Add MSVC support (#118)

Adds support for Windows builds with MSVC!

---------

Co-authored-by: Ali Hassani <alih@gatech.edu>

Author: FgoDt

.gitignore

@@ -133,3 +133,6 @@ dmypy.json
 
 # Output files
 *.out
+
+# MSVC
+.vs/

WindowsBuilder.bat

@@ -0,0 +1,50 @@
+@ECHO off
+
+if [%1] == [test] shift & goto :test
+if [%1] == [install] shift & goto :install
+if [%1] == [clean] shift & goto :clean
+
+:parseArgs
+if [%1] == [WORKERS] set NATTEN_N_WORKERS=%2 & shift & shift & goto :parseargs
+if [%1] == [CUDA_ARCH] set NATTEN_CUDA_ARCH=%2 & shift & shift & goto :parseargs
+if [%1] == [VERBOSE] set NATTEN_VERBOSE=%2 & shift & shift & goto :parseargs
+goto :installContinue
+:end
+
+:test
+echo "Testing NATTEN"
+pip install -r requirements-dev.txt
+pytest -v -x ./tests
+goto :eof
+:end
+
+:install
+goto :installStart
+:end
+
+:installStart
+goto :parseargs
+:end
+
+:installFinalize
+set NATTEN_N_WORKERS=
+set NATTEN_CUDA_ARCH=
+set NATTEN_VERBOSE=
+goto :eof
+:end
+
+:installContinue
+echo NATTEN_N_WORKERS: %NATTEN_N_WORKERS%
+echo NATTEN_CUDA_ARCH: %NATTEN_CUDA_ARCH%
+echo NATTEN_VERBOSE: %NATTEN_VERBOSE%
+pip install -r requirements.txt
+python setup.py install
+goto :installFinalize
+:end
+
+:clean
+echo Cleaning up
+echo "Removing %CD%\build"
+del %CD%\build
+goto :eof
+:end

csrc/CMakeLists.txt

@@ -97,6 +97,17 @@ set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -O3")
 set(CMAKE_CUDA_FLAGS_RELEASE "${CMAKE_CUDA_FLAGS_RELEASE} -Xcompiler -O3")
 set(CMAKE_CUDA_FLAGS_RELEASE "${CMAKE_CUDA_FLAGS_RELEASE} --use_fast_math")
 
+
+if(${NATTEN_IS_WINDOWS})
+  # Copied from xformers
+  set(CMAKE_CXX_FLAGS  "${CMAKE_CXX_FLAGS} /MP /Zc:lambda /Zc:preprocessor")
+  set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler /Zc:lambda -Xcompiler /Zc:preprocessor")
+  set(CMAKE_CXX_FLAGS_DEBUG  "${CMAKE_CXX_FLAGS_DEBUG} /MP /Zc:lambda /Zc:preprocessor")
+  set(CMAKE_CUDA_FLAGS_DEBUG "-Xcompiler /Zc:lambda -Xcompiler /Zc:preprocessor")
+  set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} /MP /Zc:lambda /Zc:preprocessor")
+  set(CMAKE_CUDA_FLAGS_RELEASE "${CMAKE_CUDA_FLAGS_RELEASE} -Xcompiler /Zc:lambda -Xcompiler /Zc:preprocessor")
+endif()
+
 message("CMAKE_CUDA_FLAGS_RELEASE: ${CMAKE_CUDA_FLAGS_RELEASE}")
 if(${NATTEN_WITH_CUDA})
   message("CUDA compiler: " ${CMAKE_CUDA_COMPILER})
@@ -179,6 +190,11 @@ include_directories(${COMMON_HEADER_DIRS})
 # Find torch lib dir so we can link with libtorch
 link_directories("${TORCH_DIR}/lib/")
 
+# Windows builds require linking with python*.lib
+if(NATTEN_IS_WINDOWS)
+  link_directories("${PY_LIB_DIR}")
+endif()
+
 add_library(natten SHARED ${ALL_SOURCES})
 
 set_target_properties(natten PROPERTIES PREFIX "" OUTPUT_NAME ${OUTPUT_FILE_NAME})
@@ -193,7 +209,7 @@ if(${NATTEN_IS_MAC})
 endif()
 
 if(${NATTEN_WITH_CUDA})
-  target_link_libraries(natten PUBLIC "-lc10 -ltorch -ltorch_cpu -ltorch_python -lcudart -lc10_cuda -ltorch_cuda")
+  target_link_libraries(natten PUBLIC c10 torch torch_cpu torch_python cudart c10_cuda torch_cuda)
 else()
-  target_link_libraries(natten PUBLIC "-lc10 -ltorch -ltorch_cpu -ltorch_python")
+  target_link_libraries(natten PUBLIC c10 torch torch_cpu torch_python)
 endif()

docs/install.md

@@ -22,13 +22,10 @@ We plan to port our naive kernels to metal soon, but we certainly welcome contri
 
 ### Windows
 
-NATTEN does not support Windows builds yet.
-This leaves Windows users with only one option: WSL.
+NATTEN does not come with Windows releases yet, but you can build it from source.
+If you're using WSL, please follow the same steps as building for Linux below.
 
-Since [switching to cmake](https://github.com/SHI-Labs/NATTEN/tree/3036259bdc5c30b7b49fe8ba17d60a0ab0d780dd) for building
-libnatten, we have not been able to support MSVC builds.
-While we welcome contributions, we are not going to be able to fix MSVC builds any time in the near future,
-so we recommend building and running NATTEN with either WSL or MinGW.
+If you're building with MSVC, please refer to [Build with MSVC](#Build-with-MSVC).
 
 ### Building from source
 In order to build from source, please make sure that you have your preferred PyTorch build installed,
@@ -62,6 +59,59 @@ You can do so by passing in the following arguments:
 make WORKERS=64
 ```
 
+It is highly recommended to run all unit tests when you build from source:
+```bash
+make test
+```
+
+#### Build with MSVC
+**NOTE: Windows builds are experimental and not regularly tested.**
+
+
+To build with MSVC, please open the "Native Tools Command Prompt for Visual Studio".
+The exact name may depend on your version of Windows, Visual Studio, and cpu architecture (in our case it was x64 Native Tools
+Command Prompt for VS".)
+
+Once in the command prompt, make sure your correct Python environment is in the system path. If you're using anaconda, you
+should be able to do `conda activate $YOUR_ENV_NAME`.
+
+Then simply confirm you have PyTorch installed, and use our Windows batch script to build:
+
+```
+WindowsBuilder.bat install 
+
+# Build with 8 parallel workers
+WindowsBuilder.bat install WORKERS=8
+
+# Build targeting SM89 (Ada Lovelace)
+WindowsBuilder.bat install CUDA_ARCH=8.9
+```
+
+Note that depending on how many workers you end up using, build time may vary, and the MSVC compiler tends to throw plenty of
+warnings at you, but as long as it does not fail and give you back the command prompt, just let it keep building.
+
+Once it's done building, it is highly recommended to run the unit tests to make sure everything went as expected:
+```
+WindowsBuilder.bat test
+```
+
+##### PyTorch issue: nvToolsExt not found
+Windows users may come across this issue when building NATTEN from source with CUDA 12.0 and newer.
+The build process fails with an error indicating "nvtoolsext" cannot be found on your system.
+This is because nvtoolsext binaries are no longer part of the CUDA toolkit for Windows starting CUDA 12.0, but the PyTorch
+cmake still looks for it (as of torch==2.2.1).
+
+The only workaround is to modify the following files:
+```
+$PATH_TO_YOUR_PYTHON_ENV\Lib\site-packages\torch\share\cmake\Torch\TorchConfig.cmake
+$PATH_TO_YOUR_PYTHON_ENV\Lib\site-packages\torch\share\cmake\Torch\Caffe2Targets.cmake
+$PATH_TO_YOUR_PYTHON_ENV\Lib\site-packages\torch\share\cmake\Caffe2\public\cuda.cmake
+```
+
+find all mentions of nvtoolsext (or nvToolsExt), and comment them out, to get past it.
+
+More information on the issue: [pytorch/pytorch#116926](https://github.com/pytorch/pytorch/pull/116926)
+
 ### NGC docker images
 NATTEN supports PyTorch builds that are built from source, an example of which is the builds that ship with
 NVIDIA's [NGC container images](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch).

requirements-dev.txt

@@ -1,5 +1,6 @@
+packaging
 torch>=2.0.0
-cmake==3.20.3
+# cmake==3.20.3
 ufmt==2.3.0
 flake8==7.0.0
 mypy==1.8.0

requirements.txt

@@ -1,2 +1,3 @@
+packaging
 torch>=2.0.0
-cmake==3.20.3
+# cmake==3.20.3

setup.py

@@ -72,12 +72,18 @@
     CUDA_VERSION = [int(x) for x in TORCH_CUDA_VERSION]
 
     assert CUDA_VERSION >= [11, 0], "NATTEN only supports CUDA 11.0 and above."
+    if CUDA_VERSION >= [12, 0] and IS_WINDOWS:
+        print(
+            "WARNING: Torch cmake will likely fail on Windows with CUDA 12.X. "
+            "Please refer to NATTEN documentation to read more about the issue "
+            "and how to get around it until the issue is fixed in torch."
+        )
 
 
-n_workers = os.environ.get("NATTEN_N_WORKERS", DEFAULT_N_WORKERS)
+n_workers = str(os.environ.get("NATTEN_N_WORKERS", DEFAULT_N_WORKERS)).strip()
 # In case the env variable is set, but to an empty string
 if n_workers == "":
-    n_workers = DEFAULT_N_WORKERS
+    n_workers = str(DEFAULT_N_WORKERS)
 
 if HAS_CUDA:
     print(f"Building NATTEN with CUDA {CUDA_TAG}")
@@ -182,6 +188,16 @@ def build_extension(self, ext):
             if max_sm >= 50:
                 cmake_args.append("-DNATTEN_WITH_CUTLASS=1")
 
+        if IS_WINDOWS:
+            python_path = sys.executable
+            assert (
+                "python.exe" in python_path
+            ), f"Expected the python executable path to end with python.exe, got {python_path}"
+            python_lib_dir = python_path.replace("python.exe", "libs").strip()
+            cmake_args.append(f"-DPY_LIB_DIR={python_lib_dir}")
+            cmake_args.append("-G Ninja")
+            cmake_args.append("-DCMAKE_BUILD_TYPE=Release")
+
         if not os.path.exists(self.build_lib):
             os.makedirs(self.build_lib)
 
@@ -193,9 +209,15 @@ def build_extension(self, ext):
         subprocess.check_call(
             ["cmake", cmake_lists_dir] + cmake_args, cwd=self.build_lib
         )
-        subprocess.check_call(
-            ["make", f"-j{n_workers}", f"VERBOSE={verbose}"], cwd=self.build_lib
-        )
+        cmake_build_args = [
+            "--build",
+            self.build_lib,
+            "-j",
+            n_workers,
+        ]
+        if verbose:
+            cmake_build_args.append("--verbose")
+        subprocess.check_call(["cmake", *cmake_build_args])
 
         # Clean up cmake files when building dist package;
         # otherwise they will get packed into the wheel.
@@ -224,7 +246,6 @@ def build_extension(self, ext):
     python_requires=">=3.8",
     install_requires=[
         "packaging",
-        "cmake==3.20.3",
         "torch>=2.0.0",
     ],
     extras_require={