build: Minor tweeks for wheel build

Signed-off-by: oliver könig <[email protected]>

Author: ko3n1g

deep_gemm/__init__.py

@@ -5,6 +5,7 @@
 try:
     # noinspection PyUnresolvedReferences
     from .envs import persistent_envs
+
     for key, value in persistent_envs.items():
         if key not in os.environ:
             os.environ[key] = value
@@ -23,19 +24,23 @@
 # Kernels
 from deep_gemm_cpp import (
     # FP8 GEMMs
-    fp8_gemm_nt, fp8_gemm_nn,
-    fp8_gemm_tn, fp8_gemm_tt,
+    fp8_gemm_nt,
+    fp8_gemm_nn,
+    fp8_gemm_tn,
+    fp8_gemm_tt,
     m_grouped_fp8_gemm_nt_contiguous,
     m_grouped_fp8_gemm_nn_contiguous,
     m_grouped_fp8_gemm_nt_masked,
     k_grouped_fp8_gemm_tn_contiguous,
     # BF16 GEMMs
-    bf16_gemm_nt, bf16_gemm_nn,
-    bf16_gemm_tn, bf16_gemm_tt,
+    bf16_gemm_nt,
+    bf16_gemm_nn,
+    bf16_gemm_tn,
+    bf16_gemm_tt,
     m_grouped_bf16_gemm_nt_contiguous,
     m_grouped_bf16_gemm_nt_masked,
     # Layout kernels
-    transform_sf_into_required_layout
+    transform_sf_into_required_layout,
 )
 
 # Some alias for legacy supports
@@ -53,22 +58,29 @@
 def _find_cuda_home() -> str:
     # TODO: reuse PyTorch API later
     # For some PyTorch versions, the original `_find_cuda_home` will initialize CUDA, which is incompatible with process forks
-    cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
+    cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH")
     if cuda_home is None:
         # noinspection PyBroadException
         try:
-            with open(os.devnull, 'w') as devnull:
-                nvcc = subprocess.check_output(['which', 'nvcc'], stderr=devnull).decode().rstrip('\r\n')
+            with open(os.devnull, "w") as devnull:
+                nvcc = (
+                    subprocess.check_output(["which", "nvcc"], stderr=devnull)
+                    .decode()
+                    .rstrip("\r\n")
+                )
                 cuda_home = os.path.dirname(os.path.dirname(nvcc))
         except Exception:
-            cuda_home = '/usr/local/cuda'
+            cuda_home = "/usr/local/cuda"
             if not os.path.exists(cuda_home):
                 cuda_home = None
     assert cuda_home is not None
     return cuda_home
 
 
 deep_gemm_cpp.init(
-    os.path.dirname(os.path.abspath(__file__)), # Library root directory path
-    _find_cuda_home()                           # CUDA home
+    os.path.dirname(os.path.abspath(__file__)),  # Library root directory path
+    _find_cuda_home(),  # CUDA home
 )
+
+
+__version__ = "2.0.0"

setup.py

@@ -1,36 +1,45 @@
 import os
 import setuptools
 import shutil
-import subprocess
 import torch
+import re
+import ast
 from setuptools import find_packages
 from setuptools.command.build_py import build_py
 from torch.utils.cpp_extension import CUDAExtension, CUDA_HOME
+from pathlib import Path
+import subprocess
+
+SKIP_CUDA_BUILD = os.getenv("DEEP_GEMM_SKIP_CUDA_BUILD", "FALSE") == "TRUE"
+
 
 current_dir = os.path.dirname(os.path.realpath(__file__))
-cxx_flags = ['-std=c++17', '-O3', '-fPIC', '-Wno-psabi', '-Wno-deprecated-declarations',
-             f'-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}']
-sources = ['csrc/python_api.cpp']
-build_include_dirs = [
-    f'{CUDA_HOME}/include',
-    f'{CUDA_HOME}/include/cccl',
-    'deep_gemm/include',
-    'third-party/cutlass/include',
-    'third-party/fmt/include',
+cxx_flags = [
+    "-std=c++17",
+    "-O3",
+    "-fPIC",
+    "-Wno-psabi",
+    "-Wno-deprecated-declarations",
+    f"-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}",
 ]
-build_libraries = ['cuda', 'cudart', 'nvrtc']
-build_library_dirs = [
-    f'{CUDA_HOME}/lib64',
-    f'{CUDA_HOME}/lib64/stubs'
+sources = ["csrc/python_api.cpp"]
+build_include_dirs = [
+    f"{CUDA_HOME}/include",
+    f"{CUDA_HOME}/include/cccl",
+    "deep_gemm/include",
+    "third-party/cutlass/include",
+    "third-party/fmt/include",
 ]
+build_libraries = ["cuda", "cudart", "nvrtc"]
+build_library_dirs = [f"{CUDA_HOME}/lib64", f"{CUDA_HOME}/lib64/stubs"]
 third_party_include_dirs = [
-    'third-party/cutlass/include/cute',
-    'third-party/cutlass/include/cutlass',
+    "third-party/cutlass/include/cute",
+    "third-party/cutlass/include/cutlass",
 ]
 
 # Use runtime API
-if int(os.environ.get('DG_JIT_USE_RUNTIME_API', '0')):
-    cxx_flags.append('-DDG_JIT_USE_RUNTIME_API')
+if int(os.environ.get("DG_JIT_USE_RUNTIME_API", "0")):
+    cxx_flags.append("-DDG_JIT_USE_RUNTIME_API")
 
 
 class CustomBuildPy(build_py):
@@ -45,22 +54,30 @@ def run(self):
         build_py.run(self)
 
     def generate_default_envs(self):
-        code = '# Pre-installed environment variables\n'
-        code += 'persistent_envs = dict()\n'
-        for name in ('DG_JIT_CACHE_DIR', 'DG_JIT_PRINT_COMPILER_COMMAND', 'DG_JIT_CPP_STANDARD'):
-            code += f"persistent_envs['{name}'] = '{os.environ[name]}'\n" if name in os.environ else ''
+        code = "# Pre-installed environment variables\n"
+        code += "persistent_envs = dict()\n"
+        for name in (
+            "DG_JIT_CACHE_DIR",
+            "DG_JIT_PRINT_COMPILER_COMMAND",
+            "DG_JIT_CPP_STANDARD",
+        ):
+            code += (
+                f"persistent_envs['{name}'] = '{os.environ[name]}'\n"
+                if name in os.environ
+                else ""
+            )
 
-        with open(os.path.join(self.build_lib, 'deep_gemm', 'envs.py'), 'w') as f:
+        with open(os.path.join(self.build_lib, "deep_gemm", "envs.py"), "w") as f:
             f.write(code)
 
     def prepare_includes(self):
         # Create temporary build directory instead of modifying package directory
-        build_include_dir = os.path.join(self.build_lib, 'deep_gemm/include')
+        build_include_dir = os.path.join(self.build_lib, "deep_gemm/include")
         os.makedirs(build_include_dir, exist_ok=True)
 
         # Copy third-party includes to the build directory
         for d in third_party_include_dirs:
-            dirname = d.split('/')[-1]
+            dirname = d.split("/")[-1]
             src_dir = os.path.join(current_dir, d)
             dst_dir = os.path.join(build_include_dir, dirname)
 
@@ -72,36 +89,53 @@ def prepare_includes(self):
             shutil.copytree(src_dir, dst_dir)
 
 
-if __name__ == '__main__':
-    # noinspection PyBroadException
-    try:
-        cmd = ['git', 'rev-parse', '--short', 'HEAD']
-        revision = '+' + subprocess.check_output(cmd).decode('ascii').rstrip()
-    except:
-        revision = ''
+if not SKIP_CUDA_BUILD:
+    ext_modules = [
+        CUDAExtension(
+            name="deep_gemm_cpp",
+            sources=sources,
+            include_dirs=build_include_dirs,
+        )
+    ]
+else:
+    ext_modules = []
+
+
+NO_LOCAL_VERSION = os.getenv("DEEP_GEMM_NO_LOCAL_VERSION", "FALSE") == "TRUE"
+
+
+def get_package_version():
+    with open(Path(current_dir) / "deep_gemm" / "__init__.py", "r") as f:
+        version_match = re.search(r"^__version__\s*=\s*(.*)$", f.read(), re.MULTILINE)
+    public_version = ast.literal_eval(version_match.group(1))
+    revision = ""
+
+    if not NO_LOCAL_VERSION:
+        try:
+            cmd = ["git", "rev-parse", "--short", "HEAD"]
+            revision = "+" + subprocess.check_output(cmd).decode("ascii").rstrip()
+        except:
+            revision = ""
+
+    return f"{public_version}{revision}"
+
 
+if __name__ == "__main__":
     # noinspection PyTypeChecker
     setuptools.setup(
-        name='deep_gemm',
-        version='2.0.0' + revision,
-        packages=find_packages('.'),
+        name="deep_gemm",
+        version=get_package_version(),
+        packages=find_packages("."),
         package_data={
-            'deep_gemm': [
-                'include/deep_gemm/**/*',
-                'include/cute/**/*',
-                'include/cutlass/**/*',
+            "deep_gemm": [
+                "include/deep_gemm/**/*",
+                "include/cute/**/*",
+                "include/cutlass/**/*",
             ]
         },
-        ext_modules=[
-            CUDAExtension(name='deep_gemm_cpp',
-                          sources=sources,
-                          include_dirs=build_include_dirs,
-                          libraries=build_libraries,
-                          library_dirs=build_library_dirs,
-                          extra_compile_args=cxx_flags)
-        ],
+        ext_modules=ext_modules,
         zip_safe=False,
         cmdclass={
-            'build_py': CustomBuildPy,
+            "build_py": CustomBuildPy,
         },
     )