[1/2] Wean off of PYBIND in favor of torch.ops.load_library (#1276)

Wean ao off of PYBIND [part 1]

Author: janeyx99

packaging/post_build_script.sh

@@ -7,15 +7,14 @@
 
 set -eux
 
-WHEEL_NAME=$(ls dist/)
+# Prepare manywheel, only for CUDA.
+# The wheel is a pure python wheel for other platforms.
+if [[ "$CU_VERSION" == cu* ]]; then
+    WHEEL_NAME=$(ls dist/)
 
-pushd dist
-# Prepare manywheel
-manylinux_plat=manylinux2014_x86_64
-if [[ "$CU_VERSION" == "xpu" ]]; then
-    manylinux_plat=manylinux_2_28_x86_64
-fi
-auditwheel repair --plat "$manylinux_plat" -w . \
+    pushd dist
+    manylinux_plat=manylinux2014_x86_64
+    auditwheel repair --plat "$manylinux_plat" -w . \
     --exclude libtorch.so \
     --exclude libtorch_python.so \
     --exclude libtorch_cuda.so \
@@ -26,10 +25,11 @@ auditwheel repair --plat "$manylinux_plat" -w . \
     --exclude libcudart.so.11.0 \
     "${WHEEL_NAME}"
 
-ls -lah .
-# Clean up the linux_x86_64 wheel
-rm "${WHEEL_NAME}"
-popd
+    ls -lah .
+    # Clean up the linux_x86_64 wheel
+    rm "${WHEEL_NAME}"
+    popd
+fi
 
 MANYWHEEL_NAME=$(ls dist/)
 # Try to install the new wheel

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",

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 (