diff --git a/python/cudnn/engines/base.py b/python/cudnn/engines/base.py index 46c523a19..907c93fa1 100644 --- a/python/cudnn/engines/base.py +++ b/python/cudnn/engines/base.py @@ -48,8 +48,6 @@ def execute(self, graph, uid_to_data, ctx): from typing import TYPE_CHECKING, Any, Dict, List from .engine_ids import PYTHON_ENGINE_ID_BASE # noqa: F401 — re-exported for engine authors -from ..datatypes import _CUDNN_TO_FROST_DTYPE_NAME -from ..frost.buffers import DeviceView if TYPE_CHECKING: from ..pygraph import pygraph diff --git a/python/cudnn/frost/buffers.py b/python/cudnn/frost/buffers.py index 003b761d3..280928e71 100644 --- a/python/cudnn/frost/buffers.py +++ b/python/cudnn/frost/buffers.py @@ -18,6 +18,7 @@ from __future__ import annotations import ctypes +import struct from cudnn import _pybind_module @@ -297,6 +298,143 @@ def memset_zero_async(ptr: int, nbytes: int, stream) -> None: raise RuntimeError(f"cudaMemsetAsync failed: {err}") +_WORD_FORMAT = {"fp32": " int: + """The 32-bit pattern that writes ``value`` to a buffer of ``dtype``. + + A memset moves bits, not numbers, so the value has to be packed as the dtype + the kernel will read it back as. int32's reduction identities are the ends + of its range and are exactly where that bites: -2**31 packed as float is + 0xcf000000 where the kernel wants 0x80000000. + """ + fmt = _WORD_FORMAT.get(dtype) + if fmt is None: + raise NotImplementedError(f"no 32-bit fill pattern for dtype {dtype!r}") + return int.from_bytes(struct.pack(fmt, value), "little") + + +def fill_word_async(ptr: int, count: int, word: int, stream) -> None: + """Stream-ordered fill of ``count`` CONTIGUOUS 32-bit words with ``word``. + + An engine that seeds a caller's buffer owns that operation itself: reaching + for ``tensor.fill_()`` works only while the buffer happens to be a torch + tensor, and queues on torch's current stream rather than the one the kernel + will run on. Every seed a reduction uses is a 32-bit pattern, so the + driver's D32 memset covers them without a kernel -- see ``init_word`` for + turning a value into one. + + :func:`fill_word_strided_async` is the same fill for a buffer that is not + one dense run. + """ + from cuda.bindings import driver as _drv + + res = _drv.cuMemsetD32Async(int(ptr), int(word), int(count), int(stream) if stream is not None else 0) + err = res[0] if isinstance(res, tuple) else res + if int(err) != 0: + raise RuntimeError(f"cuMemsetD32Async failed: {err}") + + +def _fill_word_2d_async(ptr: int, pitch_words: int, width: int, height: int, word: int, stream) -> None: + from cuda.bindings import driver as _drv + + res = _drv.cuMemsetD2D32Async(int(ptr), int(pitch_words) * 4, int(word), int(width), int(height), int(stream) if stream is not None else 0) + err = res[0] if isinstance(res, tuple) else res + if int(err) != 0: + raise RuntimeError(f"cuMemsetD2D32Async failed: {err}") + + +def collapse_layout(shape, strides) -> list: + """``(extent, stride)`` outermost first, with unit axes dropped and adjacent + axes merged where one exactly fills the other's gap. + + A padded output is usually dense underneath its declared rank -- a rank-3 + ``(1, M, 1)`` tap is one strided run, and a contiguous one is a single dense + run whatever rank it was declared at. Merging first is what keeps the fill + below down to one memset in both cases. + """ + axes = sorted(((int(d), int(s)) for d, s in zip(shape, strides) if int(d) != 1), key=lambda ds: -ds[1]) + out: list = [] + for extent, stride in axes: + if out and out[-1][1] == extent * stride: + out[-1] = (out[-1][0] * extent, stride) + else: + out.append((extent, stride)) + return out + + +def strided_fill_plan(shape, strides) -> "list | None": + """The 2D memsets that cover a strided region exactly once, or None. + + None means the region writes some element twice -- a stride of 0 over a real + extent, or an outer stride that does not clear the axis below it. That is a + write race whichever buffer it is, so it is refused rather than issued; the + caller decides how to say so. + + Returned before anything is written, which is the point: the seed's + preconditions have to be settled while the caller's buffer is still + untouched, and a plan is what lets several outputs all be checked before the + first of them is filled. + + Each entry is ``(offset, pitch, width, height)`` in ELEMENTS. The driver's + 2D memset takes a pitch, so a per-row scalar tap is one entry rather than one + per row (the reading that made this look expensive: 572 us at one memset per + row); what remains is one entry per point of whatever axis is left outside + the 2D region, which for a rank-3 output is the batch and is usually 1. + """ + if any(int(s) == 0 and int(d) != 1 for d, s in zip(shape, strides)): + return None + axes = collapse_layout(shape, strides) + # Non-overlapping iff each axis clears the whole span of the one below it. + # `pitch >= width` is this rule at the innermost pair and misses the rest: + # shape (2, 2) stride (2, 2) has width 1 and passes it, and lands both axes + # on the same element. + for (_outer_extent, outer_stride), (inner_extent, inner_stride) in zip(axes, axes[1:]): + if outer_stride < inner_extent * inner_stride: + return None + if not axes: + return [(0, 1, 1, 1)] + # The innermost run is the memset's width when it is dense; otherwise every + # element stands alone and the width is one. + width, rest = (axes[-1][0], axes[:-1]) if axes[-1][1] == 1 else (1, axes) + if not rest: + return [(0, width, width, 1)] + height, pitch = rest[-1] + offsets = [0] + for extent, stride in reversed(rest[:-1]): + offsets = [base + i * stride for base in offsets for i in range(extent)] + return [(base, pitch, width, height) for base in offsets] + + +def apply_fill_plan(ptr: int, plan, word: int, stream) -> None: + """Issue a plan from :func:`strided_fill_plan`, stream-ordered.""" + for offset, pitch, width, height in plan: + if height == 1: + fill_word_async(ptr + offset * 4, width, word, stream) + else: + _fill_word_2d_async(ptr + offset * 4, pitch, width, height, word, stream) + + +def fill_word_strided_async(ptr: int, shape, strides, elem_bytes: int, word: int, stream) -> None: + """Plan and issue in one call, for a caller with a single region to seed. + + The engine owns seeding a reduction output, and a padded one is the case + that used to send it back to the caller's ``fill_()`` -- the last place + anything here wrote through a buffer it does not own, and the reason a + perfectly legal call had to fall off the fast path. A caller with SEVERAL + regions wants :func:`strided_fill_plan` for all of them first: this one + cannot know whether the next region is refusable, so it would leave the + earlier ones filled. + """ + if elem_bytes != 4: + raise NotImplementedError(f"frost: a reduction seed is a 32-bit pattern; this output stores {elem_bytes}-byte elements") + plan = strided_fill_plan(shape, strides) + if plan is None: + raise ValueError(f"frost: a reduction output cannot write an element twice (shape {tuple(shape)} stride {tuple(strides)})") + apply_fill_plan(ptr, plan, word, stream) + + # The CuTe primitives these engines lower through landed in 4.7.0; older DSLs # fail during codegen with errors that name a missing attribute rather than the # version, so the check belongs where an engine can still decline. diff --git a/python/cudnn/frost/workspace.py b/python/cudnn/frost/workspace.py index f198cfa4e..1089bfe5f 100644 --- a/python/cudnn/frost/workspace.py +++ b/python/cudnn/frost/workspace.py @@ -125,7 +125,11 @@ def over(cls, variant_pack, required_bytes: int, owner: str, *, align: int = DEF f"{owner} requires a {required_bytes}-byte workspace but execute() received " f"none; allocate graph.get_workspace_size() bytes and pass the buffer to execute()" ) - if nbytes < required_bytes: + # 0 means the pack could not measure it, not that it is empty: a bare + # device address carries no size, and the backend takes one without + # checking either. Refusing here would make the same call depend on + # which plan ran. + if nbytes and nbytes < required_bytes: raise ValueError(f"{owner}: needs a {required_bytes}-byte workspace, got {nbytes} bytes (size it with graph.get_workspace_size())") if ptr % align != 0: raise ValueError(f"{owner}: the workspace buffer must be {align}-byte aligned; got 0x{ptr:x}") @@ -164,9 +168,14 @@ def take(self, numel: int, dtype: str) -> buffers.DeviceView: def remaining(self) -> buffers.DeviceView: """The tail no :meth:`take` has claimed, as uint8 — for a nested carver.""" + if not self._nbytes: + raise ValueError( + f"{self._owner}: the workspace was passed as a bare address, so its size is unknown " + "and the unclaimed tail cannot be measured; pass a sized buffer to execute()" + ) return buffers.DeviceView(self._ptr + self._offset, (self._nbytes - self._offset,), "uint8", self._device) def _check_span(self, offset: int, span: int) -> None: end = int(offset) + int(span) - if end > self._nbytes: + if self._nbytes and end > self._nbytes: raise ValueError(f"{self._owner}: workspace overrun — region [{offset}, {end}) exceeds the " f"{self._nbytes}-byte buffer (sizing bug)") diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index b44b18ebe..011322cf8 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -1803,13 +1803,30 @@ def _expected_output_shape(spec, chain: FusionChain, mnk) -> tuple[int, int, int return tuple(1 if spec.dim[i] == 1 else full[i] for i in range(3)) -def _initialize_reduction_outputs(chain: FusionChain, outputs) -> None: +def _initialize_reduction_outputs(chain: FusionChain, outputs, stream=None) -> None: + """Seed each reduction output with its identity, before the kernel runs. + + Through the driver rather than the buffer: an engine that reaches for + ``tensor.fill_()`` works only while the caller happened to pass a torch + tensor, and the variant pack exists so that it does not have to. + """ for spec, tensor in zip(chain.outputs, outputs): if not spec.is_reduction: continue - red_idx = int(spec.source.rsplit("_", 1)[1]) - red = chain.reductions[red_idx] - tensor.fill_(_REDUCTION_INIT_VALUE[red.compute_dtype][red.mode]) + red = chain.reductions[int(spec.source.rsplit("_", 1)[1])] + value = _REDUCTION_INIT_VALUE[red.compute_dtype][red.mode] + # The driver, on the stream the kernel will run on -- for a padded output + # too. tensor.fill_() would queue on torch's current stream instead, + # which is the same stream only by luck, and only exists at all while + # the caller happened to pass a torch tensor. The pattern is packed as + # the OUTPUT's dtype, not as float: int32's identities are the ends of + # its range. + shape, strides = tuple(tensor.shape), tuple(tensor.stride()) + word = buffers.init_word(red.compute_dtype, value) + if buffers.is_contiguous(shape, strides): + buffers.fill_word_async(tensor.data_ptr(), int(tensor.numel()), word, stream) + else: + buffers.fill_word_strided_async(tensor.data_ptr(), shape, strides, tensor.element_size(), word, stream) def _finalize_reductions(chain, out_bufs) -> None: @@ -1999,7 +2016,7 @@ def _call_positional(self, *args, stream=None): f"got A={tuple(a.shape)}, B={tuple(b.shape)}, " f"C={[tuple(ci.shape) for ci in cs]}" ) - _initialize_reduction_outputs(self.chain, cs) + _initialize_reduction_outputs(self.chain, cs, stream) if self.chain.output_specs: base_problem = (mnk[0], mnk[1], mnk[2], cs[0].shape[0]) @@ -2086,7 +2103,7 @@ def _call_multi_gemm(self, gemm_pairs, c, mnk, *aux, stream=None): for t in slots: if len(t.shape) != 3: raise ValueError(f"multi-GEMM {role} operand must be rank-3; got {tuple(t.shape)}") - _initialize_reduction_outputs(chain, cs) + _initialize_reduction_outputs(chain, cs, stream) base_problem = (mnk[0], mnk[1], mnk[2], cs[0].shape[0]) a_permuted = [t.permute(1, 2, 0) for t in a_slots] @@ -2158,7 +2175,7 @@ def _call_block_scale_multi_gemm(self, gemm_pairs, c, mnk, *aux, stream=None): raise ValueError( f"block-scale multi-GEMM output {spec.source!r} must have " f"shape {_expected_output_shape(spec, chain, mnk)}; " f"got {tuple(ci.shape)}" ) - _initialize_reduction_outputs(chain, cs) + _initialize_reduction_outputs(chain, cs, stream) base_problem = (mnk[0], mnk[1], mnk[2], cs[0].shape[0]) # Grouped by kind (all A, all B, all SFA, all SFB); single-GEMM → a,b,sfa,sfb. a_permuted = [d.permute(1, 2, 0) for d, _ in a_slots] @@ -2965,7 +2982,7 @@ def _launch_single(self, token, weight, first_token_offset, output, snke, worksp raise ValueError( f"MoE output {spec.source!r} must have shape " f"{_expected_output_shape(spec, self.chain, (S, N, K))}; " f"got {tuple(t.shape)}" ) - _initialize_reduction_outputs(self.chain, outputs) + _initialize_reduction_outputs(self.chain, outputs, stream) # num_experts = weight batch (E); num_groups = first_token_offset len # (BxE, may exceed E; group g uses expert g % E). From runtime tensors. num_experts = int(weight.shape[0]) @@ -3091,7 +3108,7 @@ def _call_multi_gemm(self, gemm_pairs, first_token_offset, output, snke, *aux, w raise ValueError( f"multi-GEMM MoE output {spec.source!r} must have shape " f"{_expected_output_shape(spec, chain, (S, N, K))}; got {tuple(ci.shape)}" ) - _initialize_reduction_outputs(chain, outs) + _initialize_reduction_outputs(chain, outs, stream) num_experts = int(b_slots[0].shape[0]) num_groups = int(first_token_offset.shape[0]) a_stride_perms = [t.permute(1, 2, 0) for t in a_slots] @@ -3347,7 +3364,7 @@ def _launch_single(self, token, weight, sfa, sfb, first_token_offset, output, sn f"{_expected_output_shape(spec, self.chain, (S, N, K))}; " f"got {tuple(t.shape)}" ) - _initialize_reduction_outputs(self.chain, outputs) + _initialize_reduction_outputs(self.chain, outputs, stream) # num_experts = weight batch (E); num_groups = first_token_offset len # (BxE, may exceed E; group g uses expert g % E). From runtime tensors. num_experts = int(weight.shape[0]) @@ -3483,7 +3500,7 @@ def _call_multi_gemm(self, gemm_pairs, first_token_offset, output, snke, *aux, w f"multi-GEMM MoE block-scale output {spec.source!r} must have shape " f"{_expected_output_shape(spec, chain, (S, N, K))}; got {tuple(ci.shape)}" ) - _initialize_reduction_outputs(chain, outs) + _initialize_reduction_outputs(chain, outs, stream) num_experts = int(b_slots[0][0].shape[0]) num_groups = int(first_token_offset.shape[0]) a0, b0 = a_slots[0][0], b_slots[0][0] diff --git a/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py index 83bf5b409..ff80aad30 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn2_prefill_f16.py @@ -2618,9 +2618,8 @@ def _build_descs( stream: cuda_driver.CUstream, ): """Build the 8 per-(batch, head) TMA-descriptor arrays (q, k, v, gate, - beta, w, o, h) into ``tensormap_workspace``. Compiled + launched - separately from the main kernel and cached by input identity in the host - bridge, so the builder launches do not recur in steady-state replay. + beta, w, o, h) into ``tensormap_workspace``. Launched on every execute: the descriptors fold cu_seqlens contents into + GLOBAL_ADDRESS and GLOBAL_DIM, which the host cannot read without a D2H sync. Each descriptor folds the sequence base + head offset into GLOBAL_ADDRESS (Int64) and caps the token GLOBAL_DIM to the sequence length, so the main kernel's coordinates are sequence-relative and tail diff --git a/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py b/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py index b78ebf950..f9fd8c670 100644 --- a/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/gdn_bprop_f16.py @@ -3958,9 +3958,8 @@ def _build_descs( ): """Build the per-(b,h) TMA-descriptor arrays (Q, K, V, dO, H loads; dQ, dK, dV stores; the io-dtype S0 loads when ``s0`` is given) into - ``tensormap_workspace``. Compiled + launched separately from the main - kernel and cached by input identity in the host bridge, so the builder - launches do not recur in steady-state replay. + ``tensormap_workspace``. Launched on every execute: the descriptors fold cu_seqlens contents into + GLOBAL_ADDRESS and GLOBAL_DIM, which the host cannot read without a D2H sync. The H descriptor is 3-D ``(dv, dk, h)`` over the packed ``[total_h, HO, DK, DV]`` H tensor; ``build_h_descs_kernel`` derives the diff --git a/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py b/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py index 235d10e38..dfd622955 100644 --- a/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py +++ b/python/cudnn/linear_attention/frost/kernel/kda_prefill_f16.py @@ -2541,9 +2541,8 @@ def _build_descs( stream: cuda_driver.CUstream, ): """Build the 6 per-(batch, head) TMA-descriptor arrays (q, k, v, gate, - o, h) into ``tensormap_workspace``. Compiled + launched separately from - the main kernel and cached by input identity in the host bridge, so the - builder launches do not recur in steady-state replay. Each descriptor + o, h) into ``tensormap_workspace``. Launched on every execute: the descriptors fold cu_seqlens contents into + GLOBAL_ADDRESS and GLOBAL_DIM, which the host cannot read without a D2H sync. Each descriptor folds the sequence base + head offset into GLOBAL_ADDRESS (Int64) and caps the token GLOBAL_DIM to the sequence length, so the main kernel's coordinates are sequence-relative and tail chunks clip in hardware. The diff --git a/python/pygraph/variant_pack.cpp b/python/pygraph/variant_pack.cpp index 8ef85175c..0c25db2b7 100644 --- a/python/pygraph/variant_pack.cpp +++ b/python/pygraph/variant_pack.cpp @@ -324,6 +324,11 @@ class VariantPackSlot { // Ownership transfers with the capsule, per DLPack: the struct carries its // own copy of the shape and stride and a deleter that frees them, so it // outlives this slot rather than aliasing storage the slot owns. + // + // Unversioned only: max_version is ignored and the capsule is always + // "dltensor". The consumer this exists for is cute's compile-time + // from_dlpack; tvm-ffi reads a slot through the exchange vtable and never + // gets here. py::capsule dlpack(py::object /*stream*/, py::object /*max_version*/) const { struct Owned { @@ -691,7 +696,7 @@ class WorkspaceCarve { std::vector out; out.reserve(protos_.size()); for (size_t i = 0; i < protos_.size(); i++) { - if (ends_[i] > nbytes) { + if (nbytes != 0 && ends_[i] > nbytes) { // 0 = size unknown (bare address) throw py::value_error(owner_ + ": workspace overrun -- region [" + std::to_string(offsets_[i]) + ", " + std::to_string(ends_[i]) + ") exceeds the " + std::to_string(nbytes) + "-byte buffer (sizing bug)"); diff --git a/test/python/gemm/frost/test_public_execute_flavors.py b/test/python/gemm/frost/test_public_execute_flavors.py new file mode 100644 index 000000000..068e159c4 --- /dev/null +++ b/test/python/gemm/frost/test_public_execute_flavors.py @@ -0,0 +1,382 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Every gemm flavor, driven through the PUBLIC ``graph.execute()``. + +The rest of this directory calls the compiled object directly with torch +tensors. That skips the whole engine wrapper -- operand binding, the variant +pack, and the buffer conversion ``execute()`` performs -- so a break that only +appears when the engine hands the kernel what the pack holds is invisible to +it. A reduction output shipped in exactly that state: the epilogue initializes +it with ``fill_`` and finalizes ``norm2`` with ``sqrt_``, both of which the +compiled object used to receive as torch tensors and now does not. + +So this file exists to exercise the same flavors the direct-call tests cover, +but through the entry point a user actually has. +""" + +import pytest +import torch + +import cudnn +from cudnn.engines import is_python_engine + +pytestmark = pytest.mark.L0 + +_GPU = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10, + reason="the FROST gemm engine claims SM100", +) + +BF16 = cudnn.data_type.BFLOAT16 +F32 = cudnn.data_type.FLOAT +M = N = 128 +K = 64 + + +@pytest.fixture(autouse=True) +def _frost_opt_in(monkeypatch): + monkeypatch.setenv("CUDNN_FRONTEND_ENABLE_FROST_ENGINES", "1") + + +def _pin_frost(g): + """Select the FROST plan, or skip when it does not claim this graph.""" + g.validate() + g.build_operation_graph() + g.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + frost = [i for i, p in enumerate(g.plans) if is_python_engine(p.engine_id)] + if not frost: + pytest.skip("no FROST plan for this graph") + g.select_plan(frost[0]) + g.check_support() + g.build_plans() + return g + + +def _operands(): + a = torch.randn(1, M, K, dtype=torch.bfloat16, device="cuda") + b = torch.randn(1, N, K, dtype=torch.bfloat16, device="cuda") + return a, b, torch.einsum("bmk,bnk->bmn", a.float(), b.float()) + + +def _run(g, data): + ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda") + g.execute(data, ws) + torch.cuda.synchronize() + + +@_GPU +def test_plain_matmul(): + """The control: if this fails the harness is wrong, not the flavor.""" + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", dim=[1, K, N], stride=[K * N, 1, K]) + C = g.matmul(A=A, B=B, name="mm") + C.set_output(True).set_data_type(BF16) + _pin_frost(g) + + a, b, ref = _operands() + c = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + _run(g, {A: a, B: b, C: c}) + torch.testing.assert_close(c, ref.to(torch.bfloat16), atol=1e-1, rtol=1e-2) + + +@_GPU +def test_epilogue_fusion(): + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", dim=[1, K, N], stride=[K * N, 1, K]) + C = g.matmul(A=A, B=B, name="mm") + Y = g.relu(input=C, name="relu") + Y.set_output(True).set_data_type(BF16) + _pin_frost(g) + + a, b, ref = _operands() + y = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + _run(g, {A: a, B: b, Y: y}) + torch.testing.assert_close(y, torch.relu(ref).to(torch.bfloat16), atol=1e-1, rtol=1e-2) + + +@_GPU +@pytest.mark.parametrize( + "mode,reference", + ( + (cudnn.reduction_mode.ADD, lambda t: t.sum().reshape(1, 1, 1)), + (cudnn.reduction_mode.AMAX, lambda t: t.abs().max().reshape(1, 1, 1)), + (cudnn.reduction_mode.MAX, lambda t: t.max().reshape(1, 1, 1)), + (cudnn.reduction_mode.MIN, lambda t: t.min().reshape(1, 1, 1)), + ), + ids=("add", "amax", "max", "min"), +) +def test_reduction_output(mode, reference): + """A reduction tap alongside the epilogue output. + + The epilogue writes the tap's initial value before the kernel runs and, for + ``norm2``, takes a square root after it. Both are device operations on a + caller buffer, which is where a buffer that is only a description rather + than a tensor shows up. + """ + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", dim=[1, K, N], stride=[K * N, 1, K]) + C = g.matmul(A=A, B=B, name="mm") + Y = g.relu(input=C, name="relu") + Y.set_output(True).set_data_type(BF16) + R = g.reduction(input=Y, mode=mode, name="red") + R.set_dim([1, 1, 1]).set_stride([1, 1, 1]) + R.set_output(True).set_data_type(F32) + _pin_frost(g) + + a, b, ref = _operands() + y = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + r = torch.empty(1, 1, 1, dtype=torch.float32, device="cuda") + _run(g, {A: a, B: b, Y: y, R: r}) + + relu = torch.relu(ref) + torch.testing.assert_close(y, relu.to(torch.bfloat16), atol=1e-1, rtol=1e-2) + torch.testing.assert_close(r, reference(relu), atol=1e-1, rtol=1e-2) + + +@_GPU +def test_int32_reduction_seed_is_packed_as_int32(): + """A memset moves bits, so the identity has to be packed as the dtype. + + int32's identities are the ends of its range, which is exactly where the + difference shows: -2**31 packed as float is 0xcf000000, and a MAX reduction + seeded with that returns -822083584 for any input below it. + """ + I32 = cudnn.data_type.INT32 + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", dim=[1, K, N], stride=[K * N, 1, K]) + bias = g.tensor(name="bias", dim=[1, 1, N], stride=[N, N, 1], data_type=I32) + C = g.matmul(A=A, B=B, name="mm") + Y = g.add(a=C, b=bias, name="add_i32", compute_data_type=I32) + Y.set_output(True).set_data_type(I32) + R = g.reduction(input=Y, mode=cudnn.reduction_mode.MAX, name="red", compute_data_type=I32) + R.set_dim([1, 1, 1]).set_stride([1, 1, 1]).set_output(True).set_data_type(I32) + _pin_frost(g) + + floor = -2_000_000_000 # below the float-packed seed, above int32's minimum + a = torch.zeros(1, M, K, dtype=torch.bfloat16, device="cuda") + b = torch.zeros(1, N, K, dtype=torch.bfloat16, device="cuda") + y = torch.empty(1, M, N, dtype=torch.int32, device="cuda") + r = torch.empty(1, 1, 1, dtype=torch.int32, device="cuda") + _run(g, {A: a, B: b, bias: torch.full((1, 1, N), floor, dtype=torch.int32, device="cuda"), Y: y, R: r}) + assert int(r.item()) == floor + + +@_GPU +def test_norm2_reduction_is_refused_at_build(): + """``norm2`` is the one reduction mode with a post-kernel finalize. + + It never reaches one: the backend refuses the reduction descriptor while + the graph is being lowered, so no plan exists and ``execute()`` is never + called. Recorded because the finalize would otherwise look like a live path + that needs a device-side square root. + """ + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", dim=[1, K, N], stride=[K * N, 1, K]) + C = g.matmul(A=A, B=B, name="mm") + Y = g.relu(input=C, name="relu") + Y.set_output(True).set_data_type(BF16) + R = g.reduction(input=Y, mode=cudnn.reduction_mode.NORM2, name="red") + R.set_dim([1, 1, 1]).set_stride([1, 1, 1]) + R.set_output(True).set_data_type(F32) + g.validate() + with pytest.raises(RuntimeError, match="NOT_SUPPORTED"): + g.build_operation_graph() + + +@_GPU +@pytest.mark.xfail( + reason=( + "the operand a bare address describes is the GRAPH's declaration, and frost reads its " + "extents by axis position from the layout a caller's buffer would report -- for a matmul " + "B those are [batch, K, N] and (batch, N, K). Broken before this branch too (the " + "geometry-less Tensor made it an IndexError); the fix is the engine recording which axis " + "is M/N/K at build, which belongs with the executor rewrite." + ), + strict=True, +) +def test_bare_address_operands(): + """The backend has always taken a raw device address; so must a python plan.""" + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", uid=1, dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", uid=2, dim=[1, K, N], stride=[K * N, 1, K]) + C = g.matmul(A=A, B=B, name="mm") + C.set_output(True).set_data_type(BF16).set_uid(3) + _pin_frost(g) + + a, b, ref = _operands() + c = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + _run(g, {1: a.data_ptr(), 2: b.data_ptr(), 3: c.data_ptr()}) + torch.testing.assert_close(c, ref.to(torch.bfloat16), atol=1e-1, rtol=1e-2) + + +@_GPU +def test_bare_address_workspace(): + """A workspace passed as a raw address has no measurable size. + + Zero means "the pack could not measure it", not "empty" -- the backend + takes a raw workspace pointer without checking either, so an engine that + needs scratch must not refuse one. This drives the unknown-capacity path + through both `Workspace.over` and the C carve's bounds check. + """ + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", dim=[1, K, N], stride=[K * N, 1, K]) + C = g.matmul(A=A, B=B, name="mm") + C.set_output(True).set_data_type(BF16) + _pin_frost(g) + + a, b, ref = _operands() + c = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + ws = torch.empty(max(g.get_workspace_size(), 1), dtype=torch.uint8, device="cuda") + g.execute({A: a, B: b, C: c}, ws.data_ptr()) + torch.cuda.synchronize() + torch.testing.assert_close(c, ref.to(torch.bfloat16), atol=1e-1, rtol=1e-2) + + +@_GPU +def test_undersized_workspace_still_rejected(): + """A workspace whose size IS known is still bounds-checked.""" + from cudnn.engines.base import VariantPack + from cudnn.frost.workspace import Workspace + + tiny = torch.empty(16, dtype=torch.uint8, device="cuda") + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", uid=1, dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", uid=2, dim=[1, K, N], stride=[K * N, 1, K]) + C = g.matmul(A=A, B=B, name="mm") + C.set_output(True).set_data_type(BF16).set_uid(3) + _pin_frost(g) + + a, b, _ = _operands() + c = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + pack = g._normalize(g._uid_to_data({1: a, 2: b, 3: c}), tiny) + assert pack.workspace_bytes == 16 + with pytest.raises(ValueError, match=r"needs a .*-byte workspace"): + Workspace.over(pack, 4096, "probe") + + +@_GPU +def test_unknown_size_workspace_refuses_to_measure_its_tail(): + """``remaining()`` cannot answer for a workspace whose size is unknown.""" + from cudnn.frost.workspace import Workspace + + ws = torch.empty(4096, dtype=torch.uint8, device="cuda") + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", uid=1, dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", uid=2, dim=[1, K, N], stride=[K * N, 1, K]) + C = g.matmul(A=A, B=B, name="mm") + C.set_output(True).set_data_type(BF16).set_uid(3) + _pin_frost(g) + + a, b, _ = _operands() + c = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + pack = g._normalize(g._uid_to_data({1: a, 2: b, 3: c}), ws.data_ptr()) + assert pack.workspace_bytes == 0 + carver = Workspace.over(pack, 1024, "probe") + with pytest.raises(ValueError, match="size is unknown"): + carver.remaining() + + +# --- seeding a padded reduction output, which the engine owns ---------------- + + +@pytest.mark.parametrize( + "shape,stride,want", + [ + ((1, 8, 4), (32, 4, 1), [(32, 1)]), # dense whatever rank it was declared at + ((1, 8, 1), (32, 4, 1), [(8, 4)]), # a per-row scalar: one strided run + ((2, 8, 4), (64, 8, 1), [(16, 8), (4, 1)]), # padded rows: batch merges into the row count + ((1, 1, 1), (1, 1, 1), []), # one element + ], +) +def test_a_layout_collapses_to_the_runs_a_memset_can_cover(shape, stride, want): + """Unit axes carry no elements and adjacent dense axes are one run. + + Collapsing first is what keeps the seed to a single memset for a contiguous + tap and one per batch for a padded one, rather than one per row. + """ + from cudnn.frost.buffers import collapse_layout + + assert collapse_layout(shape, stride) == want + + +@pytest.mark.parametrize( + "shape,stride", + [ + ((1, 8, 1), (0, 0, 1)), # stride 0 over a real extent + ((1, 4, 4), (16, 2, 1)), # rows closer together than they are wide + ((1, 2, 2), (4, 2, 2)), # two axes landing on the same elements + ], +) +def test_a_reduction_output_that_writes_an_element_twice_is_rejected(shape, stride): + """Two elements at one address is a write race, not a layout to support. + + The rule is that each axis clears the whole span of the one below it. + Checking the innermost pair alone -- pitch against width -- passes the last + case, whose width is 1 and whose two axes both land on element 2. + """ + from cudnn.frost.buffers import fill_word_strided_async, strided_fill_plan + + assert strided_fill_plan(shape, stride) is None + with pytest.raises(ValueError, match="twice"): + fill_word_strided_async(0, shape, stride, 4, 0, None) + + +@_GPU +@pytest.mark.parametrize("shape,stride", [((1, 8, 1), (32, 4, 1)), ((2, 8, 4), (64, 8, 1)), ((3, 5, 1), (7, 1, 1))]) +def test_a_padded_seed_writes_its_own_elements_and_no_others(shape, stride): + """The engine seeds a padded output itself, without the caller's ``fill_()``. + + Borrowing that method worked only while the buffer happened to be a torch + tensor -- and queued on torch's stream, not the one the kernel will run on. + What it has to get right is exactly this: every element the view covers, and + nothing between them. + """ + from cudnn.frost.buffers import fill_word_strided_async, init_word + + span = 1 + sum((d - 1) * s for d, s in zip(shape, stride)) + flat = torch.zeros(span, dtype=torch.float32, device="cuda") + view = torch.as_strided(flat, shape, stride) + fill_word_strided_async(flat.data_ptr(), shape, stride, 4, init_word("fp32", 3.5), None) + torch.cuda.synchronize() + + expected = torch.zeros(span, dtype=torch.float32, device="cuda") + torch.as_strided(expected, shape, stride).fill_(3.5) + assert torch.equal(flat, expected) + assert torch.equal(view, torch.full(shape, 3.5, device="cuda")) + + +@_GPU +def test_a_padded_reduction_output_is_seeded_on_the_kernel_s_stream(): + """The case that used to reach ``tensor.fill_()``. + + A tap declared into a padded buffer is legal and goes through the public + path like any other. Seeding it by calling ``fill_()`` on the caller's + tensor queued on torch's current stream rather than the one the kernel runs + on -- the same stream only by luck -- and only worked at all while that + buffer was a torch tensor. It is the driver's 2D memset now, so this asserts + both halves: the tap is right, and the padding it does not own is untouched. + """ + g = cudnn.pygraph(io_data_type=BF16, intermediate_data_type=F32, compute_data_type=F32) + A = g.tensor(name="A", dim=[1, M, K], stride=[M * K, K, 1]) + B = g.tensor(name="B", dim=[1, K, N], stride=[K * N, 1, K]) + C = g.matmul(A=A, B=B, name="mm") + Y = g.relu(input=C, name="relu") + Y.set_output(True).set_data_type(BF16) + R = g.reduction(input=Y, mode=cudnn.reduction_mode.ADD, name="red") + R.set_dim([1, M, 1]).set_stride([M * 4, 4, 1]).set_output(True).set_data_type(F32) + _pin_frost(g) + + a, b, ref = _operands() + y = torch.empty(1, M, N, dtype=torch.bfloat16, device="cuda") + pad = torch.full((1, M, 4), -7.0, dtype=torch.float32, device="cuda") + _run(g, {A: a, B: b, Y: y, R: pad[:, :, :1]}) + + torch.testing.assert_close(pad[:, :, 0], torch.relu(ref).sum(dim=2), atol=1.0, rtol=2e-2) + assert torch.equal(pad[:, :, 1:], torch.full((1, M, 3), -7.0, device="cuda"))