Merge pull request #19 from dzhulgakov/fix-wheel

Fix wheel building

.gitignore

@@ -5,3 +5,6 @@ build
 dist
 *.egg-info
 *.pyc
+# Third-party links created by `setup.py develop`
+deep_gemm/include/cute
+deep_gemm/include/cutlass

setup.py

@@ -2,12 +2,15 @@ import os
 import setuptools
 import shutil
 import subprocess
+from setuptools.command.build_py import build_py
 from setuptools.command.develop import develop
-from setuptools.command.install import install
 
 current_dir = os.path.dirname(os.path.realpath(__file__))
-jit_include_dirs = ('deep_gemm/include/deep_gemm', )
+jit_include_dirs = ("deep_gemm/include/deep_gemm",)
-third_party_include_dirs = ('third-party/cutlass/include/cute', 'third-party/cutlass/include/cutlass')
+third_party_include_dirs = (
+    "third-party/cutlass/include/cute",
+    "third-party/cutlass/include/cutlass",
+)
 
 
 class PostDevelopCommand(develop):
@@ -19,9 +22,9 @@ class PostDevelopCommand(develop):
     def make_jit_include_symlinks():
         # Make symbolic links of third-party include directories
         for d in third_party_include_dirs:
-            dirname = d.split('/')[-1]
+            dirname = d.split("/")[-1]
-            src_dir = f'{current_dir}/{d}'
+            src_dir = f"{current_dir}/{d}"
-            dst_dir = f'{current_dir}/deep_gemm/include/{dirname}'
+            dst_dir = f"{current_dir}/deep_gemm/include/{dirname}"
             assert os.path.exists(src_dir)
             if os.path.exists(dst_dir):
                 assert os.path.islink(dst_dir)
@@ -29,37 +32,53 @@ class PostDevelopCommand(develop):
             os.symlink(src_dir, dst_dir, target_is_directory=True)
 
 
-class PostInstallCommand(install):
+class CustomBuildPy(build_py):
     def run(self):
-        install.run(self)
+        # First, prepare the include directories
-        self.copy_jit_includes()
+        self.prepare_includes()
+        # Then run the regular build
+        build_py.run(self)
 
-    def copy_jit_includes(self):
+    def prepare_includes(self):
-        # Copy include directories needed by JIT
+        # Create temporary build directory instead of modifying package directory
-        shutil.rmtree(f'{self.build_lib}/deep_gemm/include', ignore_errors=True)
+        build_include_dir = os.path.join(self.build_lib, "deep_gemm/include")
-        os.makedirs(f'{self.build_lib}/deep_gemm/include', exist_ok=False)
+        os.makedirs(build_include_dir, exist_ok=True)
-        for d in jit_include_dirs + third_party_include_dirs:
+
-            src_dir = f'{current_dir}/{d}'
+        # Copy third-party includes to the build directory
-            dst_dir = f'{self.build_lib}/deep_gemm/include/{d.split("/")[-1]}'
+        for d in third_party_include_dirs:
-            assert os.path.exists(src_dir)
+            dirname = d.split("/")[-1]
+            src_dir = os.path.join(current_dir, d)
+            dst_dir = os.path.join(build_include_dir, dirname)
+
+            # Remove existing directory if it exists
+            if os.path.exists(dst_dir):
+                shutil.rmtree(dst_dir)
+
+            # Copy the directory
             shutil.copytree(src_dir, dst_dir)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     # noinspection PyBroadException
     try:
-        cmd = ['git', 'rev-parse', '--short', 'HEAD']
+        cmd = ["git", "rev-parse", "--short", "HEAD"]
-        revision = '+' + subprocess.check_output(cmd).decode('ascii').rstrip()
+        revision = "+" + subprocess.check_output(cmd).decode("ascii").rstrip()
     except:
-        revision = ''
+        revision = ""
 
-    # noinspection PyTypeChecker
     setuptools.setup(
-        name='deep_gemm',
+        name="deep_gemm",
-        version='1.0.0' + revision,
+        version="1.0.0" + revision,
-        packages=['deep_gemm', 'deep_gemm/jit', 'deep_gemm/jit_kernels'],
+        packages=["deep_gemm", "deep_gemm/jit", "deep_gemm/jit_kernels"],
+        package_data={
+            "deep_gemm": [
+                "include/deep_gemm/**/*",
+                "include/cute/**/*",
+                "include/cutlass/**/*",
+            ]
+        },
         cmdclass={
-            'develop': PostDevelopCommand,
+            "develop": PostDevelopCommand,
-            'install': PostInstallCommand
+            "build_py": CustomBuildPy,
-        }
+        },
     )