Wean ao off of PYBIND [part 1]

[ghstack-poisoned]
pytorch · Nov 12, 2024 · d831de1 · d831de1
1 parent addcb24
commit d831de1
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 5 deletions.
diff --git a/setup.py b/setup.py
@@ -110,6 +110,9 @@ def get_extensions():
     if use_cuda:
         sources += cuda_sources
 
+    if len(sources) == 0:
+        return None
+
     ext_modules = [
         extension(
             "torchao._C",

diff --git a/torchao/__init__.py b/torchao/__init__.py
@@ -22,10 +22,12 @@
 )
 if not _IS_FBCODE:
     try:
-        from . import _C
+        from pathlib import Path
+        so_files = list(Path(__file__).parent.glob("_C*.so"))
+        assert len(so_files) == 1, f"Expected one _C*.so file, found {len(so_files)}"
+        torch.ops.load_library(so_files[0])
         from . import ops
     except:
-        _C = None
         logging.info("Skipping import of cpp extensions")
 
 from torchao.quantization import (

diff --git a/torchao/csrc/init.cpp b/torchao/csrc/init.cpp