diff --git a/AGENTS.md b/AGENTS.md index 75b58e12f..553598023 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -45,6 +45,7 @@ Python (editable; compiles the pybind11 extension via CMake): ```bash pip install -e . # core graph API only pip install -e ".[cutedsl]" # + OSS CuTeDSL kernels (nvidia-cutlass-dsl, cuda-python, tvm-ffi; framework-neutral) +pip install -e ".[cutile]" # + the cuTile linear-attention engines (cuda-tile; needs a system tileiras) pip install --group torch # + torch for the CuTeDSL APIs (torch, torch-c-dlpack-ext) pip install --group jax # + jax for the CuTeDSL APIs (jax >= 0.5; XLA entry points via cutlass.jax) ``` diff --git a/include/cudnn_frontend/graph_interface.h b/include/cudnn_frontend/graph_interface.h index 87e65fc2a..c0a203f99 100644 --- a/include/cudnn_frontend/graph_interface.h +++ b/include/cudnn_frontend/graph_interface.h @@ -1454,6 +1454,17 @@ class Graph : public ICudnn, public INode { std::vector const &override_uids = {}, std::vector> const &override_shapes = {}, std::vector> const &override_strides = {}) const { + // The driver-API engines read the calling THREAD's context stack, and a + // thread that has done no CUDA work yet (a PyTorch autograd worker) has + // none. All roads reach here, so every execute is covered; the steady + // state is one cuCtxGetCurrent, and the stream query only runs when a + // context actually has to be established. + if (!detail::has_current_context()) { + cudaStream_t stream = nullptr; + detail::get_stream(handle, &stream); + detail::ensure_current_context(stream); + } + // Lazy init: prepare template if not done (e.g. deserialized graphs, build_plan_at_index) if (!varpack_prep_state->prepared.load(std::memory_order_acquire)) { CHECK_CUDNN_FRONTEND_ERROR(const_cast(this)->prepare_variant_pack_template()); diff --git a/include/cudnn_frontend_shim.h b/include/cudnn_frontend_shim.h index 7525690b4..bd690b30c 100644 --- a/include/cudnn_frontend_shim.h +++ b/include/cudnn_frontend_shim.h @@ -417,6 +417,140 @@ cuda_get_device(int *device) { NV_FE_CALL_TO_CUDA(cuda_get_device, cudaGetDevice, device); } +// Driver entry points resolved through the runtime, so the front end never links +// libcuda.so -- the same approach as cu_tensor_map_encode_tiled(). Returns +// nullptr when the entry point is unavailable. +inline void * +get_driver_entry_point(const char *name) { +#if defined(CUDART_VERSION) && CUDART_VERSION >= 12000 + void *pfn = nullptr; + cudaDriverEntryPointQueryResult query_result; + cudaError_t err; +#if defined NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING +#if CUDART_VERSION >= 12080 + using GetEntryPointFn = + cudaError_t (*)(const char *, void **, unsigned int, int, cudaDriverEntryPointQueryResult *); + const char *resolver = "cudaGetDriverEntryPointByVersion"; +#else + using GetEntryPointFn = + cudaError_t (*)(const char *, void **, unsigned long long, cudaDriverEntryPointQueryResult *); + const char *resolver = "cudaGetDriverEntryPoint"; +#endif + GetEntryPointFn get_entry_point = nullptr; + // get_cuda_symbol throws when the library or symbol is missing; this lookup is + // best-effort and runs inside a static initializer, so it must not escape. +#ifndef NV_CUDNN_DISABLE_EXCEPTION + try { + get_entry_point = reinterpret_cast(get_cuda_symbol(CudaLibrary::CUDART, resolver)); + } catch (...) { + return nullptr; + } +#else + get_entry_point = reinterpret_cast(get_cuda_symbol(CudaLibrary::CUDART, resolver)); +#endif + if (get_entry_point == nullptr) { + return nullptr; + } +#if CUDART_VERSION >= 12080 + err = get_entry_point(name, &pfn, 12000, cudaEnableDefault, &query_result); +#else + err = get_entry_point(name, &pfn, cudaEnableDefault, &query_result); +#endif +#elif CUDART_VERSION >= 12080 + err = cudaGetDriverEntryPointByVersion(name, &pfn, 12000, cudaEnableDefault, &query_result); +#else + err = cudaGetDriverEntryPoint(name, &pfn, cudaEnableDefault, &query_result); +#endif + if (err != cudaSuccess || query_result != cudaDriverEntryPointSuccess) { + return nullptr; + } + return pfn; +#else + (void)name; + return nullptr; +#endif +} + +// True when the calling thread already has a context. One driver call, so the +// execute path can skip the rest -- including the stream query -- in the steady +// state. +inline bool +has_current_context() { +#if defined(CUDART_VERSION) && CUDART_VERSION >= 12000 + using PfnCtxGetCurrent = CUresult(CUDAAPI *)(CUcontext *); + static const auto ctx_get_current = reinterpret_cast(get_driver_entry_point("cuCtxGetCurrent")); + if (ctx_get_current == nullptr) { + return true; // cannot tell; leave the thread alone + } + CUcontext current = nullptr; + return ctx_get_current(¤t) == CUDA_SUCCESS && current != nullptr; +#else + return true; +#endif +} + +// Bind a driver context to the calling thread when it has none: a driver-API +// launch reads that stack, and a thread that has done no CUDA work has nothing +// on it (a PyTorch autograd worker). A bound context is left alone -- it is +// process-wide and the caller chose it. A real stream names the right context; +// the default-stream handles name none, so the runtime's device decides there. +// Best-effort: what this cannot establish fails at the launch. +inline void +ensure_current_context(cudaStream_t stream) { +#if defined(CUDART_VERSION) && CUDART_VERSION >= 12000 + using PfnCtxGetCurrent = CUresult(CUDAAPI *)(CUcontext *); + using PfnCtxSetCurrent = CUresult(CUDAAPI *)(CUcontext); + using PfnStreamGetCtx = CUresult(CUDAAPI *)(CUstream, CUcontext *); + using PfnDeviceGet = CUresult(CUDAAPI *)(CUdevice *, int); + using PfnPrimaryCtxRetain = CUresult(CUDAAPI *)(CUcontext *, CUdevice); + + // Thread-safe static initialization (C++11): the lookups happen once. + static const auto ctx_get_current = reinterpret_cast(get_driver_entry_point("cuCtxGetCurrent")); + static const auto ctx_set_current = reinterpret_cast(get_driver_entry_point("cuCtxSetCurrent")); + static const auto stream_get_ctx = reinterpret_cast(get_driver_entry_point("cuStreamGetCtx")); + static const auto device_get = reinterpret_cast(get_driver_entry_point("cuDeviceGet")); + static const auto primary_retain = + reinterpret_cast(get_driver_entry_point("cuDevicePrimaryCtxRetain")); + if (ctx_get_current == nullptr || ctx_set_current == nullptr) { + return; + } + + CUcontext current = nullptr; + if (ctx_get_current(¤t) == CUDA_SUCCESS && current != nullptr) { + return; + } + + const bool names_a_context = stream != nullptr && stream != reinterpret_cast(CU_STREAM_LEGACY) && + stream != reinterpret_cast(CU_STREAM_PER_THREAD); + if (names_a_context && stream_get_ctx != nullptr) { + CUcontext stream_ctx = nullptr; + if (stream_get_ctx(reinterpret_cast(stream), &stream_ctx) == CUDA_SUCCESS && stream_ctx != nullptr) { + ctx_set_current(stream_ctx); + return; + } + } + + if (device_get == nullptr || primary_retain == nullptr) { + return; + } + int ordinal = 0; + if (cuda_get_device(&ordinal) != cudaSuccess) { + return; + } + CUdevice device = 0; + if (device_get(&device, ordinal) != CUDA_SUCCESS) { + return; + } + // Retained for the process lifetime, like the primary context itself. + CUcontext primary = nullptr; + if (primary_retain(&primary, device) == CUDA_SUCCESS) { + ctx_set_current(primary); + } +#else + (void)stream; +#endif +} + inline cudaError_t cuda_pointer_get_attributes(cudaPointerAttributes *attributes, const void *ptr) { NV_FE_CALL_TO_CUDA(cuda_pointer_get_attributes, cudaPointerGetAttributes, attributes, ptr); diff --git a/pyproject.toml b/pyproject.toml index 75a9751ed..66d51bdd0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -68,6 +68,16 @@ cutedsl = [ "cuda-python", "apache-tvm-ffi>=0.1.11", ] +cutile = [ + # The cuTile linear-attention engines. Base cuda-tile only -- its [tileiras] + # extra pins cuda-toolkit>=13.2,<13.4, and that upper bound would cap the + # whole environment's toolkit and shut out CUDA 12 entirely. Without it, + # cuda.tile falls back to a system `tileiras`, the same way this package + # already leaves GPU wheels to the user. Engines that cannot import the + # runtime decline in check_support, so a missing compiler costs those + # engines and nothing else. + "cuda-tile>=1.4; python_version >= '3.10'", +] [dependency-groups] dev = [ diff --git a/python/cudnn/_device.py b/python/cudnn/_device.py index 57b706059..d14b302bc 100644 --- a/python/cudnn/_device.py +++ b/python/cudnn/_device.py @@ -59,38 +59,76 @@ def _device_handle(device: int): return _ck(*drv.cuDeviceGet(device)) -def ensure_current_context(stream=None) -> None: - """Bind a driver context to the calling thread when none is bound. - - JIT engines talk to the driver directly, which reads the calling THREAD's - context stack — and an autograd backward runs on a worker thread where - ``cudaSetDevice`` has only moved the runtime's thread-local slot. Prefer - the stream's context, else retain the runtime-current device's primary one - (the retain ref is held for the process lifetime, like the primary context - itself). Best-effort: a context this cannot establish fails at the launch, - with the launch's own diagnostics.""" +@functools.lru_cache(maxsize=1) +def _default_streams() -> frozenset: + """Handles that name no context: ``cuStreamGetCtx`` answers for the calling + thread's current context on all of them.""" + drv = _driver() + if drv is None: + return frozenset({0}) + return frozenset({0, int(drv.CU_STREAM_LEGACY), int(drv.CU_STREAM_PER_THREAD)}) + + +@functools.lru_cache(maxsize=None) +def _primary_context(device: int): + """The retained primary context for ``device``, or ``None``. Once per + ordinal: a retain per execute would grow the usage count without bound.""" + drv = _driver() + if drv is None: + return None + err, handle = drv.cuDeviceGet(device) + if int(err) != 0: + return None + err, primary = drv.cuDevicePrimaryCtxRetain(handle) + return primary if int(err) == 0 else None + + +def _runtime_device(): + """Ordinal the CUDA *runtime* holds current on this thread, or ``None``. + The driver cannot see that slot.""" try: - drv = _driver() - if drv is None: - return - err, cur = drv.cuCtxGetCurrent() - if int(err) == 0 and int(cur) != 0: - return - if stream: - err, sctx = drv.cuStreamGetCtx(int(stream)) - if int(err) == 0: - drv.cuCtxSetCurrent(sctx) - return import cuda.bindings.runtime as rt + except ImportError: + return None + + err, device = rt.cudaGetDevice() + return int(device) if int(err) == 0 else None + - err, device = rt.cudaGetDevice() - if int(err) != 0: +def ensure_current_context(stream=None, device=None) -> None: + """Bind the context this work runs in to the calling thread. + + A driver-API launch reads the calling thread's context stack; an autograd + backward runs on a worker whose stack is empty. A real stream names the + right context; the default-stream handles name none, so ``device`` decides + there. With no ``device``, a bound context is authoritative and only a cold + thread is given one (``frost.device.ambient_device``'s rung order). + Best-effort: what this cannot establish fails at the launch.""" + drv = _driver() + if drv is None: + return + err, cur = drv.cuCtxGetCurrent() + cur = int(cur) if int(err) == 0 else 0 + stream = 0 if stream is None else int(stream) + if stream not in _default_streams(): + err, stream_ctx = drv.cuStreamGetCtx(stream) + if int(err) == 0 and int(stream_ctx) != 0: + if int(stream_ctx) != cur: + drv.cuCtxSetCurrent(stream_ctx) + return + if device is None: + if cur: + return + device = _runtime_device() + if device is None: + return + elif cur: + err, cur_device = drv.cuCtxGetDevice() + if int(err) == 0 and int(cur_device) == int(device): return - err, pctx = drv.cuDevicePrimaryCtxRetain(_device_handle(int(device))) - if int(err) == 0: - drv.cuCtxSetCurrent(pctx) - except Exception: # noqa: BLE001 - pass + primary = _primary_context(int(device)) + if primary is not None: + drv.cuCtxSetCurrent(primary) class DeviceInfo: diff --git a/python/cudnn/_pygraph.py b/python/cudnn/_pygraph.py index fef68b781..43ad718eb 100644 --- a/python/cudnn/_pygraph.py +++ b/python/cudnn/_pygraph.py @@ -1777,10 +1777,10 @@ def execute( if eng is not None: # python engine (plan id in the reserved region) h = handle if handle is not None else self._handle ctx = ExecutionContext(handle=h, stream=self._resolve_stream(h), workspace=workspace) - # A JIT engine talks to the driver directly, which reads the calling - # THREAD's context stack -- and an autograd backward runs on a worker - # thread that has none. Once here, so every python engine is covered. - ensure_current_context(ctx.stream) + # A JIT engine launches through the driver, which reads the calling + # thread's context stack; an autograd worker has none. The handle's + # device decides when the stream names no context. + ensure_current_context(ctx.stream, h.device.ordinal if h is not None else None) if self._plan_index not in self._compiled_plans: # compile with the CALLER's context (execute-supplied handle # and its stream reach the JIT build) diff --git a/test/python/test_ensure_current_context.py b/test/python/test_ensure_current_context.py new file mode 100644 index 000000000..90505b388 --- /dev/null +++ b/test/python/test_ensure_current_context.py @@ -0,0 +1,189 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""``ensure_current_context`` binds the context a plan's work runs in to the +calling thread. Two halves, and only the first used to hold: a cold thread must +end up bound at all, and a thread bound to ANOTHER GPU's context must be moved +off it. The second is not cosmetic -- a default stream under a foreign context +runs the work on that context's GPU, where the pointers are invalid.""" + +import threading + +import pytest + +from cudnn._device import _primary_context, device_count, ensure_current_context, is_available + +pytestmark = pytest.mark.L0 + + +@pytest.fixture +def drv(): + d = pytest.importorskip("cuda.bindings.driver") + if not is_available(): + pytest.skip("no CUDA device") + entry = d.cuCtxGetCurrent()[1] + yield d + d.cuCtxSetCurrent(entry) # these tests move the thread's context on purpose + + +def _current(drv): + err, ctx = drv.cuCtxGetCurrent() + return int(ctx) if int(err) == 0 else 0 + + +def _bind_primary(drv, ordinal): + """Retain ``ordinal``'s primary context and make it current on this thread.""" + dev = drv.cuDeviceGet(ordinal)[1] + ctx = drv.cuDevicePrimaryCtxRetain(dev)[1] + drv.cuCtxSetCurrent(ctx) + return int(ctx) + + +def _two_devices(): + if device_count() < 2: + pytest.skip("needs two GPUs to tell 'a context' from 'the right context'") + return 0, 1 + + +def _on_a_cold_thread(body): + """Run ``body`` on a thread that has never bound a context. Returns its dict.""" + seen = {} + worker = threading.Thread(target=body, args=(seen,)) + worker.start() + worker.join() + if "exc" in seen: + raise seen["exc"] + return seen + + +def test_binds_a_cold_thread(drv): + _bind_primary(drv, 0) # the process has a context; the worker below does not + + def body(seen): + try: + seen["before"] = _current(drv) + ensure_current_context(0, 0) + seen["after"] = _current(drv) + except BaseException as exc: # noqa: BLE001 + seen["exc"] = exc + + seen = _on_a_cold_thread(body) + assert seen["before"] == 0, "the worker was already bound, so this no longer covers the cold path" + assert seen["after"] != 0 + + +@pytest.mark.parametrize("handle", ["null", "legacy", "per_thread"]) +def test_replaces_a_context_on_another_device(drv, handle): + """No default-stream handle can name a GPU -- each resolves against the + calling thread's current context -- so ``device`` decides.""" + a, b = _two_devices() + ctx_a, ctx_b = _bind_primary(drv, a), _bind_primary(drv, b) + assert ctx_a != ctx_b + stream = {"null": 0, "legacy": int(drv.CU_STREAM_LEGACY), "per_thread": int(drv.CU_STREAM_PER_THREAD)}[handle] + assert int(drv.cuStreamGetCtx(stream)[1]) == ctx_b # follows the thread, names nothing + + drv.cuCtxSetCurrent(drv.CUcontext(ctx_a)) + ensure_current_context(stream, b) + assert _current(drv) == ctx_b, "left the thread on another GPU's context" + + +def test_follows_the_streams_context(drv): + """A real stream carries its context, and it wins over ``device``.""" + a, b = _two_devices() + ctx_a, ctx_b = _bind_primary(drv, a), _bind_primary(drv, b) + stream = drv.cuStreamCreate(0)[1] # created under ctx_b, so it belongs to it + try: + drv.cuCtxSetCurrent(drv.CUcontext(ctx_a)) + ensure_current_context(int(stream), a) # device says a, the stream says b + assert _current(drv) == ctx_b, "the stream's context did not win" + finally: + drv.cuCtxSetCurrent(drv.CUcontext(ctx_b)) + drv.cuStreamDestroy(stream) + + +def test_leaves_an_already_correct_context_alone(drv): + """Steady state is a no-op: no rebind, no primary-context churn.""" + ctx = _bind_primary(drv, 0) + ensure_current_context(0, 0) + assert _current(drv) == ctx + ensure_current_context(0, 0) + assert _current(drv) == ctx + + +def test_an_unnamed_device_does_not_override_a_bound_context(drv): + """No device named means nothing to correct: a bound context is + authoritative (``ambient_device``'s first rung).""" + a, b = _two_devices() + _bind_primary(drv, a) + ctx_b = _bind_primary(drv, b) + torch = pytest.importorskip("torch") + torch.cuda.set_device(a) # runtime slot -> a, while the driver context is b's + drv.cuCtxSetCurrent(drv.CUcontext(ctx_b)) + ensure_current_context(0, None) + assert _current(drv) == ctx_b, "overrode a bound context on the runtime's word" + + +def test_the_primary_context_is_retained_once_per_device(drv): + """A retain per execute would grow the usage count without bound.""" + a, b = _two_devices() + for _ in range(20): # alternating default-stream execution over two GPUs + ensure_current_context(0, a) + ensure_current_context(0, b) + for ordinal in (a, b): + dev = drv.cuDeviceGet(ordinal)[1] + # cuDevicePrimaryCtxGetState reports active/flags, not the count, so assert + # the cache instead: one retained handle per ordinal, reused. + assert _primary_context(ordinal) is _primary_context(ordinal) + assert int(drv.cuDevicePrimaryCtxGetState(dev)[2]) == 1 # still active, not churned + + +def test_a_backend_graph_runs_on_a_cold_thread(drv): + """The C++ execute funnel binds one too, for the same reason. + + cuDNN's runtime-compiled engines launch through the driver and read the + calling thread's stack; the precompiled ones go through the runtime and are + unaffected, so the fused ``matmul + relu + relu`` is what exercises this.""" + torch = pytest.importorskip("torch") + import cudnn + + if not torch.cuda.is_available(): + pytest.skip("needs CUDA") + if torch.cuda.get_device_capability()[0] < 8: + pytest.skip("bf16 matmul needs sm80+") + a = torch.randn(1, 128, 128, device="cuda", dtype=torch.bfloat16) + b = torch.randn(1, 128, 128, device="cuda", dtype=torch.bfloat16) + out = torch.empty(1, 128, 128, device="cuda", dtype=torch.bfloat16) + torch.cuda.synchronize() + + handle = cudnn.create_handle() + try: + g = cudnn.pygraph(handle=handle, io_data_type=cudnn.data_type.BFLOAT16, compute_data_type=cudnn.data_type.FLOAT) + ta, tb = g.tensor_like(a), g.tensor_like(b) + y = g.relu(input=g.relu(input=g.matmul(A=ta, B=tb))) + y.set_output(True).set_data_type(cudnn.data_type.BFLOAT16) + try: + g.build([cudnn.heur_mode.A]) + except cudnn.cudnnGraphNotSupportedError as exc: + pytest.skip(f"no engine for the fused graph on this arch/backend: {exc}") + ws = torch.empty(max(g.get_workspace_size(), 1), device="cuda", dtype=torch.uint8) + pack = {ta: a, tb: b, y: out} + g.execute(pack, ws, handle=handle) # warm, on this thread + torch.cuda.synchronize() + + def body(seen): + # No torch op before execute -- even out.zero_() binds the context. + seen["before"] = _current(drv) + try: + g.execute(pack, ws, handle=handle) + except BaseException as exc: # noqa: BLE001 + seen["exc"] = exc + seen["after"] = _current(drv) + + seen = _on_a_cold_thread(body) + assert seen["before"] == 0, "the worker was already bound, so this no longer covers the cold path" + assert seen["after"] != 0 + torch.cuda.synchronize() + expected = torch.relu(torch.relu(a.float() @ b.float())) + assert (out.float() - expected).abs().max().item() < 1.0 + finally: + cudnn.destroy_handle(handle)