Use cppimport instead of makefile

huggingface · Jul 24, 2023 · e0eb64b · e0eb64b
1 parent 8f08ebe
commit e0eb64b
Show file tree

Hide file tree

Showing 6 changed files with 54 additions and 81 deletions.
diff --git a/Makefile b/Makefile
@@ -64,27 +64,3 @@ doc: build_doc_docker_image
 		--version_tag_suffix "" \
 		--html \
 		--clean
-
-
-# make custom_ops
-# Builds the group_quantize_decompress custom ops
-
-CXX ?= g++
-OUT ?= build/custom_ops.so
-OBJDIR ?= $(dir $(OUT))obj
-
-CXXFLAGS = -Wall -Wno-sign-compare -std=c++17 -O2 -g -fPIC -DONNX_NAMESPACE=onnx
-LIBS = -lpoplar -lpopart -lpopops -lpoplin -lpopnn -lpoputil -lpoprand
-
-OBJECTS = $(OBJDIR)/group_quantize_decompress.o $(OBJDIR)/group_quantize_decompressx.o
-
-# Rules
-
-custom_ops: $(OUT)
-
-$(OBJECTS): $(OBJDIR)/%.o: optimum/graphcore/custom_ops/group_quantize_decompress/%.cpp
-	@mkdir -p $(@D)
-	$(CXX) -c $(CXXFLAGS) $< -o $@
-
-$(OUT): $(OBJECTS)
-	$(CXX) $(CXXFLAGS) -shared $^ -o $@ -Wl,--no-undefined $(LIBS)
diff --git a/optimum/graphcore/__init__.py b/optimum/graphcore/__init__.py
@@ -77,33 +77,33 @@
 poptorch.setLogLevel("ERR")
 
 
-# Load the custom ops
-def _load_custom_ops():
-    import ctypes
-    import pkg_resources
-    import sysconfig
-    from pathlib import Path
+# # Load the custom ops
+# def _load_custom_ops():
+#     import ctypes
+#     import pkg_resources
+#     import sysconfig
+#     from pathlib import Path
 
-    root = Path(pkg_resources.get_distribution("optimum-graphcore").location).absolute()
-    names = ["custom_ops.so", str(Path("custom_ops.so").with_suffix(sysconfig.get_config_vars()["SO"]))]
-    paths = [
-        root / "build" / names[0],
-        root / names[1],
-    ]
-    print("CUSTOM OPS", paths)
-    for path in paths:
-        if path.exists():
-            ctypes.cdll.LoadLibrary(str(path))
-            print("Loading:", path)
-            return
-    # Search recursively in build dir if not found in first level
-    for name in names:
-        for path in (root / "build").rglob(str(name)):
-            if path.exists():
-                ctypes.cdll.LoadLibrary(str(path))
-                print("Loading:", path)
-                return
-    raise ImportError(f"Cannot find extension library {name} - tried {[str(p) for p in paths]}")  # pragma: no cover
+#     root = Path(pkg_resources.get_distribution("optimum-graphcore").location).absolute()
+#     names = ["custom_ops.so", str(Path("custom_ops.so").with_suffix(sysconfig.get_config_vars()["SO"]))]
+#     paths = [
+#         root / "build" / names[0],
+#         root / names[1],
+#     ]
+#     print("CUSTOM OPS", paths)
+#     for path in paths:
+#         if path.exists():
+#             ctypes.cdll.LoadLibrary(str(path))
+#             print("Loading:", path)
+#             return
+#     # Search recursively in build dir if not found in first level
+#     for name in names:
+#         for path in (root / "build").rglob(str(name)):
+#             if path.exists():
+#                 ctypes.cdll.LoadLibrary(str(path))
+#                 print("Loading:", path)
+#                 return
+#     raise ImportError(f"Cannot find extension library {name} - tried {[str(p) for p in paths]}")  # pragma: no cover
 
 
-_load_custom_ops()
+# _load_custom_ops()
diff --git a/optimum/graphcore/custom_ops/group_quantize_decompress/group_quantize_decompress.cpp b/optimum/graphcore/custom_ops/group_quantize_decompress/group_quantize_decompress.cpp
@@ -1,3 +1,5 @@
+// cppimport
+// NOTE: the cppimport comment is necessary for dynamic compilation when loading
 // Copyright (c) 2023 Graphcore Ltd. All rights reserved.
 #include <algorithm>
 #include <cstdint>
@@ -106,3 +108,16 @@ true);
 // static popart::popx::OpxCreator<GroupQuantizeDecompressOpx> GroupQuantizeDecompressOpxCreator({GroupQuantizeDecompressOpId});
 
 } // namespace popart
+
+// -------------- cppimport --------------
+// cppimport configuration for compiling the pybind11 module.
+// clang-format off
+/*
+<%
+cfg['sources'] = ['group_quantize_decompressx.cpp']
+cfg['dependencies'] = ['common.hpp', 'group_quantize_decompress.cpp', 'group_quantize_decompressx.cpp']
+cfg['extra_compile_args'] = ['-std=c++17', '-fPIC', '-O2', '-DONNX_NAMESPACE=onnx', '-Wall', '-Wno-sign-compare']
+cfg['libraries'] = ['poplar', 'popart', 'poputil', 'popops', 'poplin', 'popnn', 'poprand']
+setup_pybind11(cfg)
+%>
+*/
diff --git a/optimum/graphcore/models/whisper/modeling_whisper.py b/optimum/graphcore/models/whisper/modeling_whisper.py
@@ -40,7 +40,6 @@
     shift_tokens_right,
     split_encoder_decoder_ipu_config,
 )
-from ...quantization.group_quantize import GroupQuantLinear
 
 
 logger = logging.get_logger(__name__)
@@ -492,7 +491,7 @@ def change_lm_head(self, restore: bool, use_cache: bool = None):
 
     def quantize_linear_layers(self, restore: bool, num_groups: int = 16):
         if not restore:
-            # Brute force way
+            from ...quantization.group_quantize import GroupQuantLinear
             for module in self.model.encoder.layers:
                 module.self_attn.q_proj = GroupQuantLinear.from_model(module.self_attn.q_proj, num_groups)
                 module.self_attn.k_proj = GroupQuantLinear.from_model(module.self_attn.k_proj, num_groups)

diff --git a/optimum/graphcore/quantization/__init__.py b/optimum/graphcore/quantization/__init__.py
@@ -0,0 +1,8 @@
+from pathlib import Path
+from ..custom_ops.utils import load_lib
+
+load_lib(Path(__file__).parent.parent / "custom_ops/group_quantize_decompress/group_quantize_decompress.cpp")
+# from optimum.graphcore.custom_ops.sdk_version_hash import sdk_version_hash
+# print(sdk_version_hash())
+
+from .group_quantize import group_quantize_compress, group_quantize_decompress, GroupQuantLinear
diff --git a/setup.py b/setup.py
@@ -15,7 +15,6 @@
 
 import re
 
-import setuptools.command.build_ext
 import subprocess
 from pathlib import Path
 from setuptools import find_namespace_packages, setup
@@ -34,6 +33,7 @@
     "transformers==4.29.2",
     "optimum==1.6.1",
     "diffusers[torch]==0.12.1",
+    "cppimport==22.8.2",
     "datasets",
     "tokenizers",
     "typeguard",
@@ -65,25 +65,6 @@
 }
 
 
-
-class make_ext(setuptools.command.build_ext.build_ext):  # type:ignore[misc]
-    def build_extension(self, ext: setuptools.Extension) -> None:
-        if ext.name == "custom_ops":
-            filename = Path(self.build_lib) / self.get_ext_filename(
-                self.get_ext_fullname(ext.name)
-            )
-            objdir = filename.with_suffix("")
-            subprocess.check_call(
-                [
-                    "make",
-                    "custom_ops",
-                    f"OUT={filename}",
-                    f"OBJDIR={objdir}",
-                ]
-            )
-        else:
-            super().build_extension(ext)
-
 setup(
     name="optimum-graphcore",
     version=__version__,
@@ -113,12 +94,6 @@ def build_extension(self, ext: setuptools.Extension) -> None:
     extras_require=EXTRA_REQUIRE,
     include_package_data=True,
     zip_safe=False,
-    ext_modules=[
-        setuptools.Extension(
-            "custom_ops",
-            sources=list(map(str, Path("optimum/graphcore/custom_ops/group_quantize_decompress").glob("*.[ch]pp"))),
-        )
-    ],
-    package_data={"optimum.graphcore": ["custom_ops/group_quantize_decompress/*_codelet_*.cpp"]},
-    cmdclass=dict(build_ext=make_ext),
+#    package_data={"optimum.graphcore": list(map(str, Path("optimum/graphcore/custom_ops/group_quantize_decompress").glob("*.[ch]pp"))) + [str(Path("optimum/graphcore/custom_ops/sdk_version_hash/sdk_version_hash_lib.cpp"))]}
+   package_data={"optimum-graphcore": [str(Path("optimum/graphcore/custom_ops/sdk_version_hash/sdk_version_hash_lib.cpp"))]}
 )