diff --git a/.github/actions/setup/action.yml b/.github/actions/setup/action.yml index af42016a8..3388149b0 100644 --- a/.github/actions/setup/action.yml +++ b/.github/actions/setup/action.yml @@ -26,6 +26,9 @@ runs: sudo rm -rf /usr/share/dotnet shell: bash + - name: Set up Windows developer command prompt + uses: ilammy/msvc-dev-cmd@v1 + - name: Install CUDA ${{ inputs.cuda-version }} if: ${{ inputs.cuda-version != 'cpu' }} run: | diff --git a/.github/workflows/building.yml b/.github/workflows/building.yml index 65ea31698..442ed1c0c 100644 --- a/.github/workflows/building.yml +++ b/.github/workflows/building.yml @@ -15,7 +15,6 @@ jobs: torch-version: [1.12.0, 1.13.0, 2.0.0, 2.1.0, 2.2.0] cuda-version: ['cpu', 'cu113', 'cu116', 'cu117', 'cu118', 'cu121'] exclude: - - os: windows-2019 # No windows support yet :( - torch-version: 1.12.0 python-version: '3.12' - torch-version: 1.13.0 @@ -112,6 +111,8 @@ jobs: source ./.github/workflows/cuda/${{ runner.os }}-env.sh ${{ matrix.cuda-version }} python setup.py bdist_wheel --dist-dir=dist shell: bash + env: + TORCH_CUDA_ARCH_LIST: "5.0+PTX;6.0;7.0;7.5;8.0;8.6" - name: Test wheel run: | @@ -121,6 +122,7 @@ jobs: python -c "import pyg_lib; print('pyg-lib:', pyg_lib.__version__)" python -c "import pyg_lib; print('CUDA:', pyg_lib.cuda_version())" cd .. + shell: bash - name: Configure AWS uses: aws-actions/configure-aws-credentials@v1 diff --git a/.github/workflows/cuda/Windows.sh b/.github/workflows/cuda/Windows.sh index 0bee4bf53..582541da0 100644 --- a/.github/workflows/cuda/Windows.sh +++ b/.github/workflows/cuda/Windows.sh @@ -1,10 +1,5 @@ #!/bin/bash -# Install NVIDIA drivers, see: -# https://github.com/pytorch/vision/blob/master/packaging/windows/internal/cuda_install.bat#L99-L102 -curl -k -L "https://drive.google.com/u/0/uc?id=1injUyo3lnarMgWyRcXqKg4UGnN0ysmuq&export=download" --output "/tmp/gpu_driver_dlls.zip" -7z x "/tmp/gpu_driver_dlls.zip" -o"/c/Windows/System32" - case ${1} in cu121) CUDA_SHORT=12.1 @@ -48,3 +43,18 @@ echo "Installing from ${CUDA_FILE}..." PowerShell -Command "Start-Process -FilePath \"${CUDA_FILE}\" -ArgumentList \"-s nvcc_${CUDA_SHORT} cuobjdump_${CUDA_SHORT} nvprune_${CUDA_SHORT} cupti_${CUDA_SHORT} cublas_dev_${CUDA_SHORT} cudart_${CUDA_SHORT} cufft_dev_${CUDA_SHORT} curand_dev_${CUDA_SHORT} cusolver_dev_${CUDA_SHORT} cusparse_dev_${CUDA_SHORT} thrust_${CUDA_SHORT} npp_dev_${CUDA_SHORT} nvrtc_dev_${CUDA_SHORT} nvml_dev_${CUDA_SHORT}\" -Wait -NoNewWindow" echo "Done!" rm -f "${CUDA_FILE}" + +# echo Installing NVIDIA drivers... +# https://github.com/pytorch/vision/blob/master/packaging/windows/internal/cuda_install.bat#L99-L102 +# curl -k -L "https://ossci-windows.s3.us-east-1.amazonaws.com/builder/additional_dlls.zip" --output "/tmp/gpu_driver_dlls.zip" +# 7z x "/tmp/gpu_driver_dlls.zip" -o"/c/Windows/System32" + +echo Installing NvToolsExt... +curl -k -L https://ossci-windows.s3.us-east-1.amazonaws.com/builder/NvToolsExt.7z --output /tmp/NvToolsExt.7z +7z x /tmp/NvToolsExt.7z -o"/tmp/NvToolsExt" +mkdir -p "/c/Program Files/NVIDIA Corporation/NvToolsExt/bin/x64" +mkdir -p "/c/Program Files/NVIDIA Corporation/NvToolsExt/include" +mkdir -p "/c/Program Files/NVIDIA Corporation/NvToolsExt/lib/x64" +cp -r /tmp/NvToolsExt/bin/x64/* "/c/Program Files/NVIDIA Corporation/NvToolsExt/bin/x64" +cp -r /tmp/NvToolsExt/include/* "/c/Program Files/NVIDIA Corporation/NvToolsExt/include" +cp -r /tmp/NvToolsExt/lib/x64/* "/c/Program Files/NVIDIA Corporation/NvToolsExt/lib/x64" diff --git a/.github/workflows/install.yml b/.github/workflows/install.yml index 1c9478ab0..809bca3c0 100644 --- a/.github/workflows/install.yml +++ b/.github/workflows/install.yml @@ -30,6 +30,9 @@ jobs: run: | source ./.github/workflows/cuda/${{ runner.os }}-env.sh ${{ matrix.cuda-version }} pip install --verbose -e . + shell: bash + env: + TORCH_CUDA_ARCH_LIST: "5.0+PTX;6.0;7.0;7.5;8.0;8.6" - name: Test imports run: | diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 1eef0b3b8..f23c6721f 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -19,7 +19,6 @@ jobs: torch-version: [1.12.0, 1.13.0, 2.0.0, 2.1.0, 2.2.0] cuda-version: ['cpu', 'cu113', 'cu116', 'cu117', 'cu118', 'cu121'] exclude: - - os: windows-2019 # No windows support yet :( - torch-version: 1.12.0 python-version: '3.12' - torch-version: 1.13.0 @@ -118,6 +117,8 @@ jobs: source ./.github/workflows/cuda/${{ runner.os }}-env.sh ${{ matrix.cuda-version }} python setup.py bdist_wheel --dist-dir=dist shell: bash + env: + TORCH_CUDA_ARCH_LIST: "5.0+PTX;6.0;7.0;7.5;8.0;8.6" - name: Test wheel run: | @@ -127,6 +128,7 @@ jobs: python -c "import pyg_lib; print('pyg-lib:', pyg_lib.__version__)" python -c "import pyg_lib; print('CUDA:', pyg_lib.cuda_version())" cd .. + shell: bash - name: Configure AWS uses: aws-actions/configure-aws-credentials@v1 diff --git a/CHANGELOG.md b/CHANGELOG.md index 4c8a5e97c..c069e6c2a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). ## [0.5.0] - 2023-MM-DD ### Added +- Added Windows support ([#315](https://github.com/pyg-team/pyg-lib/pull/315)) - Added macOS Apple Silicon support ([#310](https://github.com/pyg-team/pyg-lib/pull/310)) ### Changed ### Removed diff --git a/CMakeLists.txt b/CMakeLists.txt index d66365542..187a81073 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,6 +1,8 @@ cmake_minimum_required(VERSION 3.15) project(pyg) set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_SHARED_LIBRARY_PREFIX "lib") set(PYG_VERSION 0.4.0) option(BUILD_TEST "Enable testing" OFF) @@ -71,16 +73,24 @@ else() target_include_directories(${PROJECT_NAME} PRIVATE ${PHMAP_DIR}) endif() -set(METIS_DIR third_party/METIS) -target_include_directories(${PROJECT_NAME} PRIVATE ${METIS_DIR}/include) -set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -DIDXTYPEWIDTH=64 -DREALTYPEWIDTH=32") -set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DIDXTYPEWIDTH=64 -DREALTYPEWIDTH=32") -set(GKLIB_PATH "${METIS_DIR}/GKlib") -include(${GKLIB_PATH}/GKlibSystem.cmake) -include_directories(${GKLIB_PATH}) -include_directories("${METIS_DIR}/include") -add_subdirectory("${METIS_DIR}/libmetis") -target_link_libraries(${PROJECT_NAME} PRIVATE metis) +if (MSVC) + set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} /DIDXTYPEWIDTH=64 /DREALTYPEWIDTH=32") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /DIDXTYPEWIDTH=64 /DREALTYPEWIDTH=32") +else() + set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -DIDXTYPEWIDTH=64 -DREALTYPEWIDTH=32") + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DIDXTYPEWIDTH=64 -DREALTYPEWIDTH=32") +endif() + +if (NOT MSVC) + set(METIS_DIR third_party/METIS) + target_include_directories(${PROJECT_NAME} PRIVATE ${METIS_DIR}/include) + set(GKLIB_PATH "${METIS_DIR}/GKlib") + include(${GKLIB_PATH}/GKlibSystem.cmake) + include_directories(${GKLIB_PATH}) + include_directories("${METIS_DIR}/include") + add_subdirectory("${METIS_DIR}/libmetis") + target_link_libraries(${PROJECT_NAME} PRIVATE metis) +endif() find_package(Torch REQUIRED) target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES}) @@ -120,4 +130,6 @@ set_target_properties(${PROJECT_NAME} PROPERTIES # Cmake creates *.dylib by default, but python expects *.so by default if (APPLE) set_property(TARGET ${PROJECT_NAME} PROPERTY SUFFIX .so) +elseif (MSVC AND USE_PYTHON) + set_property(TARGET ${PROJECT_NAME} PROPERTY SUFFIX .pyd) endif() diff --git a/README.md b/README.md index 5834e651b..81ec7fb0d 100644 --- a/README.md +++ b/README.md @@ -37,31 +37,31 @@ The following combinations are supported: | PyTorch 2.2 | `cpu` | `cu102` | `cu113` | `cu116` | `cu117` | `cu118` | `cu121` | |--------------|-------|---------|---------|---------|---------|---------|---------| | **Linux** | ✅ | | | | | ✅ | ✅ | -| **Windows** | | | | | | | | +| **Windows** | ✅ | | | | | ✅ | ✅ | | **macOS** | ✅ | | | | | | | | PyTorch 2.1 | `cpu` | `cu102` | `cu113` | `cu116` | `cu117` | `cu118` | `cu121` | |--------------|-------|---------|---------|---------|---------|---------|---------| | **Linux** | ✅ | | | | | ✅ | ✅ | -| **Windows** | | | | | | | | +| **Windows** | ✅ | | | | | ✅ | ✅ | | **macOS** | ✅ | | | | | | | | PyTorch 2.0 | `cpu` | `cu102` | `cu113` | `cu116` | `cu117` | `cu118` | `cu121` | |--------------|-------|---------|---------|---------|---------|---------|---------| | **Linux** | ✅ | | | | ✅ | ✅ | | -| **Windows** | | | | | | | | +| **Windows** | ✅ | | | | ✅ | ✅ | | | **macOS** | ✅ | | | | | | | | PyTorch 1.13 | `cpu` | `cu102` | `cu113` | `cu116` | `cu117` | `cu118` | `cu121` | |--------------|-------|---------|---------|---------|---------|---------|---------| | **Linux** | ✅ | | | ✅ | ✅ | | | -| **Windows** | | | | | | | | +| **Windows** | ✅ | | | ✅ | ✅ | | | | **macOS** | ✅ | | | | | | | | PyTorch 1.12 | `cpu` | `cu102` | `cu113` | `cu116` | `cu117` | `cu118` | `cu121` | |--------------|-------|---------|---------|---------|---------|---------| --------| | **Linux** | ✅ | ✅ | ✅ | ✅ | | | | -| **Windows** | | | | | | | | +| **Windows** | ✅ | ✅ | ✅ | ✅ | | | | | **macOS** | ✅ | | | | | | | ### Form nightly diff --git a/pyg_lib/csrc/ops/cpu/matmul_kernel.cpp b/pyg_lib/csrc/ops/cpu/matmul_kernel.cpp index 06a020086..21b41d843 100644 --- a/pyg_lib/csrc/ops/cpu/matmul_kernel.cpp +++ b/pyg_lib/csrc/ops/cpu/matmul_kernel.cpp @@ -86,8 +86,7 @@ void mkl_blas_gemm_batched(const int* m_array, const int* ldc_array, const int group_count, const int* group_size) { - TORCH_INTERNAL_ASSERT(false, - "mkl_blas_gemm_batched: MKL BLAS is not supported"); + TORCH_INTERNAL_ASSERT(false, "MKL BLAS is not supported"); } void mkl_blas_gemm_batched(const int* m_array, @@ -103,8 +102,7 @@ void mkl_blas_gemm_batched(const int* m_array, const int* ldc_array, const int group_count, const int* group_size) { - TORCH_INTERNAL_ASSERT(false, - "mkl_blas_gemm_batched: MKL BLAS is not supported"); + TORCH_INTERNAL_ASSERT(false, "MKL BLAS is not supported"); } #endif @@ -205,6 +203,7 @@ void grouped_matmul_out_kernel_at_impl(const std::vector input, void grouped_matmul_out_kernel_mkl_impl(const std::vector input, const std::vector other, std::vector out) { +#if WITH_MKL_BLAS() // matrix_params using matrix_params = std::tuple; phmap::flat_hash_map> groups; @@ -276,6 +275,7 @@ void grouped_matmul_out_kernel_mkl_impl(const std::vector input, group_sizes.data()); #endif }); +#endif } std::vector grouped_matmul_kernel(const at::TensorList input, @@ -328,6 +328,7 @@ void segment_matmul_out_kernel_mkl_impl(const at::Tensor& input, const at::Tensor& other, at::Tensor& out, const at::IntArrayRef& sizes) { +#if WITH_MKL_BLAS() const int n = other.size(-1); const int k = input.size(-1); const int nk = n * k; @@ -403,6 +404,7 @@ void segment_matmul_out_kernel_mkl_impl(const at::Tensor& input, group_sizes.data()); #endif }); +#endif } at::Tensor segment_matmul_kernel(const at::Tensor& input, diff --git a/pyg_lib/csrc/partition/cpu/metis_kernel.cpp b/pyg_lib/csrc/partition/cpu/metis_kernel.cpp index df516224f..56f6da886 100644 --- a/pyg_lib/csrc/partition/cpu/metis_kernel.cpp +++ b/pyg_lib/csrc/partition/cpu/metis_kernel.cpp @@ -1,7 +1,9 @@ #include #include +#ifndef _WIN32 #include +#endif namespace pyg { namespace partition { @@ -14,6 +16,9 @@ at::Tensor metis_kernel(const at::Tensor& rowptr, const c10::optional& node_weight, const c10::optional& edge_weight, bool recursive) { +#ifdef _WIN32 + TORCH_INTERNAL_ASSERT(false, "METIS not yet supported on Windows"); +#else int64_t nvtxs = rowptr.numel() - 1; int64_t ncon = 1; auto* xadj = rowptr.data_ptr(); @@ -41,6 +46,7 @@ at::Tensor metis_kernel(const at::Tensor& rowptr, } return part; +#endif } } // namespace diff --git a/setup.py b/setup.py index 4ea063ff2..cbb4ecddd 100644 --- a/setup.py +++ b/setup.py @@ -1,11 +1,12 @@ # Environment flags to control different options # -# USE_MKL_BLAS=1 -# enables use of MKL BLAS (requires PyTorch to be built with MKL support) +# - USE_MKL_BLAS=1: +# Enables use of MKL BLAS (requires PyTorch to be built with MKL support) import importlib import os import os.path as osp +import re import subprocess import warnings @@ -19,7 +20,7 @@ class CMakeExtension(Extension): def __init__(self, name, sourcedir=''): Extension.__init__(self, name, sources=[]) - self.sourcedir = os.path.abspath(sourcedir) + self.sourcedir = osp.abspath(sourcedir) class CMakeBuild(build_ext): @@ -40,7 +41,7 @@ def build_extension(self, ext): import torch - extdir = os.path.abspath(osp.dirname(self.get_ext_fullpath(ext.name))) + extdir = osp.abspath(osp.dirname(self.get_ext_fullpath(ext.name))) self.build_type = "DEBUG" if self.debug else "RELEASE" if self.debug is None: if CMakeBuild.check_env_flag("DEBUG"): @@ -60,6 +61,7 @@ def build_extension(self, ext): '-DUSE_PYTHON=ON', f'-DWITH_CUDA={"ON" if WITH_CUDA else "OFF"}', f'-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}', + f'-DCMAKE_RUNTIME_OUTPUT_DIRECTORY={extdir}', f'-DCMAKE_BUILD_TYPE={self.build_type}', f'-DCMAKE_PREFIX_PATH={torch.utils.cmake_prefix_path}', ] @@ -85,26 +87,28 @@ def build_extension(self, ext): cwd=self.build_temp) -def maybe_append_with_mkl(dependencies): - if CMakeBuild.check_env_flag('USE_MKL_BLAS'): - import re +def mkl_dependencies(): + if not CMakeBuild.check_env_flag('USE_MKL_BLAS'): + return [] - import torch - torch_config = torch.__config__.show() - with_mkl_blas = 'BLAS_INFO=mkl' in torch_config - if torch.backends.mkl.is_available() and with_mkl_blas: - product_version = '2023.1.0' - pattern = r'oneAPI Math Kernel Library Version [0-9]{4}\.[0-9]+' - match = re.search(pattern, torch_config) - if match: - product_version = match.group(0).split(' ')[-1] + import torch + + dependencies = [] + torch_config = torch.__config__.show() + with_mkl_blas = 'BLAS_INFO=mkl' in torch_config + if torch.backends.mkl.is_available() and with_mkl_blas: + product_version = '2023.1.0' + pattern = r'oneAPI Math Kernel Library Version [0-9]{4}\.[0-9]+' + match = re.search(pattern, torch_config) + if match: + product_version = match.group(0).split(' ')[-1] + dependencies.append(f'mkl-include=={product_version}') + dependencies.append(f'mkl-static=={product_version}') - dependencies.append(f'mkl-include=={product_version}') - dependencies.append(f'mkl-static=={product_version}') + return dependencies -install_requires = [] -maybe_append_with_mkl(install_requires) +install_requires = [] + mkl_dependencies() triton_requires = [ 'triton',