From c4a131b122036fcbf163a9a375efd500c6508811 Mon Sep 17 00:00:00 2001 From: Michael Wang Date: Tue, 7 Jul 2026 10:21:07 -0700 Subject: [PATCH 1/2] Add native Windows ARM64 CUDA bindings support --- cuda_bindings/build_hooks.py | 5 ++++- cuda_bindings/cuda/bindings/_bindings/loader.cpp | 9 ++++++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/cuda_bindings/build_hooks.py b/cuda_bindings/build_hooks.py index 094d8adfbf1..247b9f2bbba 100644 --- a/cuda_bindings/build_hooks.py +++ b/cuda_bindings/build_hooks.py @@ -375,7 +375,10 @@ def _build_cuda_bindings(debug=False): os.path.dirname(sysconfig.get_path("include")), ] + include_path_list library_dirs = [sysconfig.get_path("platlib"), os.path.join(os.sys.prefix, "lib")] - cudalib_subdirs = [r"lib\x64"] if sys.platform == "win32" else ["lib64", "lib"] + if sys.platform == "win32": + cudalib_subdirs = [r"lib\arm64"] if sysconfig.get_platform() == "win-arm64" else [r"lib\x64"] + else: + cudalib_subdirs = ["lib64", "lib"] library_dirs.extend(os.path.join(cuda_path, subdir) for subdir in cudalib_subdirs) extra_compile_args = [] diff --git a/cuda_bindings/cuda/bindings/_bindings/loader.cpp b/cuda_bindings/cuda/bindings/_bindings/loader.cpp index a692eddc931..2820fd26d52 100644 --- a/cuda_bindings/cuda/bindings/_bindings/loader.cpp +++ b/cuda_bindings/cuda/bindings/_bindings/loader.cpp @@ -1,4 +1,4 @@ -// SPDX-FileCopyrightText: Copyright (c) 2021-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2021-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE #include @@ -16,6 +16,8 @@ static const size_t sysrootName64_length = (sizeof("System32") - 1); static const char* sysrootName64 = "System32"; static const size_t libcudaName64_length = (sizeof("\\nvcuda64.dll") - 1); static const char* libcudaName64 = "\\nvcuda64.dll"; +static const size_t libcudaNameArm64_length = (sizeof("\\nvcudaa64.dll") - 1); +static const char* libcudaNameArm64 = "\\nvcudaa64.dll"; static const size_t sysrootNameX86_length = (sizeof("SysWOW64") - 1); static const char* sysrootNameX86 = "SysWOW64"; static const size_t libcudaNameX86_length = (sizeof("\\nvcuda32.dll") - 1); @@ -311,8 +313,13 @@ int getCUDALibraryPath(char *libPath, bool isBit64) if (isBit64) { sysrootName_length = sysrootName64_length; sysrootName = sysrootName64; +#if defined(_M_ARM64) || defined(__aarch64__) + libcudaName_length = libcudaNameArm64_length; + libcudaName = libcudaNameArm64; +#else libcudaName_length = libcudaName64_length; libcudaName = libcudaName64; +#endif } else { sysrootName_length = sysrootNameX86_length; From 2b89547e1fd28045a226ed30d13c3c2794b8578d Mon Sep 17 00:00:00 2001 From: Michael Wang Date: Tue, 7 Jul 2026 10:21:20 -0700 Subject: [PATCH 2/2] Build cuda-core AOTI shim for Windows ARM64 --- cuda_core/setup.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/cuda_core/setup.py b/cuda_core/setup.py index bde1fe22fee..e0d745b1f21 100644 --- a/cuda_core/setup.py +++ b/cuda_core/setup.py @@ -27,7 +27,7 @@ def _ensure_compiler_initialized(compiler, plat_name): initialize(plat_name) -def _build_aoti_shim_lib(compiler): +def _build_aoti_shim_lib(compiler, plat_name): # Reuse setuptools' initialized MSVC compiler instead of rediscovering # lib.exe separately in the build backend. lib_exe = getattr(compiler, "lib", None) @@ -35,12 +35,16 @@ def _build_aoti_shim_lib(compiler): raise RuntimeError("MSVC compiler did not expose lib.exe after initialization.") _AOTI_SHIM_LIB_FILE.parent.mkdir(exist_ok=True) + machine = { + "win-amd64": "X64", + "win-arm64": "ARM64", + }.get(plat_name, "X64") compiler.spawn( [ lib_exe, f"/DEF:{_AOTI_SHIM_DEF_FILE}", f"/OUT:{_AOTI_SHIM_LIB_FILE}", - "/MACHINE:X64", + f"/MACHINE:{machine}", ] ) return str(_AOTI_SHIM_LIB_FILE) @@ -58,7 +62,7 @@ def _configure_windows_tensor_bridge(self): continue _ensure_compiler_initialized(self.compiler, self.plat_name) - shim_lib = _build_aoti_shim_lib(self.compiler) + shim_lib = _build_aoti_shim_lib(self.compiler, self.plat_name) link_args = list(ext.extra_link_args or []) if shim_lib not in link_args: ext.extra_link_args = [*link_args, shim_lib]