BLD: Initial Windows support

finch-tensor · Nov 1, 2024 · 1f8ed85 · 1f8ed85
1 parent 195eddd
commit 1f8ed85
Show file tree

Hide file tree

Showing 6 changed files with 129 additions and 30 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -1,6 +1,12 @@
 name: Wheel
 
 on:
+  pull_request:
+    branches:
+      - main  # NOTE: to remove
+  push:
+    branches:
+      - main  # NOTE: to remove
   schedule:
     # At 09:00 on Monday. (see https://crontab.guru)
     - cron: '0 9 * * 1'
@@ -17,7 +23,7 @@ jobs:
   build-wheel:
     strategy:
       matrix:
-        os: ['ubuntu-24.04', 'macos-12', 'macos-14']  # TODO: 'windows-2022'
+        os: ['windows-2022']  # TODO: , 'ubuntu-24.04', 'macos-12', 'macos-14'
         arch: ['x86_64', 'aarch64']
         python: ['3.10']
         exclude:
@@ -44,6 +50,19 @@ jobs:
           ref: '4091bc61e315f187829dca877dd908a07ba9cb91'  # Latest commit as of 2024-10-17
           path: 'llvm-project'
 
+      - uses: ilammy/[email protected]
+        if: ${{ matrix.os == 'windows-2022' }}
+
+      - name: Set up Visual Studio shell
+        if: ${{ matrix.os == 'windows-2022' }}
+        uses: egor-tensin/vs-shell@v2
+        with:
+          arch: x64
+
+      - name: MS Build
+        if: ${{ matrix.os == 'windows-2022' }}
+        uses: microsoft/[email protected]
+
       - name: Setup ccache
         uses: hendrikmuhs/ccache-action@v1
         with:
@@ -85,7 +104,18 @@ jobs:
         run: |
           echo "MACOSX_DEPLOYMENT_TARGET=11.0" | tee -a $GITHUB_ENV
 
+      # - name: set env variables - windows
+      #   if: contains(matrix.os, 'windows')
+      #   shell: pwsh
+      #   run: |
+      #     echo "SPARSE_BACKEND=MLIR" >> $env:GITHUB_ENV
+
+      # - name: Setup tmate session
+      #   uses: mxschmitt/action-tmate@v3
+
       - name: cibuildwheel run
+        env:
+          SPARSE_BACKEND: MLIR
         run: |
           mv ./Finch-mlir/setup.py .
           mv ./Finch-mlir/pyproject.toml .

diff --git a/finch-plugin/CMakeLists.txt b/finch-plugin/CMakeLists.txt
@@ -11,8 +11,10 @@ add_llvm_library(FinchPlugin
         LINK_LIBS
         MLIRFinch
         )
-target_include_directories(
-        FinchPlugin
-        PRIVATE
-        "${FINCH_BINARY_DIR}/include"
-)
+if(UNIX)
+        target_include_directories(
+                FinchPlugin
+                PRIVATE
+                ${FINCH_BINARY_DIR}/include
+        )
+endif()
diff --git a/pyproject.toml b/pyproject.toml
@@ -14,7 +14,7 @@ build = "cp310-* cp311-* cp312-* cp313-*"
 build-verbosity = 1
 test-requires = ["pytest", "pytest-cov", "PyYAML", "scipy"]
 before-test = "pip install --no-deps sparse@git+https://github.com/pydata/sparse@updated-llvm-nightly-test"
-test-command = "SPARSE_BACKEND=MLIR pytest --pyargs sparse.mlir_backend"
+test-command = "pytest --pyargs sparse.mlir_backend"
 
 [tool.cibuildwheel.linux]
 skip = ["*-manylinux_i686", "*-musllinux*"]
@@ -23,7 +23,7 @@ before-build = [
     "pip install -r {project}/Finch-mlir/requirements/requirements.txt",
     "{project}/Finch-mlir/scripts/docker_prepare_ccache.sh"
 ]
-environment-pass = ["HOST_CCACHE_DIR"]
+environment-pass = ["HOST_CCACHE_DIR", "SPARSE_BACKEND"]
 
 [tool.cibuildwheel.macos]
 environment = { PATH = "/usr/local/opt/ccache/libexec:$PATH" }
@@ -35,6 +35,7 @@ repair-wheel-command = [
 ]
 
 [tool.cibuildwheel.windows]
+build = "cp310-win_amd64 cp311-win_amd64 cp312-win_amd64 cp313-win_amd64"
 before-build = [
     "pip install delvewheel",
     "pip install -r {project}/Finch-mlir/requirements/requirements.txt",

diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt
@@ -25,37 +25,64 @@ declare_mlir_python_extension(FinchPythonSources.Extension
   ADD_TO_PARENT FinchPythonSources
   SOURCES
     FinchExtension.cpp
+  PRIVATE_LINK_LIBS
+    LLVMSupport
   EMBED_CAPI_LINK_LIBS
     FinchCAPI
+    # MLIRCAPIConversion
+    # MLIRCAPITransforms
+    # MLIRCAPIRegisterEverything
 )
 
+# EMBED_CAPI_LINK_LIBS
+#   MLIRCAPIDebug
+#   MLIRCAPIIR
+#   MLIRCAPIInterfaces
+# declare_mlir_python_extension(FinchPythonSources.RegisterWhatWeNeed
+#   MODULE_NAME _finchRegisterWhatWeNeed
+#   ADD_TO_PARENT FinchPythonSources
+#   SOURCES
+#     RegisterEverything.cpp
+#   PRIVATE_LINK_LIBS
+#     LLVMSupport
+#   EMBED_CAPI_LINK_LIBS
+#     MLIRCAPIConversion
+#     MLIRCAPITransforms
+#     MLIRCAPIRegisterEverything
+# )
+
 ################################################################################
 # Common CAPI
 ################################################################################
 
-add_mlir_python_common_capi_library(FinchPythonCAPI
+add_mlir_python_common_capi_library(FPCAPI
   INSTALL_COMPONENT FinchPythonModules
-  INSTALL_DESTINATION python_packages/finch/mlir_finch/_mlir_libs
-  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/finch/mlir_finch/_mlir_libs"
-  RELATIVE_INSTALL_ROOT "../../../.."
+  INSTALL_DESTINATION mlir_finch/_ms
+  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/mlir_finch/_ms"
+  RELATIVE_INSTALL_ROOT "../.."
   DECLARED_SOURCES
     FinchPythonSources
     # TODO: Remove this in favor of showing fine grained registration once
     # available.
     MLIRPythonExtension.ExecutionEngine
-    MLIRPythonExtension.RegisterEverything
-    MLIRPythonSources.ExecutionEngine
-    MLIRPythonSources.Dialects
-    MLIRPythonSources.Core
+    #MLIRPythonExtension.RegisterEverything
+    #MLIRPythonSources.ExecutionEngine
+    #MLIRPythonSources.Dialects
+    MLIRPythonSources #.Core
+    #LLVMSupport
+  #EMBED_CAPI_LINK_LIBS
+    #MLIRCAPIConversion
+    #MLIRCAPITransforms
+    #MLIRCAPIRegisterEverything
 )
 
 ################################################################################
 # Instantiation of all Python modules
 ################################################################################
 
 add_mlir_python_modules(FinchPythonModules
-  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/finch/mlir_finch"
-  INSTALL_PREFIX "python_packages/finch/mlir_finch"
+  ROOT_PREFIX "${MLIR_BINARY_DIR}/mlir_finch"
+  INSTALL_PREFIX "mlir_finch"
   DECLARED_SOURCES
     FinchPythonSources
     # TODO: Remove this in favor of showing fine grained registration once
@@ -66,5 +93,6 @@ add_mlir_python_modules(FinchPythonModules
     MLIRPythonExtension.RegisterEverything
     MLIRPythonSources
   COMMON_CAPI_LINK_LIBS
-    FinchPythonCAPI
+    FPCAPI
+    MLIRCAPIRegisterEverything
   )
diff --git a/python/mlir_finch/dialects/finch.py b/python/mlir_finch/dialects/finch.py
@@ -3,4 +3,4 @@
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from ._finch_ops_gen import *
-from .._mlir_libs._finchDialects.finch import *
+from .._ms._finchDialects.finch import *
diff --git a/setup.py b/setup.py
@@ -1,4 +1,5 @@
 import os
+import platform
 import shutil
 import sys
 import subprocess
@@ -23,6 +24,14 @@ def __init__(
 
 class CMakeBuild(build_ext):
     def build_extension(self, ext: CMakeExtension) -> None:
+
+        p1 = os.fspath(Path("/tmp/st").absolute())
+        shutil.move(ext.llvm_source_dir, p1)
+        ext.llvm_source_dir = os.fspath(Path("/tmp/st/llvm").absolute())
+        p2 = os.fspath(Path("/tmp/sf").absolute())
+        shutil.move(ext.finch_source_dir, p2)
+        ext.finch_source_dir = p2
+
         ext_fullpath = Path.cwd() / self.get_ext_fullpath(ext.name)
         extdir = ext_fullpath.parent.resolve()
         install_dir = extdir
@@ -33,6 +42,14 @@ def build_extension(self, ext: CMakeExtension) -> None:
         extra_flags = []
         if sys.platform.startswith("darwin"):
             extra_flags.append("-DCMAKE_OSX_DEPLOYMENT_TARGET=11.0")
+        elif platform.system() == "Windows":
+            extra_flags += [
+                "-DCMAKE_C_COMPILER=cl",
+                "-DCMAKE_CXX_COMPILER=cl",
+                "-DCMAKE_MSVC_RUNTIME_LIBRARY=MultiThreaded",
+                "-DCMAKE_C_FLAGS=/MT",
+                "-DCMAKE_CXX_FLAGS=/MT",
+            ]
 
         # BUILD LLVM
         llvm_cmake_args = [
@@ -64,12 +81,33 @@ def build_extension(self, ext: CMakeExtension) -> None:
             check=True,
         )
 
+        llvm_lit = "llvm-lit.py" if platform.system() == "Windows" else "llvm-lit"
+
+        subprocess.run(
+            ["dir", llvm_build_dir / 'bin'],
+            cwd=finch_build_dir,
+            check=True,
+        )
+
+        subprocess.run(
+            ["dir", str(llvm_build_dir)],
+            cwd=finch_build_dir,
+            check=True,
+        )
+
+        # if platform.system() == "Windows":
+        #     # fatal error LNK1170: line in command file contains 131071 or more characters
+        #     if Path("/tmp/m").exists():
+        #         shutil.rmtree("/tmp/m")
+        #     shutil.move(llvm_install_dir, "/tmp/m")
+        #     llvm_install_dir = Path("/tmp/m").absolute()
+
         # BUILD FINCH DIALECT
         dialect_cmake_args = [
             "-G Ninja",
             f"-B{finch_build_dir}",
-            f"-DMLIR_DIR={llvm_install_dir}/lib/cmake/mlir",
-            f"-DLLVM_EXTERNAL_LIT={llvm_build_dir}/bin/llvm-lit",
+            f"-DMLIR_DIR={llvm_install_dir / 'lib' / 'cmake' / 'mlir'}",
+            f"-DLLVM_EXTERNAL_LIT={llvm_build_dir / 'bin' / llvm_lit}",
             "-DCMAKE_PLATFORM_NO_VERSIONED_SONAME=ON",
             f"-DCMAKE_MAKE_PROGRAM:FILEPATH={ninja_executable_path}",
             "-DLLVM_ENABLE_ZLIB=OFF",
@@ -91,9 +129,9 @@ def build_extension(self, ext: CMakeExtension) -> None:
         )
 
         # Move Python package out of nested directories.
-        python_package_dir = install_dir / "python_packages" / "finch" / "mlir_finch"
-        shutil.copytree(python_package_dir, install_dir / "mlir_finch")
-        shutil.rmtree(install_dir / "python_packages")
+        # python_package_dir = install_dir / "python_packages" / "finch" / "mlir_finch"
+        # shutil.copytree(python_package_dir, install_dir / "mlir_finch")
+        # shutil.rmtree(install_dir / "python_packages")
 
         subprocess.run(
             [
@@ -114,15 +152,15 @@ def build_extension(self, ext: CMakeExtension) -> None:
 
 
 def create_dir(name: str) -> Path:
-    path = Path.cwd() / "build" / name
+    path = Path("/tmp").absolute() / name
     if not path.exists():
         path.mkdir(parents=True)
     return path
 
 
-llvm_build_dir = create_dir("llvm-build")
-llvm_install_dir = create_dir("llvm-install")
-finch_build_dir = create_dir("finch-build")
+llvm_build_dir = create_dir("ub")
+llvm_install_dir = create_dir("ui")
+finch_build_dir = create_dir("fb")
 
 
 setup(
@@ -134,7 +172,7 @@ def create_dir(name: str) -> Path:
     long_description_content_type="text/markdown",
     ext_modules=[CMakeExtension(
         "mlir_finch_ext",
-        llvm_source_dir=f"./llvm-project/llvm",
+        llvm_source_dir=f"./llvm-project", # /llvm
         finch_source_dir="./Finch-mlir",
     )],
     cmdclass={"build_ext": CMakeBuild},