From c6401735bec33bf07b25e9e011ed5182d60ac372 Mon Sep 17 00:00:00 2001 From: Sebastian Berg Date: Mon, 27 Jul 2026 14:29:19 +0200 Subject: [PATCH] Make sure `` is found and soft-linked In my conda setup, compiling `tensor_map` created linking issues and it seems like the reason this didn't fail here might be that it is missing when it shouldn't be. Adding the `cccl` headers breaks CI the same way, but moving things to `_tensor_map_cccl.pxd` and soft-linking it should then unbreak things again. This now uses CCCL in an *extremely* limited fashion because CCCL has some bugs/issues around validation. (Even if fixed, we still need to support all dtypes in DLPack, though.) --- cuda_core/build_hooks.py | 12 +- .../{tensor_map.cpp => tensor_map_cccl.cpp} | 68 ++++----- cuda_core/cuda/core/_cpp/tensor_map_cccl.h | 12 +- cuda_core/cuda/core/_tensor_map.pyx | 132 ++++++++++++------ cuda_core/cuda/core/_tensor_map_cccl.pxd | 33 +++++ cuda_core/cuda/core/_tensor_map_cccl.pyi | 12 ++ cuda_core/cuda/core/_tensor_map_cccl.pyx | 22 +++ cuda_core/setup.py | 55 ++++++-- cuda_core/tests/test_tensor_map.py | 21 +++ 9 files changed, 268 insertions(+), 99 deletions(-) rename cuda_core/cuda/core/_cpp/{tensor_map.cpp => tensor_map_cccl.cpp} (70%) create mode 100644 cuda_core/cuda/core/_tensor_map_cccl.pxd create mode 100644 cuda_core/cuda/core/_tensor_map_cccl.pyi create mode 100644 cuda_core/cuda/core/_tensor_map_cccl.pyx diff --git a/cuda_core/build_hooks.py b/cuda_core/build_hooks.py index dfd08d56733..dfdc20d2cb9 100644 --- a/cuda_core/build_hooks.py +++ b/cuda_core/build_hooks.py @@ -171,7 +171,17 @@ def get_sources(mod_name): return sources - all_include_dirs = [os.path.join(_get_cuda_path(), "include")] + # Match nvcc: CCCL lives under include/cccl (e.g. ). + cuda_include = os.path.join(_get_cuda_path(), "include") + cccl_include = os.path.join(cuda_include, "cccl") + all_include_dirs = [ + cccl_include, + cuda_include, + ] + # is what enables the CCCL descriptor path in _tensor_map_cccl. + # CCCL versions predating it make that path compile out silently, leaving + # tensor maps to be encoded through the driver API, so report which we get. + print("CCCL found:", os.path.exists(os.path.join(cccl_include, "cuda", "tma"))) extra_compile_args = [] extra_link_args = [] extra_cythonize_kwargs = {} diff --git a/cuda_core/cuda/core/_cpp/tensor_map.cpp b/cuda_core/cuda/core/_cpp/tensor_map_cccl.cpp similarity index 70% rename from cuda_core/cuda/core/_cpp/tensor_map.cpp rename to cuda_core/cuda/core/_cpp/tensor_map_cccl.cpp index df3f7654e54..a757f4f7c8c 100644 --- a/cuda_core/cuda/core/_cpp/tensor_map.cpp +++ b/cuda_core/cuda/core/_cpp/tensor_map_cccl.cpp @@ -6,32 +6,28 @@ #include -#include #include -#if defined(__has_include) -// Older CTK releases do not ship . When it is unavailable we keep -// the CCCL helper compiled out and fall back to the direct driver path. -# if __has_include() -# include -# define CUDA_CORE_HAS_CUDA_TMA 1 -# else -# define CUDA_CORE_HAS_CUDA_TMA 0 -# endif -# if __has_include("dlpack.h") -# include "dlpack.h" -# define CUDA_CORE_HAS_DLPACK_H 1 -# elif __has_include() -# include -# define CUDA_CORE_HAS_DLPACK_H 1 +// Ahead of : CCCL pulls DLPack in itself, and the copy included first +// wins for the whole file. +#include "dlpack.h" + +// arrived in CCCL 3.2; older CTKs use the driver path below. CCCL +// declares the make_tma_descriptor overloads we call only when it found DLPack +// for itself, so ask it rather than inferring. +#if defined(__has_include) && __has_include() +# include +# if defined(_CCCL_HAS_DLPACK) && _CCCL_HAS_DLPACK() +# define CUDA_CORE_HAS_CCCL_TMA 1 # else -# define CUDA_CORE_HAS_DLPACK_H 0 +# define CUDA_CORE_HAS_CCCL_TMA 0 # endif #else -# define CUDA_CORE_HAS_CUDA_TMA 0 -# define CUDA_CORE_HAS_DLPACK_H 0 +# define CUDA_CORE_HAS_CCCL_TMA 0 #endif +#if CUDA_CORE_HAS_CCCL_TMA + static inline void cuda_core_write_err(char* err, size_t cap, const char* msg) noexcept { if (!err || cap == 0) @@ -48,7 +44,7 @@ static inline void cuda_core_write_err(char* err, size_t cap, const char* msg) n err[n] = '\0'; } -int cuda_core_cccl_make_tma_descriptor_tiled( +static int cuda_core_cccl_make_tma_descriptor_tiled_impl( void* out_tensor_map, void* data, int device_type, @@ -68,26 +64,6 @@ int cuda_core_cccl_make_tma_descriptor_tiled( char* err, size_t err_cap) noexcept { -#if !(CUDA_CORE_HAS_CUDA_TMA && CUDA_CORE_HAS_DLPACK_H) - (void)out_tensor_map; - (void)data; - (void)device_type; - (void)device_id; - (void)ndim; - (void)shape; - (void)strides; - (void)dtype_code; - (void)dtype_bits; - (void)dtype_lanes; - (void)box_sizes; - (void)elem_strides; - (void)interleave_layout; - (void)swizzle; - (void)l2_fetch_size; - (void)oob_fill; - cuda_core_write_err(err, err_cap, "CCCL and/or not available at build time"); - return 1; -#else try { if (!out_tensor_map) @@ -140,6 +116,8 @@ int cuda_core_cccl_make_tma_descriptor_tiled( cuda_core_write_err(err, err_cap, nullptr); return 0; } + // cuda::cuda_error derives from std::runtime_error, so this covers both CCCL's + // argument validation and failures of the driver calls it makes. catch (const std::exception& e) { cuda_core_write_err(err, err_cap, e.what()); @@ -150,5 +128,13 @@ int cuda_core_cccl_make_tma_descriptor_tiled( cuda_core_write_err(err, err_cap, "unknown error while building TMA descriptor"); return 1; } -#endif } + +cuda_core_cccl_make_tma_descriptor_tiled_fn cuda_core_cccl_make_tma_descriptor_tiled = + cuda_core_cccl_make_tma_descriptor_tiled_impl; + +#else + +cuda_core_cccl_make_tma_descriptor_tiled_fn cuda_core_cccl_make_tma_descriptor_tiled = nullptr; + +#endif diff --git a/cuda_core/cuda/core/_cpp/tensor_map_cccl.h b/cuda_core/cuda/core/_cpp/tensor_map_cccl.h index 37f0e0fd8d2..0096222f90c 100644 --- a/cuda_core/cuda/core/_cpp/tensor_map_cccl.h +++ b/cuda_core/cuda/core/_cpp/tensor_map_cccl.h @@ -17,15 +17,18 @@ extern "C" { // Build a tiled CUtensorMap using CCCL's cuda::make_tma_descriptor (from ). // // Returns 0 on success; on failure returns non-zero and writes a best-effort -// human-readable message into (err, err_cap) if provided. -int cuda_core_cccl_make_tma_descriptor_tiled( +// human-readable message into (err, err_cap) if provided. CCCL rejects the same +// inputs the driver path rejects, so callers report every failure the way +// cuTensorMapEncodeTiled failures are reported. +typedef int (*cuda_core_cccl_make_tma_descriptor_tiled_fn)( void* out_tensor_map, void* data, int device_type, int device_id, int ndim, const int64_t* shape, // length ndim - const int64_t* strides, // length ndim, or NULL for contiguous + const int64_t* strides, // length ndim, in elements; required (CCCL + // rejects NULL strides for DLPack >= 1.2) uint8_t dtype_code, uint8_t dtype_bits, uint16_t dtype_lanes, @@ -38,6 +41,9 @@ int cuda_core_cccl_make_tma_descriptor_tiled( char* err, size_t err_cap) noexcept; +// NULL when a usable CCCL was not available at compile time. +extern cuda_core_cccl_make_tma_descriptor_tiled_fn cuda_core_cccl_make_tma_descriptor_tiled; + #ifdef __cplusplus } // extern "C" #endif diff --git a/cuda_core/cuda/core/_tensor_map.pyx b/cuda_core/cuda/core/_tensor_map.pyx index 7e059fe7b98..023b8405021 100644 --- a/cuda_core/cuda/core/_tensor_map.pyx +++ b/cuda_core/cuda/core/_tensor_map.pyx @@ -2,9 +2,18 @@ # # SPDX-License-Identifier: Apache-2.0 -from libc.stdint cimport intptr_t, int64_t, uint8_t, uint16_t, uint32_t, uint64_t +from libc.stdint cimport intptr_t, int64_t, uint8_t, uint16_t, uint32_t, uint64_t, uintptr_t from libc.stddef cimport size_t +from libcpp cimport bool as cpp_bool +from libcpp.atomic cimport ( + atomic as std_atomic, + memory_order_acquire, + memory_order_relaxed, + memory_order_release, +) +from cpython.pycapsule cimport PyCapsule_GetName, PyCapsule_GetPointer from cuda.bindings cimport cydriver +from cuda.core._tensor_map_cccl cimport make_tma_descriptor_tiled_t from cuda.core._utils.cuda_utils cimport HANDLE_RETURN from cuda.core._dlpack cimport kDLInt, kDLUInt, kDLFloat, kDLBfloat, _kDLCUDA @@ -15,38 +24,71 @@ from typing import TYPE_CHECKING import numpy from cuda.core._memoryview import StridedMemoryView -from cuda.core._utils.cuda_utils import check_or_create_options +from cuda.core._utils.cuda_utils import CUDAError, check_or_create_options if TYPE_CHECKING: from cuda.core._device import Device -cdef extern from "_cpp/tensor_map_cccl.h": - int cuda_core_cccl_make_tma_descriptor_tiled( - void* out_tensor_map, - void* data, - int device_type, - int device_id, - int ndim, - const int64_t* shape, - const int64_t* strides, - uint8_t dtype_code, - uint8_t dtype_bits, - uint16_t dtype_lanes, - const int* box_sizes, - const int* elem_strides, - int interleave_layout, - int swizzle, - int l2_fetch_size, - int oob_fill, - char* err, - size_t err_cap) nogil - - try: from ml_dtypes import bfloat16 as ml_bfloat16 except ImportError: ml_bfloat16 = None +# CCCL's make_tma_descriptor lives in the optional ``_tensor_map_cccl`` +# extension, which deliberately leaves ``cudaGetErrorString`` (pulled in by +# CCCL's exception formatting) unresolved. Shared cudart must therefore be in +# the process global namespace before that extension is dlopen'ed, which means +# it must not be imported at ``cuda.core`` package import time -- same +# bootstrap shape as ``_memoryview._get_tensor_bridge``. +# +# The pointer is fetched through ``__pyx_capi__`` (see +# ``_resource_handles._get_driver_fn`` / ``_cpp/DESIGN.md``). Only the getter is +# resolved by name: cimporting it would make Cython import the optional +# extension when ``_tensor_map`` loads, defeating the deferral above. +ctypedef make_tma_descriptor_tiled_t (*cccl_get_make_tma_t)() noexcept nogil + +# The resolved pointer is published to other threads the way +# Buffer._mem_attrs_inited is: stored before the flag with release ordering, +# read after the flag with acquire ordering. Threads may race through the slow +# path (resolution is idempotent), so the pointer is atomic as well. +cdef std_atomic[uintptr_t] _cccl_make_tma_fn +cdef std_atomic[cpp_bool] _cccl_tma_resolved + + +cdef make_tma_descriptor_tiled_t _resolve_cccl_make_tma_fn() except? NULL: + """Import the optional CCCL extension and return its function pointer. + + Returns NULL when the CCCL path is unusable on this system, so that callers + encode the descriptor through the driver API instead: the extension is only + built with a working CCCL ````, and importing it requires shared + cudart, which cuda.core does not otherwise depend on. + """ + cdef cccl_get_make_tma_t get_fn + from cuda.pathfinder import DynamicLibNotFoundError, load_nvidia_dynamic_lib + + try: + # pathfinder loads with RTLD_GLOBAL, putting cudaGetErrorString in the + # global namespace before the extension below is dlopen'ed. + load_nvidia_dynamic_lib("cudart") + import cuda.core._tensor_map_cccl as cccl_mod + except (DynamicLibNotFoundError, ImportError): + return NULL + + capsule = cccl_mod.__pyx_capi__["get_make_tma_descriptor_tiled"] + get_fn = PyCapsule_GetPointer( + capsule, PyCapsule_GetName(capsule)) + return get_fn() + + +cdef make_tma_descriptor_tiled_t _get_cccl_make_tma_fn() except? NULL: + """Return the CCCL helper function pointer, or NULL if unavailable.""" + if _cccl_tma_resolved.load(memory_order_acquire): + return _cccl_make_tma_fn.load(memory_order_relaxed) + cdef make_tma_descriptor_tiled_t fn = _resolve_cccl_make_tma_fn() + _cccl_make_tma_fn.store(fn, memory_order_relaxed) + _cccl_tma_resolved.store(True, memory_order_release) + return fn + class TensorMapDataType(enum.IntEnum): """Data types for tensor map descriptors. @@ -616,8 +658,8 @@ cdef class TensorMapDescriptor: cdef bint elem_strides_provided = element_strides is not None element_strides = _validate_element_strides(element_strides, rank) - # Reuse CCCL/libcu++'s DLPack -> CUtensorMap conversion when possible. - # This avoids maintaining a second, independent validation/encoding implementation. + # Reuse CCCL/libcu++'s DLPack -> CUtensorMap conversion when possible, + # otherwise encode through the driver API further below. cdef uint8_t dl_code cdef uint8_t dl_bits cdef uint16_t dl_lanes @@ -631,15 +673,23 @@ cdef class TensorMapDescriptor: cdef int i_cccl cdef int device_type cdef int c_device_id - cdef int dl_device_type - cdef int dl_device_id cdef int c_cccl_interleave_int cdef int c_cccl_swizzle_int cdef int c_cccl_l2_promotion_int cdef int c_cccl_oob_fill_int cdef int rc - if _tma_dtype_to_dlpack(tma_dt, &dl_code, &dl_bits, &dl_lanes): - c_strides_ptr = NULL + cdef int64_t contig_stride + cdef make_tma_descriptor_tiled_t make_tma = NULL + device_type, c_device_id = view.__dlpack_device__() + # Resolve the optional CCCL extension (after loading cudart). + # TODO(seberg): CCCL has various limitations as of writing, see: + # https://github.com/NVIDIA/cccl/issues/10511. + # Because of this, the path is only taken in _very_ narrow conditions. + make_tma = _get_cccl_make_tma_fn() + if (make_tma != NULL + and device_type == _kDLCUDA + and box_dim[0] == box_dim[rank - 1] + and _tma_dtype_to_dlpack(tma_dt, &dl_code, &dl_bits, &dl_lanes)): c_elem_strides_ptr = NULL errbuf[0] = 0 @@ -649,24 +699,27 @@ cdef class TensorMapDescriptor: if elem_strides_provided: c_elem_strides[i_cccl] = element_strides[i_cccl] + # CCCL rejects NULL strides; synthesize C-contiguous if omitted. if view.strides is not None: for i_cccl in range(rank): c_strides[i_cccl] = view.strides[i_cccl] - c_strides_ptr = &c_strides[0] + else: + contig_stride = 1 + for i_cccl in range(rank - 1, -1, -1): + c_strides[i_cccl] = contig_stride + contig_stride *= shape[i_cccl] + c_strides_ptr = &c_strides[0] if elem_strides_provided: c_elem_strides_ptr = &c_elem_strides[0] - dl_device_type, dl_device_id = view.__dlpack_device__() - device_type = dl_device_type - c_device_id = dl_device_id c_cccl_interleave_int = int(interleave) c_cccl_swizzle_int = int(swizzle) c_cccl_l2_promotion_int = int(l2_promotion) c_cccl_oob_fill_int = int(oob_fill) with nogil: - rc = cuda_core_cccl_make_tma_descriptor_tiled( + rc = make_tma( &desc._tensor_map, global_address, device_type, @@ -696,11 +749,12 @@ cdef class TensorMapDescriptor: } return desc + # CUDAError, as for the cuTensorMapEncodeTiled failures the driver + # path below reports: anything reaching either backend has already + # passed this class' own validation, so which one encoded the + # descriptor stays an implementation detail. msg = errbuf[:].split(b"\0", 1)[0].decode("utf-8", errors="replace") - # If CCCL isn't available at build time, fall back to the direct - # driver API path to preserve functionality on older toolchains. - if "not available at build time" not in msg: - raise ValueError(f"Failed to build TMA descriptor via CCCL: {msg}") + raise CUDAError(f"Failed to build TMA descriptor via CCCL: {msg}") cdef int elem_size = _TMA_DATA_TYPE_SIZE[tma_dt] byte_strides = _compute_byte_strides(shape, view.strides, elem_size) diff --git a/cuda_core/cuda/core/_tensor_map_cccl.pxd b/cuda_core/cuda/core/_tensor_map_cccl.pxd new file mode 100644 index 00000000000..8cd2ecf7b28 --- /dev/null +++ b/cuda_core/cuda/core/_tensor_map_cccl.pxd @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from libc.stdint cimport int64_t, uint8_t, uint16_t +from libc.stddef cimport size_t + +# Shared with ``_tensor_map.pyx``, which must cimport the typedef only: +# cimporting the getter makes Cython import ``_tensor_map_cccl`` when +# ``_tensor_map`` loads, which is exactly what the capsule lookup avoids. +ctypedef int (*make_tma_descriptor_tiled_t)( + void* out_tensor_map, + void* data, + int device_type, + int device_id, + int ndim, + const int64_t* shape, + const int64_t* strides, + uint8_t dtype_code, + uint8_t dtype_bits, + uint16_t dtype_lanes, + const int* box_sizes, + const int* elem_strides, + int interleave_layout, + int swizzle, + int l2_fetch_size, + int oob_fill, + char* err, + size_t err_cap) noexcept nogil + +# Exported via ``__pyx_capi__`` for soft-linking from ``_tensor_map.pyx``. +# Returns the CCCL implementation, or NULL when unavailable at C++ compile time. +cdef make_tma_descriptor_tiled_t get_make_tma_descriptor_tiled() noexcept nogil diff --git a/cuda_core/cuda/core/_tensor_map_cccl.pyi b/cuda_core/cuda/core/_tensor_map_cccl.pyi new file mode 100644 index 00000000000..d9b185255d3 --- /dev/null +++ b/cuda_core/cuda/core/_tensor_map_cccl.pyi @@ -0,0 +1,12 @@ +# This file was generated by stubgen-pyx v0.2.6 from cuda_core/cuda/core/_tensor_map_cccl.pyx + +"""Optional CCCL helper for TensorMapDescriptor. + +This extension is imported only after ``load_nvidia_dynamic_lib("cudart")`` so +CCCL's references to ``cudaGetErrorString`` resolve from the globally loaded +shared cudart. It must not be imported at ``cuda.core`` package import time. + +``get_make_tma_descriptor_tiled`` is exported via ``__pyx_capi__`` for +soft-linking from ``_tensor_map.pyx`` (returns real fn or NULL). +""" +from __future__ import annotations \ No newline at end of file diff --git a/cuda_core/cuda/core/_tensor_map_cccl.pyx b/cuda_core/cuda/core/_tensor_map_cccl.pyx new file mode 100644 index 00000000000..b4399e37143 --- /dev/null +++ b/cuda_core/cuda/core/_tensor_map_cccl.pyx @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Optional CCCL helper for TensorMapDescriptor. + +This extension is imported only after ``load_nvidia_dynamic_lib("cudart")`` so +CCCL's references to ``cudaGetErrorString`` resolve from the globally loaded +shared cudart. It must not be imported at ``cuda.core`` package import time. + +``get_make_tma_descriptor_tiled`` is exported via ``__pyx_capi__`` for +soft-linking from ``_tensor_map.pyx`` (returns real fn or NULL). +""" + +from cuda.core._tensor_map_cccl cimport make_tma_descriptor_tiled_t + +cdef extern from "_cpp/tensor_map_cccl.h": + make_tma_descriptor_tiled_t cuda_core_cccl_make_tma_descriptor_tiled + + +cdef make_tma_descriptor_tiled_t get_make_tma_descriptor_tiled() noexcept nogil: + return cuda_core_cccl_make_tma_descriptor_tiled diff --git a/cuda_core/setup.py b/cuda_core/setup.py index e0d745b1f21..970fc7b9742 100644 --- a/cuda_core/setup.py +++ b/cuda_core/setup.py @@ -16,6 +16,9 @@ _AOTI_SHIM_DEF_FILE = _ROOT_DIR / "cuda" / "core" / "_include" / "aoti_shim.def" _AOTI_SHIM_LIB_FILE = _ROOT_DIR / "build" / "aoti_shim.lib" _TENSOR_BRIDGE_EXT_NAME = "cuda.core._tensor_bridge" +_CUDART_SHIM_DEF_FILE = _ROOT_DIR / "build" / "cudart_shim.def" +_CUDART_SHIM_LIB_FILE = _ROOT_DIR / "build" / "cudart_shim.lib" +_TENSOR_MAP_CCCL_EXT_NAME = "cuda.core._tensor_map_cccl" def _ensure_compiler_initialized(compiler, plat_name): @@ -27,14 +30,14 @@ def _ensure_compiler_initialized(compiler, plat_name): initialize(plat_name) -def _build_aoti_shim_lib(compiler, plat_name): +def _build_stub_import_lib(compiler, plat_name, def_file, lib_file): # Reuse setuptools' initialized MSVC compiler instead of rediscovering # lib.exe separately in the build backend. lib_exe = getattr(compiler, "lib", None) if not lib_exe: raise RuntimeError("MSVC compiler did not expose lib.exe after initialization.") - _AOTI_SHIM_LIB_FILE.parent.mkdir(exist_ok=True) + lib_file.parent.mkdir(exist_ok=True) machine = { "win-amd64": "X64", "win-arm64": "ARM64", @@ -42,37 +45,59 @@ def _build_aoti_shim_lib(compiler, plat_name): compiler.spawn( [ lib_exe, - f"/DEF:{_AOTI_SHIM_DEF_FILE}", - f"/OUT:{_AOTI_SHIM_LIB_FILE}", + f"/DEF:{def_file}", + f"/OUT:{lib_file}", f"/MACHINE:{machine}", ] ) - return str(_AOTI_SHIM_LIB_FILE) + return str(lib_file) -class build_ext(_build_ext): # noqa: N801 - def _configure_windows_tensor_bridge(self): - if os.name != "nt" or getattr(self.compiler, "compiler_type", None) != "msvc": - return +def _write_cudart_shim_def(): + """Emit the .def naming the cudart symbols _tensor_map_cccl resolves lazily. + + CCCL's exception formatting calls cudaGetErrorString, so the extension is + left with that one unresolved cudart symbol. Unlike aoti_shim.def this file + is generated, because the DLL name carries the CUDA major version. + """ + cuda_major = build_hooks._determine_cuda_major_version() + _CUDART_SHIM_DEF_FILE.parent.mkdir(exist_ok=True) + _CUDART_SHIM_DEF_FILE.write_text( + f"LIBRARY cudart64_{cuda_major}.dll\nEXPORTS\n cudaGetErrorString\n", encoding="utf-8" + ) + return _CUDART_SHIM_DEF_FILE + - # _tensor_bridge imports AOTI symbols from torch_cpu.dll, which on - # Windows requires a stub import library for the MSVC linker. +class build_ext(_build_ext): # noqa: N801 + def _attach_stub_import_lib(self, ext_name, def_file, lib_file): for ext in self.extensions: - if ext.name != _TENSOR_BRIDGE_EXT_NAME: + if ext.name != ext_name: continue _ensure_compiler_initialized(self.compiler, self.plat_name) - shim_lib = _build_aoti_shim_lib(self.compiler, self.plat_name) + shim_lib = _build_stub_import_lib(self.compiler, self.plat_name, def_file, lib_file) link_args = list(ext.extra_link_args or []) if shim_lib not in link_args: ext.extra_link_args = [*link_args, shim_lib] return - raise RuntimeError(f"Failed to find extension {_TENSOR_BRIDGE_EXT_NAME!r} for Windows build.") + raise RuntimeError(f"Failed to find extension {ext_name!r} for Windows build.") + + def _configure_windows_stub_imports(self): + if os.name != "nt" or getattr(self.compiler, "compiler_type", None) != "msvc": + return + + # _tensor_bridge imports AOTI symbols from torch_cpu.dll and + # _tensor_map_cccl imports cudaGetErrorString from cudart. Both resolve + # at runtime from a DLL that is already loaded by the time the + # extension is imported, but the MSVC linker still demands an import + # library, so synthesize a minimal one from a .def file. + self._attach_stub_import_lib(_TENSOR_BRIDGE_EXT_NAME, _AOTI_SHIM_DEF_FILE, _AOTI_SHIM_LIB_FILE) + self._attach_stub_import_lib(_TENSOR_MAP_CCCL_EXT_NAME, _write_cudart_shim_def(), _CUDART_SHIM_LIB_FILE) def build_extensions(self): self.parallel = nthreads - self._configure_windows_tensor_bridge() + self._configure_windows_stub_imports() super().build_extensions() diff --git a/cuda_core/tests/test_tensor_map.py b/cuda_core/tests/test_tensor_map.py index 7abbaadb483..5de41179927 100644 --- a/cuda_core/tests/test_tensor_map.py +++ b/cuda_core/tests/test_tensor_map.py @@ -250,6 +250,27 @@ def test_from_tiled_with_oob_fill(self, dev, skip_if_no_tma): ) assert desc is not None + @pytest.mark.agent_authored(model="claude-opus-5") + @pytest.mark.parametrize( + "data_type", + [ + TensorMapDataType.FLOAT32_FTZ, + TensorMapDataType.TFLOAT32, + TensorMapDataType.TFLOAT32_FTZ, + ], + ids=lambda data_type: data_type.name, + ) + def test_from_tiled_data_types_without_dlpack_spelling(self, dev, skip_if_no_tma, data_type): + # These have no DLPack spelling, so they are encoded through the driver + # API rather than CCCL's make_tma_descriptor. + buf = dev.allocate(64 * 64 * 4, stream=dev.default_stream) + tensor = _DeviceArray(buf, (64, 64)) + desc = _as_view(tensor).as_tensor_map( + box_dim=(32, 32), + data_type=data_type, + ) + assert desc is not None + class TestTensorMapDescriptorValidation: """Test validation in TensorMapDescriptor factory methods."""