Skip to content

Commit c640173

Browse files
sebergseberg
committed
Make sure <cccl/tma> 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.)
1 parent 94e3a76 commit c640173

9 files changed

Lines changed: 268 additions & 99 deletions

File tree

‎cuda_core/build_hooks.py‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,17 @@ def get_sources(mod_name):
171171

172172
return sources
173173

174-
all_include_dirs = [os.path.join(_get_cuda_path(), "include")]
174+
# Match nvcc: CCCL lives under include/cccl (e.g. <cuda/tma>).
175+
cuda_include = os.path.join(_get_cuda_path(), "include")
176+
cccl_include = os.path.join(cuda_include, "cccl")
177+
all_include_dirs = [
178+
cccl_include,
179+
cuda_include,
180+
]
181+
# <cuda/tma> is what enables the CCCL descriptor path in _tensor_map_cccl.
182+
# CCCL versions predating it make that path compile out silently, leaving
183+
# tensor maps to be encoded through the driver API, so report which we get.
184+
print("CCCL <cuda/tma> found:", os.path.exists(os.path.join(cccl_include, "cuda", "tma")))
175185
extra_compile_args = []
176186
extra_link_args = []
177187
extra_cythonize_kwargs = {}

cuda_core/cuda/core/_cpp/tensor_map.cpp renamed to cuda_core/cuda/core/_cpp/tensor_map_cccl.cpp

Lines changed: 27 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -6,32 +6,28 @@
66

77
#include <string.h>
88

9-
#include <algorithm>
109
#include <exception>
1110

12-
#if defined(__has_include)
13-
// Older CTK releases do not ship <cuda/tma>. When it is unavailable we keep
14-
// the CCCL helper compiled out and fall back to the direct driver path.
15-
# if __has_include(<cuda/tma>)
16-
# include <cuda/tma>
17-
# define CUDA_CORE_HAS_CUDA_TMA 1
18-
# else
19-
# define CUDA_CORE_HAS_CUDA_TMA 0
20-
# endif
21-
# if __has_include("dlpack.h")
22-
# include "dlpack.h"
23-
# define CUDA_CORE_HAS_DLPACK_H 1
24-
# elif __has_include(<dlpack/dlpack.h>)
25-
# include <dlpack/dlpack.h>
26-
# define CUDA_CORE_HAS_DLPACK_H 1
11+
// Ahead of <cuda/tma>: CCCL pulls DLPack in itself, and the copy included first
12+
// wins for the whole file.
13+
#include "dlpack.h"
14+
15+
// <cuda/tma> arrived in CCCL 3.2; older CTKs use the driver path below. CCCL
16+
// declares the make_tma_descriptor overloads we call only when it found DLPack
17+
// for itself, so ask it rather than inferring.
18+
#if defined(__has_include) && __has_include(<cuda/tma>)
19+
# include <cuda/tma>
20+
# if defined(_CCCL_HAS_DLPACK) && _CCCL_HAS_DLPACK()
21+
# define CUDA_CORE_HAS_CCCL_TMA 1
2722
# else
28-
# define CUDA_CORE_HAS_DLPACK_H 0
23+
# define CUDA_CORE_HAS_CCCL_TMA 0
2924
# endif
3025
#else
31-
# define CUDA_CORE_HAS_CUDA_TMA 0
32-
# define CUDA_CORE_HAS_DLPACK_H 0
26+
# define CUDA_CORE_HAS_CCCL_TMA 0
3327
#endif
3428

29+
#if CUDA_CORE_HAS_CCCL_TMA
30+
3531
static inline void cuda_core_write_err(char* err, size_t cap, const char* msg) noexcept
3632
{
3733
if (!err || cap == 0)
@@ -48,7 +44,7 @@ static inline void cuda_core_write_err(char* err, size_t cap, const char* msg) n
4844
err[n] = '\0';
4945
}
5046

51-
int cuda_core_cccl_make_tma_descriptor_tiled(
47+
static int cuda_core_cccl_make_tma_descriptor_tiled_impl(
5248
void* out_tensor_map,
5349
void* data,
5450
int device_type,
@@ -68,26 +64,6 @@ int cuda_core_cccl_make_tma_descriptor_tiled(
6864
char* err,
6965
size_t err_cap) noexcept
7066
{
71-
#if !(CUDA_CORE_HAS_CUDA_TMA && CUDA_CORE_HAS_DLPACK_H)
72-
(void)out_tensor_map;
73-
(void)data;
74-
(void)device_type;
75-
(void)device_id;
76-
(void)ndim;
77-
(void)shape;
78-
(void)strides;
79-
(void)dtype_code;
80-
(void)dtype_bits;
81-
(void)dtype_lanes;
82-
(void)box_sizes;
83-
(void)elem_strides;
84-
(void)interleave_layout;
85-
(void)swizzle;
86-
(void)l2_fetch_size;
87-
(void)oob_fill;
88-
cuda_core_write_err(err, err_cap, "CCCL <cuda/tma> and/or <dlpack/dlpack.h> not available at build time");
89-
return 1;
90-
#else
9167
try
9268
{
9369
if (!out_tensor_map)
@@ -140,6 +116,8 @@ int cuda_core_cccl_make_tma_descriptor_tiled(
140116
cuda_core_write_err(err, err_cap, nullptr);
141117
return 0;
142118
}
119+
// cuda::cuda_error derives from std::runtime_error, so this covers both CCCL's
120+
// argument validation and failures of the driver calls it makes.
143121
catch (const std::exception& e)
144122
{
145123
cuda_core_write_err(err, err_cap, e.what());
@@ -150,5 +128,13 @@ int cuda_core_cccl_make_tma_descriptor_tiled(
150128
cuda_core_write_err(err, err_cap, "unknown error while building TMA descriptor");
151129
return 1;
152130
}
153-
#endif
154131
}
132+
133+
cuda_core_cccl_make_tma_descriptor_tiled_fn cuda_core_cccl_make_tma_descriptor_tiled =
134+
cuda_core_cccl_make_tma_descriptor_tiled_impl;
135+
136+
#else
137+
138+
cuda_core_cccl_make_tma_descriptor_tiled_fn cuda_core_cccl_make_tma_descriptor_tiled = nullptr;
139+
140+
#endif

‎cuda_core/cuda/core/_cpp/tensor_map_cccl.h‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,18 @@ extern "C" {
1717
// Build a tiled CUtensorMap using CCCL's cuda::make_tma_descriptor (from <cuda/tma>).
1818
//
1919
// Returns 0 on success; on failure returns non-zero and writes a best-effort
20-
// human-readable message into (err, err_cap) if provided.
21-
int cuda_core_cccl_make_tma_descriptor_tiled(
20+
// human-readable message into (err, err_cap) if provided. CCCL rejects the same
21+
// inputs the driver path rejects, so callers report every failure the way
22+
// cuTensorMapEncodeTiled failures are reported.
23+
typedef int (*cuda_core_cccl_make_tma_descriptor_tiled_fn)(
2224
void* out_tensor_map,
2325
void* data,
2426
int device_type,
2527
int device_id,
2628
int ndim,
2729
const int64_t* shape, // length ndim
28-
const int64_t* strides, // length ndim, or NULL for contiguous
30+
const int64_t* strides, // length ndim, in elements; required (CCCL
31+
// rejects NULL strides for DLPack >= 1.2)
2932
uint8_t dtype_code,
3033
uint8_t dtype_bits,
3134
uint16_t dtype_lanes,
@@ -38,6 +41,9 @@ int cuda_core_cccl_make_tma_descriptor_tiled(
3841
char* err,
3942
size_t err_cap) noexcept;
4043

44+
// NULL when a usable CCCL <cuda/tma> was not available at compile time.
45+
extern cuda_core_cccl_make_tma_descriptor_tiled_fn cuda_core_cccl_make_tma_descriptor_tiled;
46+
4147
#ifdef __cplusplus
4248
} // extern "C"
4349
#endif

‎cuda_core/cuda/core/_tensor_map.pyx‎

Lines changed: 93 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,18 @@
22
#
33
# SPDX-License-Identifier: Apache-2.0
44

5-
from libc.stdint cimport intptr_t, int64_t, uint8_t, uint16_t, uint32_t, uint64_t
5+
from libc.stdint cimport intptr_t, int64_t, uint8_t, uint16_t, uint32_t, uint64_t, uintptr_t
66
from libc.stddef cimport size_t
7+
from libcpp cimport bool as cpp_bool
8+
from libcpp.atomic cimport (
9+
atomic as std_atomic,
10+
memory_order_acquire,
11+
memory_order_relaxed,
12+
memory_order_release,
13+
)
14+
from cpython.pycapsule cimport PyCapsule_GetName, PyCapsule_GetPointer
715
from cuda.bindings cimport cydriver
16+
from cuda.core._tensor_map_cccl cimport make_tma_descriptor_tiled_t
817
from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
918
from cuda.core._dlpack cimport kDLInt, kDLUInt, kDLFloat, kDLBfloat, _kDLCUDA
1019

@@ -15,38 +24,71 @@ from typing import TYPE_CHECKING
1524
import numpy
1625

1726
from cuda.core._memoryview import StridedMemoryView
18-
from cuda.core._utils.cuda_utils import check_or_create_options
27+
from cuda.core._utils.cuda_utils import CUDAError, check_or_create_options
1928

2029
if TYPE_CHECKING:
2130
from cuda.core._device import Device
2231

23-
cdef extern from "_cpp/tensor_map_cccl.h":
24-
int cuda_core_cccl_make_tma_descriptor_tiled(
25-
void* out_tensor_map,
26-
void* data,
27-
int device_type,
28-
int device_id,
29-
int ndim,
30-
const int64_t* shape,
31-
const int64_t* strides,
32-
uint8_t dtype_code,
33-
uint8_t dtype_bits,
34-
uint16_t dtype_lanes,
35-
const int* box_sizes,
36-
const int* elem_strides,
37-
int interleave_layout,
38-
int swizzle,
39-
int l2_fetch_size,
40-
int oob_fill,
41-
char* err,
42-
size_t err_cap) nogil
43-
44-
4532
try:
4633
from ml_dtypes import bfloat16 as ml_bfloat16
4734
except ImportError:
4835
ml_bfloat16 = None
4936

37+
# CCCL's make_tma_descriptor lives in the optional ``_tensor_map_cccl``
38+
# extension, which deliberately leaves ``cudaGetErrorString`` (pulled in by
39+
# CCCL's exception formatting) unresolved. Shared cudart must therefore be in
40+
# the process global namespace before that extension is dlopen'ed, which means
41+
# it must not be imported at ``cuda.core`` package import time -- same
42+
# bootstrap shape as ``_memoryview._get_tensor_bridge``.
43+
#
44+
# The pointer is fetched through ``__pyx_capi__`` (see
45+
# ``_resource_handles._get_driver_fn`` / ``_cpp/DESIGN.md``). Only the getter is
46+
# resolved by name: cimporting it would make Cython import the optional
47+
# extension when ``_tensor_map`` loads, defeating the deferral above.
48+
ctypedef make_tma_descriptor_tiled_t (*cccl_get_make_tma_t)() noexcept nogil
49+
50+
# The resolved pointer is published to other threads the way
51+
# Buffer._mem_attrs_inited is: stored before the flag with release ordering,
52+
# read after the flag with acquire ordering. Threads may race through the slow
53+
# path (resolution is idempotent), so the pointer is atomic as well.
54+
cdef std_atomic[uintptr_t] _cccl_make_tma_fn
55+
cdef std_atomic[cpp_bool] _cccl_tma_resolved
56+
57+
58+
cdef make_tma_descriptor_tiled_t _resolve_cccl_make_tma_fn() except? NULL:
59+
"""Import the optional CCCL extension and return its function pointer.
60+
61+
Returns NULL when the CCCL path is unusable on this system, so that callers
62+
encode the descriptor through the driver API instead: the extension is only
63+
built with a working CCCL ``<cuda/tma>``, and importing it requires shared
64+
cudart, which cuda.core does not otherwise depend on.
65+
"""
66+
cdef cccl_get_make_tma_t get_fn
67+
from cuda.pathfinder import DynamicLibNotFoundError, load_nvidia_dynamic_lib
68+
69+
try:
70+
# pathfinder loads with RTLD_GLOBAL, putting cudaGetErrorString in the
71+
# global namespace before the extension below is dlopen'ed.
72+
load_nvidia_dynamic_lib("cudart")
73+
import cuda.core._tensor_map_cccl as cccl_mod
74+
except (DynamicLibNotFoundError, ImportError):
75+
return NULL
76+
77+
capsule = cccl_mod.__pyx_capi__["get_make_tma_descriptor_tiled"]
78+
get_fn = <cccl_get_make_tma_t>PyCapsule_GetPointer(
79+
capsule, PyCapsule_GetName(capsule))
80+
return get_fn()
81+
82+
83+
cdef make_tma_descriptor_tiled_t _get_cccl_make_tma_fn() except? NULL:
84+
"""Return the CCCL helper function pointer, or NULL if unavailable."""
85+
if _cccl_tma_resolved.load(memory_order_acquire):
86+
return <make_tma_descriptor_tiled_t>_cccl_make_tma_fn.load(memory_order_relaxed)
87+
cdef make_tma_descriptor_tiled_t fn = _resolve_cccl_make_tma_fn()
88+
_cccl_make_tma_fn.store(<uintptr_t>fn, memory_order_relaxed)
89+
_cccl_tma_resolved.store(True, memory_order_release)
90+
return fn
91+
5092

5193
class TensorMapDataType(enum.IntEnum):
5294
"""Data types for tensor map descriptors.
@@ -616,8 +658,8 @@ cdef class TensorMapDescriptor:
616658
cdef bint elem_strides_provided = element_strides is not None
617659
element_strides = _validate_element_strides(element_strides, rank)
618660

619-
# Reuse CCCL/libcu++'s DLPack -> CUtensorMap conversion when possible.
620-
# This avoids maintaining a second, independent validation/encoding implementation.
661+
# Reuse CCCL/libcu++'s DLPack -> CUtensorMap conversion when possible,
662+
# otherwise encode through the driver API further below.
621663
cdef uint8_t dl_code
622664
cdef uint8_t dl_bits
623665
cdef uint16_t dl_lanes
@@ -631,15 +673,23 @@ cdef class TensorMapDescriptor:
631673
cdef int i_cccl
632674
cdef int device_type
633675
cdef int c_device_id
634-
cdef int dl_device_type
635-
cdef int dl_device_id
636676
cdef int c_cccl_interleave_int
637677
cdef int c_cccl_swizzle_int
638678
cdef int c_cccl_l2_promotion_int
639679
cdef int c_cccl_oob_fill_int
640680
cdef int rc
641-
if _tma_dtype_to_dlpack(tma_dt, &dl_code, &dl_bits, &dl_lanes):
642-
c_strides_ptr = NULL
681+
cdef int64_t contig_stride
682+
cdef make_tma_descriptor_tiled_t make_tma = NULL
683+
device_type, c_device_id = view.__dlpack_device__()
684+
# Resolve the optional CCCL extension (after loading cudart).
685+
# TODO(seberg): CCCL has various limitations as of writing, see:
686+
# https://github.com/NVIDIA/cccl/issues/10511.
687+
# Because of this, the path is only taken in _very_ narrow conditions.
688+
make_tma = _get_cccl_make_tma_fn()
689+
if (make_tma != NULL
690+
and device_type == _kDLCUDA
691+
and box_dim[0] == box_dim[rank - 1]
692+
and _tma_dtype_to_dlpack(tma_dt, &dl_code, &dl_bits, &dl_lanes)):
643693
c_elem_strides_ptr = NULL
644694
errbuf[0] = 0
645695

@@ -649,24 +699,27 @@ cdef class TensorMapDescriptor:
649699
if elem_strides_provided:
650700
c_elem_strides[i_cccl] = <int>element_strides[i_cccl]
651701

702+
# CCCL rejects NULL strides; synthesize C-contiguous if omitted.
652703
if view.strides is not None:
653704
for i_cccl in range(rank):
654705
c_strides[i_cccl] = <int64_t>view.strides[i_cccl]
655-
c_strides_ptr = &c_strides[0]
706+
else:
707+
contig_stride = 1
708+
for i_cccl in range(rank - 1, -1, -1):
709+
c_strides[i_cccl] = contig_stride
710+
contig_stride *= <int64_t>shape[i_cccl]
711+
c_strides_ptr = &c_strides[0]
656712

657713
if elem_strides_provided:
658714
c_elem_strides_ptr = &c_elem_strides[0]
659715

660-
dl_device_type, dl_device_id = view.__dlpack_device__()
661-
device_type = dl_device_type
662-
c_device_id = dl_device_id
663716
c_cccl_interleave_int = int(interleave)
664717
c_cccl_swizzle_int = int(swizzle)
665718
c_cccl_l2_promotion_int = int(l2_promotion)
666719
c_cccl_oob_fill_int = int(oob_fill)
667720

668721
with nogil:
669-
rc = cuda_core_cccl_make_tma_descriptor_tiled(
722+
rc = make_tma(
670723
<void*>&desc._tensor_map,
671724
<void*>global_address,
672725
device_type,
@@ -696,11 +749,12 @@ cdef class TensorMapDescriptor:
696749
}
697750
return desc
698751

752+
# CUDAError, as for the cuTensorMapEncodeTiled failures the driver
753+
# path below reports: anything reaching either backend has already
754+
# passed this class' own validation, so which one encoded the
755+
# descriptor stays an implementation detail.
699756
msg = errbuf[:].split(b"\0", 1)[0].decode("utf-8", errors="replace")
700-
# If CCCL isn't available at build time, fall back to the direct
701-
# driver API path to preserve functionality on older toolchains.
702-
if "not available at build time" not in msg:
703-
raise ValueError(f"Failed to build TMA descriptor via CCCL: {msg}")
757+
raise CUDAError(f"Failed to build TMA descriptor via CCCL: {msg}")
704758

705759
cdef int elem_size = _TMA_DATA_TYPE_SIZE[tma_dt]
706760
byte_strides = _compute_byte_strides(shape, view.strides, elem_size)

‎cuda_core/cuda/core/_tensor_map_cccl.pxd‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
#
3+
# SPDX-License-Identifier: Apache-2.0
4+
5+
from libc.stdint cimport int64_t, uint8_t, uint16_t
6+
from libc.stddef cimport size_t
7+
8+
# Shared with ``_tensor_map.pyx``, which must cimport the typedef only:
9+
# cimporting the getter makes Cython import ``_tensor_map_cccl`` when
10+
# ``_tensor_map`` loads, which is exactly what the capsule lookup avoids.
11+
ctypedef int (*make_tma_descriptor_tiled_t)(
12+
void* out_tensor_map,
13+
void* data,
14+
int device_type,
15+
int device_id,
16+
int ndim,
17+
const int64_t* shape,
18+
const int64_t* strides,
19+
uint8_t dtype_code,
20+
uint8_t dtype_bits,
21+
uint16_t dtype_lanes,
22+
const int* box_sizes,
23+
const int* elem_strides,
24+
int interleave_layout,
25+
int swizzle,
26+
int l2_fetch_size,
27+
int oob_fill,
28+
char* err,
29+
size_t err_cap) noexcept nogil
30+
31+
# Exported via ``__pyx_capi__`` for soft-linking from ``_tensor_map.pyx``.
32+
# Returns the CCCL implementation, or NULL when unavailable at C++ compile time.
33+
cdef make_tma_descriptor_tiled_t get_make_tma_descriptor_tiled() noexcept nogil

0 commit comments

Comments
 (0)