Remove python version specific usage within extensions

pytorch · Nov 12, 2024 · d8b9345 · d8b9345
1 parent addcb24
commit d8b9345
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 5 deletions.
diff --git a/setup.py b/setup.py
@@ -114,6 +114,7 @@ def get_extensions():
         extension(
             "torchao._C",
             sources,
+            py_limited_api=True,
             extra_compile_args=extra_compile_args,
             extra_link_args=extra_link_args,
         )

diff --git a/torchao/__init__.py b/torchao/__init__.py
@@ -22,10 +22,14 @@
 )
 if not _IS_FBCODE:
     try:
-        from . import _C
+        from importlib.util import find_spec
+        from pathlib import Path
+        spec = find_spec("torchao")
+        assert spec is not None, "torchao python module spec is unexpectedly None"
+        SO_PATH = Path(spec.origin).parent / "_C.abi3.so"
+        torch.ops.load_library(SO_PATH)
         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