From 792cfff2a68374e033e470a54a6d1fafe55cc58c Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Wed, 5 Aug 2026 06:40:02 -0700 Subject: [PATCH 1/8] Add a channels-last-native Triton GroupNorm ScaFFold's GroupNorm was the only operator forcing a layout change: ATen's kernel launches one workgroup per (batch, group) row -- 8 of them at the benchmark's defaults -- so on a 228-CU MI300A it ran at a small fraction of achievable bandwidth, and it was the sole reason channels-last broke. This is an NDHWC-native kernel: Welford statistics, fp32/bf16/fp16 with torch's autocast contract, an int64 tile-base path for tensors past 2^31 elements, and a fused ReLU that is bit-exact against the unfused form. Registered as a custom op with fake kernels and register_autograd, so it traces under torch.compile(fullgraph=True) and composes with DistConv's DCTensor. Determinism is a contract here, not an accident: no float atomics, and the grid, split, tile and reduction order are pure functions of the shape. Verified run-to-run and across interpreters. Nothing imports this yet -- the wiring is the next commit. --- ScaFFold/unet/triton_group_norm.py | 1819 ++++++++++++++++++++++++++ tests/test_triton_group_norm.py | 868 ++++++++++++ tests/test_triton_group_norm_edge.py | 1701 ++++++++++++++++++++++++ 3 files changed, 4388 insertions(+) create mode 100644 ScaFFold/unet/triton_group_norm.py create mode 100644 tests/test_triton_group_norm.py create mode 100644 tests/test_triton_group_norm_edge.py diff --git a/ScaFFold/unet/triton_group_norm.py b/ScaFFold/unet/triton_group_norm.py new file mode 100644 index 0000000..641122e --- /dev/null +++ b/ScaFFold/unet/triton_group_norm.py @@ -0,0 +1,1819 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""Channels-last-native Triton GroupNorm (NDHWC in, NDHWC out). + +Why this exists +=============== +With ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` -- which every production ScaFFold run +sets -- every convolution in the UNet emits ``channels_last_3d`` activations, +but *every* stock GroupNorm variant (eager or Inductor-compiled) consumes them +through the *logical* NCDHW iteration order, which over a channels-last tensor +is a strided gather, and then emits a **contiguous** tensor. Measured on one +MI300A at scale 8 that costs 6.4x on GroupNorm itself (443 ms/step of GN +fwd+bwd against 69 ms/step here) *and* breaks the channels-last chain 22 times +per forward, forcing the following convolution to convert back. + +A channels-last-3d contiguous ``(N, C, D, H, W)`` tensor is *physically* a dense +``(N, S, C)`` array with ``S = D*H*W``; group ``g`` owns a contiguous run of +``C/G`` channels *inside every voxel*. The kernels below therefore let **one +program handle all groups at once** for a chunk of voxels: they read a dense +``(BLOCK_S, C)`` run, reshape the inner axis to ``(G, C/G)``, and get perfectly +coalesced loads and stores with the group axis costing nothing. Measured at +95-98% of this device's streaming roofline at the two largest UNet shapes. + +Measured on one MI300A (228 CUs), fp32, ``num_groups=8``, median of 20, the six +scale-8 UNet GroupNorm shapes, fwd / fwd+bwd in ms:: + + shape this compiled-CL compiled-CONT eager-CL fwd + [1,64,256^3] 4.23/11.39 19.51/77.01 4.77/11.85 151.2 + [1,128,128^3] 1.15/ 3.00 7.23/24.64 1.22/ 3.07 36.5 + [1,256,64^3] 0.35/ 0.97 2.27/ 7.42 0.34/ 0.84 9.1 + [1,512,32^3] 0.14/ 0.57 0.19/ 1.03 0.11/ 0.33 1.7 + [1,1024,16^3] 0.11/ 0.57 0.10/ 0.39 0.07/ 0.33 0.4 + [1,2048,8^3] 0.11/ 0.57 0.08/ 0.40 0.07/ 0.33 0.1 + +Over the 22 scale-8 call sites that is **442.8 -> 67.1 ms/step** of GroupNorm +fwd+bwd against today's production path (compiled GroupNorm on channels-last +input), i.e. **376 ms/step recovered**, and a dead heat with compiled GroupNorm +on *contiguous* input (66.4 ms/step) while additionally not breaking the +layout chain. The three smallest shapes lose on host dispatch, not on GPU +work -- see :func:`select_strategy`. + +The ``this`` column and the rollup were re-measured after the launch folding +below (2+2 kernels instead of 3+4, retuned jointly); the same measurement of +the unfused chain in the same process gives 4.26/11.59, 1.19/3.21, 0.35/1.01, +0.14/0.63, 0.13/0.63, 0.13/0.62 and a 69.5 ms/step rollup, i.e. **-3.4% over +the 22 sites** and -1.7% to -9.5% per shape. The other three columns are from +the earlier sweep and are unchanged. + +Public API +========== +``triton_group_norm(input, num_groups, weight=None, bias=None, eps=1e-5, +activation=None)`` + Drop-in for ``F.group_norm`` (plus an optionally fused ReLU) with + first-order autograd support. Accepts *anything* ``F.group_norm`` accepts; + inputs the Triton kernel cannot serve fall back to ``F.group_norm`` + internally (see "Layouts" below). + +``is_supported(input, num_groups, weight=None, bias=None, activation=None)`` + Cheap, side-effect-free predicate: ``True`` exactly when the native Triton + kernel will run. Callers that already have a good fallback (e.g. a + ``torch.compile``d GroupNorm) should test this and route rejects + themselves; ``triton_group_norm``'s own fallback is plain eager + ``F.group_norm``. + +Contract +======== +For every input ``is_supported`` accepts, the result matches ``F.group_norm`` +to within fp32 reduction-order noise, with: + +* **dtype** -- output dtype is exactly ``F.group_norm``'s. Verified + empirically on this build (torch 2.13.0+rocm7.2): without autocast the output + dtype is the input dtype (fp32/bf16/fp16); under ``torch.autocast("cuda", + ...)`` GroupNorm is an fp32-policy op, so the output is **fp32** for any + input dtype. This module reproduces that rule (see ``_autocast_out_dtype``) + without materializing the fp32 copy of the input that autocast's cast would + create: the kernels read the input at its native width and accumulate in + fp32, which is bit-for-bit the same computation as upcasting first, but reads + half the bytes. Gradients follow the same rule: ``d_input`` has the input's + dtype, ``d_weight``/``d_bias`` have the parameter's dtype. Forward time at + ``[1,64,256^3]`` / ``[1,128,128^3]`` / ``[1,256,64^3]`` in ms: fp32 + 4.28/1.20/0.35, bf16 2.42/0.68/0.22, fp16 2.43/0.66/0.22, and bf16-in with + the fp32-out autocast contract 3.08/0.84/0.28 -- so honouring autocast's + fp32 output still buys 1.4x over fp32 end to end, because only the read side + narrows. +* **statistics** -- always accumulated in fp32, never in the input dtype. +* **memory format** -- the output has the *input's* memory format. This is the + one deliberate difference from stock GroupNorm, which returns a contiguous + tensor for every input layout (measured, see the table in + ``work/gn-dctensor/triton/RESULTS.md``); preserving channels-last is the + entire point of the kernel. +* **autograd** -- ``d_input``, ``d_weight``, ``d_bias``; ``weight=None`` and/or + ``bias=None`` supported. **First order only**: the backward is itself a + custom op with no autograd formula of its own, so a second + ``torch.autograd.grad`` through this op raises ``RuntimeError: Trying to + backward through scaffold_gn.group_norm_backward.default but no autograd + formula was registered``. Stock ``F.group_norm`` *does* support double + backward, so a gradient penalty or a Hessian-vector product must route + around this kernel (``is_supported`` says nothing about second derivatives; + it is documented there too). It fails loudly rather than returning garbage. +* **device** -- the kernels run on the *input's* device, whatever device is + current, matching ATen's ``DeviceGuard`` behaviour; see ``_device_guard``. +* **determinism** -- bitwise reproducible run to run and process to process. + There are no float atomics anywhere, and the grid, split count and tile sizes + are pure functions of the shape (the tuning table is frozen in this file for + exactly that reason -- a *runtime* autotuner would break reproducibility by + changing the reduction order between runs). +* **rejections** -- every shape/dtype/parameter combination ``F.group_norm`` + raises on is one ``is_supported`` returns ``False`` for, including the + degenerate "1 value per channel" shape (``N*(C/G)*D*H*W == 1``), so a caller + that branches on ``is_supported`` never gets an answer where the op this + replaces would have raised. +* **eps** -- one deliberate divergence, at a value no run uses: for a + *subnormal* fp32 ``eps`` (``< 1.18e-38``) on a zero-variance group the GPU + flushes ``var + eps`` to zero, so ``rstd`` is ``inf`` and ``y`` is ``NaN`` + where ATen stays finite. The boundary is exactly the normal/subnormal one + (``eps=1.2e-38`` gives ``rstd=9.1e18``, ``eps=1e-38`` gives ``inf``); at + ``eps == 0`` both implementations produce non-finite output identically. + Left as is rather than clamped because clamping would perturb every + ordinary call to defend a value nine orders of magnitude below the smallest + plausible one. + +Reduction strategy +================== +Group statistics span ``S * C/G`` elements (134M at the largest UNet shape), so +one pass cannot produce them. Split-K partial reductions land at a fixed +scratch index and are combined by a fixed-order tree:: + + fwd: stats_partial -> normalize (2 kernels) + bwd: bwd_partial -> dx (2 kernels) + +Traffic (``B = numel * itemsize``): 3B forward, 5B backward. + +Each pass is **two** launches, not the three and four an unfused split-K chain +needs: the two finalize passes and the dweight/dbias row reduction are folded +into the elementwise kernel that consumes them. ``_normalize_kernel`` +re-derives ``mean``/``rstd`` from the split-K partials itself (and program 0 +stores them for the backward); ``_dx_kernel`` re-derives ``c1``/``c2`` the same +way and its first ``ceil(C/BLOCK_C)`` programs also do the dweight/dbias +reduction. Folding plus the retuning below is worth 8.0-9.5% of fwd+bwd at the +four smallest shapes, which are host-dispatch bound, 4.9-6.8% at the two middle +ones and 1.7% at the largest. + +The catch, and the reason the tuning table was re-derived rather than inherited: +**the fusion and the tiling are one problem, not two.** A fused finalize is +recomputed by every elementwise *program*, so its cost is +``nprog_elem * nsplit`` triples of redundant (L2-resident) traffic. Keeping the +unfused table's ``nsplit_target=2048`` at ``[1,64,256^3]``, whose flat +elementwise grid is 131072 programs, asks for 25.7 GB of redundant reads +against a 4.3 GB tensor and costs **+34% of fwd+bwd** (+70% of the forward). +Two things fix it, both of them in ``GNConfig``: the elementwise grid is capped +at ``elem_progs`` programs which then stride over the tiles (so the redundancy +is bounded by the *grid*, not by the tile count), and ``nsplit_target`` is +retuned per shape against that cap. With both, the same shape is 1-2% *faster* +than the unfused chain. Do not change one without re-running the other; the +coordinate-descent tuner is ``work/gn-dctensor/kernel-opt/tune.py``. + +Why not one launch per pass +--------------------------- +A device-scope software barrier (int32 atomics with volatile loads, no float +atomics, so still bitwise deterministic) collapses each pass to a single +launch and was measured at 2.3-2.6x on the four smallest shapes. It is +deliberately **not** used. A grid barrier requires every workgroup to be +co-resident, which caps the grid at the CU count (228 here); the kernel then +tops out at 0.5-0.9 TB/s against split-K's 2.7, so it loses catastrophically +the moment the shape is bandwidth-bound rather than dispatch-bound -- +**18.0 ms against 3.0 ms at [1,128,128^3]**, and it does not compile at all at +``[1,64,256^3]``. Serving both regimes therefore means shipping two kernel +families plus a crossover rule, for a whole-model gain of ~3% (65.7 -> 63.6 +ms/step at scale 8); and under CUDA-graph capture, where launch count is free, +the ten small-shape sites are already only 1.05 ms of a 64.1 ms/step total, so +the gain is zero. Hand-rolled inter-workgroup synchronisation is not a good +trade for 3% in a benchmark whose value depends on being trustworthy and +reproducible. The measurements are in +``work/gn-dctensor/triton-small/RESULTS.md``. + +Numerics: Welford, not ``E[x^2]-E[x]^2`` +======================================== +The prototype accumulated ``sum(x)`` and ``sum(x*x)`` and formed +``var = E[x^2] - E[x]^2``. That is split-friendly and cheap but cancels +catastrophically once ``mean >> std``, because it subtracts two nearly equal +large numbers to recover a small one. + +Here each tile instead produces ``(count, mean, M2)`` via a *corrected* +two-pass over registers -- ``mean0 = sum(x)/n``, then ``corr = sum(x-mean0)/n`` +to recover the digits the first sum lost, then ``M2 = sum((x-mean0-corr)^2)`` +-- and tiles and splits are merged with Chan's parallel combine. Every step is +register-only (the tile is read from HBM exactly once either way) and +atomic-free, so neither the traffic model nor determinism changes. + +Measured at ``[1,256,64^3]``, ``num_groups=8``, affine, relative error of the +*output* against a float64 reference computed from the same fp32 samples: + + x ~ N(mu, sigma) this kernel ATen fp32 E[x^2]-E[x]^2 + mu=0, sigma=1 1.6e-07 4.1e-07 1.8e-07 + mu=10, sigma=1 3.0e-07 6.9e-07 9.9e-06 + mu=100, sigma=1 8.0e-07 4.2e-06 5.6e-04 + mu=1e3, sigma=1e-2 1.1e-04 2.5e-03 2.3e+00 + +At ``mu/sigma = 1e5`` the old formulation has lost the variance outright (the +difference of the two ~1e6-sized fp32 terms is below one ulp, so ``rstd`` +saturates on ``eps`` and the output is meaningless), while this kernel is still +good to 1.1e-04 -- and is 4-23x *more* accurate than ATen's own fp32 GroupNorm +at every non-trivial mean. The residual 1.1e-04 is the fp32 representation +floor rather than an algorithm defect: a mean of 1e3 held in fp32 is quantized +to ~6e-5, which is 6e-3 of a standard deviation here, and both kernels sit on +that floor. + +Cost of the rewrite, isolated by timing the kernels alone against the +prototype's: **+0.8%** on the forward at ``[1,64,256^3]`` (the shape that +dominates the step), +3-8% at the middle shapes, and **0%** on the backward, +which does not compute a variance. A cheaper shifted-mean variant (two tile +reductions instead of three, shift taken from a peeled first tile) would +recover most of that; it was not worth the extra failure mode for ~4 ms/step. + +What the third pass (``corr``) is worth, separately +--------------------------------------------------- +The accuracy above is mostly the *two-pass* structure; the ``corr`` term is a +third reduction on top of it and deserves its own accounting. Deleting it +outright (keeping ``mean_t = mean0``, ``M2 = sum((x-mean0)^2)``) and comparing +both against float64 on the same fp32 samples, relative error of ``rstd``, +10 seeds each: + + regime with corr without ratio + one tile per group reduction (nsplit=1): + [2,64,8,4,4] G=4, mu/sigma=1e6 1.1e-07 1.6e-04 1472x + [2,64,8,4,4] G=2, mu/sigma=1e6 8.4e-08 4.9e-05 580x + [2,64,8,4,4] G=1, mu/sigma=1e6 9.8e-08 2.1e-05 216x + [2,64,8,4,4] G=4, mu/sigma=1e5 9.9e-08 8.4e-06 85x + many tiles and splits (the production configs): + [1,512,32^3] G=8, mu/sigma=1e5 1.2e-05 4.4e-05 3.8x + [1,1024,16^3] G=8, mu/sigma=1e5 8.6e-06 2.5e-05 2.9x + [1,256,24^3] G=8, mu/sigma=1e5 3.0e-05 3.7e-05 1.3x + [1,256,24^3] G=8, mu/sigma=1e7 5.8e-04 2.3e-06 0.004x + +So it is decisively load-bearing exactly where the tile mean is formed from +many large values -- up to 1472x on ``rstd`` -- and worth a steady 1.3-4x in +the multi-split configs the tuning table actually picks, at ``mu/sigma = 1e5``. +Past ``mu/sigma ~ 1e6`` with many splits it can go the *other* way (last row): +there the true spread between tile means is smaller than one ulp of the means +themselves, so Chan's between-tile term is computed from quantization noise +either way and the uncorrected version's inflated ``M2`` partly cancels it. +That regime is past fp32's floor for this computation (a mean of 1e6 held in +fp32 quantizes to 0.06, i.e. 6% of a standard deviation at sigma=1) and no +production input is near it. + +The *output* error is nearly unmoved by any of this -- at most ~1.4x in either +direction -- which is why the term looks free to delete if you only measure +``y``: the output is dominated by the fp32 representation of the mean, which +``corr`` cannot improve (``mean0 + corr`` rounds straight back to ``mean0`` +once the mean is large). It is ``rstd`` that carries the benefit. + +Price, measured the same way (median of 20 forwards, correction removed +outright rather than zeroed): **+2.3%** of the forward at ``[1,64,256^3]``, ++3.0% at ``[1,128,128^3]``, +4.7% at ``[1,256,64^3]``, and nothing measurable +(-1.2% to +0.5%, i.e. noise) at the three launch-bound shapes. At the shape +that dominates the step that is +0.10 ms of a 11.6 ms fwd+bwd, i.e. +0.9%. +Kept: a 1.3-1472x accuracy factor on the statistic the whole rewrite exists to +protect is worth ~1% of GroupNorm time. The load-bearing case is pinned by +``test_welford_correction_recovers_rstd_in_a_single_tile_reduction`` in +``tests/test_triton_group_norm_edge.py``, so deleting the term now fails the +suite instead of passing it silently. + +Layouts +======= +* ``channels_last_3d`` 5-D input -> **native Triton kernel**, channels-last + output. This is the fast path and the only one ``is_supported`` accepts. +* Plain contiguous NCDHW (and every other layout/rank) -> ``triton_group_norm`` + falls back to ``F.group_norm``, which returns a contiguous tensor, so the + input's memory format is still preserved. ``is_supported`` returns ``False`` + so that callers keep their own (probably compiled) fallback rather than + silently dropping to the eager kernel. + + This is a deliberate scope decision, not an oversight. A native NCDHW kernel + would need a *different* tiling -- with C outermost the fast axis is spatial, + so a program must own one group and stream S, rather than owning all groups + and streaming voxels -- i.e. a second family of four kernels. The payoff is + small: on contiguous input Inductor's compiled GroupNorm already reaches + 89-92% of this device's measured streaming roofline (RESULTS.md 4) -- and the + table above confirms it, 66.4 ms/step against this kernel's 67.1 -- so a + native NCDHW kernel could win ~10% there, against the 6.4x it wins on + channels-last input. If a mixed-layout model ever makes that 10% matter, the + place to add it is the strategy hook below. + +Addressing +========== +``[2, 64, 256^3]`` is *exactly* 2^31 elements, so int32 linear offsets block +batch>1 at the largest UNet shape and every shape above it. The kernels widen +only the **scalar tile base** to int64 (``INT64`` is a ``tl.constexpr``, so +shapes that fit still emit pure 32-bit code); the vector offsets inside a tile +span at most ``BLOCK_S*C + C`` elements and stay int32 either way. That is why +the wide path is free: forcing int64 on every scale-8 shape moves fwd+bwd by +-0.8% to +0.8% and forward by -3.7% to +4% (a 0.02 ms swing on the two smallest, +launch-bound shapes) -- noise in both directions. The switch is kept anyway +because it costs one constexpr and documents where the boundary is; correctness +above 2^31 elements is covered by a test at ``[2, 64, 256, 256, 257]`` +(2_155_872_256 elements, 2.5e-05 relative error on *both* batch items, i.e. the +same reduction noise a 134M-element fp32 reduction has anywhere). + +Fused activation +================ +``activation="relu"`` folds the ReLU into the forward store. In a store-bound +kernel that is free (one compare and one select) and it removes an entire 2B +streaming pass. Measured against ``F.relu(triton_group_norm(x))``: 39% off the forward +and 35% off fwd+bwd at ``[1,64,256^3]`` (6.83 -> 4.20 ms and 17.82 -> 11.61 +ms), 38%/35% at ``[1,128,128^3]``, 34%/30% at ``[1,256,64^3]``, tapering to +21%/9% at ``[1,512,32^3]`` and below, where the call is host bound and there is +less streaming pass to remove. + +The backward gates the incoming gradient on the sign of the **pre-activation** +value, which it *recomputes* from the saved ``(x, mean, rstd, weight, bias)`` +using the identical expression the forward used. Recomputation costs two FLOPs +on values already in registers and is bit-exact -- same inputs, same operation +order, same fp32 rounding -- so the sign always agrees with the forward's. The +alternative, testing ``y > 0`` on the saved output, would need the output kept +alive *in addition to* ``x`` (which the GroupNorm backward needs regardless), +and in bf16/fp16 it would also mis-gate any element whose positive +pre-activation rounded to zero on the store. + +Both the store and the gate are spelled as the *complement* of the usual test +(``tl.where(y <= 0, 0, y)``, ``tl.where(pre <= 0, 0, dy)``) rather than as +``tl.maximum(y, 0)`` / ``tl.where(pre > 0, dy, 0)``. The two are identical on +every finite value but not on NaN: ``tl.maximum`` returns the *non*-NaN operand +and ``NaN > 0`` is False, so both of the usual spellings silently map a NaN to +0.0, while ``F.relu`` propagates it and ``threshold_backward(grad, result, 0)`` +-- ReLU's real backward -- passes its gradient (``NaN <= 0`` is False too). +Matching ``F.relu`` here is not pedantry: a diverging run whose forward comes +back finite because the fused activation ate the NaN passes straight through +ScaFFold's non-finite-loss abort and checkpoints a broken model. ``+-Inf`` and +``-0.0`` are bit-identical under either spelling (``-0.0`` flushes to ``+0.0``, +as ``F.relu`` does). Cost: nil, measured -- see ``FastGroupNorm``'s tests and +``work/gn-dctensor/wiring-fixes``. + +Composition +=========== +Registered as real dispatcher ops (``scaffold_gn::group_norm`` / +``scaffold_gn::group_norm_backward``) via ``torch.library.custom_op``, with a +fake/meta kernel and ``register_autograd``. Consequences: + +* ``torch.compile(..., fullgraph=True)`` traces through without a graph break. +* Tensor subclasses that dispatch via ``__torch_dispatch__`` -- notably + DistConv's ``DCTensor`` -- intercept the op, unwrap to the local shard, run + it, and rewrap, so a DCTensor goes in and a DCTensor comes out with the graph + intact. As with the rest of DistConv today, statistics are per-shard. + +Where the host time goes, and what is left +========================================== +Composing has a price, and at the launch-bound shapes it is now the *dominant* +cost. Peeling the layers at ``[1,2048,8^3]``, steady-state wall clock per +fwd+bwd (median of 200, min of 7 rounds; GPU work is 0.030 ms):: + + kernels + this file's Python (_forward/_backward called directly) 0.145 ms + + torch.library dispatcher (both custom ops) +0.054 ms + + autograd (register_autograd node, save_for_backward, ctx) +0.338 ms + = triton_group_norm(x).backward(dy) 0.537 ms + + for scale: an *empty* python torch.autograd.Function, fwd+bwd 0.065 ms + +So **63% of the call is the autograd layer** and 10% is the dispatcher -- both +of them the cost of being a real dispatcher op that ``torch.compile`` and +``DCTensor`` can see, which is the whole point of registering it that way. Of +the 0.145 ms this file is responsible for, 0.030 ms is GPU and the remaining +0.115 ms is four launches plus the allocations, plan lookup and argument +binding around them: launching every kernel twice measures the marginal cost of +a whole invocation site at **35 us**, so the four of them are ~0.14 ms of host +work that the GPU work does not cover. + +Two things were considered for that 0.14 ms and rejected: + +* **Bypassing ``JITFunction.run`` for a cached ``CompiledKernel`` handle** + (8.68 us -> 3.94 us per launch on this node) would recover ~19 us, i.e. 3.3% + of the call. It buys that by asserting that Triton's specialization key -- + including 16-byte pointer alignment -- is a pure function of the shape. It is + not: this module accepts channels-last *views with a storage offset* and + non-contiguous affine parameters, both of which the test suite exercises, and + a stale specialization there is a wrong answer rather than a crash. 3.3% on + the host-bound shapes only is not worth a silent-miscompute failure mode. +* **Caching the scratch buffers** across calls saves ~1.7 us per allocation, + ~20 us here, and makes the buffers shared mutable state across call sites -- + correct on one stream, wrong on two, and this module has no way to know. + +What that leaves: the kernels are at 95-98% of the streaming roofline at the +two largest shapes and the fused chain is 1.7-6.8% faster than the unfused one +there, so there is no meaningful GPU headroom left. At the launch-bound +shapes the remaining 0.4 ms is torch's own plumbing, and the two ways to remove +it are both outside this file: CUDA-graph capture of the training step (which +takes the ten small scale-8 sites to ~1.05 ms/step of GPU time in total), or a +C++ autograd node. + +Triton is imported lazily, on the first call that actually reaches the kernel, +so importing this module (or running the CPU test suite) costs nothing. +""" + +import contextlib +import functools +import importlib.util +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F + +__all__ = [ + "triton_group_norm", + "is_supported", + "select_strategy", + "GNConfig", + "default_config", + "SUPPORTED_ACTIVATIONS", + "TritonKernelError", +] + + +class TritonKernelError(RuntimeError): + """A failure of the Triton kernels themselves, with the original as ``__cause__``. + + Raised in place of whatever ``_forward``/``_backward`` raised -- a missing or + mismatched ``triton``, an unwritable JIT cache, a compile error, a launch + failure, an API change between Triton releases. It exists so that a caller + with a fallback (``ScaFFold.unet.group_norm``'s ladder) can catch *exactly* + "the kernel is broken" and nothing else, instead of catching ``Exception`` + and trying to enumerate every framework mechanism that legitimately raises + through a forward -- saved-tensor pack hooks, ``torch.utils.checkpoint``'s + recompute control flow, functorch, a user's offloading hook. + + Two things are deliberately *not* tagged and therefore propagate unchanged: + + * ``torch.OutOfMemoryError``, which is a resource condition rather than a + defect (every fallback allocates an output of the same size, so retrying + one is a second, differently-shaped OOM at a call site the caller did not + ask about), and + * the ``ValueError``s ``_validate`` raises, which are contract violations by + the caller. ``is_supported`` accepts exactly what ``_validate`` accepts, + so a caller that branches on it can never see one; if one escapes, that + is a bug in this module and must be loud. + + The tagged region contains no autograd-observable work -- allocations and + kernel launches only, with ``save_for_backward`` happening in + ``_setup_context`` strictly *after* ``_forward`` returns -- so an exception + that carries this type is guaranteed to have been raised before the op saved + anything. That is what makes retrying the call on another kernel safe. + """ + + +#: The activations that may be fused into the forward store. +SUPPORTED_ACTIVATIONS = (None, "relu") + +#: Input dtypes the kernels read directly (statistics are always fp32). +SUPPORTED_DTYPES = (torch.float32, torch.bfloat16, torch.float16) + +#: Largest linear element index representable in int32. +_INT32_MAX = 2**31 - 1 + + +# --------------------------------------------------------------------------- # +# tiling configuration +# --------------------------------------------------------------------------- # +class GNConfig: + """Tiling knobs. A pure function of the shape => bitwise determinism. + + ``stats_tile``/``elem_tile`` are *element* budgets per program (the spatial + block is ``tile // channels_per_voxel``, rounded down to a power of two); + ``nsplit_target`` is the total number of split-K partials wanted across the + batch, so the per-sample split count is ``nsplit_target // N``; + ``elem_progs`` caps the elementwise grid, each program then striding over + ``ceil(nblk_elem / elem_progs)`` tiles (0 = one program per tile). + + These are **not** independent knobs, and in particular they stopped being + independent when the finalize passes were folded into the elementwise + kernels: each elementwise *program* now re-reads all ``nsplit`` split-K + partials, so the redundant traffic is ``min(nblk_elem, elem_progs) * + nsplit`` triples. Raising ``nsplit_target`` for stats-kernel occupancy and + lowering ``elem_progs`` for redundancy pull against each other and were + tuned together; see the module docstring. + """ + + __slots__ = ( + "stats_tile", + "stats_warps", + "nsplit_target", + "elem_tile", + "elem_warps", + "elem_progs", + ) + + def __init__( + self, + stats_tile=8192, + stats_warps=4, + nsplit_target=2048, + elem_tile=8192, + elem_warps=4, + elem_progs=2048, + ): + self.stats_tile = stats_tile + self.stats_warps = stats_warps + self.nsplit_target = nsplit_target + self.elem_tile = elem_tile + self.elem_warps = elem_warps + self.elem_progs = elem_progs + + def key(self): + return ( + self.stats_tile, + self.stats_warps, + self.nsplit_target, + self.elem_tile, + self.elem_warps, + self.elem_progs, + ) + + def __eq__(self, other): + return isinstance(other, GNConfig) and self.key() == other.key() + + def __hash__(self): + return hash(self.key()) + + def __repr__(self): + return ( + "GNConfig(stats_tile=%d, stats_warps=%d, nsplit_target=%d, " + "elem_tile=%d, elem_warps=%d, elem_progs=%d)" % self.key() + ) + + +#: Frozen tuning table, produced by coordinate descent on fwd+bwd time on one +#: MI300A (228 CUs) at fp32 with ``num_groups=8``, keyed by the +#: ``(num_channels, cube-root spatial extent)`` of the scale-8 ScaFFold UNet +#: GroupNorm sites plus the ``[1,4096,4^3]`` tail. Frozen -- never autotuned at +#: run time -- because the split count fixes the reduction order and therefore +#: the bits of the result. +#: +#: Re-derived for the fused (2+2 launch) kernels: every candidate was timed +#: **interleaved against the incumbent** in one process, so that anything which +#: changes on the device partway through a sweep -- a neighbour process above +#: all -- lands on both arms at once instead of on whichever ran later. The +#: table is keyed on ``(C, edge)`` and not on ``N``: ``nsplit_target`` is a target for +#: the split count *summed over the batch* (the per-sample count is +#: ``nsplit_target // N``), so the same entry serves ``N > 1`` with the same +#: total number of stats programs. Verified at ``[2,1024,16^3]``, +#: ``[4,2048,8^3]`` and ``[2,256,64^3]``. +_TUNED = { + (64, 256): GNConfig(16384, 4, 2048, 8192, 4, 2048), + (128, 128): GNConfig(16384, 4, 2048, 16384, 4, 912), + (256, 64): GNConfig(16384, 4, 512, 16384, 8, 912), + (512, 32): GNConfig(32768, 8, 4096, 16384, 4, 0), + (1024, 16): GNConfig(65536, 4, 8192, 16384, 4, 0), + (2048, 8): GNConfig(16384, 8, 1024, 8192, 8, 228), + (4096, 4): GNConfig(16384, 4, 32, 4096, 8, 0), +} + +_DEFAULT_CONFIG = GNConfig() + + +def default_config(num_channels: int, spatial: int) -> GNConfig: + """Tiling for ``num_channels`` channels and ``spatial = D*H*W`` voxels.""" + edge = round(spatial ** (1.0 / 3.0)) + if edge**3 != spatial: + edge = None + return _TUNED.get((num_channels, edge), _DEFAULT_CONFIG) + + +# --------------------------------------------------------------------------- # +# small-shape dispatch hook +# --------------------------------------------------------------------------- # +#: Every strategy name ``select_strategy`` may return. Only ``"split_k"`` is +#: implemented; anything else raises rather than silently doing the wrong +#: thing. +STRATEGIES = ("split_k",) + +#: Spatial extent (``D*H*W``) below which the split-K chain is host-dispatch +#: bound rather than bandwidth bound. Measured on MI300A: at ``[1,2048,8^3]`` +#: the kernels do 0.030 ms of GPU work behind ~0.58 ms of +#: Python/autograd/launch cost, and 0.086 ms of that is the *empty* +#: ``torch.autograd.Function`` wrapper -- i.e. 68% of the remaining call is +#: torch's plumbing, not this file's. Purely informational -- +#: ``select_strategy`` does not use it. +SMALL_SPATIAL_THRESHOLD = 4096 + + +def select_strategy(n: int, num_channels: int, spatial: int, num_groups: int) -> str: + """### SMALL-SHAPE DISPATCH HOOK ### -- the single point where a different + kernel strategy is chosen for a shape. + + Returns a name from :data:`STRATEGIES`. Today it always returns + ``"split_k"``: two forward and two backward kernels with split-K partial + reductions and the finalize passes fused into their consumers. That is + bandwidth-optimal for the large shapes and, after the fusion, within + ~0.09 ms of the floor a Python ``autograd.Function`` can reach at the small + ones -- so there is much less left here than there looks. Below roughly + ``SMALL_SPATIAL_THRESHOLD`` voxels the call is host bound, but the host cost + is now dominated by autograd and the dispatcher rather than by launches: + see "Why not one launch per pass" in the module docstring for the one + strategy that *would* cut it further and why it is not here. + + If a second strategy ever lands, add its name to :data:`STRATEGIES`, return + it from here on a rule that is a **pure function of the shape** + (determinism depends on it), and branch on it in ``_dispatch`` -- which is + the only caller, sits in front of the memoized tiling plan, and is itself + called by both ``_forward`` and ``_backward``. Nothing else in this file + needs to change. + """ + return "split_k" + + +# --------------------------------------------------------------------------- # +# planning helpers +# --------------------------------------------------------------------------- # +def _prev_pow2(x: int) -> int: + p = 1 + while p * 2 <= x: + p *= 2 + return p + + +def _next_pow2(x: int) -> int: + p = 1 + while p < x: + p *= 2 + return p + + +def _cdiv(a: int, b: int) -> int: + return -(-a // b) + + +class _Plan: + """Everything the launcher needs, derived only from the shape + config.""" + + __slots__ = ( + "n", + "channels", + "spatial", + "groups", + "group_channels", + "groups_p2", + "group_channels_p2", + "masked_c", + "int64", + "block_s_stats", + "nsplit", + "chunk", + "block_s_elem", + "nblk_elem", + "nprog_elem", + "elements_per_group", + "dwdb_rows", + "dwdb_block_c", + "dwdb_block_r", + "dwdb_progs", + "grid_dx", + "zero_dx", + "cfg", + ) + + def __init__(self, n, channels, spatial, groups, cfg, numel): + self.n = n + self.channels = channels + self.spatial = spatial + self.groups = groups + self.group_channels = channels // groups + self.groups_p2 = _next_pow2(groups) + self.group_channels_p2 = _next_pow2(self.group_channels) + # Only power-of-two group/channel counts tile the (G, C/G) axes exactly; + # anything else is rounded up and masked, which is correct but reads a + # few lanes it throws away. + self.masked_c = ( + self.groups_p2 != groups or self.group_channels_p2 != self.group_channels + ) + # int64 addressing is needed once a linear element index can exceed + # INT32_MAX. [2,64,256^3] is exactly 2^31 elements, so this is not + # hypothetical at scale. Only the *scalar* tile base is widened (see + # the kernels), which measurement showed to be free. + self.int64 = numel + channels > _INT32_MAX + self.cfg = cfg + + voxel = self.groups_p2 * self.group_channels_p2 + self.block_s_stats = max(1, _prev_pow2(cfg.stats_tile // max(1, voxel))) + ntiles = max(1, spatial // self.block_s_stats) + self.nsplit = _prev_pow2( + max(1, min(ntiles, max(1, cfg.nsplit_target // max(1, n)))) + ) + self.chunk = _cdiv(spatial, self.nsplit) + self.block_s_elem = max(1, _prev_pow2(cfg.elem_tile // max(1, voxel))) + self.nblk_elem = _cdiv(spatial, self.block_s_elem) + # Grid cap for the two elementwise kernels; each program then strides + # over its share of the tiles. Bounds the cost of the fused finalize, + # which every *program* pays once. + self.nprog_elem = ( + self.nblk_elem + if cfg.elem_progs <= 0 + else min(self.nblk_elem, cfg.elem_progs) + ) + self.elements_per_group = float(spatial * self.group_channels) + # Everything the fused dweight/dbias reduction in _dx_kernel needs. + # Precomputed rather than derived per call: the launch-bound shapes pay + # every Python statement in _backward, and _next_pow2 is a loop. + self.dwdb_rows = n * self.nsplit + self.dwdb_block_c = min(256, max(64, _next_pow2(channels))) + self.dwdb_block_r = 32 if self.dwdb_rows >= 32 else 1 + self.dwdb_progs = _cdiv(channels, self.dwdb_block_c) + # Programs past nprog_elem run no elementwise loop iterations; they + # exist only when there are more dweight/dbias blocks than tiles. + self.grid_dx = max(self.nprog_elem, self.dwdb_progs) + self.zero_dx = self.group_channels * spatial == 1 + + +@functools.lru_cache(maxsize=256) +def _plan(n, channels, spatial, groups, numel) -> _Plan: + """Memoized: a UNet presents a handful of shapes, and the three smallest + GroupNorm sites are host-dispatch bound, so rebuilding the plan (two + power-of-two loops and a dict lookup) on every call is measurable there. + Memoization cannot affect results -- the plan is a pure function of its + arguments, which is also what makes the kernels bitwise reproducible.""" + return _Plan(n, channels, spatial, groups, default_config(channels, spatial), numel) + + +def _dispatch(n, channels, spatial, groups, numel) -> _Plan: + """Consult the strategy hook, then build (or reuse) the tiling plan.""" + strategy = select_strategy(n, channels, spatial, groups) + if strategy != "split_k": + raise NotImplementedError( + f"kernel strategy {strategy!r} selected by select_strategy() is not " + f"implemented; known strategies are {STRATEGIES}" + ) + return _plan(n, channels, spatial, groups, numel) + + +# --------------------------------------------------------------------------- # +# Triton kernels (built lazily -- importing this module must not import triton) +# --------------------------------------------------------------------------- # +triton = None +tl = None +_welford_combine = None +_stats_partial_kernel = None +_normalize_kernel = None +_bwd_partial_kernel = None +_dx_kernel = None + + +_TRITON_AVAILABLE = None + + +def triton_available() -> bool: + """Whether Triton is importable, without importing it. + + ``find_spec`` on a top-level name only touches the finders, so this stays + side-effect free and is safe to call from :func:`is_supported`. The answer + is memoized in a plain global rather than an ``lru_cache`` because Dynamo + warns (loudly, once per process) when it traces through a cache wrapper. + """ + global _TRITON_AVAILABLE + if _TRITON_AVAILABLE is None: + try: + _TRITON_AVAILABLE = importlib.util.find_spec("triton") is not None + except (ImportError, ValueError): + _TRITON_AVAILABLE = False + return _TRITON_AVAILABLE + + +def _build_kernels(): + """Import Triton and install the JIT kernels into this module's globals. + + The kernels are defined inside a function purely so that ``import triton`` + is deferred to the first GPU call; they are written into ``globals()`` so + Triton's name resolution (which reads ``fn.__globals__``) sees them. + """ + global triton, tl + import triton as _triton + import triton.language as _tl + + triton = _triton + tl = _tl + + # ---------------------------------------------------------------- stats -- + @_triton.jit + def _welford_combine(cnt_a, mean_a, m2_a, cnt_b, mean_b, m2_b): + """Chan's parallel merge of two (count, mean, M2) triples. + + Exact for empty partials on either side (``cnt == 0`` leaves the other + triple untouched), which matters because the last split of a shape + whose spatial extent is not a multiple of the chunk size can be empty. + """ + cnt = cnt_a + cnt_b + denom = tl.where(cnt == 0.0, 1.0, cnt) + delta = mean_b - mean_a + mean = mean_a + delta * (cnt_b / denom) + m2 = m2_a + m2_b + delta * delta * (cnt_a * cnt_b / denom) + return cnt, mean, m2 + + @_triton.jit + def _stats_partial_kernel( + X, + PCNT, + PMEAN, + PM2, + S, + CHUNK, + C: tl.constexpr, + G: tl.constexpr, + CG: tl.constexpr, + GP: tl.constexpr, + CGP: tl.constexpr, + NSPLIT: tl.constexpr, + BLOCK_S: tl.constexpr, + MASKED_C: tl.constexpr, + INT64: tl.constexpr, + ): + """One program per ``(split, n)``: Welford partials for every group. + + Reads a dense ``(BLOCK_S, C)`` run of memory per step -- perfectly + coalesced -- and produces the statistics of all G groups at once, the + group axis being the inner channel axis reshaped to ``(G, C/G)``. + """ + sp = tl.program_id(0) + n = tl.program_id(1) + + offs_g = tl.arange(0, GP) + offs_j = tl.arange(0, CGP) + offs_s = tl.arange(0, BLOCK_S) + inner = offs_g[None, :, None] * CG + offs_j[None, None, :] + cmask = (offs_g[None, :, None] < G) & (offs_j[None, None, :] < CG) + off = offs_s[:, None, None] * C + inner + + s_begin = sp * CHUNK + s_end = tl.minimum(s_begin + CHUNK, S) + + cnt = tl.zeros((GP,), dtype=tl.float32) + mean = tl.zeros((GP,), dtype=tl.float32) + m2 = tl.zeros((GP,), dtype=tl.float32) + + for s0 in range(s_begin, s_end, BLOCK_S): + # Only the scalar tile base is ever widened to int64; the vector + # offsets stay int32 because they span at most BLOCK_S*C+C + # elements. That keeps the wide arithmetic off the hot path. + if INT64: + base = (n.to(tl.int64) * S + s0) * C + else: + base = (n * S + s0) * C + nvalid = tl.minimum(BLOCK_S, s_end - s0) + m = tl.broadcast_to((offs_s < nvalid)[:, None, None], (BLOCK_S, GP, CGP)) + if MASKED_C: + m = m & cmask + # `other` is not load-bearing for any *bounded* value: `cnt_t` + # counts only the valid lanes, so `corr` below evaluates to + # `sum_valid(x)/cnt_t - mean0` and `mean_t = mean0 + corr` is the + # true mean whatever the masked lanes contributed, while `d`/`dd` + # are re-masked before they reach `m2_t`. (Bounded: `other=1e30` + # would swamp `mean0` and the correction with it, and `other=inf` + # or `nan` would poison it outright.) 0.0 is kept because it is + # the value that survives all three of those, not because the + # cancellation is something to rely on. + x = tl.load(X + base + off, mask=m, other=0.0).to(tl.float32) + + # Corrected two-pass within the tile: the first mean loses digits + # to the magnitude of the data, `corr` puts them back, and the + # centred squares are then accurate to fp32 roundoff. Everything + # here is register traffic; the tile is read from HBM exactly once. + # `corr` is worth 1.3x to 1472x on `rstd` once the data's mean + # dominates its spread -- see "What the third pass is worth" in the + # module docstring for the measurements, for the one regime where + # it goes the other way, and for why the *output* error barely + # moves even where `rstd` improves by three orders of magnitude. + cnt_t = (nvalid * CG).to(tl.float32) + mean0 = tl.sum(tl.sum(x, 2), 0) / cnt_t + d = tl.where(m, x - mean0[None, :, None], 0.0) + corr = tl.sum(tl.sum(d, 2), 0) / cnt_t + dd = tl.where(m, d - corr[None, :, None], 0.0) + m2_t = tl.sum(tl.sum(dd * dd, 2), 0) + mean_t = mean0 + corr + + new_cnt = cnt + cnt_t + delta = mean_t - mean + mean = mean + delta * (cnt_t / new_cnt) + m2 = m2 + m2_t + delta * delta * (cnt * cnt_t / new_cnt) + cnt = new_cnt + + o = (n * NSPLIT + sp) * G + offs_g + gm = offs_g < G + tl.store(PCNT + o, cnt, mask=gm) + tl.store(PMEAN + o, mean, mask=gm) + tl.store(PM2 + o, m2, mask=gm) + + # ------------------------------------------------------------ normalize -- + @_triton.jit + def _normalize_kernel( + X, + Y, + PCNT, + PMEAN, + PM2, + MEAN, + RSTD, + W, + B, + S, + M, + eps, + C: tl.constexpr, + G: tl.constexpr, + CG: tl.constexpr, + GP: tl.constexpr, + CGP: tl.constexpr, + NSPLIT: tl.constexpr, + BLOCK_S: tl.constexpr, + NBLK: tl.constexpr, + NPROG: tl.constexpr, + RELU: tl.constexpr, + HAS_W: tl.constexpr, + HAS_B: tl.constexpr, + MASKED_C: tl.constexpr, + INT64: tl.constexpr, + ): + """Finalize the split-K statistics, then normalize NBLK/NPROG tiles. + + The finalize is *recomputed by every program* rather than round-tripped + through its own kernel launch: merging NSPLIT Welford triples is a few + KB of L2-resident traffic and a tree reduction over a ``(NSPLIT, GP)`` + tile, which is cheaper than the ~9 us launch it replaces. What it is + *not* cheap enough for is being paid once per tile at the largest + shapes, where the flat grid is 10^5 programs: the grid is therefore + capped at ``NPROG`` and each program strides over its share of the + ``NBLK`` tiles, so the redundant read costs ``NPROG * NSPLIT`` and not + ``NBLK * NSPLIT``. See :class:`GNConfig` -- ``nsplit_target``, + ``elem_tile`` and ``elem_progs`` are one joint tuning problem, not + three independent knobs. + + Every program reads the same partials with the same tile shape, so they + all get bit-identical ``mean``/``rstd``; program 0 stores them for the + backward. The loop carries nothing across iterations, so the striding + cannot affect the result. + """ + pid = tl.program_id(0) + n = tl.program_id(1) + + offs_g = tl.arange(0, GP) + offs_j = tl.arange(0, CGP) + offs_s = tl.arange(0, BLOCK_S) + inner = offs_g[None, :, None] * CG + offs_j[None, None, :] + off = offs_s[:, None, None] * C + inner + wb = offs_g[:, None] * CG + offs_j[None, :] + wbm = (offs_g[:, None] < G) & (offs_j[None, :] < CG) + + gm = offs_g < G + offs_p = tl.arange(0, NSPLIT) + pidx = (n * NSPLIT + offs_p[:, None]) * G + offs_g[None, :] + pm = tl.broadcast_to(gm[None, :], (NSPLIT, GP)) + # Padded group lanes load cnt == 0, which _welford_combine treats as the + # identity, so they merge to (0, 0, 0) and are masked off on the store. + # Reduced over axis 0 -- the *slowest* axis -- deliberately: reducing + # the fastest axis of a 2-D tile makes Triton stage the whole tile + # through LDS, which for a (G, NSPLIT) tile is 64 KB per array. + cnt_p = tl.load(PCNT + pidx, mask=pm, other=0.0) + _cnt, mu, m2 = tl.reduce( + ( + cnt_p, + tl.load(PMEAN + pidx, mask=pm, other=0.0), + tl.load(PM2 + pidx, mask=pm, other=0.0), + ), + 0, + _welford_combine, + ) + # `_cnt` equals M by construction; M is passed in so the divisor is the + # exact element count rather than a float accumulated from partials. + rs = 1.0 / tl.sqrt(m2 / M + eps) + if pid == 0: + tl.store(MEAN + n * G + offs_g, mu, mask=gm) + tl.store(RSTD + n * G + offs_g, rs, mask=gm) + mean = mu[None, :, None] + rstd = rs[None, :, None] + if HAS_W: + w = tl.load(W + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :] + else: + w = tl.full((1, GP, CGP), 1.0, tl.float32) + if HAS_B: + b = tl.load(B + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :] + else: + b = tl.zeros((1, GP, CGP), dtype=tl.float32) + + for blk in tl.range(pid, NBLK, NPROG): + s0 = blk * BLOCK_S + m = tl.broadcast_to((offs_s < S - s0)[:, None, None], (BLOCK_S, GP, CGP)) + if MASKED_C: + m = m & ((offs_g[None, :, None] < G) & (offs_j[None, None, :] < CG)) + if INT64: + base = (n.to(tl.int64) * S + s0) * C + else: + base = (n * S + s0) * C + x = tl.load(X + base + off, mask=m, other=0.0).to(tl.float32) + xhat = (x - mean) * rstd + y = xhat * w + b + if RELU: + # `tl.maximum(y, 0.0)` and `tl.where(y > 0, y, 0.0)` both map NaN + # to 0.0 (the first returns the non-NaN operand, the second + # because `NaN > 0` is False), while `F.relu` propagates it. + # Testing the *complement* keeps NaN on the pass-through side: + # `NaN <= 0` is also False, so NaN falls to `y`. Bit-identical + # to `F.relu` on NaN, +-Inf and -0.0 (which both flush to +0.0), + # for one comparison and one select -- see the module docstring. + y = tl.where(y <= 0.0, 0.0, y) + tl.store(Y + base + off, y.to(Y.dtype.element_ty), mask=m) + + # ------------------------------------------------------------- backward -- + @_triton.jit + def _bwd_partial_kernel( + X, + DY, + MEAN, + RSTD, + W, + B, + PS1, + PS2, + PDW, + PDB, + S, + CHUNK, + C: tl.constexpr, + G: tl.constexpr, + CG: tl.constexpr, + GP: tl.constexpr, + CGP: tl.constexpr, + NSPLIT: tl.constexpr, + BLOCK_S: tl.constexpr, + RELU: tl.constexpr, + HAS_W: tl.constexpr, + HAS_B: tl.constexpr, + MASKED_C: tl.constexpr, + INT64: tl.constexpr, + ): + """Partials for the two per-``(n, g)`` reductions used by dx, and for + the per-channel dweight / dbias reductions.""" + sp = tl.program_id(0) + n = tl.program_id(1) + + offs_g = tl.arange(0, GP) + offs_j = tl.arange(0, CGP) + offs_s = tl.arange(0, BLOCK_S) + inner = offs_g[None, :, None] * CG + offs_j[None, None, :] + cmask = (offs_g[None, :, None] < G) & (offs_j[None, None, :] < CG) + off = offs_s[:, None, None] * C + inner + wb = offs_g[:, None] * CG + offs_j[None, :] + wbm = (offs_g[:, None] < G) & (offs_j[None, :] < CG) + + s_begin = sp * CHUNK + s_end = tl.minimum(s_begin + CHUNK, S) + + gm = offs_g < G + mean = tl.load(MEAN + n * G + offs_g, mask=gm, other=0.0)[None, :, None] + rstd = tl.load(RSTD + n * G + offs_g, mask=gm, other=0.0)[None, :, None] + if HAS_W: + w = tl.load(W + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :] + else: + w = tl.full((1, GP, CGP), 1.0, tl.float32) + if HAS_B: + b = tl.load(B + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :] + else: + b = tl.zeros((1, GP, CGP), dtype=tl.float32) + + acc1 = tl.zeros((GP,), dtype=tl.float32) + acc2 = tl.zeros((GP,), dtype=tl.float32) + accdw = tl.zeros((GP, CGP), dtype=tl.float32) + accdb = tl.zeros((GP, CGP), dtype=tl.float32) + + for s0 in range(s_begin, s_end, BLOCK_S): + if INT64: + base = (n.to(tl.int64) * S + s0) * C + else: + base = (n * S + s0) * C + nvalid = tl.minimum(BLOCK_S, s_end - s0) + m = tl.broadcast_to((offs_s < nvalid)[:, None, None], (BLOCK_S, GP, CGP)) + if MASKED_C: + m = m & cmask + x = tl.load(X + base + off, mask=m, other=0.0).to(tl.float32) + dy = tl.load(DY + base + off, mask=m, other=0.0).to(tl.float32) + xhat = (x - mean) * rstd + if RELU: + # Identical expression (and therefore identical rounding) to + # the forward's pre-activation, so the sign test agrees with + # the forward bit for bit. Masked lanes carry dy == 0, so + # gating cannot resurrect them. Spelled as the *complement* + # (`pre <= 0` zeroes) rather than `pre > 0` passes, so that a + # NaN pre-activation passes the gradient through: that is what + # `threshold_backward(grad, result, 0)` -- ReLU's real backward + # -- does, since `NaN <= 0` is False. See the forward store. + dy = tl.where(xhat * w + b <= 0.0, 0.0, dy) + dyw = dy * w + acc1 += tl.sum(tl.sum(dyw, 2), 0) + acc2 += tl.sum(tl.sum(dyw * xhat, 2), 0) + accdw += tl.sum(dy * xhat, 0) + accdb += tl.sum(dy, 0) + + o = (n * NSPLIT + sp) * G + offs_g + tl.store(PS1 + o, acc1, mask=gm) + tl.store(PS2 + o, acc2, mask=gm) + row = (n * NSPLIT + sp) * C + wb + tl.store(PDW + row, accdw, mask=wbm) + tl.store(PDB + row, accdb, mask=wbm) + + @_triton.jit + def _dx_kernel( + X, + DY, + DX, + MEAN, + RSTD, + W, + B, + PS1, + PS2, + PDW, + PDB, + DW, + DB, + ROWS, + S, + M, + C: tl.constexpr, + G: tl.constexpr, + CG: tl.constexpr, + GP: tl.constexpr, + CGP: tl.constexpr, + NSPLIT: tl.constexpr, + BLOCK_S: tl.constexpr, + NBLK: tl.constexpr, + NPROG: tl.constexpr, + NDW: tl.constexpr, + BLOCK_C: tl.constexpr, + BLOCK_R: tl.constexpr, + RELU: tl.constexpr, + HAS_W: tl.constexpr, + HAS_B: tl.constexpr, + MASKED_C: tl.constexpr, + INT64: tl.constexpr, + ZERO_DX: tl.constexpr, + ): + """The whole backward tail: dweight/dbias, the c1/c2 finalize, and dx. + + Two reductions that used to be their own launches ride along here. The + per-channel dweight/dbias row reduction is done by the first ``NDW`` + programs of ``n == 0`` (a single pass over an ``(n*nsplit, C)`` scratch, + i.e. a few hundred KB); the per-``(n, g)`` c1/c2 finalize is recomputed + redundantly by every program, once, before the tile loop -- exactly as + in ``_normalize_kernel``, and capped the same way. The grid is + ``(max(NPROG, NDW), n)``; programs past ``NPROG`` exist only to cover + the dweight/dbias rows and run no loop iterations. + """ + pid = tl.program_id(0) + n = tl.program_id(1) + + # ---- dweight / dbias: rows of the split-K scratch, once per channel -- + if n == 0: + if pid < NDW: + offs_c = pid * BLOCK_C + tl.arange(0, BLOCK_C) + mc = offs_c < C + accw = tl.zeros((BLOCK_C,), dtype=tl.float32) + accb = tl.zeros((BLOCK_C,), dtype=tl.float32) + for r0 in range(0, ROWS, BLOCK_R): + offs_r = r0 + tl.arange(0, BLOCK_R) + rm = (offs_r[:, None] < ROWS) & mc[None, :] + roff = offs_r[:, None] * C + offs_c[None, :] + accw += tl.sum(tl.load(PDW + roff, mask=rm, other=0.0), 0) + accb += tl.sum(tl.load(PDB + roff, mask=rm, other=0.0), 0) + tl.store(DW + offs_c, accw, mask=mc) + tl.store(DB + offs_c, accb, mask=mc) + + offs_g = tl.arange(0, GP) + offs_j = tl.arange(0, CGP) + offs_s = tl.arange(0, BLOCK_S) + inner = offs_g[None, :, None] * CG + offs_j[None, None, :] + off = offs_s[:, None, None] * C + inner + wb = offs_g[:, None] * CG + offs_j[None, :] + wbm = (offs_g[:, None] < G) & (offs_j[None, :] < CG) + + if ZERO_DX: + # One element per group: mean == x and var == 0 identically, so + # xhat is the constant 0 and y does not depend on x at all -- the + # exact d_input is zero everywhere. The expression below would + # instead return rstd * (dy*w - c1), and since the compiler + # contracts that to fma(dy, w, -c1) while c1 was accumulated from + # the *rounded* product, what survives is the product's rounding + # error amplified by rstd = 1/sqrt(eps) ~ 316 (2.2e-05 at + # eps=1e-5). Answering with the exact zero costs one constexpr. + zero = tl.zeros((BLOCK_S, GP, CGP), dtype=tl.float32) + for blk in tl.range(pid, NBLK, NPROG): + s0 = blk * BLOCK_S + m = tl.broadcast_to( + (offs_s < S - s0)[:, None, None], (BLOCK_S, GP, CGP) + ) + if MASKED_C: + m = m & ((offs_g[None, :, None] < G) & (offs_j[None, None, :] < CG)) + if INT64: + base = (n.to(tl.int64) * S + s0) * C + else: + base = (n * S + s0) * C + tl.store(DX + base + off, zero.to(DX.dtype.element_ty), mask=m) + else: + gm = offs_g < G + offs_p = tl.arange(0, NSPLIT) + pidx = (n * NSPLIT + offs_p[:, None]) * G + offs_g[None, :] + pm = tl.broadcast_to(gm[None, :], (NSPLIT, GP)) + c1 = (tl.sum(tl.load(PS1 + pidx, mask=pm, other=0.0), 0) / M)[None, :, None] + c2 = (tl.sum(tl.load(PS2 + pidx, mask=pm, other=0.0), 0) / M)[None, :, None] + + mean = tl.load(MEAN + n * G + offs_g, mask=gm, other=0.0)[None, :, None] + rstd = tl.load(RSTD + n * G + offs_g, mask=gm, other=0.0)[None, :, None] + if HAS_W: + w = tl.load(W + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :] + else: + w = tl.full((1, GP, CGP), 1.0, tl.float32) + if HAS_B: + b = tl.load(B + wb, mask=wbm, other=0.0).to(tl.float32)[None, :, :] + else: + b = tl.zeros((1, GP, CGP), dtype=tl.float32) + + for blk in tl.range(pid, NBLK, NPROG): + s0 = blk * BLOCK_S + m = tl.broadcast_to( + (offs_s < S - s0)[:, None, None], (BLOCK_S, GP, CGP) + ) + if MASKED_C: + m = m & ((offs_g[None, :, None] < G) & (offs_j[None, None, :] < CG)) + if INT64: + base = (n.to(tl.int64) * S + s0) * C + else: + base = (n * S + s0) * C + x = tl.load(X + base + off, mask=m, other=0.0).to(tl.float32) + dy = tl.load(DY + base + off, mask=m, other=0.0).to(tl.float32) + xhat = (x - mean) * rstd + if RELU: + # Same complement spelling as _bwd_partial_kernel: a NaN + # pre-activation must pass the gradient, exactly as + # `threshold_backward(grad, result, 0)` does. + dy = tl.where(xhat * w + b <= 0.0, 0.0, dy) + dyw = dy * w + dx = rstd * (dyw - c1 - xhat * c2) + tl.store(DX + base + off, dx.to(DX.dtype.element_ty), mask=m) + + globals().update( + _welford_combine=_welford_combine, + _stats_partial_kernel=_stats_partial_kernel, + _normalize_kernel=_normalize_kernel, + _bwd_partial_kernel=_bwd_partial_kernel, + _dx_kernel=_dx_kernel, + ) + + +def _ensure_kernels(): + if _stats_partial_kernel is None: + _build_kernels() + + +# --------------------------------------------------------------------------- # +# python drivers +# --------------------------------------------------------------------------- # +_CL_FORMAT = torch.channels_last_3d + +#: Reused so the common (already-current device) path allocates nothing. +_NO_GUARD = contextlib.nullcontext() + + +def _device_guard(device: torch.device): + """Make ``device`` current for the kernel launches inside the ``with``. + + A Triton launch goes to whatever device is *current*, not to the device the + argument tensors live on, so without this a tensor on ``cuda:1`` while + ``cuda:0`` is current makes the kernel dereference another device's pointers + and the process dies with ``Memory access fault by GPU node-N``. ATen ops + (including ``F.group_norm``) carry a ``DeviceGuard`` and handle the same + call, so this is required for the drop-in contract, not a nicety. + + The ``current_device()`` test is not about correctness but about *cost*. + Measured on this node (median of 200k calls, torch 2.13.0+rocm7.2): + ``with torch.cuda.device(t.device)`` is **1.55 us** of host time per call, + ``with torch.cuda._DeviceGuard(t.device.index)`` **0.61 us**, and this + helper **0.51 us** when the tensor is already on the current device -- + which it is on every ScaFFold call, since ScaFFold pins one device per + rank. Two of those (forward + backward) against the 0.65 ms fwd+bwd of the + two smallest scale-8 shapes, which are host-dispatch bound, is 0.16% + instead of 0.48%. + """ + if device.index == torch.cuda.current_device(): + return _NO_GUARD + return torch.cuda.device(device) + + +def _shape_of(input: torch.Tensor): + n, channels = input.shape[0], input.shape[1] + spatial = 1 + for d in input.shape[2:]: + spatial *= d + return n, channels, spatial + + +def _tag_kernel_failures(fn): + """Re-raise anything ``fn`` raises as :class:`TritonKernelError`. + + Applied to the two functions that do nothing but import Triton, allocate + scratch and launch kernels. The region is *closed*: it runs no + autograd-observable op, so a blanket ``except Exception`` here cannot + swallow framework control flow the way one at the call site would -- there + is no pack hook, no recompute stop and no functorch layer inside it. That + closure is what lets the caller's fallback ladder use a one-element + allowlist instead of an ever-growing denylist. + + ``torch.OutOfMemoryError`` is passed through untagged; see + :class:`TritonKernelError`. + """ + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + try: + return fn(*args, **kwargs) + except torch.OutOfMemoryError: + raise + except TritonKernelError: + raise + except Exception as e: + raise TritonKernelError( + f"{fn.__name__} failed ({type(e).__name__}: {e})" + ) from e + + return wrapper + + +@_tag_kernel_failures +def _forward(input, num_groups, weight, bias, eps, activation, out_dtype): + _ensure_kernels() + n, channels, spatial = _shape_of(input) + plan = _dispatch(n, channels, spatial, num_groups, input.numel()) + groups = num_groups + device = input.device + + with _device_guard(device): + pcnt = torch.empty(n * plan.nsplit * groups, device=device, dtype=torch.float32) + pmean = torch.empty_like(pcnt) + pm2 = torch.empty_like(pcnt) + mean = torch.empty((n, groups), device=device, dtype=torch.float32) + rstd = torch.empty_like(mean) + out = torch.empty_like(input, dtype=out_dtype, memory_format=_CL_FORMAT) + + _stats_partial_kernel[(plan.nsplit, n)]( + input, + pcnt, + pmean, + pm2, + spatial, + plan.chunk, + C=channels, + G=groups, + CG=plan.group_channels, + GP=plan.groups_p2, + CGP=plan.group_channels_p2, + NSPLIT=plan.nsplit, + BLOCK_S=plan.block_s_stats, + MASKED_C=plan.masked_c, + INT64=plan.int64, + num_warps=plan.cfg.stats_warps, + ) + _normalize_kernel[(plan.nprog_elem, n)]( + input, + out, + pcnt, + pmean, + pm2, + mean, + rstd, + weight, + bias, + spatial, + plan.elements_per_group, + eps, + C=channels, + G=groups, + CG=plan.group_channels, + GP=plan.groups_p2, + CGP=plan.group_channels_p2, + NSPLIT=plan.nsplit, + BLOCK_S=plan.block_s_elem, + NBLK=plan.nblk_elem, + NPROG=plan.nprog_elem, + RELU=activation == "relu", + HAS_W=weight is not None, + HAS_B=bias is not None, + MASKED_C=plan.masked_c, + INT64=plan.int64, + num_warps=plan.cfg.elem_warps, + ) + return out, mean, rstd + + +@_tag_kernel_failures +def _backward(grad_out, input, weight, bias, mean, rstd, num_groups, activation): + _ensure_kernels() + n, channels, spatial = _shape_of(input) + plan = _dispatch(n, channels, spatial, num_groups, input.numel()) + groups = num_groups + device = input.device + + with _device_guard(device): + ps1 = torch.empty(n * plan.nsplit * groups, device=device, dtype=torch.float32) + ps2 = torch.empty_like(ps1) + pdw = torch.empty( + n * plan.nsplit * channels, device=device, dtype=torch.float32 + ) + pdb = torch.empty_like(pdw) + + _bwd_partial_kernel[(plan.nsplit, n)]( + input, + grad_out, + mean, + rstd, + weight, + bias, + ps1, + ps2, + pdw, + pdb, + spatial, + plan.chunk, + C=channels, + G=groups, + CG=plan.group_channels, + GP=plan.groups_p2, + CGP=plan.group_channels_p2, + NSPLIT=plan.nsplit, + BLOCK_S=plan.block_s_stats, + RELU=activation == "relu", + HAS_W=weight is not None, + HAS_B=bias is not None, + MASKED_C=plan.masked_c, + INT64=plan.int64, + num_warps=plan.cfg.stats_warps, + ) + + d_weight = torch.empty(channels, device=device, dtype=torch.float32) + d_bias = torch.empty_like(d_weight) + d_input = torch.empty_like(input, memory_format=_CL_FORMAT) + _dx_kernel[(plan.grid_dx, n)]( + input, + grad_out, + d_input, + mean, + rstd, + weight, + bias, + ps1, + ps2, + pdw, + pdb, + d_weight, + d_bias, + plan.dwdb_rows, + spatial, + plan.elements_per_group, + C=channels, + G=groups, + CG=plan.group_channels, + GP=plan.groups_p2, + CGP=plan.group_channels_p2, + NSPLIT=plan.nsplit, + BLOCK_S=plan.block_s_elem, + NBLK=plan.nblk_elem, + NPROG=plan.nprog_elem, + NDW=plan.dwdb_progs, + BLOCK_C=plan.dwdb_block_c, + BLOCK_R=plan.dwdb_block_r, + RELU=activation == "relu", + HAS_W=weight is not None, + HAS_B=bias is not None, + MASKED_C=plan.masked_c, + INT64=plan.int64, + ZERO_DX=plan.zero_dx, + num_warps=plan.cfg.elem_warps, + ) + return d_input, d_weight, d_bias + + +# --------------------------------------------------------------------------- # +# torch.library registration +# --------------------------------------------------------------------------- # +def _one_value_per_channel(input, num_groups: int) -> bool: + """Whether ``F.group_norm`` would reject this shape as degenerate. + + ``F.group_norm`` runs ``_verify_batch_size([N*C//G, G, *spatial])``, which + raises ``ValueError("Expected more than 1 value per channel when + training")`` exactly when ``N * (C/G) * D*H*W == 1``. All three factors are + positive, so that holds iff ``N == 1``, ``C == num_groups`` and the spatial + extent is 1 -- i.e. iff ``numel == C == num_groups``, which is the cheap + form used here (``numel`` is wanted by the caller anyway). + + Rejected rather than served: the kernel *can* compute it (it returns + ``bias``, since every group has zero variance), but a caller that branches + on :func:`is_supported` would then get a result where the op this replaces + raises, which is a worse failure than being slower. + """ + channels = input.shape[1] + return channels == num_groups and input.numel() == channels + + +def _validate(input, num_groups, weight, bias, activation): + if activation not in SUPPORTED_ACTIVATIONS: + raise ValueError( + f"activation must be one of {SUPPORTED_ACTIVATIONS}, got {activation!r}" + ) + if input.dim() != 5: + raise ValueError(f"expected a 5-D NCDHW tensor, got {tuple(input.shape)}") + if num_groups <= 0 or input.shape[1] % num_groups != 0: + raise ValueError( + f"num_channels={input.shape[1]} is not divisible by num_groups={num_groups}" + ) + if _one_value_per_channel(input, num_groups): + # Same rejection, and the same exception type, as F.group_norm's + # _verify_batch_size; see _one_value_per_channel. + raise ValueError( + f"Expected more than 1 value per channel when training, got input " + f"size {tuple(input.shape)} with num_groups={num_groups}" + ) + if input.dtype not in SUPPORTED_DTYPES: + raise ValueError(f"unsupported input dtype {input.dtype}") + if not input.is_contiguous(memory_format=_CL_FORMAT): + # Required, not converted: the fake kernel promises the *input's* + # memory format for the output, so silently converting here would make + # the traced and eager results disagree on strides. The public + # ``triton_group_norm`` routes non-channels-last input to + # ``F.group_norm`` before it ever reaches this op. + raise ValueError( + "input must be channels_last_3d-contiguous; use triton_group_norm() " + "which falls back to F.group_norm for other layouts" + ) + for name, t in (("weight", weight), ("bias", bias)): + if t is not None and t.numel() != input.shape[1]: + raise ValueError( + f"{name} has {t.numel()} elements, expected {input.shape[1]}" + ) + + +@torch.library.custom_op( + "scaffold_gn::group_norm", mutates_args=(), device_types="cuda" +) +def _group_norm_op( + input: torch.Tensor, + num_groups: int, + weight: Optional[torch.Tensor], + bias: Optional[torch.Tensor], + eps: float, + activation: Optional[str], + out_dtype: Optional[torch.dtype], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Channels-last GroupNorm forward: returns ``(output, mean, rstd)``. + + ``mean``/``rstd`` are ``(N, num_groups)`` fp32 tensors kept for the + backward; they are marked non-differentiable in ``_setup_context`` + (nothing produces a gradient for them), so they come back with + ``requires_grad=False`` and differentiating through them raises rather than + returning zeros. Callers should treat them as opaque. + """ + _validate(input, num_groups, weight, bias, activation) + weight = None if weight is None else weight.contiguous() + bias = None if bias is None else bias.contiguous() + out, mean, rstd = _forward( + input, num_groups, weight, bias, eps, activation, out_dtype or input.dtype + ) + return out, mean, rstd + + +@_group_norm_op.register_fake +def _(input, num_groups, weight, bias, eps, activation, out_dtype): + # empty_like preserves the input's memory format, which is the contract. + out = torch.empty_like(input, dtype=out_dtype or input.dtype) + mean = input.new_empty((input.shape[0], num_groups), dtype=torch.float32) + rstd = input.new_empty((input.shape[0], num_groups), dtype=torch.float32) + return out, mean, rstd + + +@torch.library.custom_op( + "scaffold_gn::group_norm_backward", mutates_args=(), device_types="cuda" +) +def _group_norm_backward_op( + grad_out: torch.Tensor, + input: torch.Tensor, + weight: Optional[torch.Tensor], + bias: Optional[torch.Tensor], + mean: torch.Tensor, + rstd: torch.Tensor, + num_groups: int, + activation: Optional[str], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Returns ``(d_input, d_weight, d_bias)``. + + ``d_weight``/``d_bias`` are zero-element tensors when the corresponding + parameter is ``None``. ``d_input`` always has the input's dtype and + channels-last memory format. + """ + if not grad_out.is_contiguous(memory_format=_CL_FORMAT): + grad_out = grad_out.contiguous(memory_format=_CL_FORMAT) + if not input.is_contiguous(memory_format=_CL_FORMAT): + input = input.contiguous(memory_format=_CL_FORMAT) + weight = None if weight is None else weight.contiguous() + bias = None if bias is None else bias.contiguous() + d_input, d_weight, d_bias = _backward( + grad_out, input, weight, bias, mean, rstd, num_groups, activation + ) + # No `d_input.to(input.dtype)`: `_backward` allocates it with + # `empty_like(input)` and `_dx_kernel` stores through `DX.dtype.element_ty`, + # so it already *is* the input's dtype. + if weight is None: + d_weight = d_weight.new_empty(0) + else: + d_weight = d_weight.to(weight.dtype) + if bias is None: + d_bias = d_bias.new_empty(0) + else: + d_bias = d_bias.to(bias.dtype) + return d_input, d_weight, d_bias + + +@_group_norm_backward_op.register_fake +def _(grad_out, input, weight, bias, mean, rstd, num_groups, activation): + channels = input.shape[1] + # channels_last_3d, *not* the input's own format: the real op relayouts a + # non-channels-last `input` and always returns a channels-last `d_input`, + # so promising `empty_like(input)` here would hand torch.compile the wrong + # strides for any contiguous NCDHW input -- silently, since eager never + # consults the fake kernel. + d_input = torch.empty_like(input, memory_format=_CL_FORMAT) + d_weight = input.new_empty( + channels if weight is not None else 0, + dtype=weight.dtype if weight is not None else torch.float32, + ) + d_bias = input.new_empty( + channels if bias is not None else 0, + dtype=bias.dtype if bias is not None else torch.float32, + ) + return d_input, d_weight, d_bias + + +def _setup_context(ctx, inputs, output): + input, num_groups, weight, bias, eps, activation, out_dtype = inputs + _out, mean, rstd = output + # Outputs 1 and 2 are backward state, not results: nothing produces a + # gradient for them. Without this they come back requiring grad, and + # differentiating through them *succeeds* -- autograd materializes an + # all-zero cotangent for the unused `out` and runs the whole backward to + # return zeros, which is a plausible wrong answer rather than an error. + ctx.mark_non_differentiable(mean, rstd) + ctx.save_for_backward(input, weight, bias, mean, rstd) + ctx.num_groups = num_groups + ctx.activation = activation + ctx.needs = ( + ctx.needs_input_grad[0], + ctx.needs_input_grad[2], + ctx.needs_input_grad[3], + ) + + +def _autograd_backward(ctx, grad_out, grad_mean, grad_rstd): + input, weight, bias, mean, rstd = ctx.saved_tensors + need_x, need_w, need_b = ctx.needs + if not (need_x or need_w or need_b): + return None, None, None, None, None, None, None + d_input, d_weight, d_bias = torch.ops.scaffold_gn.group_norm_backward( + grad_out, input, weight, bias, mean, rstd, ctx.num_groups, ctx.activation + ) + return ( + d_input if need_x else None, + None, # num_groups + d_weight if need_w else None, + d_bias if need_b else None, + None, # eps + None, # activation + None, # out_dtype + ) + + +torch.library.register_autograd( + "scaffold_gn::group_norm", _autograd_backward, setup_context=_setup_context +) + + +# --------------------------------------------------------------------------- # +# public API +# --------------------------------------------------------------------------- # +def _autocast_active(input: torch.Tensor) -> bool: + """Whether autocast is enabled for this tensor's device type.""" + try: + return bool(torch.is_autocast_enabled(input.device.type)) + except (RuntimeError, TypeError): # device type autocast does not know + return False + + +def _autocast_out_dtype(input: torch.Tensor) -> Optional[torch.dtype]: + """``F.group_norm``'s output dtype for this input, or None for "unchanged". + + ``at::group_norm`` carries autocast's ``fp32`` cast policy, so under an + enabled autocast region it upcasts its input and returns fp32 whatever came + in. Verified empirically on torch 2.13.0+rocm7.2 for fp32/bf16/fp16 input + and both bf16 and fp16 autocast dtypes. + """ + if input.dtype is not torch.float32 and _autocast_active(input): + return torch.float32 + return None + + +def is_supported( + input, + num_groups: int, + weight=None, + bias=None, + activation: Optional[str] = None, +) -> bool: + """Whether the native channels-last Triton kernel can serve this call. + + Cheap (a handful of attribute reads and one stride check) and side-effect + free -- in particular it does not import Triton, allocate, or launch. + ``False`` means "use ``F.group_norm``": the fast path needs a 5-D CUDA + tensor that is ``channels_last_3d``-contiguous, an fp32/bf16/fp16 dtype, a + channel count divisible by ``num_groups``, and affine parameters whose + dtype ``F.group_norm`` would itself accept for this input (equal to the + input's, or fp32 under autocast, which is what autocast would produce). + Shapes ``F.group_norm`` itself rejects are rejected here too, so that + branching on this predicate can never turn a stock ``ValueError`` into an + answer (see :func:`_one_value_per_channel`). + + ``True`` promises the *first* derivative only: the backward is itself a + custom op with no autograd formula, so a second ``torch.autograd.grad`` + raises where stock ``F.group_norm`` would succeed. Callers that need a + gradient penalty or a Hessian-vector product must not take this path. + + Note that this is a capability predicate, not a layout classifier: for + shapes whose spatial *and* channel extents make the contiguous and + channels-last-3d stride patterns coincide (e.g. ``(N, C, 1, 1, 1)``), a + plain contiguous tensor is accepted, correctly -- it is the same bytes. + """ + if activation not in SUPPORTED_ACTIVATIONS: + return False + if not isinstance(input, torch.Tensor): + return False + if input.device.type != "cuda" or not triton_available(): + return False + if input.dim() != 5 or input.dtype not in SUPPORTED_DTYPES: + return False + if not isinstance(num_groups, int) or num_groups <= 0: + return False + channels = input.shape[1] + if channels % num_groups != 0 or input.numel() == 0: + return False + if _one_value_per_channel(input, num_groups): + return False + if not input.is_contiguous(memory_format=_CL_FORMAT): + return False + autocast = None + for t in (weight, bias): + if t is None: + continue + if not isinstance(t, torch.Tensor): + return False + if t.dim() != 1 or t.numel() != channels: + return False + if t.device != input.device: + return False + if t.dtype is not input.dtype: + if t.dtype is not torch.float32: + return False + if autocast is None: + autocast = _autocast_active(input) + if not autocast: + # F.group_norm would raise "expected scalar type ..." here; + # reject so the caller reproduces that behaviour exactly. + return False + return True + + +def triton_group_norm( + input, + num_groups: int, + weight=None, + bias=None, + eps: float = 1e-5, + activation: Optional[str] = None, +): + """GroupNorm with an optionally fused activation, channels-last native. + + A drop-in replacement for ``F.group_norm(input, num_groups, weight, bias, + eps)`` (followed by ``F.relu`` when ``activation="relu"``). Inputs that + :func:`is_supported` rejects are served by ``F.group_norm`` itself, which + keeps this function total but means such calls get the *eager* kernel -- + callers with a faster fallback should branch on :func:`is_supported` + themselves. + + The output has the input's memory format and ``F.group_norm``'s dtype; see + the module docstring for the full contract. + """ + if activation not in SUPPORTED_ACTIVATIONS: + raise ValueError( + f"activation must be one of {SUPPORTED_ACTIVATIONS}, got {activation!r}" + ) + if not is_supported(input, num_groups, weight, bias, activation): + out = F.group_norm(input, num_groups, weight, bias, eps) + return F.relu(out) if activation == "relu" else out + out, _mean, _rstd = torch.ops.scaffold_gn.group_norm( + input, + num_groups, + weight, + bias, + float(eps), + activation, + _autocast_out_dtype(input), + ) + return out diff --git a/tests/test_triton_group_norm.py b/tests/test_triton_group_norm.py new file mode 100644 index 0000000..ff4d0ee --- /dev/null +++ b/tests/test_triton_group_norm.py @@ -0,0 +1,868 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""Tests for the channels-last Triton GroupNorm (``ScaFFold.unet.triton_group_norm``). + +The kernel replaces a stock op, so almost every test here is a *parity* test: +values and gradients against ``F.group_norm``, but with the reference computed +in **float64** rather than against another fp32 result -- an fp32-vs-fp32 +comparison cannot tell a correct kernel from one that has merely made the same +mistake, and it cannot see the variance-formula failure the Welford +implementation exists to fix (``test_welford_survives_large_mean``). Each +parity test reports the measured relative error so a regression shows up as a +number, not just a boolean. + +The other three things being pinned down: + +* the **contract** -- output dtype exactly matches ``F.group_norm``'s + (including its fp32 autocast policy), output *memory format* matches the + input's (which is where the kernel deliberately differs from stock, and the + entire reason it exists), and ``is_supported`` accepts exactly the inputs the + native kernel serves; +* **determinism** -- the same call twice is bitwise identical, forward and + backward, because the split count and tiling are pure functions of the shape; +* **composition** -- the op is a real dispatcher op, so it must survive + ``torch.compile(fullgraph=True)`` without a graph break and a ``DCTensor`` + round trip through ``__torch_dispatch__`` with the autograd graph intact. + +CPU runs never touch Triton: the module defers ``import triton`` to the first +call that reaches a kernel, which ``test_import_does_not_pull_in_triton`` +checks in a fresh interpreter. +""" + +import os +import subprocess +import sys + +import pytest +import torch +import torch.nn.functional as F + +from ScaFFold.unet import triton_group_norm as tgn +from ScaFFold.unet.triton_group_norm import is_supported, triton_group_norm + +CL = torch.channels_last_3d +GROUPS = 8 +EPS = 1e-5 + +#: Relative-error ceilings against a float64 reference, by input dtype. The +#: fp32 numbers observed on MI300A at these (small) test shapes are ~2e-07 for +#: y/dx/dweight/dbias; the ceiling leaves room for the 3e-05 that a 134M-element +#: fp32 reduction shows at the largest production shape. The low-precision +#: ceilings are set just above the output's own rounding: 2^-8 for bf16 and +#: 2^-11 for fp16, measured 3.3e-03 and 3.8e-04. +_TOL = { + torch.float32: 1e-4, + torch.bfloat16: 2e-2, + torch.float16: 3e-3, +} + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +def _rel(actual, expected): + """max|actual - expected| / max|expected|, computed in float64.""" + a = actual.detach().double() + e = expected.detach().double() + scale = e.abs().max().clamp_min(1e-30) + return ((a - e).abs().max() / scale).item() + + +def _tensors(shape, dtype, device, affine=True, seed=0, mean=0.0, std=1.0): + """Channels-last input plus (optionally) affine parameters and a cotangent.""" + gen = torch.Generator(device=device).manual_seed(seed) + x = torch.empty(shape, device=device, dtype=dtype, memory_format=CL) + x.normal_(mean, std, generator=gen) + channels = shape[1] + if affine: + weight = torch.empty(channels, device=device, dtype=dtype) + weight.normal_(1.0, 0.25, generator=gen) + bias = torch.empty(channels, device=device, dtype=dtype) + bias.normal_(0.0, 0.25, generator=gen) + else: + weight = bias = None + grad_out = torch.empty(shape, device=device, dtype=dtype, memory_format=CL) + grad_out.normal_(generator=gen) + return x, weight, bias, grad_out + + +def _run(fn, x, weight, bias, grad_out, activation=None, eps=EPS): + """Forward + backward through ``fn``, returning detached results.""" + x = x.detach().clone().requires_grad_(True) + weight = None if weight is None else weight.detach().clone().requires_grad_(True) + bias = None if bias is None else bias.detach().clone().requires_grad_(True) + out = fn(x, GROUPS, weight, bias, eps, activation) + out.backward(grad_out.to(out.dtype)) + return ( + out.detach(), + x.grad, + None if weight is None else weight.grad, + None if bias is None else bias.grad, + ) + + +def _reference(x, weight, bias, grad_out, activation=None, eps=EPS): + """``F.group_norm`` (+ optional ReLU) evaluated entirely in float64.""" + + def fn(x, groups, weight, bias, eps, activation): + out = F.group_norm(x, groups, weight, bias, eps) + return F.relu(out) if activation == "relu" else out + + return _run( + fn, + x.double(), + None if weight is None else weight.double(), + None if bias is None else bias.double(), + grad_out.double(), + activation, + eps, + ) + + +def _assert_parity(got, ref, dtype, label, tol=None): + """Compare (y, dx, dweight, dbias) against the float64 reference.""" + tol = _TOL[dtype] if tol is None else tol + errors = {} + for name, a, e in zip(("y", "dx", "dweight", "dbias"), got, ref): + if a is None: + assert e is None or True # no parameter -> no gradient to compare + continue + errors[name] = _rel(a, e) + print( + f"[{label}] " + + " ".join(f"{k}={v:.3e}" for k, v in errors.items()) + + f" (tol {tol:.1e})" + ) + for name, err in errors.items(): + assert err <= tol, f"{label}: {name} relative error {err:.3e} > {tol:.1e}" + return errors + + +def _cuda_shapes(): + """Shapes covering N>1, non-power-of-two extents and a wide channel count.""" + return [ + (1, 64, 8, 8, 8), # the canonical UNet shape, shrunk + (2, 64, 9, 7, 5), # N>1, all three extents non-power-of-two + (1, 128, 5, 6, 7), + (3, 256, 4, 4, 4), + (1, 2048, 6, 6, 6), # widest UNet channel count + ] + + +# --------------------------------------------------------------------------- +# CPU-only behaviour (no Triton, no GPU) +# --------------------------------------------------------------------------- + + +def test_import_does_not_pull_in_triton(): + """Importing the module must not import Triton. + + Run in a fresh interpreter because any earlier GPU test in this session + would already have built the kernels. The guarantee matters twice over: a + CPU-only unit run must not pay Triton's import, and the module must stay + importable on a build that has no Triton at all. + """ + script = ( + "import sys; import ScaFFold.unet.triton_group_norm as m; " + "assert m.tl is None, 'kernels built at import time'; " + "print('triton' in sys.modules)" + ) + repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + cwd=repo_root, + timeout=300, + ) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "False", ( + f"triton was imported at module import time: {result.stdout!r}" + ) + + +def test_cpu_input_falls_back_bitwise(): + """A CPU tensor is not supported, and the fallback is the stock kernel.""" + gen = torch.Generator().manual_seed(0) + x = torch.randn(1, 64, 4, 4, 4, generator=gen).requires_grad_(True) + weight = torch.randn(64, generator=gen).requires_grad_(True) + bias = torch.randn(64, generator=gen).requires_grad_(True) + assert is_supported(x, GROUPS, weight, bias) is False + + out = triton_group_norm(x, GROUPS, weight, bias, EPS) + assert torch.equal(out, F.group_norm(x, GROUPS, weight, bias, EPS)) + out.pow(2).sum().backward() + assert x.grad is not None and weight.grad is not None and bias.grad is not None + + +def test_cpu_fused_relu_falls_back_bitwise(): + gen = torch.Generator().manual_seed(1) + x = torch.randn(2, 32, 3, 4, 5, generator=gen) + weight = torch.randn(32, generator=gen) + bias = torch.randn(32, generator=gen) + got = triton_group_norm(x, 8, weight, bias, EPS, "relu") + assert torch.equal(got, F.relu(F.group_norm(x, 8, weight, bias, EPS))) + + +def test_unknown_activation_raises(): + x = torch.randn(1, 8, 2, 2, 2) + with pytest.raises(ValueError, match="activation"): + triton_group_norm(x, 2, activation="gelu") + assert is_supported(x, 2, activation="gelu") is False + + +def test_select_strategy_is_a_pure_function_of_shape(): + """The dispatch hook must be deterministic -- the reduction order, and so + the bits of the result, depend on it.""" + for args in ((1, 64, 8**3, 8), (2, 2048, 16**3, 8), (1, 128, 7 * 5 * 3, 4)): + first = tgn.select_strategy(*args) + assert first in tgn.STRATEGIES + assert all(tgn.select_strategy(*args) == first for _ in range(3)) + + +def test_tuning_table_covers_the_scale8_shapes(): + """The frozen table is what makes the kernel reproducible; keep it honest.""" + for channels, edge in ((64, 256), (128, 128), (256, 64), (512, 32), (1024, 16)): + assert tgn.default_config(channels, edge**3) is tgn._TUNED[(channels, edge)] + # An unlisted shape falls back to the generic config rather than failing. + assert tgn.default_config(96, 11**3) == tgn.GNConfig() + + +def test_plan_depends_only_on_shape(): + """Two plans for the same shape must be identical objects of identical + content, or the split count could drift between calls and break bitwise + reproducibility.""" + a = tgn._plan(2, 128, 32**3, 8, 2 * 128 * 32**3) + b = tgn._plan(2, 128, 32**3, 8, 2 * 128 * 32**3) + assert (a.nsplit, a.chunk, a.block_s_stats, a.block_s_elem, a.int64) == ( + b.nsplit, + b.chunk, + b.block_s_stats, + b.block_s_elem, + b.int64, + ) + # int64 addressing turns on exactly when a linear index can overflow int32. + small = tgn._plan(1, 64, 128**3, 8, 64 * 128**3) + big = tgn._plan(2, 64, 256**3, 8, 2 * 64 * 256**3) + assert small.int64 is False + assert big.int64 is True + + +# --------------------------------------------------------------------------- +# is_supported +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +def test_is_supported_accepts_the_fast_path(): + device = torch.device("cuda") + x = torch.empty(1, 64, 6, 6, 6, device=device, memory_format=CL).normal_() + weight = torch.randn(64, device=device) + assert is_supported(x, GROUPS, weight, weight) is True + assert is_supported(x, GROUPS) is True + assert is_supported(x, GROUPS, activation="relu") is True + + +@pytest.mark.gpu +def test_is_supported_rejections(): + """Everything ``is_supported`` rejects must be something a caller can hand + to ``F.group_norm`` instead -- so the rejections are the contract's edge.""" + device = torch.device("cuda") + cl = torch.empty(1, 64, 6, 6, 6, device=device, memory_format=CL).normal_() + cases = { + "cpu tensor": (torch.randn(1, 64, 6, 6, 6), GROUPS, None, None, None), + "contiguous (NCDHW)": ( + torch.randn(1, 64, 6, 6, 6, device=device), + GROUPS, + None, + None, + None, + ), + "float64": ( + torch.empty( + 1, 64, 6, 6, 6, device=device, dtype=torch.float64, memory_format=CL + ), + GROUPS, + None, + None, + None, + ), + "4-D": (torch.randn(1, 64, 6, 6, device=device), GROUPS, None, None, None), + "channels not divisible": (cl, 7, None, None, None), + "num_groups=0": (cl, 0, None, None, None), + "bad activation": (cl, GROUPS, None, None, "gelu"), + "weight wrong size": (cl, GROUPS, torch.randn(32, device=device), None, None), + "weight on cpu": (cl, GROUPS, torch.randn(64), None, None), + "weight dtype mismatch": ( + cl, + GROUPS, + torch.randn(64, device=device, dtype=torch.bfloat16), + None, + None, + ), + "sliced (non-contiguous)": ( + torch.empty(1, 64, 6, 6, 12, device=device, memory_format=CL)[..., ::2], + GROUPS, + None, + None, + None, + ), + "not a tensor": (None, GROUPS, None, None, None), + } + for label, args in cases.items(): + assert is_supported(*args) is False, f"{label} should be rejected" + + +@pytest.mark.gpu +def test_rejected_inputs_still_produce_stock_results(): + """``triton_group_norm`` stays total: rejects go to ``F.group_norm``.""" + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(4) + x = torch.randn(2, 64, 5, 6, 7, device=device, generator=gen).requires_grad_(True) + weight = torch.randn(64, device=device, generator=gen).requires_grad_(True) + bias = torch.randn(64, device=device, generator=gen).requires_grad_(True) + assert is_supported(x, GROUPS, weight, bias) is False + got = triton_group_norm(x, GROUPS, weight, bias, EPS) + assert torch.equal(got, F.group_norm(x, GROUPS, weight, bias, EPS)) + + +# --------------------------------------------------------------------------- +# value / gradient parity against float64 +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.parametrize("shape", _cuda_shapes()) +def test_parity_fp32(shape): + device = torch.device("cuda") + x, weight, bias, grad_out = _tensors( + shape, torch.float32, device, seed=hash(shape) % 1000 + ) + assert is_supported(x, GROUPS, weight, bias) + got = _run(triton_group_norm, x, weight, bias, grad_out) + ref = _reference(x, weight, bias, grad_out) + _assert_parity(got, ref, torch.float32, f"fp32 {shape}") + assert got[0].is_contiguous(memory_format=CL) + assert got[1].is_contiguous(memory_format=CL) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("activation", [None, "relu"]) +def test_parity_every_dtype_and_activation(dtype, activation): + device = torch.device("cuda") + shape = (2, 128, 7, 6, 5) + x, weight, bias, grad_out = _tensors(shape, dtype, device, seed=11) + assert is_supported(x, GROUPS, weight, bias, activation) + got = _run(triton_group_norm, x, weight, bias, grad_out, activation) + ref = _reference(x, weight, bias, grad_out, activation) + _assert_parity(got, ref, dtype, f"{dtype} act={activation}") + assert got[0].dtype == dtype + assert got[1].dtype == dtype + assert got[2].dtype == dtype and got[3].dtype == dtype + + +@pytest.mark.gpu +@pytest.mark.parametrize("affine", ["both", "weight_only", "bias_only", "neither"]) +def test_parity_without_affine_parameters(affine): + """``weight=None`` / ``bias=None`` are separate kernel constexpr paths.""" + device = torch.device("cuda") + shape = (2, 64, 5, 5, 5) + x, weight, bias, grad_out = _tensors(shape, torch.float32, device, seed=19) + if affine in ("bias_only", "neither"): + weight = None + if affine in ("weight_only", "neither"): + bias = None + assert is_supported(x, GROUPS, weight, bias) + got = _run(triton_group_norm, x, weight, bias, grad_out) + + reference_weight = weight + if weight is None and bias is not None: + # Upstream limitation, not a difference in this kernel: + # ``F.group_norm(x, g, None, bias).backward()`` raises "tensor does not + # have a device" on both CPU and CUDA (torch 2.13.0+rocm7.2), so the + # float64 reference has to spell the same computation with weight=1. + # This kernel handles the combination directly. + reference_weight = torch.ones_like(bias) + ref = _reference(x, reference_weight, bias, grad_out) + # ``_assert_parity`` skips outputs this configuration does not produce. + _assert_parity(got, ref, torch.float32, f"affine={affine}") + + +@pytest.mark.gpu +def test_partial_gradient_requirements(): + """Only some inputs requiring grad must not change the ones that do.""" + device = torch.device("cuda") + x, weight, bias, grad_out = _tensors( + (1, 64, 5, 5, 5), torch.float32, device, seed=23 + ) + full = _run(triton_group_norm, x, weight, bias, grad_out) + + frozen_w = weight.detach().clone() + frozen_b = bias.detach().clone() + xi = x.detach().clone().requires_grad_(True) + out = triton_group_norm(xi, GROUPS, frozen_w, frozen_b, EPS) + out.backward(grad_out) + assert torch.equal(xi.grad, full[1]) + assert frozen_w.grad is None and frozen_b.grad is None + + # ... and the mirror image: parameters only. + xn = x.detach().clone() + wn = weight.detach().clone().requires_grad_(True) + bn = bias.detach().clone().requires_grad_(True) + triton_group_norm(xn, GROUPS, wn, bn, EPS).backward(grad_out) + assert torch.equal(wn.grad, full[2]) + assert torch.equal(bn.grad, full[3]) + + +# --------------------------------------------------------------------------- +# fused activation +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_fused_relu_matches_unfused(dtype): + """The fused store and ``F.relu`` on the unfused output must agree bitwise, + forward *and* backward -- the backward recomputes the pre-activation rather + than reading it back, so this is the test that recomputation is exact.""" + device = torch.device("cuda") + x, weight, bias, grad_out = _tensors((2, 64, 6, 7, 8), dtype, device, seed=29) + fused = _run(triton_group_norm, x, weight, bias, grad_out, "relu") + + def unfused(x, groups, weight, bias, eps, _activation): + return F.relu(triton_group_norm(x, groups, weight, bias, eps)) + + separate = _run(unfused, x, weight, bias, grad_out) + for name, a, b in zip(("y", "dx", "dweight", "dbias"), fused, separate): + assert torch.equal(a, b), f"fused vs unfused+relu differ in {name}" + # A ReLU that never fires would make this test vacuous. + assert (fused[0] == 0).any() and (fused[0] > 0).any() + + +# --------------------------------------------------------------------------- +# determinism +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.parametrize("activation", [None, "relu"]) +def test_bitwise_determinism(activation): + """Same input twice => bitwise-equal output and gradients. + + Guaranteed by construction (no float atomics; grid, split count and tile + sizes are pure functions of the shape) and asserted here because a future + run-time autotuner would silently break it. + """ + device = torch.device("cuda") + x, weight, bias, grad_out = _tensors( + (2, 128, 9, 11, 13), torch.float32, device, seed=31 + ) + first = _run(triton_group_norm, x, weight, bias, grad_out, activation) + second = _run(triton_group_norm, x, weight, bias, grad_out, activation) + for name, a, b in zip(("y", "dx", "dweight", "dbias"), first, second): + assert torch.equal(a, b), f"{name} is not bitwise reproducible" + + +# --------------------------------------------------------------------------- +# numerics: Welford vs E[x^2] - E[x]^2 +# --------------------------------------------------------------------------- + + +def _naive_group_norm(x, num_groups, weight, bias, eps): + """The prototype's variance formula, reproduced in fp32 torch ops. + + ``var = E[x^2] - E[x]^2`` is split-friendly and cheap, and it is what the + kernel used before the Welford rewrite; this is the thing the test below + must show is broken so that "the new one passes" means something. + """ + n, channels = x.shape[0], x.shape[1] + flat = x.reshape(n, num_groups, -1) + mean = flat.mean(-1) + mean_sq = (flat * flat).mean(-1) + var = mean_sq - mean * mean + rstd = 1.0 / torch.sqrt(var + eps) + out = (flat - mean[..., None]) * rstd[..., None] + out = out.reshape(x.shape) + shape = (1, channels) + (1,) * (x.dim() - 2) + return out * weight.reshape(shape) + bias.reshape(shape) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "mean,std,naive_floor", + [ + (0.0, 1.0, None), # both formulations are fine here + (100.0, 1.0, 1e-4), # E[x^2]-E[x]^2 already an order of magnitude off + (1e3, 1e-2, 1e-1), # ... and here it has lost the variance entirely + ], +) +def test_welford_survives_large_mean(mean, std, naive_floor): + """Large-mean / small-variance input: the regression case for the rewrite. + + Measured here on MI300A at ``[1, 256, 24^3]`` (relative error of the output + against a float64 reference computed from the same fp32 samples), with the + production shape ``[1, 256, 64^3]`` in parentheses: + + mean=0, std=1 this 1.3e-07 (1.6e-07) E[x^2]-E[x]^2 1.4e-07 (1.8e-07) + mean=1e2, std=1 this 1.1e-06 (8.0e-07) E[x^2]-E[x]^2 9.5e-04 (5.6e-04) + mean=1e3, std=1e-2 this 4.4e-04 (1.1e-04) E[x^2]-E[x]^2 2.3e+00 (2.3e+00) + + i.e. at ``mean/std = 1e5`` the old formulation loses the variance outright + (the difference of the two ~1e6-sized fp32 terms is below one ulp, so + ``rstd`` saturates on ``eps`` and the output is meaningless) while Welford + is still correct to 4.4e-04 -- itself *5x better* than ATen's own fp32 + GroupNorm on the same input (2.2e-03), and dominated by the fp32 + representation of a mean of 1e3 rather than by anything the kernel does. + """ + device = torch.device("cuda") + x, weight, bias, grad_out = _tensors( + (1, 256, 24, 24, 24), torch.float32, device, seed=37, mean=mean, std=std + ) + got = triton_group_norm(x, GROUPS, weight, bias, EPS) + ref = F.group_norm(x.double(), GROUPS, weight.double(), bias.double(), EPS) + stock = F.group_norm(x, GROUPS, weight, bias, EPS) + naive = _naive_group_norm(x, GROUPS, weight, bias, EPS) + + err = _rel(got, ref) + err_stock = _rel(stock, ref) + err_naive = _rel(naive, ref) + print( + f"[welford mean={mean:g} std={std:g}] triton={err:.3e} " + f"aten_fp32={err_stock:.3e} naive_Ex2={err_naive:.3e}" + ) + # Never worse than ATen's own fp32 kernel by more than a small factor. + assert err <= max(4.0 * err_stock, 1e-5), ( + f"triton {err:.3e} vs aten fp32 {err_stock:.3e}" + ) + if naive_floor is not None: + assert err_naive > naive_floor, ( + "the naive formulation was expected to fail here " + f"but only reached {err_naive:.3e}" + ) + assert err < err_naive / 10.0, ( + f"Welford ({err:.3e}) is not clearly better than " + f"E[x^2]-E[x]^2 ({err_naive:.3e})" + ) + + +# --------------------------------------------------------------------------- +# dtypes, layouts, autocast +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("autocast_dtype", [None, torch.bfloat16, torch.float16]) +def test_output_dtype_matches_stock(dtype, autocast_dtype): + """The dtype contract, including autocast's fp32 policy for GroupNorm. + + Stock behaviour on this build (measured, not assumed): without autocast the + output dtype is the input dtype; inside *any* enabled CUDA autocast region + it is fp32, because ``at::group_norm`` carries the fp32 cast policy. + """ + device = torch.device("cuda") + x, weight, bias, _ = _tensors((1, 64, 5, 5, 5), dtype, device, seed=41) + if autocast_dtype is not None: + # Autocast casts the parameters itself, and production keeps them fp32. + weight = weight.float() + bias = bias.float() + ctx = ( + torch.autocast("cuda", dtype=autocast_dtype) + if autocast_dtype is not None + else torch.autocast("cuda", enabled=False) + ) + with ctx: + assert is_supported(x, GROUPS, weight, bias) + stock = F.group_norm(x, GROUPS, weight, bias, EPS) + got = triton_group_norm(x, GROUPS, weight, bias, EPS) + expected = torch.float32 if autocast_dtype is not None else dtype + assert stock.dtype == expected, "assumption about stock GroupNorm broke" + assert got.dtype == stock.dtype + # ... and the one deliberate difference: stock always returns contiguous. + assert stock.is_contiguous() and not stock.is_contiguous(memory_format=CL) + assert got.is_contiguous(memory_format=CL) + print( + f"[dtype in={dtype} autocast={autocast_dtype}] " + f"stock={stock.dtype}/CONT mine={got.dtype}/CL rel={_rel(got, stock):.3e}" + ) + assert _rel(got.float(), stock.float()) <= _TOL[stock.dtype] + + +@pytest.mark.gpu +def test_autocast_gradient_dtypes_match_stock(): + """Under autocast, ``d_input`` keeps the input's dtype and the parameter + gradients stay fp32 -- exactly what the cast nodes around stock GroupNorm + produce.""" + device = torch.device("cuda") + x, _, _, grad_out = _tensors((1, 64, 5, 5, 5), torch.bfloat16, device, seed=43) + weight = torch.randn(64, device=device, requires_grad=True) + bias = torch.randn(64, device=device, requires_grad=True) + + def run(fn): + xi = x.detach().clone().requires_grad_(True) + w = weight.detach().clone().requires_grad_(True) + b = bias.detach().clone().requires_grad_(True) + with torch.autocast("cuda", dtype=torch.bfloat16): + out = fn(xi, GROUPS, w, b, EPS) + out.backward(grad_out.to(out.dtype)) + return out, xi.grad, w.grad, b.grad + + stock = run(F.group_norm) + got = run(triton_group_norm) + for name, a, b in zip(("y", "dx", "dweight", "dbias"), got, stock): + assert a.dtype == b.dtype, f"{name}: {a.dtype} != {b.dtype}" + assert _rel(a.float(), b.float()) <= 5e-2, name + + +@pytest.mark.gpu +@pytest.mark.parametrize("layout", ["channels_last_3d", "contiguous"]) +def test_memory_format_is_preserved(layout): + """Both layouts round-trip their own format; contiguous input takes the + documented ``F.group_norm`` fallback rather than silently changing layout.""" + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(47) + x = torch.randn(2, 64, 5, 6, 7, device=device, generator=gen) + if layout == "channels_last_3d": + x = x.contiguous(memory_format=CL) + weight = torch.randn(64, device=device, generator=gen) + bias = torch.randn(64, device=device, generator=gen) + grad_out = torch.randn(2, 64, 5, 6, 7, device=device, generator=gen) + if layout == "channels_last_3d": + grad_out = grad_out.contiguous(memory_format=CL) + + assert is_supported(x, GROUPS, weight, bias) is (layout == "channels_last_3d") + got = _run(triton_group_norm, x, weight, bias, grad_out) + ref = _reference(x, weight, bias, grad_out) + _assert_parity(got, ref, torch.float32, f"layout={layout}") + if layout == "channels_last_3d": + assert got[0].is_contiguous(memory_format=CL) + assert got[1].is_contiguous(memory_format=CL) + else: + assert got[0].is_contiguous() + assert got[1].is_contiguous() + + +# --------------------------------------------------------------------------- +# int64 offsets +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +def test_int64_switch_flips_at_int32_max(): + """The switch is a pure function of the element count, so pin the boundary. + + ``[2, 64, 256^3]`` is *exactly* 2^31 elements: the shape that made an + int64 path mandatory before batch>1 or scale 16. + """ + assert tgn._plan(1, 64, 255**3, 8, 64 * 255**3).int64 is False + assert tgn._plan(2, 64, 256**3, 8, 2 * 64 * 256**3).int64 is True + + +@pytest.mark.gpu +@pytest.mark.slow +def test_correct_above_int32_max_elements(): + """Correctness at a shape whose linear element count exceeds INT32_MAX. + + ``[2, 64, 256, 256, 257]`` is 2_155_872_256 elements -- 8.4M past 2^31, and + non-power-of-two in the fastest spatial dimension so a truncated offset + cannot accidentally land on the right address. fp32 (8.03 GiB per tensor) + keeps the comparison sharp; the reference needs an NCDHW copy, so the peak + is ~48 GiB and the test skips, loudly, if the device cannot hold that. + """ + device = torch.device("cuda") + shape = (2, 64, 256, 256, 257) + numel = 1 + for dim in shape: + numel *= dim + assert numel > 2**31 - 1 + needed = 6 * numel * 4 # x, y, x_contig, reference, and slack for the diff + free, total = torch.cuda.mem_get_info() + if free < needed: + pytest.skip( + f"needs ~{needed / 2**30:.0f} GiB free, device has " + f"{free / 2**30:.0f} GiB of {total / 2**30:.0f} GiB" + ) + + gen = torch.Generator(device=device).manual_seed(53) + x = torch.empty(shape, device=device, memory_format=CL) + x.normal_(generator=gen) + weight = torch.randn(64, device=device, generator=gen) + bias = torch.randn(64, device=device, generator=gen) + + assert tgn._plan( + shape[0], shape[1], shape[2] * shape[3] * shape[4], GROUPS, numel + ).int64 + got = triton_group_norm(x, GROUPS, weight, bias, EPS) + assert got.is_contiguous(memory_format=CL) + + contiguous = x.contiguous() + del x + torch.cuda.empty_cache() + reference = F.group_norm(contiguous, GROUPS, weight, bias, EPS) + del contiguous + torch.cuda.empty_cache() + + # Compare both batch items separately: a truncated 32-bit offset wraps + # partway through, so the second half would be wrong while the first is not. + errors = [_rel(got[i], reference[i]) for i in range(shape[0])] + print(f"[int64 {shape}] per-sample relative error {errors}") + for i, err in enumerate(errors): + assert err < 1e-4, f"sample {i}: relative error {err:.3e}" + del got, reference + torch.cuda.empty_cache() + + +# --------------------------------------------------------------------------- +# composition: torch.compile and DCTensor +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +def test_custom_op_is_registered_with_a_fake_kernel(): + """A meta/fake kernel is what lets Dynamo trace the op without running it.""" + from torch._subclasses.fake_tensor import FakeTensorMode + + assert hasattr(torch.ops.scaffold_gn, "group_norm") + assert hasattr(torch.ops.scaffold_gn, "group_norm_backward") + with FakeTensorMode(): + x = torch.empty(2, 64, 5, 6, 7, device="cuda", memory_format=CL) + weight = torch.empty(64, device="cuda") + out, mean, rstd = torch.ops.scaffold_gn.group_norm( + x, GROUPS, weight, weight, EPS, "relu", None + ) + assert out.shape == x.shape and out.dtype == x.dtype + assert out.is_contiguous(memory_format=CL) + assert mean.shape == (2, GROUPS) and rstd.dtype == torch.float32 + # ... and the dtype override autocast uses. + bf16 = torch.empty( + 2, 64, 5, 6, 7, device="cuda", dtype=torch.bfloat16, memory_format=CL + ) + out32, _, _ = torch.ops.scaffold_gn.group_norm( + bf16, GROUPS, None, None, EPS, None, torch.float32 + ) + assert out32.dtype == torch.float32 + assert out32.is_contiguous(memory_format=CL) + + +@pytest.mark.gpu +@pytest.mark.parametrize("activation", [None, "relu"]) +def test_torch_compile_fullgraph(activation): + """``fullgraph=True`` raises on a graph break, so this *is* the no-break + test; the compiled result must additionally be bitwise equal to eager, + because the op is opaque to Inductor and so cannot be re-associated.""" + device = torch.device("cuda") + x, weight, bias, grad_out = _tensors( + (2, 64, 6, 6, 6), torch.float32, device, seed=59 + ) + + def fn(x, weight, bias): + return triton_group_norm(x, GROUPS, weight, bias, EPS, activation) * 2.0 + + def wrapped(x, groups, weight, bias, eps, _activation): + return fn(x, weight, bias) + + eager = _run(wrapped, x, weight, bias, grad_out) + + compiled_fn = torch.compile(fn, fullgraph=True, dynamic=False) + + def wrapped_compiled(x, groups, weight, bias, eps, _activation): + return compiled_fn(x, weight, bias) + + compiled = _run(wrapped_compiled, x, weight, bias, grad_out) + for name, a, b in zip(("y", "dx", "dweight", "dbias"), compiled, eager): + assert torch.equal(a, b), f"compiled and eager differ in {name}" + assert compiled[0].is_contiguous(memory_format=CL) + + +@pytest.fixture +def dc_cuda(): + """DistConv package plus a CUDA ParallelStrategy over a 1-rank NCCL group. + + Mirrors ``tests/test_groupnorm.py``'s fixture (``num_shards=(1, 1, 1)`` on + dims (2, 3, 4) is what worker.py builds for a single-device run) on its own + rendezvous port so the two suites can run in one session. + """ + import torch.distributed as dist + + distconv = pytest.importorskip("distconv") + + created = False + if not dist.is_initialized(): + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29519") + dist.init_process_group(backend="nccl", rank=0, world_size=1) + created = True + strategy = distconv.ParallelStrategy( + num_shards=(1, 1, 1), shard_dim=(2, 3, 4), device_type="cuda" + ) + yield distconv, strategy + if created and dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.gpu +def test_dctensor_round_trip(dc_cuda): + """A DCTensor must go in and come out, with the graph back to its producer + intact. + + The op is a real dispatcher op, so DistConv's generic + ``__torch_dispatch__`` unwraps to the local shard, runs it, and rewraps -- + no GroupNorm-specific handling needed on either side. The producer in + front matters: with a bare ``input._tensor`` read the graph would be severed + there and only GroupNorm's own parameters would see gradients. + """ + distconv, strategy = dc_cuda + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(61) + x = torch.empty(1, 64, 6, 6, 6, device=device, memory_format=CL) + x.normal_(generator=gen) + grad_out = torch.empty_like(x) + grad_out.normal_(generator=gen) + weight = torch.randn(64, device=device, generator=gen) + bias = torch.randn(64, device=device, generator=gen) + producer = torch.nn.Conv3d(64, 64, 1, bias=False).to(device) + + def run(fn, wrap): + xi = x.detach().clone().requires_grad_(True) + conv = torch.nn.Conv3d(64, 64, 1, bias=False).to(device) + with torch.no_grad(): + conv.weight.copy_(producer.weight) + w = weight.detach().clone().requires_grad_(True) + b = bias.detach().clone().requires_grad_(True) + inp = distconv.DCTensor.from_shard(xi, strategy) if wrap else xi + # The conv is the producer; the explicit channels-last conversion is + # what PYTORCH_MIOPEN_SUGGEST_NHWC=1 gives production for free. + hidden = conv(inp).contiguous(memory_format=CL) + out = fn(hidden, GROUPS, w, b, EPS, "relu") + if wrap: + assert isinstance(out, distconv.DCTensor) + assert out.is_contiguous(memory_format=CL) + out = distconv.distconv._ToTensor.apply(out) + out.backward(grad_out) + return out.detach(), xi.grad, conv.weight.grad, w.grad, b.grad + + def stock(x, groups, weight, bias, eps, _activation): + return F.relu(F.group_norm(x, groups, weight, bias, eps)) + + got = run(triton_group_norm, wrap=True) + ref = run(stock, wrap=False) + for name, a, b in zip(("y", "dx", "dconv", "dweight", "dbias"), got, ref): + assert a is not None, f"{name} never received a gradient" + err = _rel(a, b) + print(f"[dctensor] {name}={err:.3e}") + assert err <= 1e-4, f"{name}: relative error {err:.3e}" diff --git a/tests/test_triton_group_norm_edge.py b/tests/test_triton_group_norm_edge.py new file mode 100644 index 0000000..bca5e88 --- /dev/null +++ b/tests/test_triton_group_norm_edge.py @@ -0,0 +1,1701 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""Adversarial edge-case tests for the channels-last Triton GroupNorm. + +Companion to ``tests/test_triton_group_norm.py``, written independently during +an audit of the kernel. It covers the ground the author's suite does not, and +pins the divergences that audit found. See +``work/gn-dctensor/KERNEL_REVIEW.md`` for the full write-up. + +The two structural gaps this file closes: + +* **The masked channel axis is never exercised upstream.** Every GPU test in + ``test_triton_group_norm.py`` uses ``num_groups=8`` with a channel count of + 64/128/256/2048, so ``G`` and ``C/G`` are *always* powers of two and + ``_Plan.masked_c`` is always ``False``. The entire ``MASKED_C=True`` code + path -- the ``cmask``/``wbm`` predicates in all four kernels, and the + ``inner`` offsets that deliberately run past the end of a voxel -- ships + untested. :func:`test_masked_channel_axis_parity` and friends run it. + +* **Uninitialised split-K scratch is never checked.** ``_forward`` and + ``_backward`` allocate their partial buffers with ``torch.empty``, so a slot + that is read before it is written would surface as *plausible* numbers, not + as a crash. :func:`test_scratch_slots_are_all_written` poisons every + ``torch.empty`` with NaN for the duration of the call, which turns that class + of bug into a hard failure. + +The audit's six findings -- no device guard, the backward fake kernel's stride +promise, silently-differentiable ``mean``/``rstd``, accepting a shape +``F.group_norm`` rejects, a non-zero ``d_input`` for single-element groups, and +undocumented double backward -- were tested here as ``xfail(strict=True)`` +first and fixed afterwards; the tests remain, without the markers, as the +regression pins. The last section adds the coverage a mutation sweep of the +kernels found thinnest: the ``INT64=True`` branch (which a default run never +compiled), the split-K Welford merge on *unequal* split counts, ``eps`` +placement, and the tile-mean correction term. +""" + +import contextlib +import os +import subprocess +import sys +import textwrap + +import pytest +import torch +import torch.nn.functional as F + +from ScaFFold.unet import triton_group_norm as tgn +from ScaFFold.unet.triton_group_norm import is_supported, triton_group_norm + +CL = torch.channels_last_3d +EPS = 1e-5 + +#: Relative-error ceiling against the float64 reference below. fp32 parity at +#: these (small) shapes measures ~1e-07; the ceiling leaves room for the +#: reduction noise a production-sized split-K reduction shows. +FP32_TOL = 1e-4 + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +# --------------------------------------------------------------------------- +# independent float64 reference (deliberately *not* F.group_norm, and +# deliberately not the helper the author's suite uses) +# --------------------------------------------------------------------------- + + +def _ref64(x, groups, weight, bias, eps, activation=None): + """GroupNorm written from scratch in float64, in the (N, G, ...) view.""" + xd = x.double() + n, channels = xd.shape[0], xd.shape[1] + flat = xd.reshape(n, groups, -1) + mu = flat.mean(-1, keepdim=True) + var = ((flat - mu) ** 2).mean(-1, keepdim=True) + y = ((flat - mu) / torch.sqrt(var + eps)).reshape(xd.shape) + shape = (1, channels) + (1,) * (xd.dim() - 2) + if weight is not None: + y = y * weight.double().reshape(shape) + if bias is not None: + y = y + bias.double().reshape(shape) + return torch.relu(y) if activation == "relu" else y + + +def _rel(actual, expected): + a = actual.detach().double() + e = expected.detach().double() + return ((a - e).abs().max() / e.abs().max().clamp_min(1e-300)).item() + + +def _make(shape, groups, dtype=torch.float32, affine=True, seed=0, mean=0.0, std=1.0): + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(seed) + x = torch.empty(shape, device=device, dtype=dtype, memory_format=CL) + x.normal_(mean, std, generator=gen) + channels = shape[1] + if affine: + weight = torch.empty(channels, device=device, dtype=dtype) + weight.normal_(1.0, 0.25, generator=gen) + bias = torch.empty(channels, device=device, dtype=dtype) + bias.normal_(0.0, 0.25, generator=gen) + else: + weight = bias = None + grad_out = torch.empty(shape, device=device, dtype=dtype, memory_format=CL) + grad_out.normal_(generator=gen) + return x, weight, bias, grad_out + + +def _parity(shape, groups, activation=None, affine=True, seed=0, eps=EPS, label=""): + """Forward + backward against the float64 reference. Returns the errors.""" + x, weight, bias, grad_out = _make(shape, groups, affine=affine, seed=seed) + assert is_supported(x, groups, weight, bias, activation), ( + f"is_supported rejected {shape} groups={groups}" + ) + + xi = x.detach().clone().requires_grad_(True) + wi = None if weight is None else weight.detach().clone().requires_grad_(True) + bi = None if bias is None else bias.detach().clone().requires_grad_(True) + triton_group_norm(xi, groups, wi, bi, eps, activation).backward(grad_out) + + xd = x.detach().clone().double().requires_grad_(True) + wd = ( + None + if weight is None + else weight.detach().clone().double().requires_grad_(True) + ) + bd = None if bias is None else bias.detach().clone().double().requires_grad_(True) + _ref64(xd, groups, wd, bd, eps, activation).backward(grad_out.double()) + + errors = {"dx": _rel(xi.grad, xd.grad)} + if wi is not None: + errors["dw"] = _rel(wi.grad, wd.grad) + if bi is not None: + errors["db"] = _rel(bi.grad, bd.grad) + xj = x.detach().clone() + errors["y"] = _rel( + triton_group_norm(xj, groups, weight, bias, eps, activation), + _ref64(x, groups, weight, bias, eps, activation), + ) + print(f"[{label or shape}] " + " ".join(f"{k}={v:.2e}" for k, v in errors.items())) + for name, err in errors.items(): + assert err <= FP32_TOL, f"{label or shape}: {name} rel err {err:.3e}" + return errors + + +# --------------------------------------------------------------------------- +# 1. the masked channel axis (MASKED_C=True) -- never reached upstream +# --------------------------------------------------------------------------- + +#: ``(shape, num_groups)`` pairs for which ``_Plan.masked_c`` is True, i.e. +#: ``num_groups`` and/or ``num_channels // num_groups`` is not a power of two, +#: so ``GP``/``CGP`` over-cover the channel axis and every load, store and +#: reduction in all four kernels has to be predicated. +_MASKED_CASES = [ + ((1, 6, 4, 4, 4), 3), # G=3 -> GP=4, CG=2 + ((1, 15, 5, 5, 5), 3), # G=3, CG=5 -> both padded + ((2, 12, 7, 5, 3), 3), # N>1 with a padded group axis + ((1, 24, 9, 9, 9), 6), # G=6 -> GP=8, CG=4 + ((1, 20, 4, 4, 4), 5), # G=5 -> GP=8, CG=4 + ((1, 20, 4, 4, 4), 4), # G=4, CG=5 -> only the inner axis padded + ((3, 20, 3, 5, 7), 5), + ((1, 63, 5, 5, 5), 7), # G=7, CG=9 + ((2, 63, 5, 5, 5), 7), + ((1, 96, 5, 5, 5), 6), # G=6, CG=16 + ((1, 10, 3, 3, 3), 10), # G == C, neither a power of two +] + + +@pytest.mark.gpu +@pytest.mark.parametrize("shape,groups", _MASKED_CASES) +@pytest.mark.parametrize("activation", [None, "relu"]) +def test_masked_channel_axis_parity(shape, groups, activation): + """``MASKED_C=True``: the padded (G, C/G) tile must be fully predicated. + + ``inner = g * CG + j`` deliberately runs past the end of a voxel for the + padding lanes, so a wrong ``cmask``/``wbm`` predicate reads (or writes) the + *next* voxel's channels, and a wrong ``other=`` poisons the Welford sums. + Neither shows up anywhere in the author's suite, which only ever runs + ``num_groups=8`` over 64/128/256/2048 channels. + """ + plan = tgn._plan(shape[0], shape[1], shape[2] * shape[3] * shape[4], groups, 0) + assert plan.masked_c, "case is supposed to exercise the padded channel axis" + _parity(shape, groups, activation, seed=abs(hash((shape, groups))) % 997) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups", + [ + ((1, 64, 4, 4, 4), 64), # instance norm, C/G == 1 + ((1, 7, 3, 3, 3), 7), # instance norm, prime channel count + ((1, 64, 4, 4, 4), 1), # layer norm, G == 1 + ((1, 7, 3, 3, 3), 1), # layer norm, prime channel count + ((1, 1, 4, 4, 4), 1), # single channel + ((2, 1, 4, 4, 4), 1), + ((1, 3, 5, 5, 5), 3), + ], +) +def test_extreme_group_counts(shape, groups): + """``num_groups == num_channels`` (instance norm) and ``== 1`` (layer norm). + + Both collapse one axis of the ``(BLOCK_S, GP, CGP)`` tile to length 1 and + are accepted by ``is_supported``; neither appears upstream. + """ + _parity(shape, groups, seed=abs(hash((shape, groups))) % 997) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape", + [ + (1, 64, 1, 1, 1), # S == 1: a single voxel, far below one tile + (2, 64, 1, 1, 1), + (4, 64, 2, 1, 1), + (1, 64, 1, 1, 127), # just under a 128-voxel stats tile + (1, 64, 1, 1, 128), # exactly one tile + (1, 64, 1, 1, 129), # just over + (1, 64, 1, 1, 257), + (1, 64, 13, 17, 19), # three primes + (5, 64, 3, 3, 3), # N not a power of two + (1, 2048, 1, 1, 2), # widest channel count, two voxels + ], +) +def test_ragged_spatial_tails(shape): + """Spatial extents that are prime, or sit just either side of a tile edge. + + The ragged tail is where ``offs_s < S - s0`` in ``_normalize_kernel`` / + ``_dx_kernel`` and ``nvalid = min(BLOCK_S, s_end - s0)`` in the two partial + kernels have to agree; ``cnt_t = nvalid * CG`` also has to be the *valid* + lane count or the Welford mean is scaled wrong. + """ + _parity(shape, 8, seed=abs(hash(shape)) % 997) + + +# --------------------------------------------------------------------------- +# 2. split-K scratch +# --------------------------------------------------------------------------- + + +def _empty_split_count(n, channels, spatial, groups): + plan = tgn._plan(n, channels, spatial, groups, n * channels * spatial) + return sum(1 for sp in range(plan.nsplit) if sp * plan.chunk >= spatial), plan + + +#: Shapes whose ``ceil(S / nsplit)`` chunking leaves at least one split with +#: ``s_begin >= S``, i.e. a program that writes an all-zero ``(cnt, mean, M2)`` +#: partial that the finalize tree then has to absorb. Found by search over the +#: plan; ``_welford_combine``'s ``cnt == 0`` guard is what makes them harmless. +_EMPTY_SPLIT_CASES = [ + ((1, 64, 1, 1, 32775), 8), + ((2, 64, 1, 1, 32775), 8), + ((1, 128, 1, 1, 8198), 8), + ((1, 256, 1, 1, 2049), 8), + ((1, 2048, 1, 1, 33), 8), +] + + +@pytest.mark.gpu +@pytest.mark.parametrize("shape,groups", _EMPTY_SPLIT_CASES) +def test_empty_split_slots(shape, groups): + """A split whose whole chunk lies past ``S`` still has to combine cleanly. + + ``chunk = ceil(S / nsplit)`` can leave trailing splits entirely empty; that + program's loop never runs, so it stores ``(0, 0, 0)``. Chan's combine is + only exact for those because of its ``cnt == 0`` guard, and no upstream + shape produces one. + """ + empties, plan = _empty_split_count( + shape[0], shape[1], shape[2] * shape[3] * shape[4], groups + ) + assert empties > 0, ( + f"expected an empty split for {shape}; plan has nsplit={plan.nsplit} " + f"chunk={plan.chunk}" + ) + _parity(shape, groups, seed=abs(hash(shape)) % 997) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups", + [ + ((1, 64, 8, 8, 8), 8), + ((2, 64, 8, 8, 8), 8), + ((3, 15, 5, 5, 5), 3), + ((1, 64, 1, 1, 32775), 8), # has an empty split + ((2, 2048, 1, 1, 1), 8), + ], +) +def test_scratch_slots_are_all_written(shape, groups): + """Every split-K partial slot must be written before it is read. + + ``_forward``/``_backward`` allocate ``pcnt/pmean/pm2`` and + ``ps1/ps2/pdw/pdb`` with ``torch.empty``. A slot that is read but never + written would inherit whatever the caching allocator last left there -- + usually finite, plausible numbers, which no parity test can be relied on to + catch. Poisoning every ``torch.empty``/``empty_like`` with NaN for the + duration of the call turns that into a hard failure, and also proves the + output buffer itself is fully covered by the store masks. + """ + x, weight, bias, grad_out = _make(shape, groups, seed=5) + real_empty, real_empty_like = torch.empty, torch.empty_like + + def poisoned_empty(*args, **kwargs): + t = real_empty(*args, **kwargs) + return t.fill_(float("nan")) if t.is_floating_point() else t + + def poisoned_empty_like(*args, **kwargs): + t = real_empty_like(*args, **kwargs) + return t.fill_(float("nan")) if t.is_floating_point() else t + + xi = x.detach().clone().requires_grad_(True) + wi = weight.detach().clone().requires_grad_(True) + bi = bias.detach().clone().requires_grad_(True) + torch.empty, torch.empty_like = poisoned_empty, poisoned_empty_like + try: + out = triton_group_norm(xi, groups, wi, bi, EPS) + out.backward(grad_out) + finally: + torch.empty, torch.empty_like = real_empty, real_empty_like + + for name, t in (("y", out), ("dx", xi.grad), ("dw", wi.grad), ("db", bi.grad)): + assert torch.isfinite(t).all(), ( + f"{name} contains NaN with poisoned scratch: a split-K slot (or an " + f"output element) is read/returned without ever being written" + ) + ref = _ref64(x, groups, weight, bias, EPS) + assert _rel(out, ref) <= FP32_TOL + + +# --------------------------------------------------------------------------- +# 3. numerics +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.parametrize("value", [0.0, 3.0, 1e3]) +@pytest.mark.parametrize("eps", [1e-5, 1e-12]) +def test_all_equal_input_has_exactly_zero_variance(value, eps): + """Variance exactly 0 => ``rstd = 1/sqrt(eps)`` and ``xhat`` exactly 0. + + This is the sharpest possible statement of the Welford claim: with + ``weight=1, bias=0`` the output must be *identically* zero, with no + tolerance at all. ATen's fp32 GroupNorm does not manage it (it forms the + variance by cancellation and leaves ~1e-05 of noise at ``value=3`` and + ~3e-03 at ``value=1e3``), which is asserted here so the comparison stays + honest if ATen ever changes. + """ + device = torch.device("cuda") + shape = (2, 64, 8, 8, 8) + x = torch.full(shape, value, device=device).contiguous(memory_format=CL) + weight = torch.ones(64, device=device) + bias = torch.zeros(64, device=device) + + got = triton_group_norm(x, 8, weight, bias, eps) + assert torch.equal(got, torch.zeros_like(got)), ( + f"all-equal input must normalise to exactly 0, got max " + f"{got.abs().max().item():.3e}" + ) + # ... and rstd really is 1/sqrt(eps), which only the output scale can show. + _out, _mean, rstd = torch.ops.scaffold_gn.group_norm( + x, 8, None, None, eps, None, None + ) + assert torch.allclose(rstd, torch.full_like(rstd, 1.0 / eps**0.5), rtol=1e-6), ( + f"rstd={rstd.flatten()[0].item():.6e} != 1/sqrt(eps)={1.0 / eps**0.5:.6e}" + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "mean,std,naive_floor", + [(1e4, 1e-2, 1e-1), (1e6, 1.0, 1e1)], +) +def test_welford_at_extreme_mean_to_std_ratio(mean, std, naive_floor): + """Past the ratios the author's suite tests (mu/sigma up to 1e5). + + At ``mu/sigma = 1e6`` the ``E[x^2]-E[x]^2`` formulation is off by ~2e0 to + ~3e2 relative while this kernel holds ~5e-04, and ATen's own fp32 kernel is + 50-70x worse than this one. Measured on MI300A at ``[1, 256, 24^3]``. + """ + x, weight, bias, _ = _make((1, 256, 24, 24, 24), 8, seed=37, mean=mean, std=std) + got = triton_group_norm(x, 8, weight, bias, EPS) + ref = _ref64(x, 8, weight, bias, EPS) + stock = F.group_norm(x, 8, weight, bias, EPS) + + flat = x.reshape(1, 8, -1) + mu = flat.mean(-1) + var = (flat * flat).mean(-1) - mu * mu + naive = (flat - mu[..., None]) / torch.sqrt(var + EPS)[..., None] + naive = naive.reshape(x.shape) * weight.reshape(1, 256, 1, 1, 1) + bias.reshape( + 1, 256, 1, 1, 1 + ) + + err, err_stock, err_naive = _rel(got, ref), _rel(stock, ref), _rel(naive, ref) + print( + f"[mu={mean:g} sd={std:g}] triton={err:.3e} aten={err_stock:.3e} " + f"naive={err_naive:.3e}" + ) + assert err_naive > naive_floor, "the naive formulation was supposed to fail here" + assert err < err_naive / 100.0 + assert err <= err_stock, ( + f"triton {err:.3e} is worse than ATen fp32 {err_stock:.3e} at mu/sigma=" + f"{mean / std:g}" + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("eps", [1e-5, 1e-8, 1e-12]) +def test_tiny_eps_with_tiny_variance(eps): + """``eps`` far below the default with data whose std is ~1e-4. + + ``rstd = 1/sqrt(var + eps)`` reaches ~1e4 here, so any error in the + variance is amplified by that factor before it reaches the output. + """ + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(7) + x = torch.empty(1, 64, 16, 16, 16, device=device, memory_format=CL) + x.normal_(0.0, 1e-4, generator=gen) + weight = torch.ones(64, device=device) + bias = torch.zeros(64, device=device) + got = triton_group_norm(x, 8, weight, bias, eps) + assert _rel(got, _ref64(x, 8, weight, bias, eps)) <= FP32_TOL + + +# --------------------------------------------------------------------------- +# 4. autograd plumbing +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "kind", ["contiguous", "sliced", "expanded", "transposed", "channels_last"] +) +def test_grad_out_layout_variants(kind): + """A cotangent that is not channels-last-contiguous. + + ``_group_norm_backward_op`` relayouts it; the kernels index it with the + *input's* channels-last stride pattern, so a missed relayout silently + permutes the gradient rather than raising. The author's suite only ever + feeds a channels-last-contiguous cotangent to the fast path. + """ + shape = (2, 64, 5, 6, 7) + x, weight, bias, _ = _make(shape, 8, seed=71) + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(77) + base = torch.empty(shape, device=device) + base.normal_(generator=gen) + if kind == "contiguous": + grad_out = base.contiguous() + elif kind == "channels_last": + grad_out = base.contiguous(memory_format=CL) + elif kind == "sliced": + wide = torch.empty( + (shape[0], shape[1], shape[2], shape[3], shape[4] * 2), device=device + ) + wide.normal_(generator=gen) + grad_out = wide.contiguous(memory_format=CL)[..., ::2] + elif kind == "expanded": + col = torch.empty((shape[0], shape[1], shape[2], shape[3], 1), device=device) + col.normal_(generator=gen) + grad_out = col.expand(shape) + else: # transposed + swapped = torch.empty( + (shape[0], shape[1], shape[2], shape[4], shape[3]), device=device + ) + swapped.normal_(generator=gen) + grad_out = swapped.contiguous(memory_format=CL).transpose(3, 4) + + xi = x.detach().clone().requires_grad_(True) + wi = weight.detach().clone().requires_grad_(True) + bi = bias.detach().clone().requires_grad_(True) + triton_group_norm(xi, 8, wi, bi, EPS).backward(grad_out) + + xd = x.detach().clone().double().requires_grad_(True) + wd = weight.detach().clone().double().requires_grad_(True) + bd = bias.detach().clone().double().requires_grad_(True) + _ref64(xd, 8, wd, bd, EPS).backward(grad_out.double()) + + assert _rel(xi.grad, xd.grad) <= FP32_TOL + assert _rel(wi.grad, wd.grad) <= FP32_TOL + assert _rel(bi.grad, bd.grad) <= FP32_TOL + # d_input keeps the *input's* format regardless of the cotangent's. + assert xi.grad.is_contiguous(memory_format=CL) + + +@pytest.mark.gpu +def test_affine_parameters_may_be_non_contiguous_views(): + """``weight``/``bias`` sliced out of a bigger parameter tensor. + + ``is_supported`` only checks rank, numel, device and dtype, so a strided or + offset 1-D parameter reaches the op, which is why it calls ``.contiguous()`` + on both. Nothing upstream tests that. + """ + x, _weight, _bias, _ = _make((1, 64, 4, 4, 4), 8, seed=83) + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(83) + strided = torch.empty(128, device=device) + strided.normal_(1.0, 0.25, generator=gen) + weight = strided[::2] # stride 2 + pack = torch.empty(4, 64, device=device) + pack.normal_(0.0, 0.25, generator=gen) + bias = pack[2] # storage offset + assert not weight.is_contiguous() + assert is_supported(x, 8, weight, bias) + got = triton_group_norm(x, 8, weight, bias, EPS) + assert _rel(got, _ref64(x, 8, weight, bias, EPS)) <= FP32_TOL + + +@pytest.mark.gpu +@pytest.mark.parametrize("lo,hi", [(0, 2), (1, 3), (3, 4)]) +def test_channels_last_views_with_a_storage_offset(lo, hi): + """A batch slice of a bigger channels-last tensor stays channels-last + contiguous but has a non-zero storage offset -- the kernels must address + from ``data_ptr()``, not from the storage base.""" + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(89) + big = torch.empty((4, 64, 5, 6, 7), device=device, memory_format=CL) + big.normal_(generator=gen) + weight = torch.empty(64, device=device) + weight.normal_(1.0, 0.25, generator=gen) + bias = torch.empty(64, device=device) + bias.normal_(0.0, 0.25, generator=gen) + view = big[lo:hi] + assert view.is_contiguous(memory_format=CL) and is_supported(view, 8, weight, bias) + got = triton_group_norm(view, 8, weight, bias, EPS) + assert _rel(got, _ref64(view, 8, weight, bias, EPS)) <= FP32_TOL + + +@pytest.mark.gpu +def test_double_backward_raises_instead_of_returning_garbage(): + """Higher-order gradients are *not* supported, and must say so. + + ``scaffold_gn::group_norm_backward`` has no autograd formula of its own, so + a second ``torch.autograd.grad`` through the kernel raises. Stock + ``F.group_norm`` supports double backward, so this is a real (if narrow) + behavioural difference from the op it replaces -- anything that needs a + gradient penalty or a Hessian-vector product cannot use this kernel. The + test pins "raises loudly", which is the safe half of the story. + """ + x, weight, bias, grad_out = _make((1, 64, 4, 4, 4), 8, seed=97) + xi = x.detach().clone().requires_grad_(True) + wi = weight.detach().clone().requires_grad_(True) + bi = bias.detach().clone().requires_grad_(True) + y = triton_group_norm(xi, 8, wi, bi, EPS) + (gx,) = torch.autograd.grad(y, xi, grad_out, create_graph=True) + with pytest.raises(RuntimeError, match="no autograd formula was registered"): + torch.autograd.grad(gx.sum(), xi) + + # ... and stock really does support it, so this is a divergence not a law. + xr = x.detach().clone().requires_grad_(True) + yr = F.group_norm(xr, 8, weight, bias, EPS) + (gxr,) = torch.autograd.grad(yr, xr, grad_out, create_graph=True) + (ggr,) = torch.autograd.grad(gxr.sum(), xr) + assert torch.isfinite(ggr).all() + + +# --------------------------------------------------------------------------- +# 5. fake / meta kernel +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("out_dtype", [None, torch.float32]) +@pytest.mark.parametrize("has_w,has_b", [(True, True), (False, False), (True, False)]) +def test_fake_forward_matches_real_in_every_branch(dtype, out_dtype, has_w, has_b): + """The fake kernel must promise the real shape, dtype, stride *and* device. + + A meta mismatch is invisible in eager and silently corrupts + ``torch.compile``; the author's suite spot-checks two combinations, this + walks the whole cross product of dtype x out_dtype override x affine. + """ + from torch._subclasses.fake_tensor import FakeTensorMode + + device = torch.device("cuda") + shape = (2, 64, 5, 6, 7) + x = torch.empty(shape, device=device, dtype=dtype, memory_format=CL).normal_() + weight = torch.randn(64, device=device, dtype=dtype) if has_w else None + bias = torch.randn(64, device=device, dtype=dtype) if has_b else None + + real = torch.ops.scaffold_gn.group_norm(x, 8, weight, bias, EPS, "relu", out_dtype) + with FakeTensorMode() as mode: + args = [None if t is None else mode.from_tensor(t) for t in (x, weight, bias)] + fake = torch.ops.scaffold_gn.group_norm( + args[0], 8, args[1], args[2], EPS, "relu", out_dtype + ) + for i, (r, f) in enumerate(zip(real, fake)): + assert r.shape == f.shape, f"out[{i}] shape" + assert r.dtype == f.dtype, f"out[{i}] dtype {r.dtype} != {f.dtype}" + assert r.stride() == f.stride(), f"out[{i}] stride {r.stride()} != {f.stride()}" + assert r.device.type == f.device.type, f"out[{i}] device" + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "layout", ["contiguous", "channels_last", "sliced", "degenerate"] +) +@pytest.mark.parametrize("has_w,has_b", [(True, True), (False, False), (True, False)]) +def test_fake_backward_matches_real_in_every_branch(layout, has_w, has_b): + """The backward's fake kernel must promise what the real op returns. + + The real op relayouts a non-channels-last ``input`` and *always* returns a + channels-last ``d_input``; ``torch.empty_like(input)`` would instead + preserve the input's own format, so for a plain contiguous NCDHW input the + two disagree ((13440, 1, 2688, 448, 64) against (13440, 210, 42, 7, 1)). A + meta mismatch is invisible in eager and silently corrupts + ``torch.compile``, so every branch of the promise -- both layouts, a + non-contiguous view, the shape where the two formats coincide, and each + affine combination -- is checked here rather than only the CL case. + """ + from torch._subclasses.fake_tensor import FakeTensorMode + + device = torch.device("cuda") + if layout == "degenerate": + shape = (2, 64, 1, 1, 1) # contiguous *is* channels_last_3d here + x = torch.randn(shape, device=device) + else: + shape = (2, 64, 5, 6, 7) + if layout == "contiguous": + x = torch.randn(shape, device=device) + elif layout == "channels_last": + x = torch.randn(shape, device=device).contiguous(memory_format=CL) + else: # sliced: neither contiguous nor channels-last contiguous + x = torch.randn((2, 64, 5, 6, 14), device=device)[..., ::2] + grad_out = torch.randn(shape, device=device).contiguous(memory_format=CL) + weight = torch.randn(64, device=device) if has_w else None + bias = torch.randn(64, device=device) if has_b else None + mean = torch.zeros(2, 8, device=device) + rstd = torch.ones(2, 8, device=device) + + real = torch.ops.scaffold_gn.group_norm_backward( + grad_out, x, weight, bias, mean, rstd, 8, None + ) + with FakeTensorMode() as mode: + a = [ + None if t is None else mode.from_tensor(t) + for t in (grad_out, x, weight, bias, mean, rstd) + ] + fake = torch.ops.scaffold_gn.group_norm_backward( + a[0], a[1], a[2], a[3], a[4], a[5], 8, None + ) + names = ("d_input", "d_weight", "d_bias") + for name, r, f in zip(names, real, fake): + assert r.shape == f.shape, f"{name} shape {r.shape} != {f.shape}" + assert r.dtype == f.dtype, f"{name} dtype {r.dtype} != {f.dtype}" + assert r.stride() == f.stride(), f"{name} stride {r.stride()} != {f.stride()}" + assert r.device.type == f.device.type, f"{name} device" + assert real[0].is_contiguous(memory_format=CL) + + +@pytest.mark.gpu +def test_mean_and_rstd_are_not_silently_differentiable(): + """``mean``/``rstd`` are backward state, so they must refuse, not lie. + + They are documented as "not differentiable". Before they were marked as + such, they came back with ``requires_grad=True`` and differentiating + through them *succeeded*: autograd materialised an all-zero cotangent for + the unused ``out``, ran the entire backward (a full-size zeros allocation + plus four kernels) and returned zeros -- a plausible wrong answer where the + true value is ~6e-04. ``ctx.mark_non_differentiable`` turns that into an + error, which is the only safe outcome short of a real formula. + """ + device = torch.device("cuda") + shape = (2, 64, 5, 6, 7) + x = torch.empty(shape, device=device, memory_format=CL).normal_() + xi = x.clone().requires_grad_(True) + out, mean, rstd = torch.ops.scaffold_gn.group_norm( + xi, 8, None, None, EPS, None, None + ) + assert out.requires_grad, "the forward output must still be differentiable" + assert not mean.requires_grad, "mean must be marked non-differentiable" + assert not rstd.requires_grad, "rstd must be marked non-differentiable" + for name, t in (("mean", mean), ("rstd", rstd)): + with pytest.raises(RuntimeError, match="does not require grad"): + torch.autograd.grad(t.sum(), xi) + assert xi.grad is None, f"differentiating {name} left a gradient behind" + # The value that used to come back silently wrong is genuinely non-zero, + # so "returns zeros" was never defensible as an answer. + xd = x.clone().double().requires_grad_(True) + (want,) = torch.autograd.grad(xd.reshape(2, 8, -1).mean(-1).sum(), xd) + assert want.abs().max() > 0 + + +# --------------------------------------------------------------------------- +# 6. contract / drop-in divergences +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups", + [((1, 64, 1, 1, 1), 8), ((2, 64, 1, 1, 1), 8), ((1, 1, 4, 5, 6), 1)], +) +def test_is_supported_accepts_layout_ambiguous_contiguous_input(shape, groups): + """``is_supported`` is *not* simply "False for contiguous input". + + For shapes whose spatial or channel extents are all 1 the contiguous and + channels-last-3d stride patterns coincide, so a plain ``torch.randn`` + tensor is accepted by the fast path. That is benign -- the two layouts are + the same bytes -- but it means callers cannot use ``is_supported`` as a + layout *classifier*. Pinned here so the behaviour is deliberate. + """ + device = torch.device("cuda") + x = torch.randn(shape, device=device) # never asked for channels_last + assert x.is_contiguous() + assert x.is_contiguous(memory_format=CL) + assert is_supported(x, groups) is True + got = triton_group_norm(x, groups, None, None, EPS) + assert _rel(got, _ref64(x, groups, None, None, EPS)) <= FP32_TOL + + +@pytest.mark.gpu +@pytest.mark.parametrize("shape,groups", [((1, 8, 1, 1, 1), 8), ((1, 1, 1, 1, 1), 1)]) +def test_one_value_per_channel_matches_stock_rejection(shape, groups): + """``N*(C/G)*D*H*W == 1`` is a shape ``F.group_norm`` refuses to run. + + The kernel *can* compute it (every group has zero variance, so the answer + is ``bias``), and it used to: ``is_supported`` returned True and + ``triton_group_norm`` returned a value where the op it is a drop-in for + raises ``ValueError``. A caller branching on ``is_supported`` would then + get a different answer from the reference path, which is worse than being + slower, so all three of ``is_supported``, the public wrapper and the raw op + now reject it the same way stock does. + """ + device = torch.device("cuda") + x = torch.empty(shape, device=device, memory_format=CL).normal_() + with pytest.raises(ValueError, match="more than 1 value per channel"): + F.group_norm(x, groups, None, None, EPS) + assert is_supported(x, groups) is False, ( + "is_supported accepts a shape F.group_norm rejects" + ) + # The public wrapper reaches the same rejection through its fallback... + with pytest.raises(ValueError, match="more than 1 value per channel"): + triton_group_norm(x, groups, None, None, EPS) + # ... and the op itself refuses too, for anyone calling it directly. + with pytest.raises(ValueError, match="more than 1 value per channel"): + torch.ops.scaffold_gn.group_norm(x, groups, None, None, EPS, None, None) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups", + [((2, 8, 1, 1, 1), 8), ((1, 8, 2, 1, 1), 8), ((1, 16, 1, 1, 1), 8)], +) +def test_neighbours_of_the_one_value_per_channel_shape_are_still_served(shape, groups): + """The rejection must be exactly stock's, not a shape family around it. + + ``_verify_batch_size`` rejects ``N*(C/G)*spatial == 1`` and nothing else, so + bumping *any one* of N, C/G or the spatial extent to 2 has to come back to + the fast path -- including ``(2, 8, 1, 1, 1)``, which still has a single + element per group. + """ + device = torch.device("cuda") + x = torch.empty(shape, device=device, memory_format=CL).normal_() + F.group_norm(x, groups, None, None, EPS) # stock accepts it + assert is_supported(x, groups) is True + got = triton_group_norm(x, groups, None, None, EPS) + assert _rel(got, _ref64(x, groups, None, None, EPS)) <= FP32_TOL + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups,activation", + [ + ((2, 8, 1, 1, 1), 8, None), + ((2, 8, 1, 1, 1), 8, "relu"), + ((1, 8, 2, 1, 1), 8, None), # 2 elements per group: NOT the degenerate case + ((3, 16, 1, 1, 1), 16, None), + ], +) +def test_single_element_group_gradient_is_exactly_zero(shape, groups, activation): + """One element per group => y is constant in x => dx must be identically 0. + + ``mean == x`` and ``var == 0`` identically, so ``xhat`` is the constant 0 + and nothing downstream depends on ``x``. ``_dx_kernel`` used to answer + 2.2e-05 instead: the compiler contracts ``dy*w - c1`` to + ``fma(dy, w, -c1)`` while ``c1`` was accumulated from the *rounded* + product, so what survives is the product's rounding error (7.0e-08, well + under one ulp of ``dyw``), amplified by ``rstd = 1/sqrt(eps) = 316``. + ``_backward`` now recognises the degenerate case and returns the exact + zero; ATen, on the shapes where it will run at all, leaves ~3e-05 there. + + The ``(1, 8, 2, 1, 1)`` case is the control: two elements per group, so the + gradient is *not* identically zero and the kernel must not zero it. + """ + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(3) + channels = shape[1] + x = torch.empty(shape, device=device, memory_format=CL) + x.normal_(generator=gen) + weight = torch.empty(channels, device=device) + weight.normal_(1.0, 0.25, generator=gen) + bias = torch.empty(channels, device=device) + bias.normal_(0.0, 0.25, generator=gen) + grad_out = torch.empty(shape, device=device, memory_format=CL) + grad_out.normal_(generator=gen) + + assert is_supported(x, groups, weight, bias, activation) + xi = x.clone().requires_grad_(True) + wi = weight.clone().requires_grad_(True) + bi = bias.clone().requires_grad_(True) + triton_group_norm(xi, groups, wi, bi, EPS, activation).backward(grad_out) + + xd = x.clone().double().requires_grad_(True) + wd = weight.clone().double().requires_grad_(True) + bd = bias.clone().double().requires_grad_(True) + _ref64(xd, groups, wd, bd, EPS, activation).backward(grad_out.double()) + + if channels // groups * shape[2] * shape[3] * shape[4] == 1: + assert torch.equal(xi.grad, torch.zeros_like(xi.grad)), ( + f"dx should be exactly 0, got {xi.grad.abs().max().item():.3e}" + ) + assert xd.grad.abs().max() == 0, "the float64 reference disagrees" + # d_weight is exactly 0 too (xhat is exactly 0); d_bias is not. + assert torch.equal(wi.grad, torch.zeros_like(wi.grad)) + assert _rel(bi.grad, bd.grad) <= FP32_TOL + else: + assert xd.grad.abs().max() > 0, "control case is supposed to be non-trivial" + assert xi.grad.abs().max() > 0, "the kernel zeroed a non-degenerate gradient" + # A *two*-element group is merely ill-conditioned, not degenerate: + # xhat is +-1/sqrt(1+eps/var) and dx is a difference of near-equal + # terms, so every fp32 implementation loses digits here -- 4.9e-04 + # relative for this kernel and 1.4e-04 for ATen on this input. The + # bound is therefore loose against float64, and tight against ATen, + # which suffers the same cancellation. + assert _rel(xi.grad, xd.grad) <= 1e-3 + xa = x.clone().requires_grad_(True) + F.group_norm(xa, groups, weight, bias, EPS).backward(grad_out) + assert _rel(xi.grad, xa.grad) <= 1e-3 + + +# --------------------------------------------------------------------------- +# 7. composition +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.parametrize("activation", [None, "relu"]) +def test_torch_compile_with_dynamic_shapes(activation): + """``dynamic=True`` as well as the author's ``dynamic=False``. + + With dynamic shapes the fake kernel is invoked on *symbolic* sizes, so a + shape/stride promise that only happens to hold for a concrete size shows up + here and nowhere else. ``fullgraph=True`` is the no-graph-break assertion. + """ + x, weight, bias, grad_out = _make((2, 64, 6, 6, 6), 8, seed=59) + + def fn(x, weight, bias): + return triton_group_norm(x, 8, weight, bias, EPS, activation) * 2.0 + + def run(f): + xi = x.detach().clone().requires_grad_(True) + wi = weight.detach().clone().requires_grad_(True) + bi = bias.detach().clone().requires_grad_(True) + out = f(xi, wi, bi) + out.backward(grad_out) + return out.detach(), xi.grad, wi.grad, bi.grad + + torch._dynamo.reset() + eager = run(fn) + compiled = run(torch.compile(fn, fullgraph=True, dynamic=True)) + for name, a, b in zip(("y", "dx", "dweight", "dbias"), compiled, eager): + assert torch.equal(a, b), f"dynamic-shape compile differs from eager in {name}" + assert compiled[0].is_contiguous(memory_format=CL) + + +# --------------------------------------------------------------------------- +# 8. determinism, across processes +# --------------------------------------------------------------------------- + +_DETERMINISM_SCRIPT = textwrap.dedent( + """ + import hashlib, sys, torch + from ScaFFold.unet.triton_group_norm import triton_group_norm + CL = torch.channels_last_3d + + def h(t): + b = t.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes() + return hashlib.sha256(b).hexdigest() + + if sys.argv[1] == "warm": + # Different JIT order, different lru_cache occupancy, different + # allocator state and different free memory before the real work. + junk = [] + for shape, g in (((3, 128, 7, 7, 7), 8), ((1, 15, 5, 5, 5), 3)): + a = torch.empty(shape, device="cuda", memory_format=CL).normal_() + triton_group_norm(a, g, None, None, 1e-5, "relu") + junk.append(torch.empty(1 << 25, device="cuda")) + del junk + torch.cuda.empty_cache() + + for shape, groups, act, dtype in ( + ((2, 128, 9, 11, 13), 8, None, torch.float32), + ((2, 64, 6, 7, 8), 8, "relu", torch.bfloat16), + ((3, 15, 5, 5, 5), 3, None, torch.float32), + ): + gen = torch.Generator(device="cuda").manual_seed(31) + x = torch.empty(shape, device="cuda", dtype=dtype, memory_format=CL) + x.normal_(generator=gen) + w = torch.empty(shape[1], device="cuda", dtype=dtype) + w.normal_(1.0, 0.25, generator=gen) + b = torch.empty(shape[1], device="cuda", dtype=dtype) + b.normal_(0.0, 0.25, generator=gen) + go = torch.empty(shape, device="cuda", dtype=dtype, memory_format=CL) + go.normal_(generator=gen) + xi = x.clone().requires_grad_(True) + wi = w.clone().requires_grad_(True) + bi = b.clone().requires_grad_(True) + y = triton_group_norm(xi, groups, wi, bi, 1e-5, act) + y.backward(go) + print(shape, groups, act, dtype, + h(y), h(xi.grad), h(wi.grad), h(bi.grad), flush=True) + """ +) + + +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.timeout(900) +def test_bitwise_determinism_across_processes(): + """Process-to-process bitwise reproducibility, which is half the claim. + + ``test_bitwise_determinism`` upstream only calls the kernel twice in *one* + process, where the plan is already memoised and the JIT cache already warm. + This runs three fresh interpreters -- one of which first JITs other shapes, + churns the caching allocator and changes how much memory is free -- and + compares SHA-256 of the raw output bytes. Anything that made the split + count, tile size or launch geometry depend on device state rather than on + the shape would show up only here. + """ + outputs = [] + for mode in ("plain", "plain", "warm"): + result = subprocess.run( + [sys.executable, "-c", _DETERMINISM_SCRIPT, mode], + capture_output=True, + text=True, + cwd=REPO_ROOT, + timeout=600, + ) + assert result.returncode == 0, result.stderr[-3000:] + outputs.append(result.stdout) + assert outputs[0] == outputs[1], "two identical processes disagree" + assert outputs[0] == outputs[2], ( + "a process that JITted other shapes first disagrees:\n" + f"{outputs[0]}\n--- vs ---\n{outputs[2]}" + ) + assert outputs[0].count("\n") >= 3 + + +# --------------------------------------------------------------------------- +# 9. multi-device +# --------------------------------------------------------------------------- + +_DEVICE_GUARD_SCRIPT = textwrap.dedent( + """ + import sys, torch + import torch.nn.functional as F + from ScaFFold.unet.triton_group_norm import triton_group_norm + CL = torch.channels_last_3d + torch.cuda.set_device(0) # current device = 0 + other = "cuda:1" + g = torch.Generator(device=other).manual_seed(1) + x = torch.empty((1, 64, 4, 4, 4), device=other, memory_format=CL) + x.normal_(generator=g) + w = torch.empty(64, device=other); w.normal_(1.0, 0.25, generator=g) + b = torch.empty(64, device=other); b.normal_(0.0, 0.25, generator=g) + go = torch.empty((1, 64, 4, 4, 4), device=other, memory_format=CL) + go.normal_(generator=g) + + def rel(a, e): + return ((a.double() - e.double()).abs().max() + / e.double().abs().max().clamp_min(1e-30)).item() + + # ATen carries a DeviceGuard, so this is the behaviour to match. + xr = x.clone().requires_grad_(True) + wr = w.clone().requires_grad_(True) + br = b.clone().requires_grad_(True) + F.group_norm(xr, 8, wr, br, 1e-5).backward(go) + + xi = x.clone().requires_grad_(True) + wi = w.clone().requires_grad_(True) + bi = b.clone().requires_grad_(True) + y = triton_group_norm(xi, 8, wi, bi, 1e-5) # tensors on 1, current is 0 + y.backward(go) # ... and so is the backward + torch.cuda.synchronize() + assert torch.cuda.current_device() == 0, "the guard leaked the device" + for name, got, want in (("y", y, F.group_norm(xr.detach(), 8, w, b, 1e-5)), + ("dx", xi.grad, xr.grad), + ("dw", wi.grad, wr.grad), + ("db", bi.grad, br.grad)): + assert got.device == torch.device(other), f"{name} on {got.device}" + e = rel(got, want) + assert e < 1e-4, f"{name}: rel err {e}" + print("OK") + """ +) + + +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.timeout(600) +def test_kernel_runs_on_the_inputs_device_not_the_current_one(): + """Tensors on cuda:1 while cuda:0 is current, forward *and* backward. + + A Triton launch goes to whatever device is *current*, so without a device + guard the kernel dereferences another device's pointers and the process + dies with ``Memory access fault by GPU node-N``. ``F.group_norm`` carries + ATen's ``DeviceGuard`` and handles the identical call, so this is a + divergence from the op being replaced, not a PyTorch limitation. + + Run in a subprocess because the failure mode is an unrecoverable GPU memory + fault, which would take the whole pytest session with it. + """ + if torch.cuda.device_count() < 2: + pytest.skip("needs 2 visible CUDA devices") + result = subprocess.run( + [sys.executable, "-c", _DEVICE_GUARD_SCRIPT], + capture_output=True, + text=True, + cwd=REPO_ROOT, + timeout=480, + ) + assert result.returncode == 0 and "OK" in result.stdout, ( + f"returncode={result.returncode}\nstdout={result.stdout}\n" + f"stderr={result.stderr[-2000:]}" + ) + + +@pytest.mark.gpu +def test_device_guard_helper_is_a_no_op_on_the_current_device(): + """The guard must be free on the hot path and real off it. + + ``_device_guard`` skips ``torch.cuda.device`` when the tensor already lives + on the current device (1.55 us against 0.51 us of host time per call, which + is 0.5% of the two smallest scale-8 shapes' 0.65 ms fwd+bwd because they + are host-dispatch bound). This pins both halves of that shortcut so a + future edit cannot quietly turn it into "no guard at all"; the multi-device + behaviour itself is covered by the subprocess test above. + """ + device = torch.device("cuda", torch.cuda.current_device()) + guard = tgn._device_guard(device) + assert guard is tgn._NO_GUARD, "should not build a guard for the current device" + # Constructing a guard for another index does not touch that device. + other = torch.device("cuda", device.index + 1) + assert tgn._device_guard(other) is not tgn._NO_GUARD, ( + "a foreign device must get a real guard" + ) + + +# --------------------------------------------------------------------------- +# 10. addressing at the int32 boundary +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +@pytest.mark.slow +@pytest.mark.timeout(1800) +def test_int32_addressing_at_its_documented_maximum(): + """``numel = INT32_MAX - 127`` with N=2, i.e. ``plan.int64 is False``. + + The author's suite tests the shape *above* the switch + (``test_correct_above_int32_max_elements``) but never the largest shape the + **int32** path itself has to serve, which is where a missing term in the + ``numel + channels > INT32_MAX`` guard would bite. ``65 * 63 * 4097`` is + ``2^24 - 1`` voxels, so nothing about the extents is a power of two. + + Verified without materialising an NCDHW reference: the statistics are + checked against a chunked float64 reduction over the physical (N, S, C) + view, and the output against an elementwise recomputation done per batch + item (a truncated offset wraps partway through, so sample 1 would break + while sample 0 did not). + """ + device = torch.device("cuda") + shape = (2, 64, 65, 63, 4097) + n, channels = shape[0], shape[1] + spatial = shape[2] * shape[3] * shape[4] + numel = n * channels * spatial + assert numel == 2**31 - 128, numel + + plan = tgn._plan(n, channels, spatial, 8, numel) + assert plan.int64 is False, "this shape is supposed to use the int32 path" + + free, total = torch.cuda.mem_get_info() + needed = 4 * numel * 4 + if free < needed: + pytest.skip( + f"needs ~{needed / 2**30:.0f} GiB free, device has " + f"{free / 2**30:.0f} GiB of {total / 2**30:.0f} GiB" + ) + + gen = torch.Generator(device=device).manual_seed(53) + x = torch.empty(shape, device=device, memory_format=CL) + x.normal_(generator=gen) + weight = torch.randn(channels, device=device, generator=gen) + bias = torch.randn(channels, device=device, generator=gen) + out, mean, rstd = torch.ops.scaffold_gn.group_norm( + x, 8, weight, bias, EPS, None, None + ) + + group_channels = channels // 8 + flat = x.permute(0, 2, 3, 4, 1).reshape(n, spatial, channels) # no copy + chunk = 1 << 20 + for i in range(n): + acc = torch.zeros(8, dtype=torch.float64, device=device) + for s in range(0, spatial, chunk): + acc += ( + flat[i, s : s + chunk] + .double() + .reshape(-1, 8, group_channels) + .sum(dim=(0, 2)) + ) + mu = acc / (spatial * group_channels) + acc2 = torch.zeros(8, dtype=torch.float64, device=device) + for s in range(0, spatial, chunk): + d = ( + flat[i, s : s + chunk].double().reshape(-1, 8, group_channels) + - mu[None, :, None] + ) + acc2 += (d * d).sum(dim=(0, 2)) + var = acc2 / (spatial * group_channels) + assert _rel(mean[i], mu) <= 1e-5, f"sample {i} mean" + assert _rel(rstd[i], 1.0 / torch.sqrt(var + EPS)) <= 1e-5, f"sample {i} rstd" + del flat + + mv = ( + mean.reshape(n, 8, 1).expand(n, 8, group_channels).reshape(n, channels, 1, 1, 1) + ) + rv = ( + rstd.reshape(n, 8, 1).expand(n, 8, group_channels).reshape(n, channels, 1, 1, 1) + ) + for i in range(n): + recomputed = (x[i : i + 1] - mv[i : i + 1]) * rv[i : i + 1] * weight.reshape( + 1, channels, 1, 1, 1 + ) + bias.reshape(1, channels, 1, 1, 1) + # fp32 subtraction of two near-equal fp32 values is exact. + err = (out[i : i + 1] - recomputed).abs().max().item() + scale = recomputed.abs().max().item() + print(f"[int32-max sample {i}] elementwise rel err {err / scale:.3e}") + assert err / scale < 1e-5, f"sample {i}" + del recomputed + del x, out, mean, rstd, mv, rv + torch.cuda.empty_cache() + + +# --------------------------------------------------------------------------- +# 11. coverage the mutation sweep found thin +# --------------------------------------------------------------------------- + + +@contextlib.contextmanager +def _force_int64_addressing(): + """Make every plan take the int64 tile-base path, whatever the shape. + + ``_Plan`` sets ``int64 = numel + channels > _INT32_MAX``, so dropping the + threshold turns the wide path on for a shape that fits in a few MiB. The + plan cache is keyed on the shape, not on the threshold, so it has to be + cleared on the way in *and* on the way out. + """ + real = tgn._INT32_MAX + tgn._plan.cache_clear() + tgn._INT32_MAX = -1 + try: + yield + finally: + tgn._INT32_MAX = real + tgn._plan.cache_clear() + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups", + [ + ((2, 64, 5, 6, 7), 8), # ragged tail, several splits + ((1, 2048, 6, 6, 6), 8), # widest channel count + ((3, 15, 5, 5, 5), 3), # masked channel axis as well + ((1, 64, 1, 1, 32775), 8), # has an empty split + ], +) +@pytest.mark.parametrize("activation", [None, "relu"]) +def test_int64_addressing_path_is_behaviourally_correct(shape, groups, activation): + """Run the ``INT64=True`` branch of all seven kernels on a small shape. + + ``INT64`` is a ``tl.constexpr``, so the wide and narrow paths are *different + compiled kernels*; only shapes above 2^31 elements reach the wide one + naturally, and the one test that does is ``@pytest.mark.slow`` and needs + 8 GiB. In a default ``-m "not slow"`` run the int64 branch therefore has no + behavioural coverage at all -- forcing ``self.int64`` gives it some for the + price of a few MiB. + + The two paths differ only in the *type* of the scalar tile base, so the + results must be **bitwise** identical, which is a far sharper assertion than + a tolerance and would catch a widened offset that lost or duplicated a tile. + """ + x, weight, bias, grad_out = _make(shape, groups, seed=abs(hash(shape)) % 997) + + def run(): + xi = x.detach().clone().requires_grad_(True) + wi = weight.detach().clone().requires_grad_(True) + bi = bias.detach().clone().requires_grad_(True) + out = triton_group_norm(xi, groups, wi, bi, EPS, activation) + out.backward(grad_out) + return out.detach(), xi.grad, wi.grad, bi.grad + + plan32 = tgn._plan( + shape[0], shape[1], shape[2] * shape[3] * shape[4], groups, x.numel() + ) + assert plan32.int64 is False, "shape is supposed to fit the int32 path" + narrow = run() + + with _force_int64_addressing(): + plan64 = tgn._plan( + shape[0], shape[1], shape[2] * shape[3] * shape[4], groups, x.numel() + ) + assert plan64.int64 is True, "the int64 path was not forced on" + wide = run() + + for name, a, b in zip(("y", "dx", "dweight", "dbias"), wide, narrow): + assert torch.equal(a, b), f"int64 path differs from int32 in {name}" + # ... and both are actually right, not identically wrong. + ref = _ref64(x, groups, weight, bias, EPS, activation) + assert _rel(wide[0], ref) <= FP32_TOL + + +#: ``(shape, groups)`` whose split-K partials have *unequal* counts, because +#: ``chunk = ceil(S / nsplit)`` does not divide ``S``. Chan's combine weights +#: the delta by ``cnt_b / (cnt_a + cnt_b)``; with equal counts every level of +#: the reduction tree has ``cnt_a == cnt_b``, so weighting by the wrong one is +#: invisible. Only a ragged (or empty) trailing split exposes it -- which is +#: why the mutation sweep killed that bug with exactly two parametrizations of +#: one test upstream. +_UNEVEN_SPLIT_CASES = [ + ((2, 64, 9, 7, 5), 8), + ((1, 2048, 6, 6, 6), 8), + ((1, 64, 1, 1, 32775), 8), + ((1, 256, 1, 1, 2049), 8), + ((2, 128, 11, 13, 17), 8), + ((2, 15, 9, 9, 9), 3), # masked channel axis as well + ((1, 20, 17, 17, 17), 5), # 16 splits, trailing split 15 voxels short +] + + +@pytest.mark.gpu +@pytest.mark.parametrize("shape,groups", _UNEVEN_SPLIT_CASES) +@pytest.mark.parametrize("eps", [1e-5, 0.5]) +def test_group_statistics_match_float64_with_uneven_splits(shape, groups, eps): + """Assert ``mean``/``rstd`` themselves, not just the output they feed. + + Two things hide inside the output's 1e-4 tolerance and show up here: + + * **the Welford merge.** The shapes above all have at least one split with + a different element count from its neighbours, which is the only + configuration in which mis-weighting Chan's delta changes the answer. + * **where ``eps`` goes.** Every parity test in both files uses + ``eps=1e-5`` against a variance of ~1, where ``1/sqrt(var+eps)`` and + ``1/(sqrt(var)+eps)`` agree to ~1e-5 -- inside that tolerance. At + ``eps=0.5`` they are 0.816 and 0.667, a 22% difference that no tolerance + can absorb. + """ + spatial = shape[2] * shape[3] * shape[4] + plan = tgn._plan(shape[0], shape[1], spatial, groups, shape[0] * shape[1] * spatial) + counts = { + max(0, min(sp * plan.chunk + plan.chunk, spatial) - sp * plan.chunk) + for sp in range(plan.nsplit) + } + assert plan.nsplit > 1 and len(counts) > 1, ( + f"{shape} was supposed to give unequal split counts; nsplit=" + f"{plan.nsplit} chunk={plan.chunk} counts={sorted(counts)}" + ) + + x, _weight, _bias, _ = _make(shape, groups, seed=abs(hash(shape)) % 997) + _out, mean, rstd = torch.ops.scaffold_gn.group_norm( + x, groups, None, None, eps, None, None + ) + flat = x.double().reshape(shape[0], groups, -1) + mean64 = flat.mean(-1) + var64 = ((flat - mean64[..., None]) ** 2).mean(-1) + rstd64 = 1.0 / torch.sqrt(var64 + eps) + assert _rel(mean, mean64) <= 1e-5, "group mean" + assert _rel(rstd, rstd64) <= 1e-5, "group rstd (eps placement / Welford merge)" + + +@pytest.mark.gpu +@pytest.mark.parametrize("groups", [1, 2, 4]) +@pytest.mark.parametrize("seed", [0, 1, 2, 3]) +def test_welford_correction_recovers_rstd_in_a_single_tile_reduction(groups, seed): + """The third reduction pass (``corr``) is load-bearing, and here is where. + + ``mean0 = sum(x)/n`` loses digits in proportion to the tile's element count + times ``mu/sigma``; ``corr = sum(x-mean0)/n`` recovers them, and ``M2`` is + then formed around the corrected mean. The effect is largest when one tile + carries a whole group's reduction, which is this shape: ``block_s_stats`` + covers all 128 voxels and ``nsplit == 1``, so 8192/``G`` elements per group + go through a single ``mean0``. + + At ``mu/sigma = 1e6`` the correction is worth **216x** (G=1), **580x** + (G=2) and **1472x** (G=4) on the relative error of ``rstd`` -- measured by + running a copy of this module with the term deleted. Corrected lands at + ~1e-07 for every seed and group count; without it, at 2.1e-05 to 1.6e-04. + The 1e-06 ceiling below sits an order of magnitude above the first and an + order of magnitude below the second. + + The *output* is not a witness for this: ``y`` moves by at most ~1.4x with + or without the term, because it is dominated by the fp32 representation of + the mean. That is why this asserts ``rstd`` directly. + """ + device = torch.device("cuda") + shape = (2, 64, 8, 4, 4) + spatial = shape[2] * shape[3] * shape[4] + plan = tgn._plan(shape[0], shape[1], spatial, groups, shape[0] * shape[1] * spatial) + assert plan.nsplit == 1 and plan.block_s_stats >= spatial, ( + f"case is supposed to be a single-tile reduction; nsplit={plan.nsplit} " + f"block_s_stats={plan.block_s_stats} spatial={spatial}" + ) + + gen = torch.Generator(device=device).manual_seed(seed) + x = torch.empty(shape, device=device, memory_format=CL) + x.normal_(1e4, 1e-2, generator=gen) # mu/sigma = 1e6 + _out, mean, rstd = torch.ops.scaffold_gn.group_norm( + x, groups, None, None, EPS, None, None + ) + flat = x.double().reshape(shape[0], groups, -1) + mean64 = flat.mean(-1) + var64 = ((flat - mean64[..., None]) ** 2).mean(-1) + err = _rel(rstd, 1.0 / torch.sqrt(var64 + EPS)) + print(f"[corr G={groups} seed={seed}] rstd rel err {err:.3e}") + assert err <= 1e-6, ( + f"rstd rel err {err:.3e} at mu/sigma=1e6 with G={groups}: the tile mean " + f"correction is not doing its job" + ) + + +# --------------------------------------------------------------------------- +# 12. the fused finalize and the capped elementwise grid +# --------------------------------------------------------------------------- +# +# ``_stats_finalize``/``_bwd_finalize``/``_dwdb_reduce`` are no longer their own +# launches: each is recomputed inside the elementwise kernel that consumes it. +# Two consequences need pinning. +# +# * The elementwise grid is capped at ``GNConfig.elem_progs`` and each program +# *strides* over its share of the tiles, so that the fused finalize costs +# ``nprog_elem * nsplit`` and not ``nblk_elem * nsplit`` reads. No scale-8 +# shape and no shape in either suite reaches that path with the shipped +# table -- ``nprog_elem == nblk_elem`` at every small shape -- so it has to be +# reached deliberately. +# * The cap is a *performance* knob. If it could change a single bit of the +# output it would break the module's reproducibility contract, since it is +# the one plan field that does not follow from the shape alone. + + +@contextlib.contextmanager +def _forced_config(channels, spatial, **overrides): + """Temporarily install a tiling config for one ``(channels, spatial)`` key. + + ``default_config`` keys the frozen table by ``(num_channels, cube-root + spatial extent)``, so the spatial extent has to be a perfect cube here. + Restores the previous entry (or its absence) and clears the plan cache on + the way out, so no other test can see it. + """ + edge = round(spatial ** (1.0 / 3.0)) + assert edge**3 == spatial, "forced configs need a cube spatial extent" + key = (channels, edge) + cfg = tgn.GNConfig(*tgn.default_config(channels, spatial).key()) + for name, value in overrides.items(): + assert hasattr(cfg, name), name + setattr(cfg, name, value) + sentinel = object() + saved = tgn._TUNED.get(key, sentinel) + tgn._TUNED[key] = cfg + tgn._plan.cache_clear() + try: + yield cfg + finally: + if saved is sentinel: + del tgn._TUNED[key] + else: + tgn._TUNED[key] = saved + tgn._plan.cache_clear() + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups", + [ + ((1, 64, 16, 16, 16), 8), + ((2, 64, 16, 16, 16), 8), # N > 1: the stride is per (blk, n) program + ((1, 20, 8, 8, 8), 5), # capped grid *and* a padded channel axis + ], +) +@pytest.mark.parametrize("elem_progs", [1, 3, 8]) +@pytest.mark.parametrize("activation", [None, "relu"]) +def test_capped_elementwise_grid_strides_over_its_tiles( + shape, groups, elem_progs, activation +): + """Fewer elementwise programs than tiles: each must cover several tiles. + + A grid-stride loop that got its start, stride or trip count wrong leaves + part of the output (and of ``d_input``) unwritten -- which, since both are + ``torch.empty``, surfaces as plausible stale numbers rather than as a + crash. ``elem_progs=1`` is the extreme: one program per sample walks every + tile, so it also pins that the fused statistics are hoisted out of the loop + correctly rather than being recomputed per iteration from stale state. + """ + spatial = shape[2] * shape[3] * shape[4] + with _forced_config(shape[1], spatial, elem_tile=1024, elem_progs=elem_progs): + plan = tgn._plan(shape[0], shape[1], spatial, groups, 0) + assert plan.nprog_elem == min(plan.nblk_elem, elem_progs) + assert plan.nprog_elem < plan.nblk_elem, ( + f"the cap has to actually bite: nprog={plan.nprog_elem} " + f"nblk={plan.nblk_elem}" + ) + _parity( + shape, + groups, + activation, + seed=abs(hash((shape, elem_progs))) % 997, + label=f"{shape} elem_progs={elem_progs}", + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups", [((2, 64, 16, 16, 16), 8), ((1, 20, 8, 8, 8), 5)] +) +def test_elementwise_grid_cap_is_bitwise_neutral(shape, groups): + """``elem_progs`` may not change a single bit of any output. + + It is the only field of ``_Plan`` that is a free parameter rather than a + consequence of the shape, and the module promises bitwise reproducibility. + That promise holds only because the elementwise kernels carry nothing + across loop iterations: the fused finalize is computed once per program + from the *same* partials with the *same* tile shape, and the tile bodies + are pure elementwise. If tuning this knob ever moved a result, the frozen + table would have become part of the numerical contract. + """ + spatial = shape[2] * shape[3] * shape[4] + x, weight, bias, grad_out = _make(shape, groups, seed=11) + results = [] + for elem_progs in (0, 1, 5, 64, 4096): + with _forced_config(shape[1], spatial, elem_tile=1024, elem_progs=elem_progs): + xi = x.detach().clone().requires_grad_(True) + wi = weight.detach().clone().requires_grad_(True) + bi = bias.detach().clone().requires_grad_(True) + y = triton_group_norm(xi, groups, wi, bi, EPS) + y.backward(grad_out) + results.append( + (y.detach().clone(), xi.grad.clone(), wi.grad.clone(), bi.grad.clone()) + ) + for elem_progs, got in zip((1, 5, 64, 4096), results[1:]): + for name, a, b in zip(("y", "dx", "dweight", "dbias"), got, results[0]): + assert torch.equal(a, b), ( + f"elem_progs={elem_progs} changed {name} bitwise; the grid cap " + f"is supposed to be a pure performance knob" + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups,elem_progs", + [ + ((1, 64, 16, 16, 16), 8, 0), + ((1, 64, 16, 16, 16), 8, 3), + ((2, 96, 8, 8, 8), 6, 2), # padded channel axis, N > 1, capped grid + ], +) +def test_fused_finalize_publishes_the_statistics(shape, groups, elem_progs): + """``mean``/``rstd`` are published by program 0 of the *normalize* kernel. + + There is no separate finalize launch any more: every elementwise program + re-derives the statistics from the split-K Welford partials, and program 0 + is the one that stores them for the backward pass. A wrong publishing + program, a wrong partials index, or a group-mask slip in that fused + reduction would hand the backward garbage while leaving the forward -- which + uses its own locally computed copy -- perfectly correct. So check the + published tensors directly against float64. + """ + spatial = shape[2] * shape[3] * shape[4] + with _forced_config(shape[1], spatial, elem_tile=1024, elem_progs=elem_progs): + x, weight, bias, _grad = _make(shape, groups, seed=3) + _y, mean, rstd = torch.ops.scaffold_gn.group_norm( + x, groups, weight, bias, EPS, None, None + ) + flat = x.double().reshape(shape[0], groups, -1) + mean64 = flat.mean(-1) + var64 = ((flat - mean64[..., None]) ** 2).mean(-1) + assert _rel(mean, mean64) <= FP32_TOL + assert _rel(rstd, 1.0 / torch.sqrt(var64 + EPS)) <= FP32_TOL + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "shape,groups", [((2, 2048, 1, 1, 1), 8), ((1, 1024, 2, 1, 1), 8)] +) +def test_dweight_blocks_are_covered_when_there_are_more_of_them_than_tiles( + shape, groups +): + """The dweight/dbias reduction rides in ``_dx_kernel``'s first NDW programs. + + Those blocks are per-*channel*, the elementwise tiles are per-*voxel*, and + nothing makes the first outnumber the second: at ``(2, 2048, 1, 1, 1)`` + there is one elementwise tile and eight dweight blocks. The grid is + ``max(nprog_elem, dwdb_progs)`` for exactly that reason, and a grid of + ``nprog_elem`` alone would silently leave 7/8 of ``d_weight`` unwritten. + """ + spatial = shape[2] * shape[3] * shape[4] + plan = tgn._plan(shape[0], shape[1], spatial, groups, 0) + assert plan.dwdb_progs > plan.nprog_elem, ( + f"case is supposed to have more dweight blocks ({plan.dwdb_progs}) than " + f"elementwise programs ({plan.nprog_elem})" + ) + assert plan.grid_dx == plan.dwdb_progs + _parity(shape, groups, seed=13) + + +# --------------------------------------------------------------------------- +# the kernel-failure boundary +# --------------------------------------------------------------------------- + + +def test_kernel_failures_are_tagged_and_carry_their_cause(): + """Everything the launch region raises comes out as ``TritonKernelError``. + + The tag is what lets a caller with a fallback (``FastGroupNorm``'s ladder) + catch *exactly* "the kernel is broken" instead of catching ``Exception`` and + then trying to enumerate every framework mechanism -- saved-tensor pack + hooks, ``torch.utils.checkpoint``'s recompute control flow, functorch -- + that legitimately raises through a forward. The region it wraps is closed + (allocations and launches, no autograd-observable op), so a blanket catch + inside it is sound where one at the call site is not. + + The tag must survive the *type* of the original error, whatever it was: a + mismatched Triton release raises ``TypeError``/``AttributeError`` from a + changed signature, an unwritable JIT cache ``OSError``, a bad launch + ``RuntimeError``. + """ + for original in ( + RuntimeError("launch failed"), + TypeError("triton API changed"), + AttributeError("no such attribute"), + OSError("unwritable cache dir"), + ImportError("no module named triton"), + ): + + @tgn._tag_kernel_failures + def _boom(): + raise original + + with pytest.raises(tgn.TritonKernelError) as caught: + _boom() + assert caught.value.__cause__ is original + assert type(original).__name__ in str(caught.value) + + +def test_out_of_memory_is_not_tagged_as_a_kernel_failure(): + """An OOM is a resource condition, and every fallback allocates as much. + + Tagging it would make the ladder retry on a rung that is about to OOM in + the same place, and would latch a rung off for the rest of the process on a + transient, per-rank event. It has to come out unchanged. + """ + + @tgn._tag_kernel_failures + def _oom(): + raise torch.OutOfMemoryError("simulated OOM") + + with pytest.raises(torch.OutOfMemoryError): + _oom() + assert not issubclass(torch.OutOfMemoryError, tgn.TritonKernelError) + + +def test_contract_violations_are_not_tagged(): + """``_validate``'s ``ValueError``s are caller errors, and stay loud. + + ``is_supported`` accepts exactly what ``_validate`` accepts, so a caller + that branches on the predicate can never see one; if the two ever disagree, + the failure must not be laundered into "the kernel is broken" and silently + fall back. + """ + with pytest.raises(ValueError, match="activation must be one of"): + tgn._validate(torch.zeros(1, 8, 2, 2, 2), 8, None, None, "gelu") + with pytest.raises(ValueError, match="expected a 5-D"): + tgn._validate(torch.zeros(1, 8, 2, 2), 8, None, None, None) + + +@pytest.mark.gpu +def test_a_real_launch_failure_is_tagged(monkeypatch): + """End to end: break the launch and the public op raises the tagged type.""" + x = torch.randn(1, 64, 4, 4, 4, device="cuda").to(memory_format=CL) + + def _broken(*args, **kwargs): + raise RuntimeError("simulated HIP launch failure") + + tgn._ensure_kernels() + monkeypatch.setattr(tgn, "_stats_partial_kernel", _Unlaunchable(_broken)) + with pytest.raises(tgn.TritonKernelError): + torch.ops.scaffold_gn.group_norm(x, 8, None, None, EPS, None, None) + + +class _Unlaunchable: + """A stand-in for a ``triton.jit`` kernel whose launch raises.""" + + def __init__(self, fn): + self._fn = fn + + def __getitem__(self, grid): + return self._fn + + +# --------------------------------------------------------------------------- +# the fused activation on non-finite values +# --------------------------------------------------------------------------- + +#: NaN, +Inf, -Inf, -0.0 and four ordinary values. ``tl.maximum(y, 0)`` returns +#: the non-NaN operand and ``tl.where(y > 0, y, 0)`` fails ``NaN > 0``, so both +#: of the obvious spellings map NaN to 0.0 where ``F.relu`` propagates it. +_SPECIALS = [float("nan"), float("inf"), float("-inf"), -0.0, 0.0, -1.0, 1.0, 2.0] + + +def _zero_weight_case(activation, seed=3): + """A case whose pre-activation is exactly ``bias``, elementwise. + + Poisoning the *input* can only produce NaN pre-activations -- one non-finite + value makes the whole group's statistics NaN -- so the values that actually + distinguish the spellings of ReLU have to be placed directly. A zero + ``weight`` does that: ``xhat * 0 + bias == bias``. + """ + x = torch.randn( + 1, + 64, + 4, + 4, + 4, + device="cuda", + generator=torch.Generator("cuda").manual_seed(seed), + ).to(memory_format=CL) + weight = torch.zeros(64, device="cuda") + bias = torch.tensor(_SPECIALS * 8, device="cuda") + reference = F.group_norm(x, 8, weight, bias, EPS) + if activation == "relu": + reference = F.relu(reference) + return x, weight, bias, reference + + +@pytest.mark.gpu +def test_fused_relu_matches_f_relu_on_nan_inf_and_negative_zero(): + """The fused store must be ``F.relu``, bit for bit, on every special value. + + NaN in, NaN out -- and that matters beyond numerics: ScaFFold aborts a run + whose loss goes non-finite, so an activation that turns a diverging NaN into + a finite 0.0 makes the forward look healthy while the backward is still NaN, + and the run checkpoints a broken model. ``-Inf`` and both signed zeros must + come out as ``+0.0``, never ``-0.0``. + """ + x, weight, bias, reference = _zero_weight_case("relu") + out, _mean, _rstd = torch.ops.scaffold_gn.group_norm( + x, 8, weight, bias, EPS, "relu", None + ) + assert torch.equal(out.cpu().view(torch.int32), reference.cpu().view(torch.int32)) + # ... and the control: without the fusion the same values pass through. + plain, _m, _r = torch.ops.scaffold_gn.group_norm( + x, 8, weight, bias, EPS, None, None + ) + assert plain[0, 0, 0, 0, 0].isnan() and plain[0, 1, 0, 0, 0].isinf() + + +@pytest.mark.gpu +def test_fused_relu_backward_gates_like_threshold_backward(): + """ReLU's backward is ``result <= 0 ? 0 : grad``, so a NaN passes. + + The kernel recomputes the pre-activation and must gate with the same + complement: ``pre > 0 ? dy : 0`` reads identically on every finite value and + silently zeroes the NaN lane, which is the backward half of the same defect. + """ + x, weight, bias, _reference = _zero_weight_case("relu") + grad_out = torch.ones_like(x) + + weight = weight.requires_grad_(True) + bias = bias.requires_grad_(True) + xg = x.clone().requires_grad_(True) + reference = F.relu(F.group_norm(xg, 8, weight, bias, EPS)) + reference.backward(grad_out) + ref_dbias = bias.grad.detach().clone() + ref_dx = xg.grad.detach().clone() + + weight.grad = bias.grad = xg.grad = None + triton_group_norm(xg, 8, weight, bias, EPS, "relu").backward(grad_out) + + # d_bias counts exactly the elements whose gradient the gate let through. + assert ref_dbias[0].item() == 64, "the reference gated the NaN lane off" + assert torch.equal(bias.grad.cpu(), ref_dbias.cpu()) + assert torch.equal(xg.grad.cpu(), ref_dx.cpu()) From c6b6725e17562e5dffef48107de75baeb1cd5d2c Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Wed, 5 Aug 2026 06:41:46 -0700 Subject: [PATCH 2/8] Route the UNet GroupNorm through the Triton kernel FastGroupNorm becomes a three-rung ladder -- the native Triton kernel, a torch.compile'd functional, and stock eager -- with the routing decision made per call and the rejections tested in order: an explicit opt-out, a rung that has already failed in this process, an active torch.func transform, a non-CUDA tensor, a tensor subclass other than DistConv's DCTensor, a GPU the launch tables were not tuned on, and anything the kernel's own is_supported rejects. The ladder primitives live in _rungs.py, which the convolution ladder will share. The hardware guard is a preference rather than a correctness condition -- the kernel is correct anywhere Triton lowers it, and what is unknown elsewhere is only its speed -- so an explicit opt-in overrides it and nothing else on that list can be overridden at all. DoubleConv now asks the GroupNorm for its ReLU: the Triton kernel folds the activation into its forward store, removing a streaming pass worth 38% of the forward at the shapes that dominate the step. The nn.ReLU slots stay occupied by nn.Identity so nn.Sequential does not renumber its children and existing checkpoints keep loading. --- ScaFFold/unet/_rungs.py | 424 +++++ ScaFFold/unet/group_norm.py | 747 ++++++++- ScaFFold/unet/unet_parts.py | 35 +- .../rank_scripts/groupnorm_shards_2rank.py | 123 ++ tests/test_groupnorm.py | 1381 ++++++++++++++++- tests/test_groupnorm_wiring_edge.py | 1193 ++++++++++++++ tests/test_platform_guard.py | 501 ++++++ 7 files changed, 4256 insertions(+), 148 deletions(-) create mode 100644 ScaFFold/unet/_rungs.py create mode 100644 tests/helpers/rank_scripts/groupnorm_shards_2rank.py create mode 100644 tests/test_groupnorm_wiring_edge.py create mode 100644 tests/test_platform_guard.py diff --git a/ScaFFold/unet/_rungs.py b/ScaFFold/unet/_rungs.py new file mode 100644 index 0000000..acaaa8e --- /dev/null +++ b/ScaFFold/unet/_rungs.py @@ -0,0 +1,424 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""Primitives shared by the modules that run a fast kernel with a fallback. + +:mod:`ScaFFold.unet.group_norm` and :mod:`ScaFFold.unet.conv3d` both present the +same shape: a hand-written kernel tried first, the stock one behind it, a +per-rung failure latch and a per-module "this rung has served me" flag. The +pieces collected here are the ones that are *not* about either kernel and that +each carry a correction paid for once already: + +* :func:`_env_override` -- the opt-in/opt-out spelling, so both modules accept + the same words and warn about the same typos. +* :func:`_platform_declines` -- the hardware guard, which is the one routing + condition whose failure mode is a *correct* answer at the wrong speed. +* :func:`_replaying_a_forward` -- the "a backward is in flight" probe, with its + ``is_compiling()`` guard. +* :data:`_functorch_active` -- the ``torch.func`` probe, which is a routing + question and not a kernel defect. +* :func:`_dctensor_ops` -- DistConv resolved through ``sys.modules`` rather than + imported, so ScaFFold stays importable without it. +* :func:`_run_local` -- the ``DCTensor`` unwrap/rewrap, through DistConv's + autograd pair rather than a bare ``_tensor`` read. +* :func:`_warn_rung_failure` -- the fallback message, with its ``is_compiling()`` + guard. + +They live here rather than being copied because a copy drifts: every one of the +guards above was added in response to a specific measured failure, and the next +module to grow a ladder should inherit them rather than rediscover them. + +What is deliberately *not* here: the allowlist of exceptions a rung may fail +with. That is a property of the kernel behind the rung, closed at that kernel's +own boundary, and it has to be written next to the rung that uses it. +""" + +import logging +import os +import sys + +import torch + +logger = logging.getLogger(__name__) + + +def _env_override(name): + """Read boolean env var ``name``; ``None`` when unset or unparsable.""" + raw = os.environ.get(name) + if raw is None: + return None + value = raw.strip().lower() + if value in ("1", "true", "on", "yes"): + return True + if value in ("0", "false", "off", "no"): + return False + logger.warning( + f"Ignoring unrecognized {name}={raw!r}; " + "expected one of 1/0/true/false/on/off/yes/no" + ) + return None + + +# --------------------------------------------------------------------------- +# The hardware guard +# --------------------------------------------------------------------------- + +#: The GPU architecture every tuning table in both ladders was raced on, with +#: its feature suffixes stripped. ``gcnArchName`` reports +#: ``"gfx942:sramecc+:xnack-"`` here, and the suffixes are a property of how the +#: *build* was configured rather than of the silicon, so an exact string +#: comparison would decline the same chip under a different HIP build. Compared +#: against the part before the first colon for that reason. +TUNED_ARCH = "gfx942" + +#: ...and the compute-unit count, which is what makes this MI300A rather than +#: ``gfx942``. ``gfx942`` is three parts: the MI300A APU (228 CUs), the MI300X +#: (304) and the MI325X (304). The tables here were raced on the first one and +#: several of them are *written in terms of* its geometry -- +#: ``gather_gemm.candidate_configs`` defaults ``GROUP_M`` to 6 because that is +#: MI300A's XCD count, ``triton_group_norm._TUNED``'s largest entry pins a grid +#: of 228, and ``conv3d._policy_declines``'s small-``M`` rule is the shape of the +#: cliff where a GEMM stops filling 228 CUs. None of those is wrong on an +#: MI300X in the sense of returning bad numbers; all of them are simply about a +#: different machine. +#: +#: The CU count also declines a *partitioned* MI300A, which is the case an +#: arch-only test would silently accept and which is genuinely mistuned: CPX +#: mode presents one logical device per XCD, so both the 228 and the 6 above +#: become fiction while ``gcnArchName`` stays exactly the same. +TUNED_CU_COUNT = 228 + +#: Both ladders' opt-in switches, named together in the one message a user on +#: another device gets. The verdict is a fact about the *node*, not about +#: either kernel, so a user who has just discovered that the fast path is off +#: should not have to find the second switch separately. +_OVERRIDE_SWITCHES = "SCAFFOLD_CONV_TRITON=1 / SCAFFOLD_GROUPNORM_TRITON=1" + +#: ``device index -> (is the tuned platform, human description)``. Resolved on +#: the first eligible call per device and never again; see +#: :func:`_platform_verdict`. +_PLATFORM_VERDICTS = {} + +#: Device indices that have already produced their one message, per kind. +_PLATFORM_DECLINE_WARNED = set() +_PLATFORM_OVERRIDE_WARNED = set() + + +def _device_fingerprint(index): + """``(arch, cu_count, name)`` for CUDA/HIP device ``index``. + + Split out from :func:`_platform_verdict` for exactly one reason: it is the + **seam the tests replace**. Every interesting branch of this guard is the + one this node cannot take -- an MI300X, a partitioned MI300A, an NVIDIA + device, a driver that will not answer -- so the predicate has to be + injectable or those branches are shipped unexecuted. Keeping the whole + query (and nothing else) behind one function means a test substitutes a + tuple and exercises the real decision, the real caching and the real + message, rather than a parallel copy of them. + + ``gcnArchName`` exists only on a ROCm build; a CUDA build of torch has no + such attribute, so an NVIDIA device is described as an empty arch and + declines through the same clause an untuned AMD one does. That is the right + answer for a reason beyond tuning: the kernels' launch constraints are MFMA + constraints (``gather_gemm._MFMA_KDIM``, ``matrix_instr_nonkdim``) and mean + nothing on a device with no MFMA. + """ + props = torch.cuda.get_device_properties(index) + arch = getattr(props, "gcnArchName", "") or "" + return ( + arch.split(":")[0], + int(getattr(props, "multi_processor_count", 0)), + str(getattr(props, "name", "")), + ) + + +def _device_index(device): + """The integer this device is cached and named by. + + A tensor's ``device`` always carries an index, but a bare + ``torch.device("cuda")`` does not, and that one means "whichever is + current" -- which is what the kernel would launch on. + """ + return device.index if device.index is not None else torch.cuda.current_device() + + +def _platform_verdict(device): + """Whether ``device`` is the machine the tables were tuned on, cached. + + Returns ``(ok, description)`` and computes it **once per device index** for + the life of the process. Cached because it is asked on the routing path of + every convolution and every GroupNorm -- 40-odd calls per step -- and + ``get_device_properties`` is a driver query, not an attribute read; and + computed lazily rather than at import because a CPU-only run (the whole CPU + unit suite) must not initialize the GPU at all. The callers only reach here + after ``is_cuda``, so by then torch's CUDA state is already up. + + **Per device index, not per process.** ScaFFold pins one rank per GPU, so + in production this dictionary holds exactly one entry and the distinction is + invisible. It is keyed anyway because the question the guard is asking is + "will the kernel that is about to launch be running on the machine its + launch configuration was chosen for", and that is a property of the device + the tensor is on -- so on a node that exposes two different GPUs a + process-wide answer taken from device 0 would be wrong on one of them in + whichever direction hurts more. It costs a dictionary lookup on an int. + + A driver that will not answer is *not* the tuned platform: the guard's whole + job is to be sure, and "I could not find out" is not "yes". + """ + index = _device_index(device) + verdict = _PLATFORM_VERDICTS.get(index) + if verdict is None: + try: + arch, cus, name = _device_fingerprint(index) + except Exception as e: # a driver query has no correct answer to invent + arch, cus, name = "", 0, f"" + ok = arch == TUNED_ARCH and cus == TUNED_CU_COUNT + verdict = (ok, f"{name} (arch {arch or 'unknown'}, {cus} CUs)") + _PLATFORM_VERDICTS[index] = verdict + return verdict + + +def _reset_platform_cache(): + """Forget every cached verdict and every message already emitted. + + For tests only, and it exists because the cache above is process-global: + without it the first test to ask a question would fix the answer for every + later one, so a suite that checks the decline path and then the accept path + would pass while testing the first one twice. + """ + _PLATFORM_VERDICTS.clear() + _PLATFORM_DECLINE_WARNED.clear() + _PLATFORM_OVERRIDE_WARNED.clear() + + +def _platform_declines(device, override): + """``True`` when ``device`` is not the hardware these kernels were tuned for. + + What this is protecting against is **silent mistuning, not a crash**. Every + number that decides how these kernels launch was raced on one MI300A: the + convolution tile tables and ``matrix_instr_nonkdim``/``kpack``/ + ``waves_per_eu`` choices in ``triton_conv3d``, the ``GROUP_M = 6`` that is + MI300A's XCD count, and ``triton_group_norm._TUNED``, whose largest entry + names a grid of 228. Run somewhere else those are not *wrong answers*, they + are answers to a question about a different machine -- and there is no + mechanism downstream that would notice: + + * ``triton_conv3d.gather_gemm``'s docstring records that on gfx942 an + illegal MFMA configuration does not fail. It emits **zero** MFMA + instructions, drops to vector FMA, and returns correct results at a + fraction of the speed. On an architecture whose legality rules differ + from the ones ``ConvConfig.validate`` encodes, that is precisely the + failure available: a right answer, slowly, raising nothing. + * the ladders' fallback allowlist is ``triton.errors.TritonError`` and + deliberately nothing else, so there is no exception for it to catch even + in principle. + * every ``is_supported*`` predicate reads shape, dtype, layout and stride. + None of them reads the GPU, and none of them should: they are *capability* + predicates, and the kernels really are capable of computing this + convolution on other hardware. + + So the guard is a **preference, not a correctness condition**, and the code + says so by letting an explicit opt-in through. ``override`` is the caller's + tri-state ``_triton_override``: ``None`` (the default, and every production + run) means "on wherever it is safe", where this device is part of what + "safe" means; ``True`` is a human who has typed ``SCAFFOLD_CONV_TRITON=1`` + or called ``set_conv_triton_enabled(True)`` and is therefore asserting a + judgement about their own hardware, which is exactly the development case + this override exists for. ``False`` never reaches here -- the callers + decline on it first -- so the two states this function distinguishes are + "nobody said" and "somebody said yes". + + Both branches are loud, once per device: a decline that said nothing would + leave a user on an MI300X with a 1.26x regression and no thread to pull, and + an override that said nothing would leave every subsequent measurement on + that machine uncomparable with the tables it will be read against. + """ + ok, described = _platform_verdict(device) + if ok: + return False + index = _device_index(device) + if override is True: + _warn_platform_override(index, described) + return False + _warn_platform_decline(index, described) + return True + + +def _warn_platform_decline(index, described): + """The one message a user on untuned hardware gets. + + Once per device index -- not once per call, which at 40-odd routed + operations a step would be a log line every few milliseconds, and not + silence, which is the state this whole change exists to end. The message + names the device it found, the device it wanted, and both switches, because + a user who has just noticed the fast path is off is asking one question and + should not have to find the second ladder's control separately. + + ``is_compiling()`` for the reason :func:`_warn_rung_failure` documents: + Dynamo cannot trace a ``logging.Logger`` method, so a bare call here would + turn a *routing decision* into a hard error for a caller who wrapped the + forward in ``torch.compile(fullgraph=True)``. + """ + if index in _PLATFORM_DECLINE_WARNED: + return + _PLATFORM_DECLINE_WARNED.add(index) + if torch.compiler.is_compiling(): + return + logger.warning( + f"ScaFFold's Triton kernels are tuned for {TUNED_ARCH} with " + f"{TUNED_CU_COUNT} CUs (AMD Instinct MI300A); cuda:{index} is " + f"{described}. Using the fallback kernels there. The Triton kernels are " + "correct on other hardware -- every launch configuration in them was " + "chosen on that device, so what is unknown is their speed, and a " + f"mistuned launch reports nothing. Set {_OVERRIDE_SWITCHES} to use them " + "anyway." + ) + + +def _warn_platform_override(index, described): + """The message the override owes, once per device index. + + Deliberately not quiet. Every performance figure either ladder is read + against -- the 1.26x step, the block-list in ``conv3d._policy_declines``, + ``triton_group_norm``'s roofline percentages -- was measured on the device + this run is *not* on, so a number produced under this override is not + comparable with any of them, and a log line is the only place that fact can + be recovered from afterwards. + """ + if index in _PLATFORM_OVERRIDE_WARNED: + return + _PLATFORM_OVERRIDE_WARNED.add(index) + if torch.compiler.is_compiling(): + return + logger.warning( + f"Taking the Triton kernels on cuda:{index}, which is {described}, " + f"because they were explicitly enabled. They are tuned for " + f"{TUNED_ARCH} with {TUNED_CU_COUNT} CUs and nothing here has been " + "measured on this device: expect correct numbers and unknown speed, and " + "do not compare timings from this run against the tuned ones." + ) + + +def _replaying_a_forward(): + """``True`` while this thread is executing inside an autograd graph task. + + ``torch._C._current_graph_task_id()`` is ``-1`` outside a backward pass and + the running task's id inside one; it is the same signal + ``torch.utils.checkpoint`` keys its own recompute bookkeeping on + (``torch/utils/checkpoint.py``'s ``unpack_hook``). + + A module *forward* that runs while a backward is in flight is not a new + call: it is a checkpoint recompute (or a double backward) replaying a + forward that has already happened and whose saved tensors are already held. + That is the one place where quietly answering on a different rung than the + original forward used is not a fallback but a corruption -- the rungs do not + save interchangeable tensors, so the recompute's saved set no longer matches + the graph node that will consume it. For GroupNorm that shows up as + ``CheckpointError: Recomputed values ... have different metadata`` (measured, + and on one shape a GPU memory fault instead); for convolution under DistConv + it is worse, because the two rungs save tensors with *identical* shape, + dtype and device and differ only in whether slot 0 is the ``DCTensor`` + wrapper or its inner tensor -- which is precisely what + ``_default_meta_extractor`` does not compare, so the substitution succeeds + and fails later as an ``AttributeError`` from inside DistConv (measured, both + directions). See the callers. + + ``is_compiling()`` first, for the same reason :func:`_warn_rung_failure` + checks it: the probe below is a ``torch._C`` builtin returning an ``int``, + which Dynamo cannot trace ("Unsupported torch.* op returned non-Tensor"), + so a caller who wraps a ``forward`` in ``torch.compile(fullgraph=True)`` + would get a hard error where the fallback belongs. Dynamo folds it to + ``True`` at trace time, leaving ``False`` here as a constant -- which is + also the right answer: tracing is not replaying, and the recompute this + guards against runs with Dynamo disabled anyway + (``torch.utils.checkpoint``'s ``_run_fn_with_dynamo_disabled``). + """ + if torch.compiler.is_compiling(): + return False + task_id = getattr(torch._C, "_current_graph_task_id", None) + if task_id is None: # pragma: no cover - every supported torch has it + return False + return task_id() != -1 + + +#: ``True`` while a ``torch.func`` transform (``vmap``/``grad``/``jvp``) is on +#: the stack. A fast rung declines then: a functorch layer is a routing +#: question, not a kernel defect, and the stock kernel handles every transform. +#: Not merely a performance choice -- an ``is_supported``'s +#: ``is_contiguous(memory_format=...)`` raises outright under ``vmap`` +#: ("NYI: querying is_contiguous inside of vmap"), and neither hand-written op +#: has a batching rule -- so without this the modules are not the drop-in +#: replacements they claim to be for any caller using ``torch.func``. +_functorch_active = getattr(torch._C, "_are_functorch_transforms_active", lambda: False) + + +def _dctensor_ops(input): + """The ``distconv.distconv`` module when ``input`` is a DCTensor, else None. + + Resolved through ``sys.modules`` instead of an import: a DCTensor can only + exist if DistConv is already imported, and ScaFFold's model must stay + importable (and the CPU suite runnable) without DistConv installed. + """ + distconv = sys.modules.get("distconv.distconv") + if distconv is not None and isinstance(input, distconv.DCTensor): + return distconv + return None + + +def _run_local(input, distconv, kernel): + """Run ``kernel`` on a plain tensor, DCTensor in -> DCTensor out. + + ``distconv`` is ``None`` for a plain tensor, where this is just + ``kernel(input)``. For a ``DCTensor`` the unwrap goes through DistConv's + ``_ToTensor``/``_FromTensor`` autograd pair (``DCTensor.from_shard`` is the + public spelling of the latter; there is no public unwrap yet -- upstream + ask) rather than a bare ``input._tensor`` read: DistConv's own dispatch may + read ``_tensor`` directly because it runs *below* autograd, while this runs + above it and a bare read would sever the graph back to the producing + convolution. + + Note what this does *not* do: it does not consult the parallel strategy. + Whether running the kernel on the local shard alone is the same computation + DistConv's dispatch would have performed is the caller's question, and the + answer differs by operator -- see :class:`~ScaFFold.unet.conv3d.FastConv3d`, + where it is only true when nothing is actually sharded. + """ + if distconv is None: + return kernel(input) + local = distconv._ToTensor.apply(input) + return distconv.DCTensor.from_shard(kernel(local), input._parallel_strategy) + + +def _warn_rung_failure(what, error, fallback, env_var): + """Log a rung failure, without graph-breaking a compiled caller. + + Dynamo cannot trace ``logging.Logger`` methods ("Unsupported: logging.Logger + method not supported for non-export cases"), so a bare ``logger.warning`` + in a ladder's handler turns a *fallback* into a hard Dynamo error for + anyone who wraps that ``forward`` in ``torch.compile(fullgraph=True)`` -- + the one caller for whom the fallback matters most, since the failure it is + reacting to is usually a compile failure. ``is_compiling()`` is a Dynamo + intrinsic that folds to ``True`` at trace time, so the call below becomes + dead code inside a traced region and the fallback traces cleanly. The + latch itself is a global assignment, which Dynamo does replay, so the + fallback is still recorded -- only this message is dropped, and only for a + caller that is compiling the module's forward (nothing in ScaFFold does). + """ + if torch.compiler.is_compiling(): + return + logger.warning( + f"{what} failed ({type(error).__name__}: {error}); falling back to the " + f"{fallback} for modules that have not already used it. " + f"Set {env_var}=0 to skip this attempt entirely." + ) diff --git a/ScaFFold/unet/group_norm.py b/ScaFFold/unet/group_norm.py index 06d7407..fe0f510 100644 --- a/ScaFFold/unet/group_norm.py +++ b/ScaFFold/unet/group_norm.py @@ -12,44 +12,164 @@ # # SPDX-License-Identifier: (Apache-2.0) -"""GroupNorm with a ``torch.compile``d fast path on GPU. - -ATen's GroupNorm computes its per-group statistics with a kernel that launches -one workgroup per ``(batch, group)`` row. At this benchmark's defaults -(``local_batch_size=1``, ``group_norm_groups=8``) that is 8 workgroups, so on a -228-CU MI300A the normalization runs at a small fraction of achievable -bandwidth and dominates the step: measured 87 ms of a 187 ms step (47%) at -scale 7. Compiling the same functional GroupNorm hands the reduction to -Inductor, which tiles it across the whole device; the same measurement then -gives a 184.7 ms step at 100.7 ms, with GroupNorm down to ~7% of it. +"""GroupNorm with a Triton fast path and a ``torch.compile``d one behind it. + +Three kernels, tried in order, all of them producing the same numbers: + +1. **Native channels-last Triton** (:mod:`ScaFFold.unet.triton_group_norm`), + whenever that module's ``is_supported`` accepts the input. Production runs + set ``PYTORCH_MIOPEN_SUGGEST_NHWC=1``, under which every convolution emits + ``channels_last_3d`` -- and every *stock* GroupNorm (eager or Inductor) reads + that layout through the logical NCDHW order, a strided gather, and returns a + contiguous tensor, breaking the layout chain at all 22 call sites of the + forward. The Triton kernel is NDHWC in and NDHWC out and is 6.5x faster than + the compiled kernel on that input; with the ReLU fused it takes another 38% + off the forward. See that module's docstring for the measurements. +2. **``torch.compile``d ``F.group_norm``**, for inputs the Triton kernel does + not serve (contiguous NCDHW, non-5-D, unsupported dtypes) and as the landing + place if the Triton path ever raises. ATen's own kernel launches one + workgroup per ``(batch, group)`` row -- 8 of them at this benchmark's + defaults -- so on a 228-CU MI300A it runs at a small fraction of achievable + bandwidth: measured 87 ms of a 187 ms step (47%) at scale 7, against 7% of a + 184.7 ms step once Inductor tiles the reduction across the device. +3. **Stock eager ``F.group_norm``**, which is what every rejection falls back + to and what defines the semantics the other two must match. ``FastGroupNorm`` is a drop-in ``nn.GroupNorm``: same parameters, same names, same shapes, same numerics -- only the kernel differs, so checkpoints are interchangeable in both directions with any other GroupNorm-based build. The -compiled path is used only when it is safe and worthwhile, and every rejection -falls back to stock eager ``F.group_norm``: - -* non-CUDA tensors (the CPU test suite never pays compile latency), -* tensor subclasses such as DistConv's ``DCTensor``, whose ``__torch_dispatch__`` - wrapper Dynamo cannot trace, -* an already-compiled enclosing region (the functional call inlines instead), -* an explicit opt-out via ``SCAFFOLD_GROUPNORM_COMPILE=0``, -* any failure inside ``torch.compile`` -- logged once, then eager forever after. - -Determinism: the compiled kernels are bitwise reproducible. Two separate -processes running three fwd+bwd+Adam steps of the scale-7 UNet under -``more_determinism`` (``use_deterministic_algorithms(True, warn_only=True)``, -``cudnn.benchmark=False``, fixed seeds) hash identically with the compiled path, -exactly as they do with the eager one, so no determinism gate is needed. +one addition is the optional fused ``activation`` (see below), which adds no +state either. + +All three return the input's memory format, so the rungs are interchangeable in +everything a caller can observe (see :func:`_match_memory_format`). + +Routing rejections, in the order they are tested: + +* an explicit opt-out via ``SCAFFOLD_GROUPNORM_TRITON=0`` / + ``SCAFFOLD_GROUPNORM_COMPILE=0``, +* a rung that has failed in this process, for every module that has not already + had a call served by it (see "Latches" below), +* an active ``torch.func`` transform -- a ``vmap``/``grad``/``jvp`` layer is a + routing miss, not a kernel failure, and the stock kernel handles it, +* non-CUDA tensors -- the CPU test suite pays neither compile latency nor the + Triton import, +* tensor subclasses, whose ``__torch_dispatch__`` wrappers have unknown + semantics -- except DistConv's ``DCTensor``, which is unwrapped to its local + shard around both fast kernels (see ``FastGroupNorm``), +* for the Triton kernel, a GPU other than the one its launch tables were tuned + on (``gfx942`` with 228 CUs -- an MI300A). Alone on this list that one is a + *preference*: the kernel is correct anywhere Triton lowers it, and what is + unknown elsewhere is only its speed, which is why an explicit opt-in overrides + it and nothing else here can be overridden at all. See + ``_rungs._platform_declines``, +* for the Triton kernel, anything its ``is_supported`` rejects (a layout, dtype, + degenerate shape or affine-parameter dtype it does not serve); for the + compiled one, an already-compiled enclosing region (the functional call + inlines instead). + +Latches +======= +Both fast rungs are optimizations, never correctness requirements: a broken +Triton install must degrade a multi-node run, not kill it. So a *kernel* +failure is caught, logged once and retried on the next rung down. + +"A kernel failure" is an allowlist, not the absence of one. The Triton rung is +caught on ``triton_group_norm.TritonKernelError``, which that module raises for +anything its launch region produces; the compiled rung on +``torch._dynamo.exc.TorchDynamoException``, the root of every Dynamo and +Inductor compile failure, *plus* ``FailOnRecompileLimitHit``, which despite the +name derives from ``Exception`` and not from that root (see +:func:`_compiled_kernel_failures`). Everything else propagates -- saved-tensor pack +hooks, ``torch.utils.checkpoint``'s recompute control flow, a user's offloading +hook, ``torch.OutOfMemoryError``, an error from a shape the kernel mishandles +badly enough to corrupt the graph. The previous shape of this code caught +``Exception`` and re-raised a denylist of framework mechanisms, which was wrong +twice (``_StopRecomputationError``, then ``CheckpointError``): the set of things +torch may raise through a forward is open, the set of ways a kernel can be +broken is closed at its own boundary. Both allowlisted exceptions are also +raised strictly *before* their rung saves anything for backward (the Triton op +saves in ``_setup_context``, after its launch region; a Dynamo/Inductor failure +is a compile-time failure, before any execution), so the retry cannot double-fire +saved-tensor hooks. + +A failure latches the rung off **for modules that have never had a call served +by it**. A module that has already run on a rung keeps it. That is not a +performance nicety: ``torch.utils.checkpoint``'s non-reentrant recompute +compares the metadata of every tensor the recomputed forward saves against the +originals, and the three rungs intrinsically save *different tensors* -- Triton +saves ``(input, weight, bias, mean, rstd)``, the other two +``(input, weight, mean, rstd, relu_output)``. A latch that flipped between a +block's forward and its recompute would therefore kill the step with +``CheckpointError: Recomputed values ... have different metadata``, which is the +exact opposite of the contract above (measured; matching the output memory +format is *not* sufficient on its own). Keeping a proven rung pins each +module's choice for the life of the process, so forward and recompute always +agree. + +The same reasoning bounds the *fallback* itself, which the latch alone does not: +a proven module still has to answer the call its rung just failed, and answering +it eagerly is exactly the flip the paragraph above forbids -- if that call is a +checkpoint recompute. So the fallback is declined in the one case where it +would corrupt rather than degrade: a module proven on the rung, failing while an +autograd graph task is in flight (:func:`_replaying_a_forward`), re-raises. +Every other failure -- and in particular every *first* failure, which is what a +broken Triton install, an unwritable Inductor cache or a missing compiler +produce -- still degrades, which is where the "must not kill a multi-node run" +contract actually lives. + +Note that a latch is process-local: under DDP one rank can end up running a +different kernel from its peers. All three kernels agree to fp32 rounding, not +bitwise, so a rank that latches shifts that rank's gradients and therefore the +all-reduced ones -- a real (measured) change to the job's trajectory, and a +2.1x straggler besides. That is the price of degrading instead of dying, but it +is why the latch is as narrow as it is, and why ``torch.OutOfMemoryError`` -- +transient by nature, and no cheaper on any other rung -- does not latch anything +at all. :func:`set_triton_enabled` / :func:`set_compile_enabled` with ``True`` +clear the latch, which is the supported way to retry after a transient failure. + +Determinism: all three kernels are bitwise reproducible. Two separate processes +running the scale-7 UNet under ``more_determinism`` +(``use_deterministic_algorithms(True, warn_only=True)``, ``cudnn.benchmark=False``, +fixed seeds) hash identically with the Triton path, the compiled path and the +eager one alike, so no determinism gate is needed. The Triton kernel's grid, +split count and tile sizes are pure functions of the shape and it uses no float +atomics, which is what buys that. + +Fused activation +================ +Every GroupNorm in the UNet is immediately followed by a ReLU, and the Triton +kernel can fold that into its forward store for free (it is store-bound) while +removing a whole streaming pass -- 38% of the forward at the shapes that +dominate. ``FastGroupNorm(..., activation="relu")`` therefore *always* applies +the ReLU: fused inside the Triton kernel where that path is taken, and as an +explicit in-place ``F.relu`` on the compiled and eager paths. ``DoubleConv`` +consequently holds an ``nn.Identity`` where its ``nn.ReLU`` used to be, so the +positional keys of its ``nn.Sequential`` -- and therefore every checkpoint -- +are unchanged (neither module has parameters or buffers). Correctness holds on +every path; the fusion is purely an optimization inside the module. """ import logging -import os import torch import torch.nn as nn import torch.nn.functional as F +# The pieces every rung ladder needs, kept in one place so the corrections each +# of them carries cannot drift between this module and ScaFFold.unet.conv3d. +# Imported by name so they stay module attributes here, which is what the tests +# monkeypatch and what this module's own code resolves through. +from ._rungs import ( + _dctensor_ops, + _env_override, + _functorch_active, + _platform_declines, + _replaying_a_forward, + _run_local, + _warn_rung_failure, +) + logger = logging.getLogger(__name__) #: Opt-out (``0``/``false``/``off``/``no``) or explicit opt-in (``1``/``true``/ @@ -57,12 +177,31 @@ #: is safe", which is what every production run wants. COMPILE_ENV_VAR = "SCAFFOLD_GROUPNORM_COMPILE" +#: The same, for the native channels-last Triton kernel, which is tried first. +#: Same spellings, same "unset means on wherever it is safe" default -- the +#: whole point of the kernel is that production takes it. The explicit opt-in +#: is more than that default written out: it additionally overrides the hardware +#: guard, which is the one routing condition here that is a preference. See +#: :func:`set_triton_enabled`. +TRITON_ENV_VAR = "SCAFFOLD_GROUPNORM_TRITON" + +#: Activations this module can apply after normalizing. Must stay a subset of +#: ``triton_group_norm.SUPPORTED_ACTIVATIONS`` (pinned by a test); spelled out +#: here rather than imported so that constructing a module -- or running the +#: whole CPU suite -- never imports the kernel module. +SUPPORTED_ACTIVATIONS = (None, "relu") + #: Dynamo caches one entry per distinct guard set on the traced function. A #: UNet presents one entry per distinct activation shape (5 at scale 7) times #: grad-enabled/no-grad (training vs. evaluation), i.e. 10 -- above the stock #: limit of 8, which would silently drop the whole model back to eager mid-run. -#: The traced function is a single ``F.group_norm`` call, so the extra entries -#: cost only their one-time compilation. +#: ``activation_checkpointing`` on a ``DCTensor`` doubles that again: the +#: recompute reaches this module with ``__torch_function__`` subclass handling +#: *disabled* (DistConv's backward runs below it), which is part of Dynamo's +#: ``GLOBAL_STATE`` guard, so the recomputed forward misses every entry the +#: original forward built and compiles a second set beside it -- 20 for the same +#: 5 shapes (measured). The traced function is a single ``F.group_norm`` call, +#: so the extra entries cost only their one-time compilation. _MIN_RECOMPILE_LIMIT = 64 # Lazily built on the first eligible forward: importing ScaFFold must not drag @@ -76,25 +215,27 @@ # by set_compile_enabled(). _compile_override = None +# The triton_group_norm module, imported on the first CUDA forward. Importing +# it registers two dispatcher ops, and a CPU-only run must pay neither that nor +# the `triton` import the module itself defers to its first launch. +_triton_module = None -def _env_override(): - """Read ``SCAFFOLD_GROUPNORM_COMPILE``; ``None`` when unset or unparsable.""" - raw = os.environ.get(COMPILE_ENV_VAR) - if raw is None: - return None - value = raw.strip().lower() - if value in ("1", "true", "on", "yes"): - return True - if value in ("0", "false", "off", "no"): - return False - logger.warning( - f"Ignoring unrecognized {COMPILE_ENV_VAR}={raw!r}; " - "expected one of 1/0/true/false/on/off/yes/no" - ) - return None +# The ladder's two allowlists, resolved on first use of the rung they guard -- +# importing either provider (the kernel module, torch._dynamo) is exactly what +# the lazy _get_* helpers exist to avoid paying for on a CPU-only run. +_TRITON_KERNEL_FAILURES = None +_COMPILED_KERNEL_FAILURES = None + +# Set once if the Triton kernel raises; the compiled path is used from then on. +_triton_failed = False +# None = decide per tensor; True/False = forced by SCAFFOLD_GROUPNORM_TRITON or +# by set_triton_enabled(). +_triton_override = None -_compile_override = _env_override() + +_compile_override = _env_override(COMPILE_ENV_VAR) +_triton_override = _env_override(TRITON_ENV_VAR) def set_compile_enabled(enabled): @@ -103,12 +244,49 @@ def set_compile_enabled(enabled): ``None`` restores the default, which is the environment variable if set and otherwise "compile wherever it is safe". Forcing it on does not override the device and tensor-subclass checks -- those are correctness conditions, - not preferences. Returns the previous setting so callers (tests) can - restore it. + not preferences -- but it *does* clear a failure latch: an explicit "use + this rung" is the supported way to retry after a transient failure, and + leaving the latch set would make this function silently do nothing. + Returns the previous setting so callers (tests) can restore it. """ - global _compile_override + global _compile_override, _compile_failed previous = _compile_override - _compile_override = _env_override() if enabled is None else bool(enabled) + _compile_override = ( + _env_override(COMPILE_ENV_VAR) if enabled is None else bool(enabled) + ) + if _compile_override is True: + _compile_failed = False + return previous + + +def set_triton_enabled(enabled): + """Force the Triton path on (``True``) or off (``False``). + + The exact counterpart of :func:`set_compile_enabled`: ``None`` restores the + default (``SCAFFOLD_GROUPNORM_TRITON`` if set, otherwise "wherever + ``is_supported`` accepts"), and the previous setting is returned so tests + can restore it. + + Forcing it on clears any failure latch, and overrides the **hardware + guard**, and overrides nothing else -- not the device, subclass or + ``is_supported`` checks, which are correctness conditions. The guard is on + the other side of that line because its failure mode is: an untuned GPU + gives the right numbers at a speed nobody has measured, where an unsupported + input would give the wrong ones. Enabling this rung explicitly (here or as + ``SCAFFOLD_GROUPNORM_TRITON=1``) is a developer asserting a judgement about + their own hardware, which is what the override is for, and it says so in the + log. See :func:`~ScaFFold.unet._rungs._platform_declines`. + + ``None`` deliberately does *not* clear the latch: it restores a preference, + it does not assert that the kernel works again. + """ + global _triton_override, _triton_failed + previous = _triton_override + _triton_override = ( + _env_override(TRITON_ENV_VAR) if enabled is None else bool(enabled) + ) + if _triton_override is True: + _triton_failed = False return previous @@ -118,12 +296,25 @@ def _group_norm(input, num_groups, weight, bias, eps): def _raise_recompile_limit(): - """Lift Dynamo's per-function recompile cap to cover every UNet GN shape. + """Lift Dynamo's *global* recompile cap to cover every UNet GN shape. Only ever raises it, so a caller that deliberately set a larger limit keeps theirs -- but note the converse: a limit deliberately set *smaller* than ours is clobbered up to ``_MIN_RECOMPILE_LIMIT``. ``cache_size_limit`` is the older spelling of ``recompile_limit``; set whichever exists. + + This is the *portable* half of the mitigation and, on its own, not a + sufficient one: ``torch._dynamo.config`` stores user overrides in a + ``ContextVar`` (``torch/utils/_config_module.py``: "User overrides are + thread-local"), so an assignment made here is invisible to every other + thread, which keeps reading the stock default of 8. That matters because + ``torch.utils.checkpoint``'s non-reentrant recompute runs inside the + backward pass, i.e. on the autograd engine's device worker thread, and a + recompute that has to compile -- which it does on a ``DCTensor``, see + ``_MIN_RECOMPILE_LIMIT`` -- would hit 8 there no matter what this function + wrote on the main thread. :func:`_compile_group_norm` therefore also asks + ``torch.compile`` for a per-region limit, which Dynamo applies on whichever + thread is compiling. """ config = torch._dynamo.config for name in ("recompile_limit", "cache_size_limit"): @@ -132,6 +323,29 @@ def _raise_recompile_limit(): setattr(config, name, _MIN_RECOMPILE_LIMIT) +def _compile_group_norm(): + """``torch.compile`` :func:`_group_norm` with a thread-proof recompile cap. + + ``recompile_limit=`` is the per-region spelling of the cap: Dynamo applies + it with ``config.patch()`` around the compile itself, on whatever thread + that compile happens on, which is the only spelling that survives the + autograd worker thread (see :func:`_raise_recompile_limit`). Older torches + have no such keyword -- there the global assignment is all there is, and the + checkpoint-recompute case is simply out of reach. + """ + try: + return torch.compile( + _group_norm, + dynamic=False, + fullgraph=True, + recompile_limit=_MIN_RECOMPILE_LIMIT, + ) + except TypeError: + # A torch too old for the keyword: still compile, because the rung is + # worth far more than the one configuration the keyword rescues. + return torch.compile(_group_norm, dynamic=False, fullgraph=True) + + def _get_compiled_group_norm(): """Build (once) the compiled functional GroupNorm shared by every module. @@ -145,21 +359,184 @@ def _get_compiled_group_norm(): global _compiled_group_norm if _compiled_group_norm is None: _raise_recompile_limit() - _compiled_group_norm = torch.compile(_group_norm, dynamic=False, fullgraph=True) + _compiled_group_norm = _compile_group_norm() return _compiled_group_norm -def _use_compiled(input): - """Whether this particular input should take the compiled path.""" - if _compile_failed or _compile_override is False: +def _get_triton_module(): + """Import (once) :mod:`ScaFFold.unet.triton_group_norm`. + + Deferred rather than imported at the top of this file: that module registers + two dispatcher ops and builds an autograd formula at import time, and a run + that never reaches the GPU (the whole CPU unit suite) must not pay for it. + Only ever called after the input has been shown to be a CUDA tensor, which + is also what keeps ``import triton`` -- which that module defers again, to + its first kernel launch -- out of a CPU-only process entirely. + """ + global _triton_module + if _triton_module is None: + from . import triton_group_norm + + _triton_module = triton_group_norm + return _triton_module + + +def _triton_kernel_failures(): + """The ladder's allowlist for the Triton rung: exactly ``TritonKernelError``. + + The kernel module raises it for every failure of its own launch region -- + a missing or mismatched ``triton``, an unwritable JIT cache, a compile + error, a bad launch -- and for nothing else, so this catches "the kernel is + broken" without also catching the framework mechanisms that legitimately + raise through a forward. See that class's docstring for what is + deliberately left untagged (``OutOfMemoryError``, contract violations). + + Resolved separately from :func:`_get_triton_module` so the except clause is + still available when the thing that failed *is* the module lookup. An + empty tuple (no kernel module at all) means "catch nothing": the ladder + then re-raises, which is right, because with no kernel module there is + nothing that could have failed inside one. + """ + global _TRITON_KERNEL_FAILURES + if _TRITON_KERNEL_FAILURES is None: + try: + from .triton_group_norm import TritonKernelError + + _TRITON_KERNEL_FAILURES = (TritonKernelError,) + except ImportError: # pragma: no cover - the module is in-tree + _TRITON_KERNEL_FAILURES = () + return _TRITON_KERNEL_FAILURES + + +def _compiled_kernel_failures(): + """The compiled rung's allowlist: every Dynamo and Inductor compile failure. + + ``torch._dynamo.exc.TorchDynamoException`` is the root of ``Unsupported`` + (``fullgraph=True`` met something untraceable), ``BackendCompilerFailed`` + and its ``InductorError`` subclass (the backend, and therefore also an + unwritable Inductor cache or a broken C++/Triton toolchain), and + ``InternalTorchDynamoError``. + + ``FailOnRecompileLimitHit`` -- raised when a frame needs more cache entries + than the recompile limit allows, which under ``fullgraph=True`` is a hard + error rather than a drop to eager -- is *not* under that root: it derives + straight from ``Exception`` (``torch/_dynamo/exc.py``), so catching only + ``TorchDynamoException`` lets it kill the run. It is named separately + rather than assumed, and only added when it really is outside the root, so + a torch that later reparents it does not produce a duplicate entry. + + All of these are raised while *compiling*, i.e. before the compiled callable + has executed or saved anything, which is what makes the fallback safe to + retry. + + Resolved on demand and cached: importing ``torch._dynamo`` is precisely the + cost :func:`_get_compiled_group_norm` defers. An empty tuple (a torch + without the module) means "catch nothing", which fails loudly rather than + silently swallowing. + """ + global _COMPILED_KERNEL_FAILURES + if _COMPILED_KERNEL_FAILURES is None: + try: + import torch._dynamo.exc as dynamo_exc + except ImportError: # pragma: no cover - torch always ships it + _COMPILED_KERNEL_FAILURES = () + else: + failures = [dynamo_exc.TorchDynamoException] + limit_hit = getattr(dynamo_exc, "FailOnRecompileLimitHit", None) + if isinstance(limit_hit, type) and not issubclass( + limit_hit, dynamo_exc.TorchDynamoException + ): + failures.append(limit_hit) + _COMPILED_KERNEL_FAILURES = tuple(failures) + return _COMPILED_KERNEL_FAILURES + + +# Set once if a predicate raised while deciding; see _use_triton. +_predicate_warned = False + + +def _use_triton(input, num_groups, weight, bias, activation, proven=False): + """Whether this particular input should take the native Triton kernel. + + ``proven`` is the caller's "this module has already had a call served by + this rung", which keeps a proven module on it even after a *global* latch; + see the module docstring's "Latches". + + Ordered so that the cheap local tests come first and the module import last: + a CPU tensor is rejected before ``_get_triton_module`` is ever called. + """ + if _triton_override is False: + return False + if _triton_failed and not proven: + return False + if _functorch_active(): + return False + # Same policy as _use_compiled: an unknown __torch_dispatch__ wrapper has + # unknown semantics and keeps the stock kernel. is_supported() would accept + # one (it only asks isinstance), so this check is load-bearing here, not a + # copy for symmetry. DistConv's DCTensor never reaches it -- forward() + # unwraps to the local shard first. + if type(input) is not torch.Tensor: + return False + if not input.is_cuda: + return False + # ...and the GPU the kernel's launch tables were built on. ``_TUNED``'s + # entries were raced on one 228-CU MI300A and its largest one names a grid + # of exactly 228; elsewhere they are answers about a different machine, and + # nothing downstream would notice, because a mistuned launch is a *correct* + # answer at an unmeasured speed. Shared with the convolution ladder, which + # is protecting the same kind of thing -- see ``_rungs._platform_declines``, + # including why an explicit opt-in overrides this and not the checks around + # it. Cached per device: a dictionary lookup after the first call. + if _platform_declines(input.device, _triton_override): + return False + # is_supported() is cheap and side-effect free: a handful of attribute reads + # and one stride check, no allocation, no launch, no triton import. The + # broad catch is right *here* and nowhere else in this module: a predicate + # that cannot answer has a correct answer available ("no"), it has done no + # work anyone can observe, and the failure is a routing miss rather than a + # broken kernel -- so it must not latch the rung off, which is what letting + # it fall into the ladder's handler used to do. + try: + return _get_triton_module().is_supported( + input, num_groups, weight, bias, activation + ) + except Exception as e: + _warn_once_about_the_predicate(e) + return False + + +def _warn_once_about_the_predicate(error): + """Log the first ``is_supported`` failure; a repeat would log per call.""" + global _predicate_warned + if _predicate_warned: + return + _predicate_warned = True + logger.warning( + f"Triton GroupNorm routing check failed ({type(error).__name__}: " + f"{error}); using the stock kernel for inputs like this one. This is a " + "routing miss, not a kernel failure, so nothing is latched off." + ) + + +def _use_compiled(input, proven=False): + """Whether this particular input should take the compiled path. + + ``proven`` has the same meaning as in :func:`_use_triton`. + """ + if _compile_override is False: + return False + if _compile_failed and not proven: + return False + # Dynamo cannot trace a functorch layer either, and under fullgraph=True + # that is an exception rather than a graph break. + if _functorch_active(): return False - # Tensor subclasses (DistConv's DCTensor) route their ops through - # __torch_dispatch__, which Dynamo cannot trace; eager keeps the wrapper's - # semantics -- including which of its outputs come back wrapped -- exactly - # as they are today. worker.py wraps activations in DCTensor even at - # dc_num_shards=[1,1,1], so this fast path engages once that wrap is - # skipped for the unsharded case (or whenever the model is driven with - # plain tensors, as the tests and the standalone benchmarks do). + # Tensor subclasses route their ops through __torch_dispatch__, which + # Dynamo cannot trace. DistConv's DCTensor never reaches this check -- + # forward() peeks at its local shard instead -- so anything rejected here + # is an unknown wrapper, and eager keeps its semantics exactly as they + # are today. if type(input) is not torch.Tensor: return False # CPU GroupNorm is not the bottleneck and compiling it would put a @@ -172,33 +549,243 @@ def _use_compiled(input): return True +def _match_memory_format(out, reference): + """Give ``out`` ``reference``'s memory format, copying only if it differs. + + ``F.group_norm`` -- eager or Inductor-compiled -- reads a + ``channels_last_3d`` input through the logical NCDHW order and returns a + *contiguous* tensor, which is the layout break this whole module exists to + avoid: with ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` every convolution both sides + of it wants channels-last, so one fallback re-breaks the chain for the rest + of the network. One relayout of the GroupNorm output is far cheaper than + the transposes the following convolutions would otherwise insert, and it + makes the three rungs agree on everything a caller can observe rather than + only on the values. + + Free on the Triton rung (already channels-last) and on any contiguous input + (nothing to do); one copy on a fallback from a channels-last input, which is + the only case that reaches the copy at all. + """ + if reference.dim() != 5: + # is_contiguous(memory_format=channels_last_3d) is only defined for 5-D. + return out + if _functorch_active(): + # "NYI: querying is_contiguous inside of vmap for memory_format other + # than torch.contiguous_format" -- and a functorch transform has no + # layout chain to preserve anyway, since both fast rungs decline it. + return out + if not reference.is_contiguous(memory_format=torch.channels_last_3d): + return out + if out.is_contiguous(memory_format=torch.channels_last_3d): + return out + return out.contiguous(memory_format=torch.channels_last_3d) + + class FastGroupNorm(nn.GroupNorm): - """``nn.GroupNorm`` that runs its GPU forward through ``torch.compile``. + """``nn.GroupNorm`` with a Triton GPU kernel and an optional fused ReLU. Identical state: ``weight``/``bias`` of shape ``(num_channels,)``, no buffers, so state dicts are interchangeable with plain ``nn.GroupNorm`` - in both directions. + in both directions. ``activation`` is a plain Python attribute, not a + submodule or a buffer, so setting it does not add a key either. + + ``activation="relu"`` makes this module's forward *always* apply a ReLU -- + fused into the Triton kernel's store where that path is taken, and as an + explicit in-place ``F.relu`` on the compiled and eager paths. The + correctness of the model therefore does not depend on which kernel runs; + only the number of memory passes does. + + DistConv's ``DCTensor`` gets the fast kernels too, by unwrapping to the + local shard in front of them rather than by letting the op dispatch through + the wrapper. Both would work -- the Triton kernel is a real dispatcher op, + so ``DCTensor.__torch_dispatch__`` would intercept it, unwrap, run and + rewrap on its own -- but the explicit unwrap is what this module already + does for the compiled kernel, and it is better here for three reasons. + (1) It keeps the subclass policy in one place: ``is_supported`` accepts any + ``torch.Tensor`` *instance*, so relying on dispatch would silently extend + the fast path to every unknown wrapper subclass, which today keeps the + stock kernel. (2) The eligibility predicates then examine the tensor the + kernel will actually touch -- its dtype, device, strides and shape -- rather + than a wrapper's mirrored metadata. (3) The Triton and compiled paths share + one unwrap and one fallback ladder instead of needing two shapes of code, + and a Triton failure can be retried on the compiled kernel without a second + round trip through the wrapper. Semantics are unchanged either way: + DistConv's generic ``__torch_dispatch__`` has no GroupNorm-specific + handling, so statistics are per-shard and no communication happens at any + shard count, exactly as before. """ - def forward(self, input): + #: Class-level defaults, so that an instance restored from a *module* + #: pickle written before these attributes existed (``torch.save(model)`` + #: rather than a state dict) still runs. ``nn.Module.__setstate__`` + #: replaces ``__dict__`` wholesale, so anything only ever set in + #: ``__init__`` is simply missing on such an instance. + activation = None + + #: Per-module "a call has been served by this rung". A global latch does + #: not demote a module that has one, which is what keeps a checkpointed + #: block's forward and its recompute on the same rung; see the module + #: docstring's "Latches". Plain attributes, so they are not parameters, + #: buffers or state-dict keys. + _triton_ok = False + _compiled_ok = False + + def __init__( + self, + num_groups, + num_channels, + eps=1e-5, + affine=True, + device=None, + dtype=None, + activation=None, + ): + if activation not in SUPPORTED_ACTIVATIONS: + raise ValueError( + f"activation must be one of {SUPPORTED_ACTIVATIONS}, got {activation!r}" + ) + super().__init__( + num_groups, num_channels, eps=eps, affine=affine, device=device, dtype=dtype + ) + self.activation = activation + + def extra_repr(self): + base = super().extra_repr() + if self.activation is None: + return base + return f"{base}, activation={self.activation}" + + def _activate(self, out): + """Apply the activation on the two paths that cannot fuse it. + + In place, which is what the ``nn.ReLU(inplace=True)`` this module + absorbed did: ``out`` is a freshly allocated GroupNorm output with no + other consumer, and GroupNorm's backward reads its *input*, never its + output, so overwriting it is safe for autograd as well as for memory. + + Validated *here* rather than only in ``__init__``: ``activation`` is a + plain attribute, so it can be assigned after construction, and the + Triton rung would then fuse an activation this method silently skipped + -- i.e. the network's function would depend on its input's memory + format. This is also the guard that makes adding a third activation to + ``SUPPORTED_ACTIVATIONS`` a loud failure until it is implemented here. + """ + activation = self.activation + if activation is None: + return out + if activation == "relu": + return F.relu(out, inplace=True) + raise ValueError( + f"activation must be one of {SUPPORTED_ACTIVATIONS}, got " + f"{activation!r}; this rung cannot apply it" + ) + + def _triton_forward(self, local): + """The native channels-last kernel, with the activation fused in.""" + return _get_triton_module().triton_group_norm( + local, self.num_groups, self.weight, self.bias, self.eps, self.activation + ) + + def _compiled_forward(self, local): + return self._activate( + _match_memory_format( + _get_compiled_group_norm()( + local, self.num_groups, self.weight, self.bias, self.eps + ), + local, + ) + ) + + def _eager_forward(self, input): # super().forward() is the stock kernel; deferring to it keeps the eager - # path identical to nn.GroupNorm's by construction. - if not _use_compiled(input): - return super().forward(input) - global _compile_failed + # path identical to nn.GroupNorm's (plus the ReLU and the relayout) by + # construction. + return self._activate(_match_memory_format(super().forward(input), input)) + + def forward(self, input): + global _compile_failed, _triton_failed + + distconv = _dctensor_ops(input) + # The eligibility checks look at the local shard for a DCTensor (the + # peek is a plain attribute read, no autograd involvement) and at the + # tensor itself otherwise. + local_view = input._tensor if distconv is not None else input + + if _use_triton( + local_view, + self.num_groups, + self.weight, + self.bias, + self.activation, + proven=self._triton_ok, + ): + triton_failures = _triton_kernel_failures() + try: + out = _run_local(input, distconv, self._triton_forward) + except triton_failures as e: + # A broken or mismatched Triton install, an unwritable JIT cache + # or a shape the kernel mishandles must cost speed, not a + # multi-node run. GroupNorm is pure and the kernel raises this + # only from its launch region -- before it has saved anything -- + # so retrying the same call on the compiled kernel below is + # safe, and the compiled kernel, not eager, is the right landing + # place: it is still ~10x the stock one. + # + # Logged on the latch's False->True edge only. A module that + # has already used the rung keeps trying it (that is what pins + # a checkpointed block to one rung), so a persistently broken + # kernel would otherwise warn once per call for the rest of the + # run; clearing the latch re-arms the message. + first = not _triton_failed + _triton_failed = True + if first: + _warn_rung_failure( + "Triton GroupNorm", e, "compiled kernel", TRITON_ENV_VAR + ) + # ... with one exception, shared with the compiled rung below + # and explained there: a module already proven on this rung must + # not be answered from a different one while a backward is + # replaying its forward. + if self._triton_ok and _replaying_a_forward(): + raise + else: + # Only written once: nn.Module.__setattr__ is not free, and + # after the first success this reads a class attribute. + if not self._triton_ok: + self._triton_ok = True + return out + + if not _use_compiled(local_view, proven=self._compiled_ok): + return self._eager_forward(input) + compile_failures = _compiled_kernel_failures() try: - return _get_compiled_group_norm()( - input, self.num_groups, self.weight, self.bias, self.eps - ) - except Exception as e: + out = _run_local(input, distconv, self._compiled_forward) + except compile_failures as e: # Compilation is an optimization, never a correctness requirement: - # a broken Inductor/Triton install, an unwritable cache directory or - # an untraceable input must degrade to the stock kernel, not kill a - # multi-node run. GroupNorm is pure, so retrying eagerly is safe. + # a broken Inductor install, an unwritable cache directory or an + # untraceable input must degrade to the stock kernel, not kill a + # multi-node run. Every exception caught here is a *compile*-time + # one, so nothing ran and retrying eagerly is safe. Same + # once-per-latch-edge logging as the Triton rung above. + first = not _compile_failed _compile_failed = True - logger.warning( - f"torch.compile of GroupNorm failed ({type(e).__name__}: {e}); " - "falling back to the eager kernel for the rest of this run. " - f"Set {COMPILE_ENV_VAR}=0 to skip this attempt entirely." - ) - return super().forward(input) + if first: + _warn_rung_failure( + "torch.compile of GroupNorm", e, "eager kernel", COMPILE_ENV_VAR + ) + # The one call this rung must not answer eagerly: a module already + # proven on it, failing while a backward is in flight, is a + # checkpoint recompute of a forward that *did* run compiled. The + # rungs save different tensors, so handing back the eager result + # makes the recomputed saved set disagree with the saved one and + # torch rejects the step -- a `CheckpointError`, or worse (both + # measured). Degrading is for modules with nothing to contradict; + # here the honest answer is the original exception, which at least + # names the rung and the shape that could not be served. + if self._compiled_ok and _replaying_a_forward(): + raise + return self._eager_forward(input) + else: + if not self._compiled_ok: + self._compiled_ok = True + return out diff --git a/ScaFFold/unet/unet_parts.py b/ScaFFold/unet/unet_parts.py index c9e6cb0..9fffe72 100644 --- a/ScaFFold/unet/unet_parts.py +++ b/ScaFFold/unet/unet_parts.py @@ -28,18 +28,35 @@ _outconv_annotate = annotate(fmt="OutConv.{}") -def _group_norm(num_groups, num_channels): +def _group_norm(num_groups, num_channels, activation=None): if num_channels % num_groups != 0: raise ValueError( f"group_norm_groups={num_groups} must evenly divide num_channels={num_channels}" ) - # FastGroupNorm is nn.GroupNorm plus a compiled GPU kernel; it holds the - # same parameters under the same names, so checkpoints are unaffected. - return FastGroupNorm(num_groups, num_channels) + # FastGroupNorm is nn.GroupNorm plus a Triton/compiled GPU kernel; it holds + # the same parameters under the same names, and `activation` is a plain + # attribute rather than a submodule, so checkpoints are unaffected. + return FastGroupNorm(num_groups, num_channels, activation=activation) class DoubleConv(nn.Module): - """(convolution => GroupNorm => ReLU) * 2""" + """(convolution => GroupNorm => ReLU) * 2 + + The ReLU lives *inside* the GroupNorm (``activation="relu"``), because the + Triton GroupNorm kernel folds it into its forward store for free and thereby + removes an entire streaming pass -- 38% of the forward at the shapes that + dominate the step. ``FastGroupNorm`` applies the ReLU on every path, + including eager, so the network's function is unchanged; only the number of + memory passes differs. + + The ``nn.ReLU`` slots are held open by ``nn.Identity`` rather than removed: + ``nn.Sequential`` names its children by position, so deleting them would + renumber the two convolutions and the second GroupNorm and invalidate every + existing checkpoint. Neither ``nn.ReLU`` nor ``nn.Identity`` has parameters + or buffers, so with the placeholders in place the state dict is byte + identical to the pre-fusion model's (pinned by + ``tests/test_groupnorm.py::test_state_dict_matches_plain_groupnorm_model``). + """ def __init__(self, in_channels, out_channels, group_norm_groups, mid_channels=None): super().__init__() @@ -47,11 +64,11 @@ def __init__(self, in_channels, out_channels, group_norm_groups, mid_channels=No mid_channels = out_channels self.double_conv = nn.Sequential( nn.Conv3d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False), - _group_norm(group_norm_groups, mid_channels), - nn.ReLU(inplace=True), + _group_norm(group_norm_groups, mid_channels, activation="relu"), + nn.Identity(), nn.Conv3d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False), - _group_norm(group_norm_groups, out_channels), - nn.ReLU(inplace=True), + _group_norm(group_norm_groups, out_channels, activation="relu"), + nn.Identity(), ) @_doubleconv_annotate diff --git a/tests/helpers/rank_scripts/groupnorm_shards_2rank.py b/tests/helpers/rank_scripts/groupnorm_shards_2rank.py new file mode 100644 index 0000000..c347365 --- /dev/null +++ b/tests/helpers/rank_scripts/groupnorm_shards_2rank.py @@ -0,0 +1,123 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""Two-rank check that FastGroupNorm's DCTensor route is shard-count agnostic. + +Run under ``torchrun --nproc_per_node=2`` with the gloo backend (see +``tests/test_groupnorm.py``). Every other DCTensor test uses +``num_shards=(1, 1, 1)``, where sharding is a no-op and the local shard is the +whole tensor; this one shards a spatial dim across two ranks so the claim the +fast path actually rests on -- GroupNorm statistics are per-shard, and the +route does not add communication or change which elements are reduced together +-- is exercised where it can fail. + +Each rank prints one ``RESULT ...`` line plus ``DONE``; the parent asserts. +""" + +import os +import sys + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +sys.path.insert(0, os.environ.get("SCAFFOLD_ROOT", "/usr/WS1/dryden1/ScaFFold")) + +from ScaFFold.unet import group_norm as gn_mod # noqa: E402 +from ScaFFold.unet.group_norm import FastGroupNorm # noqa: E402 + +GROUPS = 8 +CHANNELS = 16 +SIZE = 8 # dim 2 is split into two shards of 4 + + +def run(): + dist.init_process_group(backend="gloo") + rank = dist.get_rank() + + import distconv + + ps = distconv.ParallelStrategy(num_shards=(2,), shard_dim=(2,), device_type="cpu") + + # The same global volume on both ranks; each takes its own slab. + generator = torch.Generator().manual_seed(41) + volume = torch.randn(1, CHANNELS, SIZE, SIZE, SIZE, generator=generator) + half = SIZE // 2 + local = volume.narrow(2, rank * half, half).contiguous() + + norm = FastGroupNorm(GROUPS, CHANNELS) + param_generator = torch.Generator().manual_seed(97) + with torch.no_grad(): + norm.weight.normal_(1.0, 0.1, generator=param_generator) + norm.bias.normal_(0.0, 0.1, generator=param_generator) + + def forward(compiled): + """Run the wrapped GroupNorm with the compiled route on or off. + + The compiled route is forced on CPU by standing in the stock functional + kernel for the compiled callable: what is under test here is the + unwrap/rewrap plumbing at shard counts > 1, not Inductor. + """ + original_use, original_get = ( + gn_mod._use_compiled, + gn_mod._get_compiled_group_norm, + ) + if compiled: + gn_mod._use_compiled = lambda t, **kw: type(t) is torch.Tensor + gn_mod._get_compiled_group_norm = lambda: F.group_norm + else: + gn_mod._use_compiled = lambda t, **kw: False + try: + norm.zero_grad(set_to_none=True) + x = local.clone().requires_grad_(True) + out = norm(distconv.DCTensor.from_shard(x, ps)) + assert isinstance(out, distconv.DCTensor), type(out) + distconv.distconv._ToTensor.apply(out).pow(2).sum().backward() + weight_grad = norm.weight.grad + if isinstance(weight_grad, distconv.DCTensor): + weight_grad = weight_grad._tensor + return out._tensor.detach().clone(), x.grad.clone(), weight_grad.clone() + finally: + gn_mod._use_compiled = original_use + gn_mod._get_compiled_group_norm = original_get + + eager = forward(compiled=False) + compiled = forward(compiled=True) + identical = all(torch.equal(a, b) for a, b in zip(eager, compiled)) + + # Per-shard statistics: this rank's output normalizes its own slab only. + with torch.no_grad(): + per_shard = F.group_norm(local, GROUPS, norm.weight, norm.bias, norm.eps) + global_slice = F.group_norm( + volume, GROUPS, norm.weight, norm.bias, norm.eps + ).narrow(2, rank * half, half) + + # No spaces in any field: torchrun interleaves the ranks' stdout and a + # line can arrive without its trailing newline, so the parent's regex has + # to be able to tell two RESULT lines apart when they run together. + shape = "x".join(str(dim) for dim in compiled[0].shape) + print( + f"RESULT rank={rank} shape={shape} " + f"identical={identical} " + f"per_shard={torch.equal(compiled[0], per_shard)} " + f"global={torch.allclose(compiled[0], global_slice, atol=1e-6)}", + flush=True, + ) + print("DONE", flush=True) + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + run() diff --git a/tests/test_groupnorm.py b/tests/test_groupnorm.py index b1e8e6a..942087d 100644 --- a/tests/test_groupnorm.py +++ b/tests/test_groupnorm.py @@ -12,18 +12,36 @@ # # SPDX-License-Identifier: (Apache-2.0) -"""Tests for the compiled GroupNorm fast path (``ScaFFold.unet.group_norm``). +"""Tests for the GroupNorm fast paths (``ScaFFold.unet.group_norm``). The optimization must be invisible everywhere except in the profile: the same state dict as a stock ``nn.GroupNorm`` model (checkpoints stay interchangeable -in both directions), the same numbers within reduction-order noise, and an -eager fallback for every input the compiled kernel cannot or should not take -(CPU, tensor subclasses such as DistConv's ``DCTensor``, a broken compiler). +in both directions), the same numbers within reduction-order noise, and a +fallback for every input the fast kernels cannot or should not take (CPU, +unknown tensor subclasses, a broken Triton or Inductor install). The ladder is +Triton -> compiled -> eager, and a failure at any rung latches that rung off and +drops to the next, never to the bottom. + +DistConv's ``DCTensor`` is not in the rejection list: ``forward`` unwraps it to +its local shard around both fast kernels, so the wrapped production path is +served too. That unwrap is *not* the bare attribute read DistConv's own +dispatch does -- dispatch runs below autograd, where a bare read is safe, while +``forward`` runs above it and must go through DistConv's +``_ToTensor``/``_FromTensor`` pair to keep the graph connected. + +The ReLU that used to follow every GroupNorm now lives inside it +(``activation="relu"``), fused into the Triton store and applied explicitly on +the other two paths. ``DoubleConv`` keeps an ``nn.Identity`` in the vacated +``nn.Sequential`` slot, so the state dict does not move by one key -- which is +what the checkpoint tests here pin. """ from __future__ import annotations import logging +import os +import re +from pathlib import Path import pytest import torch @@ -32,6 +50,7 @@ from ScaFFold.unet import group_norm as gn_mod from ScaFFold.unet.group_norm import FastGroupNorm from ScaFFold.unet.unet_model import UNet +from tests.helpers import mpi_runner _N = 16 _N_CHANNELS = 3 @@ -41,39 +60,61 @@ @pytest.fixture(autouse=True) def _restore_compile_state(): - """Keep per-test overrides of the module-level compile state contained.""" - previous = gn_mod.set_compile_enabled(None) - failed = gn_mod._compile_failed + """Keep per-test overrides of the module-level routing state contained.""" + previous_compile = gn_mod.set_compile_enabled(None) + previous_triton = gn_mod.set_triton_enabled(None) + compile_failed = gn_mod._compile_failed + triton_failed = gn_mod._triton_failed yield - gn_mod._compile_override = previous - gn_mod._compile_failed = failed + gn_mod._compile_override = previous_compile + gn_mod._triton_override = previous_triton + gn_mod._compile_failed = compile_failed + gn_mod._triton_failed = triton_failed -def _make_unet(seed: int, group_norm_cls=None): - """Build the worker.py-shaped UNet, optionally with a different norm class.""" +def _make_unet(seed: int): + """Build the worker.py-shaped UNet.""" torch.manual_seed(seed) - if group_norm_cls is None: - return UNet( - n_channels=_N_CHANNELS, - n_classes=_N_CLASSES, - trilinear=False, - layers=2, - group_norm_groups=_GROUPS, - ) - import ScaFFold.unet.unet_parts as parts + return UNet( + n_channels=_N_CHANNELS, + n_classes=_N_CLASSES, + trilinear=False, + layers=2, + group_norm_groups=_GROUPS, + ) - original = parts.FastGroupNorm - parts.FastGroupNorm = group_norm_cls - try: - return UNet( - n_channels=_N_CHANNELS, - n_classes=_N_CLASSES, - trilinear=False, - layers=2, - group_norm_groups=_GROUPS, - ) - finally: - parts.FastGroupNorm = original + +def _make_plain_unet(seed: int): + """The pre-fusion build: stock ``nn.GroupNorm`` followed by ``nn.ReLU``. + + Built by *converting* a normal UNet rather than by patching the norm class + at construction time, because the fusion moved the ReLU into the norm: a + class swap alone would leave ``DoubleConv``'s ``nn.Identity`` placeholders + in place and produce a model with no activations at all, which would make + every numeric comparison below vacuous. Converting reproduces exactly the + module graph this branch replaced -- ``nn.GroupNorm`` where the fast norm + sits, an in-place ``nn.ReLU`` where the placeholder sits -- and consumes no + RNG (``nn.GroupNorm`` initializes to ones/zeros), so the parameters are + bit-identical to what ``_make_unet(seed)`` draws. + """ + model = _make_unet(seed) + for parent in [m for m in model.modules() if isinstance(m, nn.Sequential)]: + for index, child in enumerate(list(parent)): + if isinstance(child, FastGroupNorm): + plain = nn.GroupNorm( + child.num_groups, + child.num_channels, + eps=child.eps, + affine=child.affine, + ) + if child.affine: + with torch.no_grad(): + plain.weight.copy_(child.weight) + plain.bias.copy_(child.bias) + parent[index] = plain + elif isinstance(child, nn.Identity): + parent[index] = nn.ReLU(inplace=True) + return model def _make_input(seed: int = 0, channels: int = _N_CHANNELS, size: int = _N): @@ -93,7 +134,7 @@ def test_state_dict_matches_plain_groupnorm_model(): parameter inventory of the model may not shift by even one key. """ new_model = _make_unet(seed=0) - old_model = _make_unet(seed=0, group_norm_cls=nn.GroupNorm) + old_model = _make_plain_unet(seed=0) new_sd = new_model.state_dict() old_sd = old_model.state_dict() @@ -116,7 +157,7 @@ def test_checkpoint_round_trip_both_directions(tmp_path): script). After each load the two models must agree bit for bit. """ new_model = _make_unet(seed=0) - old_model = _make_unet(seed=1, group_norm_cls=nn.GroupNorm) + old_model = _make_plain_unet(seed=1) x = _make_input(seed=3) old_path = tmp_path / "old.pth" @@ -151,6 +192,119 @@ def test_unet_uses_fast_group_norm(): assert all(isinstance(m, FastGroupNorm) for m in norms) +def test_state_dict_bytes_identical_to_plain_groupnorm_model(): + """Not just the same keys: the serialized checkpoint must be byte identical. + + ``test_state_dict_matches_plain_groupnorm_model`` compares names, shapes and + dtypes; this compares the actual bytes ``torch.save`` writes, which is the + thing that has to stay interchangeable. It is the direct guard on the + ``nn.ReLU`` -> ``nn.Identity`` swap: ``nn.Sequential`` names its children by + position, so *removing* the activation slot rather than holding it open + would renumber ``3.weight`` and ``4.weight``/``4.bias`` and silently + invalidate every checkpoint on disk. + """ + import io + + new_model = _make_unet(seed=0) + old_model = _make_plain_unet(seed=0) + + def blob(model): + buffer = io.BytesIO() + torch.save(model.state_dict(), buffer) + return buffer.getvalue() + + assert blob(new_model) == blob(old_model) + + +def test_double_conv_keeps_the_activation_slots(): + """The fused build keeps six positional slots, with nothing in the spares. + + Pins both halves of the fusion design: the ReLU is *in* the norm + (``activation == "relu"`` at positions 1 and 4) and its old slots (2 and 5) + are parameterless placeholders rather than deletions. + """ + from ScaFFold.unet.unet_parts import DoubleConv + + block = DoubleConv(3, 16, _GROUPS) + children = list(block.double_conv) + assert len(children) == 6 + for norm_index, spare_index in ((1, 2), (4, 5)): + norm = children[norm_index] + assert isinstance(norm, FastGroupNorm) + assert norm.activation == "relu" + spare = children[spare_index] + assert isinstance(spare, nn.Identity) + assert list(spare.parameters()) == [] + assert list(spare.buffers()) == [] + # And the positional key numbering is exactly the pre-fusion one. + assert list(block.state_dict().keys()) == [ + "double_conv.0.weight", + "double_conv.1.weight", + "double_conv.1.bias", + "double_conv.3.weight", + "double_conv.4.weight", + "double_conv.4.bias", + ] + + +def test_double_conv_output_matches_the_explicit_relu_build(): + """Folding the ReLU into the norm may not change a single bit of the output. + + The fused module applies the ReLU itself on every path, so on CPU (eager) + the block must reproduce ``conv -> GroupNorm -> ReLU`` exactly, gradients + included. + """ + from ScaFFold.unet.unet_parts import DoubleConv + + torch.manual_seed(4) + fused = DoubleConv(3, 16, _GROUPS) + reference = DoubleConv(3, 16, _GROUPS) + reference.load_state_dict(fused.state_dict()) + for index in (1, 4): + reference.double_conv[index].activation = None + for index in (2, 5): + reference.double_conv[index] = nn.ReLU(inplace=True) + + x_fused = _make_input(seed=8, channels=3, size=8).requires_grad_(True) + x_reference = x_fused.detach().clone().requires_grad_(True) + + out_fused = fused(x_fused) + out_reference = reference(x_reference) + assert torch.equal(out_fused, out_reference) + # A block whose activation silently vanished would still pass an + # output-equality test against another activation-free block, so assert the + # ReLU is really there. + assert (out_fused < 0).sum() == 0 + assert out_fused.max() > 0 + + out_fused.pow(2).sum().backward() + out_reference.pow(2).sum().backward() + assert torch.equal(x_fused.grad, x_reference.grad) + for (name, a), (_, b) in zip( + fused.named_parameters(), reference.named_parameters() + ): + assert torch.equal(a.grad, b.grad), name + + +def test_activation_argument_is_validated(): + """An unknown activation must fail at construction, not at the first step.""" + with pytest.raises(ValueError, match="activation"): + FastGroupNorm(_GROUPS, 16, activation="gelu") + + +def test_supported_activations_match_the_kernels(): + """The module's activation list may not drift from the kernel's. + + ``group_norm`` spells the tuple out rather than importing it (importing the + kernel module has to stay off the CPU path), so nothing but this test stops + the two copies from diverging into a runtime ``ValueError`` from inside the + custom op. + """ + from ScaFFold.unet import triton_group_norm as triton_mod + + assert set(gn_mod.SUPPORTED_ACTIVATIONS) <= set(triton_mod.SUPPORTED_ACTIVATIONS) + + # --------------------------------------------------------------------------- # CPU behavior: identical numerics, and no compilation at all # --------------------------------------------------------------------------- @@ -190,10 +344,13 @@ def _boom(*a, **kw): def test_tensor_subclass_input_stays_eager(): - """DistConv wraps activations in a ``__torch_dispatch__`` tensor subclass. + """Unknown tensor subclasses must stay on the eager path. - Dynamo cannot trace those wrappers, so the predicate must reject anything - that is not exactly ``torch.Tensor`` before a compile is attempted. + Dynamo cannot trace ``__torch_dispatch__`` wrappers, so the predicate must + reject anything that is not exactly ``torch.Tensor`` before a compile is + attempted. DistConv's ``DCTensor`` is handled separately -- ``forward`` + unwraps it before consulting the predicate -- but any other wrapper has + unknown semantics and keeps the stock kernel. """ class _Wrapper(torch.Tensor): @@ -207,14 +364,17 @@ class _Wrapper(torch.Tensor): def test_compile_failure_falls_back_to_eager(monkeypatch, caplog): """A broken compiler degrades to the stock kernel instead of killing the run. - Simulated by making the compiled callable raise; the module must return the - eager result, warn once, and stop trying for the rest of the process. + Simulated by making the compiled callable raise the real thing Dynamo and + Inductor raise (``BackendCompilerFailed``/``Unsupported`` share the + ``TorchDynamoException`` root the ladder allowlists); the module must return + the eager result, warn once, and stop trying for the rest of the process. """ + import torch._dynamo.exc def _raises(*args, **kwargs): - raise RuntimeError("simulated Inductor failure") + raise torch._dynamo.exc.Unsupported("simulated Inductor failure") - monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input: True) + monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input, **kw: True) monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _raises) gn_mod._compile_failed = False @@ -232,6 +392,148 @@ def _raises(*args, **kwargs): assert gn_mod._use_compiled(torch.randn(1, 8, 4, 4, 4)) is False +def test_triton_failure_falls_back_to_the_compiled_kernel(monkeypatch, caplog): + """A broken Triton install drops to *compiled*, not all the way to eager. + + The distinction is worth 10x on the shapes that dominate the step, so the + ladder must have three rungs and not two. Simulated by forcing the Triton + predicate on and making its kernel raise; the compiled stand-in must then be + the one that answers, exactly once, with the Triton path latched off. + """ + from ScaFFold.unet.triton_group_norm import TritonKernelError + + compiled_calls = [] + + def _raises(*args, **kwargs): + raise TritonKernelError("simulated Triton failure") + + def _recording(input, num_groups, weight, bias, eps): + compiled_calls.append(type(input)) + return nn.functional.group_norm(input, num_groups, weight, bias, eps) + + monkeypatch.setattr(gn_mod, "_use_triton", lambda *a, **kw: True) + monkeypatch.setattr(gn_mod, "_get_triton_module", _raises) + monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input, **kw: True) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _recording) + gn_mod._triton_failed = False + gn_mod._compile_failed = False + + fast = FastGroupNorm(_GROUPS, 64) + x = _make_input(seed=44, channels=64, size=8) + with caplog.at_level(logging.WARNING, logger=gn_mod.__name__): + out = fast(x) + + assert compiled_calls == [torch.Tensor], "compiled kernel was not the fallback" + assert torch.equal( + out, nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias) + ) + assert any( + "falling back to the compiled kernel" in r.message for r in caplog.records + ) + assert gn_mod._triton_failed is True + assert gn_mod._compile_failed is False + # Latched off: the predicate now refuses even a would-be eligible tensor. + monkeypatch.undo() + assert gn_mod._use_triton(torch.randn(1, 8, 4, 4, 4), 8, None, None, None) is False + + +@pytest.mark.parametrize("rung", ["triton", "compiled"]) +def test_checkpoint_recompute_stop_is_re_raised(monkeypatch, rung): + """``_StopRecomputationError`` is control flow, not a kernel failure. + + ``torch.utils.checkpoint``'s non-reentrant recompute stops itself by raising + it from a saved-tensor *pack hook*, i.e. from inside whichever op is saving + a tensor at that moment -- which, now that the ReLU is fused and nothing + follows GroupNorm in a ``DoubleConv``, is this module. A blanket + ``except Exception`` would swallow it, latch the fast kernel off and drop the + whole model to eager mid-run. (Observed exactly that on + ``test_gpu_activation_checkpointing_matches_eager`` before the re-raise.) + """ + import torch.utils.checkpoint as checkpoint_mod + + stop = checkpoint_mod._StopRecomputationError + + def _raises(*args, **kwargs): + raise stop() + + if rung == "triton": + monkeypatch.setattr(gn_mod, "_use_triton", lambda *a, **kw: True) + monkeypatch.setattr(gn_mod, "_get_triton_module", _raises) + else: + monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input, **kw: True) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _raises) + gn_mod._triton_failed = False + gn_mod._compile_failed = False + + fast = FastGroupNorm(_GROUPS, 16) + with pytest.raises(stop): + fast(_make_input(seed=45, channels=16, size=4)) + assert gn_mod._triton_failed is False + assert gn_mod._compile_failed is False + + +def test_cpu_activation_checkpointing_keeps_the_fast_path(monkeypatch): + """End-to-end version of the above, on the real model. + + The fast path is forced on for CPU tensors (with the stock kernel standing + in for the compiled one, so only the *routing* is under test) and the model + is run with activation checkpointing. Gradients must match the + non-checkpointed run and the fast path must still be live afterwards -- a + swallowed recompute-stop shows up as a latched-off kernel here. + """ + monkeypatch.setattr( + gn_mod, "_use_compiled", lambda t, **kw: type(t) is torch.Tensor + ) + monkeypatch.setattr( + gn_mod, "_get_compiled_group_norm", lambda: nn.functional.group_norm + ) + gn_mod._compile_failed = False + + x = _make_input(seed=46).requires_grad_(True) + + def grads(checkpointing): + model = _make_unet(seed=0) + if checkpointing: + model.use_checkpointing() + model.zero_grad(set_to_none=True) + model(x).pow(2).sum().backward() + return {n: p.grad.detach().clone() for n, p in model.named_parameters()} + + direct = grads(False) + checkpointed = grads(True) + assert gn_mod._compile_failed is False, "recompute-stop was swallowed" + for name in direct: + assert torch.allclose(direct[name], checkpointed[name]), name + + +def test_triton_rejects_unknown_tensor_subclasses(): + """``is_supported`` only asks ``isinstance``, so the type check lives here. + + A ``__torch_dispatch__`` wrapper other than DCTensor has unknown semantics + and must keep the stock kernel, exactly as it does for the compiled path -- + but ``triton_group_norm.is_supported`` would happily accept one, so + ``_use_triton`` has to reject it itself rather than delegating. + """ + + class _Wrapper(torch.Tensor): + pass + + plain = torch.randn(1, 8, 4, 4, 4) + assert gn_mod._use_triton(plain, 8, None, None, None) is False # CPU + assert gn_mod._use_triton(plain.as_subclass(_Wrapper), 8, None, None, None) is False + + +def test_triton_env_opt_out_skips_the_kernel_module_entirely(monkeypatch): + """``SCAFFOLD_GROUPNORM_TRITON=0`` is checked before anything is imported.""" + + def _boom(): + raise AssertionError("the kernel module must not be imported when opted out") + + monkeypatch.setattr(gn_mod, "_get_triton_module", _boom) + gn_mod.set_triton_enabled(False) + assert gn_mod._use_triton(torch.randn(1, 8, 4, 4, 4), 8, None, None, None) is False + + @pytest.mark.parametrize( "value,expected", [ @@ -259,6 +561,116 @@ def test_env_var_unset_means_auto(monkeypatch): assert gn_mod._compile_override is None +@pytest.mark.parametrize( + "value,expected", + [ + ("0", False), + ("false", False), + ("OFF", False), + ("no", False), + ("1", True), + ("true", True), + ("On", True), + ("yes", True), + ("maybe", None), + ], +) +def test_triton_env_var_controls_the_fast_path(monkeypatch, value, expected): + """``SCAFFOLD_GROUPNORM_TRITON`` is parsed exactly like its compile twin.""" + monkeypatch.setenv(gn_mod.TRITON_ENV_VAR, value) + gn_mod.set_triton_enabled(None) + assert gn_mod._triton_override is expected + + +def test_triton_env_var_unset_means_auto(monkeypatch): + """Unset means "on wherever is_supported accepts" -- the production default.""" + monkeypatch.delenv(gn_mod.TRITON_ENV_VAR, raising=False) + gn_mod.set_triton_enabled(None) + assert gn_mod._triton_override is None + + +def test_triton_env_var_garbage_warns(monkeypatch, caplog): + """An unparsable value is ignored *loudly*, like the compile variable.""" + monkeypatch.setenv(gn_mod.TRITON_ENV_VAR, "sometimes") + with caplog.at_level(logging.WARNING, logger=gn_mod.__name__): + gn_mod.set_triton_enabled(None) + assert any(gn_mod.TRITON_ENV_VAR in r.message for r in caplog.records) + assert gn_mod._triton_override is None + + +def test_set_triton_enabled_returns_the_previous_setting(monkeypatch): + """The save/restore contract tests rely on, matching set_compile_enabled.""" + monkeypatch.delenv(gn_mod.TRITON_ENV_VAR, raising=False) + gn_mod.set_triton_enabled(None) + assert gn_mod.set_triton_enabled(False) is None + assert gn_mod.set_triton_enabled(True) is False + assert gn_mod.set_triton_enabled(None) is True + assert gn_mod._triton_override is None + + +def test_cpu_never_imports_triton(fresh_python): + """A CPU-only process must not import triton, nor the kernel module. + + Two separate costs, both of which the CPU unit suite would otherwise pay on + every run: ``import triton`` (seconds, and it is not installed everywhere), + and importing ``ScaFFold.unet.triton_group_norm``, which registers two + dispatcher ops and an autograd formula at import time. ``_use_triton`` + rejects non-CUDA tensors *before* it touches the module, which is what this + pins -- run in a fresh interpreter because the test session itself has long + since imported the kernel module for the kernel's own tests. + """ + out = fresh_python( + "import sys\n" + "import torch\n" + "from ScaFFold.unet.unet_model import UNet\n" + "m = UNet(n_channels=3, n_classes=2, trilinear=False, layers=1, " + "group_norm_groups=8)\n" + "with torch.no_grad():\n" + " m(torch.randn(1, 3, 16, 16, 16))\n" + "print('triton', 'triton' in sys.modules)\n" + "print('kernel', 'ScaFFold.unet.triton_group_norm' in sys.modules)\n" + ) + assert "triton False" in out, out + assert "kernel False" in out, out + + +def test_cpu_activation_is_applied_on_the_eager_path(): + """``activation="relu"`` is a promise of the module, not of the kernel.""" + fast = FastGroupNorm(_GROUPS, 16, activation="relu") + plain = nn.GroupNorm(_GROUPS, 16) + with torch.no_grad(): + plain.weight.copy_(fast.weight) + plain.bias.copy_(fast.bias) + x = _make_input(seed=41, channels=16, size=8) + assert torch.equal(fast(x), torch.relu(plain(x))) + + +def test_cpu_activation_none_leaves_the_output_alone(): + """The default stays a bare GroupNorm -- no accidental global activation.""" + fast = FastGroupNorm(_GROUPS, 16) + x = _make_input(seed=42, channels=16, size=8) + assert torch.equal( + fast(x), nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias) + ) + + +def test_activation_does_not_allocate_a_second_output(): + """The absorbed ReLU keeps ``nn.ReLU(inplace=True)``'s memory behaviour. + + The old ``nn.Sequential`` spelling mutated the GroupNorm output in place; + an out-of-place ``F.relu`` here would add a full activation-sized allocation + at all 22 sites. Checked by handing the module a stand-in kernel whose + output we still hold: the ReLU must have rewritten *that* tensor. + """ + fast = FastGroupNorm(_GROUPS, 16, activation="relu") + x = _make_input(seed=43, channels=16, size=4) + produced = nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias, fast.eps) + assert (produced < 0).any(), "test input must have negatives to clamp" + out = fast._activate(produced) + assert out.data_ptr() == produced.data_ptr() + assert (produced < 0).sum() == 0 + + def test_recompile_limit_is_raised_never_lowered(): """Dynamo's stock cap of 8 is below what one UNet needs. @@ -285,6 +697,385 @@ def test_recompile_limit_is_raised_never_lowered(): setattr(config, name, original) +def test_the_compiled_region_carries_its_own_recompile_limit(monkeypatch): + """The global limit is thread-local, so the region must carry one too. + + ``torch._dynamo.config`` keeps user overrides in a ``ContextVar`` + ("User overrides are thread-local", ``torch/utils/_config_module.py``), so + what :func:`_raise_recompile_limit` writes is invisible from every *other* + thread -- and one of those threads matters: ``torch.utils.checkpoint``'s + non-reentrant recompute runs inside the backward pass, i.e. on the autograd + engine's device worker thread. On a ``DCTensor`` that recompute has to + compile (it reaches this module with ``__torch_function__`` subclass + handling disabled, which is part of Dynamo's ``GLOBAL_STATE`` guard, so it + misses every entry the forward built), and there the limit read the stock 8 + however large the global had been set -- ``FailOnRecompileLimitHit``, run + over. ``torch.compile``'s ``recompile_limit=`` is applied by Dynamo around + the compile itself, on whichever thread that compile happens on, which is + the only spelling that reaches the worker; this pins that we ask for it. + """ + seen = {} + + def _fake_compile(fn, **kwargs): + seen.update(kwargs) + return fn + + monkeypatch.setattr(torch, "compile", _fake_compile) + assert gn_mod._compile_group_norm() is gn_mod._group_norm + assert seen.get("recompile_limit") == gn_mod._MIN_RECOMPILE_LIMIT + assert seen.get("fullgraph") is True and seen.get("dynamic") is False + + +def test_compiling_still_works_without_a_per_region_limit(monkeypatch): + """A torch too old for ``recompile_limit=`` must still get a callable. + + The keyword is the fix for the worker thread, not a requirement for + compiling at all; dropping the whole rung on a ``TypeError`` would be a far + bigger regression than the case it addresses. + """ + calls = [] + + def _fake_compile(fn, **kwargs): + calls.append(kwargs) + if "recompile_limit" in kwargs: + raise TypeError("compile() got an unexpected keyword 'recompile_limit'") + return fn + + monkeypatch.setattr(torch, "compile", _fake_compile) + assert gn_mod._compile_group_norm() is gn_mod._group_norm + assert len(calls) == 2 and "recompile_limit" not in calls[1] + + +def test_a_recompile_limit_hit_is_a_kernel_failure_not_a_crash(monkeypatch, caplog): + """``FailOnRecompileLimitHit`` has to land in the ladder, not in the run. + + It is what ``fullgraph=True`` raises when a frame needs more cache entries + than the recompile limit allows, and -- unlike every other Dynamo failure -- + it derives straight from ``Exception`` rather than from + ``TorchDynamoException``, so an allowlist that names only the latter lets it + escape and kill the step (observed at ``5943389``). It is raised while + compiling, before the callable has run or saved anything, so the eager + retry underneath it is safe. + """ + import torch._dynamo.exc + + limit_hit = torch._dynamo.exc.FailOnRecompileLimitHit + assert not issubclass(limit_hit, torch._dynamo.exc.TorchDynamoException), ( + "naming it separately is only needed while it sits outside that root" + ) + + def _kernel(*args, **kwargs): + raise limit_hit("simulated recompile limit hit") + + monkeypatch.setattr( + gn_mod, "_use_compiled", lambda t, **kw: type(t) is torch.Tensor + ) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _kernel) + gn_mod._compile_failed = False + + fast = _seeded_norm() + x = _make_input(seed=51, channels=16, size=4) + expected = nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias, fast.eps) + + with caplog.at_level(logging.WARNING, logger=gn_mod.__name__): + out = fast(x) + + assert torch.allclose(out, expected), "the eager fallback did not run" + assert gn_mod._compile_failed is True, "the failure did not latch the rung off" + assert fast._compiled_ok is False + assert any("falling back" in record.message for record in caplog.records) + + +# --------------------------------------------------------------------------- +# DCTensor routing: unwrap -> compiled kernel -> rewrap +# --------------------------------------------------------------------------- + + +@pytest.fixture +def dc_cpu(gloo_group_1rank): + """DistConv package plus a CPU ParallelStrategy over the 1-rank group. + + ``num_shards=(1, 1, 1)`` on dims (2, 3, 4) is exactly what worker.py builds + for a single-device run; a process group must exist even then. + """ + import distconv + + ps = distconv.ParallelStrategy( + num_shards=(1, 1, 1), shard_dim=(2, 3, 4), device_type="cpu" + ) + return distconv, ps + + +def _seeded_norm(channels=16): + """A FastGroupNorm with non-default affine params (the defaults are 1/0, + which would let a kernel that drops weight/bias slip through).""" + fast = FastGroupNorm(_GROUPS, channels) + generator = torch.Generator().manual_seed(97) + with torch.no_grad(): + fast.weight.normal_(1.0, 0.1, generator=generator) + fast.bias.normal_(0.0, 0.1, generator=generator) + return fast + + +def test_dctensor_routes_through_compiled_kernel(monkeypatch, dc_cpu): + """A DCTensor input reaches the compiled callable as its plain local shard. + + Verified with a recording stand-in for the compiled callable: it must see + exactly ``torch.Tensor`` (Dynamo cannot trace the wrapper), and the caller + must get a DCTensor back with the same values the stock kernel produces. + """ + distconv, ps = dc_cpu + seen = [] + + def _recording(input, num_groups, weight, bias, eps): + seen.append(type(input)) + return nn.functional.group_norm(input, num_groups, weight, bias, eps) + + monkeypatch.setattr( + gn_mod, "_use_compiled", lambda t, **kw: type(t) is torch.Tensor + ) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _recording) + + fast = _seeded_norm() + x = _make_input(seed=21, channels=16, size=4) + out = fast(distconv.DCTensor.from_shard(x, ps)) + + assert seen == [torch.Tensor] + assert isinstance(out, distconv.DCTensor) + reference = nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias, fast.eps) + assert torch.equal(out._tensor, reference) + + +def test_dctensor_gradients_reach_the_layer_upstream(monkeypatch, dc_cpu): + """Gradients must flow past GroupNorm into the layer that produced its input. + + The unwrap has to go through DistConv's autograd pair rather than a bare + ``input._tensor`` read. The distinction is invisible when the DCTensor + wraps a leaf -- the leaf *is* ``_tensor``, so even a bare read reaches it -- + which is why this test puts a producer in front, as production does + (``conv -> GroupNorm`` in every block). With a bare read the graph is + severed there: the input and the producer's weight get no gradient at all + while GroupNorm's own weight/bias still look healthy. + """ + distconv, ps = dc_cpu + monkeypatch.setattr( + gn_mod, "_use_compiled", lambda t, **kw: type(t) is torch.Tensor + ) + monkeypatch.setattr( + gn_mod, "_get_compiled_group_norm", lambda: nn.functional.group_norm + ) + + fast = _seeded_norm() + reference = nn.GroupNorm(_GROUPS, 16) + producer = nn.Conv3d(16, 16, 1, bias=False) + reference_producer = nn.Conv3d(16, 16, 1, bias=False) + with torch.no_grad(): + reference.weight.copy_(fast.weight) + reference.bias.copy_(fast.bias) + reference_producer.weight.copy_(producer.weight) + + x_fast = _make_input(seed=22, channels=16, size=4).requires_grad_(True) + x_ref = x_fast.detach().clone().requires_grad_(True) + + # A 1x1x1 conv on a DCTensor takes DistConv's convolution path, so the + # DCTensor handed to GroupNorm is a genuine intermediate, not a leaf. + out = fast(producer(distconv.DCTensor.from_shard(x_fast, ps))) + assert isinstance(out, distconv.DCTensor) + # Unwrap the way a downstream consumer would (autograd-aware) and drive a + # scalar backward through it. + distconv.distconv._ToTensor.apply(out).pow(2).sum().backward() + reference(reference_producer(x_ref)).pow(2).sum().backward() + + assert x_fast.grad is not None, "gradient never reached the input" + assert producer.weight.grad is not None, "gradient never reached the producer" + assert torch.equal(x_fast.grad, x_ref.grad) + assert torch.equal(producer.weight.grad, reference_producer.weight.grad) + assert torch.equal(fast.weight.grad, reference.weight.grad) + assert torch.equal(fast.bias.grad, reference.bias.grad) + + +def test_dctensor_on_cpu_never_invokes_torch_compile(monkeypatch, dc_cpu): + """The unwrap route obeys the same CPU guard as plain tensors. + + On CPU the local shard fails the ``is_cuda`` check, so a DCTensor must fall + through to the stock eager dispatch -- and still come back wrapped. The + stand-in records instead of only raising: ``forward`` catches ``Exception`` + to fall back, so a raise alone would be swallowed by the very code path + under test and the assertion would never fire. + """ + distconv, ps = dc_cpu + calls = [] + + def _boom(*a, **kw): + calls.append(a) + raise AssertionError("torch.compile must not be called for CPU tensors") + + monkeypatch.setattr(torch, "compile", _boom) + monkeypatch.setattr(gn_mod, "_compiled_group_norm", None) + gn_mod.set_compile_enabled(True) # even when explicitly forced on + + fast = _seeded_norm() + x = _make_input(seed=23, channels=16, size=4) + out = fast(distconv.DCTensor.from_shard(x, ps)) + + assert not calls + assert isinstance(out, distconv.DCTensor) + reference = nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias, fast.eps) + assert torch.equal(out._tensor, reference) + + +def test_dctensor_compile_failure_falls_back_to_eager(monkeypatch, caplog, dc_cpu): + """A broken compiler degrades the wrapped path to eager, like the plain one.""" + distconv, ps = dc_cpu + import torch._dynamo.exc + + def _raises(*args, **kwargs): + raise torch._dynamo.exc.Unsupported("simulated Inductor failure") + + monkeypatch.setattr( + gn_mod, "_use_compiled", lambda t, **kw: type(t) is torch.Tensor + ) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _raises) + gn_mod._compile_failed = False + + fast = _seeded_norm() + x = _make_input(seed=24, channels=16, size=4) + with caplog.at_level(logging.WARNING, logger=gn_mod.__name__): + out = fast(distconv.DCTensor.from_shard(x, ps)) + + assert isinstance(out, distconv.DCTensor) + reference = nn.functional.group_norm(x, _GROUPS, fast.weight, fast.bias, fast.eps) + assert torch.equal(out._tensor, reference) + assert gn_mod._compile_failed is True + assert any("falling back to the eager kernel" in r.message for r in caplog.records) + + +def test_dctensor_two_shards_matches_eager_and_normalizes_per_shard(): + """Shard count > 1: same values as the eager route, still per-shard stats. + + Every other DCTensor test here runs ``num_shards=(1, 1, 1)``, where the + local shard is the whole tensor and the sharding is a no-op -- so none of + them can catch a fast path that quietly reduced over the wrong set of + elements. This one splits a spatial dim over two ranks and asserts both + halves of the claim the fast path rests on: bit-identical to the eager + wrapped route, and normalizing the local shard rather than the global + volume (which is DistConv's existing semantics, not something this change + introduces -- see the R32/R39 notes for the upstream discussion). + """ + script = ( + Path(__file__).resolve().parent + / "helpers" + / "rank_scripts" + / "groupnorm_shards_2rank.py" + ) + rc, out, err = mpi_runner.torchrun_gloo(str(script), n=2, timeout=180) + assert rc == 0, f"2-rank job failed rc={rc}\nstdout:\n{out}\nstderr:\n{err[-3000:]}" + + results = { + rank: match + for rank, *match in re.findall( + # Bounded alternatives, not \S+: the ranks' lines can arrive + # concatenated, so a greedy final field would swallow the next + # line's "RESULT". + r"RESULT rank=(\d+) shape=(\S+) identical=(True|False) " + r"per_shard=(True|False) global=(True|False)", + out, + ) + } + assert set(results) == {"0", "1"}, f"missing ranks\nstdout:\n{out}" + for rank, (shape, identical, per_shard, matches_global) in results.items(): + assert identical == "True", f"rank {rank}: compiled route differs from eager" + assert per_shard == "True", f"rank {rank}: not per-shard statistics" + assert matches_global == "False", f"rank {rank}: matched global statistics" + # Each rank holds half of the sharded dim. + assert shape == "1x16x4x8x8", f"rank {rank}: unexpected shard shape {shape}" + + +@pytest.fixture +def dc_cuda(): + """DistConv package plus a CUDA ParallelStrategy over a 1-rank NCCL group.""" + import distconv + import torch.distributed as dist + + created = False + if not dist.is_initialized(): + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29517") + dist.init_process_group(backend="nccl", rank=0, world_size=1) + created = True + ps = distconv.ParallelStrategy( + num_shards=(1, 1, 1), shard_dim=(2, 3, 4), device_type="cuda" + ) + yield distconv, ps + if created and dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.gpu +@pytest.mark.parametrize("autocast", [False, True]) +@pytest.mark.parametrize("layout", ["contiguous", "channels_last_3d"]) +def test_gpu_dctensor_matches_eager_dctensor(dc_cuda, autocast, layout): + """The compiled unwrap path matches today's eager wrapped path on GPU. + + This is the production configuration: worker.py wraps every activation in + a DCTensor (even at dc_num_shards=[1,1,1]), which used to force the eager + kernel. Values and gradients must agree within reduction-order noise, the + output must still be a DCTensor, and the compile must actually engage. + + Both layouts are covered because production requests ``channels_last_3d`` + (worker.py) and, with ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` set as it is there, + the convolutions really do hand GroupNorm channels-last activations. Only + parity is asserted, not the output layout: both routes compared here return + contiguous regardless of the input layout (the Triton kernel, which does + not, is pinned off below and covered by its own tests). + """ + distconv, ps = dc_cuda + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(31) + x = torch.randn(1, 64, 32, 32, 32, device=device, generator=generator) + if layout == "channels_last_3d": + x = x.to(memory_format=torch.channels_last_3d) + grad_out = torch.randn(*x.shape, device=device, generator=generator) + + fast = FastGroupNorm(_GROUPS, 64).to(device) + with torch.no_grad(): + fast.weight.normal_(1.0, 0.1, generator=generator) + fast.bias.normal_(0.0, 0.1, generator=generator) + + def plain(t): + return t._tensor if isinstance(t, distconv.DCTensor) else t + + def run(compiled): + # This test is about the compiled rung of the ladder, which the Triton + # one would otherwise pre-empt on the channels-last parametrization. + gn_mod.set_triton_enabled(False) + gn_mod.set_compile_enabled(compiled) + inp = x.clone().requires_grad_(True) + fast.zero_grad(set_to_none=True) + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast): + out = fast(distconv.DCTensor.from_shard(inp, ps)) + assert isinstance(out, distconv.DCTensor) + local = distconv.distconv._ToTensor.apply(out) + local.backward(grad_out.to(local.dtype)) + return ( + local.detach(), + inp.grad.detach().clone(), + plain(fast.weight.grad).detach().clone(), + plain(fast.bias.grad).detach().clone(), + ) + + eager = run(False) + compiled = run(True) + assert gn_mod._compiled_group_norm is not None, "compiled path was not taken" + assert not gn_mod._compile_failed + + _assert_close(compiled[0], eager[0], 1e-5, "output") + _assert_close(compiled[1], eager[1], 1e-4, "d_input") + _assert_close(compiled[2], eager[2], 1e-4, "d_weight") + _assert_close(compiled[3], eager[3], 1e-4, "d_bias") + assert compiled[0].dtype == eager[0].dtype + + # --------------------------------------------------------------------------- # GPU behavior: numerics, single compile, checkpointing # --------------------------------------------------------------------------- @@ -320,6 +1111,9 @@ def test_gpu_compiled_matches_eager(shape, autocast): fast.bias.normal_(0.0, 0.1, generator=generator) def run(compiled): + # Compiled-rung test: the inputs here are contiguous, which the Triton + # kernel declines anyway, but pin it off so the routing cannot drift. + gn_mod.set_triton_enabled(False) gn_mod.set_compile_enabled(compiled) inp = x.clone().requires_grad_(True) fast.zero_grad(set_to_none=True) @@ -357,6 +1151,7 @@ def test_gpu_steady_state_does_not_recompile(): """ from torch._dynamo.utils import counters + gn_mod.set_triton_enabled(False) # this is the compiled rung's guard set gn_mod.set_compile_enabled(True) device = torch.device("cuda") fast = FastGroupNorm(_GROUPS, 64).to(device) @@ -375,6 +1170,72 @@ def step(): assert counters["stats"]["unique_graphs"] == before, "recompiled in steady state" +# --------------------------------------------------------------------------- +# Whole-network gradient comparisons +# --------------------------------------------------------------------------- +# +# Two tests below run the whole UNet twice, changing only which rung serves +# GroupNorm, and ask whether the gradients agree. *How* that is asked matters +# more than it looks, and both tests used to ask it in a way that could only +# pass by luck. Measurements: ``work/triton-conv/review/GN_CHECKPOINT_TEST.md``. +# +# **Per-parameter relative L2 under bf16 autocast is not a bounded quantity for +# this model.** Against an fp64 reference, every arm -- eager, compiled and +# Triton alike -- is 12.1-12.5% off on the bottleneck's parameters, whose +# gradients are 170x smaller than the largest in the network. The difference +# between two arms is therefore the difference of two ~12% errors, and it is +# small only when they happen to cancel. Whether they cancel is settled +# *outside the source tree*: two of this model's 33 convolution problems have +# MIOpen algorithms whose benchmark times tie to within 11%, so which one wins +# is frozen into the machine-local find database, and rewriting only those two +# recorded times moves "compiled vs. eager" from 5.3e-3 to 5.4e-2 with nothing +# else changed. Three of the four algorithm combinations land at 5.3-6.7e-3 +# and the fourth at 5.4e-2 -- which is exactly the history of this file, an +# assertion that was intermittent and then, once a cache went warm, failed +# deterministically at a bit-identical value. +# +# It is not a gate either: injecting a 1e-4 relative error into the compiled +# rung moves the per-parameter figure only from 5.4e-2 to 1.1e-1, because both +# are already saturated by the bf16 floor. +# +# So rung equivalence is asserted where it is measurable -- **without +# autocast**, where the same comparison reads 1.3e-6 and that same injected +# 1e-4 error reads 1.9e-4, a 154x separation -- while the bf16 autocast run, +# which is the production combination and the one the checkpoint machinery has +# to survive, is asserted on the *aggregate* gradient. That is stable +# (5.0e-3 against a 3.2e-3 run-to-run floor, across every algorithm combination +# measured) and still catches a wrong eps (6.9e-2) or a wrong group count +# (5.2e-1). + +#: Cross-rung agreement without autocast. Measured 1.3e-6 (compiled vs. eager) +#: and 2.2e-6 (Triton vs. compiled); 1e-4 leaves ~50x headroom and still fails +#: on a 1e-4 relative kernel error, which reads 1.9e-4 here. +_FP32_RUNG_TOLERANCE = 1e-4 + +#: Aggregate agreement under bf16 autocast, and per-parameter agreement between +#: two runs *on the same rung* (where the bf16 error is common-mode and the +#: figure sits on the model's own run-to-run floor). +_BF16_TOLERANCE = 5e-2 + + +def _relative_l2(actual, expected): + """Relative L2 over the concatenated gradient of every parameter.""" + names = sorted(expected) + a = torch.cat([actual[name].float().flatten() for name in names]) + b = torch.cat([expected[name].float().flatten() for name in names]) + return ((a - b).norm() / b.norm().clamp_min(1e-12)).item() + + +def _worst_relative_l2(actual, expected): + """The largest per-parameter relative L2, with the parameter it belongs to.""" + worst = (0.0, "") + for name, reference in expected.items(): + reference = reference.float() + error = (actual[name].float() - reference).norm().item() + worst = max(worst, (error / max(reference.norm().item(), 1e-12), name)) + return worst + + @pytest.mark.gpu def test_gpu_activation_checkpointing_matches_eager(): """The compiled kernel must survive recompute under use_checkpointing(). @@ -383,39 +1244,441 @@ def test_gpu_activation_checkpointing_matches_eager(): pass; a compiled region has to produce the same activations both times or the gradients silently change. - Compared as relative L2 error per gradient tensor, because whole-network - agreement is not bitwise even without this change: with cudnn.benchmark on - and bf16 autocast, two eager runs of this model differ by ~4e-3 relative - (measured), and checkpointing on vs. off differs by the same amount. - Measured here: compiled vs. eager 6.4e-3, i.e. the same order as that noise - floor -- while a genuinely wrong kernel would be O(1). + Three assertions, each on a quantity it can actually bound -- see the + "Whole-network gradient comparisons" note above for why that distinction is + the whole point here: + + * checkpointed vs. non-checkpointed **on the same rung**, per parameter. + This is the one the test is named for, and it is well posed because the + bf16 error is common-mode: it reads 3.9e-3, the model's own run-to-run + floor, under every convolution algorithm measured. + * compiled vs. eager under bf16 autocast, on the *aggregate* gradient + (5.0e-3 measured, against the same 3.2e-3 floor). + * compiled vs. eager **without autocast**, per parameter -- the sharp one, + 1.3e-6 measured against a 1e-4 tolerance. """ device = torch.device("cuda") x = _make_input(seed=9).to(device).requires_grad_(True) - tolerance = 5e-2 - def grads(compiled, checkpointing): + def grads(compiled, checkpointing, autocast=True): + gn_mod.set_triton_enabled(False) gn_mod.set_compile_enabled(compiled) model = _make_unet(seed=0).to(device) if checkpointing: model.use_checkpointing() model.zero_grad(set_to_none=True) - with torch.autocast("cuda", dtype=torch.bfloat16): + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast): out = model(x) out.float().pow(2).mean().backward() return {n: p.grad.detach().clone() for n, p in model.named_parameters()} - def assert_agrees(actual, expected, label): - for name in expected: - reference = expected[name].float() - error = (actual[name].float() - reference).norm().item() - relative = error / max(reference.norm().item(), 1e-12) - assert relative < tolerance, f"{label} {name}: rel L2 {relative:.3e}" - eager = grads(False, True) compiled = grads(True, True) compiled_nockpt = grads(True, False) assert gn_mod._compiled_group_norm is not None, "compiled path was not taken" assert not gn_mod._compile_failed - assert_agrees(compiled, eager, "checkpointed grad") - assert_agrees(compiled_nockpt, compiled, "grad") + + relative, name = _worst_relative_l2(compiled_nockpt, compiled) + assert relative < _BF16_TOLERANCE, f"grad {name}: rel L2 {relative:.3e}" + + relative = _relative_l2(compiled, eager) + assert relative < _BF16_TOLERANCE, f"checkpointed grad: rel L2 {relative:.3e}" + + relative, name = _worst_relative_l2( + grads(True, True, autocast=False), grads(False, True, autocast=False) + ) + assert relative < _FP32_RUNG_TOLERANCE, f"fp32 grad {name}: rel L2 {relative:.3e}" + + +@pytest.mark.gpu +def test_gpu_the_recompile_limit_holds_on_a_worker_thread(): + """Past Dynamo's stock 8 entries, compiling from a thread that never set it. + + ``torch._dynamo.config``'s user overrides live in a ``ContextVar``, so the + limit :func:`_raise_recompile_limit` writes on the main thread is not the + limit another thread reads -- and the compiles that matter happen on + another thread, because ``torch.utils.checkpoint``'s recompute runs inside + the backward pass, on the autograd engine's device worker. This drives the + same shape of traffic directly: entries live on ``_group_norm``'s code + object and are shared between threads, so a worker that pushes the count + past 8 is exactly the situation the recompute creates. Before the + per-region ``recompile_limit=``, the ninth compile raised + ``FailOnRecompileLimitHit`` here. + + ``torch._dynamo.reset()`` first because those entries also accumulate + across the whole test session, which would otherwise decide the outcome. + """ + import threading + + torch._dynamo.reset() + gn_mod.set_triton_enabled(False) # this is the compiled rung's limit + gn_mod.set_compile_enabled(True) + gn_mod._compile_failed = False + + device = torch.device("cuda") + # Nine distinct channel counts: nine cache entries, one more than the stock + # limit allows, and the one that overflows must land on the worker thread. + norms = [FastGroupNorm(_GROUPS, 8 * n).to(device) for n in range(1, 10)] + failures = [] + + def run(subset): + try: + for norm in subset: + norm(torch.randn(1, norm.num_channels, 2, 2, 2, device=device)) + except BaseException as error: # noqa: BLE001 - re-raised below + failures.append(error) + + run(norms[:2]) + worker = threading.Thread(target=run, args=(norms[2:],)) + worker.start() + worker.join() + + if failures: + raise AssertionError(f"compiling off the main thread failed: {failures[0]}") + assert gn_mod._compile_failed is False, "the rung latched itself off" + assert all(norm._compiled_ok for norm in norms), ( + "some module never had a call served by the compiled rung" + ) + + +@pytest.mark.gpu +def test_gpu_checkpointed_dctensor_recompute_keeps_the_compiled_rung(dc_cuda): + """The three-way combination that used to die: ckpt + compiled rung + DCTensor. + + ``activation_checkpointing: true`` with ``SCAFFOLD_GROUPNORM_TRITON=0`` on + DistConv activations is a supported configuration and it crashed: the + recompute reaches this module with ``__torch_function__`` subclass handling + *disabled* (DistConv's backward runs below it), which is part of Dynamo's + ``GLOBAL_STATE`` guard, so it misses every cache entry the forward built and + compiles a second set beside them -- twice the shapes, past 8 -- on the + autograd worker thread, where the module's raised limit was invisible. Each + pair of the three is fine on its own; all three together raised + ``FailOnRecompileLimitHit`` (at ``5943389``) or, once the ladder caught it + and dropped a *proven* module to eager mid-recompute, ``CheckpointError``. + + Five norms is the smallest count that reproduces it: 5 forward entries plus + 5 recompute entries is 10, and the ninth compile is the one that overflows. + The convolutions are what make the block's backward run below torch-function + (a bare unwrap does not), and the loss is taken on the ``DCTensor`` for the + same reason the trainer's is. + """ + import threading + + import torch.utils.checkpoint + + distconv, ps = dc_cuda + device = torch.device("cuda") + channels = (8, 16, 24, 32, 40) + + torch._dynamo.reset() + gn_mod.set_triton_enabled(False) + gn_mod.set_compile_enabled(True) + gn_mod._compile_failed = False + + torch.manual_seed(5) + norms = [FastGroupNorm(_GROUPS, c).to(device) for c in channels] + convs = [ + nn.Conv3d(previous, c, 1, bias=False).to(device) + for previous, c in zip((1,) + channels[:-1], channels) + ] + tail = nn.Conv3d(channels[-1], 1, 1, bias=False).to(device) + + # Where each GroupNorm call happens, as Dynamo's GLOBAL_STATE guard sees it. + states = set() + + def block(t): + for conv, norm in zip(convs, norms): + states.add( + ( + threading.current_thread() is threading.main_thread(), + torch._C._is_torch_function_enabled(), + ) + ) + t = norm(conv(t)) + return tail(t) + + x = torch.randn(1, 1, 4, 4, 4, device=device) + + def step(checkpointing): + for parameter in [x] + [ + p for m in convs + norms + [tail] for p in m.parameters() + ]: + parameter.grad = None + x.requires_grad_(True) + wrapped = distconv.DCTensor.from_shard(x, ps) + if checkpointing: + out = torch.utils.checkpoint.checkpoint(block, wrapped, use_reentrant=False) + else: + out = block(wrapped) + out.float().square().mean().backward() + return [norm.weight.grad.detach().clone() for norm in norms] + + checkpointed = step(True) + step(True) # a second step must not compile anything new either + direct = step(False) + + assert (False, False) in states, ( + "the recompute did not run below torch-function off the main thread; " + "this configuration no longer reproduces the guard split it targets" + ) + assert gn_mod._compile_failed is False, "the compiled rung latched itself off" + assert all(norm._compiled_ok for norm in norms), "a norm never ran compiled" + for index, (recomputed, plain) in enumerate(zip(checkpointed, direct)): + assert torch.isfinite(recomputed).all(), index + _assert_close(recomputed, plain, 1e-4, f"norm {index} weight grad") + + +# --------------------------------------------------------------------------- +# GPU behavior: the Triton rung +# --------------------------------------------------------------------------- + + +def _channels_last(t): + return t.is_contiguous(memory_format=torch.channels_last_3d) + + +@pytest.mark.gpu +@pytest.mark.parametrize("activation", [None, "relu"]) +@pytest.mark.parametrize("autocast", [False, True]) +def test_gpu_triton_matches_eager(activation, autocast): + """The Triton kernel is the default for channels-last input and matches eager. + + ``(1, 64, 32^3)`` channels-last is the production shape family at unit-test + size. Three claims at once: the routing really picks Triton when nothing is + forced; the values and gradients match the eager reference within + reduction-order noise; and the fused activation equals an explicit ReLU on + the eager result. + + The rung is established by spying on the entry point rather than by + inspecting the output's layout: every rung now returns the *input's* memory + format, deliberately (a fallback that returned contiguous re-broke the + channels-last chain for every convolution after it), so layout no longer + distinguishes them. The layout is asserted separately, of both. + """ + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(11) + x = torch.randn(1, 64, 32, 32, 32, device=device, generator=generator).to( + memory_format=torch.channels_last_3d + ) + grad_out = torch.randn(*x.shape, device=device, generator=generator) + + fast = FastGroupNorm(_GROUPS, 64, activation=activation).to(device) + with torch.no_grad(): + fast.weight.normal_(1.0, 0.1, generator=generator) + fast.bias.normal_(0.0, 0.1, generator=generator) + + calls = [] + original_triton_forward = FastGroupNorm._triton_forward + + def run(triton): + gn_mod.set_triton_enabled(triton) + gn_mod.set_compile_enabled(False) # eager reference, not Inductor + inp = x.clone().requires_grad_(True) + fast.zero_grad(set_to_none=True) + before = len(calls) + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast): + out = fast(inp) + out.backward(grad_out.to(out.dtype)) + return ( + out.detach(), + inp.grad.detach().clone(), + fast.weight.grad.detach().clone(), + fast.bias.grad.detach().clone(), + len(calls) - before, + ) + + def spy(self, local): + calls.append(tuple(local.shape)) + return original_triton_forward(self, local) + + FastGroupNorm._triton_forward = spy + try: + eager = run(False) + triton = run(None) # None = the production default, no override at all + finally: + FastGroupNorm._triton_forward = original_triton_forward + assert not gn_mod._triton_failed + + assert triton[4] == 1, "Triton path was not taken" + # ... and the control: the reference really did *not* take it, so the + # comparison below is between two kernels and not one kernel with itself. + assert eager[4] == 0 + # Both preserve the input's channels-last layout; that is the contract now, + # not a rung signature. + assert _channels_last(triton[0]) and _channels_last(eager[0]) + _assert_close(triton[0], eager[0], 1e-5, "output") + _assert_close(triton[1], eager[1], 1e-4, "d_input") + _assert_close(triton[2], eager[2], 1e-4, "d_weight") + _assert_close(triton[3], eager[3], 1e-4, "d_bias") + # Autocast's fp32 policy for GroupNorm must survive the swap. + assert triton[0].dtype == eager[0].dtype + if activation == "relu": + assert (triton[0] < 0).sum() == 0 + assert triton[0].max() > 0 + + +@pytest.mark.gpu +@pytest.mark.parametrize("activation", [None, "relu"]) +def test_gpu_triton_dctensor_matches_eager_and_stays_wrapped(dc_cuda, activation): + """The production configuration: DCTensor in, DCTensor out, NDHWC preserved. + + worker.py wraps every activation in a DCTensor even at + ``dc_num_shards=[1,1,1]``, so this -- not the plain-tensor case -- is the + path the benchmark actually runs. A producing convolution sits in front so + that the DCTensor handed to GroupNorm is a genuine intermediate: the unwrap + has to be the autograd-aware one or the gradient never reaches the conv. + """ + distconv, ps = dc_cuda + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(53) + x = torch.randn(1, 64, 16, 16, 16, device=device, generator=generator).to( + memory_format=torch.channels_last_3d + ) + + fast = FastGroupNorm(_GROUPS, 64, activation=activation).to(device) + producer = nn.Conv3d(64, 64, 1, bias=False).to( + device, memory_format=torch.channels_last_3d + ) + with torch.no_grad(): + fast.weight.normal_(1.0, 0.1, generator=generator) + fast.bias.normal_(0.0, 0.1, generator=generator) + + def plain(t): + return t._tensor if isinstance(t, distconv.DCTensor) else t + + def run(triton): + gn_mod.set_triton_enabled(triton) + gn_mod.set_compile_enabled(False) + inp = x.clone().requires_grad_(True) + fast.zero_grad(set_to_none=True) + producer.zero_grad(set_to_none=True) + out = fast(producer(distconv.DCTensor.from_shard(inp, ps))) + assert isinstance(out, distconv.DCTensor), "DCTensor did not survive" + local = distconv.distconv._ToTensor.apply(out) + local.float().pow(2).sum().backward() + assert inp.grad is not None, "gradient never reached the input" + assert producer.weight.grad is not None, "gradient never reached the producer" + return ( + local.detach().clone(), + inp.grad.detach().clone(), + plain(producer.weight.grad).detach().clone(), + plain(fast.weight.grad).detach().clone(), + plain(fast.bias.grad).detach().clone(), + ) + + eager = run(False) + triton = run(None) + assert not gn_mod._triton_failed + assert _channels_last(triton[0]), "Triton path was not taken (output not NDHWC)" + + _assert_close(triton[0], eager[0], 1e-5, "output") + for index, what in ( + (1, "d_input"), + (2, "d_producer"), + (3, "d_weight"), + (4, "d_bias"), + ): + _assert_close(triton[index], eager[index], 1e-4, what) + + +@pytest.mark.gpu +def test_gpu_unet_keeps_the_channels_last_chain(monkeypatch): + """The whole point: GroupNorm stops breaking the layout chain in the model. + + Before this kernel, every one of the model's GroupNorms consumed + ``channels_last_3d`` and emitted contiguous, forcing the next convolution to + convert back -- 22 breaks per scale-8 forward. A hook census asserts that + every ``FastGroupNorm`` invocation now takes NDHWC in *and* hands NDHWC out. + + Needs ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` in the environment for the + convolutions to emit channels-last at all; without it there is nothing to + preserve and the test skips rather than passing vacuously. + """ + device = torch.device("cuda") + model = _make_unet(seed=0).to(device, memory_format=torch.channels_last_3d) + x = _make_input(seed=9).to(device).contiguous(memory_format=torch.channels_last_3d) + + census = [] + + def hook(module, inputs, output): + census.append((_channels_last(inputs[0]), _channels_last(output))) + + for module in model.modules(): + if isinstance(module, FastGroupNorm): + module.register_forward_hook(hook) + + gn_mod.set_triton_enabled(None) + with torch.autocast("cuda", dtype=torch.bfloat16), torch.no_grad(): + model(x) + + assert census, "no GroupNorm ran" + if not any(seen_in for seen_in, _ in census): + pytest.skip( + "convolutions did not emit channels_last_3d; set " + "PYTORCH_MIOPEN_SUGGEST_NHWC=1 (the production setting)" + ) + breaks = [i for i, (seen_in, seen_out) in enumerate(census) if seen_in != seen_out] + assert not breaks, f"GroupNorm broke the layout chain at sites {breaks}" + assert all(seen_out for _, seen_out in census) + + +@pytest.mark.gpu +def test_gpu_unet_triton_matches_the_compiled_build(monkeypatch): + """Whole-model gradients with the Triton kernel vs. without it. + + Asked twice, on the two quantities the "Whole-network gradient comparisons" + note above establishes are bounded: the *aggregate* gradient under bf16 + autocast (1.9e-3 measured, against a 3.2e-3 run-to-run floor) and the + per-parameter gradient **without** autocast, which is the sharp one -- + 2.2e-6 measured, and the same 1e-4 tolerance the compiled rung is held to. + Per-parameter under autocast, which this test used to assert, reads 1.5e-2 + against its 5e-2 tolerance here for reasons that have nothing to do with + either kernel; see the note. + + Skips (rather than passing vacuously) when the convolutions are not emitting + channels-last, since the Triton kernel would then never engage. + """ + device = torch.device("cuda") + x = _make_input(seed=9).to(device).contiguous(memory_format=torch.channels_last_3d) + + engaged = [] + original = FastGroupNorm._triton_forward + + def spy(self, local): + engaged.append(tuple(local.shape)) + return original(self, local) + + monkeypatch.setattr(FastGroupNorm, "_triton_forward", spy) + + def grads(triton, autocast=True): + gn_mod.set_triton_enabled(triton) + gn_mod.set_compile_enabled(True) + model = _make_unet(seed=0).to(device, memory_format=torch.channels_last_3d) + model.zero_grad(set_to_none=True) + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast): + out = model(x) + out.float().pow(2).mean().backward() + return {n: p.grad.detach().clone() for n, p in model.named_parameters()} + + without = grads(False) + assert not engaged + with_triton = grads(None) + if not engaged: + pytest.skip( + "Triton kernel never engaged; set PYTORCH_MIOPEN_SUGGEST_NHWC=1 " + "(the production setting) so the convolutions emit channels-last" + ) + assert not gn_mod._triton_failed + + relative = _relative_l2(with_triton, without) + assert relative < _BF16_TOLERANCE, f"grad: rel L2 {relative:.3e}" + + engaged.clear() + without_fp32 = grads(False, autocast=False) + assert not engaged + with_triton_fp32 = grads(None, autocast=False) + assert engaged, "the Triton rung declined the same input without autocast" + assert not gn_mod._triton_failed + + relative, name = _worst_relative_l2(with_triton_fp32, without_fp32) + assert relative < _FP32_RUNG_TOLERANCE, f"fp32 grad {name}: rel L2 {relative:.3e}" diff --git a/tests/test_groupnorm_wiring_edge.py b/tests/test_groupnorm_wiring_edge.py new file mode 100644 index 0000000..b1e8526 --- /dev/null +++ b/tests/test_groupnorm_wiring_edge.py @@ -0,0 +1,1193 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""Edge cases of the GroupNorm *wiring* (``FastGroupNorm``'s three-rung ladder). + +Written as an adversarial review of commit ``ca58664``; ``tests/test_groupnorm.py`` +covers the happy paths and the routing predicates, this file covers the places +where the ladder, the latches and the absorbed ReLU interact with the rest of +torch. See ``work/gn-dctensor/WIRING_REVIEW.md`` for the full write-up. + +The review left ten of these as ``xfail(strict=True)``, one per defect, each +asserting the behaviour the module *should* have. All ten are fixed and the +markers are gone; the tests stay, now as regression guards. The properties +they pin, in the order the defects were found: + +* the fused activation is bit-identical to ``F.relu`` on NaN, the infinities + and ``-0.0``, forward and backward, on all three rungs -- and a NaN produced + under it still reaches the trainer's non-finite-loss abort; +* a latch may not change the rung a checkpointed block is *recomputed* on; +* the ladder catches "the kernel is broken" and nothing else -- not the + checkpoint machinery's control flow, not a user's saved-tensor hook, not an + exception from after the rung already saved something; +* ``torch.func`` is a routing question, not a kernel failure; +* a latch can be cleared; +* ``activation`` is validated where it is used, and a module pickled before it + existed still runs. +""" + +from __future__ import annotations + +import io + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ScaFFold.unet import group_norm as gn_mod +from ScaFFold.unet.group_norm import FastGroupNorm +from ScaFFold.unet.unet_model import UNet + +_GROUPS = 8 + + +@pytest.fixture(autouse=True) +def _restore_routing_state(): + """Keep per-test overrides of the module-level routing state contained. + + Same contract as ``tests/test_groupnorm.py``'s fixture: several tests here + deliberately trip a latch, which is a process global. + """ + previous_compile = gn_mod.set_compile_enabled(None) + previous_triton = gn_mod.set_triton_enabled(None) + compile_failed = gn_mod._compile_failed + triton_failed = gn_mod._triton_failed + yield + gn_mod._compile_override = previous_compile + gn_mod._triton_override = previous_triton + gn_mod._compile_failed = compile_failed + gn_mod._triton_failed = triton_failed + + +def _cl(t): + return t.is_contiguous(memory_format=torch.channels_last_3d) + + +def _cuda_norm(channels=64, activation="relu", size=16, seed=11): + """A seeded ``FastGroupNorm`` plus a channels-last CUDA input for it.""" + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(seed) + x = torch.randn(1, channels, size, size, size, device=device, generator=generator) + x = x.to(memory_format=torch.channels_last_3d) + module = FastGroupNorm(_GROUPS, channels, activation=activation).to(device) + with torch.no_grad(): + module.weight.normal_(1.0, 0.1, generator=generator) + module.bias.normal_(0.0, 0.1, generator=generator) + return module, x + + +def _small_unet(device="cpu", channels_last=False): + torch.manual_seed(0) + model = UNet( + n_channels=3, n_classes=2, trilinear=False, layers=1, group_norm_groups=_GROUPS + ) + if channels_last: + return model.to(device, memory_format=torch.channels_last_3d) + return model.to(device) + + +def _triton_spy(monkeypatch): + """Record every call that actually reached the Triton rung.""" + calls = [] + original = FastGroupNorm._triton_forward + + def spy(self, local): + calls.append(tuple(local.shape)) + return original(self, local) + + monkeypatch.setattr(FastGroupNorm, "_triton_forward", spy) + return calls + + +# --------------------------------------------------------------------------- +# activation semantics: every rung must apply the same function +# --------------------------------------------------------------------------- + + +def test_activate_handles_every_supported_activation(): + """``_activate`` must implement every activation the module advertises. + + ``SUPPORTED_ACTIVATIONS`` is what the *constructor* accepts and what + ``is_supported`` is asked about, but the compiled and eager rungs apply it + through ``_activate``, which tests one literal string. Adding a second + activation to both tuples (the only thing + ``test_supported_activations_match_the_kernels`` checks) would fuse it into + the Triton store and silently drop it everywhere else -- i.e. the network's + function would depend on the memory format of its input. This is the guard + on that: for every non-``None`` activation, ``_activate`` has to *change* + an input that the identity would leave alone. + """ + x = torch.linspace(-2.0, 2.0, 64).reshape(1, 8, 2, 2, 2) + for activation in gn_mod.SUPPORTED_ACTIVATIONS: + module = FastGroupNorm(_GROUPS, 8, activation=activation) + out = module._activate(x.clone()) + if activation is None: + assert torch.equal(out, x) + else: + assert not torch.equal(out, x), ( + f"_activate is a no-op for activation={activation!r}: the " + "compiled and eager rungs would silently skip it while the " + "Triton rung fused it in" + ) + + +def test_affine_false_still_applies_the_activation(): + """``affine=False`` leaves ``weight``/``bias`` ``None`` on every rung.""" + module = FastGroupNorm(_GROUPS, 16, affine=False, activation="relu") + x = torch.randn(1, 16, 4, 4, 4, generator=torch.Generator().manual_seed(5)) + assert torch.equal( + module(x), F.relu(F.group_norm(x, _GROUPS, None, None, module.eps)) + ) + assert list(module.state_dict().keys()) == [] + + +def test_activation_is_not_part_of_the_state(): + """``activation`` may not become a parameter, a buffer or a state-dict key.""" + with_relu = FastGroupNorm(_GROUPS, 16, activation="relu") + without = FastGroupNorm(_GROUPS, 16) + assert list(with_relu.state_dict().keys()) == list(without.state_dict().keys()) + assert list(with_relu.buffers()) == [] + # ... and a checkpoint written by one loads into the other, strict. + result = without.load_state_dict(with_relu.state_dict(), strict=True) + assert not result.missing_keys and not result.unexpected_keys + + +def test_double_conv_bytes_match_a_hand_built_pre_fusion_block(): + """Byte-for-byte state-dict identity against an independently built block. + + ``test_state_dict_bytes_identical_to_plain_groupnorm_model`` compares against + a model produced by *converting* the fused one, which shares its + construction order by definition. This builds the pre-fusion + ``nn.Sequential`` from scratch -- ``Conv3d, GroupNorm, ReLU, Conv3d, + GroupNorm, ReLU`` -- and compares the serialized bytes, which is the + independent version of the same claim. + """ + from ScaFFold.unet.unet_parts import DoubleConv + + torch.manual_seed(17) + fused = DoubleConv(3, 16, _GROUPS) + + torch.manual_seed(17) + reference = nn.Sequential( + nn.Conv3d(3, 16, kernel_size=3, padding=1, bias=False), + nn.GroupNorm(_GROUPS, 16), + nn.ReLU(inplace=True), + nn.Conv3d(16, 16, kernel_size=3, padding=1, bias=False), + nn.GroupNorm(_GROUPS, 16), + nn.ReLU(inplace=True), + ) + + def blob(state_dict): + buffer = io.BytesIO() + torch.save(state_dict, buffer) + return buffer.getvalue() + + assert list(fused.double_conv.state_dict().keys()) == list( + reference.state_dict().keys() + ) + assert blob(fused.double_conv.state_dict()) == blob(reference.state_dict()) + + +def test_module_pickled_before_the_fusion_still_runs(): + """A whole-module pickle predates ``self.activation``; forward must cope. + + ``nn.Module.__setstate__`` replaces ``__dict__`` wholesale, so an instance + restored from a ``torch.save(model)`` written before the fusion has no + ``activation`` at all -- and none of the routing state added since either. + Every attribute this module reads outside ``__init__`` therefore needs a + class-level default. + """ + module = FastGroupNorm(_GROUPS, 16, activation="relu") + state = module.__dict__.copy() + for added_since in ("activation", "_triton_ok", "_compiled_ok"): + state.pop(added_since, None) # exactly what a pre-fusion pickle carries + + restored = FastGroupNorm.__new__(FastGroupNorm) + nn.Module.__setstate__(restored, state) + + x = torch.randn(1, 16, 4, 4, 4, generator=torch.Generator().manual_seed(6)) + out = restored(x) + # A pre-fusion pickle had no activation, so it must behave as one. + assert torch.equal( + out, F.group_norm(x, _GROUPS, restored.weight, restored.bias, restored.eps) + ) + + +def test_unsupported_activation_assigned_after_construction_is_caught(): + """``activation`` is validated where it is *used*, not only at construction. + + It is a plain attribute, so it can be assigned afterwards; ``is_supported`` + would then decline the Triton rung while ``_activate`` silently applied + nothing, i.e. the module would quietly become a bare GroupNorm. The same + hole is the forward-looking risk: adding a third entry to both + ``SUPPORTED_ACTIVATIONS`` tuples without implementing it in ``_activate`` + must not produce a network whose activation depends on its input's memory + format. Failing loudly on the rung that cannot apply it closes both. + """ + module = FastGroupNorm(_GROUPS, 16, activation="relu") + module.activation = "gelu" + x = torch.randn(1, 16, 4, 4, 4, generator=torch.Generator().manual_seed(7)) + with pytest.raises(ValueError, match="activation must be one of"): + module(x) + + +# --------------------------------------------------------------------------- +# the ladder must not swallow torch's own control flow +# --------------------------------------------------------------------------- + + +def test_base_exceptions_are_not_caught(): + """``KeyboardInterrupt``/``SystemExit`` must escape the ladder untouched.""" + for exception in (KeyboardInterrupt, SystemExit): + + def _raises(*args, **kwargs): + raise exception() + + previous = gn_mod._get_triton_module + gn_mod._get_triton_module = _raises + original_use = gn_mod._use_triton + gn_mod._use_triton = lambda *a, **kw: True + gn_mod._triton_failed = False + try: + module = FastGroupNorm(_GROUPS, 16) + with pytest.raises(exception): + module(torch.randn(1, 16, 4, 4, 4)) + assert gn_mod._triton_failed is False + finally: + gn_mod._get_triton_module = previous + gn_mod._use_triton = original_use + + +@pytest.mark.parametrize("rung", ["triton", "compiled"]) +def test_checkpoint_error_is_re_raised(monkeypatch, rung): + """``CheckpointError`` is the checkpoint machinery talking, not a kernel. + + It is raised by the recompute pack hook -- i.e. from inside whichever op is + saving a tensor -- exactly like ``_StopRecomputationError``, and it is a + ``RuntimeError`` subclass, so any handler wide enough to catch "a broken + kernel" by type catches it too. Swallowing it latches the rung off, retries + on the next one and leaves the checkpoint frame in a state the machinery + never expected. The allowlist has to be narrow enough that this propagates + untouched and nothing latches. + """ + import torch.utils.checkpoint as checkpoint_mod + + def _raises(*args, **kwargs): + raise checkpoint_mod.CheckpointError("simulated recompute mismatch") + + if rung == "triton": + monkeypatch.setattr(gn_mod, "_use_triton", lambda *a, **kw: True) + monkeypatch.setattr(gn_mod, "_get_triton_module", _raises) + else: + monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input, **kw: True) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _raises) + gn_mod._triton_failed = False + gn_mod._compile_failed = False + + module = FastGroupNorm(_GROUPS, 16) + with pytest.raises(checkpoint_mod.CheckpointError): + module(torch.randn(1, 16, 4, 4, 4)) + assert gn_mod._triton_failed is False + assert gn_mod._compile_failed is False + + +def test_a_rung_failure_does_not_re_fire_saved_tensor_hooks(monkeypatch): + """The retry has to be idempotent with respect to saved-tensor hooks. + + Under non-reentrant activation checkpointing the recompute counts pack-hook + firings and requires the count *and* the metadata to match the forward's, so + a rung that packed some tensors and then failed -- with the fallback packing + its own set on top -- corrupts the frame. A user's offloading hook has the + same problem in a less dramatic way (an offload failure retried by + offloading a second, larger set). + + The property is structural rather than defensive: the ladder catches only + failures that are raised *before* their rung saves anything (the Triton op + saves in ``_setup_context``, after the launch region its ``TritonKernelError`` + comes from; a Dynamo/Inductor error is a compile-time error, before + execution). Both halves are asserted here -- a caught failure packs exactly + what a clean fallback packs, and an exception from *after* the packing is + not the ladder's to swallow. + """ + import torch._dynamo.exc + + class _Boom(Exception): + pass + + def _fails_before_packing(input, num_groups, weight, bias, eps): + raise torch._dynamo.exc.Unsupported("failed while compiling") + + def _fails_after_packing(input, num_groups, weight, bias, eps): + F.group_norm(input, num_groups, weight, bias, eps) + raise _Boom("failed after packing") + + monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input, **kw: True) + module = FastGroupNorm(_GROUPS, 16) + x = torch.randn(1, 16, 4, 4, 4).requires_grad_(True) + + def packs_during(run): + packed = [] + with torch.autograd.graph.saved_tensors_hooks( + lambda t: (packed.append(1), t)[1], lambda t: t + ): + run() + return len(packed) + + baseline = packs_during(lambda: module._eager_forward(x)) + + monkeypatch.setattr( + gn_mod, "_get_compiled_group_norm", lambda: _fails_before_packing + ) + gn_mod._compile_failed = False + retried = packs_during(lambda: module(x)) + assert retried == baseline, ( + f"the failed rung packed {retried - baseline} extra tensors before the " + "fallback ran" + ) + + # ... and a failure that *did* have observable effects is not retried at all. + monkeypatch.setattr( + gn_mod, "_get_compiled_group_norm", lambda: _fails_after_packing + ) + gn_mod._compile_failed = False + with pytest.raises(_Boom): + module(x) + + +# --------------------------------------------------------------------------- +# latches +# --------------------------------------------------------------------------- + + +def test_forcing_a_rung_on_clears_its_failure_latch(): + """``set_*_enabled(True)`` is the documented way to retry after a failure. + + Without this a one-off failure (a transient OOM, a cache-directory hiccup) + costs the rung for the rest of the process with no recovery at all, and the + function's own docstring -- "forcing it on is overridden only by the + correctness checks" -- is false. ``None`` deliberately does *not* clear it: + that restores a preference, it does not assert that the kernel works again. + """ + for setter, latch in ( + (gn_mod.set_triton_enabled, "_triton_failed"), + (gn_mod.set_compile_enabled, "_compile_failed"), + ): + setattr(gn_mod, latch, True) + setter(True) + assert getattr(gn_mod, latch) is False + + setattr(gn_mod, latch, True) + setter(None) + assert getattr(gn_mod, latch) is True + setter(False) + assert getattr(gn_mod, latch) is True + setattr(gn_mod, latch, False) + + +@pytest.mark.parametrize("rung", ["triton", "compiled"]) +def test_out_of_memory_is_not_recorded_as_a_kernel_failure(monkeypatch, rung): + """A transient OOM must propagate, and must not latch a rung off forever. + + ``torch.OutOfMemoryError`` is a resource condition, not a defect: every + fallback allocates an output of the same size, so retrying one is a second, + differently-shaped OOM at a call site the caller never asked about. + Latching on it is worse still -- a per-rank, nondeterministic event that + permanently changes which kernel that rank runs, and therefore (measured) + the all-reduced gradients of the whole job. + """ + from ScaFFold.unet.triton_group_norm import TritonKernelError + + def _oom(*args, **kwargs): + raise torch.OutOfMemoryError("simulated OOM") + + if rung == "triton": + monkeypatch.setattr(gn_mod, "_use_triton", lambda *a, **kw: True) + monkeypatch.setattr(gn_mod, "_get_triton_module", _oom) + else: + monkeypatch.setattr(gn_mod, "_use_compiled", lambda _input, **kw: True) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _oom) + gn_mod._triton_failed = False + gn_mod._compile_failed = False + + module = FastGroupNorm(_GROUPS, 16) + with pytest.raises(torch.OutOfMemoryError): + module(torch.randn(1, 16, 4, 4, 4)) + assert gn_mod._triton_failed is False + assert gn_mod._compile_failed is False + # It is a RuntimeError, so a handler that caught the kernel's own error by + # base class would have swallowed it; the allowlist is the tagged type. + assert issubclass(torch.OutOfMemoryError, RuntimeError) + assert not issubclass(torch.OutOfMemoryError, TritonKernelError) + + +def test_a_global_latch_does_not_demote_a_module_that_already_used_the_rung( + monkeypatch, caplog +): + """The unit of the latch is the module, not the process. + + A rung that has already served a module keeps serving it; only modules that + have never used it are steered away. That is what makes a checkpointed + block's forward and its recompute agree (they save different tensors on + different rungs, so a mid-graph change is fatal), and it is why a broken + install still costs exactly one attempt per module rather than one per call. + + The corollary tested here too: because a proven module keeps retrying, the + warning has to be emitted on the latch's edge rather than per call, or a + persistently broken kernel floods the log for the rest of the run. + """ + import logging + + import torch._dynamo.exc + + calls = [] + + def _kernel(input, num_groups, weight, bias, eps): + calls.append(1) + if len(calls) > 1: + raise torch._dynamo.exc.Unsupported("simulated Inductor failure") + return F.group_norm(input, num_groups, weight, bias, eps) + + monkeypatch.setattr(gn_mod, "_use_compiled", gn_mod._use_compiled) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: _kernel) + monkeypatch.setattr( + gn_mod, + "_use_compiled", + lambda t, proven=False, **kw: (not gn_mod._compile_failed) or proven, + ) + gn_mod._compile_failed = False + + proven = FastGroupNorm(_GROUPS, 16) + fresh = FastGroupNorm(_GROUPS, 16) + x = torch.randn(1, 16, 4, 4, 4) + + proven(x) # succeeds: this module is now proven on the compiled rung + assert proven._compiled_ok is True + + with caplog.at_level(logging.WARNING, logger=gn_mod.__name__): + proven(x) # fails, latches, falls back to eager + assert gn_mod._compile_failed is True + warned = sum("falling back" in r.message for r in caplog.records) + proven(x) # ... and still *tries* the rung, because it is proven + proven(x) + assert sum("falling back" in r.message for r in caplog.records) == warned, ( + "a persistently failing rung warned once per call" + ) + assert len(calls) == 4, "the proven module stopped trying its rung" + + fresh(x) # never used it, so the global latch keeps it away entirely + assert len(calls) == 4 + assert fresh._compiled_ok is False + + +@pytest.mark.parametrize("rung", ["triton", "compiled"]) +def test_a_proven_rung_does_not_degrade_while_a_backward_replays_it(monkeypatch, rung): + """A fallback *during a recompute* corrupts rather than degrades. + + The latch already refuses to demote a module that has used a rung, but the + fallback itself sidestepped that: the failing call still got answered from + the next rung down, and if that call is ``torch.utils.checkpoint``'s + recompute of a forward that ran on the failing rung, the recomputed forward + saves a different set of tensors than the original did and torch rejects + the whole step (``CheckpointError``; on one measured shape a GPU memory + fault instead). Neither is a degradation, so this one case re-raises. + + It is narrow on purpose -- ``_replaying_a_forward()`` is false in an + ordinary forward, where ``test_a_global_latch_does_not_demote_a_module_...`` + still requires the fallback -- and the second half here pins the other side + of the narrowness: a module that has *not* used the rung degrades even + inside the backward, because the forward it is replaying went down the + ladder too and the two agree. + """ + import torch._dynamo.exc + import torch.utils.checkpoint as checkpoint_mod + + from ScaFFold.unet.triton_group_norm import TritonKernelError + + failure = TritonKernelError if rung == "triton" else torch._dynamo.exc.Unsupported + latch = "_triton_failed" if rung == "triton" else "_compile_failed" + proven_flag = "_triton_ok" if rung == "triton" else "_compiled_ok" + + def make_kernel(fail_always): + def _kernel(input, num_groups, weight, bias, eps, *activation): + if fail_always or gn_mod._replaying_a_forward(): + raise failure("simulated kernel failure") + return F.group_norm(input, num_groups, weight, bias, eps) + + return _kernel + + def run(fail_always): + gn_mod._triton_failed = False + gn_mod._compile_failed = False + kernel = make_kernel(fail_always) + if rung == "triton": + monkeypatch.setattr(gn_mod, "_use_triton", lambda *a, **kw: True) + module_stub = type("_Stub", (), {"triton_group_norm": staticmethod(kernel)}) + monkeypatch.setattr(gn_mod, "_get_triton_module", lambda: module_stub) + else: + monkeypatch.setattr(gn_mod, "_use_compiled", lambda t, **kw: True) + monkeypatch.setattr(gn_mod, "_get_compiled_group_norm", lambda: kernel) + module = FastGroupNorm(_GROUPS, 16) + x = torch.randn(1, 16, 4, 4, 4, requires_grad=True) + out = checkpoint_mod.checkpoint(module, x, use_reentrant=False) + out.pow(2).sum().backward() + return module, x + + # Proven in the forward, failing in the recompute: answering from another + # rung would be the metadata mismatch, so the failure has to come back out. + with pytest.raises(failure): + run(fail_always=False) + + # Never served by the rung: the forward already went down the ladder, so + # the recompute doing the same agrees with it and the step survives. + module, x = run(fail_always=True) + assert getattr(module, proven_flag) is False + assert getattr(gn_mod, latch) is True + assert torch.isfinite(x.grad).all() + + +def test_a_latch_flip_mid_forward_is_not_a_numerics_error(monkeypatch): + """Two rungs inside one forward still compose (values, not bits, agree).""" + monkeypatch.setattr( + gn_mod, "_use_compiled", lambda t, **kw: type(t) is torch.Tensor + ) + monkeypatch.setattr( + gn_mod, "_get_compiled_group_norm", lambda: nn.functional.group_norm + ) + gn_mod._compile_failed = False + + module = FastGroupNorm(_GROUPS, 16, activation="relu") + x = torch.randn(1, 16, 4, 4, 4).requires_grad_(True) + first = module(x) + gn_mod._compile_failed = True # latch flips between the two calls + second = module(first) + second.pow(2).sum().backward() + assert torch.isfinite(x.grad).all() + + +# --------------------------------------------------------------------------- +# predicates +# --------------------------------------------------------------------------- + + +def test_predicates_reject_a_parameter_input(): + """``nn.Parameter`` is a subclass, so both fast rungs decline it. + + Not a bug -- the model never feeds a Parameter to a norm -- but it is the + documented consequence of the ``type(input) is torch.Tensor`` policy, and a + regression that loosened it to ``isinstance`` would route real + ``__torch_dispatch__`` wrappers into the kernel. + """ + parameter = nn.Parameter(torch.randn(1, 8, 4, 4, 4)) + assert gn_mod._use_triton(parameter, _GROUPS, None, None, None) is False + assert gn_mod._use_compiled(parameter) is False + + +def test_use_triton_is_side_effect_free_for_rejected_inputs(monkeypatch): + """The predicate may look, but it may not allocate, launch or mutate.""" + module = FastGroupNorm(_GROUPS, 16) + x = torch.randn(1, 16, 4, 4, 4) + before = x.clone() + assert gn_mod._use_triton(x, _GROUPS, module.weight, module.bias, None) is False + assert torch.equal(x, before) + + +# --------------------------------------------------------------------------- +# GPU: the Triton rung in situ +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +def test_gpu_double_backward_fails_loudly(): + """The kernel is first-order only; a second derivative must *raise*. + + ``triton_group_norm``'s backward is itself a custom op with no autograd + formula, so a gradient penalty or an HVP through the wired model has to die + with a clear message rather than silently return a wrong number -- and the + failure must not be mistaken for a broken kernel and latch the rung off. + """ + module, x = _cuda_norm(activation=None) + x = x.clone().requires_grad_(True) + out = module(x) + assert _cl(out), "Triton rung was not taken; the test would be vacuous" + + (first,) = torch.autograd.grad(out.pow(2).sum(), x, create_graph=True) + with pytest.raises(RuntimeError, match="no autograd formula"): + torch.autograd.grad(first.pow(2).sum(), x) + assert gn_mod._triton_failed is False, "a caller error latched the kernel off" + + +@pytest.mark.gpu +def test_gpu_model_double_backward_fails_loudly(): + """Same, through the wired model, which is where a user would hit it. + + Either rung may raise first and the message differs, so this matches both: + ``triton_group_norm``'s backward is a custom op with **no autograd formula**, + and the Triton convolution's is marked **@once_differentiable**. Which one + the traversal reaches first is a routing detail -- before the block-list was + emptied on 2026-08-04 the ``Cin == 3`` stem was on MIOpen, which *is* twice + differentiable, so the walk got all the way to GroupNorm. Now the stem's own + backward stops it, one op earlier. + + What must not change is that it raises *at all*. Neither decorator was in + place on the convolution until that routing change, and without it the + failure was ``first`` silently arriving with no ``grad_fn`` -- a second + backward then contributing zero rather than erroring. + """ + model = _small_unet("cuda", channels_last=True) + x = ( + torch.randn(1, 3, 16, 16, 16, device="cuda") + .contiguous(memory_format=torch.channels_last_3d) + .requires_grad_(True) + ) + gn_mod.set_triton_enabled(None) + out = model(x) + (first,) = torch.autograd.grad(out.pow(2).sum(), x, create_graph=True) + assert first.grad_fn is not None, ( + "the double-backward graph was severed instead of raising; a second " + "backward would contribute zero silently" + ) + with pytest.raises( + RuntimeError, match="no autograd formula|once_differentiable" + ): + first.pow(2).sum().backward() + + +@pytest.mark.gpu +@pytest.mark.parametrize("fullgraph", [True, False]) +def test_gpu_triton_rung_inside_a_compiled_region(monkeypatch, fullgraph): + """The Triton rung is *allowed* inside ``torch.compile``; prove it works. + + ``_use_compiled`` bails out when ``torch.compiler.is_compiling()`` so the + functional GroupNorm inlines, but ``_use_triton`` has no such guard: an + enclosing compiled region traces straight into the custom op. Nothing in + ScaFFold compiles ``FastGroupNorm.forward`` today, so this was untested in + situ. Values, gradients *and* the channels-last output must survive + Dynamo/AOTAutograd unchanged. + """ + import torch._dynamo + + module, x = _cuda_norm(activation="relu") + reference_input = x.clone().requires_grad_(True) + reference = module(reference_input) + reference.pow(2).sum().backward() + reference_grad = reference_input.grad.detach().clone() + module.zero_grad(set_to_none=True) + assert _cl(reference), "Triton rung was not taken; the test would be vacuous" + + torch._dynamo.reset() + calls = _triton_spy(monkeypatch) + compiled = torch.compile(lambda t: module(t), fullgraph=fullgraph, dynamic=False) + + compiled_input = x.clone().requires_grad_(True) + out = compiled(compiled_input) + out.pow(2).sum().backward() + + assert calls, "the Triton rung was not traced inside the compiled region" + assert _cl(out), "the compiled region lost the channels-last output" + assert torch.equal(out, reference) + assert torch.equal(compiled_input.grad, reference_grad) + + +@pytest.mark.gpu +@pytest.mark.parametrize("proven", [False, True]) +def test_gpu_the_fallback_path_traces_under_fullgraph(monkeypatch, proven): + """A rung failure *while Dynamo is tracing* must still fall back, not die. + + The handler used to call ``logger.warning``, which Dynamo cannot trace + ("Unsupported: logging.Logger method not supported for non-export cases"), + so a caller compiling this forward with ``fullgraph=True`` got a hard error + instead of the fallback -- the one caller for whom the fallback matters + most, since the thing it is reacting to is usually a compile-time failure. + Nothing in ScaFFold compiles ``FastGroupNorm.forward`` today; this pins the + claim that it can. + + Both halves of the handler's guard have to trace, which is why ``proven`` + is parametrized: with ``_triton_ok`` false Dynamo folds the ``and`` away + without ever looking at ``_replaying_a_forward()``, so only the ``True`` + arm reaches it -- and a probe Dynamo cannot trace there would be the same + defect as the logging call, reintroduced. + """ + import torch._dynamo + + from ScaFFold.unet.triton_group_norm import TritonKernelError + + module, x = _cuda_norm(activation="relu") + reference = module(x).detach().clone() + assert _cl(reference), "Triton rung was not taken; the test would be vacuous" + + real = gn_mod._get_triton_module() + + class _BrokenKernelModule: + def __getattr__(self, name): + if name == "triton_group_norm": + + def _raises(*args, **kwargs): + raise TritonKernelError("simulated Triton failure") + + return _raises + return getattr(real, name) + + monkeypatch.setattr(gn_mod, "_get_triton_module", _BrokenKernelModule) + gn_mod._triton_failed = False + module._triton_ok = proven + torch._dynamo.reset() + + compiled = torch.compile(lambda t: module(t), fullgraph=True, dynamic=False) + out = compiled(x) + + assert gn_mod._triton_failed is True, "the latch was not recorded" + assert _cl(out), "the fallback rung dropped the channels-last chain" + assert (out - reference).abs().max().item() < 1e-5 + + +@pytest.mark.gpu +def test_gpu_inference_mode_takes_the_triton_rung(monkeypatch): + """``evaluate()`` runs the whole model under ``torch.inference_mode``.""" + module, x = _cuda_norm(activation="relu") + module.eval() + reference = F.relu(F.group_norm(x, _GROUPS, module.weight, module.bias, module.eps)) + calls = _triton_spy(monkeypatch) + with torch.inference_mode(): + out = module(x) + assert calls, "inference_mode fell off the Triton rung" + assert _cl(out) + assert out.is_inference() + assert (out.float() - reference.float()).abs().max().item() < 1e-5 + + +@pytest.mark.gpu +def test_gpu_evaluation_shaped_forward_matches_training_shaped_one(monkeypatch): + """``eval()`` + ``inference_mode`` + autocast is the evaluate() combination.""" + model = _small_unet("cuda", channels_last=True) + x = torch.randn(1, 3, 16, 16, 16, device="cuda").contiguous( + memory_format=torch.channels_last_3d + ) + gn_mod.set_triton_enabled(None) + model.eval() + calls = _triton_spy(monkeypatch) + with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): + out = model(x) + if not calls: + pytest.skip( + "convolutions did not emit channels_last_3d; set " + "PYTORCH_MIOPEN_SUGGEST_NHWC=1 (the production setting)" + ) + assert torch.isfinite(out.float()).all() + assert gn_mod._triton_failed is False + + +@pytest.mark.gpu +def test_gpu_vmap_over_the_module_still_works(): + """``torch.func`` is a routing question, not a kernel failure. + + ``is_supported``'s ``is_contiguous(memory_format=...)`` raises outright + under a ``vmap`` layer ("NYI"), so a predicate that reaches it -- or a + relayout helper that does -- turns a plain ``nn.GroupNorm`` drop-in into a + hard error for any caller using ``torch.func``. Both fast rungs must + decline while a transform is active and let the stock kernel answer. + """ + module, x = _cuda_norm(activation=None) + batched = torch.stack([x[0], x[0]]) + out = torch.func.vmap(lambda t: module(t.unsqueeze(0)).squeeze(0))(batched) + assert out.shape == batched.shape + assert torch.allclose(out[0], module(x[None, 0]).squeeze(0), atol=1e-5) + + +@pytest.mark.gpu +def test_gpu_a_predicate_that_cannot_answer_falls_back_without_latching( + monkeypatch, caplog +): + """``is_supported`` raising is a routing miss, and "no" is always a valid answer. + + The predicate runs *outside* the ladder's try, so anything it raises escapes + ``forward()`` -- which is how a ``torch.func`` transform used to turn a + drop-in ``nn.GroupNorm`` into a hard error. The functorch check upstream + covers the one caller known to trip it; this covers the shape of the + problem, because ``is_supported`` inspects an *arbitrary* tensor and the set + of wrappers that can make an attribute read raise is not closed. A broad + catch is right here and nowhere else in this module: the predicate has done + no work anyone can observe and a correct answer ("use the stock kernel") is + always available -- so it must fall back, and must not latch, because + nothing about the kernel has been learned. + """ + import logging + + class _Unanswerable: + def __getattr__(self, name): + if name == "is_supported": + + def _raises(*args, **kwargs): + raise RuntimeError("NYI: querying is_contiguous inside of vmap") + + return _raises + return getattr(gn_mod._get_triton_module(), name) + + module, x = _cuda_norm(activation="relu") + reference = module(x).detach().clone() + monkeypatch.setattr(gn_mod, "_get_triton_module", _Unanswerable) + monkeypatch.setattr(gn_mod, "_predicate_warned", False) + gn_mod._triton_failed = False + module._triton_ok = False + + with caplog.at_level(logging.WARNING, logger=gn_mod.__name__): + out = module(x) + assert (out - reference).abs().max().item() < 1e-5 + assert gn_mod._triton_failed is False, "a routing miss latched the kernel off" + assert gn_mod._compile_failed is False + assert any("routing check failed" in r.message for r in caplog.records) + + +@pytest.mark.gpu +def test_gpu_torch_func_grad_does_not_latch_the_rungs_off(): + """A ``torch.func`` call anywhere must not demote the whole process. + + ``torch.func.grad`` used to reach the kernel, fail, and latch *both* fast + rungs off permanently -- i.e. one transform anywhere in a process silently + dropped every GroupNorm in the model to the stock kernel for the rest of + the run. Recording a routing miss as a kernel failure is the general shape + of the bug; this pins the specific instance. + """ + module, x = _cuda_norm(activation=None) + gn_mod._triton_failed = False + gn_mod._compile_failed = False + + grad = torch.func.grad(lambda t: module(t).pow(2).sum())(x) + + assert torch.isfinite(grad).all() + assert gn_mod._triton_failed is False + assert gn_mod._compile_failed is False + # ... and the module is still on the Triton rung afterwards. + assert _cl(module(x)) + + +@pytest.mark.gpu +def test_gpu_a_triton_failure_during_a_checkpointed_step_degrades_not_dies(): + """A latch may not change the rung a checkpointed block is recomputed on. + + Non-reentrant checkpointing compares the metadata of every tensor the + recomputed forward saves against the forward's, and the three rungs save + *different* tensors -- Triton ``(input, weight, bias, mean, rstd)``, the + other two ``(input, weight, mean, rstd, relu_output)``. So a rung change + between a block's forward and its recompute kills the step with a + ``CheckpointError``, which is the exact opposite of the ladder's contract + ("a broken Triton install must degrade a multi-node run, not kill it") and + is reachable whenever the ``activation_checkpointing`` option is on + (``worker.py:230``). Matching the output memory format is *not* enough on + its own -- measured; the saved sets still differ -- so the fix is that a + global latch does not demote a module that has already used the rung. + + The same hazard predates the Triton rung: flipping ``_compile_failed`` + between forward and recompute dies too, which is why both latches are + checked here. + """ + for latch in ("_triton_failed", "_compile_failed"): + model = _small_unet("cuda", channels_last=True) + model.use_checkpointing() + x = ( + torch.randn(1, 3, 16, 16, 16, device="cuda") + .contiguous(memory_format=torch.channels_last_3d) + .requires_grad_(True) + ) + gn_mod.set_triton_enabled(None) + gn_mod._triton_failed = False + gn_mod._compile_failed = False + + out = model(x) + # A one-off failure at any *later* GroupNorm site latches the rung off + # while the blocks already run are waiting to be recomputed. + setattr(gn_mod, latch, True) + out.pow(2).sum().backward() + + assert torch.isfinite(x.grad).all(), latch + + +@pytest.mark.gpu +@pytest.mark.parametrize("channels_last", [True, False]) +def test_gpu_every_rung_returns_the_inputs_memory_format(channels_last): + """All three rungs must agree on the output layout, not just the values. + + ``F.group_norm`` -- eager or Inductor-compiled -- returns a *contiguous* + tensor whatever it was given, so a single fallback used to re-break the + channels-last chain for every convolution after it, which is the exact + thing this module exists to prevent. It also made the rungs distinguishable + to anything that inspects metadata (``torch.utils.checkpoint``, a compiled + caller's guards), which is a correctness problem rather than a speed one. + """ + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(9) + x = torch.randn(1, 64, 8, 8, 8, device=device, generator=generator) + x = x.to(memory_format=torch.channels_last_3d) if channels_last else x.contiguous() + module = FastGroupNorm(_GROUPS, 64, activation="relu").to(device) + + outputs = _run_on_every_rung(module, x) + for label, out in outputs.items(): + assert _cl(out) is channels_last, ( + f"the {label} rung returned " + f"{'channels_last_3d' if _cl(out) else 'contiguous'} for a " + f"{'channels_last_3d' if channels_last else 'contiguous'} input" + ) + reference = outputs["eager"].detach().float() + for label, out in outputs.items(): + assert (out.detach().float() - reference).abs().max().item() < 1e-5, label + + +@pytest.mark.gpu +def test_gpu_triton_rejects_a_cuda_tensor_subclass(): + """The subclass check has to be tested on a tensor that would otherwise pass. + + ``test_triton_rejects_unknown_tensor_subclasses`` hands ``_use_triton`` a + *CPU* subclass, which the ``is_cuda`` check rejects one line later -- so it + cannot tell whether the ``type(input) is torch.Tensor`` test exists at all + (a mutation deleting that line survives the whole suite). The check is + load-bearing: ``is_supported`` only asks ``isinstance``, so without it every + unknown ``__torch_dispatch__`` wrapper would be routed into the kernel. + """ + + class _Wrapper(torch.Tensor): + pass + + x = torch.randn(1, 64, 8, 8, 8, device="cuda").to( + memory_format=torch.channels_last_3d + ) + wrapped = x.as_subclass(_Wrapper) + # The control: everything *except* the subclass test accepts this tensor. + from ScaFFold.unet import triton_group_norm as triton_mod + + assert triton_mod.is_supported(wrapped, _GROUPS, None, None, None) is True + assert gn_mod._use_triton(wrapped, _GROUPS, None, None, None) is False + assert gn_mod._use_compiled(wrapped) is False + + +#: NaN, +Inf, -Inf, -0.0 and four ordinary values -- everything the fused +#: activation has to agree with ``F.relu`` on. ``tl.maximum(y, 0)`` and +#: ``tl.where(y > 0, y, 0)`` both map NaN to 0.0 (the first returns the non-NaN +#: operand, the second because ``NaN > 0`` is False); ``F.relu`` propagates it. +_SPECIAL_VALUES = [float("nan"), float("inf"), float("-inf"), -0.0, 0.0, -1.0, 1.0, 2.0] + + +def _run_on_every_rung(module, x): + """``{rung: output}`` for the same module and input on all three rungs.""" + results = {} + for label, triton, compiled in ( + ("triton", True, False), + ("compiled", False, True), + ("eager", False, False), + ): + gn_mod.set_triton_enabled(triton) + gn_mod.set_compile_enabled(compiled) + results[label] = module(x) + return results + + +def _bits(t): + return t.detach().float().cpu().contiguous().view(torch.int32) + + +@pytest.mark.gpu +@pytest.mark.parametrize("activation", ["relu", None]) +@pytest.mark.parametrize("poison", [float("nan"), float("inf"), float("-inf")]) +def test_gpu_all_rungs_agree_on_nan_and_inf(poison, activation): + """One non-finite input value must poison the same elements on every rung. + + The Triton store used ``tl.maximum(y, 0.0)``, which returns the *non*-NaN + operand, so a diverging activation came back from the Triton rung as a + finite 0.0 while ``F.relu`` on the other two kept it NaN. That is worse + than a numerics discrepancy: the forward looks finite while the backward is + still NaN, so the run sails past ScaFFold's non-finite-loss abort and + checkpoints a broken model -- and the model's output becomes a function of + its input's memory format. ``activation=None`` is the control: all three + agreed there even before the fix, which is what localizes the divergence to + the fused activation. + """ + device = torch.device("cuda") + x = torch.randn(1, 64, 4, 4, 4, device=device) + x.view(-1)[0] = poison + x = x.to(memory_format=torch.channels_last_3d) + module = FastGroupNorm(_GROUPS, 64, activation=activation).to(device) + + results = { + label: out.detach().float().cpu().contiguous().isnan() + for label, out in _run_on_every_rung(module, x).items() + } + assert int(results["eager"].sum()) > 0, "the poison did not reach the output" + assert torch.equal(results["compiled"], results["eager"]) + assert torch.equal(results["triton"], results["eager"]), ( + f"the fused activation turned {int(results['eager'].sum())} NaNs into " + f"{int(results['triton'].sum())}" + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("activation", ["relu", None]) +def test_gpu_fused_activation_is_bit_identical_to_relu(activation): + """Every special value of the *pre-activation*, on every rung, bit for bit. + + Poisoning the input can only produce NaN pre-activations (one NaN or Inf + makes the whole group's statistics NaN), so the four values that actually + distinguish the spellings of ReLU are reached the other way round: a zero + ``weight`` makes the pre-activation exactly ``bias``, elementwise, so the + bias vector chooses what the activation sees. Expected, per + ``F.relu``: NaN stays NaN, ``+Inf`` stays ``+Inf``, ``-Inf`` and both zeros + become ``+0.0`` (never ``-0.0``). + + With ``activation=None`` the zeros are normalized (``+ 0.0`` maps ``-0.0`` + to ``+0.0`` and leaves NaN, the infinities and every normal value alone) + before the same bitwise comparison: a ``-0.0`` bias survives to the output + there, and whether ``xhat * 0 + (-0.0)`` keeps the sign depends on whether + the kernel contracted the multiply-add into an FMA -- true of the Triton + *and* the Inductor rung, false of eager, and nothing to do with the + activation. ``torch.equal`` is no use for either case: it reports NaN as + unequal to itself. + """ + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(3) + x = torch.randn(1, 64, 4, 4, 4, device=device, generator=generator).to( + memory_format=torch.channels_last_3d + ) + module = FastGroupNorm(_GROUPS, 64, activation=activation).to(device) + with torch.no_grad(): + module.weight.zero_() + module.bias.copy_(torch.tensor(_SPECIAL_VALUES * 8, device=device)) + + reference = F.group_norm(x, _GROUPS, module.weight, module.bias, module.eps) + if activation == "relu": + reference = F.relu(reference) + + for label, out in _run_on_every_rung(module, x).items(): + if activation == "relu": + assert torch.equal(_bits(out), _bits(reference)), ( + f"{label} rung differs from F.relu(F.group_norm(...)) in the " + "bit pattern of at least one special value" + ) + else: + assert torch.equal(_bits(out + 0.0), _bits(reference + 0.0)), label + + +@pytest.mark.gpu +def test_gpu_fused_relu_backward_gates_like_threshold_backward(): + """ReLU's backward passes the gradient where the output is NaN, too. + + ``threshold_backward(grad, result, 0)`` zeroes where ``result <= 0``, and + ``NaN <= 0`` is False -- so a NaN pre-activation passes its gradient. The + kernel recomputes the pre-activation and must gate with the same + complement; ``pre > 0 ? dy : 0`` would silently zero it. + """ + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(4) + base = torch.randn(1, 64, 4, 4, 4, device=device, generator=generator).to( + memory_format=torch.channels_last_3d + ) + module = FastGroupNorm(_GROUPS, 64, activation="relu").to(device) + with torch.no_grad(): + module.weight.zero_() + module.bias.copy_(torch.tensor(_SPECIAL_VALUES * 8, device=device)) + + grads = {} + for label, triton in (("triton", True), ("eager", False)): + gn_mod.set_triton_enabled(triton) + gn_mod.set_compile_enabled(False) + x = base.clone().requires_grad_(True) + module.zero_grad(set_to_none=True) + module(x).sum().backward() + grads[label] = (x.grad.detach().clone(), module.bias.grad.detach().clone()) + + # d_bias is exactly the gate: one per element that passed. + assert torch.equal(_bits(grads["triton"][1]), _bits(grads["eager"][1])) + assert grads["eager"][1][0].item() > 0, "the NaN lane's gradient was gated off" + assert torch.equal( + _bits(grads["triton"][0].float()), _bits(grads["eager"][0].float()) + ) + + +@pytest.mark.gpu +def test_gpu_fused_relu_nan_still_trips_the_trainers_non_finite_guard( + monkeypatch, tiny_trainer +): + """End to end: a NaN produced under the fused path reaches the abort. + + ScaFFold aborts a run whose reduced epoch losses are non-finite, precisely + so a diverged run stops instead of overwriting ``checkpoint_last.pth`` with + NaN weights. A fused activation that ate the NaN would hand that guard a + finite loss and let the run continue on a model whose *gradients* are still + NaN. The loss below is the real one: a real UNet, on the GPU, with the + Triton rung verified to have served every GroupNorm in it. + """ + from ScaFFold.utils import trainer as trainer_mod + + model = _small_unet("cuda", channels_last=True) + poisoned = torch.randn(1, 3, 16, 16, 16, device="cuda") + poisoned.view(-1)[0] = float("nan") + x = poisoned.contiguous(memory_format=torch.channels_last_3d).requires_grad_(True) + gn_mod.set_triton_enabled(None) + calls = _triton_spy(monkeypatch) + loss = model(x).float().mean() + if not calls: + pytest.skip( + "convolutions did not emit channels_last_3d; set " + "PYTORCH_MIOPEN_SUGGEST_NHWC=1 (the production setting)" + ) + assert not torch.isfinite(loss).item(), ( + "the fused activation swallowed the NaN: the forward is finite while " + "the backward is not, which is exactly what hides divergence" + ) + + trainer = tiny_trainer(config_overrides={"checkpoint_interval": 1, "epochs": 3}) + monkeypatch.setattr( + trainer, + "_run_training_batch", + lambda batch, **kw: (1, loss.detach().cpu(), torch.tensor(0.0)), + ) + # A *finite* validation loss, so the abort can only come from the model's. + monkeypatch.setattr( + trainer_mod, "evaluate", lambda *a, **kw: (7.4e-10, 0.5, 0.5, 2, 2) + ) + trainer.cleanup_or_resume() + with pytest.raises(ValueError, match="[Nn]on-finite"): + trainer.train() + assert not trainer.checkpoint_manager.last_ckpt_path.exists() + + +@pytest.mark.gpu +def test_gpu_ddp_wrapped_model_takes_the_triton_rung(monkeypatch): + """DDP's module-tree walk must not be confused by the nn.Identity slots.""" + import torch.distributed as dist + from torch.nn.parallel import DistributedDataParallel + + created = False + if not dist.is_initialized(): + import os + + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29623") + dist.init_process_group(backend="nccl", rank=0, world_size=1) + created = True + try: + model = _small_unet("cuda", channels_last=True) + wrapped = DistributedDataParallel(model, device_ids=[0]) + x = torch.randn(1, 3, 16, 16, 16, device="cuda").contiguous( + memory_format=torch.channels_last_3d + ) + gn_mod.set_triton_enabled(None) + calls = _triton_spy(monkeypatch) + wrapped(x).pow(2).sum().backward() + if not calls: + pytest.skip( + "convolutions did not emit channels_last_3d; set " + "PYTORCH_MIOPEN_SUGGEST_NHWC=1 (the production setting)" + ) + assert all(torch.isfinite(p.grad).all() for p in model.parameters()) + finally: + if created and dist.is_initialized(): + dist.destroy_process_group() diff --git a/tests/test_platform_guard.py b/tests/test_platform_guard.py new file mode 100644 index 0000000..5e2ffc9 --- /dev/null +++ b/tests/test_platform_guard.py @@ -0,0 +1,501 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""The hardware guard: ``_rungs._platform_declines`` and both ladders' wiring. + +**Every interesting branch here is one this node cannot take.** The guard's job +is to keep the Triton rungs off hardware they were not tuned on, and the machine +running these tests is the hardware they *were* tuned on -- so a suite that only +exercised the accept path would ship the entire decline path unexecuted, which +is the opposite of what the change is for. ``_rungs._device_fingerprint`` exists +as one small seam for that reason: substituting a tuple for it poses the MI300X, +the partitioned MI300A, the MI250X and the NVIDIA questions to the *real* +predicate, the real cache and the real message, rather than to a re-implementation +of them. + +Two properties are easy to get wrong and are pinned separately. The cache is +process-global, so :func:`_clean_platform_state` clears it around every test -- +without that the first question asked would fix the answer for the rest of the +session and half these tests would silently be re-asking it. And the guard must +be evaluated *once*, which is asserted as a call count on the seam rather than as +a duration, because a timing assertion on a driver query measures the driver. + +Verified by mutation: ``work/triton-conv/bin/mutate_platform_guard.py`` breaks +each property in turn and checks that the test named there notices. +""" + +from __future__ import annotations + +import logging + +import pytest +import torch +import torch.nn as nn + +from ScaFFold.unet import _rungs +from ScaFFold.unet import conv3d as conv_mod +from ScaFFold.unet import group_norm as gn_mod +from ScaFFold.unet.conv3d import FastConv3d, FastConvTranspose3d +from ScaFFold.unet.group_norm import FastGroupNorm + +_CHANNELS_LAST = torch.channels_last_3d + +#: What this node reports, and the only fingerprint the guard accepts. +_MI300A = ("gfx942", 228, "AMD Instinct MI300A") + +#: The parts an arch-only test would wrongly accept. MI300X and MI325X are +#: ``gfx942`` too -- discrete GPUs with 304 CUs rather than an APU with 228 -- +#: and a compute-partitioned MI300A reports the same arch with one XCD's worth +#: of CUs, which makes both the 228 and ``gather_gemm``'s ``GROUP_M = 6`` +#: (MI300A's XCD count) fiction while the arch string never moves. +_UNTUNED = { + "mi300x": ("gfx942", 304, "AMD Instinct MI300X"), + "mi325x": ("gfx942", 304, "AMD Instinct MI325X"), + "mi300a-cpx": ("gfx942", 38, "AMD Instinct MI300A"), + "mi250x": ("gfx90a", 104, "AMD Instinct MI250X"), + "next-gen": ("gfx950", 256, "AMD Instinct MI355X"), + # A CUDA build of torch has no ``gcnArchName`` at all, so the fingerprint + # reports an empty arch. Declining is right for a reason beyond tuning: + # the kernels' launch rules are MFMA rules and mean nothing without MFMA. + "nvidia": ("", 132, "NVIDIA H100"), +} + + +@pytest.fixture(autouse=True) +def _clean_platform_state(): + """Clear the verdict cache and both ladders' overrides around every test. + + The cache is a process-global memo, on purpose -- the whole point is that + the question is asked once. That makes it test state: a faked MI300X left + behind would turn every later GPU test in the session into a fallback test, + and a real verdict left behind would make a decline test pass by answering + the wrong question. Cleared on both sides for that reason. + """ + _rungs._reset_platform_cache() + saved = ( + conv_mod._triton_override, + conv_mod._triton_failed, + gn_mod._triton_override, + gn_mod._triton_failed, + ) + yield + ( + conv_mod._triton_override, + conv_mod._triton_failed, + gn_mod._triton_override, + gn_mod._triton_failed, + ) = saved + _rungs._reset_platform_cache() + + +def _fake_device(monkeypatch, fingerprint, *, count=None): + """Make every device look like ``fingerprint``; optionally count the asks. + + ``count`` is a list the seam appends each asked-about index to, which is how + the "evaluated once" tests assert on a call count rather than on a duration. + + The cache is dropped here as well, and that is not tidying: swapping the + hardware out from under a memo whose whole purpose is never to ask twice + would otherwise leave the *previous* answer in place -- which on this node + is "yes, MI300A", so every decline test would quietly become another accept + test. That is exactly the failure this file exists to avoid. + """ + + def fingerprint_of(index): + if count is not None: + count.append(index) + return fingerprint + + monkeypatch.setattr(_rungs, "_device_fingerprint", fingerprint_of) + _rungs._reset_platform_cache() + + +def _cuda(index=0): + return torch.device("cuda", index) + + +# --------------------------------------------------------------------------- +# the predicate +# --------------------------------------------------------------------------- + + +def test_the_tuned_fingerprint_is_the_only_one_accepted(monkeypatch): + """Arch *and* CU count, which is what makes this MI300A and not gfx942.""" + _fake_device(monkeypatch, _MI300A) + ok, described = _rungs._platform_verdict(_cuda()) + assert ok is True + assert "MI300A" in described + + +@pytest.mark.parametrize("name", sorted(_UNTUNED)) +def test_every_other_device_is_declined(monkeypatch, name): + """Including the three an arch-only predicate would have accepted.""" + _fake_device(monkeypatch, _UNTUNED[name]) + ok, described = _rungs._platform_verdict(_cuda()) + assert ok is False, f"{name} is not the device anything here was tuned on" + assert _UNTUNED[name][2] in described + assert _rungs._platform_declines(_cuda(), None) is True + + +def test_the_arch_feature_suffixes_are_not_part_of_the_comparison(monkeypatch): + """``gcnArchName`` carries build features; exact equality would be a trap. + + This node reports ``gfx942:sramecc+:xnack-``, and those suffixes describe + how the *build* was configured rather than which silicon is present -- so a + string comparison against ``"gfx942"`` would decline the very device + everything was tuned on, and one against the full string would decline the + same chip under a different HIP build. Driven through a stub property + object so the CPU suite exercises it too; the GPU test below pins that the + real device really does carry a suffix, i.e. that this is not hypothetical. + """ + + class _Props: + gcnArchName = "gfx942:sramecc+:xnack-" + multi_processor_count = 228 + name = "AMD Instinct MI300A" + + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda index: _Props()) + assert _rungs._device_fingerprint(0) == _MI300A + assert _rungs._platform_declines(_cuda(), None) is False + + +def test_a_device_that_cannot_be_answered_for_is_not_the_tuned_one(monkeypatch): + """"I could not find out" is not "yes"; the description says which.""" + + def explode(index): + raise RuntimeError("HIP error: no device") + + monkeypatch.setattr(_rungs, "_device_fingerprint", explode) + ok, described = _rungs._platform_verdict(_cuda()) + assert ok is False + assert "RuntimeError" in described and "no device" in described + + +def test_the_verdict_is_computed_once_per_device(monkeypatch): + """A call count, not a timing: the query is a driver call on a hot path. + + Also the per-device half of the decision: device 1 gets its own answer + rather than inheriting device 0's, so a node exposing two different parts + cannot have one of them silently answered for by the other. + """ + asked = [] + _fake_device(monkeypatch, _MI300A, count=asked) + for _ in range(20): + _rungs._platform_declines(_cuda(0), None) + assert asked == [0] + for _ in range(20): + _rungs._platform_declines(_cuda(1), None) + assert asked == [0, 1] + + +def test_a_mixed_node_answers_per_device(monkeypatch): + """One MI300A and one MI300X: each device gets the routing it deserves.""" + fingerprints = {0: _MI300A, 1: _UNTUNED["mi300x"]} + monkeypatch.setattr(_rungs, "_device_fingerprint", lambda i: fingerprints[i]) + assert _rungs._platform_declines(_cuda(0), None) is False + assert _rungs._platform_declines(_cuda(1), None) is True + + +# --------------------------------------------------------------------------- +# what it says +# --------------------------------------------------------------------------- + + +def test_declining_is_quiet_but_says_so_exactly_once(monkeypatch, caplog): + """One message per device: not per call, and not silence. + + Silence is the state this change ends -- a user on an MI300X should be able + to discover why the fast path is off. Per call would be a log line every + few milliseconds, since a scale-7 step routes some forty of these. + """ + _fake_device(monkeypatch, _UNTUNED["mi300x"]) + with caplog.at_level(logging.WARNING, logger=_rungs.__name__): + for _ in range(25): + assert _rungs._platform_declines(_cuda(), None) is True + records = [r for r in caplog.records if "tuned for" in r.message] + assert len(records) == 1, [r.message for r in records] + message = records[0].message + # It has to name what it found and what it wanted, or it is not actionable. + assert "MI300X" in message and "304 CUs" in message + assert "gfx942" in message and "228" in message + assert "SCAFFOLD_CONV_TRITON=1" in message + assert "SCAFFOLD_GROUPNORM_TRITON=1" in message + + +def test_the_tuned_platform_says_nothing_at_all(monkeypatch, caplog): + """No message on the machine everything was measured on.""" + _fake_device(monkeypatch, _MI300A) + with caplog.at_level(logging.WARNING, logger=_rungs.__name__): + for _ in range(5): + assert _rungs._platform_declines(_cuda(), None) is False + assert caplog.records == [] + + +def test_a_mixed_node_names_each_declining_device(monkeypatch, caplog): + """The bound is one message per device, which is what makes it discoverable.""" + fingerprints = {0: _UNTUNED["mi300x"], 1: _UNTUNED["mi250x"]} + monkeypatch.setattr(_rungs, "_device_fingerprint", lambda i: fingerprints[i]) + with caplog.at_level(logging.WARNING, logger=_rungs.__name__): + for index in (0, 1, 0, 1, 0): + _rungs._platform_declines(_cuda(index), None) + messages = [r.message for r in caplog.records] + assert len(messages) == 2 + assert any("cuda:0" in m and "MI300X" in m for m in messages) + assert any("cuda:1" in m and "MI250X" in m for m in messages) + + +# --------------------------------------------------------------------------- +# the override +# --------------------------------------------------------------------------- + + +def test_an_explicit_opt_in_takes_the_rung_anyway_and_is_loud(monkeypatch, caplog): + """The guard is a preference, so an explicit "yes" wins -- audibly. + + Loud because every figure either ladder is read against was measured on the + other machine, so a timing produced under this override is not comparable + with any of them and the log is the only place that survives the run. + """ + _fake_device(monkeypatch, _UNTUNED["mi300x"]) + with caplog.at_level(logging.WARNING, logger=_rungs.__name__): + for _ in range(25): + assert _rungs._platform_declines(_cuda(), True) is False + records = [r for r in caplog.records if "explicitly enabled" in r.message] + assert len(records) == 1, [r.message for r in records] + assert "MI300X" in records[0].message + + +def test_the_default_is_not_an_opt_in(monkeypatch): + """``None`` means "on wherever it is safe", and this device is not that. + + The tri-state is the whole override: an unset ``SCAFFOLD_CONV_TRITON`` and + an explicit ``1`` differ *here* and nowhere else. + """ + _fake_device(monkeypatch, _UNTUNED["mi300x"]) + assert _rungs._platform_declines(_cuda(), None) is True + assert _rungs._platform_declines(_cuda(), True) is False + + +@pytest.mark.parametrize( + "module, setter", + [ + (conv_mod, "set_conv_triton_enabled"), + (gn_mod, "set_triton_enabled"), + ], +) +def test_each_ladders_setter_is_the_override_for_that_ladder( + monkeypatch, module, setter +): + """``set_*_triton_enabled(True)`` is the in-process spelling of the opt-in. + + Per ladder, deliberately: the two kernels' tuning tables are separate bodies + of work, so a developer who has satisfied themselves about one has said + nothing about the other. + """ + _fake_device(monkeypatch, _UNTUNED["mi300x"]) + getattr(module, setter)(None) + assert _rungs._platform_declines(_cuda(), module._triton_override) is True + getattr(module, setter)(True) + assert _rungs._platform_declines(_cuda(), module._triton_override) is False + getattr(module, setter)(False) + # ``False`` never reaches the guard -- both callers decline on it first -- + # but it must not be mistaken for the opt-in if it ever does. + assert _rungs._platform_declines(_cuda(), module._triton_override) is True + + +@pytest.mark.parametrize("module", [conv_mod, gn_mod]) +def test_the_env_var_opt_in_reaches_the_guard(monkeypatch, module): + """``SCAFFOLD_*_TRITON=1`` and the setter are the same statement.""" + monkeypatch.setenv(module.TRITON_ENV_VAR, "1") + monkeypatch.setattr(module, "_triton_override", _rungs._env_override( + module.TRITON_ENV_VAR + )) + _fake_device(monkeypatch, _UNTUNED["mi300x"]) + assert _rungs._platform_declines(_cuda(), module._triton_override) is False + + +# --------------------------------------------------------------------------- +# the CPU path +# --------------------------------------------------------------------------- + + +def test_a_cpu_convolution_never_asks_the_hardware(monkeypatch): + """Ordering: ``is_cuda`` is tested before the device is ever queried. + + Not a nicety. ``torch.cuda.get_device_properties`` initializes torch's CUDA + state, so asking it on the routing path of a CPU tensor would drag a GPU + context into the whole CPU unit suite -- and into any CPU-only run of + ScaFFold, which must keep working. + + Asserted as "the seam was never reached" rather than as a raising stub: the + verdict is computed inside a broad ``except``, which would swallow the stub's + own exception and let the test pass while the query happened. + """ + asked = [] + _fake_device(monkeypatch, _MI300A, count=asked) + conv = FastConv3d(16, 16, kernel_size=3, padding=1, bias=False) + x = torch.randn(1, 16, 8, 8, 8) + assert conv_mod._use_triton(conv, x, None, None) is False + assert asked == [] + torch.testing.assert_close(conv(x), nn.Conv3d.forward(conv, x)) + + +def test_a_cpu_group_norm_never_asks_the_hardware(monkeypatch): + """The same ordering in the other ladder, asserted the same way.""" + asked = [] + _fake_device(monkeypatch, _MI300A, count=asked) + module = FastGroupNorm(8, 16) + x = torch.randn(1, 16, 8, 8, 8) + assert ( + gn_mod._use_triton(x, module.num_groups, module.weight, module.bias, None) + is False + ) + assert asked == [] + torch.testing.assert_close(module(x), nn.GroupNorm.forward(module, x)) + + +# --------------------------------------------------------------------------- +# the wiring: this node, told it is another one +# --------------------------------------------------------------------------- + + +def _gpu_conv(cin=16, cout=16, **kwargs): + kwargs.setdefault("kernel_size", 3) + kwargs.setdefault("padding", 1) + kwargs.setdefault("bias", False) + torch.manual_seed(11) + conv = FastConv3d(cin, cout, **kwargs) + return conv.cuda().to(memory_format=_CHANNELS_LAST).to(torch.bfloat16) + + +def _gpu_input(shape, dtype=torch.bfloat16): + generator = torch.Generator(device="cuda").manual_seed(5) + x = torch.randn(shape, device="cuda", dtype=torch.float32, generator=generator) + return x.to(dtype).contiguous(memory_format=_CHANNELS_LAST) + + +@pytest.mark.gpu +def test_this_node_is_the_tuned_platform(caplog): + """The accept path, against the real driver rather than a stub. + + Also the check that the constants have not drifted away from the machine + every measurement in ``work/triton-conv/`` was taken on -- and that the + suffix strip is load-bearing here rather than defensive, since the raw + string this device reports really does carry ``:sramecc+:xnack-``. + """ + props = torch.cuda.get_device_properties(0) + assert ":" in props.gcnArchName, props.gcnArchName + arch, cus, _name = _rungs._device_fingerprint(0) + assert (arch, cus) == (_rungs.TUNED_ARCH, _rungs.TUNED_CU_COUNT) + with caplog.at_level(logging.WARNING, logger=_rungs.__name__): + assert _rungs._platform_declines(_cuda(), None) is False + assert caplog.records == [] + + +@pytest.mark.gpu +def test_an_untuned_device_sends_the_convolution_to_miopen(monkeypatch): + """Declining is the ordinary fallback, not an error: same answer, MIOpen.""" + conv = _gpu_conv() + x = _gpu_input((1, 16, 16, 16, 16)) + assert conv_mod._use_triton(conv, x, None, None) is True + + _fake_device(monkeypatch, _UNTUNED["mi300x"]) + assert conv_mod._use_triton(conv, x, None, None) is False + torch.testing.assert_close(conv(x), nn.Conv3d.forward(conv, x)) + + conv_mod.set_conv_triton_enabled(True) + assert conv_mod._use_triton(conv, x, None, None) is True + + +@pytest.mark.gpu +def test_an_untuned_device_sends_the_upsampler_to_miopen(monkeypatch): + """The transposed ladder shares ``_routing_declines``, so it shares this.""" + torch.manual_seed(3) + module = FastConvTranspose3d(16, 8, kernel_size=2, stride=2) + module = module.cuda().to(memory_format=_CHANNELS_LAST).to(torch.bfloat16) + x = _gpu_input((1, 16, 8, 8, 8)) + assert conv_mod._use_triton_transposed(module, x, None, None) is True + + _fake_device(monkeypatch, _UNTUNED["mi300x"]) + assert conv_mod._use_triton_transposed(module, x, None, None) is False + torch.testing.assert_close(module(x), nn.ConvTranspose3d.forward(module, x)) + + conv_mod.set_conv_triton_enabled(True) + assert conv_mod._use_triton_transposed(module, x, None, None) is True + + +@pytest.mark.gpu +def test_an_untuned_device_sends_group_norm_to_the_stock_kernel(monkeypatch): + """The second ladder, guarded by the same one predicate.""" + module = FastGroupNorm(8, 16).cuda() + x = _gpu_input((1, 16, 8, 8, 8), dtype=torch.float32) + args = (x, module.num_groups, module.weight, module.bias, None) + assert gn_mod._use_triton(*args) is True + + _fake_device(monkeypatch, _UNTUNED["mi300x"]) + assert gn_mod._use_triton(*args) is False + torch.testing.assert_close( + module(x), nn.GroupNorm.forward(module, x), rtol=1e-5, atol=1e-5 + ) + + gn_mod.set_triton_enabled(True) + assert gn_mod._use_triton(*args) is True + + +@pytest.mark.gpu +def test_both_ladders_share_one_verdict_and_ask_for_it_once(monkeypatch): + """One source of truth, one query, whichever ladder gets there first. + + Two copies of this decision would drift -- the tables they protect were + tuned in two separate sessions on the same device -- so the cheapest + available proof that there is only one is that the second ladder's routing + call does not produce a second driver query. + """ + asked = [] + _fake_device(monkeypatch, _MI300A, count=asked) + conv = _gpu_conv() + x = _gpu_input((1, 16, 16, 16, 16)) + gn = FastGroupNorm(8, 16).cuda() + for _ in range(3): + conv_mod._use_triton(conv, x, None, None) + gn_mod._use_triton( + x.float().contiguous(memory_format=_CHANNELS_LAST), + gn.num_groups, + gn.weight, + gn.bias, + None, + ) + assert asked == [0] + + +@pytest.mark.gpu +def test_a_whole_unet_step_asks_the_hardware_once(monkeypatch): + """The property that matters in production: once, not once per operation. + + A scale-7 step routes 19 convolutions, 4 upsamplers and 18 GroupNorms + through these predicates; the guard has to be a dictionary lookup after the + first of them. + """ + from ScaFFold.unet.unet_model import UNet + + asked = [] + _fake_device(monkeypatch, _MI300A, count=asked) + torch.manual_seed(0) + model = UNet(n_channels=3, n_classes=2, trilinear=False, layers=2, + group_norm_groups=8) + model = model.cuda().to(memory_format=_CHANNELS_LAST) + x = _gpu_input((1, 3, 16, 16, 16), dtype=torch.float32) + with torch.no_grad(): + model(x) + assert asked == [0] From 0f9a38c0a0f1a94cec3c1191d550dd43fcc5e4c6 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Wed, 5 Aug 2026 06:41:55 -0700 Subject: [PATCH 3/8] Add a Triton 3-D convolution package A self-contained NDHWC implicit-GEMM convolution for MI300A / gfx942, built to replace MIOpen in this benchmark and to be upstreamable to DistConv on its own terms. It imports nothing from ScaFFold. Four kernels serve seven operator-directions. Backward-data is not a kernel at all -- it is the forward contraction on a permuted weight -- and the transposed operator's backward directions are its own forward kernel and backward-weight with the operands swapped. That reuse is the main structural result. The corpus is recorded from real ScaFFold calls, and it records the *shape form*: the logical convolution, the halo'd unpadded one DistConv hands a backend, and the padded one this package's adapter actually issues. Those are different problems, and conflating them has been the most expensive mistake in this work, so ConvProblem names which is which and the benchmark driver takes --form. Backward-weight is deterministic by default: split-K with a reduction tree whose split count, tile and order are pure functions of the shape. Correctness is checked against fp64 references under a three-tier tolerance policy, with a bitwise corpus for the exactly-representable cases. Also here: the benchmark harness, which times kernels through CUDA-graph replay with 95% intervals and an online iteration count, and which can run without the MIOpen control -- 98% of a two-arm capture's wall clock was MIOpen's find, not measurement. --- triton_conv3d/__init__.py | 107 + triton_conv3d/bench/__init__.py | 14 + triton_conv3d/bench/baseline.py | 469 +++ triton_conv3d/bench/conv_bench.py | 1412 +++++++ triton_conv3d/bench/gemm_probe.py | 1009 +++++ triton_conv3d/bench/harness.py | 1068 +++++ triton_conv3d/bwd_data.py | 360 ++ triton_conv3d/gather_gemm.py | 1396 +++++++ triton_conv3d/reduce_gemm.py | 1507 +++++++ triton_conv3d/reference.py | 403 ++ triton_conv3d/scaffold_census.json | 5065 +++++++++++++++++++++++ triton_conv3d/scaffold_corpus.json | 4871 ++++++++++++++++++++++ triton_conv3d/shapes.py | 735 ++++ triton_conv3d/tests/test_bwd_data.py | 680 +++ triton_conv3d/tests/test_bwd_weight.py | 1370 ++++++ triton_conv3d/tests/test_gather_gemm.py | 1101 +++++ triton_conv3d/tests/test_infra.py | 1607 +++++++ triton_conv3d/tests/test_transposed.py | 947 +++++ triton_conv3d/transposed.py | 1149 +++++ 19 files changed, 25270 insertions(+) create mode 100644 triton_conv3d/__init__.py create mode 100644 triton_conv3d/bench/__init__.py create mode 100644 triton_conv3d/bench/baseline.py create mode 100644 triton_conv3d/bench/conv_bench.py create mode 100644 triton_conv3d/bench/gemm_probe.py create mode 100644 triton_conv3d/bench/harness.py create mode 100644 triton_conv3d/bwd_data.py create mode 100644 triton_conv3d/gather_gemm.py create mode 100644 triton_conv3d/reduce_gemm.py create mode 100644 triton_conv3d/reference.py create mode 100644 triton_conv3d/scaffold_census.json create mode 100644 triton_conv3d/scaffold_corpus.json create mode 100644 triton_conv3d/shapes.py create mode 100644 triton_conv3d/tests/test_bwd_data.py create mode 100644 triton_conv3d/tests/test_bwd_weight.py create mode 100644 triton_conv3d/tests/test_gather_gemm.py create mode 100644 triton_conv3d/tests/test_infra.py create mode 100644 triton_conv3d/tests/test_transposed.py create mode 100644 triton_conv3d/transposed.py diff --git a/triton_conv3d/__init__.py b/triton_conv3d/__init__.py new file mode 100644 index 0000000..db3807c --- /dev/null +++ b/triton_conv3d/__init__.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""Triton 3-D convolution kernels for NDHWC (``channels_last_3d``) tensors. + +This package is deliberately **self-contained**: it imports nothing from +ScaFFold or DistConv, so it can be vendored into either (or released on its own) +unchanged. ScaFFold plugs in through a thin adapter that lives on the ScaFFold +side. + +What exists today: + +- :mod:`triton_conv3d.gather_gemm` -- the forward implicit-GEMM convolution, + ``k>=1`` with ``stride=1``, bf16 / fp16 / fp32. +- :mod:`triton_conv3d.bwd_data` -- the gradient with respect to the input. It + contains no kernel of its own: at ``stride=1`` backward-data *is* the forward + contraction on a flipped, channel-transposed weight. +- :mod:`triton_conv3d.reduce_gemm` -- the gradient with respect to the weight, + which is the one direction that needs a kernel of its own: a tiny output + reduced over the whole volume, so split-K is mandatory rather than optional. + It is also where reproducibility is decided, and its deterministic path is the + default. +- :mod:`triton_conv3d.transposed` -- ``ConvTranspose3d`` at ``kernel == stride`` + and no padding, all three directions. Only its *forward* is a kernel: with + the windows tiling rather than overlapping, both backward directions are the + ordinary strided convolution seen from the other side, and the two modules + above serve them unchanged. +- :mod:`triton_conv3d.shapes` -- the convolution problems that actually occur, + extracted from real ScaFFold runs, plus synthetic edge cases. +- :mod:`triton_conv3d.reference` -- reference implementations and the tolerance + policy used to decide whether a kernel is correct. +- :mod:`triton_conv3d.bench` -- interleaved A/B timing, MIOpen baseline capture, + the ``tl.dot`` ceiling probe, and the forward benchmark. + +The public entry point takes and returns ``channels_last_3d`` tensors; the +autograd registration and the ScaFFold adapter live in a later milestone, so a +caller today drives :func:`conv3d_forward` directly and checks a gate first. + +**The gates say nothing about the GPU, deliberately.** They are *capability* +predicates -- "will this call be computed correctly here" -- and the answer to +that does not depend on which AMD part is running: the kernels compute the right +convolution wherever Triton lowers them. What *is* device-specific is every +number that decides how they launch (the tile tables, ``matrix_instr_nonkdim``, +and the ``GROUP_M`` default of 6, which is MI300A's XCD count), all of which was +raced on one MI300A -- and a launch configuration that is merely wrong for the +hardware raises nothing: see :mod:`triton_conv3d.gather_gemm`'s "Configuration +constraints are hard". Deciding whether *this* machine is one whose numbers are +trustworthy is therefore the embedder's routing question, not this package's +capability question, and putting a device allowlist inside the gates would lock +out a consumer who has retuned for their own part. ScaFFold makes that decision +in ``ScaFFold/unet/_rungs.py`` (``_platform_declines``); a consumer that has not +retuned should do the same thing there. + +**Which gate depends on what the caller will do with the answer.** The three +directions do not accept the same problems -- ``stride > 1`` is served by the +forward and by backward-weight and refused by backward-data -- so a caller that +will differentiate the result must ask :func:`is_supported_all`, which is all +three at once. :func:`is_supported` alone gates the forward alone, which is +what an inference caller wants and a training caller must not settle for: a +forward this package serves and a backward it cannot is discovered inside +``backward()``, where the caller's fallback kernel is no longer reachable. +""" + +from __future__ import annotations + +from typing import Any + +__version__ = "0.0.0.dev0" + +__all__ = ["ConvConfig", "conv3d_forward", "conv3d_backward_data", + "conv3d_backward_weight", "is_supported", "is_supported_all", + "is_supported_bwd_data", "is_supported_bwd_weight", + "conv_transpose3d_forward", "conv_transpose3d_backward_data", + "conv_transpose3d_backward_weight", "is_supported_transposed", + "is_supported_transposed_all", "is_supported_transposed_bwd_data", + "is_supported_transposed_bwd_weight", "__version__"] + +#: The public names live in :mod:`triton_conv3d.gather_gemm`, which imports +#: torch and triton. They are re-exported lazily so that +#: ``import triton_conv3d.shapes`` stays free of both: the shape and cost model +#: is pure Python by design, and it drives test parametrization at collection +#: time on machines that have no GPU and may have no triton. +_LAZY = {"ConvConfig": "gather_gemm", "conv3d_forward": "gather_gemm", + "is_supported": "gather_gemm", "is_supported_all": "gather_gemm", + "conv3d_backward_data": "bwd_data", + "is_supported_bwd_data": "bwd_data", + "conv3d_backward_weight": "reduce_gemm", + "is_supported_bwd_weight": "reduce_gemm", + "conv_transpose3d_forward": "transposed", + "conv_transpose3d_backward_data": "transposed", + "conv_transpose3d_backward_weight": "transposed", + "is_supported_transposed": "transposed", + "is_supported_transposed_all": "transposed", + "is_supported_transposed_bwd_data": "transposed", + "is_supported_transposed_bwd_weight": "transposed"} + + +def __getattr__(name: str) -> Any: + if name in _LAZY: + import importlib + + return getattr( + importlib.import_module(f".{_LAZY[name]}", __name__), name + ) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__() -> list[str]: + return sorted(__all__) diff --git a/triton_conv3d/bench/__init__.py b/triton_conv3d/bench/__init__.py new file mode 100644 index 0000000..e230d8e --- /dev/null +++ b/triton_conv3d/bench/__init__.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""Measurement infrastructure: timing, baselines and ceiling probes.""" + +from .harness import ( + Measurement, + Ratio, + flush_caches, + interleaved, + ratio, + time_callable, +) + +__all__ = ["Measurement", "Ratio", "flush_caches", "interleaved", "ratio", + "time_callable"] diff --git a/triton_conv3d/bench/baseline.py b/triton_conv3d/bench/baseline.py new file mode 100644 index 0000000..0fe1b67 --- /dev/null +++ b/triton_conv3d/bench/baseline.py @@ -0,0 +1,469 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""Capture MIOpen's time for every problem in the corpus, once. + +Every performance claim we make later is a comparison, and a comparison needs a +control taken on the same machine in the same session. Re-measuring MIOpen +alongside each experiment is both slow and fragile -- MIOpen's tuning database +warms up over a run, so the "same" baseline drifts depending on what ran before +it. Capturing it once, deliberately warmed, and storing it makes every later +number reproducible and makes database drift visible instead of silent. + +Reproducing what ScaFFold actually runs +======================================= + +Getting either of the following wrong makes MIOpen solve a *different problem*, +or solve this one *differently*, and then quietly reports a number that is not +the control. The first version of this file got both wrong and overstated +MIOpen's cost by up to 12x, which would have turned into a fabricated speedup +for every kernel measured against it. So they are enforced here, recorded in +the output header, and regression-tested in ``tests/test_infra.py``. + +1. ``torch.backends.cudnn.benchmark = True``. ScaFFold sets this at startup + unless ``more_determinism`` is on (``ScaFFold/worker.py:171``), and so does + the profiling harness the reference numbers come from + (``work/report/prof_bench.py:125``). On ROCm the flag decides whether + PyTorch asks MIOpen for an exhaustive *find* or lets it answer from its AI + heuristic. The two answers are not close. For ``conv 64->64 k3 @ 128^3``, + forward, measured with ``MIOPEN_ENABLE_LOGGING=1``:: + + benchmark=False findMode DYNAMIC_HYBRID(5), no search + DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<256,64,64,32, + Default,16,16,2,2,2,1,2,1,1,1> 12.235 ms + benchmark=True findMode NORMAL(1), GenericSearch over 23 configs + DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<256,128,64,32, + Default,32,32,2,1,8,8,8,1,1,1> 1.746 ms + + Same solver (``ConvHipImplicitGemm3DGroupFwdXdlops``), same device op -- + only the tuning config differs. The heuristic picks 16x16 MFMA tiles with + 2-element global loads; the search picks 32x32 tiles with 8-element loads. + That is the whole 7x. It is not a naive fallback, which is why it does not + look like one in a profile. + +2. The shape -- and there are **three** of it. Upstream DistConv concatenates + a ``k // 2`` halo onto every axis it manages and then sets that axis's + padding to zero, so under the MIOpen rung ScaFFold's + ``conv 64->64 k3 @ 128^3`` reaches MIOpen as a *130^3 unpadded* problem. + MIOpen's find-db key includes the padding, so these are separate problems + with separate tuning:: + + padded 128^3 pad 1, benchmark=True bwd-data 3.426 ms + halo'd 130^3 pad 0, benchmark=True bwd-data 3.324 ms <- DistConv + profiled ScaFFold (config A) bwd-data 3.483 ms + + ``--shape halo`` (the default) measures the form **DistConv** issues, which + is the form the profiled numbers this module cross-checks against were + measured in -- so it stays the default and the join stays valid. It is + *not* the shape ScaFFold's own Triton rung runs: that adapter exchanges a + halo only on genuinely split axes, so the convolution it issues is padded on + H and W at every configuration and on all three axes at one GPU. + ``--shape production`` measures **that** form, which is the one to baseline + MIOpen in if the comparison is against a Triton kernel; ``--shape unhaloed`` + measures the logical statement; ``--shape all`` measures each distinct one. + Every record says which it was, so a baseline cell can never be silently + compared against the wrong profile cell. + +Two more must be in the *environment* before the process starts, because MIOpen +reads them when it builds its handle and this module cannot set them for you: + +* ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` -- ScaFFold's production setting. Without + it ``channels_last_3d`` is inert on ROCm and MIOpen is handed NCDHW, which is + a different problem again. Refused rather than warned about, below. +* ``MIOPEN_USER_DB_PATH`` / ``MIOPEN_CUSTOM_CACHE_DIR`` pointing somewhere + persistent, so a find survives the process and ``--resume`` does not search + from scratch. Warned about; both are recorded in the output header so a + baseline taken against a cold database is identifiable after the fact. + +Cross-check +=========== + +``--cross-check`` joins each measured cell against the ``measured`` entries the +corpus carries from the profiled runs and reports the ratio, so a harness that +has drifted out of agreement says so rather than being believed. + +Results stream to the output file as they are produced, and ``--resume`` skips +what is already there. That matters because two things in this corpus do not +merely run slowly: the scale-8 backward-weight at ``128->64 @ 128x256x256`` +takes 45 s per call, and unsharded scale 8 trips an assertion inside MIOpen that +can take the process down with it. Losing 40 minutes of measurements to the +last problem in the list is avoidable, so it is avoided. + +Usage:: + + python -m triton_conv3d.bench.baseline --out baseline.json + python -m triton_conv3d.bench.baseline --out baseline.json --resume \ + --include-edge --shape both --cross-check +""" + +from __future__ import annotations + +import argparse +import json +import os +import pathlib +import sys +import time +import traceback + +import torch + +from ..shapes import DIRECTIONS, ConvProblem, Direction, edge_cases, scaffold_corpus +from .harness import format_table, interleaved + +_MEMORY_FORMAT = torch.channels_last_3d + +#: ScaFFold's default (``worker.py:171``), and the profiling harness's +#: (``prof_bench.py:125``). See the module docstring: with this off MIOpen +#: answers from its heuristic instead of searching, and the baseline is wrong +#: by up to 12x. Module-level so that importing this module is enough to put a +#: process in the configuration the recorded numbers were taken in. +REQUIRE_CUDNN_BENCHMARK = True +torch.backends.cudnn.benchmark = REQUIRE_CUDNN_BENCHMARK + + +def _key(problem: ConvProblem, direction: Direction) -> str: + return f"{problem.label}|{problem.dtype}|{direction}" + + +def _build(problem: ConvProblem, device: str, dtype: torch.dtype): + x = torch.randn(problem.input_shape, device=device, dtype=torch.float32).to(dtype) + w = torch.randn(problem.weight_shape, device=device, dtype=torch.float32).to(dtype) + x = x.contiguous(memory_format=_MEMORY_FORMAT).requires_grad_(True) + w = w.contiguous(memory_format=_MEMORY_FORMAT).requires_grad_(True) + b = ( + torch.randn(problem.cout, device=device, dtype=torch.float32).to(dtype) + if problem.bias + else None + ) + return x, w, b + + +def _callable(problem: ConvProblem, direction: Direction, device: str, dtype): + """A zero-argument closure that runs exactly the one direction, plus its shapes. + + The backward directions are isolated with ``torch.autograd.grad`` on a + pre-computed forward output so that the forward is not folded into the + measurement, and ``retain_graph`` keeps the same graph across iterations. + """ + import torch.nn.functional as F + + x, w, b = _build(problem, device, dtype) + op = F.conv_transpose3d if problem.transposed else F.conv3d + fwd = lambda: op(x, w, b, stride=problem.stride, padding=problem.padding) + + if direction == "fwd": + with torch.no_grad(): + return fwd, (x, w) + + y = fwd() + gy = torch.randn_like(y) + inputs = (x,) if direction == "bwd-data" else (w,) + fn = lambda: torch.autograd.grad(y, inputs, gy, retain_graph=True) + return fn, (x, w, y, gy) + + +def measure_one( + problem: ConvProblem, + direction: Direction, + *, + device: str = "cuda", + budget_s: float = 10.0, + target_rel: float = 0.02, + max_call_ms: float = 60_000.0, + shape_mode: str = "halo", +) -> dict: + """One (problem, direction) cell. Never raises; failures are data too.""" + if not torch.backends.cudnn.benchmark: + raise RuntimeError( + "cudnn.benchmark is off; MIOpen will answer from its heuristic " + "instead of searching and the result is not a baseline" + ) + dtype = {"bf16": torch.bfloat16, "fp32": torch.float32, + "fp16": torch.float16}[problem.dtype] + record: dict = { + "problem": problem.label, + "direction": direction, + "dtype": problem.dtype, + # Which of the three forms of this convolution was measured, and enough + # of the descriptor to tell them apart without consulting the corpus. + "shape_mode": shape_mode, + "qualified_problem": problem.qualified_label, + "padding": list(problem.padding), + "input_shape": list(problem.input_shape), + "weight_shape": list(problem.weight_shape), + "output_shape": list(problem.output_shape), + "flops": problem.flops(direction), + "bytes": problem.bytes(direction), + "arithmetic_intensity": problem.arithmetic_intensity(direction), + "roofline_tflops": problem.roofline_flops(direction) / 1e12, + "needs_int64": problem.needs_int64, + } + tensors = None + try: + torch.cuda.empty_cache() + fn, tensors = _callable(problem, direction, device, dtype) + + # One untimed call decides whether this cell is measurable at all: with + # cudnn.benchmark on, MIOpen's *find* runs on the first invocation -- + # it launches every candidate config -- and would otherwise be the whole + # measurement. Everything after that is sized by the harness, which + # re-derives the same per-call time and additionally picks the round + # count from the precision it has reached. + fn() + torch.cuda.synchronize() + t0 = time.perf_counter() + fn() + torch.cuda.synchronize() + rough_ms = (time.perf_counter() - t0) * 1e3 + + if rough_ms > max_call_ms: + record.update(ms=rough_ms, best_ms=rough_ms, iters=1, rounds=1, + spread=0.0, rel_ci=float("inf"), stop="single-call", + note="single call; exceeds max_call_ms") + else: + meas = interleaved({"miopen": fn}, warmup=None, iters=None, + rounds=None, budget_s=budget_s, + target_rel=target_rel)["miopen"] + # Both statistics, because they answer different questions. The + # median is the control -- it is what a step actually costs on a + # shared node. The best round is the diagnostic: this node has + # other tenants, and a neighbour can inflate every round at once, + # so "did MIOpen find a good kernel" has to be asked of the best + # round or it gets a flaky answer. + # + # ``spread`` is kept because every stored baseline has it, but read + # ``rel_ci`` instead: ``spread`` is a *range* and its expectation + # grows with ``rounds``, which is now chosen per cell, so two cells' + # spreads are no longer comparable to each other at all. + record.update(ms=meas.median, best_ms=meas.best, iters=meas.iters, + rounds=len(meas.rounds), spread=meas.spread, + rel_ci=meas.rel_half_width, stop=meas.stop, + group=meas.group, tax_frac=meas.tax_frac, + measure_seconds=meas.seconds) + record["tflops"] = record["flops"] / (record["ms"] * 1e-3) / 1e12 + record["pct_roofline"] = 100 * record["tflops"] / record["roofline_tflops"] + except Exception as exc: + record["error"] = f"{type(exc).__name__}: {exc}" + record["traceback"] = traceback.format_exc()[-1500:] + finally: + del tensors + torch.cuda.empty_cache() + return record + + +# --------------------------------------------------------------------------- +# Cross-check against the profiled runs +# --------------------------------------------------------------------------- + + +def cross_check(records: list[dict], problems: list[ConvProblem]) -> list[dict]: + """Join measured cells onto the profiled ScaFFold numbers they control for. + + Only the halo'd cells are joined, and the reason is narrower than it used to + be stated: the profile these numbers control for was taken with DistConv on + the path, so the calls it timed *were* the halo'd form. The other two forms + have no profiled counterpart to be compared against -- not because ScaFFold + never runs them (it runs the production form at every site, every step) but + because nobody has profiled a step in them. Joining a production-form cell + onto a DistConv-form profile figure is the bug this whole module is a + response to. The profiled figure used is the *cheapest* of the per-config + measurements, because a profiled call can be slowed by contention with the + rest of the step but cannot be sped up by it. + """ + by_key = {} + for p in problems: + halo = p.halo_variant + for d in DIRECTIONS: + hits = p.measured_for(d) + if hits: + by_key[(halo.label, d)] = hits[-1] + rows = [] + for r in records: + if r.get("shape_mode") != "halo" or "ms" not in r: + continue + hit = by_key.get((r["problem"], r["direction"])) + if hit is None: + continue + rows.append({ + "problem": r["problem"], + "direction": r["direction"], + "isolated_ms": r["ms"], + "profiled_ms": hit["ms_per_call"], + "ratio": r["ms"] / hit["ms_per_call"], + "config": hit["config"], + "profiled_solvers": hit.get("solvers", []), + }) + return rows + + +def format_cross_check(rows: list[dict], tol: float = 0.25) -> str: + table = format_table( + [ + [ + r["problem"], r["direction"], f"{r['isolated_ms']:.4f}", + f"{r['profiled_ms']:.4f}", f"{r['ratio']:.2f}x", + "ok" if abs(r["ratio"] - 1) <= tol else "MISMATCH", + ] + for r in rows + ], + ["problem", "direction", "isolated ms", "profiled ms", "ratio", ""], + aligns="llrrrl", + ) + bad = [r for r in rows if abs(r["ratio"] - 1) > tol] + return table + f"\n\n{len(rows) - len(bad)}/{len(rows)} cells within {tol:.0%}" + + +# --------------------------------------------------------------------------- + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--out", default="baseline.json") + ap.add_argument("--resume", action="store_true", + help="skip cells already present in --out") + ap.add_argument("--top", type=int, default=0, + help="only the N hottest corpus problems (0 = all)") + ap.add_argument("--skip", type=int, default=0, + help="skip the first N corpus problems (they are ordered by " + "cost, and the most expensive one takes 45 s per call)") + ap.add_argument("--include-edge", action="store_true", + help="also measure the synthetic edge cases") + ap.add_argument("--shape", + choices=("halo", "production", "unhaloed", "both", "all"), + default="halo", + help="halo: the form upstream DistConv issues (default, and " + "the form the profiled numbers measured, so the only " + "one --cross-check can join); production: the form " + "ScaFFold's own Triton adapter issues, padded on every " + "unsplit axis -- the right MIOpen baseline for a " + "Triton comparison; unhaloed: the logical statement; " + "both: halo+unhaloed, as before; all: every distinct " + "form") + ap.add_argument("--max-call-ms", type=float, default=60_000.0, + help="above this, record a single call rather than a sweep") + ap.add_argument("--budget", type=float, default=10.0, + help="wall-clock seconds per cell (default 10)") + ap.add_argument("--precision", type=float, default=0.02, + help="target relative 95%% half-width per cell") + ap.add_argument("--skip-slow", action="store_true", + help="skip cells a previous run recorded as slower than " + "--max-call-ms; useful for a quick re-capture") + ap.add_argument("--cross-check", action="store_true", + help="join the halo'd cells onto the profiled numbers") + ap.add_argument("--tolerance", type=float, default=0.25, + help="cross-check band, as a fraction of the profiled time") + args = ap.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("no GPU") + if os.environ.get("PYTORCH_MIOPEN_SUGGEST_NHWC") != "1": + raise SystemExit( + "PYTORCH_MIOPEN_SUGGEST_NHWC=1 is not set: channels_last_3d is inert " + "on ROCm without it and MIOpen would see NCDHW, which is not what " + "ScaFFold runs" + ) + if not os.environ.get("MIOPEN_USER_DB_PATH"): + print("WARNING: MIOPEN_USER_DB_PATH is unset -- every cell re-runs the " + "find from scratch and nothing is reusable afterwards", + file=sys.stderr) + + problems = list(scaffold_corpus())[args.skip:] + if args.top: + problems = problems[: args.top] + corpus_problems = list(problems) + if args.include_edge: + problems += list(edge_cases()) + + modes = {"both": ("halo", "unhaloed"), + "all": ("halo", "production", "unhaloed")}.get( + args.shape, (args.shape,)) + _forms = {"halo": lambda p: p.halo_variant, + "production": lambda p: p.production_variant, + "unhaloed": lambda p: p} + # A problem with nothing to halo is its own variant in all three forms, so + # the multi-mode runs would otherwise measure the synthetic edge cases and + # the k=1/transposed convs two or three times for nothing. Deduplicated on + # the *qualified* label, which carries the padding: two forms can share a + # ``label`` and be different problems. + variants: list[tuple[ConvProblem, str]] = [] + for p in problems: + seen: dict[str, str] = {} + for mode in modes: + q = _forms[mode](p) + if q.qualified_label in seen: + continue + seen[q.qualified_label] = mode + variants.append((q, mode)) + + out_path = pathlib.Path(args.out) + done: dict[str, dict] = {} + records: list[dict] = [] + if args.resume and out_path.exists(): + prior = json.loads(out_path.read_text()) + records = list(prior.get("records", [])) + done = {r["problem"] + "|" + r["direction"]: r for r in records} + print(f"resuming: {len(done)} cells already measured") + + props = torch.cuda.get_device_properties(0) + header = { + "device": props.name, + "torch": torch.__version__, + "miopen_suggest_nhwc": os.environ.get("PYTORCH_MIOPEN_SUGGEST_NHWC"), + "miopen_user_db_path": os.environ.get("MIOPEN_USER_DB_PATH"), + "miopen_custom_cache_dir": os.environ.get("MIOPEN_CUSTOM_CACHE_DIR"), + # The two settings that decide whether this file is a control at all. + "cudnn_benchmark": torch.backends.cudnn.benchmark, + "shape_mode": args.shape, + "memory_format": "channels_last_3d", + } + + def flush() -> None: + payload = {**header, "records": records} + if args.cross_check: + payload["cross_check"] = cross_check(records, corpus_problems) + out_path.write_text(json.dumps(payload, indent=1) + "\n") + + for problem, mode in variants: + for direction in DIRECTIONS: + key = problem.label + "|" + direction + if key in done: + continue + if args.skip_slow and done.get(key, {}).get("note"): + continue + rec = measure_one(problem, direction, max_call_ms=args.max_call_ms, + shape_mode=mode, budget_s=args.budget, + target_rel=args.precision) + records.append(rec) + flush() + if "error" in rec: + print(f" {problem.label:36s} {mode:8s} {direction:11s} " + f"ERROR {rec['error'][:80]}") + else: + print( + f" {problem.label:36s} {mode:8s} {direction:11s} " + f"{rec['ms']:10.4f} ms +-{rec.get('rel_ci', 0):5.1%} " + f"({rec.get('rounds', 0)}r/{rec.get('stop', '?')}) " + f"{rec['tflops']:7.1f} TF/s " + f"{rec['pct_roofline']:6.1f}% roofline" + + (f" [{rec['note']}]" if rec.get("note") else "") + ) + sys.stdout.flush() + + flush() + ok = [r for r in records if "error" not in r] + print(f"\n{len(ok)}/{len(records)} cells measured -> {out_path}") + if len(ok) < len(records): + print("failures:") + print(format_table( + [[r["problem"], r["direction"], r["error"][:70]] + for r in records if "error" in r], + ["problem", "direction", "error"], + )) + if args.cross_check: + rows = cross_check(records, corpus_problems) + print("\ncross-check against the profiled ScaFFold runs:") + print(format_cross_check(rows, tol=args.tolerance)) + + +if __name__ == "__main__": + main() diff --git a/triton_conv3d/bench/conv_bench.py b/triton_conv3d/bench/conv_bench.py new file mode 100644 index 0000000..4813147b --- /dev/null +++ b/triton_conv3d/bench/conv_bench.py @@ -0,0 +1,1412 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""Triton against MIOpen, per operator, per direction, on the shapes ScaFFold runs. + +**Two operators, three directions, one driver.** ``--operator`` selects the +convolution or the ``k == s`` transposed convolution; ``--direction`` selects +which of its gradients (or none). ``--operator all --direction all`` measures +every cell in the project under one methodology, in one process, from one +command. + +The transposed benchmarks used to live in a separate driver +(``work/triton-conv/bin/m5_convT_bench.py``, deleted at this commit), on the +argument that ``--direction`` selects between three *directions of one operator* +while the transposed convolution is a *different* operator, and that folding it +in would make ``_build`` branch on ``problem.transposed`` in every arm to run the +same code. The argument was right about the failure and wrong about the +factoring. Operator and direction are **two axes of one table**, not two values +on one switch: :data:`_OPERATORS` maps ``(operator, direction)`` to a builder, +and each of the six builders is a separate function that names its own operands, +its own control, its own candidate configs and its own reference. ``_build`` is +a lookup, so nothing branches on ``problem.transposed`` anywhere, and the four +things that *are* per-operator rather than per-direction -- the shape form, the +problem ordering, the config type, whether a direction is sweepable -- sit on +:class:`_Op` where a reader can see all four at once. + +What that bought, immediately: ``m5_convT_bench``'s transposed backward +directions dropped the ``config=`` argument on the floor, because the only config +object it had was the transposed *forward*'s, and passing that to a direction +served by ``conv3d_forward`` silently benchmarks a tile nothing would select. +With a builder per cell each one names the config type its own entry point +resolves, and both transposed backward directions became sweepable for free. + +Five things this driver is careful about, each because getting it wrong has +already produced a wrong answer once in this project: + +**The shape.** One ScaFFold convolution reaches a kernel in three different +shapes, and they are three different tuning problems -- MIOpen keys its find +database on the padding, ``bwd_data_config`` derives ``M`` from it, and the +kernel compiles a different ``PADDED`` body either way. (``bwd_weight_config`` +also used to change its answer on it; that clause went on 2026-08-05, and the +forms are still three different measurements without it.) ``--form`` chooses +which one a run measures, and every row records it: + +* ``distconv`` (the default, and every capture on disk): the halo'd, unpadded + form upstream DistConv hands the backend, ``130^3`` at ``padding = 0``. It is + what the profiled MIOpen baseline in ``ConvProblem.measured`` is a timing + *of*, so it is the right form for a like-for-like MIOpen comparison. +* ``adapter``: what ``ScaFFold/unet/conv3d.py`` hands the Triton kernels, which + is **the form production actually runs** -- a halo on the genuinely split axis + only, so ``128^3`` at ``padding = (1,1,1)`` unsharded and + ``130x256x256`` at ``(0,1,1)`` at two shards. Padded at every configuration. +* ``logical``: the module's own statement, unhalo'd and padded. Identical to + ``adapter`` wherever nothing is split. + +Defaulting to ``distconv`` keeps every stored capture comparable; it does *not* +mean it is the form to quote a Triton speedup in. A ``conv`` cell applies the +chosen form. **No ``convT`` cell ever differs**, and that is not an oversight: +DistConv's halo is ``k // 2`` only for an odd kernel and 0 at ``k = 2``, and the +adapter exchanges nothing there either, so a transposed site is issued in +exactly the shape the corpus records under all three names. The choice is a +field on :class:`_Op` (``form``) rather than a line inside a builder. + +**The comparison.** Never sequential. Both implementations go into one +:func:`interleaved` call so that a neighbour arriving on the device hits both +arms at once and lands in the reported interval instead of in the conclusion. +``cudnn.benchmark`` is on, because with it off MIOpen answers from a heuristic +rather than searching and reports 5-12x worse for the *same* solver -- which +would fabricate a speedup. + +**And the comparison is what a capture costs.** ``--control none`` drops the +MIOpen arm and measures the Triton kernels alone. It is not a corner-cutting +option, it is where essentially all of the wall clock is: ``cudnn.benchmark = +True`` puts MIOpen on the Find path, whose disk record cannot be replayed in a +fresh process (``review/MIOPEN_CACHE.md`` §2), so **every** cell pays a find -- +measured on this node at 92-174 s per cell against 0.3-1.2 s for the Triton +compile, the graph capture, the calibration and the timed rounds put together +(``review/HARNESS_SPEED.md`` §1). A Triton-only row therefore carries no +``miopen_*`` and no ``speedup`` key at all -- an absent measurement stays +absent -- and ``--check``, whose reference *is* MIOpen's answer, is refused with +it. + +**The control.** The MIOpen side of a backward direction is a real forward +graph plus :func:`torch.autograd.grad`, in **all four** backward cells. Never +``torch.nn.grad.conv3d_input`` / ``conv3d_weight``: those have no real operand to +pass for the tensor being differentiated, so they fabricate a zero-strided +placeholder, and ``convolution_backward`` picks its solver from that operand's +layout. At the ``k=1x1x1`` head that made MIOpen decline its own NDHWC path and +run **3.2x** slower than the same call inside a real backward, which is where a +published 4.51x came from against a true 1.39x. + +**The timed region.** See :func:`_timed_region`. The published per-shape number +is **kernel time**: Python-side dispatch, shape re-validation, tuned-table lookup +and the launcher itself are outside it, for *both* arms, because both arms are +replayed from a CUDA graph. ``--launcher include`` gives the other number. + +**The precision.** Every speedup is a paired per-round ratio with a 95% +interval, and every cell says how many rounds it took and whether it stopped +because it converged or because it ran out of budget. ``--iters``/``--rounds`` +default to 0, i.e. decided online; pass integers to pin them. + +Usage:: + + # everything, one command + python -m triton_conv3d.bench.conv_bench --operator all --direction all \\ + --top 0 --shipped --out all.json + + # a Triton-only baseline over the form production runs: no control, so no + # find, so minutes rather than hours + python -m triton_conv3d.bench.conv_bench --operator all --direction all \\ + --top 0 --shipped --control none --form adapter --out triton.json + + python -m triton_conv3d.bench.conv_bench --top 8 --out m1.json + python -m triton_conv3d.bench.conv_bench --direction bwd-data --problems 1,3,5 + python -m triton_conv3d.bench.conv_bench --operator convT --direction bwd-weight \\ + --top 0 --shipped +""" + +from __future__ import annotations + +import argparse +import dataclasses +import json +import math +import os +import pathlib +import sys +import time +from typing import Callable, Literal, Mapping + +import torch +import torch.nn.functional as F + +from ..bwd_data import bwd_data_config, conv3d_backward_data +from ..gather_gemm import ( + ConvConfig, + candidate_configs, + conv3d_forward, + select_config, +) +from ..reduce_gemm import ( + candidate_bwd_weight_configs, + bwd_weight_config, + conv3d_backward_weight, + grad_weight_empty, + split_count, + workspace_elements, +) +from ..shapes import ( + DIRECTIONS, + ConvProblem, + Direction, + census_corpus, + scaffold_corpus, +) +from ..transposed import ( + candidate_transposed_configs, + conv_transpose3d_backward_data, + conv_transpose3d_backward_weight, + conv_transpose3d_forward, + grad_transposed_weight_empty, + transposed_config, +) +from .harness import ( + CaptureError, + capture, + capture_stream, + common_chunk, + format_table, + graph_is_worthwhile, + interleaved, + on_capture_stream, + per_call_ms, + ratio, +) + +#: See the module docstring. Set at import so that merely importing this module +#: puts the process in the configuration the numbers were taken in. +torch.backends.cudnn.benchmark = True + +_TORCH_DTYPE = {"bf16": torch.bfloat16, "fp32": torch.float32, "fp16": torch.float16} + +Operator = Literal["conv", "convT"] +OPERATORS: tuple[Operator, ...] = ("conv", "convT") + +#: How many of the sweep's fastest configs go into the run-off against MIOpen. +#: One would do if the sweep were noise-free; it is not, so its winner is partly +#: whichever config drew the luckiest sample. Racing the top few restores a +#: like-for-like best-of on both sides. +_FINALISTS = 3 + +#: Split counts the backward-weight refinement pass pins. Wide, because the +#: pass is free: the split count is a runtime argument, so none of these +#: triggers a recompile. +_SPLIT_SWEEP = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024) + + +# --------------------------------------------------------------------------- +# One cell of the operator x direction table +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class _Case: + """Everything that differs between one operator-direction and another. + + The tensors are held here rather than in locals so that the caller's + ``finally`` can drop them all at once; ``triton`` is a factory rather than a + launcher because the sweep needs one launcher per config. + + Six of these exist -- two operators by three directions -- and no field is + ever computed by asking the *problem* which operator it is. That is the + whole point of the factoring: the branch happens once, in :data:`_OPERATORS`, + and never again. + """ + + triton: Callable[[ConvConfig | None], Callable[[], object]] + #: The MIOpen control, or ``None`` when the case was built with + #: ``control=False``. Optional because building it is not free and is not + #: always wanted: for a *backward* direction the control is a real forward + #: graph, and running it once costs MIOpen's find -- measured at 92-174 s + #: per cell on this node's corpus, against 0.3-1.2 s for everything else the + #: cell does (``work/triton-conv/review/HARNESS_SPEED.md`` §1). A + #: Triton-only capture that still built the control would pay all of it. + miopen: Callable[[], object] | None + #: ``Callable[[], object]`` -- the hoistable weight prep, or ``None`` where + #: the direction has none. **All six cells are now ``None``**: the + #: consuming directions read the channels-last parameter in place, and the + #: weight-gradient directions *produce* the weight, in the layout the GEMM + #: writes natively. The field stays because the reporting path is the + #: record of what a transform would have to be charged if one came back. + transform: object + #: ``Callable[[], list[ConvConfig]]`` -- the configs worth timing. + candidates: Callable[[], list[ConvConfig]] + #: ``Callable[[list[ConvConfig]], list[ConvConfig]]`` -- a second, cheap + #: pass over the finalists on an axis the first pass held fixed. + refine: Callable[[list[ConvConfig]], list[ConvConfig]] + #: The config **this cell's entry point would resolve on its own**, computed + #: once so that ``--shipped`` measures the shipped kernel without also + #: measuring the shipped table lookup. The two are not the same number: the + #: lookup is 0.0164 ms, 39% of the transposed forward kernel + #: (``HARNESS_RIGOR.md`` H13). ``test_the_shipped_config_is_what_the_entry + #: _point_resolves`` pins these six against the entry points. + shipped_config: Callable[[], ConvConfig | None] + #: ``Callable[[], tuple[Tensor, Tensor]]`` -- ``(ours, MIOpen's)`` on this + #: cell's shape, for ``--check``. ``None`` without a control, because + #: MIOpen's answer *is* the reference. + reference: Callable[[], tuple[torch.Tensor, torch.Tensor]] | None + #: The operand whose storage decides buffer-op eligibility: the gathered one. + primary: torch.Tensor + keep: tuple + + +def _randn(shape, device, dtype): + t = torch.randn(shape, device=device, dtype=torch.float32).to(dtype) + return t.contiguous(memory_format=torch.channels_last_3d) + + +def _bias(problem: ConvProblem, device, dtype): + if not problem.bias: + return None + return torch.randn(problem.cout, device=device, dtype=torch.float32).to(dtype) + + +def _gather_refine(top): + # GROUP_M is an L2 swizzle width worth a few percent; sweeping it across the + # whole grid would double a cost that is almost entirely JIT. + return [dataclasses.replace(c, GROUP_M=8) for c in top if c.GROUP_M != 8] + + +def _bwd_weight_refine(top): + # The split count is a *runtime* argument -- it changes the grid and the + # chunk length, not a constexpr -- so this second pass costs no JIT at all, + # which is why it can afford to be exhaustive where the tile pass cannot. + return [dataclasses.replace(c, SPLIT_K=sk) for c in top for sk in _SPLIT_SWEEP] + + +def _autograd_control(build_forward: Callable[[], tuple]): + """Build a real forward graph **on the capture stream** and keep it alive. + + Two things at once, and both are load-bearing. + + The graph has to be real, because ``torch.nn.grad.conv3d_*`` fabricates a + zero-strided placeholder for the operand it is differentiating and + ``convolution_backward`` chooses its solver from that operand's layout -- + measured, 3.2x slow at the ``k=1`` head, and the source of a retracted 4.51x. + + It has to be built on :func:`~triton_conv3d.bench.harness.capture_stream`, + because otherwise the autograd node records the default stream and CUDA + graph capture refuses it outright: *"During CUDA graph capture, autograd node + ``ConvolutionBackward0`` has a stale reference to the default stream"*. That + refusal is what would push a backward cell back onto the eager path -- i.e. + onto a launcher-inclusive number for both arms -- so it is worth the two + lines it costs. + """ + s = capture_stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + made = build_forward() + torch.cuda.current_stream().wait_stream(s) + torch.cuda.synchronize() + return made + + +# -- conv: the ordinary convolution ----------------------------------------- + + +def _conv_fwd(problem: ConvProblem, device: str, + control: bool = True) -> _Case: + dtype = _TORCH_DTYPE[problem.dtype] + k = tuple(problem.kernel) + x = _randn(problem.input_shape, device, dtype) + w = _randn(problem.weight_shape, device, dtype) + b = _bias(problem, device, dtype) + y = torch.empty(problem.output_shape, device=device, dtype=dtype, + memory_format=torch.channels_last_3d) + m = problem.n * math.prod(problem.out_spatial) + + def triton(cfg): + def run(): + # ``w`` itself, not a transform of it: ``_randn`` returns it + # channels-last, which is what a ScaFFold parameter is, and the + # kernel reads that layout in place. Passing ``weight_rsck=`` here + # would time a path the integration no longer takes. + conv3d_forward(x, w, b, problem.stride, problem.padding, + config=cfg, out=y) + return run + + def miopen(): + with torch.no_grad(): + F.conv3d(x, w, b, stride=problem.stride, padding=problem.padding) + + def reference(): + got = conv3d_forward(x, w, b, problem.stride, problem.padding) + with torch.no_grad(): + ref = F.conv3d(x, w, b, stride=problem.stride, padding=problem.padding) + return got, ref + + return _Case( + triton=triton, miopen=miopen if control else None, transform=None, + candidates=lambda: candidate_configs(m, problem.cin, problem.cout, dtype), + refine=_gather_refine, + shipped_config=lambda: select_config(m, problem.cin, problem.cout, k, dtype), + reference=reference if control else None, primary=x, keep=(x, w, b, y), + ) + + +def _conv_bwd_data(problem: ConvProblem, device: str, + control: bool = True) -> _Case: + dtype = _TORCH_DTYPE[problem.dtype] + k = tuple(problem.kernel) + w = _randn(problem.weight_shape, device, dtype) + b = _bias(problem, device, dtype) + gy = _randn(problem.output_shape, device, dtype) + gx = torch.empty(problem.input_shape, device=device, dtype=dtype, + memory_format=torch.channels_last_3d) + + def triton(cfg): + def run(): + # As in the forward: the channels-last parameter is the operand, and + # the tap flip and the transpose are both constexprs in the kernel. + conv3d_backward_data(gy, w, problem.input_shape, problem.stride, + problem.padding, config=cfg, out=gx) + return run + + def build(): + xg = _randn(problem.input_shape, device, dtype).requires_grad_(True) + with torch.enable_grad(): + yg = F.conv3d(xg, w, b, stride=problem.stride, padding=problem.padding) + return xg, yg + + # ``w`` does not require grad, so ``convolution_backward``'s output mask is + # ``(True, False, False)`` and no weight gradient is computed. + # + # Not built at all without a control: this line is a real ``F.conv3d``, and + # on a shape MIOpen has not found yet it is where the find is paid. + xg, yg = _autograd_control(build) if control else (None, None) + + def miopen(): + torch.autograd.grad(yg, (xg,), gy, retain_graph=True) + + def reference(): + got = conv3d_backward_data(gy, w, problem.input_shape, problem.stride, + problem.padding) + ref = torch.autograd.grad(yg, (xg,), gy, retain_graph=True)[0] + return got, ref + + return _Case( + triton=triton, miopen=miopen if control else None, transform=None, + # Swapped: backward-data reduces over Cout and its GEMM's N is Cin. + candidates=lambda: candidate_configs( + problem.n * math.prod(problem.spatial), problem.cout, problem.cin, + dtype), + refine=_gather_refine, + shipped_config=lambda: bwd_data_config( + problem.output_shape, problem.cin, k, dtype, + padding=problem.padding), + reference=reference if control else None, primary=gy, + # ``xg``/``yg`` are kept because the graph (and so the control) dies + # with them. + keep=(gy, w, b, gx, xg, yg), + ) + + +def _conv_bwd_weight(problem: ConvProblem, device: str, + control: bool = True) -> _Case: + dtype = _TORCH_DTYPE[problem.dtype] + k = tuple(problem.kernel) + x = _randn(problem.input_shape, device, dtype) + w = _randn(problem.weight_shape, device, dtype) + b = _bias(problem, device, dtype) + gy = _randn(problem.output_shape, device, dtype) + gw = grad_weight_empty(problem.cout, problem.cin, k, dtype=dtype, device=device) + k_total = problem.n * math.prod(problem.out_spatial) + padded = any(problem.padding) + + def candidates(): + return candidate_bwd_weight_configs( + problem.cout, problem.cin, k, k_total, dtype, + splits=(0,), padded=padded, + ) + + def splits_for(cfg): + return split_count(cfg, problem.cout, problem.cin, problem.tap_count, + k_total, problem.out_spatial[2])[0] + + ws = _sweep_workspace(candidates(), splits_for, problem.cout, problem.cin, + k, device) + + def triton(cfg): + def run(): + conv3d_backward_weight(x, problem.weight_shape, gy, problem.stride, + problem.padding, config=cfg, workspace=ws, + out=gw) + return run + + def build(): + wg = w.detach().clone().requires_grad_(True) + with torch.enable_grad(): + yg = F.conv3d(x, wg, b, stride=problem.stride, padding=problem.padding) + return wg, yg + + wg, yg = _autograd_control(build) if control else (None, None) + + def miopen(): + torch.autograd.grad(yg, (wg,), gy, retain_graph=True) + + def reference(): + got = conv3d_backward_weight(x, problem.weight_shape, gy, + problem.stride, problem.padding) + ref = torch.autograd.grad(yg, (wg,), gy, retain_graph=True)[0] + return got, ref + + return _Case( + triton=triton, miopen=miopen if control else None, transform=None, + candidates=candidates, refine=_bwd_weight_refine, + shipped_config=lambda: bwd_weight_config( + problem.cout, problem.cin, k, k_total, dtype, padded=padded), + reference=reference if control else None, primary=x, + keep=(x, w, b, gy, gw, ws, wg, yg), + ) + + +# -- convT: the k == s transposed convolution -------------------------------- +# +# All three of these are served by *four* entry points across three modules, and +# the channel widths swap on the way in. Each builder names the swap once, in +# the call, rather than a shared helper naming it three times differently. + + +def _convt_fwd(problem: ConvProblem, device: str, + control: bool = True) -> _Case: + dtype = _TORCH_DTYPE[problem.dtype] + k = tuple(problem.kernel) + x = _randn(problem.input_shape, device, dtype) + w = _randn(problem.weight_shape, device, dtype) + b = _bias(problem, device, dtype) + y = torch.empty(problem.output_shape, device=device, dtype=dtype, + memory_format=torch.channels_last_3d) + m = problem.n * math.prod(problem.spatial) + + def triton(cfg): + def run(): + conv_transpose3d_forward(x, w, b, k, config=cfg, out=y) + return run + + def miopen(): + with torch.no_grad(): + F.conv_transpose3d(x, w, b, stride=k) + + def reference(): + got = conv_transpose3d_forward(x, w, b, k) + with torch.no_grad(): + ref = F.conv_transpose3d(x, w, b, stride=k) + return got, ref + + return _Case( + triton=triton, miopen=miopen if control else None, transform=None, + candidates=lambda: candidate_transposed_configs( + m, problem.cin, problem.cout, problem.tap_count, dtype), + refine=_gather_refine, + shipped_config=lambda: transposed_config(m, problem.cin, problem.cout, + k, dtype), + reference=reference if control else None, primary=x, keep=(x, w, b, y), + ) + + +def _convt_bwd_data(problem: ConvProblem, device: str, + control: bool = True) -> _Case: + """``grad_input = conv3d(grad_output, w, stride=k)`` -- an ordinary forward. + + So the config that runs is :func:`~triton_conv3d.gather_gemm.select_config`'s + for the *strided* convolution, whose ``(cin, cout)`` are this operator's + ``(cout, cin)``. ``m5_convT_bench`` had no way to say that -- the only + config object it held was a ``TransposedConfig`` -- so it dropped ``config=`` + entirely and could not sweep this direction at all. + """ + dtype = _TORCH_DTYPE[problem.dtype] + k = tuple(problem.kernel) + x = _randn(problem.input_shape, device, dtype) + w = _randn(problem.weight_shape, device, dtype) + b = _bias(problem, device, dtype) + gy = _randn(problem.output_shape, device, dtype) + gx = torch.empty(problem.input_shape, device=device, dtype=dtype, + memory_format=torch.channels_last_3d) + # The strided convolution's M is this operator's *input* volume, and its + # (cin, cout) are (Cout, Cin) of the transposed operator. + m = problem.n * math.prod(problem.spatial) + + def triton(cfg): + def run(): + conv_transpose3d_backward_data(gy, w, problem.input_shape, k, + config=cfg, out=gx) + return run + + def build(): + xg = x.detach().clone().requires_grad_(True) + with torch.enable_grad(): + yg = F.conv_transpose3d(xg, w, b, stride=k) + return xg, yg + + xg, yg = _autograd_control(build) if control else (None, None) + + def miopen(): + torch.autograd.grad(yg, (xg,), gy, retain_graph=True) + + def reference(): + got = conv_transpose3d_backward_data(gy, w, problem.input_shape, k) + ref = torch.autograd.grad(yg, (xg,), gy, retain_graph=True)[0] + return got, ref + + return _Case( + triton=triton, miopen=miopen if control else None, transform=None, + candidates=lambda: candidate_configs(m, problem.cout, problem.cin, dtype), + refine=_gather_refine, + shipped_config=lambda: select_config(m, problem.cout, problem.cin, k, dtype), + reference=reference if control else None, primary=gy, + keep=(x, w, b, gy, gx, xg, yg), + ) + + +def _convt_bwd_weight(problem: ConvProblem, device: str, + control: bool = True) -> _Case: + """The same reduction ``conv3d_backward_weight`` performs, operands swapped. + + ``grad_output`` is the strided convolution's input and ``x`` is its output + gradient, so the reduction's ``(cout, cin)`` are this operator's + ``(cin, cout)`` and its ``k_total`` is this operator's *input* volume. The + workspace is sized from the swapped widths, not from the ones a reader would + name. + """ + dtype = _TORCH_DTYPE[problem.dtype] + k = tuple(problem.kernel) + x = _randn(problem.input_shape, device, dtype) + w = _randn(problem.weight_shape, device, dtype) + b = _bias(problem, device, dtype) + gy = _randn(problem.output_shape, device, dtype) + gw = grad_transposed_weight_empty(problem.cin, problem.cout, k, + dtype=dtype, device=device) + k_total = problem.n * math.prod(problem.spatial) + + def candidates(): + return candidate_bwd_weight_configs( + problem.cin, problem.cout, k, k_total, dtype, + splits=(0,), padded=False, + ) + + def splits_for(cfg): + return split_count(cfg, problem.cin, problem.cout, problem.tap_count, + k_total, problem.spatial[2])[0] + + ws = _sweep_workspace(candidates(), splits_for, problem.cin, problem.cout, + k, device) + + def triton(cfg): + def run(): + conv_transpose3d_backward_weight(x, problem.weight_shape, gy, k, + config=cfg, workspace=ws, out=gw) + return run + + def build(): + wg = w.detach().clone().requires_grad_(True) + with torch.enable_grad(): + yg = F.conv_transpose3d(x, wg, b, stride=k) + return wg, yg + + wg, yg = _autograd_control(build) if control else (None, None) + + def miopen(): + torch.autograd.grad(yg, (wg,), gy, retain_graph=True) + + def reference(): + got = conv_transpose3d_backward_weight(x, problem.weight_shape, gy, k) + ref = torch.autograd.grad(yg, (wg,), gy, retain_graph=True)[0] + return got, ref + + return _Case( + triton=triton, miopen=miopen if control else None, transform=None, + candidates=candidates, refine=_bwd_weight_refine, + shipped_config=lambda: bwd_weight_config( + problem.cin, problem.cout, k, k_total, dtype, padded=False), + reference=reference if control else None, primary=x, + keep=(x, w, b, gy, gw, ws, wg, yg), + ) + + +def _sweep_workspace(cands, splits_for, cout, cin, k, device) -> torch.Tensor: + """One workspace, sized for the largest split count the sweep can ask for. + + So that allocation is outside every timed region -- MIOpen's own time + excludes its workspace allocation too. + """ + max_splits = max( + [splits_for(c) for c in cands] + + [splits_for(dataclasses.replace(c, SPLIT_K=sk)) + for c in cands for sk in _SPLIT_SWEEP] + ) + return torch.empty(workspace_elements(max_splits, cout, cin, k), + dtype=torch.float32, device=device) + + +# --------------------------------------------------------------------------- +# The two operators +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class _Op: + """What is per-*operator* rather than per-direction, in one place. + + Four things, and each of them is a way to measure the wrong problem: + + ``form`` + Which of the module docstring's three shapes a cell measures, as a + function of the requested form name. ``conv`` honours the name; + ``convT`` must **not** be haloed under any of them, because DistConv's + halo is ``k // 2`` for an odd kernel and 0 at ``k = 2`` and the adapter + exchanges nothing there either. Haloing a transposed problem would + grow its input by two voxels per axis and measure a convolution the + model never runs. + ``order`` + The order the cells are measured in, kept per operator so a re-capture + is comparable with what is already on disk. + ``sweepable`` + Which directions have a candidate list worth racing. + ``build`` + The six-cell table itself. + """ + + name: str + selects: Callable[[ConvProblem], bool] + form: Callable[[ConvProblem, str], ConvProblem] + form_note: Callable[[str], str] + order: Callable[[ConvProblem], object] + build: Mapping[Direction, Callable[[ConvProblem, str], _Case]] + + +#: What each ``--form`` name means, in one place, so the word a user typed and +#: the shape a kernel is handed cannot drift apart. +_FORMS: dict[str, Callable[[ConvProblem], ConvProblem]] = { + "distconv": lambda p: p.halo_variant, + "adapter": lambda p: p.production_variant, + "logical": lambda p: p, +} + +_FORM_NOTES = { + "distconv": "DistConv's halo'd, unpadded form -- what the MIOpen baseline " + "was profiled in", + "adapter": "the form ScaFFold's Triton rung is handed -- what production " + "runs, padded at every configuration", + "logical": "the module's own statement, unhalo'd and padded", +} + +_OPERATORS: dict[str, _Op] = { + "conv": _Op( + name="conv", + selects=lambda p: not p.transposed, + form=lambda p, form: _FORMS[form](p), + form_note=lambda form: _FORM_NOTES[form], + # Corpus order is measured-cost order, and it is what every stored + # capture's ``--problems`` indices refer to. + order=lambda p: 0, + build={"fwd": _conv_fwd, "bwd-data": _conv_bwd_data, + "bwd-weight": _conv_bwd_weight}, + ), + "convT": _Op( + name="convT", + selects=lambda p: p.transposed, + form=lambda p, form: p, + form_note=lambda form: "as recorded (no halo in any form at k=2)", + # Cheapest first, as ``m5_convT_bench`` ran them, so a re-capture lines + # up row for row with ``m5_shipped_*.json``. + order=lambda p: math.prod(p.spatial) * p.cin, + build={"fwd": _convt_fwd, "bwd-data": _convt_bwd_data, + "bwd-weight": _convt_bwd_weight}, + ), +} + + +def operator_of(problem: ConvProblem) -> Operator: + """Which operator a problem is. The **only** place this question is asked.""" + return "convT" if problem.transposed else "conv" + + +def _build(problem: ConvProblem, direction: Direction, device: str = "cuda", + operator: Operator | None = None, control: bool = True) -> _Case: + """Operands and launchers for one problem in one direction. + + A lookup into :data:`_OPERATORS`, not a switch: the operator is resolved + once, here, and the builder it names never asks again. + + The Triton launchers exclude allocation, deliberately: MIOpen's time + excludes its own workspace allocation, so excluding ours keeps the + comparison like-for-like. They no longer exclude a weight transform, for the + stronger reason that there is not one -- the weight operand is the + channels-last parameter itself in every direction. + + ``control=False`` builds the Triton operands and **nothing else**: no MIOpen + launcher and, for a backward direction, no autograd graph -- which is the + expensive half. Building the control does not merely cost the arm's timing; + running it once costs MIOpen's *find*, which under ``cudnn.benchmark = True`` + cannot be replayed from disk (``review/MIOPEN_CACHE.md`` §2) and which + measures **92-174 s per cell** on this corpus against 0.3-1.2 s for + everything else in the cell. A Triton-only capture is therefore two orders + of magnitude cheaper than a comparison, and that is entirely MIOpen's find. + """ + op = _OPERATORS[operator or operator_of(problem)] + try: + builder = op.build[direction] + except KeyError: + raise ValueError(f"unsupported direction {direction!r}") from None + return builder(problem, device, control) + + +# --------------------------------------------------------------------------- +# What is inside the timed region +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class _Region: + """The decision about what every arm of one cell is timed with. + + One object for the whole cell, never one per arm. ``HARNESS_RIGOR.md`` H10 + is the reason: a per-arm instrument biases a ratio *even when both arms are + individually right*, and the version of that mistake available here -- + hoisting Triton's config lookup out while leaving PyTorch's dispatch inside + the MIOpen arm -- flatters us by up to 1.4x on exactly the sub-0.15 ms cells + where it is hardest to see. + """ + + #: ``"kernel"`` (both arms replayed from a graph) or ``"call"`` (both arms + #: called from Python). + kind: str + chunk: int + fns: Mapping[str, Callable[[], object]] + eager_ms: Mapping[str, float] + note: str + + @property + def excludes_launcher(self) -> bool: + return self.kind == "kernel" + + +def _timed_region(variants: Mapping[str, Callable[[], object]], + launcher: str = "exclude") -> _Region: + """Decide, for one cell, what the published number will contain. + + **Excluded** from a ``kind="kernel"`` measurement, on **every** arm: Python + call overhead, PyTorch's dispatcher, autocast and shape re-validation, the + tuned-table lookup, MIOpen's descriptor construction and find-database + probe, the autograd engine's node walk, and the launcher itself. + **Included**: the kernels, in the order and with the operands the eager call + issues them, back to back on one stream. + + That boundary is the same on both sides only because the *whole comparison* + moves together. Measured on ``convT 1024->512 @ 8^3``, per call + (``work/triton-conv/bin/launcher_symmetry.py --only census-small``): + + ============================ ======== ======= ============ =============== + arm eager kernel host launch speedup in->out + ============================ ======== ======= ============ =============== + ``convT`` fwd Triton 0.0421 0.0282 13.9 us 1.659 -> 2.420 + ``convT`` fwd MIOpen 0.0699 0.0683 **1.6 us** + ``convT`` bwd-data Triton 0.0539 0.0350 18.9 us 1.414 -> 1.130 + ``convT`` bwd-data MIOpen 0.0761 0.0395 **36.6 us** + ``convT`` bwd-weight Triton 0.0712 0.0542 17.0 us 1.176 -> 0.931 + ``convT`` bwd-weight MIOpen 0.0839 0.0504 33.4 us + ============================ ======== ======= ============ =============== + + The host costs differ by **23x** between arms of the *same* cell, so an eager + number is not a launcher-neutral number that a graph then "improves": it is a + number with a per-arm instrument in it. Note the last column runs both ways + -- excluding the launcher is not a favour to Triton. It doubles the forward + (the Triton arm was paying 13.9 us against a 28 us kernel while MIOpen paid + 1.6 us) and it turns the weight gradient from a 1.176x win into a **0.931x + loss** (the MIOpen arm was paying an autograd-engine walk worth twice the + Triton entry point's). + + Three ways this refuses to produce a mixed measurement: + + * if any arm cannot be captured, **no** arm is -- the cell falls back to + eager whole and says so in ``note``; + * ``chunk`` comes from the shortest arm's duration alone + (:func:`~triton_conv3d.bench.harness.common_chunk`), never from a per-arm + estimate of the replay cost; + * above 40 ms per call nothing is captured, because there the largest host + cost measured on this node (0.08 ms) is under 0.2% of either arm and the + eager number is already launcher-exclusive to within a fifth of the + target precision. + """ + names = list(variants) + if launcher == "include": + return _Region("call", 1, variants, {}, + "Python call: kernel + dispatch + config lookup + launcher") + if launcher != "exclude": + raise ValueError(f"unknown launcher policy {launcher!r}") + + eager = {n: per_call_ms(variants[n]) for n in names} + durations = [eager[n] for n in names] + if not graph_is_worthwhile(durations): + return _Region("call", 1, variants, eager, + f"Python call: every arm is above {min(durations):.1f} ms, " + "where the measured host cost (<=0.08 ms) is under 0.2%") + chunk = common_chunk(durations) + captured: dict[str, Callable[[], object]] = {} + try: + for n in names: + captured[n] = capture(variants[n], chunk) + except CaptureError as exc: + captured.clear() + torch.cuda.synchronize() + torch.cuda.empty_cache() + return _Region("call", 1, variants, eager, + f"Python call: {n!r} could not be captured ({exc}), so no " + "arm was -- a mixed measurement is worth up to 1.4x") + return _Region("kernel", chunk, captured, eager, + f"CUDA graph replay, {chunk} calls per graph: kernels only, " + "no dispatch, no config lookup, no launcher, on every arm") + + +# --------------------------------------------------------------------------- +# Measurement +# --------------------------------------------------------------------------- + + +def _sweep(case: _Case, configs: list[ConvConfig], + verbose: bool = False) -> list[tuple[ConvConfig, float]]: + """Time every config once, cheaply. Failures (LDS overflow, OOM) are skipped. + + Eager, deliberately: this pass only has to *rank*, it runs hundreds of + configs, and a capture per config would cost more than the ranking is worth. + The launcher cost it carries is the same for every candidate, which is the + property a ranking needs; the finalists are then re-measured in the race, + where the launcher is excluded. + """ + ranked: list[tuple[ConvConfig, float]] = [] + for cfg in configs: + run = case.triton(cfg) + try: + run() + torch.cuda.synchronize() + except Exception as exc: # noqa: BLE001 - any compile/launch failure + if verbose: + print(f" skip {cfg}: {type(exc).__name__}: {str(exc)[:70]}") + continue + # Adaptive ``iters``, fixed 3 rounds, and no tax probe. Its old + # ``iters=3`` was not neutral between the configs it was ranking: the + # first call after a synchronize pays a queue restart worth 3% at 1.4 ms + # and 42% at 0.07 ms, so a fixed small ``iters`` charges that restart to + # whichever config is fastest -- exactly the config the sweep is trying + # to find. Sizing the block by time makes the restart the same fraction + # for every candidate. + meas = interleaved({"t": run}, warmup=None, iters=None, rounds=3, + warmup_s=0.02, warmup_min=2, block_ms=5.0, + measure_tax=False)["t"] + ranked.append((cfg, meas.median)) + ranked.sort(key=lambda kv: kv[1]) + return ranked + + +def measure_problem(problem: ConvProblem, *, direction: Direction = "fwd", + operator: Operator | None = None, + max_configs: int = 0, iters: int = 0, rounds: int = 0, + shipped: bool = False, verbose: bool = False, + budget_s: float = 20.0, target_rel: float = 0.02, + launcher: str = "exclude", control: str = "miopen") -> dict: + """Sweep, then race the finalists against MIOpen in one interleaved block. + + ``control="none"`` drops the MIOpen arm and measures the Triton kernels + alone. The row then carries no ``miopen_*`` and no ``speedup*`` key -- an + absent number is absent, not zero -- and says so in ``control``. Everything + else is unchanged: the same CUDA-graph region, the same adaptive stopping, + the same 95% interval, the same ``stop`` reason. What it buys is the whole + of MIOpen's find (§1 of ``review/HARNESS_SPEED.md``: 98.3% of a three-cell + problem's wall clock on this node); what it costs is the comparison, so use + it when the baseline is the deliverable and the ratio is not. + + ``shipped`` skips the sweep and times the config **this cell's entry point + would resolve on its own** -- the tuned table plus the heuristic fallback -- + resolved once, outside the timed region. That is the kernel a caller + actually gets; it is not the same number as the *call* a caller actually + makes, which also pays 0.0164 ms of table lookup per call at the transposed + sites, and `--launcher include` is how to see that. Confirming the shipped + config's time agrees with the sweep's is what makes the sweep's numbers a + claim about the shipped kernel rather than about a config nobody will use. + + ``iters`` and ``rounds`` default to 0, meaning *decide online*: the race + grows until the paired speedup's 95% interval is inside ``target_rel`` or + ``budget_s`` of wall clock is gone, and the row records which. Pinning both + to integers restores the old fixed 10x6 exactly, for a capture that has to + be byte-comparable with an earlier one. + + A fixed 10x6 is 60 calls whatever the kernel costs. Over this corpus that + is microseconds at the transposed sites and **45 minutes** at the 2 GiB + cliff, where one backward-weight call is 45.2 s -- and no amount of + averaging is going to change a 2789x ratio. + """ + op = operator or operator_of(problem) + row: dict = { + "problem": problem.label, + "operator": op, + "direction": direction, + "cin": problem.cin, "cout": problem.cout, + "spatial": list(problem.spatial), + "kernel": list(problem.kernel), + "padding": list(problem.padding), + "dtype": problem.dtype, + "gemm": list(problem.gemm_shape(direction)), + "flops": problem.flops(direction), + "roofline_tflops": problem.roofline_flops(direction) / 1e12, + } + if control not in ("miopen", "none"): + raise ValueError(f"unknown control {control!r}") + row["control"] = control + case = None + region = None + try: + case = _build(problem, direction, operator=op, + control=(control == "miopen")) + # ``UntypedStorage.size()`` is already in *bytes*, which is what the + # specializer's ``is_within_2gb`` compares -- multiplying by the + # element size again would report every operand as ineligible. + row["x_storage_bytes"] = case.primary.untyped_storage().size() + row["buffer_ops_eligible"] = bool( + case.primary.untyped_storage().size() <= 2**31 - 1 + ) + + if shipped: + ranked: list[tuple[ConvConfig | None, float]] = [ + (case.shipped_config(), 0.0) + ] + else: + configs = case.candidates() + if max_configs: + configs = configs[:max_configs] + ranked = _sweep(case, configs, verbose=verbose) + if not ranked: + row["error"] = "no config ran" + return row + refine = case.refine([cfg for cfg, _ in ranked[: 2 * _FINALISTS]]) + ranked += _sweep(case, refine, verbose=verbose) + ranked.sort(key=lambda kv: kv[1]) + row["configs_ran"] = len(ranked) + row["sweep"] = [[str(c), ms] for c, ms in ranked] + + variants: dict[str, Callable[[], object]] = {} + if case.miopen is not None: + variants["miopen"] = case.miopen + owner: dict[str, ConvConfig | None] = {} + for i, (cfg, _) in enumerate(ranked[:_FINALISTS]): + name = f"triton#{i}" + owner[name] = cfg + variants[name] = case.triton(cfg) + # The transform a real integration would hoist out of the call. Timed + # here so its cost is a stated number rather than an assumption. + if case.transform is not None: + variants["rsck_transform"] = case.transform + + pinned = bool(iters and rounds) + # One stream for the whole cell, both policies. The MIOpen control for + # a backward direction is an autograd graph built on this stream, and + # the engine synchronizes when it is asked to run somewhere else: 35 us + # per call, on that arm only. See :class:`on_capture_stream`. + with on_capture_stream(): + region = _timed_region(variants, launcher) + meas = interleaved(region.fns, + warmup=3 if pinned else None, + iters=iters or None, rounds=rounds or None, + budget_s=budget_s, target_rel=target_rel) + # ``region.chunk`` calls sit behind one replay, so every *absolute* time + # is that many times too large. Every *relative* one -- the half-widths, + # the paired ratio, the convergence test the harness already ran -- is + # scale-invariant and needs no correction, which is why the division + # happens here and not inside the harness. + c = region.chunk + + best_name = min( + (nm for nm in meas if nm.startswith("triton")), + key=lambda nm: meas[nm].median, + ) + best = meas[best_name] + row.update( + timed_region=region.kind, + timed_region_note=region.note, + graph_chunk=c, + triton_ms=best.median / c, + triton_best_ms=best.best / c, + triton_spread=best.spread, + triton_stall=best.stall_ratio, + triton_rel_ci=best.rel_half_width, + triton_cov=best.cov, + triton_half_width_ms=best.half_width / c, + triton_tax_frac=best.tax_frac, + triton_eager_ms=region.eager_ms.get(best_name, 0.0), + triton_config=str(owner[best_name]), + triton_pct_roofline=100 * problem.efficiency(best.median / c, direction), + triton_tflops=problem.flops(direction) / (best.median / c * 1e-3) / 1e12, + measure_rounds=len(best.rounds), + measure_iters={nm: meas[nm].iters for nm in meas}, + measure_group={nm: meas[nm].group for nm in meas}, + measure_stop=best.stop, + measure_balanced=best.balanced, + measure_seconds=best.seconds, + rsck_ms=(meas["rsck_transform"].median / c + if "rsck_transform" in meas else 0.0), + ) + # Only when there *is* a control. An absent MIOpen number is left + # absent rather than written as zero: every consumer of these rows + # reads ``speedup`` straight out, and a zero would read as a 0.000x + # result instead of as "not measured here". + if "miopen" in meas: + mio = meas["miopen"] + # Paired per round, not median-over-median: the two arms of a round + # ran seconds apart under the same device state, so a common-mode + # excursion divides out of each pair before anything is reduced. + # This is also the only quantity here that comes with an interval, + # and the interval is the point -- ``speedup`` alone has been quoted + # four times in this project's history against a number that could + # not support it. + sp = ratio(mio, best) + row.update( + miopen_ms=mio.median / c, + miopen_best_ms=mio.best / c, + miopen_spread=mio.spread, + miopen_stall=mio.stall_ratio, + miopen_rel_ci=mio.rel_half_width, + miopen_tax_frac=mio.tax_frac, + miopen_eager_ms=region.eager_ms.get("miopen", 0.0), + miopen_pct_roofline=100 * problem.efficiency(mio.median / c, + direction), + miopen_tflops=problem.flops(direction) / (mio.median / c * 1e-3) / 1e12, + speedup=sp.point, + speedup_lo=sp.lo, + speedup_hi=sp.hi, + speedup_rel_ci=sp.rel_half_width, + speedup_significant=sp.significant, + ) + # What the launcher is worth, per arm. A probe estimate (one bracketed + # block, no interval) minus the measured kernel; reported because the + # *difference* between the two arms' launchers is the bias that + # excluding them removes, and a reader should be able to see it. + # + # Not a measurement, and it shows: where the launcher is already + # negligible -- above about 0.3 ms, where an event-free bracket of a + # host-paced loop already reaches kernel throughput -- this comes out at + # a few microseconds of either sign. The measured version, with + # intervals and both policies run as full races, is + # ``work/triton-conv/bin/launcher_symmetry.py --only census``. + for arm in ("triton", "miopen"): + if f"{arm}_ms" not in row: + continue + e = row[f"{arm}_eager_ms"] + row[f"{arm}_launcher_ms"] = (e - row[f"{arm}_ms"]) if e else 0.0 + row["finalists"] = { + str(owner[nm]): meas[nm].median / c for nm in owner + } + return row + except Exception as exc: # noqa: BLE001 + row["error"] = f"{type(exc).__name__}: {exc}" + return row + finally: + del case, region + torch.cuda.synchronize() + torch.cuda.empty_cache() + + +def _correctness(problem: ConvProblem, direction: Direction, + operator: Operator | None = None) -> dict: + """Error against MIOpen's own answer, on the shape just measured. + + Not a substitute for the test suite -- ``tests/`` holds the bitwise-exact + standard -- but a benchmark that reports a time without checking the result + is how a fast wrong kernel gets believed. The reference is + :attr:`_Case.reference`, so it is the same six-cell table the timing uses and + cannot drift from it. + """ + case = None + try: + case = _build(problem, direction, operator=operator, control=True) + got, ref = case.reference() + d = (got.float() - ref.float()).abs() + scale = ref.float().pow(2).mean().sqrt().item() or 1.0 + return {"max_abs_vs_miopen": d.max().item(), + "rms_rel_vs_miopen": (d.pow(2).mean().sqrt().item() / scale)} + except Exception as exc: # noqa: BLE001 + return {"correctness_error": f"{type(exc).__name__}: {exc}"} + finally: + del case + torch.cuda.empty_cache() + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def _pick(corpus, args) -> list[tuple[int, ConvProblem]]: + """``(corpus index, problem)`` pairs, in corpus order. + + Indices are into the **corpus**, for both operators, because that is the + only stable name a problem has: ``--problems 1,3,5`` means what it has + always meant, and a transposed problem is now nameable the same way + (``m5_convT_bench``'s cheapest-first index 0 is corpus index 56). Every + printed row and every stored row carries its index. + """ + if args.problems: + return [(int(i), corpus[int(i)]) for i in args.problems.split(",")] + keep = list(enumerate(corpus)) + return keep[: args.top] if args.top else keep + + +def main() -> None: + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--operator", default="all", + choices=["conv", "convT", "all"], + help="which operator; 'all' measures both (default). " + "Supersedes the old --skip-transposed, which was " + "store_true with default True and so could never be " + "turned off") + ap.add_argument("--direction", default="fwd", + choices=["fwd", "bwd-data", "bwd-weight", "all"]) + ap.add_argument("--top", type=int, default=8, + help="hottest N corpus problems (0 = all)") + ap.add_argument("--problems", default=None, + help="comma-separated corpus indices, overriding --top") + ap.add_argument("--max-configs", type=int, default=0) + ap.add_argument("--iters", type=int, default=0, + help="calls per timed block; 0 (default) sizes it online " + "from the measured per-call time") + ap.add_argument("--rounds", type=int, default=0, + help="rounds of the race; 0 (default) grows until the " + "speedup's 95%% interval is inside --precision or " + "--budget seconds are spent") + ap.add_argument("--budget", type=float, default=20.0, + help="wall-clock seconds per cell's race (default 20)") + ap.add_argument("--precision", type=float, default=0.02, + help="target relative 95%% half-width on the reported " + "speedup and on each arm's median (default 0.02)") + ap.add_argument("--launcher", default="exclude", + choices=["exclude", "include"], + help="'exclude' (default): both arms are replayed from a " + "CUDA graph, so the number is kernel time -- no " + "dispatch, no config lookup, no launcher. 'include': " + "both arms are called from Python, so the number is " + "what a caller pays today") + ap.add_argument("--shipped", action="store_true", + help="skip the sweep and time the config the entry point " + "resolves on its own -- the kernel a caller gets") + ap.add_argument("--control", default="miopen", choices=["miopen", "none"], + help="'miopen' (default): race the Triton kernel against a " + "real MIOpen control and report a paired speedup with " + "its interval. 'none': measure the Triton kernel " + "alone, emitting no miopen_* and no speedup key. The " + "control is what makes a capture expensive -- MIOpen's " + "find cannot be replayed from disk under " + "cudnn.benchmark=True and costs 92-174 s per cell on " + "this corpus, which is 98% of a cell's wall clock " + "(work/triton-conv/review/HARNESS_SPEED.md)") + ap.add_argument("--corpus", default="scaffold", + choices=["scaffold", "census"], + help="which problem list --top/--problems index into. " + "'scaffold' (default) is the 57 profiled problems, " + "cost-ordered, and is the key every stored capture " + "refers to -- its indices must not move. 'census' is " + "the 88 problems an instrumented step actually issued " + "at all four configurations, which is the only list " + "containing configuration B and the 2048-channel " + "sites; it is already in the adapter form, carries no " + "MIOpen timings and is not cost-ordered, so --form is " + "ignored for it and --top means 'the first N', not " + "'the hottest N'") + ap.add_argument("--form", default="distconv", + choices=["distconv", "adapter", "logical"], + help="which of the three shapes of a ScaFFold convolution " + "to measure. 'distconv' (default) is the halo'd, " + "unpadded form upstream DistConv issues and the form " + "every capture on disk was taken in; 'adapter' is what " + "ScaFFold's own Triton rung is handed, which is what " + "production runs and is padded everywhere; 'logical' " + "is the module's own statement. See the module " + "docstring -- these are three different tuning " + "problems, not three views of one") + ap.add_argument("--check", action="store_true", + help="also compare each result against MIOpen's") + ap.add_argument("--verbose", action="store_true") + ap.add_argument("--out", default=None) + args = ap.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("no GPU") + if os.environ.get("PYTORCH_MIOPEN_SUGGEST_NHWC") != "1": + raise SystemExit( + "PYTORCH_MIOPEN_SUGGEST_NHWC=1 is not set: channels_last_3d is inert " + "on ROCm without it, so MIOpen would be handed NCDHW" + ) + if args.check and args.control == "none": + # MIOpen's answer *is* --check's reference, and computing it costs the + # same find the run just declined to pay. Refused rather than silently + # made expensive, because a --control none run that quietly took as long + # as a comparison would be the worst of both. The bitwise standard is + # triton_conv3d/tests/, which does not need a timing run to hold. + raise SystemExit( + "--check compares against MIOpen's own answer, so it needs " + "--control miopen; with --control none it would reintroduce the " + "find the run exists to avoid. Correctness is pinned by " + "triton_conv3d/tests/ instead." + ) + + corpus = list(census_corpus() if args.corpus == "census" + else scaffold_corpus()) + picks = _pick(corpus, args) + operators = list(OPERATORS) if args.operator == "all" else [args.operator] + directions = list(DIRECTIONS) if args.direction == "all" else [args.direction] + + props = torch.cuda.get_device_properties(0) + print(f"device {props.name}, {props.multi_processor_count} CUs, " + f"torch {torch.__version__}, cudnn.benchmark={torch.backends.cudnn.benchmark}") + print(f"control: {args.control}" + + ("" if args.control == "miopen" else + " (Triton alone; no speedup is reported and none should be " + "inferred)")) + print(f"timed region: {args.launcher} launcher " + f"({'CUDA graph replay -- kernels only, both arms' if args.launcher == 'exclude' else 'Python call -- kernel + dispatch + lookup + launcher, both arms'})") + + rows: list[dict] = [] + out_path = pathlib.Path(args.out) if args.out else None + t0 = time.time() + for opname in operators: + op = _OPERATORS[opname] + mine = sorted([(i, p) for i, p in picks if op.selects(p)], + key=lambda ip: op.order(ip[1])) + if not mine: + continue + for direction in directions: + form_note = ("recorded from a real step; already the adapter form" + if args.corpus == "census" else op.form_note(args.form)) + print(f"\n== {opname} {direction} -- {len(mine)} problems, " + f"--corpus {args.corpus} --form " + f"{'adapter' if args.corpus == 'census' else args.form}: " + f"{form_note}\n") + for idx, p in mine: + # The census records the shape *as the kernel was handed it*, + # so it is already in the adapter form and carries no halo to + # re-derive one from. Re-applying a form transform would be a + # no-op today and a silent lie the day the census gains a halo + # field, so it is skipped by name rather than by luck. + hp = p if args.corpus == "census" else op.form(p, args.form) + print(f" [{idx}] {hp.qualified_label} " + f"(GEMM {hp.gemm_shape(direction)})") + sys.stdout.flush() + row = measure_problem( + hp, direction=direction, operator=opname, + max_configs=args.max_configs, iters=args.iters, + rounds=args.rounds, shipped=args.shipped, + verbose=args.verbose, budget_s=args.budget, + target_rel=args.precision, launcher=args.launcher, + control=args.control) + row["corpus_index"] = idx + row["logical_problem"] = p.label + # The form is recorded per row, not only in the header: a row + # lifted out of one capture and quoted beside another is + # precisely how a halo'd number became "what production runs". + row["shape_form"] = ("adapter" if args.corpus == "census" + else args.form) + row["corpus"] = args.corpus + row["qualified_problem"] = hp.qualified_label + row["padding"] = list(hp.padding) + row["sites"] = list(p.sites) + if args.check and "error" not in row: + row.update(_correctness(hp, direction, operator=opname)) + rows.append(row) + _print_row(row) + sys.stdout.flush() + if out_path: + out_path.write_text(json.dumps( + {"device": props.name, "torch": torch.__version__, + "operator": args.operator, "direction": args.direction, + "shape_form": ("adapter" if args.corpus == "census" + else args.form), + "corpus": args.corpus, + "launcher": args.launcher, + "control": args.control, + "cudnn_benchmark": True, "rows": rows}, indent=1) + "\n") + + ok = [r for r in rows if "error" not in r] + # Two tables, not one with empty cells: without a control there are no + # MIOpen columns to leave blank, and a blank column in a results table is + # read as a missing value rather than as an absent measurement. + if args.control == "miopen": + print("\n" + format_table( + [[ + f"{r['operator']} {r['direction']}", + r["problem"], + f"{r['gemm'][0]}x{r['gemm'][1]}x{r['gemm'][2]}", + f"{r['triton_ms']:.4f}", f"{r['triton_pct_roofline']:.0f}%", + f"{r['miopen_ms']:.4f}", f"{r['miopen_pct_roofline']:.0f}%", + f"{r['speedup']:.3f}x", + f"+-{r['speedup_rel_ci']:.1%}" + + ("" if r["speedup_significant"] else "?"), + f"{r['measure_rounds']}/{r['measure_stop'][:4]}", + f"{r['timed_region'][:4]}x{r['graph_chunk']}", + r["triton_config"], + ] for r in ok], + ["cell", "problem", "M x N x K", "triton ms", "%roof", "miopen ms", + "%roof", "speedup", "95% CI", "rounds", "timed", "best config"], + aligns="lllrrrrrrrrl", + )) + else: + print("\n" + format_table( + [[ + f"{r['operator']} {r['direction']}", + r["qualified_problem"], + f"{r['gemm'][0]}x{r['gemm'][1]}x{r['gemm'][2]}", + f"{r['triton_ms']:.4f}", + f"+-{r['triton_rel_ci']:.1%}", + f"{r['triton_cov']:.2%}", + f"{r['triton_pct_roofline']:.0f}%", + f"{r['triton_tflops']:.1f}", + f"{r['measure_rounds']}/{r['measure_stop'][:4]}", + f"{r['timed_region'][:4]}x{r['graph_chunk']}", + r["triton_config"], + ] for r in ok], + ["cell", "problem", "M x N x K", "triton ms", "95% CI", "CoV", + "%roof", "TFLOP/s", "rounds", "timed", "config"], + aligns="lllrrrrrrrl", + )) + print(f"\nelapsed {time.time() - t0:.0f} s") + if out_path: + print(f"wrote {out_path}") + + +def _print_row(row: dict) -> None: + if "error" in row: + print(f" ERROR {row['error'][:110]}") + return + def launcher(arm): + v = row.get(f"{arm}_launcher_ms", 0.0) + return f", launcher +{v:.4f}" if v else "" + out = ( + f" triton {row['triton_ms']:8.4f} +-{row['triton_rel_ci']:.1%} ms " + f"({row['triton_pct_roofline']:5.1f}% roof, stall {row['triton_stall']:.2f}x, " + f"instrument {row['triton_tax_frac']:+.1%}{launcher('triton')}) " + f"{row['triton_config']}\n" + ) + if "miopen_ms" in row: + out += ( + f" miopen {row['miopen_ms']:8.4f} +-{row['miopen_rel_ci']:.1%} ms " + f"({row['miopen_pct_roofline']:5.1f}% roof, " + f"stall {row['miopen_stall']:.2f}x, " + f"instrument {row['miopen_tax_frac']:+.1%}{launcher('miopen')})\n" + f" speedup {row['speedup']:.3f}x " + f"[{row['speedup_lo']:.3f}, {row['speedup_hi']:.3f}] " + f"{row['measure_rounds']}r/{row['measure_stop']}" + f"{'' if row['speedup_significant'] else ' NOT SIGNIFICANT'}" + ) + else: + # No control ran. Say so where the speedup would have been, rather + # than leaving a blank a reader could take for a missing win. + out += (f" no MIOpen control (--control none) " + f"{row['measure_rounds']}r/{row['measure_stop']}") + out += f" in {row['measure_seconds']:.1f}s {row['timed_region_note']}" + if row.get("rsck_ms"): + out += (f"\n weight transform {row['rsck_ms']:.4f} ms " + f"({100*row['rsck_ms']/row['triton_ms']:.1f}% of kernel)") + if "max_abs_vs_miopen" in row: + out += f"\n max_abs vs MIOpen {row['max_abs_vs_miopen']:.3e}" + print(out) + + +if __name__ == "__main__": + main() diff --git a/triton_conv3d/bench/gemm_probe.py b/triton_conv3d/bench/gemm_probe.py new file mode 100644 index 0000000..b5146be --- /dev/null +++ b/triton_conv3d/bench/gemm_probe.py @@ -0,0 +1,1009 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""M0 gate: what can ``tl.dot`` actually reach on this device, at our shapes? + +Every convolution direction reduces to a GEMM. Before writing a gather-GEMM +convolution it is worth knowing the ceiling of the thing it is built on, because +no amount of clever addressing recovers throughput the matrix core never had. + +Three measurements, in increasing specificity: + +``peak`` + A large square GEMM. Calibrates Triton against ``torch.matmul`` + (hipBLASLt) and against the 600 TFLOP/s bf16 constant, so every later + percentage has a known reference. + +``compute`` + The conv-implied ``(M, N, K)``, but with ``A``'s M-stride set to zero so + every row tile reads the same cached rows. The FLOP count is unchanged and + DRAM traffic is negligible, which isolates matrix-core throughput in the + *shape regime* our convolutions live in -- skinny ``N`` (64-512), long ``K`` + (81-27,648), enormous ``M``. This is the real ceiling for a fused kernel: + a convolution reads its input once, so it is compute-bound wherever its + arithmetic intensity exceeds the 182 FLOP/byte crossover, which is almost + everywhere in the corpus. + +``dram`` + The same shape with real strides. Always slower, and *not* a ceiling for + the convolution -- materializing im2col multiplies ``A``'s bytes by the tap + count, which is exactly the traffic a fused kernel avoids. Measured anyway, + because the gap between ``compute`` and ``dram`` is how much a fused kernel + stands to gain over the explicit-GEMM approach. + +Usage:: + + python -m triton_conv3d.bench.gemm_probe --mode peak + python -m triton_conv3d.bench.gemm_probe --top 8 --out probe.json +""" + +from __future__ import annotations + +import argparse +import dataclasses +import itertools +import json +import pathlib +import sys +import time + +import torch +import triton +import triton.language as tl + +from ..shapes import ( + DIRECTIONS, + HBM_BYTES_PER_S, + PEAK_FLOPS, + ConvProblem, + hot_corpus, +) +from .harness import flush_caches, format_table, interleaved + + +# --------------------------------------------------------------------------- +# A plain, honest GEMM +# --------------------------------------------------------------------------- + + +@triton.jit +def _gemm_kernel( + a_ptr, b_ptr, c_ptr, + M, N, K, + stride_am, stride_ak, + stride_bk, stride_bn, + stride_cm, stride_cn, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + GROUP_M: tl.constexpr, + SPLIT_K: tl.constexpr, + INT32_OFF: tl.constexpr, +): + """Textbook tiled GEMM with grouped-M ordering and optional split-K. + + Deliberately unremarkable: the point of the probe is to measure what a + competent-but-ordinary Triton GEMM achieves, so that a later convolution + kernel's number can be read as "this much of the available throughput" + rather than against an unknown. + """ + pid = tl.program_id(axis=0) + pid_k = tl.program_id(axis=1) + + grid_m = tl.cdiv(M, BLOCK_M) + grid_n = tl.cdiv(N, BLOCK_N) + + # Group consecutive programs along M so that a group shares B tiles in L2. + width = GROUP_M * grid_n + group_id = pid // width + group_size = min(grid_m - group_id * GROUP_M, GROUP_M) + pid_m = group_id * GROUP_M + (pid % group_size) + pid_n = (pid % width) // group_size + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) + + # ``INT32_OFF`` keeps every offset tensor ``i32``, which is condition 2 of + # the buffer-load fast path (briefing 1.3): ``canUseBufferOps`` bails out + # with ``if (ofstBit != 32) return false;`` before it ever looks at the + # range. The int64 form below is the safe default -- ``M`` reaches 8.4M + # here and ``M * stride`` overflows i32 for the real strides -- so the flag + # exists to *measure* what the promotion costs, not to be switched on + # blindly. + if INT32_OFF: + a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak + b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn + else: + a_ptrs = a_ptr + offs_m[:, None].to(tl.int64) * stride_am + offs_k[None, :] * stride_ak + b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :].to(tl.int64) * stride_bn + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + k_step = BLOCK_K * SPLIT_K + for k in range(pid_k * BLOCK_K, K, k_step): + k_mask = offs_k[None, :] + k - pid_k * BLOCK_K < K - k + pid_k * BLOCK_K + a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & k_mask, other=0.0) + b = tl.load( + b_ptrs, + mask=(offs_k[:, None] + k - pid_k * BLOCK_K < K - k + pid_k * BLOCK_K) + & (offs_n[None, :] < N), + other=0.0, + ) + acc = tl.dot(a, b, acc) + a_ptrs += k_step * stride_ak + b_ptrs += k_step * stride_bk + + c_ptrs = c_ptr + offs_m[:, None].to(tl.int64) * stride_cm + offs_n[None, :] * stride_cn + c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) + if SPLIT_K == 1: + tl.store(c_ptrs, acc, mask=c_mask) + else: + tl.atomic_add(c_ptrs, acc, mask=c_mask) + + +@dataclasses.dataclass(frozen=True) +class GemmConfig: + BLOCK_M: int + BLOCK_N: int + BLOCK_K: int + GROUP_M: int = 8 + SPLIT_K: int = 1 + num_warps: int = 8 + num_stages: int = 2 + #: AMD backend kernargs. ``None`` means "do not pass it at all", which is + #: not the same as passing 0: the M0 sweep never passed them, so keeping a + #: distinct sentinel lets the old numbers be reproduced exactly. + matrix_instr_nonkdim: int | None = None + kpack: int | None = None + waves_per_eu: int | None = None + int32_offsets: bool = False + + def __str__(self) -> str: + s = ( + f"{self.BLOCK_M}x{self.BLOCK_N}x{self.BLOCK_K}" + f"/g{self.GROUP_M}/sk{self.SPLIT_K}/w{self.num_warps}/s{self.num_stages}" + ) + if self.matrix_instr_nonkdim is not None: + s += f"/nk{self.matrix_instr_nonkdim}" + if self.kpack is not None: + s += f"/kp{self.kpack}" + if self.waves_per_eu: + s += f"/we{self.waves_per_eu}" + if self.int32_offsets: + s += "/i32" + return s + + def amd_kwargs(self) -> dict: + kw = {} + if self.matrix_instr_nonkdim is not None: + kw["matrix_instr_nonkdim"] = self.matrix_instr_nonkdim + if self.kpack is not None: + kw["kpack"] = self.kpack + if self.waves_per_eu is not None: + kw["waves_per_eu"] = self.waves_per_eu + return kw + + +#: Curated tile shapes rather than a full product sweep. A product over +#: (BM, BN, BK, warps, stages) is ~250 configs per shape, and every one costs a +#: JIT compile; at 57 problems x 3 directions that is hours of compilation to +#: answer a yes/no question. These are the shapes that matter on CDNA3: square +#: tiles for balanced GEMMs, tall-skinny tiles for the huge-M/small-N regime the +#: convolutions actually live in, and a couple of small tiles for the 8^3 sites. +_TILES: tuple[tuple[int, int, int], ...] = ( + (256, 128, 64), + (256, 64, 64), + (128, 256, 64), + (128, 128, 64), + (128, 128, 32), + (128, 64, 64), + (128, 64, 128), + (64, 128, 64), + (64, 64, 64), + (64, 64, 128), + (32, 64, 128), + (64, 32, 128), +) + +#: PyTorch Inductor's ROCm convolution seed grid, ``(BLOCK_M, BLOCK_N, BLOCK_K, +#: num_warps)``. Verbatim from ``torch/_inductor/heuristics/template/triton.py`` +#: (``BaseConfigHeuristic.conv_configs``, which ``ROCmConfigHeuristic`` +#: inherits); its per-config ``num_stages`` is dropped because +#: ``ROCmConfigHeuristic._filter_configs`` force-overwrites it with +#: ``get_backend_num_stages()`` == 2 on HIP. Preferred over a blind sweep +#: because these values are already tuned on ROCm. +_ROCM_CONV_TILES: tuple[tuple[int, int, int, int], ...] = ( + (64, 256, 16, 4), + (256, 64, 16, 4), + (1024, 16, 16, 8), + (128, 128, 32, 8), + (64, 64, 32, 4), + (64, 256, 32, 8), + (256, 64, 32, 8), + (128, 128, 64, 8), + (64, 128, 64, 4), + (128, 64, 64, 4), + (256, 128, 64, 8), + (128, 256, 64, 8), + (128, 128, 128, 8), + (64, 128, 128, 4), + (256, 128, 128, 8), + (128, 256, 128, 8), +) + +#: Extra tiles for the skinny-N regime, where the seed grid runs out of shapes. +#: At ``N=64`` a ``BLOCK_N=64`` tile is only four 16x16 MFMA tiles wide, so the +#: N axis cannot absorb warps; these trade N width for M depth (and, at +#: ``BLOCK_N=32``, test whether going *narrower* and taller helps at all). +_SKINNY_N_TILES: tuple[tuple[int, int, int, int], ...] = ( + (256, 64, 64, 8), + (512, 64, 64, 8), + (256, 64, 128, 8), + (512, 64, 32, 8), + (128, 64, 128, 4), + (64, 64, 64, 4), + (64, 64, 128, 4), + (256, 32, 128, 8), + (512, 32, 64, 8), + (128, 32, 128, 4), + (1024, 64, 32, 8), +) + +#: MFMA k-dimension per ``matrix_instr_nonkdim`` for bf16 on gfx942 +#: (``mfmaVersion == 3``). ``BLOCK_K`` must be a multiple of this or +#: ``chooseMfmaInstruction`` fails outright ("would introduce data duplication") +#: and the dot silently lowers to FMA. Source: Triton v3.7.0 +#: ``MfmaGroup.cpp`` (``TRITON_MFMA_v(3, 16, 16, bf16T, bf16T, +#: mfma_f32_16x16x16bf16_1k, 16, 4)`` and the 32x32x8 entry) plus the +#: ``inputKSize % kDim`` check in ``AccelerateAMDMatmul.cpp``. +_MFMA_KDIM_BF16 = {16: 16, 32: 8} + + +def _default_kpack(block_k: int) -> int: + """Inductor's arch-aware kpack default, ``get_default_kpack`` in utils.py. + + ``kWidth = kBase * kPack``; ``kpack=2`` means ``ds_read_b128``, the widest + LDS load. On gfx942 Inductor keeps it at 1 for ``BLOCK_K <= 16``, where the + wider read has nothing to read. + """ + return 1 if block_k <= 16 else 2 + + +def _amd_config(bm: int, bn: int, bk: int, warps: int, *, group_m: int, + split_k: int, nonkdim: int, kpack: int | None = None, + waves_per_eu: int = 0, int32: bool = False) -> GemmConfig | None: + """One AMD-knob config, or ``None`` if a hard constraint rejects it. + + The constraints are not preferences. ``BLOCK_M``/``BLOCK_N`` not divisible + by ``nonkdim`` is what Inductor's ``_finalize_mm_configs`` prunes; a + ``BLOCK_K`` that is not a multiple of the intrinsic's ``kDim`` is what the + Triton pass rejects. Both failure modes are *silent* -- the kernel still + runs, just on the FMA path -- so a config that violates them would quietly + contribute a meaningless number to a best-of sweep. + """ + if nonkdim and (bm % nonkdim or bn % nonkdim): + return None + if bk % _MFMA_KDIM_BF16.get(nonkdim, 16): + return None + # Each warp owns a 16x16 tile; more warps than tiles leaves warps idle. + warps = min(warps, bm * bn // 256) + if warps < 1: + return None + return GemmConfig( + bm, bn, bk, group_m, split_k, warps, num_stages=2, + matrix_instr_nonkdim=nonkdim, + kpack=_default_kpack(bk) if kpack is None else kpack, + waves_per_eu=waves_per_eu, + int32_offsets=int32, + ) + + +def _split_ks(m: int, n: int, k: int, bm: int, bn: int) -> tuple[int, ...]: + """Split-K only where the M/N grid cannot fill the 228 CUs on its own.""" + tiles = ((m + bm - 1) // bm) * ((n + bn - 1) // bn) + return (1, 4, 16) if tiles < 228 and k >= 2048 else (1,) + + +def _candidate_configs(m: int, n: int, k: int, knobs: str = "legacy") -> list[GemmConfig]: + """Configs worth trying for one shape. + + Not ``@triton.autotune``: that recompiles inside whatever is running at the + time, which is the wrong behaviour both here (it would pollute the timing) + and in production (ScaFFold's figure of merit is total wall time). + + ``knobs``: + + ``legacy`` + The M0 sweep: curated tiles, ``GROUP_M=8``, no AMD kernargs at all. + Kept verbatim so the earlier numbers stay reproducible. + + ``amd`` + Inductor's ROCm conv seed grid under the gfx942 constraints, with + ``matrix_instr_nonkdim=16`` and the arch-aware ``kpack``. ``GROUP_M`` + sweeps 6 (MI300A has 6 XCDs; AMD's L2-swizzle rule is "multiple of the + XCD count") against the MI300X-derived 8 that the M0 sweep used, so the + report can say which actually won rather than assuming. + + ``amd-wide`` + ``amd`` plus the skinny-N tiles, ``matrix_instr_nonkdim`` in + ``{16, 32}`` and ``GROUP_M`` in ``{6, 8, 12}``. For the ``N=64`` + interrogation; several times the configs, so not the default. + """ + m2 = max(32, triton.next_power_of_2(m)) + n2 = max(16, triton.next_power_of_2(n)) + k2 = max(32, triton.next_power_of_2(k)) + + def oversized(bm: int, bn: int, bk: int) -> bool: + # Skip tiles that mostly compute padding. + return bm > 2 * m2 or bn > 2 * n2 or bk > 2 * k2 + + out: list[GemmConfig] = [] + if knobs == "legacy": + for bm, bn, bk in _TILES: + if oversized(bm, bn, bk): + continue + for warps in (4, 8): + for sk in _split_ks(m, n, k, bm, bn): + out.append(GemmConfig(bm, bn, bk, 8, sk, warps, num_stages=2)) + return out + + if knobs == "amd": + tiles, nonkdims, group_ms = _ROCM_CONV_TILES, (16,), (6, 8) + elif knobs == "amd-wide": + tiles = _ROCM_CONV_TILES + _SKINNY_N_TILES + nonkdims, group_ms = (16, 32), (6, 8, 12) + else: + raise ValueError(f"unknown knob set {knobs!r}") + + seen: set[GemmConfig] = set() + for bm, bn, bk, seed_warps in tiles: + if oversized(bm, bn, bk): + continue + # Sweep both warp counts, then clamp; Inductor ships one value per tile + # but we are measuring a ceiling, not reproducing its choice. + for warps in {4, 8, seed_warps}: + for nonkdim in nonkdims: + for group_m in group_ms: + for sk in _split_ks(m, n, k, bm, bn): + cfg = _amd_config(bm, bn, bk, warps, group_m=group_m, + split_k=sk, nonkdim=nonkdim) + if cfg is not None and cfg not in seen: + seen.add(cfg) + out.append(cfg) + return out + + +def _randn(shape: tuple[int, ...], device, dtype: torch.dtype) -> torch.Tensor: + """Random operand allocated directly in ``dtype``. + + Going through fp32 and casting doubles peak memory, which for the largest + corpus shapes means a 29 GiB operand briefly needs 87 GiB and the allocator + spends minutes thrashing before it gets there. + """ + return torch.randn(shape, device=device, dtype=dtype) + + +def _launch(a, b, c, cfg: GemmConfig, *, m: int, n: int, k: int, + stride_am: int, stride_bn: int): + grid = (triton.cdiv(m, cfg.BLOCK_M) * triton.cdiv(n, cfg.BLOCK_N), cfg.SPLIT_K) + _gemm_kernel[grid]( + a, b, c, + m, n, k, + stride_am, 1, + b.stride(0), stride_bn, + c.stride(0), c.stride(1), + BLOCK_M=cfg.BLOCK_M, BLOCK_N=cfg.BLOCK_N, BLOCK_K=cfg.BLOCK_K, + GROUP_M=cfg.GROUP_M, SPLIT_K=cfg.SPLIT_K, + INT32_OFF=cfg.int32_offsets, + num_warps=cfg.num_warps, num_stages=cfg.num_stages, + **cfg.amd_kwargs(), + ) + + +#: Refuse to allocate an operand bigger than this in ``compute`` mode; above it +#: the operand is replaced by a stride-0 broadcast of the reduction axis. 1 GiB +#: comfortably exceeds MI300A's 256 MiB last level, so anything under it streams +#: from cache and anything over it would be measuring DRAM instead of the +#: matrix core. +_RESIDENT_BUDGET = 1 << 30 +#: Refuse a ``dram``-mode shape whose materialized operands would not fit. +_DRAM_BUDGET = 24 << 30 + + +def plan_operands(m: int, n: int, k: int, mode: str, elem: int) -> dict | None: + """Decide how to allocate A and B, or ``None`` if the shape cannot be run. + + In ``compute`` mode an operand's non-reduction axis gets stride 0 whenever + materializing it would spill out of cache: the same rows (or columns) are + re-read by every tile. The FLOP count is untouched, so the measured rate is + matrix-core throughput at this ``(M, N, K)`` with DRAM taken out of the + picture -- which is the ceiling a *fused* convolution kernel is entitled to + aim at, since it reads its input once rather than ``tap_count`` times. + """ + a_bytes, b_bytes, c_bytes = m * k * elem, k * n * elem, m * n * 4 + if mode == "dram": + if a_bytes + b_bytes + c_bytes > _DRAM_BUDGET: + return None + return {"a_rows": m, "stride_am": None, "b_cols": n, "stride_bn": None} + plan = {"a_rows": m, "stride_am": None, "b_cols": n, "stride_bn": None} + if a_bytes > _RESIDENT_BUDGET: + plan["a_rows"], plan["stride_am"] = min(m, 256), 0 + if b_bytes > _RESIDENT_BUDGET: + # Broadcast one column across N. Only reachable for backward-weight, + # where K is the whole volume and B is the im2col'd activation -- the + # 58 GiB tensor a fused kernel never builds. + plan["b_cols"], plan["stride_bn"] = 1, 0 + if c_bytes > _DRAM_BUDGET: + return None + return plan + + +def best_triton_gemm( + m: int, n: int, k: int, *, dtype=torch.bfloat16, mode: str = "compute", + device="cuda", max_configs: int = 0, verbose: bool = False, + knobs: str = "legacy", sink: list | None = None, +) -> tuple[float, GemmConfig | None, int]: + """Sweep configs, return ``(best ms/call, config, n_configs_that_ran)``. + + ``float('inf')`` with a ``None`` config means the shape could not be run at + all in this mode -- which for ``dram`` is itself the finding. + + ``sink``, if given, collects ``(config string, ms)`` for every config that + ran. Only the winner is reported normally, but "which knob won, and by how + much over second place" is the question the N=64 investigation asks, and it + is not answerable from a single best time. + """ + plan = plan_operands(m, n, k, mode, torch.finfo(dtype).bits // 8) + if plan is None: + return float("inf"), None, 0 + + a, b, c, stride_am, stride_bn = _alloc(plan, m, n, k, device, dtype) + configs = _candidate_configs(m, n, k, knobs) + if max_configs: + configs = configs[:max_configs] + try: + return _sweep(a, b, c, configs, m=m, n=n, k=k, stride_am=stride_am, + stride_bn=stride_bn, verbose=verbose, sink=sink) + finally: + del a, b, c + torch.cuda.empty_cache() + + +def _alloc(plan: dict, m: int, n: int, k: int, device, dtype): + a = _randn((plan["a_rows"], k), device, dtype) + b = _randn((k, plan["b_cols"]), device, dtype) + c = torch.empty((m, n), device=device, dtype=torch.float32) + return (a, b, c, + a.stride(0) if plan["stride_am"] is None else 0, + b.stride(1) if plan["stride_bn"] is None else 0) + + +def _sweep(a, b, c, configs: list[GemmConfig], *, m: int, n: int, k: int, + stride_am: int, stride_bn: int, verbose: bool = False, + sink: list | None = None) -> tuple[float, GemmConfig | None, int]: + """Time each config on already-allocated operands; return the winner.""" + best_ms, best_cfg, ran = float("inf"), None, 0 + for cfg in configs: + launch = _launcher(a, b, c, cfg, m=m, n=n, k=k, + stride_am=stride_am, stride_bn=stride_bn) + try: + if cfg.SPLIT_K > 1: + c.zero_() + launch() + torch.cuda.synchronize() + except Exception as exc: # OOM, LDS overflow, unsupported tiling + if verbose: + print(f" skip {cfg}: {type(exc).__name__}: {str(exc)[:80]}") + continue + ran += 1 + meas = interleaved({"t": launch}, warmup=2, iters=3, rounds=3)["t"] + if sink is not None: + sink.append((str(cfg), meas.median)) + if meas.median < best_ms: + best_ms, best_cfg = meas.median, cfg + return best_ms, best_cfg, ran + + +def _launcher(a, b, c, cfg: GemmConfig, **kw): + return lambda: _launch(a, b, c, cfg, **kw) + + +def torch_gemm_ms(m: int, n: int, k: int, *, dtype=torch.bfloat16, device="cuda") -> float: + """hipBLASLt's time for the same GEMM -- the library reference.""" + a = _randn((m, k), device, dtype) + b = _randn((k, n), device, dtype) + try: + meas = interleaved({"t": lambda: torch.matmul(a, b)}, warmup=5, iters=5, rounds=5) + return meas["t"].median + finally: + del a, b + torch.cuda.empty_cache() + + +# --------------------------------------------------------------------------- +# Modes +# --------------------------------------------------------------------------- + + +def run_peak(sizes=(2048, 4096, 8192), dtype=torch.bfloat16, + knobs: str = "legacy") -> list[dict]: + peak = PEAK_FLOPS["bf16"] if dtype is torch.bfloat16 else PEAK_FLOPS["fp32"] + rows = [] + for s in sizes: + flops = 2 * s * s * s + tri_ms, cfg, ran = best_triton_gemm(s, s, s, dtype=dtype, mode="dram", + knobs=knobs) + tor_ms = torch_gemm_ms(s, s, s, dtype=dtype) + rows.append({ + "size": s, + "triton_ms": tri_ms, "triton_tflops": flops / (tri_ms * 1e-3) / 1e12, + "triton_pct_peak": 100 * flops / (tri_ms * 1e-3) / peak, + "triton_config": str(cfg), "configs_ran": ran, + "torch_ms": tor_ms, "torch_tflops": flops / (tor_ms * 1e-3) / 1e12, + "torch_pct_peak": 100 * flops / (tor_ms * 1e-3) / peak, + "triton_vs_torch": tor_ms / tri_ms, + }) + print(format_table( + [[ + r["size"], + f"{r['triton_tflops']:.1f}", f"{r['triton_pct_peak']:.1f}%", + f"{r['torch_tflops']:.1f}", f"{r['torch_pct_peak']:.1f}%", + f"{r['triton_vs_torch']:.2f}x", r["triton_config"], + ] for r in rows[-1:]], + ["MNK", "triton TF/s", "%peak", "torch TF/s", "%peak", "tri/torch", "config"], + aligns="rrrrrrl", + ) if len(rows) == 1 else " " + " ".join([ + str(rows[-1]["size"]), + f"{rows[-1]['triton_tflops']:.1f}", f"{rows[-1]['triton_pct_peak']:.1f}%", + f"{rows[-1]['torch_tflops']:.1f}", f"{rows[-1]['torch_pct_peak']:.1f}%", + f"{rows[-1]['triton_vs_torch']:.2f}x", rows[-1]["triton_config"], + ])) + sys.stdout.flush() + return rows + + +def run_peak_compare(sizes=(2048, 4096, 8192), knob_sets=("legacy", "amd"), + dtype=torch.bfloat16, device="cuda") -> list[dict]: + """Peak calibration where the knob sets are compared *against each other*. + + :func:`run_peak` sweeps one knob set and reports its winner, which is fine + for a single number but useless for a before/after: two sweeps run minutes + apart are two different machines. Measured here, three passes over the same + ``legacy`` grid spanned 386-428 TF/s at 8192 -- an 11% band with no code + change at all, which is larger than most knob effects we are looking for. + + So: sweep each knob set to find its own champion, then put the champions + (and hipBLASLt) into a single :func:`interleaved` call. Drift then hits + every variant equally and lands in the reported spread instead of in the + conclusion. The sweep-time numbers are kept alongside as ``*_sweep_tflops`` + precisely so the size of that effect stays visible. + """ + peak = PEAK_FLOPS["bf16"] if dtype is torch.bfloat16 else PEAK_FLOPS["fp32"] + rows = [] + for s in sizes: + flops = 2 * s * s * s + plan = plan_operands(s, s, s, "dram", torch.finfo(dtype).bits // 8) + assert plan is not None + a, b, c, stride_am, stride_bn = _alloc(plan, s, s, s, device, dtype) + try: + variants: dict = {} + owner: dict = {} + row: dict = {"size": s} + for ks in knob_sets: + configs = _candidate_configs(s, s, s, ks) + sink: list = [] + ms, cfg, ran = _sweep(a, b, c, configs, m=s, n=s, k=s, + stride_am=stride_am, stride_bn=stride_bn, + sink=sink) + row[f"{ks}_config"] = str(cfg) + row[f"{ks}_configs_ran"] = ran + row[f"{ks}_sweep_tflops"] = flops / (ms * 1e-3) / 1e12 + by_str = {str(x): x for x in configs} + for i, (name, _) in enumerate( + sorted(sink, key=lambda kv: kv[1])[:_FINALISTS]): + owner[f"{ks}#{i}"] = (ks, by_str[name]) + variants[f"{ks}#{i}"] = _launcher( + a, b, c, by_str[name], m=s, n=s, k=s, + stride_am=stride_am, stride_bn=stride_bn) + owner["torch"] = ("torch", None) + variants["torch"] = lambda: torch.matmul(a, b) + meas = interleaved(variants, warmup=5, iters=5, + rounds=2 * len(variants)) + for name, m_ in meas.items(): + ks, cfg = owner[name] + if m_.median >= row.get(f"{ks}_ms", float("inf")): + continue + row[f"{ks}_ms"] = m_.median + row[f"{ks}_tflops"] = flops / (m_.median * 1e-3) / 1e12 + row[f"{ks}_pct_peak"] = 100 * flops / (m_.median * 1e-3) / peak + row[f"{ks}_spread"] = m_.spread + # Launch-gap inflation, from ``harness._time_block``. A value + # far above 1 means the GPU idled between launches and the + # number is not kernel time; recorded so a reader can see + # whether the device was actually ours for the duration. + row[f"{ks}_stall"] = m_.stall_ratio + if cfg is not None: + row[f"{ks}_config"] = str(cfg) + finally: + del a, b, c + torch.cuda.empty_cache() + rows.append(row) + names = list(knob_sets) + ["torch"] + print(" " + " ".join( + [f"MNK={s:<5d}"] + + [f"{name}={row[f'{name}_tflops']:6.1f} TF/s" + f" ({row[f'{name}_pct_peak']:5.1f}%, spread {row[f'{name}_spread']:.1%}," + f" stall {row[f'{name}_stall']:.2f}x)" + for name in names if f"{name}_tflops" in row] + )) + for ks in knob_sets: + print(f" {ks:8s} winner {row[f'{ks}_config']} " + f"[{row[f'{ks}_configs_ran']} configs, " + f"{row[f'{ks}_sweep_tflops']:.1f} TF/s during the sweep]") + sys.stdout.flush() + return rows + + +#: Named knob-set *comparisons*. A bare knob-set name sweeps that set alone; +#: these run several sets over the same operands and finish with a single +#: interleaved head-to-head between their champions, which is the only honest +#: way to state a before/after: two sweeps minutes apart need not have shared +#: the device with the same neighbours. +_KNOB_SETS: dict[str, tuple[str, ...]] = { + "compare": ("legacy", "amd"), + "compare-wide": ("legacy", "amd", "amd-wide"), +} + + +#: How many of each knob set's fastest configs go into the run-off. +#: One would be enough if the sweep were noise-free. It is not: a sweep reports +#: ``min`` over ~50 noisy samples, so its winner is partly the config that got +#: the luckiest sample, and re-timing that one config reproduces the luck only +#: sometimes. Racing the top few and taking each set's best restores the +#: like-for-like comparison -- both sides get a best-of, in the same interleaved +#: measurement. +_FINALISTS = 3 + + +def _measure_cell(m: int, n: int, k: int, *, mode: str, dtype, knob_sets, + keep_all: bool, max_configs: int, device="cuda") -> dict | None: + """Sweep each knob set on one shape, then race the finalists. + + Returns ``None`` when the shape does not fit this mode's budget. Operands + are allocated once and shared by every knob set, so the comparison is not + confounded by a different allocation or a differently warmed cache. + """ + plan = plan_operands(m, n, k, mode, torch.finfo(dtype).bits // 8) + if plan is None: + return None + a, b, c, stride_am, stride_bn = _alloc(plan, m, n, k, device, dtype) + try: + out: dict = {"winners": {}, "ran": {}, "sweep": {}, "sinks": {}, + "headtohead": {}, "spread": {}, "stall": {}} + finalists: dict[str, list[GemmConfig]] = {} + for ks in knob_sets: + configs = _candidate_configs(m, n, k, ks) + if max_configs: + configs = configs[:max_configs] + sink: list = [] + ms, cfg, ran = _sweep(a, b, c, configs, m=m, n=n, k=k, + stride_am=stride_am, stride_bn=stride_bn, + sink=sink) + if cfg is None: + continue + out["winners"][ks], out["ran"][ks], out["sweep"][ks] = cfg, ran, ms + ranked = sorted(sink, key=lambda kv: kv[1]) + if keep_all: + out["sinks"][ks] = ranked + by_str = {str(x): x for x in configs} + finalists[ks] = [by_str[name] for name, _ in ranked[:_FINALISTS]] + + if len(out["winners"]) > 1: + variants, owner = {}, {} + for ks, cfgs in finalists.items(): + for i, cfg in enumerate(cfgs): + name = f"{ks}#{i}" + owner[name] = ks + variants[name] = _launcher(a, b, c, cfg, m=m, n=n, k=k, + stride_am=stride_am, + stride_bn=stride_bn) + if cfg.SPLIT_K > 1: + c.zero_() + # Rounds a multiple of the variant count, so each variant occupies + # each slot the same number of times -- with 2 variants over 5 + # rounds one of them gets the post-warmup slot three times and the + # other twice, which is worth several percent here. + meas = interleaved(variants, warmup=3, iters=5, + rounds=2 * len(variants)) + out["h2h_best"] = {} + for name, m_ in meas.items(): + ks = owner[name] + out["h2h_best"][ks] = min(out["h2h_best"].get(ks, float("inf")), + m_.best) + if m_.median < out["headtohead"].get(ks, float("inf")): + out["headtohead"][ks] = m_.median + out["spread"][ks] = m_.spread + # See ``Measurement.stall_ratio``: how much of this cell's + # elapsed time was the GPU waiting for the host. Carried + # into the JSON so "was the node quiet" is answerable from + # the artifact instead of from a contemporaneous rocm-smi. + out["stall"][ks] = m_.stall_ratio + out["winners"][ks] = finalists[ks][int(name.split("#")[1])] + else: + out["headtohead"] = dict(out["sweep"]) + out["h2h_best"] = dict(out["sweep"]) + out["spread"] = {ks: 0.0 for ks in out["winners"]} + out["stall"] = {ks: float("nan") for ks in out["winners"]} + return out + finally: + del a, b, c + torch.cuda.empty_cache() + + +def run_shapes(problems: list[ConvProblem], modes=("compute", "dram"), + dtype=torch.bfloat16, max_configs: int = 0, + knobs: str = "legacy", only_n: int | None = None, + keep_all: bool = False, + prior: dict[tuple[str, str], dict] | None = None, + flush=None) -> list[dict]: + """Sweep every (problem, direction) cell. + + ``prior`` supplies cells a previous run already measured, keyed by + ``(label, direction)``; they are carried through untouched. ``flush``, if + given, is called with the row list after every cell. Both exist because a + full ``compare-wide`` pass over the corpus is ~90 minutes and the JSON used + to be written only at the end -- a run interrupted at cell 53 of 60 left + nothing behind but a log, twice. + """ + peak = PEAK_FLOPS["bf16"] + rows = [] + for p in problems: + for direction in DIRECTIONS: + m, n, k = p.gemm_shape(direction) + if only_n is not None and n != only_n: + continue + if prior and (p.label, direction) in prior: + rows.append(prior[(p.label, direction)]) + print(f" {p.label:34s} {direction:11s} [resumed]") + sys.stdout.flush() + if flush: + flush(rows) + continue + flops = 2 * m * n * k + row = { + "problem": p.label, "direction": direction, + "M": m, "N": n, "K": k, + "conv_flops": p.flops(direction), + "conv_ai": p.arithmetic_intensity(direction), + "conv_roofline_tflops": p.roofline_flops(direction) / 1e12, + } + best = p.measured_for(direction, config="A") or p.measured_for(direction) + if best: + row["miopen_ms"] = best[0]["ms_per_call"] + row["miopen_pct_roofline"] = best[0]["pct_roofline"] + row["miopen_solver"] = best[0]["solvers"][0] + knob_sets = _KNOB_SETS.get(knobs, (knobs,)) + for mode in modes: + try: + cell = _measure_cell(m, n, k, mode=mode, dtype=dtype, + knob_sets=knob_sets, keep_all=keep_all, + max_configs=max_configs) + except torch.OutOfMemoryError: + cell = None + if cell is None or not cell["winners"]: + # Shape not runnable in this mode. For ``dram`` that is the + # finding: materialized im2col does not fit in 128 GiB. + row[f"{mode}_ms"] = None + row[f"{mode}_skipped"] = "operands exceed budget" + continue + + def record(prefix: str, ms: float) -> None: + rate = flops / (ms * 1e-3) + row[f"{prefix}_ms"] = ms + row[f"{prefix}_tflops"] = rate / 1e12 + row[f"{prefix}_pct_peak"] = 100 * rate / peak + # What the convolution would take at this FLOP rate. + row[f"{prefix}_implied_conv_ms"] = p.flops(direction) / rate * 1e3 + row[f"{prefix}_implied_pct_roofline"] = ( + 100 * rate / p.roofline_flops(direction) + ) + + for ks in cell["winners"]: + if keep_all: + row[f"{mode}_{ks}_all_configs"] = cell["sinks"][ks] + if len(knob_sets) > 1: + record(f"{mode}_{ks}", cell["headtohead"][ks]) + row[f"{mode}_{ks}_config"] = str(cell["winners"][ks]) + row[f"{mode}_{ks}_configs_ran"] = cell["ran"][ks] + # Kept because it is the number the M0 run reported, and + # a large gap between the two is itself a finding about + # how long this device holds its clocks. + row[f"{mode}_{ks}_sweep_tflops"] = ( + flops / (cell["sweep"][ks] * 1e-3) / 1e12 + ) + # Min over rounds. This node is shared: a neighbouring + # job on another die drags whole rounds down by 10x and + # shows up as a spread in the hundreds of percent. The + # median is then a measure of the neighbour, and the + # minimum is the best available estimate of what the + # kernel can actually do. + row[f"{mode}_{ks}_best_tflops"] = ( + flops / (cell["h2h_best"][ks] * 1e-3) / 1e12 + ) + row[f"{mode}_{ks}_spread"] = cell["spread"][ks] + row[f"{mode}_{ks}_stall"] = cell["stall"].get(ks) + # The headline ``{mode}_*`` keys carry the *last* knob set, which + # is the tuned one in a comparison run and the only one in a + # single-set run. Keeps the JSON shape compatible with + # ``probe.json`` so before/after can be diffed key for key. + head = knob_sets[-1] + record(mode, (cell["headtohead"] if len(knob_sets) > 1 + else cell["sweep"])[head]) + row[f"{mode}_config"] = str(cell["winners"][head]) + row[f"{mode}_configs_ran"] = cell["ran"][head] + if len(knob_sets) > 1: + row[f"{mode}_headtohead"] = list(knob_sets) + row[f"{mode}_spread"] = cell["spread"][head] + row[f"{mode}_stall"] = cell["stall"].get(head) + rows.append(row) + print( + f" {p.label:34s} {direction:11s} " + f"M={m:<9d} N={n:<5d} K={k:<6d} " + + " ".join( + f"{mode}={row.get(f'{mode}_tflops', float('nan')):6.1f} TF/s" + f" ({row.get(f'{mode}_pct_peak', float('nan')):5.1f}% peak)" + for mode in modes + ) + + (f" miopen={row['miopen_ms']:.3f} ms" + f" ({row['miopen_pct_roofline']:.0f}%)" if "miopen_ms" in row else "") + ) + if len(knob_sets) > 1: + for mode in modes: + if row.get(f"{mode}_ms") is None: + continue + print(" " + " | ".join( + f"{ks}: {row[f'{mode}_{ks}_pct_peak']:5.1f}% " + f"(sweep {row[f'{mode}_{ks}_sweep_tflops']:.0f} TF/s, " + f"stall {row[f'{mode}_{ks}_stall']:.2f}x) " + f"{row[f'{mode}_{ks}_config']}" + for ks in knob_sets if f"{mode}_{ks}_pct_peak" in row + )) + sys.stdout.flush() + if flush: + flush(rows) + return rows + + +def run_isa(spec: str, *, m: int, n: int, k: int, device="cuda", + dtype=torch.bfloat16) -> None: + """Compile and launch exactly one config, so its ISA can be inspected. + + Run under ``AMDGCN_ENABLE_DUMP=1`` with a cold ``TRITON_CACHE_DIR`` -- a + cache hit skips the compile and therefore the dump, which is an easy way to + conclude "no MFMA" from an empty grep. Neither ``matrix_instr_nonkdim`` nor + ``kpack`` is validated at the Python level: an illegal value falls back to + FMA with only an MLIR remark, so this check is not optional before trusting + a number. + """ + parts = spec.split(",") + if len(parts) != 8: + raise SystemExit( + "--isa expects BM,BN,BK,GROUP_M,warps,nonkdim,kpack,int32 " + f"(got {spec!r})" + ) + bm, bn, bk, gm, warps, nonkdim, kpack, int32 = (int(x) for x in parts) + cfg = GemmConfig(bm, bn, bk, gm, 1, warps, num_stages=2, + matrix_instr_nonkdim=nonkdim, kpack=kpack, waves_per_eu=0, + int32_offsets=bool(int32)) + a = _randn((m, k), device, dtype) + b = _randn((k, n), device, dtype) + c = torch.empty((m, n), device=device, dtype=torch.float32) + _launch(a, b, c, cfg, m=m, n=n, k=k, + stride_am=a.stride(0), stride_bn=b.stride(1)) + torch.cuda.synchronize() + print(f"ISA-DUMP-CONFIG {cfg} M={m} N={n} K={k} " + f"a_storage={a.untyped_storage().size()} " + f"b_storage={b.untyped_storage().size()} " + f"c_storage={c.untyped_storage().size()}") + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--mode", choices=("peak", "shapes", "both", "isa"), + default="both") + ap.add_argument("--knobs", + choices=("legacy", "amd", "amd-wide", + "compare", "compare-wide"), + default="legacy", + help="config generator. 'legacy' reproduces the M0 sweep " + "(no AMD kernargs). 'amd' is Inductor's ROCm conv " + "seed grid under the gfx942 MFMA constraints. " + "'amd-wide' adds skinny-N tiles and sweeps " + "matrix_instr_nonkdim. 'compare[-wide]' runs several " + "sets and races their champions in one interleaved " + "measurement, which is the only before/after immune " + "to a change of tenancy on the device.") + ap.add_argument("--only-n", type=int, default=None, + help="restrict to cells whose GEMM N equals this") + ap.add_argument("--keep-all-configs", action="store_true", + help="record every config's time, not just the winner") + ap.add_argument("--isa", default="128,128,64,6,8,16,2,0", + help="BM,BN,BK,GROUP_M,warps,nonkdim,kpack,int32 for --mode isa") + ap.add_argument("--isa-mnk", default="4096,4096,4096", + help="M,N,K for --mode isa") + ap.add_argument("--modes", default="compute", + help="comma-separated subset of compute,dram. 'dram' is " + "informative but slow -- it materializes im2col, which " + "for the corpus's larger shapes means a 15-30 GiB " + "operand per config -- and it is not the ceiling a " + "fused kernel is measured against.") + ap.add_argument("--top", type=int, default=6, help="hottest N corpus problems") + ap.add_argument("--max-configs", type=int, default=0, + help="cap the config sweep (0 = no cap)") + ap.add_argument("--out", default=None) + ap.add_argument("--resume", action="store_true", + help="reuse cells (and the peak block) already present in " + "--out instead of re-measuring them") + args = ap.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("no GPU") + if args.mode == "isa": + m, n, k = (int(x) for x in args.isa_mnk.split(",")) + run_isa(args.isa, m=m, n=n, k=k) + return + props = torch.cuda.get_device_properties(0) + print(f"device: {props.name}, {props.multi_processor_count} CUs, " + f"torch {torch.__version__}, triton {triton.__version__}") + print(f"roofline constants: {PEAK_FLOPS['bf16']/1e12:.0f} TFLOP/s bf16, " + f"{HBM_BYTES_PER_S/1e12:.1f} TB/s HBM " + f"(crossover at {PEAK_FLOPS['bf16']/HBM_BYTES_PER_S:.0f} FLOP/byte)\n") + + result: dict = {"device": props.name, "torch": torch.__version__, + "triton": triton.__version__, "knobs": args.knobs} + prior: dict[tuple[str, str], dict] = {} + out_path = pathlib.Path(args.out) if args.out else None + if args.resume and out_path and out_path.exists(): + old = json.loads(out_path.read_text()) + prior = {(r["problem"], r["direction"]): r for r in old.get("shapes", [])} + if old.get("peak"): + result["peak"] = old["peak"] + print(f"resuming: {len(prior)} cells already measured" + + (", peak block reused" if "peak" in result else "")) + + def write(rows=None) -> None: + if not out_path: + return + payload = dict(result) + if rows is not None: + payload["shapes"] = rows + out_path.write_text(json.dumps(payload, indent=1) + "\n") + + print(f"knob set: {args.knobs}" + + (f", restricted to N={args.only_n}" if args.only_n else "")) + t0 = time.time() + if args.mode in ("peak", "both") and "peak" not in result: + print("== peak: square bf16 GEMM, Triton vs hipBLASLt ==") + if args.knobs in _KNOB_SETS: + result["peak"] = run_peak_compare(knob_sets=_KNOB_SETS[args.knobs]) + else: + result["peak"] = run_peak(knobs=args.knobs) + write() + print() + if args.mode in ("shapes", "both"): + print(f"== shapes: conv-implied GEMMs, top {args.top} corpus problems ==") + result["shapes"] = run_shapes(list(hot_corpus(args.top)), + modes=tuple(args.modes.split(",")), + max_configs=args.max_configs, + knobs=args.knobs, only_n=args.only_n, + keep_all=args.keep_all_configs, + prior=prior, flush=write) + result["elapsed_s"] = time.time() - t0 + print(f"\nelapsed {result['elapsed_s']:.0f} s") + + if out_path: + write(result.get("shapes")) + print(f"wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/triton_conv3d/bench/harness.py b/triton_conv3d/bench/harness.py new file mode 100644 index 0000000..0a94131 --- /dev/null +++ b/triton_conv3d/bench/harness.py @@ -0,0 +1,1068 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""Timing that survives this machine -- and says how well it survived it. + +**What this node does.** Measured on GPU 2 on 2026-08-03 +(``work/triton-conv/bin/drift_probe.py``: 47,686 consecutive 10-call blocks of +``conv 64->64 k3 @ 130^3`` over 14 minutes, with ``rocm-smi`` sampled +concurrently from a sibling process): + +* steady-state dispersion is **CoV 0.56%**, IQR/median 0.39%, p99/p01 = 1.022; +* the *sequential* protocol -- measure A for a while, then B -- simulated out of + that stream, where both windows are the same kernel and the true ratio is + therefore exactly 1, has a **worst error of 0.32%** over 206 replications at + 2 s separation, and 0.14 / 0.24 / 0.18% at 5 / 15 / 60 s; +* a deliberate **90-second idle gap** moves the measured time by **0.13%**, so + there is no thermal ramp to catch; +* ``sclk`` under load stays in 1396-1451 MHz and junction temperature rises + 49 -> 66 C, and neither correlates with the time (r = -0.10 and -0.05). The + device is power-capped at 550 W and clocks *stably* at that cap. + +So the clock is not what threatens a comparison here, and neither is elapsed +time between two measurements in one process. What does threaten one is a +**neighbour**: a foreign tenant on the device inflated three cells **2.6x**, +with per-cell spreads 20-100x their quiet-device values +(``work/triton-conv/BASELINE_NOTES.md``, "cause 3"). That is common-mode +within a round -- every arm pays it at once -- so :func:`interleaved` survives it +where a between-run comparison does not. It costs nothing, so it stays. It is +not, however, a substitute for an error bar. + +**What actually goes wrong here is the instrument, twice.** + +1. Below about 0.15 ms the per-iteration event pair is a material part of what + it reports. Measured: ``hipEventRecord`` costs **9.5 us of host time**, and + a block with an event between every iteration reports **16-26% more** than + the same kernel's wall-clock throughput (``convT 1024->512 @ 8^3`` forward: + 0.0718 ms with events, 0.0617 bracketed, 0.0568 by wall clock). The tax is + per-arm, not common-mode, so it does *not* cancel in a ratio. This harness + now widens the event interval to cover several calls whenever the kernel is + short enough for that to matter -- one common width for every arm of a + comparison, from a rule that depends only on the measured duration -- and + reports the residual it measures against an event-free bracket. Checked + against a wall-clock throughput reference at 512^2 / 1024^2 / 2048^2 / 4096^2 + bf16 GEMM: the harness now reads 0.96-1.08x of it across three orders of + magnitude, where at the smallest of those it used to read **1.52x**. +2. ``iters`` was never a neutral knob. The first call after a synchronize pays + a queue restart, so ``iters=1`` over-reports by 3% at 1.4 ms and by **42%** + at 0.07 ms. The adaptive path picks ``iters`` from the measured duration and + never leaves it at 1 for a small kernel. + +**What the timed region contains, and what it does not.** Below about 0.15 ms +the *host* is the pacer: a launch costs 2-37 us of Python and dispatch on this +node, against kernels of 28-68 us, so an event-timed loop of ``fn()`` measures +the launcher and the kernel together and cannot separate them. :func:`capture` +puts ``chunk`` back-to-back calls behind one CUDA graph, which contains the +device work and none of the host work, so replaying it measures the kernel +alone. Measured, per call, on ``convT 1024->512 @ 8^3`` +(``work/triton-conv/bin/launcher_symmetry.py --only census-small``): + +============================= ======= ======== ============= +arm eager kernel host launch +============================= ======= ======== ============= +``convT`` fwd, Triton 0.0421 0.0282 13.9 us +``convT`` fwd, MIOpen 0.0699 0.0683 **1.6 us** +``convT`` bwd-data, Triton 0.0539 0.0350 18.9 us +``convT`` bwd-data, MIOpen 0.0761 0.0395 **36.6 us** +``convT`` bwd-weight, Triton 0.0712 0.0542 17.0 us +``convT`` bwd-weight, MIOpen 0.0839 0.0504 33.4 us +============================= ======= ======== ============= + +Those six host costs differ by **23x**, so *leaving* the launcher in is as much +a per-arm instrument as taking it out asymmetrically would be: the MIOpen +control for a backward direction pays an autograd-engine walk the Triton arm +does not, and the Triton arm pays a tuned-table lookup MIOpen does not. At +these sizes the two do not cancel -- the same three cells read 1.659x / 1.414x / +1.176x with the launcher in and **2.420x / 1.130x / 0.931x** without it, i.e. +the launcher-inclusive number is 1.5x too *low* on the forward and 1.3x too +*high* on the weight gradient. Hence the rule this module enforces: a graph is +chosen for the whole comparison or for none of it, and ``chunk`` -- like +``group`` -- is a function of the shortest arm's duration alone, never of a +per-arm measurement. Above ~40 ms per call the host cost is under 0.2% of +either arm and no graph is used. + +**And every number now carries its precision.** ``rounds`` and ``iters`` are +chosen online from what has already been observed, stopping when the reported +statistic reaches a stated relative precision or when a wall-clock budget is +exhausted, and :class:`Measurement` says which of the two happened. A cell that +stopped on the budget with a wide interval is visibly different from one that +converged. That distinction is worth more than the time saved: four of this +project's five retracted results were numbers quoted without one. + +Everything here is in-process. Sub-process benchmarking adds interpreter start, +allocator state and MIOpen database warmth as confounders, none of which the +kernel controls. +""" + +from __future__ import annotations + +import dataclasses +import math +import statistics +import time +import warnings +from typing import Callable, Iterable, Mapping, Sequence + +import torch + +#: Big enough to evict MI300A's 256 MiB infinity cache. +_FLUSH_BYTES = 512 * 1024 * 1024 +_flush_buffer: torch.Tensor | None = None + +#: Fraction of a kernel's own time the event instrument is allowed to add before +#: :func:`interleaved` starts grouping iterations behind one event interval. +#: 2% because the smallest published per-cell differences are around 3%. +_TAX_BUDGET = 0.02 + +#: Wall time one timed block should aim for. Large enough that the queue-restart +#: transient on the first iteration is a small fraction of the block, small +#: enough that a round is cheap. +_BLOCK_TARGET_MS = 15.0 +_MAX_ITERS = 512 + +#: Student-t 97.5th percentile by degrees of freedom, so an interval can be +#: quoted without a scipy dependency. Index 0 is unused. +_T975 = ( + math.nan, 12.706, 4.303, 3.182, 2.776, 2.571, 2.447, 2.365, 2.306, 2.262, + 2.228, 2.201, 2.179, 2.160, 2.145, 2.131, 2.120, 2.110, 2.101, 2.093, + 2.086, 2.080, 2.074, 2.069, 2.064, 2.060, 2.056, 2.052, 2.048, 2.045, +) + +#: Asymptotic ratio of the standard error of a sample median to that of the +#: sample mean, for normally distributed data: ``sqrt(pi/2)``. Round values are +#: already medians over ``iters`` calls, so they are close to normal; where they +#: are heavier-tailed this factor is conservative (the median's true SE is then +#: smaller than the formula says), which is the direction to err in. +_MEDIAN_SE_FACTOR = 1.2533141373155003 + + +def _t975(n: int) -> float: + if n < 2: + return math.inf + return _T975[n - 1] if n - 1 < len(_T975) else 1.96 + + +def _half_width(values: Sequence[float]) -> float: + """95% half-width for the *median* of ``values``, in the same units. + + Closed form rather than a bootstrap: it is deterministic (this project + re-runs its numbers and compares them), it needs no RNG seed to be + reproducible, and at the round counts in use (4-16) a percentile bootstrap + of a median cannot produce an interval wider than the observed range, which + understates exactly when it matters most. + """ + n = len(values) + if n < 2: + return math.inf + sd = statistics.stdev(values) + return _t975(n) * _MEDIAN_SE_FACTOR * sd / math.sqrt(n) + + +def flush_caches(device: torch.device | str = "cuda") -> None: + """Evict the cache hierarchy so a measurement starts cold. + + Matters for the memory-bound directions: a 16 MiB working set measured hot + reports bandwidth the same kernel will never see inside a real step, where + everything upstream has already flushed it. + + Two things to know before trusting it. It **works** -- the first iteration + after a flush is 1.50-1.54x the hot time at ``2048^2`` bf16 GEMM and at the + transposed sites -- but a caller that flushes once per block and then reports + the *median* over ``iters`` calls throws that one cold sample away: measured + ``median_moved_by_flush`` is 0.99-1.01 at every real workload. Use + ``iters=1`` (what the adaptive path does when ``flush=True``) or read + :attr:`Measurement.cold`, which this module records for exactly this reason. + """ + global _flush_buffer + want = torch.device(device) + if want.index is None and want.type == "cuda": + # ``torch.device("cuda")`` carries no index but a tensor created on it + # does, so a naive ``!=`` is always true and this function used to + # reallocate 512 MiB on every call -- one extra 512 MiB block live at a + # time, and an allocation on the critical path of every timed round. + want = torch.device("cuda", torch.cuda.current_device()) + if _flush_buffer is None or _flush_buffer.device != want: + _flush_buffer = torch.empty(_FLUSH_BYTES, dtype=torch.uint8, device=want) + _flush_buffer.zero_() + + +@dataclasses.dataclass(frozen=True) +class Measurement: + """Per-round times for one variant, in milliseconds per call. + + The headline statistic is still :attr:`median`. What is new is that it + comes with :attr:`half_width` -- a 95% interval -- and with :attr:`stop`, + which says whether the measurement reached its precision target or ran out + of wall clock. Print it; do not quote the median alone. + """ + + name: str + rounds: tuple[float, ...] + #: Per round, ``block mean / per-iteration median``. See :func:`_time_block`: + #: greater than 1 means the host failed to keep the queue full and the GPU + #: idled between launches. Reported rather than hidden, because the previous + #: version of this harness folded that idle time into the kernel time and + #: produced spreads of up to 2363% that were then misdiagnosed twice. + #: + #: It is a *skew* statistic and it is blind to the uniform case: if every + #: iteration is inflated by the same launch gap it reads exactly 1.00. Use + #: :attr:`tax_frac`, which is measured against an event-free bracket and is + #: therefore independent, for that question. + stalls: tuple[float, ...] = () + #: Calls per timed block, and calls per event interval within it. + iters: int = 0 + group: int = 1 + #: Per round, the *first* iteration of the block. With ``flush=True`` that + #: is the only cold sample there is. + firsts: tuple[float, ...] = () + #: ``per-iteration-event time - event-free bracket time``, ms per call, + #: measured for this variant during calibration. The instrument's own cost. + tax_ms: float = 0.0 + #: Why the measurement stopped: ``fixed`` (caller pinned ``rounds``), + #: ``converged``, ``budget`` or ``max_rounds``. + stop: str = "fixed" + #: True when every variant occupied every position, and every ordered + #: adjacency occurred, equally often -- i.e. ``rounds`` was a multiple of + #: ``2 * len(variants)``. + balanced: bool = True + seconds: float = 0.0 + + @property + def median(self) -> float: + return statistics.median(self.rounds) + + @property + def best(self) -> float: + return min(self.rounds) + + @property + def cold(self) -> float: + """Median first-iteration time; the cold number when ``flush`` is on.""" + return statistics.median(self.firsts) if self.firsts else self.median + + @property + def spread(self) -> float: + """Relative range across rounds. **Grows with** ``rounds`` by construction. + + Kept because every recorded result JSON has it, but it is not a measure + of how much the machine moved: the expected range of ``n`` samples grows + like ``d2(n)`` even on a perfectly stationary device. Measured on this + node with one kernel held constant for 14 minutes, the median of this + statistic runs 0.23% at 2 rounds, 0.63% at 6, 0.98% at 20 and 2.70% at + 100 -- all of it arithmetic. Compare :attr:`rel_half_width` instead, + which is an interval and does not have that defect. + """ + return (max(self.rounds) - min(self.rounds)) / self.median if self.rounds else 0.0 + + @property + def cov(self) -> float: + """Coefficient of variation across rounds; comparable between runs.""" + if len(self.rounds) < 2: + return 0.0 + return statistics.stdev(self.rounds) / statistics.fmean(self.rounds) + + @property + def half_width(self) -> float: + """95% half-width on :attr:`median`, in ms.""" + return _half_width(self.rounds) + + @property + def rel_half_width(self) -> float: + m = self.median + return self.half_width / m if m > 0 else math.inf + + @property + def converged(self) -> bool: + return self.stop in ("converged", "fixed") + + @property + def tax_frac(self) -> float: + """Instrument cost as a fraction of the reported time. + + Independent of the numbers the median came from -- it is the gap between + an event-per-iteration block and an event-free bracket of the same + kernel -- which is precisely what :attr:`stall_ratio` cannot see. + """ + m = self.median + return self.tax_ms / m if m > 0 else 0.0 + + @property + def stall_ratio(self) -> float: + """Worst launch-gap inflation seen in any round; 1.0 is a clean queue.""" + return max(self.stalls) if self.stalls else 1.0 + + def __str__(self) -> str: + stall = f", stall {self.stall_ratio:.2f}x" if self.stall_ratio > 1.05 else "" + tax = f", instrument {self.tax_frac:+.1%}" if abs(self.tax_frac) > 0.02 else "" + mark = "" if self.converged else f" [{self.stop}]" + return ( + f"{self.name}: {self.median:.4f} +-{self.rel_half_width:.1%} ms " + f"({len(self.rounds)}x{self.iters}{mark}, best {self.best:.4f}" + f"{stall}{tax})" + ) + + +@dataclasses.dataclass(frozen=True) +class Ratio: + """A paired ratio of two variants, with the interval that makes it a claim. + + Paired **per round**, not median-over-median: the two arms of a round were + measured seconds apart under the same device state, so a common-mode + excursion divides out of every pair before anything is averaged. That is + the property :func:`interleaved` exists to buy, and taking a ratio of two + independently-reduced medians throws it away. + """ + + numerator: str + denominator: str + point: float + lo: float + hi: float + n: int + + @property + def rel_half_width(self) -> float: + return (self.hi - self.lo) / (2 * self.point) if self.point > 0 else math.inf + + @property + def significant(self) -> bool: + """Does the interval exclude 1.0? If not, there is no measured win.""" + return self.lo > 1.0 or self.hi < 1.0 + + def __str__(self) -> str: + star = "" if self.significant else " (consistent with no difference)" + return (f"{self.numerator}/{self.denominator} = {self.point:.3f}x " + f"[{self.lo:.3f}, {self.hi:.3f}], n={self.n}{star}") + + +def ratio(numerator: Measurement, denominator: Measurement) -> Ratio: + """Paired ratio ``numerator / denominator`` with a 95% interval.""" + n = min(len(numerator.rounds), len(denominator.rounds)) + pairs = [numerator.rounds[i] / denominator.rounds[i] for i in range(n) + if denominator.rounds[i] > 0] + if not pairs: + return Ratio(numerator.name, denominator.name, math.nan, math.nan, + math.nan, 0) + logs = [math.log(p) for p in pairs] + point = math.exp(statistics.median(logs)) + hw = _half_width(logs) + if not math.isfinite(hw): + return Ratio(numerator.name, denominator.name, point, 0.0, math.inf, + len(pairs)) + return Ratio(numerator.name, denominator.name, point, + point * math.exp(-hw), point * math.exp(hw), len(pairs)) + + +def _time_block(fn: Callable[[], object], iters: int, + group: int = 1) -> tuple[float, float]: + """Return ``(median ms per call, stall ratio)`` for ``iters`` calls. + + Events rather than the wall clock: the launch is asynchronous, so a wall + clock measures the host's ability to enqueue until something forces a + synchronize, and the forced synchronize is then part of the measurement. + + But bracketing the *whole block* with two events, as this used to, has the + same disease one level up: if the host cannot keep the queue full the GPU + goes idle between launches, and that idle time is silently attributed to the + kernel. On a contended node it produced per-round spreads of 250-2363% that + were diagnosed as host jitter, then as a rogue tenant, before turning out to + be a duplicate driver process of our own. + + So time each iteration separately and return the **median**, which rejects a + stalled launch instead of averaging it in, alongside the ratio of the old + block-mean to that median. A ratio near 1 means the queue stayed full and + the two agree; a large ratio means the measurement is launch-bound and the + number should not be read as kernel time. + + ``group`` widens the event interval to ``group`` calls. An event costs + ~9.5 us of host time and shows up in the reported number below ~0.15 ms of + kernel, where it inflated ``convT`` sites by 16-26%; grouping divides that + by ``group`` while keeping enough samples per block for the median to still + reject a stall. ``group=1`` is the historical behaviour and is what a large + kernel gets, because there the tax is already under a tenth of a percent. + """ + marks = _blocked_events(fn, iters, group) + n = len(marks) - 1 + per_iter = [marks[i].elapsed_time(marks[i + 1]) / group for i in range(n)] + median = statistics.median(per_iter) + block_mean = marks[0].elapsed_time(marks[n]) / (n * group) + return median, (block_mean / median if median > 0 else 1.0) + + +def _blocked_events(fn: Callable[[], object], iters: int, group: int) -> list: + groups = max(1, iters // group) + marks = [torch.cuda.Event(enable_timing=True) for _ in range(groups + 1)] + torch.cuda.synchronize() + for g in range(groups): + marks[g].record() + for _ in range(group): + fn() + marks[groups].record() + torch.cuda.synchronize() + return marks + + +def _time_block_full(fn: Callable[[], object], iters: int, + group: int) -> tuple[float, float, float]: + """``(median, stall ratio, first sample)`` -- the first is the cold one.""" + marks = _blocked_events(fn, iters, group) + n = len(marks) - 1 + per = [marks[i].elapsed_time(marks[i + 1]) / group for i in range(n)] + median = statistics.median(per) + block_mean = marks[0].elapsed_time(marks[n]) / (n * group) + return median, (block_mean / median if median > 0 else 1.0), per[0] + + +def _bracket(fn: Callable[[], object], iters: int) -> float: + """Per-call ms with two events around the whole block and none inside. + + The event-free control the instrument tax is measured against. + """ + a = torch.cuda.Event(enable_timing=True) + b = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize() + a.record() + for _ in range(iters): + fn() + b.record() + torch.cuda.synchronize() + return a.elapsed_time(b) / iters + + +@dataclasses.dataclass(frozen=True) +class Plan: + """What calibration decided for one variant, so it can be reported.""" + + per_call_ms: float + iters: int + group: int + warmup: int + tax_ms: float + + +#: Above this per-call time the event instrument is provably under 1%, so the +#: tax probe -- ten extra blocks -- is not worth running. It would cost 450 s +#: on the 45 s cliff cell alone. +_TAX_PROBE_MAX_MS = 1.0 + +#: Floor on the cost of one event interval, in ms. Measured on this device with +#: ``torch.cuda._sleep`` at four durations spanning 10 us to 285 us per call: +#: **2.85, 2.85, 2.85 and 3.09 us**, i.e. flat in the kernel size, which is what +#: makes it a property of the instrument rather than of the workload. A real +#: kernel can cost *more* than this, because the host also pays ~9.5 us per +#: ``record()`` and a launch path expensive enough to make the host the pacer +#: turns that into device idle -- so this is a floor, and the per-arm probe +#: raises it. It exists so that a probe which happens to measure near zero +#: cannot leave a tiny kernel ungrouped. +_EVENT_INTERVAL_MS = 0.00285 + +#: Group only once the instrument is worth more than twice its budget, i.e. more +#: than 4% of the kernel by default. The band below that is left alone because +#: grouping is not free: it averages ``group`` calls behind one event interval, +#: so it also *reduces* the median's ability to reject a stalled launch. At +#: 0.08 ms, forcing a group of 2 moved the reported time 7% the wrong way. +_GROUP_TRIGGER = 2.0 + + +def _measure_tax(fn: Callable[[], object], iters: int, group: int = 1, + reps: int = 5) -> float: + """Per-call cost of the event instrument: events minus an event-free bracket. + + Measured as a *paired* difference -- ev, br, ev, br, ... -- rather than as a + difference of two separately-collected medians. The two arms of each pair + are adjacent in time, so a device-wide excursion cancels inside the pair + instead of landing in the estimate; taken unpaired, this estimate was noisy + enough to pick wildly different groups for byte-identical work. + """ + diffs = [] + for _ in range(reps): + ev = _time_block(fn, iters, group)[0] + br = _bracket(fn, iters) + diffs.append(ev - br) + return statistics.median(diffs) + + +def _probe(fn: Callable[[], object], *, pinned_warmup: int | None, + warmup_s: float, warmup_min: int, warmup_max: int, + warmup_hard_s: float, block_ms: float, max_iters: int, + need_duration: bool) -> tuple[float, int]: + """Warm one variant and return ``(per-call ms, warmup calls issued)``. + + The first call absorbs whatever one-off the variant has -- MIOpen's find is + **8.2 s** on ``conv 64->64 k3 @ 130^3`` against a 1.65 ms steady state, and + Triton's JIT is a compile -- so it is never the call that decides anything. + + A pinned ``warmup`` issues exactly that many calls and nothing else, so a + caller that also pins ``iters`` gets precisely the pre-adaptive call + sequence and a re-capture stays comparable with what is on disk. + """ + if pinned_warmup is not None: + for _ in range(pinned_warmup): + fn() + torch.cuda.synchronize() + if not need_duration: + return 0.0, pinned_warmup + probe = max(4, min(max_iters, 8)) + return max(_bracket(fn, probe), 1e-6), pinned_warmup + + fn() + torch.cuda.synchronize() + t0 = time.perf_counter() + fn() + torch.cuda.synchronize() + rough = max((time.perf_counter() - t0) * 1e3, 1e-4) + + # Warm to a settled state. Measured: Triton needs ~5 calls to come within + # 0.5% of steady, MIOpen ~4 after its find; both are cheap when the kernel + # is small and unaffordable when it is 45 s, which is what the hard cap is. + n = min(warmup_max, max(warmup_min, int(warmup_s * 1e3 / rough))) + if n * rough > warmup_hard_s * 1e3: + n = max(0, int(warmup_hard_s * 1e3 / rough)) + for _ in range(n): + fn() + torch.cuda.synchronize() + + # ``rough`` is one call bracketed by two synchronizes and it over-states a + # small kernel badly -- 0.37 ms for a 0.017 ms GEMM, because the sync and + # the queue restart are most of it. Sizing ``iters`` off that number gave + # blocks 20x too short and left the instrument tax at 6.9% instead of 2%. + # So re-estimate from an event-free bracket, which is the same quantity the + # tax is measured against. + d = rough + if rough < block_ms: + probe = max(4, min(max_iters, round(block_ms / rough))) + d = max(_bracket(fn, probe), 1e-6) + return d, n + 2 + + +def per_call_ms(fn: Callable[[], object], *, warmup_s: float = 0.05, + warmup_min: int = 5, warmup_max: int = 200, + warmup_hard_s: float = 2.0, + block_ms: float = _BLOCK_TARGET_MS, + max_iters: int = _MAX_ITERS, + settle_s: float = 0.25, settle_calls: int = 5) -> float: + """One warmed, event-free estimate of a callable's per-call time. + + The same probe :func:`interleaved` runs internally, exposed because a caller + that is about to decide *how* to time something -- whether to put it in a + graph, and how many calls to put in one -- has to know roughly what it costs + first, and that decision has to be made from a quantity measured the same + way for every arm. + + With one difference, and it cost an afternoon: :func:`_probe` absorbs *one* + one-off call before it estimates anything, and one is not always enough. + The first ``256^2`` bf16 ``matmul`` in a fresh process measured **653 ms** on + its *second* call, because rocBLAS/hipBLASLt loads its kernel library + lazily and does it after the first launch. A decision made from that number + puts a 0.02 ms kernel on the eager path with ``iters=1``. So this settles + first, bounded by time rather than by a call count so that a 45-second + kernel is called once and a 20-microsecond one five times. + + It is a probe, not a measurement: one bracketed block, no rounds, no + interval. Do not publish it. + """ + t0 = time.perf_counter() + for _ in range(settle_calls): + fn() + torch.cuda.synchronize() + if time.perf_counter() - t0 > settle_s: + break + return _probe(fn, pinned_warmup=None, warmup_s=warmup_s, + warmup_min=warmup_min, warmup_max=warmup_max, + warmup_hard_s=warmup_hard_s, block_ms=block_ms, + max_iters=max_iters, need_duration=True)[0] + + +def _common_group(durations: Sequence[float], *, tax_budget: float, + max_iters: int) -> int: + """One event-interval width for **every** arm of a comparison. + + Two things are load-bearing here and both were bought with a wrong answer. + + *The group is a function of the duration alone*, not of a per-arm + measurement of the instrument tax. A group derived from a per-arm probe + made two byte-identical arms pick 12 and 2 in the same call -- a 4% + difference in residual instrument cost, and therefore a 4% bias in a ratio + whose true value was exactly 1.000. That is the failure this whole + exercise is about, reintroduced by the fix for it; + ``test_a_paired_ratio_of_two_identical_arms_covers_one`` caught it. + + *And it is common to the whole call*, taken from the **shortest** arm, so + that even when two arms are far apart in duration -- or merely far enough + apart to straddle a power-of-two boundary, which byte-identical arms did -- + they are measured with the same ruler. ``iters`` stays per-arm, because + that is what the five orders of magnitude in this corpus need; the *width + of the event interval* is what has to match. + """ + d = min(durations) + need = _EVENT_INTERVAL_MS / (tax_budget * d) + if need <= _GROUP_TRIGGER: + return 1 + return max(1, min(max_iters // 4, 1 << math.ceil(math.log2(need)))) + + +def _size_block(d: float, group: int, block_ms: float, max_iters: int) -> int: + """Calls per timed block: about ``block_ms`` of work, at least 4 samples.""" + if group == 1: + return max(1, min(max_iters, round(block_ms / d))) + samples = max(4, min(max_iters // group, round(block_ms / (d * group)))) + return group * samples + + +def _williams(n: int) -> list[int]: + """``0, 1, n-1, 2, n-2, ...`` -- the first row of a Williams square. + + Its successive differences are ``1, -2, 3, -4, ...`` mod ``n``, which are + all distinct, and that is what makes the rotations of this row + *row-complete* for even ``n``: every ordered pair of variants occurs + adjacent equally often instead of only the cyclically adjacent ones. + """ + out, lo, hi = [], 0, n + while lo < hi: + out.append(lo) + lo += 1 + if lo < hi: + hi -= 1 + out.append(hi) + return out + + +def _order(names: list[str], r: int) -> list[str]: + """A position- and adjacency-balanced order for round ``r``. + + The old rule rotated by one position per round. That balances *positions* + but it preserves *adjacency*: with three or more variants, B always ran + immediately after A, so whatever A left in the caches was a constant charged + to B and averaged out of nothing. Measured on the adversarial case -- a + 1 GiB cache-polluting arm plus two arms doing byte-identical work, 40 + replications -- cyclic rotation reported the two identical arms **2.8% + apart** (sd 4.3%), while this rule reported them 0.2% apart and a uniformly + random order 0.7%. 2.8% is larger than several of the per-cell differences + this project publishes. + + Rotating a Williams row every *second* round and reversing it on odd rounds + gives, over ``2 * len(names)`` rounds, every variant in every position + exactly twice and every ordered adjacency equally often -- and for an even + number of variants (``conv_bench`` runs four) *every* ordered pair occurs, + not just the cyclic ones. Deterministic, so a capture is reproducible. + """ + n = len(names) + base = _williams(n) + k = (r // 2) % n + seq = [(i + k) % n for i in base] + if r % 2: + seq = seq[::-1] + return [names[i] for i in seq] + + +# --------------------------------------------------------------------------- +# Taking the launcher out of the timed region +# --------------------------------------------------------------------------- + +#: Cost of one ``cudaGraphLaunch``, in ms per replay, measured on this device by +#: fitting ``per_call(chunk) = kernel + cost / chunk`` to graphs of 1, 2, 4, 8, +#: 16 and 32 calls at ``convT 1024->512 @ 8^3`` +#: (``work/triton-conv/bin/graph_probe2.py``): +#: +#: =========================== ============ +#: arm fitted cost +#: =========================== ============ +#: ``convT`` fwd, Triton 3.9 us +#: ``convT`` fwd, MIOpen **12.8 us** +#: ``convT`` bwd-data, Triton 4.4 us +#: ``convT`` bwd-data, MIOpen 4.1 us +#: =========================== ============ +#: +#: The constant below is the **worst** of those, not their mean, because the +#: quantity that has to be bounded is the residual on whichever arm pays most -- +#: and because a per-arm estimate is exactly the mistake :func:`_common_group` +#: exists to document (H10: two byte-identical arms picked different instruments +#: and read 4% apart). It is a property of the launcher, not of the workload. +_REPLAY_COST_MS = 0.0128 + +#: Fraction of the *shortest* arm's per-call time the residual replay cost is +#: allowed to reach. 1% because the smallest published per-cell differences are +#: around 3% and the target precision is 2%. +_REPLAY_BUDGET = 0.01 + +#: A graph holding this many calls is already 128 x 0.0128 us = 0.1 us per call, +#: below the event instrument's own floor; more would only cost capture time. +_MAX_CHUNK = 128 + +#: Above this per-call time no graph is used: the largest host launch cost +#: measured on this node is 0.08 ms (the autograd engine's, on the MIOpen +#: backward control), which at 40 ms per call is 0.2% of either arm -- smaller +#: than the harness's own target precision, so the exclusion is not worth the +#: capture. Below it the host cost reaches 190% of the kernel and decides the +#: answer. +_GRAPH_MAX_MS = 40.0 + +_capture_stream: torch.cuda.Stream | None = None + + +class CaptureError(RuntimeError): + """A callable could not be put in a CUDA graph, or the graph came out empty. + + Raised rather than swallowed: a caller that silently fell back to eager for + *one* arm would be comparing a launcher-exclusive number against a + launcher-inclusive one, which at these sizes is worth up to 1.4x. The + decision to fall back belongs to the comparison, not to an arm. + """ + + +def capture_stream() -> torch.cuda.Stream: + """The one side stream every capture in this process uses. + + It has to be shared, and callers have to build their autograd graphs on it. + ``torch.autograd.grad`` refuses to be captured with *"autograd node + ``ConvolutionBackward0`` has a stale reference to the default stream"* when + the forward that created the node ran on the default stream while the + capture runs on another -- so the MIOpen control for a backward direction, + which is a real forward plus :func:`torch.autograd.grad`, has to have its + forward built here. Handing every caller the same stream is what makes that + possible without each of them inventing one. + """ + global _capture_stream + if _capture_stream is None: + _capture_stream = torch.cuda.Stream() + return _capture_stream + + +class on_capture_stream: # noqa: N801 - a context manager, spelled like one + """Run every timed call of one comparison on :func:`capture_stream`. + + Not a nicety. The MIOpen control for a backward direction is + :func:`torch.autograd.grad` over a forward graph that had to be built on the + capture stream (see :func:`capture_stream`), and the autograd engine + synchronizes when the node's recorded stream is not the caller's current + one. Measured, ``convT`` backward at ``1024->512 @ 8^3``: the MIOpen arm + read **0.148 ms** when timed from the default stream and **0.113 ms** from + the stream its graph was built on -- a **35 us** cross-stream tax that the + Triton arm, which has no autograd graph, does not pay. At those sizes that + alone is worth 1.3x, and it is per-arm. + + So the stream is a property of the *comparison*, exactly like ``group`` and + ``chunk``: one stream, entered once, for every arm and both launcher + policies. Above ~0.3 ms the tax is under 1% and the two agree, which is why + it was invisible until the transposed cells. + """ + + def __enter__(self): + self._stream = capture_stream() + self._ctx = torch.cuda.stream(self._stream) + self._stream.wait_stream(torch.cuda.current_stream()) + self._ctx.__enter__() + return self._stream + + def __exit__(self, *exc): + self._ctx.__exit__(*exc) + torch.cuda.current_stream().wait_stream(self._stream) + torch.cuda.synchronize() + return False + + +@dataclasses.dataclass(frozen=True) +class Captured: + """``chunk`` back-to-back calls of one callable, behind one graph replay. + + Calling this runs ``chunk`` calls' worth of device work and *no* host work + beyond one ``cudaGraphLaunch``, so a harness that times it is timing the + kernel. :attr:`chunk` is the divisor a caller needs to get back to ms per + call -- deliberately not hidden, because every relative quantity the harness + reports (``rel_half_width``, ``cov``, :func:`ratio`) is scale-invariant and + only the absolute times need it. + """ + + graph: object + chunk: int + #: Per-call ms of the eager callable, as measured before capture. Kept so a + #: caller can report what the launcher was worth. + eager_ms: float = 0.0 + + def __call__(self) -> None: + self.graph.replay() + + +def capture(fn: Callable[[], object], chunk: int = 1, *, + warmup: int = 3) -> Captured: + """Put ``chunk`` calls of ``fn`` in a CUDA graph, or raise :class:`CaptureError`. + + Everything -- the warmup and the capture -- runs on :func:`capture_stream`, + for the autograd reason given there. + + The warmup is not optional and it is not only PyTorch's lazy-init + requirement: with ``cudnn.benchmark`` on, MIOpen's find is **8.2 s** and + happens on the first call, and a find inside a capture would synchronize and + abort it. Three calls are enough (measured: MIOpen settles by call ~4, + Triton by ~5, and the JIT compile is on call 1). + + An *empty* graph is treated as a failure. PyTorch only warns -- "The CUDA + Graph is empty. This usually means that the graph was attempted to be + captured on wrong device or stream" -- and a caller that ignored the warning + would publish the cost of ``cudaGraphLaunch`` as a kernel time, which is the + fastest wrong answer available. + """ + if chunk < 1: + raise ValueError(f"chunk must be >= 1, got {chunk}") + s = capture_stream() + try: + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(warmup): + fn() + torch.cuda.current_stream().wait_stream(s) + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + with torch.cuda.graph(g, stream=s): + for _ in range(chunk): + fn() + torch.cuda.synchronize() + except CaptureError: + raise + except Exception as exc: # noqa: BLE001 - any capture refusal + torch.cuda.synchronize() + raise CaptureError(f"{type(exc).__name__}: {exc}") from exc + for w in caught: + if "graph is empty" in str(w.message).lower(): + raise CaptureError("the captured graph is empty") + warnings.warn_explicit(w.message, w.category, w.filename, w.lineno) + return Captured(g, chunk) + + +def common_chunk(durations: Sequence[float], *, cost_ms: float = _REPLAY_COST_MS, + budget: float = _REPLAY_BUDGET, + max_chunk: int = _MAX_CHUNK) -> int: + """Calls per graph, for **every** arm of a comparison. + + Same shape and same reasoning as :func:`_common_group`, one level up: a + function of the shortest arm's measured duration alone, so that two arms are + never measured with two different rulers. A graph replay costs 3.9-12.8 us + of device time whatever is inside it, which is 45% of a 0.028 ms kernel at + ``chunk = 1``; dividing it by ``chunk`` brings it under ``budget`` of the + arm it hurts most. + """ + d = min(durations) + if d <= 0: + return 1 + need = cost_ms / (budget * d) + if need <= 1.0: + return 1 + return max(1, min(max_chunk, 1 << math.ceil(math.log2(need)))) + + +def graph_is_worthwhile(durations: Sequence[float], + max_ms: float = _GRAPH_MAX_MS) -> bool: + """Is the host launch cost big enough to be worth capturing away? + + Above :data:`_GRAPH_MAX_MS` the measured worst-case host launch cost + (0.08 ms, the autograd engine's) is under 0.2% of either arm, so eager + timing is already launcher-exclusive to within a fifth of the target + precision and the capture would only cost wall clock. + """ + return min(durations) <= max_ms + + +def time_callable( + fn: Callable[[], object], + *, + warmup: int | None = None, + iters: int | None = None, + rounds: int | None = None, + flush: bool = False, + **kwargs, +) -> Measurement: + """Time a single callable. Prefer :func:`interleaved` for comparisons.""" + return interleaved({"fn": fn}, warmup=warmup, iters=iters, rounds=rounds, + flush=flush, **kwargs)["fn"] + + +def interleaved( + variants: Mapping[str, Callable[[], object]], + *, + warmup: int | None = None, + iters: int | None = None, + rounds: int | None = None, + flush: bool = False, + target_rel: float = 0.02, + budget_s: float = 20.0, + hard_budget_s: float = 300.0, + min_rounds: int = 4, + floor_rounds: int = 3, + max_rounds: int = 64, + block_ms: float = _BLOCK_TARGET_MS, + max_iters: int = _MAX_ITERS, + tax_budget: float = _TAX_BUDGET, + warmup_s: float = 0.05, + warmup_min: int = 5, + warmup_max: int = 200, + warmup_hard_s: float = 2.0, + measure_tax: bool = True, +) -> dict[str, Measurement]: + """Time several variants against each other, and say how precisely. + + Each round runs every variant once, in a position- *and* adjacency-balanced + order (see :func:`_order`), and the arms are compared round by round so that + a neighbour process -- which inflates every arm of a round at once -- divides + out of the comparison instead of deciding it. + + **Online sizing.** With ``iters`` and ``rounds`` left at ``None`` this + picks both from what it has already measured, per variant: + + * ``iters`` so a block lasts about ``block_ms``. The corpus spans five + orders of magnitude -- 0.06 ms at the transposed sites, 45,241 ms at the + 2 GiB cliff -- and a fixed ``iters=10, rounds=6`` means 60 calls either + way: microseconds for one cell and **45 minutes** for another. + * ``group``, the number of calls behind one event interval, so the + instrument costs under ``tax_budget`` of the *shortest* arm -- **one width + for every arm**, never per-arm; see :func:`_common_group` for the ratio + that got biased by 4% when it was per-arm. + * ``rounds``, growing until the 95% half-width on every reported quantity -- + each variant's median and each variant's paired ratio against the first -- + is within ``target_rel``, or until ``budget_s`` of wall clock is spent, or + ``max_rounds``. It only ever stops on a **balanced block boundary** + (a multiple of ``2 * len(variants)`` rounds) so that stopping early cannot + reintroduce the position bias the ordering exists to remove -- except + against ``hard_budget_s``, which is checked every round because a cell + whose single round costs minutes must be able to stop without first + completing a design it cannot afford. ``Measurement.balanced`` then says + the design did not close. + + **Arms are never stopped individually.** If one arm's interval tightens + first, it keeps running: dropping it would leave the other arm measured + against a different stretch of wall clock, which is precisely the sequential + comparison this function exists to avoid. Stopping is a property of the + block, not of an arm. + + Passing ``iters`` and/or ``rounds`` as integers pins them exactly, which is + what every existing caller does and what a deliberately reproducible capture + should keep doing. ``warmup`` likewise. + """ + names = list(variants) + if not names: + return {} + + # Warm and size in two passes, because the second decision belongs to the + # comparison rather than to any one arm: every variant is probed first, and + # only then is one common event-interval width chosen for all of them. + need_duration = iters is None and not flush + probes = { + name: _probe( + variants[name], pinned_warmup=warmup, warmup_s=warmup_s, + warmup_min=warmup_min, warmup_max=warmup_max, + warmup_hard_s=warmup_hard_s, block_ms=block_ms, + max_iters=max_iters, need_duration=need_duration, + ) + for name in names + } + if flush: + # One flush per block reaches only the block's first sample, so a cold + # measurement has to have exactly one sample per block. + group = 1 + sizes = {name: iters or 1 for name in names} + elif iters is not None: + group = 1 + sizes = {name: iters for name in names} + else: + group = _common_group([probes[n][0] for n in names], + tax_budget=tax_budget, max_iters=max_iters) + sizes = {name: _size_block(probes[name][0], group, block_ms, max_iters) + for name in names} + plans: dict[str, Plan] = {} + for name in names: + d, warmed = probes[name] + tax = 0.0 + if measure_tax and sizes[name] >= 4 and 0 < d <= _TAX_PROBE_MAX_MS: + tax = _measure_tax(variants[name], sizes[name], group) + plans[name] = Plan(d, sizes[name], group, warmed, tax) + torch.cuda.synchronize() + + times: dict[str, list[float]] = {name: [] for name in names} + stalls: dict[str, list[float]] = {name: [] for name in names} + firsts: dict[str, list[float]] = {name: [] for name in names} + + block = 2 * len(names) + fixed = rounds is not None + limit = rounds if fixed else max_rounds + stop = "fixed" if fixed else "max_rounds" + t0 = time.perf_counter() + r = 0 + while r < limit: + order = _order(names, r) + for name in order: + if flush: + flush_caches() + ms, stall, first = _time_block_full(variants[name], plans[name].iters, + plans[name].group) + times[name].append(ms) + stalls[name].append(stall) + firsts[name].append(first) + r += 1 + if fixed: + continue + elapsed = time.perf_counter() - t0 + # The hard cap is checked every round, not only on a block boundary: + # a cell whose single round costs minutes must be able to stop without + # first completing a balanced design it cannot afford. When it does, + # ``balanced`` says so. + if elapsed > hard_budget_s and r >= floor_rounds: + stop = "budget" + break + if r % block: + continue + if r >= min_rounds and _precise_enough(names, times, target_rel): + stop = "converged" + break + if r >= floor_rounds and elapsed + elapsed / r * block > budget_s: + stop = "budget" + break + + seconds = time.perf_counter() - t0 + # With one variant there is no order to balance, so the question does not + # arise; with more, the design only closes on a multiple of ``2 * n``. + balanced = len(names) < 2 or (r % block) == 0 + return { + name: Measurement( + name, tuple(times[name]), tuple(stalls[name]), + iters=plans[name].iters, group=plans[name].group, + firsts=tuple(firsts[name]), tax_ms=plans[name].tax_ms, + stop=stop, balanced=balanced, seconds=seconds, + ) + for name in names + } + + +def _precise_enough(names, times, target_rel: float) -> bool: + """Every reported quantity within ``target_rel``: the medians and the ratios.""" + for name in names: + vals = times[name] + med = statistics.median(vals) + if med <= 0 or _half_width(vals) / med > target_rel: + return False + ref = names[0] + for name in names[1:]: + logs = [math.log(a / b) for a, b in zip(times[name], times[ref]) if b > 0] + if not logs or _half_width(logs) > target_rel: + return False + return True + + +def format_table( + rows: Iterable[Sequence[object]], + headers: Sequence[str], + aligns: str | None = None, +) -> str: + """A plain fixed-width table; the reports are read in a terminal.""" + rows = [[str(c) for c in row] for row in rows] + widths = [ + max(len(h), *(len(r[i]) for r in rows)) if rows else len(h) + for i, h in enumerate(headers) + ] + aligns = aligns or "l" * len(headers) + def fmt(cells: Sequence[str]) -> str: + return " ".join( + c.rjust(w) if a == "r" else c.ljust(w) + for c, w, a in zip(cells, widths, aligns) + ) + lines = [fmt(headers), " ".join("-" * w for w in widths)] + lines += [fmt(r) for r in rows] + return "\n".join(lines) diff --git a/triton_conv3d/bwd_data.py b/triton_conv3d/bwd_data.py new file mode 100644 index 0000000..fd29d61 --- /dev/null +++ b/triton_conv3d/bwd_data.py @@ -0,0 +1,360 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""Backward-data: the forward kernel, run on a transformed weight. + +At ``stride == 1`` the gradient with respect to the input is the *same* +contraction as the forward pass. Starting from the forward, + + y[n, oc, o] = sum_{ic, t} x[n, ic, o*1 - p + t*dil] * w[oc, ic, t] + +each ``x`` voxel contributes to every ``y`` voxel whose window covers it, so + + gx[n, ic, i] = sum_{oc, t} gy[n, oc, i + p - t*dil] * w[oc, ic, t] + +which is a gather with a *negative* tap stride. Substituting ``t' = k-1-t`` +flips it back:: + + i + p - t*dil = i - (dil*(k-1) - p) + t'*dil + +and that is exactly the forward gather with padding ``dil*(k-1) - p``. So + + grad_input = conv3d_forward(grad_output, flip_taps(swap_channels(w)), + padding = dil*(k-1) - padding) + +with the output spatial extent working out to the input's on its own: +``OD + 2p' - dil*(k-1) == ID`` identically, for every ``p`` and ``dil``. + +Consequences, all of which the code below leans on: + +* **There is no ``@triton.jit`` in this module, deliberately.** + ``grep -c '^@triton\.jit' triton_conv3d/*.py`` reporting 1 for + ``gather_gemm.py`` and 0 for everything else is the claim this file is making + -- anchored at column 0 so that this very sentence does not satisfy the grep, + which the first version of it did. It is also why the tests here can reuse + the forward's correctness standards unchanged. +* **There is no weight transform either, and there used to be.** ``to_bwd_rsck`` + materialized ``(kd, kh, kw, Cout, Cin)`` with the taps flipped, once per layer + per optimizer step -- 0.531 ms/step over one configuration's 19 Conv3d sites, + which no caching removes because the optimizer dirties every parameter every + step. Both halves of it are now free, and neither cost what it looked like it + would. The flip is a constexpr index (``taps - 1 - dij``: flipping all three + kernel axes is the complement of a mixed-radix index). The transpose is *not + performed at all* -- the kernel addresses the weight through its strides, and + a ``permute`` supplies those, so "transposed" is a matter of which stride is + which. A ``channels_last_3d`` parameter is ``[Cout][tap][Cin]``, and this + direction's N is ``Cin``, so the parameter is read here with the *same* + contiguous-N tile the forward gets from a materialized RSCK buffer. Measured + against the ``to_bwd_rsck`` path it replaced, on the eight hottest sites of + config A: 0.98-1.02x, i.e. free. +* The tuning does **not** transfer from the forward. The effective GEMM has + the channel widths swapped: forward ``128 -> 64`` is ``N=64, K=128*27``, + and its backward-data is ``N=128, K=64*27``. Same kernel, different corner + of the tuning surface, so :data:`_TUNED_BWD` is its own table. +* ``PADDED`` is always true for ``k > 1``, whatever the forward's padding was. + The equivalent forward has ``p' = d*(k-1) - p``, which is 2 for an unpadded + ``k = 3`` and 1 for a "same"-padded one -- both non-zero, so the six-compare + boundary predicate is unavoidable here even in the one case the forward + compiles it away. That is a property of the mathematics, not of the reuse, + and it is why this direction is the one place the halo'd/padded distinction + does *not* change which kernel body is compiled. + +Restrictions beyond the forward's +================================= + +``stride > 1`` is refused. The substitution above needs ``o = i + p - t*dil`` +to have a solution for every ``i``, which at ``stride s`` it has only when +``s`` divides ``i + p - t*dil``; the backward is then a *dilated scatter* into +a strided sub-lattice, which the forward kernel's addressing cannot express. +``padding > dil*(k-1)`` is refused for the mirror-image reason: ``p'`` would be +negative, i.e. the backward would have to *crop*, which is again not something +the forward gather does. Neither occurs in ScaFFold's corpus. +""" + +from __future__ import annotations + +import math +from typing import Sequence + +import torch + +from .gather_gemm import ( + ConvConfig, + _MFMA_KDIM, + _check_weight_rsck, + _triple, + conv3d_forward, + select_config, + tune_key, +) + +__all__ = [ + "conv3d_backward_data", + "is_supported_bwd_data", + "bwd_data_padding", +] + + +def bwd_data_padding(padding, dilation, kernel) -> tuple[int, int, int]: + """The forward padding that reproduces backward-data: ``dil*(k-1) - p``.""" + p = _triple(padding, "padding") + d = _triple(dilation, "dilation") + k = _triple(kernel, "kernel") + return tuple(d[i] * (k[i] - 1) - p[i] for i in range(3)) # type: ignore[return-value] + + +# --------------------------------------------------------------------------- +# Tuning +# --------------------------------------------------------------------------- + + +def _tuned(bm: int, bn: int, bk: int, warps: int, group_m: int = 6) -> ConvConfig: + return ConvConfig(BLOCK_M=bm, BLOCK_N=bn, BLOCK_K=bk, GROUP_M=group_m, + num_warps=warps, num_stages=2, matrix_instr_nonkdim=16, + kpack=1 if bk <= 16 else 2) + + +#: Measured backward-data winners, keyed by the **forward** problem's +#: ``(dtype, Cin, Cout, kernel)`` -- i.e. by the convolution a reader would name +#: -- even though the GEMM that runs has those two widths swapped. Keying it +#: the other way round would make ``512 -> 256`` in this table and ``512 -> 256`` +#: in the forward's mean different things, which is a trap not worth setting. +#: +#: Source: ``work/triton-conv/m2_bwd.json`` -- 21 problems, 1468 timed +#: configurations, analysed by ``work/triton-conv/m2_analyze.py``. Only +#: channel pairs that were actually timed appear; a miss falls to the heuristic +#: on the effective widths, which is a real gap and not an extrapolation dressed +#: up as a measurement. +#: +#: Three things this table records that the forward's does *not*: +#: +#: * ``BLOCK_N`` runs to **256**. The GEMM's N is ``Cin``, so the decoder +#: convolutions whose forward is skinny (``Cout=64``) are the widest ones +#: here, and the tile follows. +#: * ``BLOCK_K`` of **32** wins twice. The reduction is ``Cout * 27``, which for +#: ``Cout=64`` is only 1728, and a deep K-tile then wastes the tail. +#: ``GROUP_M`` and ``matrix_instr_nonkdim`` are pinned at 6 and 16, as in the +#: forward, and both were re-checked here rather than inherited. ``nonkdim=32`` +#: won exactly one pair (``64 -> 64``) and forcing 16 there costs **0.3%**. +#: ``GROUP_M=8`` won 8 of 21 problems with a geometric mean of 1.006x, and that +#: count is biased in its favour -- the sweep only tries 8 on each problem's +#: finalists -- so 6 (MI300A's XCD count) is used throughout. Selecting each +#: entry from the ``g6``/``nk16`` arm alone, which is the arm every configuration +#: was timed in, costs a geometric mean of **1.9%** against the per-problem best. +_TUNED_BWD: dict[tuple, ConvConfig] = { + **{ + tune_key(torch.bfloat16, cin, cout, (3, 3, 3)): cfg + for (cin, cout), cfg in { + (64, 64): _tuned(256, 64, 32, 4), + (64, 128): _tuned(128, 64, 64, 4), + (128, 64): _tuned(128, 128, 32, 4), + (128, 128): _tuned(128, 128, 64, 4), + (128, 256): _tuned(128, 128, 64, 4), + (256, 128): _tuned(128, 256, 64, 8), + (256, 256): _tuned(128, 256, 64, 8), + (256, 512): _tuned(128, 256, 64, 8), + (512, 256): _tuned(128, 256, 64, 8), + (512, 512): _tuned(128, 128, 128, 8), + (512, 1024): _tuned(64, 64, 128, 4), + (1024, 512): _tuned(128, 256, 64, 8), + (1024, 1024): _tuned(128, 128, 128, 8), + }.items() + }, + # The segmentation head. ``k=1`` means the backward has no gather at all + # (``p' = 0``) and a reduction of just ``Cout = 6``, so it is a different + # regime from every entry above and gets its own key. + # + # ``num_warps = 2``, not the 4 M2 shipped, and that is the only axis of this + # entry that moved: M2's sweep drew ``num_warps`` from ``{4, 8, seed}``, so + # 1 and 2 were never timed at any site in this project. Raced at all three + # head volumes (``work/triton-conv/tune/warps_bwdd.json``) 2 is 1.099x, + # 1.057x and 1.101x over 4, and 8 is 0.89-0.90x. The reduction here is + # ``Cout * taps = 6``, one MFMA fragment deep, so a second pair of waves has + # nothing to reduce and only replicates the addressing. 1 warp is within + # noise of 2 at this site and is 0.245x at ``128 -> 128 @ 66^3``, so the + # narrow regime is where this stops, not a direction-wide rule. + tune_key(torch.bfloat16, 64, 6, (1, 1, 1)): _tuned(256, 64, 16, 2), +} + + +def register_tuned_bwd_data(dtype, cin, cout, kernel, config: ConvConfig) -> None: + _TUNED_BWD[tune_key(dtype, cin, cout, kernel)] = config + + +def bwd_data_config( + grad_output_shape: Sequence[int], cin: int, kernel: Sequence[int], + dtype: torch.dtype = torch.bfloat16, *, padding=0, dilation=1, +) -> ConvConfig: + """The config :func:`conv3d_backward_data` would pick for this problem. + + Exposed because the benchmark and the ISA gate both need to know what the + shipped path chooses without having to reconstruct the effective GEMM. + """ + n, cout, *out_sp = (int(v) for v in grad_output_shape) + k = _triple(kernel, "kernel") + d = _triple(dilation, "dilation") + p = _triple(padding, "padding") + in_sp = [o + 2 * (d[i] * (k[i] - 1) - p[i]) - d[i] * (k[i] - 1) + for i, o in enumerate(out_sp)] + m = n * math.prod(in_sp) + return select_config( + m, cout, cin, k, dtype, + table=_TUNED_BWD, key=tune_key(dtype, cin, cout, k), + ) + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + + +def is_supported_bwd_data( + grad_output: torch.Tensor, + weight: torch.Tensor, + input_shape: Sequence[int], + stride=1, + padding=0, + dilation=1, + groups: int = 1, +) -> bool: + """Whether :func:`conv3d_backward_data` will serve this call. + + Same asymmetry as :func:`~triton_conv3d.gather_gemm.is_supported`: the + caller's fallback is MIOpen, which is correct everywhere, so a false + negative costs a little speed and a false positive returns a wrong gradient. + + The two checks that are *not* in the forward's predicate are the two the + module docstring derives: ``stride == 1``, and ``padding <= dil*(k-1)``. + They make this the narrowest of the three gates -- the forward and + backward-weight both serve a stride this one refuses -- so a caller that + will differentiate must ask + :func:`~triton_conv3d.gather_gemm.is_supported_all` rather than assume the + forward's ``True`` covers this direction. + """ + if groups != 1: + return False + if grad_output.dim() != 5 or weight.dim() != 5 or len(tuple(input_shape)) != 5: + return False + if grad_output.dtype != weight.dtype or grad_output.dtype not in _MFMA_KDIM: + return False + # Same device, not merely both on *a* device: Triton launches on the current + # one and dereferences the foreign pointer regardless, and on a node with + # peer access enabled -- which is how ScaFFold runs its four GPUs -- that + # reads another rank's memory instead of faulting. A wrong gradient, not a + # crash. + if (not grad_output.is_cuda or not weight.is_cuda + or weight.device != grad_output.device): + return False + try: + s = _triple(stride, "stride") + p = _triple(padding, "padding") + d = _triple(dilation, "dilation") + except ValueError: + return False + k = tuple(int(v) for v in weight.shape[2:]) + if any(v < 1 for v in d) or any(v < 0 for v in p): + return False + if s != (1, 1, 1): + return False + if any(p[i] > d[i] * (k[i] - 1) for i in range(3)): + return False + n, cin, *in_sp = (int(v) for v in input_shape) + if int(weight.shape[0]) != int(grad_output.shape[1]): + return False + if int(weight.shape[1]) != cin: + return False + if int(grad_output.shape[0]) != n: + return False + # The gradient's own shape has to be the one this ``grad_output`` came from, + # or the caller has mixed up two problems and the kernel would happily write + # a differently-shaped answer into a buffer sized for the other one. + for i in range(3): + if int(grad_output.shape[2 + i]) != in_sp[i] + 2 * p[i] - d[i] * (k[i] - 1): + return False + # ``n`` alongside the spatial extents, for the same reason: this predicate's + # own "every output voxel must exist" argument excludes an empty batch, and + # a gate that answers ``True`` for a problem with no voxels in it is stating + # something it has not checked. Costs a fallback to MIOpen on a call that + # has nothing to compute. + if n < 1 or any(v < 1 for v in in_sp): + return False + return True + + +def conv3d_backward_data( + grad_output: torch.Tensor, + weight: torch.Tensor, + input_shape: Sequence[int], + stride=1, + padding=0, + dilation=1, + groups: int = 1, + *, + config: ConvConfig | None = None, + weight_rsck: torch.Tensor | None = None, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Gradient of a 3-D convolution with respect to its input. + + ``grad_output`` and the returned gradient are ``channels_last_3d``. + ``input_shape`` is PyTorch's ``(N, Cin, D, H, W)``; it is redundant at + ``stride == 1`` (the derivation recovers it exactly) and is required anyway, + both to match ``torch.nn.grad.conv3d_input``'s signature and because a + mismatch is the cheapest available check that the caller has not paired a + ``grad_output`` with the wrong problem. + + ``weight_rsck`` is the **forward's** RSCK buffer, ``(kd, kh, kw, Cin, + Cout)`` -- the same tensor + :func:`~triton_conv3d.gather_gemm.conv3d_forward` takes, not a second one + transformed for this direction. It is optional and, on a + ``channels_last_3d`` parameter, pointless: pass the parameter as ``weight`` + and the kernel reads it in place. + """ + if not is_supported_bwd_data(grad_output, weight, input_shape, stride, + padding, dilation, groups): + raise NotImplementedError( + f"unsupported: grad_output={tuple(grad_output.shape)}/" + f"{grad_output.dtype} w={tuple(weight.shape)} " + f"input_shape={tuple(input_shape)} stride={stride} " + f"padding={padding} dilation={dilation} groups={groups}" + ) + k = tuple(int(v) for v in weight.shape[2:]) + pad = bwd_data_padding(padding, dilation, k) + + n, cin, *in_sp = (int(v) for v in input_shape) + cout = int(grad_output.shape[1]) + if config is None: + config = select_config( + n * math.prod(in_sp), cout, cin, k, grad_output.dtype, + table=_TUNED_BWD, + key=tune_key(grad_output.dtype, cin, cout, k), + ) + + # A *view* in both branches, never a copy. The effective convolution's + # channel widths are the real one's swapped, and ``permute`` is exactly that + # relabelling: the forward's :func:`~triton_conv3d.gather_gemm._weight_plan` + # reads the resulting strides and picks the load orientation off them, so + # both of these are addressed in place. + # + # * the parameter itself becomes ``(Cin, Cout, kd, kh, kw)``. Channels-last + # makes its ``Cin`` contiguous, which is this GEMM's N -- the same + # ``W_ORDER == 0`` load the forward has always used, at a different stride. + # * the forward's RSCK buffer becomes the same shape with ``Cout`` + # contiguous, i.e. this GEMM's K, so its N is strided and it takes the + # general load. Tap ``t`` of this direction is tap ``flip(t)`` of the + # weight and its matrix is the transpose; both are addressing, and neither + # is a copy or a register shuffle. + # + # ``out=`` is still validated by the forward rather than here, and that is + # exact rather than approximate: the effective forward's output shape + # ``(n, Cin_eff, out_d, out_h, out_w)`` *is* ``input_shape``. + if weight_rsck is None: + w_view = weight.permute(1, 0, 2, 3, 4) + else: + # Checked here rather than by the forward: what the forward would check + # it against is the *effective* problem's RSCK shape, and this is the + # real problem's, with the two channel widths the other way round. + _check_weight_rsck(weight_rsck, (*k, cin, cout), grad_output) + w_view = weight_rsck.permute(3, 4, 0, 1, 2) + return conv3d_forward( + grad_output, w_view, None, 1, pad, dilation, 1, + config=config, weight_flip=True, out=out, + ) diff --git a/triton_conv3d/gather_gemm.py b/triton_conv3d/gather_gemm.py new file mode 100644 index 0000000..34aff56 --- /dev/null +++ b/triton_conv3d/gather_gemm.py @@ -0,0 +1,1396 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""Forward 3-D convolution as a fused implicit GEMM over NDHWC tensors. + +The convolution is evaluated as a single GEMM whose ``A`` operand is gathered +rather than materialized:: + + M = N * OD * OH * OW output voxels, a flat linear index + N = Cout output channels + K = kd * kh * kw * Cin taps x input channels + +Nothing is written to memory between the gather and the matrix core: for each +tap the kernel re-reads the input at a *constant* voxel shift, which is what +makes NDHWC the right layout. ``Cin`` is the fastest-varying axis of the input +and is also the GEMM's reduction axis, so a K-tile is a contiguous vector load, +and moving from one tap to the next is a scalar addend on the row offset rather +than per-element index arithmetic. + +Provenance +========== + +The tiling is PyTorch Inductor's ``conv3d_template`` +(``torch/_inductor/kernel/conv.py``), not a fresh derivation. Three things are +taken from it unchanged because they are already right: + +* the M-unravel of a fused ``ndhw`` linear index by successive ``%`` / ``//``; +* the fused ``dijk`` reduction loop with **channel blocks innermost and taps + outermost**, which keeps the contiguous ``C`` axis fast-varying so the loads + vectorize (Inductor's own comment records that the nested-loop form is + slightly slower); +* halo handling as pure predication -- no shared-memory staging, no im2col. + +What is *not* taken from it is the address arithmetic. The template is written +against NCDHW; every offset here is re-derived for NDHWC. What survives that +re-derivation, deliberately, is the pointer *shape*: every global access is +``splat(scalar_base) + offset_tensor``. That is condition 1 of the AMD backend's +``canUseBufferOps``, and a tensor-of-pointers formulation loses buffer-op +lowering outright -- the documented reason naive Triton convolutions are slow. + +The weight is read where it lies +================================ + +There is no weight transform on the shipped path, in any direction, and the +reason is worth stating because the obvious design has one. The B tile wants +``(BLOCK_K, BLOCK_N) = (Cin, Cout)`` per tap, PyTorch stores neither channel +axis in that position, and the natural fix -- materialize +``(kd, kh, kw, Cin, Cout)`` once and reuse it -- costs **0.786 ms/step** for the +forward and 0.531 for backward-data across one configuration's 19 Conv3d sites. +Not per call: per *optimizer step*, because the optimizer dirties every +parameter every step, so no cache removes it. + +So the kernel addresses the weight through its strides instead, and which axis +is unit-stride selects the load (``W_ORDER``): + +=========================== ================== ==================== ========== +weight layout forward backward-data copies +=========================== ================== ==================== ========== +``channels_last_3d`` gathered columns contiguous rows **no** +RSCK-strided or ``rsck=`` contiguous rows gathered columns **no** +PyTorch default -- -- yes +=========================== ================== ==================== ========== + +A ScaFFold model is entirely in the first row: ``worker.py`` moves it to +``channels_last_3d`` at construction, which makes every conv weight +``[Cout][kd][kh][kw][Cin]``. The two rows that copy nothing are within 0.8% of +each other over the eight hottest config-A sites, so the choice between them is +not a performance question; the third is 4.8-9.0x slower if addressed in place, +because neither tile axis is dense, and is therefore copied. + +Two designs that look better and measure worse, both raced per site +(``work/triton-conv/review/wt/``): loading the tile coalesced and transposing +it in registers (1.05-1.55x, the transpose is the cost); and holding the +parameter in RSCK order, which is 0.159 ms/step better in the kernels and +**2.6 ms/step worse in the optimizer**, since the gradient this package produces +is channels-last and the elementwise update would then be strided. + +Configuration constraints are hard +================================== + +On gfx942 an illegal MFMA configuration does not fail. It emits **zero** MFMA +instructions, falls back to vector FMA, and returns correct results at a +fraction of the speed. So :class:`ConvConfig` refuses rather than deprioritises: +``BLOCK_M``/``BLOCK_N`` must be multiples of ``matrix_instr_nonkdim`` and +``BLOCK_K`` a multiple of the intrinsic's ``kDim`` (16 at nonkdim 16, 8 at 32, +for bf16). :func:`verify_isa` exists because the only way to know the +constraints were met is to read the emitted ISA. +""" + +from __future__ import annotations + +import dataclasses +from typing import Sequence + +import torch +import triton +import triton.language as tl + +__all__ = [ + "ConvConfig", + "conv3d_forward", + "default_config", + "candidate_configs", + "is_supported", + "is_supported_all", + "select_config", + "to_rsck", + "tune_key", +] + + +# --------------------------------------------------------------------------- +# The kernel +# --------------------------------------------------------------------------- + + +@triton.jit +def _conv3d_fwd_kernel( + X, W, Y, BIAS, + # Sizes. ``M_TOTAL`` is ``BATCH * OUT_D * OUT_H * OUT_W``. + BATCH, IN_D, IN_H, IN_W, + OUT_D, OUT_H, OUT_W, + CIN, COUT, M_TOTAL, + # Element strides. The channel stride of X and Y is 1 by construction -- + # that is what NDHWC means -- so it is not passed and not multiplied by. + stride_xn, stride_xd, stride_xh, stride_xw, + # The weight, described by three strides over the *effective* GEMM's axes: + # the fused tap index, the reduction axis K (Cin), and the output axis N + # (Cout). Which of the two channel strides is 1 is a constexpr (``W_ORDER``) + # rather than a runtime fact, because it decides how the tile is loaded. + stride_wt, stride_wk, stride_wn, + stride_yn, stride_yd, stride_yh, stride_yw, + KD: tl.constexpr, KH: tl.constexpr, KW: tl.constexpr, + SD: tl.constexpr, SH: tl.constexpr, SW: tl.constexpr, + PD: tl.constexpr, PH: tl.constexpr, PW: tl.constexpr, + DD: tl.constexpr, DH: tl.constexpr, DW: tl.constexpr, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + BLOCK_K_COUNT: tl.constexpr, GROUP_M: tl.constexpr, + HAS_BIAS: tl.constexpr, EVEN_K: tl.constexpr, EVEN_N: tl.constexpr, + PADDED: tl.constexpr, INDEX_DTYPE: tl.constexpr, + INPUT_PRECISION: tl.constexpr, + W_ORDER: tl.constexpr, W_FLIP: tl.constexpr, +): + # -- which output tile this program owns ------------------------------ + # + # A flat program id with grouped-M ordering rather than a 2-D grid: the + # group width is the L2 swizzle, and on MI300A it wants to be a multiple of + # the 6 XCDs rather than MI300X's 8. Programs in a group share their B + # tiles, which for a convolution is the whole weight -- small and hot. + pid = tl.program_id(0) + grid_m = tl.cdiv(M_TOTAL, BLOCK_M) + grid_n = tl.cdiv(COUT, BLOCK_N) + width = GROUP_M * grid_n + group_id = pid // width + group_size = min(grid_m - group_id * GROUP_M, GROUP_M) + pid_m = group_id * GROUP_M + (pid % group_size) + pid_n = (pid % width) // group_size + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + + # -- unravel the fused ndhw index ------------------------------------- + # + # Done once, outside the reduction. The divisions are expensive and the + # whole point of hoisting them is that the tap shift below is then a scalar. + idx_w = offs_m % OUT_W + tmp = offs_m // OUT_W + idx_h = tmp % OUT_H + tmp = tmp // OUT_H + idx_d = tmp % OUT_D + idx_n = tmp // OUT_D + + # Input coordinate of tap (0,0,0); tap (d,i,j) is this plus a scalar. + src_d = idx_d * SD - PD + src_h = idx_h * SH - PH + src_w = idx_w * SW - PW + + # The row offset of the A operand. Cast per term rather than after the sum: + # ``idx_n * stride_xn`` alone overflows int32 for a batched scale-8 volume, + # and the sum would then be wrong before the widening ever happened. + x_row = ( + idx_n.to(INDEX_DTYPE) * stride_xn + + src_d.to(INDEX_DTYPE) * stride_xd + + src_h.to(INDEX_DTYPE) * stride_xh + + src_w.to(INDEX_DTYPE) * stride_xw + ) + m_valid = offs_m < M_TOTAL + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + + # -- reduction: taps outermost, channel blocks innermost --------------- + for dijk in range(KD * KH * KW * BLOCK_K_COUNT): + k = (dijk % BLOCK_K_COUNT) * BLOCK_K + dij = dijk // BLOCK_K_COUNT + j = dij % KW + di = dij // KW + i = di % KH + d = di // KH + + offs_k = k + tl.arange(0, BLOCK_K) + + # A: one voxel shift. The addend is a scalar, so this is a uniform + # bump of the row offset rather than a recomputed gather. + tap_off = (d * DD) * stride_xd + (i * DH) * stride_xh + (j * DW) * stride_xw + x_ptrs = X + (x_row + tap_off.to(INDEX_DTYPE))[:, None] + offs_k[None, :] + + if PADDED: + in_d = src_d + d * DD + in_h = src_h + i * DH + in_w = src_w + j * DW + row_ok = ( + m_valid + & (in_d >= 0) & (in_d < IN_D) + & (in_h >= 0) & (in_h < IN_H) + & (in_w >= 0) & (in_w < IN_W) + ) + else: + # Unpadded: every tap of an in-range output voxel is in range, so + # the six compares above are dead. Worth compiling out -- they run + # 27 times per K sweep. This is the *rarer* arm at a ScaFFold site: + # the adapter halos only the split axis, so a k>1 production + # convolution compiles the PADDED branch above at every + # configuration. The transposed upsamplers and the k=1 head land + # here. + row_ok = m_valid + mask_x = tl.broadcast_to(row_ok[:, None], (BLOCK_M, BLOCK_K)) + if not EVEN_K: + mask_x = mask_x & (offs_k < CIN)[None, :] + a = tl.load(x_ptrs, mask=mask_x, other=0.0) + + # B: the weight tile, in whatever layout the weight arrived in. + # + # W_ORDER == 0 Cout is contiguous, so the tile is a run of BLOCK_N + # elements per row -- the RSCK buffer :func:`to_rsck` + # materializes, and also a channels-last *parameter* seen + # from backward-data, whose N is Cin. + # W_ORDER == 1 Cout is *not* contiguous, so each column of the tile is + # addressed on its own. A channels-last parameter seen + # from the *forward* is this: Cin is unit-stride, and Cin + # is K. + # + # An uncoalesced B tile sounds like it should be much worse than a + # vectorized one and is not -- 0.96-1.05x per site, winning 6 of the 8 + # hottest forward cells outright (``1024->512 @ 18^3``: 0.384 vs 0.408 ms + # before the transform it avoids is charged at all). What matters is not + # that the *lanes* are contiguous but that the tile's K-run is: at + # ``BLOCK_K`` consecutive unit-stride elements each column costs two + # cache lines, the weight is small and stays hot, and B is not where the + # bandwidth goes. Take that away -- a weight where neither channel axis + # is unit-stride -- and the same instruction sequence is 4.8-9.0x + # slower, which is why :func:`_weight_plan` refuses it rather than + # compiling it. + # + # Two alternatives were implemented and measured worse, both on the eight + # hot config-A sites (``work/triton-conv/review/wt/``): loading the tile + # coalesced along K and transposing it in registers into the + # ``(BLOCK_K, BLOCK_N)`` the dot wants costs 1.05-1.55x, and holding the + # parameter in RSCK order costs the *backward* direction the same, since + # the axis RSCK makes contiguous is backward-data's reduction axis. + # + # ``W_FLIP`` reverses the fused tap index. Flipping all three kernel + # axes is the complement of a mixed-radix index, i.e. exactly + # ``taps - 1 - dij``, so backward-data's tap flip is this scalar rather + # than a materialized copy of the weight. + # + # Offsets are widened by the same ``INDEX_DTYPE`` as A. The widest + # weight in this project is 28.3 M elements, 80x below int32, but + # ``taps * Cin * Cout`` is bounded by nothing a caller cannot exceed: at + # 2.30e9 elements the truncated offset goes *negative* and faults the + # GPU. Cast per term rather than after the sum -- ``dij * stride_wt`` is + # the term that overflows on its own. On the int32 path the casts are + # frontend no-ops, so the operand keeps the buffer-load eligibility it + # has today; it only loses it at sizes where PLAN.md 3.5 says the storage + # has already lost it. + dij_w = (KD * KH * KW - 1 - dij) if W_FLIP else dij + w_row = ( + dij_w.to(INDEX_DTYPE) * stride_wt + offs_k.to(INDEX_DTYPE) * stride_wk + ) + if W_ORDER == 0: + w_ptrs = W + w_row[:, None] + offs_n[None, :] + else: + w_ptrs = W + w_row[:, None] + offs_n[None, :].to(INDEX_DTYPE) * stride_wn + if EVEN_K and EVEN_N: + b = tl.load(w_ptrs) + elif EVEN_K: + b = tl.load(w_ptrs, mask=(offs_n < COUT)[None, :], other=0.0) + elif EVEN_N: + b = tl.load(w_ptrs, mask=(offs_k < CIN)[:, None], other=0.0) + else: + b = tl.load( + w_ptrs, + mask=(offs_k < CIN)[:, None] & (offs_n < COUT)[None, :], + other=0.0, + ) + + # ``input_precision`` only bites for fp32 operands, where the backend's + # default splits the dot into reduced-precision pieces. bf16 already + # accumulates in fp32 and is unaffected; fp32 is the ``more_determinism`` + # path and has to actually be fp32, so it is asked for explicitly. + acc = tl.dot(a, b, acc, input_precision=INPUT_PRECISION) + + if HAS_BIAS: + bias = tl.load(BIAS + offs_n, mask=offs_n < COUT, other=0.0) + acc += bias[None, :].to(tl.float32) + + y_row = ( + idx_n.to(INDEX_DTYPE) * stride_yn + + idx_d.to(INDEX_DTYPE) * stride_yd + + idx_h.to(INDEX_DTYPE) * stride_yh + + idx_w.to(INDEX_DTYPE) * stride_yw + ) + y_ptrs = Y + y_row[:, None] + offs_n[None, :] + mask_y = tl.broadcast_to(m_valid[:, None], (BLOCK_M, BLOCK_N)) + if not EVEN_N: + mask_y = mask_y & (offs_n < COUT)[None, :] + tl.store(y_ptrs, acc.to(Y.dtype.element_ty), mask=mask_y) + + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + + +#: MFMA reduction depth per ``matrix_instr_nonkdim`` on gfx942, by operand +#: dtype. ``BLOCK_K`` must be a multiple of this or ``chooseMfmaInstruction`` +#: rejects the shape ("would introduce data duplication") and the dot silently +#: lowers to FMA. Source: Triton v3.7.0 ``MfmaGroup.cpp`` plus the +#: ``inputKSize % kDim`` check in ``AccelerateAMDMatmul.cpp``. +_MFMA_KDIM = { + torch.bfloat16: {16: 16, 32: 8}, + torch.float16: {16: 16, 32: 8}, + torch.float32: {16: 4, 32: 2}, +} + + +@dataclasses.dataclass(frozen=True) +class ConvConfig: + """One launch configuration, with the gfx942 constraints enforced. + + ``validate`` is not advisory. Every constraint here has a *silent* failure + mode: an illegal ``matrix_instr_nonkdim``, or a ``BLOCK_K`` that is not a + multiple of the intrinsic's reduction depth, produces a kernel that runs and + returns the right answer with no MFMA instruction in it at all. A config + generator that merely ranked such a config last would still be feeding + meaningless entries into a best-of sweep. + """ + + BLOCK_M: int = 128 + BLOCK_N: int = 64 + BLOCK_K: int = 64 + GROUP_M: int = 6 + num_warps: int = 4 + num_stages: int = 2 + matrix_instr_nonkdim: int = 16 + kpack: int = 2 + waves_per_eu: int = 0 + + def __str__(self) -> str: + return ( + f"{self.BLOCK_M}x{self.BLOCK_N}x{self.BLOCK_K}" + f"/g{self.GROUP_M}/w{self.num_warps}/s{self.num_stages}" + f"/nk{self.matrix_instr_nonkdim}/kp{self.kpack}" + + (f"/we{self.waves_per_eu}" if self.waves_per_eu else "") + ) + + def validate(self, dtype: torch.dtype) -> str | None: + """``None`` if the config can reach the matrix core, else the reason.""" + kdims = _MFMA_KDIM.get(dtype) + if kdims is None: + return f"unsupported operand dtype {dtype}" + if self.matrix_instr_nonkdim not in kdims: + return ( + f"matrix_instr_nonkdim={self.matrix_instr_nonkdim} is not one of " + f"{sorted(kdims)}; anything else falls back to FMA" + ) + nk = self.matrix_instr_nonkdim + if self.BLOCK_M % nk or self.BLOCK_N % nk: + return f"BLOCK_M/BLOCK_N must be multiples of nonkdim={nk}" + if self.BLOCK_K % kdims[nk]: + return f"BLOCK_K must be a multiple of {kdims[nk]} at nonkdim={nk}" + if self.num_warps < 1 or self.num_warps & (self.num_warps - 1): + return "num_warps must be a positive power of two" + if self.num_warps > self.BLOCK_M * self.BLOCK_N // 256: + return "more warps than 16x16 tiles in the output block" + if self.num_stages < 2: + # Block-pingpong is on by default for gfx942 and needs > 1. + return "num_stages must be at least 2 on gfx942" + if self.GROUP_M < 1: + # The swizzle divides by ``GROUP_M * grid_n`` and takes ``pid % + # group_size``; at 0 that is a division by zero inside the kernel, + # which on gfx942 is not a trap but a garbage ``pid_m`` and a memory + # access fault. A negative value reaches the kernel just as far. + # Every legal value is fine, including ones that do not divide + # ``grid_m`` and ones far larger than it. + return "GROUP_M must be at least 1" + return None + + def lds_bytes(self, dtype: torch.dtype) -> int: + """Shared memory the two operand tiles need, in bytes. + + Measured rather than assumed: across 22 tile/dtype combinations the + compiler's own ``metadata.shared`` equalled + ``(BLOCK_M*BLOCK_K + BLOCK_K*BLOCK_N) * itemsize`` exactly, with no + double-buffering factor -- Triton's gfx942 pipeliner keeps one LDS + buffer at ``num_stages=2``. So this is the real number and not a bound. + """ + elem = torch.empty((), dtype=dtype).element_size() + return (self.BLOCK_M * self.BLOCK_K + self.BLOCK_K * self.BLOCK_N) * elem + + def launch_kwargs(self) -> dict: + return { + "num_warps": self.num_warps, + "num_stages": self.num_stages, + "matrix_instr_nonkdim": self.matrix_instr_nonkdim, + "kpack": self.kpack, + "waves_per_eu": self.waves_per_eu, + } + + +def _pow2_at_most(x: int, cap: int) -> int: + return max(16, min(cap, 1 << max(0, (max(1, x)).bit_length() - 1))) + + +#: gfx942's shared memory per workgroup. Exceeding it raises ``OutOfResources`` +#: at launch -- loudly, unlike the MFMA constraints, which is why M1 left it to +#: fail rather than guarding it statically. That is the right call for a *sweep* +#: candidate and the wrong one for the config the entry point picks on its own, +#: which is what :func:`_fit_to_lds` is for. +_LDS_BYTES = 64 * 1024 + + +def _fit_to_lds(cfg: ConvConfig, dtype: torch.dtype) -> ConvConfig: + """Shrink a tile until its operands fit in LDS, then legalize the warps. + + This exists because of an fp32 hole that M2's tests fell into: the block + sizes in :func:`default_config` were chosen against bf16, and fp32 operands + are twice the bytes, so ``Cin >= 512`` in fp32 asked for 128 KiB and the + *shipped* configuration raised. ``more_determinism`` runs the model in + fp32, so that was reachable from a real ScaFFold configuration. + + ``BLOCK_K`` is halved first: it is the reduction depth, so shortening it + costs some reuse but changes neither the grid nor the parallelism, whereas + halving ``BLOCK_M`` doubles the program count. The loud failure remains as + the backstop for an explicitly supplied ``config=``. + """ + nk = cfg.matrix_instr_nonkdim + kdim = _MFMA_KDIM.get(dtype, {}).get(nk) + if kdim is None: + return cfg + while cfg.lds_bytes(dtype) > _LDS_BYTES: + half_k, half_m, half_n = cfg.BLOCK_K // 2, cfg.BLOCK_M // 2, cfg.BLOCK_N // 2 + if half_k >= kdim and half_k % kdim == 0: + cfg = dataclasses.replace(cfg, BLOCK_K=half_k, + kpack=1 if half_k <= 16 else cfg.kpack) + elif half_m >= nk and half_m % nk == 0: + cfg = dataclasses.replace(cfg, BLOCK_M=half_m) + elif half_n >= nk and half_n % nk == 0: + cfg = dataclasses.replace(cfg, BLOCK_N=half_n) + else: + break # nothing left to shrink; let the launch say so + warps = max(1, min(cfg.num_warps, cfg.BLOCK_M * cfg.BLOCK_N // 256)) + return dataclasses.replace(cfg, num_warps=1 << (warps.bit_length() - 1)) + + +#: Below this many programs the grid cannot fill MI300A's 228 CUs, and a +#: narrower ``BLOCK_M`` buys more parallelism than it loses in reuse. Half a +#: wave rather than a whole one: measured, the ``1024 -> 1024`` bottleneck at +#: ``M = 512`` wants 128 programs and is *slower* when pushed to 256. +_MIN_PROGRAMS = 114 + + +def default_config( + m: int, cin: int, cout: int, dtype: torch.dtype = torch.bfloat16 +) -> ConvConfig: + """A config that is legal for any shape and close to tuned for most. + + Measured, not guessed: over 15 corpus shapes and ~1000 timed configurations + the surface is remarkably flat and almost entirely determined by the channel + widths. ``BLOCK_M=128`` won 14 of 15, ``BLOCK_N`` tracks ``Cout`` up to 128, + and ``BLOCK_K`` is 64 below ``Cin=512`` and 128 above. + + Two things this does *not* do, both because the measurement said not to: + + * It does not use ``matrix_instr_nonkdim=32``. The M0 ceiling probe found 32 + winning 9 of 10 ``N=64`` cells on a plain GEMM and worth 1.21x there, which + is why the M1 brief said to sweep it. Swept here on the real convolution, + **16 won all 15 cells**, by 0.6-13% (geometric mean 5.8%). The GEMM result + does not transfer: the convolution's inner loop carries a per-tap boundary + predicate and a 27x longer reduction, so it is not the same instruction + mix the ceiling probe measured. + * It does not scale ``BLOCK_M`` with ``M``. A tall tile only pays while the + grid still fills the device, which at these shapes it always does; the one + place it does not is handled by :data:`_MIN_PROGRAMS`. + """ + block_n = _pow2_at_most(cout, 128) + block_k = 128 if cin >= 512 else _pow2_at_most(cin, 64) + block_m = _pow2_at_most(m, 128) + nonkdim = 16 + kdim = _MFMA_KDIM[dtype][nonkdim] + block_k = max(kdim, block_k - block_k % kdim) + return _fit_to_lds( + _fit_to_grid( + ConvConfig( + BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=block_k, GROUP_M=6, + num_warps=8 if block_k >= 128 or block_n >= 256 else 4, + num_stages=2, matrix_instr_nonkdim=nonkdim, + kpack=1 if block_k <= 16 else 2, + ), + m, cout, + ), + dtype, + ) + + +def _fit_to_grid(cfg: ConvConfig, m: int, cout: int) -> ConvConfig: + """Shrink ``BLOCK_M`` until the grid can fill the device, then legalize warps. + + Applied to tuned entries as well as to the heuristic, because a tuned entry + is keyed on the channel widths and so can be reused at an ``M`` far smaller + than the one it was measured at -- which is exactly where a 128-row tile + stops being a good idea. + """ + nk = cfg.matrix_instr_nonkdim + while ( + cfg.BLOCK_M > max(16, nk) + and (cfg.BLOCK_M // 2) % nk == 0 + and -(-m // cfg.BLOCK_M) * -(-cout // cfg.BLOCK_N) < _MIN_PROGRAMS + ): + cfg = dataclasses.replace(cfg, BLOCK_M=cfg.BLOCK_M // 2) + warps = max(1, min(cfg.num_warps, cfg.BLOCK_M * cfg.BLOCK_N // 256)) + warps = 1 << (warps.bit_length() - 1) + return dataclasses.replace(cfg, num_warps=warps) + + +#: PyTorch Inductor's ROCm convolution seed grid, ``(BLOCK_M, BLOCK_N, BLOCK_K, +#: num_warps)``. Preferred over a blind sweep because these values are already +#: tuned on ROCm; its per-config ``num_stages`` is dropped because +#: ``ROCmConfigHeuristic._filter_configs`` overwrites it with 2 on HIP anyway. +_SEED_TILES: tuple[tuple[int, int, int, int], ...] = ( + (64, 256, 16, 4), + (256, 64, 16, 4), + (128, 128, 32, 8), + (64, 64, 32, 4), + (64, 256, 32, 8), + (256, 64, 32, 8), + (128, 128, 64, 8), + (64, 128, 64, 4), + (128, 64, 64, 4), + (256, 128, 64, 8), + (128, 256, 64, 8), + (128, 128, 128, 8), + (64, 128, 128, 4), + (256, 128, 128, 8), + (128, 256, 128, 8), +) + +#: Extra tiles for the skinny-N regime. ``Cout=64`` is the model's most common +#: output width and the seed grid has little there; M0 found every winner in +#: this band was tall in M with ``BLOCK_N=64``. +_SKINNY_N_TILES: tuple[tuple[int, int, int, int], ...] = ( + (256, 64, 64, 8), + (512, 64, 64, 8), + (256, 64, 128, 8), + (512, 64, 32, 8), + (128, 64, 128, 4), + (64, 64, 64, 4), + (64, 64, 128, 4), + (1024, 64, 32, 8), +) + +#: Tiles for the *narrow*-N regime, ``Cout <= 16``. Only the segmentation head +#: (``64 -> 6``) reaches it in ScaFFold, and until 2026-08-03 the head had no +#: tile evidence at all: ``candidate_configs`` prunes on ``bn > 2 * n2``, which +#: at ``Cout = 6`` gives ``n2 = 16`` and removes every entry of both grids +#: above, so the generator fell through to ``default_config`` and the "sweep" +#: recorded for that site timed the shipped config twice with two ``GROUP_M``. +#: +#: The ``num_warps`` column is the point of this grid, not the tile. A +#: ``BLOCK_N`` of 16 is one MFMA fragment wide, so there is no N work to hand a +#: second wave; four warps each take a quarter of ``BLOCK_M`` and replicate the +#: whole per-K-tile address computation for a fragment that is 10/16 padding. +#: Measured at all three head volumes, one warp is 1.02-1.22x over the shipped +#: four (``work/triton-conv/tune/k1_fwd*.json``). +#: +#: Gated to ``Cout <= 16`` in :func:`candidate_configs` rather than added to the +#: grids above, because a 16-column tile at ``Cout = 512`` is 32x padding and +#: would only lengthen every other site's sweep -- and because ``num_warps=1`` +#: is a **catastrophe** outside this regime: raced on the shipped tile it is +#: 0.166x on the ``128 -> 128 @ 66^3`` forward and 0.245x on its backward-data. +_NARROW_N_TILES: tuple[tuple[int, int, int, int], ...] = ( + (128, 16, 64, 1), + (64, 16, 64, 1), + (128, 16, 64, 2), + (256, 16, 64, 1), + (128, 16, 32, 1), + (64, 16, 32, 1), + (128, 32, 64, 1), +) + + +def candidate_configs( + m: int, cin: int, cout: int, dtype: torch.dtype = torch.bfloat16, + *, group_ms: Sequence[int] = (6,), nonkdims: Sequence[int] = (16, 32), +) -> list[ConvConfig]: + """Configs worth timing for one shape, already pruned to legal ones. + + ``matrix_instr_nonkdim`` is *swept* over {16, 32} rather than fixed at 16. + Fixing it at 16 is what Inductor does and what AMD's guidance says, and at + ``Cout=64`` the M0 probe measured that advice costing 11-18%. + + ``GROUP_M`` defaults to 6 alone -- MI300A's XCD count, and the value that won + the M0 square-GEMM probe -- because sweeping it doubles a list whose cost is + almost entirely JIT compilation. The caller refines it on the finalists + instead, which is where a few percent of L2 locality is actually decidable. + """ + m2 = max(16, triton.next_power_of_2(m)) + n2 = max(16, triton.next_power_of_2(cout)) + k2 = max(16, triton.next_power_of_2(cin)) + out: list[ConvConfig] = [] + seen: set[ConvConfig] = set() + tiles = _SEED_TILES + _SKINNY_N_TILES + if n2 <= 16: + tiles += _NARROW_N_TILES + for bm, bn, bk, seed_warps in tiles: + # Skip tiles that would mostly compute padding. BLOCK_K is capped at + # the channel count rather than twice it because the reduction is + # per-tap: a BLOCK_K above Cin wastes a whole tap's worth of MFMA. + if bm > 2 * m2 or bn > 2 * n2 or bk > k2: + continue + for warps in {4, 8, seed_warps}: + for nonkdim in nonkdims: + for group_m in group_ms: + cfg = ConvConfig( + BLOCK_M=bm, BLOCK_N=bn, BLOCK_K=bk, GROUP_M=group_m, + num_warps=warps, num_stages=2, + matrix_instr_nonkdim=nonkdim, + kpack=1 if bk <= 16 else 2, + ) + # LDS overflow is pruned rather than shrunk: shrinking would + # fold two seed tiles onto one entry and silently + # double-count it in the sweep. Only configs that could not + # have run at all are removed, so no measured winner is lost. + if (cfg.validate(dtype) is not None + or cfg.lds_bytes(dtype) > _LDS_BYTES + or cfg in seen): + continue + seen.add(cfg) + out.append(cfg) + if not out: + out.append(default_config(m, cin, cout, dtype)) + return out + + +def tune_key(dtype: torch.dtype, cin: int, cout: int, + kernel: tuple[int, ...]) -> tuple: + return (str(dtype), cin, cout, tuple(kernel)) + + +def _tuned(bm: int, bn: int, bk: int, warps: int, group_m: int = 6, + nk: int = 16) -> ConvConfig: + """One measured row. + + ``nk`` defaults to 16 because that is what the forward measured -- 16 won + all 15 cells of M1's 1090-config sweep, and :func:`default_config` says why. + It is a parameter at all for exactly one row, the ``3 -> 64`` stem, where 32 + is not chosen for its own sake: ``_MFMA_KDIM[bf16][32] = 8`` is what makes + ``BLOCK_K = 8`` legal, and ``BLOCK_K = 8`` is the whole effect. Raced as a + control, ``nonkdim=32`` on the *shipped* ``128x64x16`` tile measures 0.9682 + ms against 0.9554 -- i.e. nothing. Spelled the same way + :func:`~triton_conv3d.reduce_gemm._tuned` spells it, and for the same + reason: a per-row knob whose default carries the rule. + """ + return ConvConfig(BLOCK_M=bm, BLOCK_N=bn, BLOCK_K=bk, GROUP_M=group_m, + num_warps=warps, num_stages=2, matrix_instr_nonkdim=nk, + kpack=1 if bk <= 16 else 2) + + +#: Measured winners, keyed by ``(dtype, Cin, Cout, kernel)``; a miss falls back to +#: :func:`default_config`, and every hit still goes through :func:`_fit_to_grid`. +#: +#: Keyed on the channel widths and *not* on the spatial extent because that is +#: what the measurement showed: where a channel pair occurs at more than one +#: volume in the corpus -- ``128 -> 64`` at three, ``512 -> 512`` at two -- +#: the same tile won at each. ``GROUP_M`` is the exception; it flips between 6 +#: and 8 across volumes but is worth under 1% either way at the shapes where it +#: flips, so 6 (MI300A's XCD count) is used throughout. +#: +#: Deliberately a table and not ``@triton.autotune``: ScaFFold's figure of merit +#: is total wall time, so a recompile inside a training step is a direct loss. +#: Source: ``work/triton-conv/m1_fwd.json``, 15 problems, ~1050 timed configs. +#: Only channel pairs that were actually timed appear here. The unmeasured +#: pairs -- ``64 -> 128``, ``128 -> 256``, ``256 -> 512``, ``512 -> 1024``, all +#: encoder-side -- fall to the heuristic on purpose: an extrapolated entry in a +#: table called "measured winners" is worse than no entry, because it cannot be +#: told apart from one. Priced since: served by the heuristic those eight cells +#: span 0.80x to 1.32x and are worth **+0.02 ms/step at config A and +0.12 at +#: C** -- a wash, so tuning them is not the action. Three of them are *losses* +#: (``256 -> 512 @ 18^3`` 0.84x, ``512 -> 1024 @ 10^3`` 0.80x, ``512 -> 1024 @ +#: 6x18x18`` 0.90x) and belong on the adapter's block-list instead. (An earlier +#: version of this comment priced the four pairs at "12.7 ms/step"; that figure +#: is config B's *all-directions* total at those pairs, not the forward time an +#: absent row here governs. See ``STATUS.md`` §4.3.) +_TUNED: dict[tuple, ConvConfig] = { + # The segmentation head, and the one entry here that is *not* from + # ``m1_fwd.json``: source ``work/triton-conv/tune/k1_fwd_confirm.json``. + # ``Cout = 6`` prunes every seed tile, so M1's sweep never timed anything + # but ``default_config`` at this site and the tile below is that same + # ``128x16x64`` with **one warp instead of four** -- see + # :data:`_NARROW_N_TILES` for why one, and why only here. Raced against the + # shipped four at all three head volumes in one interleaved block: + # 1.024x @ 128^3, 1.137x @ 64x256^2, 1.217x @ 128x256^2, i.e. 1.19-1.21x of + # MIOpen where the shipped config was 1.16x, 1.04x and **0.99x**. The + # runner-up ``64x16x64/w1`` wins the two smaller volumes by 2-3% and loses + # the largest by 8%, so it is not shipped; the choice between them is worth + # under 0.01 ms/step either way. + tune_key(torch.bfloat16, 64, 6, (1, 1, 1)): _tuned(128, 16, 64, 1), + **{ + tune_key(torch.bfloat16, cin, cout, (3, 3, 3)): cfg + for (cin, cout), cfg in { + # The UNet stem, and the row that refutes ``OPTIMIZATION.md`` D4's + # "``conv 3->64`` is genuinely hopeless, leave it on MIOpen". It was + # never hopeless and it was never a matrix-core feeding problem: the + # reduction axis of this kernel's ``tl.dot`` is ``Cin`` **alone** + # (``BLOCK_K_COUNT = cdiv(Cin, BLOCK_K)``, taps outermost), and + # ``BLOCK_K`` is floored both by the MFMA intrinsic's reduction depth + # and by ``_pow2_at_most``'s own floor of 16 -- so at ``Cin = 3`` + # every dot had **3 live columns of 16** and 81% of the matrix-core + # work multiplied padding this kernel put there itself. + # ``SQ_INSTS_MFMA`` measures 14,155,776 per call at ``130^3``, + # exactly 5.333x the useful FLOPs, against MIOpen's 3,145,728 + # (2.370x -- CK contracts over the merged ``(Z,Y,X,C)`` axis, dense + # 81 padded once to 96). On *issued* MFMA the old config already ran + # at 23.7% of the measured ``tl.dot`` ceiling against CK's 16.2%; it + # just issued 2.25x more of it. + # + # ``BLOCK_K = 8`` is the fix and ``nonkdim=32`` is only how it is + # spelled -- see :func:`_tuned`. Raced against the heuristic's + # ``128x64x16/nk16/w4`` and MIOpen, one interleaved block per volume, + # kernel-only, at every volume the corpus has for this pair: + # 130^3 0.5236 ms vs MIOpen 0.6236 -- 1.193x [1.190,1.196] + # 130x258^2 2.0522 ms vs MIOpen 2.4907 -- 1.214x [1.211,1.217] + # 66x258^2 1.0364 ms vs MIOpen 1.2468 -- 1.203x [1.201,1.206] + # i.e. 1.83-1.86x over the config it replaces, which was 0.651-0.653x + # of MIOpen. Flat across the pair's 4.0x span of volume, which is the + # property this table has twice been burned by not checking. + # + # Bitwise **identical** to the config it replaces on random operands + # (0 of 134 M elements differ), because at ``Cin = 3`` only 3 products + # per tap are non-zero whatever ``BLOCK_K`` is and the 27 taps are + # still visited in order. So no determinism baseline moves. + # + # ``kpack = 1`` is not a rounding detail here: ``kp2`` on the same + # tile is 0.6085 ms, 1.17x worse. Taller than 512 turns over + # (``1024x64x8/w16`` 0.6042); ``GROUP_M`` is inert at this site + # (g1/g6/g12 within 1%). ``work/triton-conv/review/LOSS_OPTIMIZATION.md`` + # §1.2 has the full race and the counters. + (3, 64): _tuned(512, 64, 8, 8, nk=32), + (64, 64): _tuned(128, 64, 64, 4), + (128, 64): _tuned(128, 64, 64, 4), + (128, 128): _tuned(128, 128, 64, 4), + (256, 128): _tuned(128, 128, 64, 4), + (256, 256): _tuned(128, 128, 64, 4), + # The one place the heuristic's "Cin >= 512 wants BLOCK_K=128" rule + # is wrong: here BLOCK_K=64 is 7% faster. Cout=256 rather than 512 + # is what distinguishes it, on one data point, so the rule stands. + (512, 256): _tuned(128, 128, 64, 4), + (512, 512): _tuned(128, 128, 128, 8), + (1024, 512): _tuned(128, 128, 128, 8), + (1024, 1024): _tuned(64, 64, 128, 8), + # The two 2048-channel bottleneck pairs. They are scale-8 sites + # that appeared in no corpus until ``SHAPE_AUDIT.md`` 2.3 found + # them -- the corpus's scale-8 model was a *four*-layer network and + # the harness runs a five-layer one -- so until now they fell to + # the heuristic, and against fresh MIOpen the forward **lost** + # (``TRITON_BASELINES.md`` 5: 0.952x and 0.977x). + # + # ``128x64x128`` wins at **every volume both pairs occur at**, + # which is the property this table has twice been burned by not + # checking, and it is one row rather than three because + # ``_fit_to_grid`` walks ``BLOCK_M`` down 128 -> 64 -> 32 as ``M`` + # falls 512 -> 256 -> 128. Raced against the heuristic it + # replaces, quiet node, **four independently allocated operand sets + # per cell**, both arms sharing each build's operands so the pair is + # immune to the placement effect below; kernel-only, median over + # builds with the worst build in brackets + # (``review/opt3/row_2048_quiet.json``): + # (1024,2048) 8^3 0.3196 vs 0.3224 ms -- 1.010x [1.008] + # (1024,2048) 6x8^2 0.2208 vs 0.2769 -- 1.253x [1.251] + # (1024,2048) 4x8^2 0.1760 vs 0.2304 -- 1.308x [1.306] + # (2048,2048) 8^3 0.6663 vs 0.7321 -- 1.099x [1.097] + # (2048,2048) 6x8^2 0.4819 vs 0.6354 -- 1.318x [1.315] + # (2048,2048) 4x8^2 0.3844 vs 0.5066 -- 1.317x [1.303] + # + # The builds are not ceremony, and one caveat has to travel with + # this row. ``(2048, 2048)`` is the one site in this project whose + # time depends on **where its weight lands**: at 216 MiB against a + # 256 MiB MALL the heuristic is bimodal, two tight states up to + # 14.7% apart and fixed for the life of the allocation + # (``OPTIMIZATION_ROUND3.md`` 4). Measured solo -- one config, one + # operand set, six rebuilds, which is what a caller actually sees -- + # this row is **stable**: 0.5% / 0.7% / 1.4% spread at the three + # volumes against the heuristic's 15% / 5% / 20%. But at ``8^3`` + # its 0.6681 ms sits *between* the heuristic's two states (0.6472 + # and 0.7453), so it beats the unlucky allocation by 1.12x and + # **loses to the lucky one by 0.97x**. It is shipped because the + # expected value and both sharded volumes are clear wins and the + # variance goes away, not because it dominates. + # + # ``(2048, 1024)`` is deliberately **absent**, and that is a result + # rather than an omission: ``128x256x64`` is **1.433x** at ``16^3`` + # and **0.804x** and **0.504x** at the two sharded volumes of the + # same pair (``review/opt3/row_2048to1024.json``, two builds each, + # every interval tight). The discriminator is ``BLOCK_K`` against + # ``M`` -- 64 wins at ``M = 4096`` and 128 wins at 2048 and 1024 -- + # this table is keyed per channel pair and cannot say that, and an + # entry would trade config B's 0.68 ms saving for config C and D's + # 0.23 and 0.16 ms losses. ``_fit_to_grid`` already walks + # ``BLOCK_M`` with ``M``; making it walk ``BLOCK_K`` too is the + # change this measurement argues for, and it is not made here + # because one channel pair is not enough evidence to move a rule + # every pair goes through. + (1024, 2048): _tuned(128, 64, 128, 8), + (2048, 2048): _tuned(128, 64, 128, 8), + }.items() + }, +} + + +def register_tuned(dtype, cin, cout, kernel, config: ConvConfig) -> None: + _TUNED[tune_key(dtype, cin, cout, kernel)] = config + + +def select_config( + m: int, cin: int, cout: int, kernel: Sequence[int], dtype: torch.dtype, + *, table: dict | None = None, key: tuple | None = None, +) -> ConvConfig: + """The config the kernel will run: tuned entry if there is one, else heuristic. + + ``m``/``cin``/``cout`` always describe the GEMM that will actually be issued + -- ``(M, N, K) = (m, cout, cin * prod(kernel))`` -- because that is what + :func:`_fit_to_grid` has to reason about. ``table`` and ``key`` are separate + so that :mod:`triton_conv3d.bwd_data`, whose effective GEMM has the channel + widths *swapped*, can keep a table keyed on the problem a reader recognises + while the tile is still fitted to the grid it will really launch on. + """ + table = _TUNED if table is None else table + if key is None: + key = tune_key(dtype, cin, cout, tuple(kernel)) + tuned = table.get(key) + if tuned is not None: + return _fit_to_lds(_fit_to_grid(tuned, m, cout), dtype) + return default_config(m, cin, cout, dtype) + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + + +def _triple(v, name: str) -> tuple[int, int, int]: + if isinstance(v, int): + return (v, v, v) + t = tuple(int(x) for x in v) + if len(t) != 3: + raise ValueError(f"{name} must be an int or a length-3 sequence, got {v!r}") + return t # type: ignore[return-value] + + +def to_rsck(w: torch.Tensor) -> torch.Tensor: + """PyTorch's ``(Cout, Cin, kd, kh, kw)`` weight as ``(kd, kh, kw, Cin, Cout)``. + + A B tile whose row is a contiguous run wants Cout fastest-varying, which is + what this produces. **It is no longer on any shipped path.** The kernel + reads the parameter wherever it lies, and this copy ran once per layer per + *optimizer step* -- 0.786 ms/step over the 19 Conv3d sites of one + configuration, 8.4x the single-copy floor, because 19 small strided + ``permute().contiguous()`` launches are latency-bound rather than + bandwidth-bound, and no caching could remove it because the optimizer + dirties every parameter every step. Measured against reading a + channels-last parameter in place, materializing this buffer is *slower* on 6 + of the 8 hottest forward sites before the copy is charged at all. + + It is kept, and still supported through ``weight_rsck=``, for the weights + :func:`_weight_plan` refuses -- chiefly PyTorch's *default* layout, in which + neither channel axis is unit-stride and the gathered load is 4.8-9.0x slower + than copying. A ScaFFold parameter is never in it, because the model is + moved to ``channels_last_3d`` at construction. + """ + return w.permute(2, 3, 4, 1, 0).contiguous() + + +#: How the kernel's B operand is laid out -- the values of the kernel's +#: ``W_ORDER`` constexpr. ``_W_GENERAL`` costs nothing measurable against +#: ``_W_N_CONTIG`` (0.96-1.05x per site), which is why there is no third value: +#: a "coalesce along K and ``tl.trans``" order was implemented and measured at +#: 1.05-1.55x, i.e. a real loss, and deleted. +_W_N_CONTIG = 0 +_W_GENERAL = 1 + + +def _weight_plan(w: torch.Tensor) -> tuple[int, int, int, int] | None: + """``(W_ORDER, stride_wt, stride_wk, stride_wn)`` for ``w``, or ``None``. + + ``w`` is the weight *as this GEMM sees it*: ``(Cout, Cin, kd, kh, kw)``, + where for backward-data the two channel widths are the real convolution's + swapped and ``w`` is a permuted view. Strides, not memory format, are what + the kernel needs, so this is a stride computation and not a + ``is_contiguous(memory_format=...)`` test -- the backward-data view is + neither contiguous nor channels-last and is still perfectly addressable. + + ``None`` means materialize :func:`to_rsck` instead, for one of two reasons: + the three kernel axes are not one fused axis of constant stride, which is + what the kernel's single ``dij * stride_wt`` assumes (a weight sliced along a + kernel axis; nothing in this project produces one), or *neither* channel axis + is unit-stride, which is a correctness-neutral but 4.8-9.0x performance + cliff -- see the comment below. + + Extents of 1 carry no observable stride, so they constrain nothing and are + skipped -- ``k=1x1x1`` is a real corpus shape (the segmentation head), and + demanding ``stride(4) == 1`` of it would reject the weights of every model + that has one. + """ + cout, cin, kd, kh, kw = (int(v) for v in w.shape) + s = tuple(int(v) for v in w.stride()) + if kw > 1: + st = s[4] + elif kh > 1: + st = s[3] + elif kd > 1: + st = s[2] + else: + st = 0 # one tap: ``dij`` is always 0, so any stride is the right one + if ((kw > 1 and s[4] != st) + or (kh > 1 and s[3] != st * kw) + or (kd > 1 and s[2] != st * kw * kh)): + return None + if cout == 1 or s[0] == 1: + return (_W_N_CONTIG, st, s[1], 1) + if cin == 1 or s[1] == 1: + return (_W_GENERAL, st, s[1], s[0]) + # Neither channel axis is unit-stride -- PyTorch's *default* weight layout, + # where the only dense axis is the 27-element tap axis, which is not a tile + # axis. Every element of the B tile is then its own cache line: measured + # **4.8-9.0x** slower than materializing RSCK across the eight hottest + # forward sites, the copy charged to every call (9.31 ms against 1.15 at + # ``256->128 @ 66^3``), and 2.2-6.2x on backward-data. So this one really + # does have to be copied, and it is the only layout left that does. + return None + + +def is_supported( + x: torch.Tensor, + w: torch.Tensor, + bias: torch.Tensor | None = None, + stride=1, + padding=0, + dilation=1, + groups: int = 1, +) -> bool: + """Whether :func:`conv3d_forward` will serve this call. + + Deliberately conservative: the caller's fallback is MIOpen, which is correct + everywhere, so a false negative costs a little speed and a false positive + costs a wrong answer. + + It is also **total**. This is the gate of a Triton -> MIOpen rung ladder, so + an argument it cannot interpret has to be a ``False`` and not an exception: + ``padding=None`` and ``padding=1.5`` are ``TypeError`` out of :func:`_triple` + and would otherwise take down a caller that was only asking a question. + + **This gates the forward and nothing else, and the three gates do not + agree.** A ``stride > 1`` call is served here and by + :func:`~triton_conv3d.reduce_gemm.is_supported_bwd_weight`, and *refused* by + :func:`~triton_conv3d.bwd_data.is_supported_bwd_data`, whose kernel-free + formulation (backward-data as the forward contraction on a flipped weight) + only holds at unit stride. A caller that will differentiate the result must + therefore ask :func:`is_supported_all` instead: a ``True`` from this function + alone builds a graph node whose backward this package cannot answer, and by + then the caller's fallback is gone. A forward-only caller (inference) should + keep asking this one -- the stride support is real, and the combined gate + would take it away. + """ + if groups != 1: + return False + if x.dim() != 5 or w.dim() != 5: + return False + if x.dtype != w.dtype or x.dtype not in _MFMA_KDIM: + return False + # Same device, not merely both on *a* device. Triton launches on the current + # device and dereferences the other pointer anyway; ScaFFold runs four GPUs + # per node, where peer access turns that into another rank's data rather than + # a fault. + if not x.is_cuda or not w.is_cuda or w.device != x.device: + return False + if bias is not None: + # The kernel masks the bias load against ``Cout``, which says nothing + # about how long the bias actually is, and indexes it with an element + # stride of 1. So a short bias reads past the end -- whatever is in + # memory there becomes the bias, ``nan`` if you are lucky -- and a + # stride-2 view of the right length silently applies every other value. + # ``torch.conv3d`` rejects both; so does this. + if (bias.dim() != 1 or int(bias.shape[0]) != int(w.shape[0]) + or bias.dtype != x.dtype or not bias.is_cuda + or bias.device != x.device or bias.stride(0) != 1): + return False + if x.shape[1] != w.shape[1]: + return False + try: + s = _triple(stride, "stride") + p = _triple(padding, "padding") + d = _triple(dilation, "dilation") + except (ValueError, TypeError): + return False + k = tuple(w.shape[2:]) + if any(v < 1 for v in s + d) or any(v < 0 for v in p): + return False + # Degenerate extents. Each of these clears the output-voxel test below and + # then disagrees with torch, which is the asymmetry this predicate exists to + # prevent: a zero-length spatial axis with padding returns a volume of pure + # padding where torch raises; a zero-size kernel returns an output *larger* + # than the input, because ``(in + 2p - d(k-1) - 1)//s + 1`` gains one at + # ``k = 0``; and ``Cin = 0`` returns ``Cout`` channels of zeros where torch + # returns a tensor with no channels at all -- a different shape. ``N = 0`` + # is not here: it agrees with torch (an empty grid, an empty result). + if any(v < 1 for v in x.shape[2:]) or any(v < 1 for v in k): + return False + if w.shape[0] < 1 or w.shape[1] < 1: + return False + # Every output voxel must exist: a kernel wider than the padded input has + # an empty output, which the M-unravel cannot express. + for i in range(3): + eff = d[i] * (k[i] - 1) + 1 + if x.shape[2 + i] + 2 * p[i] < eff: + return False + return True + + +def is_supported_all( + x: torch.Tensor, + w: torch.Tensor, + bias: torch.Tensor | None = None, + stride=1, + padding=0, + dilation=1, + groups: int = 1, +) -> bool: + """Whether **every** direction of this convolution will be served. + + The gate for a caller that is going to differentiate: :func:`is_supported` + and ``bwd_data.is_supported_bwd_data`` and + ``reduce_gemm.is_supported_bwd_weight``, asked about the one call the caller + has in hand and about the gradient it does not have yet. + + It exists because the three direction gates genuinely disagree and the + disagreement is a trap. ``stride > 1`` is supported by the forward (its + output-voxel unravel simply steps by ``s``) and by backward-weight (the + reduction is indexed by the *output* voxel, so a stride is three extra + multiplies), and is not supported by backward-data, which has no kernel of + its own: at unit stride it *is* the forward contraction on a flipped, + channel-transposed weight, and a stride turns that into a scatter into a + sub-lattice. So a training caller that asks only the forward gate gets a + ``True``, builds a graph node, and discovers at ``backward()`` -- when its + own fallback is no longer reachable, because the node is already in the + graph -- that the gradient cannot be computed. + + The direction gates are deliberately left as they are. Narrowing the + forward's to the intersection would agree with backward-data by taking a + capability away from inference, which asks only the forward and for which + strided convolution works today; and there is no single "the backward" + answer to agree with anyway, since backward-weight accepts the stride the + forward does. The asymmetry is a fact about the three kernels; what was + wrong was that a caller had to know it. Now it can ask. + + Total for the same reason :func:`is_supported` is: an argument that cannot + be interpreted is a ``False``, never an exception. The forward's gate runs + first and validates the triples, so the arithmetic below is reached only + with arguments it has already accepted. + + The gradient is passed as a **metadata-only stand-in**: all three predicates + read rank, shape, dtype, device and ``is_cuda`` and never a stride, a value + or a contiguity, so a one-element allocation expanded to the output shape + answers exactly as the real gradient would. ``expand`` gives every dim a + stride of 0, so if a predicate ever grows a stride test it will see those + zeros and answer ``False`` -- a fallback to the caller's other kernel, which + is the safe direction. + """ + if not is_supported(x, w, bias, stride, padding, dilation, groups): + return False + s = _triple(stride, "stride") + p = _triple(padding, "padding") + d = _triple(dilation, "dilation") + k = tuple(int(v) for v in w.shape[2:]) + grad_shape = (int(x.shape[0]), int(w.shape[0])) + tuple( + (int(x.shape[2 + i]) + 2 * p[i] - d[i] * (k[i] - 1) - 1) // s[i] + 1 + for i in range(3) + ) + grad = x.new_empty((1, 1, 1, 1, 1)).expand(grad_shape) + + # Imported here rather than at module scope: both backward modules import + # this one, so a top-level import would be a cycle. By the time this runs + # they are ordinary already-initialized modules. + from .bwd_data import is_supported_bwd_data + from .reduce_gemm import is_supported_bwd_weight + + if not is_supported_bwd_data( + grad, w, tuple(x.shape), stride, padding, dilation, groups + ): + return False + return bool( + is_supported_bwd_weight( + x, tuple(w.shape), grad, stride, padding, dilation, groups + ) + ) + + +def _check_out( + y: torch.Tensor, shape: tuple[int, ...], like: torch.Tensor +) -> None: + """Reject an ``out=`` the kernel would write outside of, or write wrongly. + + Nothing downstream catches either failure. The grid is sized from the + *problem* and not from ``out``, and the store addresses come from + ``out.stride(0/2/3/4)`` with a channel stride of 1 assumed -- so an + undersized buffer is an out-of-bounds device write (1920 elements into a + 128-element allocation, observed, with no error), and an NCDHW buffer is a + full-rate kernel that returns a scrambled answer. + + The shape is compared explicitly rather than inferred from the strides. + ``reduce_gemm._layout_ok`` checks strides alone and cannot see ``Cout`` -- + none of the five channels-last strides depends on it -- so a buffer built + for a different output width has byte-identical strides and passes. + """ + if tuple(y.shape) != tuple(shape): + raise ValueError(f"out= must have shape {tuple(shape)}, got {tuple(y.shape)}") + if y.dtype != like.dtype: + raise ValueError(f"out= must have dtype {like.dtype}, got {y.dtype}") + if y.device != like.device: + raise ValueError(f"out= must be on {like.device}, got {y.device}") + if not y.is_contiguous(memory_format=torch.channels_last_3d): + raise ValueError( + "out= must have channels_last_3d strides -- the store addressing " + f"assumes a channel stride of 1; got {tuple(y.stride())}" + ) + + +def _check_weight_rsck( + wr: torch.Tensor, shape: tuple[int, ...], like: torch.Tensor +) -> None: + """Reject a hoisted weight that is not the one this call needs. + + ``weight_rsck`` supplies every weight *value* the kernel reads -- ``w`` is + then consulted only for its shape -- so a wrong one is a smooth, correctly + shaped, entirely wrong result. That is a live hazard rather than a + "you asked for it": the transform is meant to be cached across calls, and a + cache keyed on the parameter's version is exactly the thing that can go + stale without changing shape. + + Checked against this tensor's own shape, never against ``w``'s strides: + :mod:`~triton_conv3d.bwd_data` deliberately passes a permuted *view* as + ``w`` and supplies the values through here. + """ + if tuple(wr.shape) != tuple(shape): + raise ValueError( + f"weight_rsck= must have shape {tuple(shape)} (kd, kh, kw, Cin, " + f"Cout), got {tuple(wr.shape)}" + ) + if wr.dtype != like.dtype: + raise ValueError(f"weight_rsck= must have dtype {like.dtype}, got {wr.dtype}") + if wr.device != like.device: + raise ValueError(f"weight_rsck= must be on {like.device}, got {wr.device}") + if not wr.is_contiguous(): + raise ValueError( + "weight_rsck= must be contiguous -- the B tile is loaded as a " + f"contiguous vector along Cout; got strides {tuple(wr.stride())}" + ) + + +def _index_dtype(*operands: torch.Tensor): + """``tl.int64`` offsets, and only for the shapes that need them. + + They are not free -- the AMD backend's buffer-load path requires an i32 + offset tensor -- but triton 3.7.1 narrows i64 offsets it can prove safe, so + the cost is paid only where the storage really is over 2 GiB. See PLAN.md + 3.5: storage size, not offset dtype, is the lever, and it is not one the + kernel controls. + + *Every* operand the kernel indexes has to be passed here, the weight + included -- it is the one this decision used to omit, on an assumption + ("weights are never that large") that nothing enforced. ``numel`` is the + right quantity for each: the largest element offset a contiguous operand + sees is ``numel - 1``, and offsets computed for masked-off lanes can exceed + it but are never dereferenced. + """ + return tl.int64 if max(t.numel() for t in operands) > 2**31 - 1 else tl.int32 + + +def conv3d_forward( + x: torch.Tensor, + w: torch.Tensor, + bias: torch.Tensor | None = None, + stride=1, + padding=0, + dilation=1, + groups: int = 1, + *, + config: ConvConfig | None = None, + weight_rsck: torch.Tensor | None = None, + weight_flip: bool = False, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Forward 3-D convolution. Input and output are ``channels_last_3d``. + + **The weight is read where it lies**, decided from its strides. A + ``channels_last_3d`` parameter -- which is what a ScaFFold model's weights + already are, since ``worker.py`` moves the whole model to that format -- + costs *no* weight transform at all. That matters because the transform was + per optimizer step rather than per call: the optimizer dirties every + parameter every step, so no amount of caching removed it. A weight in + PyTorch's *default* layout is still copied, and has to be; see + :func:`_weight_plan`. + + ``weight_rsck`` remains for a caller who has an RSCK buffer already, and is + a wash against reading the parameter (0.96-1.10x per site, both directions). + It is checked rather than trusted, since it supplies every weight value the + kernel reads and ``w`` is then consulted only for its shape. + + ``weight_flip`` consumes the taps in reverse. It exists for + :mod:`~triton_conv3d.bwd_data`, whose gather is the forward's with the taps + flipped: doing it with a constexpr index rather than a ``torch.flip`` copy is + what lets backward-data share the forward's weight buffer instead of + materializing a second one. + + ``out=`` is checked rather than trusted for the same reason as + ``weight_rsck``: the kernel writes it with addressing derived from *this* + call's shapes, so a mismatched one is an out-of-bounds write. See + :func:`_check_out` and :func:`_check_weight_rsck`. + """ + if not is_supported(x, w, bias, stride, padding, dilation, groups): + raise NotImplementedError( + f"unsupported: x={tuple(x.shape)}/{x.dtype} w={tuple(w.shape)} " + f"stride={stride} padding={padding} dilation={dilation} groups={groups}" + ) + sd, sh, sw = _triple(stride, "stride") + pd, ph, pw = _triple(padding, "padding") + dd, dh, dw = _triple(dilation, "dilation") + kd, kh, kw = (int(v) for v in w.shape[2:]) + + # NDHWC is not a preference here, it is the layout the addressing assumes. + x = x.contiguous(memory_format=torch.channels_last_3d) + n, cin, in_d, in_h, in_w = (int(v) for v in x.shape) + cout = int(w.shape[0]) + out_d = (in_d + 2 * pd - dd * (kd - 1) - 1) // sd + 1 + out_h = (in_h + 2 * ph - dh * (kh - 1) - 1) // sh + 1 + out_w = (in_w + 2 * pw - dw * (kw - 1) - 1) // sw + 1 + + y_shape = (n, cout, out_d, out_h, out_w) + if out is None: + # One allocation, already in the layout the kernel stores into. Spelling + # it ``torch.empty(shape).contiguous(memory_format=...)`` allocates NCDHW + # and then copies the whole thing: 2.82 ms against 0.012 ms on a 256 MiB + # output, 235x, on a path a training step takes about 19 times. + y = torch.empty( + y_shape, device=x.device, dtype=x.dtype, + memory_format=torch.channels_last_3d, + ) + else: + y = out + _check_out(y, y_shape, x) + if weight_rsck is not None: + wr = weight_rsck + _check_weight_rsck(wr, (kd, kh, kw, cin, cout), x) + # RSCK is contiguous by the check above, so the strides are exactly these. + plan = (_W_N_CONTIG, wr.stride(2), wr.stride(3), 1) + else: + plan = _weight_plan(w) + if plan is None: + # The only path left that copies the weight; see :func:`_weight_plan`. + wr = to_rsck(w) + plan = (_W_N_CONTIG, wr.stride(2), wr.stride(3), 1) + else: + wr = w + + m_total = n * out_d * out_h * out_w + if config is None: + config = select_config(m_total, cin, cout, (kd, kh, kw), x.dtype) + why = config.validate(x.dtype) + if why is not None: + raise ValueError(f"illegal config {config}: {why}") + + index_dtype = _index_dtype(x, y, wr) + + block_k_count = triton.cdiv(cin, config.BLOCK_K) + grid = ( + triton.cdiv(m_total, config.BLOCK_M) * triton.cdiv(cout, config.BLOCK_N), + ) + _conv3d_fwd_kernel[grid]( + x, wr, y, bias, + n, in_d, in_h, in_w, + out_d, out_h, out_w, + cin, cout, m_total, + x.stride(0), x.stride(2), x.stride(3), x.stride(4), + plan[1], plan[2], plan[3], + y.stride(0), y.stride(2), y.stride(3), y.stride(4), + KD=kd, KH=kh, KW=kw, + SD=sd, SH=sh, SW=sw, + PD=pd, PH=ph, PW=pw, + DD=dd, DH=dh, DW=dw, + BLOCK_M=config.BLOCK_M, BLOCK_N=config.BLOCK_N, BLOCK_K=config.BLOCK_K, + BLOCK_K_COUNT=block_k_count, GROUP_M=config.GROUP_M, + HAS_BIAS=bias is not None, + EVEN_K=(cin % config.BLOCK_K == 0), + EVEN_N=(cout % config.BLOCK_N == 0), + PADDED=(pd > 0 or ph > 0 or pw > 0), + INDEX_DTYPE=index_dtype, + INPUT_PRECISION="ieee", + W_ORDER=plan[0], W_FLIP=bool(weight_flip), + **config.launch_kwargs(), + ) + return y + + +# --------------------------------------------------------------------------- +# ISA verification +# --------------------------------------------------------------------------- + + +def verify_isa( + problem_shape: Sequence[int] | None = None, + direction: str = "fwd", + config: "ConvConfig | None" = None, + padding: int = 1, + kernel: int = 3, + weight_layout: str = "channels_last", +) -> None: # pragma: no cover + """Compile and launch one configuration so its ISA can be inspected. + + Run under ``AMDGCN_ENABLE_DUMP=1`` with a **cold** ``TRITON_CACHE_DIR``: a + cache hit skips the compile and therefore the dump, and an empty grep then + looks exactly like a kernel with no MFMA in it. The other trap is the + mnemonic -- the emitted instruction is ``v_mfma_f32_16x16x16_bf16`` with no + ``_1k`` suffix even though Triton's internal table entry is named ``_1k``, so + grepping for ``_1k`` reports zero on a healthy kernel. + + ``direction="bwd-data"`` runs the same kernel through + :func:`~triton_conv3d.bwd_data.conv3d_backward_data`. It is the *same* + ``@triton.jit`` function, so a reader could reasonably ask why it needs + checking again: because the constexprs differ. Backward-data's ``PADDED`` + is true where the halo'd forward's is false, its ``EVEN_K``/``EVEN_N`` are + computed from the swapped channel widths, and its tile comes from a + different table -- and every one of those changes the code that is emitted. + + ``weight_layout`` selects which of the three B loads is compiled, and it has + to be gated separately for the same reason: ``W_ORDER`` is a constexpr, and + ``channels_last`` (the shipped path, a transposing load) and ``rsck`` (a + hoisted buffer, a straight load) emit different instructions for the operand + that feeds the matrix core. + """ + n, cin, cout, d, h, wd = problem_shape or (1, 64, 64, 32, 64, 64) + k = (kernel, kernel, kernel) + w = torch.randn((cout, cin, *k), device="cuda", dtype=torch.bfloat16) + if weight_layout == "channels_last": + w = w.contiguous(memory_format=torch.channels_last_3d) + elif weight_layout not in ("rsck", "contiguous"): + raise ValueError(f"unknown weight_layout {weight_layout!r}") + rsck = to_rsck(w) if weight_layout == "rsck" else None + if direction == "fwd": + x = torch.randn( + (n, cin, d, h, wd), device="cuda", dtype=torch.bfloat16 + ).contiguous(memory_format=torch.channels_last_3d) + cfg = config or default_config(n * d * h * wd, cin, cout, torch.bfloat16) + y = conv3d_forward(x, w, padding=padding, config=cfg, weight_rsck=rsck) + big = x + elif direction == "bwd-data": + # Local import: bwd_data imports this module, so a top-level import here + # would be a cycle. It is a wrapper over this file's kernel, not a peer. + from .bwd_data import bwd_data_config, conv3d_backward_data + + out = tuple(v + 2 * padding - (kernel - 1) for v in (d, h, wd)) + gy = torch.randn( + (n, cout, *out), device="cuda", dtype=torch.bfloat16 + ).contiguous(memory_format=torch.channels_last_3d) + cfg = config or bwd_data_config( + gy.shape, cin, k, torch.bfloat16, padding=padding + ) + y = conv3d_backward_data( + gy, w, (n, cin, d, h, wd), padding=padding, config=cfg, + weight_rsck=rsck, + ) + big = gy + else: + raise ValueError(f"unknown direction {direction!r}") + torch.cuda.synchronize() + print( + f"ISA-DUMP-CONFIG [{direction}/{weight_layout}] {cfg} cin={cin} cout={cout} " + f"spatial={(d, h, wd)} k={kernel} pad={padding} " + f"x_storage={big.untyped_storage().size()} " + f"y_storage={y.untyped_storage().size()}" + ) diff --git a/triton_conv3d/reduce_gemm.py b/triton_conv3d/reduce_gemm.py new file mode 100644 index 0000000..7e14a14 --- /dev/null +++ b/triton_conv3d/reduce_gemm.py @@ -0,0 +1,1507 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""Backward-weight: a split-K reduction GEMM over the whole output volume. + + dW[co, ci, kd, kh, kw] = sum_{n,d,h,w} dY[n,d,h,w,co] * X[n, d*s+kd*dil-p, ..., ci] + +As a GEMM that is ``M = Cout``, ``N = taps * Cin``, ``K = N*OD*OH*OW`` -- a *tiny* +output reduced over an enormous K (8.4 M at config B's largest site). That is the +transpose of the situation the forward and backward-data face, and it is why this +is the one direction that needs a kernel of its own. + +Why the forward kernel cannot serve this +======================================== + +It *can*, algebraically, and checking that is what justifies the new file. +Swapping the batch and channel axes of both activations turns backward-weight +into a forward convolution:: + + dW^T (Cin, Cout, KD, KH, KW) = conv3d(X^T (Cin, N, ID, IH, IW), + weight = dY^T (Cout, N, OD, OH, OW)) + +with ``N`` as the channel count and the *output volume* as the kernel extent. +``test_bwd_weight.py::test_the_forward_kernel_can_express_backward_weight`` runs +exactly that and checks it bitwise, so this is a measurement and not an argument. + +It is also unusable. ScaFFold runs ``N = 1``, so the reused kernel's ``Cin`` is +1: ``BLOCK_K`` would have to be 16 to reach the matrix core and 15 of every 16 +lanes would be padding. Worse, the forward's reduction loop runs +``KD*KH*KW * ceil(Cin/BLOCK_K)`` iterations, and ``KD*KH*KW`` here is the output +volume -- **8.4 million** trip counts of a six-compare boundary predicate at +config B's ``dec3`` site, with no split-K anywhere. The reuse is correct and +about four orders of magnitude too slow. + +Shape of the kernel +=================== + +The contraction is a "TN" GEMM: both operands have the reduction axis (the +output voxel) *slowest* and the GEMM's M / N axes contiguous, because NDHWC puts +the channel last. So the A tile is loaded ``(BLOCK_K, BLOCK_M)`` and +transposed, which costs one ``tl.trans`` and keeps both global loads +contiguous -- the alternative, addressing A as ``(BLOCK_M, BLOCK_K)``, strides by +``Cout`` down the fast axis and devectorizes every load. + +**Split-K is mandatory, not optional.** ``Cout`` alone is one or two ``BLOCK_M`` +rows, so an unsplit grid is ``taps * ceil(Cin/BLOCK_N)`` programs -- 27 at the +``64 -> 64`` stem, on a device with 228 CUs. :func:`split_count` derives the +split count from the shape. + +**Several taps per tile**, which is the part that took a measurement to find. +The obvious tiling gives each program one tap, so that the tap's spatial shift +stays a scalar addend on the row offset (which is what the forward does, and +what makes NDHWC pay). Written that way, and with its own split count tuned, +the kernel's best configuration at the ``64 -> 64 @ 130^3`` stem is 21% of +roofline and **0.86x of MIOpen**, against 38% and 1.53x for the tiling below. + +The reason is arithmetic intensity: per reduction element a tile loads +``BLOCK_M + BLOCK_N`` values and does ``2*BLOCK_M*BLOCK_N`` flops, so its +intensity is ``BLOCK_M*BLOCK_N/(BLOCK_M+BLOCK_N)`` flops per byte -- 32 at that +stem, where ``Cout`` caps ``BLOCK_M`` at 64 and one tap caps ``BLOCK_N`` at 64. +Every operand is then re-read once per tap: the one-tap kernel issues 14.5 GB of +load requests for a 549 MB working set and runs them at 2.4 TB/s, which is HBM +speed rather than cache speed. + +Widening ``BLOCK_N`` across ``TAP_BLOCK`` taps fixes both halves of that at once. +The upstream gradient is read ``taps/TAP_BLOCK`` times instead of ``taps``, and +the ``TAP_BLOCK`` shifted reads of the input land in the same instruction stream +on overlapping cache lines instead of in unrelated programs on different XCDs. +The cost is that the tap shift is no longer a scalar: it becomes a per-column +addend, hoisted out of the reduction loop, and -- only when the convolution is +padded -- a two-dimensional boundary predicate instead of a one-dimensional one. + +**Padding used to veto the wide tile, and no longer does.** Until 2026-08-05 +:func:`default_bwd_weight_config` dropped to ``TAP_BLOCK=1`` on a padded problem +and :func:`bwd_weight_config` declined a tuned row with ``TAP_BLOCK > 1`` there, +on the strength of that two-dimensional predicate. Both clauses were written +believing they could not fire: the docstring claimed DistConv hands every real +convolution to the backend unpadded, "so the expensive case is the one that does +not occur". DistConv does do that -- but ScaFFold does not route its +convolutions through DistConv any more (``ScaFFold/unet/conv3d.py`` performs the +halo exchange itself, and only on axes that are *genuinely split*), so the +kernel is handed ``padding = (1,1,1)`` at one GPU and ``(0,1,1)`` at two or +four. Measured inside running steps at all four configurations +(``work/triton-conv/review/SHAPE_AUDIT.md``): every ``k = 3`` site is padded, +and the veto fired at **eight of them**. + +It was then raced rather than argued about. On the padded production form of +all 18 affected cells, the wide tile against the pinned one, one interleaved +block per cell with 95% intervals +(``work/triton-conv/review/OPTIMIZATION_ROUND3.md`` §1): the wide tile wins +**18 of 18**, geometric mean **1.946x**, range 1.137x-5.336x. The +two-dimensional predicate is real and it costs something; it costs far less +than the arithmetic intensity the wide tile buys. Both clauses are gone. + +Correctness was never what they protected, and that too is measured rather than +assumed: forcing every tuned ``TAP_BLOCK > 1`` row onto the padded form of its +own channel pair is bitwise exact against an fp64 reference, in both paddings, +including the ``PADDED and ROW_ALIGNED`` corner, in bf16 and fp32 +(``work/triton-conv/bin/padded_tap_probe.py``). It does not overflow LDS and it +does not move the workspace bound. + +Determinism +=========== + +This direction is where ScaFFold's reproducibility is decided. MIOpen serves it +with ``kernel_batched_gemm_xdlops_bwd_weight``, a split-K GEMM using **float +atomics**, which is why the default configuration is not bitwise reproducible +today and why ``more_determinism`` -- which fixes it by disabling far more than +convolution -- costs 640x. + +Two paths, one kernel, one ``ATOMIC`` constexpr: + +* **deterministic** (the default): each split writes its own slice of an fp32 + workspace ``[splits, Cout, taps*Cin]``, and :func:`_reduce_partials_kernel` + sums the splits in index order. No float atomics, a grid and a split count + that are pure functions of the shape, and accumulation order fixed by the + compiled code. This is the mechanism ``ScaFFold/unet/triton_group_norm.py`` + already ships. The claim is the same one that file makes: **bitwise identical + run to run and process to process, for the same input, dtype, shape, device + and tuning config** -- not bitwise against MIOpen, and not across configs. +* **atomic**: ``tl.atomic_add`` into a zeroed fp32 accumulator, i.e. what CK + does. It exists solely to price determinism and is never selected on its own; + a caller has to ask for ``deterministic=False``. + +:func:`split_count` deliberately does not consult free memory, occupancy or a +runtime autotuner. Anything that lets the split count vary between two runs of +the same shape breaks the claim above, and it would break it *intermittently*, +which is worse than breaking it outright. +""" + +from __future__ import annotations + +import dataclasses +import math +from typing import Sequence + +import torch +import triton +import triton.language as tl + +from .gather_gemm import ( + ConvConfig, + _LDS_BYTES, + _MFMA_KDIM, + _pow2_at_most, + _triple, + tune_key, +) + +__all__ = [ + "BwdWeightConfig", + "conv3d_backward_weight", + "is_supported_bwd_weight", + "bwd_weight_config", + "default_bwd_weight_config", + "candidate_bwd_weight_configs", + "split_count", + "workspace_elements", + "grad_weight_empty", + "register_tuned_bwd_weight", +] + + +# --------------------------------------------------------------------------- +# The kernel +# --------------------------------------------------------------------------- + + +@triton.jit +def _conv3d_bwd_weight_kernel( + X, GY, OUT, + # Sizes. ``K_TOTAL`` is ``BATCH * OUT_D * OUT_H * OUT_W``, the reduction + # length; ``K_CHUNK`` is how much of it one split owns. + IN_D, IN_H, IN_W, + OUT_D, OUT_H, OUT_W, + CIN, COUT, K_TOTAL, K_CHUNK, GRID, + # Element strides. The channel stride of X and GY is 1 by construction -- + # that is what NDHWC means -- so it is neither passed nor multiplied by. + stride_xn, stride_xd, stride_xh, stride_xw, + stride_gn, stride_gd, stride_gh, stride_gw, + # Destination: ``[split][Cout][tap][Cin]`` with the last three contiguous, + # so one output channel's whole gradient is ``stride_wo`` long. + stride_ws, stride_wo, + NUM_M: tl.constexpr, NUM_CI: tl.constexpr, NUM_TG: tl.constexpr, + TAPS: tl.constexpr, TAP_BLOCK: tl.constexpr, BLOCK_NC: tl.constexpr, + KD: tl.constexpr, KH: tl.constexpr, KW: tl.constexpr, + SD: tl.constexpr, SH: tl.constexpr, SW: tl.constexpr, + PD: tl.constexpr, PH: tl.constexpr, PW: tl.constexpr, + DD: tl.constexpr, DH: tl.constexpr, DW: tl.constexpr, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_K: tl.constexpr, + PADDED: tl.constexpr, ROW_ALIGNED: tl.constexpr, ATOMIC: tl.constexpr, + NUM_XCD: tl.constexpr, INDEX_DTYPE: tl.constexpr, + INPUT_PRECISION: tl.constexpr, +): + # -- which tile of dW, and which slice of the reduction ------------------ + # + # Split slowest, tiles fastest. Every program in one split reads the *same* + # range of output voxels, so the taps and Cin blocks of a chunk are + # co-resident and their overlapping reads of X hit cache rather than HBM. + # The opposite order (splits fastest) spreads concurrent programs over the + # whole volume and has no reuse at all. + pid = tl.program_id(0) + if NUM_XCD > 1: + # MI300A dispatches workgroups round-robin over its six XCDs, each with + # its own 4 MiB L2, so the tiles of one split -- which read the *same* + # chunk of both activations -- land on six different caches and share + # nothing but the MALL. Remapping the id so that consecutive logical + # tiles are consecutive *within* an XCD puts them back together. This + # is the same fact about this device that makes ``GROUP_M`` want to be a + # multiple of 6 in the gather kernel. + per = GRID // NUM_XCD + rem = GRID % NUM_XCD + xcd = pid % NUM_XCD + seq = pid // NUM_XCD + pid = tl.where(xcd < rem, xcd * (per + 1), + rem * (per + 1) + (xcd - rem) * per) + seq + num_tiles = NUM_M * NUM_CI * NUM_TG + split = pid // num_tiles + tile = pid % num_tiles + pid_m = tile % NUM_M + rest = tile // NUM_M + pid_ci = rest % NUM_CI + tap0 = (rest // NUM_CI) * TAP_BLOCK + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) # output channels + + # -- the N axis: TAP_BLOCK taps x BLOCK_NC input channels --------------- + # + # All of this is hoisted out of the reduction: the column decomposition + # depends on the tile, not on the voxel. ``BLOCK_NC``, ``KH`` and ``KW`` are + # constexpr, so the divisions fold away. + col = tl.arange(0, BLOCK_N) + t_local = col // BLOCK_NC + offs_n = pid_ci * BLOCK_NC + (col % BLOCK_NC) # input channels + # ``tl.arange`` needs a power of two, so ``TAP_BLOCK`` is one and cannot + # divide 27. Rather than mask the load for the ragged last group -- which + # would cost a predicate on every B tile of every group -- the tap is + # *clamped*: those columns read a real, in-bounds tap, compute a value + # nobody wants, and are dropped by the store mask. The waste is one tap in + # 28 at ``TAP_BLOCK=4``; the alternative is a masked load everywhere. + tap_ok = (tap0 + t_local) < TAPS + tap = tl.minimum(tap0 + t_local, TAPS - 1) + kd = tap // (KH * KW) + khw = tap % (KH * KW) + kh = khw // KW + kw = khw % KW + # The column part of the X address: the tap's spatial shift plus the channel. + x_col = ( + (kd * DD).to(INDEX_DTYPE) * stride_xd + + (kh * DH).to(INDEX_DTYPE) * stride_xh + + (kw * DW).to(INDEX_DTYPE) * stride_xw + + offs_n.to(INDEX_DTYPE) + ) + col_ok = (offs_n < CIN) & tap_ok + + k_begin = split * K_CHUNK + k_end = min(k_begin + K_CHUNK, K_TOTAL) + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + + for k0 in range(k_begin, k_end, BLOCK_K): + offs_k = k0 + tl.arange(0, BLOCK_K) + + # -- unravel the fused ndhw voxel index --------------------------- + # + # Unlike the forward, this cannot be hoisted out of the reduction: here + # the *reduction* axis is the volume. ROW_ALIGNED is what keeps it + # cheap. When BLOCK_K divides OUT_W and the chunk is row-aligned a + # K-tile lies inside a single row of the output, so the whole unravel is + # scalar (four SALU divisions) and the only vector term is ``ow``. Every + # real ScaFFold volume has a power-of-two output extent, so this is the + # path that runs in production; the general branch below exists for the + # 8^3 bottleneck (where BLOCK_K > OUT_W) and for the test shapes. + if ROW_ALIGNED: + row = k0 // OUT_W + ow = (k0 - row * OUT_W) + tl.arange(0, BLOCK_K) + oh = row % OUT_H + tmp = row // OUT_H + od = tmp % OUT_D + idn = tmp // OUT_D + else: + ow = offs_k % OUT_W + tmp = offs_k // OUT_W + oh = tmp % OUT_H + tmp = tmp // OUT_H + od = tmp % OUT_D + idn = tmp // OUT_D + + # -- A: the upstream gradient, (BLOCK_K, BLOCK_M) ----------------- + # + # Contiguous along Cout, which is the GEMM's M. Cast per term rather + # than after the sum: at scale 8 a single term overflows int32 and the + # sum would already be wrong before any widening. + g_row = ( + idn.to(INDEX_DTYPE) * stride_gn + + od.to(INDEX_DTYPE) * stride_gd + + oh.to(INDEX_DTYPE) * stride_gh + + ow.to(INDEX_DTYPE) * stride_gw + ) + a_ptrs = GY + g_row[:, None] + offs_m[None, :] + if EVEN_K and EVEN_M: + a = tl.load(a_ptrs) + elif EVEN_K: + a = tl.load(a_ptrs, mask=(offs_m < COUT)[None, :], other=0.0) + elif EVEN_M: + a = tl.load(a_ptrs, mask=(offs_k < k_end)[:, None], other=0.0) + else: + a = tl.load( + a_ptrs, + mask=(offs_k < k_end)[:, None] & (offs_m < COUT)[None, :], + other=0.0, + ) + + # -- B: the input, (BLOCK_K, BLOCK_N) ----------------------------- + # + # The row part is the voxel, the column part is (tap shift, channel). + src_d = od * SD - PD + src_h = oh * SH - PH + src_w = ow * SW - PW + x_row = ( + idn.to(INDEX_DTYPE) * stride_xn + + src_d.to(INDEX_DTYPE) * stride_xd + + src_h.to(INDEX_DTYPE) * stride_xh + + src_w.to(INDEX_DTYPE) * stride_xw + ) + b_ptrs = X + x_row[:, None] + x_col[None, :] + if PADDED: + # Two-dimensional, because the tap now varies down the columns. + # Unpadded, every tap of an in-range output voxel is in range and + # all of this compiles out -- but that is the *rarer* case in + # production, not the common one: ScaFFold's adapter halos only the + # split axis, so every k>1 convolution it issues arrives here with + # PADDED true. See the module docstring. + in_d = src_d[:, None] + (kd * DD)[None, :] + in_h = src_h[:, None] + (kh * DH)[None, :] + in_w = src_w[:, None] + (kw * DW)[None, :] + mask_b = ( + (in_d >= 0) & (in_d < IN_D) + & (in_h >= 0) & (in_h < IN_H) + & (in_w >= 0) & (in_w < IN_W) + ) + if not EVEN_N: + mask_b = mask_b & (offs_n < CIN)[None, :] + if not EVEN_K: + mask_b = mask_b & (offs_k < k_end)[:, None] + b = tl.load(b_ptrs, mask=mask_b, other=0.0) + elif EVEN_K and EVEN_N: + b = tl.load(b_ptrs) + elif EVEN_K: + b = tl.load(b_ptrs, mask=(offs_n < CIN)[None, :], other=0.0) + elif EVEN_N: + b = tl.load(b_ptrs, mask=(offs_k < k_end)[:, None], other=0.0) + else: + b = tl.load( + b_ptrs, + mask=(offs_k < k_end)[:, None] & (offs_n < CIN)[None, :], + other=0.0, + ) + + # ``tl.trans`` rather than a strided A load: see the module docstring. + # ``input_precision`` only bites for fp32 operands, where the backend + # default splits the dot into reduced-precision pieces; bf16 already + # accumulates in fp32. ``more_determinism`` runs in fp32 and has to + # actually be fp32, so it is asked for explicitly. + acc = tl.dot(tl.trans(a), b, acc, input_precision=INPUT_PRECISION) + + # -- epilogue --------------------------------------------------------- + # + # One expression serves three destinations. With one split ``stride_ws`` is + # 0 and ``OUT`` is the real gradient in its own dtype, so the workspace and + # the reduction pass disappear entirely; with several it is an fp32 slice; + # with ATOMIC it is a single fp32 accumulator every split adds into. + out_ptrs = ( + OUT + + split.to(INDEX_DTYPE) * stride_ws + + offs_m.to(INDEX_DTYPE)[:, None] * stride_wo + + (tap * CIN + offs_n)[None, :] + ) + if EVEN_M: + mask_o = tl.broadcast_to(col_ok[None, :], (BLOCK_M, BLOCK_N)) + else: + mask_o = (offs_m < COUT)[:, None] & col_ok[None, :] + if ATOMIC: + tl.atomic_add(out_ptrs, acc, mask=mask_o, sem="relaxed") + else: + tl.store(out_ptrs, acc.to(OUT.dtype.element_ty), mask=mask_o) + + +@triton.jit +def _reduce_partials_kernel( + PARTIAL, OUT, N_ELEM, SPLITS, + BLOCK: tl.constexpr, BLOCK_S: tl.constexpr, +): + """Sum the split-K partials in index order and cast to the output dtype. + + The determinism of the whole direction rests on this loop. ``SPLITS`` is a + runtime argument rather than a constexpr (one compile serves every shape), + but the loop is sequential, its bound is a pure function of the problem, and + ``tl.sum`` over a fixed tile shape is a fixed order -- so two runs of the + same problem add the same numbers in the same order. A tree reduction is + equally reproducible; a ``tl.atomic_add`` is not, which is the whole point. + + ``BLOCK_S`` splits are read at a time rather than one, and that is not a + micro-optimization: the gradient can be *smaller* than one program's tile + (the ``k=1`` head is 384 elements), and a one-split-at-a-time loop is then a + single workgroup paying 900 dependent memory latencies in series. Measured, + that alone made the deterministic path 2.0x slower than the atomic one at + that site -- the only shape where determinism cost anything at all. + """ + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < N_ELEM + acc = tl.zeros((BLOCK,), dtype=tl.float32) + base = offs.to(tl.int64) + stride = N_ELEM.to(tl.int64) + offs_s = tl.arange(0, BLOCK_S) + for s0 in range(0, SPLITS, BLOCK_S): + s = s0 + offs_s + tile = tl.load( + PARTIAL + s.to(tl.int64)[:, None] * stride + base[None, :], + mask=(s < SPLITS)[:, None] & mask[None, :], other=0.0, + ) + acc += tl.sum(tile, axis=0) + tl.store(OUT + offs, acc.to(OUT.dtype.element_ty), mask=mask) + + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class BwdWeightConfig(ConvConfig): + """A launch configuration with the two knobs only this direction has. + + A subclass rather than two more fields on :class:`ConvConfig`, because the + gather directions would never set either and a config printed in a forward + sweep should not grow suffixes it cannot use. Everything else -- the gfx942 + legality rules, the measured LDS model, ``launch_kwargs`` -- is inherited + unchanged, and ``BLOCK_N`` keeps its meaning as the *full* tile width so + that both of those stay correct. + """ + + #: Split-K partial count. ``0`` means "derive it from the shape", which is + #: what the shipped path does; a non-zero value pins it so that a sweep can + #: see the shape of the curve rather than one point on it. + SPLIT_K: int = 0 + #: How many taps one tile spans. ``BLOCK_N = TAP_BLOCK * BLOCK_NC``. + TAP_BLOCK: int = 1 + #: XCD count to swizzle the program id for; 0 or 1 disables it. MI300A has + #: six, and it is the same device fact that makes ``GROUP_M`` want to be a + #: multiple of six in the gather kernel. **Measured at 1.006x** -- inside + #: the round-to-round spread -- so it buys nothing; it is on by default only + #: because every number in M3_RESULTS was taken with it on, and turning it + #: off would make those numbers describe a kernel nobody ran. + NUM_XCD: int = 6 + + @property + def BLOCK_NC(self) -> int: + """Input channels per tap in the tile.""" + return self.BLOCK_N // self.TAP_BLOCK + + def __str__(self) -> str: + return ( + super().__str__() + + (f"/tb{self.TAP_BLOCK}" if self.TAP_BLOCK != 1 else "") + + (f"/x{self.NUM_XCD}" if self.NUM_XCD != 6 else "") + + (f"/sk{self.SPLIT_K}" if self.SPLIT_K else "") + ) + + def validate(self, dtype: torch.dtype) -> str | None: + why = super().validate(dtype) + if why is not None: + return why + if self.SPLIT_K < 0: + return "SPLIT_K must be non-negative (0 means derive from shape)" + if self.TAP_BLOCK < 1: + return "TAP_BLOCK must be at least 1" + if self.BLOCK_N % self.TAP_BLOCK: + return f"BLOCK_N must be a multiple of TAP_BLOCK={self.TAP_BLOCK}" + if self.NUM_XCD < 0: + return "NUM_XCD must be non-negative" + return None + + +#: MI300A's compute units. It appears here rather than being read from the +#: device on purpose: the split count has to be a pure function of the *problem* +#: for the determinism claim to hold, so a device with a different CU count gets +#: a differently-tuned kernel rather than a differently-ordered reduction. +_CU_COUNT = 228 + +#: Waves of programs the split count aims for. Four, measured -- but the whole +#: *number* matters more than the value. Every program in this kernel does the +#: same amount of work, so a grid of 4.5 waves runs five and idles through half +#: of the last one, and the measured split-count curve is dominated by that +#: quantization rather than by anything about the reduction. In one interleaved +#: block at ``64 -> 64 @ 130x258x258``: 512 splits (8.98 waves) 8.05 ms, 596 +#: (10.46 waves) 8.79 ms, 384 (6.70) 8.92 ms, 256 (4.49) 9.50 ms -- monotone in +#: how fractional the wave count is and not in the parallelism. 228 splits, +#: which is 4.00 waves exactly, measures 7.66-7.73 ms in two further runs. +#: Hence the snap in :func:`split_count`, which is worth more than the target. +_SPLIT_TARGET_WAVES = 4 + +#: A split must be worth its epilogue. Every split writes a full +#: ``BLOCK_M x BLOCK_N`` fp32 tile and the reduction reads every one of them +#: back, so at a short reduction the partials become the dominant traffic: the +#: rule is that they stay under this fraction of the main loop's. Without it, +#: ``512 -> 1024 @ 10x18x18`` (a 2048-voxel reduction into a 14 M-element +#: gradient) asks for three splits and pays 1.25x for them. +_MAX_EPILOGUE_FRACTION = 10 + +#: And a split must be at least a few K-tiles long, or the loop's own prologue +#: is most of it. +_MIN_K_TILES_PER_SPLIT = 4 + +#: Ceiling on the fp32 partial workspace. The product to watch is +#: ``splits * Cout * taps * Cin * 4``: one split at ``1024 -> 1024`` is 113 MiB, +#: though that site has a tiny K and needs no splitting at all. Bounded here so +#: that no shape in the corpus can ask for an allocation that fails at step 400. +_WORKSPACE_BYTES = 256 * 1024 * 1024 + + +def _fit_bwd_weight_to_lds(cfg: BwdWeightConfig, dtype: torch.dtype) -> BwdWeightConfig: + """Shrink ``BLOCK_K`` until the operand tiles fit in LDS. + + Only ``BLOCK_K`` moves. ``BLOCK_M`` and ``BLOCK_N`` are already bounded by + ``Cout`` and ``TAP_BLOCK * Cin`` -- shrinking either throws away the + arithmetic intensity this kernel is short of -- whereas ``BLOCK_K`` is the + reduction depth and costs only reuse. That is the same ordering + :func:`~triton_conv3d.gather_gemm._fit_to_lds` uses and the same reason. + """ + kdim = _MFMA_KDIM.get(dtype, {}).get(cfg.matrix_instr_nonkdim) + if kdim is None: + return cfg + while cfg.lds_bytes(dtype) > _LDS_BYTES and cfg.BLOCK_K // 2 >= kdim: + cfg = dataclasses.replace( + cfg, BLOCK_K=cfg.BLOCK_K // 2, + kpack=1 if cfg.BLOCK_K // 2 <= 16 else cfg.kpack, + ) + warps = max(1, min(cfg.num_warps, cfg.BLOCK_M * cfg.BLOCK_N // 256)) + return dataclasses.replace(cfg, num_warps=1 << (warps.bit_length() - 1)) + + +def default_bwd_weight_config( + cout: int, cin: int, kernel: Sequence[int], k_total: int, + dtype: torch.dtype = torch.bfloat16, *, padded: bool = False, +) -> BwdWeightConfig: + """A config that is legal for any shape and close to tuned for most. + + The tile is bounded by the *problem*, not chosen freely: ``BLOCK_M`` cannot + usefully exceed ``Cout``, and one tap's N extent is exactly ``Cin``. That is + the opposite of the forward, where M is the volume and a tall tile is always + available -- and it is why ``TAP_BLOCK`` exists. + + ``TAP_BLOCK`` brings the tile *width* up to a target that the sweep put at + 256 columns -- or 512 when ``Cout`` is 64 or less, where the tile has no + height to trade against. Two measured facts sit behind that rule. Where + ``Cin >= 256`` the channels reach 256 columns unaided and adding taps on top + **costs** 5-24%; where ``Cout = 64`` the taps are the only way to get there + at all and they are worth 1.7-1.9x. + + ``TAP_BLOCK`` used to be forced to 1 whenever the convolution was padded, + because the boundary predicate becomes two-dimensional then (see the module + docstring). **That was the path production took**, at every site with + ``k > 1`` and every configuration: ScaFFold's adapter exchanges a halo only + on the axis it actually splits, so H and W keep the module's ``padding = 1`` + and an unsharded run keeps all three. The clause was documented as + unreachable ("no real ScaFFold convolution is padded -- DistConv halos them + all"), which was true of the MIOpen rung and false of the shipped one; it + cost a step-level projection a factor of three before anyone checked + (``work/triton-conv/review/SHAPE_AUDIT.md``). + + It is gone as of 2026-08-05, on a measurement rather than on the + observation that its premise was false. Raced on the padded production + form of the six channel pairs that reach it, the widened tile against the + pinned one, one interleaved block per cell with 95% intervals + (``work/triton-conv/review/opt3/defwide_adapter.json``): widening wins + **6 of 6**, 1.263x-2.084x, every interval clear of 1.000. At + ``64 -> 128`` the widened *heuristic* beats that pair's tuned row as well + (2.084x against 1.867x at ``66x128^2``), which is why the table now carries + the wider tile there. + + ``padded`` is therefore accepted and **not consulted**. It is kept in the + signature because it describes the problem rather than the policy, every + caller already computes it, and the next rule that wants it should not have + to re-thread it through eight call sites -- but nothing here branches on it + today, and a reader should not have to run the function to learn that. + """ + k = _triple(kernel, "kernel") + taps = math.prod(k) + # 256, not 128: at ``Cout >= 256`` a 128-row tile leaves half the available + # M on the floor, and the tile's arithmetic intensity + # ``BLOCK_M*BLOCK_N/(BLOCK_M+BLOCK_N)`` goes from 85 flops/byte at + # ``128x256`` to 128 at ``256x256``. Measured 1.11-1.56x over the shipped + # 128-row tile at every ``Cout >= 256`` channel pair in the corpus bar two; + # see ``_TUNED_BWD_W``. ``_pow2_at_most`` keeps ``BLOCK_M <= Cout``, which + # genuinely binds at ``Cout = 128`` -- ``128x512`` and ``128x1024`` were both + # tried there and are dead ends (1.07x and 0.13-0.24x). + block_m = _pow2_at_most(cout, 256) + block_nc = _pow2_at_most(cin, 256) + # ``TAP_BLOCK`` is a power of two because ``BLOCK_N`` has to be one + # (``tl.arange`` refuses anything else), so it never divides 27 exactly; the + # ragged last group is handled by clamping in the kernel. + target_width = 512 if block_m <= 64 else 256 + tap_block = 1 + while tap_block * 2 <= taps and block_nc * tap_block * 2 <= target_width: + tap_block *= 2 + block_n = block_nc * tap_block + nonkdim = 16 + kdim = _MFMA_KDIM[dtype][nonkdim] + # The deepest K-tile that still leaves the operands inside LDS. Deeper is + # better for the loop's own overhead and does nothing for the intensity, so + # it is the axis that gives way -- and it is the only one that can, since + # BLOCK_M and BLOCK_N are pinned to the problem above. + itemsize = torch.empty((), dtype=dtype).element_size() + # ...but only to 32 once the tile is 256 rows tall, and that cap is + # measured, not a guess about registers. ``256x256x64`` is a **cliff** at + # the long-reduction sites: ``512 -> 256 @ 34x66x66`` runs 8.00 ms against + # the 128-row tile's 3.57 (0.45x) and ``512 -> 256 @ 18x66x66`` 4.21 against + # 1.76 (0.42x), with an identical program count, so it is not parallelism. + # ``256x256x32`` is never worse than 0.93x anywhere I measured and is + # 1.11-1.20x where the tall tile pays. The heuristic derives ``BLOCK_K`` + # from the LDS budget, which at ``256 + 256`` columns lands on exactly 64 in + # bf16 -- i.e. straight into the cliff -- so the cap has to be explicit. + # The tuned table is free to ship 64 where a measurement says so. + block_k = _pow2_at_most(_LDS_BYTES // (itemsize * (block_m + block_n)), + 32 if block_m >= 256 else 64) + block_k = max(kdim, block_k - block_k % kdim) + # A K axis shorter than one tile is not an error, only waste; shrink so the + # tiny synthetic shapes do not run a mostly-masked reduction. + while block_k > kdim and block_k > k_total: + block_k //= 2 + return _fit_bwd_weight_to_lds( + BwdWeightConfig( + BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=block_k, GROUP_M=6, + num_warps=min(8, max(1, block_m * block_n // 256)), + num_stages=2, matrix_instr_nonkdim=nonkdim, + kpack=1 if block_k <= 16 else 2, TAP_BLOCK=tap_block, + ), + dtype, + ) + + +def _row_aligned(block_k: int, out_w: int) -> bool: + """Whether a K-tile is guaranteed to lie inside one row of the output.""" + return out_w % block_k == 0 + + +def split_count( + cfg: BwdWeightConfig, cout: int, cin: int, taps: int, k_total: int, + out_w: int, +) -> tuple[int, int]: + """``(splits, chunk)``: how the reduction axis is divided, and by how much. + + A **pure function of the shape and the config**, which is what makes the + deterministic path reproducible process to process. Two things pull it up + and three cap it, and both directions are measured rather than assumed: + + * up: enough programs to fill the device four times over, which is the whole + reason split-K is here; + * down: the epilogue must stay a minority of the traffic + (:data:`_MAX_EPILOGUE_FRACTION`), a split must be a few K-tiles long, and + the workspace ceiling must hold -- the last is what stops ``1024 -> 1024`` + (113 MiB per split) from asking for gigabytes, and that site needs no + splitting anyway, so it is a guard rather than a compromise. The first + two bound the split count the *target* asks for; only the workspace + ceiling survives the wave snap below, and the comment there says why; + * and then the result is **snapped to a whole number of waves**, which is + the step that actually matters. Every program here does the same work, so + a grid of 4.5 waves runs five and idles through half of the last; the + measured split-count curve is mostly that sawtooth, and reading it as "the + cache prefers short chunks" -- which is what it looks like if the snap is + missing -- leads to picking 2.6x more splits than the shape wants. + + ``chunk`` is rounded up to a whole number of output *rows* when the tile is + row-aligned, because the kernel's cheap scalar unravel needs every K-tile to + stay inside one row; otherwise to a whole number of K-tiles. + """ + tiles = (-(-cout // cfg.BLOCK_M) + * -(-cin // cfg.BLOCK_NC) + * -(-taps // cfg.TAP_BLOCK)) + k_tiles = -(-k_total // cfg.BLOCK_K) + per_split = cout * taps * cin * 4 + ceiling = max(1, min(k_tiles // _MIN_K_TILES_PER_SPLIT, + _WORKSPACE_BYTES // per_split)) + + if cfg.SPLIT_K: + want = min(cfg.SPLIT_K, ceiling) + else: + want = -(-_SPLIT_TARGET_WAVES * _CU_COUNT // tiles) + # The partials are written once and read once, in fp32; the loop reads + # ``BLOCK_M + BLOCK_N`` operand elements per reduction element per tile. + # Where the gradient is huge and the volume small -- ``512 -> 1024 @ + # 10x18x18`` is a 2048-voxel reduction into 14 M elements -- this is the + # bound that bites, and without it that site pays 1.25x for splits it + # cannot use. + loop_elems = tiles * k_total * (cfg.BLOCK_M + cfg.BLOCK_N) + epilogue_elems = cout * taps * cin * 2 * 2 + epilogue_bound = max( + 1, loop_elems // (_MAX_EPILOGUE_FRACTION * max(1, epilogue_elems))) + want = max(1, min(want, epilogue_bound, ceiling)) + # The snap goes last and **outranks the epilogue bound**, which is a + # deliberate ordering and not the oversight it looks like. ``round`` + # here can only move ``want`` up: where ``tiles * want`` is under half a + # wave it gives 0, ``max(1, ...)`` forces one whole wave, and ``want`` + # becomes ``_CU_COUNT // tiles``, which can be several times the + # epilogue-bounded value. That is the right trade, because the two + # costs are not the same size. Re-applying the epilogue bound after the + # snap was implemented and measured, and it loses: at + # ``128 -> 256 @ 34^3`` it takes 16 splits to 7, i.e. a 98-program grid + # on 228 CUs, and the site goes 0.2536 -> 0.4563 ms (**1.80x**); at + # ``256 -> 512 @ 10x34x34`` 108 programs against 216 and 0.2731 -> + # 0.6553 ms. Half an idle device costs more than a doubled epilogue + # whenever the grid is that small -- and the shapes where the snap + # overrides the bound are exactly the shapes where the grid is that + # small, because that is the condition under which ``round`` rounds to + # zero. The **workspace** ceiling is different in kind (an allocation + # that fails is not a slow kernel) and is re-applied. + waves = max(1, round(tiles * want / _CU_COUNT)) + want = max(1, min(waves * _CU_COUNT // tiles, ceiling)) + want = max(1, want) + + align = out_w if _row_aligned(cfg.BLOCK_K, out_w) else cfg.BLOCK_K + chunk = -(-(-(-k_total // want)) // align) * align + return -(-k_total // chunk), chunk + + +#: Seed tiles, ``(BLOCK_M, BLOCK_NC, TAP_BLOCK, BLOCK_K, num_warps)``. Not the +#: forward's grid: there M is the volume and the useful tiles are tall, here M is +#: ``Cout``, one tap of N is ``Cin``, and the volume lives in ``BLOCK_K`` and in +#: the split count. So the axes worth sweeping are ``TAP_BLOCK`` and +#: ``BLOCK_K``, neither of which the forward's seed grid varies at all. +_SEED_TILES: tuple[tuple[int, int, int, int, int], ...] = ( + (64, 64, 1, 64, 4), + (64, 64, 1, 128, 4), + (64, 64, 1, 256, 4), + (64, 64, 2, 64, 4), + (64, 64, 2, 128, 8), + (64, 64, 4, 32, 4), + (64, 64, 4, 64, 8), + (64, 64, 8, 16, 4), + (64, 64, 8, 32, 8), + (64, 128, 1, 64, 4), + (64, 128, 1, 128, 4), + (64, 128, 2, 32, 4), + (64, 128, 2, 64, 8), + (64, 128, 4, 32, 8), + (128, 64, 1, 64, 4), + (128, 64, 1, 128, 8), + (128, 64, 2, 32, 4), + (128, 64, 2, 64, 8), + (128, 64, 4, 32, 8), + (128, 128, 1, 64, 8), + (128, 128, 1, 128, 8), + (128, 128, 2, 32, 8), + (128, 128, 2, 64, 8), + (128, 256, 1, 64, 8), + (256, 128, 1, 64, 8), + (256, 64, 1, 64, 4), + # Tall *and* wide. The grid used to top out at 128 columns for every + # ``BLOCK_M=256`` entry, so at the ``Cout >= 256`` sites -- where the M axis + # has the room -- the tile that wins was never timed at all and the sweep + # reported a tie it had not actually measured. Both ``BLOCK_K`` are here + # because the choice between them is not a preference: 64 is 1.11-1.56x at + # the short-reduction sites and 0.42-0.45x at the long ones. + (256, 256, 1, 32, 8), + (256, 256, 1, 64, 8), + (32, 64, 1, 128, 4), + (32, 64, 4, 64, 4), + (16, 64, 1, 128, 4), + (16, 64, 4, 64, 4), +) + +#: Split counts worth trying. 0 means "let :func:`split_count` decide", which is +#: what the shipped path does. +_SEED_SPLITS: tuple[int, ...] = (0, 1, 4, 16, 64, 256) + + +def candidate_bwd_weight_configs( + cout: int, cin: int, kernel: Sequence[int], k_total: int, + dtype: torch.dtype = torch.bfloat16, + *, splits: Sequence[int] = _SEED_SPLITS, padded: bool = False, +) -> list[BwdWeightConfig]: + """Configs worth timing for one shape, already pruned to legal ones. + + Pruned rather than shrunk, for the reason M2 gives: shrinking an oversized + tile folds two seed entries onto one config and silently double-counts it in + a best-of sweep. + """ + taps = math.prod(_triple(kernel, "kernel")) + m2 = max(16, triton.next_power_of_2(cout)) + c2 = max(16, triton.next_power_of_2(cin)) + k2 = max(16, triton.next_power_of_2(k_total)) + out: list[BwdWeightConfig] = [] + seen: set[BwdWeightConfig] = set() + for bm, bnc, tb, bk, seed_warps in _SEED_TILES: + # A BLOCK_M past Cout is pure padding (M is Cout, not a volume), a + # BLOCK_NC past Cin is padding for the same reason, and a BLOCK_K past + # the whole reduction is padding too. + if bm > m2 or bnc > c2 or bk > k2 or tb > taps: + continue + for warps in {4, 8, seed_warps}: + for sk in splits: + cfg = BwdWeightConfig( + BLOCK_M=bm, BLOCK_N=bnc * tb, BLOCK_K=bk, GROUP_M=6, + num_warps=warps, num_stages=2, matrix_instr_nonkdim=16, + kpack=1 if bk <= 16 else 2, SPLIT_K=sk, TAP_BLOCK=tb, + ) + if (cfg.validate(dtype) is not None + or cfg.lds_bytes(dtype) > _LDS_BYTES + or cfg in seen): + continue + seen.add(cfg) + out.append(cfg) + if not out: + out.append(default_bwd_weight_config(cout, cin, kernel, k_total, dtype, + padded=padded)) + return out + + +def _tuned(bm: int, bnc: int, tb: int, bk: int, warps: int, + sk: int = 0, nk: int = 32) -> BwdWeightConfig: + """One measured row. + + ``nk`` defaults to **32 here and to 16 everywhere else in the package**, and + that asymmetry is the whole point of it. See :data:`_TUNED_BWD_W`. + """ + return BwdWeightConfig(BLOCK_M=bm, BLOCK_N=bnc * tb, BLOCK_K=bk, GROUP_M=6, + num_warps=warps, num_stages=2, + matrix_instr_nonkdim=nk, + kpack=1 if bk <= 16 else 2, SPLIT_K=sk, TAP_BLOCK=tb) + + +#: Measured backward-weight winners, keyed by ``(dtype, Cin, Cout, kernel)`` -- +#: the convolution a reader would name, as in the other two tables, even though +#: this direction's GEMM has ``Cout`` on M and ``taps * Cin`` on N. +#: +#: Source: ``work/triton-conv/m3_bwd.json``. Only channel pairs that were +#: actually timed appear; a miss falls to :func:`default_bwd_weight_config` plus +#: :func:`split_count`, which is a real gap and not an extrapolation dressed up +#: as a measurement. +#: +#: ``SPLIT_K`` is left at 0 -- "derive from the shape" -- in every entry, and +#: that is a finding rather than an omission. The same channel pair occurs at +#: volumes three orders of magnitude apart, the split count is the one knob that +#: genuinely has to follow the volume, and pinning a sweep's winner would carry +#: one volume's answer to every other. See M3_RESULTS 5. +#: +#: **``matrix_instr_nonkdim`` is 32 in this table and 16 in every other table in +#: the package.** That is not an inconsistency, it is the measurement. This +#: direction's ``tl.dot(tl.trans(a), b)`` lowers on gfx942 to an *element-wise* +#: transpose of A through LDS -- 128 two-byte ``ds_read_u16`` per loop body +#: against 16 ``ds_read_b128`` for the untransposed operand, at a structural +#: 50.00% bank-conflict rate -- which leaves the LDS pipe ~65% busy while the +#: matrix core idles at ~32% (``work/triton-conv/review/BWDW_OPTIMIZATION.md`` §1, counters). +#: The 32x32x8 fragment halves the MFMA instruction count for identical FLOPs +#: and so stops the MFMA stream competing with that transpose for issue slots. +#: It does *not* remove the transpose: the ``nk32`` ISA still emits 128 +#: ``ds_read_u16``. The forward has no transposed operand, and ``nonkdim=16`` +#: won 15 of 15 there over a 1090-config sweep; backward-data reuses the forward +#: kernel and 32 won exactly one pair by 0.3%. **The rule is per direction and +#: does not generalise** -- which this table then has to say a second time, +#: because it does not generalise across channel pair either. +#: +#: Raced per *volume* against the identical tile at ``nonkdim=16``, one +#: interleaved block per site, 27 sites covering every ``k=3, Cin >= 64`` cell +#: in the corpus bar the 2 GiB cliff (``work/triton-conv/tune/nk_bwdw.json``): +#: +#: ============ ========================================= ====== +#: pair gain by volume shipped +#: ============ ========================================= ====== +#: (64,64) 1.042 @130^3, 1.026 @66x258^2, 1.023 @130x258^2 nk32 +#: (128,64) 1.015 @130^3, 1.012 @66x258^2 nk32 +#: (64,128) 0.992 @66^3, 1.012 @34x130^2, 1.012 @66x130^2 nk32 +#: (128,128) 1.105 @66^3, 1.124 @34x130^2, 1.155 @66x130^2 nk32 +#: (256,128) 1.056 @66^3, 1.061 @34x130^2, 1.065 @66x130^2 nk32 +#: (128,256) 1.051 @34^3, 1.101 @34x66^2, 1.096 @18x66^2 nk32 +#: (256,256) 1.083 @34^3, 1.086 @34x66^2, 1.083 @18x66^2 nk32 +#: (512,256) 1.031 @34^3, 1.054 @34x66^2, 1.052 @18x66^2 nk32 +#: (256,512) 1.012 @18^3, 1.043 @18x34^2, 1.036 @10x34^2 nk32 +#: (512,512) 1.028 @18^3, 1.045 @18x34^2, 1.040 @10x34^2 nk32 +#: (1024,512) 1.016 @18^3, 1.038 @18x34^2, 1.026 @10x34^2 nk32 +#: (512,1024) 1.088 @10^3, **0.935** @6x18^2, **0.969** @10x18^2 nk16 +#: (1024,1024) 1.069 @10^3, **0.963** @6x18^2, 0.997 @10x18^2 nk16 +#: ============ ========================================= ====== +#: +#: The last two rows are why this was raced per volume rather than per pair. +#: ``BWDW_OPTIMIZATION.md`` §3.1 reports ``nk32`` winning **13 of 13** and +#: "never losing"; every one of those 13 was a ``Cout <= 512`` site, and at +#: ``Cout = 1024`` the sign reverses at two of the three volumes each pair has. +#: Both pairs keep ``nonkdim=16``. The transposed operator's weight gradient +#: runs this same kernel and was raced too (7 sites, +#: ``work/triton-conv/tune/nk_convT.json``): 0.934-1.056x, no consistent sign, +#: so :func:`default_bwd_weight_config` -- which is the only thing serving it, +#: and also serves fp32 and every untuned pair, including the four +#: 2048-channel backward-weight sites of a scale-8 run -- **stays at 16**. +#: (Until 2026-08-05 it also served the eight padded production sites whose +#: tuned row widens ``TAP_BLOCK``, because the resolver declined those rows; +#: it no longer does, and those eight now run the table.) A heuristic is the +#: path with no measurement +#: behind it; putting a knob +#: there whose sign is shape-dependent is exactly the extrapolation this table +#: refuses to make elsewhere. +#: +#: Results are **not** bitwise identical to the ``nonkdim=16`` kernel -- a +#: different MFMA fragment sums the same products in a different order. They +#: are still bitwise *reproducible*, which is what the determinism claim says: +#: the split-K partition, the reduction's tiling and its ``tl.sum`` order are +#: all unchanged, and ``work/triton-conv/review/bwd/det_verify.py`` was re-run on this table. +_TUNED_BWD_W: dict[tuple, BwdWeightConfig] = { + # The segmentation head, which until now had no row at all and no ``nk`` + # question either: ``BLOCK_M`` is ``Cout = 6`` rounded up to 16, so + # ``nonkdim=32`` is illegal here and the axis above simply does not apply. + # What does apply is ``num_warps``, which no sweep in this project has ever + # taken below 4 (``candidate_bwd_weight_configs`` draws it from + # ``{4, 8, seed}``). Raced at all three head volumes, 1 warp is 1.044x, + # 1.052x and 1.033x over the heuristic's 4 + # (``work/triton-conv/tune/warps_bwdw.json``). It is still a **loss** + # against MIOpen at two of those three volumes (0.73x, 0.93x, 1.19x); the + # row is here so the number the adapter's block-list is built from is the + # best this kernel can do, not the best it happened to be doing. + tune_key(torch.bfloat16, 64, 6, (1, 1, 1)): _tuned(16, 64, 1, 64, 1, nk=16), + # The one transposed row. ``conv_transpose3d_backward_weight`` calls this + # module with the operator's widths **swapped** -- it passes the transposed + # weight's own ``(Cin, Cout, k, k, k)`` shape, whose first axis is this + # reduction's M -- so the key below reads ``(64, 128)`` and the module a + # reader would name is ``ConvTranspose3d(128, 64, 2, stride=2)``. + # + # It is the *only* one of the four transposed channel pairs where + # ``nonkdim=32`` wins, and it wins at all three of that pair's volumes: + # 1.128x @ 64^3, 1.112x @ 32x128^2, 1.069x @ 64x128^2. The other three + # pairs measure 0.926-1.056x with no consistent sign and stay on + # ``default_bwd_weight_config``, i.e. at 16 + # (``work/triton-conv/tune/nk_convT.json``). The split is the mechanism, + # not luck: this pair is the only transposed site whose tile is + # ``BLOCK_M = 128``; the other three reach ``BLOCK_M = 256``, which already + # amortises the transpose, and that is the same boundary the ``k=3`` rows + # show. Everything except ``nonkdim`` here restates what the heuristic + # already picks, so the row cannot drift away from it silently. + tune_key(torch.bfloat16, 64, 128, (2, 2, 2)): _tuned(128, 64, 4, 64, 8), + **{ + tune_key(torch.bfloat16, cin, cout, (3, 3, 3)): cfg + for (cin, cout), cfg in { + # The UNet stem, and the row that refutes ``OPTIMIZATION.md`` D4. + # It is the forward's disease one axis over: ``BLOCK_NC = + # _pow2_at_most(Cin, 256) = 16`` against ``Cin = 3``, times + # ``TAP_BLOCK = 16`` covering 27 taps in two groups, is **512 issued + # columns for 81 useful -- 6.32x**. ``BLOCK_NC = 4`` (which is what + # ``bnc=4, tb=16`` spells, keeping ``BLOCK_N = 64``) is 3 live of 4 + # and 128 columns for 81, **1.58x**. That is the structural optimum + # for this axis and no kernel change can beat it: a dense N would + # need ``BLOCK_N = 96``, which ``tl.arange`` cannot express, and + # ``BLOCK_N = 128`` dense is the same 1.58x. + # + # Raced against the heuristic's ``64x256x64/nk16/w8/tb16`` and + # MIOpen, one interleaved block per volume, kernel-only: + # 130^3 0.3279 ms vs MIOpen 0.6034 -- 1.842x [1.838,1.846] + # 130x258^2 1.3738 ms vs MIOpen 2.5105 -- 1.825x [1.817,1.834] + # 66x258^2 0.6980 ms vs MIOpen 1.2924 -- 1.854x [1.846,1.863] + # i.e. 3.21-3.43x over the config it replaces, which was + # 0.536-0.577x of MIOpen. + # + # ``nk=16``, not this table's 32, and measured rather than inherited: + # every ``nk32`` twin is 1-5% behind. ``num_warps=1`` beats 2 by + # 1.08x and 4 by 1.04x -- no sweep in this project had ever taken the + # backward-weight warps below 4. ``BLOCK_M`` below ``Cout`` costs + # 1.84x (``32x64x64/tb16`` 0.6031 ms) and a deeper K is a 2.3x loss + # (``64x64x256/tb16`` 0.7616); ``64x128x64/tb16``, i.e. ``BLOCK_NC = + # 8``, gets only half the win at 0.5071. + # + # Bitwise **identical** to the config it replaces on random operands + # (0 of 5184 gradient elements differ): ``split_count`` returns the + # same 456 splits and the same 9.018 MiB workspace at all three + # volumes, so the reduction tree does not move, and ``BLOCK_N`` / + # ``BLOCK_NC`` / ``TAP_BLOCK`` only decide which columns a program + # owns, never the order a column is summed in. So no determinism + # baseline moves and the 168.8 MiB workspace bound is unchanged. + # ``work/triton-conv/review/LOSS_OPTIMIZATION.md`` §1.2. + (3, 64): _tuned(64, 4, 16, 64, 1, nk=16), + # Cout = 64 is where TAP_BLOCK earns its existence: the tile can only + # get to 512 columns through the taps, and getting there is worth + # 1.7-1.9x against the one-tap form. + (64, 64): _tuned(64, 64, 8, 16, 4), + # ``(64, 128)`` is the one pair of the three whose ``Cout`` has room + # for a 128-row tile, and it wants one. Raced against the + # ``64x512x16/tb8`` row it replaces, one interleaved block per cell, + # kernel-only, 95% intervals -- **both** forms, because this table + # is shared and the padded form is not the only caller + # (``review/opt3/r64x128_{adapter,distconv}.json``): + # padded 66x128^2 1.6928 vs 1.9896 ms -- 1.177x [1.165,1.189] + # padded 34x128^2 0.8760 vs 1.0036 -- 1.145x [1.140,1.151] + # padded 64^3 0.4277 vs 0.5059 -- 1.183x [1.180,1.185] + # unpadded 66x130^2 1.7061 vs 1.8609 -- 1.088x [1.083,1.093] + # unpadded 34x130^2 0.9087 vs 0.9563 -- 1.052x [1.046,1.057] + # unpadded 66^3 0.4448 vs 0.5052 -- 1.136x [1.133,1.140] + # 6 of 6, every interval clear of 1.000. ``nk=32`` is worth a + # further 1.05-1.10x over the ``nk16`` twin of the same tile, which + # is this table's usual sign; ``BLOCK_K=32`` is 0.91-0.97x and is + # not taken. **This row moves the unpadded form too**, so a stored + # DistConv number for ``64->128`` backward-weight is superseded. + (64, 128): _tuned(128, 64, 4, 64, 8), + (128, 64): _tuned(64, 64, 8, 16, 4), + (128, 128): _tuned(128, 128, 2, 64, 8), + (128, 256): _tuned(128, 128, 2, 64, 8), + # From Cin = 256 up, the channels alone reach a 256-column tile and the + # taps are not needed for it; TAP_BLOCK > 1 then *costs* 5-24%, because + # a wider tile past 256 buys less than the register pressure takes. + (256, 128): _tuned(128, 256, 1, 64, 8), + (512, 256): _tuned(128, 256, 1, 64, 8), + # ``BLOCK_M = 256`` wherever ``Cout`` has the room. Every entry below + # was raced against the 128-row tile it replaces at *every* volume the + # corpus has for that channel pair, interleaved, on an idle GPU 1 -- + # per-pair, not per-site, because that is what this table is keyed on. + # Gains, worst volume first: + # (256,256) 1.11x @ 34x66x66, 1.15x @ 18x66x66, 1.20x @ 34^3 + # (256,512) 1.13x @ 18x34x34, 1.14x @ 10x34x34, 1.17x @ 18^3 + # (512,512) 1.14x @ 10x34x34, 1.16x @ 18x34x34, 1.21x @ 18^3 + # (1024,512) 1.11x @ 18x34x34, 1.21x @ 10x34x34, 1.56x @ 18^3 + # (512,1024) 1.04x @ 10^3, 1.14x @ 10x18x18, 1.16x @ 6x18x18 + # (1024,1024)1.18x @ 10^3, 1.33x @ 6x18x18 and 10x18x18 + # The two ``10^3`` cells are the thin ones: at 0.09 and 0.19 ms they are + # the smallest sites in the corpus, the tile has barely a wave of work to + # do, and a 16-deep K-tile is 1.10x and 1.07x better still than the entry + # shipped here. They are not tuned for, because the pair's other two + # volumes are 3-7x larger and want 64. + # ``(1024, 512)`` was ``_tuned(64, 64, 4, 32, 4)`` -- ``BLOCK_M = 64`` at + # ``Cout = 512``, a plain table bug and the reason M3 recorded that site + # as its worst k=3 cell at 0.72x of MIOpen. It is 1.11x of MIOpen now. + # + # ``BLOCK_K`` is 32 at ``(256, 256)`` and 64 elsewhere, and that is a + # measurement rather than an oversight: at ``(256, 256)`` the 64-deep + # tile is 4-8% behind at all three volumes, while at the sites below 64 + # is 5-15% ahead. See ``default_bwd_weight_config`` for the cliff that + # makes the heuristic refuse 64 at this width. + (256, 256): _tuned(256, 256, 1, 32, 8), + (256, 512): _tuned(256, 256, 1, 64, 8), + (512, 512): _tuned(256, 256, 1, 64, 8), + (1024, 512): _tuned(256, 256, 1, 64, 8), + # ``nk=16``: the two ``Cout = 1024`` pairs, and the only rows here that + # keep the forward's value. See the table above -- 0.935x and 0.963x at + # their ``6x18x18`` volumes. + (512, 1024): _tuned(256, 256, 1, 64, 8, nk=16), + (1024, 1024): _tuned(256, 256, 1, 64, 8, nk=16), + }.items() + }, +} +# ``(128, 256)`` and ``(512, 256)`` keep their 128-row tile, and that is a +# result rather than an omission -- both were raced with ``BLOCK_M = 256`` at +# ``BLOCK_K`` 16, 32 and 64 at every volume the corpus has for them: +# +# * ``(128, 256)``: 1.00x, 1.00x, 0.99x. A genuine tie. ``Cin = 128`` forces +# ``TAP_BLOCK = 2`` to reach 256 columns, and the tap traffic eats the +# intensity the taller tile buys. +# * ``(512, 256)``: 1.19x at ``34^3``, but **0.93x and 0.94x** at the two +# ``out_w = 66`` volumes -- which are 3.5x and 1.8x larger, so the pair is a +# net loss, and this table is keyed per channel pair and cannot split one by +# volume. ``256x256x64`` there is the cliff named in +# :func:`default_bwd_weight_config`: 0.45x and 0.42x. + + +def register_tuned_bwd_weight(dtype, cin, cout, kernel, config) -> None: + _TUNED_BWD_W[tune_key(dtype, cin, cout, tuple(kernel))] = config + + +def bwd_weight_config( + cout: int, cin: int, kernel: Sequence[int], k_total: int, + dtype: torch.dtype = torch.bfloat16, *, padded: bool = False, +) -> BwdWeightConfig: + """The config :func:`conv3d_backward_weight` would pick for this problem. + + ``padded`` no longer selects a different row, and that is a measurement. + + Until 2026-08-05 this function declined a tuned row whose ``TAP_BLOCK`` was + greater than 1 whenever the convolution was padded, on the argument that + the boundary predicate becomes two-dimensional there. ``padded`` is not a + rare argument -- every ScaFFold convolution with ``k > 1`` reaches here with + ``padded=True``, at every configuration -- so that clause fired at **eight + sites over six channel pairs**, and what ran instead was + :func:`default_bwd_weight_config` with ``TAP_BLOCK`` pinned to 1. + + Raced arm against arm on the padded production form of all 18 affected + cells (``work/triton-conv/review/OPTIMIZATION_ROUND3.md`` §1; the tuned row + forced against the config the decline produced, in one interleaved block + per cell, CUDA-graph replay, 95% intervals): the tuned row is faster in + **18 of 18**, geometric mean **1.946x**, range 1.137x-5.336x, every + interval clear of 1.000. Worst cell was the stem, ``3->64 k3 @ 130x256^2``, + at 7.9505 ms declined against 1.4910 ms with the row. + + The two-dimensional predicate is real and it is not free -- it is simply + much cheaper than the arithmetic intensity ``TAP_BLOCK`` buys, which is + 1.7-1.9x at ``Cout = 64`` where the tile has no height to trade. + + That race forces the row; what the *shipped* resolver then delivers was + measured separately, over the whole corpus in the production shape, at the + commit before this clause was removed and at the commit after + (``work/triton-conv/review/HEAD_MEASUREMENT.md`` §3). **1.813x** geometric + mean over the 21 backward-weight cells of those six channel pairs, + 105.5 ms summed to 57.5 ms, worst cell the stem at 8.0000 ms to 1.523 ms. + The forced-row 1.946x above is the counterfactual and this is the delivered + figure; quote this one for the shipped kernel. + + The set the clause never reached is the control, and it is flat: the other + 36 backward-weight cells move 1.0015x, the forward 1.0043x and + backward-data 1.0032x, against a floor of 0.9999x measured by repeating one + 171-cell capture back to back on an idle device. So the effect is a + backward-weight change on six channel pairs and nothing else. + + Two consequences worth carrying. **The production shape no longer costs + anything in this direction**: measured against the halo'd, unpadded shape + the same problems take on the MIOpen rung, production backward-weight was + 2.019x over those cells and is now **1.021x** (1.3444x to **1.0037x** over + all 42 form-sensitive cells). And **the step falls at every + configuration**: +9.89 ms [9.01, 10.77] at scale 7 on one GPU, and +83.85 + [82.19, 85.51], +45.53 [38.50, 52.55] and +20.69 [14.36, 27.03] at scale 8 + on one, two and four -- paired alternating arms on an idle node, 5/5/8/10 + pairs, all four significant. + + Determinism is unaffected in the sense the package claims it: a different + config is a different reduction order, so results are **not** bitwise equal + to what the declined path produced, but they remain bitwise reproducible + run to run and process to process, because :func:`split_count` is still a + pure function of the shape and the config. See the module docstring. + + That inequality is **measured**, not inferred, and it is wider than the + split count suggests (``work/triton-conv/review/verify3/ulp_bf16.json``, one + ordinary draw per cell, the six pairs at both production paddings). In bf16 + the old and new configs disagree at **10 of 12 cells**, but + :func:`split_count` moves at only 5 of them: the other five move because the + row changes ``matrix_instr_nonkdim`` from 16 to 32, and a different MFMA + fragment sums the same products in a different order within one ``BLOCK_K``. + The one pair that stays bitwise identical, ``3 -> 64``, is the one whose new + row keeps both ``nonkdim`` and ``BLOCK_K``. So the thing to check before + asserting two configs agree is not the split count alone. The disagreement + is small: at most 1 ULP at eight of the ten cells, 4 and 20 ULP at the other + two, on 6-52 elements of 110k-885k, and every result stays inside + :func:`~triton_conv3d.reference.error_bound`. + """ + k = _triple(kernel, "kernel") + tuned = _TUNED_BWD_W.get(tune_key(dtype, cin, cout, tuple(k))) + if tuned is not None: + return _fit_bwd_weight_to_lds(tuned, dtype) + return default_bwd_weight_config(cout, cin, k, k_total, dtype, padded=padded) + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + + +def workspace_elements( + splits: int, cout: int, cin: int, kernel: Sequence[int] +) -> int: + """fp32 elements :func:`conv3d_backward_weight` needs for ``splits`` splits.""" + return splits * cout * cin * math.prod(_triple(kernel, "kernel")) + + +def grad_weight_empty( + cout: int, cin: int, kernel: Sequence[int], *, dtype, device +) -> torch.Tensor: + """An empty gradient in the layout the kernel writes: ``[Cout][tap][Cin]``. + + That is exactly ``channels_last_3d`` for a ``(Cout, Cin, kd, kh, kw)`` + tensor -- its memory order is ``Cout, kd, kh, kw, Cin`` -- so the natural + output of this GEMM is already the memory format ScaFFold runs in. The + forward and backward-data both need a weight transform on the way *in*; this + direction gets the layout for free on the way out. + + ``memory_format=`` on the allocation rather than ``.contiguous(...)`` after + it, and that is not cosmetic: ``torch.empty(shape).contiguous(memory_format= + channels_last_3d)`` allocates the tensor in NCDHW and then runs a permuting + device copy to reach the layout it was going to be asked for anyway. The + contents are undefined either way, so the copy transports nothing; measured + on the forward's identically-shaped defect it is **235x** the cost of the + one-shot allocation. + """ + k = _triple(kernel, "kernel") + return torch.empty( + (cout, cin, *k), dtype=dtype, device=device, + memory_format=torch.channels_last_3d, + ) + + +def _validate_out( + gw: torch.Tensor, cout: int, cin: int, k: tuple[int, int, int], + dtype: torch.dtype, device: torch.device, +) -> str | None: + """``None`` if ``gw`` can be written as this problem's gradient, else why not. + + Shaped like :meth:`ConvConfig.validate` because it is the same kind of + guard: nothing downstream looks at ``gw`` again. The reduction pass is + launched with ``n_elem = Cout * taps * Cin`` derived from ``weight_shape`` + and stores that many elements into ``gw`` whatever ``gw`` actually is, and + the one-split fast path hands the tile kernel ``gw``'s data pointer + directly. + + All four clauses have teeth: + + * **shape**, which the stride comparison below *cannot* see. None of the + five strides depends on ``Cout``, so a gradient allocated for ``Cout=8`` + with the same ``Cin`` and kernel is stride-identical to one allocated for + ``Cout=64``. Passing the small one to the large problem used to be + accepted, and wrote 55 296 elements into a 6 912-element allocation -- + past the end of the buffer, into whatever the caching allocator handed out + next. No fault and no exception; some other live tensor is simply wrong + later. + * **strides**: the reduction pass treats both the workspace and the + destination as flat ``[Cout][tap][Cin]`` arrays, so an ``out=`` in the + default contiguous layout would be filled with a correctly shaped, + *transposed* answer. + * **device**: the kernel launches on the current device and dereferences + whatever pointer it is given. ScaFFold runs four GPUs per node, and with + peer access enabled a foreign pointer does not fault -- it scribbles on + another rank. + * **dtype**: the epilogue casts to ``OUT.dtype.element_ty``, so a + mismatched ``out=`` returns a gradient in a dtype the caller's optimizer + is not expecting rather than raising, and at an integer dtype the cast + truncates. + """ + want_shape = (cout, cin, *k) + if tuple(gw.shape) != want_shape: + return ( + f"shape must be {want_shape} (Cout x Cin x kernel); got " + f"{tuple(gw.shape)} -- the strides alone cannot tell these apart, " + "because none of them depends on Cout" + ) + if gw.device != device: + return f"device must be {device}; got {gw.device}" + if gw.dtype != dtype: + return f"dtype must match the operands' {dtype}; got {gw.dtype}" + taps = k[0] * k[1] * k[2] + want = (taps * cin, 1, k[1] * k[2] * cin, k[2] * cin, cin) + got = tuple(gw.stride()) + # An extent of 1 makes its stride unobservable, so only compare the ones + # that can be told apart -- ``k=1`` is a real corpus shape. + if not all(g == w for g, w, n in zip(got, want, want_shape) if n > 1): + return ( + "must have channels_last_3d strides -- the reduction pass treats it " + f"as [Cout][tap][Cin]; want {want}, got {got}" + ) + return None + + +def is_supported_bwd_weight( + input: torch.Tensor, + weight_shape: Sequence[int], + grad_output: torch.Tensor, + stride=1, + padding=0, + dilation=1, + groups: int = 1, +) -> bool: + """Whether :func:`conv3d_backward_weight` will serve this call. + + Same asymmetry as the other two predicates: the caller's fallback is MIOpen, + which is correct everywhere, so a false negative costs a little speed and a + false positive returns a wrong gradient. + + Unlike backward-data this direction has **no** stride restriction. The + reduction axis is the *output* voxel and the input coordinate + ``o*s + t*dil - p`` is a function of it, so a stride is three extra + multiplies rather than a scatter into a sub-lattice. It is supported and + tested, though ScaFFold's corpus never uses one on a non-transposed + convolution. + """ + if groups != 1: + return False + if input.dim() != 5 or grad_output.dim() != 5 or len(tuple(weight_shape)) != 5: + return False + if input.dtype != grad_output.dtype or input.dtype not in _MFMA_KDIM: + return False + # Same device, not merely both on *a* device. Triton launches on the current + # device and dereferences the other pointer anyway; ScaFFold runs four GPUs + # per node, where peer access turns that into another rank's data rather than + # a fault -- i.e. a plausible wrong gradient instead of a crash. The same + # clause is in ``gather_gemm.is_supported``; the two gates are kept symmetric + # deliberately, because the caller picks between them by direction and a hole + # in one of them is a hole in the ladder. + if (not input.is_cuda or not grad_output.is_cuda + or grad_output.device != input.device): + return False + try: + s = _triple(stride, "stride") + p = _triple(padding, "padding") + d = _triple(dilation, "dilation") + except ValueError: + return False + cout, cin, *k = (int(v) for v in weight_shape) + if any(v < 1 for v in s + d + tuple(k)) or any(v < 0 for v in p): + return False + if cout < 1 or cin < 1: + return False + n, in_c, *in_sp = (int(v) for v in input.shape) + if in_c != cin or int(grad_output.shape[1]) != cout: + return False + if int(grad_output.shape[0]) != n: + return False + for i in range(3): + eff = d[i] * (k[i] - 1) + 1 + if in_sp[i] + 2 * p[i] < eff: + return False + if int(grad_output.shape[2 + i]) != (in_sp[i] + 2 * p[i] - eff) // s[i] + 1: + return False + return True + + +def conv3d_backward_weight( + input: torch.Tensor, + weight_shape: Sequence[int], + grad_output: torch.Tensor, + stride=1, + padding=0, + dilation=1, + groups: int = 1, + *, + deterministic: bool = True, + config: BwdWeightConfig | None = None, + workspace: torch.Tensor | None = None, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Gradient of a 3-D convolution with respect to its weight. + + ``input`` and ``grad_output`` are ``channels_last_3d``; the returned + gradient is ``channels_last_3d`` too, which for a weight is the layout this + GEMM produces natively (see :func:`grad_weight_empty`). The argument order + mirrors ``torch.nn.grad.conv3d_weight``. + + ``deterministic`` defaults to **True**, which is a deliberate divergence + from PLAN.md 3.4's ``deterministic=None`` following + ``torch.are_deterministic_algorithms_enabled()``. The atomic path exists to + price determinism, not to be selected; if it were the default whenever + torch's flag is off then every ScaFFold run would silently get the + nonreproducible one, which is the state this milestone exists to end. A + caller who wants it has to say so. + + ``workspace`` hoists the fp32 partial buffer out of the call, the way + ``weight_rsck`` hoists the weight transform in the other two directions. + :func:`split_count` and :func:`workspace_elements` say how big it must be. + """ + if not is_supported_bwd_weight(input, weight_shape, grad_output, stride, + padding, dilation, groups): + raise NotImplementedError( + f"unsupported: input={tuple(input.shape)}/{input.dtype} " + f"weight_shape={tuple(weight_shape)} " + f"grad_output={tuple(grad_output.shape)}/{grad_output.dtype} " + f"stride={stride} padding={padding} dilation={dilation} " + f"groups={groups}" + ) + sd, sh, sw = _triple(stride, "stride") + pd, ph, pw = _triple(padding, "padding") + dd, dh, dw = _triple(dilation, "dilation") + cout, cin, *k = (int(v) for v in weight_shape) + kd, kh, kw = k + taps = kd * kh * kw + padded = pd > 0 or ph > 0 or pw > 0 + + # NDHWC is not a preference, it is the layout the addressing assumes. + x = input.contiguous(memory_format=torch.channels_last_3d) + gy = grad_output.contiguous(memory_format=torch.channels_last_3d) + n, _, in_d, in_h, in_w = (int(v) for v in x.shape) + out_d, out_h, out_w = (int(v) for v in gy.shape[2:]) + k_total = n * out_d * out_h * out_w + + if out is None: + gw = grad_weight_empty(cout, cin, k, dtype=x.dtype, device=x.device) + else: + gw = out + why = _validate_out(gw, cout, cin, (kd, kh, kw), x.dtype, x.device) + if why is not None: + raise ValueError(f"out= is not usable for this problem: {why}") + + if config is None: + config = bwd_weight_config(cout, cin, k, k_total, x.dtype, padded=padded) + why = config.validate(x.dtype) + if why is not None: + raise ValueError(f"illegal config {config}: {why}") + + splits, chunk = split_count(config, cout, cin, taps, k_total, out_w) + num_m = triton.cdiv(cout, config.BLOCK_M) + num_ci = triton.cdiv(cin, config.BLOCK_NC) + num_tg = triton.cdiv(taps, config.TAP_BLOCK) + grid = (num_m * num_ci * num_tg * splits,) + + n_elem = cout * taps * cin + atomic = not deterministic + if atomic or splits > 1: + need = n_elem * (1 if atomic else splits) + if workspace is None: + ws = torch.empty(need, dtype=torch.float32, device=x.device) + else: + if workspace.numel() < need or workspace.dtype is not torch.float32: + # Say the size, not just that this one is wrong: a caller who + # hoists the workspace sizes it once, out of the step, from a + # number in a document -- and the number in the document was + # understated by 1.46x for a year, so the first thing they see + # is this message at step 1. ``workspace_elements(splits, ...)`` + # with ``split_count``'s own ``splits`` is the supported way to + # get here; the arithmetic is repeated in the text so that a + # traceback alone is enough to fix the call. + raise ValueError( + f"workspace must be at least {need} float32 elements " + f"({need * 4 / 2**20:.1f} MiB) -- {splits} splits x " + f"{cout} Cout x {taps} taps x {cin} Cin; got " + f"{workspace.numel()} of {workspace.dtype} " + f"({workspace.numel() * workspace.element_size() / 2**20:.1f}" + " MiB)" + ) + ws = workspace + if atomic: + # CK's own shape: an fp32 accumulator every split adds into, zeroed + # first. The zeroing and the cast below are part of the atomic + # path's cost and are timed as such -- excluding them would price + # determinism against a variant that does not exist. + ws[:n_elem].zero_() + dest, stride_ws = ws, (0 if atomic else n_elem) + else: + # One split: the kernel writes the answer straight out in its own dtype + # and neither the workspace nor the reduction pass exists. Worth the + # branch -- at ``1024 -> 1024`` a round trip through fp32 partials is + # 113 MiB read plus 57 MiB written against a 0.4 ms kernel. + dest, stride_ws = gw, 0 + + big = max(x.numel(), gy.numel()) > 2**31 - 1 + index_dtype = tl.int64 if big else tl.int32 + + _conv3d_bwd_weight_kernel[grid]( + x, gy, dest, + in_d, in_h, in_w, + out_d, out_h, out_w, + cin, cout, k_total, chunk, grid[0], + x.stride(0), x.stride(2), x.stride(3), x.stride(4), + gy.stride(0), gy.stride(2), gy.stride(3), gy.stride(4), + stride_ws, taps * cin, + NUM_M=num_m, NUM_CI=num_ci, NUM_TG=num_tg, + TAPS=taps, TAP_BLOCK=config.TAP_BLOCK, BLOCK_NC=config.BLOCK_NC, + KD=kd, KH=kh, KW=kw, + SD=sd, SH=sh, SW=sw, + PD=pd, PH=ph, PW=pw, + DD=dd, DH=dh, DW=dw, + BLOCK_M=config.BLOCK_M, BLOCK_N=config.BLOCK_N, BLOCK_K=config.BLOCK_K, + EVEN_M=(cout % config.BLOCK_M == 0), + EVEN_N=(cin % config.BLOCK_NC == 0 and taps % config.TAP_BLOCK == 0), + EVEN_K=(k_total % config.BLOCK_K == 0 and chunk % config.BLOCK_K == 0), + PADDED=padded, + ROW_ALIGNED=_row_aligned(config.BLOCK_K, out_w), + ATOMIC=atomic, + NUM_XCD=config.NUM_XCD, + INDEX_DTYPE=index_dtype, + INPUT_PRECISION="ieee", + **config.launch_kwargs(), + ) + if atomic or splits > 1: + # A narrower tile when the whole gradient is small, so that the grid is + # not one program: at ``Cout=6, k=1`` the gradient is 384 elements. + block = min(1024, max(64, triton.next_power_of_2(n_elem))) + _reduce_partials_kernel[(triton.cdiv(n_elem, block),)]( + ws, gw, n_elem, 1 if atomic else splits, + BLOCK=block, BLOCK_S=8, num_warps=4, + ) + return gw + + +# --------------------------------------------------------------------------- +# ISA verification +# --------------------------------------------------------------------------- + + +def verify_isa_bwd_weight( + problem_shape: Sequence[int] | None = None, + config: BwdWeightConfig | None = None, + padding: int = 1, + kernel: int = 3, + deterministic: bool = True, +) -> None: # pragma: no cover + """Compile and launch one configuration so its ISA can be inspected. + + Run under ``AMDGCN_ENABLE_DUMP=1`` with a **cold** ``TRITON_CACHE_DIR``; a + cache hit skips the compile and the empty grep that follows is + indistinguishable from a kernel with no MFMA in it. Grep ``v_mfma``, not + ``v_mfma.*_1k``: the emitted mnemonic has no ``_1k`` suffix. + + ``padding`` defaults to **1**, matching :func:`~triton_conv3d.gather_gemm. + verify_isa`: ``PADDED`` is a ``constexpr``, so it selects a different kernel + body, and every production ScaFFold convolution with ``k > 1`` compiles the + padded one. The default used to be 0, which inspected the ISA of a kernel + no ScaFFold site launches. + """ + n, cin, cout, d, h, w = problem_shape or (1, 64, 64, 32, 64, 64) + k = (kernel, kernel, kernel) + out = tuple(v + 2 * padding - (kernel - 1) for v in (d, h, w)) + x = torch.randn((n, cin, d, h, w), device="cuda", dtype=torch.bfloat16 + ).contiguous(memory_format=torch.channels_last_3d) + gy = torch.randn((n, cout, *out), device="cuda", dtype=torch.bfloat16 + ).contiguous(memory_format=torch.channels_last_3d) + k_total = n * out[0] * out[1] * out[2] + cfg = config or bwd_weight_config(cout, cin, k, k_total, torch.bfloat16, + padded=padding > 0) + splits, chunk = split_count(cfg, cout, cin, kernel ** 3, k_total, out[2]) + gw = conv3d_backward_weight(x, (cout, cin, *k), gy, padding=padding, + config=cfg, deterministic=deterministic) + torch.cuda.synchronize() + print( + f"ISA-DUMP-CONFIG [bwd-weight] {cfg} cin={cin} cout={cout} " + f"spatial={(d, h, w)} k={kernel} pad={padding} splits={splits} " + f"chunk={chunk} det={deterministic} " + f"row_aligned={_row_aligned(cfg.BLOCK_K, out[2])} " + f"x_storage={x.untyped_storage().size()} " + f"gw_storage={gw.untyped_storage().size()}" + ) diff --git a/triton_conv3d/reference.py b/triton_conv3d/reference.py new file mode 100644 index 0000000..e36cc0e --- /dev/null +++ b/triton_conv3d/reference.py @@ -0,0 +1,403 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""Reference implementations and the tolerance policy. + +The hard part of testing a reduced-precision kernel is deciding what "correct" +means. Three standards are used here, in decreasing order of strictness: + +1. **Exact.** Inputs drawn so that every partial sum is exactly representable in + the working dtype (small integers, bounded reduction length). The kernel must + then match the reference *bitwise*. This is the standard that catches + indexing, masking and boundary bugs, which a tolerance would hide -- a kernel + that reads the wrong voxel usually reads a plausible one. + +2. **No worse than the incumbent.** Error against an fp64 reference must not + exceed MIOpen's error on the same problem by more than a small factor. This + is the honest bar for a replacement: we do not have to be better, but we must + not be worse, and it adapts automatically to shape and reduction length. + +3. **Absolute tolerance.** A dtype- and K-derived bound, used where an fp64 + reference is impractical. Weakest, and only a backstop. + +Everything takes and returns NCDHW tensors in PyTorch's usual convention; the +NDHWC memory format is a layout question, not a semantic one, and is handled by +``contiguous(memory_format=...)`` at the boundary. +""" + +from __future__ import annotations + +import dataclasses +import math + +import torch +import torch.nn.functional as F + +from .shapes import ConvProblem, Direction + +_TORCH_DTYPE = { + "fp32": torch.float32, + "bf16": torch.bfloat16, + "fp16": torch.float16, +} + +#: Mantissa bits, including the implicit leading one. +_MANTISSA_BITS = {torch.float64: 53, torch.float32: 24, torch.bfloat16: 8, + torch.float16: 11} + + +def torch_dtype(problem: ConvProblem) -> torch.dtype: + return _TORCH_DTYPE[problem.dtype] + + +def unit_roundoff(dtype: torch.dtype) -> float: + """One half ulp, relative -- the classic ``u`` of error analysis.""" + return 2.0 ** -_MANTISSA_BITS[dtype] + + +# --------------------------------------------------------------------------- +# Operand construction +# --------------------------------------------------------------------------- + + +def make_inputs( + problem: ConvProblem, + device: torch.device | str = "cuda", + *, + seed: int = 0, + exact: bool = False, + channels_last: bool = True, + dtype: torch.dtype | None = None, + density: float | None = None, +) -> dict[str, torch.Tensor]: + """Input, weight, bias and upstream gradient for one problem. + + With ``exact=True`` the values are small integers chosen so that every + partial sum of the contraction is exactly representable in ``dtype``; see + :func:`is_exactly_representable` for when that is possible and + :func:`exact_density` for the knob that makes it possible at real widths. + + ``density`` thins the *activations* -- ``input`` and ``grad_output`` -- to + that fraction of nonzeros, and only has an effect under ``exact=True``. It + exists because at a real ScaFFold channel width the dense ``{-1,0,1}`` draw + is not exactly representable in bf16 at all: the forward reduces over + ``Cin * taps``, which is 27 648 terms at ``Cin = 1024`` regardless of how + small the volume is made, and a sum of that many random signs runs to a few + hundred while bf16 holds integers only to 256. Thinning is the one lever + that shortens the *realized* reduction without touching the shape, so the + channel widths, the tile selection and the 512-byte row strides under test + all stay exactly as ScaFFold runs them. + + The weight is deliberately left dense. Every one of the ``K`` gather + addresses then contributes to every output element, so a wrong address is + masked only by the sparsity of the value it happens to read -- independently + per element, over millions of them. Thinning the weight instead would + multiply whole ``(tap, Cin)`` rows by an exact zero for a whole output + channel, which is a hole in precisely the coverage this draw exists for. + """ + dtype = dtype or torch_dtype(problem) + device = torch.device(device) + gen = torch.Generator(device=device).manual_seed(seed) + thin = exact and density is not None and density < 1.0 + + def draw(shape: tuple[int, ...], offset: int, + activation: bool = False) -> torch.Tensor: + g = torch.Generator(device=device).manual_seed(seed + offset) + if exact: + # {-1, 0, 1}: products are exact and sums stay small. + t = torch.randint(-1, 2, shape, generator=g, device=device, + dtype=torch.int8).to(dtype) + if thin and activation: + # A separate stream, offset far enough that it cannot collide + # with any operand's *value* stream: those are seed + 0..3, and + # a mask drawn from one of them would correlate the zeros with + # the signs of another tensor. + gm = torch.Generator(device=device).manual_seed( + seed + offset + (1 << 20)) + t = t * (torch.rand(shape, generator=gm, device=device) + < density).to(dtype) + else: + t = torch.randn(shape, generator=g, device=device, dtype=torch.float32) + t = t.to(dtype) + return t + + del gen + fmt = torch.channels_last_3d if channels_last else torch.contiguous_format + out: dict[str, torch.Tensor] = { + "input": draw(problem.input_shape, 0, True).contiguous(memory_format=fmt), + "weight": draw(problem.weight_shape, 1).contiguous(memory_format=fmt), + "grad_output": draw(problem.output_shape, 2, True).contiguous(memory_format=fmt), + } + out["bias"] = draw((problem.cout,), 3) if problem.bias else None + return out + + +def exact_density( + problem: ConvProblem, + direction: Direction = "fwd", + *, + dtype: torch.dtype | None = None, + headroom: float = 4.0, +) -> float: + """Activation density that keeps the realized result inside the mantissa. + + The arithmetic. Draw the activations from ``{-1,0,1}`` and then zero all + but a fraction ``q`` of them, against a dense ``{-1,0,1}`` weight. Each + product then has variance ``(4/9) q``, so a reduction over ``K`` terms has + standard deviation ``(2/3) sqrt(qK)``, and the largest of ``M*N`` such sums + is about ``sqrt(2 ln(M*N))`` deviations out. Setting that equal to + ``2**mantissa / headroom`` and solving for ``q`` gives what is returned. + + Two things fall out that are worth stating. ``q*K`` -- the number of terms + that actually contribute to an output element -- comes out at a few hundred + and is nearly independent of ``K``, so the draw is not "mostly zeros" in the + sense that matters: every output element is still a sum of hundreds of + genuine gathers, and 96% of them are nonzero. And the shape is untouched, + which is the whole point -- this is what lets the forward's bitwise standard + run at ``Cin = 1024`` instead of skipping there, which is where it had no + coverage at all. + + ``headroom`` is against the order-statistic estimate, which is an estimate: + measured over the eleven corpus forward shapes the realized maxima land at + 58-71 against bf16's limit of 256, so 4.0 buys a genuine 3.6-4.4x rather + than a nominal 4x. Returns 1.0 -- no thinning at all -- wherever the dense + draw already fits, so a caller can pass this unconditionally. + """ + dtype = dtype or torch_dtype(problem) + m, n, k = problem.gemm_shape(direction) + if k <= 0 or m * n <= 0: + return 1.0 + limit = 2 ** _MANTISSA_BITS[dtype] / headroom + spread = math.sqrt(2.0 * math.log(max(m * n, 2))) + return min(1.0, (1.5 * limit / spread) ** 2 / k) + + +def is_exactly_representable(result: torch.Tensor, dtype: torch.dtype) -> bool: + """Whether every value in an fp64 reference survives ``dtype`` unchanged. + + With operands in ``{-1, 0, 1}`` every product is exact and every partial sum + is an integer, so the only question is whether the *realized* magnitudes fit + in the mantissa. Asking that of the actual result rather than of the + worst-case reduction length matters a great deal: bf16 has 8 mantissa bits, + so a worst-case bound rejects any reduction longer than 256 and would skip + almost the whole corpus, while a sum of a few hundred random signs is in + practice tens. The bitwise standard is the only one that reliably catches an + off-by-one gather, so it is worth keeping applicable. + """ + finite = result[torch.isfinite(result)] + if finite.numel() == 0: + return True + integral = torch.equal(finite, finite.round()) + return bool(integral and finite.abs().max().item() < 2 ** _MANTISSA_BITS[dtype]) + + +# --------------------------------------------------------------------------- +# References +# --------------------------------------------------------------------------- + + +def _conv(problem: ConvProblem, x, w, b): + op = F.conv_transpose3d if problem.transposed else F.conv3d + return op(x, w, b, stride=problem.stride, padding=problem.padding) + + +def reference( + problem: ConvProblem, + operands: dict[str, torch.Tensor], + direction: Direction = "fwd", + *, + dtype: torch.dtype = torch.float64, + device: torch.device | str | None = None, +) -> torch.Tensor: + """The trusted answer, computed in ``dtype`` (fp64 by default). + + fp64 3-D convolution has no fast path on any backend, so this is slow by + construction -- it is for correctness, not for benchmarking. Callers that + need it at ScaFFold's real sizes should reach for ``direction``-specific + tiling or simply use a smaller problem; every bug we are trying to catch + reproduces at small sizes. + """ + device = torch.device(device) if device is not None else operands["input"].device + x = operands["input"].to(device=device, dtype=dtype) + w = operands["weight"].to(device=device, dtype=dtype) + b = operands["bias"] + b = b.to(device=device, dtype=dtype) if b is not None else None + + if direction == "fwd": + return _conv(problem, x, w, b) + + gy = operands["grad_output"].to(device=device, dtype=dtype) + # Ask for only the gradient wanted. fp64 convolution has no fast path on + # any backend, so the unwanted one is not a rounding error in the test + # suite's runtime -- differentiating both roughly doubled it. + x = x.detach().requires_grad_(direction == "bwd-data") + w = w.detach().requires_grad_(direction != "bwd-data") + y = _conv(problem, x, w, b) + (grad,) = torch.autograd.grad(y, (x if direction == "bwd-data" else w,), gy) + return grad + + +def incumbent( + problem: ConvProblem, + operands: dict[str, torch.Tensor], + direction: Direction = "fwd", +) -> torch.Tensor: + """What MIOpen produces today -- the thing we have to be no worse than.""" + x = operands["input"] + w = operands["weight"] + b = operands["bias"] + if direction == "fwd": + return _conv(problem, x, w, b) + gy = operands["grad_output"] + x = x.detach().requires_grad_(True) + w = w.detach().requires_grad_(True) + y = _conv(problem, x, w, b) + grad_x, grad_w = torch.autograd.grad(y, (x, w), gy) + return grad_x if direction == "bwd-data" else grad_w + + +# --------------------------------------------------------------------------- +# Comparison +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class ErrorReport: + max_abs: float + max_rel: float + rms_rel: float + #: Number of elements differing at all, and the total. + n_diff: int + n_total: int + + @property + def bitwise(self) -> bool: + return self.n_diff == 0 + + def __str__(self) -> str: + return ( + f"max_abs={self.max_abs:.3e} max_rel={self.max_rel:.3e} " + f"rms_rel={self.rms_rel:.3e} diff={self.n_diff}/{self.n_total}" + ) + + +def compare(actual: torch.Tensor, expected: torch.Tensor) -> ErrorReport: + """Error of ``actual`` against a higher-precision ``expected``. + + Relative error is normalized by the RMS of ``expected`` rather than + elementwise, because a convolution output legitimately contains + near-cancellations whose elementwise relative error is unbounded and + uninformative. + """ + a = actual.detach().to(torch.float64) + e = expected.detach().to(device=a.device, dtype=torch.float64) + if a.shape != e.shape: + raise ValueError(f"shape mismatch: {tuple(a.shape)} vs {tuple(e.shape)}") + diff = (a - e).abs() + scale = e.pow(2).mean().sqrt().item() + scale = scale if scale > 0 else 1.0 + return ErrorReport( + max_abs=diff.max().item(), + max_rel=(diff.max() / scale).item(), + rms_rel=(diff.pow(2).mean().sqrt() / scale).item(), + n_diff=int((a != e).sum().item()), + n_total=a.numel(), + ) + + +def error_bound( + problem: ConvProblem, + expected: torch.Tensor, + direction: Direction = "fwd", + *, + roundings: float = 1.0, +) -> float: + """Absolute bound on ``max |actual - expected|``. + + Two error sources, and they do not scale with the same quantity: + + - **Accumulation** in fp32 over ``K`` terms. Rounding there behaves like a + random walk rather than a worst case, so ``u * sqrt(K)``, and it scales + with the *typical* magnitude of the result -- its RMS. A random walk is + an average, not a bound, so this term carries the 8x safety factor. + - **The final store** down to bf16, which is up to one ulp of each element + and so scales with the *largest* element, not the typical one. + + Conflating the two is a real trap, and one this code fell into: measuring + error relative to the RMS while bounding it in per-element ulps understates + the bound by the tensor's peak-to-RMS ratio, which for a convolution result + is comfortably 5x. MIOpen itself failed that bound on the transposed + backward-weight, which is how the mistake surfaced -- the tolerance was + wrong, not the incumbent. + + The 8x used to sit on *both* terms, and that was the opposite mistake. The + store is a single deterministic rounding, bounded by half an ulp of the + element and so by ``u_dtype * peak`` outright -- there is no walk to take a + safety factor against, and charging four ulps of the peak for it made the + static bound 12-17x MIOpen's measured error and left the "no worse than the + incumbent" clause of :func:`assert_close` dead in 46 of 48 cells. Measured + here over 78 (problem, direction) cells, MIOpen's own ``max_abs`` in bf16 + and fp16 lands at **0.24-0.66 ulps of the peak** in the forward and + backward-data -- both of which are bitwise reproducible, i.e. genuinely one + rounding -- so ``roundings=1`` (a full ulp of the peak) covers a + single-store kernel with 1.5-4x to spare. fp32 is the exception and is + covered by the other term: there ``u_dtype`` is 2**16 smaller, the + accumulation dominates, and MIOpen sits at 7-8 ulps of a very small ulp. + + ``roundings`` is the number of times a value is rounded into the working + dtype on its way out, and it is a knob because the incumbent is not always + 1: MIOpen's backward-weight reduces with atomics, so two identical calls + differ bitwise (verified) and its error *wanders* -- over eight calls on one + cell it ranged 0.61-1.05 ulps of the peak, disagreeing with itself by 0.72, + where a single rounding would repeat exactly. Our backward-weight reduces + its split-K partials in fp32 and stores once, so it stays at 1. + + What this bound still does not cover, and no tolerance can: a ``tl.dot`` + silently running at ~10-11 mantissa bits sits *under* one bf16 ulp of the + peak and passes. Only the bitwise standard rejects that. + """ + dt = _TORCH_DTYPE[problem.dtype] + k = problem.gemm_shape(direction)[2] + e = expected.detach().to(torch.float64) + rms = e.pow(2).mean().sqrt().item() + peak = e.abs().max().item() + accum = unit_roundoff(torch.float32) * math.sqrt(k) * rms + store = unit_roundoff(dt) * peak + return 8.0 * accum + 2.0 * roundings * store + + +def assert_close( + actual: torch.Tensor, + expected: torch.Tensor, + problem: ConvProblem, + direction: Direction = "fwd", + *, + incumbent_error: ErrorReport | None = None, + margin: float = 4.0, + roundings: float = 1.0, +) -> ErrorReport: + """Apply the strictest standard the situation supports. + + If ``incumbent_error`` is supplied, the bar is "no worse than MIOpen by more + than ``margin``"; otherwise :func:`error_bound` applies. The two are combined + with ``max`` so that a shape where MIOpen happens to be unusually accurate + cannot make the test stricter than the numerics justify. + + That ``max`` is only worth writing if both arms can win, and for a long time + only one could: with four ulps of the peak charged for the final store the + static bound won 46 of 48 cells and the documented standard was never the + one applied. With the store term at one ulp (see :func:`error_bound`) + ``margin * incumbent`` is operative in 71 of the 75 bf16/fp16 cells + measured. Which arm wins is closest to a coin toss in fp32, where + ``u_dtype`` is 2**16 smaller: the store term stops dominating, the bound + collapses onto MIOpen's own accumulation error, and the two arms come out + within about 1.2x of each other in either direction. + """ + report = compare(actual, expected) + bound = error_bound(problem, expected, direction, roundings=roundings) + if incumbent_error is not None: + bound = max(bound, margin * incumbent_error.max_abs) + if not (report.max_abs <= bound): + raise AssertionError( + f"{problem.label} [{direction}]: {report}, bound max_abs <= {bound:.3e}" + + (f" (incumbent {incumbent_error})" if incumbent_error else "") + ) + return report diff --git a/triton_conv3d/scaffold_census.json b/triton_conv3d/scaffold_census.json new file mode 100644 index 0000000..456bf8e --- /dev/null +++ b/triton_conv3d/scaffold_census.json @@ -0,0 +1,5065 @@ +{ + "source": "work/triton-conv/bin/conv_census.py, three training steps per configuration, wrapped entry points", + "form": "adapter", + "note": "The shape and padding the kernel was handed, read off a real call. Every k>1 convolution here is padded: the ScaFFold adapter exchanges a halo only on genuinely split axes.", + "configs": [ + { + "tag": "A", + "desc": "scale 7, 1 GPU, shards (1,1,1)", + "capture": "cens_A.census.json" + }, + { + "tag": "B", + "desc": "scale 8, 1 GPU, shards (1,1,1)", + "capture": "cens_B.census.json" + }, + { + "tag": "C", + "desc": "scale 8, 2 GPUs, shards (2,1,1)", + "capture": "cens_C.r0.census.json" + }, + { + "tag": "D", + "desc": "scale 8, 4 GPUs, shards (4,1,1)", + "capture": "cens_D.r0.census.json" + } + ], + "n_problems": 88, + "problems": [ + { + "op": "Conv3d", + "weight_shape": [ + 64, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 256, + 256, + 256 + ], + "out_shape": [ + 1, + 64, + 256, + 256, + 256 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "B:up_list.4.conv.double_conv.0" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-up_list.4.conv.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 64, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 130, + 256, + 256 + ], + "out_shape": [ + 1, + 64, + 128, + 256, + 256 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "C:up_list.4.conv.double_conv.0" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-up_list.4.conv.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 64, + 64, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 64, + 256, + 256, + 256 + ], + "out_shape": [ + 1, + 64, + 256, + 256, + 256 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "B:down_list.0.double_conv.3", + "B:up_list.4.conv.double_conv.3" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-down_list.0.double_conv.3" + }, + { + "op": "Conv3d", + "weight_shape": [ + 3, + 64, + 1, + 1, + 1 + ], + "in_shape": [ + 1, + 64, + 256, + 256, + 256 + ], + "out_shape": [ + 1, + 3, + 256, + 256, + 256 + ], + "kernel": [ + 1, + 1, + 1 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "B:up_list.5.conv" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-up_list.5.conv" + }, + { + "op": "Conv3d", + "weight_shape": [ + 64, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 66, + 256, + 256 + ], + "out_shape": [ + 1, + 64, + 64, + 256, + 256 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "D:up_list.4.conv.double_conv.0" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-up_list.4.conv.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 64, + 64, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 64, + 130, + 256, + 256 + ], + "out_shape": [ + 1, + 64, + 128, + 256, + 256 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "C:down_list.0.double_conv.3", + "C:up_list.4.conv.double_conv.3" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-down_list.0.double_conv.3" + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 128, + 128, + 128 + ], + "out_shape": [ + 1, + 128, + 128, + 128, + 128 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "B:up_list.3.conv.double_conv.0" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-up_list.3.conv.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 3, + 64, + 1, + 1, + 1 + ], + "in_shape": [ + 1, + 64, + 128, + 256, + 256 + ], + "out_shape": [ + 1, + 3, + 128, + 256, + 256 + ], + "kernel": [ + 1, + 1, + 1 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "C:up_list.5.conv" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-up_list.5.conv" + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 66, + 128, + 128 + ], + "out_shape": [ + 1, + 128, + 64, + 128, + 128 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "C:up_list.3.conv.double_conv.0" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-up_list.3.conv.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 64, + 64, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 64, + 66, + 256, + 256 + ], + "out_shape": [ + 1, + 64, + 64, + 256, + 256 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "D:down_list.0.double_conv.3", + "D:up_list.4.conv.double_conv.3" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-down_list.0.double_conv.3" + }, + { + "op": "Conv3d", + "weight_shape": [ + 64, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 128, + 128, + 128 + ], + "out_shape": [ + 1, + 64, + 128, + 128, + 128 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "A:up_list.3.conv.double_conv.0" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "name": "A-up_list.3.conv.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 128, + 128, + 128 + ], + "out_shape": [ + 1, + 128, + 128, + 128, + 128 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "B:down_list.1.maxpool_conv.1.double_conv.3", + "B:up_list.3.conv.double_conv.3" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-down_list.1.maxpool_conv.1.double_conv.3" + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 128, + 64, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 128, + 128, + 128, + 128 + ], + "out_shape": [ + 1, + 64, + 256, + 256, + 256 + ], + "kernel": [ + 2, + 2, + 2 + ], + "stride": [ + 2, + 2, + 2 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "B:up_list.4.up" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-up_list.4.up" + }, + { + "op": "Conv3d", + "weight_shape": [ + 3, + 64, + 1, + 1, + 1 + ], + "in_shape": [ + 1, + 64, + 64, + 256, + 256 + ], + "out_shape": [ + 1, + 3, + 64, + 256, + 256 + ], + "kernel": [ + 1, + 1, + 1 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "D:up_list.5.conv" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-up_list.5.conv" + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 34, + 128, + 128 + ], + "out_shape": [ + 1, + 128, + 32, + 128, + 128 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "D:up_list.3.conv.double_conv.0" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-up_list.3.conv.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 66, + 128, + 128 + ], + "out_shape": [ + 1, + 128, + 64, + 128, + 128 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "C:down_list.1.maxpool_conv.1.double_conv.3", + "C:up_list.3.conv.double_conv.3" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-down_list.1.maxpool_conv.1.double_conv.3" + }, + { + "op": "Conv3d", + "weight_shape": [ + 64, + 64, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 64, + 128, + 128, + 128 + ], + "out_shape": [ + 1, + 64, + 128, + 128, + 128 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "A:down_list.0.double_conv.3", + "A:up_list.3.conv.double_conv.3" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "name": "A-down_list.0.double_conv.3" + }, + { + "op": "Conv3d", + "weight_shape": [ + 3, + 64, + 1, + 1, + 1 + ], + "in_shape": [ + 1, + 64, + 128, + 128, + 128 + ], + "out_shape": [ + 1, + 3, + 128, + 128, + 128 + ], + "kernel": [ + 1, + 1, + 1 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "A:up_list.4.conv" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "name": "A-up_list.4.conv" + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 64, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 64, + 128, + 128, + 128 + ], + "out_shape": [ + 1, + 128, + 128, + 128, + 128 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "B:down_list.1.maxpool_conv.1.double_conv.0" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-down_list.1.maxpool_conv.1.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 64, + 64, + 64 + ], + "out_shape": [ + 1, + 256, + 64, + 64, + 64 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "B:up_list.2.conv.double_conv.0" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-up_list.2.conv.double_conv.0" + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 128, + 64, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 128, + 64, + 128, + 128 + ], + "out_shape": [ + 1, + 64, + 128, + 256, + 256 + ], + "kernel": [ + 2, + 2, + 2 + ], + "stride": [ + 2, + 2, + 2 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "C:up_list.4.up" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-up_list.4.up" + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 34, + 64, + 64 + ], + "out_shape": [ + 1, + 256, + 32, + 64, + 64 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "C:up_list.2.conv.double_conv.0" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-up_list.2.conv.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 34, + 128, + 128 + ], + "out_shape": [ + 1, + 128, + 32, + 128, + 128 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "D:down_list.1.maxpool_conv.1.double_conv.3", + "D:up_list.3.conv.double_conv.3" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-down_list.1.maxpool_conv.1.double_conv.3" + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 64, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 64, + 66, + 128, + 128 + ], + "out_shape": [ + 1, + 128, + 64, + 128, + 128 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "C:down_list.1.maxpool_conv.1.double_conv.0" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-down_list.1.maxpool_conv.1.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 64, + 64, + 64 + ], + "out_shape": [ + 1, + 128, + 64, + 64, + 64 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "A:up_list.2.conv.double_conv.0" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "name": "A-up_list.2.conv.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 64, + 64, + 64 + ], + "out_shape": [ + 1, + 256, + 64, + 64, + 64 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "B:down_list.2.maxpool_conv.1.double_conv.3", + "B:up_list.2.conv.double_conv.3" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-down_list.2.maxpool_conv.1.double_conv.3" + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 256, + 128, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 256, + 64, + 64, + 64 + ], + "out_shape": [ + 1, + 128, + 128, + 128, + 128 + ], + "kernel": [ + 2, + 2, + 2 + ], + "stride": [ + 2, + 2, + 2 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "B:up_list.3.up" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-up_list.3.up" + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 128, + 64, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 128, + 32, + 128, + 128 + ], + "out_shape": [ + 1, + 64, + 64, + 256, + 256 + ], + "kernel": [ + 2, + 2, + 2 + ], + "stride": [ + 2, + 2, + 2 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "D:up_list.4.up" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-up_list.4.up" + }, + { + "op": "Conv3d", + "weight_shape": [ + 64, + 3, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 3, + 256, + 256, + 256 + ], + "out_shape": [ + 1, + 64, + 256, + 256, + 256 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "B:down_list.0.double_conv.0" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-down_list.0.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 18, + 64, + 64 + ], + "out_shape": [ + 1, + 256, + 16, + 64, + 64 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "D:up_list.2.conv.double_conv.0" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-up_list.2.conv.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 34, + 64, + 64 + ], + "out_shape": [ + 1, + 256, + 32, + 64, + 64 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "C:down_list.2.maxpool_conv.1.double_conv.3", + "C:up_list.2.conv.double_conv.3" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-down_list.2.maxpool_conv.1.double_conv.3" + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 64, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 64, + 34, + 128, + 128 + ], + "out_shape": [ + 1, + 128, + 32, + 128, + 128 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "D:down_list.1.maxpool_conv.1.double_conv.0" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-down_list.1.maxpool_conv.1.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 64, + 64, + 64 + ], + "out_shape": [ + 1, + 128, + 64, + 64, + 64 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "A:down_list.1.maxpool_conv.1.double_conv.3", + "A:up_list.2.conv.double_conv.3" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "name": "A-down_list.1.maxpool_conv.1.double_conv.3" + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 128, + 64, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 128, + 64, + 64, + 64 + ], + "out_shape": [ + 1, + 64, + 128, + 128, + 128 + ], + "kernel": [ + 2, + 2, + 2 + ], + "stride": [ + 2, + 2, + 2 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "A:up_list.3.up" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "name": "A-up_list.3.up" + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 64, + 64, + 64 + ], + "out_shape": [ + 1, + 256, + 64, + 64, + 64 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "B:down_list.2.maxpool_conv.1.double_conv.0" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-down_list.2.maxpool_conv.1.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 1024, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 1024, + 32, + 32, + 32 + ], + "out_shape": [ + 1, + 512, + 32, + 32, + 32 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "B:up_list.1.conv.double_conv.0" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-up_list.1.conv.double_conv.0" + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 256, + 128, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 256, + 32, + 64, + 64 + ], + "out_shape": [ + 1, + 128, + 64, + 128, + 128 + ], + "kernel": [ + 2, + 2, + 2 + ], + "stride": [ + 2, + 2, + 2 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "C:up_list.3.up" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-up_list.3.up" + }, + { + "op": "Conv3d", + "weight_shape": [ + 64, + 3, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 3, + 130, + 256, + 256 + ], + "out_shape": [ + 1, + 64, + 128, + 256, + 256 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "C:down_list.0.double_conv.0" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-down_list.0.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 1024, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 1024, + 18, + 32, + 32 + ], + "out_shape": [ + 1, + 512, + 16, + 32, + 32 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "C:up_list.1.conv.double_conv.0" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-up_list.1.conv.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 18, + 64, + 64 + ], + "out_shape": [ + 1, + 256, + 16, + 64, + 64 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "D:down_list.2.maxpool_conv.1.double_conv.3", + "D:up_list.2.conv.double_conv.3" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-down_list.2.maxpool_conv.1.double_conv.3" + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 34, + 64, + 64 + ], + "out_shape": [ + 1, + 256, + 32, + 64, + 64 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "C:down_list.2.maxpool_conv.1.double_conv.0" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-down_list.2.maxpool_conv.1.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 64, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 64, + 64, + 64, + 64 + ], + "out_shape": [ + 1, + 128, + 64, + 64, + 64 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "A:down_list.1.maxpool_conv.1.double_conv.0" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "name": "A-down_list.1.maxpool_conv.1.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 32, + 32, + 32 + ], + "out_shape": [ + 1, + 256, + 32, + 32, + 32 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "A:up_list.1.conv.double_conv.0" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "name": "A-up_list.1.conv.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 32, + 32, + 32 + ], + "out_shape": [ + 1, + 512, + 32, + 32, + 32 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "B:down_list.3.maxpool_conv.1.double_conv.3", + "B:up_list.1.conv.double_conv.3" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-down_list.3.maxpool_conv.1.double_conv.3" + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 512, + 256, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 512, + 32, + 32, + 32 + ], + "out_shape": [ + 1, + 256, + 64, + 64, + 64 + ], + "kernel": [ + 2, + 2, + 2 + ], + "stride": [ + 2, + 2, + 2 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "B:up_list.2.up" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-up_list.2.up" + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 256, + 128, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 256, + 16, + 64, + 64 + ], + "out_shape": [ + 1, + 128, + 32, + 128, + 128 + ], + "kernel": [ + 2, + 2, + 2 + ], + "stride": [ + 2, + 2, + 2 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "D:up_list.3.up" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-up_list.3.up" + }, + { + "op": "Conv3d", + "weight_shape": [ + 64, + 3, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 3, + 66, + 256, + 256 + ], + "out_shape": [ + 1, + 64, + 64, + 256, + 256 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": true, + "sites": [ + "D:down_list.0.double_conv.0" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-down_list.0.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 1024, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 1024, + 10, + 32, + 32 + ], + "out_shape": [ + 1, + 512, + 8, + 32, + 32 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "D:up_list.1.conv.double_conv.0" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-up_list.1.conv.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 18, + 32, + 32 + ], + "out_shape": [ + 1, + 512, + 16, + 32, + 32 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "C:down_list.3.maxpool_conv.1.double_conv.3", + "C:up_list.1.conv.double_conv.3" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-down_list.3.maxpool_conv.1.double_conv.3" + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 18, + 64, + 64 + ], + "out_shape": [ + 1, + 256, + 16, + 64, + 64 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "D:down_list.2.maxpool_conv.1.double_conv.0" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-down_list.2.maxpool_conv.1.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 32, + 32, + 32 + ], + "out_shape": [ + 1, + 256, + 32, + 32, + 32 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "A:down_list.2.maxpool_conv.1.double_conv.3", + "A:up_list.1.conv.double_conv.3" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "name": "A-down_list.2.maxpool_conv.1.double_conv.3" + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 256, + 128, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 256, + 32, + 32, + 32 + ], + "out_shape": [ + 1, + 128, + 64, + 64, + 64 + ], + "kernel": [ + 2, + 2, + 2 + ], + "stride": [ + 2, + 2, + 2 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "A:up_list.2.up" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "name": "A-up_list.2.up" + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 32, + 32, + 32 + ], + "out_shape": [ + 1, + 512, + 32, + 32, + 32 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "B:down_list.3.maxpool_conv.1.double_conv.0" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-down_list.3.maxpool_conv.1.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 1024, + 2048, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 2048, + 16, + 16, + 16 + ], + "out_shape": [ + 1, + 1024, + 16, + 16, + 16 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "B:up_list.0.conv.double_conv.0" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-up_list.0.conv.double_conv.0" + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 512, + 256, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 512, + 16, + 32, + 32 + ], + "out_shape": [ + 1, + 256, + 32, + 64, + 64 + ], + "kernel": [ + 2, + 2, + 2 + ], + "stride": [ + 2, + 2, + 2 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "C:up_list.2.up" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-up_list.2.up" + }, + { + "op": "Conv3d", + "weight_shape": [ + 64, + 3, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 3, + 128, + 128, + 128 + ], + "out_shape": [ + 1, + 64, + 128, + 128, + 128 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "A:down_list.0.double_conv.0" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "name": "A-down_list.0.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 1024, + 2048, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 2048, + 10, + 16, + 16 + ], + "out_shape": [ + 1, + 1024, + 8, + 16, + 16 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "C:up_list.0.conv.double_conv.0" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-up_list.0.conv.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 10, + 32, + 32 + ], + "out_shape": [ + 1, + 512, + 8, + 32, + 32 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "D:down_list.3.maxpool_conv.1.double_conv.3", + "D:up_list.1.conv.double_conv.3" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-down_list.3.maxpool_conv.1.double_conv.3" + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 18, + 32, + 32 + ], + "out_shape": [ + 1, + 512, + 16, + 32, + 32 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "C:down_list.3.maxpool_conv.1.double_conv.0" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-down_list.3.maxpool_conv.1.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 32, + 32, + 32 + ], + "out_shape": [ + 1, + 256, + 32, + 32, + 32 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "A:down_list.2.maxpool_conv.1.double_conv.0" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "name": "A-down_list.2.maxpool_conv.1.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 1024, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 1024, + 16, + 16, + 16 + ], + "out_shape": [ + 1, + 512, + 16, + 16, + 16 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "A:up_list.0.conv.double_conv.0" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "name": "A-up_list.0.conv.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 1024, + 1024, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 1024, + 16, + 16, + 16 + ], + "out_shape": [ + 1, + 1024, + 16, + 16, + 16 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "B:down_list.4.maxpool_conv.1.double_conv.3", + "B:up_list.0.conv.double_conv.3" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-down_list.4.maxpool_conv.1.double_conv.3" + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 1024, + 512, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 1024, + 16, + 16, + 16 + ], + "out_shape": [ + 1, + 512, + 32, + 32, + 32 + ], + "kernel": [ + 2, + 2, + 2 + ], + "stride": [ + 2, + 2, + 2 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "B:up_list.1.up" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-up_list.1.up" + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 512, + 256, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 512, + 8, + 32, + 32 + ], + "out_shape": [ + 1, + 256, + 16, + 64, + 64 + ], + "kernel": [ + 2, + 2, + 2 + ], + "stride": [ + 2, + 2, + 2 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "D:up_list.2.up" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-up_list.2.up" + }, + { + "op": "Conv3d", + "weight_shape": [ + 1024, + 2048, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 2048, + 6, + 16, + 16 + ], + "out_shape": [ + 1, + 1024, + 4, + 16, + 16 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "D:up_list.0.conv.double_conv.0" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-up_list.0.conv.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 1024, + 1024, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 1024, + 10, + 16, + 16 + ], + "out_shape": [ + 1, + 1024, + 8, + 16, + 16 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "C:down_list.4.maxpool_conv.1.double_conv.3", + "C:up_list.0.conv.double_conv.3" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-down_list.4.maxpool_conv.1.double_conv.3" + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 10, + 32, + 32 + ], + "out_shape": [ + 1, + 512, + 8, + 32, + 32 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "D:down_list.3.maxpool_conv.1.double_conv.0" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-down_list.3.maxpool_conv.1.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 16, + 16, + 16 + ], + "out_shape": [ + 1, + 512, + 16, + 16, + 16 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "A:down_list.3.maxpool_conv.1.double_conv.3", + "A:up_list.0.conv.double_conv.3" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "name": "A-down_list.3.maxpool_conv.1.double_conv.3" + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 512, + 256, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 512, + 16, + 16, + 16 + ], + "out_shape": [ + 1, + 256, + 32, + 32, + 32 + ], + "kernel": [ + 2, + 2, + 2 + ], + "stride": [ + 2, + 2, + 2 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "A:up_list.1.up" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "name": "A-up_list.1.up" + }, + { + "op": "Conv3d", + "weight_shape": [ + 1024, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 16, + 16, + 16 + ], + "out_shape": [ + 1, + 1024, + 16, + 16, + 16 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "B:down_list.4.maxpool_conv.1.double_conv.0" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-down_list.4.maxpool_conv.1.double_conv.0" + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 1024, + 512, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 1024, + 8, + 16, + 16 + ], + "out_shape": [ + 1, + 512, + 16, + 32, + 32 + ], + "kernel": [ + 2, + 2, + 2 + ], + "stride": [ + 2, + 2, + 2 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "C:up_list.1.up" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-up_list.1.up" + }, + { + "op": "Conv3d", + "weight_shape": [ + 1024, + 1024, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 1024, + 6, + 16, + 16 + ], + "out_shape": [ + 1, + 1024, + 4, + 16, + 16 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "D:down_list.4.maxpool_conv.1.double_conv.3", + "D:up_list.0.conv.double_conv.3" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-down_list.4.maxpool_conv.1.double_conv.3" + }, + { + "op": "Conv3d", + "weight_shape": [ + 1024, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 10, + 16, + 16 + ], + "out_shape": [ + 1, + 1024, + 8, + 16, + 16 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "C:down_list.4.maxpool_conv.1.double_conv.0" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-down_list.4.maxpool_conv.1.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 16, + 16, + 16 + ], + "out_shape": [ + 1, + 512, + 16, + 16, + 16 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "A:down_list.3.maxpool_conv.1.double_conv.0" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "name": "A-down_list.3.maxpool_conv.1.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 2048, + 2048, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 2048, + 8, + 8, + 8 + ], + "out_shape": [ + 1, + 2048, + 8, + 8, + 8 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "B:down_list.5.maxpool_conv.1.double_conv.3" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-down_list.5.maxpool_conv.1.double_conv.3" + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 2048, + 1024, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 2048, + 8, + 8, + 8 + ], + "out_shape": [ + 1, + 1024, + 16, + 16, + 16 + ], + "kernel": [ + 2, + 2, + 2 + ], + "stride": [ + 2, + 2, + 2 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "B:up_list.0.up" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-up_list.0.up" + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 1024, + 512, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 1024, + 4, + 16, + 16 + ], + "out_shape": [ + 1, + 512, + 8, + 32, + 32 + ], + "kernel": [ + 2, + 2, + 2 + ], + "stride": [ + 2, + 2, + 2 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "D:up_list.1.up" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-up_list.1.up" + }, + { + "op": "Conv3d", + "weight_shape": [ + 2048, + 2048, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 2048, + 6, + 8, + 8 + ], + "out_shape": [ + 1, + 2048, + 4, + 8, + 8 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "C:down_list.5.maxpool_conv.1.double_conv.3" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-down_list.5.maxpool_conv.1.double_conv.3" + }, + { + "op": "Conv3d", + "weight_shape": [ + 1024, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 6, + 16, + 16 + ], + "out_shape": [ + 1, + 1024, + 4, + 16, + 16 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "D:down_list.4.maxpool_conv.1.double_conv.0" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-down_list.4.maxpool_conv.1.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 1024, + 1024, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 1024, + 8, + 8, + 8 + ], + "out_shape": [ + 1, + 1024, + 8, + 8, + 8 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "A:down_list.4.maxpool_conv.1.double_conv.3" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "name": "A-down_list.4.maxpool_conv.1.double_conv.3" + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 1024, + 512, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 1024, + 8, + 8, + 8 + ], + "out_shape": [ + 1, + 512, + 16, + 16, + 16 + ], + "kernel": [ + 2, + 2, + 2 + ], + "stride": [ + 2, + 2, + 2 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "A:up_list.0.up" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "name": "A-up_list.0.up" + }, + { + "op": "Conv3d", + "weight_shape": [ + 2048, + 1024, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 1024, + 8, + 8, + 8 + ], + "out_shape": [ + 1, + 2048, + 8, + 8, + 8 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "B:down_list.5.maxpool_conv.1.double_conv.0" + ], + "configs": [ + "scale 8, 1 GPU, shards (1,1,1)" + ], + "name": "B-down_list.5.maxpool_conv.1.double_conv.0" + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 2048, + 1024, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 2048, + 4, + 8, + 8 + ], + "out_shape": [ + 1, + 1024, + 8, + 16, + 16 + ], + "kernel": [ + 2, + 2, + 2 + ], + "stride": [ + 2, + 2, + 2 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "C:up_list.0.up" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-up_list.0.up" + }, + { + "op": "Conv3d", + "weight_shape": [ + 2048, + 2048, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 2048, + 4, + 8, + 8 + ], + "out_shape": [ + 1, + 2048, + 2, + 8, + 8 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "D:down_list.5.maxpool_conv.1.double_conv.3" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-down_list.5.maxpool_conv.1.double_conv.3" + }, + { + "op": "Conv3d", + "weight_shape": [ + 2048, + 1024, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 1024, + 6, + 8, + 8 + ], + "out_shape": [ + 1, + 2048, + 4, + 8, + 8 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "C:down_list.5.maxpool_conv.1.double_conv.0" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "name": "C-down_list.5.maxpool_conv.1.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 1024, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 8, + 8, + 8 + ], + "out_shape": [ + 1, + 1024, + 8, + 8, + 8 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 1, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "A:down_list.4.maxpool_conv.1.double_conv.0" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "name": "A-down_list.4.maxpool_conv.1.double_conv.0" + }, + { + "op": "Conv3d", + "weight_shape": [ + 2048, + 1024, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 1024, + 4, + 8, + 8 + ], + "out_shape": [ + 1, + 2048, + 2, + 8, + 8 + ], + "kernel": [ + 3, + 3, + 3 + ], + "stride": [ + 1, + 1, + 1 + ], + "padding": [ + 0, + 1, + 1 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": false, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "D:down_list.5.maxpool_conv.1.double_conv.0" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-down_list.5.maxpool_conv.1.double_conv.0" + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 2048, + 1024, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 2048, + 2, + 8, + 8 + ], + "out_shape": [ + 1, + 1024, + 4, + 16, + 16 + ], + "kernel": [ + 2, + 2, + 2 + ], + "stride": [ + 2, + 2, + 2 + ], + "padding": [ + 0, + 0, + 0 + ], + "dilation": [ + 1, + 1, + 1 + ], + "groups": 1, + "bias": true, + "dtype": "bf16", + "memory_format": "channels_last_3d", + "dctensor": true, + "large": false, + "sites": [ + "D:up_list.0.up" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "name": "D-up_list.0.up" + } + ] +} diff --git a/triton_conv3d/scaffold_corpus.json b/triton_conv3d/scaffold_corpus.json new file mode 100644 index 0000000..1170655 --- /dev/null +++ b/triton_conv3d/scaffold_corpus.json @@ -0,0 +1,4871 @@ +{ + "source": "model-analysis/unet_shapes.py via work/report/data", + "configs": [ + { + "tag": "A", + "desc": "scale 7, 1 GPU, shards (1,1,1)" + }, + { + "tag": "B", + "desc": "scale 8, 2 GPUs, shards (2,1,1)" + }, + { + "tag": "C", + "desc": "scale 8, 4 GPUs, shards (4,1,1)" + } + ], + "n_problems": 57, + "problems": [ + { + "op": "Conv3d", + "weight_shape": [ + 64, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 128, + 256, + 256 + ], + "out_shape": [ + 1, + 64, + 128, + 256, + 256 + ], + "halo_in_shape": [ + 1, + 128, + 130, + 258, + 258 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "B:dec3 (up4)#61" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "measured": [ + { + "config": "B", + "direction": "bwd-weight", + "ms_per_step": 45401.1315, + "ms_per_call": 45401.1315, + "calls": 1, + "pct_roofline": 0.0, + "solvers": [ + "unknown" + ] + }, + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 8199.9801, + "ms_per_call": 8199.9801, + "calls": 1, + "pct_roofline": 0.1, + "solvers": [ + "unknown" + ] + }, + { + "config": "B", + "direction": "fwd", + "ms_per_step": 18.861, + "ms_per_call": 18.861, + "calls": 1, + "pct_roofline": 32.8, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_d_grouped_gemm_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 64, + 64, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 64, + 128, + 256, + 256 + ], + "out_shape": [ + 1, + 64, + 128, + 256, + 256 + ], + "halo_in_shape": [ + 1, + 64, + 130, + 258, + 258 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "B:enc0 (inc)#4", + "B:dec3 (up4)#64" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "measured": [ + { + "config": "B", + "direction": "bwd-weight", + "ms_per_step": 25.9167, + "ms_per_call": 12.9584, + "calls": 2, + "pct_roofline": 23.9, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 25.4267, + "ms_per_call": 12.7133, + "calls": 2, + "pct_roofline": 24.3, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "B", + "direction": "fwd", + "ms_per_step": 13.9525, + "ms_per_call": 6.9763, + "calls": 2, + "pct_roofline": 44.3, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_abd_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 64, + 64, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 64, + 64, + 256, + 256 + ], + "out_shape": [ + 1, + 64, + 64, + 256, + 256 + ], + "halo_in_shape": [ + 1, + 64, + 66, + 258, + 258 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "C:enc0 (inc)#4", + "C:dec3 (up4)#64" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "measured": [ + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 13.072, + "ms_per_call": 6.536, + "calls": 2, + "pct_roofline": 23.7, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "C", + "direction": "bwd-weight", + "ms_per_step": 12.0946, + "ms_per_call": 6.0473, + "calls": 2, + "pct_roofline": 25.6, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "C", + "direction": "fwd", + "ms_per_step": 7.5631, + "ms_per_call": 3.7815, + "calls": 2, + "pct_roofline": 40.9, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_abd_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 64, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 64, + 256, + 256 + ], + "out_shape": [ + 1, + 64, + 64, + 256, + 256 + ], + "halo_in_shape": [ + 1, + 128, + 66, + 258, + 258 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "C:dec3 (up4)#61" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "measured": [ + { + "config": "C", + "direction": "bwd-weight", + "ms_per_step": 12.2945, + "ms_per_call": 12.2945, + "calls": 1, + "pct_roofline": 25.2, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "C", + "direction": "fwd", + "ms_per_step": 8.2749, + "ms_per_call": 8.2749, + "calls": 1, + "pct_roofline": 37.4, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_abd_xdl_cshuffle" + ] + }, + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 8.1346, + "ms_per_call": 8.1346, + "calls": 1, + "pct_roofline": 38.0, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 64, + 128, + 128 + ], + "out_shape": [ + 1, + 128, + 64, + 128, + 128 + ], + "halo_in_shape": [ + 1, + 128, + 66, + 130, + 130 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "B:enc1 (down1)#11", + "B:dec2 (up3)#56" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "measured": [ + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 8.5547, + "ms_per_call": 4.2774, + "calls": 2, + "pct_roofline": 36.1, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "B", + "direction": "bwd-weight", + "ms_per_step": 7.0896, + "ms_per_call": 3.5448, + "calls": 2, + "pct_roofline": 43.6, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "B", + "direction": "fwd", + "ms_per_step": 5.066, + "ms_per_call": 2.533, + "calls": 2, + "pct_roofline": 61.0, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_abd_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 64, + 128, + 128 + ], + "out_shape": [ + 1, + 128, + 64, + 128, + 128 + ], + "halo_in_shape": [ + 1, + 256, + 66, + 130, + 130 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "B:dec2 (up3)#53" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "measured": [ + { + "config": "B", + "direction": "bwd-weight", + "ms_per_step": 7.412, + "ms_per_call": 7.412, + "calls": 1, + "pct_roofline": 41.7, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 6.3417, + "ms_per_call": 6.3417, + "calls": 1, + "pct_roofline": 48.8, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "B", + "direction": "fwd", + "ms_per_step": 6.2742, + "ms_per_call": 6.2742, + "calls": 1, + "pct_roofline": 49.3, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_abd_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 64, + 64, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 64, + 128, + 128, + 128 + ], + "out_shape": [ + 1, + 64, + 128, + 128, + 128 + ], + "halo_in_shape": [ + 1, + 64, + 130, + 130, + 130 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "A:enc0 (inc)#4", + "A:dec3 (up4)#64" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "measured": [ + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 6.9658, + "ms_per_call": 3.4829, + "calls": 2, + "pct_roofline": 22.2, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "A", + "direction": "bwd-weight", + "ms_per_step": 6.7362, + "ms_per_call": 3.3681, + "calls": 2, + "pct_roofline": 23.0, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "A", + "direction": "fwd", + "ms_per_step": 3.8699, + "ms_per_call": 1.9349, + "calls": 2, + "pct_roofline": 40.0, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_abd_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 64, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 128, + 128, + 128 + ], + "out_shape": [ + 1, + 64, + 128, + 128, + 128 + ], + "halo_in_shape": [ + 1, + 128, + 130, + 130, + 130 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "A:dec3 (up4)#61" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "measured": [ + { + "config": "A", + "direction": "bwd-weight", + "ms_per_step": 5.8648, + "ms_per_call": 5.8648, + "calls": 1, + "pct_roofline": 26.4, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "A", + "direction": "fwd", + "ms_per_step": 4.3772, + "ms_per_call": 4.3772, + "calls": 1, + "pct_roofline": 35.3, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_abd_xdl_cshuffle" + ] + }, + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 3.9948, + "ms_per_call": 3.9948, + "calls": 1, + "pct_roofline": 38.7, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 32, + 128, + 128 + ], + "out_shape": [ + 1, + 128, + 32, + 128, + 128 + ], + "halo_in_shape": [ + 1, + 256, + 34, + 130, + 130 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "C:dec2 (up3)#53" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "measured": [ + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 3.8658, + "ms_per_call": 3.8658, + "calls": 1, + "pct_roofline": 40.0, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "C", + "direction": "bwd-weight", + "ms_per_step": 3.7909, + "ms_per_call": 3.7909, + "calls": 1, + "pct_roofline": 40.8, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "C", + "direction": "fwd", + "ms_per_step": 3.7217, + "ms_per_call": 3.7217, + "calls": 1, + "pct_roofline": 41.5, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_abd_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 32, + 64, + 64 + ], + "out_shape": [ + 1, + 256, + 32, + 64, + 64 + ], + "halo_in_shape": [ + 1, + 512, + 34, + 66, + 66 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "B:dec1 (up2)#45" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "measured": [ + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 3.9987, + "ms_per_call": 3.9987, + "calls": 1, + "pct_roofline": 38.7, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "B", + "direction": "fwd", + "ms_per_step": 3.7667, + "ms_per_call": 3.7667, + "calls": 1, + "pct_roofline": 41.0, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_abd_xdl_cshuffle" + ] + }, + { + "config": "B", + "direction": "bwd-weight", + "ms_per_step": 3.4461, + "ms_per_call": 3.4461, + "calls": 1, + "pct_roofline": 44.9, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 32, + 128, + 128 + ], + "out_shape": [ + 1, + 128, + 32, + 128, + 128 + ], + "halo_in_shape": [ + 1, + 128, + 34, + 130, + 130 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "C:enc1 (down1)#11", + "C:dec2 (up3)#56" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "measured": [ + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 4.3724, + "ms_per_call": 2.1862, + "calls": 2, + "pct_roofline": 35.4, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "C", + "direction": "bwd-weight", + "ms_per_step": 3.7598, + "ms_per_call": 1.8799, + "calls": 2, + "pct_roofline": 41.1, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "C", + "direction": "fwd", + "ms_per_step": 2.8026, + "ms_per_call": 1.4013, + "calls": 2, + "pct_roofline": 55.2, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_abd_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 32, + 64, + 64 + ], + "out_shape": [ + 1, + 256, + 32, + 64, + 64 + ], + "halo_in_shape": [ + 1, + 256, + 34, + 66, + 66 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "B:enc2 (down2)#18", + "B:dec1 (up2)#48" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "measured": [ + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 4.7582, + "ms_per_call": 2.3791, + "calls": 2, + "pct_roofline": 32.5, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "B", + "direction": "bwd-weight", + "ms_per_step": 3.4186, + "ms_per_call": 1.7093, + "calls": 2, + "pct_roofline": 45.2, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "B", + "direction": "fwd", + "ms_per_step": 2.6846, + "ms_per_call": 1.3423, + "calls": 2, + "pct_roofline": 57.6, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 64, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 64, + 64, + 128, + 128 + ], + "out_shape": [ + 1, + 128, + 64, + 128, + 128 + ], + "halo_in_shape": [ + 1, + 64, + 66, + 130, + 130 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "B:enc1 (down1)#8" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "measured": [ + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 3.1462, + "ms_per_call": 3.1462, + "calls": 1, + "pct_roofline": 24.6, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "B", + "direction": "bwd-weight", + "ms_per_step": 1.8767, + "ms_per_call": 1.8767, + "calls": 1, + "pct_roofline": 41.2, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "B", + "direction": "fwd", + "ms_per_step": 1.2644, + "ms_per_call": 1.2644, + "calls": 1, + "pct_roofline": 61.1, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_abd_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 1024, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 1024, + 16, + 32, + 32 + ], + "out_shape": [ + 1, + 512, + 16, + 32, + 32 + ], + "halo_in_shape": [ + 1, + 1024, + 18, + 34, + 34 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "B:dec0 (up1)#37" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "measured": [ + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 2.3372, + "ms_per_call": 2.3372, + "calls": 1, + "pct_roofline": 33.1, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "B", + "direction": "fwd", + "ms_per_step": 2.2222, + "ms_per_call": 2.2222, + "calls": 1, + "pct_roofline": 34.8, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + }, + { + "config": "B", + "direction": "bwd-weight", + "ms_per_step": 1.7249, + "ms_per_call": 1.7249, + "calls": 1, + "pct_roofline": 44.8, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 16, + 32, + 32 + ], + "out_shape": [ + 1, + 512, + 16, + 32, + 32 + ], + "halo_in_shape": [ + 1, + 512, + 18, + 34, + 34 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "B:enc3 (down3)#25", + "B:dec0 (up1)#40" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "measured": [ + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 2.7845, + "ms_per_call": 1.3922, + "calls": 2, + "pct_roofline": 27.8, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "B", + "direction": "bwd-weight", + "ms_per_step": 1.8518, + "ms_per_call": 0.9259, + "calls": 2, + "pct_roofline": 41.7, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "B", + "direction": "fwd", + "ms_per_step": 1.6448, + "ms_per_call": 0.8224, + "calls": 2, + "pct_roofline": 47.0, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 16, + 64, + 64 + ], + "out_shape": [ + 1, + 256, + 16, + 64, + 64 + ], + "halo_in_shape": [ + 1, + 512, + 18, + 66, + 66 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "C:dec1 (up2)#45" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "measured": [ + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 2.4556, + "ms_per_call": 2.4556, + "calls": 1, + "pct_roofline": 31.5, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "C", + "direction": "fwd", + "ms_per_step": 2.0744, + "ms_per_call": 2.0744, + "calls": 1, + "pct_roofline": 37.3, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + }, + { + "config": "C", + "direction": "bwd-weight", + "ms_per_step": 1.7406, + "ms_per_call": 1.7406, + "calls": 1, + "pct_roofline": 44.4, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 16, + 64, + 64 + ], + "out_shape": [ + 1, + 256, + 16, + 64, + 64 + ], + "halo_in_shape": [ + 1, + 256, + 18, + 66, + 66 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "C:enc2 (down2)#18", + "C:dec1 (up2)#48" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "measured": [ + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 2.5593, + "ms_per_call": 1.2796, + "calls": 2, + "pct_roofline": 30.2, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "C", + "direction": "bwd-weight", + "ms_per_step": 1.9686, + "ms_per_call": 0.9843, + "calls": 2, + "pct_roofline": 39.3, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "C", + "direction": "fwd", + "ms_per_step": 1.5206, + "ms_per_call": 0.7603, + "calls": 2, + "pct_roofline": 50.8, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 64, + 3, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 3, + 128, + 256, + 256 + ], + "out_shape": [ + 1, + 64, + 128, + 256, + 256 + ], + "halo_in_shape": [ + 1, + 3, + 130, + 258, + 258 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "B:enc0 (inc)#1" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "measured": [ + { + "config": "B", + "direction": "bwd-weight", + "ms_per_step": 3.172, + "ms_per_call": 3.172, + "calls": 1, + "pct_roofline": 10.8, + "solvers": [ + "kernel_grouped_conv_bwd_weight_xdl_cshuffle_v3" + ] + }, + { + "config": "B", + "direction": "fwd", + "ms_per_step": 2.4064, + "ms_per_call": 2.4064, + "calls": 1, + "pct_roofline": 14.2, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_abd_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 64, + 64, + 64 + ], + "out_shape": [ + 1, + 128, + 64, + 64, + 64 + ], + "halo_in_shape": [ + 1, + 256, + 66, + 66, + 66 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "A:dec2 (up3)#53" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "measured": [ + { + "config": "A", + "direction": "bwd-weight", + "ms_per_step": 1.9356, + "ms_per_call": 1.9356, + "calls": 1, + "pct_roofline": 39.9, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "A", + "direction": "fwd", + "ms_per_step": 1.9287, + "ms_per_call": 1.9287, + "calls": 1, + "pct_roofline": 40.1, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_abd_xdl_cshuffle" + ] + }, + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 1.655, + "ms_per_call": 1.655, + "calls": 1, + "pct_roofline": 46.7, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 64, + 64, + 64 + ], + "out_shape": [ + 1, + 128, + 64, + 64, + 64 + ], + "halo_in_shape": [ + 1, + 128, + 66, + 66, + 66 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "A:enc1 (down1)#11", + "A:dec2 (up3)#56" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "measured": [ + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 2.3401, + "ms_per_call": 1.1701, + "calls": 2, + "pct_roofline": 33.0, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "A", + "direction": "bwd-weight", + "ms_per_step": 1.8785, + "ms_per_call": 0.9392, + "calls": 2, + "pct_roofline": 41.2, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "A", + "direction": "fwd", + "ms_per_step": 1.132, + "ms_per_call": 0.566, + "calls": 2, + "pct_roofline": 68.3, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 8, + 32, + 32 + ], + "out_shape": [ + 1, + 512, + 8, + 32, + 32 + ], + "halo_in_shape": [ + 1, + 512, + 10, + 34, + 34 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "C:enc3 (down3)#25", + "C:dec0 (up1)#40" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "measured": [ + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 1.6378, + "ms_per_call": 0.8189, + "calls": 2, + "pct_roofline": 23.6, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "C", + "direction": "fwd", + "ms_per_step": 1.17, + "ms_per_call": 0.585, + "calls": 2, + "pct_roofline": 33.0, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + }, + { + "config": "C", + "direction": "bwd-weight", + "ms_per_step": 1.0142, + "ms_per_call": 0.5071, + "calls": 2, + "pct_roofline": 38.1, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 1024, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 1024, + 8, + 32, + 32 + ], + "out_shape": [ + 1, + 512, + 8, + 32, + 32 + ], + "halo_in_shape": [ + 1, + 1024, + 10, + 34, + 34 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "C:dec0 (up1)#37" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "measured": [ + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 1.4805, + "ms_per_call": 1.4805, + "calls": 1, + "pct_roofline": 26.1, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "C", + "direction": "fwd", + "ms_per_step": 1.2633, + "ms_per_call": 1.2633, + "calls": 1, + "pct_roofline": 30.6, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + }, + { + "config": "C", + "direction": "bwd-weight", + "ms_per_step": 0.9532, + "ms_per_call": 0.9532, + "calls": 1, + "pct_roofline": 40.6, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + } + ] + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 128, + 64, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 128, + 64, + 128, + 128 + ], + "out_shape": [ + 1, + 64, + 128, + 256, + 256 + ], + "halo_in_shape": [ + 1, + 128, + 64, + 128, + 128 + ], + "halo_dhw": [ + 0, + 0, + 0 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 2, + "stride": 2, + "padding": 0, + "bias": true, + "sites": [ + "B:dec3 (up4)#59" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "measured": [ + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 1.6491, + "ms_per_call": 1.6491, + "calls": 1, + "pct_roofline": 24.7, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "B", + "direction": "bwd-weight", + "ms_per_step": 1.2135, + "ms_per_call": 1.2135, + "calls": 1, + "pct_roofline": 33.5, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 0.7098, + "ms_per_call": 0.7098, + "calls": 1, + "pct_roofline": 57.3, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_d_grouped_gemm_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 64, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 64, + 32, + 128, + 128 + ], + "out_shape": [ + 1, + 128, + 32, + 128, + 128 + ], + "halo_in_shape": [ + 1, + 64, + 34, + 130, + 130 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "C:enc1 (down1)#8" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "measured": [ + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 1.6875, + "ms_per_call": 1.6875, + "calls": 1, + "pct_roofline": 22.9, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "C", + "direction": "bwd-weight", + "ms_per_step": 0.9152, + "ms_per_call": 0.9152, + "calls": 1, + "pct_roofline": 42.2, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "C", + "direction": "fwd", + "ms_per_step": 0.5931, + "ms_per_call": 0.5931, + "calls": 1, + "pct_roofline": 65.2, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_d_grouped_gemm_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 32, + 64, + 64 + ], + "out_shape": [ + 1, + 256, + 32, + 64, + 64 + ], + "halo_in_shape": [ + 1, + 128, + 34, + 66, + 66 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "B:enc2 (down2)#15" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "measured": [ + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 1.58, + "ms_per_call": 1.58, + "calls": 1, + "pct_roofline": 24.5, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "B", + "direction": "bwd-weight", + "ms_per_step": 0.9241, + "ms_per_call": 0.9241, + "calls": 1, + "pct_roofline": 41.8, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "B", + "direction": "fwd", + "ms_per_step": 0.6599, + "ms_per_call": 0.6599, + "calls": 1, + "pct_roofline": 58.6, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_d_grouped_gemm_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 32, + 32, + 32 + ], + "out_shape": [ + 1, + 256, + 32, + 32, + 32 + ], + "halo_in_shape": [ + 1, + 256, + 34, + 34, + 34 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "A:enc2 (down2)#18", + "A:dec1 (up2)#48" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "measured": [ + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 1.2779, + "ms_per_call": 0.6389, + "calls": 2, + "pct_roofline": 30.2, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "A", + "direction": "bwd-weight", + "ms_per_step": 1.016, + "ms_per_call": 0.508, + "calls": 2, + "pct_roofline": 38.0, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "A", + "direction": "fwd", + "ms_per_step": 0.7613, + "ms_per_call": 0.3806, + "calls": 2, + "pct_roofline": 50.8, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3_2lds" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 32, + 32, + 32 + ], + "out_shape": [ + 1, + 256, + 32, + 32, + 32 + ], + "halo_in_shape": [ + 1, + 512, + 34, + 34, + 34 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "A:dec1 (up2)#45" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "measured": [ + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 1.2046, + "ms_per_call": 1.2046, + "calls": 1, + "pct_roofline": 32.1, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "A", + "direction": "bwd-weight", + "ms_per_step": 0.9492, + "ms_per_call": 0.9492, + "calls": 1, + "pct_roofline": 40.7, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "A", + "direction": "fwd", + "ms_per_step": 0.8416, + "ms_per_call": 0.8416, + "calls": 1, + "pct_roofline": 45.9, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3_2lds" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 64, + 3, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 3, + 64, + 256, + 256 + ], + "out_shape": [ + 1, + 64, + 64, + 256, + 256 + ], + "halo_in_shape": [ + 1, + 3, + 66, + 258, + 258 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "C:enc0 (inc)#1" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "measured": [ + { + "config": "C", + "direction": "bwd-weight", + "ms_per_step": 1.4547, + "ms_per_call": 1.4547, + "calls": 1, + "pct_roofline": 11.7, + "solvers": [ + "kernel_grouped_conv_bwd_weight_xdl_cshuffle_v3" + ] + }, + { + "config": "C", + "direction": "fwd", + "ms_per_step": 1.2133, + "ms_per_call": 1.2133, + "calls": 1, + "pct_roofline": 14.1, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_abd_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 1024, + 1024, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 1024, + 8, + 16, + 16 + ], + "out_shape": [ + 1, + 1024, + 8, + 16, + 16 + ], + "halo_in_shape": [ + 1, + 1024, + 10, + 18, + 18 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "B:bottleneck (down4)#32" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "measured": [ + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 1.5751, + "ms_per_call": 1.5751, + "calls": 1, + "pct_roofline": 12.3, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "B", + "direction": "bwd-weight", + "ms_per_step": 0.5376, + "ms_per_call": 0.5376, + "calls": 1, + "pct_roofline": 35.9, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "B", + "direction": "fwd", + "ms_per_step": 0.442, + "ms_per_call": 0.442, + "calls": 1, + "pct_roofline": 43.7, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 16, + 16, + 16 + ], + "out_shape": [ + 1, + 512, + 16, + 16, + 16 + ], + "halo_in_shape": [ + 1, + 512, + 18, + 18, + 18 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "A:enc3 (down3)#25", + "A:dec0 (up1)#40" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "measured": [ + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 1.2806, + "ms_per_call": 0.6403, + "calls": 2, + "pct_roofline": 15.1, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "A", + "direction": "bwd-weight", + "ms_per_step": 0.5551, + "ms_per_call": 0.2776, + "calls": 2, + "pct_roofline": 34.8, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "A", + "direction": "fwd", + "ms_per_step": 0.4488, + "ms_per_call": 0.2244, + "calls": 2, + "pct_roofline": 43.1, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 6, + 64, + 1, + 1, + 1 + ], + "in_shape": [ + 1, + 64, + 128, + 256, + 256 + ], + "out_shape": [ + 1, + 6, + 128, + 256, + 256 + ], + "halo_in_shape": [ + 1, + 64, + 128, + 256, + 256 + ], + "halo_dhw": [ + 0, + 0, + 0 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 1, + "stride": 1, + "padding": 0, + "bias": true, + "sites": [ + "B:outc#67" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "measured": [ + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 1.7305, + "ms_per_call": 1.7305, + "calls": 1, + "pct_roofline": 20.6, + "solvers": [ + "kernel_batched_gemm_xdl_cshuffle_v3_multi_d", + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "B", + "direction": "fwd", + "ms_per_step": 0.5188, + "ms_per_call": 0.5188, + "calls": 1, + "pct_roofline": 68.6, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 128, + 64, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 128, + 32, + 128, + 128 + ], + "out_shape": [ + 1, + 64, + 64, + 256, + 256 + ], + "halo_in_shape": [ + 1, + 128, + 32, + 128, + 128 + ], + "halo_dhw": [ + 0, + 0, + 0 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 2, + "stride": 2, + "padding": 0, + "bias": true, + "sites": [ + "C:dec3 (up4)#59" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "measured": [ + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 0.8154, + "ms_per_call": 0.8154, + "calls": 1, + "pct_roofline": 24.9, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "C", + "direction": "bwd-weight", + "ms_per_step": 0.5538, + "ms_per_call": 0.5538, + "calls": 1, + "pct_roofline": 36.7, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 0.5073, + "ms_per_call": 0.5073, + "calls": 1, + "pct_roofline": 40.1, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 1024, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 1024, + 16, + 16, + 16 + ], + "out_shape": [ + 1, + 512, + 16, + 16, + 16 + ], + "halo_in_shape": [ + 1, + 1024, + 18, + 18, + 18 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "A:dec0 (up1)#37" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "measured": [ + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 0.7961, + "ms_per_call": 0.7961, + "calls": 1, + "pct_roofline": 24.3, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "A", + "direction": "fwd", + "ms_per_step": 0.5311, + "ms_per_call": 0.5311, + "calls": 1, + "pct_roofline": 36.4, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + }, + { + "config": "A", + "direction": "bwd-weight", + "ms_per_step": 0.4595, + "ms_per_call": 0.4595, + "calls": 1, + "pct_roofline": 42.1, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 16, + 64, + 64 + ], + "out_shape": [ + 1, + 256, + 16, + 64, + 64 + ], + "halo_in_shape": [ + 1, + 128, + 18, + 66, + 66 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "C:enc2 (down2)#15" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "measured": [ + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 0.7486, + "ms_per_call": 0.7486, + "calls": 1, + "pct_roofline": 25.8, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "C", + "direction": "bwd-weight", + "ms_per_step": 0.5563, + "ms_per_call": 0.5563, + "calls": 1, + "pct_roofline": 34.7, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "C", + "direction": "fwd", + "ms_per_step": 0.4106, + "ms_per_call": 0.4106, + "calls": 1, + "pct_roofline": 47.1, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_abd_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 16, + 32, + 32 + ], + "out_shape": [ + 1, + 512, + 16, + 32, + 32 + ], + "halo_in_shape": [ + 1, + 256, + 18, + 34, + 34 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "B:enc3 (down3)#22" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "measured": [ + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 0.7796, + "ms_per_call": 0.7796, + "calls": 1, + "pct_roofline": 24.8, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "B", + "direction": "bwd-weight", + "ms_per_step": 0.4954, + "ms_per_call": 0.4954, + "calls": 1, + "pct_roofline": 39.0, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "B", + "direction": "fwd", + "ms_per_step": 0.4096, + "ms_per_call": 0.4096, + "calls": 1, + "pct_roofline": 47.2, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3_2lds" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 128, + 64, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 64, + 64, + 64, + 64 + ], + "out_shape": [ + 1, + 128, + 64, + 64, + 64 + ], + "halo_in_shape": [ + 1, + 64, + 66, + 66, + 66 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "A:enc1 (down1)#8" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "measured": [ + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 0.8504, + "ms_per_call": 0.8504, + "calls": 1, + "pct_roofline": 22.7, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "A", + "direction": "bwd-weight", + "ms_per_step": 0.4766, + "ms_per_call": 0.4766, + "calls": 1, + "pct_roofline": 40.6, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "A", + "direction": "fwd", + "ms_per_step": 0.2942, + "ms_per_call": 0.2942, + "calls": 1, + "pct_roofline": 65.7, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 1024, + 1024, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 1024, + 4, + 16, + 16 + ], + "out_shape": [ + 1, + 1024, + 4, + 16, + 16 + ], + "halo_in_shape": [ + 1, + 1024, + 6, + 18, + 18 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "C:bottleneck (down4)#32" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "measured": [ + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 0.9524, + "ms_per_call": 0.9524, + "calls": 1, + "pct_roofline": 10.1, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "C", + "direction": "bwd-weight", + "ms_per_step": 0.3211, + "ms_per_call": 0.3211, + "calls": 1, + "pct_roofline": 30.1, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "C", + "direction": "fwd", + "ms_per_step": 0.3151, + "ms_per_call": 0.3151, + "calls": 1, + "pct_roofline": 30.7, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 64, + 3, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 3, + 128, + 128, + 128 + ], + "out_shape": [ + 1, + 64, + 128, + 128, + 128 + ], + "halo_in_shape": [ + 1, + 3, + 130, + 130, + 130 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "A:enc0 (inc)#1" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "measured": [ + { + "config": "A", + "direction": "bwd-weight", + "ms_per_step": 0.6971, + "ms_per_call": 0.6971, + "calls": 1, + "pct_roofline": 12.2, + "solvers": [ + "kernel_grouped_conv_bwd_weight_xdl_cshuffle_v3" + ] + }, + { + "config": "A", + "direction": "fwd", + "ms_per_step": 0.5992, + "ms_per_call": 0.5992, + "calls": 1, + "pct_roofline": 14.2, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_abd_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 1024, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 8, + 16, + 16 + ], + "out_shape": [ + 1, + 1024, + 8, + 16, + 16 + ], + "halo_in_shape": [ + 1, + 512, + 10, + 18, + 18 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "B:bottleneck (down4)#29" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "measured": [ + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 0.7379, + "ms_per_call": 0.7379, + "calls": 1, + "pct_roofline": 13.1, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "B", + "direction": "bwd-weight", + "ms_per_step": 0.3073, + "ms_per_call": 0.3073, + "calls": 1, + "pct_roofline": 31.4, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "B", + "direction": "fwd", + "ms_per_step": 0.2272, + "ms_per_call": 0.2272, + "calls": 1, + "pct_roofline": 42.5, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 256, + 128, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 256, + 32, + 64, + 64 + ], + "out_shape": [ + 1, + 128, + 64, + 128, + 128 + ], + "halo_in_shape": [ + 1, + 256, + 32, + 64, + 64 + ], + "halo_dhw": [ + 0, + 0, + 0 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 2, + "stride": 2, + "padding": 0, + "bias": true, + "sites": [ + "B:dec2 (up3)#51" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "measured": [ + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 0.4457, + "ms_per_call": 0.4457, + "calls": 1, + "pct_roofline": 25.7, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "B", + "direction": "bwd-weight", + "ms_per_step": 0.4407, + "ms_per_call": 0.4407, + "calls": 1, + "pct_roofline": 26.0, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 0.3786, + "ms_per_call": 0.3786, + "calls": 1, + "pct_roofline": 30.3, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_abd_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 8, + 32, + 32 + ], + "out_shape": [ + 1, + 512, + 8, + 32, + 32 + ], + "halo_in_shape": [ + 1, + 256, + 10, + 34, + 34 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "C:enc3 (down3)#22" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "measured": [ + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 0.7077, + "ms_per_call": 0.7077, + "calls": 1, + "pct_roofline": 13.7, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "C", + "direction": "bwd-weight", + "ms_per_step": 0.2882, + "ms_per_call": 0.2882, + "calls": 1, + "pct_roofline": 33.5, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "C", + "direction": "fwd", + "ms_per_step": 0.2191, + "ms_per_call": 0.2191, + "calls": 1, + "pct_roofline": 44.1, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 6, + 64, + 1, + 1, + 1 + ], + "in_shape": [ + 1, + 64, + 64, + 256, + 256 + ], + "out_shape": [ + 1, + 6, + 64, + 256, + 256 + ], + "halo_in_shape": [ + 1, + 64, + 64, + 256, + 256 + ], + "halo_dhw": [ + 0, + 0, + 0 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 1, + "stride": 1, + "padding": 0, + "bias": true, + "sites": [ + "C:outc#67" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "measured": [ + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 0.8846, + "ms_per_call": 0.8846, + "calls": 1, + "pct_roofline": 20.1, + "solvers": [ + "kernel_batched_gemm_xdl_cshuffle_v3_multi_d", + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "C", + "direction": "fwd", + "ms_per_step": 0.279, + "ms_per_call": 0.279, + "calls": 1, + "pct_roofline": 63.8, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 1024, + 1024, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 1024, + 8, + 8, + 8 + ], + "out_shape": [ + 1, + 1024, + 8, + 8, + 8 + ], + "halo_in_shape": [ + 1, + 1024, + 10, + 10, + 10 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "A:bottleneck (down4)#32" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "measured": [ + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 0.6021, + "ms_per_call": 0.6021, + "calls": 1, + "pct_roofline": 8.0, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "A", + "direction": "bwd-weight", + "ms_per_step": 0.2279, + "ms_per_call": 0.2279, + "calls": 1, + "pct_roofline": 21.2, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "A", + "direction": "fwd", + "ms_per_step": 0.2267, + "ms_per_call": 0.2267, + "calls": 1, + "pct_roofline": 21.3, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 256, + 128, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 128, + 32, + 32, + 32 + ], + "out_shape": [ + 1, + 256, + 32, + 32, + 32 + ], + "halo_in_shape": [ + 1, + 128, + 34, + 34, + 34 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "A:enc2 (down2)#15" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "measured": [ + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 0.4657, + "ms_per_call": 0.4657, + "calls": 1, + "pct_roofline": 20.7, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "A", + "direction": "bwd-weight", + "ms_per_step": 0.3011, + "ms_per_call": 0.3011, + "calls": 1, + "pct_roofline": 32.1, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "A", + "direction": "fwd", + "ms_per_step": 0.2263, + "ms_per_call": 0.2263, + "calls": 1, + "pct_roofline": 42.7, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3_2lds" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 1024, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 4, + 16, + 16 + ], + "out_shape": [ + 1, + 1024, + 4, + 16, + 16 + ], + "halo_in_shape": [ + 1, + 512, + 6, + 18, + 18 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 1, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "C:bottleneck (down4)#29" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "measured": [ + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 0.5983, + "ms_per_call": 0.5983, + "calls": 1, + "pct_roofline": 8.1, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "C", + "direction": "bwd-weight", + "ms_per_step": 0.1656, + "ms_per_call": 0.1656, + "calls": 1, + "pct_roofline": 29.2, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "C", + "direction": "fwd", + "ms_per_step": 0.1532, + "ms_per_call": 0.1532, + "calls": 1, + "pct_roofline": 31.5, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 128, + 64, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 128, + 64, + 64, + 64 + ], + "out_shape": [ + 1, + 64, + 128, + 128, + 128 + ], + "halo_in_shape": [ + 1, + 128, + 64, + 64, + 64 + ], + "halo_dhw": [ + 0, + 0, + 0 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 2, + "stride": 2, + "padding": 0, + "bias": true, + "sites": [ + "A:dec3 (up4)#59" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "measured": [ + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 0.4025, + "ms_per_call": 0.4025, + "calls": 1, + "pct_roofline": 25.3, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "A", + "direction": "bwd-weight", + "ms_per_step": 0.2411, + "ms_per_call": 0.2411, + "calls": 1, + "pct_roofline": 42.2, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 0.2042, + "ms_per_call": 0.2042, + "calls": 1, + "pct_roofline": 49.8, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_abd_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 1024, + 512, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 512, + 8, + 8, + 8 + ], + "out_shape": [ + 1, + 1024, + 8, + 8, + 8 + ], + "halo_in_shape": [ + 1, + 512, + 10, + 10, + 10 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "A:bottleneck (down4)#29" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "measured": [ + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 0.4172, + "ms_per_call": 0.4172, + "calls": 1, + "pct_roofline": 5.8, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "A", + "direction": "bwd-weight", + "ms_per_step": 0.117, + "ms_per_call": 0.117, + "calls": 1, + "pct_roofline": 20.7, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "A", + "direction": "fwd", + "ms_per_step": 0.1112, + "ms_per_call": 0.1112, + "calls": 1, + "pct_roofline": 21.7, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 256, + 128, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 256, + 16, + 64, + 64 + ], + "out_shape": [ + 1, + 128, + 32, + 128, + 128 + ], + "halo_in_shape": [ + 1, + 256, + 16, + 64, + 64 + ], + "halo_dhw": [ + 0, + 0, + 0 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 2, + "stride": 2, + "padding": 0, + "bias": true, + "sites": [ + "C:dec2 (up3)#51" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "measured": [ + { + "config": "C", + "direction": "bwd-weight", + "ms_per_step": 0.2422, + "ms_per_call": 0.2422, + "calls": 1, + "pct_roofline": 23.6, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 0.2274, + "ms_per_call": 0.2274, + "calls": 1, + "pct_roofline": 25.2, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 0.1592, + "ms_per_call": 0.1592, + "calls": 1, + "pct_roofline": 36.0, + "solvers": [ + "kernel_grouped_conv_fwd_multiple_abd_xdl_cshuffle" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 512, + 256, + 3, + 3, + 3 + ], + "in_shape": [ + 1, + 256, + 16, + 16, + 16 + ], + "out_shape": [ + 1, + 512, + 16, + 16, + 16 + ], + "halo_in_shape": [ + 1, + 256, + 18, + 18, + 18 + ], + "halo_dhw": [ + 1, + 1, + 1 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 3, + "stride": 1, + "padding": 1, + "bias": false, + "sites": [ + "A:enc3 (down3)#22" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "measured": [ + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 0.3144, + "ms_per_call": 0.3144, + "calls": 1, + "pct_roofline": 15.4, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "A", + "direction": "bwd-weight", + "ms_per_step": 0.1559, + "ms_per_call": 0.1559, + "calls": 1, + "pct_roofline": 31.0, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "A", + "direction": "fwd", + "ms_per_step": 0.1251, + "ms_per_call": 0.1251, + "calls": 1, + "pct_roofline": 38.6, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 512, + 256, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 512, + 16, + 32, + 32 + ], + "out_shape": [ + 1, + 256, + 32, + 64, + 64 + ], + "halo_in_shape": [ + 1, + 512, + 16, + 32, + 32 + ], + "halo_dhw": [ + 0, + 0, + 0 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 2, + "stride": 2, + "padding": 0, + "bias": true, + "sites": [ + "B:dec1 (up2)#43" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "measured": [ + { + "config": "B", + "direction": "bwd-weight", + "ms_per_step": 0.2355, + "ms_per_call": 0.2355, + "calls": 1, + "pct_roofline": 24.3, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 0.1843, + "ms_per_call": 0.1843, + "calls": 1, + "pct_roofline": 31.1, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 0.1434, + "ms_per_call": 0.1434, + "calls": 1, + "pct_roofline": 39.9, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3_2lds" + ] + } + ] + }, + { + "op": "Conv3d", + "weight_shape": [ + 6, + 64, + 1, + 1, + 1 + ], + "in_shape": [ + 1, + 64, + 128, + 128, + 128 + ], + "out_shape": [ + 1, + 6, + 128, + 128, + 128 + ], + "halo_in_shape": [ + 1, + 64, + 128, + 128, + 128 + ], + "halo_dhw": [ + 0, + 0, + 0 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 1, + "stride": 1, + "padding": 0, + "bias": true, + "sites": [ + "A:outc#67" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "measured": [ + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 0.4321, + "ms_per_call": 0.4321, + "calls": 1, + "pct_roofline": 20.6, + "solvers": [ + "kernel_batched_gemm_xdl_cshuffle_v3_multi_d", + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "A", + "direction": "fwd", + "ms_per_step": 0.0981, + "ms_per_call": 0.0981, + "calls": 1, + "pct_roofline": 90.7, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 256, + 128, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 256, + 32, + 32, + 32 + ], + "out_shape": [ + 1, + 128, + 64, + 64, + 64 + ], + "halo_in_shape": [ + 1, + 256, + 32, + 32, + 32 + ], + "halo_dhw": [ + 0, + 0, + 0 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 2, + "stride": 2, + "padding": 0, + "bias": true, + "sites": [ + "A:dec2 (up3)#51" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "measured": [ + { + "config": "A", + "direction": "bwd-weight", + "ms_per_step": 0.175, + "ms_per_call": 0.175, + "calls": 1, + "pct_roofline": 16.4, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 0.1207, + "ms_per_call": 0.1207, + "calls": 1, + "pct_roofline": 23.7, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 0.0907, + "ms_per_call": 0.0907, + "calls": 1, + "pct_roofline": 31.6, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3_2lds" + ] + } + ] + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 512, + 256, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 512, + 8, + 32, + 32 + ], + "out_shape": [ + 1, + 256, + 16, + 64, + 64 + ], + "halo_in_shape": [ + 1, + 512, + 8, + 32, + 32 + ], + "halo_dhw": [ + 0, + 0, + 0 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 2, + "stride": 2, + "padding": 0, + "bias": true, + "sites": [ + "C:dec1 (up2)#43" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "measured": [ + { + "config": "C", + "direction": "bwd-weight", + "ms_per_step": 0.136, + "ms_per_call": 0.136, + "calls": 1, + "pct_roofline": 21.1, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 0.1238, + "ms_per_call": 0.1238, + "calls": 1, + "pct_roofline": 23.1, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 0.0872, + "ms_per_call": 0.0872, + "calls": 1, + "pct_roofline": 32.8, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 1024, + 512, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 1024, + 8, + 16, + 16 + ], + "out_shape": [ + 1, + 512, + 16, + 32, + 32 + ], + "halo_in_shape": [ + 1, + 1024, + 8, + 16, + 16 + ], + "halo_dhw": [ + 0, + 0, + 0 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 2, + "stride": 2, + "padding": 0, + "bias": true, + "sites": [ + "B:dec0 (up1)#35" + ], + "configs": [ + "scale 8, 2 GPUs, shards (2,1,1)" + ], + "measured": [ + { + "config": "B", + "direction": "bwd-weight", + "ms_per_step": 0.1488, + "ms_per_call": 0.1488, + "calls": 1, + "pct_roofline": 19.2, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 0.1153, + "ms_per_call": 0.1153, + "calls": 1, + "pct_roofline": 24.8, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "B", + "direction": "bwd-data", + "ms_per_step": 0.0816, + "ms_per_call": 0.0816, + "calls": 1, + "pct_roofline": 35.1, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 1024, + 512, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 1024, + 4, + 16, + 16 + ], + "out_shape": [ + 1, + 512, + 8, + 32, + 32 + ], + "halo_in_shape": [ + 1, + 1024, + 4, + 16, + 16 + ], + "halo_dhw": [ + 0, + 0, + 0 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 2, + "stride": 2, + "padding": 0, + "bias": true, + "sites": [ + "C:dec0 (up1)#35" + ], + "configs": [ + "scale 8, 4 GPUs, shards (4,1,1)" + ], + "measured": [ + { + "config": "C", + "direction": "bwd-weight", + "ms_per_step": 0.0815, + "ms_per_call": 0.0815, + "calls": 1, + "pct_roofline": 17.6, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 0.0733, + "ms_per_call": 0.0733, + "calls": 1, + "pct_roofline": 19.5, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "C", + "direction": "bwd-data", + "ms_per_step": 0.0611, + "ms_per_call": 0.0611, + "calls": 1, + "pct_roofline": 23.4, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 512, + 256, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 512, + 16, + 16, + 16 + ], + "out_shape": [ + 1, + 256, + 32, + 32, + 32 + ], + "halo_in_shape": [ + 1, + 512, + 16, + 16, + 16 + ], + "halo_dhw": [ + 0, + 0, + 0 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 2, + "stride": 2, + "padding": 0, + "bias": true, + "sites": [ + "A:dec1 (up2)#43" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "measured": [ + { + "config": "A", + "direction": "bwd-weight", + "ms_per_step": 0.0821, + "ms_per_call": 0.0821, + "calls": 1, + "pct_roofline": 17.4, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 0.065, + "ms_per_call": 0.065, + "calls": 1, + "pct_roofline": 22.0, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 0.0458, + "ms_per_call": 0.0458, + "calls": 1, + "pct_roofline": 31.3, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + }, + { + "op": "ConvTranspose3d", + "weight_shape": [ + 1024, + 512, + 2, + 2, + 2 + ], + "in_shape": [ + 1, + 1024, + 8, + 8, + 8 + ], + "out_shape": [ + 1, + 512, + 16, + 16, + 16 + ], + "halo_in_shape": [ + 1, + 1024, + 8, + 8, + 8 + ], + "halo_dhw": [ + 0, + 0, + 0 + ], + "shard_halo_dhw": [ + 0, + 0, + 0 + ], + "k": 2, + "stride": 2, + "padding": 0, + "bias": true, + "sites": [ + "A:dec0 (up1)#35" + ], + "configs": [ + "scale 7, 1 GPU, shards (1,1,1)" + ], + "measured": [ + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 0.056, + "ms_per_call": 0.056, + "calls": 1, + "pct_roofline": 12.8, + "solvers": [ + "kernel_grouped_conv_bwd_data_multiple_d_xdl_cshuffle" + ] + }, + { + "config": "A", + "direction": "bwd-weight", + "ms_per_step": 0.05, + "ms_per_call": 0.05, + "calls": 1, + "pct_roofline": 14.3, + "solvers": [ + "kernel_batched_gemm_xdlops_bwd_weight" + ] + }, + { + "config": "A", + "direction": "bwd-data", + "ms_per_step": 0.043, + "ms_per_call": 0.043, + "calls": 1, + "pct_roofline": 16.6, + "solvers": [ + "kernel_grouped_conv_fwd_xdl_cshuffle_v3" + ] + } + ] + } + ] +} diff --git a/triton_conv3d/shapes.py b/triton_conv3d/shapes.py new file mode 100644 index 0000000..0cd3ff1 --- /dev/null +++ b/triton_conv3d/shapes.py @@ -0,0 +1,735 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""The convolution problems the kernels have to serve. + +Three sources: + +``scaffold_corpus()`` + The 57 distinct convolutions that occur in real ScaFFold runs, extracted + from ``model-analysis/unet_shapes.py`` dumps for the three profiled + configurations (scale 7 on 1 GPU, scale 8 sharded over 2 and over 4 GPUs) + and joined with the measured MIOpen cost and roofline efficiency of each. + This is the tuning target and the priority list. + +``census_corpus()`` + Every convolution an **instrumented ScaFFold training step actually + issued**, at the four configurations the current benchmark harness runs. + Recorded by wrapping the entry points inside a running step + (``work/triton-conv/bin/conv_census.py``), so it is a measurement rather + than a model. It carries no MIOpen timings; it exists to say what the + shapes *are*. + +``edge_cases()`` + Synthetic problems chosen to break addressing, masking and tiling + assumptions: channel counts that are not multiples of the MFMA granularity, + prime spatial extents, anisotropic volumes, batches, and volumes whose + linear element index exceeds 2**31. + +All return :class:`ConvProblem`, which knows how to derive the tensor shapes, +the FLOP and byte counts, and the implicit-GEMM shape for each of the three +directions. Nothing here imports torch, so it is cheap to introspect and can +drive test parametrization at collection time. + +One problem, three forms +======================== + +A single ScaFFold convolution reaches a kernel in three different shapes +depending on *who* issues it, and they are three different tuning problems -- +MIOpen keys its find database on the padding, this package's ``bwd_data_config`` +derives ``M`` from it, and every one of them is a different *measurement*. +Confusing them has already cost this project one wrong projection, so +:attr:`ConvProblem.form` names which one an instance is and every corpus +accessor says which it returns. (``bwd_weight_config`` used to change its +answer on the padding as well; that clause went on 2026-08-05 and the forms are +still three problems without it.) + +``"logical"`` + The convolution as ``unet_parts.py`` states it: the local shard at the + module's own padding, ``k // 2`` on every axis. This is what + :func:`scaffold_corpus` stores. + +``"distconv"`` + What upstream DistConv hands the backend: it concatenates a ``k // 2`` halo + on **every dimension listed in ``dc_shard_dims``**, split or not, and zeroes + the padding there. :attr:`ConvProblem.halo_variant`. This is the form the + MIOpen baseline in ``measured`` was profiled in, and it is the incumbent's + problem. + +``"adapter"`` + What ``ScaFFold/unet/conv3d.py`` -- the shipped Triton rung -- hands the + kernel: it exchanges a halo only on axes that are *genuinely split*, so at + ``dc_num_shards = (1, 1, 1)`` nothing is exchanged and the convolution is + padded on all three axes, and at ``(2, 1, 1)`` or ``(4, 1, 1)`` only D is + halo'd while H and W stay padded. :attr:`ConvProblem.production_variant`. + **This is the form production actually runs today**, at every + configuration, and it is padded at every one of them. +""" + +from __future__ import annotations + +import dataclasses +import functools +import json +import math +import pathlib +from typing import Iterator, Literal, Sequence + +_CORPUS_PATH = pathlib.Path(__file__).resolve().parent / "scaffold_corpus.json" +_CENSUS_PATH = pathlib.Path(__file__).resolve().parent / "scaffold_census.json" + +Direction = Literal["fwd", "bwd-data", "bwd-weight"] +DIRECTIONS: tuple[Direction, ...] = ("fwd", "bwd-data", "bwd-weight") + +#: Which of the three statements of a problem an instance is. See the module +#: docstring; the short version is that they differ in the padding and in the +#: input extent, and that ``"adapter"`` is the one production runs. +Form = Literal["logical", "distconv", "adapter"] +FORMS: tuple[Form, ...] = ("logical", "distconv", "adapter") + +#: Empirical MI300A roofline constants, supplied rather than derived. They are +#: measured ceilings: a kernel that exceeds one is a fact about the measurement, +#: not necessarily an error. +HBM_BYTES_PER_S = 3.3e12 +PEAK_FLOPS = {"fp32": 82.6e12, "bf16": 600e12, "fp16": 600e12} +ELEM_BYTES = {"fp32": 4, "bf16": 2, "fp16": 2} + +#: Largest linear element index representable in int32. +INT32_MAX = 2**31 - 1 + +#: Largest *storage*, in bytes, that an AMD buffer instruction can address. +#: Triton's runtime specializer tags a pointer argument ``tt.pointer_range = 32`` +#: -- the flag that lets the backend emit ``buffer_load_dwordx4`` rather than +#: ``global_load_dwordx4`` -- from ``arg.untyped_storage().size() <= 2**31 - 1``. +#: That check reads the whole *storage* and it counts *bytes*, so it is a +#: different question from whether an element index overflows int32, and the two +#: answers differ by ``elem_bytes`` (``PLAN.md`` §3.5). +BUFFER_OP_MAX_BYTES = 2**31 - 1 + + +def _prod(xs: Sequence[int]) -> int: + return math.prod(xs) + + +@dataclasses.dataclass(frozen=True) +class ConvProblem: + """One convolution, in the form a kernel is tuned for. + + ``spatial`` is always the *input* volume, ``(D, H, W)``. For the transposed + case ``cin``/``cout`` keep their logical meaning (the operator maps ``cin`` + channels to ``cout``), which is the transpose of how PyTorch stores the + weight -- :meth:`weight_shape` accounts for that. + """ + + name: str + cin: int + cout: int + spatial: tuple[int, int, int] + kernel: tuple[int, int, int] = (3, 3, 3) + stride: tuple[int, int, int] = (1, 1, 1) + padding: tuple[int, int, int] = (1, 1, 1) + n: int = 1 + transposed: bool = False + bias: bool = False + dtype: str = "bf16" + #: Where this problem came from: ScaFFold site tags, or "synthetic". + sites: tuple[str, ...] = () + #: Measured MIOpen results, if any: one dict per (config, direction). + measured: tuple[dict, ...] = () + #: Set for problems that need a lot of memory or a long time; opt-in. + large: bool = False + #: **Upstream DistConv's** halo width per spatial dim: ``k // 2`` on every + #: dim listed in ``dc_shard_dims``, whether or not that dim is actually + #: split. Non-zero means the *MIOpen* rung -- which goes through + #: ``distconv_forward`` -- runs the convolution unpadded at a larger extent; + #: see :meth:`halo_variant`. It says nothing about the Triton rung, which + #: exchanges only what :attr:`shard_halo` records. Zero for synthetic + #: problems. + halo: tuple[int, int, int] = (0, 0, 0) + #: **The ScaFFold Triton adapter's** halo width per spatial dim: ``k // 2`` + #: on the dims that are genuinely split (``dc_num_shards > 1``) and zero + #: elsewhere, because ``ScaFFold/unet/conv3d.py``'s ``_halo_plan`` skips an + #: unsplit axis and leaves the module's own padding on it. This is what + #: separates the production form from DistConv's -- see + #: :meth:`production_variant`. + shard_halo: tuple[int, int, int] = (0, 0, 0) + #: Which of the three statements of the problem this instance is. Set by + #: :meth:`halo_variant` and :meth:`production_variant`; ``"logical"`` + #: otherwise. Carried so that a table row, a benchmark cell and a JSON + #: record all say which shape they measured instead of leaving it to be + #: inferred from the padding. + form: Form = "logical" + + # -- derived shapes --------------------------------------------------- + + @functools.cached_property + def out_spatial(self) -> tuple[int, int, int]: + if self.transposed: + return tuple( + (i - 1) * s - 2 * p + k + for i, k, s, p in zip( + self.spatial, self.kernel, self.stride, self.padding + ) + ) + return tuple( + (i + 2 * p - k) // s + 1 + for i, k, s, p in zip(self.spatial, self.kernel, self.stride, self.padding) + ) + + @property + def input_shape(self) -> tuple[int, ...]: + return (self.n, self.cin, *self.spatial) + + @property + def output_shape(self) -> tuple[int, ...]: + return (self.n, self.cout, *self.out_spatial) + + @property + def weight_shape(self) -> tuple[int, ...]: + """PyTorch's storage order, which differs between the two operators.""" + if self.transposed: + return (self.cin, self.cout, *self.kernel) + return (self.cout, self.cin, *self.kernel) + + @property + def elem_bytes(self) -> int: + return ELEM_BYTES[self.dtype] + + @property + def halo_variant(self) -> "ConvProblem": + """The same convolution in the form **upstream DistConv** issues it. + + ``distconv_forward`` never lets PyTorch pad a convolution on a dimension + it manages. It concatenates a halo slab of width ``k // 2`` onto both + faces -- neighbour data, or zeros at the mesh boundary and at one shard + -- and then sets that dimension's padding to zero (``distconv.py``). + It does this for **every dim in ``dc_shard_dims``**, including dims with + a single shard, where the slab is provably zeros; that is why + :attr:`halo` is ``(1, 1, 1)`` on every ``k = 3`` corpus problem even at + one GPU. So the tensor MIOpen sees is two voxels larger per listed axis + and the convolution is unpadded. + + **This is the incumbent's form, not the shipped one.** ScaFFold routes + its convolutions through ``ScaFFold/unet/conv3d.py``, which performs the + exchange itself and only on axes that are genuinely split -- see + :meth:`production_variant`. The MIOpen numbers in :attr:`measured` were + profiled through DistConv, so they *are* timings of this form, and it is + the right form to compare an MIOpen baseline in; it is not the form the + Triton kernels are handed. + + This distinction is not cosmetic. MIOpen keys its find database on the + whole problem descriptor, padding included: ``64-128-128-128-...-1x1x1`` + and ``64-130-130-130-...-0x0x0`` are two different problems that tune + independently and can land on different kernels. This package used to + change its own answer on it too -- :func:`~triton_conv3d.reduce_gemm. + bwd_weight_config` declined a tuned row with ``TAP_BLOCK > 1`` on a + padded convolution until 2026-08-05 -- and no longer does; but the two + forms remain different *problems*, they are timed differently, and + :func:`~triton_conv3d.bwd_data.bwd_data_config` still reads the padding + because it derives ``M`` from it. + + The cost model follows the shape rather than being special-cased: the + halo'd form genuinely reads a slightly larger input and genuinely + produces a slightly larger input gradient, and :meth:`flops` and + :meth:`bytes` say so because they are derived from ``spatial``. + Returns ``self`` when there is no halo -- the three forms genuinely + coincide there, and returning an unequal copy would only make a + distinction the problem does not have. + """ + if not any(self.halo): + return self + return dataclasses.replace( + self, + name=f"{self.name}+halo" if self.name else "halo", + spatial=tuple(s + 2 * h for s, h in zip(self.spatial, self.halo)), + padding=tuple(0 if h else p for h, p in zip(self.halo, self.padding)), + halo=(0, 0, 0), + shard_halo=(0, 0, 0), + form="distconv", + ) + + @property + def production_variant(self) -> "ConvProblem": + """The same convolution in the form **ScaFFold runs it today**. + + ``ScaFFold/unet/conv3d.py``'s ``_halo_plan`` walks the parallel + strategy and ``continue``s past any axis with a single shard, so it + exchanges a halo *only* on axes that are genuinely split and leaves the + module's own padding on every other one. ScaFFold ships + ``dc_shard_dims: [2, 3, 4]`` with ``dc_num_shards`` of ``[1,1,1]``, + ``[2,1,1]`` or ``[4,1,1]``, so: + + * unsharded, nothing is exchanged and the convolution reaches the kernel + at its logical extent with ``padding = (1, 1, 1)``; + * sharded, D is halo'd and H and W are still padded -- + ``padding = (0, 1, 1)`` at ``(D_loc + 2, H, W)``. + + **Every production convolution with ``k > 1`` is therefore padded, at + every configuration.** Measured inside running steps at all four + (``work/triton-conv/review/SHAPE_AUDIT.md``), not inferred: 18 of the 19 + distinct ordinary convolutions at scale 7 on one GPU arrive with + ``padding = (1, 1, 1)``, and the nineteenth is the ``k = 1`` head, which + has no padding to begin with. + + Dropping the zero slabs on the unsplit axes is a deliberate and + separately verified decision -- ``cat(zeros, x, zeros)`` at + ``padding = 0`` is the same arithmetic as ``padding = k // 2`` on ``x``, + and it is measured bitwise identical through these kernels -- so this is + not a divergence to be repaired but the shape to be tuned for. + + Returns ``self`` when nothing is split, for the same reason + :meth:`halo_variant` does: unsharded, the logical statement *is* what + the adapter issues, and there is no distinction to record. + """ + if not any(self.shard_halo): + return self + return dataclasses.replace( + self, + name=f"{self.name}+shard" if self.name else "shard", + spatial=tuple(s + 2 * h + for s, h in zip(self.spatial, self.shard_halo)), + padding=tuple(0 if h else p + for h, p in zip(self.shard_halo, self.padding)), + halo=(0, 0, 0), + shard_halo=(0, 0, 0), + form="adapter", + ) + + # -- cost model ------------------------------------------------------- + + @property + def tap_count(self) -> int: + return _prod(self.kernel) + + def flops(self, direction: Direction = "fwd") -> int: + """Multiply-accumulate count x2. + + All three directions perform the same contraction with different operands + held fixed, so the count differs only in which volume indexes it. For a + forward convolution each *output* voxel gathers ``taps`` contributions; + backward-data is the same contraction over the *input* volume. + + The transposed operator scatters instead of gathering, so every direction + is indexed by the *input* volume -- and with ``kernel == stride`` that + makes the tap factor illusory: each output voxel receives exactly one + contribution, because the windows tile rather than overlap. + """ + if self.transposed: + vol = _prod(self.spatial) + else: + vol = _prod(self.spatial if direction == "bwd-data" else self.out_spatial) + return 2 * self.n * vol * self.cin * self.cout * self.tap_count + + def bytes(self, direction: Direction = "fwd") -> int: + """Compulsory traffic: each tensor the direction touches, read once. + + This is the denominator of the memory roof. It credits the kernel with + perfect reuse -- no im2col materialization, no partial spilling -- which + is exactly the standard a fused implicit-GEMM kernel should be held to. + """ + eb = self.elem_bytes + x = self.n * self.cin * _prod(self.spatial) * eb + y = self.n * self.cout * _prod(self.out_spatial) * eb + w = self.cin * self.cout * self.tap_count * eb + return {"fwd": x + w + y, "bwd-data": y + w + x, "bwd-weight": y + x + w}[ + direction + ] + + def arithmetic_intensity(self, direction: Direction = "fwd") -> float: + return self.flops(direction) / self.bytes(direction) + + def roofline_flops(self, direction: Direction = "fwd") -> float: + """Attainable FLOP/s: whichever of compute and bandwidth binds first.""" + return min( + PEAK_FLOPS[self.dtype], + self.arithmetic_intensity(direction) * HBM_BYTES_PER_S, + ) + + def efficiency(self, ms: float, direction: Direction = "fwd") -> float: + """Fraction of the roofline achieved by a measured time in milliseconds.""" + return (self.flops(direction) / (ms * 1e-3)) / self.roofline_flops(direction) + + # -- implicit-GEMM decomposition ------------------------------------- + + def gemm_shape(self, direction: Direction = "fwd") -> tuple[int, int, int]: + """``(M, N, K)`` of the GEMM this direction reduces to. + + Forward and backward-data tile over a volume with the channel count as N + and the taps folded into K. Backward-weight is the transpose of that + situation: a tiny output reduced over the whole volume, which is why it + needs split-K and why determinism is a live question there. + + The transposed operator with ``kernel == stride`` and no padding is a + special case -- a pointwise GEMM producing ``cout * taps`` channels, + followed by a voxel shuffle -- so its forward K carries no tap factor. + """ + taps = self.tap_count + if self.transposed: + if self.kernel != self.stride or set(self.padding) != {0}: + raise NotImplementedError( + "transposed convolutions are only decomposed for " + f"kernel == stride and no padding; got kernel={self.kernel}, " + f"stride={self.stride}, padding={self.padding}" + ) + in_vol = self.n * _prod(self.spatial) + if direction == "fwd": + return (in_vol, self.cout * taps, self.cin) + if direction == "bwd-data": + return (in_vol, self.cin, self.cout * taps) + return (self.cin, self.cout * taps, in_vol) + out_vol = self.n * _prod(self.out_spatial) + in_vol = self.n * _prod(self.spatial) + if direction == "fwd": + return (out_vol, self.cout, self.cin * taps) + if direction == "bwd-data": + return (in_vol, self.cin, self.cout * taps) + return (self.cout, self.cin * taps, out_vol) + + # -- indexing --------------------------------------------------------- + + @property + def max_elements(self) -> int: + """Element count of the larger activation. + + The largest linear index a pointer into it will see is therefore + ``max_elements - 1``. + """ + return max( + self.n * self.cin * _prod(self.spatial), + self.n * self.cout * _prod(self.out_spatial), + ) + + @property + def max_activation_bytes(self) -> int: + """Storage of the larger activation, in bytes. + + This -- not :attr:`max_elements` -- is the quantity the AMD backend + cares about, and it is ``elem_bytes`` times larger. + """ + return self.max_elements * self.elem_bytes + + @property + def index_exceeds_int32(self) -> bool: + """The kernel's *element* offsets must be widened to int64. + + Counted in elements because that is what a Triton offset holds: the + largest one is ``max_elements - 1``, so the boundary sits at ``2**31`` + elements and not at ``INT32_MAX``. (The old form compared + ``max_elements > INT32_MAX``, which fires one element early -- harmless, + but it made the predicate hard to reason about at the boundary the + edge cases exist to pin.) + + This is emphatically **not** the 2 GiB cliff. No corpus problem + reaches it in either shape mode, including the cliff cell itself: + ``conv 128->64 k3 @ 130x258x258`` holds 1.108e9 elements -- half of + int32's range -- in 2.22 GiB of storage. What that shape loses is + buffer ops, which is :attr:`buffer_ops_eligible`. + """ + return self.max_elements - 1 > INT32_MAX + + #: The name ``bench/baseline.py`` records this predicate under, and so the + #: name it carries in every row of ``baseline.json``. Kept as an alias + #: rather than renamed in place, because the field is published data. + needs_int64 = index_exceeds_int32 + + @property + def buffer_ops_eligible(self) -> bool: + """The larger activation still fits the buffer-load fast path. + + False costs about 4.5% on the shapes we measured it on (M1), and it is + the property that separates the corpus's one cliff cell from the rest of + it: a 2.22 GiB activation is 3.2% over the byte limit while being + nowhere near the *element* limit. Modelled from the shape, so it + assumes a freshly allocated tensor -- a narrowed view keeps its parent's + storage and ``conv_bench`` therefore measures the same predicate off + ``untyped_storage().size()`` instead. + """ + return self.max_activation_bytes <= BUFFER_OP_MAX_BYTES + + # -- reporting -------------------------------------------------------- + + @property + def label(self) -> str: + k = "x".join(map(str, self.kernel)) + s = "x".join(map(str, self.spatial)) + op = "convT" if self.transposed else "conv" + return f"{op} {self.cin}->{self.cout} k{k} @ {s}" + + @property + def qualified_label(self) -> str: + """:attr:`label` plus the two things that make it a *different problem*. + + ``label`` names the operator, the channels, the kernel and the extent, + and every published table is keyed on it -- so it stays exactly as it + is. It does not name the padding, and the padding is what separates the + three forms of the module docstring: ``conv 64->64 k3 @ 128x128x128`` + alone does not say whether it is the padded convolution ScaFFold runs or + an unpadded one, and MIOpen and this package both answer differently on + that. Use this wherever a reader could otherwise take a halo'd cell for + a production one. + """ + p = ",".join(map(str, self.padding)) + return f"{self.label} p{p} [{self.form}]" + + def measured_for(self, direction: Direction, config: str | None = None): + """The MIOpen measurements for one direction, most expensive first.""" + hits = [ + m + for m in self.measured + if m["direction"] == direction and (config is None or m["config"] == config) + ] + return sorted(hits, key=lambda m: -m["ms_per_call"]) + + +# --------------------------------------------------------------------------- +# The ScaFFold corpus +# --------------------------------------------------------------------------- + + +@functools.lru_cache(maxsize=1) +def scaffold_corpus() -> tuple[ConvProblem, ...]: + """Every distinct convolution in the three profiled ScaFFold configurations. + + Ordered by measured cost, so truncating the list keeps the problems that + matter. Regenerate with ``work/triton-conv/make_corpus.py``. + """ + raw = json.loads(_CORPUS_PATH.read_text()) + problems = [] + for entry in raw["problems"]: + w = entry["weight_shape"] + transposed = entry["op"] == "ConvTranspose3d" + cin, cout = (w[0], w[1]) if transposed else (w[1], w[0]) + n, _, *spatial = entry["in_shape"] + k = tuple(w[2:5]) + problems.append( + ConvProblem( + name="+".join(s.split(":", 1)[1] for s in entry["sites"][:1]), + cin=cin, + cout=cout, + spatial=tuple(spatial), + kernel=k, + stride=(entry["stride"],) * 3, + padding=(entry["padding"],) * 3, + n=n, + transposed=transposed, + bias=entry["bias"], + dtype="bf16", + sites=tuple(entry["sites"]), + measured=tuple(entry.get("measured", ())), + halo=tuple(entry.get("halo_dhw") or (0, 0, 0)), + shard_halo=tuple(entry.get("shard_halo_dhw") or (0, 0, 0)), + ) + ) + return tuple(problems) + + +@functools.lru_cache(maxsize=1) +def halo_corpus() -> tuple[ConvProblem, ...]: + """The corpus as **upstream DistConv** issues it: halo'd input, no padding. + + This -- not :func:`scaffold_corpus` -- is what the ``measured`` MIOpen + timings in the corpus are timings *of*, because the profile that produced + them ran through ``distconv_forward``. It is therefore the right form to + hold an *MIOpen baseline* in, and the wrong one to hold a Triton result in: + the shipped Triton rung is handed :func:`production_corpus`'s form instead. + :func:`scaffold_corpus` keeps the logical, unhaloed statement of each + problem because that is what the shape dump records and what the FLOP model + is naturally expressed in; the three differ by + :meth:`ConvProblem.halo_variant` and :meth:`ConvProblem.production_variant`. + """ + return tuple(p.halo_variant for p in scaffold_corpus()) + + +@functools.lru_cache(maxsize=1) +def production_corpus() -> tuple[ConvProblem, ...]: + """The corpus in the form **ScaFFold runs today** -- padded, mostly. + + :meth:`ConvProblem.production_variant` of every corpus problem: the local + shard with a halo on the genuinely split axis only, and the module's own + padding still in place on the others. At the unsharded configuration this + is identical to :func:`scaffold_corpus`; at the sharded ones it is a third + shape, in neither :func:`scaffold_corpus` nor :func:`halo_corpus`. + + Verified against an instrumented run rather than asserted -- see + ``test_infra.py::test_the_production_variant_matches_the_measured_census``, + which joins this against :func:`census_corpus`. + """ + return tuple(p.production_variant for p in scaffold_corpus()) + + +@functools.lru_cache(maxsize=1) +def census_corpus() -> tuple[ConvProblem, ...]: + """Every convolution an instrumented ScaFFold step actually issued. + + Recorded by ``work/triton-conv/bin/conv_census.py``, which wraps + ``FastConv3d``/``FastConvTranspose3d`` and the six kernel entry points and + runs three real training steps at each of the four configurations the + benchmark harness uses (A = scale 7 / 1 GPU, B = scale 8 / 1 GPU, C = scale + 8 / 2 GPUs, D = scale 8 / 4 GPUs). Every problem here is in + :attr:`ConvProblem.form` ``"adapter"`` by construction: it is the shape and + padding the kernel was handed, read off the call. + + Why this exists beside :func:`scaffold_corpus`, rather than being folded + into it: + + * it covers a configuration the profiled corpus does not (scale 8 on one + GPU), and a *network depth* the profiled corpus does not -- the shape + dumps behind :func:`scaffold_corpus` were taken at + ``unet_bottleneck_dim = 4`` at scale 8, giving a four-layer model topping + out at 1024 channels, while every step-level measurement in + ``work/triton-conv/minibatch/`` runs the shipped default of 3, i.e. a + five-layer model topping out at 2048; + * it carries no ``measured`` MIOpen data and no cost ordering, so it is not + a priority list and must not be used as one; + * and :func:`scaffold_corpus`'s ordering, indices and contents are the key + every stored capture in this project refers to, so they do not move. + + ``large`` is set from the activation size, so a caller that iterates this + without opting in does not try to allocate the 2 GiB scale-8 unsharded + activations. + """ + if not _CENSUS_PATH.exists(): # pragma: no cover - shipped with the package + return () + raw = json.loads(_CENSUS_PATH.read_text()) + out = [] + for entry in raw["problems"]: + w = entry["weight_shape"] + transposed = entry["op"] == "ConvTranspose3d" + cin, cout = (w[0], w[1]) if transposed else (w[1], w[0]) + n, _, *spatial = entry["in_shape"] + out.append( + ConvProblem( + name=entry.get("name", ""), + cin=cin, + cout=cout, + spatial=tuple(spatial), + kernel=tuple(w[2:5]), + stride=tuple(entry["stride"]), + padding=tuple(entry["padding"]), + n=n, + transposed=transposed, + bias=entry["bias"], + dtype=entry.get("dtype", "bf16"), + sites=tuple(entry["sites"]), + large=bool(entry.get("large")), + form="adapter", + ) + ) + return tuple(out) + + +def hot_corpus(top: int = 12) -> tuple[ConvProblem, ...]: + """The most expensive distinct problems -- the fast loop during development.""" + return scaffold_corpus()[:top] + + +# --------------------------------------------------------------------------- +# Synthetic edge cases +# --------------------------------------------------------------------------- + + +def edge_cases(include_large: bool = False) -> tuple[ConvProblem, ...]: + """Problems chosen to break assumptions rather than to be fast. + + Each one targets a specific way an implicit-GEMM kernel goes wrong: tile + remainders in every dimension, masking at volume faces, anisotropy, batching, + and the int32 offset overflow that MIOpen itself gets wrong. + """ + cases: list[ConvProblem] = [ + # Channel counts that are not multiples of any plausible BLOCK_K. + ConvProblem("cin_tiny", 3, 64, (16, 16, 16), sites=("synthetic",)), + ConvProblem("cin_odd", 5, 32, (8, 8, 8), sites=("synthetic",)), + ConvProblem("cin_prime", 17, 24, (8, 8, 8), sites=("synthetic",)), + ConvProblem("cout_tiny", 64, 6, (8, 8, 8), (1, 1, 1), padding=(0, 0, 0), + bias=True, sites=("synthetic",)), + ConvProblem("cout_odd", 32, 7, (8, 8, 8), sites=("synthetic",)), + # Spatial extents that do not divide any plausible tile. + ConvProblem("spatial_prime", 32, 32, (13, 13, 13), sites=("synthetic",)), + ConvProblem("spatial_one", 32, 32, (1, 8, 8), sites=("synthetic",)), + ConvProblem("spatial_thin", 32, 32, (2, 31, 3), sites=("synthetic",)), + ConvProblem("spatial_aniso", 64, 64, (5, 40, 96), sites=("synthetic",)), + # Smaller than the kernel in one axis: every tap is masked somewhere. + ConvProblem("smaller_than_kernel", 16, 16, (2, 2, 2), sites=("synthetic",)), + # Padding variants: unpadded shrinks the output, k=1 removes the gather. + ConvProblem("unpadded", 32, 32, (16, 16, 16), padding=(0, 0, 0), + sites=("synthetic",)), + ConvProblem("pointwise", 64, 6, (16, 16, 16), (1, 1, 1), padding=(0, 0, 0), + bias=True, sites=("synthetic",)), + ConvProblem("kernel_aniso", 32, 32, (8, 8, 8), (1, 3, 3), padding=(0, 1, 1), + sites=("synthetic",)), + # The padding a *sharded* ScaFFold convolution actually reaches the + # kernel with: a symmetric ``k = 3`` with the split axis halo'd (so + # ``p = 0`` there) and H and W still padded. Anisotropic padding under + # an isotropic kernel is a combination nothing else here produces -- + # ``kernel_aniso`` gets its zero from ``kd = 1``, where the boundary + # predicate on D is dead for a different reason -- and it is the form + # every k=3 site runs at ``dc_num_shards = (2,1,1)`` or ``(4,1,1)``. + ConvProblem("shard_padded", 32, 32, (8, 8, 8), padding=(0, 1, 1), + sites=("synthetic",)), + # Batch > 1: ScaFFold never does this, but the M decomposition must. + ConvProblem("batched", 32, 32, (8, 8, 8), n=3, sites=("synthetic",)), + # The transposed upsample, at a size that is quick to check. + ConvProblem("transposed", 64, 32, (8, 8, 8), (2, 2, 2), (2, 2, 2), (0, 0, 0), + transposed=True, bias=True, sites=("synthetic",)), + # fp32, for more_determinism and for exact-arithmetic tests. + ConvProblem("fp32", 32, 32, (8, 8, 8), dtype="fp32", sites=("synthetic",)), + ConvProblem("fp16", 32, 32, (8, 8, 8), dtype="fp16", sites=("synthetic",)), + ] + if include_large: + # The 2**31 *element* boundary, bracketed. ``1 x 128 x 258^3`` = + # 2.198e9 elements is the unsharded scale-8 activation that makes MIOpen + # assert; ``255^3`` is the largest volume of the same shape family that + # still fits an int32 index, at 2.122e9 elements. The pair differs only + # in spatial extent so that what it brackets is the boundary and not a + # change of channel width or kernel as well. + # + # The previous ``int32_below`` was ``64 -> 64 @ 512^3``: 8.59e9 + # elements, four times *above* the boundary it was meant to sit below, + # so the pair pinned nothing and the low case needed 16 GiB per + # activation. Both of these are 4.2-4.4 GiB in bf16 and both are past + # the buffer-op byte limit -- see :attr:`ConvProblem.buffer_ops_eligible` + # for why that is a different question from this one. + cases += [ + ConvProblem("int32_below", 128, 64, (255, 255, 255), large=True, + sites=("synthetic",)), + ConvProblem("int32_above", 128, 64, (258, 258, 258), large=True, + sites=("synthetic",)), + ] + return tuple(cases) + + +def all_problems(include_large: bool = False) -> Iterator[ConvProblem]: + yield from scaffold_corpus() + yield from edge_cases(include_large=include_large) + + +def problems_in_form(form: Form) -> tuple[ConvProblem, ...]: + """The corpus in one of the three forms, chosen by name. + + A driver that takes a ``--form`` flag wants exactly this, and wants it in + one place: the mapping from the word a user typed to the shape a kernel is + handed is the thing this whole distinction exists to keep honest. + """ + return { + "logical": scaffold_corpus, + "distconv": halo_corpus, + "adapter": production_corpus, + }[form]() + + +if __name__ == "__main__": # pragma: no cover - a human-readable dump + hdr = (f"{'ms/step':>9} {'logical':38s} {'adapter (production)':46s} " + f"{'AI':>7} {'i64':>4}") + print(hdr) + print("-" * len(hdr)) + for p in scaffold_corpus(): + ms = sum(m["ms_per_step"] for m in p.measured) + print( + f"{ms:9.3f} {p.label:38s} {p.production_variant.qualified_label:46s} " + f"{p.arithmetic_intensity():7.0f} {'yes' if p.needs_int64 else '':>4}" + ) + padded = sum(1 for p in production_corpus() if any(p.padding)) + print(f"\n{len(scaffold_corpus())} ScaFFold problems, " + f"{len(edge_cases(include_large=True))} synthetic edge cases, " + f"{len(census_corpus())} measured by census") + print(f"{padded}/{len(production_corpus())} of the production forms are " + f"padded; {sum(1 for p in halo_corpus() if any(p.padding))} of the " + f"DistConv forms are") diff --git a/triton_conv3d/tests/test_bwd_data.py b/triton_conv3d/tests/test_bwd_data.py new file mode 100644 index 0000000..f94c36c --- /dev/null +++ b/triton_conv3d/tests/test_bwd_data.py @@ -0,0 +1,680 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""Tests for backward-data, which is the forward kernel on a transformed weight. + +Because no new kernel is introduced, these tests are not re-testing the gather +-- ``test_gather_gemm.py`` does that. What they test is the *transform*, and +the transform is the part with a uniquely nasty failure mode: it flips the tap +axes and swaps the two channel axes, and getting either half wrong produces a +gradient that is the right shape, the right magnitude, smooth, and wrong. A +tolerance test cannot see that. Two of the tests here exist purely to prove the +bitwise standard is not vacuous: + +* :func:`test_bitwise_standard_rejects_a_shifted_gather` -- shift the upstream + gradient by one voxel and the comparison must fail; +* :func:`test_an_unflipped_weight_is_detected` -- omit the tap flip and the + comparison must fail. This is the specific bug the whole module could have, + and without this test a passing suite would not rule it out. + +The other thing these tests cover that the forward's do not is that +backward-data's effective convolution is **always padded** for ``k > 1``, even +when the forward was not: DistConv issues an unpadded ``130^3`` convolution and +its backward-data has ``p' = 2``. So the halo'd corpus is parametrized here in +its own right rather than only in its logical, padded form. +""" + +from __future__ import annotations + +import math + +import pytest +import torch +import triton + +from triton_conv3d import reference +from triton_conv3d.bwd_data import ( + bwd_data_config, + bwd_data_padding, + conv3d_backward_data, + is_supported_bwd_data, +) +from triton_conv3d.gather_gemm import candidate_configs, default_config, to_rsck +from triton_conv3d.shapes import ConvProblem, edge_cases, scaffold_corpus + +requires_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") + +#: The synthetic corpus, minus the transposed upsample (a later milestone). +EDGE = [p for p in edge_cases() if not p.transposed] + + +def _corpus_channel_pairs() -> list[ConvProblem]: + """Every distinct ``(Cin, Cout, kernel)`` in the corpus, at a testable volume. + + The forward suite selects corpus problems by *volume* -- small enough for an + fp64 reference -- and that works there. It does not work here, and the way + it fails is worth recording because the first version of this file shipped + with it: backward-data reduces over ``Cout * taps``, so the surviving + problems are exactly the deep, wide ones, and a sum of ``27648`` random + signs runs to about 500 while bf16 holds integers only to 256. Every single + corpus case then hit ``is_exactly_representable`` and skipped, and the file + reported "89 passed" with zero real-shape coverage. + + Restating each channel pair at ``6x7x8`` instead keeps what the corpus is + *for* -- the channel widths, and with them ``EVEN_K``/``EVEN_N``, the tile + selection and the 512-byte row strides -- while making the reference cheap. + + **All three paddings** are generated, because ScaFFold issues all three and + they are three different problems (``shapes.py``'s module docstring): + + * ``p = (1,1,1)`` -- what the adapter hands the kernel at one GPU, and the + module's own statement everywhere; + * ``p = (0,1,1)`` -- what it hands the kernel at two or four GPUs, where D + is halo'd and H and W are not. Anisotropic, which no other case in this + file is: the backward's ``p'`` is then ``(2,1,1)``, so one axis reads a + two-voxel boundary shell and the other two read one; + * ``p = (0,0,0)`` -- what upstream DistConv hands MIOpen, and the form every + published baseline was measured in. + + None of the three subsumes another, and the middle one is the one that used + to be missing. + """ + seen: set[tuple] = set() + out: list[ConvProblem] = [] + for p in scaffold_corpus(): + if p.transposed or (p.cin, p.cout, p.kernel) in seen: + continue + seen.add((p.cin, p.cout, p.kernel)) + shard = tuple(0 if i == 0 else v for i, v in enumerate(p.padding)) + forms = [(p.padding, ""), ((0, 0, 0), "-halo")] + if shard != p.padding and shard != (0, 0, 0): + forms.insert(1, (shard, "-shard")) + for pad, tag in forms: + out.append(ConvProblem( + f"{p.cin}to{p.cout}{tag}", p.cin, p.cout, (6, 7, 8), p.kernel, + padding=pad, sites=("corpus-pair",), + )) + return out + + +#: See :func:`_corpus_channel_pairs`. +CORPUS_PAIRS = _corpus_channel_pairs() + +#: Real ScaFFold shapes, at their real volumes, small enough to reference in +#: fp64. Used only for the fp32 test below: their bf16 references are never +#: exactly representable, which is the whole point of the note above. +CORPUS_SMALL = [ + p for p in scaffold_corpus() + if not p.transposed + and math.prod(p.halo_variant.spatial) * max(p.cin, p.cout) <= 1 << 22 +] +CORPUS_SMALL += [p.halo_variant for p in CORPUS_SMALL] + + +def _ids(problems): + return [p.name or p.label for p in problems] + + +def _run(problem: ConvProblem, ops: dict, **kwargs) -> torch.Tensor: + return conv3d_backward_data( + ops["grad_output"], ops["weight"], problem.input_shape, + problem.stride, problem.padding, **kwargs, + ) + + +# --------------------------------------------------------------------------- +# The algebra, before any GPU is involved +# --------------------------------------------------------------------------- + + +def test_the_padding_identity_is_the_one_the_derivation_claims(): + """``p' = dil*(k-1) - p``, and the output extent then lands on the input's. + + Stated as a test rather than left in a docstring because every other file in + this module depends on it and it is one sign error away from producing a + gradient of the wrong *shape* -- which at least fails loudly -- or, at + ``k=3, p=1``, the right shape and the wrong answer, which does not. + """ + assert bwd_data_padding(1, 1, 3) == (1, 1, 1) + assert bwd_data_padding(0, 1, 3) == (2, 2, 2) # the halo'd form + assert bwd_data_padding(0, 1, 1) == (0, 0, 0) # k=1: no gather at all + assert bwd_data_padding((0, 1, 1), 1, (1, 3, 3)) == (0, 1, 1) + assert bwd_data_padding(1, 2, 3) == (3, 3, 3) # dilation widens the reach + + # And the extent identity: OD + 2p' - dil*(k-1) == ID, for every combination. + for k in (1, 2, 3, 5): + for dil in (1, 2, 3): + for p in range(0, dil * (k - 1) + 1): + for in_d in (1, 4, 17): + out_d = in_d + 2 * p - dil * (k - 1) + if out_d < 1: + continue + pp = bwd_data_padding(p, dil, k)[0] + assert out_d + 2 * pp - dil * (k - 1) == in_d, (k, dil, p, in_d) + + +def test_flipping_every_tap_axis_is_complementing_the_fused_index(): + """The identity the kernel's ``taps - 1 - dij`` rests on. + + The transform this replaced -- ``permute(2,3,4,0,1).flip((0,1,2))`` -- + materialized a whole second copy of every weight, once per optimizer step, + to express a *reindexing*. The kernel now flips by walking the fused tap + index backwards, which is only the same thing because the fused index is a + mixed-radix number and complementing every digit complements the number. + That is exactly the sort of claim that is obvious, load-bearing and one + off-by-one away from a silently wrong gradient, so it is checked over + anisotropic kernels rather than argued. + + ``k=(1,3,1)``-shaped cases are in the list on purpose: an axis of extent 1 + contributes ``0`` to both sides, which is where a formula that got the radix + order wrong would still look right. + """ + for kd, kh, kw in [(3, 3, 3), (1, 1, 1), (2, 3, 4), (1, 3, 1), (5, 1, 2)]: + taps = kd * kh * kw + for d in range(kd): + for i in range(kh): + for j in range(kw): + flipped = (((kd - 1 - d) * kh + (kh - 1 - i)) * kw + + (kw - 1 - j)) + fused = (d * kh + i) * kw + j + assert flipped == taps - 1 - fused, (kd, kh, kw, d, i, j) + + +# --------------------------------------------------------------------------- +# Configuration legality -- the failure mode is silent, so it is checked apart +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("problem", EDGE + CORPUS_PAIRS, + ids=_ids(EDGE + CORPUS_PAIRS)) +def test_selected_config_is_legal_for_every_shape(problem: ConvProblem): + """The config picked for backward-data must still reach the matrix core. + + Not implied by the forward's version of this test: the effective GEMM has + ``Cin`` and ``Cout`` swapped, so a shape whose forward tile is legal can + have a backward tile that is not -- ``Cout=6`` becomes ``BLOCK_K`` rather + than ``BLOCK_N``, and ``BLOCK_K`` is the one with the hard MFMA constraint. + """ + dtype = reference.torch_dtype(problem) + cfg = bwd_data_config(problem.output_shape, problem.cin, problem.kernel, dtype, + padding=problem.padding, dilation=(1, 1, 1)) + assert cfg.validate(dtype) is None, f"{problem.label}: {cfg} -> {cfg.validate(dtype)}" + + +@pytest.mark.parametrize("problem", CORPUS_PAIRS, ids=_ids(CORPUS_PAIRS)) +def test_every_backward_candidate_config_is_legal(problem: ConvProblem): + """The sweep that produced ``_TUNED_BWD`` must not contain an FMA kernel. + + Same reasoning as the forward's: an illegal config runs and returns the + right answer slowly, so the sweep's reported winner could be one. The + argument order is what differs -- the candidate list is generated for the + *effective* widths. + """ + dtype = reference.torch_dtype(problem) + m = problem.n * math.prod(problem.spatial) + cfgs = candidate_configs(m, problem.cout, problem.cin, dtype) + assert cfgs + for cfg in cfgs: + assert cfg.validate(dtype) is None, f"{cfg}: {cfg.validate(dtype)}" + + +# --------------------------------------------------------------------------- +# Support predicate +# --------------------------------------------------------------------------- + + +@requires_gpu +def test_is_supported_declines_what_the_algebra_cannot_express(): + """Two of these refusals are backward-only and both are load-bearing. + + ``stride > 1`` makes the backward a scatter into a sub-lattice, and + ``padding > dil*(k-1)`` makes ``p'`` negative -- a crop. Neither is a + forward gather, and neither raises anything by itself: the kernel would run + and write a plausible, wrong gradient. + """ + gy = torch.empty((1, 8, 4, 4, 4), device="cuda", dtype=torch.bfloat16) + w = torch.empty((8, 8, 3, 3, 3), device="cuda", dtype=torch.bfloat16) + shape = (1, 8, 4, 4, 4) + assert is_supported_bwd_data(gy, w, shape, padding=1) + + assert not is_supported_bwd_data(gy, w, shape, stride=2, padding=1) + assert not is_supported_bwd_data(gy, w, shape, padding=3) # p > dil*(k-1) + assert not is_supported_bwd_data(gy, w, shape, padding=1, groups=2) + assert not is_supported_bwd_data(gy, w.float(), shape, padding=1) + # Cin of the weight must match the gradient being asked for ... + assert not is_supported_bwd_data(gy, w, (1, 4, 4, 4, 4), padding=1) + # ... Cout of the weight must match grad_output ... + assert not is_supported_bwd_data( + gy, torch.empty((4, 8, 3, 3, 3), device="cuda", dtype=torch.bfloat16), + shape, padding=1, + ) + # ... and grad_output's spatial extent must be the one this problem produces. + assert not is_supported_bwd_data(gy, w, (1, 8, 6, 4, 4), padding=1) + assert is_supported_bwd_data(gy, w, (1, 8, 6, 6, 6), padding=0) + + +@requires_gpu +def test_is_supported_declines_an_empty_batch(): + """The predicate bounded every spatial extent below and not ``N``. + + Degenerate rather than dangerous -- the grid comes out empty -- but this + gate's own "every output voxel must exist" reasoning excludes a batch with no + samples in it, and a ``True`` here is the gate asserting something it never + looked at. The cost of declining is one call's worth of MIOpen on a problem + that has nothing to compute. + """ + gy = torch.empty((0, 8, 4, 4, 4), device="cuda", dtype=torch.bfloat16) + w = torch.empty((8, 8, 3, 3, 3), device="cuda", dtype=torch.bfloat16) + assert not is_supported_bwd_data(gy, w, (0, 8, 4, 4, 4), padding=1) + with pytest.raises(NotImplementedError): + conv3d_backward_data(gy, w, (0, 8, 4, 4, 4), padding=1) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two GPUs") +def test_is_supported_declines_operands_on_different_devices(): + """Both operands on *a* GPU is not both on the *same* GPU. + + Triton launches on the current device and dereferences the foreign pointer + regardless. ScaFFold runs four ranks to a node, and with peer access enabled + that reads another rank's weights rather than faulting -- a wrong gradient + with no symptom at all. Skipped, not absent, on a single-GPU box. + """ + gy = torch.empty((1, 8, 4, 4, 4), device="cuda:0", dtype=torch.bfloat16) + w = torch.empty((8, 8, 3, 3, 3), device="cuda:0", dtype=torch.bfloat16) + assert is_supported_bwd_data(gy, w, (1, 8, 4, 4, 4), padding=1) + assert not is_supported_bwd_data(gy, w.to("cuda:1"), (1, 8, 4, 4, 4), padding=1) + + +# --------------------------------------------------------------------------- +# Correctness: the bitwise standard +# --------------------------------------------------------------------------- + + +@requires_gpu +@pytest.mark.parametrize("problem", EDGE, ids=_ids(EDGE)) +def test_exact_operands_match_bitwise(problem: ConvProblem): + """Bitwise against ``torch.autograd.grad`` in fp64, on the nasty shapes. + + The same synthetic corpus as the forward, and it earns its place twice over + here: ``Cin=3`` and ``Cout=6`` land on the GEMM's *N* rather than its K, and + ``smaller_than_kernel`` is where the flipped gather masks every tap + somewhere. + """ + ops = reference.make_inputs(problem, seed=3, exact=True) + expected = reference.reference(problem, ops, "bwd-data") + dtype = reference.torch_dtype(problem) + if not reference.is_exactly_representable(expected, dtype): + pytest.skip("realized magnitudes exceed the mantissa in this dtype") + actual = _run(problem, ops) + report = reference.compare(actual, expected.to(dtype)) + assert report.bitwise, f"{problem.label}: {report}" + + +@requires_gpu +@pytest.mark.parametrize("problem", CORPUS_PAIRS, ids=_ids(CORPUS_PAIRS)) +def test_corpus_channel_pairs_match_bitwise(problem: ConvProblem): + """Every channel pair ScaFFold runs, in all three paddings, bitwise in bf16. + + The three arms are the three forms of the same site, and each has something + the others do not. ``p=0`` (DistConv's) has a backward padding of 2, so it + reads a boundary shell two voxels thick, which no forward convolution in + this project ever does. ``p=(0,1,1)`` (the adapter's, sharded) is + anisotropic: ``p'`` is ``(2,1,1)`` and the two shell thicknesses coexist in + one kernel. ``p=1`` (the adapter's, unsharded) is the ordinary one. + Running only one of them would leave a shape ScaFFold actually issues + untested; the middle one was the one missing until 2026-08-04. + + The five deepest pairs skip here and are picked up by the fp32 test below; + see :func:`test_the_bitwise_corpus_is_not_entirely_skipped` for why that is + checked rather than assumed. + """ + ops = reference.make_inputs(problem, seed=5, exact=True) + expected = reference.reference(problem, ops, "bwd-data") + dtype = reference.torch_dtype(problem) + if not reference.is_exactly_representable(expected, dtype): + pytest.skip("realized magnitudes exceed the mantissa in this dtype") + actual = _run(problem, ops) + assert reference.compare(actual, expected.to(dtype)).bitwise + + +@requires_gpu +def test_the_bitwise_corpus_is_not_entirely_skipped(): + """A regression guard on this file, not on the kernel. + + ``is_exactly_representable`` declining is the correct behaviour, but if it + declines for *every* parametrized case the suite reports a wall of passes + and tests nothing. That is exactly what the first version of this file did. + So pin a floor: most of the corpus's channel pairs must actually reach the + bitwise comparison in bf16. + """ + exact = 0 + for problem in CORPUS_PAIRS: + ops = reference.make_inputs(problem, seed=5, exact=True) + expected = reference.reference(problem, ops, "bwd-data") + exact += reference.is_exactly_representable( + expected, reference.torch_dtype(problem) + ) + assert exact >= len(CORPUS_PAIRS) // 2, ( + f"only {exact}/{len(CORPUS_PAIRS)} corpus pairs are bf16-exact; the " + "bitwise corpus test is close to vacuous" + ) + + +@requires_gpu +@pytest.mark.parametrize( + "problem", CORPUS_PAIRS + CORPUS_SMALL, ids=_ids(CORPUS_PAIRS + CORPUS_SMALL) +) +def test_deep_corpus_shapes_match_bitwise_in_fp32(problem: ConvProblem): + """The shapes bf16 cannot express exactly, at their real widths and volumes. + + ``Cout >= 512`` makes backward-data's reduction 13824 or 27648 terms long, + and no choice of ``{-1,0,1}`` operands keeps that inside bf16's 8-bit + mantissa -- it is a property of the arithmetic, not of the test. fp32 has + 24 bits, which covers it with room to spare, and the addressing under test + is dtype-independent: what changes is the MFMA intrinsic and therefore the + legal ``BLOCK_K``, so this is also the only bitwise coverage the fp32 tile + selection gets at real widths. + + This test does not skip. If the fp32 reference is ever not exact either, + that is a fact worth failing on rather than stepping around. + """ + ops = reference.make_inputs(problem, seed=7, exact=True, dtype=torch.float32) + expected = reference.reference(problem, ops, "bwd-data") + assert reference.is_exactly_representable(expected, torch.float32) + actual = _run(problem, ops) + assert actual.dtype is torch.float32 + assert reference.compare(actual, expected.to(torch.float32)).bitwise + + +@requires_gpu +def test_bitwise_standard_rejects_a_shifted_gather(): + """Prove the comparison discriminates: a one-voxel shift must fail it. + + Same argument as the forward's version -- ``{-1,0,1}`` operands could in + principle make everything agree, and this project has shipped two vacuous + exact tests before. + """ + problem = ConvProblem("shift", 16, 16, (6, 6, 6)) + ops = reference.make_inputs(problem, seed=11, exact=True) + actual = _run(problem, ops) + correct = reference.reference(problem, ops, "bwd-data").to(torch.bfloat16) + assert reference.compare(actual, correct).bitwise + + shifted = torch.roll(ops["grad_output"], shifts=1, dims=-1) + wrong = reference.reference( + problem, {**ops, "grad_output": shifted}, "bwd-data" + ).to(torch.bfloat16) + assert not reference.compare(actual, wrong).bitwise, ( + "a one-voxel shift of the upstream gradient produced a bitwise-identical " + "result; the comparison is not discriminating" + ) + + +@requires_gpu +def test_an_unflipped_weight_is_detected(): + """The one bug this module can uniquely have, pinned. + + The gather reads tap ``t`` of this direction from tap ``taps-1-t`` of the + weight. Omitting that -- the single most plausible mistake in the whole + file, and now a constexpr in the kernel rather than a ``torch.flip``, which + makes it easier to get wrong and no easier to see -- still produces a + correctly shaped, correctly scaled, smooth gradient, and would pass every + tolerance test written. So construct exactly that wrong answer and require + a mismatch. + + ``padding=1`` with ``k=3`` is deliberate: it is the case where flipped and + unflipped agree on the *shape*, so nothing else catches it. + """ + problem = ConvProblem("flip", 16, 16, (6, 7, 8)) + ops = reference.make_inputs(problem, seed=13, exact=True) + actual = _run(problem, ops) + correct = reference.reference(problem, ops, "bwd-data").to(torch.bfloat16) + assert reference.compare(actual, correct).bitwise + + unflipped = reference.reference( + problem, {**ops, "weight": ops["weight"].flip(2, 3, 4)}, "bwd-data" + ).to(torch.bfloat16) + assert unflipped.shape == actual.shape + assert not reference.compare(actual, unflipped).bitwise, ( + "a weight with the taps un-flipped gave a bitwise-identical gradient; " + "the kernel's W_FLIP is untested by this suite" + ) + + +@requires_gpu +@pytest.mark.parametrize("problem", EDGE, ids=_ids(EDGE)) +def test_every_config_gives_the_same_answer(problem: ConvProblem): + """The tuning surface, not one point on it. + + ``_TUNED_BWD`` is free to pick any of these, and the backward's boundary + shell is two voxels thick rather than one -- so a mask that is right when + ``BLOCK_M`` divides the row length and wrong when it does not has more room + to hide here than in the forward. That argument applies with most force to + the cases this stopped short of when it swept only ``EDGE[:8]``: ``batched`` + (the only ``n > 1`` shape), ``kernel_aniso``, ``smaller_than_kernel``, + ``unpadded`` -- whose backward is padded where its forward is not -- and both + non-bf16 dtypes, which move the MFMA reduction depth and so the set of legal + ``BLOCK_K`` values. + """ + ops = reference.make_inputs(problem, seed=2, exact=True) + expected = reference.reference(problem, ops, "bwd-data") + dtype = reference.torch_dtype(problem) + if not reference.is_exactly_representable(expected, dtype): + pytest.skip("realized magnitudes exceed the mantissa in this dtype") + expected = expected.to(dtype) + m = problem.n * math.prod(problem.spatial) + # Effective widths: the reduction is over Cout and the GEMM's N is Cin. + cfgs = candidate_configs(m, problem.cout, problem.cin, dtype, group_ms=(6, 8)) + cfgs = list(dict.fromkeys( + cfgs + [default_config(m, problem.cout, problem.cin, dtype)] + )) + ran = 0 + for cfg in cfgs: + try: + actual = _run(problem, ops, config=cfg) + except triton.runtime.errors.OutOfResources: + # Operands that do not fit in 64 KiB of LDS. A loud failure, so no + # static guard is wanted -- the sweep skips it and so does this. + continue + ran += 1 + assert reference.compare(actual, expected).bitwise, f"{problem.label} {cfg}" + assert ran, "no candidate configuration was runnable" + + +# --------------------------------------------------------------------------- +# Correctness: the tolerance standards +# --------------------------------------------------------------------------- + + +@requires_gpu +@pytest.mark.parametrize("problem", EDGE + CORPUS_PAIRS, + ids=_ids(EDGE + CORPUS_PAIRS)) +def test_no_worse_than_miopen(problem: ConvProblem): + """The honest bar at realistic magnitudes, against MIOpen on the same data. + + Worth stating separately from the forward's because backward-data's + reduction is over ``Cout * taps`` rather than ``Cin * taps``, so on the + asymmetric decoder convolutions the two directions accumulate over + different lengths and inherit different error. + """ + ops = reference.make_inputs(problem, seed=17) + expected = reference.reference(problem, ops, "bwd-data") + incumbent_err = reference.compare( + reference.incumbent(problem, ops, "bwd-data"), expected + ) + actual = _run(problem, ops) + reference.assert_close(actual, expected, problem, "bwd-data", + incumbent_error=incumbent_err) + + +@requires_gpu +def test_fp32_accumulates_in_fp32(): + """``more_determinism`` runs the model in fp32, and the backward too. + + A tf32-style split dot would pass any bf16-sized tolerance, so the bound is + fp32-sized and held against fp64. + """ + problem = ConvProblem("fp32", 48, 32, (7, 9, 5), dtype="fp32") + ops = reference.make_inputs(problem, seed=23) + expected = reference.reference(problem, ops, "bwd-data") + actual = _run(problem, ops) + assert actual.dtype is torch.float32 + report = reference.compare(actual, expected) + peak = expected.abs().max().item() + assert report.max_abs < 1e-4 * peak, f"looks like a reduced-precision dot: {report}" + + +# --------------------------------------------------------------------------- +# Entry-point behaviour +# --------------------------------------------------------------------------- + + +@requires_gpu +def test_out_buffer_is_written_in_place_and_is_validated(): + """``out=`` is forwarded straight to the forward entry point, unexamined. + + Handing a preallocated gradient buffer to the backward is exactly what a + DistConv integration does, and both halves of the hole were reachable from + here: an undersized buffer wrote 10752 elements into a 256-element + allocation with no error, and an NCDHW buffer returned a gradient with + ``max_abs = 99.0``. ``reduce_gemm`` validated the same parameter and this + direction did not. + + The check lives in the forward, and that is exact rather than approximate: + the effective forward's output shape *is* ``input_shape``. This test is what + says so. + """ + problem = ConvProblem("out", 16, 24, (4, 5, 6)) + ops = reference.make_inputs(problem, seed=67, exact=True) + expected = _run(problem, ops) + assert tuple(expected.shape) == problem.input_shape + + buf = torch.empty_like(expected) + got = _run(problem, ops, out=buf) + assert got.data_ptr() == buf.data_ptr(), "out= was allocated over, not written" + assert torch.equal(got, expected) + + bf16 = torch.bfloat16 + with pytest.raises(ValueError): + _run(problem, ops, out=torch.empty((1, 16, 2, 2, 2), device="cuda", dtype=bf16)) + with pytest.raises(ValueError): # right shape, NCDHW + _run(problem, ops, + out=torch.empty(problem.input_shape, device="cuda", dtype=bf16)) + with pytest.raises(ValueError): + _run(problem, ops, out=torch.empty_like(expected, dtype=torch.float32)) + + +@requires_gpu +def test_hoisted_weight_buffer_is_validated(): + """``weight_rsck`` supplies every weight value; ``weight`` supplies a shape. + + So a buffer belonging to another parameter -- a stale cache entry is the + realistic way to get one -- is a smooth, correctly shaped, entirely wrong + gradient. This direction has its own trap on top of the forward's: the + buffer it takes is the **forward's** ``(kd, kh, kw, Cin, Cout)``, so the + transposed spelling, which is what a reader who knows backward-data reduces + over ``Cout`` would reach for, has to be rejected rather than quietly + transposing the answer. + """ + problem = ConvProblem("wr", 16, 24, (4, 5, 6)) + ops = reference.make_inputs(problem, seed=67, exact=True) + good = to_rsck(ops["weight"]) + assert torch.equal(_run(problem, ops, weight_rsck=good), _run(problem, ops)) + + other = torch.randn((24, 16, 1, 1, 1), device="cuda", dtype=torch.bfloat16) + with pytest.raises(ValueError): + _run(problem, ops, weight_rsck=to_rsck(other)) + with pytest.raises(ValueError): + _run(problem, ops, weight_rsck=good.float()) + # (kd, kh, kw, Cout, Cin) -- the layout the deleted ``to_bwd_rsck`` produced. + with pytest.raises(ValueError): + _run(problem, ops, weight_rsck=good.transpose(3, 4).contiguous()) + + +@requires_gpu +def test_every_weight_layout_gives_the_same_gradient(): + """The parameter is read where it lies, so its strides pick the B load. + + Three layouts reach three different ``W_ORDER``/copy decisions and must not + reach three different answers. Bitwise, not close: they are the same + multiply-accumulate in the same order, and anything less would mean the + layout had leaked into the arithmetic. + + The RSCK-strided case is the one worth having a test for. It is a weight + with PyTorch's shape and this kernel's storage order, which is what an + integration that wanted the forward's B tile contiguous would allocate; here + it is the layout in which ``weight_rsck`` and ``weight`` are the *same + tensor*, so it is also the case that would hide a mix-up between them. + """ + problem = ConvProblem("layouts", 32, 48, (4, 5, 6)) + ops = reference.make_inputs(problem, seed=53, exact=True) + w = ops["weight"] + cout, cin, kd, kh, kw = w.shape + layouts = { + "channels_last": w.contiguous(memory_format=torch.channels_last_3d), + "contiguous": w.contiguous(), + "rsck_strided": (w.permute(2, 3, 4, 1, 0).contiguous() + .permute(4, 3, 0, 1, 2)), + } + ref = _run(problem, ops) + for name, wl in layouts.items(): + assert torch.equal(wl, w), name # same values, different strides + got = _run(problem, {**ops, "weight": wl}) + assert torch.equal(ref, got), name + # And the hoisted buffer, which is a fourth spelling of the same values. + assert torch.equal(ref, _run(problem, ops, weight_rsck=to_rsck(w))) + + +@requires_gpu +def test_output_is_channels_last_and_matches_torch_grad_shape(): + problem = ConvProblem("shape", 16, 40, (3, 11, 5)) + ops = reference.make_inputs(problem, seed=61) + gx = _run(problem, ops) + ref = torch.nn.grad.conv3d_input( + problem.input_shape, ops["weight"], ops["grad_output"], + stride=problem.stride, padding=problem.padding, + ) + assert gx.shape == ref.shape + assert gx.is_contiguous(memory_format=torch.channels_last_3d) + + +@requires_gpu +def test_ncdhw_grad_output_is_converted_rather_than_misread(): + """The addressing assumes ``stride_c == 1`` on the upstream gradient. + + An NCDHW ``grad_output`` read with NDHWC strides gives a full-rate kernel + and a completely wrong gradient. ScaFFold's own backward can hand us either + layout depending on what produced the gradient, so this is not hypothetical. + """ + problem = ConvProblem("layout", 24, 16, (5, 6, 7)) + ops = reference.make_inputs(problem, seed=41, exact=True) + ndhwc = _run(problem, ops) + nc = {k: (v.contiguous() if torch.is_tensor(v) else v) for k, v in ops.items()} + assert nc["grad_output"].stride(1) != 1 + ncdhw = _run(problem, nc) + assert torch.equal(ndhwc, ncdhw) + + +@requires_gpu +def test_repeated_calls_are_bitwise_reproducible(): + """MIOpen's backward-data is not; this is the direction where that is fixed. + + ScaFFold's default configuration is nonreproducible today, and backward-data + is one of the contributors. Stating the property as a test is what stops a + later split-K variant from quietly giving it up. + """ + problem = ConvProblem("determinism", 64, 64, (8, 12, 10)) + ops = reference.make_inputs(problem, seed=71) + first = _run(problem, ops) + for _ in range(4): + assert torch.equal(first, _run(problem, ops)) + + +@requires_gpu +def test_unsupported_calls_raise_rather_than_return_garbage(): + gy = torch.randn((1, 8, 4, 4, 4), device="cuda", dtype=torch.bfloat16) + w = torch.randn((8, 8, 3, 3, 3), device="cuda", dtype=torch.bfloat16) + with pytest.raises(NotImplementedError): + conv3d_backward_data(gy, w, (1, 8, 4, 4, 4), stride=2, padding=1) + with pytest.raises(NotImplementedError): + conv3d_backward_data(gy, w, (1, 8, 4, 4, 4), padding=3) + with pytest.raises(NotImplementedError): + conv3d_backward_data(gy, w, (1, 8, 5, 4, 4), padding=1) diff --git a/triton_conv3d/tests/test_bwd_weight.py b/triton_conv3d/tests/test_bwd_weight.py new file mode 100644 index 0000000..65a11a6 --- /dev/null +++ b/triton_conv3d/tests/test_bwd_weight.py @@ -0,0 +1,1370 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""Tests for backward-weight, the one direction with a kernel of its own. + +Three things are tested here that the other two directions do not have: + +* **the split-K decomposition**, which is not an optimization but the only way + this GEMM fills the device -- ``M = Cout`` is one or two tile rows. Its + correctness lives in :func:`split_count` agreeing with what the kernel and the + reduction pass assume about each other, so the tests pin the arithmetic (every + voxel in exactly one split) as well as the answer; +* **determinism**, stated as the claim PLAN.md 3.4 makes and tested the way that + claim is worded: bitwise identical run to run *and process to process*, for the + same input, dtype, shape, device and tuning config. There are three tests -- + in-process repetition, three separate interpreters, and a negative control on + the atomic path -- because a determinism test that cannot fail is the most + comfortable kind to write and the least useful; +* **the reuse that was checked and rejected.** + :func:`test_the_forward_kernel_can_express_backward_weight` runs the algebraic + identity that would have made this file unnecessary, and passes; the reason it + is not used is a trip count, which the same test asserts. + +The bitwise standard has the same teeth as elsewhere in this suite and the same +two guards against being vacuous: a shifted operand must fail the comparison, and +the specific bug this module can uniquely have -- writing a tap to the wrong slot +of the ``[Cout][tap][Cin]`` output -- is constructed and required to fail. +""" + +from __future__ import annotations + +import math +import pathlib +import subprocess +import sys +import textwrap + +import pytest +import torch +import triton + +from triton_conv3d import reference +from triton_conv3d.gather_gemm import conv3d_forward, default_config +from triton_conv3d.reduce_gemm import ( + _CU_COUNT, + _MAX_EPILOGUE_FRACTION, + _SPLIT_TARGET_WAVES, + _WORKSPACE_BYTES, + BwdWeightConfig, + _row_aligned, + bwd_weight_config, + candidate_bwd_weight_configs, + conv3d_backward_weight, + default_bwd_weight_config, + grad_weight_empty, + is_supported_bwd_weight, + split_count, + workspace_elements, +) +from triton_conv3d.shapes import ConvProblem, edge_cases, scaffold_corpus + +requires_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") + +#: The synthetic corpus, minus the transposed upsample (a later milestone). +EDGE = [p for p in edge_cases() if not p.transposed] + + +def _corpus_channel_pairs() -> list[ConvProblem]: + """Every distinct ``(Cin, Cout, kernel)`` in the corpus, at a testable volume. + + Same construction as ``test_bwd_data.py``, and it is needed for the same + reason with a different arithmetic: backward-weight reduces over the whole + *output volume*, so a sum of ``{-1,0,1}`` products at a real ScaFFold shape + runs to about ``sqrt(2.1e6) = 1450`` while bf16 holds integers only to 256. + Restating each pair at ``6x7x8`` keeps what the corpus is for -- the channel + widths, and with them ``EVEN_M``/``EVEN_N``, the tile selection and the + 512-byte row strides -- and brings the reduction down to 336 terms, which + bf16 does hold. + + **All three paddings** are generated, because ScaFFold issues all three + (``shapes.py``'s module docstring): ``p=(1,1,1)`` is what the adapter hands + the kernel unsharded and the module's own statement everywhere, + ``p=(0,1,1)`` is what it hands the kernel at two or four shards, and + ``p=(0,0,0)`` is what upstream DistConv hands MIOpen. They do not differ in + the *predicate* the kernel compiles the way backward-data's do -- this + direction reads X at ``o + t - p``, and ``PADDED`` is on for any non-zero + padding -- but they differ in the output extent and therefore in the + reduction length, the split count and whether ``BLOCK_K`` divides a row, and + the anisotropic one is the only case where the ``d`` half of the boundary + predicate is dead while the ``h``/``w`` halves are live. + """ + seen: set[tuple] = set() + out: list[ConvProblem] = [] + for p in scaffold_corpus(): + if p.transposed or (p.cin, p.cout, p.kernel) in seen: + continue + seen.add((p.cin, p.cout, p.kernel)) + shard = tuple(0 if i == 0 else v for i, v in enumerate(p.padding)) + forms = [(p.padding, ""), ((0, 0, 0), "-halo")] + if shard != p.padding and shard != (0, 0, 0): + forms.insert(1, (shard, "-shard")) + for pad, tag in forms: + out.append(ConvProblem( + f"{p.cin}to{p.cout}{tag}", p.cin, p.cout, (6, 7, 8), p.kernel, + padding=pad, sites=("corpus-pair",), + )) + return out + + +#: See :func:`_corpus_channel_pairs`. +CORPUS_PAIRS = _corpus_channel_pairs() + +#: Real ScaFFold shapes, at their real volumes, small enough to reference in +#: fp64. Used only for the fp32 test below. +CORPUS_SMALL = [ + p for p in scaffold_corpus() + if not p.transposed + and math.prod(p.halo_variant.spatial) * max(p.cin, p.cout) <= 1 << 22 +] +CORPUS_SMALL += [p.halo_variant for p in CORPUS_SMALL] + + +def _ids(problems): + return [p.name or p.label for p in problems] + + +def _run(problem: ConvProblem, ops: dict, **kwargs) -> torch.Tensor: + return conv3d_backward_weight( + ops["input"], problem.weight_shape, ops["grad_output"], + problem.stride, problem.padding, **kwargs, + ) + + +# --------------------------------------------------------------------------- +# The split decomposition, before any GPU is involved +# --------------------------------------------------------------------------- + + +def _cfg(**kw) -> BwdWeightConfig: + base = dict(BLOCK_M=64, BLOCK_N=64, BLOCK_K=64, num_warps=4) + base.update(kw) + return BwdWeightConfig(**base) + + +def test_the_splits_partition_the_reduction_exactly_once(): + """Every output voxel lands in exactly one split, for every shape. + + This is the invariant the whole direction rests on: the kernel clamps its + last tile with ``k_end = min(k_begin + chunk, K)`` and the reduction adds + every split unconditionally, so a chunk arithmetic that overlapped would + double-count silently and one that fell short would drop voxels at the end + of the volume -- both of which produce a plausible gradient. + """ + for out_w in (8, 16, 128, 256, 7, 13): + for k_total in (out_w, out_w * 3, out_w * 4096, out_w * 4097): + for bk in (16, 32, 64, 128): + for sk in (0, 1, 3, 64, 1000): + cfg = _cfg(BLOCK_K=bk, SPLIT_K=sk) + splits, chunk = split_count(cfg, 64, 64, 27, k_total, out_w) + assert splits >= 1 and chunk >= 1 + assert (splits - 1) * chunk < k_total <= splits * chunk + # The kernel's cheap scalar unravel is only valid when a + # K-tile cannot straddle a row, which needs the chunk to be + # row-aligned as well as the tile. + if out_w % bk == 0: + assert chunk % out_w == 0 + else: + assert chunk % bk == 0 + + +def test_the_grid_lands_on_whole_waves(): + """The snap that the measured split-count curve turned out to be about. + + Every program in this kernel does the same amount of work, so a grid of 4.5 + waves runs five and idles through half the last one. At + ``64 -> 64 @ 130x258x258`` that sawtooth is an 18% effect -- 596 splits is + 16% more parallelism than 512 and 9% slower -- and it is easy to misread as + a statement about cache footprint. So the property is pinned here rather + than left to the constant that happens to produce it. + """ + for cout, cin, taps, k_total, out_w in ( + (64, 64, 27, 8_388_608, 256), (64, 128, 27, 2_097_152, 128), + (128, 128, 27, 1_048_576, 128), (256, 256, 27, 131_072, 64), + (512, 512, 27, 16_384, 32), (6, 64, 1, 2_097_152, 256), + ): + cfg = bwd_weight_config(cout, cin, (3, 3, 3) if taps > 1 else (1, 1, 1), + k_total, torch.bfloat16) + splits, _ = split_count(cfg, cout, cin, taps, k_total, out_w) + tiles = (-(-cout // cfg.BLOCK_M) * -(-cin // cfg.BLOCK_NC) + * -(-taps // cfg.TAP_BLOCK)) + progs = tiles * splits + if progs <= _CU_COUNT: + continue # one wave or less: nothing to snap + waste = (-(-progs // _CU_COUNT) * _CU_COUNT - progs) / progs + assert waste < 0.10, (cout, cin, k_total, splits, tiles, progs, waste) + assert _SPLIT_TARGET_WAVES >= 1 + + +def test_the_split_count_is_a_pure_function_of_the_shape(): + """Determinism starts here: no clock, no free memory, no autotuner. + + Stated as a test because it is an easy thing to give away later -- a split + count that adapted to the device's current occupancy would be a perfectly + reasonable optimization and would silently end the reproducibility claim. + """ + cfg = _cfg() + first = split_count(cfg, 128, 256, 27, 2_097_152, 128) + for _ in range(4): + assert split_count(cfg, 128, 256, 27, 2_097_152, 128) == first + # And it responds to the shape, so the property above is not vacuous. + assert split_count(cfg, 128, 256, 27, 4096, 16)[0] < first[0] + + +#: The largest fp32 partial workspace any corpus problem asks for, in MiB, on +#: the path :func:`conv3d_backward_weight` actually takes. Quoted to anyone +#: sizing a hoisted ``workspace=`` once and out of the step, so it is pinned by +#: a test rather than recorded in a document: the previously published figure +#: (111 MiB) was the maximum over the ten shapes of the determinism table +#: measured on the *heuristic* config, and understated the real corpus maximum +#: by 1.46x. An integration that had sized from it would have taken a +#: ``ValueError`` mid-run. +def _every_form(problems): + """Each non-transposed problem in all three of the forms ScaFFold issues. + + Bounds like the workspace ceiling have to hold on the shape the *kernel* + is handed, and there are three of those: the module's own padded statement, + the adapter's (padded on every unsplit axis) and upstream DistConv's (halo'd + and unpadded). Iterating only the last of them -- which this file did until + 2026-08-04 -- bounds the one form production never issues. Deduplicated on + the qualified label, since the three coincide wherever nothing is split. + """ + seen: set[str] = set() + for p in problems: + if p.transposed: + continue # a later milestone; ``weight_shape`` is swapped too + for q in (p, p.production_variant, p.halo_variant): + if q.qualified_label in seen: + continue + seen.add(q.qualified_label) + yield q + + +_WORST_WORKSPACE_MIB = 216 + + +def test_the_partial_workspace_is_bounded_across_the_whole_corpus(): + """``splits * Cout * taps * Cin * 4`` bytes, on every problem ScaFFold runs. + + The number to watch is the *product*: one split at ``1024 -> 1024`` is + 113 MiB, and a split count picked for a shallow site would ask for + gigabytes of it. The two bounds pull against each other -- the sites that + want many splits are the ones with a small ``Cout * taps * Cin`` -- but that + is an observation about this corpus, not a theorem, so it is checked. + + Two things this test used to get wrong, both of which made it unable to + fail: + + * it built its config with :func:`default_bwd_weight_config` while the + entry point uses :func:`bwd_weight_config`, which prefers the tuned table + and therefore a different ``BLOCK_M``/``BLOCK_NC``/``TAP_BLOCK``, a + different tile count and a different split count. The two disagree by up + to 1.5x in practice, so the bound it certified was not the shipped one; + * ``mib <= _WORKSPACE_BYTES`` is *trivially* true -- ``split_count``'s own + ``ceiling`` is ``_WORKSPACE_BYTES // per_split``, so no config it returns + can violate it. The assertion that carries the weight is the pinned + maximum, which is a number people size allocations from. + """ + worst, worst_label = 0.0, "" + for hp in _every_form(list(scaffold_corpus()) + EDGE): + k_total = hp.n * math.prod(hp.out_spatial) + # Both, and the worst of the two: a caller who passes no ``config=`` + # gets the resolver's answer, and one who builds a config from + # :func:`default_bwd_weight_config` gets the heuristic's, which at an + # untuned pair is what production launches. Since 2026-08-05 the two + # agree on padded problems that they used to disagree on -- the + # ``TAP_BLOCK`` decline is gone -- so the worst of the pair is a smaller + # set than it was, and it is still the number to size an allocation + # from. + for cfg in ( + bwd_weight_config(hp.cout, hp.cin, hp.kernel, k_total, + torch.bfloat16, padded=any(hp.padding)), + default_bwd_weight_config(hp.cout, hp.cin, hp.kernel, k_total, + torch.bfloat16, padded=any(hp.padding)), + ): + splits, _ = split_count(cfg, hp.cout, hp.cin, hp.tap_count, k_total, + hp.out_spatial[2]) + mib = workspace_elements(splits, hp.cout, hp.cin, + hp.kernel) * 4 / 2**20 + assert mib <= _WORKSPACE_BYTES / 2**20, f"{hp.label}: {mib:.0f} MiB" + if mib > worst: + worst, worst_label = mib, f"{hp.label} {cfg}" + # A ceiling nothing approaches would not be a useful test either. + assert worst > 1.0 + assert round(worst) == _WORST_WORKSPACE_MIB, ( + f"the corpus workspace maximum moved to {worst:.1f} MiB at " + f"{worst_label}, from the {_WORST_WORKSPACE_MIB} MiB pinned here. " + "That number is what an integration sizes a hoisted workspace= from, " + "so update it deliberately -- do not widen this assertion" + ) + + +def test_the_wave_snap_outranks_the_epilogue_bound_and_only_below_one_wave(): + """The bound the docstring states is the bound the code applies. + + :func:`split_count` clamps its target against three ceilings and *then* + snaps to a whole number of waves, and the snap can push the result back + above the epilogue bound. That reads like an oversight and is not: the + alternative was implemented and raced, and it loses badly. Re-applying the + epilogue bound after the snap takes ``128 -> 256 @ 34^3`` from 16 splits to + 7 -- a 98-program grid on 228 CUs -- and the site from 0.2536 to 0.4563 ms + (**1.80x**); ``256 -> 512 @ 10x34x34`` goes 0.2731 -> 0.6553 ms. Half an + idle device costs more than a doubled epilogue, and the shapes where the + snap overrides the bound are *exactly* the shapes with a sub-wave grid, + because that is the condition under which ``round`` rounds to zero. + + So what is pinned here is the ordering itself, in both directions: + + * the snap may exceed the epilogue bound only by *rounding the grid to the + nearest whole wave* -- at most half a wave of extra programs, or one + whole wave where the bounded grid does not fill even that. Anything + beyond that would mean the bound had stopped constraining anything; + * the **workspace** ceiling is different in kind (a failed allocation at + step 400 is not a slow kernel) and is re-applied after the snap, so it is + never exceeded. + + Both halves have failed at some point in this function's history, in + opposite directions. + """ + def bounds(cfg, cout, cin, taps, k_total): + tiles = (-(-cout // cfg.BLOCK_M) * -(-cin // cfg.BLOCK_NC) + * -(-taps // cfg.TAP_BLOCK)) + loop_elems = tiles * k_total * (cfg.BLOCK_M + cfg.BLOCK_N) + epi = max(1, loop_elems + // (_MAX_EPILOGUE_FRACTION * max(1, cout * taps * cin * 4))) + return tiles, epi + + checked = overridden = 0 + for hp in _every_form(list(scaffold_corpus()) + EDGE): + k_total = hp.n * math.prod(hp.out_spatial) + for cfg in ( + bwd_weight_config(hp.cout, hp.cin, hp.kernel, k_total, + torch.bfloat16, padded=any(hp.padding)), + default_bwd_weight_config(hp.cout, hp.cin, hp.kernel, k_total, + torch.bfloat16, padded=any(hp.padding)), + ): + splits, _ = split_count(cfg, hp.cout, hp.cin, hp.tap_count, k_total, + hp.out_spatial[2]) + tiles, epi = bounds(cfg, hp.cout, hp.cin, hp.tap_count, k_total) + checked += 1 + if splits <= epi: + continue + overridden += 1 + # Over the epilogue bound is allowed, but only by the rounding the + # snap does: to the *nearest* whole wave, so at most half a wave of + # extra programs -- or one whole wave where the bounded grid does + # not fill even one. + assert tiles * splits <= max(_CU_COUNT, + tiles * epi + _CU_COUNT // 2), ( + f"{hp.label} {cfg}: {splits} splits against an epilogue bound " + f"of {epi} is {tiles * splits} programs, more than a wave past " + f"the bounded grid's {tiles * epi}" + ) + # And the workspace ceiling still holds, which is the one bound the + # snap is *not* allowed to escape. + mib = workspace_elements(splits, hp.cout, hp.cin, + hp.kernel) * 4 / 2**20 + assert mib <= _WORKSPACE_BYTES / 2**20, f"{hp.label}: {mib:.0f} MiB" + assert checked > 40 + + assert overridden, ( + "no shape in the corpus reaches the sub-wave regime any more, so this " + "test no longer covers the ordering it exists to pin" + ) + + +@pytest.mark.parametrize("problem", EDGE + CORPUS_PAIRS, + ids=_ids(EDGE + CORPUS_PAIRS)) +def test_selected_config_is_legal_for_every_shape(problem: ConvProblem): + """The config picked for backward-weight must reach the matrix core. + + Not implied by the other two directions' versions of this test: here + ``BLOCK_M`` is bounded by ``Cout`` rather than by a volume, so a shape whose + forward tile is legal can select a tile here that is not -- and an illegal + MFMA configuration on gfx942 runs, returns the right answer, and emits no + matrix instruction at all. + """ + dtype = reference.torch_dtype(problem) + hp = problem.halo_variant + k_total = hp.n * math.prod(hp.out_spatial) + cfg = bwd_weight_config(hp.cout, hp.cin, hp.kernel, k_total, dtype, + padded=any(hp.padding)) + assert cfg.validate(dtype) is None, f"{hp.label}: {cfg} -> {cfg.validate(dtype)}" + assert cfg.lds_bytes(dtype) <= 64 * 1024, f"{hp.label}: {cfg}" + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32], + ids=["bf16", "fp16", "fp32"]) +def test_default_config_fits_in_lds_in_every_dtype(dtype): + """fp32 operands are twice the bytes, and ``more_determinism`` runs in fp32. + + M2 found exactly this hole in the *forward*'s shipped heuristic, where + ``128x128x128`` is 64 KiB in bf16 and 128 KiB in fp32 and the shipped + configuration raised ``OutOfResources``. This direction's tiles are wider + still -- ``TAP_BLOCK`` multiplies ``BLOCK_N`` -- so the same trap is closer, + not further away. + """ + for cout in (6, 64, 128, 256, 512, 1024): + for cin in (3, 64, 128, 256, 512, 1024): + for k in ((1, 1, 1), (3, 3, 3)): + cfg = default_bwd_weight_config(cout, cin, k, 1 << 20, dtype) + assert cfg.validate(dtype) is None, cfg + assert cfg.lds_bytes(dtype) <= 64 * 1024, (cout, cin, k, cfg) + + +@pytest.mark.parametrize("problem", CORPUS_PAIRS, ids=_ids(CORPUS_PAIRS)) +def test_every_backward_weight_candidate_config_is_legal(problem: ConvProblem): + """The sweep that produced the tuned table must not contain an FMA kernel. + + Same reasoning as the other two directions': an illegal config runs and + returns the right answer slowly, so a best-of sweep that merely ranked it + last would still be reporting a meaningless winner. + """ + dtype = reference.torch_dtype(problem) + k_total = problem.n * math.prod(problem.out_spatial) + cfgs = candidate_bwd_weight_configs(problem.cout, problem.cin, + problem.kernel, k_total, dtype) + assert cfgs + for cfg in cfgs: + assert cfg.validate(dtype) is None, f"{cfg}: {cfg.validate(dtype)}" + assert cfg.lds_bytes(dtype) <= 64 * 1024, cfg + assert cfg.BLOCK_N == cfg.BLOCK_NC * cfg.TAP_BLOCK + + +def test_config_validate_refuses_the_two_knobs_this_direction_adds(): + bf16 = torch.bfloat16 + assert BwdWeightConfig().validate(bf16) is None + assert BwdWeightConfig(SPLIT_K=-1).validate(bf16) + assert BwdWeightConfig(TAP_BLOCK=0).validate(bf16) + # BLOCK_N is the *full* tile width, so it has to divide into whole taps -- + # otherwise BLOCK_NC is a truncated integer and the column decode silently + # addresses the wrong channels. + assert BwdWeightConfig(BLOCK_N=64, TAP_BLOCK=3).validate(bf16) + assert BwdWeightConfig(BLOCK_N=192, TAP_BLOCK=3).validate(bf16) is None + # And the inherited gfx942 rules still apply. + assert BwdWeightConfig(BLOCK_K=8).validate(bf16) + + +# --------------------------------------------------------------------------- +# Support predicate +# --------------------------------------------------------------------------- + + +@requires_gpu +def test_is_supported_declines_what_the_kernel_cannot_express(): + """Note what is *not* refused: ``stride > 1``. + + Backward-data has to refuse a stride because its substitution turns into a + scatter into a sub-lattice. This direction does not: its reduction axis is + the output voxel and the input coordinate ``o*s + t*dil - p`` is a function + of it, so a stride is three extra multiplies. The asymmetry is real and is + pinned here so that a later reader does not "fix" it by symmetry. + """ + x = torch.empty((1, 8, 6, 6, 6), device="cuda", dtype=torch.bfloat16) + gy = torch.empty((1, 8, 6, 6, 6), device="cuda", dtype=torch.bfloat16) + ws = (8, 8, 3, 3, 3) + assert is_supported_bwd_weight(x, ws, gy, padding=1) + + strided = torch.empty((1, 8, 3, 3, 3), device="cuda", dtype=torch.bfloat16) + assert is_supported_bwd_weight(x, ws, strided, stride=2, padding=1) + + assert not is_supported_bwd_weight(x, ws, gy, padding=1, groups=2) + assert not is_supported_bwd_weight(x, ws, gy.float(), padding=1) + + # Both operands on *the same* device, not merely both on a device. Triton + # launches on the current device and dereferences the other pointer anyway, + # and ScaFFold runs four GPUs per node: with peer access enabled a foreign + # pointer does not fault, it reads another rank's activations and returns a + # plausible wrong gradient. ``gather_gemm.is_supported`` refuses the same + # thing; the two gates sit behind one rung ladder and a hole in either is a + # hole in the ladder. + assert not is_supported_bwd_weight(x, ws, gy.cpu(), padding=1) + assert not is_supported_bwd_weight(x.cpu(), ws, gy, padding=1) + if torch.cuda.device_count() >= 2: + # The clause above ``is_cuda`` cannot reach: two *CUDA* devices. Only + # runnable on a multi-GPU node -- this suite is normally run with one + # device pinned -- so the CPU cases above stay unconditional rather than + # letting the whole check disappear behind the guard. + assert not is_supported_bwd_weight(x, ws, gy.to("cuda:1"), padding=1) + # ...and the same-device pair is still accepted, so none of this is a + # predicate that has simply started refusing everything. + assert is_supported_bwd_weight(x, ws, gy, padding=1) + assert not is_supported_bwd_weight(x, (8, 4, 3, 3, 3), gy, padding=1) + assert not is_supported_bwd_weight(x, (4, 8, 3, 3, 3), gy, padding=1) + # grad_output's extent has to be the one this problem produces, or the + # reduction would run over a volume the input does not have. + assert not is_supported_bwd_weight(x, ws, gy, padding=0) + assert not is_supported_bwd_weight( + x, ws, torch.empty((1, 8, 4, 6, 6), device="cuda", dtype=torch.bfloat16), + padding=1, + ) + + +@requires_gpu +def test_unsupported_calls_raise_rather_than_return_garbage(): + x = torch.randn((1, 8, 6, 6, 6), device="cuda", dtype=torch.bfloat16) + gy = torch.randn((1, 8, 6, 6, 6), device="cuda", dtype=torch.bfloat16) + with pytest.raises(NotImplementedError): + conv3d_backward_weight(x, (8, 8, 3, 3, 3), gy, padding=1, groups=2) + with pytest.raises(NotImplementedError): + conv3d_backward_weight(x, (8, 8, 3, 3, 3), gy, padding=0) + # An out= in the wrong layout is refused rather than filled transposed. + with pytest.raises(ValueError): + conv3d_backward_weight( + x, (8, 8, 3, 3, 3), gy, padding=1, + out=torch.empty((8, 8, 3, 3, 3), device="cuda", dtype=torch.bfloat16), + ) + + +# --------------------------------------------------------------------------- +# Correctness: the bitwise standard +# --------------------------------------------------------------------------- + + +@requires_gpu +@pytest.mark.parametrize("problem", EDGE, ids=_ids(EDGE)) +def test_exact_operands_match_bitwise(problem: ConvProblem): + """Bitwise against ``torch.autograd.grad`` in fp64, on the nasty shapes. + + The synthetic corpus earns its place here differently than in the other two + directions: ``Cout=6`` and ``Cout=7`` land on the GEMM's *M*, which is the + axis this kernel has least of, and ``spatial_thin`` (2x31x3) gives an output + volume of 12 -- a reduction shorter than one ``BLOCK_K``. + """ + ops = reference.make_inputs(problem, seed=3, exact=True) + expected = reference.reference(problem, ops, "bwd-weight") + dtype = reference.torch_dtype(problem) + if not reference.is_exactly_representable(expected, dtype): + pytest.skip("realized magnitudes exceed the mantissa in this dtype") + actual = _run(problem, ops) + report = reference.compare(actual, expected.to(dtype)) + assert report.bitwise, f"{problem.label}: {report}" + + +@requires_gpu +@pytest.mark.parametrize("problem", CORPUS_PAIRS, ids=_ids(CORPUS_PAIRS)) +def test_corpus_channel_pairs_match_bitwise(problem: ConvProblem): + """Every channel pair ScaFFold runs, in both paddings, bitwise in bf16.""" + ops = reference.make_inputs(problem, seed=5, exact=True) + expected = reference.reference(problem, ops, "bwd-weight") + dtype = reference.torch_dtype(problem) + if not reference.is_exactly_representable(expected, dtype): + pytest.skip("realized magnitudes exceed the mantissa in this dtype") + actual = _run(problem, ops) + assert reference.compare(actual, expected.to(dtype)).bitwise + + +@requires_gpu +def test_the_bitwise_corpus_is_not_entirely_skipped(): + """A regression guard on this file, not on the kernel. + + ``is_exactly_representable`` declining is correct behaviour, but if it + declines for every parametrized case the suite reports a wall of passes and + tests nothing. That is what the first version of ``test_bwd_data.py`` did. + """ + exact = sum( + reference.is_exactly_representable( + reference.reference( + p, reference.make_inputs(p, seed=5, exact=True), "bwd-weight" + ), + reference.torch_dtype(p), + ) + for p in CORPUS_PAIRS + ) + assert exact >= len(CORPUS_PAIRS) // 2, ( + f"only {exact}/{len(CORPUS_PAIRS)} corpus pairs are bf16-exact; the " + "bitwise corpus test is close to vacuous" + ) + + +@requires_gpu +@pytest.mark.parametrize( + "problem", CORPUS_PAIRS + CORPUS_SMALL, ids=_ids(CORPUS_PAIRS + CORPUS_SMALL) +) +def test_deep_corpus_shapes_match_bitwise_in_fp32(problem: ConvProblem): + """The shapes bf16 cannot express exactly, at their real widths and volumes. + + A reduction over a real ScaFFold output volume runs to about ``sqrt(K)`` in + ``{-1,0,1}`` arithmetic -- 1450 at the 128^3 sites -- which bf16's 8-bit + mantissa provably cannot hold, as a property of the arithmetic and not of + the test. fp32 has 24 bits, which covers it, and the addressing under test + is dtype-independent: what changes is the MFMA intrinsic and therefore the + legal ``BLOCK_K``, so this is also the only bitwise coverage the fp32 tile + selection gets at real widths. + + This test does not skip. If the fp32 reference is ever not exact either, + that is a fact worth failing on rather than stepping around. + """ + ops = reference.make_inputs(problem, seed=7, exact=True, dtype=torch.float32) + expected = reference.reference(problem, ops, "bwd-weight") + assert reference.is_exactly_representable(expected, torch.float32) + actual = _run(problem, ops) + assert actual.dtype is torch.float32 + assert reference.compare(actual, expected.to(torch.float32)).bitwise + + +#: Shapes that compile the ``PADDED and ROW_ALIGNED`` pair of constexprs. Every +#: other padded shape in this file has ``out_w < BLOCK_K``, so ``_row_aligned`` +#: is False at all of them and this combination had never been compiled by the +#: suite at all. See the test below for why that is worth fixing. +_PADDED_ROW_ALIGNED = [ + # ``IN_D = IN_H = 1`` under ``padding=1``: ``src_d`` is -1 at every voxel and + # the three taps land at -1, 0 and +1, so both the low and the high ``d``/ + # ``h`` boundaries fire on every K-tile rather than only at the volume's + # edge. 16 reduction terms, so bf16 holds the result exactly. + ConvProblem("pad-rowaligned-thin", 16, 16, (1, 1, 16)), + # The logical (non-halo'd) form of a real corpus site: ``256->128 k3 @ + # 64x128x128, padding=1`` has ``out_w = 128`` against ``BLOCK_K = 64``. Same + # shape of predicate at a width the corpus actually produces; fp32 because + # a 2048-term reduction is past bf16's mantissa. + ConvProblem("pad-rowaligned-corpus", 32, 32, (4, 4, 128), dtype="fp32"), +] + + +@requires_gpu +@pytest.mark.parametrize("problem", _PADDED_ROW_ALIGNED, ids=_ids(_PADDED_ROW_ALIGNED)) +def test_the_padded_row_aligned_corner_is_compiled_and_correct(problem): + """The one ``constexpr`` pair nothing else in this suite reaches. + + ``PADDED`` and ``ROW_ALIGNED`` are independent, and they interact. In the + ``ROW_ALIGNED`` branch ``row``, ``od``, ``oh`` and ``idn`` collapse to + **rank-0 scalars** -- the whole point of that branch is that the unravel + becomes four SALU divisions -- so the padded branch's boundary predicate + ``src_d[:, None] + (kd*DD)[None, :]`` is a different expression there than + in the general branch: broadcast from a scalar rather than from a + ``BLOCK_K`` vector, and collapsed to one row of the mask instead of + ``BLOCK_K`` of them. It is the right predicate, because within a + row-aligned K-tile ``od`` and ``oh`` really are constant -- but "it is + correct" and "it is tested" are different claims, and a bug planted in the + ``d`` or ``h`` half of it passed the entire suite. + + **This is a production branch, not a hypothetical one.** It used to be + documented as reachable only by "a caller who bypasses DistConv", on the + premise that every ScaFFold convolution is issued halo'd and unpadded. That + premise is false: the shipped adapter halos only the split axis, so + ``256->128 k3 @ 64x128x128`` arrives padded with ``out_w = 128`` against + ``BLOCK_K = 64`` -- exactly this branch -- every step. + """ + hp = problem + k_total = hp.n * math.prod(hp.out_spatial) + dtype = reference.torch_dtype(problem) + cfg = bwd_weight_config(hp.cout, hp.cin, hp.kernel, k_total, dtype, + padded=any(hp.padding)) + # The two constexprs, asserted rather than hoped for: this test's whole + # value is that it compiles a branch, so it has to fail loudly if a config + # change ever stops it reaching that branch. + assert any(hp.padding), "PADDED would be False" + assert _row_aligned(cfg.BLOCK_K, hp.out_spatial[2]), ( + f"ROW_ALIGNED is False: BLOCK_K={cfg.BLOCK_K} does not divide " + f"out_w={hp.out_spatial[2]}" + ) + + ops = reference.make_inputs(problem, seed=97, exact=True, dtype=dtype) + expected = reference.reference(problem, ops, "bwd-weight") + assert reference.is_exactly_representable(expected, dtype) + assert reference.compare(_run(problem, ops), expected.to(dtype)).bitwise + + +#: The channel pairs whose tuned backward-weight row widens ``TAP_BLOCK``. +#: Until 2026-08-05 these were exactly the rows *declined* on a padded problem, +#: i.e. at every production ScaFFold site; the decline is gone and this set is +#: now the rows that must survive the padding. Resolved from the table rather +#: than listed, so a retune moves this set instead of stranding it. +def _tap_widened_pairs() -> list[tuple[ConvProblem, BwdWeightConfig]]: + from triton_conv3d.reduce_gemm import ( + _TUNED_BWD_W, _fit_bwd_weight_to_lds, tune_key, + ) + + out, seen = [], set() + for p in scaffold_corpus(): + if p.transposed or (p.cin, p.cout, p.kernel) in seen: + continue + seen.add((p.cin, p.cout, p.kernel)) + row = _TUNED_BWD_W.get( + tune_key(torch.bfloat16, p.cin, p.cout, tuple(p.kernel))) + if row is not None and row.TAP_BLOCK > 1: + out.append((p, _fit_bwd_weight_to_lds(row, torch.bfloat16))) + return out + + +_TAP_WIDENED = _tap_widened_pairs() + + +def test_a_tuned_tap_block_row_survives_the_padding(): + """The replacement for ``..._is_declined_when_padded``, and why it flipped. + + Until 2026-08-05 ``bwd_weight_config`` refused a tuned row with + ``TAP_BLOCK > 1`` whenever the convolution was padded, and + ``default_bwd_weight_config`` refused to widen ``TAP_BLOCK`` there at all. + Both clauses were written believing they could not fire -- "no real ScaFFold + convolution is padded, DistConv halos them all" -- and that was false: + ScaFFold's own adapter halos only the split axis, so every ``k > 1`` site + arrives padded and **eight sites over six channel pairs** took the decline + at every configuration. + + The old test asserted the decline and asked whoever relaxed it to replace + the assertion with a measurement. That is what happened. Raced on the + padded production form of all 18 affected cells, the tuned row against the + config the decline produced, one interleaved block per cell with 95% + intervals (``work/triton-conv/review/OPTIMIZATION_ROUND3.md`` §1): the + tuned row wins **18 of 18**, geometric mean **1.946x**, range + 1.137x-5.336x, worst cell 7.9505 ms declined against 1.4910 ms with the + row. The heuristic's half was raced separately on the six pairs that reach + it and widening wins **6 of 6**, 1.263x-2.084x. + + So this test now pins the opposite property, and it is the one that matters + for production: the tuned row must be what a *padded* problem resolves, + because a padded problem is the only kind ScaFFold issues. + """ + assert _TAP_WIDENED, ( + "no tuned backward-weight row widens TAP_BLOCK any more; this test and " + "the behaviour it pins are both about a table that has changed" + ) + for p, row in _TAP_WIDENED: + k_total = p.n * math.prod(p.out_spatial) + padded = bwd_weight_config(p.cout, p.cin, p.kernel, k_total, + torch.bfloat16, padded=True) + unpadded = bwd_weight_config(p.cout, p.cin, p.kernel, k_total, + torch.bfloat16, padded=False) + assert unpadded == row, ( + f"{p.cin}->{p.cout}: the tuned row is not selected even unpadded" + ) + assert padded == row, ( + f"{p.cin}->{p.cout}: a padded problem resolved {padded} instead of " + f"the tuned row {row}. Production issues nothing but padded " + "convolutions, so this is the whole of what the table buys -- read " + "this test's docstring before accepting it" + ) + assert padded.TAP_BLOCK > 1 + + +def test_the_heuristic_widens_tap_block_under_padding_too(): + """The other half of the same predicate, pinned separately. + + :func:`default_bwd_weight_config` used to pin ``TAP_BLOCK`` to 1 on a padded + problem. It no longer does, and the two are now the same config: padding + changes the boundary predicate inside the kernel and nothing about the tile + the host picks. Kept apart from the test above because this one governs + every channel pair the tuned table does *not* list, which is where a new + ScaFFold site lands. + """ + for p, _row in _TAP_WIDENED: + k_total = p.n * math.prod(p.out_spatial) + wide = default_bwd_weight_config(p.cout, p.cin, p.kernel, k_total, + torch.bfloat16, padded=False) + padded = default_bwd_weight_config(p.cout, p.cin, p.kernel, k_total, + torch.bfloat16, padded=True) + assert padded == wide, ( + f"{p.cin}->{p.cout}: the heuristic still answers differently under " + f"padding ({padded} vs {wide})" + ) + assert wide.TAP_BLOCK > 1, ( + f"{p.cin}->{p.cout}: the heuristic did not widen TAP_BLOCK at all; " + "this test is about a rule that has changed" + ) + + +@requires_gpu +@pytest.mark.parametrize( + "problem,cfg", + [(ConvProblem(f"{p.cin}to{p.cout}-padded", p.cin, p.cout, (6, 7, 8), + p.kernel, padding=p.padding, sites=("tap-widened",)), c) + for p, c in _TAP_WIDENED], + ids=[f"{p.cin}to{p.cout}" for p, _ in _TAP_WIDENED], +) +def test_a_padded_tap_block_row_is_still_bitwise_correct(problem, cfg): + """The gradient a widened row produces on a padded problem, bitwise. + + Written while the row was still *declined* on a padded problem, to establish + that what the decline protected was a performance argument and not a + correctness one. Since 2026-08-05 the decline is gone and this is no longer + a hypothetical: it is the gradient every ``k = 3`` ScaFFold site computes, + so a failure here is a wrong weight gradient in production rather than a + reason not to relax a clause. + """ + assert cfg.TAP_BLOCK > 1 and any(problem.padding) + ops = reference.make_inputs(problem, seed=1234, exact=True) + expected = reference.reference(problem, ops, "bwd-weight") + dtype = reference.torch_dtype(problem) + assert reference.is_exactly_representable(expected, dtype) + actual = _run(problem, ops, config=cfg) + assert reference.compare(actual, expected.to(dtype)).bitwise, ( + f"{problem.label} with {cfg} (TAP_BLOCK>1 on a padded convolution) is " + "not bitwise correct" + ) + + +#: The triple ``PADDED and ROW_ALIGNED and TAP_BLOCK > 1``, and -- in the last +#: entry -- the *quintuple* production actually launches. ``_PADDED_ROW_ +#: ALIGNED`` above reaches the first two but not the third, so until 2026-08-05 +#: the combination was compiled nowhere; it is now compiled at every ``k = 3`` +#: site of every configuration. ``out_w`` is chosen equal to ``BLOCK_K`` so a +#: K-tile is exactly one output row. +#: +#: ``block_nc`` is carried per case because the stem needs it. ``3 -> 64`` +#: resolves ``64x64x64/tb16``, i.e. ``BLOCK_NC = 4`` against ``Cin = 3``, so it +#: adds two raggednesses to the triple -- a partial channel group *and* a +#: partial tap group (27 taps in blocks of 16) -- inside the ``ROW_ALIGNED`` +#: branch where ``src_d``/``src_h`` collapse to scalars. Nothing else in the +#: suite compiles that: the three cases above hold ``Cin = BLOCK_NC = 32``, and +#: the ragged-``Cin`` tests elsewhere are not row-aligned. It is also the +#: largest single win in the round (5.3x), which is a poor thing to have +#: untested. +_PADDED_ROW_ALIGNED_TAPS = [ + (ConvProblem("triple-tb8", 32, 32, (2, 3, 16)), 8, 16, 32), + (ConvProblem("triple-tb2", 32, 32, (2, 2, 64), dtype="fp32"), 2, 64, 32), + (ConvProblem("triple-tb16", 32, 64, (2, 2, 32)), 16, 32, 32), + # The sharded production padding, which is anisotropic: the ``d`` half of + # the boundary predicate is dead and the ``h``/``w`` halves are live, inside + # the branch where ``src_d`` is a rank-0 scalar. + (ConvProblem("triple-shardpad", 32, 32, (4, 4, 32), padding=(0, 1, 1)), + 4, 32, 32), + # The stem, in both of its production paddings. ``BLOCK_M`` is 32 here + # rather than the shipped 64 only because this test fixes it; every other + # constexpr is the one the resolver returns. + (ConvProblem("quintuple-stem", 3, 64, (2, 2, 64)), 16, 64, 4), + (ConvProblem("quintuple-stem-shardpad", 3, 64, (4, 4, 64), + padding=(0, 1, 1)), 16, 64, 4), +] + + +@requires_gpu +@pytest.mark.parametrize( + "problem,tap_block,block_k,block_nc", _PADDED_ROW_ALIGNED_TAPS, + ids=[p.name for p, _, _, _ in _PADDED_ROW_ALIGNED_TAPS], +) +def test_the_padded_row_aligned_tap_widened_corner_is_correct( + problem, tap_block, block_k, block_nc +): + """Three independent ``constexpr`` at once, which nothing else compiles. + + ``PADDED`` selects a two-dimensional boundary predicate; ``ROW_ALIGNED`` + collapses ``od``/``oh``/``idn`` to rank-0 scalars; ``TAP_BLOCK > 1`` makes + the tap vary down the *columns*. Together the predicate is a scalar + broadcast against a per-column tap shift, and since 2026-08-05 it is the + shape production launches at every ``k = 3`` site with a widened row -- see + :func:`test_a_tuned_tap_block_row_survives_the_padding`. It was written + while the combination was still unreachable, which is why it forces the + constexpr triple by hand rather than going through the resolver. + + The last two cases add the stem's two raggednesses on top, which is the + combination the shipped ``3 -> 64`` row launches and which nothing else + reaches; the assertions below say which case is which so a failure names the + axis rather than the tile. + """ + cfg = BwdWeightConfig(BLOCK_M=32, BLOCK_N=block_nc * tap_block, + BLOCK_K=block_k, TAP_BLOCK=tap_block, num_warps=4, + matrix_instr_nonkdim=16, kpack=1) + assert cfg.BLOCK_NC == block_nc + assert any(problem.padding), "PADDED would be False" + assert _row_aligned(cfg.BLOCK_K, problem.out_spatial[2]), ( + f"ROW_ALIGNED is False: BLOCK_K={cfg.BLOCK_K} does not divide " + f"out_w={problem.out_spatial[2]}" + ) + # 27 taps never divide by a power of two, so *every* case here has a ragged + # last tap group; the stem cases add a ragged channel group on top, and that + # is the axis the four original cases do not reach. + assert math.prod(problem.kernel) % cfg.TAP_BLOCK != 0 + assert (problem.cin % cfg.BLOCK_NC != 0) == (problem.cin == 3), ( + "the stem cases are the ragged-Cin ones; the others must not be" + ) + dtype = reference.torch_dtype(problem) + ops = reference.make_inputs(problem, seed=31, exact=True, dtype=dtype) + expected = reference.reference(problem, ops, "bwd-weight") + assert reference.is_exactly_representable(expected, dtype) + assert reference.compare(_run(problem, ops, config=cfg), + expected.to(dtype)).bitwise + + +@requires_gpu +def test_bitwise_standard_rejects_a_shifted_input(): + """Prove the comparison discriminates: a one-voxel shift must fail it.""" + problem = ConvProblem("shift", 16, 16, (6, 6, 6)) + ops = reference.make_inputs(problem, seed=11, exact=True) + actual = _run(problem, ops) + correct = reference.reference(problem, ops, "bwd-weight").to(torch.bfloat16) + assert reference.compare(actual, correct).bitwise + + shifted = torch.roll(ops["input"], shifts=1, dims=-1) + wrong = reference.reference( + problem, {**ops, "input": shifted}, "bwd-weight" + ).to(torch.bfloat16) + assert not reference.compare(actual, wrong).bitwise, ( + "a one-voxel shift of the input produced a bitwise-identical gradient; " + "the comparison is not discriminating" + ) + + +@requires_gpu +def test_a_permuted_tap_axis_is_detected(): + """The bug this module can uniquely have, pinned. + + The kernel's N axis is ``(tap, Cin)`` and its output offset is + ``co*taps*Cin + tap*Cin + ci``. Getting the tap ordering wrong -- reversing + it, or transposing (kd,kh,kw) -- produces a correctly shaped, correctly + scaled, entirely plausible weight gradient, and would pass every tolerance + test one could write. At ``k=3`` with a symmetric volume nothing else in + this file would catch it, so the wrong answer is constructed and required to + differ. + """ + problem = ConvProblem("taps", 16, 16, (6, 7, 8)) + ops = reference.make_inputs(problem, seed=13, exact=True) + actual = _run(problem, ops) + correct = reference.reference(problem, ops, "bwd-weight").to(torch.bfloat16) + assert reference.compare(actual, correct).bitwise + + for wrong in (correct.flip(2, 3, 4), correct.transpose(2, 4).contiguous()): + assert wrong.shape == actual.shape + assert not reference.compare(actual, wrong).bitwise, ( + "a permuted tap axis gave a bitwise-identical gradient; the " + "[Cout][tap][Cin] output ordering is untested by this suite" + ) + + +@requires_gpu +@pytest.mark.parametrize("problem", EDGE[:8], ids=_ids(EDGE[:8])) +def test_every_config_gives_the_same_answer(problem: ConvProblem): + """The tuning surface, not one point on it. + + This matters more here than in the other two directions because the + candidate list varies ``TAP_BLOCK`` and ``SPLIT_K``, and both change the + *decomposition* rather than only the tiling: a wrong tap-column decode shows + up only at ``TAP_BLOCK > 1``, and an off-by-one in the chunk arithmetic only + at split counts the shipped heuristic happens not to pick. + """ + ops = reference.make_inputs(problem, seed=2, exact=True) + expected = reference.reference(problem, ops, "bwd-weight") + dtype = reference.torch_dtype(problem) + if not reference.is_exactly_representable(expected, dtype): + pytest.skip("realized magnitudes exceed the mantissa in this dtype") + expected = expected.to(dtype) + k_total = problem.n * math.prod(problem.out_spatial) + cfgs = candidate_bwd_weight_configs( + problem.cout, problem.cin, problem.kernel, k_total, dtype, + splits=(0, 1, 3, 64), + ) + ran = 0 + for cfg in cfgs: + try: + actual = _run(problem, ops, config=cfg) + except triton.runtime.errors.OutOfResources: + continue # a loud failure; the sweep skips these too + ran += 1 + assert reference.compare(actual, expected).bitwise, f"{problem.label} {cfg}" + assert ran, "no candidate configuration was runnable" + + +@requires_gpu +def test_the_atomic_path_agrees_with_the_deterministic_one(): + """Same answer, different summation order -- so *not* bitwise, but close. + + The atomic path exists only to price determinism, and the price is only + meaningful if the two compute the same thing. The bar is the fp64 reference + rather than each other, because "equal to the wrong answer" is exactly what + a shared bug would look like. + """ + problem = ConvProblem("atomic", 64, 64, (10, 12, 16), padding=(0, 0, 0)) + ops = reference.make_inputs(problem, seed=19, exact=True) + expected = reference.reference(problem, ops, "bwd-weight") + assert reference.is_exactly_representable(expected, torch.bfloat16) + expected = expected.to(torch.bfloat16) + assert reference.compare(_run(problem, ops), expected).bitwise + assert reference.compare(_run(problem, ops, deterministic=False), expected).bitwise + + +@requires_gpu +def test_the_forward_kernel_can_express_backward_weight(): + """The reuse M2 got for free, checked here and then rejected on a trip count. + + Swapping the batch and channel axes of both activations turns + backward-weight into a forward convolution whose kernel extent is the + *output volume*. It is a real identity and the forward kernel really + computes it, which is what this half of the test shows. + + The other half is why ``reduce_gemm.py`` exists anyway. At config B's + ``dec3`` site that convolution has 8.4 million taps and a channel count of + ``N = 1``, so the forward's reduction loop -- ``taps * ceil(Cin/BLOCK_K)`` + iterations, each carrying a six-compare boundary predicate -- runs 8.4 + million times with 15 of every 16 ``BLOCK_K`` lanes masked off, and there is + no split-K anywhere in it. Both numbers are asserted rather than described, + because "too slow" is the kind of claim that rots. + """ + problem = ConvProblem("reuse", 4, 5, (4, 5, 6), padding=(0, 0, 0)) + ops = reference.make_inputs(problem, seed=29, exact=True) + expected = reference.reference(problem, ops, "bwd-weight") + assert reference.is_exactly_representable(expected, torch.bfloat16) + + # (Cin, N, ID, IH, IW) convolved with (Cout, N, OD, OH, OW) -> (Cin, Cout, k) + as_conv = conv3d_forward( + ops["input"].transpose(0, 1).contiguous(memory_format=torch.channels_last_3d), + ops["grad_output"].transpose(0, 1).contiguous( + memory_format=torch.channels_last_3d), + padding=0, + ) + assert tuple(as_conv.shape) == (problem.cin, problem.cout, *problem.kernel) + assert reference.compare( + as_conv.transpose(0, 1), expected.to(torch.bfloat16) + ).bitwise + + # And the shape of that same reuse at a real site. + big = ConvProblem("dec3", 128, 64, (130, 258, 258), padding=(0, 0, 0)) + reused_taps = math.prod(big.out_spatial) + cfg = default_config(big.cin * reused_taps, 1, big.cout, torch.bfloat16) + assert reused_taps == 8_388_608 + assert cfg.BLOCK_K >= 16 and big.n == 1, ( + "the reused kernel's reduction is Cin=N=1 deep but BLOCK_K cannot go " + "below the MFMA's kDim" + ) + assert reused_taps * triton.cdiv(big.n, cfg.BLOCK_K) > 8e6 + + +# --------------------------------------------------------------------------- +# Correctness: the tolerance standards +# --------------------------------------------------------------------------- + + +@requires_gpu +@pytest.mark.parametrize("problem", EDGE + CORPUS_PAIRS, + ids=_ids(EDGE + CORPUS_PAIRS)) +def test_no_worse_than_miopen(problem: ConvProblem): + """The honest bar at realistic magnitudes, against MIOpen on the same data.""" + ops = reference.make_inputs(problem, seed=17) + expected = reference.reference(problem, ops, "bwd-weight") + incumbent_err = reference.compare( + reference.incumbent(problem, ops, "bwd-weight"), expected + ) + actual = _run(problem, ops) + reference.assert_close(actual, expected, problem, "bwd-weight", + incumbent_error=incumbent_err) + + +@requires_gpu +def test_split_k_is_more_accurate_than_miopen_at_a_long_reduction(): + """A claim worth making in the other direction, for once. + + Splitting a 32k-term fp32 reduction into fixed chunks and summing the + partials is not just reproducible, it is *more accurate* than one long + accumulation -- the error of a sum of ``K`` terms grows like ``sqrt(K)`` and + a two-level sum trades that for ``sqrt(K/S) + sqrt(S)``. Measured here so + that "deterministic" is not read as "at some cost in accuracy": Triton's + error lands at the bf16 rounding limit of the output, and MIOpen's is + several times larger. + """ + problem = ConvProblem("acc", 64, 64, (34, 34, 34), padding=(0, 0, 0)) + ops = reference.make_inputs(problem, seed=31) + expected = reference.reference(problem, ops, "bwd-weight") + mine = reference.compare(_run(problem, ops), expected) + theirs = reference.compare( + reference.incumbent(problem, ops, "bwd-weight"), expected + ) + assert mine.max_abs < theirs.max_abs, f"triton {mine} vs miopen {theirs}" + + +@requires_gpu +def test_fp32_accumulates_in_fp32(): + """``more_determinism`` runs the model in fp32, and the backward too. + + A tf32-style split dot would pass any bf16-sized tolerance, so the bound is + fp32-sized and held against fp64. + """ + problem = ConvProblem("fp32", 48, 32, (7, 9, 5), dtype="fp32") + ops = reference.make_inputs(problem, seed=23) + expected = reference.reference(problem, ops, "bwd-weight") + actual = _run(problem, ops) + assert actual.dtype is torch.float32 + report = reference.compare(actual, expected) + peak = expected.abs().max().item() + assert report.max_abs < 1e-4 * peak, f"looks like a reduced-precision dot: {report}" + + +@requires_gpu +@pytest.mark.parametrize("stride,padding", [(2, 1), (2, 0), (3, 2)]) +def test_a_strided_convolution_is_served_correctly(stride, padding): + """Backward-data refuses a stride; this direction does not, so it is tested. + + ScaFFold's corpus has no strided non-transposed convolution, so nothing else + in this suite would exercise the ``o*s`` term at all, and an unexercised + multiply that is *also* not refused by ``is_supported`` is the combination + that returns a wrong gradient silently. + """ + problem = ConvProblem("strided", 16, 24, (9, 11, 13), + stride=(stride,) * 3, padding=(padding,) * 3) + ops = reference.make_inputs(problem, seed=37, exact=True) + expected = reference.reference(problem, ops, "bwd-weight") + assert reference.is_exactly_representable(expected, torch.bfloat16) + assert reference.compare(_run(problem, ops), + expected.to(torch.bfloat16)).bitwise + + +# --------------------------------------------------------------------------- +# Determinism -- the property this milestone exists for +# --------------------------------------------------------------------------- + + +#: A shape whose split count is well above 1, so that the deterministic path is +#: actually exercising the workspace and the reduction pass rather than the +#: single-split shortcut that trivially cannot disagree with itself. +_DET = ConvProblem("determinism", 64, 64, (18, 34, 34), padding=(0, 0, 0)) + + +@requires_gpu +def test_repeated_calls_are_bitwise_reproducible_in_process(): + problem = _DET + ops = reference.make_inputs(problem, seed=71) + cfg = bwd_weight_config(problem.cout, problem.cin, problem.kernel, + math.prod(problem.out_spatial), torch.bfloat16) + assert split_count(cfg, problem.cout, problem.cin, problem.tap_count, + math.prod(problem.out_spatial), + problem.out_spatial[2])[0] > 1, "not exercising split-K" + first = _run(problem, ops) + for _ in range(4): + assert torch.equal(first, _run(problem, ops)) + + +#: The ``k=1`` segmentation head, at a volume that splits ~800 ways, in **fp32**. +#: The dtype is the entire point -- see the negative-control test below. +_DET_K1 = ConvProblem("determinism-k1", 64, 6, (64, 64, 64), (1, 1, 1), + padding=(0, 0, 0), dtype="fp32") + + +@requires_gpu +@pytest.mark.parametrize("problem", [_DET, _DET_K1], ids=["k3-bf16", "k1-fp32"]) +def test_the_atomic_path_is_not_bitwise_reproducible(problem: ConvProblem): + """The negative control, and the reason the default is not the atomic one. + + Without this the reproducibility test above could pass on a kernel that was + reproducible for some unrelated reason -- a grid too small to race, say -- + and the claim would be about the shape rather than about the mechanism. + Float addition is not associative and ``tl.atomic_add`` fixes no order, so + at 100-odd racing splits a repeat that agrees bitwise every time would mean + the atomic path is not doing what it says. + + **The second cell is the interesting one, and it is the reason this test is + parametrized at all.** The ``k=1`` head at ``64 -> 6 @ 128^3`` was recorded + elsewhere in this project as a shape where the atomic control "reproduced by + scheduling accident". That is not the mechanism. The atomic accumulator is + fp32 and the *result* is bf16, so a reordering that perturbs the sum at the + fp32 ulp is simply invisible after the cast -- measured, the perturbation + there is about 300x below one bf16 ulp of the output. The splits are racing + the whole time; the race is under the resolution of the dtype it is being + observed in. Run the identical shape in fp32 and the control fires every + single time (15/15, ~400 ulps of the fp32 output). + + That distinction matters because it says what the control *can* certify: it + is informative wherever the reordering is resolvable in the output dtype, + and it certifies nothing on a short-``Cout`` bf16 shape -- a change that + made the deterministic path non-deterministic at the ``k=1`` head would be + invisible in a bf16 cell. So the ``k=1`` head is covered here in the dtype + where the control has teeth. + + If this ever goes flaky it is worth reading as a result rather than as a + flake: it would mean the splits stopped racing. + """ + dtype = reference.torch_dtype(problem) + k_total = problem.n * math.prod(problem.out_spatial) + cfg = bwd_weight_config(problem.cout, problem.cin, problem.kernel, k_total, + dtype) + splits = split_count(cfg, problem.cout, problem.cin, problem.tap_count, + k_total, problem.out_spatial[2])[0] + assert splits > 8, f"{splits} splits: too few writers to contend" + + ops = reference.make_inputs(problem, seed=71, dtype=dtype) + first = _run(problem, ops, deterministic=False) + differed = any( + not torch.equal(first, _run(problem, ops, deterministic=False)) + for _ in range(15) + ) + assert differed, ( + f"16 runs of the atomic path agreed bitwise at {splits} splits; either " + "the splits are not racing, or the reordering is below one ulp of " + f"{dtype} and this cell certifies nothing" + ) + + +_CHILD = textwrap.dedent( + """ + import hashlib, sys, torch + sys.path.insert(0, {repo!r}) + from triton_conv3d import reference + from triton_conv3d.reduce_gemm import conv3d_backward_weight + from triton_conv3d.shapes import ConvProblem + p = ConvProblem("determinism", 64, 64, (18, 34, 34), padding=(0, 0, 0)) + ops = reference.make_inputs(p, seed=71) + gw = conv3d_backward_weight(ops["input"], p.weight_shape, + ops["grad_output"], p.stride, p.padding) + # bf16 has no numpy dtype; widening to fp32 is exact, so the digest still + # answers the bitwise question. + print(hashlib.sha256(gw.float().cpu().numpy().tobytes()).hexdigest()) + """ +) + + +@requires_gpu +def test_three_separate_processes_agree_bitwise(): + """Process to process, which is the half of the claim a loop cannot test. + + An in-process repeat shares the allocator state, the JIT cache and the + module-level tuning table, so it would still pass if any of those were what + fixed the reduction order. Separate interpreters share none of it, which is + what makes this the test that the *shape* determines the split count. + """ + repo = str(pathlib.Path(__file__).resolve().parents[2]) + digests = [] + for _ in range(3): + proc = subprocess.run([sys.executable, "-c", _CHILD.format(repo=repo)], + capture_output=True, text=True, timeout=900) + assert proc.returncode == 0, proc.stderr[-2000:] + digests.append(proc.stdout.strip().splitlines()[-1]) + assert len(set(digests)) == 1, digests + + +# --------------------------------------------------------------------------- +# Entry-point behaviour +# --------------------------------------------------------------------------- + + +@requires_gpu +def test_the_output_is_a_channels_last_weight_of_the_right_shape(): + """The GEMM writes ``[Cout][tap][Cin]``, which *is* channels_last_3d. + + Worth asserting rather than assuming: it is the reason this direction needs + no layout transform at all, and a future change to the epilogue that + produced a contiguous weight instead would still pass every value test in + this file while costing the integration a permute per parameter per step. + """ + problem = ConvProblem("shape", 16, 40, (3, 11, 5)) + ops = reference.make_inputs(problem, seed=61) + gw = _run(problem, ops) + ref = torch.nn.grad.conv3d_weight( + ops["input"], problem.weight_shape, ops["grad_output"], + stride=problem.stride, padding=problem.padding, + ) + assert gw.shape == ref.shape + assert gw.is_contiguous(memory_format=torch.channels_last_3d) + assert gw.stride(1) == 1 + + +@requires_gpu +def test_an_out_the_kernel_would_overrun_is_refused(): + """The ``Cout`` extent is invisible to a stride check, and it is the extent. + + ``[Cout][kd][kh][kw][Cin]`` strides are + ``(taps*Cin, 1, kh*kw*Cin, kw*Cin, Cin)`` -- **not one of them mentions + Cout**. So a gradient allocated for ``Cout=8`` is stride-identical to one + allocated for ``Cout=64`` with the same ``Cin`` and kernel, and the + reduction pass takes its element count from ``weight_shape`` rather than + from ``gw``: passing the small one used to be accepted and wrote 55 296 + elements into a 6 912-element allocation. No fault and no exception -- the + write lands in whatever the caching allocator has next, and some other live + tensor is wrong later. + + The other three clauses are here for the same reason they are in the + function: a foreign device is a pointer this kernel will happily + dereference (ScaFFold runs four ranks per node), and a mismatched dtype + silently changes the dtype of the gradient the caller gets back. + """ + k = (3, 3, 3) + x = torch.randn((1, 32, 6, 6, 6), device="cuda", dtype=torch.bfloat16 + ).contiguous(memory_format=torch.channels_last_3d) + gy = torch.randn((1, 64, 6, 6, 6), device="cuda", dtype=torch.bfloat16 + ).contiguous(memory_format=torch.channels_last_3d) + ws = (64, 32, *k) + + right = grad_weight_empty(64, 32, k, dtype=torch.bfloat16, device="cuda") + small = grad_weight_empty(8, 32, k, dtype=torch.bfloat16, device="cuda") + # The trap, stated: the guard that used to be here could not tell these two + # apart, because the only thing that differs is an extent. + assert small.stride() == right.stride() + assert small.numel() * 8 == right.numel() + + for bad, what in ( + (small, "shape"), + (grad_weight_empty(64, 32, k, dtype=torch.float32, device="cuda"), + "dtype"), + (grad_weight_empty(64, 32, k, dtype=torch.bfloat16, device="cpu"), + "device"), + (torch.empty(ws, device="cuda", dtype=torch.bfloat16), "strides"), + ): + with pytest.raises(ValueError, match="out="): + conv3d_backward_weight(x, ws, gy, padding=1, out=bad) + # ...and the buffer that *is* right is still accepted, so the guard is not + # simply refusing everything. + assert conv3d_backward_weight(x, ws, gy, padding=1, + out=right).data_ptr() == right.data_ptr() + + +@requires_gpu +def test_the_gradient_buffer_is_allocated_in_the_layout_it_is_used_in(): + """One allocation, no copy. + + ``torch.empty(shape).contiguous(memory_format=channels_last_3d)`` allocates + NCDHW and then runs a permuting device copy to reach the layout it was + always going to be asked for. The contents are undefined either way, so + the copy transports nothing -- it is pure waste on a buffer this direction + allocates once per parameter per step, and the identical defect in the + forward measured 235x the cost of the one-shot allocation. + """ + from torch.utils._python_dispatch import TorchDispatchMode + + seen: list[str] = [] + + class _Record(TorchDispatchMode): + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + seen.append(str(func)) + return func(*args, **(kwargs or {})) + + with _Record(): + gw = grad_weight_empty(64, 32, (3, 3, 3), dtype=torch.bfloat16, + device="cuda") + assert gw.shape == (64, 32, 3, 3, 3) + assert gw.is_contiguous(memory_format=torch.channels_last_3d) + assert not [op for op in seen if "copy" in op or "clone" in op], seen + assert len(seen) == 1, f"expected one allocation and nothing else: {seen}" + + +@requires_gpu +def test_hoisted_workspace_and_out_are_equivalent(): + """Both are optimizations, so both must change nothing.""" + problem = ConvProblem("hoist", 32, 48, (10, 12, 16), padding=(0, 0, 0)) + ops = reference.make_inputs(problem, seed=53, exact=True) + inline = _run(problem, ops) + + k_total = math.prod(problem.out_spatial) + cfg = bwd_weight_config(problem.cout, problem.cin, problem.kernel, k_total, + torch.bfloat16) + splits, _ = split_count(cfg, problem.cout, problem.cin, problem.tap_count, + k_total, problem.out_spatial[2]) + ws = torch.empty( + workspace_elements(splits, problem.cout, problem.cin, problem.kernel), + dtype=torch.float32, device="cuda", + ) + gw = grad_weight_empty(problem.cout, problem.cin, problem.kernel, + dtype=torch.bfloat16, device="cuda") + hoisted = _run(problem, ops, workspace=ws, out=gw) + assert hoisted.data_ptr() == gw.data_ptr() + assert torch.equal(inline, hoisted) + + # An undersized workspace has to say *how big* it needed to be. A hoisted + # workspace is sized once, out of the step, from a number someone read + # somewhere -- and the number that was published for this corpus was + # understated by 1.46x, so the first thing that caller sees is this + # exception at step 1 with no way to compute the right size from it. + need = workspace_elements(splits, problem.cout, problem.cin, problem.kernel) + with pytest.raises(ValueError, match=rf"at least {need} float32 elements"): + _run(problem, ops, workspace=ws[:8]) + with pytest.raises(ValueError, match=rf"at least {need} float32 elements"): + _run(problem, ops, workspace=ws.double()) + + +@requires_gpu +def test_ncdhw_operands_are_converted_rather_than_misread(): + """The addressing assumes ``stride_c == 1`` on both activations. + + An NCDHW operand read with NDHWC strides gives a full-rate kernel and a + completely wrong gradient. ScaFFold's own backward can hand us either + layout depending on what produced the tensor, so this is not hypothetical. + """ + problem = ConvProblem("layout", 24, 16, (5, 6, 7)) + ops = reference.make_inputs(problem, seed=41, exact=True) + ndhwc = _run(problem, ops) + nc = {k: (v.contiguous() if torch.is_tensor(v) else v) for k, v in ops.items()} + assert nc["input"].stride(1) != 1 + assert torch.equal(ndhwc, _run(problem, nc)) diff --git a/triton_conv3d/tests/test_gather_gemm.py b/triton_conv3d/tests/test_gather_gemm.py new file mode 100644 index 0000000..55da17c --- /dev/null +++ b/triton_conv3d/tests/test_gather_gemm.py @@ -0,0 +1,1101 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""Tests for the forward gather-GEMM convolution. + +The organising idea is that a convolution kernel fails by *reading the wrong +voxel*, and a wrong voxel holds a plausible number. A tolerance-based test +waves that through: swap a tap, drop a boundary compare, or transpose two +spatial strides and the result is still smooth, still the right magnitude, and +still passes ``allclose``. So the primary standard here is **bitwise**, made +attainable by drawing operands from ``{-1, 0, 1}`` -- every product is exact and +every partial sum is a small integer, so the reference and the kernel must agree +exactly or the kernel is wrong. :func:`test_bitwise_standard_rejects_a_shifted_gather` +exists to prove that standard has teeth rather than being vacuously satisfied. + +The tolerance-based tests are still here, but as the *second* line: they are the +ones that hold at real magnitudes and real reduction lengths, where exactness is +not available. + +Everything that needs a GPU is skipped without one; the configuration-legality +tests are pure Python and always run, which matters because an illegal MFMA +config on gfx942 does not raise -- it silently emits no matrix instructions and +returns correct results at a fraction of the speed. +""" + +from __future__ import annotations + +import itertools +import math + +import pytest +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from triton_conv3d import reference +from triton_conv3d.gather_gemm import ( + ConvConfig, + candidate_configs, + conv3d_forward, + default_config, + is_supported, + is_supported_all, + to_rsck, +) +from triton_conv3d.shapes import ConvProblem, edge_cases, scaffold_corpus + +requires_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") + +#: The synthetic corpus, minus the transposed upsample (a later milestone). +EDGE = [p for p in edge_cases() if not p.transposed] + +#: Real ScaFFold shapes small enough to test against an fp64 reference. The hot +#: ones are 2 GiB activations; correctness does not need them, and every bug +#: these tests are looking for reproduces at 16^3. +#: +#: Selecting by ``volume * channels`` has a consequence worth stating, because +#: it is not obvious and it is what made the bitwise test on this list skip 11 +#: of 11 for two milestones: the problems that survive the filter are the +#: *widest* ones -- 256->512 up to 1024->1024, K = 6912 to 27648 -- because +#: those are the ones ScaFFold runs at a small spatial extent. So this list is +#: precisely the regime where the forward's reduction is longest, which is the +#: regime the exactness question is hardest in. See +#: :func:`test_corpus_shapes_match_bitwise`. +CORPUS_SMALL = [ + p for p in scaffold_corpus() + if not p.transposed and math.prod(p.spatial) * max(p.cin, p.cout) <= 1 << 22 +] + + +def _ids(problems): + return [p.name or p.label for p in problems] + + +def _run(problem: ConvProblem, ops: dict, **kwargs) -> torch.Tensor: + return conv3d_forward( + ops["input"], ops["weight"], ops["bias"], + problem.stride, problem.padding, **kwargs, + ) + + +# --------------------------------------------------------------------------- +# Configuration legality -- no GPU needed, and the failure mode is silent +# --------------------------------------------------------------------------- + + +def test_validate_rejects_the_silent_mfma_failures(): + """Each of these produces a working kernel with zero MFMA instructions. + + Verified by negative control in M0: ``BLOCK_K=8`` at ``nonkdim=16`` and an + illegal ``nonkdim=64`` both compile, run, and return correct results with the + dot lowered to vector FMA. Nothing raises and nothing warns, so a config + generator that merely ranked them last would still feed meaningless entries + into a best-of sweep -- which is why :meth:`ConvConfig.validate` refuses. + """ + bf16 = torch.bfloat16 + assert ConvConfig(BLOCK_K=8, matrix_instr_nonkdim=16).validate(bf16) + assert ConvConfig(matrix_instr_nonkdim=64).validate(bf16) + assert ConvConfig(BLOCK_M=24, matrix_instr_nonkdim=16).validate(bf16) + assert ConvConfig(BLOCK_M=64, BLOCK_N=16, num_warps=8).validate(bf16) + assert ConvConfig(num_warps=3).validate(bf16) + assert ConvConfig(num_stages=1).validate(bf16) + assert ConvConfig().validate(torch.float64) + # And the default is legal, or none of the above means anything. + assert ConvConfig().validate(bf16) is None + + +def test_validate_rejects_a_group_m_that_faults_the_gpu(): + """``GROUP_M`` was the one config field ``validate`` did not look at. + + Its failure mode is not the silent FMA fallback the others have, which is + why it is stated apart from them: the swizzle computes ``width = GROUP_M * + grid_n`` and then ``pid // width``, so ``GROUP_M = 0`` divides by zero inside + the kernel -- on gfx942 that is a garbage ``pid_m`` and a memory access + fault, not a trap -- and ``GROUP_M = -3`` reaches the kernel just as far. + + The other half of the pin is that every *legal* value is accepted, including + values that do not divide ``grid_m`` and values far larger than it; the + swizzle is a bijection for all of them and rejecting them would cost tuning + range for nothing. + """ + bf16 = torch.bfloat16 + assert ConvConfig(GROUP_M=0).validate(bf16) + assert ConvConfig(GROUP_M=-3).validate(bf16) + for group_m in (1, 5, 6, 7, 8, 4096): + assert ConvConfig(GROUP_M=group_m).validate(bf16) is None, group_m + + +def test_the_index_width_decision_covers_every_operand_including_the_weight(): + """The predicate behind ``INDEX_DTYPE``, checked without allocating 4 GiB. + + It used to be ``max(x.numel(), y.numel())``, and the weight's absence from it + was an assumption ("weights are never large enough") that nothing enforced: + at ``taps * Cin * Cout = 2.30e9`` the int32 offset wraps negative and the GPU + faults, with ``is_supported`` returning ``True``. Meta tensors carry a + ``numel`` and no storage, so this end of the fix costs nothing to state. + """ + from triton_conv3d.gather_gemm import _index_dtype + + small = torch.empty((1 << 20,), device="meta") + huge = torch.empty((1 << 31,), device="meta") # numel > 2**31 - 1 by one + assert _index_dtype(small, small, small) == tl.int32 + assert _index_dtype(huge, small, small) == tl.int64 + assert _index_dtype(small, huge, small) == tl.int64 + assert _index_dtype(small, small, huge) == tl.int64 + + +def test_block_k_constraint_follows_the_intrinsic_not_the_tile(): + """``BLOCK_K`` is constrained by the MFMA's reduction depth, which moves. + + bf16 on gfx942 has one intrinsic per shape: ``16x16x16`` at nonkdim 16 and + ``32x32x8`` at 32. So a ``BLOCK_K`` of 8 is legal at nonkdim 32 and illegal + at 16 -- the constraint is not a property of the block size alone, and the + briefing's claim that ``BLOCK_K=16`` rows get pruned at nonkdim 16 is simply + wrong arithmetic (16 % 16 == 0). + """ + ok32 = ConvConfig(BLOCK_M=32, BLOCK_N=32, BLOCK_K=8, matrix_instr_nonkdim=32, + num_warps=4, kpack=1) + assert ok32.validate(torch.bfloat16) is None + assert ConvConfig(BLOCK_M=32, BLOCK_N=32, BLOCK_K=8, + matrix_instr_nonkdim=16, num_warps=4).validate(torch.bfloat16) + assert ConvConfig(BLOCK_K=16, matrix_instr_nonkdim=16, + kpack=1).validate(torch.bfloat16) is None + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32], + ids=["bf16", "fp16", "fp32"]) +def test_default_config_fits_in_lds_in_every_dtype(dtype): + """The block sizes were chosen against bf16; fp32 operands are twice the bytes. + + This is a regression test for a hole M2 fell into rather than a hypothetical: + ``default_config`` returned ``128x128x128`` for ``Cin >= 512``, which is 64 + KiB in bf16 and 128 KiB in fp32, and the *shipped* configuration therefore + raised ``OutOfResources`` on any wide fp32 convolution. ``more_determinism`` + runs the model in fp32, so a real ScaFFold configuration reached it. + + An explicitly supplied ``config=`` is still allowed to overflow and still + fails loudly; what must never overflow is the one the entry point picks by + itself. + """ + from triton_conv3d.gather_gemm import _LDS_BYTES + + for cin in (3, 6, 64, 128, 256, 512, 1024): + for cout in (6, 64, 128, 256, 512, 1024): + for m in (512, 4096, 2 << 20): + cfg = default_config(m, cin, cout, dtype) + assert cfg.validate(dtype) is None, f"{cin}->{cout} m={m}: {cfg}" + assert cfg.lds_bytes(dtype) <= _LDS_BYTES, ( + f"{cin}->{cout} m={m}: {cfg} needs " + f"{cfg.lds_bytes(dtype)} B of LDS" + ) + + +@pytest.mark.parametrize( + "problem", EDGE + CORPUS_SMALL, ids=_ids(EDGE + CORPUS_SMALL) +) +def test_default_config_is_legal_for_every_shape(problem: ConvProblem): + """The heuristic must never hand back a config that loses the matrix core. + + It is the config used when no tuned entry exists, which is most of the time, + and it derives block sizes from the shape -- so the tiny synthetic problems + are exactly where it can round itself into an illegal combination. + """ + dtype = reference.torch_dtype(problem) + m = problem.n * math.prod(problem.out_spatial) + cfg = default_config(m, problem.cin, problem.cout, dtype) + assert cfg.validate(dtype) is None, f"{problem.label}: {cfg} -> {cfg.validate(dtype)}" + + +@pytest.mark.parametrize("problem", CORPUS_SMALL[:6], ids=_ids(CORPUS_SMALL[:6])) +def test_every_candidate_config_is_legal(problem: ConvProblem): + """The sweep must not contain a config that cannot reach the matrix core. + + Otherwise the sweep's *reported* winner could be an FMA kernel that happened + to beat the others, and the whole tuning surface would be measuring the + wrong thing. + """ + dtype = reference.torch_dtype(problem) + m = problem.n * math.prod(problem.out_spatial) + cfgs = candidate_configs(m, problem.cin, problem.cout, dtype) + assert cfgs + for cfg in cfgs: + assert cfg.validate(dtype) is None, f"{cfg}: {cfg.validate(dtype)}" + + +# --------------------------------------------------------------------------- +# Support predicate +# --------------------------------------------------------------------------- + + +def test_is_supported_declines_what_the_kernel_cannot_do(): + """A false positive returns a wrong answer; a false negative costs speed. + + The caller's fallback is MIOpen, which is correct everywhere, so the + predicate is deliberately asymmetric and this test pins that asymmetry. + """ + x = torch.empty((1, 8, 4, 4, 4), device="meta", dtype=torch.bfloat16) + w = torch.empty((8, 8, 3, 3, 3), device="meta", dtype=torch.bfloat16) + # Meta tensors are not on a device, so the real predicate rejects them; the + # checks below are about everything *except* device placement. + assert not is_supported(x, w, padding=1) + + if not torch.cuda.is_available(): + pytest.skip("the remaining branches need a real device") + x = torch.empty((1, 8, 4, 4, 4), device="cuda", dtype=torch.bfloat16) + w = torch.empty((8, 8, 3, 3, 3), device="cuda", dtype=torch.bfloat16) + assert is_supported(x, w, padding=1) + assert not is_supported(x, w, padding=1, groups=2) + assert not is_supported(x, w.float(), padding=1) + assert not is_supported(x, torch.empty((8, 4, 3, 3, 3), device="cuda", + dtype=torch.bfloat16), padding=1) + # A kernel wider than the padded input has no output voxels at all, which + # the M-unravel cannot express. + tiny = torch.empty((1, 8, 1, 4, 4), device="cuda", dtype=torch.bfloat16) + assert not is_supported(tiny, w, padding=0) + assert is_supported(tiny, w, padding=1) + + +@requires_gpu +def test_is_supported_never_raises_on_an_argument_it_cannot_parse(): + """A gate that throws is not a gate. + + This predicate is the first rung of a Triton -> MIOpen ladder, so a caller + asking "will you serve this?" about a ``padding`` it holds in a variable must + get an answer. ``_triple`` raises ``TypeError`` for anything neither ``int`` + nor iterable, and only ``ValueError`` was caught -- so ``padding=None`` and + ``padding=1.5`` came back out of the *predicate* as exceptions while + ``padding=(1, 1)`` and ``padding='same'`` correctly returned ``False``. + """ + x = torch.empty((1, 8, 4, 4, 4), device="cuda", dtype=torch.bfloat16) + w = torch.empty((8, 8, 3, 3, 3), device="cuda", dtype=torch.bfloat16) + assert is_supported(x, w, padding=1) + for bad in (None, 1.5, "same", (1, 1), [1, 2, 3, 4], object()): + assert not is_supported(x, w, padding=bad), bad + assert not is_supported(x, w, stride=bad), bad + assert not is_supported(x, w, dilation=bad), bad + + +@requires_gpu +def test_is_supported_declines_a_bias_torch_itself_rejects(): + """The kernel masks the bias load against ``Cout`` and assumes stride 1. + + Neither of those is a property of the bias, so both failures are silent: a + short bias reads whatever is in memory past its end (``nan`` if you are + lucky, a plausible finite number if you are not), and a stride-2 view of the + right length applies ``[0,1,2,...]`` where the caller passed ``[0,2,4,...]``. + ``torch.conv3d`` refuses the first outright; this predicate now refuses both. + """ + bf16 = torch.bfloat16 + x = torch.empty((1, 8, 4, 5, 6), device="cuda", dtype=bf16) + w = torch.empty((32, 8, 3, 3, 3), device="cuda", dtype=bf16) + bias = torch.empty(32, device="cuda", dtype=bf16) + assert is_supported(x, w, bias, padding=1) + + assert not is_supported(x, w, bias[:4], padding=1) # too short + assert not is_supported(x, w, torch.empty(64, device="cuda", + dtype=bf16)[::2], padding=1) + assert not is_supported(x, w, bias.float(), padding=1) + assert not is_supported(x, w, bias.cpu(), padding=1) + assert not is_supported(x, w, bias.view(1, 32), padding=1) + # And the entry point declines rather than running on it. + with pytest.raises(NotImplementedError): + conv3d_forward(x, w, bias[:4], padding=1) + + +@requires_gpu +def test_the_forward_gate_alone_is_a_trap_for_a_caller_that_differentiates(): + """``stride=2``: served forward, served backward-weight, refused backward-data. + + The disagreement is real and each side of it is deliberate -- the forward's + M-unravel simply steps by ``s``, backward-weight is indexed by the *output* + voxel so a stride is three extra multiplies, and backward-data has no kernel + of its own and is the forward contraction on a flipped weight, which is only + the right contraction at unit stride. What was wrong was that nothing said + so: a training caller who asked :func:`is_supported`, got ``True`` and built + a graph node found out at ``backward()``, where its own MIOpen fallback is no + longer reachable because the node is already in the graph. + + So this pins both halves: the trap still exists at the direction gates (they + describe their own kernels and must keep doing so), and + :func:`is_supported_all` is the one question that closes it. + """ + from triton_conv3d.bwd_data import conv3d_backward_data, is_supported_bwd_data + from triton_conv3d.reduce_gemm import is_supported_bwd_weight + + bf16 = torch.bfloat16 + x = torch.empty((1, 8, 8, 8, 8), device="cuda", dtype=bf16) + w = torch.empty((16, 8, 3, 3, 3), device="cuda", dtype=bf16) + gy = torch.empty((1, 16, 4, 4, 4), device="cuda", dtype=bf16) + args = dict(stride=2, padding=1) + + assert is_supported(x, w, **args) + assert is_supported_bwd_weight(x, w.shape, gy, **args) + assert not is_supported_bwd_data(gy, w, x.shape, **args) + assert not is_supported_all(x, w, **args) + + # The trap itself, run: the forward serves the call and the gradient this + # very forward produces cannot be turned back into an input gradient. + y = conv3d_forward( + x.contiguous(memory_format=torch.channels_last_3d), + w.contiguous(memory_format=torch.channels_last_3d), + **args, + ) + assert tuple(y.shape) == (1, 16, 4, 4, 4) + with pytest.raises(NotImplementedError): + conv3d_backward_data( + y.contiguous(memory_format=torch.channels_last_3d), + w.contiguous(memory_format=torch.channels_last_3d), + x.shape, + **args, + ) + + # And the same problem at unit stride, where all three do agree, is not + # collateral damage: the combined gate must still say yes. + assert is_supported_all(x, w, padding=1) + + +@requires_gpu +def test_is_supported_all_is_exactly_the_three_gates_conjoined(): + """The combined gate against the conjunction it stands for, term by term. + + Two things could rot here and neither would fail loudly. The gradient is + passed to the two backward predicates as a metadata-only stand-in -- a + one-element allocation expanded to the output shape -- which is sound only + while those predicates read metadata and nothing else, so it is compared + against the answer a *real* gradient gets. And the output shape is computed + here rather than by the caller, so a wrong one would be a gate answering + about a different problem than the one it was asked about. + """ + from triton_conv3d.bwd_data import is_supported_bwd_data + from triton_conv3d.reduce_gemm import is_supported_bwd_weight + + bf16 = torch.bfloat16 + cases = [ + # (x shape, w shape, kwargs) + ((1, 8, 8, 8, 8), (16, 8, 3, 3, 3), dict(padding=1)), # all yes + ((1, 8, 8, 8, 8), (16, 8, 3, 3, 3), dict(stride=2, padding=1)), # bwd-data no + ((1, 8, 8, 8, 8), (16, 8, 3, 3, 3), dict(padding=1, groups=2)), # fwd no + ((1, 8, 8, 8, 8), (16, 8, 3, 3, 3), dict(padding=0)), # all yes + ((1, 8, 8, 8, 8), (16, 8, 1, 1, 1), dict(padding=0)), # k=1 + ((1, 8, 8, 8, 8), (16, 8, 3, 3, 3), dict(padding=2)), # p > dil*(k-1) + ((1, 8, 1, 8, 8), (16, 8, 3, 3, 3), dict(padding=1)), # thin D + ((1, 8, 8, 8, 8), (16, 8, 3, 3, 3), dict(dilation=2, padding=2)), + ] + for x_shape, w_shape, kwargs in cases: + x = torch.empty(x_shape, device="cuda", dtype=bf16) + w = torch.empty(w_shape, device="cuda", dtype=bf16) + s = kwargs.get("stride", 1) + p = kwargs.get("padding", 0) + d = kwargs.get("dilation", 1) + out = tuple( + (x_shape[2 + i] + 2 * p - d * (w_shape[2 + i] - 1) - 1) // s + 1 + for i in range(3) + ) + gy = torch.empty((x_shape[0], w_shape[0]) + out, device="cuda", dtype=bf16) + expected = ( + is_supported(x, w, **kwargs) + and is_supported_bwd_data(gy, w, x.shape, **kwargs) + and is_supported_bwd_weight(x, w.shape, gy, **kwargs) + ) + assert is_supported_all(x, w, **kwargs) is expected, (x_shape, w_shape, kwargs) + + +@requires_gpu +def test_is_supported_all_never_raises_on_an_argument_it_cannot_parse(): + """Total, for the same reason :func:`is_supported` is: it is a gate. + + The forward's predicate runs first and refuses everything unparsable, so the + output-shape arithmetic below it is never reached with an argument that would + make it throw -- but the caller's contract is "you get an answer", and that + has to be checked and not argued. + """ + x = torch.empty((1, 8, 4, 4, 4), device="cuda", dtype=torch.bfloat16) + w = torch.empty((8, 8, 3, 3, 3), device="cuda", dtype=torch.bfloat16) + assert is_supported_all(x, w, padding=1) + for bad in (None, 1.5, "same", (1, 1), [1, 2, 3, 4], object()): + assert not is_supported_all(x, w, padding=bad), bad + assert not is_supported_all(x, w, stride=bad), bad + assert not is_supported_all(x, w, dilation=bad), bad + + +@requires_gpu +def test_is_supported_declines_degenerate_extents(): + """Three shapes where the kernel returned something ``torch.conv3d`` does not. + + Each clears the "every output voxel must exist" test and then diverges, which + is the asymmetry the predicate exists to prevent -- the MIOpen fallback would + have raised on all three and the Triton path silently did not. + + ``N = 0`` is deliberately *not* in the rejection list: it agrees with torch, + both in the shape it returns and in doing no work to return it, so declining + it would be a false negative with nothing behind it. + """ + bf16 = torch.bfloat16 + x = torch.empty((1, 8, 4, 5, 6), device="cuda", dtype=bf16) + w = torch.empty((16, 8, 3, 3, 3), device="cuda", dtype=bf16) + + # A zero-length spatial axis: returned a volume of pure padding, where torch + # raises "Only zero batch or zero channel inputs are supported". + assert not is_supported(torch.empty((1, 8, 0, 5, 6), device="cuda", + dtype=bf16), w, padding=2) + # A zero-size kernel: ``(in + 2p - d(k-1) - 1)//s + 1`` gains one at k=0, so + # the returned output was *larger* than the input. + assert not is_supported(x, torch.empty((16, 8, 0, 0, 0), device="cuda", + dtype=bf16), padding=0) + # Cin = 0: returned Cout channels of zeros where torch returns a tensor with + # no channels at all -- a different shape, not a different value. + assert not is_supported(torch.empty((1, 0, 4, 5, 6), device="cuda", dtype=bf16), + torch.empty((16, 0, 3, 3, 3), device="cuda", dtype=bf16), + padding=1) + + empty_batch = torch.empty((0, 8, 4, 5, 6), device="cuda", + dtype=bf16).contiguous(memory_format=torch.channels_last_3d) + assert is_supported(empty_batch, w, padding=1) + assert tuple(conv3d_forward(empty_batch, w, padding=1).shape) == (0, 16, 4, 5, 6) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two GPUs") +def test_is_supported_declines_operands_on_different_devices(): + """Both operands on *a* GPU is not the same as both on the *same* GPU. + + Triton launches on the current device and dereferences the foreign pointer + regardless. ScaFFold runs four ranks per node, and with peer access enabled + that reads another rank's activations instead of faulting -- a wrong answer + with no symptom. Skipped, not absent, on a single-GPU box. + """ + bf16 = torch.bfloat16 + x = torch.empty((1, 8, 4, 4, 4), device="cuda:0", dtype=bf16) + w = torch.empty((8, 8, 3, 3, 3), device="cuda:0", dtype=bf16) + bias = torch.empty(8, device="cuda:0", dtype=bf16) + assert is_supported(x, w, bias, padding=1) + assert not is_supported(x, w.to("cuda:1"), padding=1) + assert not is_supported(x, w, bias.to("cuda:1"), padding=1) + + +# --------------------------------------------------------------------------- +# Correctness: the bitwise standard +# --------------------------------------------------------------------------- + + +@requires_gpu +@pytest.mark.parametrize("problem", EDGE, ids=_ids(EDGE)) +def test_exact_operands_match_bitwise(problem: ConvProblem): + """The strictest standard, on the problems chosen to break addressing. + + The corpus covers channel counts that are not multiples of any plausible + ``BLOCK_K``, prime and unit spatial extents, anisotropic volumes and kernels, + ``N > 1``, and a volume smaller than the kernel in every axis -- where every + tap is masked somewhere and the boundary predicate is the whole computation. + """ + ops = reference.make_inputs(problem, seed=3, exact=True) + expected = reference.reference(problem, ops, "fwd") + dtype = reference.torch_dtype(problem) + if not reference.is_exactly_representable(expected, dtype): + pytest.skip("realized magnitudes exceed the mantissa in this dtype") + actual = _run(problem, ops) + report = reference.compare(actual, expected.to(dtype)) + assert report.bitwise, f"{problem.label}: {report}" + + +@requires_gpu +def test_bitwise_standard_rejects_a_shifted_gather(): + """Prove the bitwise test has teeth: a one-voxel shift must fail it. + + Without this, ``test_exact_operands_match_bitwise`` could be passing because + ``{-1,0,1}`` operands happen to make everything agree -- a vacuous test, and + this project has already shipped two of those in the GroupNorm suite. So + compare the kernel's output against a reference computed for a *deliberately + wrong* padding and require a mismatch. + """ + problem = ConvProblem("shift", 16, 16, (6, 6, 6)) + ops = reference.make_inputs(problem, seed=11, exact=True) + actual = _run(problem, ops) + correct = reference.reference(problem, ops, "fwd").to(torch.bfloat16) + assert reference.compare(actual, correct).bitwise + + # Same problem, but the reference gathers from one voxel further along W. + shifted = torch.roll(ops["input"], shifts=1, dims=-1) + wrong = reference.reference( + problem, {**ops, "input": shifted}, "fwd" + ).to(torch.bfloat16) + assert not reference.compare(actual, wrong).bitwise, ( + "a one-voxel shift of the input produced a bitwise-identical result; " + "the comparison is not discriminating" + ) + + +@requires_gpu +@pytest.mark.parametrize("problem", CORPUS_SMALL, ids=_ids(CORPUS_SMALL)) +def test_corpus_shapes_match_bitwise(problem: ConvProblem): + """Same standard, on the shapes ScaFFold actually runs. + + The synthetic cases break addressing; these check that nothing about the + real channel widths -- 64 through 1024, all multiples of 256 bytes in bf16 + and so all candidates for the stride hazard -- changes the answer. + + This test used to skip 11 of 11 cases, every run, for two milestones, and + the mechanism is worth stating because the obvious fix does not work here. + A dense ``{-1,0,1}`` draw is exactly representable only while the realized + sums stay inside the mantissa, and bf16 holds integers only to 256; the + forward reduces over ``Cin * taps``, which is 27 648 terms at ``Cin = 1024`` + and gives sums running to several hundred. ``test_bwd_data.py`` and + ``test_bwd_weight.py`` hit the same wall and got out of it by restating each + channel pair at ``6x7x8`` -- their reductions run over ``Cout * taps`` and + over the *output volume*, so a smaller volume shortens them. The forward's + does not depend on the volume at all, so no shape substitution can help it: + the only lever is the operands. + + So the activations are thinned to :func:`reference.exact_density` -- about + 1-4% here, which leaves ~300 live terms per output element and a realized + maximum of 58-71 against the limit of 256 -- while the shape, the channel + widths and the weight stay exactly as they are. It is the gather that is + under test, not the arithmetic. Two things then have to be asserted rather + than assumed, because a thinned draw is exactly how one would accidentally + build a test that compares zeros against zeros: that the answer is mostly + nonzero, and that the comparison still rejects a one-voxel shift. + + And it no longer skips. If the draw is ever not exact, that is a fact worth + failing on rather than stepping around -- the skip is what hid the hole. + """ + density = reference.exact_density(problem, "fwd") + ops = reference.make_inputs(problem, seed=5, exact=True, density=density) + expected = reference.reference(problem, ops, "fwd") + dtype = reference.torch_dtype(problem) + assert reference.is_exactly_representable(expected, dtype), ( + f"{problem.label}: the thinned draw at density {density:.4g} still " + f"realizes |max| = {expected.abs().max().item():g}, which " + f"{problem.dtype} cannot hold exactly" + ) + nonzero = (expected != 0).to(torch.float64).mean().item() + assert nonzero > 0.5, ( + f"{problem.label}: only {nonzero:.1%} of the reference is nonzero; the " + "thinning has gone far enough to make the comparison vacuous" + ) + + actual = _run(problem, ops) + assert reference.compare(actual, expected.to(dtype)).bitwise + + # The negative control, per case rather than once: run the kernel over an + # input shifted by one voxel and require the *same* comparison to reject it. + # A tolerance would wave that through, and so would a draw thinned until + # everything it touches is zero -- this is what says the passing assertion + # above is coverage rather than a coincidence, at this width. + shifted = _run(problem, {**ops, "input": torch.roll(ops["input"], 1, dims=-1)}) + assert not reference.compare(shifted, expected.to(dtype)).bitwise, ( + f"{problem.label}: a one-voxel shift of the input produced a " + "bitwise-identical result; the comparison is not discriminating" + ) + + +def test_the_bitwise_corpus_is_not_entirely_skipped(): + """A regression guard on this file, not on the kernel. + + ``test_bwd_data.py`` and ``test_bwd_weight.py`` both carry a guard of this + name because this suite has already once reported "89 passed" while + skipping 100% of its real-shape cases. The forward had no such guard, which + is how it went two milestones with 11 of 11 skipping and nothing saying so. + + The shape of the guard differs from theirs, because the fix here differs: + :func:`test_corpus_shapes_match_bitwise` has no skip branch left, so what + needs pinning is not "enough cases are representable" but the two ways the + test could still stop meaning anything -- the parametrization collapsing to + nothing or to only narrow shapes, and the thinning going so far that every + output element is a sum of nothing. Both are pure arithmetic, so this runs + without a GPU, which is the other half of the point: a guard that skips with + the thing it guards is not a guard. + """ + assert len(CORPUS_SMALL) >= 8, CORPUS_SMALL + # The widths are the reason this list exists. A filter that quietly stopped + # selecting the deep encoder problems would leave the forward tested bitwise + # only at the synthetic sizes again. + assert max(p.cin for p in CORPUS_SMALL) >= 1024 + assert max(p.cout for p in CORPUS_SMALL) >= 1024 + for problem in CORPUS_SMALL: + k = problem.gemm_shape("fwd")[2] + density = reference.exact_density(problem, "fwd") + assert 0.0 < density <= 1.0, f"{problem.label}: density {density}" + # Live terms per output element: the reduction that actually happens. + # At 64 a wrong gather still has dozens of independent chances to show + # up in every element; below it the draw would be approaching a test of + # whether zero equals zero. + assert density * k >= 64.0, ( + f"{problem.label}: only {density * k:.1f} of {k} terms contribute; " + "the thinned draw is close to vacuous" + ) + + +@requires_gpu +@pytest.mark.parametrize("problem", EDGE, ids=_ids(EDGE)) +def test_every_config_gives_the_same_answer(problem: ConvProblem): + """Tiling must not be observable in the result. + + A boundary bug usually only shows up at one tile shape: a mask that is right + when ``BLOCK_M`` divides ``OUT_W`` and wrong when it does not, or a ``BLOCK_K`` + remainder that is only exercised when ``Cin`` is not a multiple of the tile. + Sweeping the whole candidate list against a bitwise reference tests the + *tuning surface* rather than one point on it, which matters because the + tuned table is free to pick any of them. + + Over all of ``EDGE`` rather than its first eight. The eight are the channel + and spatial oddities, and stopping there left ``batched`` (the only ``n > 1`` + case), ``kernel_aniso``, ``smaller_than_kernel``, ``unpadded``, ``pointwise`` + and both non-bf16 dtypes checked at the *default* config alone -- so a tiling + bug that needed ``BLOCK_M=256`` with ``n > 1``, or fp32 at ``nonkdim=32``, + had nowhere to show up. fp32 and fp16 matter here in their own right: the + dtype moves the MFMA intrinsic's reduction depth and therefore which + ``BLOCK_K`` values are even legal. + """ + ops = reference.make_inputs(problem, seed=2, exact=True) + expected = reference.reference(problem, ops, "fwd") + dtype = reference.torch_dtype(problem) + if not reference.is_exactly_representable(expected, dtype): + pytest.skip("realized magnitudes exceed the mantissa in this dtype") + expected = expected.to(dtype) + m = problem.n * math.prod(problem.out_spatial) + # Plus the shipped default, which for a shape too small for any seed tile + # (``Cout=6``) is the only candidate there is. + cfgs = candidate_configs(m, problem.cin, problem.cout, dtype, + group_ms=(6, 8)) + cfgs = list(dict.fromkeys(cfgs + [default_config(m, problem.cin, + problem.cout, dtype)])) + ran = 0 + for cfg in cfgs: + try: + actual = _run(problem, ops, config=cfg) + except triton.runtime.errors.OutOfResources: + # A tile whose operands do not fit in 64 KiB of LDS. Unlike the + # MFMA constraints this failure is *loud*: Triton refuses at compile + # time and says so, so it needs no static guard -- the sweep skips + # it and so does this test. + continue + ran += 1 + assert reference.compare(actual, expected).bitwise, f"{problem.label} {cfg}" + assert ran, "no candidate configuration was runnable" + + +# --------------------------------------------------------------------------- +# Correctness: the tolerance standards +# --------------------------------------------------------------------------- + + +@requires_gpu +@pytest.mark.parametrize("problem", EDGE + CORPUS_SMALL, + ids=_ids(EDGE + CORPUS_SMALL)) +def test_no_worse_than_miopen(problem: ConvProblem): + """The honest bar for a replacement: not better than MIOpen, but not worse. + + Held against an fp64 reference with MIOpen measured on the same operands, so + the bar adapts to shape and reduction length instead of being a constant + somebody picked. Random operands rather than ``{-1,0,1}`` because this is + the standard that has to hold at realistic magnitudes, where the reduction + genuinely does lose bits. + """ + ops = reference.make_inputs(problem, seed=17) + expected = reference.reference(problem, ops, "fwd") + incumbent_err = reference.compare( + reference.incumbent(problem, ops, "fwd"), expected + ) + actual = _run(problem, ops) + reference.assert_close(actual, expected, problem, "fwd", + incumbent_error=incumbent_err) + + +@requires_gpu +def test_fp32_accumulates_in_fp32(): + """fp32 in, fp32 out, and no silent demotion to a reduced-precision dot. + + ``more_determinism`` runs the model in fp32, and on this backend it is not + obvious whether ``tl.dot`` on fp32 operands uses the exact ``f32`` MFMA or a + tf32-style split. A tf32 dot would still pass a bf16-sized tolerance, so the + check is against fp64 with an fp32-sized bound. + """ + problem = ConvProblem("fp32", 48, 32, (7, 9, 5), dtype="fp32") + ops = reference.make_inputs(problem, seed=23) + expected = reference.reference(problem, ops, "fwd") + actual = _run(problem, ops) + assert actual.dtype is torch.float32 + report = reference.compare(actual, expected) + # tf32 keeps 10 explicit mantissa bits; fp32 keeps 23. A bound between the + # two separates them, which a dtype-generic tolerance would not. + peak = expected.abs().max().item() + assert report.max_abs < 1e-4 * peak, f"looks like a reduced-precision dot: {report}" + + +# --------------------------------------------------------------------------- +# Entry-point behaviour +# --------------------------------------------------------------------------- + + +@requires_gpu +def test_bias_is_added_once_and_broadcast_over_channels(): + problem = ConvProblem("bias", 32, 24, (5, 6, 7), (1, 1, 1), + padding=(0, 0, 0), bias=True) + ops = reference.make_inputs(problem, seed=31, exact=True) + with_bias = _run(problem, ops) + without = conv3d_forward(ops["input"], ops["weight"], None, + problem.stride, problem.padding) + delta = (with_bias.float() - without.float()) + # The difference must be exactly the bias, in every voxel. + expected = ops["bias"].float().view(1, -1, 1, 1, 1).expand_as(delta) + assert torch.equal(delta, expected) + + +@requires_gpu +def test_ncdhw_input_is_converted_rather_than_misread(): + """A contiguous NCDHW input must give the same answer, not a transposed one. + + The addressing assumes ``stride_xc == 1``. Silently reading an NCDHW tensor + with NDHWC strides produces a full-rate kernel and a completely wrong result, + so the entry point converts; this pins that it converts rather than assumes. + """ + problem = ConvProblem("layout", 24, 16, (5, 6, 7)) + ops = reference.make_inputs(problem, seed=41, exact=True) + ndhwc = _run(problem, ops) + nc = {k: (v.contiguous() if torch.is_tensor(v) else v) for k, v in ops.items()} + assert nc["input"].stride(1) != 1 + ncdhw = _run(problem, nc) + assert torch.equal(ndhwc, ncdhw) + + +@requires_gpu +def test_out_buffer_is_written_in_place_and_is_validated(): + """``out=`` had no check of any kind, and nothing downstream can catch one. + + The grid is sized from the problem rather than from the buffer, so an + undersized ``out=`` is an out-of-bounds *device write* -- 1920 elements into + a 128-element allocation, observed, with no error and no fault, surviving + only because the allocator slab happened to be bigger. An NCDHW buffer is + the other half: the store addressing writes NDHWC strides into it and returns + a scrambled answer at full speed. + + Handing a preallocated gradient buffer to the backward is precisely what the + ``nn.Module`` adapter will do, so this is the parameter that most needs the + check and had none. + """ + problem = ConvProblem("out", 16, 24, (4, 5, 6)) + ops = reference.make_inputs(problem, seed=67, exact=True) + expected = _run(problem, ops) + + buf = torch.empty_like(expected) + got = _run(problem, ops, out=buf) + assert got.data_ptr() == buf.data_ptr(), "out= was allocated over, not written" + assert torch.equal(got, expected) + + shape = tuple(expected.shape) + bf16 = torch.bfloat16 + # Undersized, right layout: the write ran off the end. + with pytest.raises(ValueError): + _run(problem, ops, out=torch.empty((1, 24, 2, 2, 2), device="cuda", dtype=bf16)) + # Right shape, NCDHW: read with NDHWC strides. + with pytest.raises(ValueError): + _run(problem, ops, out=torch.empty(shape, device="cuda", dtype=bf16)) + with pytest.raises(ValueError): + _run(problem, ops, out=torch.empty_like(expected, dtype=torch.float32)) + with pytest.raises(ValueError): + _run(problem, ops, out=torch.empty(shape, dtype=bf16)) # on the CPU + + +@requires_gpu +def test_the_output_is_allocated_directly_in_channels_last(): + """One allocation in the final layout, not an NCDHW one plus a full copy. + + ``torch.empty(shape).contiguous(memory_format=channels_last_3d)`` is a + correct way to spell an expensive thing: it allocates NCDHW and then copies + the whole tensor, which measured **2.82 ms against 0.012 ms** on a 256 MiB + output -- 235x, on a path a training step takes about 19 times. The copy is + invisible in the result, so what pins it is the peak allocation: the wrong + form needs two output-sized buffers live at once, the right form needs one. + """ + x = torch.randn((1, 64, 64, 64, 64), device="cuda", dtype=torch.bfloat16 + ).contiguous(memory_format=torch.channels_last_3d) + w = torch.randn((64, 64, 3, 3, 3), device="cuda", dtype=torch.bfloat16) + wr = to_rsck(w) + conv3d_forward(x, w, padding=1, weight_rsck=wr) # warm the JIT out of the way + + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + y = conv3d_forward(x, w, padding=1, weight_rsck=wr) + peak = torch.cuda.max_memory_allocated() - base + want = y.numel() * y.element_size() + assert peak < 1.5 * want, ( + f"the call peaked at {peak} B for a {want} B output; that is the " + "allocate-then-copy form, not the one-shot one" + ) + + +def test_the_layout_conversion_is_a_no_op_only_where_stride_c_is_moot(): + """The kernel's unstated ``stride_c == 1`` rests on a PyTorch detail. + + ``contiguous(memory_format=channels_last_3d)`` is a *no-op* on an + NCDHW-contiguous tensor whenever enough dims are size 1 that the two formats + cannot be told apart -- PyTorch skips size-1 dims in its format predicate. + The entry point converts unconditionally, so in those shapes it converts + nothing and the kernel reads NCDHW strides as if they were NDHWC. + + That is safe, but for a reason outside this code: every such shape either has + ``stride(1) == 1`` outright (all three spatial extents are 1) or has + ``Cin == 1``, which makes the channel stride unobservable because the only + channel index the kernel ever dereferences is 0. 87 of the 243 shapes over + ``{1,2,3}^5`` are ambiguous and all 87 land in one of those two cases -- a + property of PyTorch's predicate rather than of ours, so it is pinned rather + than assumed. + """ + ambiguous = 0 + for shape in itertools.product((1, 2, 3), repeat=5): + t = torch.empty(shape) + if not (t.is_contiguous() + and t.is_contiguous(memory_format=torch.channels_last_3d)): + continue + ambiguous += 1 + assert t.stride(1) == 1 or shape[1] == 1, shape + assert ambiguous, "no shape was ambiguous; the enumeration is vacuous" + + +@requires_gpu +def test_an_ambiguous_layout_still_gives_the_right_answer(): + """One of the shapes above, end to end: ``Cin = 1``, where nothing converts. + + The conversion is a no-op, the strides the kernel is handed are NCDHW's, and + the result still has to be the reference's -- which it is only because the + one stride that differs is the one a single-channel input never uses. + """ + x = torch.randint(-1, 2, (1, 1, 4, 5, 6), device="cuda", dtype=torch.int8 + ).to(torch.bfloat16) + w = torch.randint(-1, 2, (8, 1, 3, 3, 3), device="cuda", dtype=torch.int8 + ).to(torch.bfloat16) + assert x.is_contiguous() + assert x.contiguous(memory_format=torch.channels_last_3d).data_ptr() == x.data_ptr() + assert torch.equal(conv3d_forward(x, w, padding=1), F.conv3d(x, w, padding=1)) + + +@requires_gpu +def test_hoisted_weight_transform_is_validated(): + """``weight_rsck`` supplies every weight *value* the kernel reads. + + ``w`` is then consulted only for its shape, so a hoisted transform of the + wrong parameter runs and returns a smooth, correctly shaped, entirely wrong + result -- measured ``max_abs = 60.0``. That is a live hazard rather than a + "you asked for it": the transform exists to be cached across calls, and a + cache keyed on the parameter's version is exactly the thing that goes stale. + """ + problem = ConvProblem("wr", 16, 24, (4, 5, 6)) + ops = reference.make_inputs(problem, seed=67, exact=True) + good = to_rsck(ops["weight"]) + assert torch.equal(_run(problem, ops, weight_rsck=good), _run(problem, ops)) + + other = torch.randn((24, 16, 1, 1, 1), device="cuda", dtype=torch.bfloat16) + with pytest.raises(ValueError): + _run(problem, ops, weight_rsck=to_rsck(other)) # a different kernel + with pytest.raises(ValueError): + _run(problem, ops, weight_rsck=good.float()) + # Right shape, wrong layout: the B tile load assumes Cout is contiguous. + with pytest.raises(ValueError): + _run(problem, ops, + weight_rsck=good.transpose(3, 4).contiguous().transpose(3, 4)) + + +@requires_gpu +def test_every_weight_layout_gives_the_same_answer(): + """The weight is read where it lies, so its strides pick the B load. + + Three layouts take three different decisions -- ``channels_last_3d`` is + addressed in place with a gathered tile, PyTorch's default is copied because + a gathered tile is 5-8x slower when *neither* channel axis is unit-stride, + and an RSCK-strided weight is addressed in place with a contiguous one -- + and they must not produce three answers. Bitwise, not close: it is the same + multiply-accumulate in the same order, and anything less would mean the + layout had leaked into the arithmetic. + + The RSCK-strided case is the one a test is really needed for. It has + PyTorch's shape over this kernel's storage order, which is what an + integration would allocate to make the B tile contiguous, and it is the case + where ``to_rsck`` is a no-op and ``weight`` and ``weight_rsck`` are the same + tensor -- so it is also where a mix-up between them would hide. + """ + problem = ConvProblem("layouts", 32, 48, (4, 5, 6)) + ops = reference.make_inputs(problem, seed=53, exact=True) + w = ops["weight"] + layouts = { + "channels_last": w.contiguous(memory_format=torch.channels_last_3d), + "contiguous": w.contiguous(), + "rsck_strided": (w.permute(2, 3, 4, 1, 0).contiguous() + .permute(4, 3, 0, 1, 2)), + } + ref = _run(problem, ops) + for name, wl in layouts.items(): + assert torch.equal(wl, w), name # same values, different strides + assert torch.equal(ref, _run(problem, {**ops, "weight": wl})), name + assert torch.equal(ref, _run(problem, ops, weight_rsck=to_rsck(w))) + # ``to_rsck`` of an already-RSCK-strided weight must not copy: that is what + # makes the layout free for a caller who chooses it, and ``.contiguous()`` + # returning ``self`` is the whole mechanism. + assert to_rsck(layouts["rsck_strided"]).data_ptr() == \ + layouts["rsck_strided"].data_ptr() + + +@requires_gpu +def test_hoisted_weight_transform_is_equivalent(): + """``weight_rsck`` is an optimization, so it must change nothing observable. + + A caller that hoists the transform out of a training step must get the + identical result to one that lets the entry point decide. + """ + problem = ConvProblem("hoist", 32, 48, (4, 5, 6)) + ops = reference.make_inputs(problem, seed=53, exact=True) + inline = _run(problem, ops) + hoisted = _run(problem, ops, weight_rsck=to_rsck(ops["weight"])) + assert torch.equal(inline, hoisted) + + +@requires_gpu +def test_output_is_channels_last_and_matches_torch_shape(): + problem = ConvProblem("shape", 16, 40, (3, 11, 5)) + ops = reference.make_inputs(problem, seed=61) + y = _run(problem, ops) + ref = F.conv3d(ops["input"], ops["weight"], ops["bias"], + stride=problem.stride, padding=problem.padding) + assert y.shape == ref.shape + assert y.is_contiguous(memory_format=torch.channels_last_3d) + + +@requires_gpu +def test_unsupported_calls_raise_rather_than_return_garbage(): + x = torch.randn((1, 8, 4, 4, 4), device="cuda", dtype=torch.bfloat16) + w = torch.randn((8, 4, 3, 3, 3), device="cuda", dtype=torch.bfloat16) + with pytest.raises(NotImplementedError): + conv3d_forward(x, w, padding=1, groups=2) + good_w = torch.randn((8, 8, 3, 3, 3), device="cuda", dtype=torch.bfloat16) + with pytest.raises(ValueError): + conv3d_forward(x, good_w, padding=1, + config=ConvConfig(BLOCK_K=8, matrix_instr_nonkdim=16)) + + +@requires_gpu +def test_repeated_calls_are_bitwise_reproducible(): + """No float atomics, fixed grid, fixed accumulation order. + + ScaFFold's default configuration is *not* bitwise reproducible today because + MIOpen's backward-weight uses atomics. The forward has no reason to inherit + that, and stating the property as a test is what stops a later split-K + variant from quietly giving it up. + """ + problem = ConvProblem("determinism", 64, 64, (8, 12, 10)) + ops = reference.make_inputs(problem, seed=71) + first = _run(problem, ops) + for _ in range(4): + assert torch.equal(first, _run(problem, ops)) + + +@requires_gpu +@pytest.mark.slow +def test_indices_beyond_int32_are_addressed_correctly(): + """A 2.2 GiB activation: the offsets must widen, and the far end must be read. + + Unsharded scale 8 is ``1 x 128 x 258^3 = 2.20e9`` elements, which is where + MIOpen itself asserts (upstream bug, reproducer filed) and where the + buffer-load fast path is lost because ``is_within_2gb`` reads the whole + storage. Losing buffer loads is a performance question; getting the *index* + wrong is a correctness one, and only a tensor this size asks it. + + The check is placed at the far end deliberately: a truncated 32-bit offset + aliases back to the start of the tensor, so a spot check near the end catches + it while a check of the mean would not. + """ + free, _ = torch.cuda.mem_get_info() + if free < 12 << 30: + pytest.skip("needs ~12 GiB free") + cin, cout, sp = 128, 16, (258, 258, 258) + x = torch.zeros((1, cin, *sp), device="cuda", dtype=torch.bfloat16).contiguous( + memory_format=torch.channels_last_3d + ) + w = torch.zeros((cout, cin, 1, 1, 1), device="cuda", dtype=torch.bfloat16) + # One channel of one weight, so the output is a copy of one input channel. + w[0, 0, 0, 0, 0] = 1.0 + x[0, 0, -1, -1, -1] = 3.0 + x[0, 0, 0, 0, 0] = 5.0 + y = conv3d_forward(x, w, padding=0) + assert y[0, 0, -1, -1, -1].item() == 3.0 + assert y[0, 0, 0, 0, 0].item() == 5.0 + assert y.sum().item() == 8.0 + + +@requires_gpu +@pytest.mark.slow +def test_indices_beyond_int32_are_addressed_correctly_with_taps_and_padding(): + """The same widening where ``tap_off`` and the ``PADDED`` predicate are live. + + The test above uses a ``1x1x1`` weight and no padding, so the widened row + offset is never bumped by a tap and the six boundary compares are compiled + out entirely -- a change that widened only the pointwise path would pass it. + This is the same ``1 x 128 x 258^3`` activation (2.20e9 elements, unsharded + scale 8) at ``k=3, padding=1``. + + One tap, the last one, so that ``tap_off`` is at its maximum and the far + corner of the output reads the far corner of the input: a truncated 32-bit + offset aliases back towards the start, which a spot check at the end catches + and a check of the mean does not. + """ + free, _ = torch.cuda.mem_get_info() + if free < 12 << 30: + pytest.skip("needs ~12 GiB free") + cin, cout, sp = 128, 16, (258, 258, 258) + x = torch.empty((1, cin, *sp), device="cuda", dtype=torch.bfloat16, + memory_format=torch.channels_last_3d).zero_() + w = torch.zeros((cout, cin, 3, 3, 3), device="cuda", dtype=torch.bfloat16) + # Tap (2,2,2) of channel 0 alone. At padding 1 that is y[o] = x[o + 1]. + w[0, 0, 2, 2, 2] = 1.0 + x[0, 0, -1, -1, -1] = 3.0 + x[0, 0, 1, 1, 1] = 5.0 + y = conv3d_forward(x, w, padding=1) + assert y[0, 0, -2, -2, -2].item() == 3.0 + assert y[0, 0, 0, 0, 0].item() == 5.0 + assert y.sum().item() == 8.0 + + +@requires_gpu +@pytest.mark.slow +def test_a_weight_beyond_int32_is_addressed_correctly(): + """``taps * Cin * Cout`` over ``2**31``, which used to fault the GPU. + + The weight offset was int32 unconditionally, on a stated assumption + ("weights are never large enough") that ``is_supported`` did not enforce: the + reviewer's matched pair at ``k=13`` differing only in ``Cout`` ran clean at + 1.15e9 weight elements and took a memory access fault at 2.30e9. Not + reachable from this model -- its widest weight is 28.3 M elements, 80x below + -- but the package is meant to be lifted into DistConv and released, and a + kernel whose reason for existing is MIOpen's int32 overflow should not have + one of its own. + + Shaped for the *offset* and not for the arithmetic: 2 taps, ``M = 1``, and + the widths chosen so the GEMM stays trivial while the row offset does not. + What has to overflow is ``dij * stride_wt + offs_k * stride_wc``, because + ``offs_n`` is a *second* ``addptr`` and is sign-extended on its own -- a first + attempt at this test put the excess there and passed against the unfixed + kernel. So the quantity to push past ``2**31`` is ``taps*Cin*Cout - Cout``, + which here is 65,537 elements over. + + ``w`` is an expanded view: ``weight_rsck`` supplies the values and ``w`` is + read only for its shape, which keeps this to one 4.29 GiB allocation rather + than two. + """ + free, _ = torch.cuda.mem_get_info() + if free < 8 << 30: + pytest.skip("needs ~8 GiB free") + bf16, cin, cout = torch.bfloat16, 16385, 65536 + wr = torch.zeros((2, 1, 1, cin, cout), device="cuda", dtype=bf16) + assert (2 * cin - 1) * cout > 2**31 - 1 # the largest row offset + w = torch.zeros((), device="cuda", dtype=bf16).expand(cout, cin, 2, 1, 1) + x = torch.zeros((1, cin, 2, 1, 1), device="cuda", dtype=bf16).contiguous( + memory_format=torch.channels_last_3d) + + wr[1, 0, 0, cin - 1, cout - 1] = 1.0 # the last element of the weight + x[0, cin - 1, 1, 0, 0] = 3.0 + wr[0, 0, 0, 0, 0] = 1.0 # and the first + x[0, 0, 0, 0, 0] = 5.0 + + y = conv3d_forward(x, w, padding=0, weight_rsck=wr) + assert tuple(y.shape) == (1, cout, 1, 1, 1) + assert y[0, cout - 1, 0, 0, 0].item() == 3.0 + assert y[0, 0, 0, 0, 0].item() == 5.0 + assert y.sum().item() == 8.0 diff --git a/triton_conv3d/tests/test_infra.py b/triton_conv3d/tests/test_infra.py new file mode 100644 index 0000000..5a125c5 --- /dev/null +++ b/triton_conv3d/tests/test_infra.py @@ -0,0 +1,1607 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""Tests for the measurement infrastructure itself. + +No kernels exist yet. What exists is a shape model, a cost model, a reference +and a timing harness, and every performance claim we make later is only as good +as those. So they get tested first, and mostly by cross-checking them against +PyTorch rather than against my own arithmetic: :func:`test_output_shape_matches_torch` +and :func:`test_flops_match_gemm_decomposition` between them caught a real error +in the transposed-convolution FLOP count, where the tap factor was applied twice. + +The GPU tests are skipped without a device; the shape and cost model tests are +pure Python and always run. +""" + +from __future__ import annotations + +import dataclasses +import json +import math + +import pytest +import torch +import torch.nn.functional as F + +from triton_conv3d import reference +from triton_conv3d.shapes import ( + _CORPUS_PATH, + BUFFER_OP_MAX_BYTES, + DIRECTIONS, + INT32_MAX, + ConvProblem, + edge_cases, + scaffold_corpus, +) + +requires_gpu = pytest.mark.skipif( + not torch.cuda.is_available(), reason="needs a GPU" +) + +#: ``include_large=True`` because the shape and cost model tests below are meta +#: tensors and integer arithmetic -- a 4 GiB activation costs nothing here, and +#: until this call existed the two int32-boundary cases were never instantiated +#: by anything at all. The GPU tests parametrize over ``SMALL`` instead. +ALL = list(scaffold_corpus()) + list(edge_cases(include_large=True)) +SMALL = [p for p in edge_cases() if math.prod(p.spatial) * p.cin <= 1 << 16] + + +def _ids(problems): + return [p.name or p.label for p in problems] + + +# --------------------------------------------------------------------------- +# Shape model +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("problem", ALL, ids=_ids(ALL)) +def test_output_shape_matches_torch(problem: ConvProblem): + """The derived output shape is what PyTorch actually produces. + + Cross-checking the whole corpus against the operator it models is what makes + the extracted shapes trustworthy; an off-by-one in the padding arithmetic + would otherwise propagate silently into every FLOP count and every roofline. + Run on meta tensors so a 2 GiB scale-8 activation costs nothing. + """ + x = torch.empty(problem.input_shape, device="meta") + w = torch.empty(problem.weight_shape, device="meta") + op = F.conv_transpose3d if problem.transposed else F.conv3d + y = op(x, w, None, stride=problem.stride, padding=problem.padding) + assert tuple(y.shape) == problem.output_shape + + +@pytest.mark.parametrize("problem", ALL, ids=_ids(ALL)) +@pytest.mark.parametrize("direction", DIRECTIONS) +def test_flops_match_gemm_decomposition(problem: ConvProblem, direction): + """``flops()`` and ``gemm_shape()`` must describe the same contraction. + + They are derived independently -- one from the convolution's definition, one + from the implicit-GEMM decomposition the kernels will use -- so agreement is + a real check rather than a tautology. + """ + m, n, k = problem.gemm_shape(direction) + assert 2 * m * n * k == problem.flops(direction) + + +@pytest.mark.parametrize("problem", ALL, ids=_ids(ALL)) +def test_bytes_counts_each_tensor_once(problem: ConvProblem): + """Compulsory traffic is the three tensors, nothing more and nothing less.""" + x = math.prod(problem.input_shape) * problem.elem_bytes + y = math.prod(problem.output_shape) * problem.elem_bytes + w = math.prod(problem.weight_shape) * problem.elem_bytes + for direction in DIRECTIONS: + assert problem.bytes(direction) == x + y + w + + +def test_transposed_flops_have_no_phantom_tap_factor(): + """A ``k == s`` transposed convolution does one MAC per output voxel. + + With kernel equal to stride the scatter windows tile the output rather than + overlapping, so each output voxel receives exactly one contribution. Scaling + the input volume by the tap count *and* keeping the tap factor would inflate + the count eightfold, which is exactly the bug this pins. + """ + p = ConvProblem("t", 64, 32, (8, 8, 8), (2, 2, 2), (2, 2, 2), (0, 0, 0), + transposed=True) + assert p.out_spatial == (16, 16, 16) + macs = math.prod(p.out_spatial) * p.cin * p.cout + assert p.flops("fwd") == 2 * macs + + +def test_the_int32_edge_cases_bracket_the_element_boundary(): + """The pair has to sit either side of 2**31 elements, or it pins nothing. + + It did not: ``int32_below`` was ``64 -> 64 @ 512^3``, 8.59e9 elements -- + four times *above* the boundary it is named for, so both cases were above + it and the transition was unbracketed. Asserted on the element count + rather than on the predicate so that this fails if either the shape or the + predicate moves. + """ + cases = {p.name: p for p in edge_cases(include_large=True)} + below, above = cases["int32_below"], cases["int32_above"] + assert below.max_elements < 2**31 <= above.max_elements + assert not below.index_exceeds_int32 + assert above.index_exceeds_int32 + # Same channel pair and kernel: the only thing that differs is the volume. + assert (below.cin, below.cout, below.kernel) == (above.cin, above.cout, + above.kernel) + # And the boundary is where the largest index -- not the count -- crosses. + assert below.max_elements - 1 <= INT32_MAX < above.max_elements - 1 + + small = ConvProblem("small", 32, 32, (8, 8, 8)) + assert not small.index_exceeds_int32 + #: ``bench/baseline.py`` records the predicate under its old name. + assert small.needs_int64 is small.index_exceeds_int32 + + +def test_the_2gib_cliff_is_a_byte_problem_and_not_an_index_problem(): + """The two int32 predicates are about different quantities, in both senses. + + ``conv 128->64 @ 130x258x258`` is the shape behind the project's two + largest numbers (769x and 2789x): DistConv's halo pushes its activation + 3.2% past 2 GiB, MIOpen falls off its solver database, and Triton does not. + That shape holds 1.11e9 elements -- about half of int32's range -- so an + element-counting predicate says nothing about it, and ``needs_int64`` + used to be read as though it did. What it exceeds is the *byte* limit on + the whole storage, which is what decides buffer-op eligibility + (``PLAN.md`` §3.5). + """ + cliff = next(p.halo_variant for p in scaffold_corpus() + if p.halo_variant.label == "conv 128->64 k3x3x3 @ 130x258x258") + assert not cliff.index_exceeds_int32 + assert cliff.max_elements / 2**31 < 0.55 # half int32's range + assert not cliff.buffer_ops_eligible + assert cliff.max_activation_bytes / BUFFER_OP_MAX_BYTES == pytest.approx( + 1.032, abs=0.002 # 3.2% past 2 GiB + ) + + # It is the only corpus shape on either side of that line, in either shape + # mode -- and *no* corpus shape needs a 64-bit element index. A test or a + # dispatch rule parametrized on the index predicate selects nothing. + over = [p.halo_variant.label for p in scaffold_corpus() + if not p.halo_variant.buffer_ops_eligible] + assert over == ["conv 128->64 k3x3x3 @ 130x258x258"] + assert not any(p.index_exceeds_int32 or p.halo_variant.index_exceeds_int32 + for p in scaffold_corpus()) + + +def test_corpus_covers_the_three_scaffold_configurations(): + corpus = scaffold_corpus() + assert len(corpus) > 40 + kinds = {(p.kernel, p.stride, p.padding, p.transposed) for p in corpus} + assert kinds == { + ((3, 3, 3), (1, 1, 1), (1, 1, 1), False), + ((2, 2, 2), (2, 2, 2), (0, 0, 0), True), + ((1, 1, 1), (1, 1, 1), (0, 0, 0), False), + } + # Ordered by measured cost, so truncation keeps what matters. + costs = [sum(m["ms_per_step"] for m in p.measured) for p in corpus] + assert costs == sorted(costs, reverse=True) + + +def test_halo_variant_is_the_shape_distconv_actually_issues(): + """The halo'd form is derived, not guessed, and it matches the shape dump. + + Upstream DistConv concatenates a ``k // 2`` halo and zeroes the padding on + every axis it manages -- including unsplit ones, where the slab is provably + zeros -- so a convolution routed through it reaches MIOpen two voxels larger + per axis and unpadded. ``halo_variant`` reconstructs that from ``halo`` + alone; this pins the reconstruction against ``halo_in_shape``, which the + shape dump recorded independently. If they ever disagree, every profiled + number in ``measured`` is attached to the wrong problem. + + This is the *incumbent's* form. What ScaFFold's own Triton rung issues is + :meth:`ConvProblem.production_variant`, pinned separately below against a + census of real calls. + """ + raw = json.loads(_CORPUS_PATH.read_text())["problems"] + corpus = scaffold_corpus() + assert len(raw) == len(corpus) + for entry, problem in zip(raw, corpus): + halo = problem.halo_variant + assert list(halo.input_shape) == entry["halo_in_shape"] + # The halo changes the input, never the output -- that is what makes it + # a halo and not a padding change. + assert halo.output_shape == problem.output_shape + # Padding is dropped exactly on the axes that gained a halo, and left + # alone elsewhere -- that swap is the whole transformation. + assert halo.padding == tuple( + 0 if h else p for h, p in zip(problem.halo, problem.padding) + ) + assert halo.flops("fwd") == problem.flops("fwd") + + +def test_halo_variant_is_a_distinct_miopen_problem(): + """The two forms must not collide in any table keyed by label. + + MIOpen keys its find database on the full descriptor including padding, so + ``64-128-128-128-...-1x1x1`` and ``64-130-130-130-...-0x0x0`` tune + separately and can land on different kernels. A baseline that labelled + them the same would let a cell measured on one be compared against a + profile of the other -- silently, and in whichever direction flatters the + kernel under test. + """ + hot = [p for p in scaffold_corpus() if any(p.halo)] + assert hot, "corpus has no halo'd problems; the dump lost halo_dhw" + for p in hot: + assert p.halo_variant.label != p.label + # And a problem with no halo is its own variant, so "both" shape modes do + # not measure the transposed upsamples and 1x1x1 convs twice. + for p in scaffold_corpus(): + if not any(p.halo): + assert p.halo_variant is p + + +def test_the_production_variant_is_what_a_real_step_issues(): + """The one that was wrong, pinned against a measurement. + + Every published ``conv`` cell in this project measures + :meth:`ConvProblem.halo_variant`, on the premise that "DistConv halos them + all". ScaFFold does not route its convolutions through DistConv: the + adapter in ``ScaFFold/unet/conv3d.py`` exchanges a halo only on axes with + more than one shard, so H and W keep ``padding = 1`` at every configuration + and all three axes do at one GPU. A projection built on the halo'd cells + over-credited a tuning commit by 3x before anyone checked. + + Checked here against ``census_corpus()`` -- a recording of the shapes and + paddings real ``FastConv3d`` calls handed the kernels at all four + configurations -- rather than against the same arithmetic twice. The + segmentation head is excluded on ``cout``: its output channel count is + ``n_categories + 1``, a dataset knob, and the corpus and the census were + taken with different values of it (6 and 3). + """ + from triton_conv3d.shapes import census_corpus, production_corpus + + def key(p): + return (p.transposed, p.cin, p.cout, tuple(p.kernel), + tuple(p.spatial), tuple(p.padding), tuple(p.stride), p.n) + + census = {key(p) for p in census_corpus()} + assert len(census) > 60, "the census is missing; nothing is being checked" + missing = [p for p in production_corpus() + if key(p) not in census and p.kernel != (1, 1, 1)] + assert not missing, ( + "the corpus's production form does not match what a real step issued: " + + ", ".join(p.qualified_label for p in missing) + ) + # And the census is in the form it claims: every k>1 convolution padded. + unpadded = [p for p in census_corpus() + if p.kernel == (3, 3, 3) and not any(p.padding)] + assert not unpadded, ( + "a k=3 production convolution arrived unpadded, which would mean the " + "adapter's halo plan changed: " + + ", ".join(p.qualified_label for p in unpadded) + ) + + +def test_the_three_forms_are_told_apart_by_the_qualified_label(): + """A cell must never be quotable as a form it is not. + + ``label`` carries the extent but not the padding, and the two sharded forms + of one site differ in *both* -- while the unsharded production form differs + from the DistConv one in the padding alone at some extents. Any table that + mixes forms therefore has to key on ``qualified_label``. + """ + sharded = [p for p in scaffold_corpus() if any(p.shard_halo)] + assert sharded, "corpus has no sharded problems; shard_halo_dhw was lost" + for p in sharded: + forms = {p.qualified_label, p.production_variant.qualified_label, + p.halo_variant.qualified_label} + assert len(forms) == 3, f"{p.label}: forms collide -> {forms}" + # The adapter halos D and leaves H and W padded; DistConv does neither. + assert p.production_variant.padding == (0, *p.padding[1:]) + assert p.halo_variant.padding == (0, 0, 0) + assert p.production_variant.spatial == (p.spatial[0] + 2, *p.spatial[1:]) + # Unsharded, the adapter form *is* the logical one and says so by identity. + for p in scaffold_corpus(): + if not any(p.shard_halo): + assert p.production_variant is p + + +def test_the_production_corpus_is_padded_where_the_halo_corpus_is_not(): + """The headline of the whole distinction, as a number. + + If this ever reads "0 padded" again, either the adapter has started haloing + every axis or ``shard_halo`` has been confused with ``halo`` -- and the + consequence is that every backward-weight kernel silently stops compiling + the ``PADDED`` body it compiles at every ``k = 3`` site today, so every + number in this project's adapter-form tables would describe a kernel + production no longer launches. (Until 2026-08-05 the consequence was + larger still: ``bwd_weight_config`` declined a tuned ``TAP_BLOCK > 1`` row + on a padded problem, so this count decided which *tile* eight sites ran.) + """ + from triton_conv3d.shapes import halo_corpus, production_corpus + + padded_prod = [p for p in production_corpus() if any(p.padding)] + padded_halo = [p for p in halo_corpus() if any(p.padding)] + assert len(padded_prod) == 42, len(padded_prod) + assert padded_halo == [] + # Every one of them is a k=3 convolution; the k=1 head and the k=2 + # upsamplers are genuinely unpadded in every form. + assert {p.kernel for p in padded_prod} == {(3, 3, 3)} + + +def test_stored_efficiency_agrees_with_the_cost_model(): + """The corpus's ``pct_roofline`` must be what ``efficiency(ms_per_call)`` says. + + They are computed by different code -- one by ``make_corpus.py`` out of the + profile's own FLOP and byte counts, one here out of the shape -- so agreement + is a real cross-check, and it caught a real error. ``make_corpus.py`` used + to divide the profile's *per-step* FLOP count by the *per-call* time, which + multiplies the efficiency by the number of call sites. Every affected + problem is a symmetric ``C -> C`` convolution occurring at two sites, so the + artifact read as "MIOpen is excellent on symmetric convolutions and poor on + asymmetric ones" and produced the three forward points that appeared to + exceed 100% of roofline. With it fixed, MIOpen's forward spans 21-68% + everywhere and the three impossible points are gone. + """ + for problem in scaffold_corpus(): + for m in problem.measured: + got = 100 * problem.efficiency(m["ms_per_call"], m["direction"]) + assert got == pytest.approx(m["pct_roofline"], abs=0.06, rel=0.01), ( + f"{problem.label} [{m['direction']}, config {m['config']}]: " + f"stored {m['pct_roofline']}%, cost model {got:.3f}%" + ) + # And no forward cell exceeds the roof, which is what the artifact implied. + fwd = [m["pct_roofline"] for p in scaffold_corpus() for m in p.measured + if m["direction"] == "fwd"] + assert fwd and max(fwd) < 100 + + +def test_roofline_switches_at_the_crossover(): + """Below ~182 FLOP/byte the memory roof binds; above it, compute does.""" + memory_bound = ConvProblem("thin", 3, 8, (16, 16, 16)) + compute_bound = ConvProblem("fat", 512, 512, (16, 16, 16)) + assert memory_bound.arithmetic_intensity() < 182 + assert memory_bound.roofline_flops() < 600e12 + assert compute_bound.arithmetic_intensity() > 182 + assert compute_bound.roofline_flops() == 600e12 + + +# --------------------------------------------------------------------------- +# Reference and tolerance policy +# --------------------------------------------------------------------------- + + +@requires_gpu +@pytest.mark.parametrize("problem", SMALL, ids=_ids(SMALL)) +def test_reference_agrees_with_miopen_within_tolerance(problem: ConvProblem): + """MIOpen itself must pass the bar we intend to hold our kernel to. + + If the incumbent failed this, the tolerance would be wrong rather than + MIOpen -- so this is a test of the policy, not of MIOpen. + """ + ops = reference.make_inputs(problem, seed=7) + for direction in DIRECTIONS: + expected = reference.reference(problem, ops, direction) + actual = reference.incumbent(problem, ops, direction) + # MIOpen's backward-weight is the one direction that is not a single + # rounding: it reduces with atomics, so two identical calls differ + # bitwise and the result carries several roundings rather than the one + # ``error_bound`` charges by default. Measured at ``conv 32->32 k3x3x3 + # @ 8x8x8``, its error wanders over 0.61-1.05 ulps of the peak from call + # to call while the forward sits at a fixed 0.284 -- and a one-ulp bound + # therefore does not merely fail it, it fails it *intermittently*, which + # is the worse outcome. ``convT 64->32 k2x2x2 @ 8x8x8`` is the other + # cell that reaches past one ulp. The nondeterminism and the size of + # the excess are both pinned by + # :func:`test_the_incumbents_extra_roundings_are_the_atomic_ones`, so + # this is a measured allowance rather than a tolerance nudged until the + # test passed. Only this direction and only the incumbent get it: our + # own backward-weight reduces its split-K partials in fp32 and stores + # once, so it is held to ``roundings=1`` like everything else. + reference.assert_close(actual, expected, problem, direction, + roundings=2 if direction == "bwd-weight" else 1) + + +@requires_gpu +@pytest.mark.parametrize("problem", SMALL, ids=_ids(SMALL)) +@pytest.mark.parametrize("direction", DIRECTIONS) +def test_exact_inputs_give_a_bitwise_reference(problem: ConvProblem, direction): + """With ``{-1,0,1}`` operands the contraction is exact, so equality holds. + + This is the standard that catches indexing and masking bugs: a kernel that + reads a neighbouring voxel still produces a plausible number, and only an + exact comparison rejects it. MIOpen passing it is what establishes that the + standard is attainable rather than aspirational. + """ + ops = reference.make_inputs(problem, seed=3, exact=True) + expected = reference.reference(problem, ops, direction) + if not reference.is_exactly_representable(expected, reference.torch_dtype(problem)): + pytest.skip("realized magnitudes exceed the mantissa in this dtype") + actual = reference.incumbent(problem, ops, direction) + report = reference.compare(actual, expected) + assert report.bitwise, f"{problem.label} [{direction}]: {report}" + + +def test_error_bound_grows_with_reduction_length_and_shrinks_with_precision(): + expected = torch.randn(4096, dtype=torch.float64) + short = ConvProblem("short", 8, 8, (8, 8, 8)) + long_ = ConvProblem("long", 1024, 1024, (8, 8, 8)) + assert reference.error_bound(long_, expected) > reference.error_bound(short, expected) + fp32 = ConvProblem("fp32", 64, 64, (8, 8, 8), dtype="fp32") + bf16 = ConvProblem("bf16", 64, 64, (8, 8, 8), dtype="bf16") + assert reference.error_bound(fp32, expected) < reference.error_bound(bf16, expected) + + +def test_error_bound_tracks_peak_not_just_rms(): + """A tensor with a big outlier gets a proportionally bigger absolute bound. + + This is the property whose absence made the bound too tight: the final + rounding to bf16 costs an ulp of the *largest* element, so a spiky tensor + legitimately admits more absolute error than a flat one of the same RMS. + """ + problem = ConvProblem("p", 64, 64, (8, 8, 8)) + flat = torch.ones(4096, dtype=torch.float64) + spiky = flat.clone() + spiky[0] = 100.0 + assert reference.error_bound(problem, spiky) > 10 * reference.error_bound(problem, flat) + + +def test_the_store_term_is_charged_as_one_rounding_not_four(): + """The safety factor belongs on the walk, not on the deterministic store. + + ``error_bound`` used to be ``8 * (accum + store)``, four ulps of the peak + for a rounding that is bounded by half an ulp of the *element* outright. + There is no random walk in a single store to take a factor against, and the + consequence was not academic: the static arm then won + :func:`reference.assert_close`'s ``max()`` in 46 of 48 measured cells, so + ``test_no_worse_than_miopen`` in all three kernel files was not holding the + kernel to the standard its name and docstring claim. + + Pinned arithmetically rather than by measurement so that it fails on the + formula rather than on a GPU: at ``K`` short enough that the accumulation + term is negligible, the bound must be one ulp of the peak per rounding. + """ + problem = ConvProblem("p", 8, 8, (4, 4, 4)) # K = 216 + peak = torch.zeros(4096, dtype=torch.float64) + peak[0] = 64.0 + ulp = 2.0 * reference.unit_roundoff(torch.bfloat16) * 64.0 + assert reference.error_bound(problem, peak) == pytest.approx(ulp, rel=1e-3) + assert reference.error_bound(problem, peak, roundings=2) == pytest.approx( + 2 * ulp, rel=1e-3) + + +@requires_gpu +def test_the_incumbents_extra_roundings_are_the_atomic_ones(): + """Why the incumbent gets ``roundings=2`` in exactly one direction. + + The store term models one deterministic rounding into the working dtype. + That is what the forward and backward-data do -- both are bitwise + reproducible here, and their error lands under one ulp of the peak. + MIOpen's backward-weight is not: it reduces with atomics, two identical + calls differ, and the extra roundings can carry it past one ulp. Without + this the allowance in + :func:`test_reference_agrees_with_miopen_within_tolerance` looks like a + tolerance that was widened until the test passed. + + ``conv 32->32 k3x3x3 @ 8x8x8`` because it is the cell that measures the + excess most clearly; the nondeterminism is a property of the direction, not + of the shape. What is asserted is the *call-to-call spread*, not the error + on any one call, and that distinction is the finding: the error itself + wanders (0.61 to 1.05 ulps of the peak over eight calls) precisely because + the reduction order does, so an assertion on a single call would be as + intermittent as the bound it is defending. A single rounding has a spread + of exactly zero, which is what the other two directions measure. + """ + problem = ConvProblem("atomic", 32, 32, (8, 8, 8)) + ops = reference.make_inputs(problem, seed=7) + ulp = 2.0 * reference.unit_roundoff(reference.torch_dtype(problem)) + + def probe(direction, repeats=6): + expected = reference.reference(problem, ops, direction) + scale = ulp * expected.abs().max().item() + runs = [reference.incumbent(problem, ops, direction) for _ in range(repeats)] + errs = [reference.compare(r, expected).max_abs / scale for r in runs] + spread = max((a - b).abs().max().item() for a in runs for b in runs) / scale + return max(errs), spread + + deterministic = 0.0 + for direction in ("fwd", "bwd-data"): + err, spread = probe(direction) + assert spread == 0.0, f"{direction}: MIOpen disagreed with itself by {spread}" + assert err < 1.0, f"{direction}: {err:.3f} ulps of the peak" + deterministic = max(deterministic, err) + + err, spread = probe("bwd-weight") + assert spread > 0.25, ( + f"MIOpen's backward-weight agreed with itself to {spread:.3f} ulps of " + "the peak; if it has stopped reducing with atomics then the roundings=2 " + "allowance it is given has lost its reason and should be dropped" + ) + assert err > deterministic, ( + f"backward-weight ({err:.3f} ulps) is no worse than the directions that " + f"round once ({deterministic:.3f}); the allowance is unmotivated" + ) + # And the allowance is an envelope, not a blank cheque: two roundings must + # still be enough. If this trips, the right response is to find out how + # many partials MIOpen is accumulating, not to raise the number. + assert err < 2.0, f"backward-weight needs more than two roundings: {err:.3f}" + + +@requires_gpu +def test_the_incumbent_clause_binds_more_often_than_the_static_bound(): + """The anti-vacuity guard on ``assert_close``'s ``max()``. + + A ``max()`` is only worth writing if both arms can win. Under the old + four-ulp store term the static arm won essentially always and the "no worse + than MIOpen by more than ``margin``" standard was dead code -- documented, + named in three test functions, and never applied. So pin the property that + made it live: over these cells the incumbent arm must be the operative one + more often than not. + + A floor rather than a per-cell assertion because which arm wins is a real + measurement and does move: it is the incumbent in 12 of these 13 cells, and + the one that goes the other way is a shape where MIOpen happens to be + unusually accurate -- exactly the case the ``max()`` exists to stop from + tightening the test beyond what the numerics justify. Not parametrized, + because a per-case fixture cannot state a floor over the set and this file's + whole reason for existing is that a test which reports a pass without + testing anything is worse than no test. + """ + binds = [] + for problem in SMALL: + ops = reference.make_inputs(problem, seed=7) + expected = reference.reference(problem, ops, "fwd") + err = reference.compare(reference.incumbent(problem, ops, "fwd"), expected) + binds.append(4.0 * err.max_abs + > reference.error_bound(problem, expected, "fwd")) + assert sum(binds) > len(binds) // 2, ( + f"the incumbent clause bound only {sum(binds)}/{len(binds)} cells; the " + "static bound has drifted back to swallowing it" + ) + + +@requires_gpu +def test_assert_close_rejects_a_wrong_answer(): + """The policy has to fail when it should; a tolerance nobody can trip is not one. + + A one-voxel shift is the realistic failure mode for a gather kernel, and it + is the one a loose elementwise tolerance would wave through. + """ + problem = ConvProblem("shift", 16, 16, (8, 8, 8)) + ops = reference.make_inputs(problem, seed=11) + expected = reference.reference(problem, ops, "fwd") + shifted = reference.incumbent(problem, ops, "fwd").roll(1, dims=-1) + with pytest.raises(AssertionError): + reference.assert_close(shifted, expected, problem, "fwd") + + +@requires_gpu +def test_channels_last_is_preserved_by_make_inputs(): + problem = ConvProblem("cl", 32, 32, (8, 8, 8)) + ops = reference.make_inputs(problem) + assert ops["input"].is_contiguous(memory_format=torch.channels_last_3d) + assert ops["grad_output"].is_contiguous(memory_format=torch.channels_last_3d) + + +# --------------------------------------------------------------------------- +# Timing harness +# --------------------------------------------------------------------------- + + +@requires_gpu +def test_interleaved_rotates_variants_and_reports_spread(): + """Every variant occupies every slot, so no one of them owns the fast one.""" + from triton_conv3d.bench.harness import interleaved + + a = torch.randn(512, 512, device="cuda") + seen: dict[str, list[int]] = {"x": [], "y": [], "z": []} + order: list[str] = [] + + def make(name): + def fn(): + order.append(name) + return a @ a + return fn + + result = interleaved({k: make(k) for k in seen}, warmup=1, iters=1, rounds=3) + assert set(result) == set(seen) + assert all(len(m.rounds) == 3 for m in result.values()) + # Rotation: the first variant of each round differs from round to round. + starts = {order[i] for i in range(0, len(order), 1) if i % 3 == 0} + assert len(starts) > 1, "rounds did not rotate" + # The old assertion here was ``m.spread >= 0``, which is true by + # construction of ``(max - min) / median`` and could not fail. What is + # worth pinning is that pinning ``warmup``/``iters``/``rounds`` still runs + # exactly the calls it says: 1 warmup and 3 rounds of 1 iteration each, per + # variant, with no calibration probe smuggled in. + assert len(order) == 3 * (1 + 3 * 1) + assert all(m.iters == 1 and m.group == 1 and m.stop == "fixed" + for m in result.values()) + + +def test_the_round_order_is_position_and_adjacency_balanced(): + """Rotating by one position per round de-biases slots but not neighbours. + + Under the old rule -- ``names[r % n:] + names[:r % n]`` -- variant B ran + immediately after variant A in *every* round, so whatever A left in the + caches was a constant charged to B and averaged out of nothing. Measured on + the adversarial case (a 1 GiB cache-polluting arm plus two arms doing + byte-identical work, 40 replications): cyclic rotation reported the two + identical arms **2.8% apart**, this rule 0.2% apart, a random order 0.7%. + 2.8% is larger than several of the per-cell differences this project + publishes, so the design property is worth asserting rather than trusting. + + Pure Python and exhaustive, so it fails on the *rule* rather than on a + measurement: over ``2 * n`` rounds every variant must occupy every position + equally often **and** every ordered adjacent pair must occur equally often. + Reverting :func:`_order` to the cyclic rotation fails the second clause at + every ``n >= 3`` (it makes the count of ``(A, B)`` equal to the number of + rounds and the count of ``(B, A)`` zero). + """ + from triton_conv3d.bench.harness import _order + + for n in range(1, 7): + names = [chr(ord("A") + i) for i in range(n)] + rounds = 2 * n + positions = {x: [0] * n for x in names} + adjacency: dict[tuple[str, str], int] = {} + for r in range(rounds): + got = _order(names, r) + assert sorted(got) == sorted(names), f"{n}: {got} is not a permutation" + for slot, x in enumerate(got): + positions[x][slot] += 1 + for pair in zip(got, got[1:]): + adjacency[pair] = adjacency.get(pair, 0) + 1 + for x in names: + assert len(set(positions[x])) == 1, ( + f"n={n}: {x} occupied positions unevenly: {positions[x]}" + ) + if n >= 2: + assert len(set(adjacency.values())) == 1, ( + f"n={n}: adjacency is not balanced: {adjacency}" + ) + assert len(adjacency) == n * (n - 1), ( + f"n={n}: only {len(adjacency)} of {n*(n-1)} ordered pairs occur" + ) + + +def test_spread_is_a_range_statistic_and_the_interval_is_not(): + """Why ``spread`` cannot support a claim about how much the machine moved. + + ``(max - min) / median`` is a *range*, and the expected range of ``n`` + samples grows like ``d2(n)`` even on a perfectly stationary device. + Measured on this node with one kernel held constant for 14 minutes + (47,686 blocks, ``work/triton-conv/bin/drift_probe.py``) the median of + this statistic runs 0.23% at 2 rounds, 0.63% at 6, 0.98% at 20 and 2.70% at + 100 -- all of it arithmetic, none of it the machine. Since ``rounds`` is + now chosen per cell, two cells' spreads are not comparable to each other at + all, and the replacement has to be an interval. + + Pinned on a fixed draw so it tests the formulae, not the GPU. + """ + import random + + from triton_conv3d.bench.harness import Measurement + + rng = random.Random(20260803) + def draw(n): + return tuple(1.0 + 0.01 * rng.gauss(0, 1) for _ in range(n)) + + short = Measurement("short", draw(4)) + long_ = Measurement("long", draw(1000)) + # Same underlying dispersion, by construction. + assert abs(long_.cov - 0.01) < 0.002 + # The range grows with n ... + assert long_.spread > 2.5 * short.spread + # ... while the interval, which is the thing to quote, shrinks. + assert long_.rel_half_width < 0.2 * short.rel_half_width + assert short.rel_half_width > 0.005 + + +@requires_gpu +def test_a_paired_ratio_of_two_identical_arms_covers_one(): + """The anti-vacuity guard on the interval: it must be right *and* narrow. + + Two arms that are the same callable have a true ratio of exactly 1, so an + interval that misses 1 is too narrow and one that spans a factor of two is + useless. Both failures are live: an interval computed on the *mean* of + per-iteration times rather than on the round medians is too narrow, and one + taken over two rounds is too wide. + + This also puts a number on what a published ratio has to beat. At a + 0.08 ms kernel two identical arms measured the old way -- 6 rounds of 10 -- + came out **0.941x to 1.058x** over 40 replications (sd 2.5%), so a "1.02x" + at that size was never a measurement. Nothing here asserts that; it is why + the interval exists. + """ + from triton_conv3d.bench.harness import interleaved, ratio + + a = torch.randn(2048, 2048, device="cuda", dtype=torch.bfloat16) + fn = lambda: a @ a # noqa: E731 + meas = interleaved({"x": fn, "y": fn}, target_rel=0.02, budget_s=15.0) + r = ratio(meas["y"], meas["x"]) + assert r.lo <= 1.0 <= r.hi, f"interval missed the truth: {r}" + assert not r.significant, f"identical arms declared different: {r}" + assert r.rel_half_width < 0.06, f"interval uselessly wide: {r}" + assert abs(r.point - 1.0) < 0.05, f"identical arms differ by {r.point:.4f}" + + +@requires_gpu +def test_the_block_is_sized_from_the_measured_duration(): + """``iters`` is chosen online, and it has to move with the kernel. + + The corpus spans five orders of magnitude -- 0.06 ms at the transposed sites + against 45,241 ms for one call at the 2 GiB cliff -- and a fixed + ``iters=10, rounds=6`` is 60 calls either way: microseconds for one cell and + 45 minutes for the other. + + ``torch.cuda._sleep`` rather than a real kernel: it consumes a stated number + of device cycles with no memory traffic and no tuning database, so the test + asserts the *sizing rule* and cannot fail because MIOpen picked a different + solver today. + """ + from triton_conv3d.bench.harness import time_callable + + fast = time_callable(lambda: torch.cuda._sleep(200_000), budget_s=5.0) + slow = time_callable(lambda: torch.cuda._sleep(60_000_000), budget_s=5.0) + assert slow.median > 20 * fast.median, "the two probes are not far apart" + assert slow.iters <= 2, f"a {slow.median:.1f} ms call got iters={slow.iters}" + assert fast.iters >= 10 * slow.iters, ( + f"iters did not track duration: {fast.iters} at {fast.median:.4f} ms " + f"vs {slow.iters} at {slow.median:.2f} ms" + ) + # And the block lands near its target rather than anywhere at all. + assert 0.1 <= fast.iters * fast.median / 15.0 <= 10.0 + + +@requires_gpu +def test_a_slow_kernel_stops_on_the_budget_and_says_so(): + """The ceiling, and the flag that makes a loose measurement visible. + + With an unreachable precision target the only way out is the wall clock, so + this pins both that the budget is honoured and that ``stop`` reports it. + Without the budget check the same call runs to ``max_rounds`` -- 64 rounds + of a ~0.5 s kernel, half a minute -- which is what the assertion on elapsed + time detects. + """ + import time + + from triton_conv3d.bench.harness import time_callable + + t0 = time.perf_counter() + m = time_callable(lambda: torch.cuda._sleep(1_000_000_000), + budget_s=1.0, target_rel=1e-9) + elapsed = time.perf_counter() - t0 + assert m.stop == "budget", f"stopped for the wrong reason: {m.stop}" + assert not m.converged + assert len(m.rounds) <= 8, f"{len(m.rounds)} rounds against a 1 s budget" + assert elapsed < 15.0, f"budget not honoured: {elapsed:.1f} s" + assert m.rel_half_width > 0, "a budget-stopped cell must still report a width" + + +@requires_gpu +def test_the_instrument_tax_is_measured_and_grouped_away(): + """The sub-0.15 ms regime, with its own negative control. + + An ``hipEventRecord`` costs ~9.5 us of host time, and at a 0.017 ms kernel a + block with an event between every iteration reports **1.5x** what the same + kernel's wall-clock throughput does. That is not noise, it is not the node, + and it does not cancel in a ratio because it is per-arm: measured on + ``convT 1024->512 @ 8^3``, the Triton forward pays 10.1 us and the MIOpen + weight-gradient control 11.9 us on times of 0.057 and 0.067 ms. + + The grouped block is checked against an event-free wall-clock measurement of + the same callable, and against the *ungrouped* harness in the same run. The + second is the control: if grouping ever stops working, the two agree and + this fails, rather than both drifting together unnoticed. + """ + import statistics + import time + + from triton_conv3d.bench.harness import time_callable + + a = torch.randn(512, 512, device="cuda", dtype=torch.bfloat16) + fn = lambda: a @ a # noqa: E731 + for _ in range(50): + fn() + torch.cuda.synchronize() + + def wall(n=4000): + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(n): + fn() + torch.cuda.synchronize() + return (time.perf_counter() - t0) * 1e3 / n + + # The *minimum* of several wall-clock runs, not the median. This reference + # has no events in it at all, which is the point, but it is therefore + # host-throughput-bound: it can only be inflated by contention, never + # deflated. Taking the median made this test fail once inside the full + # suite -- reference 0.0214 ms against 0.0167 ms in isolation, while the + # harness's own number moved by 9% -- which is the harness's stall rejection + # working and the reference's absence of it showing. + reference = min(wall() for _ in range(5)) + grouped = time_callable(fn, budget_s=10.0) + # ``tax_budget`` above 1.0 can never be exceeded, which disables grouping + # and reproduces the historical instrument exactly. + ungrouped = time_callable(fn, budget_s=10.0, tax_budget=10.0) + + assert grouped.group > 1, "a 0.02 ms kernel was left at one event per call" + assert ungrouped.group == 1 + assert 0.80 <= grouped.median / reference <= 1.25, ( + f"grouped block disagrees with the event-free reference: " + f"{grouped.median:.5f} vs {reference:.5f} ms" + ) + assert ungrouped.median > 1.2 * grouped.median, ( + f"the instrument tax has vanished on its own ({ungrouped.median:.5f} " + f"vs {grouped.median:.5f} ms); if that is real this test's premise is " + "gone and the grouping can be removed, but check the ruler first" + ) + + +@requires_gpu +def test_flush_caches_reuses_one_buffer_and_reaches_only_the_first_sample(): + """Two defects in one small function, both of which had teeth. + + ``torch.device("cuda")`` carries no index and a tensor made on it does, so + the guard ``_flush_buffer.device != torch.device(device)`` was *always* + true: every flush allocated a fresh 512 MiB tensor while the old one was + still live, on the critical path of every timed round. + + And a flush before a block reaches only the block's **first** call, while + the block reports the median over ``iters`` of them -- so at ``iters=10`` + the one cold sample is precisely the one the median throws away. Measured + ``median_moved_by_flush`` is 0.99-1.01 at every real workload while the + first iteration moves 1.02-1.54x. The adaptive path therefore measures + ``iters=1`` when ``flush`` is on, and :attr:`Measurement.cold` records the + first sample either way. + """ + from triton_conv3d.bench import harness as H + + H.flush_caches() + first = H._flush_buffer.data_ptr() + for _ in range(5): + H.flush_caches() + assert H._flush_buffer.data_ptr() == first, ( + "flush_caches reallocated its buffer; the device comparison is wrong " + "again" + ) + assert H._flush_buffer.numel() == H._FLUSH_BYTES + + a = torch.randn(1024, 1024, device="cuda", dtype=torch.bfloat16) + cold = H.time_callable(lambda: a @ a, flush=True, rounds=4) + assert cold.iters == 1, ( + f"flush=True measured {cold.iters} calls per block, so {cold.iters-1} " + "of them are hot and the median reports a hot number" + ) + assert cold.cold == cold.median # with one sample per block they coincide + + +@requires_gpu +def test_pinning_iters_and_rounds_reproduces_the_fixed_protocol(): + """Backward compatibility, asserted on the call count rather than assumed. + + Every existing driver and every stored result JSON was produced by pinned + ``warmup``/``iters``/``rounds``. Those callers must keep issuing exactly + the calls they always did -- no calibration probe, no warmup of its own, no + grouping -- or a re-capture is not comparable with what is on disk. + """ + from triton_conv3d.bench.harness import time_callable + + a = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16) + calls = {"n": 0} + + def fn(): + calls["n"] += 1 + return a @ a + + m = time_callable(fn, warmup=3, iters=10, rounds=6) + assert calls["n"] == 3 + 6 * 10, f"issued {calls['n']} calls, expected 63" + assert len(m.rounds) == 6 and m.iters == 10 and m.group == 1 + assert m.stop == "fixed" and m.converged + + +# --------------------------------------------------------------------------- +# The baseline is the control for every later claim, so it gets its own guards +# --------------------------------------------------------------------------- + + +def test_importing_the_baseline_turns_on_the_miopen_find(): + """A baseline taken with ``cudnn.benchmark`` off is not a baseline. + + On ROCm that flag decides whether PyTorch asks MIOpen to *search* for a + tuning config or to answer from its AI heuristic. The heuristic's answer + for the corpus' hottest problems is 5-12x slower than the searched one -- + same solver, same device op, just 16x16 MFMA tiles with 2-element global + loads instead of 32x32 with 8-element loads. The first version of this + harness left the flag at its default and overstated MIOpen by up to 12x, + which would have become a fabricated speedup for every kernel measured + against it. ScaFFold itself sets it (``worker.py:171``) and so does the + profiler the reference numbers come from (``prof_bench.py:125``). + """ + from triton_conv3d.bench import baseline + + assert baseline.REQUIRE_CUDNN_BENCHMARK is True + assert torch.backends.cudnn.benchmark is True, ( + "importing the baseline module must leave the process in the " + "configuration its recorded numbers were taken in" + ) + + +def test_measure_one_refuses_to_report_a_heuristic_time(): + """The guard has to be at the measurement, not only at import. + + Anything may flip ``cudnn.benchmark`` between import and use -- a + determinism experiment, another test, a notebook cell. Refusing loudly is + the only outcome that cannot end up in a JSON file that looks like a + control. + """ + from triton_conv3d.bench.baseline import measure_one + + problem = ConvProblem("guard", 8, 8, (4, 4, 4)) + previous = torch.backends.cudnn.benchmark + try: + torch.backends.cudnn.benchmark = False + with pytest.raises(RuntimeError, match="cudnn.benchmark is off"): + measure_one(problem, "fwd") + finally: + torch.backends.cudnn.benchmark = previous + + +#: Corpus cells the harness is anchored to. Chosen because their isolated and +#: profiled shapes genuinely match: the halo'd input is 281 MiB, well under the +#: 2 GiB threshold above which MIOpen abandons its tuned solvers for the naive +#: non-packed ones and the isolated and profiled numbers legitimately diverge +#: (``work/nonpacked-conv/RESULTS.md``). The profiled time is read from the +#: corpus rather than copied here so there is one source of truth for it. +ANCHOR_CELLS = ((6, "fwd"), (6, "bwd-data")) + + +@requires_gpu +@pytest.mark.slow +@pytest.mark.timeout(900) +@pytest.mark.parametrize("index, direction", ANCHOR_CELLS) +def test_baseline_reproduces_the_profiled_scaffold_conv(index, direction): + """An end-to-end anchor: the harness lands on the profiled number. + + The two tests above check the settings; this one checks the thing the + settings are for, and would still fail if the harness went wrong in a way + nobody anticipated -- a memory-format regression, a dtype regression, a + future PyTorch that stops honouring ``benchmark`` on ROCm. + + The band is asymmetric and generous. Isolated *should* come out a little + faster than profiled -- no contention for bandwidth, no other kernel in + flight, no allocator pressure -- but never much faster, and a run that is + slower than the profile has lost the find. + + Judged on the *best* round, not the median: this is a shared node and a + neighbouring job can inflate all five rounds at once (observed once while + writing this, at 2.1x). That is a fact about the node, not about the + harness, and a test that fails on it teaches people to ignore it. The + failure this test is for -- a lost find, the wrong shape, an inert memory + format -- is 5-12x and survives taking the minimum easily. + """ + from triton_conv3d.bench.baseline import measure_one + + logical = scaffold_corpus()[index] + problem = logical.halo_variant + profiled_ms = logical.measured_for(direction)[-1]["ms_per_call"] + record = measure_one(problem, direction) + assert "error" not in record, record.get("error") + ratio = record["best_ms"] / profiled_ms + assert 0.5 <= ratio <= 1.4, ( + f"{problem.label} {direction}: {record['best_ms']:.3f} ms isolated " + f"(best of {record['rounds']}) vs {profiled_ms:.3f} ms profiled " + f"({ratio:.2f}x). Off by this much means MIOpen is not solving the " + f"problem ScaFFold solves -- check cudnn.benchmark, " + f"PYTORCH_MIOPEN_SUGGEST_NHWC and the halo shape." + ) + + +@requires_gpu +def test_sporadic_host_stall_is_rejected_from_the_median_and_flagged(): + """An occasional slow launch must not be charged to the kernel. + + This is the harness bug that produced last session's 250-2363% spreads, + which were then misdiagnosed twice -- first as host jitter, then as a rogue + tenant on the GPU -- before turning out to be a duplicate driver process of + our own. The old ``_time_block`` bracketed a whole block of iterations with + two events, so any launch gap inside it was silently added to kernel time. + + One stalled launch in ten is the realistic shape of the problem: contention + is intermittent, so a mean absorbs it and a median rejects it. The stall + ratio exists so that rejecting it is not the same as hiding it. + """ + import time + + from triton_conv3d.bench.harness import interleaved + + a = torch.randn(1024, 1024, device="cuda", dtype=torch.bfloat16) + + # The control needs a quiet host and this test cannot assume one: the node + # is shared, and a neighbouring job stalls our launches exactly as well as + # the duplicate driver of our own did. When that happens the diagnostic is + # firing *correctly* and it is the premise -- "this run is clean" -- that is + # false. Observed at 12.87 against a threshold of 2.0 while a sibling job + # was running, passing three times in a row on the same tree once the node + # went idle. So take the quietest of several attempts, and if none of them + # is quiet, say the host was loaded rather than assert something this run + # cannot decide -- failing here would re-enact the original misdiagnosis in + # test form, blaming the measurement for observing real contention. + clean = min( + (interleaved({"g": lambda: a @ a}, warmup=3, iters=20, rounds=3)["g"] + for _ in range(5)), + key=lambda m: m.stall_ratio, + ) + if clean.stall_ratio >= 2.0: + pytest.skip( + f"host too loaded for a quiet control (best stall ratio " + f"{clean.stall_ratio:.2f} over 5 attempts); the diagnostic is " + "reporting real contention, so this test cannot separate a false " + "positive from a true one" + ) + + calls = {"n": 0} + + def sporadic(): + calls["n"] += 1 + if calls["n"] % 10 == 0: + time.sleep(0.02) + return a @ a + + stalled = interleaved({"g": sporadic}, warmup=3, iters=20, rounds=3)["g"] + + # The reported time is still the kernel's, not the kernel plus the gap. + assert stalled.median < 2.0 * clean.median, ( + f"stall leaked into the median: {stalled.median:.4f} vs {clean.median:.4f}" + ) + # And the gap is visible rather than absorbed. + assert stalled.stall_ratio > 3.0, f"stall not flagged: {stalled.stall_ratio:.2f}" + # The converse -- that a quiet run is *not* flagged -- is established by the + # skip above rather than here, because on a loaded node it is not true and + # should not be asserted. + assert clean.stall_ratio < 2.0, f"clean run falsely flagged: {clean.stall_ratio:.2f}" + + +# --------------------------------------------------------------------------- +# What is inside the timed region +# --------------------------------------------------------------------------- +# +# The published per-shape number is *kernel* time: the Python-side dispatch, +# the tuned-table lookup and the launcher in front of the kernel are outside it. +# That is a decision about what to measure, and it has exactly one way to go +# wrong -- taking the launcher out of one arm and not the other, which at these +# sizes is worth up to 1.4x in the direction that flatters us. These tests are +# the guard on that, and each of them is verified by breaking what it tests +# (``work/triton-conv/bin/mutate_bench.py``). + + +def test_the_graph_chunk_is_one_ruler_for_every_arm(): + """``chunk`` is a function of the shortest arm's duration, and nothing else. + + A CUDA graph replay costs 3.9-12.8 us of device time whatever is inside it + -- measured by fitting ``per_call(chunk) = kernel + cost / chunk`` to graphs + of 1, 2, 4, 8, 16 and 32 calls on four real arms. At ``chunk = 1`` that is + 45% of a 0.028 ms kernel and only 19% of a 0.068 ms one, so a per-arm chunk + would be a per-arm instrument: exactly the failure ``_common_group`` already + documents, one level up (H10, where two byte-identical arms picked different + event groups and read 4% apart). + + Hence: the rule reads only ``min(durations)``, so two arms of the same call + can never be given different rulers. + """ + from triton_conv3d.bench.harness import _REPLAY_COST_MS, common_chunk + + # Only the minimum matters: a slow second arm cannot loosen the ruler. + assert common_chunk([0.03, 0.03]) == common_chunk([0.03, 3.0]) + assert common_chunk([0.03, 3.0]) == common_chunk([3.0, 0.03]) + # Monotone: a shorter kernel needs a wider graph. + chunks = [common_chunk([d]) for d in (0.01, 0.03, 0.1, 0.3, 1.0, 10.0)] + assert chunks == sorted(chunks, reverse=True), chunks + # And the residual really is inside the budget it claims. + for d in (0.01, 0.02, 0.05, 0.1, 0.5): + c = common_chunk([d]) + residual = _REPLAY_COST_MS / c / d + assert residual <= 0.011 or c == 128, ( + f"at {d} ms the chunk {c} leaves {residual:.1%} of replay cost in" + ) + # A kernel long enough not to care is left alone. + assert common_chunk([5.0]) == 1 + + +def test_no_graph_where_the_launcher_is_already_negligible(): + """Above 40 ms per call the exclusion is not worth the capture. + + The largest host launch cost measured on this node is 0.08 ms -- the + autograd engine's, on the MIOpen backward control. At 40 ms per call that + is 0.2% of either arm, a fifth of the harness's own 2% target, so both arms + stay eager and the exclusion is negligible *for both* rather than applied to + one. Below it the same 0.08 ms reaches 190% of the kernel and decides the + answer. + """ + from triton_conv3d.bench.harness import graph_is_worthwhile + + assert graph_is_worthwhile([0.03]) + assert graph_is_worthwhile([0.03, 5000.0]), "the shortest arm decides" + assert not graph_is_worthwhile([100.0]) + assert not graph_is_worthwhile([45241.0]), "the 2 GiB cliff cell" + + +@requires_gpu +def test_an_empty_graph_is_a_capture_failure(): + """PyTorch only *warns* when a capture caught nothing. + + "The CUDA Graph is empty. This usually means that the graph was attempted + to be captured on wrong device or stream." It is a ``UserWarning``, and a + caller that ignored it would publish the cost of ``cudaGraphLaunch`` -- a + few microseconds -- as a kernel time. That is the fastest wrong answer + available and it looks like a spectacular win, so the warning is promoted to + a refusal. + + It is not hypothetical: the first version of this work built the MIOpen + backward control's forward graph on the default stream, captured on another, + and got an empty graph plus this warning on two of six cells. + """ + from triton_conv3d.bench.harness import CaptureError, capture + + with pytest.raises(CaptureError, match="empty"): + capture(lambda: None, 1) + + +@requires_gpu +def test_a_captured_ratio_of_two_identical_arms_covers_one(): + """The null experiment for the launcher-exclusion boundary. + + Two arms doing byte-identical work have a true ratio of exactly 1.000, so + anything else is the instrument. Measured over 12 replications on + ``convT 1024->512 @ 8^3`` through the shipped decision path + (``work/triton-conv/bin/launcher_symmetry.py --only null``): under + ``exclude`` the median is 0.9996, the range 0.9982-1.0021, and **12 of 12** + intervals cover 1.000. + + The sibling test for the *event* instrument is + :func:`test_a_paired_ratio_of_two_identical_arms_covers_one`; this one is + for the graph. + """ + from triton_conv3d.bench.conv_bench import _timed_region + from triton_conv3d.bench.harness import interleaved, ratio + + # 512, not a "nicer" 256 or 384: on this torch/ROCm build a bf16 + # ``a @ a`` is **~600 ms** at 128, 192, 256, 320, 384, 448, 640 and + # 768, and 0.019 ms at 512 and 1024. That is the ``torch.mm`` bf16 + # pathology ``STATUS.md`` §7 already owes upstream, measured here from + # a second direction; a test that picked one of the slow sizes would + # be timing a 600 ms kernel and would correctly be told it does not + # need a graph. + a = torch.randn(512, 512, device="cuda", dtype=torch.bfloat16) + fn = lambda: a @ a # noqa: E731 + region = _timed_region({"x": fn, "y": fn}, "exclude") + assert region.kind == "kernel", region.note + out = interleaved(region.fns, budget_s=10.0) + r = ratio(out["y"], out["x"]) + # A tolerance, not the interval, and deliberately. At 0.019 ms with a + # 64-call graph one race converges to a *within-race* half-width of ~0.04%, + # which is narrower than the between-race scatter of the same pair (sd + # 0.13%, range 0.9982-1.0021 over 12 replications) -- so an interval that + # misses 1 by 0.2% here is the same residual the sequential protocol has + # (0.32%, HARNESS_RIGOR Part 1), not a biased instrument. What a biased + # instrument looks like is 4% (H10, per-arm event groups) or 45% (a one-call + # graph), and 1% catches both. The coverage claim is the 12-replication + # experiment, where 12 of 12 intervals contained 1. + assert abs(r.point - 1.0) < 0.01, ( + f"two byte-identical arms read {r} under the kernel-time definition" + ) + assert r.rel_half_width < 0.05, f"interval uselessly wide: {r}" + + +@requires_gpu +def test_an_inflated_launcher_does_not_move_the_reported_kernel_time(): + """The whole point of the exclusion, stated as a property. + + Three arms run the **same kernel** with deliberately different launchers. + Under ``exclude`` they must be indistinguishable; under ``include`` they + must not be, or the experiment proves nothing and the exclusion is + measuring something that was not there. + + That negative control is deliberate. Measured on the real thing + (``launcher_symmetry.py --only inflate``, ``convT 1024->512 @ 8^3``): the + entry point's own per-call table lookup reads 1.0003x of the hoisted config + under ``exclude`` and **1.404x** under ``include``, and 500 us of Python in + front of the launch reads 0.9993x and **13.12x**. + """ + import time + + from triton_conv3d.bench.conv_bench import _timed_region + from triton_conv3d.bench.harness import interleaved, ratio + + # 512, not a "nicer" 256 or 384: on this torch/ROCm build a bf16 + # ``a @ a`` is **~600 ms** at 128, 192, 256, 320, 384, 448, 640 and + # 768, and 0.019 ms at 512 and 1024. That is the ``torch.mm`` bf16 + # pathology ``STATUS.md`` §7 already owes upstream, measured here from + # a second direction; a test that picked one of the slow sizes would + # be timing a 600 ms kernel and would correctly be told it does not + # need a graph. + a = torch.randn(512, 512, device="cuda", dtype=torch.bfloat16) + + def plain(): + return a @ a + + def slow_launcher(): + end = time.perf_counter() + 300e-6 + while time.perf_counter() < end: + pass + return a @ a + + variants = {"plain": plain, "slow": slow_launcher} + kern = _timed_region(dict(variants), "exclude") + assert kern.kind == "kernel", kern.note + k = interleaved(kern.fns, budget_s=10.0) + rk = ratio(k["slow"], k["plain"]) + # 1%, for the reason given in + # ``test_a_captured_ratio_of_two_identical_arms_covers_one``. 300 us in + # front of a 0.019 ms kernel is a 16x effect if it is inside the timed + # region, so 1% is not a generous threshold here. + assert abs(rk.point - 1.0) < 0.01, ( + f"300 us of host work moved the kernel time: {rk}" + ) + + eager = _timed_region(dict(variants), "include") + assert eager.kind == "call" + e = interleaved(eager.fns, budget_s=10.0) + re_ = ratio(e["slow"], e["plain"]) + assert re_.point > 2.0, ( + f"the negative control did not fire: the same host work read {re_} " + "under the launcher-inclusive definition, so this test would pass " + "against a version that excludes nothing" + ) + + +@requires_gpu +def test_the_replay_cost_is_amortized_by_the_chunk(): + """With its own negative control, like the event-tax test. + + One graph replay costs up to 12.8 us of device time whatever is in it, so a + one-call graph is 45% instrument at a 0.028 ms kernel. The chunk divides + that away. The control is the *same* kernel measured at ``chunk = 1`` in + the same run: if the replay ever becomes free, the two agree and this fails + rather than both drifting together unnoticed. + """ + from triton_conv3d.bench.harness import capture, common_chunk, interleaved + + # 512, not a "nicer" 256 or 384: on this torch/ROCm build a bf16 + # ``a @ a`` is **~600 ms** at 128, 192, 256, 320, 384, 448, 640 and + # 768, and 0.019 ms at 512 and 1024. That is the ``torch.mm`` bf16 + # pathology ``STATUS.md`` §7 already owes upstream, measured here from + # a second direction; a test that picked one of the slow sizes would + # be timing a 600 ms kernel and would correctly be told it does not + # need a graph. + a = torch.randn(512, 512, device="cuda", dtype=torch.bfloat16) + fn = lambda: a @ a # noqa: E731 + for _ in range(50): + fn() + torch.cuda.synchronize() + + one = capture(fn, 1) + chunk = common_chunk([0.02]) + assert chunk >= 8, chunk + many = capture(fn, chunk) + out = interleaved({"one": one, "many": many}, budget_s=10.0) + per_one = out["one"].median + per_many = out["many"].median / chunk + assert per_many < per_one, ( + f"a {chunk}-call graph is not cheaper per call ({per_many:.5f}) than a " + f"one-call graph ({per_one:.5f}); the replay cost has vanished and this " + "test's premise with it -- check the ruler before deleting the chunk" + ) + assert per_one - per_many < 0.05, "implausible replay cost; something else moved" + + +@requires_gpu +def test_a_capture_failure_takes_the_whole_cell_back_to_eager(): + """Never a mixed measurement. + + If one arm cannot be captured, the other must not be either: comparing a + launcher-exclusive number against a launcher-inclusive one is worth 1.4x at + the transposed sites and 3.0x on the backward controls. So the fallback is + a property of the *cell*, and ``_Region`` is one object for all of its arms. + """ + from triton_conv3d.bench.conv_bench import _timed_region + from triton_conv3d.bench.harness import Captured + + # 512, not a "nicer" 256 or 384: on this torch/ROCm build a bf16 + # ``a @ a`` is **~600 ms** at 128, 192, 256, 320, 384, 448, 640 and + # 768, and 0.019 ms at 512 and 1024. That is the ``torch.mm`` bf16 + # pathology ``STATUS.md`` §7 already owes upstream, measured here from + # a second direction; a test that picked one of the slow sizes would + # be timing a 600 ms kernel and would correctly be told it does not + # need a graph. + a = torch.randn(512, 512, device="cuda", dtype=torch.bfloat16) + + def fine(): + return a @ a + + def uncapturable(): + # A device-to-host read inside a capture is illegal, and this is how a + # real arm fails: any control that peeks at a value does it. + return (a @ a).sum().item() + + ok = _timed_region({"x": fine, "y": fine}, "exclude") + assert ok.kind == "kernel" + assert all(isinstance(f, Captured) for f in ok.fns.values()) + + mixed = _timed_region({"x": fine, "y": uncapturable}, "exclude") + assert mixed.kind == "call", "a cell with an uncapturable arm was captured" + assert not any(isinstance(f, Captured) for f in mixed.fns.values()), ( + "one arm kept its graph while the other fell back -- that is the " + "asymmetry the whole exclusion exists to avoid" + ) + assert "could not be captured" in mixed.note + + +# --------------------------------------------------------------------------- +# The operator x direction table +# --------------------------------------------------------------------------- + + +def test_the_operator_direction_table_is_complete_and_has_six_distinct_cells(): + """Two operators, three directions, six builders, no sharing. + + The driver these replaced argued that a transposed convolution could not be + a fourth value of ``--direction`` because it is a different *operator*. It + was right, and the answer is a second axis rather than a fourth case: what + is per-operator (the shape form, the ordering) lives on ``_Op``, what is + per-cell (operands, control, candidates, shipped config, reference) lives in + one function per cell, and nothing is shared by accident. + """ + from triton_conv3d.bench.conv_bench import _OPERATORS, OPERATORS + + assert set(_OPERATORS) == set(OPERATORS) == {"conv", "convT"} + builders = [] + for op in _OPERATORS.values(): + assert set(op.build) == set(DIRECTIONS), op.name + builders += list(op.build.values()) + assert len(set(builders)) == 6, "two cells share a builder" + + +def test_no_builder_asks_a_problem_which_operator_it_is(): + """The design property, asserted rather than trusted. + + The objection to folding the transposed driver in was that ``_build`` would + "branch on ``problem.transposed`` in every arm to run the same code". It + does not: the operator is resolved **once**, in ``operator_of``, and the + builder it selects never asks again. If a future edit puts the question + back inside a builder, this fails. + """ + import inspect + + from triton_conv3d.bench.conv_bench import _OPERATORS, operator_of + + assert ".transposed" in inspect.getsource(operator_of) + for op in _OPERATORS.values(): + for direction, builder in op.build.items(): + src = inspect.getsource(builder) + assert ".transposed" not in src, ( + f"{op.name}/{direction} branches on the operator inside the " + "builder; that is the switch this factoring exists to remove" + ) + + +def test_a_backward_control_is_never_a_fabricated_operand(): + """``torch.nn.grad.conv3d_*`` must appear nowhere in this driver. + + It has no real tensor for the operand being differentiated, so it fabricates + ``grad_output.new_empty(1).expand(input_size)`` -- zero-strided, and + therefore not channels-last. ``convolution_backward`` picks its solver from + that operand's layout, so at the ``k=1x1x1`` head MIOpen declined its own + NDHWC path and ran **3.2x** slower than the same call inside a real + backward (0.9649 vs 0.2972 ms). A published 4.51x for that head came from + the fabricated control; against the real one the cell is 1.39x. + + Both drivers that existed before this one used it somewhere, which is why + the rule is a grep and not a convention. + """ + import inspect + + from triton_conv3d.bench import conv_bench + + for op in conv_bench._OPERATORS.values(): + for direction, builder in op.build.items(): + if direction == "fwd": + continue + bsrc = inspect.getsource(builder) + assert "torch.nn.grad" not in bsrc, ( + f"{op.name}/{direction} uses a fabricated operand for its " + "MIOpen control" + ) + assert "torch.autograd.grad" in bsrc, ( + f"{op.name}/{direction} has no real autograd control" + ) + + +def test_the_transposed_problems_are_never_haloed_and_the_others_follow_the_form(): + """The one shape decision that is per-operator, and it is silent when wrong. + + Upstream DistConv concatenates a ``k // 2`` halo onto every axis it manages + and zeroes that axis's padding, so an ordinary convolution reaches MIOpen at + ``130^3`` unpadded rather than ``128^3`` padded -- two problems MIOpen tunes + independently. At ``k = 2`` the halo is ``2 // 2 = 1``... which is why the + *corpus* is the authority and not the arithmetic: every transposed problem + in it records ``halo = (0, 0, 0)``, because ScaFFold's transposed sites are + not sharded convolutions at all. Applying ``halo_variant`` to them anyway + would silently grow the input by two voxels per axis and measure a different + problem -- **under any of the three ``--form`` names**, which is what this + test pins now that there is more than one. + """ + from triton_conv3d.bench.conv_bench import _FORMS, _OPERATORS + + conv, convt = _OPERATORS["conv"], _OPERATORS["convT"] + assert set(_FORMS) == {"distconv", "adapter", "logical"} + # The *function*, not just its effect on today's corpus: every transposed + # problem happens to record a zero halo, so a form that called + # ``halo_variant`` would be indistinguishable from the right one until the + # day one of them did not. Pin the rule instead. + haloed = ConvProblem("would_halo", 32, 16, (4, 4, 4), (2, 2, 2), (2, 2, 2), + (0, 0, 0), transposed=True, halo=(1, 1, 1), + shard_halo=(1, 1, 1)) + for name in _FORMS: + assert convt.form(haloed, name) is haloed, ( + f"--form {name} gave a transposed problem a halo; at k=2 there is " + "none, and adding one measures a convolution the model never runs" + ) + plain = dataclasses.replace(haloed, transposed=False) + assert conv.form(plain, "distconv").spatial == (6, 6, 6) + assert conv.form(plain, "adapter").spatial == (6, 6, 6) + assert conv.form(plain, "logical").spatial == (4, 4, 4) + + seen = {"conv": 0, "convT": 0} + for p in scaffold_corpus(): + if p.transposed: + assert convt.selects(p) and not conv.selects(p) + for name in _FORMS: + assert convt.form(p, name) is p, f"{p.label} was haloed" + assert p.halo == (0, 0, 0) + seen["convT"] += 1 + else: + assert conv.selects(p) and not convt.selects(p) + assert conv.form(p, "distconv") == p.halo_variant + assert conv.form(p, "adapter") == p.production_variant + assert conv.form(p, "logical") is p + seen["conv"] += 1 + assert seen["convT"] == 12 and seen["conv"] > 0, seen + + +@requires_gpu +def test_the_shipped_config_is_the_one_the_entry_point_resolves(): + """``--shipped`` must measure the shipped *kernel*, not a lookalike. + + The launcher-exclusive definition means the config cannot be resolved inside + the timed region, so the driver resolves it outside and passes it in. That + is only honest if the two agree, and nothing except this test makes them: + the six cells reach four different resolvers across three modules, with the + channel widths swapped on three of them. + + Checked by spying on the resolver each entry point actually calls, rather + than by re-deriving the answer here -- which would be the same arithmetic + twice and would agree with itself while both were wrong. + """ + from triton_conv3d.bench.conv_bench import _OPERATORS, _build + from triton_conv3d import bwd_data, gather_gemm, reduce_gemm, transposed + + problems = { + "conv": ConvProblem("t", 32, 16, (8, 8, 8)), + "convT": ConvProblem("tt", 32, 16, (4, 4, 4), (2, 2, 2), (2, 2, 2), + (0, 0, 0), transposed=True), + } + spied = [] + + def spy(mod, name): + real = getattr(mod, name) + def wrapper(*a, **kw): + cfg = real(*a, **kw) + spied.append(cfg) + return cfg + return real, wrapper + + patched = [(gather_gemm, "select_config"), (bwd_data, "select_config"), + (reduce_gemm, "bwd_weight_config"), + (transposed, "transposed_config")] + originals = {} + for mod, name in patched: + real, wrapper = spy(mod, name) + originals[(mod, name)] = real + setattr(mod, name, wrapper) + try: + for opname, op in _OPERATORS.items(): + for direction in DIRECTIONS: + case = _build(problems[opname], direction, operator=opname) + declared = case.shipped_config() + spied.clear() + case.triton(None)() + torch.cuda.synchronize() + assert spied, f"{opname}/{direction}: no resolver was called" + assert declared in spied, ( + f"{opname}/{direction}: --shipped would time {declared}, " + f"but the entry point resolves {spied}" + ) + del case + torch.cuda.empty_cache() + finally: + for (mod, name), real in originals.items(): + setattr(mod, name, real) + + +@requires_gpu +def test_the_published_time_is_per_call_and_never_exceeds_the_eager_call(): + """``chunk`` calls sit behind one replay; the row must report one call. + + The division happens in the driver rather than in the harness, because every + *relative* quantity the harness computes -- the half-widths, the convergence + test, the paired ratio -- is scale-invariant and only the absolute times need + it. That is easy to forget, and forgetting it multiplies every published + time by up to 128 while leaving every interval and every speedup looking + perfectly healthy. + + The invariant that catches it: kernel time is the eager call *minus* its + launcher, so it can never exceed the eager call. + """ + from triton_conv3d.bench.conv_bench import measure_problem + + p = ConvProblem("tt", 64, 32, (4, 4, 4), (2, 2, 2), (2, 2, 2), (0, 0, 0), + transposed=True) + row = measure_problem(p, direction="fwd", shipped=True, budget_s=5.0) + assert "error" not in row, row.get("error") + assert row["timed_region"] == "kernel", row["timed_region_note"] + assert row["graph_chunk"] > 1 + for arm in ("triton", "miopen"): + kernel, eager = row[f"{arm}_ms"], row[f"{arm}_eager_ms"] + assert 0.0 < kernel <= 1.05 * eager, ( + f"{arm}: reported kernel time {kernel:.5f} ms exceeds the eager " + f"call it is part of ({eager:.5f} ms) -- the chunk divisor is " + "missing or wrong" + ) + assert row[f"{arm}_launcher_ms"] > 0 + + +@requires_gpu +def test_a_control_free_row_omits_the_control_rather_than_zeroing_it(): + """``--control none`` must leave MIOpen *absent*, not present and zero. + + Two failures this pins, and they are opposite ones. + + A row that carried ``miopen_ms = 0.0`` and ``speedup = 0.0`` would be read + by every consumer of these captures -- the report generator, the aggregate + scripts, a human scanning a table -- as a measured 0.000x result rather than + as "no control ran here". Absence has to be representable. + + And a case built with ``control=False`` must not construct the control + either. For a backward direction the control is a real ``F.conv3d`` forward + graph, and *running* it once is where MIOpen's find is paid -- 92-174 s per + cell on this corpus (``work/triton-conv/review/HARNESS_SPEED.md`` 1). + Dropping the arm from the timing while still building it would save the + timing and none of the cost, which is the whole point of the flag. + """ + from triton_conv3d.bench.conv_bench import _build, measure_problem + + p = ConvProblem("tt", 64, 32, (4, 4, 4), (2, 2, 2), (2, 2, 2), (0, 0, 0), + transposed=True) + + for direction in DIRECTIONS: + case = _build(p, direction, control=False) + assert case.miopen is None, f"{direction}: a control was built anyway" + assert case.reference is None, f"{direction}: a reference was built" + del case + torch.cuda.empty_cache() + + row = measure_problem(p, direction="bwd-weight", shipped=True, + budget_s=5.0, control="none") + assert "error" not in row, row.get("error") + assert row["control"] == "none" + # The Triton half is unchanged: same region, same interval, same stop rule. + assert row["timed_region"] == "kernel", row["timed_region_note"] + assert row["triton_ms"] > 0.0 + # Present and finite, not strictly positive: on a kernel this small every + # round can read the same value to the last bit, and ``stdev`` of identical + # samples is exactly 0. A zero half-width there is the honest answer, not a + # missing one -- what would be wrong is the key being absent or ``inf``. + assert math.isfinite(row["triton_rel_ci"]) and row["triton_rel_ci"] >= 0.0 + assert row["measure_stop"] in ("converged", "budget", "max_rounds") + # The MIOpen half is gone, not zeroed. + for key in ("miopen_ms", "miopen_rel_ci", "miopen_eager_ms", "speedup", + "speedup_lo", "speedup_hi", "speedup_significant"): + assert key not in row, ( + f"{key} is present in a --control none row; an absent measurement " + "must stay absent, because a zero here reads as a result" + ) diff --git a/triton_conv3d/tests/test_transposed.py b/triton_conv3d/tests/test_transposed.py new file mode 100644 index 0000000..2b99d2c --- /dev/null +++ b/triton_conv3d/tests/test_transposed.py @@ -0,0 +1,947 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""Tests for the ``kernel == stride`` transposed convolution. + +Three directions, one new kernel, so these tests split unevenly on purpose. + +**The forward** is a new ``@triton.jit`` function with a *scatter* store, which +is the one addressing pattern nothing else in this package has. Its failure +mode is a permutation: write tap ``(kd,kh,kw)`` into the wrong sub-lattice and +the result is the right shape, the right magnitude, smooth, and wrong -- a +tolerance test cannot see it, and neither can a test that only checks sums. So +the bar here is bitwise, and two tests exist purely to prove that bar is not +vacuous (:func:`test_a_transposed_tap_permutation_is_detected` and +:func:`test_bitwise_standard_rejects_a_shifted_scatter`), because this project +has shipped a vacuous exact test before. + +**Both backward directions** are re-expressions: backward-data is +``conv3d_forward`` at ``stride = k`` and backward-weight is +``conv3d_backward_weight`` with the two activations swapped. There is no new +arithmetic in either, so what is tested is the *re-expression* -- above all the +swap, which is the single most plausible mistake in the file and which produces +a correctly shaped gradient when it is wrong (:func:`test_backward_weight_ +operand_swap_is_not_reversible`). + +**The FLOP count** is checked in its own right. ``k == s`` makes the per-tap +factor illusory (the windows tile rather than overlap) and this project once +counted it anyway, 8x too high. ``shapes.py`` has it right; here it is checked +against the elementary MAC count of the reference implementation rather than +against another formula. +""" + +from __future__ import annotations + +import dataclasses +import math + +import pytest +import torch +import torch.nn.functional as F + +from triton_conv3d import reference +from triton_conv3d.gather_gemm import ConvConfig, is_supported +from triton_conv3d.shapes import ConvProblem, scaffold_corpus +from triton_conv3d.transposed import ( + TransposedConfig, + candidate_transposed_configs, + conv_transpose3d_backward_data, + conv_transpose3d_backward_weight, + conv_transpose3d_forward, + default_transposed_config, + grad_transposed_weight_empty, + is_supported_transposed, + is_supported_transposed_all, + is_supported_transposed_bwd_data, + is_supported_transposed_bwd_weight, + to_tkn, + transposed_config, +) + +requires_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") + + +def _problem(name, cin, cout, spatial, k=(2, 2, 2), *, bias=False, + dtype="bf16", n=1) -> ConvProblem: + return ConvProblem(name, cin, cout, spatial, k, k, (0, 0, 0), n=n, + transposed=True, bias=bias, dtype=dtype, + sites=("synthetic",)) + + +#: Synthetic problems chosen to break the scatter rather than to be fast. Each +#: one targets a specific way a tiling store goes wrong. +EDGE: list[ConvProblem] = [ + # The four channel pairs the model contains, at a volume an fp64 reference + # can afford. The pairs matter because they decide BLOCK_NC and TAP_BLOCK. + _problem("128to64", 128, 64, (4, 4, 4), bias=True), + _problem("256to128", 256, 128, (4, 4, 4), bias=True), + _problem("512to256", 512, 256, (2, 2, 2), bias=True), + _problem("1024to512", 1024, 512, (2, 2, 2), bias=True), + # Channel counts that divide no plausible tile: Cout below the MFMA + # granularity is the interesting one, since BLOCK_NC cannot go under 16 and + # the column mask is then the only thing keeping the store in bounds. + _problem("cout_tiny", 64, 6, (4, 5, 6), bias=True), + _problem("cout_odd", 32, 7, (3, 4, 5)), + _problem("cin_prime", 17, 24, (3, 5, 7), bias=True), + _problem("cin_one", 1, 32, (4, 4, 4)), + _problem("cout_one", 32, 1, (4, 4, 4), bias=True), + # Spatial extents that do not divide BLOCK_M, so the M-unravel's rows wrap + # the W axis inside a tile and the scatter's rows are no longer a dense run. + _problem("spatial_prime", 32, 32, (13, 11, 7)), + _problem("spatial_one", 32, 32, (1, 5, 8)), + _problem("spatial_thin", 32, 32, (2, 31, 3)), + # Batch > 1: ScaFFold never does this, but the M decomposition must. + _problem("batched", 32, 32, (3, 4, 5), n=3, bias=True), + # Kernels other than 2. ``k=3`` gives 27 taps, which no power of two + # divides, so TAP_BLOCK must fall back to 1; the anisotropic ones check that + # the fused tap index is unpacked in the right radix order. + _problem("k3", 32, 32, (3, 4, 5), (3, 3, 3), bias=True), + _problem("k_aniso", 32, 32, (3, 4, 5), (1, 2, 4), bias=True), + _problem("k_aniso2", 32, 32, (4, 3, 2), (4, 2, 1)), + _problem("k1", 32, 32, (4, 5, 6), (1, 1, 1), bias=True), + # fp32 (more_determinism) and fp16, which change the LDS budget and the + # MFMA reduction depth. + _problem("fp32", 64, 64, (4, 4, 4), dtype="fp32", bias=True), + _problem("fp16", 64, 64, (4, 4, 4), dtype="fp16", bias=True), +] + +#: The corpus's real transposed problems, restated at a volume an fp64 reference +#: can afford. What survives the restatement is what matters: the channel +#: widths, and with them ``EVEN_N``, ``TAP_BLOCK`` and the tile selection. The +#: extents are deliberately not powers of two so ``BLOCK_M`` does not divide +#: ``M`` and the store's rows wrap. +CORPUS_PAIRS: list[ConvProblem] = [ + _problem(f"{p.cin}to{p.cout}-corpus", p.cin, p.cout, (3, 4, 5), + tuple(p.kernel), bias=p.bias) + for p in {(q.cin, q.cout, tuple(q.kernel), q.bias): q + for q in scaffold_corpus() if q.transposed}.values() +] + + +def _ids(problems): + return [p.name or p.label for p in problems] + + +def _ops(problem: ConvProblem, seed: int = 0, direction="fwd") -> dict: + """Operands drawn so the *realized* sums stay inside the mantissa. + + ``exact_density`` rather than a dense ``{-1,0,1}`` draw, for the reason it + documents: at ``Cin = 1024`` the forward reduces over 1024 terms and a sum + of that many random signs runs past bf16's integer limit of 256, so a dense + draw would skip every wide problem -- which is where the coverage is needed. + The shape is untouched, so the channel widths and the tile selection under + test stay exactly what ScaFFold runs. + """ + dtype = reference.torch_dtype(problem) + return reference.make_inputs( + problem, seed=seed, exact=True, + density=reference.exact_density(problem, direction, dtype=dtype), + ) + + +def _reference(problem: ConvProblem, ops: dict, direction: str) -> torch.Tensor: + return reference.reference(problem, ops, direction) + + +def _fwd(problem: ConvProblem, ops: dict, **kw) -> torch.Tensor: + return conv_transpose3d_forward( + ops["input"], ops["weight"], ops["bias"], problem.stride, **kw + ) + + +def _bwd_data(problem: ConvProblem, ops: dict, **kw) -> torch.Tensor: + return conv_transpose3d_backward_data( + ops["grad_output"], ops["weight"], problem.input_shape, problem.stride, + **kw, + ) + + +def _bwd_weight(problem: ConvProblem, ops: dict, **kw) -> torch.Tensor: + return conv_transpose3d_backward_weight( + ops["input"], problem.weight_shape, ops["grad_output"], problem.stride, + **kw, + ) + + +# --------------------------------------------------------------------------- +# The algebra, before any GPU is involved +# --------------------------------------------------------------------------- + + +def test_the_windows_tile_the_output_exactly_once(): + """The identity the whole module rests on, checked by counting. + + At ``k == s`` every output voxel must be written by exactly one ``(input + voxel, tap)`` pair. A ``k != s`` case is included as the negative control: + there the count is not 1 everywhere, which is precisely why this module + refuses it rather than generalising. + """ + for k in ((2, 2, 2), (3, 3, 3), (1, 2, 4), (4, 2, 1)): + extents = (3, 4, 5) + hits = torch.zeros(tuple(e * kk for e, kk in zip(extents, k))) + for d in range(extents[0]): + for h in range(extents[1]): + for w in range(extents[2]): + for kd in range(k[0]): + for kh in range(k[1]): + for kw in range(k[2]): + hits[d * k[0] + kd, h * k[1] + kh, + w * k[2] + kw] += 1 + assert torch.equal(hits, torch.ones_like(hits)), k + + # ``k=3, s=2`` overlaps: the windows cover some voxels twice. If this ever + # stops being true the gate could be widened; it is here so that widening it + # by accident is impossible. + k, s, extents = 3, 2, (4, 1, 1) + hits = torch.zeros((extents[0] - 1) * s + k) + for d in range(extents[0]): + for kd in range(k): + hits[d * s + kd] += 1 + assert hits.max() > 1 + + +def test_transposed_flops_have_no_phantom_tap_factor(): + """``flops()`` against the elementary MAC count, not against another formula. + + The trap: the general transposed FLOP count carries a per-tap factor, and at + ``k == s`` it does not apply, because the windows tile rather than overlap. + Applying it anyway overstates the count by ``taps`` -- 8x at ``k=2`` -- which + this project did once, and a wrong FLOP count is invisible: it produces a + plausible roofline percentage and a wrong conclusion about where the + opportunity is. + + So the count is derived here from first principles: one MAC per (output + voxel, output channel, input channel), times two. + """ + for cin, cout, spatial, k in [(128, 64, (4, 5, 6), (2, 2, 2)), + (32, 32, (3, 3, 3), (3, 3, 3)), + (16, 8, (2, 3, 4), (1, 2, 4))]: + p = _problem("f", cin, cout, spatial, k) + out_vol = math.prod(p.out_spatial) + macs = out_vol * cout * cin + assert p.flops("fwd") == 2 * macs, p.label + # Every direction performs the same contraction, so all three agree. + assert p.flops("bwd-data") == 2 * macs + assert p.flops("bwd-weight") == 2 * macs + # And the GEMM decomposition has to describe the same contraction: + # M*N*K must equal the MAC count, with K = Cin and no taps in it. + m, n, kk = p.gemm_shape("fwd") + assert m * n * kk == macs, (p.label, (m, n, kk)) + assert kk == cin, "the forward's K carries a tap factor it should not" + + +def test_the_gemm_decomposition_matches_the_kernels_grid(): + """``gemm_shape`` and the launch have to agree about what N is. + + ``gemm_shape`` reports ``N = Cout * taps`` and the kernel tiles that as + ``(taps // TAP_BLOCK)`` groups of ``TAP_BLOCK * BLOCK_NC`` columns. If the + two ever disagree the cost model is describing a different kernel from the + one that runs, which is the class of error that produced this project's + largest published mistake. + """ + for p in CORPUS_PAIRS + EDGE: + m, n, k = p.gemm_shape("fwd") + taps = p.tap_count + assert n == p.cout * taps + assert m == p.n * math.prod(p.spatial) + cfg = transposed_config(m, p.cin, p.cout, p.kernel, + reference.torch_dtype(p)) + assert taps % cfg.TAP_BLOCK == 0, (p.label, cfg) + columns = (taps // cfg.TAP_BLOCK) * cfg.TAP_BLOCK * ( + -(-p.cout // cfg.BLOCK_NC) * cfg.BLOCK_NC + ) + assert columns >= n, (p.label, cfg) + + +# --------------------------------------------------------------------------- +# Configuration legality -- the failure mode is silent, so it is checked apart +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("problem", EDGE + CORPUS_PAIRS, + ids=_ids(EDGE + CORPUS_PAIRS)) +def test_selected_config_is_legal_for_every_shape(problem: ConvProblem): + """An illegal MFMA config does not fail, it silently emits vector FMA. + + So legality is asserted rather than discovered from a timing. ``TAP_BLOCK`` + adds two constraints the other kernels do not have -- it must divide the tap + count, and it multiplies ``BLOCK_NC`` into ``BLOCK_N`` -- and both are + checked here at every shape the module can be handed, including ``k=3`` + (27 taps, which no power of two divides) and ``Cout=1``. + """ + dtype = reference.torch_dtype(problem) + m = problem.n * math.prod(problem.spatial) + cfg = transposed_config(m, problem.cin, problem.cout, problem.kernel, dtype) + assert cfg.validate(dtype) is None, (problem.label, cfg) + assert cfg.lds_bytes(dtype) <= 64 * 1024, (problem.label, cfg) + assert problem.tap_count % cfg.TAP_BLOCK == 0, (problem.label, cfg) + assert cfg.BLOCK_N == cfg.TAP_BLOCK * cfg.BLOCK_NC + + +@pytest.mark.parametrize("problem", CORPUS_PAIRS, ids=_ids(CORPUS_PAIRS)) +def test_every_candidate_config_is_legal(problem: ConvProblem): + """The sweep must never time a config that cannot reach the matrix core. + + A ranked-last illegal config still pollutes a best-of: it runs, it is + correct, and it is slow for a reason that has nothing to do with the tile. + """ + dtype = reference.torch_dtype(problem) + m = problem.n * math.prod(problem.spatial) + cands = candidate_transposed_configs(m, problem.cin, problem.cout, + problem.tap_count, dtype) + assert cands + for cfg in cands: + assert cfg.validate(dtype) is None, cfg + assert cfg.lds_bytes(dtype) <= 64 * 1024, cfg + assert problem.tap_count % cfg.TAP_BLOCK == 0, cfg + + +def test_the_fp32_config_fits_lds(): + """fp32 operands are twice the bytes, and that hole has bitten before. + + ``more_determinism`` runs the model in fp32, and the gather kernel shipped a + ``default_config`` that asked for 128 KiB there -- reachable from a real + ScaFFold configuration. This kernel's tile is *wider* than that one's + (``TAP_BLOCK`` multiplies the column count), so the same hole is closer. + """ + for cin, cout, taps in [(1024, 512, 8), (512, 256, 8), (256, 128, 8), + (128, 64, 8), (64, 64, 27), (2048, 1024, 8)]: + for dtype in (torch.float32, torch.bfloat16, torch.float16): + cfg = default_transposed_config(1 << 16, cin, cout, taps, dtype) + assert cfg.validate(dtype) is None, (cin, cout, dtype, cfg) + assert cfg.lds_bytes(dtype) <= 64 * 1024, (cin, cout, dtype, cfg) + + +def test_transposed_config_is_a_pure_function_of_its_arguments(): + """No device state, no clock, no allocator: two calls must agree. + + The same property ``split_count`` needs and for a weaker but related reason + -- a tuning choice that varied between two runs of the same shape would make + the kernel's own reproducibility claim untestable. + """ + args = (1 << 20, 128, 64, (2, 2, 2), torch.bfloat16) + assert transposed_config(*args) == transposed_config(*args) + + +# --------------------------------------------------------------------------- +# The gates +# --------------------------------------------------------------------------- + + +def test_is_supported_declines_what_the_tiling_argument_does_not_cover(): + """Everything outside ``kernel == stride, p = 0, output_padding = 0, dil = 1``. + + Each of these breaks the bijection ``(d, kd) -> d*k + kd`` in a different + way, and the failure is not a crash: ``k != s`` would write some output + voxels twice and others never, which is a smooth, plausible, wrong answer. + """ + x = torch.zeros(1, 8, 4, 4, 4) + w = torch.zeros(8, 4, 2, 2, 2) + # The shape checks run on CPU tensors, so a True is impossible here; what is + # asserted is that each of these is refused *before* the device test, which + # is why the positive control below is on the GPU. + assert not is_supported_transposed(x, w, None, 2, 0, 0, 1, 1) # not cuda + for stride, padding, output_padding, dilation, groups in [ + (1, 0, 0, 1, 1), # k != s: the windows overlap + (3, 0, 0, 1, 1), # k != s: the windows leave gaps + (2, 1, 0, 1, 1), # padding crops the tiled result + (2, 0, 1, 1, 1), # output_padding extends it asymmetrically + (2, 0, 0, 2, 1), # dilation interleaves the window with holes + (2, 0, 0, 1, 2), # groups + ((2, 2, 1), 0, 0, 1, 1), # anisotropic mismatch on one axis only + ]: + assert not is_supported_transposed( + x.cuda() if torch.cuda.is_available() else x, + w.cuda() if torch.cuda.is_available() else w, + None, stride, padding, output_padding, dilation, groups + ), (stride, padding, output_padding, dilation, groups) + + +def test_the_gates_are_total(): + """An argument the gate cannot interpret is a ``False``, never an exception. + + This is the gate of a Triton -> MIOpen rung ladder. A caller that is only + asking a question must not be taken down by the answer, and ``_triple`` + raises ``TypeError`` on ``None`` and ``ValueError`` on a bad length. + """ + x = torch.zeros(1, 8, 4, 4, 4) + w = torch.zeros(8, 4, 2, 2, 2) + for bad in (None, 1.5, "2", (2, 2), (2, 2, 2, 2), object()): + assert is_supported_transposed(x, w, None, bad, 0, 0, 1, 1) is False + assert is_supported_transposed(x, w, None, 2, bad, 0, 1, 1) is False + assert is_supported_transposed(x, w, None, 2, 0, bad, 1, 1) is False + assert is_supported_transposed(x, w, None, 2, 0, 0, bad, 1) is False + assert is_supported_transposed_all(x, w, None, bad, 0, 0, 1, 1) is False + assert is_supported_transposed_bwd_data( + x, w, (1, 8, 4, 4, 4), bad, 0, 0, 1, 1) is False + assert is_supported_transposed_bwd_weight( + x, (8, 4, 2, 2, 2), x, bad, 0, 0, 1, 1) is False + # A malformed ``input_shape`` / ``weight_shape`` is the same kind of + # question and gets the same kind of answer. + for bad in (None, (1, 8, 4, 4), "abcde", 5): + assert is_supported_transposed_bwd_data(x, w, bad, 2, 0, 0, 1, 1) is False + assert is_supported_transposed_bwd_weight(x, bad, x, 2, 0, 0, 1, 1) is False + + +def test_is_supported_declines_degenerate_extents(): + """Zero-length axes and empty channel counts, which torch handles otherwise. + + Each clears the tiling argument and then disagrees with torch: a zero-length + spatial axis gives an output the M-unravel has no rows to index, and + ``Cin = 0`` returns ``Cout`` channels of zeros where torch returns a tensor + with no channels at all -- a different *shape*, not a different value. + """ + good_x = torch.zeros(1, 8, 4, 4, 4) + good_w = torch.zeros(8, 4, 2, 2, 2) + assert not is_supported_transposed(torch.zeros(1, 8, 0, 4, 4), good_w, + None, 2, 0, 0, 1, 1) + assert not is_supported_transposed(good_x, torch.zeros(0, 4, 2, 2, 2), + None, 2, 0, 0, 1, 1) + assert not is_supported_transposed(good_x, torch.zeros(8, 0, 2, 2, 2), + None, 2, 0, 0, 1, 1) + assert not is_supported_transposed(good_x, torch.zeros(8, 4, 0, 2, 2), + None, (0, 2, 2), 0, 0, 1, 1) + + +@requires_gpu +def test_is_supported_reads_the_transposed_weight_convention(): + """``(Cin, Cout, k, k, k)``, not ``(Cout, Cin, k, k, k)``. + + The two operators store their weights with the channel axes the other way + round. A gate that read ``nn.Conv3d``'s convention would accept a weight + whose channel counts happen to match and compute a transposed answer -- the + right shape, the wrong numbers, silently. + """ + x = torch.zeros(1, 8, 4, 4, 4, device="cuda", dtype=torch.bfloat16) + assert is_supported_transposed( + x, torch.zeros(8, 4, 2, 2, 2, device="cuda", dtype=torch.bfloat16), + None, 2, 0, 0, 1, 1) + # Same tensor read the other way round: Cin=4 does not match x's 8 channels. + assert not is_supported_transposed( + x, torch.zeros(4, 8, 2, 2, 2, device="cuda", dtype=torch.bfloat16), + None, 2, 0, 0, 1, 1) + # A bias is ``Cout`` = w.shape[1] long, not w.shape[0]. + w = torch.zeros(8, 4, 2, 2, 2, device="cuda", dtype=torch.bfloat16) + assert is_supported_transposed( + x, w, torch.zeros(4, device="cuda", dtype=torch.bfloat16), + None or 2, 0, 0, 1, 1) + assert not is_supported_transposed( + x, w, torch.zeros(8, device="cuda", dtype=torch.bfloat16), 2, 0, 0, 1, 1) + # A stride-2 view of the right length applies every other value; the kernel + # indexes the bias with an element stride of 1 and cannot see this. + long_bias = torch.zeros(8, device="cuda", dtype=torch.bfloat16) + assert not is_supported_transposed(x, w, long_bias[::2], 2, 0, 0, 1, 1) + + +@requires_gpu +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two GPUs") +def test_is_supported_declines_operands_on_different_devices(): + """Triton launches on the current device and dereferences the other pointer. + + ScaFFold runs four GPUs per node with peer access, so a foreign pointer does + not fault -- it reads another rank's memory, which is a plausible wrong + answer rather than a crash. + """ + x = torch.zeros(1, 8, 4, 4, 4, device="cuda:0", dtype=torch.bfloat16) + w = torch.zeros(8, 4, 2, 2, 2, device="cuda:1", dtype=torch.bfloat16) + assert not is_supported_transposed(x, w, None, 2, 0, 0, 1, 1) + assert not is_supported_transposed_all(x, w, None, 2, 0, 0, 1, 1) + + +@requires_gpu +@pytest.mark.parametrize("problem", EDGE + CORPUS_PAIRS, + ids=_ids(EDGE + CORPUS_PAIRS)) +def test_all_three_gates_accept_every_problem_this_module_serves(problem): + """``is_supported_transposed_all`` must not be narrower than the forward. + + Unlike the ordinary convolution -- whose three gates genuinely disagree + about ``stride > 1``, which is a trap the package documents -- all three + directions of this operator accept the same problems, because both backward + directions are the same ``k == s`` convolution seen from the other side. + That is an argument, and the adapter needs a fact: a site only leaves the + block-list if the *combined* gate says yes, so it is asked for real here at + every shape the module claims. + """ + dtype = reference.torch_dtype(problem) + x = torch.zeros(problem.input_shape, device="cuda", dtype=dtype) + w = torch.zeros(problem.weight_shape, device="cuda", dtype=dtype) + b = (torch.zeros(problem.cout, device="cuda", dtype=dtype) + if problem.bias else None) + args = (problem.stride, 0, 0, 1, 1) + assert is_supported_transposed(x, w, b, *args), problem.label + assert is_supported_transposed_all(x, w, b, *args), problem.label + + +def test_the_ordinary_forward_gate_would_not_have_served_these(): + """Why this module exists at all, stated as a test. + + The ordinary ``is_supported`` takes no ``transposed`` parameter, so a caller + holding a ``ConvTranspose3d`` has no way to ask it the right question: it + answers about the *non*-transposed convolution with the same tensors, whose + output shape is 8x smaller. Asking it and believing the answer is precisely + the bug the adapter's ``module.transposed`` check exists to prevent. + """ + x = torch.zeros(1, 128, 4, 4, 4) + w = torch.zeros(128, 64, 2, 2, 2) + # It answers -- about a 128 -> 64 strided convolution, not about the + # upsample -- and the answer says nothing about this operator. + assert is_supported(x, w, None, 2, 0, 1, 1) in (True, False) + p = _problem("t", 128, 64, (4, 4, 4)) + assert p.out_spatial == (8, 8, 8) + non_transposed = (4 + 2 * 0 - 2) // 2 + 1 + assert non_transposed == 2 != 8 + + +# --------------------------------------------------------------------------- +# Correctness: bitwise, because a permuted scatter is invisible to a tolerance +# --------------------------------------------------------------------------- + + +@requires_gpu +@pytest.mark.parametrize("problem", EDGE, ids=_ids(EDGE)) +def test_forward_matches_bitwise(problem: ConvProblem): + ops = _ops(problem, seed=3, direction="fwd") + expected = _reference(problem, ops, "fwd") + dtype = reference.torch_dtype(problem) + assert reference.is_exactly_representable(expected, dtype), ( + f"{problem.label}: the draw is not exact, so this case tests nothing" + ) + actual = _fwd(problem, ops) + assert reference.compare(actual, expected.to(dtype)).bitwise, problem.label + + +@requires_gpu +@pytest.mark.parametrize("problem", EDGE, ids=_ids(EDGE)) +def test_backward_data_matches_bitwise(problem: ConvProblem): + ops = _ops(problem, seed=5, direction="bwd-data") + expected = _reference(problem, ops, "bwd-data") + dtype = reference.torch_dtype(problem) + assert reference.is_exactly_representable(expected, dtype), problem.label + actual = _bwd_data(problem, ops) + assert reference.compare(actual, expected.to(dtype)).bitwise, problem.label + + +@requires_gpu +@pytest.mark.parametrize("problem", EDGE, ids=_ids(EDGE)) +def test_backward_weight_matches_bitwise(problem: ConvProblem): + ops = _ops(problem, seed=7, direction="bwd-weight") + expected = _reference(problem, ops, "bwd-weight") + dtype = reference.torch_dtype(problem) + assert reference.is_exactly_representable(expected, dtype), problem.label + actual = _bwd_weight(problem, ops) + assert reference.compare(actual, expected.to(dtype)).bitwise, problem.label + + +@requires_gpu +@pytest.mark.parametrize("problem", CORPUS_PAIRS, ids=_ids(CORPUS_PAIRS)) +@pytest.mark.parametrize("direction", ["fwd", "bwd-data", "bwd-weight"]) +def test_corpus_channel_pairs_match_bitwise(problem: ConvProblem, direction: str): + """Every ``ConvTranspose3d`` channel pair ScaFFold runs, bitwise in bf16. + + ``exact_density`` is what makes this reachable at ``Cin = 1024``: it thins + the activations so the *realized* sums stay inside bf16's mantissa while the + shape -- and so the tile, ``TAP_BLOCK`` and the 512-byte row strides -- is + exactly what the model runs. Asserted rather than skipped, so this cannot + quietly become a wall of passes that tests nothing, which is how a sibling + file lost its whole real-shape coverage once. + """ + ops = _ops(problem, seed=11, direction=direction) + expected = _reference(problem, ops, direction) + dtype = reference.torch_dtype(problem) + assert reference.is_exactly_representable(expected, dtype), problem.label + actual = {"fwd": _fwd, "bwd-data": _bwd_data, + "bwd-weight": _bwd_weight}[direction](problem, ops) + assert reference.compare(actual, expected.to(dtype)).bitwise, problem.label + + +@requires_gpu +def test_a_transposed_tap_permutation_is_detected(): + """The bug this kernel can uniquely have, pinned. + + Each output voxel takes its value from one tap, and which tap is decided by + ``(D % kd, H % kh, W % kw)``. Unpack the fused tap index in the wrong radix + order -- swap ``kd`` and ``kw``, the single most plausible mistake in the + epilogue -- and every value written is a value that *belongs* somewhere in + the output, just not there. The norms are identical, the histogram is + identical, and every tolerance test ever written passes. + + So construct exactly that wrong answer, from the same operands, and require + a bitwise mismatch. An anisotropic volume is used so that the permutation + cannot coincide with a symmetry of the data. + """ + problem = _problem("perm", 32, 32, (3, 4, 5)) + ops = _ops(problem, seed=13) + actual = _fwd(problem, ops) + correct = _reference(problem, ops, "fwd").to(torch.bfloat16) + assert reference.compare(actual, correct).bitwise + + # The same convolution with the weight's three kernel axes permuted, which + # is exactly what unpacking the fused index in the wrong order computes. + permuted = _reference( + problem, {**ops, "weight": ops["weight"].permute(0, 1, 4, 3, 2) + .contiguous(memory_format=torch.channels_last_3d)}, "fwd", + ).to(torch.bfloat16) + assert permuted.shape == actual.shape + assert not reference.compare(actual, permuted).bitwise, ( + "a kd/kw-swapped tap unpacking gave a bitwise-identical answer; the " + "scatter's radix order is untested by this suite" + ) + # And the sums agree, which is the point: nothing weaker than bitwise sees it. + assert torch.allclose(actual.double().sum(), permuted.double().sum()) + + +@requires_gpu +def test_bitwise_standard_rejects_a_shifted_scatter(): + """Prove the comparison discriminates at all. + + ``{-1,0,1}`` operands could in principle make two different answers agree, + and this project has shipped two vacuous exact tests before. A one-voxel + roll of the input is a different convolution and must be rejected. + """ + problem = _problem("shift", 16, 16, (3, 4, 5)) + ops = _ops(problem, seed=17) + actual = _fwd(problem, ops) + correct = _reference(problem, ops, "fwd").to(torch.bfloat16) + assert reference.compare(actual, correct).bitwise + + rolled = torch.roll(ops["input"], shifts=1, dims=-1) + wrong = _reference(problem, {**ops, "input": rolled}, "fwd").to(torch.bfloat16) + assert not reference.compare(actual, wrong).bitwise + + +@requires_gpu +def test_backward_weight_operand_swap_is_not_reversible(): + """The one mistake the backward-weight re-expression can make. + + ``conv_transpose3d_backward_weight`` hands ``grad_output`` to + ``conv3d_backward_weight``'s *input* slot and ``x`` to its *grad_output* + slot. The swap is checked twice, because it has two regimes and only one of + them is dangerous: + + * at ``k > 1`` the swap is **not shape-legal** -- the strided convolution's + input is the ``k``-times-larger volume, so ``is_supported_bwd_weight`` + refuses it. That is worth pinning as a fact rather than assumed: it is + the reason the swap cannot silently produce a wrong gradient at any real + ScaFFold site. + * at ``k == 1`` the two activations have the *same* shape, the gate cannot + tell them apart, and the swap returns a correctly shaped, transposed + gradient. That is the case where only the operand order stands between a + right and a wrong answer, so it is constructed and required to differ. + """ + from triton_conv3d.reduce_gemm import conv3d_backward_weight + + problem = _problem("swap", 32, 32, (3, 4, 5)) + ops = _ops(problem, seed=19, direction="bwd-weight") + actual = _bwd_weight(problem, ops) + expected = _reference(problem, ops, "bwd-weight").to(torch.bfloat16) + assert reference.compare(actual, expected).bitwise + with pytest.raises(NotImplementedError): + conv3d_backward_weight(ops["input"], problem.weight_shape, + ops["grad_output"], problem.stride, 0, 1, 1) + + # ``k=1``: same shapes, so nothing but the argument order decides. + flat = _problem("swap1", 32, 32, (3, 4, 5), (1, 1, 1)) + ops1 = _ops(flat, seed=19, direction="bwd-weight") + got = _bwd_weight(flat, ops1) + want = _reference(flat, ops1, "bwd-weight").to(torch.bfloat16) + assert reference.compare(got, want).bitwise + swapped = conv3d_backward_weight(ops1["input"], flat.weight_shape, + ops1["grad_output"], flat.stride, 0, 1, 1) + assert swapped.shape == got.shape + assert not reference.compare(got, swapped).bitwise, ( + "swapping the two activations gave the same gradient; the operand " + "order of the backward-weight re-expression is untested" + ) + + +@requires_gpu +def test_backward_data_is_the_strided_convolution_it_claims_to_be(): + """The re-expression, stated as an identity and checked bitwise. + + ``grad_input = conv3d(grad_output, w, stride=k)`` with ``w`` *unpermuted*. + If a permute were needed the two would differ, and the difference would be a + transposed gradient of the right shape whenever ``Cin == Cout``. + """ + problem = _problem("bd", 64, 32, (3, 4, 5)) + ops = _ops(problem, seed=23, direction="bwd-data") + from triton_conv3d.gather_gemm import conv3d_forward + direct = conv3d_forward(ops["grad_output"], ops["weight"], None, + problem.stride, 0, 1, 1) + assert torch.equal(direct, _bwd_data(problem, ops)) + + +# --------------------------------------------------------------------------- +# The tuning surface, the layouts, and the entry-point contract +# --------------------------------------------------------------------------- + + +@requires_gpu +@pytest.mark.parametrize("problem", EDGE[:8], ids=_ids(EDGE[:8])) +def test_every_config_gives_the_same_answer(problem: ConvProblem): + """The whole tuning surface, not the one point the table happens to pick. + + ``TAP_BLOCK`` is the axis that matters here: it changes how many taps share + an accumulator and therefore the column decomposition of the store, so a + mask that is right at ``TAP_BLOCK=1`` and wrong at 8 would be invisible to + a test that only ran the shipped config. Capped at four configs per shape + to keep the JIT cost bounded; they are chosen to span ``TAP_BLOCK``. + """ + dtype = reference.torch_dtype(problem) + ops = _ops(problem, seed=29) + expected = _reference(problem, ops, "fwd") + assert reference.is_exactly_representable(expected, dtype) + m = problem.n * math.prod(problem.spatial) + cands = candidate_transposed_configs(m, problem.cin, problem.cout, + problem.tap_count, dtype) + by_tb: dict[int, TransposedConfig] = {} + for cfg in cands: + by_tb.setdefault(cfg.TAP_BLOCK, cfg) + chosen = list(by_tb.values())[:4] + assert chosen, problem.label + assert len({c.TAP_BLOCK for c in chosen}) == len(chosen) + for cfg in chosen: + actual = _fwd(problem, ops, config=cfg) + assert reference.compare(actual, expected.to(dtype)).bitwise, ( + f"{problem.label} with {cfg}" + ) + + +@requires_gpu +def test_every_weight_layout_gives_the_same_answer(): + """Channels-last, the materialized ``(t, K, N)`` buffer, and PyTorch's default. + + Three layouts, one answer. The middle one is the copy + :func:`~triton_conv3d.transposed.to_tkn` makes for a weight the plan + refuses, and the last one is what the plan refuses -- a weight where neither + channel axis is unit-stride. Getting the stride plan wrong is a *silent* + wrong answer, because the kernel will happily read whatever the strides say. + """ + problem = _problem("layout", 64, 32, (3, 4, 5), bias=True) + ops = _ops(problem, seed=31) + expected = _reference(problem, ops, "fwd").to(torch.bfloat16) + cl = ops["weight"] + assert cl.is_contiguous(memory_format=torch.channels_last_3d) + plain = cl.contiguous() + assert not plain.is_contiguous(memory_format=torch.channels_last_3d) + for w in (cl, plain, to_tkn(cl).permute(3, 4, 0, 1, 2)): + got = conv_transpose3d_forward(ops["input"], w, ops["bias"], + problem.stride) + assert reference.compare(got, expected).bitwise, tuple(w.stride()) + + +@requires_gpu +def test_out_buffer_is_written_in_place_and_is_validated(): + """``out=`` is checked rather than trusted, and nothing downstream catches it. + + The grid is sized from the *problem*, not from ``out``, and the store + addressing assumes a channel stride of 1 -- so an undersized buffer is an + out-of-bounds device write with no error and an NCDHW one is a full-rate + kernel returning a scrambled answer. + """ + problem = _problem("outbuf", 32, 16, (3, 4, 5), bias=True) + ops = _ops(problem, seed=37) + expected = _reference(problem, ops, "fwd").to(torch.bfloat16) + y = torch.empty(problem.output_shape, device="cuda", dtype=torch.bfloat16, + memory_format=torch.channels_last_3d) + got = _fwd(problem, ops, out=y) + assert got.data_ptr() == y.data_ptr() + assert reference.compare(y, expected).bitwise + + small = torch.empty((1, 16, 2, 2, 2), device="cuda", dtype=torch.bfloat16, + memory_format=torch.channels_last_3d) + with pytest.raises(ValueError, match="shape"): + _fwd(problem, ops, out=small) + ncdhw = torch.empty(problem.output_shape, device="cuda", + dtype=torch.bfloat16) + with pytest.raises(ValueError, match="channels_last_3d"): + _fwd(problem, ops, out=ncdhw) + wrong_dtype = torch.empty(problem.output_shape, device="cuda", + dtype=torch.float32, + memory_format=torch.channels_last_3d) + with pytest.raises(ValueError, match="dtype"): + _fwd(problem, ops, out=wrong_dtype) + + +@requires_gpu +def test_an_illegal_tap_block_is_refused_rather_than_run(): + """A ``TAP_BLOCK`` that does not divide the tap count. + + The kernel's ``pid % (taps // TAP_BLOCK)`` would then address a tap group + that runs off the end of the weight -- a wrong answer, not a fault, because + the offsets stay inside the allocation for small kernels. Refused at the + entry point, loudly, since it can only arrive through an explicit + ``config=``. + """ + problem = _problem("tb", 32, 32, (3, 4, 5), (3, 3, 3)) + ops = _ops(problem, seed=41) + bad = TransposedConfig(BLOCK_M=64, BLOCK_N=64, BLOCK_K=32, TAP_BLOCK=2) + with pytest.raises(ValueError, match="TAP_BLOCK"): + _fwd(problem, ops, config=bad) + # And an outright illegal MFMA config is refused by the inherited rules. + with pytest.raises(ValueError, match="nonkdim"): + _fwd(problem, ops, + config=TransposedConfig(BLOCK_M=64, BLOCK_N=64, BLOCK_K=32, + matrix_instr_nonkdim=64)) + + +@requires_gpu +def test_ncdhw_input_is_converted_rather_than_misread(): + """A plain-contiguous input is relayouted, not read as if it were NDHWC. + + The addressing assumes a channel stride of 1. Reading an NCDHW tensor with + it would produce a full-rate kernel and a scrambled answer, which is why the + entry point calls ``contiguous(memory_format=...)`` rather than asserting. + """ + problem = _problem("ncdhw", 32, 16, (3, 4, 5), bias=True) + ops = _ops(problem, seed=43) + expected = _reference(problem, ops, "fwd").to(torch.bfloat16) + plain = ops["input"].contiguous() + assert not plain.is_contiguous(memory_format=torch.channels_last_3d) + got = conv_transpose3d_forward(plain, ops["weight"], ops["bias"], + problem.stride) + assert got.is_contiguous(memory_format=torch.channels_last_3d) + assert reference.compare(got, expected).bitwise + + +@requires_gpu +def test_output_matches_torchs_shape_and_layout(): + """Shape and memory format against ``F.conv_transpose3d``, at every kernel.""" + for k in ((2, 2, 2), (3, 3, 3), (1, 2, 4)): + problem = _problem("shape", 32, 16, (3, 4, 5), k, bias=True) + ops = _ops(problem, seed=47) + got = _fwd(problem, ops) + want = F.conv_transpose3d(ops["input"], ops["weight"], ops["bias"], + stride=k) + assert got.shape == want.shape, k + assert got.is_contiguous(memory_format=torch.channels_last_3d) + assert tuple(got.shape[2:]) == problem.out_spatial + + +@requires_gpu +def test_no_worse_than_miopen(): + """Error against fp64, held to the incumbent's own error where possible. + + ``assert_close``'s policy, unchanged and not reinvented: it once failed on + MIOpen's *transposed* backward-weight, and the resolution was that the + tolerance was wrong -- it charged the final store like an accumulation. + ``roundings`` is 2 for MIOpen's backward-weight because that direction + reduces with atomics and disagrees with itself bitwise between two calls. + """ + for problem in [_problem("mi", 64, 32, (4, 5, 6), bias=True), + _problem("mi3", 32, 32, (3, 4, 5), (3, 3, 3))]: + for direction in ("fwd", "bwd-data", "bwd-weight"): + ops = reference.make_inputs(problem, seed=53) + expected = _reference(problem, ops, direction) + incumbent = reference.compare( + reference.incumbent(problem, ops, direction), expected) + actual = {"fwd": _fwd, "bwd-data": _bwd_data, + "bwd-weight": _bwd_weight}[direction](problem, ops) + reference.assert_close(actual, expected, problem, direction, + incumbent_error=incumbent) + + +@requires_gpu +def test_repeated_calls_are_bitwise_reproducible(): + """The whole operator, run twice, must agree bitwise in all three directions. + + Backward-weight is the one at risk: it is ``conv3d_backward_weight``, whose + deterministic split-K path is the default and whose atomic path is not + reproducible. Nothing here asks for the atomic path, and this test is what + says so. + """ + problem = _problem("repro", 64, 32, (4, 5, 6), bias=True) + ops = reference.make_inputs(problem, seed=59) + for run in (_fwd, _bwd_data, _bwd_weight): + first = run(problem, ops) + for _ in range(3): + assert torch.equal(first, run(problem, ops)) + + +@requires_gpu +def test_grad_weight_buffer_has_the_transposed_shape(): + """``(Cin, Cout, k, k, k)``, channels-last -- the parameter's own layout. + + The ordinary ``grad_weight_empty`` allocates ``(Cout, Cin, ...)``. Passing + that here is a correctly-strided buffer of the wrong shape, which the + reduction's ``out=`` check catches only because it compares the shape + explicitly -- none of the five channels-last strides depends on the first + dimension. + """ + from triton_conv3d.reduce_gemm import grad_weight_empty + gw = grad_transposed_weight_empty(128, 64, (2, 2, 2), + dtype=torch.bfloat16, device="cuda") + assert tuple(gw.shape) == (128, 64, 2, 2, 2) + assert gw.is_contiguous(memory_format=torch.channels_last_3d) + + problem = _problem("gw", 32, 16, (3, 4, 5)) + ops = _ops(problem, seed=61, direction="bwd-weight") + out = grad_transposed_weight_empty(32, 16, (2, 2, 2), + dtype=torch.bfloat16, device="cuda") + got = _bwd_weight(problem, ops, out=out) + assert got.data_ptr() == out.data_ptr() + expected = _reference(problem, ops, "bwd-weight").to(torch.bfloat16) + assert reference.compare(got, expected).bitwise + + wrong = grad_weight_empty(32, 16, (2, 2, 2), dtype=torch.bfloat16, + device="cuda") + assert tuple(wrong.shape) == (32, 16, 2, 2, 2) + other = _problem("gw2", 16, 32, (3, 4, 5)) + ops2 = _ops(other, seed=61, direction="bwd-weight") + with pytest.raises(ValueError, match="shape"): + _bwd_weight(other, ops2, out=wrong) + + +@requires_gpu +def test_unsupported_calls_raise_rather_than_return_garbage(): + """Each entry point re-asks its own gate and refuses, never guesses.""" + problem = _problem("raise", 32, 16, (3, 4, 5)) + ops = _ops(problem, seed=67) + with pytest.raises(NotImplementedError): + conv_transpose3d_forward(ops["input"], ops["weight"], None, 3) + with pytest.raises(NotImplementedError): + conv_transpose3d_backward_data(ops["grad_output"], ops["weight"], + problem.input_shape, 3) + with pytest.raises(NotImplementedError): + conv_transpose3d_backward_weight(ops["input"], problem.weight_shape, + ops["grad_output"], 3) + # And a padding, which is the one a caller is most likely to pass by habit. + with pytest.raises(NotImplementedError): + conv_transpose3d_forward(ops["input"], ops["weight"], None, + problem.stride, 1) + + +@requires_gpu +def test_fp32_accumulates_in_fp32(): + """``more_determinism`` runs the model in fp32 and it has to really be fp32. + + The backend's default ``input_precision`` splits an fp32 dot into + reduced-precision pieces, which is a ~10-bit mantissa and passes every + tolerance this package has. Only a bitwise test over a long reduction sees + it, so the reduction here is long enough to matter. + """ + problem = _problem("fp32acc", 512, 64, (2, 3, 4), dtype="fp32") + ops = _ops(problem, seed=71) + expected = _reference(problem, ops, "fwd") + assert reference.is_exactly_representable(expected, torch.float32) + actual = _fwd(problem, ops) + assert actual.dtype is torch.float32 + assert reference.compare(actual, expected.to(torch.float32)).bitwise + + +@requires_gpu +def test_bias_is_per_channel_and_not_per_column(): + """One bias value per output channel, shared by all ``taps`` sub-lattices. + + In the kernel the bias is indexed by ``offs_n`` and not by the column, which + is the difference between a bias and a per-tap offset. Indexing it by the + column would read ``TAP_BLOCK * Cout`` values from a ``Cout``-long tensor -- + past the end for every tap but the first, and wrong even where it is in + bounds. Checked by making the bias the only nonzero operand, so the answer + *is* the bias broadcast over the upsampled volume. + """ + problem = _problem("bias", 64, 48, (3, 4, 5), bias=True) + ops = _ops(problem, seed=73) + ops = {**ops, "weight": torch.zeros_like(ops["weight"]), + "bias": torch.arange(1, 49, device="cuda", + dtype=torch.bfloat16)} + got = _fwd(problem, ops) + want = ops["bias"].view(1, 48, 1, 1, 1).expand(got.shape) + assert torch.equal(got, want.to(got.dtype)) diff --git a/triton_conv3d/transposed.py b/triton_conv3d/transposed.py new file mode 100644 index 0000000..1f2d383 --- /dev/null +++ b/triton_conv3d/transposed.py @@ -0,0 +1,1149 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""Transposed 3-D convolution at ``kernel == stride``, on NDHWC tensors. + +ScaFFold's decoder upsamples with four ``nn.ConvTranspose3d(k=2, s=2, p=0)`` +sites. That case is *much* simpler than a general transposed convolution, and +the whole design here follows from one observation: + + at ``kernel == stride`` and no padding the scatter windows **tile** the + output rather than overlapping, so every output voxel receives exactly one + contribution. + +Concretely, with ``k = s`` the map ``(d, kd) -> D = d*k + kd`` is a bijection +onto ``[0, k*ID)`` -- it is just base-``k`` positional notation -- so + + y[n, oc, d*KD+kd, h*KH+kh, w*KW+kw] = sum_ic x[n, ic, d, h, w] * w[ic, oc, kd, kh, kw] + +with **no sum over taps at all**. There is no accumulation across windows and +no overlap-add: the operator is a pointwise GEMM from ``Cin`` to ``Cout * taps`` +channels, followed by an interleave of those ``taps`` groups into the ``taps`` +sub-lattices of the output volume -- a 3-D pixel shuffle. + +Three directions, one new kernel +================================ + +Only the forward needs a kernel. Both backward directions are the *ordinary* +strided convolution this operator is the transpose of, which this package +already serves: + + let C(u, w) = conv3d(u, weight=w, stride=k, padding=0) + with u an NDHWC tensor of Cout channels and w read as (Cin, Cout, kd, kh, kw) + -- i.e. PyTorch's ConvTranspose3d weight *unpermuted*, whose dim 0 is the + convolution's output-channel axis and whose dim 1 is its input-channel axis. + + C(u, w)[n, ic, d, h, w'] = sum_{oc, t} u[n, oc, d*k+kd, ...] * w[ic, oc, t] + +Comparing that with the display above: + +* **backward-data is exactly** ``C(grad_output, w)``. A strided forward + convolution, so :func:`~triton_conv3d.gather_gemm.conv3d_forward` serves it + with no new code and no permute of the parameter -- the transposed operator's + weight already *is* the shape a ``Cout -> Cin`` convolution wants. +* **backward-weight is exactly** ``C``'s backward-weight, with ``grad_output`` + in the "input" slot and ``x`` in the "grad_output" slot. Backward-weight has + no stride restriction (its reduction is indexed by the output voxel), so + :func:`~triton_conv3d.reduce_gemm.conv3d_backward_weight` serves it unchanged, + and it produces the gradient in ``channels_last_3d`` -- which for a + ``(Cin, Cout, k, k, k)`` parameter is the layout ScaFFold's optimizer wants. +* **the forward is** ``C``'s backward-*data*, which + :mod:`~triton_conv3d.bwd_data` refuses: its kernel-free formulation (the + forward contraction on a flipped weight) holds only at unit stride, and at + ``stride > 1`` the gather becomes a scatter into a sub-lattice. That scatter + is what :func:`_convT3d_fwd_kernel` below is. + +So this module adds one ``@triton.jit`` function and two thin re-expressions. +The package's structural claim becomes: **four operators, three kernels.** + +The FLOP count has no per-tap factor +==================================== +``2 * in_vol * Cin * Cout * taps`` looks like the general transposed formula and +is not: the ``taps`` here is the *output/input volume ratio*, not a per-tap +gather. Each output voxel takes ``Cin`` MACs per output channel and there are +``taps * in_vol`` of them. Applying both factors at once overstates the count +by ``taps`` -- 8x at ``k=2`` -- which this project did once; +``shapes.ConvProblem.flops`` has it right and +``test_shapes.py::test_transposed_flops_have_no_phantom_tap_factor`` pins it. +Check any new arithmetic against :meth:`ConvProblem.gemm_shape`, which reports +``(in_vol, Cout*taps, Cin)`` for the forward -- one ``K = Cin``, no taps in it. + +Why the taps go in N and not in M +================================= +The GEMM is ``M = N*ID*IH*IW`` input voxels, ``N = Cout * taps``, ``K = Cin``, +and the only real design question is which axis carries the taps. + +Putting them in **M** -- tiling the *output* volume, so each row is one output +voxel and N is plain ``Cout`` -- gives a perfectly coalesced store and a +correctly-once-written output, and then dies: the weight column a row needs +depends on that row's tap, and ``kw`` alternates between adjacent rows along W. +The B operand would have to vary down the M axis, which is not a GEMM. + +Putting them in **N** keeps B constant per tile, and the tile's tap group is a +function of the program id alone. The store is then a scatter -- but a +*structured* one: within one tap the columns are consecutive output channels at +one voxel, i.e. contiguous, and the row-to-row step is ``k`` voxels. With +``TAP_BLOCK`` covering the ``kw`` pair and ``BLOCK_NC == Cout`` the tile is one +dense run of memory outright. + +``TAP_BLOCK`` is why the tap axis has to be in the tile rather than in the grid. +Every tap of a given input voxel reads the *same* A row, so ``taps`` separate +programs would read the input ``taps`` times -- 8x at ``k=2``, against an output +that is only 4x the input at ``128 -> 64``, i.e. more traffic than the answer. +One program spanning ``TAP_BLOCK`` taps loads A once for all of them, in exactly +the way :mod:`~triton_conv3d.reduce_gemm` widens its N across taps and for +exactly the same reason. +""" + +from __future__ import annotations + +import dataclasses +import math +from typing import Sequence + +import torch +import triton +import triton.language as tl + +from .gather_gemm import ( + ConvConfig, + _LDS_BYTES, + _MFMA_KDIM, + _check_out, + _index_dtype, + _pow2_at_most, + _triple, + conv3d_forward, + is_supported as _is_supported_fwd, +) +from .reduce_gemm import conv3d_backward_weight, is_supported_bwd_weight + +__all__ = [ + "TransposedConfig", + "conv_transpose3d_forward", + "conv_transpose3d_backward_data", + "conv_transpose3d_backward_weight", + "default_transposed_config", + "is_supported_transposed", + "is_supported_transposed_all", + "is_supported_transposed_bwd_data", + "is_supported_transposed_bwd_weight", + "candidate_transposed_configs", + "grad_transposed_weight_empty", + "register_tuned_transposed", + "to_tkn", + "transposed_config", + "verify_isa_transposed", +] + + +# --------------------------------------------------------------------------- +# The kernel +# --------------------------------------------------------------------------- + + +@triton.jit +def _convT3d_fwd_kernel( + X, W, Y, BIAS, + # Sizes. ``M_TOTAL`` is ``BATCH * IN_D * IN_H * IN_W`` -- the *input* + # volume, because that is what a scatter is indexed by. + BATCH, IN_D, IN_H, IN_W, + CIN, COUT, M_TOTAL, + # Element strides. The channel stride of X and Y is 1 by construction -- + # that is what NDHWC means -- so it is neither passed nor multiplied by. + stride_xn, stride_xd, stride_xh, stride_xw, + # The weight over the effective GEMM's axes: the fused tap index, the + # reduction axis K (Cin), and the output axis N (Cout). Which of the two + # channel strides is 1 is a constexpr (``W_ORDER``), as in the forward. + stride_wt, stride_wk, stride_wn, + stride_yn, stride_yd, stride_yh, stride_yw, + KD: tl.constexpr, KH: tl.constexpr, KW: tl.constexpr, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_NC: tl.constexpr, + BLOCK_K: tl.constexpr, BLOCK_K_COUNT: tl.constexpr, + TAP_BLOCK: tl.constexpr, GROUP_M: tl.constexpr, + HAS_BIAS: tl.constexpr, EVEN_K: tl.constexpr, EVEN_N: tl.constexpr, + INDEX_DTYPE: tl.constexpr, INPUT_PRECISION: tl.constexpr, + W_ORDER: tl.constexpr, +): + # -- which tile this program owns -------------------------------------- + # + # The tap group is the *fastest*-varying part of the id, which is a cache + # decision rather than a cosmetic one: the programs that share an A tile are + # the ones differing only in tap group, and consecutive ids are dispatched + # together, so the second read of a row lands while the first is still in + # L2/MALL. With the tap group slowest, every tap group would sweep the + # whole volume before the next one started and each sweep would come from + # HBM. Within a tap group the ordinary grouped-M swizzle applies. + pid = tl.program_id(0) + grid_m = tl.cdiv(M_TOTAL, BLOCK_M) + grid_nc = tl.cdiv(COUT, BLOCK_NC) + # Every operand is constexpr, so this is a compile-time constant and the + # ``%`` / ``//`` below fold into shifts at ``TAP_BLOCK`` a power of two. + grid_t = (KD * KH * KW) // TAP_BLOCK + pid_t = pid % grid_t + pid_mn = pid // grid_t + width = GROUP_M * grid_nc + group_id = pid_mn // width + group_size = min(grid_m - group_id * GROUP_M, GROUP_M) + pid_m = group_id * GROUP_M + (pid_mn % group_size) + pid_nc = (pid_mn % width) // group_size + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + m_valid = offs_m < M_TOTAL + + # -- unravel the fused ndhw index over the INPUT volume ---------------- + idx_w = offs_m % IN_W + tmp = offs_m // IN_W + idx_h = tmp % IN_H + tmp = tmp // IN_H + idx_d = tmp % IN_D + idx_n = tmp // IN_D + + # A's row. Cast per term rather than after the sum: ``idx_n * stride_xn`` + # alone overflows int32 on a batched scale-8 volume and the sum would then + # already be wrong before the widening happened. + x_row = ( + idx_n.to(INDEX_DTYPE) * stride_xn + + idx_d.to(INDEX_DTYPE) * stride_xd + + idx_h.to(INDEX_DTYPE) * stride_xh + + idx_w.to(INDEX_DTYPE) * stride_xw + ) + # The destination row: the *corner* of this input voxel's output window. + # The tap's offset within the window is a per-column addend below, so the + # scatter costs one vector add in the epilogue and nothing in the loop. + y_row = ( + idx_n.to(INDEX_DTYPE) * stride_yn + + (idx_d * KD).to(INDEX_DTYPE) * stride_yd + + (idx_h * KH).to(INDEX_DTYPE) * stride_yh + + (idx_w * KW).to(INDEX_DTYPE) * stride_yw + ) + + # -- the N axis: TAP_BLOCK taps x BLOCK_NC output channels -------------- + # + # All hoisted out of the reduction: the column decomposition depends on the + # tile and not on the reduction index. ``BLOCK_NC``, ``KH`` and ``KW`` are + # constexpr, so the divisions fold away. No tap needs clamping here (unlike + # ``reduce_gemm``, whose 27 taps cannot be divided by a power of two): + # ``TAP_BLOCK`` is required to divide ``taps`` exactly, so every column + # addresses a real tap. + col = tl.arange(0, BLOCK_N) + tap = pid_t * TAP_BLOCK + col // BLOCK_NC + offs_n = pid_nc * BLOCK_NC + (col % BLOCK_NC) + kd = tap // (KH * KW) + khw = tap % (KH * KW) + kh = khw // KW + kw = khw % KW + col_ok = offs_n < COUT + + # Where this column lands in the output: the tap's corner offset inside the + # window, plus the channel. ``stride_y*`` are element strides of a + # channels-last tensor, so this is the same expression the forward's + # ``x_row`` uses, read in the other direction. + y_col = ( + kd.to(INDEX_DTYPE) * stride_yd + + kh.to(INDEX_DTYPE) * stride_yh + + kw.to(INDEX_DTYPE) * stride_yw + + offs_n.to(INDEX_DTYPE) + ) + # B's column. ``W_ORDER == 0`` means Cout is unit-stride, which is what a + # ``channels_last_3d`` ConvTranspose3d parameter is: its memory order is + # ``[Cin][kd][kh][kw][Cout]``, i.e. this GEMM's ``[K][tap][N]`` with N dense. + # That is the *good* case for this direction and it needs no transform, + # unlike the ordinary forward, where the same parameter layout puts the + # reduction axis in the contiguous slot. + if W_ORDER == 0: + w_col = tap.to(INDEX_DTYPE) * stride_wt + offs_n.to(INDEX_DTYPE) + else: + w_col = ( + tap.to(INDEX_DTYPE) * stride_wt + + offs_n.to(INDEX_DTYPE) * stride_wn + ) + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + + # -- the reduction: over Cin alone. There is no tap loop -------------- + for k0 in range(BLOCK_K_COUNT): + offs_k = k0 * BLOCK_K + tl.arange(0, BLOCK_K) + + # A: no gather and no boundary predicate. Every input voxel of an + # in-range row contributes to every one of its taps, so the six compares + # the ordinary forward runs per tap do not exist here -- which is the + # whole reason ``kernel == stride`` is worth a kernel of its own. + x_ptrs = X + x_row[:, None] + offs_k[None, :] + if EVEN_K: + a = tl.load(x_ptrs, mask=m_valid[:, None], other=0.0) + else: + a = tl.load( + x_ptrs, mask=m_valid[:, None] & (offs_k < CIN)[None, :], other=0.0 + ) + + w_ptrs = W + (offs_k.to(INDEX_DTYPE) * stride_wk)[:, None] + w_col[None, :] + if EVEN_K and EVEN_N: + b = tl.load(w_ptrs) + elif EVEN_K: + b = tl.load(w_ptrs, mask=col_ok[None, :], other=0.0) + elif EVEN_N: + b = tl.load(w_ptrs, mask=(offs_k < CIN)[:, None], other=0.0) + else: + b = tl.load( + w_ptrs, mask=(offs_k < CIN)[:, None] & col_ok[None, :], other=0.0 + ) + + # ``input_precision`` only bites for fp32 operands, where the backend's + # default splits the dot into reduced-precision pieces. bf16 already + # accumulates in fp32 and is unaffected; fp32 is the ``more_determinism`` + # path and has to actually be fp32, so it is asked for explicitly. + acc = tl.dot(a, b, acc, input_precision=INPUT_PRECISION) + + if HAS_BIAS: + # Indexed by the output *channel*, not by the column: the same bias + # value serves every tap, which is what makes ``ConvTranspose3d``'s bias + # a per-channel constant over the upsampled volume. + bias = tl.load(BIAS + offs_n, mask=col_ok, other=0.0) + acc += bias[None, :].to(tl.float32) + + y_ptrs = Y + y_row[:, None] + y_col[None, :] + mask_y = tl.broadcast_to(m_valid[:, None], (BLOCK_M, BLOCK_N)) + if not EVEN_N: + mask_y = mask_y & col_ok[None, :] + tl.store(y_ptrs, acc.to(Y.dtype.element_ty), mask=mask_y) + + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class TransposedConfig(ConvConfig): + """A launch configuration with the one knob only this direction has. + + A subclass rather than another field on :class:`ConvConfig`, for the reason + :class:`~triton_conv3d.reduce_gemm.BwdWeightConfig` is one: the gather + directions have no tap axis in their tile and a config printed in a forward + sweep should not grow a suffix it cannot use. ``BLOCK_N`` keeps its meaning + as the *full* tile width, so the inherited gfx942 legality rules and the + measured LDS model stay correct unchanged. + """ + + #: How many taps one tile spans. ``BLOCK_N = TAP_BLOCK * BLOCK_NC``. + TAP_BLOCK: int = 1 + + @property + def BLOCK_NC(self) -> int: + """Output channels per tap in the tile.""" + return self.BLOCK_N // self.TAP_BLOCK + + def __str__(self) -> str: + return super().__str__() + ( + f"/tb{self.TAP_BLOCK}" if self.TAP_BLOCK != 1 else "" + ) + + def validate(self, dtype: torch.dtype) -> str | None: + why = super().validate(dtype) + if why is not None: + return why + if self.TAP_BLOCK < 1: + return "TAP_BLOCK must be at least 1" + if self.BLOCK_N % self.TAP_BLOCK: + return f"BLOCK_N must be a multiple of TAP_BLOCK={self.TAP_BLOCK}" + if self.BLOCK_NC < 1: + return "BLOCK_N // TAP_BLOCK must be at least 1" + return None + + +#: Below this many programs the grid cannot fill MI300A's 228 CUs. The same +#: value the gather kernel uses, restated rather than imported so that a change +#: there is a deliberate change here too -- the two kernels have different +#: occupancies and there is no measurement saying they should track. +_MIN_PROGRAMS = 114 + + +def _fit_transposed( + cfg: TransposedConfig, m: int, cout: int, taps: int, dtype: torch.dtype +) -> TransposedConfig: + """Shrink a tile until it fits LDS *and* the grid fills the device. + + Two shrinks, in this order, for the same reasons the gather kernel's + :func:`~triton_conv3d.gather_gemm._fit_to_lds` and ``_fit_to_grid`` give: + ``BLOCK_K`` first, because it changes neither the grid nor the parallelism; + then ``BLOCK_M``. ``BLOCK_N`` is shrunk last and only down to + ``TAP_BLOCK * nonkdim``, because halving it below that would drop + ``BLOCK_NC`` under the MFMA granularity and the tile would be mostly + padding. + + The grid clause differs from the gather kernel's in one term and it matters: + this grid has ``taps // TAP_BLOCK`` tap groups in it, so a problem whose M + and N alone look too small for 228 CUs may already fill them. Leaving that + factor out would halve ``BLOCK_M`` at every decoder site and lose the reuse + for nothing. + """ + nk = cfg.matrix_instr_nonkdim + kdim = _MFMA_KDIM.get(dtype, {}).get(nk) + if kdim is None: + return cfg + while cfg.lds_bytes(dtype) > _LDS_BYTES: + half_k, half_m, half_n = cfg.BLOCK_K // 2, cfg.BLOCK_M // 2, cfg.BLOCK_N // 2 + if half_k >= kdim and half_k % kdim == 0: + cfg = dataclasses.replace(cfg, BLOCK_K=half_k, + kpack=1 if half_k <= 16 else cfg.kpack) + elif half_m >= nk and half_m % nk == 0: + cfg = dataclasses.replace(cfg, BLOCK_M=half_m) + elif half_n >= cfg.TAP_BLOCK * nk and half_n % nk == 0: + cfg = dataclasses.replace(cfg, BLOCK_N=half_n) + else: + break # nothing left to shrink; let the launch say so + while ( + cfg.BLOCK_M > max(16, nk) + and (cfg.BLOCK_M // 2) % nk == 0 + and (-(-m // cfg.BLOCK_M) * -(-cout // cfg.BLOCK_NC) + * (taps // cfg.TAP_BLOCK)) < _MIN_PROGRAMS + ): + cfg = dataclasses.replace(cfg, BLOCK_M=cfg.BLOCK_M // 2) + warps = max(1, min(cfg.num_warps, cfg.BLOCK_M * cfg.BLOCK_N // 256)) + return dataclasses.replace(cfg, num_warps=1 << (warps.bit_length() - 1)) + + +def _largest_pow2_divisor(n: int, cap: int) -> int: + """The largest power of two that divides ``n`` and is at most ``cap``.""" + d = 1 + while d * 2 <= cap and n % (d * 2) == 0: + d *= 2 + return d + + +def default_transposed_config( + m: int, cin: int, cout: int, taps: int, dtype: torch.dtype = torch.bfloat16 +) -> TransposedConfig: + """A config that is legal for any shape this module accepts. + + ``TAP_BLOCK`` is the only choice here that is not the gather kernel's, and + it is chosen to cut the A traffic rather than to fill a tile: every tap of + an input voxel reads the same row, so a program spanning ``TAP_BLOCK`` taps + reads the input ``taps / TAP_BLOCK`` times instead of ``taps`` times. It is + capped so the tile stays 256 columns wide -- past that the accumulator alone + is 128 registers per lane at four warps and occupancy collapses. + """ + block_nc = _pow2_at_most(cout, 128) + tap_block = _largest_pow2_divisor(taps, max(1, 256 // block_nc)) + block_k = 128 if cin >= 512 else _pow2_at_most(cin, 64) + block_m = _pow2_at_most(m, 128) + nonkdim = 16 + kdim = _MFMA_KDIM[dtype][nonkdim] + block_k = max(kdim, block_k - block_k % kdim) + block_n = tap_block * block_nc + return _fit_transposed( + TransposedConfig( + BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=block_k, GROUP_M=6, + num_warps=8 if block_n >= 256 or block_k >= 128 else 4, + num_stages=2, matrix_instr_nonkdim=nonkdim, + kpack=1 if block_k <= 16 else 2, TAP_BLOCK=tap_block, + ), + m, cout, taps, dtype, + ) + + +def _tuned(bm: int, bnc: int, tb: int, bk: int, warps: int, + group_m: int = 6) -> TransposedConfig: + return TransposedConfig( + BLOCK_M=bm, BLOCK_N=tb * bnc, BLOCK_K=bk, GROUP_M=group_m, + num_warps=warps, num_stages=2, matrix_instr_nonkdim=16, + kpack=1 if bk <= 16 else 2, TAP_BLOCK=tb, + ) + + +def transposed_tune_key(dtype: torch.dtype, cin: int, cout: int, + kernel: tuple[int, ...]) -> tuple: + return (str(dtype), cin, cout, tuple(kernel)) + + +#: Measured winners for the transposed forward, keyed by ``(dtype, Cin, Cout, +#: kernel)``. Source: ``work/triton-conv/m5_fwd.json`` -- the four +#: ``ConvTranspose3d`` channel pairs the model contains, swept over the tile and +#: ``TAP_BLOCK`` grid of :func:`candidate_transposed_configs` and then raced +#: against MIOpen. A miss falls back to :func:`default_transposed_config`. +#: +#: Keyed on the channel widths and not the volume, as the gather kernel's table +#: is -- and with the same caveat, which this project has now paid for twice: a +#: *speedup ratio* does not transfer across volume even when the winning tile +#: does. So only the *tile* is claimed to transfer, and only where it was +#: measured winning at every volume the pair occurs at. Each pair below occurs +#: at three volumes (one per profiled configuration) and the entry named won all +#: three; the speedups they produce differ by up to 2.3x between those volumes, +#: are recorded per volume in ``work/triton-conv/m5_shipped_*.json``, and are +#: never averaged. +#: +#: **Two of the four pairs are deliberately absent.** ``512 -> 256`` and +#: ``1024 -> 512`` were swept just as thoroughly and +#: :func:`default_transposed_config` picked the winner or a tie at every volume +#: (within 0.4-8%, and the sweep's nominal best flipped tile between volumes), +#: so an entry would restate the heuristic while claiming to have improved on +#: it. An absent row here means "measured, and the heuristic was right", which +#: is a different statement from "never measured" -- the gather kernel's table +#: had to learn that distinction the hard way. +_TUNED_T: dict[tuple, TransposedConfig] = { + transposed_tune_key(torch.bfloat16, cin, cout, (2, 2, 2)): cfg + for (cin, cout), cfg in { + # Both winners are ``BLOCK_NC = 64`` with ``TAP_BLOCK = 4``, i.e. a + # 256-column tile spanning half the taps, against the heuristic's + # ``BLOCK_NC = Cout, TAP_BLOCK = 2``. Same column count, twice the tap + # reuse: the input is read twice instead of four times, which is what + # this operator is short of at these channel widths. Worth 1.12-1.23x + # over the heuristic and it is the whole gap between them. + (128, 64): _tuned(256, 64, 4, 64, 8), + (256, 128): _tuned(128, 64, 4, 64, 8), + }.items() +} + + +def register_tuned_transposed(dtype, cin, cout, kernel, config) -> None: + _TUNED_T[transposed_tune_key(dtype, cin, cout, kernel)] = config + + +def transposed_config( + m: int, cin: int, cout: int, kernel: Sequence[int], + dtype: torch.dtype = torch.bfloat16, +) -> TransposedConfig: + """The config :func:`conv_transpose3d_forward` would pick for this problem.""" + k = _triple(kernel, "kernel") + taps = k[0] * k[1] * k[2] + tuned = _TUNED_T.get(transposed_tune_key(dtype, cin, cout, tuple(k))) + if tuned is not None: + return _fit_transposed(tuned, m, cout, taps, dtype) + return default_transposed_config(m, cin, cout, taps, dtype) + + +#: Seed tiles for a sweep, ``(BLOCK_M, BLOCK_NC, BLOCK_K, num_warps)``. Narrower +#: than the gather kernel's grid because this GEMM's K is ``Cin`` alone -- there +#: is no tap factor in it -- so a ``BLOCK_K`` above ``Cin`` is pure padding, and +#: because ``TAP_BLOCK`` multiplies the column count on top of ``BLOCK_NC``. +_SEED_TILES: tuple[tuple[int, int, int, int], ...] = ( + (64, 64, 32, 4), + (64, 64, 64, 4), + (128, 64, 64, 4), + (256, 64, 64, 8), + (64, 128, 64, 4), + (128, 128, 64, 8), + (64, 64, 128, 4), + (128, 64, 128, 8), + (64, 128, 128, 8), + (32, 64, 64, 4), + (32, 128, 64, 4), +) + + +def candidate_transposed_configs( + m: int, cin: int, cout: int, taps: int, + dtype: torch.dtype = torch.bfloat16, + *, tap_blocks: Sequence[int] = (1, 2, 4, 8), nonkdims: Sequence[int] = (16, 32), +) -> list[TransposedConfig]: + """Configs worth timing for one transposed problem, pruned to legal ones.""" + n2 = max(16, triton.next_power_of_2(cout)) + k2 = max(16, triton.next_power_of_2(cin)) + m2 = max(16, triton.next_power_of_2(m)) + out: list[TransposedConfig] = [] + seen: set[TransposedConfig] = set() + for bm, bnc, bk, seed_warps in _SEED_TILES: + if bm > 2 * m2 or bnc > 2 * n2 or bk > k2: + continue + for tb in tap_blocks: + if taps % tb: + continue + for warps in {4, 8, seed_warps}: + for nonkdim in nonkdims: + cfg = TransposedConfig( + BLOCK_M=bm, BLOCK_N=tb * bnc, BLOCK_K=bk, GROUP_M=6, + num_warps=warps, num_stages=2, + matrix_instr_nonkdim=nonkdim, + kpack=1 if bk <= 16 else 2, TAP_BLOCK=tb, + ) + if (cfg.validate(dtype) is not None + or cfg.lds_bytes(dtype) > _LDS_BYTES + or cfg in seen): + continue + seen.add(cfg) + out.append(cfg) + if not out: + out.append(default_transposed_config(m, cin, cout, taps, dtype)) + return out + + +# --------------------------------------------------------------------------- +# Host side +# --------------------------------------------------------------------------- + + +_W_N_CONTIG = 0 +_W_GENERAL = 1 + + +def to_tkn(w: torch.Tensor) -> torch.Tensor: + """A ``(Cin, Cout, kd, kh, kw)`` transposed weight as ``(kd, kh, kw, Cin, Cout)``. + + The B tile wants ``[tap][K=Cin][N=Cout]`` with N dense, which is what this + produces. **It is off the shipped path**: a ``channels_last_3d`` parameter + -- which is what ``worker.py`` makes every 5-D parameter -- already has + ``Cout`` unit-stride and the three kernel axes fused, so + :func:`_transposed_weight_plan` addresses it in place and this copy never + runs. Note that this is the opposite of the ordinary forward's situation, + where the same layout puts the *reduction* axis in the dense slot and the + tile has to be gathered. + + Kept for the layouts the plan refuses -- chiefly PyTorch's default, where + neither channel axis is unit-stride and every element of the B tile is its + own cache line. + """ + return w.permute(2, 3, 4, 0, 1).contiguous() + + +def _transposed_weight_plan(w: torch.Tensor) -> tuple[int, int, int, int] | None: + """``(W_ORDER, stride_wt, stride_wk, stride_wn)`` for ``w``, or ``None``. + + ``w`` is the weight as PyTorch stores it for ``ConvTranspose3d``: + ``(Cin, Cout, kd, kh, kw)``, i.e. dim 0 is this GEMM's reduction axis and + dim 1 is its N. That is the transpose of the ordinary convolution's + convention, which is why this cannot simply call + :func:`~triton_conv3d.gather_gemm._weight_plan`; everything else about it is + the same computation, including why it is a stride test rather than a + ``is_contiguous(memory_format=...)`` one. + + ``None`` means materialize :func:`to_tkn` instead, for one of two reasons: + the three kernel axes are not one fused axis of constant stride, which is + what the kernel's single ``tap * stride_wt`` assumes; or neither channel + axis is unit-stride, which is correctness-neutral and a large performance + cliff, since every element of the B tile is then its own cache line. + + Extents of 1 carry no observable stride, so they constrain nothing and are + skipped -- the same reason the gather kernel's plan skips them. + """ + cin, cout, kd, kh, kw = (int(v) for v in w.shape) + s = tuple(int(v) for v in w.stride()) + if kw > 1: + st = s[4] + elif kh > 1: + st = s[3] + elif kd > 1: + st = s[2] + else: + st = 0 # one tap: ``tap`` is always 0, so any stride is the right one + if ((kw > 1 and s[4] != st) + or (kh > 1 and s[3] != st * kw) + or (kd > 1 and s[2] != st * kw * kh)): + return None + if cout == 1 or s[1] == 1: + return (_W_N_CONTIG, st, s[0], 1) + if cin == 1 or s[0] == 1: + return (_W_GENERAL, st, s[0], s[1]) + return None + + +def _transposed_out_spatial( + in_spatial: Sequence[int], kernel: tuple[int, int, int] +) -> tuple[int, int, int]: + """``k * extent`` per axis -- the only output shape this module produces. + + Spelled from ``kernel`` rather than from PyTorch's general + ``(i-1)*s - 2p + d*(k-1) + 1 + output_padding`` because the gate has already + pinned ``s == k``, ``p == 0``, ``d == 1`` and ``output_padding == 0``, at + which point that formula collapses to exactly this. Writing the general one + here would suggest the module served the general case. + """ + return tuple(int(i) * k for i, k in zip(in_spatial, kernel)) # type: ignore[return-value] + + +def _transposed_shape_ok( + x: torch.Tensor, w: torch.Tensor, stride, padding, output_padding, dilation, + groups: int, +) -> tuple[int, int, int] | None: + """The kernel triple if this is a ``kernel == stride`` upsample, else ``None``. + + Total, like the gates that call it: an argument it cannot interpret is a + ``None`` and never an exception, because these are the predicates of a + Triton -> MIOpen rung ladder and a caller asking a question must not be + taken down by the answer. + """ + if groups != 1: + return None + if x.dim() != 5 or w.dim() != 5: + return None + try: + s = _triple(stride, "stride") + p = _triple(padding, "padding") + op = _triple(output_padding, "output_padding") + d = _triple(dilation, "dilation") + except (ValueError, TypeError): + return None + k = tuple(int(v) for v in w.shape[2:]) + # The whole of this module's mathematics: windows that tile rather than + # overlap. ``k != s`` overlaps (or leaves gaps), a padding crops the + # result, an ``output_padding`` extends it asymmetrically and a dilation + # interleaves the window with holes -- each one breaks the bijection + # ``(d, kd) -> d*k + kd`` that makes every output voxel a single + # contribution, and none of them occurs in ScaFFold. + if s != k or p != (0, 0, 0) or op != (0, 0, 0) or d != (1, 1, 1): + return None + if any(v < 1 for v in k): + return None + # Degenerate extents, refused for the same reason ``is_supported`` refuses + # them: each clears every other test here and then disagrees with torch. A + # zero-length spatial axis produces an empty output where the M-unravel has + # no rows to index; ``Cin = 0`` returns ``Cout`` channels of zeros where + # torch returns a tensor with no channels at all. + if any(int(v) < 1 for v in x.shape[2:]): + return None + if int(w.shape[0]) < 1 or int(w.shape[1]) < 1: + return None + return k + + +def is_supported_transposed( + x: torch.Tensor, + w: torch.Tensor, + bias: torch.Tensor | None = None, + stride=1, + padding=0, + output_padding=0, + dilation=1, + groups: int = 1, +) -> bool: + """Whether :func:`conv_transpose3d_forward` will serve this call. + + Deliberately conservative and **total**, for the reasons + :func:`~triton_conv3d.gather_gemm.is_supported` gives: the caller's fallback + is MIOpen, which is correct everywhere, and an argument this cannot + interpret has to be a ``False`` rather than an exception. + + ``w`` is PyTorch's ``ConvTranspose3d`` weight, ``(Cin, Cout, kd, kh, kw)`` + -- the channel axes the other way round from ``nn.Conv3d``'s. That is not a + detail: passing a ``Conv3d`` weight here would be accepted whenever the two + channel counts happen to match and would compute a transposed answer. + + **This gates the forward alone.** Unlike the ordinary convolution, all three + of this operator's directions accept exactly the same problems -- both + backward directions are the *same* ``k == s`` convolution seen from the + other side -- so :func:`is_supported_transposed_all` should agree with this + on every input. It exists anyway, and asks all three for real, because + "should agree" is an argument and the ladder needs a fact: the three gates + of the ordinary convolution were also expected to agree until ``stride > 1`` + showed that they do not. + """ + k = _transposed_shape_ok(x, w, stride, padding, output_padding, dilation, + groups) + if k is None: + return False + if x.dtype != w.dtype or x.dtype not in _MFMA_KDIM: + return False + # Same device, not merely both on *a* device. Triton launches on the current + # device and dereferences the other pointer anyway; ScaFFold runs four GPUs + # per node, where peer access turns that into another rank's data rather + # than a fault. + if not x.is_cuda or not w.is_cuda or w.device != x.device: + return False + if int(x.shape[1]) != int(w.shape[0]): + return False + if bias is not None: + # The kernel masks the bias load against ``Cout`` -- which says nothing + # about how long the bias actually is -- and indexes it with an element + # stride of 1. A short bias reads past the end and a stride-2 view of + # the right length silently applies every other value. + # ``torch.conv_transpose3d`` rejects both; so does this. ``Cout`` is + # ``w.shape[1]`` here, not ``w.shape[0]``. + if (bias.dim() != 1 or int(bias.shape[0]) != int(w.shape[1]) + or bias.dtype != x.dtype or not bias.is_cuda + or bias.device != x.device or bias.stride(0) != 1): + return False + return True + + +def is_supported_transposed_bwd_data( + grad_output: torch.Tensor, + w: torch.Tensor, + input_shape: Sequence[int], + stride=1, + padding=0, + output_padding=0, + dilation=1, + groups: int = 1, +) -> bool: + """Whether :func:`conv_transpose3d_backward_data` will serve this call. + + Asks the *ordinary* forward's gate about the strided convolution this + direction actually is -- ``conv3d(grad_output, w, stride=k)`` -- rather than + re-deriving a predicate, so the two can never drift apart. The extra checks + on top are the ones that gate cannot see: that the problem is a + ``kernel == stride`` upsample at all, and that ``input_shape`` is the shape + this ``grad_output`` came from. + """ + k = _transposed_shape_ok(grad_output, w, stride, padding, output_padding, + dilation, groups) + if k is None: + return False + try: + shape = tuple(int(v) for v in input_shape) + except (TypeError, ValueError): + return False + if len(shape) != 5: + return False + n, cin, *in_sp = shape + if n < 1 or any(v < 1 for v in in_sp): + return False + if int(w.shape[0]) != cin or int(grad_output.shape[1]) != int(w.shape[1]): + return False + if int(grad_output.shape[0]) != n: + return False + if tuple(int(v) for v in grad_output.shape[2:]) != _transposed_out_spatial( + in_sp, k + ): + return False + # The effective convolution, asked of the gate that will actually serve it. + return bool(_is_supported_fwd(grad_output, w, None, k, 0, 1, 1)) + + +def is_supported_transposed_bwd_weight( + x: torch.Tensor, + weight_shape: Sequence[int], + grad_output: torch.Tensor, + stride=1, + padding=0, + output_padding=0, + dilation=1, + groups: int = 1, +) -> bool: + """Whether :func:`conv_transpose3d_backward_weight` will serve this call. + + As with backward-data, this asks the incumbent gate about the problem that + will really run -- ``conv3d_backward_weight`` on the strided convolution, + with ``grad_output`` in the input slot and ``x`` in the gradient slot -- and + adds only what that gate cannot see. + """ + try: + ws = tuple(int(v) for v in weight_shape) + except (TypeError, ValueError): + return False + if len(ws) != 5 or groups != 1: + return False + if x.dim() != 5 or grad_output.dim() != 5: + return False + try: + s = _triple(stride, "stride") + p = _triple(padding, "padding") + op = _triple(output_padding, "output_padding") + d = _triple(dilation, "dilation") + except (ValueError, TypeError): + return False + k = ws[2:] + if s != k or p != (0, 0, 0) or op != (0, 0, 0) or d != (1, 1, 1): + return False + if any(v < 1 for v in k) or ws[0] < 1 or ws[1] < 1: + return False + if int(x.shape[1]) != ws[0] or int(grad_output.shape[1]) != ws[1]: + return False + if any(int(v) < 1 for v in x.shape[2:]): + return False + return bool( + is_supported_bwd_weight(grad_output, (ws[0], ws[1], *k), x, k, 0, 1, 1) + ) + + +def is_supported_transposed_all( + x: torch.Tensor, + w: torch.Tensor, + bias: torch.Tensor | None = None, + stride=1, + padding=0, + output_padding=0, + dilation=1, + groups: int = 1, +) -> bool: + """Whether **every** direction of this transposed convolution will be served. + + The gate for a caller that is going to differentiate, and the counterpart of + :func:`~triton_conv3d.gather_gemm.is_supported_all`. A forward this package + serves and a backward it cannot is discovered inside ``backward()``, where + the caller's fallback kernel is no longer reachable -- so a training caller + must ask this one. + + The gradient is passed as a **metadata-only stand-in**: all three predicates + read rank, shape, dtype, device and ``is_cuda`` and never a stride, a value + or a contiguity, so a one-element allocation expanded to the output shape + answers exactly as the real gradient would. ``expand`` gives every dim a + stride of 0, so if a predicate ever grows a stride test it will see those + zeros and answer ``False`` -- a fallback to the caller's other kernel, which + is the safe direction. + """ + if not is_supported_transposed(x, w, bias, stride, padding, output_padding, + dilation, groups): + return False + k = tuple(int(v) for v in w.shape[2:]) + grad_shape = (int(x.shape[0]), int(w.shape[1])) + _transposed_out_spatial( + tuple(int(v) for v in x.shape[2:]), k + ) + grad = x.new_empty((1, 1, 1, 1, 1)).expand(grad_shape) + if not is_supported_transposed_bwd_data( + grad, w, tuple(x.shape), stride, padding, output_padding, dilation, groups + ): + return False + return bool( + is_supported_transposed_bwd_weight( + x, tuple(w.shape), grad, stride, padding, output_padding, dilation, + groups, + ) + ) + + +def conv_transpose3d_forward( + x: torch.Tensor, + w: torch.Tensor, + bias: torch.Tensor | None = None, + stride=1, + padding=0, + output_padding=0, + dilation=1, + groups: int = 1, + *, + config: TransposedConfig | None = None, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Transposed 3-D convolution at ``kernel == stride``. NDHWC in and out. + + ``w`` is PyTorch's ``ConvTranspose3d`` weight, ``(Cin, Cout, kd, kh, kw)``, + and is read where it lies: a ``channels_last_3d`` parameter has ``Cout`` + unit-stride, which is this GEMM's N, so no transform runs. A weight in + PyTorch's *default* layout is copied and has to be; see + :func:`_transposed_weight_plan`. + + ``out=`` is checked rather than trusted, for the reason + :func:`~triton_conv3d.gather_gemm._check_out` gives: the store addressing is + derived from *this* call's shapes, so a mismatched buffer is an + out-of-bounds device write with no error and an NCDHW one is a full-rate + kernel returning a scrambled answer. + """ + if not is_supported_transposed(x, w, bias, stride, padding, output_padding, + dilation, groups): + raise NotImplementedError( + f"unsupported: x={tuple(x.shape)}/{x.dtype} w={tuple(w.shape)} " + f"stride={stride} padding={padding} output_padding={output_padding} " + f"dilation={dilation} groups={groups}" + ) + kd, kh, kw = (int(v) for v in w.shape[2:]) + taps = kd * kh * kw + + # NDHWC is not a preference here, it is the layout the addressing assumes. + x = x.contiguous(memory_format=torch.channels_last_3d) + n, cin, in_d, in_h, in_w = (int(v) for v in x.shape) + cout = int(w.shape[1]) + out_d, out_h, out_w = _transposed_out_spatial((in_d, in_h, in_w), + (kd, kh, kw)) + + y_shape = (n, cout, out_d, out_h, out_w) + if out is None: + # One allocation, already in the layout the kernel stores into. Spelling + # it ``torch.empty(shape).contiguous(memory_format=...)`` allocates NCDHW + # and then copies the whole thing -- 235x, measured on the gather + # kernel's identically-shaped defect. + y = torch.empty(y_shape, device=x.device, dtype=x.dtype, + memory_format=torch.channels_last_3d) + else: + y = out + _check_out(y, y_shape, x) + + plan = _transposed_weight_plan(w) + if plan is None: + # The only path that copies the weight; see the docstring of + # :func:`to_tkn`. Contiguous ``(kd, kh, kw, Cin, Cout)``, so the three + # strides are exactly these. + wt = to_tkn(w) + plan = (_W_N_CONTIG, cin * cout, cout, 1) + else: + wt = w + + m_total = n * in_d * in_h * in_w + if config is None: + config = transposed_config(m_total, cin, cout, (kd, kh, kw), x.dtype) + why = config.validate(x.dtype) + if why is not None: + raise ValueError(f"illegal config {config}: {why}") + if taps % config.TAP_BLOCK: + raise ValueError( + f"illegal config {config}: TAP_BLOCK must divide the tap count " + f"{taps}; a ragged last group would address a tap that is not there" + ) + + index_dtype = _index_dtype(x, y, wt) + grid = ( + triton.cdiv(m_total, config.BLOCK_M) + * triton.cdiv(cout, config.BLOCK_NC) + * (taps // config.TAP_BLOCK), + ) + _convT3d_fwd_kernel[grid]( + x, wt, y, bias, + n, in_d, in_h, in_w, + cin, cout, m_total, + x.stride(0), x.stride(2), x.stride(3), x.stride(4), + plan[1], plan[2], plan[3], + y.stride(0), y.stride(2), y.stride(3), y.stride(4), + KD=kd, KH=kh, KW=kw, + BLOCK_M=config.BLOCK_M, BLOCK_N=config.BLOCK_N, + BLOCK_NC=config.BLOCK_NC, BLOCK_K=config.BLOCK_K, + BLOCK_K_COUNT=triton.cdiv(cin, config.BLOCK_K), + TAP_BLOCK=config.TAP_BLOCK, GROUP_M=config.GROUP_M, + HAS_BIAS=bias is not None, + EVEN_K=(cin % config.BLOCK_K == 0), + EVEN_N=(cout % config.BLOCK_NC == 0), + INDEX_DTYPE=index_dtype, + INPUT_PRECISION="ieee", + W_ORDER=plan[0], + **config.launch_kwargs(), + ) + return y + + +def conv_transpose3d_backward_data( + grad_output: torch.Tensor, + w: torch.Tensor, + input_shape: Sequence[int], + stride=1, + padding=0, + output_padding=0, + dilation=1, + groups: int = 1, + *, + config: ConvConfig | None = None, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Gradient of a ``k == s`` transposed convolution with respect to its input. + + **This is an ordinary strided forward convolution**, and the module + docstring derives it: ``grad_input = conv3d(grad_output, w, stride=k)``, + with ``w`` passed *unpermuted*. PyTorch stores a ``ConvTranspose3d`` weight + as ``(Cin, Cout, k, k, k)``, and that already is the ``(out_channels, + in_channels, k, k, k)`` an ordinary ``Cout -> Cin`` convolution wants -- the + transpose is in the storage convention, so it costs nothing here. + + No bias term: the bias is added to the forward's output, so its gradient is + a reduction of ``grad_output`` and not part of this direction at all. + """ + if not is_supported_transposed_bwd_data( + grad_output, w, input_shape, stride, padding, output_padding, dilation, + groups + ): + raise NotImplementedError( + f"unsupported: grad_output={tuple(grad_output.shape)}/" + f"{grad_output.dtype} w={tuple(w.shape)} " + f"input_shape={tuple(input_shape)} stride={stride} " + f"padding={padding} output_padding={output_padding} " + f"dilation={dilation} groups={groups}" + ) + k = tuple(int(v) for v in w.shape[2:]) + return conv3d_forward(grad_output, w, None, k, 0, 1, 1, config=config, + out=out) + + +def conv_transpose3d_backward_weight( + x: torch.Tensor, + weight_shape: Sequence[int], + grad_output: torch.Tensor, + stride=1, + padding=0, + output_padding=0, + dilation=1, + groups: int = 1, + *, + config=None, + workspace: torch.Tensor | None = None, + out: torch.Tensor | None = None, + deterministic: bool = True, +) -> torch.Tensor: + """Gradient of a ``k == s`` transposed convolution with respect to its weight. + + The *same* reduction ``conv3d_backward_weight`` already performs, with the + two activations in the slots the strided convolution of the module docstring + puts them in: ``grad_output`` is that convolution's input and ``x`` is its + output gradient. Reading the call and expecting ``x`` first is the one way + to misuse this function, which is why the argument order still matches + ``conv3d_backward_weight``'s -- the swap happens inside, once, here. + + The returned gradient is ``(Cin, Cout, k, k, k)`` in ``channels_last_3d``, + which is both the parameter's own shape and the layout ``worker.py`` puts it + in, so the optimizer's elementwise update is contiguous. + """ + if not is_supported_transposed_bwd_weight( + x, weight_shape, grad_output, stride, padding, output_padding, dilation, + groups + ): + raise NotImplementedError( + f"unsupported: x={tuple(x.shape)}/{x.dtype} " + f"weight_shape={tuple(weight_shape)} " + f"grad_output={tuple(grad_output.shape)} stride={stride} " + f"padding={padding} output_padding={output_padding} " + f"dilation={dilation} groups={groups}" + ) + ws = tuple(int(v) for v in weight_shape) + k = ws[2:] + return conv3d_backward_weight( + grad_output, ws, x, k, 0, 1, 1, + config=config, workspace=workspace, out=out, + deterministic=deterministic, + ) + + +def grad_transposed_weight_empty( + cin: int, cout: int, kernel: Sequence[int], *, dtype, device +) -> torch.Tensor: + """An empty transposed-weight gradient in the layout the kernel writes. + + ``(Cin, Cout, kd, kh, kw)`` in ``channels_last_3d``, i.e. memory order + ``Cin, kd, kh, kw, Cout``. The ``Cin``/``Cout`` order is the only thing + that differs from + :func:`~triton_conv3d.reduce_gemm.grad_weight_empty`, and it differs because + a ``ConvTranspose3d`` parameter is stored the other way round; passing the + ordinary one here allocates a correctly-strided buffer of the wrong shape, + which ``conv3d_backward_weight``'s ``out=`` check catches. + """ + k = _triple(kernel, "kernel") + return torch.empty((cin, cout, *k), dtype=dtype, device=device, + memory_format=torch.channels_last_3d) + + +# --------------------------------------------------------------------------- +# ISA verification +# --------------------------------------------------------------------------- + + +def verify_isa_transposed( + problem_shape: Sequence[int] | None = None, + config: "TransposedConfig | None" = None, + kernel: int = 2, + weight_layout: str = "channels_last", +) -> None: # pragma: no cover + """Compile and launch one configuration so its ISA can be inspected. + + Run under ``AMDGCN_ENABLE_DUMP=1`` with a **cold** ``TRITON_CACHE_DIR``: a + cache hit skips the compile and therefore the dump, and an empty grep then + looks exactly like a kernel with no MFMA in it. The emitted mnemonic is + ``v_mfma_f32_16x16x16_bf16`` with **no** ``_1k`` suffix, despite Triton's + internal table entry being named ``_1k``. + + ``weight_layout`` selects which of the two B loads is compiled, for the same + reason the gather kernel's does: ``W_ORDER`` is a constexpr and the two + orders emit different instructions for the operand that feeds the matrix + core. + """ + n, cin, cout, d, h, wd = problem_shape or (1, 128, 64, 64, 64, 64) + k = (kernel, kernel, kernel) + w = torch.randn((cin, cout, *k), device="cuda", dtype=torch.bfloat16) + if weight_layout == "channels_last": + w = w.contiguous(memory_format=torch.channels_last_3d) + elif weight_layout != "tkn": + raise ValueError(f"unknown weight_layout {weight_layout!r}") + x = torch.randn((n, cin, d, h, wd), device="cuda", + dtype=torch.bfloat16).contiguous( + memory_format=torch.channels_last_3d) + cfg = config or transposed_config(n * d * h * wd, cin, cout, k, + torch.bfloat16) + y = conv_transpose3d_forward(x, w, None, k, config=cfg) + torch.cuda.synchronize() + print( + f"ISA-DUMP-CONFIG [convT/{weight_layout}] {cfg} cin={cin} cout={cout} " + f"spatial={(d, h, wd)} k={kernel} " + f"x_storage={x.untyped_storage().size()} " + f"y_storage={y.untyped_storage().size()}" + ) From fcb696499b61c33a53d4fdb53712c2a0c839ed25 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Wed, 5 Aug 2026 06:42:05 -0700 Subject: [PATCH 4/8] Route every convolution through the Triton rung FastConv3d and FastConvTranspose3d mirror FastGroupNorm: drop-in nn.Conv3d and nn.ConvTranspose3d, same parameters under the same names, no buffers, a rung ladder sharing _rungs.py, and MIOpen underneath everything the kernel declines. The transposed operator gets a factory of its own rather than a flag, because it is a different operator -- the weight's channel axes are the other way round and a different set of kernels sits behind it. The adapter performs the halo exchange itself, above autograd, rather than leaving it to the one DistConv does below. That has a consequence for the shape the kernel sees, and it is the thing to know when reading any number from this work: only the split axis is halo'd, so padding=1 survives on the other two and every k=3 convolution in the network is padded at every configuration -- unsharded there is nothing to halo at all. The benchmark corpus calls that the adapter form, and it is what production issues. --- ScaFFold/unet/conv3d.py | 1871 +++++++++++++++++++++++++++++++++++ ScaFFold/unet/unet_parts.py | 112 ++- tests/test_conv3d.py | 1620 ++++++++++++++++++++++++++++++ tests/test_unet.py | 141 +++ 4 files changed, 3738 insertions(+), 6 deletions(-) create mode 100644 ScaFFold/unet/conv3d.py create mode 100644 tests/test_conv3d.py diff --git a/ScaFFold/unet/conv3d.py b/ScaFFold/unet/conv3d.py new file mode 100644 index 0000000..83dd933 --- /dev/null +++ b/ScaFFold/unet/conv3d.py @@ -0,0 +1,1871 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""3-D convolution with a Triton fast path and MIOpen behind it. + +Two kernels, tried in order: + +1. **Native channels-last Triton** (:mod:`triton_conv3d`, a self-contained + package that imports nothing from ScaFFold), whenever its ``is_supported_all`` + gate -- all three direction predicates, which do not accept the same problems + -- accepts the call. It is an implicit-GEMM forward, + a re-strided view of that forward for backward-data, and a deterministic + split-K reduction for backward-weight; it reads a ``channels_last_3d`` weight + where it lies, which is the layout ``worker.py`` already puts every + convolution parameter in. Measured at **1.26x of the MIOpen step** at + ScaFFold's scale-7 single-GPU configuration, with the backward-weight + direction reproducible at no cost -- see ``work/triton-conv/STATUS.md``. +2. **MIOpen**, via ``nn.Conv3d.forward``, which is what every rejection falls + back to and what defines the semantics the Triton rung must match. + +``FastConv3d`` is a drop-in ``nn.Conv3d``: same parameters, same names, same +shapes, no buffers and no ``state_dict`` keys of its own, so checkpoints are +interchangeable in both directions with any other ``nn.Conv3d``-based build. + +Two ladders +=========== +:class:`FastConvTranspose3d` is the same construction for the decoder's four +``nn.ConvTranspose3d(k=2, s=2)`` upsamplers, which ``triton_conv3d.transposed`` +serves in all three directions. It shares with :class:`FastConv3d` everything +that is genuinely one mechanism -- the failure allowlist, the latch, the +autocast reproduction, the ``DCTensor`` unwrap, the routing conditions that are +not about which operator this is (:func:`_routing_declines`) -- and shares +nothing else, because the rest reads different numbers off different tensors: + +* the two parameters store their channel axes in opposite orders, so a + block-list or a gate written for one **accepts** the other whenever the two + channel counts happen to match, and computes a different operator quietly; +* the transposed forward's GEMM has one row per *input* voxel where the + ordinary one has one per output voxel -- 8x apart at ``k = s = 2``; +* and the halo, which is the whole difficulty below, does not exist for + ``kernel == stride``: every output voxel reads exactly one input voxel. So + that ladder has no :class:`_Halo3d` in it, and :func:`_transposed_halo_plan` + is where "there is nothing to exchange" is *checked* -- against DistConv's own + ``halo = k // 2``, which is 0 for an even kernel -- rather than assumed. + +Sharding is the whole difficulty +================================ +GroupNorm's fast path unwraps DistConv's ``DCTensor`` to its local shard, runs, +and rewraps. That is semantically inert *for GroupNorm*, whose statistics are +per-shard at every shard count. **It is not inert for convolution.** A +convolution needs its neighbours' boundary voxels, and DistConv supplies them +*below* autograd, inside its interception of ``aten.convolution.default`` +(``distconv.py``, ``distconv_forward``): it concatenates a halo slab of width +``k // 2`` onto both faces of every sharded dimension and then zeroes that +dimension's padding. A module-level adapter that unwraps and calls a kernel +directly never reaches that code, so it would silently drop the halo and return +a wrong answer at every shard boundary -- 55% relative error, measured, so any +defect here is loud once it is looked for. + +So this module performs the exchange itself, *above* autograd, as its own +``autograd.Function`` (:class:`_Halo3d`) stacked over the convolution's +(:class:`_TritonConv3dFn`). The forward concatenates the neighbours' boundary +slabs onto the local shard and hands the kernel a padding with those dims +zeroed; the backward sends the outer boundary gradient slabs back and +accumulates what arrives into the inner region, which is +``distconv.backward_halo_exchange`` exactly. ``grad_weight`` is computed +against the *halo'd* input, as ``distconv_backward`` does, so each rank counts +every filter tap that straddles its boundary once and ``DistConvDDP``'s +``world_size / ddp_ranks`` rescaling turns DDP's average back into the sum the +spatial shards need. + +**Only the dims that are actually split are exchanged.** ScaFFold ships +``dc_shard_dims: [2, 3, 4]`` with ``dc_num_shards`` of ``[1,1,1]``, ``[2,1,1]`` +or ``[4,1,1]``, so D is the only axis ever divided and DistConv's halo on H and +W is two ``cat`` copies of a slab that is provably zeros. ``cat(zeros, x, +zeros)`` at ``padding = 0`` is the same arithmetic as ``padding = k // 2`` on +``x`` -- and *measured bitwise identical* through these kernels in fp32 and +bf16 at 2 and 4 shards, which is what makes dropping it a decision with no +tolerance argument in it. At ``dc_num_shards = (1, 1, 1)`` every dim is unsplit, +:class:`_Halo3d` is not applied at all, and the rung is stage 1's exactly: +``forward_halo_exchange`` allocates its receive buffers with ``zeros_like`` and +posts nothing when ``shard_ind`` is 0 on every axis, so the 3.797 ms/step of +zero-slab ``cat`` copies ``work/upstream-repros/DISTCONV_ACTIONS.md`` prices at one shard just +disappear. + +Every fact that argument rests on is checked *positively* before the rung is +taken -- see :func:`_halo_plan`, which returns ``None`` for anything it could not +read. "I could not find evidence of a problem" is not the same statement as "I +have checked that there is none", and only the second one is safe: the failure +mode of getting this wrong is a plausible-looking wrong gradient at scale, not a +crash. + +The exchange sends **plain-contiguous** buffers, which +``forward_halo_exchange`` does not: it sends +``.contiguous(memory_format=channels_last_3d)``, and a ``channels_last_3d`` +tensor is not plain-contiguous for ``C > 1``. That, not the narrow, is why +DistConv's spatial sharding does not run under gloo. Packing NDHWC into a +plain-contiguous slab costs one copy of 1.6% of the activation and makes the +sharded path runnable on a machine with no RCCL -- see :func:`_packed_slab`. + +Autocast +======== +ScaFFold trains inside ``torch.autocast(device_type="cuda", +dtype=torch.bfloat16)`` (``trainer.py``'s ``_autocast_kwargs``, on by default via +``torch_amp: 1``). Autocast's cast for ``aten::convolution`` -- and for +``aten::conv_transpose3d``, which carries the same ``lower_precision_fp`` policy +(measured: an fp32 activation and an fp32 parameter produce a bf16 output, and +an fp64 pair is left alone) -- happens *in the +dispatcher*, so an adapter that calls a kernel directly bypasses it: the +convolution's operands arrive as the fp32 tensors GroupNorm produced and the +fp32 parameters the model holds, and running them as-is would silently execute +the whole network's convolutions in fp32 -- a different computation from the one +the benchmark defines, several times slower, and nothing fails. +:func:`_autocast_dtype` and :func:`_cast_operand` reproduce ATen's rule +(``cached_cast``: cast a floating tensor to the autocast dtype *only* if it is +exactly fp32), outside the autograd node so the cast's backward returns the fp32 +parameter gradient exactly as autocast's own does. +``unet_parts._consumer_dtype`` is the same reasoning applied to ``torch.cat``'s +``promote`` policy. + +The cast is applied *before* the halo exchange, which is both what DistConv does +-- its ``__torch_dispatch__`` runs below the autocast key, so the tensor it +concatenates is already bf16 -- and what halves the bytes on the wire. + +Hardware +======== +The rung is taken only on the GPU the kernels were tuned on -- ``gfx942`` with +228 CUs, i.e. an MI300A -- and declines quietly to MIOpen anywhere else. That +guard is not about capability: ``triton_conv3d``'s kernels compute the right +convolution wherever Triton lowers them, and its ``is_supported*`` predicates +are right to say nothing about the device. It is about the fact that *every* +number deciding how they launch was raced on one machine, and that a launch +configuration which is merely wrong for the hardware raises nothing at all -- +``gather_gemm``'s docstring records the concrete mechanism, an illegal MFMA +configuration that emits zero MFMA instructions and returns correct results at a +fraction of the speed. A ladder whose fallback allowlist is +``triton.errors.TritonError`` has nothing to catch there. So the check is a +routing decision and lives with the others, in +:func:`~ScaFFold.unet._rungs._platform_declines`; the decision it encodes, and +why the explicit opt-in overrides it while none of the correctness conditions +can be overridden at all, is argued there. + +Latches +======= +The Triton rung is an optimization, never a correctness requirement: a broken +Triton install must degrade a multi-node run, not kill it. So a *kernel* +failure is caught, logged once and retried on MIOpen. + +The retry has two shapes, and which one applies is decided by whether this +rank has already put a halo slab on the wire. With nothing sent -- the +unsharded case, and every call in a one-shard run -- the whole call is re-run +from the top on ``_miopen_forward``, which is also the only rung that can be +handed a ``DCTensor``. Once a dim really is split, that route is closed: it +goes through ``distconv_forward``, which would exchange a *second* time, giving +this rank one more collective than a peer whose kernel compiled and hanging the +mesh (or cross-pairing two convolutions' slabs) instead of merely slowing down. +So the exchange is performed once, *above* the retried region, and only the +kernel call is retried: MIOpen is handed the already-exchanged tensor at +``plan.padding`` and swaps nothing but the kernel. Otherwise a Triton compile +failure would be **fatal** at ``num_shards > 1`` while costing only speed at 1, +which is exactly backwards for a ladder whose reason to exist is degrading +instead of dying. + +"A kernel failure" is an allowlist, not the absence of one: exactly +``triton.errors.TritonError`` (see :func:`_triton_kernel_failures`). +``triton_conv3d`` has no exception type of its own -- unlike +``triton_group_norm``, it does not tag its launch region -- so this allowlist is +drawn at Triton's boundary instead of at the package's. It is nonetheless +closed and small: that one root covers ``OutOfResources``, ``CompilationError``, +``CompileTimeAssertionFailure``, ``UnsupportedLanguageConstruct``, ``PTXASError`` +and ``AutotunerError``, every one of which is raised while compiling or sizing a +launch, i.e. before any device work, which is what makes the retry safe. +Everything else propagates, and each exclusion is load-bearing: + +* ``ValueError`` -- every one the package raises is a caller-contract violation + (``_check_out``, ``_check_weight_rsck``, ``_triple``, ``ConvConfig.validate``). + Catching them would turn the one bug class that produces silently wrong + numbers into a quiet performance regression. +* ``NotImplementedError`` -- raised by all three entry points when their own + ``is_supported*`` says no. Unreachable because this module branches on those + predicates first; if it fires, the predicate and the entry point disagree and + that has to be seen. +* ``torch.OutOfMemoryError`` and ``torch.AcceleratorError`` -- both + ``RuntimeError`` subclasses. The first is a resource condition, and MIOpen + needs *more* memory than the Triton path at the shapes where it bites; the + second means the HIP context is already poisoned and no fallback can succeed. +* bare ``RuntimeError`` -- excluded precisely because the three above are all + subclasses of it. + +Not covered by any allowlist, and worth stating: a kernel that stores out of +bounds takes the process down with a HSA "Memory access fault" and SIGABRT +without raising anything Python can see (measured on this MI300A). Robustness +against that lives in ``triton_conv3d``'s own argument checks, not here. + +A failure latches the rung off **for modules that have never had a call served +by it**; a module that has already run on it keeps it. That is not a +performance nicety. ``torch.utils.checkpoint``'s non-reentrant recompute +substitutes recomputed tensors positionally into the original graph node, and +under DistConv the two rungs save tensors that are *metadata-identical* -- +same shape, dtype and device -- differing only in whether slot 0 holds the +``DCTensor`` wrapper (MIOpen, which saves below DistConv's dispatch) or its +inner tensor (Triton, which saves above it). ``_default_meta_extractor`` +compares shape, dtype and device and nothing else, so a rung flip across a +recompute passes torch's own check and fails later inside DistConv with an +``AttributeError`` about ``_parallel_strategy`` or ``_is_periodic`` (measured, +both directions). Once anything is actually sharded the Triton rung saves the +*halo'd* input, whose extent differs, so a flip there is caught by torch's own +metadata check as a plain ``CheckpointError`` (measured) -- but the unsharded +configuration is the shipped one and it is the silent case, so the latch is what +the correctness argument rests on either way. Pinning each module instance's +choice for the life of the +process is what makes forward and recompute agree; :func:`_replaying_a_forward` +bounds the *fallback* by the same argument, and +:meth:`_TritonConv3dFn.backward` carries a cheap type check that turns the +``AttributeError`` into an actionable message if it ever does happen. + +A latch is process-local, so under DDP one rank can end up on a different kernel +from its peers. The two rungs agree to fp32 rounding, not bitwise, so a rank +that latches shifts its gradients and therefore the all-reduced ones. That is +the price of degrading instead of dying, and it is why the latch is as narrow as +it is and why ``torch.OutOfMemoryError`` latches nothing. + +Determinism +=========== +``conv3d_backward_weight`` reduces its split-K partials in fp32 and stores once; +its ``deterministic=True`` default is both reproducible and *faster* than the +atomic path, so nothing here plumbs ``more_determinism`` into the kernel choice. +MIOpen's backward-weight, by contrast, reduces with atomics and disagrees with +itself bitwise between two identical calls. +""" + +import logging + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from torch.autograd.function import once_differentiable + +from ._rungs import ( + _dctensor_ops, + _env_override, + _functorch_active, + _platform_declines, + _replaying_a_forward, + _run_local, + _warn_rung_failure, +) + +logger = logging.getLogger(__name__) + +#: Opt-out (``0``/``false``/``off``/``no``) or explicit opt-in (``1``/``true``/ +#: ``on``/``yes``) for the Triton convolution path. Unset means "on wherever it +#: is safe", matching ``SCAFFOLD_GROUPNORM_TRITON``: the whole point of the +#: kernel is that production takes it, and every configuration it must not serve +#: is refused by a correctness check rather than by this default. +#: +#: The explicit opt-in is *not* just that default written out: it additionally +#: overrides the hardware guard, which is the only routing condition here that +#: is a preference rather than a correctness condition. See +#: :func:`set_conv_triton_enabled`. +TRITON_ENV_VAR = "SCAFFOLD_CONV_TRITON" + +# The triton_conv3d package, imported on the first eligible forward. Its +# __init__ is lazy in turn -- the entry points live in submodules that import +# torch and triton -- so a CPU-only run pays neither. +_triton_module = None + +# The ladder's allowlist, resolved on first use of the rung it guards. +_TRITON_KERNEL_FAILURES = None + +# Set once if the Triton kernel raises; MIOpen is used from then on, except by +# modules the rung has already served. +_triton_failed = False + +# None = decide per tensor; True/False = forced by SCAFFOLD_CONV_TRITON or by +# set_conv_triton_enabled(). +_triton_override = _env_override(TRITON_ENV_VAR) + +# Set once if a predicate raised while deciding; see _use_triton. +_predicate_warned = False + +#: One-element tensors, one per ``(dtype, device)``, from which +#: :func:`_metadata_probe` expands. Bounded by the number of dtype/device pairs +#: a process actually convolves in, which for ScaFFold is one. +_PROBE_BASES = {} + + +def set_conv_triton_enabled(enabled): + """Force the Triton convolution path on (``True``) or off (``False``). + + The counterpart of ``group_norm.set_triton_enabled``: ``None`` restores the + default (``SCAFFOLD_CONV_TRITON`` if set, otherwise "wherever it is safe"), + and the previous setting is returned so tests can restore it. + + Forcing it on clears any failure latch, and overrides the **hardware + guard**, and overrides nothing else -- not the device, subclass, sharding or + ``is_supported`` checks, which are correctness conditions rather than + preferences. The hardware guard is on the other side of that line and the + difference is the failure mode: an unsupported shape or an unknown tensor + subclass would make this rung compute the *wrong thing*, while an untuned + GPU makes it compute the right thing at a speed nobody has measured (see + :func:`~ScaFFold.unet._rungs._platform_declines`). The second is a judgement + a developer is entitled to make about their own machine, so ``True`` here -- + and ``SCAFFOLD_CONV_TRITON=1``, which is the same statement spelled in the + environment -- takes the kernels on hardware the default declines, and says + so in the log. Note that this makes an explicit ``1`` mean something + stronger than leaving the variable unset, which is the one place the + tri-state is not merely "the default, written out". + + ``None`` deliberately does *not* clear the latch: it restores a preference, + it does not assert that the kernel works again. + """ + global _triton_override, _triton_failed + previous = _triton_override + _triton_override = ( + _env_override(TRITON_ENV_VAR) if enabled is None else bool(enabled) + ) + if _triton_override is True: + _triton_failed = False + return previous + + +def _get_triton_module(): + """Import (once) the :mod:`triton_conv3d` package. + + Deferred rather than imported at the top of this file for the same reason + ``group_norm`` defers its kernel module: a run that never reaches the GPU + (the whole CPU unit suite) must not pay for ``triton``. The package's own + ``__init__`` re-exports its entry points lazily, so even this import does not + pull in torch's Triton stack until a predicate is actually asked. + """ + global _triton_module + if _triton_module is None: + import triton_conv3d + + _triton_module = triton_conv3d + return _triton_module + + +def _triton_kernel_failures(): + """The ladder's allowlist: exactly ``triton.errors.TritonError``. + + That single root is the parent of ``OutOfResources``, + ``CompilationError``, ``CompileTimeAssertionFailure``, + ``UnsupportedLanguageConstruct``, ``PTXASError``, ``AutotunerError`` and + ``InterpreterError``; ``triton.runtime.errors.TritonError`` and + ``triton.compiler.errors.TritonError`` are the same object. Every one of + them is raised while compiling a kernel or sizing its launch -- a missing or + mismatched ``triton``, an unwritable JIT cache, a compile error, a tile that + does not fit in LDS -- so nothing has executed and retrying the same call on + MIOpen is safe. + + The module docstring lists what is deliberately left out and why. Resolved + on demand and cached, so a CPU-only run never imports ``triton``. An empty + tuple (no ``triton`` at all) means "catch nothing": the ladder then + re-raises, which is right, because with no Triton there is nothing that + could have failed inside one -- and the predicate would have declined the + rung long before. + """ + global _TRITON_KERNEL_FAILURES + if _TRITON_KERNEL_FAILURES is None: + try: + from triton.errors import TritonError + + _TRITON_KERNEL_FAILURES = (TritonError,) + except ImportError: # pragma: no cover - triton ships it + _TRITON_KERNEL_FAILURES = () + return _TRITON_KERNEL_FAILURES + + +def _latch_rung_failure(error, what="Triton conv3d"): + """Latch the Triton rung off, logging only on the ``False -> True`` edge. + + Both handlers that catch an allowlisted kernel failure in a *forward* need + exactly this, so it is written once: the one in :meth:`FastConv3d.forward`, + which re-runs the whole call on MIOpen, and the one in + :meth:`FastConv3d._triton_forward`, which cannot (its halo is already on the + wire) and instead hands the exchanged tensor to MIOpen itself. A module that + has already used the rung keeps trying it -- that is what pins a checkpointed + block to one rung -- so without the edge test a persistently broken kernel + would warn once per call for the rest of the run. Clearing the latch re-arms + the message. + + ``what`` names the ladder in the message and nothing else: **one latch + serves both**. Every allowlisted failure is a property of the *install* -- + a missing or mismatched ``triton``, an unwritable JIT cache, a compile error + -- and those break both ladders at once; the exception is ``OutOfResources``, + which is per-launch, and there the over-latch costs little because the + encoder's ordinary convolutions all run (and become ``proven``) before the + decoder reaches its first transposed one. + """ + global _triton_failed + first = not _triton_failed + _triton_failed = True + if first: + _warn_rung_failure(what, error, "MIOpen kernel", TRITON_ENV_VAR) + + +def _warn_once_about_the_predicate(error): + """Log the first ``is_supported*`` failure; a repeat would log per call.""" + global _predicate_warned + if _predicate_warned: + return + _predicate_warned = True + logger.warning( + f"Triton conv3d routing check failed ({type(error).__name__}: {error}); " + "using MIOpen for calls like this one. This is a routing miss, not a " + "kernel failure, so nothing is latched off." + ) + + +# --------------------------------------------------------------------------- +# Autocast +# --------------------------------------------------------------------------- + + +def _autocast_dtype(tensor): + """The dtype autocast would cast this call's operands to, or ``None``. + + ``aten::convolution`` carries the ``lower_precision_fp`` cast policy, so + inside an enabled autocast region for this device the answer is autocast's + dtype. The cast is applied by the dispatcher, *below* this module, which is + exactly why calling a kernel directly has to reproduce it -- see the module + docstring. + + ``None`` means "no cast", which is what a run with ``torch_amp: 0``, an + evaluation outside the autocast region, or a device autocast does not know + about all get. + """ + device_type = tensor.device.type + try: + if not torch.is_autocast_enabled(device_type): + return None + return torch.get_autocast_dtype(device_type) + except (RuntimeError, TypeError): # a device type autocast does not know + return None + + +def _cast_operand(tensor, dtype): + """Apply :func:`_autocast_dtype`'s answer the way ATen's ``cached_cast`` does. + + Only an exactly-fp32 tensor is cast: ATen's ``is_eligible`` requires + ``scalar_type() == kFloat``, so an operand already in the lower precision is + left alone and an fp64 one is *not* narrowed. Getting that wrong in either + direction changes which computation the benchmark performs. + + Done outside :class:`_TritonConv3dFn` so the cast is an ordinary autograd + node: the weight gradient then arrives back at the fp32 parameter through + it, exactly as it does on the MIOpen rung. + """ + if tensor is None or dtype is None or tensor.dtype is not torch.float32: + return tensor + return tensor.to(dtype) + + +def _cast_dtype(tensor, dtype): + """The dtype :func:`_cast_operand` would leave ``tensor`` in.""" + if tensor is None: + return None + if dtype is None or tensor.dtype is not torch.float32: + return tensor.dtype + return dtype + + +# --------------------------------------------------------------------------- +# Eligibility +# --------------------------------------------------------------------------- + + +def _metadata_probe(shape, dtype, device): + """A 5-D tensor with this shape, dtype and device, over one element. + + The ``is_supported*`` predicates read metadata only -- rank, shape, dtype, + device, ``is_cuda`` -- and never a stride, a value or a contiguity. That + lets the gate ask about the bf16 operands autocast will produce without + materializing them, which would cost a full-size copy and be wasted whenever + the answer is no. ``is_supported_all`` uses the same shortcut one level + down, for the *gradient* this forward will later be handed and which does not + exist yet -- a forward served by Triton whose backward cannot be is a trap, + not a fallback, so that question has to be asked here and now. + + ``expand`` gives every dimension a stride of 0, so the result is safe to + read and useless to compute with; nothing computes with it. If a predicate + ever grows a stride or contiguity test it will see those zeros and answer + ``False``, routing the call to MIOpen -- the conservative direction. + ``tests/test_conv3d.py::test_metadata_probe_answers_like_a_real_tensor`` + pins the agreement so the shortcut cannot rot silently. + """ + key = (dtype, device) + base = _PROBE_BASES.get(key) + if base is None: + base = torch.empty((1, 1, 1, 1, 1), dtype=dtype, device=device) + _PROBE_BASES[key] = base + return base.expand(tuple(int(v) for v in shape)) + + +def _out_spatial(in_spatial, kernel, stride, padding, dilation): + """PyTorch's output extents for a non-transposed convolution. + + No caller since the block-list was emptied on 2026-08-04 -- its only user was + :func:`_policy_declines`'s small-``M`` rule. Kept because a future entry + there is the reason that function still takes ``stride``/``padding``/ + ``dilation``, and because two docstrings warn about what this returns for a + *transposed* operator (the input volume over 8 at ``k == s == 2``), a warning + that needs the thing it warns about to exist. + """ + return tuple( + (i + 2 * p - d * (k - 1) - 1) // s + 1 + for i, k, s, p, d in zip(in_spatial, kernel, stride, padding, dilation) + ) + + +class _HaloPlan: + """Which dims this call exchanges, and what the kernel is left holding. + + Built once per forward by :func:`_halo_plan` and read by everything + downstream -- the ``is_supported*`` probes, the policy block-list, + :class:`_Halo3d` and :class:`_TritonConv3dFn` -- so that exactly one function + decides which dims are split and nothing re-derives it. An empty + ``exchanges`` is the unsharded case: no node is applied and the kernel sees + the module's own padding, which is stage 1 unchanged. + """ + + __slots__ = ("strategy", "exchanges", "padding", "input_shape") + + def __init__(self, strategy, exchanges, padding, input_shape): + #: The ``ParallelStrategy``; :class:`_Halo3d` reads ``shard_ind``, + #: ``num_shards`` and ``shard_to_rank`` off it. + self.strategy = strategy + #: ``(dim_index, dim, halo)`` per dim actually split, in ``shard_dim`` + #: order -- the order ``distconv_forward`` exchanges in and the order + #: ``distconv_backward`` folds back down. + self.exchanges = exchanges + #: The module's padding with each exchanged dim's entry zeroed, exactly + #: as ``distconv_forward`` mutates the caller's list. + self.padding = padding + #: The shape the kernel sees: the local shard, plus ``2 * halo`` on each + #: exchanged dim. + self.input_shape = input_shape + + +def _halo_plan(dc_input, strategy, x, weight, stride, padding, dilation): + """The halo this ``DCTensor`` needs, or ``None`` if it must go to MIOpen. + + Returns a plan only when every fact the exchange rests on has been *read and + checked*, and ``None`` for anything it could not read. The asymmetry is + deliberate: a false negative costs a convolution its fast kernel, a false + positive returns a wrong gradient at every shard boundary of a large run, + which is the worst outcome available here. + + The strategy-wide facts, each matching a line of ``distconv.py``: + + * ``shard_dim`` has the same length as ``num_shards``, because + ``distconv_forward`` indexes ``num_shards`` and ``_is_periodic`` by + position in ``shard_dim`` and a mismatch means the two disagree about which + axis is which. ``shard_ind`` likewise, since the exchange indexes it the + same way. And no axis may be named twice: two entries for one dim would + exchange it twice and count the neighbour's slab twice. + * no axis is periodic. Periodicity is the one case where even a single + shard exchanges: ``shard_ind == 0 and is_periodic`` posts a send and a + receive to itself, so the halo is the tensor's own opposite face rather + than zeros, and the padding becomes ``_periodic_shard_padding`` instead of + 0. ScaFFold never sets it (it calls ``F.pad`` in the default constant + mode), which is a reason to check rather than a reason to assume. + * the shapes and the padding are the ordinary 5-D triples this module can + reason about at all. + + Then, per axis. An axis with ``num_shards == 1`` is **skipped**: its + ``shard_ind`` is 0, neither ``shard_ind > 0`` nor + ``shard_ind < num_shards - 1`` holds, no ``P2POp`` is posted and the + ``zeros_like`` receive buffers stay zero -- so DistConv's ``cat`` there is + ``cat(zeros, x, zeros)``, which at ``padding = 0`` is the same arithmetic as + the module's own ``padding = k // 2``, and measured bitwise so. It keeps its + padding and costs nothing. + + An axis with ``num_shards > 1`` is exchanged, and only after checking what + ``check_is_distconv_supported`` checks plus what this spelling of the + exchange additionally needs: + + * ``2 <= dim < 5``: a spatial dim of a 5-D tensor. ``ParallelStrategy`` + already rejects 0 and 1, but ``padding[dim - 2]`` is indexed here. + * the kernel extent on that dim is **odd**. An even kernel gives + ``halo_size == 0`` in DistConv, which is only correct for the strided + tiling ``check_is_distconv_supported`` then insists on; that is not a case + this module has reasoned about, so it declines rather than guesses. + * the padding on that dim is exactly ``k // 2`` ("same"), and the stride and + dilation are 1. Those three are what make "the halo'd extent at padding 0" + equal to "the global volume's slice for this shard": with ``k = 2h + 1`` + the halo'd input is ``D_loc + 2h`` long and produces exactly ``D_loc`` + outputs at zero padding, aligned with the shard's global offset. + * the shard is at least ``2 * halo`` thick, so the backward's two + accumulation regions do not overlap. + * ``shard_to_rank`` is callable, since the exchange has to name its + neighbours. + + What is *not* checked here is whether a process group exists: that is a + routing question rather than a property of the strategy, and it belongs with + the rest of them in :func:`_use_triton`. + """ + num_shards = getattr(strategy, "num_shards", None) + shard_dim = getattr(strategy, "shard_dim", None) + shard_ind = getattr(strategy, "shard_ind", None) + if not isinstance(num_shards, (tuple, list)) or not isinstance( + shard_dim, (tuple, list) + ): + return None + if not num_shards or len(shard_dim) != len(num_shards): + return None + if not isinstance(shard_ind, (tuple, list)) or len(shard_ind) != len(num_shards): + return None + if len(set(shard_dim)) != len(shard_dim): + return None + for count in num_shards: + if not isinstance(count, int) or count < 1: + return None + periodic = getattr(dc_input, "_is_periodic", None) + if not isinstance(periodic, (tuple, list)) or len(periodic) != len(shard_dim): + return None + if any(periodic): + return None + if x.dim() != 5 or weight.dim() != 5: + return None + for triple in (stride, padding, dilation): + if not isinstance(triple, (tuple, list)) or len(triple) != 3: + return None + + exchanges = [] + plan_padding = list(int(p) for p in padding) + plan_shape = list(int(s) for s in x.shape) + for i, dim in enumerate(shard_dim): + if num_shards[i] == 1: + continue + if not isinstance(dim, int) or not 2 <= dim < 5: + return None + index = shard_ind[i] + if not isinstance(index, int) or not 0 <= index < num_shards[i]: + return None + kernel = int(weight.shape[dim]) + if kernel % 2 == 0: + return None + halo = kernel // 2 + if plan_padding[dim - 2] != halo: + return None + if int(stride[dim - 2]) != 1 or int(dilation[dim - 2]) != 1: + return None + if halo == 0: # k == 1: no neighbour voxel is ever read + continue + if int(x.shape[dim]) < 2 * halo: + return None + exchanges.append((i, dim, halo)) + plan_padding[dim - 2] = 0 + plan_shape[dim] += 2 * halo + + if exchanges and not callable(getattr(strategy, "shard_to_rank", None)): + return None + return _HaloPlan(strategy, exchanges, tuple(plan_padding), tuple(plan_shape)) + + +def _policy_declines(x_shape, w_shape, stride, padding, dilation): + """Shapes the Triton rung is slower on: **none, today.** + + Empty since 2026-08-04, and the arguments are kept so a future entry has + somewhere to go. It held three rules, and re-measuring every one of them + against the shipped kernels retired all three + (``work/triton-conv/review/CONV_REMEASURE.md``; per-site sums over all + three directions, kernel time with 95% intervals): + + * ``Cin == 3``, the stem, blocked at a quoted **0.53x**. That figure was + backward-weight alone. Summed over the three directions the site is + **0.93x**, because backward-data runs 1.58-1.70x the other way, and 0.93x + on this site is **+0.13% to +0.19% of a step**. + * ``k == 1``, the ``64 -> 6`` head, blocked because "the forward is 1.22x but + backward-weight runs 0.73-1.14x". Backward-data is **1.94-2.00x** and was + never in that accounting: the site is **1.40x**. + * ``M <= 4096 and Cout >= 512``, the small-``M`` class, blocked on two + measured losses and extended by predicate to nine problems. All nine now + favour the rung, **1.17-1.65x**, worth **2.28% of a config-A step**. + + The rule that mattered was the last one, and its error was structural rather + than stale: ``M`` is the *forward* GEMM's row count, but a decline keeps the + whole site on MIOpen including its gradients, and backward-data is the + forward contraction on a *permuted weight* -- a different GEMM, which wins + 1.31-2.47x on every problem the rule blocked. A rule derived from one + direction was deciding three. **Anything added here inherits that trap**: + either measure all three directions, or block per direction. + + Emptying it also removes the convolutions as a source of run-to-run + variation, since MIOpen's backward-weight is not bitwise reproducible and + the deterministic split-K path is the rung's default. + + **This function is for the non-transposed operator only.** Every term a rule + might use reads a different quantity for the other one: ``w_shape``'s channel + axes are the other way round, and ``M`` is the *input* volume rather than + ``_out_spatial``'s -- 8x larger at ``k == s == 2``. See + :func:`_transposed_policy_declines`, which is a separate function for + exactly that reason. + """ + return False + + +# --------------------------------------------------------------------------- +# The halo exchange +# --------------------------------------------------------------------------- + + +def _packed_slab(shape, dtype, device, zero=False): + """A plain-contiguous NDHWC allocation, returned as ``(base, NCDHW view)``. + + ``base`` is what goes on the wire. It matters that it is *plain* contiguous: + ``forward_halo_exchange`` sends + ``inner_halo_plus.contiguous(memory_format=channels_last_3d)``, which for + ``C > 1`` is not plain-contiguous, and ``ProcessGroupGloo``'s send/recv + rejects that -- which is the whole of "DistConv spatial sharding does not + work under gloo". The cause is the memory format, not the narrow, so + allocating the wire buffer ourselves fixes it and lets the sharded suite run + on a machine with no RCCL. + + ``view`` is the same storage addressed as ``(N, C, ...)``, and a permute of a + contiguous ``(N, D, H, W, C)`` gives exactly ``channels_last_3d``'s strides + -- so copying a shard's boundary slab into it is a straight copy of 1.6% of + the activation and not a transpose. + """ + n, c = int(shape[0]), int(shape[1]) + spatial = tuple(int(v) for v in shape[2:]) + allocate = torch.zeros if zero else torch.empty + base = allocate((n, *spatial, c), dtype=dtype, device=device) + return base, base.permute(0, 4, 1, 2, 3) + + +def _neighbour_ranks(strategy, dim_index): + """The ranks holding the shards either side of this one along ``dim_index``. + + ``shard_to_rank`` on a copy of ``shard_ind``, which is what + ``forward_halo_exchange`` does; only ever asked for a neighbour that exists, + so its wrap-around branches (which exist for periodicity) are not reached. + """ + minus = list(strategy.shard_ind) + minus[dim_index] -= 1 + plus = list(strategy.shard_ind) + plus[dim_index] += 1 + return strategy.shard_to_rank(minus), strategy.shard_to_rank(plus) + + +def _exchange_forward(x, strategy, dim_index, dim, halo): + """``distconv.forward_halo_exchange`` for one dim, on plain-contiguous wires. + + Same sends, same receives, same posting order, and the same result: the + local shard with each neighbour's ``halo``-thick boundary slab concatenated + onto the matching face, and zeros where there is no neighbour. The + differences are both deliberate: the wire buffers are plain-contiguous (see + :func:`_packed_slab`), and the output is allocated in the layout the kernel + wants rather than left to ``torch.cat``'s memory-format inference. + """ + shard_ind = strategy.shard_ind[dim_index] + num_shards = strategy.num_shards[dim_index] + minus_rank, plus_rank = _neighbour_ranks(strategy, dim_index) + + slab_shape = list(x.shape) + slab_shape[dim] = halo + recv_minus, recv_minus_view = _packed_slab( + slab_shape, x.dtype, x.device, zero=True + ) + recv_plus, recv_plus_view = _packed_slab(slab_shape, x.dtype, x.device, zero=True) + + ops = [] + if shard_ind > 0: + send_minus, view = _packed_slab(slab_shape, x.dtype, x.device) + view.copy_(x.narrow(dim, 0, halo)) + ops += [ + dist.P2POp(dist.irecv, recv_minus, minus_rank), + dist.P2POp(dist.isend, send_minus, minus_rank), + ] + if shard_ind < num_shards - 1: + send_plus, view = _packed_slab(slab_shape, x.dtype, x.device) + view.copy_(x.narrow(dim, x.size(dim) - halo, halo)) + ops += [ + dist.P2POp(dist.isend, send_plus, plus_rank), + dist.P2POp(dist.irecv, recv_plus, plus_rank), + ] + if ops: + for request in dist.batch_isend_irecv(ops): + request.wait() + + halo_shape = list(x.shape) + halo_shape[dim] = x.size(dim) + 2 * halo + _, out = _packed_slab(halo_shape, x.dtype, x.device) + out.narrow(dim, halo, x.size(dim)).copy_(x) + out.narrow(dim, 0, halo).copy_(recv_minus_view) + out.narrow(dim, out.size(dim) - halo, halo).copy_(recv_plus_view) + return out + + +def _exchange_backward(grad, strategy, dim_index, dim, halo): + """``distconv.backward_halo_exchange`` for one dim, on plain-contiguous wires. + + The transpose of :func:`_exchange_forward`: the gradient of a value this rank + borrowed belongs to the rank it was borrowed from, so each outer boundary + slab of ``grad`` goes back to the neighbour that supplied it and is + *accumulated* into that neighbour's inner region. ``grad`` is mutated in + place and a narrowed view of it is returned, which is what + ``distconv_backward`` already hands back today -- safe because the only + producer of this gradient is :meth:`_TritonConv3dFn.backward`, whose output + has no other consumer. + """ + shard_ind = strategy.shard_ind[dim_index] + num_shards = strategy.num_shards[dim_index] + minus_rank, plus_rank = _neighbour_ranks(strategy, dim_index) + + slab_shape = list(grad.shape) + slab_shape[dim] = halo + recv_minus, recv_minus_view = _packed_slab( + slab_shape, grad.dtype, grad.device, zero=True + ) + recv_plus, recv_plus_view = _packed_slab( + slab_shape, grad.dtype, grad.device, zero=True + ) + + ops = [] + if shard_ind > 0: + send_minus, view = _packed_slab(slab_shape, grad.dtype, grad.device) + view.copy_(grad.narrow(dim, 0, halo)) + ops += [ + dist.P2POp(dist.irecv, recv_minus, minus_rank), + dist.P2POp(dist.isend, send_minus, minus_rank), + ] + if shard_ind < num_shards - 1: + send_plus, view = _packed_slab(slab_shape, grad.dtype, grad.device) + view.copy_(grad.narrow(dim, grad.size(dim) - halo, halo)) + ops += [ + dist.P2POp(dist.isend, send_plus, plus_rank), + dist.P2POp(dist.irecv, recv_plus, plus_rank), + ] + if ops: + for request in dist.batch_isend_irecv(ops): + request.wait() + + inner = grad.narrow(dim, halo, grad.size(dim) - 2 * halo) + inner.narrow(dim, 0, halo).add_(recv_minus_view) + inner.narrow(dim, inner.size(dim) - halo, halo).add_(recv_plus_view) + return inner + + +class _Halo3d(torch.autograd.Function): + """The halo exchange, as an autograd node of its own above the kernel's. + + Two ``Function``s rather than one, deliberately. The convolution node is + then shard-agnostic and *identical* whether or not anything is sharded -- it + is handed an input and a padding and knows nothing about either -- so the + unsharded path stays exactly what it was, this node is separately testable + against ``distconv.forward_halo_exchange``/``backward_halo_exchange``, and + the saved-tensor set stays uniform, which is what the per-rung latch's + argument needs. + + **Nothing is saved.** ``ctx`` carries the :class:`_HaloPlan`, which is + Python objects: the strategy, and which dims were exchanged at what width. + Applied only when there is something to exchange, so at + ``dc_num_shards = (1, 1, 1)`` it is not in the graph at all. + """ + + @staticmethod + def forward(ctx, x, plan): + ctx.plan = plan + for dim_index, dim, halo in plan.exchanges: + x = _exchange_forward(x, plan.strategy, dim_index, dim, halo) + return x + + @staticmethod + def backward(ctx, grad_output): + grad = grad_output + # Same order as the forward, which is also the order distconv_backward + # uses -- it iterates shard_dim forwards in both directions. + for dim_index, dim, halo in ctx.plan.exchanges: + grad = _exchange_backward(grad, ctx.plan.strategy, dim_index, dim, halo) + return grad, None + + +# --------------------------------------------------------------------------- +# Eligibility, continued +# --------------------------------------------------------------------------- + + +def _routing_declines(x, dc_input, plan, proven): + """The conditions *both* ladders decline on, in one place. + + ``True`` means "this call must not take a Triton rung", for a reason that is + not about which convolution it is: the override, the latch, a functorch + layer, an unknown tensor subclass, the device, a sharding this module could + not prove it can serve, and the layout. Everything operator-specific -- + the module's own attributes, the performance block-list and which + ``is_supported*`` to ask -- stays with the caller, because + :class:`FastConv3d` and :class:`FastConvTranspose3d` differ on every one of + them: they read the weight's channel axes in opposite orders, their measured + losses are different shapes of a different kernel, and their gates take + different arguments. + + Written once rather than copied because each clause below is a correction + with a measured failure behind it (the module docstring says which), and a + copy would drift away from them. The one thing the split cost is ordering: + the layout test used to run *after* ``FastConv3d``'s module-level checks and + now runs before them. Both are pure predicates returning ``False``, so only + the cost of reaching the answer moved, and it moved towards the cheaper end. + """ + if _triton_override is False: + return True + if _triton_failed and not proven: + return True + if _functorch_active(): + return True + # An unknown __torch_dispatch__ wrapper has unknown semantics and keeps + # MIOpen. DistConv's DCTensor never reaches this check -- forward() has + # already unwrapped to the local shard, and its sharding is checked through + # ``plan`` -- so anything rejected here is a subclass nobody has reasoned + # about. Note the identity test, not isinstance: is_supported only asks + # isinstance and would accept any of them. + if type(x) is not torch.Tensor: + return True + if not x.is_cuda: + return True + # ...and not merely *a* CUDA device: the one every launch configuration in + # both packages was raced on. Unlike every other clause here, the thing + # this one prevents is a right answer at an unknown speed -- see + # :func:`~ScaFFold.unet._rungs._platform_declines`, which is also where the + # "explicit opt-in wins" decision is argued. Cached per device, so this is + # a dictionary lookup after the first call. + if _platform_declines(x.device, _triton_override): + return True + if dc_input is not None and plan is None: + return True + # An exchange needs a process group. A real ``ParallelStrategy`` cannot be + # constructed without one, so this only fires for a hand-built strategy or a + # torch built without distributed -- but it is the call that would otherwise + # fail inside ``dist`` rather than routing to MIOpen. A transposed plan + # never has an exchange in it, so this clause is the non-transposed ladder's + # alone; it is here because it is a property of the *plan*, not of the + # operator. + if plan is not None and plan.exchanges: + if not (dist.is_available() and dist.is_initialized()): + return True + # A relayout is a correctness no-op -- every entry point calls + # ``.contiguous(memory_format=channels_last_3d)`` itself -- but it is a + # full-size hidden copy, and the whole point of the rung is that ScaFFold's + # activations are already in that layout. Declining is cheaper than paying + # it silently. Note that DistConv's narrowed ``_tensor`` stops being + # channels-last at ``local_batch_size > 1``, which is a supported config key, + # so this is a live branch and not a formality. + if x.dim() != 5 or not x.is_contiguous(memory_format=torch.channels_last_3d): + return True + return False + + +def _use_triton(module, x, dc_input, plan, proven=False): + """Whether this particular call should take the Triton rung. + + ``x`` is the tensor the halo would be added to -- a ``DCTensor``'s local + shard, or the input itself -- read as a plain attribute, so the tests examine + real strides and dtypes rather than a wrapper's mirrored metadata. + ``dc_input`` is the wrapper or ``None``, and ``plan`` is what + :func:`_halo_plan` made of its strategy: ``None`` for a plain tensor, and + ``None`` *also* for a ``DCTensor`` whose sharding this module could not + prove it can serve -- which is why the two are passed separately. + ``proven`` is the caller's "this module has already had a call served by this + rung", which keeps a proven module on it even after a global latch; see the + module docstring's "Latches". + + **The predicates are asked about the tensor the kernel will actually see**, + which once a dim is split is the halo'd one at the reduced padding, not the + local shard at the module's own padding. The two differ in extent and in + output shape, so asking about the wrong one would gate on a call that never + happens. + + Ordered so the cheap local tests come first and the package import last, and + so that nothing is allocated or cast before the answer is known. + """ + if _routing_declines(x, dc_input, plan, proven): + return False + # Module-level conditions the kernels have no argument for at all. A + # transposed convolution is served by :class:`FastConvTranspose3d` and by a + # different set of entry points (``is_supported`` does not even take a + # ``transposed`` parameter, and passing a ``(Cin, Cout, k, k, k)`` weight to + # the gates below would be *accepted* whenever the two channel counts match + # and would compute a different operator), and a non-zeros padding mode is an + # F.pad that happens in nn.Conv3d._conv_forward and that this rung would + # skip. Neither can occur for a FastConv3d built by ``unet_parts``; both are + # checked because the class is a public drop-in. + if module.transposed or module.output_padding != (0, 0, 0): + return False + if module.padding_mode != "zeros" or not isinstance(module.padding, tuple): + return False + weight = module.weight + # What the kernel is handed: the halo'd extent and the padding left over + # after each exchanged dim's was zeroed. Identical to ``x.shape`` and + # ``module.padding`` whenever nothing is split -- which is every call at + # ``dc_num_shards = (1,1,1)`` and no call at (2,1,1) or (4,1,1), where D is + # exchanged and H and W keep ``padding = 1``. So the kernel sees + # ``(1,1,1)`` unsharded and ``(0,1,1)`` sharded, and is padded either way. + kernel_shape = plan.input_shape if plan is not None else tuple(x.shape) + kernel_padding = plan.padding if plan is not None else module.padding + if _policy_declines( + kernel_shape, weight.shape, module.stride, kernel_padding, module.dilation + ): + return False + + # The predicates are cheap and side-effect free: attribute reads and integer + # arithmetic, no allocation, no launch. The broad catch is right *here* and + # nowhere else in this module: a predicate that cannot answer has a correct + # answer available ("no"), it has done no work anyone can observe, and the + # failure is a routing miss rather than a broken kernel -- so it must not + # latch the rung off. + try: + conv = _get_triton_module() + dtype = _autocast_dtype(x) + x_probe = _metadata_probe(kernel_shape, _cast_dtype(x, dtype), x.device) + w_probe = _metadata_probe( + weight.shape, _cast_dtype(weight, dtype), weight.device + ) + bias = module.bias + # The bias is Cout elements; casting it for real is cheaper than + # explaining a probe that also has to have stride 1. + bias_probe = _cast_operand(bias, dtype) if bias is not None else None + args = (module.stride, kernel_padding, module.dilation, module.groups) + # ``is_supported_all``, not ``is_supported``: every direction the + # backward will need, asked once. The forward's own gate does not imply + # the other two and the package says so -- ``stride > 1`` is served by + # the forward and by ``is_supported_bwd_weight`` and refused by + # ``is_supported_bwd_data``, whose kernel-free formulation only holds at + # unit stride. Taking the rung on the forward's answer alone would build + # a graph node whose backward ``triton_conv3d`` cannot answer, and by + # then MIOpen is no longer an option for it. + return bool(conv.is_supported_all(x_probe, w_probe, bias_probe, *args)) + except Exception as e: + _warn_once_about_the_predicate(e) + return False + + +# --------------------------------------------------------------------------- +# Autograd +# --------------------------------------------------------------------------- + + +def _aten_backward( + grad_output, + x, + weight, + stride, + padding, + dilation, + mask, + has_bias, + transposed=False, + output_padding=(0, 0, 0), +): + """MIOpen's ``convolution_backward``, for the callers that need it. + + Used as either backward rung's fallback and, in tests, as the reference the + Triton gradients are compared against. ``bias_sizes`` is required even when + the mask says no bias gradient is wanted, so it is always supplied -- and it + is the *output* channel count, which is ``weight.shape[0]`` for an ordinary + convolution and ``weight.shape[1]`` for a transposed one, because PyTorch + stores the two parameters with their channel axes the other way round. + + ``transposed`` and ``output_padding`` are the aten op's own arguments passed + through rather than a mode flag: this is one call to one operator, and both + ladders want it with the arguments their module holds. + """ + cout = int(weight.shape[1 if transposed else 0]) + return torch.ops.aten.convolution_backward( + grad_output, + x, + weight, + [cout] if has_bias else None, + list(stride), + list(padding), + list(dilation), + bool(transposed), + list(output_padding), + 1, # groups + list(mask), + ) + + +class _TritonConv3dFn(torch.autograd.Function): + """The Triton rung's autograd node. + + A plain ``torch.autograd.Function`` rather than + ``torch.library.custom_op`` + ``register_autograd``, which is the shape + ``triton_group_norm`` uses. That module measured the custom-op route at + +0.054 ms of dispatcher and +0.338 ms of autograd node against 0.065 ms for + an empty ``Function``, and paid it so its op would compose with + ``torch.compile`` and with ``DCTensor``'s ``__torch_dispatch__``. Neither + reason applies here: nothing in ScaFFold compiles the convolutions, and + ``DCTensor`` dispatch is precisely what this rung is bypassing -- a real + dispatcher op would be intercepted by DistConv's generic unwrap, which is + the wrong behaviour, because that path has no halo. + + Operands arrive already cast (see :func:`_cast_operand`), so this node is + dtype-transparent and its gradients flow back to the fp32 parameters through + the cast's own backward. + """ + + @staticmethod + def forward(ctx, x, weight, bias, stride, padding, dilation): + conv = _get_triton_module() + y = conv.conv3d_forward(x, weight, bias, stride, padding, dilation, 1) + # Saved *after* the launch, deliberately. The ladder retries a failed + # call on MIOpen, and that is only safe while the failing region has + # done nothing autograd can observe; ``triton_group_norm`` gets the same + # property from saving in ``_setup_context``, which runs after its + # forward returns. Nothing between here and the return can raise. + ctx.save_for_backward(x, weight) + ctx.conv_args = (stride, padding, dilation, bias is not None) + return y + + # The gradients below are computed by hand from the saved operands, not by + # composing differentiable ops, so nothing here can be differentiated again. + # Without this decorator that is *silent*: ``create_graph=True`` returns a + # gradient with no ``grad_fn``, and a second backward through it contributes + # zero instead of raising -- which is a wrong number rather than an error, in + # a project where those are the defects that keep nearly shipping. With it, + # the second differentiation says so. MIOpen's rung *is* twice + # differentiable, so this is a real difference between the rungs and it is + # why the block-list emptying (2026-08-04) is what surfaced it: the stem is + # the site whose input gradient a double backward actually reaches. + @staticmethod + @once_differentiable + def backward(ctx, grad_output): + global _triton_failed + x, weight = ctx.saved_tensors + if type(x) is not torch.Tensor: + # The one detector available for a rung flip across a checkpoint + # recompute. ``_default_meta_extractor`` compares shape, dtype and + # device, which are identical between the rungs, so torch's own + # check passes and a DCTensor lands in this slot instead of the + # plain tensor this forward saved. Without this the next line dies + # inside DistConv with an AttributeError about ``_parallel_strategy`` + # that names neither checkpointing nor the rung. One type check per + # backward. + raise RuntimeError( + "FastConv3d: the tensor saved for backward is a " + f"{type(x).__name__}, not a torch.Tensor. The Triton rung and " + "the MIOpen rung save different things, so this module's " + "forward and its checkpoint recompute were served by different " + "rungs. Set SCAFFOLD_CONV_TRITON=0 to pin the whole run to " + "MIOpen." + ) + stride, padding, dilation, has_bias = ctx.conv_args + needs_x, needs_w, needs_b = ctx.needs_input_grad[:3] + # One relayout for both directions: conv3d_backward_data and + # conv3d_backward_weight each call ``.contiguous(channels_last_3d)`` on + # it, and doing it here makes the second one free. + grad_output = grad_output.contiguous(memory_format=torch.channels_last_3d) + + conv = _get_triton_module() + try: + grad_x = ( + conv.conv3d_backward_data( + grad_output, weight, x.shape, stride, padding, dilation, 1 + ) + if needs_x + else None + ) + grad_w = ( + conv.conv3d_backward_weight( + x, weight.shape, grad_output, stride, padding, dilation, 1 + ) + if needs_w + else None + ) + except _triton_kernel_failures() as e: + # A backward-direction failure *can* be answered by MIOpen, and the + # argument that forbids it in the forward does not reach here. That + # argument is about the saved set: a proven module answering a + # checkpoint recompute from the other rung saves different tensors + # than the graph node holds. This code is not a recompute -- it is + # the node itself, running once, consuming exactly the tensors it + # was given. Nothing downstream can tell which kernel produced the + # gradients, so degrading is a fallback here in the full sense, and + # it matters: the backward-weight kernel is a separate compilation + # from the forward's, so it can raise OutOfResources on a call whose + # forward compiled cleanly. The rung is still latched, so no module + # that has not used it will try it again. + first = not _triton_failed + _triton_failed = True + if first: + _warn_rung_failure( + "Triton conv3d backward", e, "MIOpen kernel", TRITON_ENV_VAR + ) + grad_x, grad_w, grad_b = _aten_backward( + grad_output, + x, + weight, + stride, + padding, + dilation, + (needs_x, needs_w, has_bias and needs_b), + has_bias, + ) + return grad_x, grad_w, grad_b, None, None, None + + # d(bias) is the sum of grad_output over every axis but the channel one, + # whatever the forward kernel was. Only the segmentation head has a + # bias and it is on the policy block-list, so this is unreachable in + # ScaFFold today -- but ``is_supported`` accepts a bias, and the class is + # a public drop-in, so it has to be right. + grad_b = grad_output.sum(dim=(0, 2, 3, 4)) if (has_bias and needs_b) else None + return grad_x, grad_w, grad_b, None, None, None + + +class FastConv3d(nn.Conv3d): + """``nn.Conv3d`` with a Triton GPU kernel and MIOpen behind it. + + Identical state: ``weight`` of shape ``(Cout, Cin, kd, kh, kw)`` and the + optional ``bias`` of shape ``(Cout,)``, both from ``nn.Conv3d.__init__``, + no buffers and no extra attributes that are parameters or ``state_dict`` + keys. ``__init__`` is not overridden at all, so the constructor signature is + ``nn.Conv3d``'s by construction and state dicts are interchangeable in both + directions with a plain ``nn.Conv3d`` model. + + **The weight is used exactly as it lies.** There is no transform, no cache + and no stride contract: ``triton_conv3d`` addresses the parameter through its + strides, and ``worker.py``'s ``model.to(device, + memory_format=channels_last_3d)`` already puts every 5-D parameter in the + layout the kernel wants. An earlier design held the parameter in RSCK order + behind ``state_dict`` hooks; it was measured 0.159 ms/step better in the + kernels and 2.615 ms/step *worse* in the optimizer, and deleted. + + DistConv's ``DCTensor`` takes the fast kernel by being unwrapped to its local + shard in front of it, with the halo this module exchanges itself where a dim + is actually split -- which, unlike GroupNorm, is a real computation and not a + formality. See the module docstring, and :func:`_halo_plan` for the check + that decides it. + """ + + #: Per-module "a call has been answered by this ladder" -- by the Triton + #: kernel, or by the MIOpen fallback :meth:`_triton_forward` runs on an + #: already-exchanged tensor, which saves the same set. A global latch does + #: not demote a module that has one, which is what keeps a checkpointed + #: block's forward and its recompute on the same rung. A plain class + #: attribute, so it is not a parameter, a buffer or a state-dict key, and it + #: costs nothing until the first success -- the instance attribute is only + #: written on the False -> True edge, because nn.Module.__setattr__ is not + #: free. + _triton_ok = False + + def _triton_forward(self, local, plan=None): + """Run the Triton rung -- and its fallback -- on an unwrapped tensor. + + ``plan`` is :func:`_halo_plan`'s answer, or ``None`` for a tensor that + is not sharded at all. The cast comes first, so the exchange carries the + bf16 tensor autocast's dispatcher would have produced rather than the + fp32 one the model holds -- which is both what DistConv does and half the + bytes on the wire. + + **The exchange is above the retry; only the kernel call is inside it.** + That split is what makes a Triton failure survivable once a dim is + actually split. ``_Halo3d`` runs before the kernel compiles, so by the + time a ``TritonError`` arrives this rank has already posted the sends and + receives its peers are matched against. Re-running the call from the top + would take it to :meth:`_miopen_forward` and therefore through + ``distconv_forward``, which exchanges *again*: one more collective on + this rank than on a peer whose kernel compiled, which desynchronises the + mesh and hangs, or pairs this convolution's slabs with the next one's. + So the fallback here consumes the tensor that has already been exchanged + -- ``x``, at ``plan.padding``, which is precisely the pair + :func:`_halo_plan` proved equal to the module's own padding on the + unexchanged shard -- and swaps nothing but the kernel. Without it a + broken Triton install is *fatal* at ``num_shards > 1`` where it is merely + slow at 1, and the whole point of the ladder is the opposite. + + At one shard nothing has been sent and the exception is re-raised + unchanged, so :meth:`forward`'s handler re-runs the whole call on the + rung that defines the semantics -- which is also the only rung that can + be handed a ``DCTensor``. + """ + dtype = _autocast_dtype(local) + x = _cast_operand(local, dtype) + weight = _cast_operand(self.weight, dtype) + bias = _cast_operand(self.bias, dtype) + padding = self.padding + exchanged = plan is not None and bool(plan.exchanges) + if exchanged: + x = _Halo3d.apply(x, plan) + padding = plan.padding + try: + return _TritonConv3dFn.apply( + x, weight, bias, self.stride, padding, self.dilation + ) + except _triton_kernel_failures() as e: + if not exchanged: + raise + _latch_rung_failure(e) + # The one call this fallback must not answer either, for the same + # reason :meth:`forward`'s does not: a module already proven on the + # rung, failing while a backward is in flight, is a checkpoint + # recompute of a forward that ran on Triton, and the honest answer is + # the original exception rather than a differently-produced tensor + # substituted into a graph node that already holds one. + if self._triton_ok and _replaying_a_forward(): + raise + # ``nn.Conv3d.forward``'s own body for ``padding_mode="zeros"``, + # which :func:`_use_triton` has already checked, on the halo'd + # operands. A plain tensor, so DistConv's ``__torch_dispatch__`` + # does not see it and no second exchange happens; ``grad_x`` still + # flows back through :class:`_Halo3d` and ``grad_weight`` is still + # computed against the halo'd input, which is what + # ``distconv_backward`` does. + return F.conv3d( + x, weight, bias, self.stride, padding, self.dilation, self.groups + ) + + def _miopen_forward(self, input): + """The semantics-defining rung. + + ``nn.Conv3d.forward`` unchanged, which is also the only rung that takes + a ``DCTensor`` as it stands: DistConv's ``__torch_dispatch__`` intercepts + the ``aten::convolution`` underneath and performs the halo exchange. + """ + return super().forward(input) + + def forward(self, input): + distconv = _dctensor_ops(input) + # The eligibility checks look at the local shard for a DCTensor (the + # peek is a plain attribute read, no autograd involvement) and at the + # tensor itself otherwise. + local_view = input._tensor if distconv is not None else input + # None for a plain tensor, and None also for a DCTensor whose strategy + # this module cannot prove it can serve -- _use_triton tells the two + # apart from ``dc_input``. + plan = ( + _halo_plan( + input, + input._parallel_strategy, + local_view, + self.weight, + self.stride, + self.padding, + self.dilation, + ) + if distconv is not None + else None + ) + + if _use_triton( + self, + local_view, + input if distconv is not None else None, + plan, + proven=self._triton_ok, + ): + triton_failures = _triton_kernel_failures() + try: + out = _run_local( + input, distconv, lambda local: self._triton_forward(local, plan) + ) + except triton_failures as e: + # A broken or mismatched Triton install, an unwritable JIT cache + # or a tile that does not fit must cost speed, not a multi-node + # run. Every allowlisted exception is raised while compiling or + # sizing a launch -- before the node saved anything -- so + # retrying this same call on MIOpen is safe. + _latch_rung_failure(e) + # The one call this rung must not answer from MIOpen: a module + # already proven on it, failing while a backward is in flight, is + # a checkpoint recompute of a forward that *did* run on Triton. + # The rungs save metadata-identical but structurally different + # tensors under DistConv, so handing back MIOpen's result + # substitutes a DCTensor into a slot holding a plain one and + # fails later, inside DistConv, with a message about neither. + # Degrading is for modules with nothing to contradict; here the + # honest answer is the original exception. + if self._triton_ok and _replaying_a_forward(): + raise + # The second call it must not answer from MIOpen: one whose halo + # has already been exchanged. ``_miopen_forward`` goes through + # ``distconv_forward``, which exchanges *again* -- one more + # collective on this rank than on a peer whose kernel compiled, + # which desynchronises the whole mesh and hangs, or silently + # pairs this convolution's slabs with the next one's. Such a + # call is normally answered inside ``_triton_forward``, on the + # tensor that was already exchanged, and never reaches here; this + # is the backstop for a failure raised *outside* that try block + # (the cast, the exchange itself, ``DCTensor.from_shard``), none + # of which runs Triton today. It is unconditional because the + # invariant is: once this rank has sent a slab, it does not enter + # a code path that sends another. + if plan is not None and plan.exchanges: + raise + else: + # Set for a call ``_triton_forward`` answered from MIOpen on the + # exchanged tensor too, which is not a slip: the flag pins the + # module to *this ladder*, and what a checkpoint recompute has to + # agree about is the saved set, not which kernel produced the + # values. Both of this ladder's answers save the halo'd input, + # unwrapped, above autograd; ``_miopen_forward``'s saves what + # DistConv's dispatch saves. Those are the two that must not be + # mixed. + if not self._triton_ok: + self._triton_ok = True + return out + + return self._miopen_forward(input) + + +# --------------------------------------------------------------------------- +# The transposed ladder +# --------------------------------------------------------------------------- + + +def _transposed_policy_declines(x_shape, w_shape): + """Shapes the transposed Triton rung is *slower* on: **none, today**. + + The peer of :func:`_policy_declines`, and a separate function rather than a + flag on that one, because two of its three entries would read the wrong + number here and the third is a measurement of a different kernel: + + * ``w_shape`` is a ``ConvTranspose3d`` parameter, ``(Cin, Cout, kd, kh, + kw)``. Its channel axes are the other way round, so + ``_policy_declines``'s ``cout, cin = w_shape[0], w_shape[1]`` reads each + one as the other -- at ``up1`` it would test ``Cin == 3`` against 512 and + ``Cout >= 512`` against 1024. + * ``M``, the forward GEMM's row count, is ``N * prod(in_spatial)``: this + operator's windows *tile* the output, so its GEMM has one row per **input** + voxel and the taps are in N. ``_out_spatial`` computes the non-transposed + extents, which at ``k == s == 2`` are half the input's per axis -- an ``M`` + 8x too small. + * and the small-``M`` cliff those numbers describe is a property of + ``gather_gemm``'s tuning, not of this kernel, whose N axis carries + ``Cout * taps`` and is 8x wider for it. + + So the four sites were measured directly rather than inherited + (``work/triton-conv/m5_shipped_{fwd,bwd-data,bwd-weight}.json``, config A, + Triton over MIOpen per direction): + + =========================== ===== ======== ========== + site fwd bwd-data bwd-weight + =========================== ===== ======== ========== + up1 ``1024 -> 512`` @ 8^3 1.26x 1.43x 0.94x + up2 ``512 -> 256`` @ 16^3 1.66x 1.54x 1.12x + up3 ``256 -> 128`` @ 32^3 2.44x 1.36x 1.22x + up4 ``128 -> 64`` @ 64^3 3.31x 1.00x 1.33x + =========================== ===== ======== ========== + + Every site wins overall, so nothing is blocked. Worth stating explicitly: + the small-``M`` rule would have blocked ``up1`` even with ``M`` read + correctly (``M = 512``, ``Cout = 512``), and the measurement says it should + not. The one sub-unity cell is ``up1``'s backward-weight at 0.94x, against + +1.26x and +1.43x on the same site's other two directions. + + Checked again at the level that decides it, because those are *kernel* times + and this ladder's per-call cost is not zero: on the real model at config A, + interleaved, 12 steps per arm + (``work/triton-conv/bin/convT_config_a.py``), the whole ladder is worth + **0.51 ms/step** and a variant that keeps only ``up3``/``up4`` on the rung is + worth 0.65 -- a 0.14 ms difference inside a 1.5 ms spread, i.e. not a + measurement, so there is nothing here to write an entry from. What *is* + measurable is that the gap between 0.51 and the 0.9 ms the kernel times + project is per-call Python and launch overhead (~0.10 ms per forward and + ~0.15 ms per backward, measured pipelined), which is a property of the + adapter rather than of any shape -- so it belongs in a profile, not in a + block-list. + + Both arguments are unused today and are taken anyway: they are what an entry + would be written in terms of, and taking them keeps this call site the same + shape as the other ladder's. ``tests/test_conv3d.py`` asserts that the + answer is ``False`` at all four sites, so an entry added here cannot silently + turn the rung off. + """ + return False + + +def _transposed_halo_plan(dc_input, strategy, x, weight, padding): + """The halo a ``k == s`` transposed convolution needs: **none, ever**. + + Returns a :class:`_HaloPlan` with an empty ``exchanges``, or ``None`` if any + fact it rests on could not be read -- the same asymmetry :func:`_halo_plan` + documents, and for the same reason. + + Where :func:`_halo_plan` has to *reproduce* an exchange, this one has to + prove there is not one, and the proof has two halves: + + * **The operator needs no neighbour voxel.** At ``kernel == stride`` and no + padding the map ``(d, kd) -> d*k + kd`` is a bijection, so output voxel + ``d*k + kd`` reads input voxel ``d`` and nothing else. A shard that holds + a contiguous block of input voxels therefore holds everything its own + output block needs, at any shard count. The gate + (``is_supported_transposed``) pins ``k == s``, ``padding == 0``, + ``output_padding == 0`` and ``dilation == 1``, which is exactly that case. + * **DistConv agrees, so the two rungs compute the same thing.** Its + ``distconv_forward`` sets ``halo_size = kernel_size // 2 if odd else 0``, + and ``forward_halo_exchange``/``backward_halo_exchange`` both return their + argument unchanged at ``halo_size == 0``. With ``k = 2`` on every axis the + MIOpen rung therefore also runs on the bare local shard -- no ``cat``, no + ``P2POp``, no padding rewrite -- at *every* shard count, which is what + makes the fallback and the fast rung the same computation rather than + merely both defensible. + + Hence the parity test below. An **odd** kernel on a split dim is declined, + even though the bijection above holds for it too: DistConv would want a + ``k // 2`` halo there and then refuse the problem outright in + ``check_is_distconv_supported`` ("when kernel size is odd, padding must be + equivalent to same", and this operator's padding is 0), so there would be no + incumbent to agree with, no measurement, and a silent change from "the run + raises" to "the run answers". ScaFFold's four sites are all ``k = 2``. + + An axis with ``num_shards == 1`` is skipped before that test, as in + :func:`_halo_plan`: it is not split, so nothing about it can matter. + """ + num_shards = getattr(strategy, "num_shards", None) + shard_dim = getattr(strategy, "shard_dim", None) + shard_ind = getattr(strategy, "shard_ind", None) + if not isinstance(num_shards, (tuple, list)) or not isinstance( + shard_dim, (tuple, list) + ): + return None + if not num_shards or len(shard_dim) != len(num_shards): + return None + if not isinstance(shard_ind, (tuple, list)) or len(shard_ind) != len(num_shards): + return None + if len(set(shard_dim)) != len(shard_dim): + return None + for count in num_shards: + if not isinstance(count, int) or count < 1: + return None + # Periodicity is the one case where even a single shard exchanges, and it + # rewrites the padding as well. ScaFFold never sets it; this is a check + # rather than an assumption for the same reason it is one in _halo_plan. + periodic = getattr(dc_input, "_is_periodic", None) + if not isinstance(periodic, (tuple, list)) or len(periodic) != len(shard_dim): + return None + if any(periodic): + return None + if x.dim() != 5 or weight.dim() != 5: + return None + if not isinstance(padding, (tuple, list)) or len(padding) != 3: + return None + + for i, dim in enumerate(shard_dim): + if num_shards[i] == 1: + continue + if not isinstance(dim, int) or not 2 <= dim < 5: + return None + index = shard_ind[i] + if not isinstance(index, int) or not 0 <= index < num_shards[i]: + return None + if int(weight.shape[dim]) % 2 != 0: + return None + + # Empty exchanges, and the module's own padding: the kernel is handed the + # local shard exactly as it stands. Carrying a plan at all -- rather than a + # bare boolean -- is what lets :func:`_routing_declines` read "a DCTensor + # this module could not prove it can serve" the same way for both ladders. + return _HaloPlan(strategy, (), tuple(int(p) for p in padding), tuple(x.shape)) + + +def _use_triton_transposed(module, x, dc_input, plan, proven=False): + """Whether this transposed call should take the Triton rung. + + :func:`_use_triton`'s peer, with the same arguments and the same contract. + What differs after :func:`_routing_declines` has answered the shared half: + the module-level conditions are the mirror image (``transposed`` must be + *true* here), the block-list is :func:`_transposed_policy_declines`, and the + gate is ``is_supported_transposed_all`` -- all three directions at once, for + the reason its own docstring gives. Nothing about a halo appears, because a + transposed plan never has an exchange in it. + """ + if _routing_declines(x, dc_input, plan, proven): + return False + # The mirror of the check that sends a transposed module here in the first + # place. ``FastConvTranspose3d`` cannot be built any other way, but the + # class is a public drop-in and ``transposed`` is a plain attribute: were it + # false, ``is_supported_transposed`` would read a ``(Cout, Cin, k, k, k)`` + # weight as ``(Cin, Cout, ...)`` and accept it whenever the two channel + # counts match, computing a different operator without a word. + if not module.transposed: + return False + if module.padding_mode != "zeros": + return False + for triple in (module.padding, module.output_padding, module.stride): + if not isinstance(triple, tuple) or len(triple) != 3: + return False + weight = module.weight + if _transposed_policy_declines(tuple(x.shape), weight.shape): + return False + + # As in _use_triton: the predicates are attribute reads and integer + # arithmetic, and one that cannot answer has a correct answer available + # ("no") and has done no observable work -- a routing miss, not a kernel + # failure, so it must not latch the rung off. + try: + conv = _get_triton_module() + dtype = _autocast_dtype(x) + x_probe = _metadata_probe(x.shape, _cast_dtype(x, dtype), x.device) + w_probe = _metadata_probe( + weight.shape, _cast_dtype(weight, dtype), weight.device + ) + bias = module.bias + # Cast for real rather than probed: ``is_supported_transposed`` reads + # the bias's *stride*, which an expanded probe reports as 0, so a probe + # would answer "no" for every biased call -- which is all four of them. + bias_probe = _cast_operand(bias, dtype) if bias is not None else None + return bool( + conv.is_supported_transposed_all( + x_probe, + w_probe, + bias_probe, + module.stride, + module.padding, + module.output_padding, + module.dilation, + module.groups, + ) + ) + except Exception as e: + _warn_once_about_the_predicate(e) + return False + + +class _TritonConvTranspose3dFn(torch.autograd.Function): + """The transposed rung's autograd node. + + :class:`_TritonConv3dFn`'s peer; the reasons for a plain + ``autograd.Function`` and for saving *after* the launch are that class's and + unchanged. Two things are genuinely different: + + * **``grad_bias`` is on the path.** All four transposed sites have one + (``nn.ConvTranspose3d`` defaults to ``bias=True`` and ``unet_parts`` does + not turn it off), where the only biased ordinary convolution is the + segmentation head, which the block-list keeps on MIOpen. So the reduction + below is reached on every step of every run, not only by a test. + * **The rung-flip hazard is wider.** ``DCTensor`` mirrors its local shard's + size, strides, dtype and device exactly, and this ladder never adds a halo + -- so the tensor this node saves and the one ``nn.ConvTranspose3d.forward`` + saves under DistConv are metadata-identical at *every* shard count, not + only at one. ``FastConv3d`` gets a loud ``CheckpointError`` once a dim is + split, because its Triton rung saves the halo'd input and the extent + differs; here there is nothing to differ, so the per-module latch is the + whole of the defence and the type check below is the only detector. + """ + + @staticmethod + def forward(ctx, x, weight, bias, stride, padding, output_padding, dilation): + conv = _get_triton_module() + y = conv.conv_transpose3d_forward( + x, weight, bias, stride, padding, output_padding, dilation, 1 + ) + # After the launch: the ladder retries a failed call on MIOpen, which is + # only safe while the failing region has done nothing autograd can + # observe. Nothing between here and the return can raise. + ctx.save_for_backward(x, weight) + ctx.conv_args = (stride, padding, output_padding, dilation, bias is not None) + return y + + # See :meth:`_TritonConv3dFn.backward`: hand-computed gradients, so a second + # differentiation must raise rather than silently contribute zero. + @staticmethod + @once_differentiable + def backward(ctx, grad_output): + global _triton_failed + x, weight = ctx.saved_tensors + if type(x) is not torch.Tensor: + # See the class docstring: the two rungs save tensors that agree on + # everything ``_default_meta_extractor`` compares, so a rung flip + # across a checkpoint recompute passes torch's own check and dies + # inside DistConv with an AttributeError about ``_parallel_strategy`` + # that names neither checkpointing nor the rung. One type check per + # backward buys an actionable message instead. + raise RuntimeError( + "FastConvTranspose3d: the tensor saved for backward is a " + f"{type(x).__name__}, not a torch.Tensor. The Triton rung and " + "the MIOpen rung save different things, so this module's " + "forward and its checkpoint recompute were served by different " + "rungs. Set SCAFFOLD_CONV_TRITON=0 to pin the whole run to " + "MIOpen." + ) + stride, padding, output_padding, dilation, has_bias = ctx.conv_args + needs_x, needs_w, needs_b = ctx.needs_input_grad[:3] + # One relayout for both directions, as in _TritonConv3dFn: each entry + # point would call ``.contiguous(channels_last_3d)`` itself. + grad_output = grad_output.contiguous(memory_format=torch.channels_last_3d) + + conv = _get_triton_module() + try: + grad_x = ( + conv.conv_transpose3d_backward_data( + grad_output, + weight, + x.shape, + stride, + padding, + output_padding, + dilation, + 1, + ) + if needs_x + else None + ) + grad_w = ( + conv.conv_transpose3d_backward_weight( + x, + weight.shape, + grad_output, + stride, + padding, + output_padding, + dilation, + 1, + ) + if needs_w + else None + ) + except _triton_kernel_failures() as e: + # Degrading is a genuine fallback here, for the reason + # _TritonConv3dFn.backward gives: this is the node itself running + # once, not a recompute, and nothing downstream can tell which + # kernel produced the gradients. Both backward directions are + # separate compilations from the forward's, so either can raise + # OutOfResources on a call whose forward compiled cleanly. + first = not _triton_failed + _triton_failed = True + if first: + _warn_rung_failure( + "Triton conv_transpose3d backward", + e, + "MIOpen kernel", + TRITON_ENV_VAR, + ) + grad_x, grad_w, grad_b = _aten_backward( + grad_output, + x, + weight, + stride, + padding, + dilation, + (needs_x, needs_w, has_bias and needs_b), + has_bias, + transposed=True, + output_padding=output_padding, + ) + return grad_x, grad_w, grad_b, None, None, None, None + + # d(bias) is the sum of grad_output over every axis but the channel one, + # whatever the forward kernel was -- the same expression ATen's + # convolution backward uses, and reached on every step here. + grad_b = grad_output.sum(dim=(0, 2, 3, 4)) if (has_bias and needs_b) else None + return grad_x, grad_w, grad_b, None, None, None, None + + +class FastConvTranspose3d(nn.ConvTranspose3d): + """``nn.ConvTranspose3d`` with a Triton GPU kernel and MIOpen behind it. + + :class:`FastConv3d`'s peer, for the four ``k = 2, s = 2`` upsamplers in the + decoder. Identical state: ``weight`` of shape ``(Cin, Cout, kd, kh, kw)`` + and ``bias`` of shape ``(Cout,)``, both from ``nn.ConvTranspose3d.__init__``, + which is not overridden -- so the constructor signature is the stock one by + construction and state dicts are interchangeable in both directions with a + plain ``nn.ConvTranspose3d`` model. + + **Sharding is not the difficulty it is for the ordinary convolution.** At + ``kernel == stride`` every output voxel reads exactly one input voxel, so + there is no halo to exchange and no ``_Halo3d`` in this ladder; DistConv + reaches the same conclusion by a different route (``halo = k // 2`` is 0 for + an even kernel) and also runs on the bare local shard, at every shard count. + :func:`_transposed_halo_plan` is where that is checked rather than assumed. + + What does *not* get weaker here: the per-module latch, which is if anything + more load-bearing than it is for ``FastConv3d`` -- see + :class:`_TritonConvTranspose3dFn` -- and the autocast reproduction, since + ``conv_transpose3d`` carries the same ``lower_precision_fp`` cast policy as + ``convolution`` (verified: an fp32 pair under ``autocast(bf16)`` produces a + bf16 output, an fp64 pair is not narrowed). + """ + + #: Per-module "a call has been answered by the Triton rung"; see + #: :attr:`FastConv3d._triton_ok`, which this mirrors exactly. + _triton_ok = False + + def _triton_forward(self, local): + """Run the Triton rung on an unwrapped tensor. + + No halo, so no exchange, so no retry-below-the-exchange split: a + ``TritonError`` here has put nothing on the wire and :meth:`forward` can + re-run the whole call on MIOpen. The cast comes first for the same + reason it does in ``FastConv3d._triton_forward`` -- it reproduces what + the dispatcher would have done below this module. + """ + dtype = _autocast_dtype(local) + return _TritonConvTranspose3dFn.apply( + _cast_operand(local, dtype), + _cast_operand(self.weight, dtype), + _cast_operand(self.bias, dtype), + self.stride, + self.padding, + self.output_padding, + self.dilation, + ) + + def _miopen_forward(self, input, output_size=None): + """The semantics-defining rung: ``nn.ConvTranspose3d.forward``. + + Also the only rung that takes a ``DCTensor`` as it stands, since + DistConv's ``__torch_dispatch__`` intercepts the ``aten::convolution`` + underneath it. + """ + return super().forward(input, output_size) + + def forward(self, input, output_size=None): + # ``output_size`` re-derives ``output_padding`` inside + # ``_output_padding``, so the call the kernel would be gated on is not + # the call that would run. It is not used anywhere in ScaFFold (``Up`` + # calls ``self.up(x1)``), so the stock rung answers it rather than this + # ladder growing a second way to compute a padding. + if output_size is not None: + return self._miopen_forward(input, output_size) + + distconv = _dctensor_ops(input) + local_view = input._tensor if distconv is not None else input + plan = ( + _transposed_halo_plan( + input, + input._parallel_strategy, + local_view, + self.weight, + self.padding, + ) + if distconv is not None + else None + ) + + if _use_triton_transposed( + self, + local_view, + input if distconv is not None else None, + plan, + proven=self._triton_ok, + ): + try: + out = _run_local(input, distconv, self._triton_forward) + except _triton_kernel_failures() as e: + # Every allowlisted exception is raised while compiling or + # sizing a launch -- before the node saved anything -- and this + # ladder never puts a halo slab on the wire, so re-running the + # whole call on MIOpen is safe at every shard count. That is + # the clause ``FastConv3d`` needs a second fallback for and this + # one does not. + _latch_rung_failure(e, "Triton conv_transpose3d") + # The one call the fallback must not answer: a module already + # proven on the rung, failing while a backward is in flight, is + # a checkpoint recompute of a forward that ran on Triton. The + # two rungs save metadata-identical but structurally different + # tensors under DistConv, so MIOpen's answer would substitute a + # DCTensor into a slot holding a plain one and fail later, + # inside DistConv, with a message about neither. + if self._triton_ok and _replaying_a_forward(): + raise + else: + if not self._triton_ok: + self._triton_ok = True + return out + + return self._miopen_forward(input) diff --git a/ScaFFold/unet/unet_parts.py b/ScaFFold/unet/unet_parts.py index 9fffe72..a49f72c 100644 --- a/ScaFFold/unet/unet_parts.py +++ b/ScaFFold/unet/unet_parts.py @@ -20,6 +20,7 @@ from ScaFFold.utils.perf_measure import annotate +from .conv3d import FastConv3d, FastConvTranspose3d from .group_norm import FastGroupNorm _doubleconv_annotate = annotate(fmt="DoubleConv.{}") @@ -39,6 +40,87 @@ def _group_norm(num_groups, num_channels, activation=None): return FastGroupNorm(num_groups, num_channels, activation=activation) +def _conv3d(in_channels, out_channels, **kwargs): + """The model's non-transposed convolutions, in one place. + + ``FastConv3d`` is ``nn.Conv3d`` plus a Triton GPU kernel; it holds the same + parameters under the same names and adds no buffers, so checkpoints are + unaffected in either direction. It falls back to MIOpen for anything the + kernel does not serve, and it serves the sharded configurations too: it + performs the halo exchange itself, above autograd, rather than leaving it to + the one DistConv does below. Note what that means for the *shape* the + kernel sees -- only the split axis is halo'd, so ``padding=1`` survives on + the other two and these convolutions are padded at every configuration; see + :mod:`ScaFFold.unet.conv3d`. + + The four ``nn.ConvTranspose3d`` in ``Up`` do not come through here: they are + a different operator, with the weight's channel axes the other way round and + a different set of kernels behind them, so they have a factory of their own + (:func:`_conv_transpose3d`) rather than a flag on this one. + """ + return FastConv3d(in_channels, out_channels, **kwargs) + + +def _conv_transpose3d(in_channels, out_channels, **kwargs): + """The model's transposed convolutions -- the decoder's upsamplers. + + ``FastConvTranspose3d`` is ``nn.ConvTranspose3d`` plus a Triton GPU kernel; + it holds the same parameters (``weight`` *and* ``bias``, which these sites + have and the ordinary convolutions mostly do not) under the same names and + adds no buffers, so checkpoints are unaffected in either direction. It falls + back to MIOpen for anything the kernel does not serve, which is everything + except the ``kernel == stride``, no-padding upsample built below. + """ + return FastConvTranspose3d(in_channels, out_channels, **kwargs) + + +def _consumer_dtype(*tensors): + """The dtype the convolution consuming a concatenation will actually see. + + Inside an enabled autocast region the answer is autocast's dtype, because + ``aten::convolution`` carries the ``lower_precision_fp`` cast policy and + casts whatever it is handed. Producing that dtype from the concatenation + is *bitwise identical* to producing ATen's promoted dtype and letting the + convolution narrow it -- the promoted tensor holds exact widenings of both + sources, so narrowing before or after the copy rounds the same values once + -- while writing and reading back half the bytes. + + Outside autocast the answer is ``torch.cat``'s ordinary promotion, so eval, + ``inference_mode`` and pure-fp32 runs are unchanged. + """ + dtype = tensors[0].dtype + for tensor in tensors[1:]: + dtype = torch.promote_types(dtype, tensor.dtype) + device_type = tensors[0].device.type + try: + if not torch.is_autocast_enabled(device_type): + return dtype + autocast_dtype = torch.get_autocast_dtype(device_type) + except (RuntimeError, TypeError): # a device type autocast does not know + return dtype + # Only ever narrow: if autocast's dtype is the wider of the two, keep the + # promotion ATen would have done. + if torch.promote_types(autocast_dtype, dtype) is autocast_dtype: + return dtype + return autocast_dtype + + +def _skip_concat(skip, upsampled): + """``torch.cat([skip, upsampled], dim=1)`` at the consumer's dtype. + + Under ``torch.autocast`` the two halves do not share a dtype: the skip + comes from a GroupNorm, an fp32-policy op, while the upsampled half comes + from a ``ConvTranspose3d`` and is bf16. ``torch.cat`` carries the + ``promote`` policy, so it widens the bf16 half to fp32, concatenates at + fp32, and the following convolution narrows the whole double-width result + straight back down -- three full-resolution passes to deliver one. Casting + the inputs first collapses that to one, and the convolution reads the same + bits either way (see :func:`_consumer_dtype`). + """ + dtype = _consumer_dtype(skip, upsampled) + return torch.cat([skip.to(dtype), upsampled.to(dtype)], dim=1) + + class DoubleConv(nn.Module): """(convolution => GroupNorm => ReLU) * 2 @@ -63,10 +145,10 @@ def __init__(self, in_channels, out_channels, group_norm_groups, mid_channels=No if not mid_channels: mid_channels = out_channels self.double_conv = nn.Sequential( - nn.Conv3d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False), + _conv3d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False), _group_norm(group_norm_groups, mid_channels, activation="relu"), nn.Identity(), - nn.Conv3d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False), + _conv3d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False), _group_norm(group_norm_groups, out_channels, activation="relu"), nn.Identity(), ) @@ -92,7 +174,23 @@ def forward(self, x): class Up(nn.Module): - """Upscaling then double conv""" + """Upscaling then double conv + + The skip concatenation goes through :func:`_skip_concat` rather than + ``torch.cat`` directly, so that it emits the dtype the following + convolution will use instead of ``torch.cat``'s promoted one. The tensor + that convolution reads is bitwise unchanged either way; it is written and + read back at half the width. This rests on ``self.conv`` beginning with a + convolution, which the constructor below guarantees on either branch. + + Measured at scale 7: 1.09 ms of a 92.8 ms step, and 0.50 GiB of peak + memory. A channels-last-native Triton concatenation kernel was built and + measured too -- ``cat``'s *backward* is a narrowed view that consumers force + contiguous, at 51-63% of this device's streaming roofline against the + kernel's 90-103% -- but it was worth a further 0.08 ms of the step, which + did not justify a second hand-written kernel in a benchmark other people + have to trust. ``work/skip-path/RESULTS.md`` has the numbers. + """ def __init__(self, in_channels, out_channels, group_norm_groups, trilinear=True): super().__init__() @@ -107,7 +205,7 @@ def __init__(self, in_channels, out_channels, group_norm_groups, trilinear=True) in_channels // 2, ) else: - self.up = nn.ConvTranspose3d( + self.up = _conv_transpose3d( in_channels, in_channels // 2, kernel_size=2, stride=2 ) self.conv = DoubleConv(in_channels, out_channels, group_norm_groups) @@ -135,14 +233,16 @@ def forward(self, x1, x2): # if you have padding issues, see # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd - x = torch.cat([x2, x1], dim=1) + # torch.cat([x2, x1], dim=1) with the dtype and the layout the + # convolution below actually wants; see the class docstring. + x = _skip_concat(x2, x1) return self.conv(x) class OutConv(nn.Module): def __init__(self, in_channels, out_channels): super(OutConv, self).__init__() - self.conv = nn.Conv3d(in_channels, out_channels, kernel_size=1) + self.conv = _conv3d(in_channels, out_channels, kernel_size=1) @_outconv_annotate def forward(self, x): diff --git a/tests/test_conv3d.py b/tests/test_conv3d.py new file mode 100644 index 0000000..b62e0ee --- /dev/null +++ b/tests/test_conv3d.py @@ -0,0 +1,1620 @@ +# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC. +# Produced at the Lawrence Livermore National Laboratory. +# Written by the LBANN Research Team (B. Van Essen, et al.) listed in +# the CONTRIBUTORS file. See the top-level LICENSE file for details. +# +# LLNL-CODE-697807. +# All rights reserved. +# +# This file is part of LBANN: Livermore Big Artificial Neural Network +# Toolkit. For details, see http://software.llnl.gov/LBANN or +# https://github.com/LBANN and https://github.com/LBANN/ScaFFold. +# +# SPDX-License-Identifier: (Apache-2.0) + +"""``FastConv3d``: the rung ladder, the sharding gate, and the numerics. + +The single most important test in this file is +:func:`test_halo_plan_refuses_a_split_dim_whose_arithmetic_it_has_not_checked`. +Every other property here fails loudly; that one fails silently, as a plausible +wrong gradient at every shard boundary of a large run, because the halo DistConv +adds below autograd is invisible to a module-level adapter and the halo this one +adds instead is only right where the plan says it is. The exchange itself +needs real ranks and lives in +``work/triton-conv/review/halo/stage2_dist.py``. + +Tolerances come from ``triton_conv3d.reference``'s policy (an fp64 reference and +a dtype/K-derived bound, or MIOpen's own error where that is looser). Nothing +here invents one. +""" + +from __future__ import annotations + +import logging + +import pytest +import torch +import torch.nn as nn + +from ScaFFold.unet import conv3d as conv_mod +from ScaFFold.unet.conv3d import FastConv3d, FastConvTranspose3d +from ScaFFold.unet.unet_model import UNet +from ScaFFold.unet.unet_parts import DoubleConv, OutConv, Up + +_CHANNELS_LAST = torch.channels_last_3d + + +# --------------------------------------------------------------------------- +# helpers +# --------------------------------------------------------------------------- + + +class _StubStrategy: + """The attributes :func:`_halo_plan` reads off a strategy. + + A real ``distconv.ParallelStrategy`` calls ``dist.get_rank()`` and builds a + device mesh, so it cannot describe a >1 shard count in a one-rank test + process at all -- ``ddp_ranks`` would be ``1 // 2 == 0``. The gate reads + exactly ``num_shards``, ``shard_dim`` and ``shard_ind``, so a stub can pose + the sharded question that the environment otherwise cannot. + + ``shard_ind`` is here so that the *MIOpen* rung -- which takes a DCTensor + through DistConv's own dispatch -- still runs, which is what makes "the + sharded call went to the other rung" an assertion about routing rather than + about which stub attribute is missing. + """ + + def __init__(self, num_shards, shard_dim=(2, 3, 4)): + self.num_shards = num_shards + self.shard_dim = shard_dim + self.shard_ind = [0] * (len(num_shards) if num_shards else 0) + + def shard_to_rank(self, shard_ind): + """Every shard is this rank: there is only one, and nothing is sent.""" + return 0 + + +def _dc(tensor, num_shards=(1, 1, 1), shard_dim=(2, 3, 4)): + """A ``DCTensor`` over ``tensor`` with a stubbed strategy.""" + import distconv + + return distconv.DCTensor(tensor, _StubStrategy(num_shards, shard_dim)) + + +def _seeded_conv(cin=16, cout=32, kernel_size=3, padding=1, bias=False, **kwargs): + """A ``FastConv3d`` whose weights are not the ones a bug would guess.""" + conv = FastConv3d( + cin, cout, kernel_size=kernel_size, padding=padding, bias=bias, **kwargs + ) + generator = torch.Generator().manual_seed(1234) + with torch.no_grad(): + conv.weight.normal_(0.0, 0.1, generator=generator) + if conv.bias is not None: + conv.bias.normal_(0.0, 0.1, generator=generator) + return conv + + +def _gpu_conv(cin=16, cout=32, dtype=torch.bfloat16, **kwargs): + """The same, on GPU and in the layout ``worker.py`` puts the model in. + + ``dtype`` defaults to bf16 because outside an autocast region the operands + have to agree: an fp32 parameter against a bf16 activation is a call neither + rung serves. The autocast tests pass fp32 on purpose -- that is the state + ``worker.py`` actually leaves the model in, and reproducing the dispatcher's + cast is what makes it work. + """ + conv = _seeded_conv(cin, cout, **kwargs).cuda().to(memory_format=_CHANNELS_LAST) + return conv.to(dtype) + + +def _gpu_input(shape, dtype=torch.bfloat16, seed=7): + generator = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn(shape, device="cuda", dtype=torch.float32, generator=generator) + return x.to(dtype).contiguous(memory_format=_CHANNELS_LAST) + + +@pytest.fixture(autouse=True) +def _clean_module_state(): + """Reset the process-global latch and override around every test.""" + saved = (conv_mod._triton_override, conv_mod._triton_failed) + yield + conv_mod._triton_override, conv_mod._triton_failed = saved + + +# --------------------------------------------------------------------------- +# the sharding gate +# --------------------------------------------------------------------------- + + +def _plan(dc_input, strategy, x=None, kernel=(3, 3, 3), padding=(1, 1, 1)): + """:func:`conv_mod._halo_plan` with the operands a 3x3x3 "same" conv has.""" + if x is None: + x = torch.empty(1, 8, 8, 8, 8) + weight = torch.empty(8, 8, *kernel) + return conv_mod._halo_plan( + dc_input, strategy, x, weight, (1, 1, 1), padding, (1, 1, 1) + ) + + +def test_halo_plan_refuses_every_strategy_it_cannot_read(): + """Every branch of the gate, including the ones it cannot read. + + The asymmetry under test is the whole safety argument: a plan is produced + only when every fact has been checked, and ``None`` -- which routes the call + to MIOpen and DistConv -- for anything else. + """ + unsharded = _StubStrategy((1, 1, 1)) + + class _Input: + _is_periodic = (False, False, False) + + ok = _Input() + plan = _plan(ok, unsharded) + assert plan is not None and plan.exchanges == [] + assert plan.padding == (1, 1, 1), "an unsplit dim keeps the module's padding" + assert plan.input_shape == (1, 8, 8, 8, 8), "nothing was added to the extent" + + # An axis whose count is 1 but which is not named in shard_dim: distconv + # indexes num_shards by position in shard_dim, so a length mismatch means + # the two disagree about which axis is which and the argument is unmade. + assert _plan(ok, _StubStrategy((1, 1, 1), shard_dim=(2,))) is None + # One axis named twice would be exchanged twice. + assert _plan(ok, _StubStrategy((1, 2, 1), shard_dim=(2, 2, 4))) is None + + # Nothing readable at all. + assert _plan(ok, None) is None + assert _plan(ok, _StubStrategy(None)) is None + assert _plan(ok, _StubStrategy(())) is None + # A count that is not an int is a count this gate has not understood. + assert _plan(ok, _StubStrategy((1, 1, "1"))) is None + # A shard index outside its own axis is a strategy that does not describe a + # mesh this exchange can address. + bad_index = _StubStrategy((2, 1, 1)) + bad_index.shard_ind = [2, 0, 0] + assert _plan(ok, bad_index) is None + missing_index = _StubStrategy((2, 1, 1)) + missing_index.shard_ind = None + assert _plan(ok, missing_index) is None + + # Periodicity: one shard still exchanges with itself, so the halo is the + # opposite face rather than zeros, and the padding becomes + # _periodic_shard_padding instead of 0. + class _Periodic: + _is_periodic = (False, True, False) + + assert _plan(_Periodic(), unsharded) is None + + class _NoPeriodicAttr: + pass + + assert _plan(_NoPeriodicAttr(), unsharded) is None + + class _WrongLength: + _is_periodic = (False, False) + + assert _plan(_WrongLength(), unsharded) is None + + +def test_halo_plan_exchanges_only_the_dims_that_are_actually_split(): + """The move stage 2 exists for, and the shapes it hands the kernel. + + ScaFFold ships ``dc_shard_dims: [2, 3, 4]`` with only D ever divided, so + DistConv's halo on H and W is two ``cat`` copies of a slab that is provably + zeros. Dropping it is measured bitwise inert; this pins that the plan does + drop it, and that the split dim -- and only the split dim -- trades its + padding for a wider extent. + """ + + class _Input: + _is_periodic = (False, False, False) + + plan = _plan(_Input(), _StubStrategy((2, 1, 1))) + assert plan.exchanges == [(0, 2, 1)], "H and W were exchanged, or D was not" + assert plan.padding == (0, 1, 1), "H/W lost their ordinary padding" + assert plan.input_shape == (1, 8, 10, 8, 8) + + plan = _plan(_Input(), _StubStrategy((2, 2, 1))) + assert plan.exchanges == [(0, 2, 1), (1, 3, 1)] + assert plan.padding == (0, 0, 1) + assert plan.input_shape == (1, 8, 10, 10, 8) + + # A 5x5x5 kernel wants two rows from each neighbour. + plan = _plan( + _Input(), _StubStrategy((2, 1, 1)), kernel=(5, 5, 5), padding=(2, 2, 2) + ) + assert plan.exchanges == [(0, 2, 2)] + assert plan.input_shape == (1, 8, 12, 8, 8) + + +def test_halo_plan_refuses_a_split_dim_whose_arithmetic_it_has_not_checked(): + """The per-axis conditions, each of which would give a wrong answer. + + An unsplit dim is exempt from all of them -- its halo is zeros either way -- + which is what makes the block-list narrow enough to be worth having. + """ + + class _Input: + _is_periodic = (False, False, False) + + ok = _Input() + split = _StubStrategy((2, 1, 1)) + + # Padding that is not "same" on the split dim: the halo'd extent at padding + # 0 would not be the shard's slice of the global volume. + assert _plan(ok, split, padding=(0, 1, 1)) is None + # An even kernel gives DistConv halo_size 0 and a strided-tiling contract + # this module has not reasoned about. + assert _plan(ok, split, kernel=(2, 3, 3), padding=(0, 1, 1)) is None + # A shard thinner than the halo it must give away. + assert _plan(ok, split, x=torch.empty(1, 8, 1, 8, 8)) is None + # Stride and dilation on the split dim. + weight = torch.empty(8, 8, 3, 3, 3) + x = torch.empty(1, 8, 8, 8, 8) + assert ( + conv_mod._halo_plan(ok, split, x, weight, (2, 1, 1), (1, 1, 1), (1, 1, 1)) + is None + ) + assert ( + conv_mod._halo_plan(ok, split, x, weight, (1, 1, 1), (1, 1, 1), (2, 1, 1)) + is None + ) + # ... but the same stride on an *unsplit* dim is nothing to do with the halo. + assert ( + conv_mod._halo_plan(ok, split, x, weight, (1, 2, 1), (1, 1, 1), (1, 1, 1)) + is not None + ) + + # k == 1 on the split dim reads no neighbour voxel at all, so there is + # nothing to exchange and padding 0 is already right. + plan = _plan(ok, split, kernel=(1, 3, 3), padding=(0, 1, 1)) + assert plan is not None and plan.exchanges == [] + + # A strategy that cannot name its neighbours cannot be exchanged with. + nameless = _StubStrategy((2, 1, 1)) + nameless.shard_to_rank = None + assert _plan(ok, nameless) is None + + +@pytest.mark.gpu +def test_the_gate_asks_the_predicates_about_the_tensor_the_kernel_will_see(): + """Widened, but still checked from both ends. + + Both halves matter. A test that only checked the refusal would pass just as + well against a gate that refuses everything -- and the sharding check sits + behind the ``is_cuda`` test, so on CPU it is never even reached. So the + same module and the same tensor are asked twice, differing only in + ``num_shards``. + + The sharded answer is ``False`` here for one reason and one reason only: + this process has no process group to exchange over. The plan is made, and + the predicates are asked about the halo'd extent -- ``8 -> 10`` on D -- at + the padding the exchange leaves behind. The multi-rank half of this lives + in ``work/triton-conv/review/halo/stage2_dist.py``. + """ + conv = _gpu_conv() + x = _gpu_input((1, 16, 8, 8, 8)) + + unsharded = _dc(x, num_shards=(1, 1, 1)) + sharded = _dc(x, num_shards=(2, 1, 1)) + plans = { + name: conv_mod._halo_plan( + dc, + dc._parallel_strategy, + x, + conv.weight, + conv.stride, + conv.padding, + conv.dilation, + ) + for name, dc in (("unsharded", unsharded), ("sharded", sharded)) + } + + assert conv_mod._use_triton(conv, x, unsharded, plans["unsharded"]) is True + assert plans["sharded"] is not None, "the gate refused a strategy it can serve" + assert plans["sharded"].input_shape == (1, 16, 10, 8, 8) + assert plans["sharded"].padding == (0, 1, 1) + assert not torch.distributed.is_initialized() + assert conv_mod._use_triton(conv, x, sharded, plans["sharded"]) is False + + +@pytest.mark.gpu +def test_a_sharded_dctensor_forward_goes_to_miopen_without_a_process_group(monkeypatch): + """End to end, not just the predicate: the rung must not fire. + + Routing is asserted from both ends -- the Triton rung is not entered *and* + DistConv's halo exchange is, which is the path that supplies the neighbours' + voxels the Triton rung would otherwise have to supply itself. + ``forward_halo_exchange`` is stubbed to the identity so the MIOpen rung + completes without a process group; it is the call count that is being + measured, not the values. + """ + import distconv.distconv as dc + + conv = _gpu_conv() + x = _gpu_input((1, 16, 8, 8, 8)) + fast_calls, halo_calls = [], [] + + original = FastConv3d._triton_forward + monkeypatch.setattr( + FastConv3d, + "_triton_forward", + lambda self, local, plan=None: fast_calls.append(local) + or original(self, local, plan), + ) + + def _local_halo(tensor, halo_size, strategy, dim_index, is_periodic=False): + """What the real exchange does when nothing has to be received. + + Concatenating zero slabs is exactly ``forward_halo_exchange``'s + behaviour at one shard; spelling it out here lets the MIOpen rung run to + completion for a *sharded* strategy too, without a process group. + """ + halo_calls.append(dim_index) + if halo_size == 0: + return tensor + dim = strategy.shard_dim[dim_index] + slab = torch.zeros_like(tensor.narrow(dim, 0, halo_size)) + return torch.cat([slab, tensor, slab], dim=dim) + + monkeypatch.setattr(dc, "forward_halo_exchange", _local_halo) + + conv(_dc(x, num_shards=(1, 1, 1))) + assert len(fast_calls) == 1, "the unsharded control did not take the Triton rung" + assert halo_calls == [], "the Triton rung still paid for a halo exchange" + + conv(_dc(x, num_shards=(1, 2, 1))) + assert len(fast_calls) == 1, "a sharded DCTensor reached the Triton rung" + assert len(halo_calls) == 3, "the sharded call did not go through DistConv" + + +@pytest.mark.gpu +def test_a_kernel_failure_after_the_halo_falls_back_without_exchanging_twice( + monkeypatch, +): + """A Triton failure at 2 shards must cost speed, not the run. + + The halo goes on the wire before the kernel compiles, so a ``TritonError`` + arrives with this rank's sends and receives already matched against its + peers'. Re-running the whole call would take it to ``_miopen_forward`` and + therefore through ``distconv_forward``, which exchanges *again* -- one more + collective on this rank than on a peer whose kernel compiled, which hangs the + mesh or pairs this convolution's slabs with the next one's. Raising instead + made a broken Triton install **fatal** at ``num_shards > 1`` while costing + only speed at 1. + + So the count is the assertion, not the absence of an exception: exactly one + exchange, the adapter's, and none of DistConv's. Both are stubbed to the + "nothing to receive" form so a one-rank process can run a two-shard strategy; + it is which of them is *called* that is being measured. And because + ``cat(zeros, x, zeros)`` at padding 0 is the same arithmetic as the module's + own padding on the unexchanged shard, the answer has an independent + reference: what ``nn.Conv3d`` computes on the original input. + + The multi-rank half, with real slabs on a real mesh, is + ``work/triton-conv/review/halo/stage2_dist.py``'s ``fallback`` stage. + """ + import distconv + import distconv.distconv as dc + from triton.errors import TritonError + + conv = _gpu_conv() + x = _gpu_input((1, 16, 8, 8, 8)) + local = x.detach().clone().requires_grad_(True) + # ``from_shard`` rather than the bare constructor: it is DistConv's + # autograd-connected wrap, so the gradient below really does have to travel + # back through the exchange to reach ``local``. + strategy = _StubStrategy((2, 1, 1)) + sharded = distconv.DCTensor.from_shard(local, strategy) + plan = conv_mod._halo_plan( + sharded, + strategy, + local, + conv.weight, + conv.stride, + conv.padding, + conv.dilation, + ) + assert plan is not None and len(plan.exchanges) == 1, "D must be the split dim" + + mine, theirs = [], [] + + def _adapter_exchange(tensor, strategy, dim_index, dim, halo): + mine.append(dim) + slab = torch.zeros_like(tensor.narrow(dim, 0, halo)) + return torch.cat([slab, tensor, slab], dim=dim).contiguous( + memory_format=_CHANNELS_LAST + ) + + def _adapter_backward(grad, strategy, dim_index, dim, halo): + return grad.narrow(dim, halo, grad.size(dim) - 2 * halo) + + def _distconv_exchange(tensor, halo_size, strategy, dim_index, is_periodic=False): + theirs.append(dim_index) + if halo_size == 0: + return tensor + dim = strategy.shard_dim[dim_index] + slab = torch.zeros_like(tensor.narrow(dim, 0, halo_size)) + return torch.cat([slab, tensor, slab], dim=dim) + + monkeypatch.setattr(conv_mod, "_exchange_forward", _adapter_exchange) + monkeypatch.setattr(conv_mod, "_exchange_backward", _adapter_backward) + monkeypatch.setattr(dc, "forward_halo_exchange", _distconv_exchange) + # The gate declines a sharded plan in a process with no group, which is a + # routing condition and not the one under test. + monkeypatch.setattr(conv_mod, "_use_triton", lambda *a, **kw: True) + monkeypatch.setattr(conv_mod, "_triton_failed", False) + + def _boom(*args, **kwargs): + raise TritonError("forced: this tile does not fit in LDS") + + monkeypatch.setattr(conv_mod._get_triton_module(), "conv3d_forward", _boom) + + out = conv(sharded) + + assert isinstance(out, distconv.DCTensor), "the fallback lost the wrapper" + assert mine == [2], f"the halo was exchanged {len(mine)} times, not once" + assert theirs == [], "MIOpen was reached through DistConv, which exchanges again" + assert conv_mod._triton_failed is True, "the kernel failure did not latch" + torch.testing.assert_close(out._tensor, nn.Conv3d.forward(conv, x)) + + # And the graph the fallback built is the halo'd one: the gradient reaches + # the shard through ``_Halo3d``, so it must match the unsharded gradient. + gy = _gpu_input((1, 32, 8, 8, 8), seed=41) + out.backward(distconv.DCTensor.from_shard(gy, strategy)) + plain_x = x.detach().clone().requires_grad_(True) + nn.Conv3d.forward(conv, plain_x).backward(gy) + assert local.grad is not None, "the fallback severed the graph at the halo" + torch.testing.assert_close(local.grad.float(), plain_x.grad.float(), + rtol=2e-2, atol=2e-2) + + +# --------------------------------------------------------------------------- +# routing +# --------------------------------------------------------------------------- + + +def test_cpu_input_never_reaches_the_triton_rung(): + conv = _seeded_conv() + x = torch.randn(1, 16, 8, 8, 8) + assert conv_mod._use_triton(conv, x, None, None) is False + torch.testing.assert_close(conv(x), nn.Conv3d.forward(conv, x)) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "kwargs, why", + [ + ({"cin": 3}, "the stem: 0.93x over three directions, +0.19% of a step"), + ({"kernel_size": 1, "padding": 0, "bias": True}, "the k=1 head: 1.40x"), + ({"cout": 512}, "small M with Cout >= 512: 1.17-1.65x"), + ], +) +def test_the_block_list_is_empty_at_the_shapes_it_used_to_hold(kwargs, why): + """Each of these was kept on MIOpen until 2026-08-04; all three now route. + + Parametrized on the three retired rules rather than asserting + ``_policy_declines`` is empty, because what matters is the *routing* answer: + a rule could return ``False`` while some other clause of :func:`_use_triton` + still declined, and then the block would be gone in name only. + + ``cout=512`` also covers the small-``M`` predicate's real defect -- it read + the forward GEMM's row count and then kept all three directions on MIOpen, + including a backward-data that wins 1.31-2.47x. See + :func:`~ScaFFold.unet.conv3d._policy_declines`. + """ + conv = _gpu_conv(**kwargs) + x = _gpu_input((1, kwargs.get("cin", 16), 8, 8, 8)) + assert conv_mod._use_triton(conv, x, None, None) is True, why + + +@pytest.mark.gpu +def test_ladder_falls_back_on_a_shape_the_kernel_does_not_serve(): + """``stride=2``: the forward predicate accepts it, both backwards reject it. + + This is the concrete witness for why all three directions are gated, not + just the forward: taking the rung here would build a graph node whose + backward ``triton_conv3d`` cannot answer, and by then MIOpen is no longer an + option for it. + """ + conv = _gpu_conv(kernel_size=3, padding=1, stride=2) + x = _gpu_input((1, 16, 8, 8, 8)) + triton_conv3d = conv_mod._get_triton_module() + + probe = conv_mod._metadata_probe((1, 16, 8, 8, 8), torch.bfloat16, x.device) + w_probe = conv_mod._metadata_probe( + tuple(conv.weight.shape), torch.bfloat16, x.device + ) + assert ( + triton_conv3d.is_supported(probe, w_probe, None, (2, 2, 2), (1, 1, 1)) is True + ) + assert conv_mod._use_triton(conv, x, None, None) is False + + out = conv(x) + torch.testing.assert_close(out, nn.Conv3d.forward(conv, x)) + + +@pytest.mark.gpu +def test_metadata_probe_answers_like_a_real_tensor(): + """The stand-in shortcut, pinned against the tensors it stands in for. + + ``_metadata_probe`` exists so the gate can ask about operands that do not + exist yet (the gradient) or that would cost a full-size copy to build (the + bf16 cast of an fp32 activation). It is only sound while the predicates + read metadata and nothing else, which is a property of a package this + module does not own. + """ + triton_conv3d = conv_mod._get_triton_module() + x = _gpu_input((1, 16, 8, 8, 8)) + w = _gpu_conv().weight.detach().to(torch.bfloat16) + gy = _gpu_input((1, 32, 8, 8, 8), seed=11) + args = ((1, 1, 1), (1, 1, 1), (1, 1, 1), 1) + + px = conv_mod._metadata_probe(x.shape, x.dtype, x.device) + pw = conv_mod._metadata_probe(w.shape, w.dtype, w.device) + pgy = conv_mod._metadata_probe(gy.shape, gy.dtype, gy.device) + + assert triton_conv3d.is_supported( + px, pw, None, *args + ) == triton_conv3d.is_supported(x, w, None, *args) + assert triton_conv3d.is_supported_bwd_data( + pgy, pw, x.shape, *args + ) == triton_conv3d.is_supported_bwd_data(gy, w, x.shape, *args) + assert triton_conv3d.is_supported_bwd_weight( + px, w.shape, pgy, *args + ) == triton_conv3d.is_supported_bwd_weight(x, w.shape, gy, *args) + + +# --------------------------------------------------------------------------- +# latch / proven / opt-in +# --------------------------------------------------------------------------- + + +def test_env_var_off_declines_before_anything_else(monkeypatch): + monkeypatch.setattr(conv_mod, "_triton_override", False) + conv = _seeded_conv() + assert conv_mod._use_triton(conv, torch.randn(1, 16, 4, 4, 4), None, None) is False + + +def test_set_conv_triton_enabled_round_trips_and_clears_the_latch(monkeypatch): + monkeypatch.setattr(conv_mod, "_triton_failed", True) + previous = conv_mod.set_conv_triton_enabled(True) + try: + assert conv_mod._triton_override is True + assert conv_mod._triton_failed is False, "an explicit opt-in must re-arm" + # None restores the env default and deliberately does NOT clear a latch. + conv_mod._triton_failed = True + conv_mod.set_conv_triton_enabled(None) + assert conv_mod._triton_failed is True + assert conv_mod._triton_override is conv_mod._env_override( + conv_mod.TRITON_ENV_VAR + ) + finally: + conv_mod.set_conv_triton_enabled(previous) + + +def test_latch_spares_a_proven_module_and_demotes_the_others(monkeypatch, caplog): + """A failure latches the rung off for modules that have never used it.""" + + class _Boom(Exception): + pass + + proven = _seeded_conv() + fresh = _seeded_conv() + proven._triton_ok = True + attempts = [] + + def _fails(self, local, plan=None): + attempts.append(self) + raise _Boom("kernel is broken") + + monkeypatch.setattr(conv_mod, "_triton_kernel_failures", lambda: (_Boom,)) + monkeypatch.setattr(FastConv3d, "_triton_forward", _fails) + monkeypatch.setattr(conv_mod, "_use_triton", lambda module, *a, **kw: True) + monkeypatch.setattr(conv_mod, "_triton_failed", False) + + x = torch.randn(1, 16, 4, 4, 4) + with caplog.at_level(logging.WARNING): + torch.testing.assert_close(proven(x), nn.Conv3d.forward(proven, x)) + assert attempts == [proven] + assert conv_mod._triton_failed is True, "the failure did not latch" + assert any("Triton conv3d failed" in r.message for r in caplog.records) + + # The latch is consulted by the real predicate, which the stub above + # replaced. Restore it and ask directly: a module that has never used the + # rung is now declined, and one that has is not. + monkeypatch.undo() + monkeypatch.setattr(conv_mod, "_triton_failed", True) + assert conv_mod._use_triton(fresh, x, None, None, proven=False) is False + # `proven=True` gets past the latch and is only declined further down, on + # the CPU check -- which is what the second half of the claim needs. + assert conv_mod._triton_failed is True + + +@pytest.mark.gpu +def test_a_proven_module_re_raises_rather_than_flipping_rungs_mid_backward(monkeypatch): + """The fallback is declined where it would corrupt instead of degrade. + + A module already proven on the rung, failing while an autograd graph task is + in flight, is answering a checkpoint recompute of a forward that ran on + Triton. Handing back MIOpen's result puts a differently-structured tensor + into a slot the graph node already holds; the honest answer is the original + exception. + """ + + class _Boom(Exception): + pass + + conv = _gpu_conv() + conv._triton_ok = True + monkeypatch.setattr(conv_mod, "_triton_kernel_failures", lambda: (_Boom,)) + monkeypatch.setattr(conv_mod, "_use_triton", lambda *a, **kw: True) + monkeypatch.setattr( + FastConv3d, + "_triton_forward", + lambda self, local, plan=None: (_ for _ in ()).throw(_Boom()), + ) + + x = _gpu_input((1, 16, 8, 8, 8)).float().requires_grad_(True) + seen = {} + + class _Probe(torch.autograd.Function): + @staticmethod + def forward(ctx, t): + return t.clone() + + @staticmethod + def backward(ctx, g): + # Inside a graph task: this is where a recompute would run. + try: + conv(g) + except _Boom: + seen["raised"] = True + return g + + _Probe.apply(x).sum().backward() + assert seen.get("raised") is True + + +# --------------------------------------------------------------------------- +# state dict / model wiring +# --------------------------------------------------------------------------- + + +def test_state_dict_matches_a_plain_conv3d_model(): + """No new keys, no renamed keys, no buffers -- checkpoints are unaffected.""" + fast = UNet( + n_channels=3, n_classes=4, trilinear=False, layers=1, group_norm_groups=2 + ) + fast_keys = list(fast.state_dict()) + + original = nn.Conv3d + try: + # Build the same model with stock convolutions by making the factories' + # classes the stock ones for the duration. Both of them: a transposed + # parameter that changed name or shape would be just as invisible here + # as an ordinary one. + import ScaFFold.unet.unet_parts as parts + + parts.FastConv3d = nn.Conv3d + parts.FastConvTranspose3d = nn.ConvTranspose3d + plain = UNet( + n_channels=3, n_classes=4, trilinear=False, layers=1, group_norm_groups=2 + ) + finally: + parts.FastConv3d = FastConv3d + parts.FastConvTranspose3d = FastConvTranspose3d + assert original is nn.Conv3d + + assert fast_keys == list(plain.state_dict()) + for key, value in fast.state_dict().items(): + assert value.shape == plain.state_dict()[key].shape + + +def test_checkpoint_round_trips_between_fast_and_plain_convolutions(): + fast = _seeded_conv(cin=8, cout=8) + plain = nn.Conv3d(8, 8, kernel_size=3, padding=1, bias=False) + plain.load_state_dict(fast.state_dict()) + torch.testing.assert_close(plain.weight, fast.weight) + + back = FastConv3d(8, 8, kernel_size=3, padding=1, bias=False) + back.load_state_dict(plain.state_dict()) + torch.testing.assert_close(back.weight, fast.weight) + + x = torch.randn(1, 8, 6, 6, 6) + torch.testing.assert_close(back(x), plain(x)) + + +def test_every_upsampler_in_the_model_is_a_fastconvtranspose3d(): + """The census: all four decoder sites, and none of them a plain module. + + A rung that is wired in but never reached is the failure this pins -- it + costs nothing, breaks nothing and shows up only as a benchmark that did not + get faster. ``layers=4`` is the shipped depth, so four is the real count. + """ + up = Up(16, 8, group_norm_groups=2, trilinear=False) + assert type(up.up) is FastConvTranspose3d + assert not isinstance(up.up, FastConv3d), "the two ladders are separate classes" + + model = UNet( + n_channels=3, n_classes=6, trilinear=False, layers=4, group_norm_groups=8 + ) + transposed = [ + m for m in model.modules() if isinstance(m, nn.modules.conv._ConvTransposeNd) + ] + assert len(transposed) == 4 + assert all(type(m) is FastConvTranspose3d for m in transposed) + # Every one of them has a bias, which is why ``grad_bias`` is a live path in + # this ladder and a test-only one in the other. + assert all(m.bias is not None for m in transposed) + + +def test_every_plain_convolution_in_the_model_is_a_fastconv3d(): + model = UNet( + n_channels=3, n_classes=6, trilinear=False, layers=4, group_norm_groups=8 + ) + plain = [ + m + for m in model.modules() + if isinstance(m, nn.Conv3d) and not isinstance(m, FastConv3d) + ] + assert plain == [] + assert sum(isinstance(m, FastConv3d) for m in model.modules()) == 19 + assert isinstance(DoubleConv(4, 4, 2).double_conv[0], FastConv3d) + assert isinstance(OutConv(4, 2).conv, FastConv3d) + + +@pytest.mark.gpu +def test_a_transposed_module_would_be_declined_even_if_one_were_wrapped(): + """The class is a public drop-in, so it checks rather than assumes.""" + conv = _gpu_conv() + conv.transposed = True + assert conv_mod._use_triton(conv, _gpu_input((1, 16, 8, 8, 8)), None, None) is False + + +# --------------------------------------------------------------------------- +# numerics +# --------------------------------------------------------------------------- + + +def _problem(cin, cout, spatial, kernel=(3, 3, 3), padding=(1, 1, 1), bias=False): + from triton_conv3d.shapes import ConvProblem + + return ConvProblem( + name=f"{cin}->{cout} k{kernel[0]} {spatial}", + cin=cin, + cout=cout, + spatial=spatial, + kernel=kernel, + padding=padding, + bias=bias, + dtype="bf16", + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize( + "cin, cout, spatial", [(16, 32, (8, 8, 8)), (64, 64, (6, 10, 10))] +) +def test_forward_and_gradients_match_nn_conv3d(cin, cout, spatial): + """All three directions against an fp64 reference, at MIOpen's own standard. + + ``assert_close`` applies the stricter of ``triton_conv3d``'s dtype/K-derived + bound and "no worse than MIOpen by more than 4x"; the incumbent's error is + measured here from the ``nn.Conv3d`` route this module replaces, which is + exactly the comparison the wiring has to survive. + """ + from triton_conv3d import reference as ref + + problem = _problem(cin, cout, spatial) + conv = _gpu_conv(cin, cout) + x = _gpu_input((1, cin, *spatial)) + gy = _gpu_input((1, cout, *spatial), seed=23) + + fast_x = x.detach().clone().requires_grad_(True) + plain_x = x.detach().clone().requires_grad_(True) + plain = nn.Conv3d(cin, cout, kernel_size=3, padding=1, bias=False).cuda() + plain = plain.to(torch.bfloat16).to(memory_format=_CHANNELS_LAST) + with torch.no_grad(): + plain.weight.copy_(conv.weight) + + assert conv_mod._use_triton(conv, fast_x, None, None) is True + y = conv(fast_x) + y_plain = plain(plain_x) + y.backward(gy) + y_plain.backward(gy) + + operands = { + "input": x, + "weight": conv.weight.detach(), + "bias": None, + "grad_output": gy, + } + for direction, actual, incumbent in ( + ("fwd", y, y_plain), + ("bwd-data", fast_x.grad, plain_x.grad), + ("bwd-weight", conv.weight.grad, plain.weight.grad), + ): + expected = ref.reference(problem, operands, direction) + incumbent_error = ref.compare(incumbent, expected) + ref.assert_close( + actual, expected, problem, direction, incumbent_error=incumbent_error + ) + + +@pytest.mark.gpu +def test_bias_gradient_is_correct_even_though_the_head_is_blocklisted(): + """``is_supported`` accepts a bias, so the node has to produce its gradient.""" + from triton_conv3d import reference as ref + + cin, cout, spatial = 16, 32, (8, 8, 8) + problem = _problem(cin, cout, spatial, bias=True) + conv = _gpu_conv(cin, cout, bias=True) + x = _gpu_input((1, cin, *spatial)).requires_grad_(True) + gy = _gpu_input((1, cout, *spatial), seed=31) + + conv(x).backward(gy) + + operands = { + "input": x.detach(), + "weight": conv.weight.detach(), + "bias": conv.bias.detach(), + "grad_output": gy, + } + expected = ref.reference(problem, operands, "fwd") + del expected # the forward is covered above; here only d(bias) is new. + expected_gb = gy.to(torch.float64).sum(dim=(0, 2, 3, 4)) + torch.testing.assert_close( + conv.bias.grad.to(torch.float64), expected_gb, rtol=2e-2, atol=2e-2 + ) + + +@pytest.mark.gpu +def test_autocast_runs_the_kernel_at_the_dtype_aten_would_have_chosen(): + """The cast ATen does in the dispatcher, reproduced above it. + + Without this the module's fp32 parameters and GroupNorm's fp32 output would + be handed straight to the kernel and the whole network's convolutions would + quietly run in fp32 -- a different computation from the benchmark's, and a + much slower one. + """ + conv = _gpu_conv(16, 32, dtype=torch.float32) # as worker.py builds them + x = _gpu_input((1, 16, 8, 8, 8), dtype=torch.float32).requires_grad_(True) + + seen = {} + original = conv_mod._TritonConv3dFn.apply + + def _spy(x_, w_, b_, *rest): + seen["x"] = x_.dtype + seen["w"] = w_.dtype + return original(x_, w_, b_, *rest) + + conv_mod._TritonConv3dFn.apply = staticmethod(_spy) + try: + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + y = conv(x) + finally: + conv_mod._TritonConv3dFn.apply = original + + assert seen == {"x": torch.bfloat16, "w": torch.bfloat16} + assert y.dtype is torch.bfloat16 + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + expected = nn.functional.conv3d(x, conv.weight, None, 1, 1) + assert expected.dtype is torch.bfloat16 + + y.sum().backward() + # The cast is an ordinary autograd node, so the parameter gradient comes + # back at the parameter's own dtype, exactly as it does on the MIOpen rung. + assert conv.weight.grad.dtype is torch.float32 + assert x.grad.dtype is torch.float32 + + +@pytest.mark.gpu +@pytest.mark.parametrize("guard", ["no_grad", "inference_mode"]) +def test_the_rung_serves_the_evaluation_path(guard): + """``evaluate`` runs the whole model under ``@torch.inference_mode()``. + + That is a different autograd state from training -- ``Function.apply`` never + builds a node and the tensors it produces are inference tensors -- and it is + every validation epoch of every run, so it is not an edge case. + """ + conv = _gpu_conv(16, 32, dtype=torch.float32) + x = _gpu_input((1, 16, 8, 8, 8), dtype=torch.float32) + with ( + getattr(torch, guard)(), + torch.autocast(device_type="cuda", dtype=torch.bfloat16), + ): + y = conv(x) + expected = nn.functional.conv3d(x, conv.weight, None, 1, 1) + assert conv._triton_ok is True, "the evaluation path did not take the rung" + assert y.shape == expected.shape and y.dtype is expected.dtype + torch.testing.assert_close(y.float(), expected.float(), rtol=2e-2, atol=2e-2) + + +@pytest.mark.gpu +def test_a_checkpointed_block_recomputes_on_the_same_rung(): + """``activation_checkpointing`` is a shipped config key. + + The recompute runs inside the backward pass and its saved tensors are + compared against the original forward's. What this pins is that a module + stays on one rung across the two, which is what the ``proven`` flag exists + for: a flip is invisible to torch's metadata check and fails later, inside + DistConv, with a message about neither checkpointing nor the rung. + """ + import torch.utils.checkpoint as cp + + conv = _gpu_conv() + x = _gpu_input((1, 16, 8, 8, 8)).requires_grad_(True) + plain_x = x.detach().clone().requires_grad_(True) + + y = cp.checkpoint(conv, x, use_reentrant=False) + y.sum().backward() + assert conv._triton_ok is True + + reference = nn.Conv3d.forward(conv, plain_x) + reference.sum().backward() + torch.testing.assert_close( + x.grad.float(), plain_x.grad.float(), rtol=2e-2, atol=2e-2 + ) + + +@pytest.mark.gpu +def test_dctensor_forward_and_backward_agree_with_the_distconv_route(): + """The unwrap/rewrap at one shard, against DistConv's own halo path. + + This is the equivalence the whole gate rests on: at ``num_shards=(1,1,1)`` + the halo slabs are provably zeros, so running the kernel on the local shard + at the module's own padding must reproduce what DistConv's dispatch computes + on the halo'd tensor at zero padding. + """ + import distconv + + from triton_conv3d import reference as ref + + cin, cout, spatial = 16, 32, (8, 8, 8) + problem = _problem(cin, cout, spatial) + conv = _gpu_conv(cin, cout) + plain = nn.Conv3d(cin, cout, kernel_size=3, padding=1, bias=False) + plain = plain.cuda().to(torch.bfloat16).to(memory_format=_CHANNELS_LAST) + with torch.no_grad(): + plain.weight.copy_(conv.weight) + + x = _gpu_input((1, cin, *spatial)) + gy = _gpu_input((1, cout, *spatial), seed=41) + + fast_x = x.detach().clone().requires_grad_(True) + plain_x = x.detach().clone().requires_grad_(True) + strategy = _StubStrategy((1, 1, 1)) + + y = conv(distconv.DCTensor.from_shard(fast_x, strategy)) + assert isinstance(y, distconv.DCTensor) + assert conv._triton_ok is True, "the fast rung did not serve the DCTensor" + y_plain = plain(distconv.DCTensor.from_shard(plain_x, strategy)) + + y.backward(distconv.DCTensor.from_shard(gy, strategy)) + y_plain.backward(distconv.DCTensor.from_shard(gy, strategy)) + + operands = { + "input": x, + "weight": conv.weight.detach(), + "bias": None, + "grad_output": gy, + } + for direction, actual, incumbent in ( + ("fwd", y._tensor, y_plain._tensor), + ("bwd-data", fast_x.grad, plain_x.grad), + ("bwd-weight", conv.weight.grad, plain.weight.grad), + ): + expected = ref.reference(problem, operands, direction) + ref.assert_close( + actual, + expected, + problem, + direction, + incumbent_error=ref.compare(incumbent, expected), + ) + + +@pytest.mark.gpu +def test_backward_falls_back_to_miopen_when_the_kernel_direction_fails(monkeypatch): + """A backward-direction failure degrades; the saved set cannot change.""" + from triton_conv3d import reference as ref + + class _Boom(Exception): + pass + + cin, cout, spatial = 16, 32, (8, 8, 8) + problem = _problem(cin, cout, spatial) + conv = _gpu_conv(cin, cout) + x = _gpu_input((1, cin, *spatial)) + gy = _gpu_input((1, cout, *spatial), seed=53) + fast_x = x.detach().clone().requires_grad_(True) + + y = conv(fast_x) + module = conv_mod._get_triton_module() + monkeypatch.setattr(conv_mod, "_triton_kernel_failures", lambda: (_Boom,)) + monkeypatch.setattr( + module, + "conv3d_backward_data", + lambda *a, **kw: (_ for _ in ()).throw(_Boom()), + raising=False, + ) + monkeypatch.setattr(conv_mod, "_triton_failed", False) + y.backward(gy) + + assert conv_mod._triton_failed is True + operands = { + "input": x, + "weight": conv.weight.detach(), + "bias": None, + "grad_output": gy, + } + expected = ref.reference(problem, operands, "bwd-data") + ref.assert_close(fast_x.grad, expected, problem, "bwd-data") + + +class _FakeCtx: + """Just enough ``ctx`` to call ``_TritonConv3dFn.backward`` directly.""" + + def __init__(self, saved): + self.saved_tensors = saved + self.conv_args = ((1, 1, 1), (1, 1, 1), (1, 1, 1), False) + self.needs_input_grad = (True, True, False, False, False, False) + + +def test_backward_names_a_rung_flip_instead_of_dying_inside_distconv(): + """The only detector for a flip torch's checkpoint metadata check misses. + + ``_default_meta_extractor`` compares shape, dtype and device, all of which + the two rungs agree on, so a subclass landing in slot 0 gets through torch's + own check and then fails somewhere else entirely. + """ + + class _Wrapper(torch.Tensor): + pass + + x = torch.randn(1, 8, 4, 4, 4) + weight = torch.randn(8, 8, 3, 3, 3) + with pytest.raises(RuntimeError, match="served by different"): + conv_mod._TritonConv3dFn.backward( + _FakeCtx((_Wrapper(x), weight)), torch.randn(1, 8, 4, 4, 4) + ) + + +@pytest.mark.gpu +def test_the_triton_rung_performs_no_halo_exchange_at_one_shard(): + """Not merely correct without the halo -- it must not pay for one either. + + ``forward_halo_exchange`` has no ``num_shards == 1`` early-out, so today + every convolution concatenates two zero slabs onto each of the three + sharded dims: 54 calls and 3.797 ms/step of pure copying at + ``dc_num_shards=(1,1,1)`` (``work/upstream-repros/DISTCONV_ACTIONS.md``). Taking the + Triton rung removes all of them, and leaves the caller's ``_tensor`` + un-narrowed and still channels-last for the consumers downstream of it. + """ + import distconv + import distconv.distconv as dc + + conv = _gpu_conv() + x = _gpu_input((1, 16, 8, 8, 8)).requires_grad_(True) + dc_input = distconv.DCTensor.from_shard(x, _StubStrategy((1, 1, 1))) + + calls = [] + original = dc.forward_halo_exchange + dc.forward_halo_exchange = lambda *a, **kw: calls.append(a) or original(*a, **kw) + try: + out = conv(dc_input) + finally: + dc.forward_halo_exchange = original + + assert conv._triton_ok is True + assert calls == [], "the Triton rung went through DistConv's halo exchange" + assert dc_input._tensor_with_halo is None + assert dc_input._tensor.is_contiguous(memory_format=_CHANNELS_LAST) + assert isinstance(out, distconv.DCTensor) + + +# --------------------------------------------------------------------------- +# the transposed ladder +# --------------------------------------------------------------------------- + + +def _seeded_convT(cin=16, cout=8, kernel_size=2, stride=2, **kwargs): + """A ``FastConvTranspose3d`` whose weights are not the ones a bug would guess. + + ``bias`` is left at ``nn.ConvTranspose3d``'s default of ``True``, which is + what the four decoder sites have and what puts ``grad_bias`` on the live + path. + """ + conv = FastConvTranspose3d(cin, cout, kernel_size=kernel_size, stride=stride, + **kwargs) + generator = torch.Generator().manual_seed(4321) + with torch.no_grad(): + conv.weight.normal_(0.0, 0.1, generator=generator) + if conv.bias is not None: + conv.bias.normal_(0.0, 0.1, generator=generator) + return conv + + +def _gpu_convT(cin=16, cout=8, dtype=torch.bfloat16, **kwargs): + conv = _seeded_convT(cin, cout, **kwargs).cuda().to(memory_format=_CHANNELS_LAST) + return conv.to(dtype) + + +def _stock_like(conv, dtype=torch.bfloat16): + """A stock ``nn.ConvTranspose3d`` holding the same parameters.""" + cin, cout = int(conv.weight.shape[0]), int(conv.weight.shape[1]) + plain = nn.ConvTranspose3d( + cin, cout, kernel_size=conv.kernel_size, stride=conv.stride, + bias=conv.bias is not None, + ) + plain = plain.cuda().to(dtype).to(memory_format=_CHANNELS_LAST) + with torch.no_grad(): + plain.weight.copy_(conv.weight) + if conv.bias is not None: + plain.bias.copy_(conv.bias) + return plain + + +def _transposed_problem(cin, cout, spatial, kernel=(2, 2, 2), bias=True): + from triton_conv3d.shapes import ConvProblem + + return ConvProblem( + name=f"convT {cin}->{cout} k{kernel[0]} {spatial}", + cin=cin, + cout=cout, + spatial=spatial, + kernel=kernel, + stride=kernel, + padding=(0, 0, 0), + transposed=True, + bias=bias, + dtype="bf16", + ) + + +#: ``(x_shape, weight_shape)`` of the four decoder upsamplers at config A +#: (scale 7, 128^3, ``layers=4``), in the order the decoder runs them. +_UPSAMPLER_SITES = [ + ((1, 1024, 8, 8, 8), (1024, 512, 2, 2, 2)), + ((1, 512, 16, 16, 16), (512, 256, 2, 2, 2)), + ((1, 256, 32, 32, 32), (256, 128, 2, 2, 2)), + ((1, 128, 64, 64, 64), (128, 64, 2, 2, 2)), +] + + +def test_the_transposed_block_list_is_empty_at_every_decoder_site(): + """No decoder site is blocked, in either ladder. + + Both block-lists are empty as of 2026-08-04, so this asserts the routing + answer rather than the shape of a rule. It is still worth a test: emptiness + is a claim about *measurements*, and a future entry in either function has to + re-establish it here. + + The reason the two functions stay separate survives the emptying, and is + recorded in :func:`~ScaFFold.unet.conv3d._transposed_policy_declines`: every + term the ordinary rule used reads a different quantity for this operator -- + ``w_shape``'s channel axes are reversed, and ``M`` from ``_out_spatial`` is + the input volume over 8 at ``k == s == 2``. The retired small-``M`` rule + answered ``True`` for ``up1`` on numbers that do not describe it. + """ + for x_shape, w_shape in _UPSAMPLER_SITES: + assert conv_mod._transposed_policy_declines(x_shape, w_shape) is False + assert conv_mod._policy_declines( + x_shape, w_shape, (2, 2, 2), (0, 0, 0), (1, 1, 1) + ) is False + + +@pytest.mark.gpu +def test_the_transposed_gate_answers_for_the_four_sites_and_refuses_the_rest(): + """The gate fires where it must, and declines what the kernels do not serve. + + Every refusal below is a condition ``triton_conv3d.transposed`` states in + its own gate; asking through ``_use_triton_transposed`` is what pins that + this module *asks* -- with the module's real ``stride``, ``padding``, + ``output_padding``, ``dilation`` and ``groups``, in the right slots. + """ + conv = _gpu_convT(16, 8) + x = _gpu_input((1, 16, 8, 8, 8)) + assert conv_mod._use_triton_transposed(conv, x, None, None) is True + + # k != s: the windows overlap and the bijection this module rests on is + # gone. The gate must not read the module's kernel_size as its stride. + assert ( + conv_mod._use_triton_transposed(_gpu_convT(16, 8, kernel_size=3), x, None, None) + is False + ) + # A padding crops the result and an output_padding extends it + # asymmetrically; both break the tiling. + assert ( + conv_mod._use_triton_transposed(_gpu_convT(16, 8, padding=1), x, None, None) + is False + ) + assert ( + conv_mod._use_triton_transposed( + _gpu_convT(16, 8, stride=3, output_padding=1), _gpu_input((1, 16, 8, 8, 8)), + None, None, + ) + is False + ) + # groups > 1 has no coverage in any direction. + assert ( + conv_mod._use_triton_transposed( + _gpu_convT(16, 8, groups=2), x, None, None + ) + is False + ) + # NCDHW would be a full-size hidden relayout, which is the cost the rung + # exists to avoid. + assert ( + conv_mod._use_triton_transposed(conv, x.contiguous(), None, None) is False + ) + # And the CPU, where there is no kernel at all. + assert conv_mod._use_triton_transposed(_seeded_convT(), torch.randn(1, 16, 4, 4, 4), + None, None) is False + + +@pytest.mark.gpu +def test_the_transposed_gate_asked_is_the_one_that_covers_the_backward(monkeypatch): + """``is_supported_transposed_all``, not the forward's gate alone. + + The three transposed predicates accept the same problems today, so no shape + can tell them apart -- which is exactly why *which one is called* has to be + pinned directly. A forward this package serves and a backward it cannot is + discovered inside ``backward()``, where MIOpen is no longer reachable, and + the ordinary convolution has a live witness for that (``stride > 1``). + """ + conv = _gpu_convT(16, 8) + x = _gpu_input((1, 16, 8, 8, 8)) + module = conv_mod._get_triton_module() + + assert conv_mod._use_triton_transposed(conv, x, None, None) is True + # Patching the package attribute reaches the module's call and not the one + # ``is_supported_transposed_all`` makes internally, so this distinguishes + # "asked the combined gate" from "asked the forward's and got the same + # answer" -- which is the only way to tell them apart while they agree. + monkeypatch.setattr(module, "is_supported_transposed_all", lambda *a, **kw: False, + raising=False) + assert conv_mod._use_triton_transposed(conv, x, None, None) is False + + +@pytest.mark.gpu +def test_a_non_transposed_module_would_be_declined_by_the_transposed_gate(): + """The mirror of ``test_a_transposed_module_would_be_declined``. + + ``is_supported_transposed`` reads ``w.shape[0]`` as ``Cin`` and + ``w.shape[1]`` as ``Cout``; an ``nn.Conv3d`` weight stores them the other way + round, so a square one would be *accepted* and would compute the wrong + operator without raising. This module checks ``transposed`` rather than + trusting its own construction, because the class is a public drop-in. + """ + conv = _gpu_convT(16, 16) + conv.transposed = False + assert conv_mod._use_triton_transposed(conv, _gpu_input((1, 16, 8, 8, 8)), None, + None) is False + + +@pytest.mark.gpu +@pytest.mark.parametrize("cin, cout, spatial", [(16, 8, (8, 8, 8)), (32, 16, (4, 6, 6))]) +def test_transposed_forward_and_gradients_match_nn_convtranspose3d(cin, cout, spatial): + """All three directions against an fp64 reference, at MIOpen's own standard. + + ``assert_close`` applies the stricter of ``triton_conv3d``'s dtype/K-derived + bound and "no worse than MIOpen by more than 4x", with the incumbent's error + measured from the stock ``nn.ConvTranspose3d`` route this module replaces. + """ + from triton_conv3d import reference as ref + + problem = _transposed_problem(cin, cout, spatial) + conv = _gpu_convT(cin, cout) + plain = _stock_like(conv) + out_spatial = tuple(2 * s for s in spatial) + + x = _gpu_input((1, cin, *spatial)) + gy = _gpu_input((1, cout, *out_spatial), seed=23) + fast_x = x.detach().clone().requires_grad_(True) + plain_x = x.detach().clone().requires_grad_(True) + + assert conv_mod._use_triton_transposed(conv, fast_x, None, None) is True + y = conv(fast_x) + y_plain = plain(plain_x) + assert conv._triton_ok is True, "the rung did not serve the call" + y.backward(gy) + y_plain.backward(gy) + + operands = { + "input": x, + "weight": conv.weight.detach(), + "bias": conv.bias.detach(), + "grad_output": gy, + } + for direction, actual, incumbent in ( + ("fwd", y, y_plain), + ("bwd-data", fast_x.grad, plain_x.grad), + ("bwd-weight", conv.weight.grad, plain.weight.grad), + ): + expected = ref.reference(problem, operands, direction) + ref.assert_close( + actual, + expected, + problem, + direction, + incumbent_error=ref.compare(incumbent, expected), + ) + # grad_bias is not one of ``reference``'s directions -- it is not a + # convolution -- so it gets the exact answer and the incumbent's own error + # as its bar, which is the same standard by a different route. + expected_gb = gy.to(torch.float64).sum(dim=(0, 2, 3, 4)) + incumbent_gb = (plain.bias.grad.to(torch.float64) - expected_gb).abs().max().item() + actual_gb = (conv.bias.grad.to(torch.float64) - expected_gb).abs().max().item() + assert conv.bias.grad.dtype is conv.bias.dtype + assert actual_gb <= max(4.0 * incumbent_gb, 2.0**-8 * expected_gb.abs().max().item()) + + +@pytest.mark.gpu +def test_autocast_runs_the_transposed_kernel_at_the_dtype_aten_would_have_chosen(): + """``conv_transpose3d`` carries the same ``lower_precision_fp`` policy. + + Verified two ways here: that the operands reaching the node are bf16 when the + module's own parameters are fp32 (which is the state ``worker.py`` leaves the + model in), and that the stock op under the same region produces the same + dtype. Without this the four upsamplers would quietly run in fp32 -- a + different computation from the benchmark's, several times slower, and + nothing failing. + """ + conv = _gpu_convT(16, 8, dtype=torch.float32) # as worker.py builds them + x = _gpu_input((1, 16, 8, 8, 8), dtype=torch.float32).requires_grad_(True) + + seen = {} + original = conv_mod._TritonConvTranspose3dFn.apply + + def _spy(x_, w_, b_, *rest): + seen["x"], seen["w"], seen["b"] = x_.dtype, w_.dtype, b_.dtype + return original(x_, w_, b_, *rest) + + conv_mod._TritonConvTranspose3dFn.apply = staticmethod(_spy) + try: + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + y = conv(x) + finally: + conv_mod._TritonConvTranspose3dFn.apply = original + + assert seen == {"x": torch.bfloat16, "w": torch.bfloat16, "b": torch.bfloat16} + assert y.dtype is torch.bfloat16 + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + expected = nn.functional.conv_transpose3d(x, conv.weight, conv.bias, 2, 0) + assert expected.dtype is torch.bfloat16, "the op's autocast policy changed" + torch.testing.assert_close(y.float(), expected.float(), rtol=2e-2, atol=2e-2) + + y.sum().backward() + # The casts are ordinary autograd nodes, so every gradient comes back at its + # parameter's own dtype, exactly as it does on the MIOpen rung. + assert conv.weight.grad.dtype is torch.float32 + assert conv.bias.grad.dtype is torch.float32 + assert x.grad.dtype is torch.float32 + + +@pytest.mark.gpu +@pytest.mark.parametrize("guard", ["no_grad", "inference_mode"]) +def test_the_transposed_rung_serves_the_evaluation_path(guard): + """``evaluate`` runs the whole model under ``@torch.inference_mode()``.""" + conv = _gpu_convT(16, 8, dtype=torch.float32) + x = _gpu_input((1, 16, 8, 8, 8), dtype=torch.float32) + with ( + getattr(torch, guard)(), + torch.autocast(device_type="cuda", dtype=torch.bfloat16), + ): + y = conv(x) + expected = nn.functional.conv_transpose3d(x, conv.weight, conv.bias, 2, 0) + assert conv._triton_ok is True, "the evaluation path did not take the rung" + assert y.shape == expected.shape and y.dtype is expected.dtype + torch.testing.assert_close(y.float(), expected.float(), rtol=2e-2, atol=2e-2) + + +@pytest.mark.gpu +def test_a_checkpointed_upsampler_recomputes_on_the_same_rung(): + """``activation_checkpointing`` is a shipped config key. + + The hazard is wider here than for ``FastConv3d``: this ladder never adds a + halo, so the tensor the Triton rung saves and the ``DCTensor`` the MIOpen + rung saves agree on shape, dtype and device at *every* shard count, and a + flip would pass ``_default_meta_extractor``'s check silently. + """ + import torch.utils.checkpoint as cp + + conv = _gpu_convT() + x = _gpu_input((1, 16, 8, 8, 8)).requires_grad_(True) + plain_x = x.detach().clone().requires_grad_(True) + + y = cp.checkpoint(conv, x, use_reentrant=False) + y.sum().backward() + assert conv._triton_ok is True + + nn.ConvTranspose3d.forward(conv, plain_x).sum().backward() + torch.testing.assert_close( + x.grad.float(), plain_x.grad.float(), rtol=2e-2, atol=2e-2 + ) + + +def test_transposed_backward_names_a_rung_flip_instead_of_dying_inside_distconv(): + """The only detector for a flip torch's checkpoint metadata check misses.""" + + class _Wrapper(torch.Tensor): + pass + + class _Ctx: + saved_tensors = (_Wrapper(torch.randn(1, 8, 4, 4, 4)), torch.randn(8, 8, 2, 2, 2)) + conv_args = ((2, 2, 2), (0, 0, 0), (0, 0, 0), (1, 1, 1), True) + needs_input_grad = (True, True, True, False, False, False, False) + + with pytest.raises(RuntimeError, match="served by different"): + conv_mod._TritonConvTranspose3dFn.backward( + _Ctx(), torch.randn(1, 8, 8, 8, 8) + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("direction", ["conv_transpose3d_backward_data", + "conv_transpose3d_backward_weight"]) +def test_transposed_backward_falls_back_to_miopen_when_a_direction_fails( + monkeypatch, direction +): + """A backward-direction failure degrades; the saved set cannot change. + + Both directions, because they are separate compilations from the forward's + and from each other, so either can raise on a call whose forward compiled. + """ + from triton_conv3d import reference as ref + + class _Boom(Exception): + pass + + cin, cout, spatial = 16, 8, (8, 8, 8) + problem = _transposed_problem(cin, cout, spatial) + conv = _gpu_convT(cin, cout) + plain = _stock_like(conv) + x = _gpu_input((1, cin, *spatial)) + gy = _gpu_input((1, cout, 16, 16, 16), seed=53) + fast_x = x.detach().clone().requires_grad_(True) + plain_x = x.detach().clone().requires_grad_(True) + + y = conv(fast_x) + module = conv_mod._get_triton_module() + monkeypatch.setattr(conv_mod, "_triton_kernel_failures", lambda: (_Boom,)) + monkeypatch.setattr( + module, direction, lambda *a, **kw: (_ for _ in ()).throw(_Boom()), + raising=False, + ) + monkeypatch.setattr(conv_mod, "_triton_failed", False) + y.backward(gy) + plain(plain_x).backward(gy) + + assert conv_mod._triton_failed is True + operands = { + "input": x, + "weight": conv.weight.detach(), + "bias": conv.bias.detach(), + "grad_output": gy, + } + # The incumbent's own error is the bar, and it has to be: MIOpen's + # transposed backward-weight exceeds the static bound at this shape (0.76 + # against 0.64), which is the observation ``error_bound``'s docstring + # records. What this direction is being asked is whether the fallback ran + # the right operator with the right operands, not whether MIOpen is accurate. + for name, actual, incumbent in ( + ("bwd-data", fast_x.grad, plain_x.grad), + ("bwd-weight", conv.weight.grad, plain.weight.grad), + ): + expected = ref.reference(problem, operands, name) + ref.assert_close(actual, expected, problem, name, + incumbent_error=ref.compare(incumbent, expected)) + # The bias gradient comes from the aten call on this path, not from the sum + # above it, so it is the one this test would otherwise never look at. + torch.testing.assert_close( + conv.bias.grad.to(torch.float64), + gy.to(torch.float64).sum(dim=(0, 2, 3, 4)), + rtol=2e-2, atol=2e-2, + ) + + +def test_the_transposed_halo_plan_exchanges_nothing_and_refuses_what_it_cannot_read(): + """The plan is "exchange nothing" -- at every shard count, or not at all. + + The first assertion is the one that matters at scale: ``num_shards > 1`` + must produce a plan with an *empty* ``exchanges``, not merely a plan. If it + silently produced one at 1 shard and ``None`` at 2, the four sites would go + back to MIOpen on every multi-GPU run and nothing would say so. + """ + x = torch.empty(1, 8, 8, 8, 8) + weight = torch.empty(8, 4, 2, 2, 2) + + def plan(num_shards, shard_dim=(2, 3, 4), w=weight, padding=(0, 0, 0)): + strategy = _StubStrategy(num_shards, shard_dim) + return conv_mod._transposed_halo_plan( + _dc(x, num_shards, shard_dim), strategy, x, w, padding + ) + + for num_shards in ((1, 1, 1), (2, 1, 1), (4, 1, 1), (2, 2, 2)): + made = plan(num_shards) + assert made is not None, num_shards + assert made.exchanges == (), num_shards + assert made.padding == (0, 0, 0) + assert made.input_shape == tuple(x.shape) + + # An odd kernel on a split dim: DistConv would want a k//2 halo there and + # then refuse the problem outright, so there is no incumbent to agree with. + assert plan((2, 1, 1), w=torch.empty(8, 4, 3, 3, 3)) is None + # ...but only on a dim that is actually split. + assert plan((1, 2, 1), w=torch.empty(8, 4, 3, 2, 2)) is not None + # Everything the plan could not read. + assert plan((2, 1, 1), shard_dim=(0, 1, 2)) is None + assert plan((2, 1, 1), shard_dim=(2, 2, 2)) is None + assert plan((2, 1), shard_dim=(2, 3, 4)) is None + assert plan((2, 1, 1), padding=(1, 1)) is None + assert conv_mod._transposed_halo_plan(None, _StubStrategy((2, 1, 1)), x, weight, + (0, 0, 0)) is None + periodic = _dc(x, (2, 1, 1)) + periodic._is_periodic = (True, False, False) + assert conv_mod._transposed_halo_plan( + periodic, _StubStrategy((2, 1, 1)), x, weight, (0, 0, 0) + ) is None + + +@pytest.mark.gpu +@pytest.mark.parametrize("num_shards", [(1, 1, 1), (2, 1, 1)]) +def test_the_transposed_rung_serves_a_dctensor_without_any_halo_exchange(num_shards): + """The rung takes a sharded ``DCTensor``, and posts nothing to do it. + + At ``k = 2`` DistConv's own ``halo_size`` is ``k // 2 == 0``, so + ``forward_halo_exchange`` returns its argument unchanged and the MIOpen rung + also runs on the bare local shard. That is what makes the two rungs the same + computation at more than one shard, and it is why this ladder has no + ``_Halo3d`` in it. ``_StubStrategy`` reports ``shard_ind = 0``, so the + MIOpen comparison arm is runnable in a one-rank process. + """ + import distconv + import distconv.distconv as dc + + from triton_conv3d import reference as ref + + cin, cout, spatial = 16, 8, (8, 8, 8) + problem = _transposed_problem(cin, cout, spatial) + conv = _gpu_convT(cin, cout) + plain = _stock_like(conv) + x = _gpu_input((1, cin, *spatial)) + gy = _gpu_input((1, cout, 16, 16, 16), seed=41) + fast_x = x.detach().clone().requires_grad_(True) + plain_x = x.detach().clone().requires_grad_(True) + strategy = _StubStrategy(num_shards) + + calls = [] + original = dc.forward_halo_exchange + dc.forward_halo_exchange = lambda *a, **kw: calls.append(a) or original(*a, **kw) + try: + y = conv(distconv.DCTensor.from_shard(fast_x, strategy)) + finally: + dc.forward_halo_exchange = original + + assert isinstance(y, distconv.DCTensor) + assert conv._triton_ok is True, "the fast rung did not serve the DCTensor" + assert calls == [], "the Triton rung went through DistConv's halo exchange" + + y_plain = plain(distconv.DCTensor.from_shard(plain_x, strategy)) + y.backward(distconv.DCTensor.from_shard(gy, strategy)) + y_plain.backward(distconv.DCTensor.from_shard(gy, strategy)) + + operands = { + "input": x, + "weight": conv.weight.detach(), + "bias": conv.bias.detach(), + "grad_output": gy, + } + for name, actual, incumbent in ( + ("fwd", y._tensor, y_plain._tensor), + ("bwd-data", fast_x.grad, plain_x.grad), + ("bwd-weight", conv.weight.grad, plain.weight.grad), + ): + expected = ref.reference(problem, operands, name) + ref.assert_close(actual, expected, problem, name, + incumbent_error=ref.compare(incumbent, expected)) + torch.testing.assert_close(conv.bias.grad.float(), plain.bias.grad.float(), + rtol=2e-2, atol=2e-2) diff --git a/tests/test_unet.py b/tests/test_unet.py index b61ae21..fa5f47b 100644 --- a/tests/test_unet.py +++ b/tests/test_unet.py @@ -198,3 +198,144 @@ def counting_pad(tensor, pad, *args, **kwargs): f"Guard not yet in place: {exact_match_pad_calls} pad calls with diffs=0 " f"(expected 0 when fixed). This is the RED baseline." ) + + +# --------------------------------------------------------------------------- # +# The decoder skip concatenation. +# +# ``Up.forward`` no longer calls ``torch.cat`` directly; it goes through +# ``unet_parts._skip_concat``, which may legitimately emit a narrower dtype +# than ``torch.cat`` would when autocast is on (see the ``Up`` docstring). +# Everything below pins the part that must NOT change: outside autocast the +# block is bitwise what it was, the ``F.pad`` path still works, and both the +# ``trilinear`` and ``ConvTranspose3d`` branches agree with an explicit +# ``torch.cat`` reference. +# --------------------------------------------------------------------------- # +@pytest.mark.parametrize("trilinear", [False, True]) +def test_up_matches_an_explicit_torch_cat_reference(trilinear): + """``Up.forward`` must equal ``conv(cat([x2, up(x1)]))``, bitwise, on CPU.""" + from ScaFFold.unet.unet_parts import Up + + up = Up(in_channels=32, out_channels=16, group_norm_groups=8, trilinear=trilinear) + up.eval() + generator = torch.Generator().manual_seed(11) + # Either branch must hand ``self.conv`` ``in_channels`` channels: the + # transposed convolution halves 32 -> 16, while ``nn.Upsample`` changes no + # channels, so its input already carries 16. + x1 = torch.randn(1, 16 if trilinear else 32, 8, 8, 8, generator=generator) + x2 = torch.randn(1, 16, 16, 16, 16, generator=generator) + + with torch.no_grad(): + got = up(x1, x2) + reference = up.conv(torch.cat([x2, up.up(x1)], dim=1)) + + assert got.shape == reference.shape + assert torch.equal(got, reference), ( + "the skip concatenation must be bitwise torch.cat outside autocast" + ) + + +def test_up_still_pads_and_concatenates_when_shapes_disagree(): + """The non-power-of-two path: ``F.pad`` fires and the result still matches.""" + import torch.nn.functional as F + + from ScaFFold.unet.unet_parts import Up + + up = Up(in_channels=32, out_channels=16, group_norm_groups=8, trilinear=False) + up.eval() + generator = torch.Generator().manual_seed(12) + x1 = torch.randn(1, 32, 7, 7, 7, generator=generator) # -> 14^3 after up + x2 = torch.randn(1, 16, 16, 16, 16, generator=generator) # 16^3: diff = 2 + + with torch.no_grad(): + got = up(x1, x2) + padded = F.pad(up.up(x1), [1, 1, 1, 1, 1, 1]) + reference = up.conv(torch.cat([x2, padded], dim=1)) + + assert got.shape == (1, 16, 16, 16, 16) + assert torch.equal(got, reference) + + +def test_up_gradients_match_an_explicit_torch_cat_reference(): + """Backward through the skip concatenation, bitwise, on CPU.""" + from ScaFFold.unet.unet_parts import Up + + up = Up(in_channels=32, out_channels=16, group_norm_groups=8, trilinear=False) + generator = torch.Generator().manual_seed(13) + x1 = torch.randn(1, 32, 8, 8, 8, generator=generator) + x2 = torch.randn(1, 16, 16, 16, 16, generator=generator) + + a, b = x1.clone().requires_grad_(True), x2.clone().requires_grad_(True) + up.zero_grad(set_to_none=True) + up(a, b).pow(2).sum().backward() + got = (a.grad.clone(), b.grad.clone()) + got_params = {n: p.grad.clone() for n, p in up.named_parameters()} + + c, d = x1.clone().requires_grad_(True), x2.clone().requires_grad_(True) + up.zero_grad(set_to_none=True) + up.conv(torch.cat([d, up.up(c)], dim=1)).pow(2).sum().backward() + + assert torch.equal(got[0], c.grad) + assert torch.equal(got[1], d.grad) + for name, param in up.named_parameters(): + assert torch.equal(got_params[name], param.grad), name + + +def test_up_concatenation_keeps_channels_last(): + """The concatenation must not break the layout chain it exists to preserve. + + Asserted on ``_skip_concat`` with two channels-last halves rather than on a + whole ``Up`` block: on CPU ``nn.ConvTranspose3d`` returns a *contiguous* + tensor whatever it is handed, so the block's own inputs to the + concatenation are not both channels-last there and the block-level + assertion would be measuring the convolution's layout policy, not this + one's. On GPU with ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` -- the production + configuration -- both halves are channels-last and this is the property + ``Up`` relies on. + """ + from ScaFFold.unet.unet_parts import _skip_concat as skip_concat + + generator = torch.Generator().manual_seed(14) + x1 = torch.randn(1, 16, 16, 16, 16, generator=generator).contiguous( + memory_format=torch.channels_last_3d + ) + x2 = torch.randn(1, 16, 16, 16, 16, generator=generator).contiguous( + memory_format=torch.channels_last_3d + ) + out = skip_concat(x2, x1) + assert out.shape == (1, 32, 16, 16, 16) + assert out.is_contiguous(memory_format=torch.channels_last_3d) + assert torch.equal(out, torch.cat([x2, x1], dim=1)) + + +def test_whole_model_forward_and_backward_still_agree_with_a_cat_based_up(): + """End to end: swapping the concatenation back must change nothing on CPU.""" + import torch as _torch + + from ScaFFold.unet import unet_parts + + def cat_forward(self, x1, x2): + x1 = self.up(x1) + return self.conv(_torch.cat([x2, x1], dim=1)) + + model = UNet( + n_channels=_N_CHANNELS, n_classes=_N_CLASSES, trilinear=False, layers=2 + ) + x = _make_input(seed=15).requires_grad_(True) + + model.zero_grad(set_to_none=True) + model(x).pow(2).sum().backward() + grads = {n: p.grad.clone() for n, p in model.named_parameters()} + x_grad = x.grad.clone() + + original = unet_parts.Up.forward + try: + unet_parts.Up.forward = cat_forward + x2 = _make_input(seed=15).requires_grad_(True) + model.zero_grad(set_to_none=True) + model(x2).pow(2).sum().backward() + for name, param in model.named_parameters(): + assert torch.equal(grads[name], param.grad), name + assert torch.equal(x_grad, x2.grad) + finally: + unet_parts.Up.forward = original From b5fe8532cb0ef44eca744bf15b6752fa72289be0 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Wed, 5 Aug 2026 10:21:43 -0700 Subject: [PATCH 5/8] Cite findings rather than the scratch tree that produced them Every comment and docstring in the committed tree stood on its own path into the untracked work/ scratch directory: 70 references across 20 files, all of them dangling for anyone who clones this branch. The measurement each one supported is kept and stated as a result -- the number, the direction of the effect, and the "this was tried and it loses" warnings that stop a closed question being re-litigated -- while the path, the capture filename, the section number of an unshipped document and the blow-by-blow methodology go. Three pointers were dead even with work/ present and are simply gone: the transposed benchmark driver the conv_bench docstring narrated a refactor away from, a tuned table's named source capture, and the review commit SHA in the GroupNorm wiring tests. The corpus JSON keeps model-analysis/unet_shapes.py as its provenance; nothing parses either file's "source" field. --- ScaFFold/unet/conv3d.py | 19 ++--- ScaFFold/unet/triton_group_norm.py | 21 ++--- ScaFFold/unet/unet_parts.py | 2 +- tests/test_conv3d.py | 12 ++- tests/test_groupnorm.py | 4 +- tests/test_groupnorm_wiring_edge.py | 4 +- tests/test_platform_guard.py | 10 +-- tests/test_triton_group_norm_edge.py | 3 +- triton_conv3d/bench/baseline.py | 10 +-- triton_conv3d/bench/conv_bench.py | 82 ++++++++----------- triton_conv3d/bench/harness.py | 20 ++--- triton_conv3d/bwd_data.py | 19 ++--- triton_conv3d/gather_gemm.py | 70 ++++++++-------- triton_conv3d/reduce_gemm.py | 109 ++++++++++++------------- triton_conv3d/scaffold_census.json | 2 +- triton_conv3d/scaffold_corpus.json | 2 +- triton_conv3d/shapes.py | 30 ++++--- triton_conv3d/tests/test_bwd_weight.py | 7 +- triton_conv3d/tests/test_infra.py | 52 ++++++------ triton_conv3d/transposed.py | 11 ++- 20 files changed, 223 insertions(+), 266 deletions(-) diff --git a/ScaFFold/unet/conv3d.py b/ScaFFold/unet/conv3d.py index 83dd933..2178dfe 100644 --- a/ScaFFold/unet/conv3d.py +++ b/ScaFFold/unet/conv3d.py @@ -25,7 +25,7 @@ where it lies, which is the layout ``worker.py`` already puts every convolution parameter in. Measured at **1.26x of the MIOpen step** at ScaFFold's scale-7 single-GPU configuration, with the backward-weight - direction reproducible at no cost -- see ``work/triton-conv/STATUS.md``. + direction reproducible at no cost. 2. **MIOpen**, via ``nn.Conv3d.forward``, which is what every rejection falls back to and what defines the semantics the Triton rung must match. @@ -91,8 +91,7 @@ :class:`_Halo3d` is not applied at all, and the rung is stage 1's exactly: ``forward_halo_exchange`` allocates its receive buffers with ``zeros_like`` and posts nothing when ``shard_ind`` is 0 on every axis, so the 3.797 ms/step of -zero-slab ``cat`` copies ``work/upstream-repros/DISTCONV_ACTIONS.md`` prices at one shard just -disappear. +zero-slab ``cat`` copies measured at one shard just disappear. Every fact that argument rests on is checked *positively* before the rung is taken -- see :func:`_halo_plan`, which returns ``None`` for anything it could not @@ -673,8 +672,7 @@ def _policy_declines(x_shape, w_shape, stride, padding, dilation): Empty since 2026-08-04, and the arguments are kept so a future entry has somewhere to go. It held three rules, and re-measuring every one of them - against the shipped kernels retired all three - (``work/triton-conv/review/CONV_REMEASURE.md``; per-site sums over all + against the shipped kernels retired all three (per-site sums over all three directions, kernel time with 95% intervals): * ``Cin == 3``, the stem, blocked at a quoted **0.53x**. That figure was @@ -1438,9 +1436,8 @@ def _transposed_policy_declines(x_shape, w_shape): ``gather_gemm``'s tuning, not of this kernel, whose N axis carries ``Cout * taps`` and is 8x wider for it. - So the four sites were measured directly rather than inherited - (``work/triton-conv/m5_shipped_{fwd,bwd-data,bwd-weight}.json``, config A, - Triton over MIOpen per direction): + So the four sites were measured directly rather than inherited (config A, + Triton over MIOpen, per direction): =========================== ===== ======== ========== site fwd bwd-data bwd-weight @@ -1459,10 +1456,8 @@ def _transposed_policy_declines(x_shape, w_shape): Checked again at the level that decides it, because those are *kernel* times and this ladder's per-call cost is not zero: on the real model at config A, - interleaved, 12 steps per arm - (``work/triton-conv/bin/convT_config_a.py``), the whole ladder is worth - **0.51 ms/step** and a variant that keeps only ``up3``/``up4`` on the rung is - worth 0.65 -- a 0.14 ms difference inside a 1.5 ms spread, i.e. not a + interleaved, 12 steps per arm, the whole ladder is worth **0.51 ms/step** + and a variant that keeps only ``up3``/``up4`` on the rung is worth 0.65 -- a 0.14 ms difference inside a 1.5 ms spread, i.e. not a measurement, so there is nothing here to write an entry from. What *is* measurable is that the gap between 0.51 and the 0.9 ms the kernel times project is per-call Python and launch overhead (~0.10 ms per forward and diff --git a/ScaFFold/unet/triton_group_norm.py b/ScaFFold/unet/triton_group_norm.py index 641122e..2dd28e9 100644 --- a/ScaFFold/unet/triton_group_norm.py +++ b/ScaFFold/unet/triton_group_norm.py @@ -97,9 +97,8 @@ * **statistics** -- always accumulated in fp32, never in the input dtype. * **memory format** -- the output has the *input's* memory format. This is the one deliberate difference from stock GroupNorm, which returns a contiguous - tensor for every input layout (measured, see the table in - ``work/gn-dctensor/triton/RESULTS.md``); preserving channels-last is the - entire point of the kernel. + tensor for every input layout; preserving channels-last is the entire point + of the kernel. * **autograd** -- ``d_input``, ``d_weight``, ``d_bias``; ``weight=None`` and/or ``bias=None`` supported. **First order only**: the backward is itself a custom op with no autograd formula of its own, so a second @@ -163,8 +162,8 @@ at ``elem_progs`` programs which then stride over the tiles (so the redundancy is bounded by the *grid*, not by the tile count), and ``nsplit_target`` is retuned per shape against that cap. With both, the same shape is 1-2% *faster* -than the unfused chain. Do not change one without re-running the other; the -coordinate-descent tuner is ``work/gn-dctensor/kernel-opt/tune.py``. +than the unfused chain. The two were tuned jointly by coordinate descent, so +do not change one without re-running the other. Why not one launch per pass --------------------------- @@ -182,8 +181,7 @@ the ten small-shape sites are already only 1.05 ms of a 64.1 ms/step total, so the gain is zero. Hand-rolled inter-workgroup synchronisation is not a good trade for 3% in a benchmark whose value depends on being trustworthy and -reproducible. The measurements are in -``work/gn-dctensor/triton-small/RESULTS.md``. +reproducible. Numerics: Welford, not ``E[x^2]-E[x]^2`` ======================================== @@ -287,9 +285,9 @@ so a program must own one group and stream S, rather than owning all groups and streaming voxels -- i.e. a second family of four kernels. The payoff is small: on contiguous input Inductor's compiled GroupNorm already reaches - 89-92% of this device's measured streaming roofline (RESULTS.md 4) -- and the - table above confirms it, 66.4 ms/step against this kernel's 67.1 -- so a - native NCDHW kernel could win ~10% there, against the 6.4x it wins on + 89-92% of this device's measured streaming roofline -- and the table above + confirms it, 66.4 ms/step against this kernel's 67.1 -- so a native NCDHW + kernel could win ~10% there, against the 6.4x it wins on channels-last input. If a mixed-layout model ever makes that 10% matter, the place to add it is the strategy hook below. @@ -339,8 +337,7 @@ back finite because the fused activation ate the NaN passes straight through ScaFFold's non-finite-loss abort and checkpoints a broken model. ``+-Inf`` and ``-0.0`` are bit-identical under either spelling (``-0.0`` flushes to ``+0.0``, -as ``F.relu`` does). Cost: nil, measured -- see ``FastGroupNorm``'s tests and -``work/gn-dctensor/wiring-fixes``. +as ``F.relu`` does). Cost: nil, measured -- see ``FastGroupNorm``'s tests. Composition =========== diff --git a/ScaFFold/unet/unet_parts.py b/ScaFFold/unet/unet_parts.py index a49f72c..9483b17 100644 --- a/ScaFFold/unet/unet_parts.py +++ b/ScaFFold/unet/unet_parts.py @@ -189,7 +189,7 @@ class Up(nn.Module): contiguous, at 51-63% of this device's streaming roofline against the kernel's 90-103% -- but it was worth a further 0.08 ms of the step, which did not justify a second hand-written kernel in a benchmark other people - have to trust. ``work/skip-path/RESULTS.md`` has the numbers. + have to trust. """ def __init__(self, in_channels, out_channels, group_norm_groups, trilinear=True): diff --git a/tests/test_conv3d.py b/tests/test_conv3d.py index b62e0ee..cd9991b 100644 --- a/tests/test_conv3d.py +++ b/tests/test_conv3d.py @@ -20,8 +20,7 @@ wrong gradient at every shard boundary of a large run, because the halo DistConv adds below autograd is invisible to a module-level adapter and the halo this one adds instead is only right where the plan says it is. The exchange itself -needs real ranks and lives in -``work/triton-conv/review/halo/stage2_dist.py``. +needs real ranks and is exercised by a separate multi-rank harness. Tolerances come from ``triton_conv3d.reference``'s policy (an fp64 reference and a dtype/K-derived bound, or MIOpen's own error where that is looser). Nothing @@ -289,7 +288,7 @@ def test_the_gate_asks_the_predicates_about_the_tensor_the_kernel_will_see(): this process has no process group to exchange over. The plan is made, and the predicates are asked about the halo'd extent -- ``8 -> 10`` on D -- at the padding the exchange leaves behind. The multi-rank half of this lives - in ``work/triton-conv/review/halo/stage2_dist.py``. + in a separate harness that needs real ranks. """ conv = _gpu_conv() x = _gpu_input((1, 16, 8, 8, 8)) @@ -390,8 +389,8 @@ def test_a_kernel_failure_after_the_halo_falls_back_without_exchanging_twice( own padding on the unexchanged shard, the answer has an independent reference: what ``nn.Conv3d`` computes on the original input. - The multi-rank half, with real slabs on a real mesh, is - ``work/triton-conv/review/halo/stage2_dist.py``'s ``fallback`` stage. + The multi-rank half, with real slabs on a real mesh, needs real ranks and + lives in a separate harness. """ import distconv import distconv.distconv as dc @@ -1087,8 +1086,7 @@ def test_the_triton_rung_performs_no_halo_exchange_at_one_shard(): ``forward_halo_exchange`` has no ``num_shards == 1`` early-out, so today every convolution concatenates two zero slabs onto each of the three sharded dims: 54 calls and 3.797 ms/step of pure copying at - ``dc_num_shards=(1,1,1)`` (``work/upstream-repros/DISTCONV_ACTIONS.md``). Taking the - Triton rung removes all of them, and leaves the caller's ``_tensor`` + ``dc_num_shards=(1,1,1)``. Taking the Triton rung removes all of them, and leaves the caller's ``_tensor`` un-narrowed and still channels-last for the consumers downstream of it. """ import distconv diff --git a/tests/test_groupnorm.py b/tests/test_groupnorm.py index 942087d..4b1cb05 100644 --- a/tests/test_groupnorm.py +++ b/tests/test_groupnorm.py @@ -960,7 +960,7 @@ def test_dctensor_two_shards_matches_eager_and_normalizes_per_shard(): halves of the claim the fast path rests on: bit-identical to the eager wrapped route, and normalizing the local shard rather than the global volume (which is DistConv's existing semantics, not something this change - introduces -- see the R32/R39 notes for the upstream discussion). + introduces). """ script = ( Path(__file__).resolve().parent @@ -1177,7 +1177,7 @@ def step(): # Two tests below run the whole UNet twice, changing only which rung serves # GroupNorm, and ask whether the gradients agree. *How* that is asked matters # more than it looks, and both tests used to ask it in a way that could only -# pass by luck. Measurements: ``work/triton-conv/review/GN_CHECKPOINT_TEST.md``. +# pass by luck. # # **Per-parameter relative L2 under bf16 autocast is not a bounded quantity for # this model.** Against an fp64 reference, every arm -- eager, compiled and diff --git a/tests/test_groupnorm_wiring_edge.py b/tests/test_groupnorm_wiring_edge.py index b1e8526..0d513e4 100644 --- a/tests/test_groupnorm_wiring_edge.py +++ b/tests/test_groupnorm_wiring_edge.py @@ -14,10 +14,10 @@ """Edge cases of the GroupNorm *wiring* (``FastGroupNorm``'s three-rung ladder). -Written as an adversarial review of commit ``ca58664``; ``tests/test_groupnorm.py`` +Written as an adversarial review of the wiring: ``tests/test_groupnorm.py`` covers the happy paths and the routing predicates, this file covers the places where the ladder, the latches and the absorbed ReLU interact with the rest of -torch. See ``work/gn-dctensor/WIRING_REVIEW.md`` for the full write-up. +torch. The review left ten of these as ``xfail(strict=True)``, one per defect, each asserting the behaviour the module *should* have. All ten are fixed and the diff --git a/tests/test_platform_guard.py b/tests/test_platform_guard.py index 5e2ffc9..8c4c048 100644 --- a/tests/test_platform_guard.py +++ b/tests/test_platform_guard.py @@ -31,8 +31,8 @@ be evaluated *once*, which is asserted as a call count on the seam rather than as a duration, because a timing assertion on a driver query measures the driver. -Verified by mutation: ``work/triton-conv/bin/mutate_platform_guard.py`` breaks -each property in turn and checks that the test named there notices. +Every property here was verified by mutation: breaking it in the guard and +checking that the test named alongside it notices. """ from __future__ import annotations @@ -390,9 +390,9 @@ def test_this_node_is_the_tuned_platform(caplog): """The accept path, against the real driver rather than a stub. Also the check that the constants have not drifted away from the machine - every measurement in ``work/triton-conv/`` was taken on -- and that the - suffix strip is load-bearing here rather than defensive, since the raw - string this device reports really does carry ``:sramecc+:xnack-``. + every measurement in this project was taken on -- and that the suffix strip + is load-bearing here rather than defensive, since the raw string this device + reports really does carry ``:sramecc+:xnack-``. """ props = torch.cuda.get_device_properties(0) assert ":" in props.gcnArchName, props.gcnArchName diff --git a/tests/test_triton_group_norm_edge.py b/tests/test_triton_group_norm_edge.py index bca5e88..4cc91ef 100644 --- a/tests/test_triton_group_norm_edge.py +++ b/tests/test_triton_group_norm_edge.py @@ -16,8 +16,7 @@ Companion to ``tests/test_triton_group_norm.py``, written independently during an audit of the kernel. It covers the ground the author's suite does not, and -pins the divergences that audit found. See -``work/gn-dctensor/KERNEL_REVIEW.md`` for the full write-up. +pins the divergences that audit found. The two structural gaps this file closes: diff --git a/triton_conv3d/bench/baseline.py b/triton_conv3d/bench/baseline.py index 0fe1b67..5610851 100644 --- a/triton_conv3d/bench/baseline.py +++ b/triton_conv3d/bench/baseline.py @@ -20,11 +20,11 @@ 1. ``torch.backends.cudnn.benchmark = True``. ScaFFold sets this at startup unless ``more_determinism`` is on (``ScaFFold/worker.py:171``), and so does - the profiling harness the reference numbers come from - (``work/report/prof_bench.py:125``). On ROCm the flag decides whether - PyTorch asks MIOpen for an exhaustive *find* or lets it answer from its AI - heuristic. The two answers are not close. For ``conv 64->64 k3 @ 128^3``, - forward, measured with ``MIOPEN_ENABLE_LOGGING=1``:: + the profiling harness the reference numbers come from. On ROCm the flag + decides whether PyTorch asks MIOpen for an exhaustive *find* or lets it + answer from its AI heuristic. The two answers are not close. For + ``conv 64->64 k3 @ 128^3``, forward, measured with + ``MIOPEN_ENABLE_LOGGING=1``:: benchmark=False findMode DYNAMIC_HYBRID(5), no search DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<256,64,64,32, diff --git a/triton_conv3d/bench/conv_bench.py b/triton_conv3d/bench/conv_bench.py index 4813147b..40c3f91 100644 --- a/triton_conv3d/bench/conv_bench.py +++ b/triton_conv3d/bench/conv_bench.py @@ -7,27 +7,22 @@ every cell in the project under one methodology, in one process, from one command. -The transposed benchmarks used to live in a separate driver -(``work/triton-conv/bin/m5_convT_bench.py``, deleted at this commit), on the -argument that ``--direction`` selects between three *directions of one operator* -while the transposed convolution is a *different* operator, and that folding it -in would make ``_build`` branch on ``problem.transposed`` in every arm to run the -same code. The argument was right about the failure and wrong about the -factoring. Operator and direction are **two axes of one table**, not two values -on one switch: :data:`_OPERATORS` maps ``(operator, direction)`` to a builder, -and each of the six builders is a separate function that names its own operands, -its own control, its own candidate configs and its own reference. ``_build`` is -a lookup, so nothing branches on ``problem.transposed`` anywhere, and the four +Operator and direction are **two axes of one table**, not two values on one +switch: :data:`_OPERATORS` maps ``(operator, direction)`` to a builder, and each +of the six builders is a separate function that names its own operands, its own +control, its own candidate configs and its own reference. ``_build`` is a +lookup, so nothing branches on ``problem.transposed`` anywhere, and the four things that *are* per-operator rather than per-direction -- the shape form, the problem ordering, the config type, whether a direction is sweepable -- sit on :class:`_Op` where a reader can see all four at once. -What that bought, immediately: ``m5_convT_bench``'s transposed backward -directions dropped the ``config=`` argument on the floor, because the only config -object it had was the transposed *forward*'s, and passing that to a direction -served by ``conv3d_forward`` silently benchmarks a tile nothing would select. -With a builder per cell each one names the config type its own entry point -resolves, and both transposed backward directions became sweepable for free. +That factoring is what makes the transposed backward directions sweepable at +all. When the transposed benchmarks lived in a driver of their own, the only +config object that driver had was the transposed *forward*'s, and passing it to +a direction served by ``conv3d_forward`` silently benchmarks a tile nothing +would select -- so those directions dropped ``config=`` on the floor instead. +With a builder per cell, each one names the config type its own entry point +resolves and the mistake is not expressible. Five things this driver is careful about, each because getting it wrong has already produced a wrong answer once in this project: @@ -70,12 +65,11 @@ MIOpen arm and measures the Triton kernels alone. It is not a corner-cutting option, it is where essentially all of the wall clock is: ``cudnn.benchmark = True`` puts MIOpen on the Find path, whose disk record cannot be replayed in a -fresh process (``review/MIOPEN_CACHE.md`` §2), so **every** cell pays a find -- -measured on this node at 92-174 s per cell against 0.3-1.2 s for the Triton -compile, the graph capture, the calibration and the timed rounds put together -(``review/HARNESS_SPEED.md`` §1). A Triton-only row therefore carries no -``miopen_*`` and no ``speedup`` key at all -- an absent measurement stays -absent -- and ``--check``, whose reference *is* MIOpen's answer, is refused with +fresh process, so **every** cell pays a find -- measured on this node at +92-174 s per cell against 0.3-1.2 s for the Triton compile, the graph capture, +the calibration and the timed rounds put together. A Triton-only row therefore +carries no ``miopen_*`` and no ``speedup`` key at all -- an absent measurement +stays absent -- and ``--check``, whose reference *is* MIOpen's answer, is refused with it. **The control.** The MIOpen side of a backward direction is a real forward @@ -218,8 +212,8 @@ class _Case: #: always wanted: for a *backward* direction the control is a real forward #: graph, and running it once costs MIOpen's find -- measured at 92-174 s #: per cell on this node's corpus, against 0.3-1.2 s for everything else the - #: cell does (``work/triton-conv/review/HARNESS_SPEED.md`` §1). A - #: Triton-only capture that still built the control would pay all of it. + #: cell does. A Triton-only capture that still built the control would pay + #: all of it. miopen: Callable[[], object] | None #: ``Callable[[], object]`` -- the hoistable weight prep, or ``None`` where #: the direction has none. **All six cells are now ``None``**: the @@ -236,8 +230,8 @@ class _Case: #: The config **this cell's entry point would resolve on its own**, computed #: once so that ``--shipped`` measures the shipped kernel without also #: measuring the shipped table lookup. The two are not the same number: the - #: lookup is 0.0164 ms, 39% of the transposed forward kernel - #: (``HARNESS_RIGOR.md`` H13). ``test_the_shipped_config_is_what_the_entry + #: lookup is 0.0164 ms, 39% of the transposed forward kernel. + #: ``test_the_shipped_config_is_what_the_entry #: _point_resolves`` pins these six against the entry points. shipped_config: Callable[[], ConvConfig | None] #: ``Callable[[], tuple[Tensor, Tensor]]`` -- ``(ours, MIOpen's)`` on this @@ -737,9 +731,8 @@ def _build(problem: ConvProblem, direction: Direction, device: str = "cuda", launcher and, for a backward direction, no autograd graph -- which is the expensive half. Building the control does not merely cost the arm's timing; running it once costs MIOpen's *find*, which under ``cudnn.benchmark = True`` - cannot be replayed from disk (``review/MIOPEN_CACHE.md`` §2) and which - measures **92-174 s per cell** on this corpus against 0.3-1.2 s for - everything else in the cell. A Triton-only capture is therefore two orders + cannot be replayed from disk and which measures **92-174 s per cell** on + this corpus against 0.3-1.2 s for everything else in the cell. A Triton-only capture is therefore two orders of magnitude cheaper than a comparison, and that is entirely MIOpen's find. """ op = _OPERATORS[operator or operator_of(problem)] @@ -759,12 +752,11 @@ def _build(problem: ConvProblem, direction: Direction, device: str = "cuda", class _Region: """The decision about what every arm of one cell is timed with. - One object for the whole cell, never one per arm. ``HARNESS_RIGOR.md`` H10 - is the reason: a per-arm instrument biases a ratio *even when both arms are - individually right*, and the version of that mistake available here -- - hoisting Triton's config lookup out while leaving PyTorch's dispatch inside - the MIOpen arm -- flatters us by up to 1.4x on exactly the sub-0.15 ms cells - where it is hardest to see. + One object for the whole cell, never one per arm. A per-arm instrument + biases a ratio *even when both arms are individually right*, and the version + of that mistake available here -- hoisting Triton's config lookup out while + leaving PyTorch's dispatch inside the MIOpen arm -- flatters us by up to + 1.4x on exactly the sub-0.15 ms cells where it is hardest to see. """ #: ``"kernel"`` (both arms replayed from a graph) or ``"call"`` (both arms @@ -792,8 +784,7 @@ def _timed_region(variants: Mapping[str, Callable[[], object]], issues them, back to back on one stream. That boundary is the same on both sides only because the *whole comparison* - moves together. Measured on ``convT 1024->512 @ 8^3``, per call - (``work/triton-conv/bin/launcher_symmetry.py --only census-small``): + moves together. Measured on ``convT 1024->512 @ 8^3``, per call: ============================ ======== ======= ============ =============== arm eager kernel host launch speedup in->out @@ -910,9 +901,9 @@ def measure_problem(problem: ConvProblem, *, direction: Direction = "fwd", absent number is absent, not zero -- and says so in ``control``. Everything else is unchanged: the same CUDA-graph region, the same adaptive stopping, the same 95% interval, the same ``stop`` reason. What it buys is the whole - of MIOpen's find (§1 of ``review/HARNESS_SPEED.md``: 98.3% of a three-cell - problem's wall clock on this node); what it costs is the comparison, so use - it when the baseline is the deliverable and the ratio is not. + of MIOpen's find -- 98.3% of a three-cell problem's wall clock on this node; + what it costs is the comparison, so use it when the baseline is the + deliverable and the ratio is not. ``shipped`` skips the sweep and times the config **this cell's entry point would resolve on its own** -- the tuned table plus the heuristic fallback -- @@ -1082,9 +1073,9 @@ def measure_problem(problem: ConvProblem, *, direction: Direction = "fwd", # Not a measurement, and it shows: where the launcher is already # negligible -- above about 0.3 ms, where an event-free bracket of a # host-paced loop already reaches kernel throughput -- this comes out at - # a few microseconds of either sign. The measured version, with - # intervals and both policies run as full races, is - # ``work/triton-conv/bin/launcher_symmetry.py --only census``. + # a few microseconds of either sign. Taking it seriously means running + # both launcher policies as full races with intervals, which is a + # separate experiment and not this row. for arm in ("triton", "miopen"): if f"{arm}_ms" not in row: continue @@ -1195,8 +1186,7 @@ def main() -> None: "control is what makes a capture expensive -- MIOpen's " "find cannot be replayed from disk under " "cudnn.benchmark=True and costs 92-174 s per cell on " - "this corpus, which is 98% of a cell's wall clock " - "(work/triton-conv/review/HARNESS_SPEED.md)") + "this corpus, which is 98% of a cell's wall clock") ap.add_argument("--corpus", default="scaffold", choices=["scaffold", "census"], help="which problem list --top/--problems index into. " diff --git a/triton_conv3d/bench/harness.py b/triton_conv3d/bench/harness.py index 0a94131..14831b0 100644 --- a/triton_conv3d/bench/harness.py +++ b/triton_conv3d/bench/harness.py @@ -1,10 +1,9 @@ # SPDX-License-Identifier: (Apache-2.0) """Timing that survives this machine -- and says how well it survived it. -**What this node does.** Measured on GPU 2 on 2026-08-03 -(``work/triton-conv/bin/drift_probe.py``: 47,686 consecutive 10-call blocks of -``conv 64->64 k3 @ 130^3`` over 14 minutes, with ``rocm-smi`` sampled -concurrently from a sibling process): +**What this node does.** Measured on GPU 2 on 2026-08-03 -- 47,686 consecutive +10-call blocks of ``conv 64->64 k3 @ 130^3`` over 14 minutes, with ``rocm-smi`` +sampled concurrently from a sibling process: * steady-state dispersion is **CoV 0.56%**, IQR/median 0.39%, p99/p01 = 1.022; * the *sequential* protocol -- measure A for a while, then B -- simulated out of @@ -20,8 +19,7 @@ So the clock is not what threatens a comparison here, and neither is elapsed time between two measurements in one process. What does threaten one is a **neighbour**: a foreign tenant on the device inflated three cells **2.6x**, -with per-cell spreads 20-100x their quiet-device values -(``work/triton-conv/BASELINE_NOTES.md``, "cause 3"). That is common-mode +with per-cell spreads 20-100x their quiet-device values. That is common-mode within a round -- every arm pays it at once -- so :func:`interleaved` survives it where a between-run comparison does not. It costs nothing, so it stays. It is not, however, a substitute for an error bar. @@ -52,8 +50,7 @@ the launcher and the kernel together and cannot separate them. :func:`capture` puts ``chunk`` back-to-back calls behind one CUDA graph, which contains the device work and none of the host work, so replaying it measures the kernel -alone. Measured, per call, on ``convT 1024->512 @ 8^3`` -(``work/triton-conv/bin/launcher_symmetry.py --only census-small``): +alone. Measured, per call, on ``convT 1024->512 @ 8^3``: ============================= ======= ======== ============= arm eager kernel host launch @@ -659,8 +656,7 @@ def _order(names: list[str], r: int) -> list[str]: #: Cost of one ``cudaGraphLaunch``, in ms per replay, measured on this device by #: fitting ``per_call(chunk) = kernel + cost / chunk`` to graphs of 1, 2, 4, 8, -#: 16 and 32 calls at ``convT 1024->512 @ 8^3`` -#: (``work/triton-conv/bin/graph_probe2.py``): +#: 16 and 32 calls at ``convT 1024->512 @ 8^3``: #: #: =========================== ============ #: arm fitted cost @@ -674,8 +670,8 @@ def _order(names: list[str], r: int) -> list[str]: #: The constant below is the **worst** of those, not their mean, because the #: quantity that has to be bounded is the residual on whichever arm pays most -- #: and because a per-arm estimate is exactly the mistake :func:`_common_group` -#: exists to document (H10: two byte-identical arms picked different instruments -#: and read 4% apart). It is a property of the launcher, not of the workload. +#: exists to document: two byte-identical arms picked different instruments and +#: read 4% apart. It is a property of the launcher, not of the workload. _REPLAY_COST_MS = 0.0128 #: Fraction of the *shortest* arm's per-call time the residual replay cost is diff --git a/triton_conv3d/bwd_data.py b/triton_conv3d/bwd_data.py index fd29d61..07d8078 100644 --- a/triton_conv3d/bwd_data.py +++ b/triton_conv3d/bwd_data.py @@ -118,11 +118,10 @@ def _tuned(bm: int, bn: int, bk: int, warps: int, group_m: int = 6) -> ConvConfi #: the other way round would make ``512 -> 256`` in this table and ``512 -> 256`` #: in the forward's mean different things, which is a trap not worth setting. #: -#: Source: ``work/triton-conv/m2_bwd.json`` -- 21 problems, 1468 timed -#: configurations, analysed by ``work/triton-conv/m2_analyze.py``. Only -#: channel pairs that were actually timed appear; a miss falls to the heuristic -#: on the effective widths, which is a real gap and not an extrapolation dressed -#: up as a measurement. +#: Drawn from a backward-data sweep over 21 problems and 1468 timed +#: configurations. Only channel pairs that were actually timed appear; a miss +#: falls to the heuristic on the effective widths, which is a real gap and not +#: an extrapolation dressed up as a measurement. #: #: Three things this table records that the forward's does *not*: #: @@ -162,11 +161,11 @@ def _tuned(bm: int, bn: int, bk: int, warps: int, group_m: int = 6) -> ConvConfi # (``p' = 0``) and a reduction of just ``Cout = 6``, so it is a different # regime from every entry above and gets its own key. # - # ``num_warps = 2``, not the 4 M2 shipped, and that is the only axis of this - # entry that moved: M2's sweep drew ``num_warps`` from ``{4, 8, seed}``, so - # 1 and 2 were never timed at any site in this project. Raced at all three - # head volumes (``work/triton-conv/tune/warps_bwdd.json``) 2 is 1.099x, - # 1.057x and 1.101x over 4, and 8 is 0.89-0.90x. The reduction here is + # ``num_warps = 2``, not the 4 the original sweep shipped, and that is the + # only axis of this entry that moved: that sweep drew ``num_warps`` from + # ``{4, 8, seed}``, so 1 and 2 were never timed at any site in this + # project. Raced at all three head volumes, 2 is 1.099x, 1.057x and 1.101x + # over 4, and 8 is 0.89-0.90x. The reduction here is # ``Cout * taps = 6``, one MFMA fragment deep, so a second pair of waves has # nothing to reduce and only replicates the addressing. 1 warp is within # noise of 2 at this site and is 0.245x at ``128 -> 128 @ 66^3``, so the diff --git a/triton_conv3d/gather_gemm.py b/triton_conv3d/gather_gemm.py index 34aff56..66fcc38 100644 --- a/triton_conv3d/gather_gemm.py +++ b/triton_conv3d/gather_gemm.py @@ -66,9 +66,9 @@ not a performance question; the third is 4.8-9.0x slower if addressed in place, because neither tile axis is dense, and is therefore copied. -Two designs that look better and measure worse, both raced per site -(``work/triton-conv/review/wt/``): loading the tile coalesced and transposing -it in registers (1.05-1.55x, the transpose is the cost); and holding the +Two designs that look better and measure worse, both raced per site: loading the +tile coalesced and transposing it in registers (1.05-1.55x, the transpose is the +cost); and holding the parameter in RSCK order, which is 0.159 ms/step better in the kernels and **2.6 ms/step worse in the optimizer**, since the gradient this package produces is channels-last and the elementwise update would then be strided. @@ -250,9 +250,9 @@ def _conv3d_fwd_kernel( # compiling it. # # Two alternatives were implemented and measured worse, both on the eight - # hot config-A sites (``work/triton-conv/review/wt/``): loading the tile - # coalesced along K and transposing it in registers into the - # ``(BLOCK_K, BLOCK_N)`` the dot wants costs 1.05-1.55x, and holding the + # hot config-A sites: loading the tile coalesced along K and + # transposing it in registers into the ``(BLOCK_K, BLOCK_N)`` the dot + # wants costs 1.05-1.55x, and holding the # parameter in RSCK order costs the *backward* direction the same, since # the axis RSCK makes contiguous is backward-data's reduction axis. # @@ -268,8 +268,9 @@ def _conv3d_fwd_kernel( # GPU. Cast per term rather than after the sum -- ``dij * stride_wt`` is # the term that overflows on its own. On the int32 path the casts are # frontend no-ops, so the operand keeps the buffer-load eligibility it - # has today; it only loses it at sizes where PLAN.md 3.5 says the storage - # has already lost it. + # has today; it only loses it at sizes where the *storage* is already + # over the buffer-op limit and has lost it anyway (see + # :data:`~triton_conv3d.shapes.BUFFER_OP_MAX_BYTES`). dij_w = (KD * KH * KW - 1 - dij) if W_FLIP else dij w_row = ( dij_w.to(INDEX_DTYPE) * stride_wt + offs_k.to(INDEX_DTYPE) * stride_wk @@ -578,7 +579,7 @@ def _fit_to_grid(cfg: ConvConfig, m: int, cout: int) -> ConvConfig: #: second wave; four warps each take a quarter of ``BLOCK_M`` and replicate the #: whole per-K-tile address computation for a fragment that is 10/16 padding. #: Measured at all three head volumes, one warp is 1.02-1.22x over the shipped -#: four (``work/triton-conv/tune/k1_fwd*.json``). +#: four. #: #: Gated to ``Cout <= 16`` in :func:`candidate_configs` rather than added to the #: grids above, because a 16-column tile at ``Cout = 512`` is 32x padding and @@ -685,7 +686,7 @@ def _tuned(bm: int, bn: int, bk: int, warps: int, group_m: int = 6, #: #: Deliberately a table and not ``@triton.autotune``: ScaFFold's figure of merit #: is total wall time, so a recompile inside a training step is a direct loss. -#: Source: ``work/triton-conv/m1_fwd.json``, 15 problems, ~1050 timed configs. +#: Drawn from a forward sweep of 15 problems and ~1050 timed configs. #: Only channel pairs that were actually timed appear here. The unmeasured #: pairs -- ``64 -> 128``, ``128 -> 256``, ``256 -> 512``, ``512 -> 1024``, all #: encoder-side -- fall to the heuristic on purpose: an extrapolated entry in a @@ -697,11 +698,11 @@ def _tuned(bm: int, bn: int, bk: int, warps: int, group_m: int = 6, #: 6x18x18`` 0.90x) and belong on the adapter's block-list instead. (An earlier #: version of this comment priced the four pairs at "12.7 ms/step"; that figure #: is config B's *all-directions* total at those pairs, not the forward time an -#: absent row here governs. See ``STATUS.md`` §4.3.) +#: absent row here governs.) _TUNED: dict[tuple, ConvConfig] = { - # The segmentation head, and the one entry here that is *not* from - # ``m1_fwd.json``: source ``work/triton-conv/tune/k1_fwd_confirm.json``. - # ``Cout = 6`` prunes every seed tile, so M1's sweep never timed anything + # The segmentation head, and the one entry here that is *not* from the main + # forward sweep but from a follow-up race at this site alone. + # ``Cout = 6`` prunes every seed tile, so that sweep never timed anything # but ``default_config`` at this site and the tile below is that same # ``128x16x64`` with **one warp instead of four** -- see # :data:`_NARROW_N_TILES` for why one, and why only here. Raced against the @@ -715,7 +716,7 @@ def _tuned(bm: int, bn: int, bk: int, warps: int, group_m: int = 6, **{ tune_key(torch.bfloat16, cin, cout, (3, 3, 3)): cfg for (cin, cout), cfg in { - # The UNet stem, and the row that refutes ``OPTIMIZATION.md`` D4's + # The UNet stem, and the row that refutes the standing verdict that # "``conv 3->64`` is genuinely hopeless, leave it on MIOpen". It was # never hopeless and it was never a matrix-core feeding problem: the # reduction axis of this kernel's ``tl.dot`` is ``Cin`` **alone** @@ -750,8 +751,7 @@ def _tuned(bm: int, bn: int, bk: int, warps: int, group_m: int = 6, # ``kpack = 1`` is not a rounding detail here: ``kp2`` on the same # tile is 0.6085 ms, 1.17x worse. Taller than 512 turns over # (``1024x64x8/w16`` 0.6042); ``GROUP_M`` is inert at this site - # (g1/g6/g12 within 1%). ``work/triton-conv/review/LOSS_OPTIMIZATION.md`` - # §1.2 has the full race and the counters. + # (g1/g6/g12 within 1%). (3, 64): _tuned(512, 64, 8, 8, nk=32), (64, 64): _tuned(128, 64, 64, 4), (128, 64): _tuned(128, 64, 64, 4), @@ -766,11 +766,11 @@ def _tuned(bm: int, bn: int, bk: int, warps: int, group_m: int = 6, (1024, 512): _tuned(128, 128, 128, 8), (1024, 1024): _tuned(64, 64, 128, 8), # The two 2048-channel bottleneck pairs. They are scale-8 sites - # that appeared in no corpus until ``SHAPE_AUDIT.md`` 2.3 found - # them -- the corpus's scale-8 model was a *four*-layer network and - # the harness runs a five-layer one -- so until now they fell to - # the heuristic, and against fresh MIOpen the forward **lost** - # (``TRITON_BASELINES.md`` 5: 0.952x and 0.977x). + # that appeared in no corpus until a shape census of running steps + # found them -- the corpus's scale-8 model was a *four*-layer + # network and the harness runs a five-layer one -- so until now + # they fell to the heuristic, and against fresh MIOpen the forward + # **lost**, 0.952x and 0.977x. # # ``128x64x128`` wins at **every volume both pairs occur at**, # which is the property this table has twice been burned by not @@ -780,8 +780,7 @@ def _tuned(bm: int, bn: int, bk: int, warps: int, group_m: int = 6, # replaces, quiet node, **four independently allocated operand sets # per cell**, both arms sharing each build's operands so the pair is # immune to the placement effect below; kernel-only, median over - # builds with the worst build in brackets - # (``review/opt3/row_2048_quiet.json``): + # builds with the worst build in brackets: # (1024,2048) 8^3 0.3196 vs 0.3224 ms -- 1.010x [1.008] # (1024,2048) 6x8^2 0.2208 vs 0.2769 -- 1.253x [1.251] # (1024,2048) 4x8^2 0.1760 vs 0.2304 -- 1.308x [1.306] @@ -793,11 +792,11 @@ def _tuned(bm: int, bn: int, bk: int, warps: int, group_m: int = 6, # this row. ``(2048, 2048)`` is the one site in this project whose # time depends on **where its weight lands**: at 216 MiB against a # 256 MiB MALL the heuristic is bimodal, two tight states up to - # 14.7% apart and fixed for the life of the allocation - # (``OPTIMIZATION_ROUND3.md`` 4). Measured solo -- one config, one - # operand set, six rebuilds, which is what a caller actually sees -- - # this row is **stable**: 0.5% / 0.7% / 1.4% spread at the three - # volumes against the heuristic's 15% / 5% / 20%. But at ``8^3`` + # 14.7% apart and fixed for the life of the allocation. Measured + # solo -- one config, one operand set, six rebuilds, which is what + # a caller actually sees -- this row is **stable**: 0.5% / 0.7% / + # 1.4% spread at the three volumes against the heuristic's 15% / 5% + # / 20%. But at ``8^3`` # its 0.6681 ms sits *between* the heuristic's two states (0.6472 # and 0.7453), so it beats the unlucky allocation by 1.12x and # **loses to the lucky one by 0.97x**. It is shipped because the @@ -807,10 +806,10 @@ def _tuned(bm: int, bn: int, bk: int, warps: int, group_m: int = 6, # ``(2048, 1024)`` is deliberately **absent**, and that is a result # rather than an omission: ``128x256x64`` is **1.433x** at ``16^3`` # and **0.804x** and **0.504x** at the two sharded volumes of the - # same pair (``review/opt3/row_2048to1024.json``, two builds each, - # every interval tight). The discriminator is ``BLOCK_K`` against - # ``M`` -- 64 wins at ``M = 4096`` and 128 wins at 2048 and 1024 -- - # this table is keyed per channel pair and cannot say that, and an + # same pair (two builds each, every interval tight). The + # discriminator is ``BLOCK_K`` against ``M`` -- 64 wins at ``M = + # 4096`` and 128 wins at 2048 and 1024 -- this table is keyed per + # channel pair and cannot say that, and an # entry would trade config B's 0.68 ms saving for config C and D's # 0.23 and 0.16 ms losses. ``_fit_to_grid`` already walks # ``BLOCK_M`` with ``M``; making it walk ``BLOCK_K`` too is the @@ -1178,9 +1177,8 @@ def _index_dtype(*operands: torch.Tensor): They are not free -- the AMD backend's buffer-load path requires an i32 offset tensor -- but triton 3.7.1 narrows i64 offsets it can prove safe, so - the cost is paid only where the storage really is over 2 GiB. See PLAN.md - 3.5: storage size, not offset dtype, is the lever, and it is not one the - kernel controls. + the cost is paid only where the storage really is over 2 GiB. Storage size, + not offset dtype, is the lever, and it is not one the kernel controls. *Every* operand the kernel indexes has to be passed here, the weight included -- it is the one this decision used to omit, on an assumption diff --git a/triton_conv3d/reduce_gemm.py b/triton_conv3d/reduce_gemm.py index 7e14a14..ac01440 100644 --- a/triton_conv3d/reduce_gemm.py +++ b/triton_conv3d/reduce_gemm.py @@ -78,24 +78,22 @@ convolutions through DistConv any more (``ScaFFold/unet/conv3d.py`` performs the halo exchange itself, and only on axes that are *genuinely split*), so the kernel is handed ``padding = (1,1,1)`` at one GPU and ``(0,1,1)`` at two or -four. Measured inside running steps at all four configurations -(``work/triton-conv/review/SHAPE_AUDIT.md``): every ``k = 3`` site is padded, -and the veto fired at **eight of them**. +four. A shape census taken inside running steps at all four configurations +settled it: every ``k = 3`` site is padded, and the veto fired at **eight of +them**. It was then raced rather than argued about. On the padded production form of all 18 affected cells, the wide tile against the pinned one, one interleaved -block per cell with 95% intervals -(``work/triton-conv/review/OPTIMIZATION_ROUND3.md`` §1): the wide tile wins -**18 of 18**, geometric mean **1.946x**, range 1.137x-5.336x. The -two-dimensional predicate is real and it costs something; it costs far less -than the arithmetic intensity the wide tile buys. Both clauses are gone. +block per cell with 95% intervals: the wide tile wins **18 of 18**, geometric +mean **1.946x**, range 1.137x-5.336x. The two-dimensional predicate is real +and it costs something; it costs far less than the arithmetic intensity the +wide tile buys. Both clauses are gone. Correctness was never what they protected, and that too is measured rather than assumed: forcing every tuned ``TAP_BLOCK > 1`` row onto the padded form of its own channel pair is bitwise exact against an fp64 reference, in both paddings, -including the ``PADDED and ROW_ALIGNED`` corner, in bf16 and fp32 -(``work/triton-conv/bin/padded_tap_probe.py``). It does not overflow LDS and it -does not move the workspace bound. +including the ``PADDED and ROW_ALIGNED`` corner, in bf16 and fp32. It does not +overflow LDS and it does not move the workspace bound. Determinism =========== @@ -450,8 +448,8 @@ class BwdWeightConfig(ConvConfig): #: six, and it is the same device fact that makes ``GROUP_M`` want to be a #: multiple of six in the gather kernel. **Measured at 1.006x** -- inside #: the round-to-round spread -- so it buys nothing; it is on by default only - #: because every number in M3_RESULTS was taken with it on, and turning it - #: off would make those numbers describe a kernel nobody ran. + #: because every number in this direction's sweep was taken with it on, and + #: turning it off would make those numbers describe a kernel nobody ran. NUM_XCD: int = 6 @property @@ -565,16 +563,15 @@ def default_bwd_weight_config( on the axis it actually splits, so H and W keep the module's ``padding = 1`` and an unsharded run keeps all three. The clause was documented as unreachable ("no real ScaFFold convolution is padded -- DistConv halos them - all"), which was true of the MIOpen rung and false of the shipped one; it - cost a step-level projection a factor of three before anyone checked - (``work/triton-conv/review/SHAPE_AUDIT.md``). + all"), which was true of the MIOpen rung and false of the shipped one; a + shape census inside running steps caught it, but only after it had cost a + step-level projection a factor of three. It is gone as of 2026-08-05, on a measurement rather than on the observation that its premise was false. Raced on the padded production form of the six channel pairs that reach it, the widened tile against the - pinned one, one interleaved block per cell with 95% intervals - (``work/triton-conv/review/opt3/defwide_adapter.json``): widening wins - **6 of 6**, 1.263x-2.084x, every interval clear of 1.000. At + pinned one, one interleaved block per cell with 95% intervals: widening + wins **6 of 6**, 1.263x-2.084x, every interval clear of 1.000. At ``64 -> 128`` the widened *heuristic* beats that pair's tuned row as well (2.084x against 1.867x at ``66x128^2``), which is why the table now carries the wider tile there. @@ -834,24 +831,24 @@ def _tuned(bm: int, bnc: int, tb: int, bk: int, warps: int, #: the convolution a reader would name, as in the other two tables, even though #: this direction's GEMM has ``Cout`` on M and ``taps * Cin`` on N. #: -#: Source: ``work/triton-conv/m3_bwd.json``. Only channel pairs that were -#: actually timed appear; a miss falls to :func:`default_bwd_weight_config` plus -#: :func:`split_count`, which is a real gap and not an extrapolation dressed up -#: as a measurement. +#: Drawn from a backward-weight sweep over the corpus. Only channel pairs that +#: were actually timed appear; a miss falls to +#: :func:`default_bwd_weight_config` plus :func:`split_count`, which is a real +#: gap and not an extrapolation dressed up as a measurement. #: #: ``SPLIT_K`` is left at 0 -- "derive from the shape" -- in every entry, and #: that is a finding rather than an omission. The same channel pair occurs at #: volumes three orders of magnitude apart, the split count is the one knob that #: genuinely has to follow the volume, and pinning a sweep's winner would carry -#: one volume's answer to every other. See M3_RESULTS 5. +#: one volume's answer to every other. #: #: **``matrix_instr_nonkdim`` is 32 in this table and 16 in every other table in #: the package.** That is not an inconsistency, it is the measurement. This #: direction's ``tl.dot(tl.trans(a), b)`` lowers on gfx942 to an *element-wise* #: transpose of A through LDS -- 128 two-byte ``ds_read_u16`` per loop body #: against 16 ``ds_read_b128`` for the untransposed operand, at a structural -#: 50.00% bank-conflict rate -- which leaves the LDS pipe ~65% busy while the -#: matrix core idles at ~32% (``work/triton-conv/review/BWDW_OPTIMIZATION.md`` §1, counters). +#: 50.00% bank-conflict rate -- which by hardware counters leaves the LDS pipe +#: ~65% busy while the matrix core idles at ~32%. #: The 32x32x8 fragment halves the MFMA instruction count for identical FLOPs #: and so stops the MFMA stream competing with that transpose for issue slots. #: It does *not* remove the transpose: the ``nk32`` ISA still emits 128 @@ -863,7 +860,7 @@ def _tuned(bm: int, bnc: int, tb: int, bk: int, warps: int, #: #: Raced per *volume* against the identical tile at ``nonkdim=16``, one #: interleaved block per site, 27 sites covering every ``k=3, Cin >= 64`` cell -#: in the corpus bar the 2 GiB cliff (``work/triton-conv/tune/nk_bwdw.json``): +#: in the corpus bar the 2 GiB cliff: #: #: ============ ========================================= ====== #: pair gain by volume shipped @@ -884,15 +881,14 @@ def _tuned(bm: int, bnc: int, tb: int, bk: int, warps: int, #: ============ ========================================= ====== #: #: The last two rows are why this was raced per volume rather than per pair. -#: ``BWDW_OPTIMIZATION.md`` §3.1 reports ``nk32`` winning **13 of 13** and -#: "never losing"; every one of those 13 was a ``Cout <= 512`` site, and at -#: ``Cout = 1024`` the sign reverses at two of the three volumes each pair has. +#: An earlier per-pair race had ``nk32`` winning **13 of 13** and never losing; +#: every one of those 13 was a ``Cout <= 512`` site, and at ``Cout = 1024`` the +#: sign reverses at two of the three volumes each pair has. #: Both pairs keep ``nonkdim=16``. The transposed operator's weight gradient -#: runs this same kernel and was raced too (7 sites, -#: ``work/triton-conv/tune/nk_convT.json``): 0.934-1.056x, no consistent sign, -#: so :func:`default_bwd_weight_config` -- which is the only thing serving it, -#: and also serves fp32 and every untuned pair, including the four -#: 2048-channel backward-weight sites of a scale-8 run -- **stays at 16**. +#: runs this same kernel and was raced too (7 sites): 0.934-1.056x, no +#: consistent sign, so :func:`default_bwd_weight_config` -- which is the only +#: thing serving it, and also serves fp32 and every untuned pair, including the +#: four 2048-channel backward-weight sites of a scale-8 run -- **stays at 16**. #: (Until 2026-08-05 it also served the eight padded production sites whose #: tuned row widens ``TAP_BLOCK``, because the resolver declined those rows; #: it no longer does, and those eight now run the table.) A heuristic is the @@ -905,7 +901,7 @@ def _tuned(bm: int, bnc: int, tb: int, bk: int, warps: int, #: different MFMA fragment sums the same products in a different order. They #: are still bitwise *reproducible*, which is what the determinism claim says: #: the split-K partition, the reduction's tiling and its ``tl.sum`` order are -#: all unchanged, and ``work/triton-conv/review/bwd/det_verify.py`` was re-run on this table. +#: all unchanged, and the run-to-run determinism check was re-run on this table. _TUNED_BWD_W: dict[tuple, BwdWeightConfig] = { # The segmentation head, which until now had no row at all and no ``nk`` # question either: ``BLOCK_M`` is ``Cout = 6`` rounded up to 16, so @@ -913,8 +909,7 @@ def _tuned(bm: int, bnc: int, tb: int, bk: int, warps: int, # What does apply is ``num_warps``, which no sweep in this project has ever # taken below 4 (``candidate_bwd_weight_configs`` draws it from # ``{4, 8, seed}``). Raced at all three head volumes, 1 warp is 1.044x, - # 1.052x and 1.033x over the heuristic's 4 - # (``work/triton-conv/tune/warps_bwdw.json``). It is still a **loss** + # 1.052x and 1.033x over the heuristic's 4. It is still a **loss** # against MIOpen at two of those three volumes (0.73x, 0.93x, 1.19x); the # row is here so the number the adapter's block-list is built from is the # best this kernel can do, not the best it happened to be doing. @@ -929,8 +924,7 @@ def _tuned(bm: int, bnc: int, tb: int, bk: int, warps: int, # ``nonkdim=32`` wins, and it wins at all three of that pair's volumes: # 1.128x @ 64^3, 1.112x @ 32x128^2, 1.069x @ 64x128^2. The other three # pairs measure 0.926-1.056x with no consistent sign and stay on - # ``default_bwd_weight_config``, i.e. at 16 - # (``work/triton-conv/tune/nk_convT.json``). The split is the mechanism, + # ``default_bwd_weight_config``, i.e. at 16. The split is the mechanism, # not luck: this pair is the only transposed site whose tile is # ``BLOCK_M = 128``; the other three reach ``BLOCK_M = 256``, which already # amortises the transpose, and that is the same boundary the ``k=3`` rows @@ -940,7 +934,8 @@ def _tuned(bm: int, bnc: int, tb: int, bk: int, warps: int, **{ tune_key(torch.bfloat16, cin, cout, (3, 3, 3)): cfg for (cin, cout), cfg in { - # The UNet stem, and the row that refutes ``OPTIMIZATION.md`` D4. + # The UNet stem, and the row that refutes the standing verdict that + # ``conv 3->64`` is hopeless and belongs on MIOpen. # It is the forward's disease one axis over: ``BLOCK_NC = # _pow2_at_most(Cin, 256) = 16`` against ``Cin = 3``, times # ``TAP_BLOCK = 16`` covering 27 taps in two groups, is **512 issued @@ -974,7 +969,6 @@ def _tuned(bm: int, bnc: int, tb: int, bk: int, warps: int, # ``BLOCK_NC`` / ``TAP_BLOCK`` only decide which columns a program # owns, never the order a column is summed in. So no determinism # baseline moves and the 168.8 MiB workspace bound is unchanged. - # ``work/triton-conv/review/LOSS_OPTIMIZATION.md`` §1.2. (3, 64): _tuned(64, 4, 16, 64, 1, nk=16), # Cout = 64 is where TAP_BLOCK earns its existence: the tile can only # get to 512 columns through the taps, and getting there is worth @@ -984,8 +978,7 @@ def _tuned(bm: int, bnc: int, tb: int, bk: int, warps: int, # for a 128-row tile, and it wants one. Raced against the # ``64x512x16/tb8`` row it replaces, one interleaved block per cell, # kernel-only, 95% intervals -- **both** forms, because this table - # is shared and the padded form is not the only caller - # (``review/opt3/r64x128_{adapter,distconv}.json``): + # is shared and the padded form is not the only caller: # padded 66x128^2 1.6928 vs 1.9896 ms -- 1.177x [1.165,1.189] # padded 34x128^2 0.8760 vs 1.0036 -- 1.145x [1.140,1.151] # padded 64^3 0.4277 vs 0.5059 -- 1.183x [1.180,1.185] @@ -1078,10 +1071,9 @@ def bwd_weight_config( :func:`default_bwd_weight_config` with ``TAP_BLOCK`` pinned to 1. Raced arm against arm on the padded production form of all 18 affected - cells (``work/triton-conv/review/OPTIMIZATION_ROUND3.md`` §1; the tuned row - forced against the config the decline produced, in one interleaved block - per cell, CUDA-graph replay, 95% intervals): the tuned row is faster in - **18 of 18**, geometric mean **1.946x**, range 1.137x-5.336x, every + cells (the tuned row forced against the config the decline produced, in one + interleaved block per cell, CUDA-graph replay, 95% intervals): the tuned row + is faster in **18 of 18**, geometric mean **1.946x**, range 1.137x-5.336x, every interval clear of 1.000. Worst cell was the stem, ``3->64 k3 @ 130x256^2``, at 7.9505 ms declined against 1.4910 ms with the row. @@ -1090,10 +1082,9 @@ def bwd_weight_config( 1.7-1.9x at ``Cout = 64`` where the tile has no height to trade. That race forces the row; what the *shipped* resolver then delivers was - measured separately, over the whole corpus in the production shape, at the - commit before this clause was removed and at the commit after - (``work/triton-conv/review/HEAD_MEASUREMENT.md`` §3). **1.813x** geometric - mean over the 21 backward-weight cells of those six channel pairs, + measured separately, over the whole corpus in the production shape, before + this clause was removed and after. **1.813x** geometric mean over the 21 + backward-weight cells of those six channel pairs, 105.5 ms summed to 57.5 ms, worst cell the stem at 8.0000 ms to 1.523 ms. The forced-row 1.946x above is the counterfactual and this is the delivered figure; quote this one for the shipped kernel. @@ -1121,12 +1112,12 @@ def bwd_weight_config( pure function of the shape and the config. See the module docstring. That inequality is **measured**, not inferred, and it is wider than the - split count suggests (``work/triton-conv/review/verify3/ulp_bf16.json``, one - ordinary draw per cell, the six pairs at both production paddings). In bf16 - the old and new configs disagree at **10 of 12 cells**, but - :func:`split_count` moves at only 5 of them: the other five move because the - row changes ``matrix_instr_nonkdim`` from 16 to 32, and a different MFMA - fragment sums the same products in a different order within one ``BLOCK_K``. + split count suggests (one ordinary draw per cell, the six pairs at both + production paddings). In bf16 the old and new configs disagree at **10 of + 12 cells**, but :func:`split_count` moves at only 5 of them: the other five + move because the row changes ``matrix_instr_nonkdim`` from 16 to 32, and a + different MFMA fragment sums the same products in a different order within + one ``BLOCK_K``. The one pair that stays bitwise identical, ``3 -> 64``, is the one whose new row keeps both ``nonkdim`` and ``BLOCK_K``. So the thing to check before asserting two configs agree is not the split count alone. The disagreement @@ -1324,7 +1315,7 @@ def conv3d_backward_weight( mirrors ``torch.nn.grad.conv3d_weight``. ``deterministic`` defaults to **True**, which is a deliberate divergence - from PLAN.md 3.4's ``deterministic=None`` following + from the original design, where ``deterministic=None`` was to follow ``torch.are_deterministic_algorithms_enabled()``. The atomic path exists to price determinism, not to be selected; if it were the default whenever torch's flag is off then every ScaFFold run would silently get the diff --git a/triton_conv3d/scaffold_census.json b/triton_conv3d/scaffold_census.json index 456bf8e..fbb9939 100644 --- a/triton_conv3d/scaffold_census.json +++ b/triton_conv3d/scaffold_census.json @@ -1,5 +1,5 @@ { - "source": "work/triton-conv/bin/conv_census.py, three training steps per configuration, wrapped entry points", + "source": "instrumented training run, three steps per configuration, wrapped entry points", "form": "adapter", "note": "The shape and padding the kernel was handed, read off a real call. Every k>1 convolution here is padded: the ScaFFold adapter exchanges a halo only on genuinely split axes.", "configs": [ diff --git a/triton_conv3d/scaffold_corpus.json b/triton_conv3d/scaffold_corpus.json index 1170655..8b6595c 100644 --- a/triton_conv3d/scaffold_corpus.json +++ b/triton_conv3d/scaffold_corpus.json @@ -1,5 +1,5 @@ { - "source": "model-analysis/unet_shapes.py via work/report/data", + "source": "model-analysis/unet_shapes.py", "configs": [ { "tag": "A", diff --git a/triton_conv3d/shapes.py b/triton_conv3d/shapes.py index 0cd3ff1..84fba46 100644 --- a/triton_conv3d/shapes.py +++ b/triton_conv3d/shapes.py @@ -13,10 +13,9 @@ ``census_corpus()`` Every convolution an **instrumented ScaFFold training step actually issued**, at the four configurations the current benchmark harness runs. - Recorded by wrapping the entry points inside a running step - (``work/triton-conv/bin/conv_census.py``), so it is a measurement rather - than a model. It carries no MIOpen timings; it exists to say what the - shapes *are*. + Recorded by wrapping the entry points inside a running step, so it is a + measurement rather than a model. It carries no MIOpen timings; it exists + to say what the shapes *are*. ``edge_cases()`` Synthetic problems chosen to break addressing, masking and tiling @@ -101,7 +100,7 @@ #: ``global_load_dwordx4`` -- from ``arg.untyped_storage().size() <= 2**31 - 1``. #: That check reads the whole *storage* and it counts *bytes*, so it is a #: different question from whether an element index overflows int32, and the two -#: answers differ by ``elem_bytes`` (``PLAN.md`` §3.5). +#: answers differ by ``elem_bytes``. BUFFER_OP_MAX_BYTES = 2**31 - 1 @@ -263,11 +262,10 @@ def production_variant(self) -> "ConvProblem": ``padding = (0, 1, 1)`` at ``(D_loc + 2, H, W)``. **Every production convolution with ``k > 1`` is therefore padded, at - every configuration.** Measured inside running steps at all four - (``work/triton-conv/review/SHAPE_AUDIT.md``), not inferred: 18 of the 19 - distinct ordinary convolutions at scale 7 on one GPU arrive with - ``padding = (1, 1, 1)``, and the nineteenth is the ``k = 1`` head, which - has no padding to begin with. + every configuration.** Measured inside running steps at all four, not + inferred: 18 of the 19 distinct ordinary convolutions at scale 7 on + one GPU arrive with ``padding = (1, 1, 1)``, and the nineteenth is the + ``k = 1`` head, which has no padding to begin with. Dropping the zero slabs on the unsplit axes is a deliberate and separately verified decision -- ``cat(zeros, x, zeros)`` at @@ -489,7 +487,8 @@ def scaffold_corpus() -> tuple[ConvProblem, ...]: """Every distinct convolution in the three profiled ScaFFold configurations. Ordered by measured cost, so truncating the list keeps the problems that - matter. Regenerate with ``work/triton-conv/make_corpus.py``. + matter. Loaded from ``scaffold_corpus.json``, which is generated from the + profiled shape dumps rather than written by hand. """ raw = json.loads(_CORPUS_PATH.read_text()) problems = [] @@ -559,8 +558,8 @@ def production_corpus() -> tuple[ConvProblem, ...]: def census_corpus() -> tuple[ConvProblem, ...]: """Every convolution an instrumented ScaFFold step actually issued. - Recorded by ``work/triton-conv/bin/conv_census.py``, which wraps - ``FastConv3d``/``FastConvTranspose3d`` and the six kernel entry points and + Recorded by a census harness that wraps ``FastConv3d`` / + ``FastConvTranspose3d`` and the six kernel entry points and runs three real training steps at each of the four configurations the benchmark harness uses (A = scale 7 / 1 GPU, B = scale 8 / 1 GPU, C = scale 8 / 2 GPUs, D = scale 8 / 4 GPUs). Every problem here is in @@ -574,9 +573,8 @@ def census_corpus() -> tuple[ConvProblem, ...]: GPU), and a *network depth* the profiled corpus does not -- the shape dumps behind :func:`scaffold_corpus` were taken at ``unet_bottleneck_dim = 4`` at scale 8, giving a four-layer model topping - out at 1024 channels, while every step-level measurement in - ``work/triton-conv/minibatch/`` runs the shipped default of 3, i.e. a - five-layer model topping out at 2048; + out at 1024 channels, while every step-level measurement in this project + runs the shipped default of 3, i.e. a five-layer model topping out at 2048; * it carries no ``measured`` MIOpen data and no cost ordering, so it is not a priority list and must not be used as one; * and :func:`scaffold_corpus`'s ordering, indices and contents are the key diff --git a/triton_conv3d/tests/test_bwd_weight.py b/triton_conv3d/tests/test_bwd_weight.py index 65a11a6..d9fe2b8 100644 --- a/triton_conv3d/tests/test_bwd_weight.py +++ b/triton_conv3d/tests/test_bwd_weight.py @@ -8,8 +8,8 @@ correctness lives in :func:`split_count` agreeing with what the kernel and the reduction pass assume about each other, so the tests pin the arithmetic (every voxel in exactly one split) as well as the answer; -* **determinism**, stated as the claim PLAN.md 3.4 makes and tested the way that - claim is worded: bitwise identical run to run *and process to process*, for the +* **determinism**, tested the way the package's determinism claim is worded: + bitwise identical run to run *and process to process*, for the same input, dtype, shape, device and tuning config. There are three tests -- in-process repetition, three separate interpreters, and a negative control on the atomic path -- because a determinism test that cannot fail is the most @@ -704,8 +704,7 @@ def test_a_tuned_tap_block_row_survives_the_padding(): the assertion with a measurement. That is what happened. Raced on the padded production form of all 18 affected cells, the tuned row against the config the decline produced, one interleaved block per cell with 95% - intervals (``work/triton-conv/review/OPTIMIZATION_ROUND3.md`` §1): the - tuned row wins **18 of 18**, geometric mean **1.946x**, range + intervals: the tuned row wins **18 of 18**, geometric mean **1.946x**, range 1.137x-5.336x, worst cell 7.9505 ms declined against 1.4910 ms with the row. The heuristic's half was raced separately on the six pairs that reach it and widening wins **6 of 6**, 1.263x-2.084x. diff --git a/triton_conv3d/tests/test_infra.py b/triton_conv3d/tests/test_infra.py index 5a125c5..825a7b4 100644 --- a/triton_conv3d/tests/test_infra.py +++ b/triton_conv3d/tests/test_infra.py @@ -143,8 +143,7 @@ def test_the_2gib_cliff_is_a_byte_problem_and_not_an_index_problem(): That shape holds 1.11e9 elements -- about half of int32's range -- so an element-counting predicate says nothing about it, and ``needs_int64`` used to be read as though it did. What it exceeds is the *byte* limit on - the whole storage, which is what decides buffer-op eligibility - (``PLAN.md`` §3.5). + the whole storage, which is what decides buffer-op eligibility. """ cliff = next(p.halo_variant for p in scaffold_corpus() if p.halo_variant.label == "conv 128->64 k3x3x3 @ 130x258x258") @@ -662,12 +661,12 @@ def test_spread_is_a_range_statistic_and_the_interval_is_not(): ``(max - min) / median`` is a *range*, and the expected range of ``n`` samples grows like ``d2(n)`` even on a perfectly stationary device. - Measured on this node with one kernel held constant for 14 minutes - (47,686 blocks, ``work/triton-conv/bin/drift_probe.py``) the median of - this statistic runs 0.23% at 2 rounds, 0.63% at 6, 0.98% at 20 and 2.70% at - 100 -- all of it arithmetic, none of it the machine. Since ``rounds`` is - now chosen per cell, two cells' spreads are not comparable to each other at - all, and the replacement has to be an interval. + Measured on this node with one kernel held constant for 14 minutes (47,686 + blocks) the median of this statistic runs 0.23% at 2 rounds, 0.63% at 6, + 0.98% at 20 and 2.70% at 100 -- all of it arithmetic, none of it the + machine. Since ``rounds`` is now chosen per cell, two cells' spreads are + not comparable to each other at all, and the replacement has to be an + interval. Pinned on a fixed draw so it tests the formulae, not the GPU. """ @@ -944,9 +943,9 @@ def test_measure_one_refuses_to_report_a_heuristic_time(): #: Corpus cells the harness is anchored to. Chosen because their isolated and #: profiled shapes genuinely match: the halo'd input is 281 MiB, well under the #: 2 GiB threshold above which MIOpen abandons its tuned solvers for the naive -#: non-packed ones and the isolated and profiled numbers legitimately diverge -#: (``work/nonpacked-conv/RESULTS.md``). The profiled time is read from the -#: corpus rather than copied here so there is one source of truth for it. +#: non-packed ones and the isolated and profiled numbers legitimately diverge. +#: The profiled time is read from the corpus rather than copied here so there +#: is one source of truth for it. ANCHOR_CELLS = ((6, "fwd"), (6, "bwd-data")) @@ -1065,8 +1064,8 @@ def sporadic(): # That is a decision about what to measure, and it has exactly one way to go # wrong -- taking the launcher out of one arm and not the other, which at these # sizes is worth up to 1.4x in the direction that flatters us. These tests are -# the guard on that, and each of them is verified by breaking what it tests -# (``work/triton-conv/bin/mutate_bench.py``). +# the guard on that, and each of them was verified by mutation -- breaking the +# thing it tests and confirming it fails. def test_the_graph_chunk_is_one_ruler_for_every_arm(): @@ -1077,8 +1076,8 @@ def test_the_graph_chunk_is_one_ruler_for_every_arm(): of 1, 2, 4, 8, 16 and 32 calls on four real arms. At ``chunk = 1`` that is 45% of a 0.028 ms kernel and only 19% of a 0.068 ms one, so a per-arm chunk would be a per-arm instrument: exactly the failure ``_common_group`` already - documents, one level up (H10, where two byte-identical arms picked different - event groups and read 4% apart). + documents, one level up, where two byte-identical arms picked different + event groups and read 4% apart. Hence: the rule reads only ``min(durations)``, so two arms of the same call can never be given different rulers. @@ -1147,8 +1146,7 @@ def test_a_captured_ratio_of_two_identical_arms_covers_one(): Two arms doing byte-identical work have a true ratio of exactly 1.000, so anything else is the instrument. Measured over 12 replications on - ``convT 1024->512 @ 8^3`` through the shipped decision path - (``work/triton-conv/bin/launcher_symmetry.py --only null``): under + ``convT 1024->512 @ 8^3`` through the shipped decision path: under ``exclude`` the median is 0.9996, the range 0.9982-1.0021, and **12 of 12** intervals cover 1.000. @@ -1162,7 +1160,7 @@ def test_a_captured_ratio_of_two_identical_arms_covers_one(): # 512, not a "nicer" 256 or 384: on this torch/ROCm build a bf16 # ``a @ a`` is **~600 ms** at 128, 192, 256, 320, 384, 448, 640 and # 768, and 0.019 ms at 512 and 1024. That is the ``torch.mm`` bf16 - # pathology ``STATUS.md`` §7 already owes upstream, measured here from + # pathology this project already owes upstream, measured here from # a second direction; a test that picked one of the slow sizes would # be timing a 600 ms kernel and would correctly be told it does not # need a graph. @@ -1177,9 +1175,9 @@ def test_a_captured_ratio_of_two_identical_arms_covers_one(): # which is narrower than the between-race scatter of the same pair (sd # 0.13%, range 0.9982-1.0021 over 12 replications) -- so an interval that # misses 1 by 0.2% here is the same residual the sequential protocol has - # (0.32%, HARNESS_RIGOR Part 1), not a biased instrument. What a biased - # instrument looks like is 4% (H10, per-arm event groups) or 45% (a one-call - # graph), and 1% catches both. The coverage claim is the 12-replication + # (0.32%), not a biased instrument. What a biased instrument looks like + # is 4% (per-arm event groups) or 45% (a one-call graph), and 1% catches + # both. The coverage claim is the 12-replication # experiment, where 12 of 12 intervals contained 1. assert abs(r.point - 1.0) < 0.01, ( f"two byte-identical arms read {r} under the kernel-time definition" @@ -1210,7 +1208,7 @@ def test_an_inflated_launcher_does_not_move_the_reported_kernel_time(): # 512, not a "nicer" 256 or 384: on this torch/ROCm build a bf16 # ``a @ a`` is **~600 ms** at 128, 192, 256, 320, 384, 448, 640 and # 768, and 0.019 ms at 512 and 1024. That is the ``torch.mm`` bf16 - # pathology ``STATUS.md`` §7 already owes upstream, measured here from + # pathology this project already owes upstream, measured here from # a second direction; a test that picked one of the slow sizes would # be timing a 600 ms kernel and would correctly be told it does not # need a graph. @@ -1264,7 +1262,7 @@ def test_the_replay_cost_is_amortized_by_the_chunk(): # 512, not a "nicer" 256 or 384: on this torch/ROCm build a bf16 # ``a @ a`` is **~600 ms** at 128, 192, 256, 320, 384, 448, 640 and # 768, and 0.019 ms at 512 and 1024. That is the ``torch.mm`` bf16 - # pathology ``STATUS.md`` §7 already owes upstream, measured here from + # pathology this project already owes upstream, measured here from # a second direction; a test that picked one of the slow sizes would # be timing a 600 ms kernel and would correctly be told it does not # need a graph. @@ -1304,7 +1302,7 @@ def test_a_capture_failure_takes_the_whole_cell_back_to_eager(): # 512, not a "nicer" 256 or 384: on this torch/ROCm build a bf16 # ``a @ a`` is **~600 ms** at 128, 192, 256, 320, 384, 448, 640 and # 768, and 0.019 ms at 512 and 1024. That is the ``torch.mm`` bf16 - # pathology ``STATUS.md`` §7 already owes upstream, measured here from + # pathology this project already owes upstream, measured here from # a second direction; a test that picked one of the slow sizes would # be timing a 600 ms kernel and would correctly be told it does not # need a graph. @@ -1569,9 +1567,9 @@ def test_a_control_free_row_omits_the_control_rather_than_zeroing_it(): And a case built with ``control=False`` must not construct the control either. For a backward direction the control is a real ``F.conv3d`` forward graph, and *running* it once is where MIOpen's find is paid -- 92-174 s per - cell on this corpus (``work/triton-conv/review/HARNESS_SPEED.md`` 1). - Dropping the arm from the timing while still building it would save the - timing and none of the cost, which is the whole point of the flag. + cell on this corpus. Dropping the arm from the timing while still building + it would save the timing and none of the cost, which is the whole point of + the flag. """ from triton_conv3d.bench.conv_bench import _build, measure_problem diff --git a/triton_conv3d/transposed.py b/triton_conv3d/transposed.py index 1f2d383..1b1681b 100644 --- a/triton_conv3d/transposed.py +++ b/triton_conv3d/transposed.py @@ -456,10 +456,10 @@ def transposed_tune_key(dtype: torch.dtype, cin: int, cout: int, #: Measured winners for the transposed forward, keyed by ``(dtype, Cin, Cout, -#: kernel)``. Source: ``work/triton-conv/m5_fwd.json`` -- the four -#: ``ConvTranspose3d`` channel pairs the model contains, swept over the tile and -#: ``TAP_BLOCK`` grid of :func:`candidate_transposed_configs` and then raced -#: against MIOpen. A miss falls back to :func:`default_transposed_config`. +#: kernel)``: the four ``ConvTranspose3d`` channel pairs the model contains, +#: swept over the tile and ``TAP_BLOCK`` grid of +#: :func:`candidate_transposed_configs` and then raced against MIOpen. A miss +#: falls back to :func:`default_transposed_config`. #: #: Keyed on the channel widths and not the volume, as the gather kernel's table #: is -- and with the same caveat, which this project has now paid for twice: a @@ -468,8 +468,7 @@ def transposed_tune_key(dtype: torch.dtype, cin: int, cout: int, #: measured winning at every volume the pair occurs at. Each pair below occurs #: at three volumes (one per profiled configuration) and the entry named won all #: three; the speedups they produce differ by up to 2.3x between those volumes, -#: are recorded per volume in ``work/triton-conv/m5_shipped_*.json``, and are -#: never averaged. +#: and were therefore recorded per volume rather than averaged. #: #: **Two of the four pairs are deliberately absent.** ``512 -> 256`` and #: ``1024 -> 512`` were swept just as thoroughly and From 77cf50d8b9ce033515ba6e8cbacf4f56febe2f6b Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Wed, 5 Aug 2026 15:24:50 -0700 Subject: [PATCH 6/8] Add helper script to run tests. --- scripts/run-tests.sh | 120 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 120 insertions(+) create mode 100755 scripts/run-tests.sh diff --git a/scripts/run-tests.sh b/scripts/run-tests.sh new file mode 100755 index 0000000..eb93fed --- /dev/null +++ b/scripts/run-tests.sh @@ -0,0 +1,120 @@ +#!/bin/bash +# Run the ScaFFold test suites with the environment they require. +# +# scripts/run-tests.sh # both suites +# scripts/run-tests.sh scaffold # tests/ only +# scripts/run-tests.sh triton # triton_conv3d/tests/ only +# scripts/run-tests.sh scaffold -x -k gn # extra args go to pytest +# +# Set PYTHON to choose an interpreter; otherwise the first virtualenv under +# .venvs/ is used, falling back to python3 on PATH. +# +# Everything this script exports is here because omitting it changes the +# result, not because it seemed prudent. See the comments at each one. + +set -euo pipefail + +cd "$(dirname "${BASH_SOURCE[0]}")/.." + +# --- interpreter ------------------------------------------------------------ +# There is no editable install: the packages are importable from the repo root +# and nowhere else, hence the cd above and `python -m pytest` rather than a +# bare `pytest` (which would run from wherever its console script resolves). +if [ -z "${PYTHON:-}" ]; then + for _venv in .venvs/*/bin/python; do + [ -x "$_venv" ] && PYTHON="$_venv" && break + done + PYTHON="${PYTHON:-python3}" +fi +# --- required: channels-last has to reach MIOpen ---------------------------- +# Without this, channels_last_3d is inert on ROCm and MIOpen is silently handed +# NCDHW -- a different problem than the one under test. This is not a tuning +# preference: both parametrizations of +# tests/test_groupnorm.py::test_gpu_triton_dctensor_matches_eager_and_stays_wrapped +# fail deterministically when it is unset. It is also what production runs set. +export PYTORCH_MIOPEN_SUGGEST_NHWC=1 + +# --- required: ROCm needs a writable TMPDIR --------------------------------- +# ROCm aborts the process (SIGABRT, no Python traceback) when it cannot write +# to TMPDIR, so an unwritable one reads as a crashed test run rather than as a +# configuration error. Check it here, where the message can say so. +_tmp="${TMPDIR:-/tmp}" +if ! ( : > "$_tmp/.scaffold-write-probe.$$" ) 2>/dev/null; then + echo "error: TMPDIR ($_tmp) is not writable; ROCm will abort the run." >&2 + echo " Set TMPDIR to a writable directory and re-run." >&2 + exit 1 +fi +rm -f "$_tmp/.scaffold-write-probe.$$" +export TMPDIR="$_tmp" + +# --- if set, these caches have to be writable ------------------------------- +# Neither Triton nor MIOpen fails when it cannot write its cache; both just +# redo the work every time. Nothing reports it, so the suite reads as hung +# rather than as misconfigured. Validate whatever the caller has set. +for _var in TRITON_CACHE_DIR MIOPEN_USER_DB_PATH MIOPEN_CUSTOM_CACHE_DIR; do + _dir="${!_var:-}" + [ -n "$_dir" ] || continue + if ! mkdir -p "$_dir" 2>/dev/null || + ! ( : > "$_dir/.scaffold-write-probe.$$" ) 2>/dev/null; then + echo "error: $_var ($_dir) is not writable; unset it or point it somewhere else." >&2 + exit 1 + fi + rm -f "$_dir/.scaffold-write-probe.$$" +done + +# --- runtime warning: a cold MIOpen find database dominates the run --------- +# The tests compare against MIOpen, and with an empty find database MIOpen +# searches for an algorithm per convolution problem instead of looking one up. +# Measured on triton_conv3d/tests/test_bwd_data.py (305 tests): 322 s cold +# against 14.7 s with a populated database -- 22x, and it is all search, not +# test work. MIOPEN_USER_DB_PATH defaults to ~/.config/miopen; point it at a +# warm database to avoid paying this on every run. +_miopen_db="${MIOPEN_USER_DB_PATH:-$HOME/.config/miopen}" +if ! ls "$_miopen_db"/*.ufdb.txt >/dev/null 2>&1; then + echo "note: MIOpen find database ($_miopen_db) is cold, so this run will be" >&2 + echo " slow -- ~22x on the convolution tests, all of it algorithm search." >&2 + echo " Set MIOPEN_USER_DB_PATH to a warm database to skip it." >&2 +fi + +# --- interpreter check + coverage warning ----------------------------------- +# One import for both: torch is slow to load, and this is the only thing the +# script needs from it. Runs after the exports above so it inherits them. +# +# The cross-device tests skip themselves when only one device is visible, so a +# one-device run reports a healthy pass count with those clauses never +# exercised. Warn rather than fail: a one-device run is still worth doing, it +# is just not the full one. +_devices=$("$PYTHON" - <<'PY' 2>/dev/null +import ScaFFold, torch # noqa: F401 -- import is the check +print(torch.cuda.device_count() if torch.cuda.is_available() else 0) +PY +) || { + echo "error: $PYTHON cannot import ScaFFold and torch." >&2 + echo " Set PYTHON to the right interpreter, or run from the repo root." >&2 + exit 1 +} +if [ "$_devices" -lt 2 ]; then + echo "warning: $_devices GPU(s) visible; the cross-device tests in" >&2 + echo " test_gather_gemm.py and test_bwd_weight.py will skip." >&2 + echo " Two or more devices are needed for full coverage." >&2 +fi + +# --- run -------------------------------------------------------------------- +# The mpi-marked tests skip themselves when no launcher is present, so they need +# no deselection here. +_suite="${1:-all}" +case "$_suite" in + scaffold|triton|all) shift || true ;; + *) _suite=all ;; +esac + +_status=0 +if [ "$_suite" = all ] || [ "$_suite" = scaffold ]; then + echo "== ScaFFold suite ==" + "$PYTHON" -m pytest tests -q "$@" || _status=$? +fi +if { [ "$_suite" = all ] || [ "$_suite" = triton ]; } && [ -d triton_conv3d/tests ]; then + echo "== triton_conv3d suite ==" + "$PYTHON" -m pytest triton_conv3d/tests -q "$@" || _status=$? +fi +exit $_status From 981e7afa0c7048f5f1d7c7b7e57acb03f11c4fe6 Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Wed, 5 Aug 2026 15:50:17 -0700 Subject: [PATCH 7/8] Satisfy the style checks CI actually runs The style workflow runs `ruff format --diff .` and `ruff check .`, and both were failing: 22 files would be reformatted and there were 19 lint errors, all of the latter in the Triton package. Most of this is whitespace. Three fixes are not mechanical: `gemm_probe` built two closures over `a` and `b` in functions whose `finally` deletes both names, so each was correct only because the harness happens to call it before the cleanup runs. They now bind the tensors as default arguments, which captures at definition time and does not depend on call order. This is also what ruff was reporting as F821. `baseline._callable` assigned two lambdas to names, now plain functions. Six imports were unused and are gone. `ScaFFold/viz/standard_viz.py` is not part of this branch's work -- it fails the format check on round2-fixes too, and is reformatted here only so the check can pass. Both suites re-run afterwards, since several tests read their subject's source: triton_conv3d 1425 passed / 16 skipped, ScaFFold 742 passed / 8 skipped / 1 xfailed. Reformatting changes Triton's JIT cache key, so the first run after this recompiles every kernel and takes ~5x longer. --- ScaFFold/unet/conv3d.py | 4 +- ScaFFold/viz/standard_viz.py | 16 +- tests/test_conv3d.py | 127 ++-- tests/test_groupnorm_wiring_edge.py | 4 +- tests/test_platform_guard.py | 13 +- triton_conv3d/__init__.py | 59 +- triton_conv3d/bench/__init__.py | 10 +- triton_conv3d/bench/baseline.py | 246 +++++--- triton_conv3d/bench/conv_bench.py | 807 ++++++++++++++++-------- triton_conv3d/bench/gemm_probe.py | 646 +++++++++++++------ triton_conv3d/bench/harness.py | 211 +++++-- triton_conv3d/bwd_data.py | 69 +- triton_conv3d/gather_gemm.py | 245 ++++--- triton_conv3d/reduce_gemm.py | 333 +++++++--- triton_conv3d/reference.py | 30 +- triton_conv3d/shapes.py | 107 +++- triton_conv3d/tests/test_bwd_data.py | 95 ++- triton_conv3d/tests/test_bwd_weight.py | 368 +++++++---- triton_conv3d/tests/test_gather_gemm.py | 199 +++--- triton_conv3d/tests/test_infra.py | 197 ++++-- triton_conv3d/tests/test_transposed.py | 325 ++++++---- triton_conv3d/transposed.py | 285 ++++++--- 22 files changed, 2973 insertions(+), 1423 deletions(-) diff --git a/ScaFFold/unet/conv3d.py b/ScaFFold/unet/conv3d.py index 2178dfe..ec1a4ed 100644 --- a/ScaFFold/unet/conv3d.py +++ b/ScaFFold/unet/conv3d.py @@ -767,9 +767,7 @@ def _exchange_forward(x, strategy, dim_index, dim, halo): slab_shape = list(x.shape) slab_shape[dim] = halo - recv_minus, recv_minus_view = _packed_slab( - slab_shape, x.dtype, x.device, zero=True - ) + recv_minus, recv_minus_view = _packed_slab(slab_shape, x.dtype, x.device, zero=True) recv_plus, recv_plus_view = _packed_slab(slab_shape, x.dtype, x.device, zero=True) ops = [] diff --git a/ScaFFold/viz/standard_viz.py b/ScaFFold/viz/standard_viz.py index d9f2dda..36fa6a5 100644 --- a/ScaFFold/viz/standard_viz.py +++ b/ScaFFold/viz/standard_viz.py @@ -49,7 +49,9 @@ def main(config: RunConfig): if "val_loss_avg" in row: val_loss.append(float(row["val_loss_avg"])) - plot_title = f"v={config.vol_size}, c={config.n_categories}, u={config.unet_layers}" + plot_title = ( + f"v={config.vol_size}, c={config.n_categories}, u={config.unet_layers}" + ) line_thickness = 2 fontsize = 20 tick_fontsize = 14 @@ -64,7 +66,9 @@ def main(config: RunConfig): plt.tick_params(axis="both", which="major", labelsize=tick_fontsize) plt.yscale("log") plt.title(plot_title, fontsize=12) - plt.legend(loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize) + plt.legend( + loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize + ) plt.grid(True, axis="y") plt.savefig(figures_path / "train_loss.png", dpi=300, bbox_inches="tight") plt.close(figures[-1]) @@ -76,7 +80,9 @@ def main(config: RunConfig): plt.ylabel("Val dice score", fontsize=fontsize) plt.tick_params(axis="both", which="major", labelsize=tick_fontsize) plt.title(plot_title, fontsize=12) - plt.legend(loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize) + plt.legend( + loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize + ) plt.grid(True, axis="y") plt.savefig(figures_path / "val_dice.png", dpi=300, bbox_inches="tight") plt.close(figures[-1]) @@ -89,7 +95,9 @@ def main(config: RunConfig): plt.ylabel("Val loss", fontsize=fontsize) plt.tick_params(axis="both", which="major", labelsize=tick_fontsize) plt.title(plot_title, fontsize=12) - plt.legend(loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize) + plt.legend( + loc="upper left", bbox_to_anchor=legend_loc, fontsize=legend_fontsize + ) plt.grid(True, axis="y") plt.savefig(figures_path / "val_loss.png", dpi=300, bbox_inches="tight") plt.close(figures[-1]) diff --git a/tests/test_conv3d.py b/tests/test_conv3d.py index cd9991b..9a27170 100644 --- a/tests/test_conv3d.py +++ b/tests/test_conv3d.py @@ -337,8 +337,9 @@ def test_a_sharded_dctensor_forward_goes_to_miopen_without_a_process_group(monke monkeypatch.setattr( FastConv3d, "_triton_forward", - lambda self, local, plan=None: fast_calls.append(local) - or original(self, local, plan), + lambda self, local, plan=None: ( + fast_calls.append(local) or original(self, local, plan) + ), ) def _local_halo(tensor, halo_size, strategy, dim_index, is_periodic=False): @@ -463,8 +464,9 @@ def _boom(*args, **kwargs): plain_x = x.detach().clone().requires_grad_(True) nn.Conv3d.forward(conv, plain_x).backward(gy) assert local.grad is not None, "the fallback severed the graph at the halo" - torch.testing.assert_close(local.grad.float(), plain_x.grad.float(), - rtol=2e-2, atol=2e-2) + torch.testing.assert_close( + local.grad.float(), plain_x.grad.float(), rtol=2e-2, atol=2e-2 + ) # --------------------------------------------------------------------------- @@ -1123,8 +1125,9 @@ def _seeded_convT(cin=16, cout=8, kernel_size=2, stride=2, **kwargs): what the four decoder sites have and what puts ``grad_bias`` on the live path. """ - conv = FastConvTranspose3d(cin, cout, kernel_size=kernel_size, stride=stride, - **kwargs) + conv = FastConvTranspose3d( + cin, cout, kernel_size=kernel_size, stride=stride, **kwargs + ) generator = torch.Generator().manual_seed(4321) with torch.no_grad(): conv.weight.normal_(0.0, 0.1, generator=generator) @@ -1142,7 +1145,10 @@ def _stock_like(conv, dtype=torch.bfloat16): """A stock ``nn.ConvTranspose3d`` holding the same parameters.""" cin, cout = int(conv.weight.shape[0]), int(conv.weight.shape[1]) plain = nn.ConvTranspose3d( - cin, cout, kernel_size=conv.kernel_size, stride=conv.stride, + cin, + cout, + kernel_size=conv.kernel_size, + stride=conv.stride, bias=conv.bias is not None, ) plain = plain.cuda().to(dtype).to(memory_format=_CHANNELS_LAST) @@ -1197,9 +1203,10 @@ def test_the_transposed_block_list_is_empty_at_every_decoder_site(): """ for x_shape, w_shape in _UPSAMPLER_SITES: assert conv_mod._transposed_policy_declines(x_shape, w_shape) is False - assert conv_mod._policy_declines( - x_shape, w_shape, (2, 2, 2), (0, 0, 0), (1, 1, 1) - ) is False + assert ( + conv_mod._policy_declines(x_shape, w_shape, (2, 2, 2), (0, 0, 0), (1, 1, 1)) + is False + ) @pytest.mark.gpu @@ -1229,26 +1236,28 @@ def test_the_transposed_gate_answers_for_the_four_sites_and_refuses_the_rest(): ) assert ( conv_mod._use_triton_transposed( - _gpu_convT(16, 8, stride=3, output_padding=1), _gpu_input((1, 16, 8, 8, 8)), - None, None, + _gpu_convT(16, 8, stride=3, output_padding=1), + _gpu_input((1, 16, 8, 8, 8)), + None, + None, ) is False ) # groups > 1 has no coverage in any direction. assert ( - conv_mod._use_triton_transposed( - _gpu_convT(16, 8, groups=2), x, None, None - ) + conv_mod._use_triton_transposed(_gpu_convT(16, 8, groups=2), x, None, None) is False ) # NCDHW would be a full-size hidden relayout, which is the cost the rung # exists to avoid. + assert conv_mod._use_triton_transposed(conv, x.contiguous(), None, None) is False + # And the CPU, where there is no kernel at all. assert ( - conv_mod._use_triton_transposed(conv, x.contiguous(), None, None) is False + conv_mod._use_triton_transposed( + _seeded_convT(), torch.randn(1, 16, 4, 4, 4), None, None + ) + is False ) - # And the CPU, where there is no kernel at all. - assert conv_mod._use_triton_transposed(_seeded_convT(), torch.randn(1, 16, 4, 4, 4), - None, None) is False @pytest.mark.gpu @@ -1270,8 +1279,9 @@ def test_the_transposed_gate_asked_is_the_one_that_covers_the_backward(monkeypat # ``is_supported_transposed_all`` makes internally, so this distinguishes # "asked the combined gate" from "asked the forward's and got the same # answer" -- which is the only way to tell them apart while they agree. - monkeypatch.setattr(module, "is_supported_transposed_all", lambda *a, **kw: False, - raising=False) + monkeypatch.setattr( + module, "is_supported_transposed_all", lambda *a, **kw: False, raising=False + ) assert conv_mod._use_triton_transposed(conv, x, None, None) is False @@ -1287,12 +1297,16 @@ def test_a_non_transposed_module_would_be_declined_by_the_transposed_gate(): """ conv = _gpu_convT(16, 16) conv.transposed = False - assert conv_mod._use_triton_transposed(conv, _gpu_input((1, 16, 8, 8, 8)), None, - None) is False + assert ( + conv_mod._use_triton_transposed(conv, _gpu_input((1, 16, 8, 8, 8)), None, None) + is False + ) @pytest.mark.gpu -@pytest.mark.parametrize("cin, cout, spatial", [(16, 8, (8, 8, 8)), (32, 16, (4, 6, 6))]) +@pytest.mark.parametrize( + "cin, cout, spatial", [(16, 8, (8, 8, 8)), (32, 16, (4, 6, 6))] +) def test_transposed_forward_and_gradients_match_nn_convtranspose3d(cin, cout, spatial): """All three directions against an fp64 reference, at MIOpen's own standard. @@ -1345,7 +1359,9 @@ def test_transposed_forward_and_gradients_match_nn_convtranspose3d(cin, cout, sp incumbent_gb = (plain.bias.grad.to(torch.float64) - expected_gb).abs().max().item() actual_gb = (conv.bias.grad.to(torch.float64) - expected_gb).abs().max().item() assert conv.bias.grad.dtype is conv.bias.dtype - assert actual_gb <= max(4.0 * incumbent_gb, 2.0**-8 * expected_gb.abs().max().item()) + assert actual_gb <= max( + 4.0 * incumbent_gb, 2.0**-8 * expected_gb.abs().max().item() + ) @pytest.mark.gpu @@ -1440,19 +1456,21 @@ class _Wrapper(torch.Tensor): pass class _Ctx: - saved_tensors = (_Wrapper(torch.randn(1, 8, 4, 4, 4)), torch.randn(8, 8, 2, 2, 2)) + saved_tensors = ( + _Wrapper(torch.randn(1, 8, 4, 4, 4)), + torch.randn(8, 8, 2, 2, 2), + ) conv_args = ((2, 2, 2), (0, 0, 0), (0, 0, 0), (1, 1, 1), True) needs_input_grad = (True, True, True, False, False, False, False) with pytest.raises(RuntimeError, match="served by different"): - conv_mod._TritonConvTranspose3dFn.backward( - _Ctx(), torch.randn(1, 8, 8, 8, 8) - ) + conv_mod._TritonConvTranspose3dFn.backward(_Ctx(), torch.randn(1, 8, 8, 8, 8)) @pytest.mark.gpu -@pytest.mark.parametrize("direction", ["conv_transpose3d_backward_data", - "conv_transpose3d_backward_weight"]) +@pytest.mark.parametrize( + "direction", ["conv_transpose3d_backward_data", "conv_transpose3d_backward_weight"] +) def test_transposed_backward_falls_back_to_miopen_when_a_direction_fails( monkeypatch, direction ): @@ -1479,7 +1497,9 @@ class _Boom(Exception): module = conv_mod._get_triton_module() monkeypatch.setattr(conv_mod, "_triton_kernel_failures", lambda: (_Boom,)) monkeypatch.setattr( - module, direction, lambda *a, **kw: (_ for _ in ()).throw(_Boom()), + module, + direction, + lambda *a, **kw: (_ for _ in ()).throw(_Boom()), raising=False, ) monkeypatch.setattr(conv_mod, "_triton_failed", False) @@ -1503,14 +1523,20 @@ class _Boom(Exception): ("bwd-weight", conv.weight.grad, plain.weight.grad), ): expected = ref.reference(problem, operands, name) - ref.assert_close(actual, expected, problem, name, - incumbent_error=ref.compare(incumbent, expected)) + ref.assert_close( + actual, + expected, + problem, + name, + incumbent_error=ref.compare(incumbent, expected), + ) # The bias gradient comes from the aten call on this path, not from the sum # above it, so it is the one this test would otherwise never look at. torch.testing.assert_close( conv.bias.grad.to(torch.float64), gy.to(torch.float64).sum(dim=(0, 2, 3, 4)), - rtol=2e-2, atol=2e-2, + rtol=2e-2, + atol=2e-2, ) @@ -1548,13 +1574,20 @@ def plan(num_shards, shard_dim=(2, 3, 4), w=weight, padding=(0, 0, 0)): assert plan((2, 1, 1), shard_dim=(2, 2, 2)) is None assert plan((2, 1), shard_dim=(2, 3, 4)) is None assert plan((2, 1, 1), padding=(1, 1)) is None - assert conv_mod._transposed_halo_plan(None, _StubStrategy((2, 1, 1)), x, weight, - (0, 0, 0)) is None + assert ( + conv_mod._transposed_halo_plan( + None, _StubStrategy((2, 1, 1)), x, weight, (0, 0, 0) + ) + is None + ) periodic = _dc(x, (2, 1, 1)) periodic._is_periodic = (True, False, False) - assert conv_mod._transposed_halo_plan( - periodic, _StubStrategy((2, 1, 1)), x, weight, (0, 0, 0) - ) is None + assert ( + conv_mod._transposed_halo_plan( + periodic, _StubStrategy((2, 1, 1)), x, weight, (0, 0, 0) + ) + is None + ) @pytest.mark.gpu @@ -1612,7 +1645,13 @@ def test_the_transposed_rung_serves_a_dctensor_without_any_halo_exchange(num_sha ("bwd-weight", conv.weight.grad, plain.weight.grad), ): expected = ref.reference(problem, operands, name) - ref.assert_close(actual, expected, problem, name, - incumbent_error=ref.compare(incumbent, expected)) - torch.testing.assert_close(conv.bias.grad.float(), plain.bias.grad.float(), - rtol=2e-2, atol=2e-2) + ref.assert_close( + actual, + expected, + problem, + name, + incumbent_error=ref.compare(incumbent, expected), + ) + torch.testing.assert_close( + conv.bias.grad.float(), plain.bias.grad.float(), rtol=2e-2, atol=2e-2 + ) diff --git a/tests/test_groupnorm_wiring_edge.py b/tests/test_groupnorm_wiring_edge.py index 0d513e4..86a6597 100644 --- a/tests/test_groupnorm_wiring_edge.py +++ b/tests/test_groupnorm_wiring_edge.py @@ -658,9 +658,7 @@ def test_gpu_model_double_backward_fails_loudly(): "the double-backward graph was severed instead of raising; a second " "backward would contribute zero silently" ) - with pytest.raises( - RuntimeError, match="no autograd formula|once_differentiable" - ): + with pytest.raises(RuntimeError, match="no autograd formula|once_differentiable"): first.pow(2).sum().backward() diff --git a/tests/test_platform_guard.py b/tests/test_platform_guard.py index 8c4c048..36e6a13 100644 --- a/tests/test_platform_guard.py +++ b/tests/test_platform_guard.py @@ -171,7 +171,7 @@ class _Props: def test_a_device_that_cannot_be_answered_for_is_not_the_tuned_one(monkeypatch): - """"I could not find out" is not "yes"; the description says which.""" + """ "I could not find out" is not "yes"; the description says which.""" def explode(index): raise RuntimeError("HIP error: no device") @@ -318,9 +318,9 @@ def test_each_ladders_setter_is_the_override_for_that_ladder( def test_the_env_var_opt_in_reaches_the_guard(monkeypatch, module): """``SCAFFOLD_*_TRITON=1`` and the setter are the same statement.""" monkeypatch.setenv(module.TRITON_ENV_VAR, "1") - monkeypatch.setattr(module, "_triton_override", _rungs._env_override( - module.TRITON_ENV_VAR - )) + monkeypatch.setattr( + module, "_triton_override", _rungs._env_override(module.TRITON_ENV_VAR) + ) _fake_device(monkeypatch, _UNTUNED["mi300x"]) assert _rungs._platform_declines(_cuda(), module._triton_override) is False @@ -492,8 +492,9 @@ def test_a_whole_unet_step_asks_the_hardware_once(monkeypatch): asked = [] _fake_device(monkeypatch, _MI300A, count=asked) torch.manual_seed(0) - model = UNet(n_channels=3, n_classes=2, trilinear=False, layers=2, - group_norm_groups=8) + model = UNet( + n_channels=3, n_classes=2, trilinear=False, layers=2, group_norm_groups=8 + ) model = model.cuda().to(memory_format=_CHANNELS_LAST) x = _gpu_input((1, 3, 16, 16, 16), dtype=torch.float32) with torch.no_grad(): diff --git a/triton_conv3d/__init__.py b/triton_conv3d/__init__.py index db3807c..795175a 100644 --- a/triton_conv3d/__init__.py +++ b/triton_conv3d/__init__.py @@ -65,41 +65,54 @@ __version__ = "0.0.0.dev0" -__all__ = ["ConvConfig", "conv3d_forward", "conv3d_backward_data", - "conv3d_backward_weight", "is_supported", "is_supported_all", - "is_supported_bwd_data", "is_supported_bwd_weight", - "conv_transpose3d_forward", "conv_transpose3d_backward_data", - "conv_transpose3d_backward_weight", "is_supported_transposed", - "is_supported_transposed_all", "is_supported_transposed_bwd_data", - "is_supported_transposed_bwd_weight", "__version__"] +__all__ = [ + "ConvConfig", + "conv3d_forward", + "conv3d_backward_data", + "conv3d_backward_weight", + "is_supported", + "is_supported_all", + "is_supported_bwd_data", + "is_supported_bwd_weight", + "conv_transpose3d_forward", + "conv_transpose3d_backward_data", + "conv_transpose3d_backward_weight", + "is_supported_transposed", + "is_supported_transposed_all", + "is_supported_transposed_bwd_data", + "is_supported_transposed_bwd_weight", + "__version__", +] #: The public names live in :mod:`triton_conv3d.gather_gemm`, which imports #: torch and triton. They are re-exported lazily so that #: ``import triton_conv3d.shapes`` stays free of both: the shape and cost model #: is pure Python by design, and it drives test parametrization at collection #: time on machines that have no GPU and may have no triton. -_LAZY = {"ConvConfig": "gather_gemm", "conv3d_forward": "gather_gemm", - "is_supported": "gather_gemm", "is_supported_all": "gather_gemm", - "conv3d_backward_data": "bwd_data", - "is_supported_bwd_data": "bwd_data", - "conv3d_backward_weight": "reduce_gemm", - "is_supported_bwd_weight": "reduce_gemm", - "conv_transpose3d_forward": "transposed", - "conv_transpose3d_backward_data": "transposed", - "conv_transpose3d_backward_weight": "transposed", - "is_supported_transposed": "transposed", - "is_supported_transposed_all": "transposed", - "is_supported_transposed_bwd_data": "transposed", - "is_supported_transposed_bwd_weight": "transposed"} +_LAZY = { + "ConvConfig": "gather_gemm", + "conv3d_forward": "gather_gemm", + "is_supported": "gather_gemm", + "is_supported_all": "gather_gemm", + "conv3d_backward_data": "bwd_data", + "is_supported_bwd_data": "bwd_data", + "conv3d_backward_weight": "reduce_gemm", + "is_supported_bwd_weight": "reduce_gemm", + "conv_transpose3d_forward": "transposed", + "conv_transpose3d_backward_data": "transposed", + "conv_transpose3d_backward_weight": "transposed", + "is_supported_transposed": "transposed", + "is_supported_transposed_all": "transposed", + "is_supported_transposed_bwd_data": "transposed", + "is_supported_transposed_bwd_weight": "transposed", +} def __getattr__(name: str) -> Any: if name in _LAZY: import importlib - return getattr( - importlib.import_module(f".{_LAZY[name]}", __name__), name - ) + return getattr(importlib.import_module(f".{_LAZY[name]}", __name__), name) raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/triton_conv3d/bench/__init__.py b/triton_conv3d/bench/__init__.py index e230d8e..9cca128 100644 --- a/triton_conv3d/bench/__init__.py +++ b/triton_conv3d/bench/__init__.py @@ -10,5 +10,11 @@ time_callable, ) -__all__ = ["Measurement", "Ratio", "flush_caches", "interleaved", "ratio", - "time_callable"] +__all__ = [ + "Measurement", + "Ratio", + "flush_caches", + "interleaved", + "ratio", + "time_callable", +] diff --git a/triton_conv3d/bench/baseline.py b/triton_conv3d/bench/baseline.py index 5610851..d7b94ce 100644 --- a/triton_conv3d/bench/baseline.py +++ b/triton_conv3d/bench/baseline.py @@ -148,7 +148,9 @@ def _callable(problem: ConvProblem, direction: Direction, device: str, dtype): x, w, b = _build(problem, device, dtype) op = F.conv_transpose3d if problem.transposed else F.conv3d - fwd = lambda: op(x, w, b, stride=problem.stride, padding=problem.padding) + + def fwd(): + return op(x, w, b, stride=problem.stride, padding=problem.padding) if direction == "fwd": with torch.no_grad(): @@ -157,7 +159,10 @@ def _callable(problem: ConvProblem, direction: Direction, device: str, dtype): y = fwd() gy = torch.randn_like(y) inputs = (x,) if direction == "bwd-data" else (w,) - fn = lambda: torch.autograd.grad(y, inputs, gy, retain_graph=True) + + def fn(): + return torch.autograd.grad(y, inputs, gy, retain_graph=True) + return fn, (x, w, y, gy) @@ -177,8 +182,9 @@ def measure_one( "cudnn.benchmark is off; MIOpen will answer from its heuristic " "instead of searching and the result is not a baseline" ) - dtype = {"bf16": torch.bfloat16, "fp32": torch.float32, - "fp16": torch.float16}[problem.dtype] + dtype = {"bf16": torch.bfloat16, "fp32": torch.float32, "fp16": torch.float16}[ + problem.dtype + ] record: dict = { "problem": problem.label, "direction": direction, @@ -216,13 +222,25 @@ def measure_one( rough_ms = (time.perf_counter() - t0) * 1e3 if rough_ms > max_call_ms: - record.update(ms=rough_ms, best_ms=rough_ms, iters=1, rounds=1, - spread=0.0, rel_ci=float("inf"), stop="single-call", - note="single call; exceeds max_call_ms") + record.update( + ms=rough_ms, + best_ms=rough_ms, + iters=1, + rounds=1, + spread=0.0, + rel_ci=float("inf"), + stop="single-call", + note="single call; exceeds max_call_ms", + ) else: - meas = interleaved({"miopen": fn}, warmup=None, iters=None, - rounds=None, budget_s=budget_s, - target_rel=target_rel)["miopen"] + meas = interleaved( + {"miopen": fn}, + warmup=None, + iters=None, + rounds=None, + budget_s=budget_s, + target_rel=target_rel, + )["miopen"] # Both statistics, because they answer different questions. The # median is the control -- it is what a step actually costs on a # shared node. The best round is the diagnostic: this node has @@ -234,11 +252,18 @@ def measure_one( # ``rel_ci`` instead: ``spread`` is a *range* and its expectation # grows with ``rounds``, which is now chosen per cell, so two cells' # spreads are no longer comparable to each other at all. - record.update(ms=meas.median, best_ms=meas.best, iters=meas.iters, - rounds=len(meas.rounds), spread=meas.spread, - rel_ci=meas.rel_half_width, stop=meas.stop, - group=meas.group, tax_frac=meas.tax_frac, - measure_seconds=meas.seconds) + record.update( + ms=meas.median, + best_ms=meas.best, + iters=meas.iters, + rounds=len(meas.rounds), + spread=meas.spread, + rel_ci=meas.rel_half_width, + stop=meas.stop, + group=meas.group, + tax_frac=meas.tax_frac, + measure_seconds=meas.seconds, + ) record["tflops"] = record["flops"] / (record["ms"] * 1e-3) / 1e12 record["pct_roofline"] = 100 * record["tflops"] / record["roofline_tflops"] except Exception as exc: @@ -283,15 +308,17 @@ def cross_check(records: list[dict], problems: list[ConvProblem]) -> list[dict]: hit = by_key.get((r["problem"], r["direction"])) if hit is None: continue - rows.append({ - "problem": r["problem"], - "direction": r["direction"], - "isolated_ms": r["ms"], - "profiled_ms": hit["ms_per_call"], - "ratio": r["ms"] / hit["ms_per_call"], - "config": hit["config"], - "profiled_solvers": hit.get("solvers", []), - }) + rows.append( + { + "problem": r["problem"], + "direction": r["direction"], + "isolated_ms": r["ms"], + "profiled_ms": hit["ms_per_call"], + "ratio": r["ms"] / hit["ms_per_call"], + "config": hit["config"], + "profiled_solvers": hit.get("solvers", []), + } + ) return rows @@ -299,8 +326,11 @@ def format_cross_check(rows: list[dict], tol: float = 0.25) -> str: table = format_table( [ [ - r["problem"], r["direction"], f"{r['isolated_ms']:.4f}", - f"{r['profiled_ms']:.4f}", f"{r['ratio']:.2f}x", + r["problem"], + r["direction"], + f"{r['isolated_ms']:.4f}", + f"{r['profiled_ms']:.4f}", + f"{r['ratio']:.2f}x", "ok" if abs(r["ratio"] - 1) <= tol else "MISMATCH", ] for r in rows @@ -316,42 +346,79 @@ def format_cross_check(rows: list[dict], tol: float = 0.25) -> str: def main() -> None: - ap = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) ap.add_argument("--out", default="baseline.json") - ap.add_argument("--resume", action="store_true", - help="skip cells already present in --out") - ap.add_argument("--top", type=int, default=0, - help="only the N hottest corpus problems (0 = all)") - ap.add_argument("--skip", type=int, default=0, - help="skip the first N corpus problems (they are ordered by " - "cost, and the most expensive one takes 45 s per call)") - ap.add_argument("--include-edge", action="store_true", - help="also measure the synthetic edge cases") - ap.add_argument("--shape", - choices=("halo", "production", "unhaloed", "both", "all"), - default="halo", - help="halo: the form upstream DistConv issues (default, and " - "the form the profiled numbers measured, so the only " - "one --cross-check can join); production: the form " - "ScaFFold's own Triton adapter issues, padded on every " - "unsplit axis -- the right MIOpen baseline for a " - "Triton comparison; unhaloed: the logical statement; " - "both: halo+unhaloed, as before; all: every distinct " - "form") - ap.add_argument("--max-call-ms", type=float, default=60_000.0, - help="above this, record a single call rather than a sweep") - ap.add_argument("--budget", type=float, default=10.0, - help="wall-clock seconds per cell (default 10)") - ap.add_argument("--precision", type=float, default=0.02, - help="target relative 95%% half-width per cell") - ap.add_argument("--skip-slow", action="store_true", - help="skip cells a previous run recorded as slower than " - "--max-call-ms; useful for a quick re-capture") - ap.add_argument("--cross-check", action="store_true", - help="join the halo'd cells onto the profiled numbers") - ap.add_argument("--tolerance", type=float, default=0.25, - help="cross-check band, as a fraction of the profiled time") + ap.add_argument( + "--resume", action="store_true", help="skip cells already present in --out" + ) + ap.add_argument( + "--top", + type=int, + default=0, + help="only the N hottest corpus problems (0 = all)", + ) + ap.add_argument( + "--skip", + type=int, + default=0, + help="skip the first N corpus problems (they are ordered by " + "cost, and the most expensive one takes 45 s per call)", + ) + ap.add_argument( + "--include-edge", + action="store_true", + help="also measure the synthetic edge cases", + ) + ap.add_argument( + "--shape", + choices=("halo", "production", "unhaloed", "both", "all"), + default="halo", + help="halo: the form upstream DistConv issues (default, and " + "the form the profiled numbers measured, so the only " + "one --cross-check can join); production: the form " + "ScaFFold's own Triton adapter issues, padded on every " + "unsplit axis -- the right MIOpen baseline for a " + "Triton comparison; unhaloed: the logical statement; " + "both: halo+unhaloed, as before; all: every distinct " + "form", + ) + ap.add_argument( + "--max-call-ms", + type=float, + default=60_000.0, + help="above this, record a single call rather than a sweep", + ) + ap.add_argument( + "--budget", + type=float, + default=10.0, + help="wall-clock seconds per cell (default 10)", + ) + ap.add_argument( + "--precision", + type=float, + default=0.02, + help="target relative 95%% half-width per cell", + ) + ap.add_argument( + "--skip-slow", + action="store_true", + help="skip cells a previous run recorded as slower than " + "--max-call-ms; useful for a quick re-capture", + ) + ap.add_argument( + "--cross-check", + action="store_true", + help="join the halo'd cells onto the profiled numbers", + ) + ap.add_argument( + "--tolerance", + type=float, + default=0.25, + help="cross-check band, as a fraction of the profiled time", + ) args = ap.parse_args() if not torch.cuda.is_available(): @@ -363,23 +430,28 @@ def main() -> None: "ScaFFold runs" ) if not os.environ.get("MIOPEN_USER_DB_PATH"): - print("WARNING: MIOPEN_USER_DB_PATH is unset -- every cell re-runs the " - "find from scratch and nothing is reusable afterwards", - file=sys.stderr) + print( + "WARNING: MIOPEN_USER_DB_PATH is unset -- every cell re-runs the " + "find from scratch and nothing is reusable afterwards", + file=sys.stderr, + ) - problems = list(scaffold_corpus())[args.skip:] + problems = list(scaffold_corpus())[args.skip :] if args.top: problems = problems[: args.top] corpus_problems = list(problems) if args.include_edge: problems += list(edge_cases()) - modes = {"both": ("halo", "unhaloed"), - "all": ("halo", "production", "unhaloed")}.get( - args.shape, (args.shape,)) - _forms = {"halo": lambda p: p.halo_variant, - "production": lambda p: p.production_variant, - "unhaloed": lambda p: p} + modes = { + "both": ("halo", "unhaloed"), + "all": ("halo", "production", "unhaloed"), + }.get(args.shape, (args.shape,)) + _forms = { + "halo": lambda p: p.halo_variant, + "production": lambda p: p.production_variant, + "unhaloed": lambda p: p, + } # A problem with nothing to halo is its own variant in all three forms, so # the multi-mode runs would otherwise measure the synthetic edge cases and # the k=1/transposed convs two or three times for nothing. Deduplicated on @@ -430,14 +502,21 @@ def flush() -> None: continue if args.skip_slow and done.get(key, {}).get("note"): continue - rec = measure_one(problem, direction, max_call_ms=args.max_call_ms, - shape_mode=mode, budget_s=args.budget, - target_rel=args.precision) + rec = measure_one( + problem, + direction, + max_call_ms=args.max_call_ms, + shape_mode=mode, + budget_s=args.budget, + target_rel=args.precision, + ) records.append(rec) flush() if "error" in rec: - print(f" {problem.label:36s} {mode:8s} {direction:11s} " - f"ERROR {rec['error'][:80]}") + print( + f" {problem.label:36s} {mode:8s} {direction:11s} " + f"ERROR {rec['error'][:80]}" + ) else: print( f" {problem.label:36s} {mode:8s} {direction:11s} " @@ -454,11 +533,16 @@ def flush() -> None: print(f"\n{len(ok)}/{len(records)} cells measured -> {out_path}") if len(ok) < len(records): print("failures:") - print(format_table( - [[r["problem"], r["direction"], r["error"][:70]] - for r in records if "error" in r], - ["problem", "direction", "error"], - )) + print( + format_table( + [ + [r["problem"], r["direction"], r["error"][:70]] + for r in records + if "error" in r + ], + ["problem", "direction", "error"], + ) + ) if args.cross_check: rows = cross_check(records, corpus_problems) print("\ncross-check against the profiled ScaFFold runs:") diff --git a/triton_conv3d/bench/conv_bench.py b/triton_conv3d/bench/conv_bench.py index 40c3f91..5a5df07 100644 --- a/triton_conv3d/bench/conv_bench.py +++ b/triton_conv3d/bench/conv_bench.py @@ -131,8 +131,8 @@ select_config, ) from ..reduce_gemm import ( - candidate_bwd_weight_configs, bwd_weight_config, + candidate_bwd_weight_configs, conv3d_backward_weight, grad_weight_empty, split_count, @@ -297,15 +297,18 @@ def _autograd_control(build_forward: Callable[[], tuple]): # -- conv: the ordinary convolution ----------------------------------------- -def _conv_fwd(problem: ConvProblem, device: str, - control: bool = True) -> _Case: +def _conv_fwd(problem: ConvProblem, device: str, control: bool = True) -> _Case: dtype = _TORCH_DTYPE[problem.dtype] k = tuple(problem.kernel) x = _randn(problem.input_shape, device, dtype) w = _randn(problem.weight_shape, device, dtype) b = _bias(problem, device, dtype) - y = torch.empty(problem.output_shape, device=device, dtype=dtype, - memory_format=torch.channels_last_3d) + y = torch.empty( + problem.output_shape, + device=device, + dtype=dtype, + memory_format=torch.channels_last_3d, + ) m = problem.n * math.prod(problem.out_spatial) def triton(cfg): @@ -314,8 +317,8 @@ def run(): # channels-last, which is what a ScaFFold parameter is, and the # kernel reads that layout in place. Passing ``weight_rsck=`` here # would time a path the integration no longer takes. - conv3d_forward(x, w, b, problem.stride, problem.padding, - config=cfg, out=y) + conv3d_forward(x, w, b, problem.stride, problem.padding, config=cfg, out=y) + return run def miopen(): @@ -329,30 +332,45 @@ def reference(): return got, ref return _Case( - triton=triton, miopen=miopen if control else None, transform=None, + triton=triton, + miopen=miopen if control else None, + transform=None, candidates=lambda: candidate_configs(m, problem.cin, problem.cout, dtype), refine=_gather_refine, shipped_config=lambda: select_config(m, problem.cin, problem.cout, k, dtype), - reference=reference if control else None, primary=x, keep=(x, w, b, y), + reference=reference if control else None, + primary=x, + keep=(x, w, b, y), ) -def _conv_bwd_data(problem: ConvProblem, device: str, - control: bool = True) -> _Case: +def _conv_bwd_data(problem: ConvProblem, device: str, control: bool = True) -> _Case: dtype = _TORCH_DTYPE[problem.dtype] k = tuple(problem.kernel) w = _randn(problem.weight_shape, device, dtype) b = _bias(problem, device, dtype) gy = _randn(problem.output_shape, device, dtype) - gx = torch.empty(problem.input_shape, device=device, dtype=dtype, - memory_format=torch.channels_last_3d) + gx = torch.empty( + problem.input_shape, + device=device, + dtype=dtype, + memory_format=torch.channels_last_3d, + ) def triton(cfg): def run(): # As in the forward: the channels-last parameter is the operand, and # the tap flip and the transpose are both constexprs in the kernel. - conv3d_backward_data(gy, w, problem.input_shape, problem.stride, - problem.padding, config=cfg, out=gx) + conv3d_backward_data( + gy, + w, + problem.input_shape, + problem.stride, + problem.padding, + config=cfg, + out=gx, + ) + return run def build(): @@ -372,30 +390,33 @@ def miopen(): torch.autograd.grad(yg, (xg,), gy, retain_graph=True) def reference(): - got = conv3d_backward_data(gy, w, problem.input_shape, problem.stride, - problem.padding) + got = conv3d_backward_data( + gy, w, problem.input_shape, problem.stride, problem.padding + ) ref = torch.autograd.grad(yg, (xg,), gy, retain_graph=True)[0] return got, ref return _Case( - triton=triton, miopen=miopen if control else None, transform=None, + triton=triton, + miopen=miopen if control else None, + transform=None, # Swapped: backward-data reduces over Cout and its GEMM's N is Cin. candidates=lambda: candidate_configs( - problem.n * math.prod(problem.spatial), problem.cout, problem.cin, - dtype), + problem.n * math.prod(problem.spatial), problem.cout, problem.cin, dtype + ), refine=_gather_refine, shipped_config=lambda: bwd_data_config( - problem.output_shape, problem.cin, k, dtype, - padding=problem.padding), - reference=reference if control else None, primary=gy, + problem.output_shape, problem.cin, k, dtype, padding=problem.padding + ), + reference=reference if control else None, + primary=gy, # ``xg``/``yg`` are kept because the graph (and so the control) dies # with them. keep=(gy, w, b, gx, xg, yg), ) -def _conv_bwd_weight(problem: ConvProblem, device: str, - control: bool = True) -> _Case: +def _conv_bwd_weight(problem: ConvProblem, device: str, control: bool = True) -> _Case: dtype = _TORCH_DTYPE[problem.dtype] k = tuple(problem.kernel) x = _randn(problem.input_shape, device, dtype) @@ -408,22 +429,42 @@ def _conv_bwd_weight(problem: ConvProblem, device: str, def candidates(): return candidate_bwd_weight_configs( - problem.cout, problem.cin, k, k_total, dtype, - splits=(0,), padded=padded, + problem.cout, + problem.cin, + k, + k_total, + dtype, + splits=(0,), + padded=padded, ) def splits_for(cfg): - return split_count(cfg, problem.cout, problem.cin, problem.tap_count, - k_total, problem.out_spatial[2])[0] - - ws = _sweep_workspace(candidates(), splits_for, problem.cout, problem.cin, - k, device) + return split_count( + cfg, + problem.cout, + problem.cin, + problem.tap_count, + k_total, + problem.out_spatial[2], + )[0] + + ws = _sweep_workspace( + candidates(), splits_for, problem.cout, problem.cin, k, device + ) def triton(cfg): def run(): - conv3d_backward_weight(x, problem.weight_shape, gy, problem.stride, - problem.padding, config=cfg, workspace=ws, - out=gw) + conv3d_backward_weight( + x, + problem.weight_shape, + gy, + problem.stride, + problem.padding, + config=cfg, + workspace=ws, + out=gw, + ) + return run def build(): @@ -438,17 +479,23 @@ def miopen(): torch.autograd.grad(yg, (wg,), gy, retain_graph=True) def reference(): - got = conv3d_backward_weight(x, problem.weight_shape, gy, - problem.stride, problem.padding) + got = conv3d_backward_weight( + x, problem.weight_shape, gy, problem.stride, problem.padding + ) ref = torch.autograd.grad(yg, (wg,), gy, retain_graph=True)[0] return got, ref return _Case( - triton=triton, miopen=miopen if control else None, transform=None, - candidates=candidates, refine=_bwd_weight_refine, + triton=triton, + miopen=miopen if control else None, + transform=None, + candidates=candidates, + refine=_bwd_weight_refine, shipped_config=lambda: bwd_weight_config( - problem.cout, problem.cin, k, k_total, dtype, padded=padded), - reference=reference if control else None, primary=x, + problem.cout, problem.cin, k, k_total, dtype, padded=padded + ), + reference=reference if control else None, + primary=x, keep=(x, w, b, gy, gw, ws, wg, yg), ) @@ -460,20 +507,24 @@ def reference(): # the call, rather than a shared helper naming it three times differently. -def _convt_fwd(problem: ConvProblem, device: str, - control: bool = True) -> _Case: +def _convt_fwd(problem: ConvProblem, device: str, control: bool = True) -> _Case: dtype = _TORCH_DTYPE[problem.dtype] k = tuple(problem.kernel) x = _randn(problem.input_shape, device, dtype) w = _randn(problem.weight_shape, device, dtype) b = _bias(problem, device, dtype) - y = torch.empty(problem.output_shape, device=device, dtype=dtype, - memory_format=torch.channels_last_3d) + y = torch.empty( + problem.output_shape, + device=device, + dtype=dtype, + memory_format=torch.channels_last_3d, + ) m = problem.n * math.prod(problem.spatial) def triton(cfg): def run(): conv_transpose3d_forward(x, w, b, k, config=cfg, out=y) + return run def miopen(): @@ -487,18 +538,23 @@ def reference(): return got, ref return _Case( - triton=triton, miopen=miopen if control else None, transform=None, + triton=triton, + miopen=miopen if control else None, + transform=None, candidates=lambda: candidate_transposed_configs( - m, problem.cin, problem.cout, problem.tap_count, dtype), + m, problem.cin, problem.cout, problem.tap_count, dtype + ), refine=_gather_refine, - shipped_config=lambda: transposed_config(m, problem.cin, problem.cout, - k, dtype), - reference=reference if control else None, primary=x, keep=(x, w, b, y), + shipped_config=lambda: transposed_config( + m, problem.cin, problem.cout, k, dtype + ), + reference=reference if control else None, + primary=x, + keep=(x, w, b, y), ) -def _convt_bwd_data(problem: ConvProblem, device: str, - control: bool = True) -> _Case: +def _convt_bwd_data(problem: ConvProblem, device: str, control: bool = True) -> _Case: """``grad_input = conv3d(grad_output, w, stride=k)`` -- an ordinary forward. So the config that runs is :func:`~triton_conv3d.gather_gemm.select_config`'s @@ -513,16 +569,22 @@ def _convt_bwd_data(problem: ConvProblem, device: str, w = _randn(problem.weight_shape, device, dtype) b = _bias(problem, device, dtype) gy = _randn(problem.output_shape, device, dtype) - gx = torch.empty(problem.input_shape, device=device, dtype=dtype, - memory_format=torch.channels_last_3d) + gx = torch.empty( + problem.input_shape, + device=device, + dtype=dtype, + memory_format=torch.channels_last_3d, + ) # The strided convolution's M is this operator's *input* volume, and its # (cin, cout) are (Cout, Cin) of the transposed operator. m = problem.n * math.prod(problem.spatial) def triton(cfg): def run(): - conv_transpose3d_backward_data(gy, w, problem.input_shape, k, - config=cfg, out=gx) + conv_transpose3d_backward_data( + gy, w, problem.input_shape, k, config=cfg, out=gx + ) + return run def build(): @@ -542,17 +604,19 @@ def reference(): return got, ref return _Case( - triton=triton, miopen=miopen if control else None, transform=None, + triton=triton, + miopen=miopen if control else None, + transform=None, candidates=lambda: candidate_configs(m, problem.cout, problem.cin, dtype), refine=_gather_refine, shipped_config=lambda: select_config(m, problem.cout, problem.cin, k, dtype), - reference=reference if control else None, primary=gy, + reference=reference if control else None, + primary=gy, keep=(x, w, b, gy, gx, xg, yg), ) -def _convt_bwd_weight(problem: ConvProblem, device: str, - control: bool = True) -> _Case: +def _convt_bwd_weight(problem: ConvProblem, device: str, control: bool = True) -> _Case: """The same reduction ``conv3d_backward_weight`` performs, operands swapped. ``grad_output`` is the strided convolution's input and ``x`` is its output @@ -567,27 +631,42 @@ def _convt_bwd_weight(problem: ConvProblem, device: str, w = _randn(problem.weight_shape, device, dtype) b = _bias(problem, device, dtype) gy = _randn(problem.output_shape, device, dtype) - gw = grad_transposed_weight_empty(problem.cin, problem.cout, k, - dtype=dtype, device=device) + gw = grad_transposed_weight_empty( + problem.cin, problem.cout, k, dtype=dtype, device=device + ) k_total = problem.n * math.prod(problem.spatial) def candidates(): return candidate_bwd_weight_configs( - problem.cin, problem.cout, k, k_total, dtype, - splits=(0,), padded=False, + problem.cin, + problem.cout, + k, + k_total, + dtype, + splits=(0,), + padded=False, ) def splits_for(cfg): - return split_count(cfg, problem.cin, problem.cout, problem.tap_count, - k_total, problem.spatial[2])[0] - - ws = _sweep_workspace(candidates(), splits_for, problem.cin, problem.cout, - k, device) + return split_count( + cfg, + problem.cin, + problem.cout, + problem.tap_count, + k_total, + problem.spatial[2], + )[0] + + ws = _sweep_workspace( + candidates(), splits_for, problem.cin, problem.cout, k, device + ) def triton(cfg): def run(): - conv_transpose3d_backward_weight(x, problem.weight_shape, gy, k, - config=cfg, workspace=ws, out=gw) + conv_transpose3d_backward_weight( + x, problem.weight_shape, gy, k, config=cfg, workspace=ws, out=gw + ) + return run def build(): @@ -607,11 +686,16 @@ def reference(): return got, ref return _Case( - triton=triton, miopen=miopen if control else None, transform=None, - candidates=candidates, refine=_bwd_weight_refine, + triton=triton, + miopen=miopen if control else None, + transform=None, + candidates=candidates, + refine=_bwd_weight_refine, shipped_config=lambda: bwd_weight_config( - problem.cin, problem.cout, k, k_total, dtype, padded=False), - reference=reference if control else None, primary=x, + problem.cin, problem.cout, k, k_total, dtype, padded=False + ), + reference=reference if control else None, + primary=x, keep=(x, w, b, gy, gw, ws, wg, yg), ) @@ -624,11 +708,15 @@ def _sweep_workspace(cands, splits_for, cout, cin, k, device) -> torch.Tensor: """ max_splits = max( [splits_for(c) for c in cands] - + [splits_for(dataclasses.replace(c, SPLIT_K=sk)) - for c in cands for sk in _SPLIT_SWEEP] + + [ + splits_for(dataclasses.replace(c, SPLIT_K=sk)) + for c in cands + for sk in _SPLIT_SWEEP + ] + ) + return torch.empty( + workspace_elements(max_splits, cout, cin, k), dtype=torch.float32, device=device ) - return torch.empty(workspace_elements(max_splits, cout, cin, k), - dtype=torch.float32, device=device) # --------------------------------------------------------------------------- @@ -677,9 +765,9 @@ class _Op: _FORM_NOTES = { "distconv": "DistConv's halo'd, unpadded form -- what the MIOpen baseline " - "was profiled in", + "was profiled in", "adapter": "the form ScaFFold's Triton rung is handed -- what production " - "runs, padded at every configuration", + "runs, padded at every configuration", "logical": "the module's own statement, unhalo'd and padded", } @@ -692,8 +780,11 @@ class _Op: # Corpus order is measured-cost order, and it is what every stored # capture's ``--problems`` indices refer to. order=lambda p: 0, - build={"fwd": _conv_fwd, "bwd-data": _conv_bwd_data, - "bwd-weight": _conv_bwd_weight}, + build={ + "fwd": _conv_fwd, + "bwd-data": _conv_bwd_data, + "bwd-weight": _conv_bwd_weight, + }, ), "convT": _Op( name="convT", @@ -703,8 +794,11 @@ class _Op: # Cheapest first, as ``m5_convT_bench`` ran them, so a re-capture lines # up row for row with ``m5_shipped_*.json``. order=lambda p: math.prod(p.spatial) * p.cin, - build={"fwd": _convt_fwd, "bwd-data": _convt_bwd_data, - "bwd-weight": _convt_bwd_weight}, + build={ + "fwd": _convt_fwd, + "bwd-data": _convt_bwd_data, + "bwd-weight": _convt_bwd_weight, + }, ), } @@ -714,8 +808,13 @@ def operator_of(problem: ConvProblem) -> Operator: return "convT" if problem.transposed else "conv" -def _build(problem: ConvProblem, direction: Direction, device: str = "cuda", - operator: Operator | None = None, control: bool = True) -> _Case: +def _build( + problem: ConvProblem, + direction: Direction, + device: str = "cuda", + operator: Operator | None = None, + control: bool = True, +) -> _Case: """Operands and launchers for one problem in one direction. A lookup into :data:`_OPERATORS`, not a switch: the operator is resolved @@ -772,8 +871,9 @@ def excludes_launcher(self) -> bool: return self.kind == "kernel" -def _timed_region(variants: Mapping[str, Callable[[], object]], - launcher: str = "exclude") -> _Region: +def _timed_region( + variants: Mapping[str, Callable[[], object]], launcher: str = "exclude" +) -> _Region: """Decide, for one cell, what the published number will contain. **Excluded** from a ``kind="kernel"`` measurement, on **every** arm: Python @@ -820,17 +920,27 @@ def _timed_region(variants: Mapping[str, Callable[[], object]], """ names = list(variants) if launcher == "include": - return _Region("call", 1, variants, {}, - "Python call: kernel + dispatch + config lookup + launcher") + return _Region( + "call", + 1, + variants, + {}, + "Python call: kernel + dispatch + config lookup + launcher", + ) if launcher != "exclude": raise ValueError(f"unknown launcher policy {launcher!r}") eager = {n: per_call_ms(variants[n]) for n in names} durations = [eager[n] for n in names] if not graph_is_worthwhile(durations): - return _Region("call", 1, variants, eager, - f"Python call: every arm is above {min(durations):.1f} ms, " - "where the measured host cost (<=0.08 ms) is under 0.2%") + return _Region( + "call", + 1, + variants, + eager, + f"Python call: every arm is above {min(durations):.1f} ms, " + "where the measured host cost (<=0.08 ms) is under 0.2%", + ) chunk = common_chunk(durations) captured: dict[str, Callable[[], object]] = {} try: @@ -840,12 +950,22 @@ def _timed_region(variants: Mapping[str, Callable[[], object]], captured.clear() torch.cuda.synchronize() torch.cuda.empty_cache() - return _Region("call", 1, variants, eager, - f"Python call: {n!r} could not be captured ({exc}), so no " - "arm was -- a mixed measurement is worth up to 1.4x") - return _Region("kernel", chunk, captured, eager, - f"CUDA graph replay, {chunk} calls per graph: kernels only, " - "no dispatch, no config lookup, no launcher, on every arm") + return _Region( + "call", + 1, + variants, + eager, + f"Python call: {n!r} could not be captured ({exc}), so no " + "arm was -- a mixed measurement is worth up to 1.4x", + ) + return _Region( + "kernel", + chunk, + captured, + eager, + f"CUDA graph replay, {chunk} calls per graph: kernels only, " + "no dispatch, no config lookup, no launcher, on every arm", + ) # --------------------------------------------------------------------------- @@ -853,8 +973,9 @@ def _timed_region(variants: Mapping[str, Callable[[], object]], # --------------------------------------------------------------------------- -def _sweep(case: _Case, configs: list[ConvConfig], - verbose: bool = False) -> list[tuple[ConvConfig, float]]: +def _sweep( + case: _Case, configs: list[ConvConfig], verbose: bool = False +) -> list[tuple[ConvConfig, float]]: """Time every config once, cheaply. Failures (LDS overflow, OOM) are skipped. Eager, deliberately: this pass only has to *rank*, it runs hundreds of @@ -880,20 +1001,36 @@ def _sweep(case: _Case, configs: list[ConvConfig], # whichever config is fastest -- exactly the config the sweep is trying # to find. Sizing the block by time makes the restart the same fraction # for every candidate. - meas = interleaved({"t": run}, warmup=None, iters=None, rounds=3, - warmup_s=0.02, warmup_min=2, block_ms=5.0, - measure_tax=False)["t"] + meas = interleaved( + {"t": run}, + warmup=None, + iters=None, + rounds=3, + warmup_s=0.02, + warmup_min=2, + block_ms=5.0, + measure_tax=False, + )["t"] ranked.append((cfg, meas.median)) ranked.sort(key=lambda kv: kv[1]) return ranked -def measure_problem(problem: ConvProblem, *, direction: Direction = "fwd", - operator: Operator | None = None, - max_configs: int = 0, iters: int = 0, rounds: int = 0, - shipped: bool = False, verbose: bool = False, - budget_s: float = 20.0, target_rel: float = 0.02, - launcher: str = "exclude", control: str = "miopen") -> dict: +def measure_problem( + problem: ConvProblem, + *, + direction: Direction = "fwd", + operator: Operator | None = None, + max_configs: int = 0, + iters: int = 0, + rounds: int = 0, + shipped: bool = False, + verbose: bool = False, + budget_s: float = 20.0, + target_rel: float = 0.02, + launcher: str = "exclude", + control: str = "miopen", +) -> dict: """Sweep, then race the finalists against MIOpen in one interleaved block. ``control="none"`` drops the MIOpen arm and measures the Triton kernels @@ -930,7 +1067,8 @@ def measure_problem(problem: ConvProblem, *, direction: Direction = "fwd", "problem": problem.label, "operator": op, "direction": direction, - "cin": problem.cin, "cout": problem.cout, + "cin": problem.cin, + "cout": problem.cout, "spatial": list(problem.spatial), "kernel": list(problem.kernel), "padding": list(problem.padding), @@ -945,8 +1083,7 @@ def measure_problem(problem: ConvProblem, *, direction: Direction = "fwd", case = None region = None try: - case = _build(problem, direction, operator=op, - control=(control == "miopen")) + case = _build(problem, direction, operator=op, control=(control == "miopen")) # ``UntypedStorage.size()`` is already in *bytes*, which is what the # specializer's ``is_within_2gb`` compares -- multiplying by the # element size again would report every operand as ineligible. @@ -993,10 +1130,14 @@ def measure_problem(problem: ConvProblem, *, direction: Direction = "fwd", # per call, on that arm only. See :class:`on_capture_stream`. with on_capture_stream(): region = _timed_region(variants, launcher) - meas = interleaved(region.fns, - warmup=3 if pinned else None, - iters=iters or None, rounds=rounds or None, - budget_s=budget_s, target_rel=target_rel) + meas = interleaved( + region.fns, + warmup=3 if pinned else None, + iters=iters or None, + rounds=rounds or None, + budget_s=budget_s, + target_rel=target_rel, + ) # ``region.chunk`` calls sit behind one replay, so every *absolute* time # is that many times too large. Every *relative* one -- the half-widths, # the paired ratio, the convergence test the harness already ran -- is @@ -1031,8 +1172,9 @@ def measure_problem(problem: ConvProblem, *, direction: Direction = "fwd", measure_stop=best.stop, measure_balanced=best.balanced, measure_seconds=best.seconds, - rsck_ms=(meas["rsck_transform"].median / c - if "rsck_transform" in meas else 0.0), + rsck_ms=( + meas["rsck_transform"].median / c if "rsck_transform" in meas else 0.0 + ), ) # Only when there *is* a control. An absent MIOpen number is left # absent rather than written as zero: every consumer of these rows @@ -1056,8 +1198,7 @@ def measure_problem(problem: ConvProblem, *, direction: Direction = "fwd", miopen_rel_ci=mio.rel_half_width, miopen_tax_frac=mio.tax_frac, miopen_eager_ms=region.eager_ms.get("miopen", 0.0), - miopen_pct_roofline=100 * problem.efficiency(mio.median / c, - direction), + miopen_pct_roofline=100 * problem.efficiency(mio.median / c, direction), miopen_tflops=problem.flops(direction) / (mio.median / c * 1e-3) / 1e12, speedup=sp.point, speedup_lo=sp.lo, @@ -1081,9 +1222,7 @@ def measure_problem(problem: ConvProblem, *, direction: Direction = "fwd", continue e = row[f"{arm}_eager_ms"] row[f"{arm}_launcher_ms"] = (e - row[f"{arm}_ms"]) if e else 0.0 - row["finalists"] = { - str(owner[nm]): meas[nm].median / c for nm in owner - } + row["finalists"] = {str(owner[nm]): meas[nm].median / c for nm in owner} return row except Exception as exc: # noqa: BLE001 row["error"] = f"{type(exc).__name__}: {exc}" @@ -1094,8 +1233,9 @@ def measure_problem(problem: ConvProblem, *, direction: Direction = "fwd", torch.cuda.empty_cache() -def _correctness(problem: ConvProblem, direction: Direction, - operator: Operator | None = None) -> dict: +def _correctness( + problem: ConvProblem, direction: Direction, operator: Operator | None = None +) -> dict: """Error against MIOpen's own answer, on the shape just measured. Not a substitute for the test suite -- ``tests/`` holds the bitwise-exact @@ -1110,8 +1250,10 @@ def _correctness(problem: ConvProblem, direction: Direction, got, ref = case.reference() d = (got.float() - ref.float()).abs() scale = ref.float().pow(2).mean().sqrt().item() or 1.0 - return {"max_abs_vs_miopen": d.max().item(), - "rms_rel_vs_miopen": (d.pow(2).mean().sqrt().item() / scale)} + return { + "max_abs_vs_miopen": d.max().item(), + "rms_rel_vs_miopen": (d.pow(2).mean().sqrt().item() / scale), + } except Exception as exc: # noqa: BLE001 return {"correctness_error": f"{type(exc).__name__}: {exc}"} finally: @@ -1143,76 +1285,117 @@ def main() -> None: ap = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter ) - ap.add_argument("--operator", default="all", - choices=["conv", "convT", "all"], - help="which operator; 'all' measures both (default). " - "Supersedes the old --skip-transposed, which was " - "store_true with default True and so could never be " - "turned off") - ap.add_argument("--direction", default="fwd", - choices=["fwd", "bwd-data", "bwd-weight", "all"]) - ap.add_argument("--top", type=int, default=8, - help="hottest N corpus problems (0 = all)") - ap.add_argument("--problems", default=None, - help="comma-separated corpus indices, overriding --top") + ap.add_argument( + "--operator", + default="all", + choices=["conv", "convT", "all"], + help="which operator; 'all' measures both (default). " + "Supersedes the old --skip-transposed, which was " + "store_true with default True and so could never be " + "turned off", + ) + ap.add_argument( + "--direction", default="fwd", choices=["fwd", "bwd-data", "bwd-weight", "all"] + ) + ap.add_argument( + "--top", type=int, default=8, help="hottest N corpus problems (0 = all)" + ) + ap.add_argument( + "--problems", + default=None, + help="comma-separated corpus indices, overriding --top", + ) ap.add_argument("--max-configs", type=int, default=0) - ap.add_argument("--iters", type=int, default=0, - help="calls per timed block; 0 (default) sizes it online " - "from the measured per-call time") - ap.add_argument("--rounds", type=int, default=0, - help="rounds of the race; 0 (default) grows until the " - "speedup's 95%% interval is inside --precision or " - "--budget seconds are spent") - ap.add_argument("--budget", type=float, default=20.0, - help="wall-clock seconds per cell's race (default 20)") - ap.add_argument("--precision", type=float, default=0.02, - help="target relative 95%% half-width on the reported " - "speedup and on each arm's median (default 0.02)") - ap.add_argument("--launcher", default="exclude", - choices=["exclude", "include"], - help="'exclude' (default): both arms are replayed from a " - "CUDA graph, so the number is kernel time -- no " - "dispatch, no config lookup, no launcher. 'include': " - "both arms are called from Python, so the number is " - "what a caller pays today") - ap.add_argument("--shipped", action="store_true", - help="skip the sweep and time the config the entry point " - "resolves on its own -- the kernel a caller gets") - ap.add_argument("--control", default="miopen", choices=["miopen", "none"], - help="'miopen' (default): race the Triton kernel against a " - "real MIOpen control and report a paired speedup with " - "its interval. 'none': measure the Triton kernel " - "alone, emitting no miopen_* and no speedup key. The " - "control is what makes a capture expensive -- MIOpen's " - "find cannot be replayed from disk under " - "cudnn.benchmark=True and costs 92-174 s per cell on " - "this corpus, which is 98% of a cell's wall clock") - ap.add_argument("--corpus", default="scaffold", - choices=["scaffold", "census"], - help="which problem list --top/--problems index into. " - "'scaffold' (default) is the 57 profiled problems, " - "cost-ordered, and is the key every stored capture " - "refers to -- its indices must not move. 'census' is " - "the 88 problems an instrumented step actually issued " - "at all four configurations, which is the only list " - "containing configuration B and the 2048-channel " - "sites; it is already in the adapter form, carries no " - "MIOpen timings and is not cost-ordered, so --form is " - "ignored for it and --top means 'the first N', not " - "'the hottest N'") - ap.add_argument("--form", default="distconv", - choices=["distconv", "adapter", "logical"], - help="which of the three shapes of a ScaFFold convolution " - "to measure. 'distconv' (default) is the halo'd, " - "unpadded form upstream DistConv issues and the form " - "every capture on disk was taken in; 'adapter' is what " - "ScaFFold's own Triton rung is handed, which is what " - "production runs and is padded everywhere; 'logical' " - "is the module's own statement. See the module " - "docstring -- these are three different tuning " - "problems, not three views of one") - ap.add_argument("--check", action="store_true", - help="also compare each result against MIOpen's") + ap.add_argument( + "--iters", + type=int, + default=0, + help="calls per timed block; 0 (default) sizes it online " + "from the measured per-call time", + ) + ap.add_argument( + "--rounds", + type=int, + default=0, + help="rounds of the race; 0 (default) grows until the " + "speedup's 95%% interval is inside --precision or " + "--budget seconds are spent", + ) + ap.add_argument( + "--budget", + type=float, + default=20.0, + help="wall-clock seconds per cell's race (default 20)", + ) + ap.add_argument( + "--precision", + type=float, + default=0.02, + help="target relative 95%% half-width on the reported " + "speedup and on each arm's median (default 0.02)", + ) + ap.add_argument( + "--launcher", + default="exclude", + choices=["exclude", "include"], + help="'exclude' (default): both arms are replayed from a " + "CUDA graph, so the number is kernel time -- no " + "dispatch, no config lookup, no launcher. 'include': " + "both arms are called from Python, so the number is " + "what a caller pays today", + ) + ap.add_argument( + "--shipped", + action="store_true", + help="skip the sweep and time the config the entry point " + "resolves on its own -- the kernel a caller gets", + ) + ap.add_argument( + "--control", + default="miopen", + choices=["miopen", "none"], + help="'miopen' (default): race the Triton kernel against a " + "real MIOpen control and report a paired speedup with " + "its interval. 'none': measure the Triton kernel " + "alone, emitting no miopen_* and no speedup key. The " + "control is what makes a capture expensive -- MIOpen's " + "find cannot be replayed from disk under " + "cudnn.benchmark=True and costs 92-174 s per cell on " + "this corpus, which is 98% of a cell's wall clock", + ) + ap.add_argument( + "--corpus", + default="scaffold", + choices=["scaffold", "census"], + help="which problem list --top/--problems index into. " + "'scaffold' (default) is the 57 profiled problems, " + "cost-ordered, and is the key every stored capture " + "refers to -- its indices must not move. 'census' is " + "the 88 problems an instrumented step actually issued " + "at all four configurations, which is the only list " + "containing configuration B and the 2048-channel " + "sites; it is already in the adapter form, carries no " + "MIOpen timings and is not cost-ordered, so --form is " + "ignored for it and --top means 'the first N', not " + "'the hottest N'", + ) + ap.add_argument( + "--form", + default="distconv", + choices=["distconv", "adapter", "logical"], + help="which of the three shapes of a ScaFFold convolution " + "to measure. 'distconv' (default) is the halo'd, " + "unpadded form upstream DistConv issues and the form " + "every capture on disk was taken in; 'adapter' is what " + "ScaFFold's own Triton rung is handed, which is what " + "production runs and is padded everywhere; 'logical' " + "is the module's own statement. See the module " + "docstring -- these are three different tuning " + "problems, not three views of one", + ) + ap.add_argument( + "--check", action="store_true", help="also compare each result against MIOpen's" + ) ap.add_argument("--verbose", action="store_true") ap.add_argument("--out", default=None) args = ap.parse_args() @@ -1237,38 +1420,51 @@ def main() -> None: "triton_conv3d/tests/ instead." ) - corpus = list(census_corpus() if args.corpus == "census" - else scaffold_corpus()) + corpus = list(census_corpus() if args.corpus == "census" else scaffold_corpus()) picks = _pick(corpus, args) operators = list(OPERATORS) if args.operator == "all" else [args.operator] directions = list(DIRECTIONS) if args.direction == "all" else [args.direction] props = torch.cuda.get_device_properties(0) - print(f"device {props.name}, {props.multi_processor_count} CUs, " - f"torch {torch.__version__}, cudnn.benchmark={torch.backends.cudnn.benchmark}") - print(f"control: {args.control}" - + ("" if args.control == "miopen" else - " (Triton alone; no speedup is reported and none should be " - "inferred)")) - print(f"timed region: {args.launcher} launcher " - f"({'CUDA graph replay -- kernels only, both arms' if args.launcher == 'exclude' else 'Python call -- kernel + dispatch + lookup + launcher, both arms'})") + print( + f"device {props.name}, {props.multi_processor_count} CUs, " + f"torch {torch.__version__}, cudnn.benchmark={torch.backends.cudnn.benchmark}" + ) + print( + f"control: {args.control}" + + ( + "" + if args.control == "miopen" + else " (Triton alone; no speedup is reported and none should be inferred)" + ) + ) + print( + f"timed region: {args.launcher} launcher " + f"({'CUDA graph replay -- kernels only, both arms' if args.launcher == 'exclude' else 'Python call -- kernel + dispatch + lookup + launcher, both arms'})" + ) rows: list[dict] = [] out_path = pathlib.Path(args.out) if args.out else None t0 = time.time() for opname in operators: op = _OPERATORS[opname] - mine = sorted([(i, p) for i, p in picks if op.selects(p)], - key=lambda ip: op.order(ip[1])) + mine = sorted( + [(i, p) for i, p in picks if op.selects(p)], key=lambda ip: op.order(ip[1]) + ) if not mine: continue for direction in directions: - form_note = ("recorded from a real step; already the adapter form" - if args.corpus == "census" else op.form_note(args.form)) - print(f"\n== {opname} {direction} -- {len(mine)} problems, " - f"--corpus {args.corpus} --form " - f"{'adapter' if args.corpus == 'census' else args.form}: " - f"{form_note}\n") + form_note = ( + "recorded from a real step; already the adapter form" + if args.corpus == "census" + else op.form_note(args.form) + ) + print( + f"\n== {opname} {direction} -- {len(mine)} problems, " + f"--corpus {args.corpus} --form " + f"{'adapter' if args.corpus == 'census' else args.form}: " + f"{form_note}\n" + ) for idx, p in mine: # The census records the shape *as the kernel was handed it*, # so it is already in the adapter form and carries no halo to @@ -1276,23 +1472,30 @@ def main() -> None: # no-op today and a silent lie the day the census gains a halo # field, so it is skipped by name rather than by luck. hp = p if args.corpus == "census" else op.form(p, args.form) - print(f" [{idx}] {hp.qualified_label} " - f"(GEMM {hp.gemm_shape(direction)})") + print( + f" [{idx}] {hp.qualified_label} (GEMM {hp.gemm_shape(direction)})" + ) sys.stdout.flush() row = measure_problem( - hp, direction=direction, operator=opname, - max_configs=args.max_configs, iters=args.iters, - rounds=args.rounds, shipped=args.shipped, - verbose=args.verbose, budget_s=args.budget, - target_rel=args.precision, launcher=args.launcher, - control=args.control) + hp, + direction=direction, + operator=opname, + max_configs=args.max_configs, + iters=args.iters, + rounds=args.rounds, + shipped=args.shipped, + verbose=args.verbose, + budget_s=args.budget, + target_rel=args.precision, + launcher=args.launcher, + control=args.control, + ) row["corpus_index"] = idx row["logical_problem"] = p.label # The form is recorded per row, not only in the header: a row # lifted out of one capture and quoted beside another is # precisely how a halo'd number became "what production runs". - row["shape_form"] = ("adapter" if args.corpus == "census" - else args.form) + row["shape_form"] = "adapter" if args.corpus == "census" else args.form row["corpus"] = args.corpus row["qualified_problem"] = hp.qualified_label row["padding"] = list(hp.padding) @@ -1303,58 +1506,106 @@ def main() -> None: _print_row(row) sys.stdout.flush() if out_path: - out_path.write_text(json.dumps( - {"device": props.name, "torch": torch.__version__, - "operator": args.operator, "direction": args.direction, - "shape_form": ("adapter" if args.corpus == "census" - else args.form), - "corpus": args.corpus, - "launcher": args.launcher, - "control": args.control, - "cudnn_benchmark": True, "rows": rows}, indent=1) + "\n") + out_path.write_text( + json.dumps( + { + "device": props.name, + "torch": torch.__version__, + "operator": args.operator, + "direction": args.direction, + "shape_form": ( + "adapter" if args.corpus == "census" else args.form + ), + "corpus": args.corpus, + "launcher": args.launcher, + "control": args.control, + "cudnn_benchmark": True, + "rows": rows, + }, + indent=1, + ) + + "\n" + ) ok = [r for r in rows if "error" not in r] # Two tables, not one with empty cells: without a control there are no # MIOpen columns to leave blank, and a blank column in a results table is # read as a missing value rather than as an absent measurement. if args.control == "miopen": - print("\n" + format_table( - [[ - f"{r['operator']} {r['direction']}", - r["problem"], - f"{r['gemm'][0]}x{r['gemm'][1]}x{r['gemm'][2]}", - f"{r['triton_ms']:.4f}", f"{r['triton_pct_roofline']:.0f}%", - f"{r['miopen_ms']:.4f}", f"{r['miopen_pct_roofline']:.0f}%", - f"{r['speedup']:.3f}x", - f"+-{r['speedup_rel_ci']:.1%}" - + ("" if r["speedup_significant"] else "?"), - f"{r['measure_rounds']}/{r['measure_stop'][:4]}", - f"{r['timed_region'][:4]}x{r['graph_chunk']}", - r["triton_config"], - ] for r in ok], - ["cell", "problem", "M x N x K", "triton ms", "%roof", "miopen ms", - "%roof", "speedup", "95% CI", "rounds", "timed", "best config"], - aligns="lllrrrrrrrrl", - )) + print( + "\n" + + format_table( + [ + [ + f"{r['operator']} {r['direction']}", + r["problem"], + f"{r['gemm'][0]}x{r['gemm'][1]}x{r['gemm'][2]}", + f"{r['triton_ms']:.4f}", + f"{r['triton_pct_roofline']:.0f}%", + f"{r['miopen_ms']:.4f}", + f"{r['miopen_pct_roofline']:.0f}%", + f"{r['speedup']:.3f}x", + f"+-{r['speedup_rel_ci']:.1%}" + + ("" if r["speedup_significant"] else "?"), + f"{r['measure_rounds']}/{r['measure_stop'][:4]}", + f"{r['timed_region'][:4]}x{r['graph_chunk']}", + r["triton_config"], + ] + for r in ok + ], + [ + "cell", + "problem", + "M x N x K", + "triton ms", + "%roof", + "miopen ms", + "%roof", + "speedup", + "95% CI", + "rounds", + "timed", + "best config", + ], + aligns="lllrrrrrrrrl", + ) + ) else: - print("\n" + format_table( - [[ - f"{r['operator']} {r['direction']}", - r["qualified_problem"], - f"{r['gemm'][0]}x{r['gemm'][1]}x{r['gemm'][2]}", - f"{r['triton_ms']:.4f}", - f"+-{r['triton_rel_ci']:.1%}", - f"{r['triton_cov']:.2%}", - f"{r['triton_pct_roofline']:.0f}%", - f"{r['triton_tflops']:.1f}", - f"{r['measure_rounds']}/{r['measure_stop'][:4]}", - f"{r['timed_region'][:4]}x{r['graph_chunk']}", - r["triton_config"], - ] for r in ok], - ["cell", "problem", "M x N x K", "triton ms", "95% CI", "CoV", - "%roof", "TFLOP/s", "rounds", "timed", "config"], - aligns="lllrrrrrrrl", - )) + print( + "\n" + + format_table( + [ + [ + f"{r['operator']} {r['direction']}", + r["qualified_problem"], + f"{r['gemm'][0]}x{r['gemm'][1]}x{r['gemm'][2]}", + f"{r['triton_ms']:.4f}", + f"+-{r['triton_rel_ci']:.1%}", + f"{r['triton_cov']:.2%}", + f"{r['triton_pct_roofline']:.0f}%", + f"{r['triton_tflops']:.1f}", + f"{r['measure_rounds']}/{r['measure_stop'][:4]}", + f"{r['timed_region'][:4]}x{r['graph_chunk']}", + r["triton_config"], + ] + for r in ok + ], + [ + "cell", + "problem", + "M x N x K", + "triton ms", + "95% CI", + "CoV", + "%roof", + "TFLOP/s", + "rounds", + "timed", + "config", + ], + aligns="lllrrrrrrrl", + ) + ) print(f"\nelapsed {time.time() - t0:.0f} s") if out_path: print(f"wrote {out_path}") @@ -1364,9 +1615,11 @@ def _print_row(row: dict) -> None: if "error" in row: print(f" ERROR {row['error'][:110]}") return + def launcher(arm): v = row.get(f"{arm}_launcher_ms", 0.0) return f", launcher +{v:.4f}" if v else "" + out = ( f" triton {row['triton_ms']:8.4f} +-{row['triton_rel_ci']:.1%} ms " f"({row['triton_pct_roofline']:5.1f}% roof, stall {row['triton_stall']:.2f}x, " @@ -1387,12 +1640,16 @@ def launcher(arm): else: # No control ran. Say so where the speedup would have been, rather # than leaving a blank a reader could take for a missing win. - out += (f" no MIOpen control (--control none) " - f"{row['measure_rounds']}r/{row['measure_stop']}") + out += ( + f" no MIOpen control (--control none) " + f"{row['measure_rounds']}r/{row['measure_stop']}" + ) out += f" in {row['measure_seconds']:.1f}s {row['timed_region_note']}" if row.get("rsck_ms"): - out += (f"\n weight transform {row['rsck_ms']:.4f} ms " - f"({100*row['rsck_ms']/row['triton_ms']:.1f}% of kernel)") + out += ( + f"\n weight transform {row['rsck_ms']:.4f} ms " + f"({100 * row['rsck_ms'] / row['triton_ms']:.1f}% of kernel)" + ) if "max_abs_vs_miopen" in row: out += f"\n max_abs vs MIOpen {row['max_abs_vs_miopen']:.3e}" print(out) diff --git a/triton_conv3d/bench/gemm_probe.py b/triton_conv3d/bench/gemm_probe.py index b5146be..dfc1f6a 100644 --- a/triton_conv3d/bench/gemm_probe.py +++ b/triton_conv3d/bench/gemm_probe.py @@ -39,7 +39,6 @@ import argparse import dataclasses -import itertools import json import pathlib import sys @@ -56,8 +55,7 @@ ConvProblem, hot_corpus, ) -from .harness import flush_caches, format_table, interleaved - +from .harness import format_table, interleaved # --------------------------------------------------------------------------- # A plain, honest GEMM @@ -66,11 +64,18 @@ @triton.jit def _gemm_kernel( - a_ptr, b_ptr, c_ptr, - M, N, K, - stride_am, stride_ak, - stride_bk, stride_bn, - stride_cm, stride_cn, + a_ptr, + b_ptr, + c_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, @@ -113,8 +118,16 @@ def _gemm_kernel( a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn else: - a_ptrs = a_ptr + offs_m[:, None].to(tl.int64) * stride_am + offs_k[None, :] * stride_ak - b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :].to(tl.int64) * stride_bn + a_ptrs = ( + a_ptr + + offs_m[:, None].to(tl.int64) * stride_am + + offs_k[None, :] * stride_ak + ) + b_ptrs = ( + b_ptr + + offs_k[:, None] * stride_bk + + offs_n[None, :].to(tl.int64) * stride_bn + ) acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) k_step = BLOCK_K * SPLIT_K @@ -131,7 +144,9 @@ def _gemm_kernel( a_ptrs += k_step * stride_ak b_ptrs += k_step * stride_bk - c_ptrs = c_ptr + offs_m[:, None].to(tl.int64) * stride_cm + offs_n[None, :] * stride_cn + c_ptrs = ( + c_ptr + offs_m[:, None].to(tl.int64) * stride_cm + offs_n[None, :] * stride_cn + ) c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) if SPLIT_K == 1: tl.store(c_ptrs, acc, mask=c_mask) @@ -267,9 +282,19 @@ def _default_kpack(block_k: int) -> int: return 1 if block_k <= 16 else 2 -def _amd_config(bm: int, bn: int, bk: int, warps: int, *, group_m: int, - split_k: int, nonkdim: int, kpack: int | None = None, - waves_per_eu: int = 0, int32: bool = False) -> GemmConfig | None: +def _amd_config( + bm: int, + bn: int, + bk: int, + warps: int, + *, + group_m: int, + split_k: int, + nonkdim: int, + kpack: int | None = None, + waves_per_eu: int = 0, + int32: bool = False, +) -> GemmConfig | None: """One AMD-knob config, or ``None`` if a hard constraint rejects it. The constraints are not preferences. ``BLOCK_M``/``BLOCK_N`` not divisible @@ -288,7 +313,13 @@ def _amd_config(bm: int, bn: int, bk: int, warps: int, *, group_m: int, if warps < 1: return None return GemmConfig( - bm, bn, bk, group_m, split_k, warps, num_stages=2, + bm, + bn, + bk, + group_m, + split_k, + warps, + num_stages=2, matrix_instr_nonkdim=nonkdim, kpack=_default_kpack(bk) if kpack is None else kpack, waves_per_eu=waves_per_eu, @@ -302,7 +333,9 @@ def _split_ks(m: int, n: int, k: int, bm: int, bn: int) -> tuple[int, ...]: return (1, 4, 16) if tiles < 228 and k >= 2048 else (1,) -def _candidate_configs(m: int, n: int, k: int, knobs: str = "legacy") -> list[GemmConfig]: +def _candidate_configs( + m: int, n: int, k: int, knobs: str = "legacy" +) -> list[GemmConfig]: """Configs worth trying for one shape. Not ``@triton.autotune``: that recompiles inside whatever is running at the @@ -363,8 +396,15 @@ def oversized(bm: int, bn: int, bk: int) -> bool: for nonkdim in nonkdims: for group_m in group_ms: for sk in _split_ks(m, n, k, bm, bn): - cfg = _amd_config(bm, bn, bk, warps, group_m=group_m, - split_k=sk, nonkdim=nonkdim) + cfg = _amd_config( + bm, + bn, + bk, + warps, + group_m=group_m, + split_k=sk, + nonkdim=nonkdim, + ) if cfg is not None and cfg not in seen: seen.add(cfg) out.append(cfg) @@ -381,19 +421,31 @@ def _randn(shape: tuple[int, ...], device, dtype: torch.dtype) -> torch.Tensor: return torch.randn(shape, device=device, dtype=dtype) -def _launch(a, b, c, cfg: GemmConfig, *, m: int, n: int, k: int, - stride_am: int, stride_bn: int): +def _launch( + a, b, c, cfg: GemmConfig, *, m: int, n: int, k: int, stride_am: int, stride_bn: int +): grid = (triton.cdiv(m, cfg.BLOCK_M) * triton.cdiv(n, cfg.BLOCK_N), cfg.SPLIT_K) _gemm_kernel[grid]( - a, b, c, - m, n, k, - stride_am, 1, - b.stride(0), stride_bn, - c.stride(0), c.stride(1), - BLOCK_M=cfg.BLOCK_M, BLOCK_N=cfg.BLOCK_N, BLOCK_K=cfg.BLOCK_K, - GROUP_M=cfg.GROUP_M, SPLIT_K=cfg.SPLIT_K, + a, + b, + c, + m, + n, + k, + stride_am, + 1, + b.stride(0), + stride_bn, + c.stride(0), + c.stride(1), + BLOCK_M=cfg.BLOCK_M, + BLOCK_N=cfg.BLOCK_N, + BLOCK_K=cfg.BLOCK_K, + GROUP_M=cfg.GROUP_M, + SPLIT_K=cfg.SPLIT_K, INT32_OFF=cfg.int32_offsets, - num_warps=cfg.num_warps, num_stages=cfg.num_stages, + num_warps=cfg.num_warps, + num_stages=cfg.num_stages, **cfg.amd_kwargs(), ) @@ -437,9 +489,17 @@ def plan_operands(m: int, n: int, k: int, mode: str, elem: int) -> dict | None: def best_triton_gemm( - m: int, n: int, k: int, *, dtype=torch.bfloat16, mode: str = "compute", - device="cuda", max_configs: int = 0, verbose: bool = False, - knobs: str = "legacy", sink: list | None = None, + m: int, + n: int, + k: int, + *, + dtype=torch.bfloat16, + mode: str = "compute", + device="cuda", + max_configs: int = 0, + verbose: bool = False, + knobs: str = "legacy", + sink: list | None = None, ) -> tuple[float, GemmConfig | None, int]: """Sweep configs, return ``(best ms/call, config, n_configs_that_ran)``. @@ -460,8 +520,19 @@ def best_triton_gemm( if max_configs: configs = configs[:max_configs] try: - return _sweep(a, b, c, configs, m=m, n=n, k=k, stride_am=stride_am, - stride_bn=stride_bn, verbose=verbose, sink=sink) + return _sweep( + a, + b, + c, + configs, + m=m, + n=n, + k=k, + stride_am=stride_am, + stride_bn=stride_bn, + verbose=verbose, + sink=sink, + ) finally: del a, b, c torch.cuda.empty_cache() @@ -471,19 +542,35 @@ def _alloc(plan: dict, m: int, n: int, k: int, device, dtype): a = _randn((plan["a_rows"], k), device, dtype) b = _randn((k, plan["b_cols"]), device, dtype) c = torch.empty((m, n), device=device, dtype=torch.float32) - return (a, b, c, - a.stride(0) if plan["stride_am"] is None else 0, - b.stride(1) if plan["stride_bn"] is None else 0) + return ( + a, + b, + c, + a.stride(0) if plan["stride_am"] is None else 0, + b.stride(1) if plan["stride_bn"] is None else 0, + ) -def _sweep(a, b, c, configs: list[GemmConfig], *, m: int, n: int, k: int, - stride_am: int, stride_bn: int, verbose: bool = False, - sink: list | None = None) -> tuple[float, GemmConfig | None, int]: +def _sweep( + a, + b, + c, + configs: list[GemmConfig], + *, + m: int, + n: int, + k: int, + stride_am: int, + stride_bn: int, + verbose: bool = False, + sink: list | None = None, +) -> tuple[float, GemmConfig | None, int]: """Time each config on already-allocated operands; return the winner.""" best_ms, best_cfg, ran = float("inf"), None, 0 for cfg in configs: - launch = _launcher(a, b, c, cfg, m=m, n=n, k=k, - stride_am=stride_am, stride_bn=stride_bn) + launch = _launcher( + a, b, c, cfg, m=m, n=n, k=k, stride_am=stride_am, stride_bn=stride_bn + ) try: if cfg.SPLIT_K > 1: c.zero_() @@ -506,12 +593,19 @@ def _launcher(a, b, c, cfg: GemmConfig, **kw): return lambda: _launch(a, b, c, cfg, **kw) -def torch_gemm_ms(m: int, n: int, k: int, *, dtype=torch.bfloat16, device="cuda") -> float: +def torch_gemm_ms( + m: int, n: int, k: int, *, dtype=torch.bfloat16, device="cuda" +) -> float: """hipBLASLt's time for the same GEMM -- the library reference.""" a = _randn((m, k), device, dtype) b = _randn((k, n), device, dtype) try: - meas = interleaved({"t": lambda: torch.matmul(a, b)}, warmup=5, iters=5, rounds=5) + # Bound as defaults, not captured: the ``finally`` below deletes both + # names, and a closure over a deleted name is only safe by accident of + # when it happens to be called. + meas = interleaved( + {"t": lambda a=a, b=b: torch.matmul(a, b)}, warmup=5, iters=5, rounds=5 + ) return meas["t"].median finally: del a, b @@ -523,45 +617,80 @@ def torch_gemm_ms(m: int, n: int, k: int, *, dtype=torch.bfloat16, device="cuda" # --------------------------------------------------------------------------- -def run_peak(sizes=(2048, 4096, 8192), dtype=torch.bfloat16, - knobs: str = "legacy") -> list[dict]: +def run_peak( + sizes=(2048, 4096, 8192), dtype=torch.bfloat16, knobs: str = "legacy" +) -> list[dict]: peak = PEAK_FLOPS["bf16"] if dtype is torch.bfloat16 else PEAK_FLOPS["fp32"] rows = [] for s in sizes: flops = 2 * s * s * s - tri_ms, cfg, ran = best_triton_gemm(s, s, s, dtype=dtype, mode="dram", - knobs=knobs) + tri_ms, cfg, ran = best_triton_gemm( + s, s, s, dtype=dtype, mode="dram", knobs=knobs + ) tor_ms = torch_gemm_ms(s, s, s, dtype=dtype) - rows.append({ - "size": s, - "triton_ms": tri_ms, "triton_tflops": flops / (tri_ms * 1e-3) / 1e12, - "triton_pct_peak": 100 * flops / (tri_ms * 1e-3) / peak, - "triton_config": str(cfg), "configs_ran": ran, - "torch_ms": tor_ms, "torch_tflops": flops / (tor_ms * 1e-3) / 1e12, - "torch_pct_peak": 100 * flops / (tor_ms * 1e-3) / peak, - "triton_vs_torch": tor_ms / tri_ms, - }) - print(format_table( - [[ - r["size"], - f"{r['triton_tflops']:.1f}", f"{r['triton_pct_peak']:.1f}%", - f"{r['torch_tflops']:.1f}", f"{r['torch_pct_peak']:.1f}%", - f"{r['triton_vs_torch']:.2f}x", r["triton_config"], - ] for r in rows[-1:]], - ["MNK", "triton TF/s", "%peak", "torch TF/s", "%peak", "tri/torch", "config"], - aligns="rrrrrrl", - ) if len(rows) == 1 else " " + " ".join([ - str(rows[-1]["size"]), - f"{rows[-1]['triton_tflops']:.1f}", f"{rows[-1]['triton_pct_peak']:.1f}%", - f"{rows[-1]['torch_tflops']:.1f}", f"{rows[-1]['torch_pct_peak']:.1f}%", - f"{rows[-1]['triton_vs_torch']:.2f}x", rows[-1]["triton_config"], - ])) + rows.append( + { + "size": s, + "triton_ms": tri_ms, + "triton_tflops": flops / (tri_ms * 1e-3) / 1e12, + "triton_pct_peak": 100 * flops / (tri_ms * 1e-3) / peak, + "triton_config": str(cfg), + "configs_ran": ran, + "torch_ms": tor_ms, + "torch_tflops": flops / (tor_ms * 1e-3) / 1e12, + "torch_pct_peak": 100 * flops / (tor_ms * 1e-3) / peak, + "triton_vs_torch": tor_ms / tri_ms, + } + ) + print( + format_table( + [ + [ + r["size"], + f"{r['triton_tflops']:.1f}", + f"{r['triton_pct_peak']:.1f}%", + f"{r['torch_tflops']:.1f}", + f"{r['torch_pct_peak']:.1f}%", + f"{r['triton_vs_torch']:.2f}x", + r["triton_config"], + ] + for r in rows[-1:] + ], + [ + "MNK", + "triton TF/s", + "%peak", + "torch TF/s", + "%peak", + "tri/torch", + "config", + ], + aligns="rrrrrrl", + ) + if len(rows) == 1 + else " " + + " ".join( + [ + str(rows[-1]["size"]), + f"{rows[-1]['triton_tflops']:.1f}", + f"{rows[-1]['triton_pct_peak']:.1f}%", + f"{rows[-1]['torch_tflops']:.1f}", + f"{rows[-1]['torch_pct_peak']:.1f}%", + f"{rows[-1]['triton_vs_torch']:.2f}x", + rows[-1]["triton_config"], + ] + ) + ) sys.stdout.flush() return rows -def run_peak_compare(sizes=(2048, 4096, 8192), knob_sets=("legacy", "amd"), - dtype=torch.bfloat16, device="cuda") -> list[dict]: +def run_peak_compare( + sizes=(2048, 4096, 8192), + knob_sets=("legacy", "amd"), + dtype=torch.bfloat16, + device="cuda", +) -> list[dict]: """Peak calibration where the knob sets are compared *against each other*. :func:`run_peak` sweeps one knob set and reports its winner, which is fine @@ -590,23 +719,41 @@ def run_peak_compare(sizes=(2048, 4096, 8192), knob_sets=("legacy", "amd"), for ks in knob_sets: configs = _candidate_configs(s, s, s, ks) sink: list = [] - ms, cfg, ran = _sweep(a, b, c, configs, m=s, n=s, k=s, - stride_am=stride_am, stride_bn=stride_bn, - sink=sink) + ms, cfg, ran = _sweep( + a, + b, + c, + configs, + m=s, + n=s, + k=s, + stride_am=stride_am, + stride_bn=stride_bn, + sink=sink, + ) row[f"{ks}_config"] = str(cfg) row[f"{ks}_configs_ran"] = ran row[f"{ks}_sweep_tflops"] = flops / (ms * 1e-3) / 1e12 by_str = {str(x): x for x in configs} for i, (name, _) in enumerate( - sorted(sink, key=lambda kv: kv[1])[:_FINALISTS]): + sorted(sink, key=lambda kv: kv[1])[:_FINALISTS] + ): owner[f"{ks}#{i}"] = (ks, by_str[name]) variants[f"{ks}#{i}"] = _launcher( - a, b, c, by_str[name], m=s, n=s, k=s, - stride_am=stride_am, stride_bn=stride_bn) + a, + b, + c, + by_str[name], + m=s, + n=s, + k=s, + stride_am=stride_am, + stride_bn=stride_bn, + ) owner["torch"] = ("torch", None) - variants["torch"] = lambda: torch.matmul(a, b) - meas = interleaved(variants, warmup=5, iters=5, - rounds=2 * len(variants)) + # Defaults rather than a closure, for the reason in ``torch_gemm_ms``. + variants["torch"] = lambda a=a, b=b: torch.matmul(a, b) + meas = interleaved(variants, warmup=5, iters=5, rounds=2 * len(variants)) for name, m_ in meas.items(): ks, cfg = owner[name] if m_.median >= row.get(f"{ks}_ms", float("inf")): @@ -627,17 +774,25 @@ def run_peak_compare(sizes=(2048, 4096, 8192), knob_sets=("legacy", "amd"), torch.cuda.empty_cache() rows.append(row) names = list(knob_sets) + ["torch"] - print(" " + " ".join( - [f"MNK={s:<5d}"] - + [f"{name}={row[f'{name}_tflops']:6.1f} TF/s" - f" ({row[f'{name}_pct_peak']:5.1f}%, spread {row[f'{name}_spread']:.1%}," - f" stall {row[f'{name}_stall']:.2f}x)" - for name in names if f"{name}_tflops" in row] - )) + print( + " " + + " ".join( + [f"MNK={s:<5d}"] + + [ + f"{name}={row[f'{name}_tflops']:6.1f} TF/s" + f" ({row[f'{name}_pct_peak']:5.1f}%, spread {row[f'{name}_spread']:.1%}," + f" stall {row[f'{name}_stall']:.2f}x)" + for name in names + if f"{name}_tflops" in row + ] + ) + ) for ks in knob_sets: - print(f" {ks:8s} winner {row[f'{ks}_config']} " - f"[{row[f'{ks}_configs_ran']} configs, " - f"{row[f'{ks}_sweep_tflops']:.1f} TF/s during the sweep]") + print( + f" {ks:8s} winner {row[f'{ks}_config']} " + f"[{row[f'{ks}_configs_ran']} configs, " + f"{row[f'{ks}_sweep_tflops']:.1f} TF/s during the sweep]" + ) sys.stdout.flush() return rows @@ -663,8 +818,18 @@ def run_peak_compare(sizes=(2048, 4096, 8192), knob_sets=("legacy", "amd"), _FINALISTS = 3 -def _measure_cell(m: int, n: int, k: int, *, mode: str, dtype, knob_sets, - keep_all: bool, max_configs: int, device="cuda") -> dict | None: +def _measure_cell( + m: int, + n: int, + k: int, + *, + mode: str, + dtype, + knob_sets, + keep_all: bool, + max_configs: int, + device="cuda", +) -> dict | None: """Sweep each knob set on one shape, then race the finalists. Returns ``None`` when the shape does not fit this mode's budget. Operands @@ -676,17 +841,33 @@ def _measure_cell(m: int, n: int, k: int, *, mode: str, dtype, knob_sets, return None a, b, c, stride_am, stride_bn = _alloc(plan, m, n, k, device, dtype) try: - out: dict = {"winners": {}, "ran": {}, "sweep": {}, "sinks": {}, - "headtohead": {}, "spread": {}, "stall": {}} + out: dict = { + "winners": {}, + "ran": {}, + "sweep": {}, + "sinks": {}, + "headtohead": {}, + "spread": {}, + "stall": {}, + } finalists: dict[str, list[GemmConfig]] = {} for ks in knob_sets: configs = _candidate_configs(m, n, k, ks) if max_configs: configs = configs[:max_configs] sink: list = [] - ms, cfg, ran = _sweep(a, b, c, configs, m=m, n=n, k=k, - stride_am=stride_am, stride_bn=stride_bn, - sink=sink) + ms, cfg, ran = _sweep( + a, + b, + c, + configs, + m=m, + n=n, + k=k, + stride_am=stride_am, + stride_bn=stride_bn, + sink=sink, + ) if cfg is None: continue out["winners"][ks], out["ran"][ks], out["sweep"][ks] = cfg, ran, ms @@ -702,22 +883,30 @@ def _measure_cell(m: int, n: int, k: int, *, mode: str, dtype, knob_sets, for i, cfg in enumerate(cfgs): name = f"{ks}#{i}" owner[name] = ks - variants[name] = _launcher(a, b, c, cfg, m=m, n=n, k=k, - stride_am=stride_am, - stride_bn=stride_bn) + variants[name] = _launcher( + a, + b, + c, + cfg, + m=m, + n=n, + k=k, + stride_am=stride_am, + stride_bn=stride_bn, + ) if cfg.SPLIT_K > 1: c.zero_() # Rounds a multiple of the variant count, so each variant occupies # each slot the same number of times -- with 2 variants over 5 # rounds one of them gets the post-warmup slot three times and the # other twice, which is worth several percent here. - meas = interleaved(variants, warmup=3, iters=5, - rounds=2 * len(variants)) + meas = interleaved(variants, warmup=3, iters=5, rounds=2 * len(variants)) out["h2h_best"] = {} for name, m_ in meas.items(): ks = owner[name] - out["h2h_best"][ks] = min(out["h2h_best"].get(ks, float("inf")), - m_.best) + out["h2h_best"][ks] = min( + out["h2h_best"].get(ks, float("inf")), m_.best + ) if m_.median < out["headtohead"].get(ks, float("inf")): out["headtohead"][ks] = m_.median out["spread"][ks] = m_.spread @@ -738,12 +927,17 @@ def _measure_cell(m: int, n: int, k: int, *, mode: str, dtype, knob_sets, torch.cuda.empty_cache() -def run_shapes(problems: list[ConvProblem], modes=("compute", "dram"), - dtype=torch.bfloat16, max_configs: int = 0, - knobs: str = "legacy", only_n: int | None = None, - keep_all: bool = False, - prior: dict[tuple[str, str], dict] | None = None, - flush=None) -> list[dict]: +def run_shapes( + problems: list[ConvProblem], + modes=("compute", "dram"), + dtype=torch.bfloat16, + max_configs: int = 0, + knobs: str = "legacy", + only_n: int | None = None, + keep_all: bool = False, + prior: dict[tuple[str, str], dict] | None = None, + flush=None, +) -> list[dict]: """Sweep every (problem, direction) cell. ``prior`` supplies cells a previous run already measured, keyed by @@ -769,8 +963,11 @@ def run_shapes(problems: list[ConvProblem], modes=("compute", "dram"), continue flops = 2 * m * n * k row = { - "problem": p.label, "direction": direction, - "M": m, "N": n, "K": k, + "problem": p.label, + "direction": direction, + "M": m, + "N": n, + "K": k, "conv_flops": p.flops(direction), "conv_ai": p.arithmetic_intensity(direction), "conv_roofline_tflops": p.roofline_flops(direction) / 1e12, @@ -783,9 +980,16 @@ def run_shapes(problems: list[ConvProblem], modes=("compute", "dram"), knob_sets = _KNOB_SETS.get(knobs, (knobs,)) for mode in modes: try: - cell = _measure_cell(m, n, k, mode=mode, dtype=dtype, - knob_sets=knob_sets, keep_all=keep_all, - max_configs=max_configs) + cell = _measure_cell( + m, + n, + k, + mode=mode, + dtype=dtype, + knob_sets=knob_sets, + keep_all=keep_all, + max_configs=max_configs, + ) except torch.OutOfMemoryError: cell = None if cell is None or not cell["winners"]: @@ -835,8 +1039,10 @@ def record(prefix: str, ms: float) -> None: # single-set run. Keeps the JSON shape compatible with # ``probe.json`` so before/after can be diffed key for key. head = knob_sets[-1] - record(mode, (cell["headtohead"] if len(knob_sets) > 1 - else cell["sweep"])[head]) + record( + mode, + (cell["headtohead"] if len(knob_sets) > 1 else cell["sweep"])[head], + ) row[f"{mode}_config"] = str(cell["winners"][head]) row[f"{mode}_configs_ran"] = cell["ran"][head] if len(knob_sets) > 1: @@ -852,28 +1058,37 @@ def record(prefix: str, ms: float) -> None: f" ({row.get(f'{mode}_pct_peak', float('nan')):5.1f}% peak)" for mode in modes ) - + (f" miopen={row['miopen_ms']:.3f} ms" - f" ({row['miopen_pct_roofline']:.0f}%)" if "miopen_ms" in row else "") + + ( + f" miopen={row['miopen_ms']:.3f} ms" + f" ({row['miopen_pct_roofline']:.0f}%)" + if "miopen_ms" in row + else "" + ) ) if len(knob_sets) > 1: for mode in modes: if row.get(f"{mode}_ms") is None: continue - print(" " + " | ".join( - f"{ks}: {row[f'{mode}_{ks}_pct_peak']:5.1f}% " - f"(sweep {row[f'{mode}_{ks}_sweep_tflops']:.0f} TF/s, " - f"stall {row[f'{mode}_{ks}_stall']:.2f}x) " - f"{row[f'{mode}_{ks}_config']}" - for ks in knob_sets if f"{mode}_{ks}_pct_peak" in row - )) + print( + " " + + " | ".join( + f"{ks}: {row[f'{mode}_{ks}_pct_peak']:5.1f}% " + f"(sweep {row[f'{mode}_{ks}_sweep_tflops']:.0f} TF/s, " + f"stall {row[f'{mode}_{ks}_stall']:.2f}x) " + f"{row[f'{mode}_{ks}_config']}" + for ks in knob_sets + if f"{mode}_{ks}_pct_peak" in row + ) + ) sys.stdout.flush() if flush: flush(rows) return rows -def run_isa(spec: str, *, m: int, n: int, k: int, device="cuda", - dtype=torch.bfloat16) -> None: +def run_isa( + spec: str, *, m: int, n: int, k: int, device="cuda", dtype=torch.bfloat16 +) -> None: """Compile and launch exactly one config, so its ISA can be inspected. Run under ``AMDGCN_ENABLE_DUMP=1`` with a cold ``TRITON_CACHE_DIR`` -- a @@ -886,63 +1101,90 @@ def run_isa(spec: str, *, m: int, n: int, k: int, device="cuda", parts = spec.split(",") if len(parts) != 8: raise SystemExit( - "--isa expects BM,BN,BK,GROUP_M,warps,nonkdim,kpack,int32 " - f"(got {spec!r})" + f"--isa expects BM,BN,BK,GROUP_M,warps,nonkdim,kpack,int32 (got {spec!r})" ) bm, bn, bk, gm, warps, nonkdim, kpack, int32 = (int(x) for x in parts) - cfg = GemmConfig(bm, bn, bk, gm, 1, warps, num_stages=2, - matrix_instr_nonkdim=nonkdim, kpack=kpack, waves_per_eu=0, - int32_offsets=bool(int32)) + cfg = GemmConfig( + bm, + bn, + bk, + gm, + 1, + warps, + num_stages=2, + matrix_instr_nonkdim=nonkdim, + kpack=kpack, + waves_per_eu=0, + int32_offsets=bool(int32), + ) a = _randn((m, k), device, dtype) b = _randn((k, n), device, dtype) c = torch.empty((m, n), device=device, dtype=torch.float32) - _launch(a, b, c, cfg, m=m, n=n, k=k, - stride_am=a.stride(0), stride_bn=b.stride(1)) + _launch(a, b, c, cfg, m=m, n=n, k=k, stride_am=a.stride(0), stride_bn=b.stride(1)) torch.cuda.synchronize() - print(f"ISA-DUMP-CONFIG {cfg} M={m} N={n} K={k} " - f"a_storage={a.untyped_storage().size()} " - f"b_storage={b.untyped_storage().size()} " - f"c_storage={c.untyped_storage().size()}") + print( + f"ISA-DUMP-CONFIG {cfg} M={m} N={n} K={k} " + f"a_storage={a.untyped_storage().size()} " + f"b_storage={b.untyped_storage().size()} " + f"c_storage={c.untyped_storage().size()}" + ) def main() -> None: - ap = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) - ap.add_argument("--mode", choices=("peak", "shapes", "both", "isa"), - default="both") - ap.add_argument("--knobs", - choices=("legacy", "amd", "amd-wide", - "compare", "compare-wide"), - default="legacy", - help="config generator. 'legacy' reproduces the M0 sweep " - "(no AMD kernargs). 'amd' is Inductor's ROCm conv " - "seed grid under the gfx942 MFMA constraints. " - "'amd-wide' adds skinny-N tiles and sweeps " - "matrix_instr_nonkdim. 'compare[-wide]' runs several " - "sets and races their champions in one interleaved " - "measurement, which is the only before/after immune " - "to a change of tenancy on the device.") - ap.add_argument("--only-n", type=int, default=None, - help="restrict to cells whose GEMM N equals this") - ap.add_argument("--keep-all-configs", action="store_true", - help="record every config's time, not just the winner") - ap.add_argument("--isa", default="128,128,64,6,8,16,2,0", - help="BM,BN,BK,GROUP_M,warps,nonkdim,kpack,int32 for --mode isa") - ap.add_argument("--isa-mnk", default="4096,4096,4096", - help="M,N,K for --mode isa") - ap.add_argument("--modes", default="compute", - help="comma-separated subset of compute,dram. 'dram' is " - "informative but slow -- it materializes im2col, which " - "for the corpus's larger shapes means a 15-30 GiB " - "operand per config -- and it is not the ceiling a " - "fused kernel is measured against.") + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--mode", choices=("peak", "shapes", "both", "isa"), default="both") + ap.add_argument( + "--knobs", + choices=("legacy", "amd", "amd-wide", "compare", "compare-wide"), + default="legacy", + help="config generator. 'legacy' reproduces the M0 sweep " + "(no AMD kernargs). 'amd' is Inductor's ROCm conv " + "seed grid under the gfx942 MFMA constraints. " + "'amd-wide' adds skinny-N tiles and sweeps " + "matrix_instr_nonkdim. 'compare[-wide]' runs several " + "sets and races their champions in one interleaved " + "measurement, which is the only before/after immune " + "to a change of tenancy on the device.", + ) + ap.add_argument( + "--only-n", + type=int, + default=None, + help="restrict to cells whose GEMM N equals this", + ) + ap.add_argument( + "--keep-all-configs", + action="store_true", + help="record every config's time, not just the winner", + ) + ap.add_argument( + "--isa", + default="128,128,64,6,8,16,2,0", + help="BM,BN,BK,GROUP_M,warps,nonkdim,kpack,int32 for --mode isa", + ) + ap.add_argument("--isa-mnk", default="4096,4096,4096", help="M,N,K for --mode isa") + ap.add_argument( + "--modes", + default="compute", + help="comma-separated subset of compute,dram. 'dram' is " + "informative but slow -- it materializes im2col, which " + "for the corpus's larger shapes means a 15-30 GiB " + "operand per config -- and it is not the ceiling a " + "fused kernel is measured against.", + ) ap.add_argument("--top", type=int, default=6, help="hottest N corpus problems") - ap.add_argument("--max-configs", type=int, default=0, - help="cap the config sweep (0 = no cap)") + ap.add_argument( + "--max-configs", type=int, default=0, help="cap the config sweep (0 = no cap)" + ) ap.add_argument("--out", default=None) - ap.add_argument("--resume", action="store_true", - help="reuse cells (and the peak block) already present in " - "--out instead of re-measuring them") + ap.add_argument( + "--resume", + action="store_true", + help="reuse cells (and the peak block) already present in " + "--out instead of re-measuring them", + ) args = ap.parse_args() if not torch.cuda.is_available(): @@ -952,14 +1194,22 @@ def main() -> None: run_isa(args.isa, m=m, n=n, k=k) return props = torch.cuda.get_device_properties(0) - print(f"device: {props.name}, {props.multi_processor_count} CUs, " - f"torch {torch.__version__}, triton {triton.__version__}") - print(f"roofline constants: {PEAK_FLOPS['bf16']/1e12:.0f} TFLOP/s bf16, " - f"{HBM_BYTES_PER_S/1e12:.1f} TB/s HBM " - f"(crossover at {PEAK_FLOPS['bf16']/HBM_BYTES_PER_S:.0f} FLOP/byte)\n") - - result: dict = {"device": props.name, "torch": torch.__version__, - "triton": triton.__version__, "knobs": args.knobs} + print( + f"device: {props.name}, {props.multi_processor_count} CUs, " + f"torch {torch.__version__}, triton {triton.__version__}" + ) + print( + f"roofline constants: {PEAK_FLOPS['bf16'] / 1e12:.0f} TFLOP/s bf16, " + f"{HBM_BYTES_PER_S / 1e12:.1f} TB/s HBM " + f"(crossover at {PEAK_FLOPS['bf16'] / HBM_BYTES_PER_S:.0f} FLOP/byte)\n" + ) + + result: dict = { + "device": props.name, + "torch": torch.__version__, + "triton": triton.__version__, + "knobs": args.knobs, + } prior: dict[tuple[str, str], dict] = {} out_path = pathlib.Path(args.out) if args.out else None if args.resume and out_path and out_path.exists(): @@ -967,8 +1217,10 @@ def main() -> None: prior = {(r["problem"], r["direction"]): r for r in old.get("shapes", [])} if old.get("peak"): result["peak"] = old["peak"] - print(f"resuming: {len(prior)} cells already measured" - + (", peak block reused" if "peak" in result else "")) + print( + f"resuming: {len(prior)} cells already measured" + + (", peak block reused" if "peak" in result else "") + ) def write(rows=None) -> None: if not out_path: @@ -978,8 +1230,10 @@ def write(rows=None) -> None: payload["shapes"] = rows out_path.write_text(json.dumps(payload, indent=1) + "\n") - print(f"knob set: {args.knobs}" - + (f", restricted to N={args.only_n}" if args.only_n else "")) + print( + f"knob set: {args.knobs}" + + (f", restricted to N={args.only_n}" if args.only_n else "") + ) t0 = time.time() if args.mode in ("peak", "both") and "peak" not in result: print("== peak: square bf16 GEMM, Triton vs hipBLASLt ==") @@ -991,12 +1245,16 @@ def write(rows=None) -> None: print() if args.mode in ("shapes", "both"): print(f"== shapes: conv-implied GEMMs, top {args.top} corpus problems ==") - result["shapes"] = run_shapes(list(hot_corpus(args.top)), - modes=tuple(args.modes.split(",")), - max_configs=args.max_configs, - knobs=args.knobs, only_n=args.only_n, - keep_all=args.keep_all_configs, - prior=prior, flush=write) + result["shapes"] = run_shapes( + list(hot_corpus(args.top)), + modes=tuple(args.modes.split(",")), + max_configs=args.max_configs, + knobs=args.knobs, + only_n=args.only_n, + keep_all=args.keep_all_configs, + prior=prior, + flush=write, + ) result["elapsed_s"] = time.time() - t0 print(f"\nelapsed {result['elapsed_s']:.0f} s") diff --git a/triton_conv3d/bench/harness.py b/triton_conv3d/bench/harness.py index 14831b0..e5cc8d0 100644 --- a/triton_conv3d/bench/harness.py +++ b/triton_conv3d/bench/harness.py @@ -118,9 +118,36 @@ #: Student-t 97.5th percentile by degrees of freedom, so an interval can be #: quoted without a scipy dependency. Index 0 is unused. _T975 = ( - math.nan, 12.706, 4.303, 3.182, 2.776, 2.571, 2.447, 2.365, 2.306, 2.262, - 2.228, 2.201, 2.179, 2.160, 2.145, 2.131, 2.120, 2.110, 2.101, 2.093, - 2.086, 2.080, 2.074, 2.069, 2.064, 2.060, 2.056, 2.052, 2.048, 2.045, + math.nan, + 12.706, + 4.303, + 3.182, + 2.776, + 2.571, + 2.447, + 2.365, + 2.306, + 2.262, + 2.228, + 2.201, + 2.179, + 2.160, + 2.145, + 2.131, + 2.120, + 2.110, + 2.101, + 2.093, + 2.086, + 2.080, + 2.074, + 2.069, + 2.064, + 2.060, + 2.056, + 2.052, + 2.048, + 2.045, ) #: Asymptotic ratio of the standard error of a sample median to that of the @@ -247,7 +274,9 @@ def spread(self) -> float: 100 -- all of it arithmetic. Compare :attr:`rel_half_width` instead, which is an interval and does not have that defect. """ - return (max(self.rounds) - min(self.rounds)) / self.median if self.rounds else 0.0 + return ( + (max(self.rounds) - min(self.rounds)) / self.median if self.rounds else 0.0 + ) @property def cov(self) -> float: @@ -326,30 +355,40 @@ def significant(self) -> bool: def __str__(self) -> str: star = "" if self.significant else " (consistent with no difference)" - return (f"{self.numerator}/{self.denominator} = {self.point:.3f}x " - f"[{self.lo:.3f}, {self.hi:.3f}], n={self.n}{star}") + return ( + f"{self.numerator}/{self.denominator} = {self.point:.3f}x " + f"[{self.lo:.3f}, {self.hi:.3f}], n={self.n}{star}" + ) def ratio(numerator: Measurement, denominator: Measurement) -> Ratio: """Paired ratio ``numerator / denominator`` with a 95% interval.""" n = min(len(numerator.rounds), len(denominator.rounds)) - pairs = [numerator.rounds[i] / denominator.rounds[i] for i in range(n) - if denominator.rounds[i] > 0] + pairs = [ + numerator.rounds[i] / denominator.rounds[i] + for i in range(n) + if denominator.rounds[i] > 0 + ] if not pairs: - return Ratio(numerator.name, denominator.name, math.nan, math.nan, - math.nan, 0) + return Ratio(numerator.name, denominator.name, math.nan, math.nan, math.nan, 0) logs = [math.log(p) for p in pairs] point = math.exp(statistics.median(logs)) hw = _half_width(logs) if not math.isfinite(hw): - return Ratio(numerator.name, denominator.name, point, 0.0, math.inf, - len(pairs)) - return Ratio(numerator.name, denominator.name, point, - point * math.exp(-hw), point * math.exp(hw), len(pairs)) - - -def _time_block(fn: Callable[[], object], iters: int, - group: int = 1) -> tuple[float, float]: + return Ratio(numerator.name, denominator.name, point, 0.0, math.inf, len(pairs)) + return Ratio( + numerator.name, + denominator.name, + point, + point * math.exp(-hw), + point * math.exp(hw), + len(pairs), + ) + + +def _time_block( + fn: Callable[[], object], iters: int, group: int = 1 +) -> tuple[float, float]: """Return ``(median ms per call, stall ratio)`` for ``iters`` calls. Events rather than the wall clock: the launch is asynchronous, so a wall @@ -397,8 +436,9 @@ def _blocked_events(fn: Callable[[], object], iters: int, group: int) -> list: return marks -def _time_block_full(fn: Callable[[], object], iters: int, - group: int) -> tuple[float, float, float]: +def _time_block_full( + fn: Callable[[], object], iters: int, group: int +) -> tuple[float, float, float]: """``(median, stall ratio, first sample)`` -- the first is the cold one.""" marks = _blocked_events(fn, iters, group) n = len(marks) - 1 @@ -459,8 +499,9 @@ class Plan: _GROUP_TRIGGER = 2.0 -def _measure_tax(fn: Callable[[], object], iters: int, group: int = 1, - reps: int = 5) -> float: +def _measure_tax( + fn: Callable[[], object], iters: int, group: int = 1, reps: int = 5 +) -> float: """Per-call cost of the event instrument: events minus an event-free bracket. Measured as a *paired* difference -- ev, br, ev, br, ... -- rather than as a @@ -477,10 +518,18 @@ def _measure_tax(fn: Callable[[], object], iters: int, group: int = 1, return statistics.median(diffs) -def _probe(fn: Callable[[], object], *, pinned_warmup: int | None, - warmup_s: float, warmup_min: int, warmup_max: int, - warmup_hard_s: float, block_ms: float, max_iters: int, - need_duration: bool) -> tuple[float, int]: +def _probe( + fn: Callable[[], object], + *, + pinned_warmup: int | None, + warmup_s: float, + warmup_min: int, + warmup_max: int, + warmup_hard_s: float, + block_ms: float, + max_iters: int, + need_duration: bool, +) -> tuple[float, int]: """Warm one variant and return ``(per-call ms, warmup calls issued)``. The first call absorbs whatever one-off the variant has -- MIOpen's find is @@ -530,12 +579,18 @@ def _probe(fn: Callable[[], object], *, pinned_warmup: int | None, return d, n + 2 -def per_call_ms(fn: Callable[[], object], *, warmup_s: float = 0.05, - warmup_min: int = 5, warmup_max: int = 200, - warmup_hard_s: float = 2.0, - block_ms: float = _BLOCK_TARGET_MS, - max_iters: int = _MAX_ITERS, - settle_s: float = 0.25, settle_calls: int = 5) -> float: +def per_call_ms( + fn: Callable[[], object], + *, + warmup_s: float = 0.05, + warmup_min: int = 5, + warmup_max: int = 200, + warmup_hard_s: float = 2.0, + block_ms: float = _BLOCK_TARGET_MS, + max_iters: int = _MAX_ITERS, + settle_s: float = 0.25, + settle_calls: int = 5, +) -> float: """One warmed, event-free estimate of a callable's per-call time. The same probe :func:`interleaved` runs internally, exposed because a caller @@ -562,14 +617,22 @@ def per_call_ms(fn: Callable[[], object], *, warmup_s: float = 0.05, torch.cuda.synchronize() if time.perf_counter() - t0 > settle_s: break - return _probe(fn, pinned_warmup=None, warmup_s=warmup_s, - warmup_min=warmup_min, warmup_max=warmup_max, - warmup_hard_s=warmup_hard_s, block_ms=block_ms, - max_iters=max_iters, need_duration=True)[0] - - -def _common_group(durations: Sequence[float], *, tax_budget: float, - max_iters: int) -> int: + return _probe( + fn, + pinned_warmup=None, + warmup_s=warmup_s, + warmup_min=warmup_min, + warmup_max=warmup_max, + warmup_hard_s=warmup_hard_s, + block_ms=block_ms, + max_iters=max_iters, + need_duration=True, + )[0] + + +def _common_group( + durations: Sequence[float], *, tax_budget: float, max_iters: int +) -> int: """One event-interval width for **every** arm of a comparison. Two things are load-bearing here and both were bought with a wrong answer. @@ -777,8 +840,7 @@ def __call__(self) -> None: self.graph.replay() -def capture(fn: Callable[[], object], chunk: int = 1, *, - warmup: int = 3) -> Captured: +def capture(fn: Callable[[], object], chunk: int = 1, *, warmup: int = 3) -> Captured: """Put ``chunk`` calls of ``fn`` in a CUDA graph, or raise :class:`CaptureError`. Everything -- the warmup and the capture -- runs on :func:`capture_stream`, @@ -825,9 +887,13 @@ def capture(fn: Callable[[], object], chunk: int = 1, *, return Captured(g, chunk) -def common_chunk(durations: Sequence[float], *, cost_ms: float = _REPLAY_COST_MS, - budget: float = _REPLAY_BUDGET, - max_chunk: int = _MAX_CHUNK) -> int: +def common_chunk( + durations: Sequence[float], + *, + cost_ms: float = _REPLAY_COST_MS, + budget: float = _REPLAY_BUDGET, + max_chunk: int = _MAX_CHUNK, +) -> int: """Calls per graph, for **every** arm of a comparison. Same shape and same reasoning as :func:`_common_group`, one level up: a @@ -846,8 +912,9 @@ def common_chunk(durations: Sequence[float], *, cost_ms: float = _REPLAY_COST_MS return max(1, min(max_chunk, 1 << math.ceil(math.log2(need)))) -def graph_is_worthwhile(durations: Sequence[float], - max_ms: float = _GRAPH_MAX_MS) -> bool: +def graph_is_worthwhile( + durations: Sequence[float], max_ms: float = _GRAPH_MAX_MS +) -> bool: """Is the host launch cost big enough to be worth capturing away? Above :data:`_GRAPH_MAX_MS` the measured worst-case host launch cost @@ -868,8 +935,9 @@ def time_callable( **kwargs, ) -> Measurement: """Time a single callable. Prefer :func:`interleaved` for comparisons.""" - return interleaved({"fn": fn}, warmup=warmup, iters=iters, rounds=rounds, - flush=flush, **kwargs)["fn"] + return interleaved( + {"fn": fn}, warmup=warmup, iters=iters, rounds=rounds, flush=flush, **kwargs + )["fn"] def interleaved( @@ -943,10 +1011,15 @@ def interleaved( need_duration = iters is None and not flush probes = { name: _probe( - variants[name], pinned_warmup=warmup, warmup_s=warmup_s, - warmup_min=warmup_min, warmup_max=warmup_max, - warmup_hard_s=warmup_hard_s, block_ms=block_ms, - max_iters=max_iters, need_duration=need_duration, + variants[name], + pinned_warmup=warmup, + warmup_s=warmup_s, + warmup_min=warmup_min, + warmup_max=warmup_max, + warmup_hard_s=warmup_hard_s, + block_ms=block_ms, + max_iters=max_iters, + need_duration=need_duration, ) for name in names } @@ -959,10 +1032,13 @@ def interleaved( group = 1 sizes = {name: iters for name in names} else: - group = _common_group([probes[n][0] for n in names], - tax_budget=tax_budget, max_iters=max_iters) - sizes = {name: _size_block(probes[name][0], group, block_ms, max_iters) - for name in names} + group = _common_group( + [probes[n][0] for n in names], tax_budget=tax_budget, max_iters=max_iters + ) + sizes = { + name: _size_block(probes[name][0], group, block_ms, max_iters) + for name in names + } plans: dict[str, Plan] = {} for name in names: d, warmed = probes[name] @@ -987,8 +1063,9 @@ def interleaved( for name in order: if flush: flush_caches() - ms, stall, first = _time_block_full(variants[name], plans[name].iters, - plans[name].group) + ms, stall, first = _time_block_full( + variants[name], plans[name].iters, plans[name].group + ) times[name].append(ms) stalls[name].append(stall) firsts[name].append(first) @@ -1018,10 +1095,16 @@ def interleaved( balanced = len(names) < 2 or (r % block) == 0 return { name: Measurement( - name, tuple(times[name]), tuple(stalls[name]), - iters=plans[name].iters, group=plans[name].group, - firsts=tuple(firsts[name]), tax_ms=plans[name].tax_ms, - stop=stop, balanced=balanced, seconds=seconds, + name, + tuple(times[name]), + tuple(stalls[name]), + iters=plans[name].iters, + group=plans[name].group, + firsts=tuple(firsts[name]), + tax_ms=plans[name].tax_ms, + stop=stop, + balanced=balanced, + seconds=seconds, ) for name in names } @@ -1054,11 +1137,13 @@ def format_table( for i, h in enumerate(headers) ] aligns = aligns or "l" * len(headers) + def fmt(cells: Sequence[str]) -> str: return " ".join( c.rjust(w) if a == "r" else c.ljust(w) for c, w, a in zip(cells, widths, aligns) ) + lines = [fmt(headers), " ".join("-" * w for w in widths)] lines += [fmt(r) for r in rows] return "\n".join(lines) diff --git a/triton_conv3d/bwd_data.py b/triton_conv3d/bwd_data.py index 07d8078..6f18e05 100644 --- a/triton_conv3d/bwd_data.py +++ b/triton_conv3d/bwd_data.py @@ -77,8 +77,8 @@ import torch from .gather_gemm import ( - ConvConfig, _MFMA_KDIM, + ConvConfig, _check_weight_rsck, _triple, conv3d_forward, @@ -107,9 +107,16 @@ def bwd_data_padding(padding, dilation, kernel) -> tuple[int, int, int]: def _tuned(bm: int, bn: int, bk: int, warps: int, group_m: int = 6) -> ConvConfig: - return ConvConfig(BLOCK_M=bm, BLOCK_N=bn, BLOCK_K=bk, GROUP_M=group_m, - num_warps=warps, num_stages=2, matrix_instr_nonkdim=16, - kpack=1 if bk <= 16 else 2) + return ConvConfig( + BLOCK_M=bm, + BLOCK_N=bn, + BLOCK_K=bk, + GROUP_M=group_m, + num_warps=warps, + num_stages=2, + matrix_instr_nonkdim=16, + kpack=1 if bk <= 16 else 2, + ) #: Measured backward-data winners, keyed by the **forward** problem's @@ -179,8 +186,13 @@ def register_tuned_bwd_data(dtype, cin, cout, kernel, config: ConvConfig) -> Non def bwd_data_config( - grad_output_shape: Sequence[int], cin: int, kernel: Sequence[int], - dtype: torch.dtype = torch.bfloat16, *, padding=0, dilation=1, + grad_output_shape: Sequence[int], + cin: int, + kernel: Sequence[int], + dtype: torch.dtype = torch.bfloat16, + *, + padding=0, + dilation=1, ) -> ConvConfig: """The config :func:`conv3d_backward_data` would pick for this problem. @@ -191,12 +203,19 @@ def bwd_data_config( k = _triple(kernel, "kernel") d = _triple(dilation, "dilation") p = _triple(padding, "padding") - in_sp = [o + 2 * (d[i] * (k[i] - 1) - p[i]) - d[i] * (k[i] - 1) - for i, o in enumerate(out_sp)] + in_sp = [ + o + 2 * (d[i] * (k[i] - 1) - p[i]) - d[i] * (k[i] - 1) + for i, o in enumerate(out_sp) + ] m = n * math.prod(in_sp) return select_config( - m, cout, cin, k, dtype, - table=_TUNED_BWD, key=tune_key(dtype, cin, cout, k), + m, + cout, + cin, + k, + dtype, + table=_TUNED_BWD, + key=tune_key(dtype, cin, cout, k), ) @@ -239,8 +258,11 @@ def is_supported_bwd_data( # peer access enabled -- which is how ScaFFold runs its four GPUs -- that # reads another rank's memory instead of faulting. A wrong gradient, not a # crash. - if (not grad_output.is_cuda or not weight.is_cuda - or weight.device != grad_output.device): + if ( + not grad_output.is_cuda + or not weight.is_cuda + or weight.device != grad_output.device + ): return False try: s = _triple(stride, "stride") @@ -307,8 +329,9 @@ def conv3d_backward_data( ``channels_last_3d`` parameter, pointless: pass the parameter as ``weight`` and the kernel reads it in place. """ - if not is_supported_bwd_data(grad_output, weight, input_shape, stride, - padding, dilation, groups): + if not is_supported_bwd_data( + grad_output, weight, input_shape, stride, padding, dilation, groups + ): raise NotImplementedError( f"unsupported: grad_output={tuple(grad_output.shape)}/" f"{grad_output.dtype} w={tuple(weight.shape)} " @@ -322,7 +345,11 @@ def conv3d_backward_data( cout = int(grad_output.shape[1]) if config is None: config = select_config( - n * math.prod(in_sp), cout, cin, k, grad_output.dtype, + n * math.prod(in_sp), + cout, + cin, + k, + grad_output.dtype, table=_TUNED_BWD, key=tune_key(grad_output.dtype, cin, cout, k), ) @@ -354,6 +381,14 @@ def conv3d_backward_data( _check_weight_rsck(weight_rsck, (*k, cin, cout), grad_output) w_view = weight_rsck.permute(3, 4, 0, 1, 2) return conv3d_forward( - grad_output, w_view, None, 1, pad, dilation, 1, - config=config, weight_flip=True, out=out, + grad_output, + w_view, + None, + 1, + pad, + dilation, + 1, + config=config, + weight_flip=True, + out=out, ) diff --git a/triton_conv3d/gather_gemm.py b/triton_conv3d/gather_gemm.py index 66fcc38..fe29f51 100644 --- a/triton_conv3d/gather_gemm.py +++ b/triton_conv3d/gather_gemm.py @@ -114,30 +114,63 @@ @triton.jit def _conv3d_fwd_kernel( - X, W, Y, BIAS, + X, + W, + Y, + BIAS, # Sizes. ``M_TOTAL`` is ``BATCH * OUT_D * OUT_H * OUT_W``. - BATCH, IN_D, IN_H, IN_W, - OUT_D, OUT_H, OUT_W, - CIN, COUT, M_TOTAL, + BATCH, + IN_D, + IN_H, + IN_W, + OUT_D, + OUT_H, + OUT_W, + CIN, + COUT, + M_TOTAL, # Element strides. The channel stride of X and Y is 1 by construction -- # that is what NDHWC means -- so it is not passed and not multiplied by. - stride_xn, stride_xd, stride_xh, stride_xw, + stride_xn, + stride_xd, + stride_xh, + stride_xw, # The weight, described by three strides over the *effective* GEMM's axes: # the fused tap index, the reduction axis K (Cin), and the output axis N # (Cout). Which of the two channel strides is 1 is a constexpr (``W_ORDER``) # rather than a runtime fact, because it decides how the tile is loaded. - stride_wt, stride_wk, stride_wn, - stride_yn, stride_yd, stride_yh, stride_yw, - KD: tl.constexpr, KH: tl.constexpr, KW: tl.constexpr, - SD: tl.constexpr, SH: tl.constexpr, SW: tl.constexpr, - PD: tl.constexpr, PH: tl.constexpr, PW: tl.constexpr, - DD: tl.constexpr, DH: tl.constexpr, DW: tl.constexpr, - BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, - BLOCK_K_COUNT: tl.constexpr, GROUP_M: tl.constexpr, - HAS_BIAS: tl.constexpr, EVEN_K: tl.constexpr, EVEN_N: tl.constexpr, - PADDED: tl.constexpr, INDEX_DTYPE: tl.constexpr, + stride_wt, + stride_wk, + stride_wn, + stride_yn, + stride_yd, + stride_yh, + stride_yw, + KD: tl.constexpr, + KH: tl.constexpr, + KW: tl.constexpr, + SD: tl.constexpr, + SH: tl.constexpr, + SW: tl.constexpr, + PD: tl.constexpr, + PH: tl.constexpr, + PW: tl.constexpr, + DD: tl.constexpr, + DH: tl.constexpr, + DW: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_K_COUNT: tl.constexpr, + GROUP_M: tl.constexpr, + HAS_BIAS: tl.constexpr, + EVEN_K: tl.constexpr, + EVEN_N: tl.constexpr, + PADDED: tl.constexpr, + INDEX_DTYPE: tl.constexpr, INPUT_PRECISION: tl.constexpr, - W_ORDER: tl.constexpr, W_FLIP: tl.constexpr, + W_ORDER: tl.constexpr, + W_FLIP: tl.constexpr, ): # -- which output tile this program owns ------------------------------ # @@ -208,9 +241,12 @@ def _conv3d_fwd_kernel( in_w = src_w + j * DW row_ok = ( m_valid - & (in_d >= 0) & (in_d < IN_D) - & (in_h >= 0) & (in_h < IN_H) - & (in_w >= 0) & (in_w < IN_W) + & (in_d >= 0) + & (in_d < IN_D) + & (in_h >= 0) + & (in_h < IN_H) + & (in_w >= 0) + & (in_w < IN_W) ) else: # Unpadded: every tap of an in-range output voxel is in range, so @@ -272,9 +308,7 @@ def _conv3d_fwd_kernel( # over the buffer-op limit and has lost it anyway (see # :data:`~triton_conv3d.shapes.BUFFER_OP_MAX_BYTES`). dij_w = (KD * KH * KW - 1 - dij) if W_FLIP else dij - w_row = ( - dij_w.to(INDEX_DTYPE) * stride_wt + offs_k.to(INDEX_DTYPE) * stride_wk - ) + w_row = dij_w.to(INDEX_DTYPE) * stride_wt + offs_k.to(INDEX_DTYPE) * stride_wk if W_ORDER == 0: w_ptrs = W + w_row[:, None] + offs_n[None, :] else: @@ -449,8 +483,9 @@ def _fit_to_lds(cfg: ConvConfig, dtype: torch.dtype) -> ConvConfig: while cfg.lds_bytes(dtype) > _LDS_BYTES: half_k, half_m, half_n = cfg.BLOCK_K // 2, cfg.BLOCK_M // 2, cfg.BLOCK_N // 2 if half_k >= kdim and half_k % kdim == 0: - cfg = dataclasses.replace(cfg, BLOCK_K=half_k, - kpack=1 if half_k <= 16 else cfg.kpack) + cfg = dataclasses.replace( + cfg, BLOCK_K=half_k, kpack=1 if half_k <= 16 else cfg.kpack + ) elif half_m >= nk and half_m % nk == 0: cfg = dataclasses.replace(cfg, BLOCK_M=half_m) elif half_n >= nk and half_n % nk == 0: @@ -500,12 +535,17 @@ def default_config( return _fit_to_lds( _fit_to_grid( ConvConfig( - BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=block_k, GROUP_M=6, + BLOCK_M=block_m, + BLOCK_N=block_n, + BLOCK_K=block_k, + GROUP_M=6, num_warps=8 if block_k >= 128 or block_n >= 256 else 4, - num_stages=2, matrix_instr_nonkdim=nonkdim, + num_stages=2, + matrix_instr_nonkdim=nonkdim, kpack=1 if block_k <= 16 else 2, ), - m, cout, + m, + cout, ), dtype, ) @@ -598,8 +638,13 @@ def _fit_to_grid(cfg: ConvConfig, m: int, cout: int) -> ConvConfig: def candidate_configs( - m: int, cin: int, cout: int, dtype: torch.dtype = torch.bfloat16, - *, group_ms: Sequence[int] = (6,), nonkdims: Sequence[int] = (16, 32), + m: int, + cin: int, + cout: int, + dtype: torch.dtype = torch.bfloat16, + *, + group_ms: Sequence[int] = (6,), + nonkdims: Sequence[int] = (16, 32), ) -> list[ConvConfig]: """Configs worth timing for one shape, already pruned to legal ones. @@ -630,8 +675,12 @@ def candidate_configs( for nonkdim in nonkdims: for group_m in group_ms: cfg = ConvConfig( - BLOCK_M=bm, BLOCK_N=bn, BLOCK_K=bk, GROUP_M=group_m, - num_warps=warps, num_stages=2, + BLOCK_M=bm, + BLOCK_N=bn, + BLOCK_K=bk, + GROUP_M=group_m, + num_warps=warps, + num_stages=2, matrix_instr_nonkdim=nonkdim, kpack=1 if bk <= 16 else 2, ) @@ -639,9 +688,11 @@ def candidate_configs( # fold two seed tiles onto one entry and silently # double-count it in the sweep. Only configs that could not # have run at all are removed, so no measured winner is lost. - if (cfg.validate(dtype) is not None - or cfg.lds_bytes(dtype) > _LDS_BYTES - or cfg in seen): + if ( + cfg.validate(dtype) is not None + or cfg.lds_bytes(dtype) > _LDS_BYTES + or cfg in seen + ): continue seen.add(cfg) out.append(cfg) @@ -650,13 +701,13 @@ def candidate_configs( return out -def tune_key(dtype: torch.dtype, cin: int, cout: int, - kernel: tuple[int, ...]) -> tuple: +def tune_key(dtype: torch.dtype, cin: int, cout: int, kernel: tuple[int, ...]) -> tuple: return (str(dtype), cin, cout, tuple(kernel)) -def _tuned(bm: int, bn: int, bk: int, warps: int, group_m: int = 6, - nk: int = 16) -> ConvConfig: +def _tuned( + bm: int, bn: int, bk: int, warps: int, group_m: int = 6, nk: int = 16 +) -> ConvConfig: """One measured row. ``nk`` defaults to 16 because that is what the forward measured -- 16 won @@ -669,9 +720,16 @@ def _tuned(bm: int, bn: int, bk: int, warps: int, group_m: int = 6, :func:`~triton_conv3d.reduce_gemm._tuned` spells it, and for the same reason: a per-row knob whose default carries the rule. """ - return ConvConfig(BLOCK_M=bm, BLOCK_N=bn, BLOCK_K=bk, GROUP_M=group_m, - num_warps=warps, num_stages=2, matrix_instr_nonkdim=nk, - kpack=1 if bk <= 16 else 2) + return ConvConfig( + BLOCK_M=bm, + BLOCK_N=bn, + BLOCK_K=bk, + GROUP_M=group_m, + num_warps=warps, + num_stages=2, + matrix_instr_nonkdim=nk, + kpack=1 if bk <= 16 else 2, + ) #: Measured winners, keyed by ``(dtype, Cin, Cout, kernel)``; a miss falls back to @@ -828,8 +886,14 @@ def register_tuned(dtype, cin, cout, kernel, config: ConvConfig) -> None: def select_config( - m: int, cin: int, cout: int, kernel: Sequence[int], dtype: torch.dtype, - *, table: dict | None = None, key: tuple | None = None, + m: int, + cin: int, + cout: int, + kernel: Sequence[int], + dtype: torch.dtype, + *, + table: dict | None = None, + key: tuple | None = None, ) -> ConvConfig: """The config the kernel will run: tuned entry if there is one, else heuristic. @@ -927,9 +991,11 @@ def _weight_plan(w: torch.Tensor) -> tuple[int, int, int, int] | None: st = s[2] else: st = 0 # one tap: ``dij`` is always 0, so any stride is the right one - if ((kw > 1 and s[4] != st) - or (kh > 1 and s[3] != st * kw) - or (kd > 1 and s[2] != st * kw * kh)): + if ( + (kw > 1 and s[4] != st) + or (kh > 1 and s[3] != st * kw) + or (kd > 1 and s[2] != st * kw * kh) + ): return None if cout == 1 or s[0] == 1: return (_W_N_CONTIG, st, s[1], 1) @@ -996,9 +1062,14 @@ def is_supported( # memory there becomes the bias, ``nan`` if you are lucky -- and a # stride-2 view of the right length silently applies every other value. # ``torch.conv3d`` rejects both; so does this. - if (bias.dim() != 1 or int(bias.shape[0]) != int(w.shape[0]) - or bias.dtype != x.dtype or not bias.is_cuda - or bias.device != x.device or bias.stride(0) != 1): + if ( + bias.dim() != 1 + or int(bias.shape[0]) != int(w.shape[0]) + or bias.dtype != x.dtype + or not bias.is_cuda + or bias.device != x.device + or bias.stride(0) != 1 + ): return False if x.shape[1] != w.shape[1]: return False @@ -1110,9 +1181,7 @@ def is_supported_all( ) -def _check_out( - y: torch.Tensor, shape: tuple[int, ...], like: torch.Tensor -) -> None: +def _check_out(y: torch.Tensor, shape: tuple[int, ...], like: torch.Tensor) -> None: """Reject an ``out=`` the kernel would write outside of, or write wrongly. Nothing downstream catches either failure. The grid is sized from the @@ -1256,7 +1325,9 @@ def conv3d_forward( # and then copies the whole thing: 2.82 ms against 0.012 ms on a 256 MiB # output, 235x, on a path a training step takes about 19 times. y = torch.empty( - y_shape, device=x.device, dtype=x.dtype, + y_shape, + device=x.device, + dtype=x.dtype, memory_format=torch.channels_last_3d, ) else: @@ -1286,30 +1357,58 @@ def conv3d_forward( index_dtype = _index_dtype(x, y, wr) block_k_count = triton.cdiv(cin, config.BLOCK_K) - grid = ( - triton.cdiv(m_total, config.BLOCK_M) * triton.cdiv(cout, config.BLOCK_N), - ) + grid = (triton.cdiv(m_total, config.BLOCK_M) * triton.cdiv(cout, config.BLOCK_N),) _conv3d_fwd_kernel[grid]( - x, wr, y, bias, - n, in_d, in_h, in_w, - out_d, out_h, out_w, - cin, cout, m_total, - x.stride(0), x.stride(2), x.stride(3), x.stride(4), - plan[1], plan[2], plan[3], - y.stride(0), y.stride(2), y.stride(3), y.stride(4), - KD=kd, KH=kh, KW=kw, - SD=sd, SH=sh, SW=sw, - PD=pd, PH=ph, PW=pw, - DD=dd, DH=dh, DW=dw, - BLOCK_M=config.BLOCK_M, BLOCK_N=config.BLOCK_N, BLOCK_K=config.BLOCK_K, - BLOCK_K_COUNT=block_k_count, GROUP_M=config.GROUP_M, + x, + wr, + y, + bias, + n, + in_d, + in_h, + in_w, + out_d, + out_h, + out_w, + cin, + cout, + m_total, + x.stride(0), + x.stride(2), + x.stride(3), + x.stride(4), + plan[1], + plan[2], + plan[3], + y.stride(0), + y.stride(2), + y.stride(3), + y.stride(4), + KD=kd, + KH=kh, + KW=kw, + SD=sd, + SH=sh, + SW=sw, + PD=pd, + PH=ph, + PW=pw, + DD=dd, + DH=dh, + DW=dw, + BLOCK_M=config.BLOCK_M, + BLOCK_N=config.BLOCK_N, + BLOCK_K=config.BLOCK_K, + BLOCK_K_COUNT=block_k_count, + GROUP_M=config.GROUP_M, HAS_BIAS=bias is not None, EVEN_K=(cin % config.BLOCK_K == 0), EVEN_N=(cout % config.BLOCK_N == 0), PADDED=(pd > 0 or ph > 0 or pw > 0), INDEX_DTYPE=index_dtype, INPUT_PRECISION="ieee", - W_ORDER=plan[0], W_FLIP=bool(weight_flip), + W_ORDER=plan[0], + W_FLIP=bool(weight_flip), **config.launch_kwargs(), ) return y @@ -1379,7 +1478,11 @@ def verify_isa( gy.shape, cin, k, torch.bfloat16, padding=padding ) y = conv3d_backward_data( - gy, w, (n, cin, d, h, wd), padding=padding, config=cfg, + gy, + w, + (n, cin, d, h, wd), + padding=padding, + config=cfg, weight_rsck=rsck, ) big = gy diff --git a/triton_conv3d/reduce_gemm.py b/triton_conv3d/reduce_gemm.py index ac01440..cd4991e 100644 --- a/triton_conv3d/reduce_gemm.py +++ b/triton_conv3d/reduce_gemm.py @@ -135,9 +135,9 @@ import triton.language as tl from .gather_gemm import ( - ConvConfig, _LDS_BYTES, _MFMA_KDIM, + ConvConfig, _pow2_at_most, _triple, tune_key, @@ -164,29 +164,65 @@ @triton.jit def _conv3d_bwd_weight_kernel( - X, GY, OUT, + X, + GY, + OUT, # Sizes. ``K_TOTAL`` is ``BATCH * OUT_D * OUT_H * OUT_W``, the reduction # length; ``K_CHUNK`` is how much of it one split owns. - IN_D, IN_H, IN_W, - OUT_D, OUT_H, OUT_W, - CIN, COUT, K_TOTAL, K_CHUNK, GRID, + IN_D, + IN_H, + IN_W, + OUT_D, + OUT_H, + OUT_W, + CIN, + COUT, + K_TOTAL, + K_CHUNK, + GRID, # Element strides. The channel stride of X and GY is 1 by construction -- # that is what NDHWC means -- so it is neither passed nor multiplied by. - stride_xn, stride_xd, stride_xh, stride_xw, - stride_gn, stride_gd, stride_gh, stride_gw, + stride_xn, + stride_xd, + stride_xh, + stride_xw, + stride_gn, + stride_gd, + stride_gh, + stride_gw, # Destination: ``[split][Cout][tap][Cin]`` with the last three contiguous, # so one output channel's whole gradient is ``stride_wo`` long. - stride_ws, stride_wo, - NUM_M: tl.constexpr, NUM_CI: tl.constexpr, NUM_TG: tl.constexpr, - TAPS: tl.constexpr, TAP_BLOCK: tl.constexpr, BLOCK_NC: tl.constexpr, - KD: tl.constexpr, KH: tl.constexpr, KW: tl.constexpr, - SD: tl.constexpr, SH: tl.constexpr, SW: tl.constexpr, - PD: tl.constexpr, PH: tl.constexpr, PW: tl.constexpr, - DD: tl.constexpr, DH: tl.constexpr, DW: tl.constexpr, - BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, - EVEN_M: tl.constexpr, EVEN_N: tl.constexpr, EVEN_K: tl.constexpr, - PADDED: tl.constexpr, ROW_ALIGNED: tl.constexpr, ATOMIC: tl.constexpr, - NUM_XCD: tl.constexpr, INDEX_DTYPE: tl.constexpr, + stride_ws, + stride_wo, + NUM_M: tl.constexpr, + NUM_CI: tl.constexpr, + NUM_TG: tl.constexpr, + TAPS: tl.constexpr, + TAP_BLOCK: tl.constexpr, + BLOCK_NC: tl.constexpr, + KD: tl.constexpr, + KH: tl.constexpr, + KW: tl.constexpr, + SD: tl.constexpr, + SH: tl.constexpr, + SW: tl.constexpr, + PD: tl.constexpr, + PH: tl.constexpr, + PW: tl.constexpr, + DD: tl.constexpr, + DH: tl.constexpr, + DW: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + EVEN_M: tl.constexpr, + EVEN_N: tl.constexpr, + EVEN_K: tl.constexpr, + PADDED: tl.constexpr, + ROW_ALIGNED: tl.constexpr, + ATOMIC: tl.constexpr, + NUM_XCD: tl.constexpr, + INDEX_DTYPE: tl.constexpr, INPUT_PRECISION: tl.constexpr, ): # -- which tile of dW, and which slice of the reduction ------------------ @@ -209,8 +245,10 @@ def _conv3d_bwd_weight_kernel( rem = GRID % NUM_XCD xcd = pid % NUM_XCD seq = pid // NUM_XCD - pid = tl.where(xcd < rem, xcd * (per + 1), - rem * (per + 1) + (xcd - rem) * per) + seq + pid = ( + tl.where(xcd < rem, xcd * (per + 1), rem * (per + 1) + (xcd - rem) * per) + + seq + ) num_tiles = NUM_M * NUM_CI * NUM_TG split = pid // num_tiles tile = pid % num_tiles @@ -219,7 +257,7 @@ def _conv3d_bwd_weight_kernel( pid_ci = rest % NUM_CI tap0 = (rest // NUM_CI) * TAP_BLOCK - offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) # output channels + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) # output channels # -- the N axis: TAP_BLOCK taps x BLOCK_NC input channels --------------- # @@ -228,7 +266,7 @@ def _conv3d_bwd_weight_kernel( # constexpr, so the divisions fold away. col = tl.arange(0, BLOCK_N) t_local = col // BLOCK_NC - offs_n = pid_ci * BLOCK_NC + (col % BLOCK_NC) # input channels + offs_n = pid_ci * BLOCK_NC + (col % BLOCK_NC) # input channels # ``tl.arange`` needs a power of two, so ``TAP_BLOCK`` is one and cannot # divide 27. Rather than mask the load for the ragged last group -- which # would cost a predicate on every B tile of every group -- the tap is @@ -332,9 +370,12 @@ def _conv3d_bwd_weight_kernel( in_h = src_h[:, None] + (kh * DH)[None, :] in_w = src_w[:, None] + (kw * DW)[None, :] mask_b = ( - (in_d >= 0) & (in_d < IN_D) - & (in_h >= 0) & (in_h < IN_H) - & (in_w >= 0) & (in_w < IN_W) + (in_d >= 0) + & (in_d < IN_D) + & (in_h >= 0) + & (in_h < IN_H) + & (in_w >= 0) + & (in_w < IN_W) ) if not EVEN_N: mask_b = mask_b & (offs_n < CIN)[None, :] @@ -385,8 +426,12 @@ def _conv3d_bwd_weight_kernel( @triton.jit def _reduce_partials_kernel( - PARTIAL, OUT, N_ELEM, SPLITS, - BLOCK: tl.constexpr, BLOCK_S: tl.constexpr, + PARTIAL, + OUT, + N_ELEM, + SPLITS, + BLOCK: tl.constexpr, + BLOCK_S: tl.constexpr, ): """Sum the split-K partials in index order and cast to the output dtype. @@ -415,7 +460,8 @@ def _reduce_partials_kernel( s = s0 + offs_s tile = tl.load( PARTIAL + s.to(tl.int64)[:, None] * stride + base[None, :], - mask=(s < SPLITS)[:, None] & mask[None, :], other=0.0, + mask=(s < SPLITS)[:, None] & mask[None, :], + other=0.0, ) acc += tl.sum(tile, axis=0) tl.store(OUT + offs, acc.to(OUT.dtype.element_ty), mask=mask) @@ -531,7 +577,8 @@ def _fit_bwd_weight_to_lds(cfg: BwdWeightConfig, dtype: torch.dtype) -> BwdWeigh return cfg while cfg.lds_bytes(dtype) > _LDS_BYTES and cfg.BLOCK_K // 2 >= kdim: cfg = dataclasses.replace( - cfg, BLOCK_K=cfg.BLOCK_K // 2, + cfg, + BLOCK_K=cfg.BLOCK_K // 2, kpack=1 if cfg.BLOCK_K // 2 <= 16 else cfg.kpack, ) warps = max(1, min(cfg.num_warps, cfg.BLOCK_M * cfg.BLOCK_N // 256)) @@ -539,8 +586,13 @@ def _fit_bwd_weight_to_lds(cfg: BwdWeightConfig, dtype: torch.dtype) -> BwdWeigh def default_bwd_weight_config( - cout: int, cin: int, kernel: Sequence[int], k_total: int, - dtype: torch.dtype = torch.bfloat16, *, padded: bool = False, + cout: int, + cin: int, + kernel: Sequence[int], + k_total: int, + dtype: torch.dtype = torch.bfloat16, + *, + padded: bool = False, ) -> BwdWeightConfig: """A config that is legal for any shape and close to tuned for most. @@ -619,8 +671,9 @@ def default_bwd_weight_config( # from the LDS budget, which at ``256 + 256`` columns lands on exactly 64 in # bf16 -- i.e. straight into the cliff -- so the cap has to be explicit. # The tuned table is free to ship 64 where a measurement says so. - block_k = _pow2_at_most(_LDS_BYTES // (itemsize * (block_m + block_n)), - 32 if block_m >= 256 else 64) + block_k = _pow2_at_most( + _LDS_BYTES // (itemsize * (block_m + block_n)), 32 if block_m >= 256 else 64 + ) block_k = max(kdim, block_k - block_k % kdim) # A K axis shorter than one tile is not an error, only waste; shrink so the # tiny synthetic shapes do not run a mostly-masked reduction. @@ -628,10 +681,15 @@ def default_bwd_weight_config( block_k //= 2 return _fit_bwd_weight_to_lds( BwdWeightConfig( - BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=block_k, GROUP_M=6, + BLOCK_M=block_m, + BLOCK_N=block_n, + BLOCK_K=block_k, + GROUP_M=6, num_warps=min(8, max(1, block_m * block_n // 256)), - num_stages=2, matrix_instr_nonkdim=nonkdim, - kpack=1 if block_k <= 16 else 2, TAP_BLOCK=tap_block, + num_stages=2, + matrix_instr_nonkdim=nonkdim, + kpack=1 if block_k <= 16 else 2, + TAP_BLOCK=tap_block, ), dtype, ) @@ -643,7 +701,11 @@ def _row_aligned(block_k: int, out_w: int) -> bool: def split_count( - cfg: BwdWeightConfig, cout: int, cin: int, taps: int, k_total: int, + cfg: BwdWeightConfig, + cout: int, + cin: int, + taps: int, + k_total: int, out_w: int, ) -> tuple[int, int]: """``(splits, chunk)``: how the reduction axis is divided, and by how much. @@ -672,13 +734,14 @@ def split_count( row-aligned, because the kernel's cheap scalar unravel needs every K-tile to stay inside one row; otherwise to a whole number of K-tiles. """ - tiles = (-(-cout // cfg.BLOCK_M) - * -(-cin // cfg.BLOCK_NC) - * -(-taps // cfg.TAP_BLOCK)) + tiles = ( + -(-cout // cfg.BLOCK_M) * -(-cin // cfg.BLOCK_NC) * -(-taps // cfg.TAP_BLOCK) + ) k_tiles = -(-k_total // cfg.BLOCK_K) per_split = cout * taps * cin * 4 - ceiling = max(1, min(k_tiles // _MIN_K_TILES_PER_SPLIT, - _WORKSPACE_BYTES // per_split)) + ceiling = max( + 1, min(k_tiles // _MIN_K_TILES_PER_SPLIT, _WORKSPACE_BYTES // per_split) + ) if cfg.SPLIT_K: want = min(cfg.SPLIT_K, ceiling) @@ -693,7 +756,8 @@ def split_count( loop_elems = tiles * k_total * (cfg.BLOCK_M + cfg.BLOCK_N) epilogue_elems = cout * taps * cin * 2 * 2 epilogue_bound = max( - 1, loop_elems // (_MAX_EPILOGUE_FRACTION * max(1, epilogue_elems))) + 1, loop_elems // (_MAX_EPILOGUE_FRACTION * max(1, epilogue_elems)) + ) want = max(1, min(want, epilogue_bound, ceiling)) # The snap goes last and **outranks the epilogue bound**, which is a # deliberate ordering and not the oversight it looks like. ``round`` @@ -773,9 +837,14 @@ def split_count( def candidate_bwd_weight_configs( - cout: int, cin: int, kernel: Sequence[int], k_total: int, + cout: int, + cin: int, + kernel: Sequence[int], + k_total: int, dtype: torch.dtype = torch.bfloat16, - *, splits: Sequence[int] = _SEED_SPLITS, padded: bool = False, + *, + splits: Sequence[int] = _SEED_SPLITS, + padded: bool = False, ) -> list[BwdWeightConfig]: """Configs worth timing for one shape, already pruned to legal ones. @@ -798,33 +867,52 @@ def candidate_bwd_weight_configs( for warps in {4, 8, seed_warps}: for sk in splits: cfg = BwdWeightConfig( - BLOCK_M=bm, BLOCK_N=bnc * tb, BLOCK_K=bk, GROUP_M=6, - num_warps=warps, num_stages=2, matrix_instr_nonkdim=16, - kpack=1 if bk <= 16 else 2, SPLIT_K=sk, TAP_BLOCK=tb, + BLOCK_M=bm, + BLOCK_N=bnc * tb, + BLOCK_K=bk, + GROUP_M=6, + num_warps=warps, + num_stages=2, + matrix_instr_nonkdim=16, + kpack=1 if bk <= 16 else 2, + SPLIT_K=sk, + TAP_BLOCK=tb, ) - if (cfg.validate(dtype) is not None - or cfg.lds_bytes(dtype) > _LDS_BYTES - or cfg in seen): + if ( + cfg.validate(dtype) is not None + or cfg.lds_bytes(dtype) > _LDS_BYTES + or cfg in seen + ): continue seen.add(cfg) out.append(cfg) if not out: - out.append(default_bwd_weight_config(cout, cin, kernel, k_total, dtype, - padded=padded)) + out.append( + default_bwd_weight_config(cout, cin, kernel, k_total, dtype, padded=padded) + ) return out -def _tuned(bm: int, bnc: int, tb: int, bk: int, warps: int, - sk: int = 0, nk: int = 32) -> BwdWeightConfig: +def _tuned( + bm: int, bnc: int, tb: int, bk: int, warps: int, sk: int = 0, nk: int = 32 +) -> BwdWeightConfig: """One measured row. ``nk`` defaults to **32 here and to 16 everywhere else in the package**, and that asymmetry is the whole point of it. See :data:`_TUNED_BWD_W`. """ - return BwdWeightConfig(BLOCK_M=bm, BLOCK_N=bnc * tb, BLOCK_K=bk, GROUP_M=6, - num_warps=warps, num_stages=2, - matrix_instr_nonkdim=nk, - kpack=1 if bk <= 16 else 2, SPLIT_K=sk, TAP_BLOCK=tb) + return BwdWeightConfig( + BLOCK_M=bm, + BLOCK_N=bnc * tb, + BLOCK_K=bk, + GROUP_M=6, + num_warps=warps, + num_stages=2, + matrix_instr_nonkdim=nk, + kpack=1 if bk <= 16 else 2, + SPLIT_K=sk, + TAP_BLOCK=tb, + ) #: Measured backward-weight winners, keyed by ``(dtype, Cin, Cout, kernel)`` -- @@ -1055,8 +1143,13 @@ def register_tuned_bwd_weight(dtype, cin, cout, kernel, config) -> None: def bwd_weight_config( - cout: int, cin: int, kernel: Sequence[int], k_total: int, - dtype: torch.dtype = torch.bfloat16, *, padded: bool = False, + cout: int, + cin: int, + kernel: Sequence[int], + k_total: int, + dtype: torch.dtype = torch.bfloat16, + *, + padded: bool = False, ) -> BwdWeightConfig: """The config :func:`conv3d_backward_weight` would pick for this problem. @@ -1137,9 +1230,7 @@ def bwd_weight_config( # --------------------------------------------------------------------------- -def workspace_elements( - splits: int, cout: int, cin: int, kernel: Sequence[int] -) -> int: +def workspace_elements(splits: int, cout: int, cin: int, kernel: Sequence[int]) -> int: """fp32 elements :func:`conv3d_backward_weight` needs for ``splits`` splits.""" return splits * cout * cin * math.prod(_triple(kernel, "kernel")) @@ -1165,14 +1256,20 @@ def grad_weight_empty( """ k = _triple(kernel, "kernel") return torch.empty( - (cout, cin, *k), dtype=dtype, device=device, + (cout, cin, *k), + dtype=dtype, + device=device, memory_format=torch.channels_last_3d, ) def _validate_out( - gw: torch.Tensor, cout: int, cin: int, k: tuple[int, int, int], - dtype: torch.dtype, device: torch.device, + gw: torch.Tensor, + cout: int, + cin: int, + k: tuple[int, int, int], + dtype: torch.dtype, + device: torch.device, ) -> str | None: """``None`` if ``gw`` can be written as this problem's gradient, else why not. @@ -1265,8 +1362,11 @@ def is_supported_bwd_weight( # clause is in ``gather_gemm.is_supported``; the two gates are kept symmetric # deliberately, because the caller picks between them by direction and a hole # in one of them is a hole in the ladder. - if (not input.is_cuda or not grad_output.is_cuda - or grad_output.device != input.device): + if ( + not input.is_cuda + or not grad_output.is_cuda + or grad_output.device != input.device + ): return False try: s = _triple(stride, "stride") @@ -1326,8 +1426,9 @@ def conv3d_backward_weight( ``weight_rsck`` hoists the weight transform in the other two directions. :func:`split_count` and :func:`workspace_elements` say how big it must be. """ - if not is_supported_bwd_weight(input, weight_shape, grad_output, stride, - padding, dilation, groups): + if not is_supported_bwd_weight( + input, weight_shape, grad_output, stride, padding, dilation, groups + ): raise NotImplementedError( f"unsupported: input={tuple(input.shape)}/{input.dtype} " f"weight_shape={tuple(weight_shape)} " @@ -1413,20 +1514,51 @@ def conv3d_backward_weight( index_dtype = tl.int64 if big else tl.int32 _conv3d_bwd_weight_kernel[grid]( - x, gy, dest, - in_d, in_h, in_w, - out_d, out_h, out_w, - cin, cout, k_total, chunk, grid[0], - x.stride(0), x.stride(2), x.stride(3), x.stride(4), - gy.stride(0), gy.stride(2), gy.stride(3), gy.stride(4), - stride_ws, taps * cin, - NUM_M=num_m, NUM_CI=num_ci, NUM_TG=num_tg, - TAPS=taps, TAP_BLOCK=config.TAP_BLOCK, BLOCK_NC=config.BLOCK_NC, - KD=kd, KH=kh, KW=kw, - SD=sd, SH=sh, SW=sw, - PD=pd, PH=ph, PW=pw, - DD=dd, DH=dh, DW=dw, - BLOCK_M=config.BLOCK_M, BLOCK_N=config.BLOCK_N, BLOCK_K=config.BLOCK_K, + x, + gy, + dest, + in_d, + in_h, + in_w, + out_d, + out_h, + out_w, + cin, + cout, + k_total, + chunk, + grid[0], + x.stride(0), + x.stride(2), + x.stride(3), + x.stride(4), + gy.stride(0), + gy.stride(2), + gy.stride(3), + gy.stride(4), + stride_ws, + taps * cin, + NUM_M=num_m, + NUM_CI=num_ci, + NUM_TG=num_tg, + TAPS=taps, + TAP_BLOCK=config.TAP_BLOCK, + BLOCK_NC=config.BLOCK_NC, + KD=kd, + KH=kh, + KW=kw, + SD=sd, + SH=sh, + SW=sw, + PD=pd, + PH=ph, + PW=pw, + DD=dd, + DH=dh, + DW=dw, + BLOCK_M=config.BLOCK_M, + BLOCK_N=config.BLOCK_N, + BLOCK_K=config.BLOCK_K, EVEN_M=(cout % config.BLOCK_M == 0), EVEN_N=(cin % config.BLOCK_NC == 0 and taps % config.TAP_BLOCK == 0), EVEN_K=(k_total % config.BLOCK_K == 0 and chunk % config.BLOCK_K == 0), @@ -1443,8 +1575,13 @@ def conv3d_backward_weight( # not one program: at ``Cout=6, k=1`` the gradient is 384 elements. block = min(1024, max(64, triton.next_power_of_2(n_elem))) _reduce_partials_kernel[(triton.cdiv(n_elem, block),)]( - ws, gw, n_elem, 1 if atomic else splits, - BLOCK=block, BLOCK_S=8, num_warps=4, + ws, + gw, + n_elem, + 1 if atomic else splits, + BLOCK=block, + BLOCK_S=8, + num_warps=4, ) return gw @@ -1477,16 +1614,20 @@ def verify_isa_bwd_weight( n, cin, cout, d, h, w = problem_shape or (1, 64, 64, 32, 64, 64) k = (kernel, kernel, kernel) out = tuple(v + 2 * padding - (kernel - 1) for v in (d, h, w)) - x = torch.randn((n, cin, d, h, w), device="cuda", dtype=torch.bfloat16 - ).contiguous(memory_format=torch.channels_last_3d) - gy = torch.randn((n, cout, *out), device="cuda", dtype=torch.bfloat16 - ).contiguous(memory_format=torch.channels_last_3d) + x = torch.randn((n, cin, d, h, w), device="cuda", dtype=torch.bfloat16).contiguous( + memory_format=torch.channels_last_3d + ) + gy = torch.randn((n, cout, *out), device="cuda", dtype=torch.bfloat16).contiguous( + memory_format=torch.channels_last_3d + ) k_total = n * out[0] * out[1] * out[2] - cfg = config or bwd_weight_config(cout, cin, k, k_total, torch.bfloat16, - padded=padding > 0) - splits, chunk = split_count(cfg, cout, cin, kernel ** 3, k_total, out[2]) - gw = conv3d_backward_weight(x, (cout, cin, *k), gy, padding=padding, - config=cfg, deterministic=deterministic) + cfg = config or bwd_weight_config( + cout, cin, k, k_total, torch.bfloat16, padded=padding > 0 + ) + splits, chunk = split_count(cfg, cout, cin, kernel**3, k_total, out[2]) + gw = conv3d_backward_weight( + x, (cout, cin, *k), gy, padding=padding, config=cfg, deterministic=deterministic + ) torch.cuda.synchronize() print( f"ISA-DUMP-CONFIG [bwd-weight] {cfg} cin={cin} cout={cout} " diff --git a/triton_conv3d/reference.py b/triton_conv3d/reference.py index e36cc0e..120fb12 100644 --- a/triton_conv3d/reference.py +++ b/triton_conv3d/reference.py @@ -40,8 +40,12 @@ } #: Mantissa bits, including the implicit leading one. -_MANTISSA_BITS = {torch.float64: 53, torch.float32: 24, torch.bfloat16: 8, - torch.float16: 11} +_MANTISSA_BITS = { + torch.float64: 53, + torch.float32: 24, + torch.bfloat16: 8, + torch.float16: 11, +} def torch_dtype(problem: ConvProblem) -> torch.dtype: @@ -98,22 +102,26 @@ def make_inputs( gen = torch.Generator(device=device).manual_seed(seed) thin = exact and density is not None and density < 1.0 - def draw(shape: tuple[int, ...], offset: int, - activation: bool = False) -> torch.Tensor: + def draw( + shape: tuple[int, ...], offset: int, activation: bool = False + ) -> torch.Tensor: g = torch.Generator(device=device).manual_seed(seed + offset) if exact: # {-1, 0, 1}: products are exact and sums stay small. - t = torch.randint(-1, 2, shape, generator=g, device=device, - dtype=torch.int8).to(dtype) + t = torch.randint( + -1, 2, shape, generator=g, device=device, dtype=torch.int8 + ).to(dtype) if thin and activation: # A separate stream, offset far enough that it cannot collide # with any operand's *value* stream: those are seed + 0..3, and # a mask drawn from one of them would correlate the zeros with # the signs of another tensor. gm = torch.Generator(device=device).manual_seed( - seed + offset + (1 << 20)) - t = t * (torch.rand(shape, generator=gm, device=device) - < density).to(dtype) + seed + offset + (1 << 20) + ) + t = t * (torch.rand(shape, generator=gm, device=device) < density).to( + dtype + ) else: t = torch.randn(shape, generator=g, device=device, dtype=torch.float32) t = t.to(dtype) @@ -124,7 +132,9 @@ def draw(shape: tuple[int, ...], offset: int, out: dict[str, torch.Tensor] = { "input": draw(problem.input_shape, 0, True).contiguous(memory_format=fmt), "weight": draw(problem.weight_shape, 1).contiguous(memory_format=fmt), - "grad_output": draw(problem.output_shape, 2, True).contiguous(memory_format=fmt), + "grad_output": draw(problem.output_shape, 2, True).contiguous( + memory_format=fmt + ), } out["bias"] = draw((problem.cout,), 3) if problem.bias else None return out diff --git a/triton_conv3d/shapes.py b/triton_conv3d/shapes.py index 84fba46..50e90cf 100644 --- a/triton_conv3d/shapes.py +++ b/triton_conv3d/shapes.py @@ -282,10 +282,8 @@ def production_variant(self) -> "ConvProblem": return dataclasses.replace( self, name=f"{self.name}+shard" if self.name else "shard", - spatial=tuple(s + 2 * h - for s, h in zip(self.spatial, self.shard_halo)), - padding=tuple(0 if h else p - for h, p in zip(self.shard_halo, self.padding)), + spatial=tuple(s + 2 * h for s, h in zip(self.spatial, self.shard_halo)), + padding=tuple(0 if h else p for h, p in zip(self.shard_halo, self.padding)), halo=(0, 0, 0), shard_halo=(0, 0, 0), form="adapter", @@ -636,8 +634,16 @@ def edge_cases(include_large: bool = False) -> tuple[ConvProblem, ...]: ConvProblem("cin_tiny", 3, 64, (16, 16, 16), sites=("synthetic",)), ConvProblem("cin_odd", 5, 32, (8, 8, 8), sites=("synthetic",)), ConvProblem("cin_prime", 17, 24, (8, 8, 8), sites=("synthetic",)), - ConvProblem("cout_tiny", 64, 6, (8, 8, 8), (1, 1, 1), padding=(0, 0, 0), - bias=True, sites=("synthetic",)), + ConvProblem( + "cout_tiny", + 64, + 6, + (8, 8, 8), + (1, 1, 1), + padding=(0, 0, 0), + bias=True, + sites=("synthetic",), + ), ConvProblem("cout_odd", 32, 7, (8, 8, 8), sites=("synthetic",)), # Spatial extents that do not divide any plausible tile. ConvProblem("spatial_prime", 32, 32, (13, 13, 13), sites=("synthetic",)), @@ -647,12 +653,28 @@ def edge_cases(include_large: bool = False) -> tuple[ConvProblem, ...]: # Smaller than the kernel in one axis: every tap is masked somewhere. ConvProblem("smaller_than_kernel", 16, 16, (2, 2, 2), sites=("synthetic",)), # Padding variants: unpadded shrinks the output, k=1 removes the gather. - ConvProblem("unpadded", 32, 32, (16, 16, 16), padding=(0, 0, 0), - sites=("synthetic",)), - ConvProblem("pointwise", 64, 6, (16, 16, 16), (1, 1, 1), padding=(0, 0, 0), - bias=True, sites=("synthetic",)), - ConvProblem("kernel_aniso", 32, 32, (8, 8, 8), (1, 3, 3), padding=(0, 1, 1), - sites=("synthetic",)), + ConvProblem( + "unpadded", 32, 32, (16, 16, 16), padding=(0, 0, 0), sites=("synthetic",) + ), + ConvProblem( + "pointwise", + 64, + 6, + (16, 16, 16), + (1, 1, 1), + padding=(0, 0, 0), + bias=True, + sites=("synthetic",), + ), + ConvProblem( + "kernel_aniso", + 32, + 32, + (8, 8, 8), + (1, 3, 3), + padding=(0, 1, 1), + sites=("synthetic",), + ), # The padding a *sharded* ScaFFold convolution actually reaches the # kernel with: a symmetric ``k = 3`` with the split axis halo'd (so # ``p = 0`` there) and H and W still padded. Anisotropic padding under @@ -660,13 +682,24 @@ def edge_cases(include_large: bool = False) -> tuple[ConvProblem, ...]: # ``kernel_aniso`` gets its zero from ``kd = 1``, where the boundary # predicate on D is dead for a different reason -- and it is the form # every k=3 site runs at ``dc_num_shards = (2,1,1)`` or ``(4,1,1)``. - ConvProblem("shard_padded", 32, 32, (8, 8, 8), padding=(0, 1, 1), - sites=("synthetic",)), + ConvProblem( + "shard_padded", 32, 32, (8, 8, 8), padding=(0, 1, 1), sites=("synthetic",) + ), # Batch > 1: ScaFFold never does this, but the M decomposition must. ConvProblem("batched", 32, 32, (8, 8, 8), n=3, sites=("synthetic",)), # The transposed upsample, at a size that is quick to check. - ConvProblem("transposed", 64, 32, (8, 8, 8), (2, 2, 2), (2, 2, 2), (0, 0, 0), - transposed=True, bias=True, sites=("synthetic",)), + ConvProblem( + "transposed", + 64, + 32, + (8, 8, 8), + (2, 2, 2), + (2, 2, 2), + (0, 0, 0), + transposed=True, + bias=True, + sites=("synthetic",), + ), # fp32, for more_determinism and for exact-arithmetic tests. ConvProblem("fp32", 32, 32, (8, 8, 8), dtype="fp32", sites=("synthetic",)), ConvProblem("fp16", 32, 32, (8, 8, 8), dtype="fp16", sites=("synthetic",)), @@ -686,10 +719,22 @@ def edge_cases(include_large: bool = False) -> tuple[ConvProblem, ...]: # the buffer-op byte limit -- see :attr:`ConvProblem.buffer_ops_eligible` # for why that is a different question from this one. cases += [ - ConvProblem("int32_below", 128, 64, (255, 255, 255), large=True, - sites=("synthetic",)), - ConvProblem("int32_above", 128, 64, (258, 258, 258), large=True, - sites=("synthetic",)), + ConvProblem( + "int32_below", + 128, + 64, + (255, 255, 255), + large=True, + sites=("synthetic",), + ), + ConvProblem( + "int32_above", + 128, + 64, + (258, 258, 258), + large=True, + sites=("synthetic",), + ), ] return tuple(cases) @@ -714,8 +759,10 @@ def problems_in_form(form: Form) -> tuple[ConvProblem, ...]: if __name__ == "__main__": # pragma: no cover - a human-readable dump - hdr = (f"{'ms/step':>9} {'logical':38s} {'adapter (production)':46s} " - f"{'AI':>7} {'i64':>4}") + hdr = ( + f"{'ms/step':>9} {'logical':38s} {'adapter (production)':46s} " + f"{'AI':>7} {'i64':>4}" + ) print(hdr) print("-" * len(hdr)) for p in scaffold_corpus(): @@ -725,9 +772,13 @@ def problems_in_form(form: Form) -> tuple[ConvProblem, ...]: f"{p.arithmetic_intensity():7.0f} {'yes' if p.needs_int64 else '':>4}" ) padded = sum(1 for p in production_corpus() if any(p.padding)) - print(f"\n{len(scaffold_corpus())} ScaFFold problems, " - f"{len(edge_cases(include_large=True))} synthetic edge cases, " - f"{len(census_corpus())} measured by census") - print(f"{padded}/{len(production_corpus())} of the production forms are " - f"padded; {sum(1 for p in halo_corpus() if any(p.padding))} of the " - f"DistConv forms are") + print( + f"\n{len(scaffold_corpus())} ScaFFold problems, " + f"{len(edge_cases(include_large=True))} synthetic edge cases, " + f"{len(census_corpus())} measured by census" + ) + print( + f"{padded}/{len(production_corpus())} of the production forms are " + f"padded; {sum(1 for p in halo_corpus() if any(p.padding))} of the " + f"DistConv forms are" + ) diff --git a/triton_conv3d/tests/test_bwd_data.py b/triton_conv3d/tests/test_bwd_data.py index f94c36c..ed64b25 100644 --- a/triton_conv3d/tests/test_bwd_data.py +++ b/triton_conv3d/tests/test_bwd_data.py @@ -88,10 +88,17 @@ def _corpus_channel_pairs() -> list[ConvProblem]: if shard != p.padding and shard != (0, 0, 0): forms.insert(1, (shard, "-shard")) for pad, tag in forms: - out.append(ConvProblem( - f"{p.cin}to{p.cout}{tag}", p.cin, p.cout, (6, 7, 8), p.kernel, - padding=pad, sites=("corpus-pair",), - )) + out.append( + ConvProblem( + f"{p.cin}to{p.cout}{tag}", + p.cin, + p.cout, + (6, 7, 8), + p.kernel, + padding=pad, + sites=("corpus-pair",), + ) + ) return out @@ -102,7 +109,8 @@ def _corpus_channel_pairs() -> list[ConvProblem]: #: fp64. Used only for the fp32 test below: their bf16 references are never #: exactly representable, which is the whole point of the note above. CORPUS_SMALL = [ - p for p in scaffold_corpus() + p + for p in scaffold_corpus() if not p.transposed and math.prod(p.halo_variant.spatial) * max(p.cin, p.cout) <= 1 << 22 ] @@ -115,8 +123,12 @@ def _ids(problems): def _run(problem: ConvProblem, ops: dict, **kwargs) -> torch.Tensor: return conv3d_backward_data( - ops["grad_output"], ops["weight"], problem.input_shape, - problem.stride, problem.padding, **kwargs, + ops["grad_output"], + ops["weight"], + problem.input_shape, + problem.stride, + problem.padding, + **kwargs, ) @@ -134,10 +146,10 @@ def test_the_padding_identity_is_the_one_the_derivation_claims(): ``k=3, p=1``, the right shape and the wrong answer, which does not. """ assert bwd_data_padding(1, 1, 3) == (1, 1, 1) - assert bwd_data_padding(0, 1, 3) == (2, 2, 2) # the halo'd form - assert bwd_data_padding(0, 1, 1) == (0, 0, 0) # k=1: no gather at all + assert bwd_data_padding(0, 1, 3) == (2, 2, 2) # the halo'd form + assert bwd_data_padding(0, 1, 1) == (0, 0, 0) # k=1: no gather at all assert bwd_data_padding((0, 1, 1), 1, (1, 3, 3)) == (0, 1, 1) - assert bwd_data_padding(1, 2, 3) == (3, 3, 3) # dilation widens the reach + assert bwd_data_padding(1, 2, 3) == (3, 3, 3) # dilation widens the reach # And the extent identity: OD + 2p' - dil*(k-1) == ID, for every combination. for k in (1, 2, 3, 5): @@ -172,8 +184,7 @@ def test_flipping_every_tap_axis_is_complementing_the_fused_index(): for d in range(kd): for i in range(kh): for j in range(kw): - flipped = (((kd - 1 - d) * kh + (kh - 1 - i)) * kw - + (kw - 1 - j)) + flipped = ((kd - 1 - d) * kh + (kh - 1 - i)) * kw + (kw - 1 - j) fused = (d * kh + i) * kw + j assert flipped == taps - 1 - fused, (kd, kh, kw, d, i, j) @@ -183,8 +194,7 @@ def test_flipping_every_tap_axis_is_complementing_the_fused_index(): # --------------------------------------------------------------------------- -@pytest.mark.parametrize("problem", EDGE + CORPUS_PAIRS, - ids=_ids(EDGE + CORPUS_PAIRS)) +@pytest.mark.parametrize("problem", EDGE + CORPUS_PAIRS, ids=_ids(EDGE + CORPUS_PAIRS)) def test_selected_config_is_legal_for_every_shape(problem: ConvProblem): """The config picked for backward-data must still reach the matrix core. @@ -194,9 +204,17 @@ def test_selected_config_is_legal_for_every_shape(problem: ConvProblem): than ``BLOCK_N``, and ``BLOCK_K`` is the one with the hard MFMA constraint. """ dtype = reference.torch_dtype(problem) - cfg = bwd_data_config(problem.output_shape, problem.cin, problem.kernel, dtype, - padding=problem.padding, dilation=(1, 1, 1)) - assert cfg.validate(dtype) is None, f"{problem.label}: {cfg} -> {cfg.validate(dtype)}" + cfg = bwd_data_config( + problem.output_shape, + problem.cin, + problem.kernel, + dtype, + padding=problem.padding, + dilation=(1, 1, 1), + ) + assert cfg.validate(dtype) is None, ( + f"{problem.label}: {cfg} -> {cfg.validate(dtype)}" + ) @pytest.mark.parametrize("problem", CORPUS_PAIRS, ids=_ids(CORPUS_PAIRS)) @@ -236,15 +254,17 @@ def test_is_supported_declines_what_the_algebra_cannot_express(): assert is_supported_bwd_data(gy, w, shape, padding=1) assert not is_supported_bwd_data(gy, w, shape, stride=2, padding=1) - assert not is_supported_bwd_data(gy, w, shape, padding=3) # p > dil*(k-1) + assert not is_supported_bwd_data(gy, w, shape, padding=3) # p > dil*(k-1) assert not is_supported_bwd_data(gy, w, shape, padding=1, groups=2) assert not is_supported_bwd_data(gy, w.float(), shape, padding=1) # Cin of the weight must match the gradient being asked for ... assert not is_supported_bwd_data(gy, w, (1, 4, 4, 4, 4), padding=1) # ... Cout of the weight must match grad_output ... assert not is_supported_bwd_data( - gy, torch.empty((4, 8, 3, 3, 3), device="cuda", dtype=torch.bfloat16), - shape, padding=1, + gy, + torch.empty((4, 8, 3, 3, 3), device="cuda", dtype=torch.bfloat16), + shape, + padding=1, ) # ... and grad_output's spatial extent must be the one this problem produces. assert not is_supported_bwd_data(gy, w, (1, 8, 6, 4, 4), padding=1) @@ -463,9 +483,9 @@ def test_every_config_gives_the_same_answer(problem: ConvProblem): m = problem.n * math.prod(problem.spatial) # Effective widths: the reduction is over Cout and the GEMM's N is Cin. cfgs = candidate_configs(m, problem.cout, problem.cin, dtype, group_ms=(6, 8)) - cfgs = list(dict.fromkeys( - cfgs + [default_config(m, problem.cout, problem.cin, dtype)] - )) + cfgs = list( + dict.fromkeys(cfgs + [default_config(m, problem.cout, problem.cin, dtype)]) + ) ran = 0 for cfg in cfgs: try: @@ -485,8 +505,7 @@ def test_every_config_gives_the_same_answer(problem: ConvProblem): @requires_gpu -@pytest.mark.parametrize("problem", EDGE + CORPUS_PAIRS, - ids=_ids(EDGE + CORPUS_PAIRS)) +@pytest.mark.parametrize("problem", EDGE + CORPUS_PAIRS, ids=_ids(EDGE + CORPUS_PAIRS)) def test_no_worse_than_miopen(problem: ConvProblem): """The honest bar at realistic magnitudes, against MIOpen on the same data. @@ -501,8 +520,9 @@ def test_no_worse_than_miopen(problem: ConvProblem): reference.incumbent(problem, ops, "bwd-data"), expected ) actual = _run(problem, ops) - reference.assert_close(actual, expected, problem, "bwd-data", - incumbent_error=incumbent_err) + reference.assert_close( + actual, expected, problem, "bwd-data", incumbent_error=incumbent_err + ) @requires_gpu @@ -555,9 +575,12 @@ def test_out_buffer_is_written_in_place_and_is_validated(): bf16 = torch.bfloat16 with pytest.raises(ValueError): _run(problem, ops, out=torch.empty((1, 16, 2, 2, 2), device="cuda", dtype=bf16)) - with pytest.raises(ValueError): # right shape, NCDHW - _run(problem, ops, - out=torch.empty(problem.input_shape, device="cuda", dtype=bf16)) + with pytest.raises(ValueError): # right shape, NCDHW + _run( + problem, + ops, + out=torch.empty(problem.input_shape, device="cuda", dtype=bf16), + ) with pytest.raises(ValueError): _run(problem, ops, out=torch.empty_like(expected, dtype=torch.float32)) @@ -611,12 +634,11 @@ def test_every_weight_layout_gives_the_same_gradient(): layouts = { "channels_last": w.contiguous(memory_format=torch.channels_last_3d), "contiguous": w.contiguous(), - "rsck_strided": (w.permute(2, 3, 4, 1, 0).contiguous() - .permute(4, 3, 0, 1, 2)), + "rsck_strided": (w.permute(2, 3, 4, 1, 0).contiguous().permute(4, 3, 0, 1, 2)), } ref = _run(problem, ops) for name, wl in layouts.items(): - assert torch.equal(wl, w), name # same values, different strides + assert torch.equal(wl, w), name # same values, different strides got = _run(problem, {**ops, "weight": wl}) assert torch.equal(ref, got), name # And the hoisted buffer, which is a fourth spelling of the same values. @@ -629,8 +651,11 @@ def test_output_is_channels_last_and_matches_torch_grad_shape(): ops = reference.make_inputs(problem, seed=61) gx = _run(problem, ops) ref = torch.nn.grad.conv3d_input( - problem.input_shape, ops["weight"], ops["grad_output"], - stride=problem.stride, padding=problem.padding, + problem.input_shape, + ops["weight"], + ops["grad_output"], + stride=problem.stride, + padding=problem.padding, ) assert gx.shape == ref.shape assert gx.is_contiguous(memory_format=torch.channels_last_3d) diff --git a/triton_conv3d/tests/test_bwd_weight.py b/triton_conv3d/tests/test_bwd_weight.py index d9fe2b8..399c3aa 100644 --- a/triton_conv3d/tests/test_bwd_weight.py +++ b/triton_conv3d/tests/test_bwd_weight.py @@ -98,10 +98,17 @@ def _corpus_channel_pairs() -> list[ConvProblem]: if shard != p.padding and shard != (0, 0, 0): forms.insert(1, (shard, "-shard")) for pad, tag in forms: - out.append(ConvProblem( - f"{p.cin}to{p.cout}{tag}", p.cin, p.cout, (6, 7, 8), p.kernel, - padding=pad, sites=("corpus-pair",), - )) + out.append( + ConvProblem( + f"{p.cin}to{p.cout}{tag}", + p.cin, + p.cout, + (6, 7, 8), + p.kernel, + padding=pad, + sites=("corpus-pair",), + ) + ) return out @@ -111,7 +118,8 @@ def _corpus_channel_pairs() -> list[ConvProblem]: #: Real ScaFFold shapes, at their real volumes, small enough to reference in #: fp64. Used only for the fp32 test below. CORPUS_SMALL = [ - p for p in scaffold_corpus() + p + for p in scaffold_corpus() if not p.transposed and math.prod(p.halo_variant.spatial) * max(p.cin, p.cout) <= 1 << 22 ] @@ -124,8 +132,12 @@ def _ids(problems): def _run(problem: ConvProblem, ops: dict, **kwargs) -> torch.Tensor: return conv3d_backward_weight( - ops["input"], problem.weight_shape, ops["grad_output"], - problem.stride, problem.padding, **kwargs, + ops["input"], + problem.weight_shape, + ops["grad_output"], + problem.stride, + problem.padding, + **kwargs, ) @@ -177,18 +189,25 @@ def test_the_grid_lands_on_whole_waves(): than left to the constant that happens to produce it. """ for cout, cin, taps, k_total, out_w in ( - (64, 64, 27, 8_388_608, 256), (64, 128, 27, 2_097_152, 128), - (128, 128, 27, 1_048_576, 128), (256, 256, 27, 131_072, 64), - (512, 512, 27, 16_384, 32), (6, 64, 1, 2_097_152, 256), + (64, 64, 27, 8_388_608, 256), + (64, 128, 27, 2_097_152, 128), + (128, 128, 27, 1_048_576, 128), + (256, 256, 27, 131_072, 64), + (512, 512, 27, 16_384, 32), + (6, 64, 1, 2_097_152, 256), ): - cfg = bwd_weight_config(cout, cin, (3, 3, 3) if taps > 1 else (1, 1, 1), - k_total, torch.bfloat16) + cfg = bwd_weight_config( + cout, cin, (3, 3, 3) if taps > 1 else (1, 1, 1), k_total, torch.bfloat16 + ) splits, _ = split_count(cfg, cout, cin, taps, k_total, out_w) - tiles = (-(-cout // cfg.BLOCK_M) * -(-cin // cfg.BLOCK_NC) - * -(-taps // cfg.TAP_BLOCK)) + tiles = ( + -(-cout // cfg.BLOCK_M) + * -(-cin // cfg.BLOCK_NC) + * -(-taps // cfg.TAP_BLOCK) + ) progs = tiles * splits if progs <= _CU_COUNT: - continue # one wave or less: nothing to snap + continue # one wave or less: nothing to snap waste = (-(-progs // _CU_COUNT) * _CU_COUNT - progs) / progs assert waste < 0.10, (cout, cin, k_total, splits, tiles, progs, waste) assert _SPLIT_TARGET_WAVES >= 1 @@ -230,7 +249,7 @@ def _every_form(problems): seen: set[str] = set() for p in problems: if p.transposed: - continue # a later milestone; ``weight_shape`` is swapped too + continue # a later milestone; ``weight_shape`` is swapped too for q in (p, p.production_variant, p.halo_variant): if q.qualified_label in seen: continue @@ -275,15 +294,27 @@ def test_the_partial_workspace_is_bounded_across_the_whole_corpus(): # set than it was, and it is still the number to size an allocation # from. for cfg in ( - bwd_weight_config(hp.cout, hp.cin, hp.kernel, k_total, - torch.bfloat16, padded=any(hp.padding)), - default_bwd_weight_config(hp.cout, hp.cin, hp.kernel, k_total, - torch.bfloat16, padded=any(hp.padding)), + bwd_weight_config( + hp.cout, + hp.cin, + hp.kernel, + k_total, + torch.bfloat16, + padded=any(hp.padding), + ), + default_bwd_weight_config( + hp.cout, + hp.cin, + hp.kernel, + k_total, + torch.bfloat16, + padded=any(hp.padding), + ), ): - splits, _ = split_count(cfg, hp.cout, hp.cin, hp.tap_count, k_total, - hp.out_spatial[2]) - mib = workspace_elements(splits, hp.cout, hp.cin, - hp.kernel) * 4 / 2**20 + splits, _ = split_count( + cfg, hp.cout, hp.cin, hp.tap_count, k_total, hp.out_spatial[2] + ) + mib = workspace_elements(splits, hp.cout, hp.cin, hp.kernel) * 4 / 2**20 assert mib <= _WORKSPACE_BYTES / 2**20, f"{hp.label}: {mib:.0f} MiB" if mib > worst: worst, worst_label = mib, f"{hp.label} {cfg}" @@ -324,25 +355,43 @@ def test_the_wave_snap_outranks_the_epilogue_bound_and_only_below_one_wave(): Both halves have failed at some point in this function's history, in opposite directions. """ + def bounds(cfg, cout, cin, taps, k_total): - tiles = (-(-cout // cfg.BLOCK_M) * -(-cin // cfg.BLOCK_NC) - * -(-taps // cfg.TAP_BLOCK)) + tiles = ( + -(-cout // cfg.BLOCK_M) + * -(-cin // cfg.BLOCK_NC) + * -(-taps // cfg.TAP_BLOCK) + ) loop_elems = tiles * k_total * (cfg.BLOCK_M + cfg.BLOCK_N) - epi = max(1, loop_elems - // (_MAX_EPILOGUE_FRACTION * max(1, cout * taps * cin * 4))) + epi = max( + 1, loop_elems // (_MAX_EPILOGUE_FRACTION * max(1, cout * taps * cin * 4)) + ) return tiles, epi checked = overridden = 0 for hp in _every_form(list(scaffold_corpus()) + EDGE): k_total = hp.n * math.prod(hp.out_spatial) for cfg in ( - bwd_weight_config(hp.cout, hp.cin, hp.kernel, k_total, - torch.bfloat16, padded=any(hp.padding)), - default_bwd_weight_config(hp.cout, hp.cin, hp.kernel, k_total, - torch.bfloat16, padded=any(hp.padding)), + bwd_weight_config( + hp.cout, + hp.cin, + hp.kernel, + k_total, + torch.bfloat16, + padded=any(hp.padding), + ), + default_bwd_weight_config( + hp.cout, + hp.cin, + hp.kernel, + k_total, + torch.bfloat16, + padded=any(hp.padding), + ), ): - splits, _ = split_count(cfg, hp.cout, hp.cin, hp.tap_count, k_total, - hp.out_spatial[2]) + splits, _ = split_count( + cfg, hp.cout, hp.cin, hp.tap_count, k_total, hp.out_spatial[2] + ) tiles, epi = bounds(cfg, hp.cout, hp.cin, hp.tap_count, k_total) checked += 1 if splits <= epi: @@ -352,16 +401,14 @@ def bounds(cfg, cout, cin, taps, k_total): # snap does: to the *nearest* whole wave, so at most half a wave of # extra programs -- or one whole wave where the bounded grid does # not fill even one. - assert tiles * splits <= max(_CU_COUNT, - tiles * epi + _CU_COUNT // 2), ( + assert tiles * splits <= max(_CU_COUNT, tiles * epi + _CU_COUNT // 2), ( f"{hp.label} {cfg}: {splits} splits against an epilogue bound " f"of {epi} is {tiles * splits} programs, more than a wave past " f"the bounded grid's {tiles * epi}" ) # And the workspace ceiling still holds, which is the one bound the # snap is *not* allowed to escape. - mib = workspace_elements(splits, hp.cout, hp.cin, - hp.kernel) * 4 / 2**20 + mib = workspace_elements(splits, hp.cout, hp.cin, hp.kernel) * 4 / 2**20 assert mib <= _WORKSPACE_BYTES / 2**20, f"{hp.label}: {mib:.0f} MiB" assert checked > 40 @@ -371,8 +418,7 @@ def bounds(cfg, cout, cin, taps, k_total): ) -@pytest.mark.parametrize("problem", EDGE + CORPUS_PAIRS, - ids=_ids(EDGE + CORPUS_PAIRS)) +@pytest.mark.parametrize("problem", EDGE + CORPUS_PAIRS, ids=_ids(EDGE + CORPUS_PAIRS)) def test_selected_config_is_legal_for_every_shape(problem: ConvProblem): """The config picked for backward-weight must reach the matrix core. @@ -385,14 +431,18 @@ def test_selected_config_is_legal_for_every_shape(problem: ConvProblem): dtype = reference.torch_dtype(problem) hp = problem.halo_variant k_total = hp.n * math.prod(hp.out_spatial) - cfg = bwd_weight_config(hp.cout, hp.cin, hp.kernel, k_total, dtype, - padded=any(hp.padding)) + cfg = bwd_weight_config( + hp.cout, hp.cin, hp.kernel, k_total, dtype, padded=any(hp.padding) + ) assert cfg.validate(dtype) is None, f"{hp.label}: {cfg} -> {cfg.validate(dtype)}" assert cfg.lds_bytes(dtype) <= 64 * 1024, f"{hp.label}: {cfg}" -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32], - ids=["bf16", "fp16", "fp32"]) +@pytest.mark.parametrize( + "dtype", + [torch.bfloat16, torch.float16, torch.float32], + ids=["bf16", "fp16", "fp32"], +) def test_default_config_fits_in_lds_in_every_dtype(dtype): """fp32 operands are twice the bytes, and ``more_determinism`` runs in fp32. @@ -420,8 +470,9 @@ def test_every_backward_weight_candidate_config_is_legal(problem: ConvProblem): """ dtype = reference.torch_dtype(problem) k_total = problem.n * math.prod(problem.out_spatial) - cfgs = candidate_bwd_weight_configs(problem.cout, problem.cin, - problem.kernel, k_total, dtype) + cfgs = candidate_bwd_weight_configs( + problem.cout, problem.cin, problem.kernel, k_total, dtype + ) assert cfgs for cfg in cfgs: assert cfg.validate(dtype) is None, f"{cfg}: {cfg.validate(dtype)}" @@ -493,7 +544,9 @@ def test_is_supported_declines_what_the_kernel_cannot_express(): # reduction would run over a volume the input does not have. assert not is_supported_bwd_weight(x, ws, gy, padding=0) assert not is_supported_bwd_weight( - x, ws, torch.empty((1, 8, 4, 6, 6), device="cuda", dtype=torch.bfloat16), + x, + ws, + torch.empty((1, 8, 4, 6, 6), device="cuda", dtype=torch.bfloat16), padding=1, ) @@ -509,7 +562,10 @@ def test_unsupported_calls_raise_rather_than_return_garbage(): # An out= in the wrong layout is refused rather than filled transposed. with pytest.raises(ValueError): conv3d_backward_weight( - x, (8, 8, 3, 3, 3), gy, padding=1, + x, + (8, 8, 3, 3, 3), + gy, + padding=1, out=torch.empty((8, 8, 3, 3, 3), device="cuda", dtype=torch.bfloat16), ) @@ -646,8 +702,9 @@ def test_the_padded_row_aligned_corner_is_compiled_and_correct(problem): hp = problem k_total = hp.n * math.prod(hp.out_spatial) dtype = reference.torch_dtype(problem) - cfg = bwd_weight_config(hp.cout, hp.cin, hp.kernel, k_total, dtype, - padded=any(hp.padding)) + cfg = bwd_weight_config( + hp.cout, hp.cin, hp.kernel, k_total, dtype, padded=any(hp.padding) + ) # The two constexprs, asserted rather than hoped for: this test's whole # value is that it compiles a branch, so it has to fail loudly if a config # change ever stops it reaching that branch. @@ -670,7 +727,9 @@ def test_the_padded_row_aligned_corner_is_compiled_and_correct(problem): #: than listed, so a retune moves this set instead of stranding it. def _tap_widened_pairs() -> list[tuple[ConvProblem, BwdWeightConfig]]: from triton_conv3d.reduce_gemm import ( - _TUNED_BWD_W, _fit_bwd_weight_to_lds, tune_key, + _TUNED_BWD_W, + _fit_bwd_weight_to_lds, + tune_key, ) out, seen = [], set() @@ -678,8 +737,7 @@ def _tap_widened_pairs() -> list[tuple[ConvProblem, BwdWeightConfig]]: if p.transposed or (p.cin, p.cout, p.kernel) in seen: continue seen.add((p.cin, p.cout, p.kernel)) - row = _TUNED_BWD_W.get( - tune_key(torch.bfloat16, p.cin, p.cout, tuple(p.kernel))) + row = _TUNED_BWD_W.get(tune_key(torch.bfloat16, p.cin, p.cout, tuple(p.kernel))) if row is not None and row.TAP_BLOCK > 1: out.append((p, _fit_bwd_weight_to_lds(row, torch.bfloat16))) return out @@ -719,10 +777,12 @@ def test_a_tuned_tap_block_row_survives_the_padding(): ) for p, row in _TAP_WIDENED: k_total = p.n * math.prod(p.out_spatial) - padded = bwd_weight_config(p.cout, p.cin, p.kernel, k_total, - torch.bfloat16, padded=True) - unpadded = bwd_weight_config(p.cout, p.cin, p.kernel, k_total, - torch.bfloat16, padded=False) + padded = bwd_weight_config( + p.cout, p.cin, p.kernel, k_total, torch.bfloat16, padded=True + ) + unpadded = bwd_weight_config( + p.cout, p.cin, p.kernel, k_total, torch.bfloat16, padded=False + ) assert unpadded == row, ( f"{p.cin}->{p.cout}: the tuned row is not selected even unpadded" ) @@ -747,10 +807,12 @@ def test_the_heuristic_widens_tap_block_under_padding_too(): """ for p, _row in _TAP_WIDENED: k_total = p.n * math.prod(p.out_spatial) - wide = default_bwd_weight_config(p.cout, p.cin, p.kernel, k_total, - torch.bfloat16, padded=False) - padded = default_bwd_weight_config(p.cout, p.cin, p.kernel, k_total, - torch.bfloat16, padded=True) + wide = default_bwd_weight_config( + p.cout, p.cin, p.kernel, k_total, torch.bfloat16, padded=False + ) + padded = default_bwd_weight_config( + p.cout, p.cin, p.kernel, k_total, torch.bfloat16, padded=True + ) assert padded == wide, ( f"{p.cin}->{p.cout}: the heuristic still answers differently under " f"padding ({padded} vs {wide})" @@ -764,9 +826,21 @@ def test_the_heuristic_widens_tap_block_under_padding_too(): @requires_gpu @pytest.mark.parametrize( "problem,cfg", - [(ConvProblem(f"{p.cin}to{p.cout}-padded", p.cin, p.cout, (6, 7, 8), - p.kernel, padding=p.padding, sites=("tap-widened",)), c) - for p, c in _TAP_WIDENED], + [ + ( + ConvProblem( + f"{p.cin}to{p.cout}-padded", + p.cin, + p.cout, + (6, 7, 8), + p.kernel, + padding=p.padding, + sites=("tap-widened",), + ), + c, + ) + for p, c in _TAP_WIDENED + ], ids=[f"{p.cin}to{p.cout}" for p, _ in _TAP_WIDENED], ) def test_a_padded_tap_block_row_is_still_bitwise_correct(problem, cfg): @@ -814,20 +888,24 @@ def test_a_padded_tap_block_row_is_still_bitwise_correct(problem, cfg): # The sharded production padding, which is anisotropic: the ``d`` half of # the boundary predicate is dead and the ``h``/``w`` halves are live, inside # the branch where ``src_d`` is a rank-0 scalar. - (ConvProblem("triple-shardpad", 32, 32, (4, 4, 32), padding=(0, 1, 1)), - 4, 32, 32), + (ConvProblem("triple-shardpad", 32, 32, (4, 4, 32), padding=(0, 1, 1)), 4, 32, 32), # The stem, in both of its production paddings. ``BLOCK_M`` is 32 here # rather than the shipped 64 only because this test fixes it; every other # constexpr is the one the resolver returns. (ConvProblem("quintuple-stem", 3, 64, (2, 2, 64)), 16, 64, 4), - (ConvProblem("quintuple-stem-shardpad", 3, 64, (4, 4, 64), - padding=(0, 1, 1)), 16, 64, 4), + ( + ConvProblem("quintuple-stem-shardpad", 3, 64, (4, 4, 64), padding=(0, 1, 1)), + 16, + 64, + 4, + ), ] @requires_gpu @pytest.mark.parametrize( - "problem,tap_block,block_k,block_nc", _PADDED_ROW_ALIGNED_TAPS, + "problem,tap_block,block_k,block_nc", + _PADDED_ROW_ALIGNED_TAPS, ids=[p.name for p, _, _, _ in _PADDED_ROW_ALIGNED_TAPS], ) def test_the_padded_row_aligned_tap_widened_corner_is_correct( @@ -849,9 +927,15 @@ def test_the_padded_row_aligned_tap_widened_corner_is_correct( reaches; the assertions below say which case is which so a failure names the axis rather than the tile. """ - cfg = BwdWeightConfig(BLOCK_M=32, BLOCK_N=block_nc * tap_block, - BLOCK_K=block_k, TAP_BLOCK=tap_block, num_warps=4, - matrix_instr_nonkdim=16, kpack=1) + cfg = BwdWeightConfig( + BLOCK_M=32, + BLOCK_N=block_nc * tap_block, + BLOCK_K=block_k, + TAP_BLOCK=tap_block, + num_warps=4, + matrix_instr_nonkdim=16, + kpack=1, + ) assert cfg.BLOCK_NC == block_nc assert any(problem.padding), "PADDED would be False" assert _row_aligned(cfg.BLOCK_K, problem.out_spatial[2]), ( @@ -869,8 +953,7 @@ def test_the_padded_row_aligned_tap_widened_corner_is_correct( ops = reference.make_inputs(problem, seed=31, exact=True, dtype=dtype) expected = reference.reference(problem, ops, "bwd-weight") assert reference.is_exactly_representable(expected, dtype) - assert reference.compare(_run(problem, ops, config=cfg), - expected.to(dtype)).bitwise + assert reference.compare(_run(problem, ops, config=cfg), expected.to(dtype)).bitwise @requires_gpu @@ -883,9 +966,9 @@ def test_bitwise_standard_rejects_a_shifted_input(): assert reference.compare(actual, correct).bitwise shifted = torch.roll(ops["input"], shifts=1, dims=-1) - wrong = reference.reference( - problem, {**ops, "input": shifted}, "bwd-weight" - ).to(torch.bfloat16) + wrong = reference.reference(problem, {**ops, "input": shifted}, "bwd-weight").to( + torch.bfloat16 + ) assert not reference.compare(actual, wrong).bitwise, ( "a one-voxel shift of the input produced a bitwise-identical gradient; " "the comparison is not discriminating" @@ -937,7 +1020,11 @@ def test_every_config_gives_the_same_answer(problem: ConvProblem): expected = expected.to(dtype) k_total = problem.n * math.prod(problem.out_spatial) cfgs = candidate_bwd_weight_configs( - problem.cout, problem.cin, problem.kernel, k_total, dtype, + problem.cout, + problem.cin, + problem.kernel, + k_total, + dtype, splits=(0, 1, 3, 64), ) ran = 0 @@ -994,8 +1081,9 @@ def test_the_forward_kernel_can_express_backward_weight(): # (Cin, N, ID, IH, IW) convolved with (Cout, N, OD, OH, OW) -> (Cin, Cout, k) as_conv = conv3d_forward( ops["input"].transpose(0, 1).contiguous(memory_format=torch.channels_last_3d), - ops["grad_output"].transpose(0, 1).contiguous( - memory_format=torch.channels_last_3d), + ops["grad_output"] + .transpose(0, 1) + .contiguous(memory_format=torch.channels_last_3d), padding=0, ) assert tuple(as_conv.shape) == (problem.cin, problem.cout, *problem.kernel) @@ -1021,8 +1109,7 @@ def test_the_forward_kernel_can_express_backward_weight(): @requires_gpu -@pytest.mark.parametrize("problem", EDGE + CORPUS_PAIRS, - ids=_ids(EDGE + CORPUS_PAIRS)) +@pytest.mark.parametrize("problem", EDGE + CORPUS_PAIRS, ids=_ids(EDGE + CORPUS_PAIRS)) def test_no_worse_than_miopen(problem: ConvProblem): """The honest bar at realistic magnitudes, against MIOpen on the same data.""" ops = reference.make_inputs(problem, seed=17) @@ -1031,8 +1118,9 @@ def test_no_worse_than_miopen(problem: ConvProblem): reference.incumbent(problem, ops, "bwd-weight"), expected ) actual = _run(problem, ops) - reference.assert_close(actual, expected, problem, "bwd-weight", - incumbent_error=incumbent_err) + reference.assert_close( + actual, expected, problem, "bwd-weight", incumbent_error=incumbent_err + ) @requires_gpu @@ -1084,13 +1172,13 @@ def test_a_strided_convolution_is_served_correctly(stride, padding): multiply that is *also* not refused by ``is_supported`` is the combination that returns a wrong gradient silently. """ - problem = ConvProblem("strided", 16, 24, (9, 11, 13), - stride=(stride,) * 3, padding=(padding,) * 3) + problem = ConvProblem( + "strided", 16, 24, (9, 11, 13), stride=(stride,) * 3, padding=(padding,) * 3 + ) ops = reference.make_inputs(problem, seed=37, exact=True) expected = reference.reference(problem, ops, "bwd-weight") assert reference.is_exactly_representable(expected, torch.bfloat16) - assert reference.compare(_run(problem, ops), - expected.to(torch.bfloat16)).bitwise + assert reference.compare(_run(problem, ops), expected.to(torch.bfloat16)).bitwise # --------------------------------------------------------------------------- @@ -1108,11 +1196,24 @@ def test_a_strided_convolution_is_served_correctly(stride, padding): def test_repeated_calls_are_bitwise_reproducible_in_process(): problem = _DET ops = reference.make_inputs(problem, seed=71) - cfg = bwd_weight_config(problem.cout, problem.cin, problem.kernel, - math.prod(problem.out_spatial), torch.bfloat16) - assert split_count(cfg, problem.cout, problem.cin, problem.tap_count, - math.prod(problem.out_spatial), - problem.out_spatial[2])[0] > 1, "not exercising split-K" + cfg = bwd_weight_config( + problem.cout, + problem.cin, + problem.kernel, + math.prod(problem.out_spatial), + torch.bfloat16, + ) + assert ( + split_count( + cfg, + problem.cout, + problem.cin, + problem.tap_count, + math.prod(problem.out_spatial), + problem.out_spatial[2], + )[0] + > 1 + ), "not exercising split-K" first = _run(problem, ops) for _ in range(4): assert torch.equal(first, _run(problem, ops)) @@ -1120,8 +1221,9 @@ def test_repeated_calls_are_bitwise_reproducible_in_process(): #: The ``k=1`` segmentation head, at a volume that splits ~800 ways, in **fp32**. #: The dtype is the entire point -- see the negative-control test below. -_DET_K1 = ConvProblem("determinism-k1", 64, 6, (64, 64, 64), (1, 1, 1), - padding=(0, 0, 0), dtype="fp32") +_DET_K1 = ConvProblem( + "determinism-k1", 64, 6, (64, 64, 64), (1, 1, 1), padding=(0, 0, 0), dtype="fp32" +) @requires_gpu @@ -1159,10 +1261,15 @@ def test_the_atomic_path_is_not_bitwise_reproducible(problem: ConvProblem): """ dtype = reference.torch_dtype(problem) k_total = problem.n * math.prod(problem.out_spatial) - cfg = bwd_weight_config(problem.cout, problem.cin, problem.kernel, k_total, - dtype) - splits = split_count(cfg, problem.cout, problem.cin, problem.tap_count, - k_total, problem.out_spatial[2])[0] + cfg = bwd_weight_config(problem.cout, problem.cin, problem.kernel, k_total, dtype) + splits = split_count( + cfg, + problem.cout, + problem.cin, + problem.tap_count, + k_total, + problem.out_spatial[2], + )[0] assert splits > 8, f"{splits} splits: too few writers to contend" ops = reference.make_inputs(problem, seed=71, dtype=dtype) @@ -1208,8 +1315,12 @@ def test_three_separate_processes_agree_bitwise(): repo = str(pathlib.Path(__file__).resolve().parents[2]) digests = [] for _ in range(3): - proc = subprocess.run([sys.executable, "-c", _CHILD.format(repo=repo)], - capture_output=True, text=True, timeout=900) + proc = subprocess.run( + [sys.executable, "-c", _CHILD.format(repo=repo)], + capture_output=True, + text=True, + timeout=900, + ) assert proc.returncode == 0, proc.stderr[-2000:] digests.append(proc.stdout.strip().splitlines()[-1]) assert len(set(digests)) == 1, digests @@ -1233,8 +1344,11 @@ def test_the_output_is_a_channels_last_weight_of_the_right_shape(): ops = reference.make_inputs(problem, seed=61) gw = _run(problem, ops) ref = torch.nn.grad.conv3d_weight( - ops["input"], problem.weight_shape, ops["grad_output"], - stride=problem.stride, padding=problem.padding, + ops["input"], + problem.weight_shape, + ops["grad_output"], + stride=problem.stride, + padding=problem.padding, ) assert gw.shape == ref.shape assert gw.is_contiguous(memory_format=torch.channels_last_3d) @@ -1261,10 +1375,12 @@ def test_an_out_the_kernel_would_overrun_is_refused(): silently changes the dtype of the gradient the caller gets back. """ k = (3, 3, 3) - x = torch.randn((1, 32, 6, 6, 6), device="cuda", dtype=torch.bfloat16 - ).contiguous(memory_format=torch.channels_last_3d) - gy = torch.randn((1, 64, 6, 6, 6), device="cuda", dtype=torch.bfloat16 - ).contiguous(memory_format=torch.channels_last_3d) + x = torch.randn((1, 32, 6, 6, 6), device="cuda", dtype=torch.bfloat16).contiguous( + memory_format=torch.channels_last_3d + ) + gy = torch.randn((1, 64, 6, 6, 6), device="cuda", dtype=torch.bfloat16).contiguous( + memory_format=torch.channels_last_3d + ) ws = (64, 32, *k) right = grad_weight_empty(64, 32, k, dtype=torch.bfloat16, device="cuda") @@ -1276,18 +1392,18 @@ def test_an_out_the_kernel_would_overrun_is_refused(): for bad, what in ( (small, "shape"), - (grad_weight_empty(64, 32, k, dtype=torch.float32, device="cuda"), - "dtype"), - (grad_weight_empty(64, 32, k, dtype=torch.bfloat16, device="cpu"), - "device"), + (grad_weight_empty(64, 32, k, dtype=torch.float32, device="cuda"), "dtype"), + (grad_weight_empty(64, 32, k, dtype=torch.bfloat16, device="cpu"), "device"), (torch.empty(ws, device="cuda", dtype=torch.bfloat16), "strides"), ): with pytest.raises(ValueError, match="out="): conv3d_backward_weight(x, ws, gy, padding=1, out=bad) # ...and the buffer that *is* right is still accepted, so the guard is not # simply refusing everything. - assert conv3d_backward_weight(x, ws, gy, padding=1, - out=right).data_ptr() == right.data_ptr() + assert ( + conv3d_backward_weight(x, ws, gy, padding=1, out=right).data_ptr() + == right.data_ptr() + ) @requires_gpu @@ -1311,8 +1427,7 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None): return func(*args, **(kwargs or {})) with _Record(): - gw = grad_weight_empty(64, 32, (3, 3, 3), dtype=torch.bfloat16, - device="cuda") + gw = grad_weight_empty(64, 32, (3, 3, 3), dtype=torch.bfloat16, device="cuda") assert gw.shape == (64, 32, 3, 3, 3) assert gw.is_contiguous(memory_format=torch.channels_last_3d) assert not [op for op in seen if "copy" in op or "clone" in op], seen @@ -1327,16 +1442,25 @@ def test_hoisted_workspace_and_out_are_equivalent(): inline = _run(problem, ops) k_total = math.prod(problem.out_spatial) - cfg = bwd_weight_config(problem.cout, problem.cin, problem.kernel, k_total, - torch.bfloat16) - splits, _ = split_count(cfg, problem.cout, problem.cin, problem.tap_count, - k_total, problem.out_spatial[2]) + cfg = bwd_weight_config( + problem.cout, problem.cin, problem.kernel, k_total, torch.bfloat16 + ) + splits, _ = split_count( + cfg, + problem.cout, + problem.cin, + problem.tap_count, + k_total, + problem.out_spatial[2], + ) ws = torch.empty( workspace_elements(splits, problem.cout, problem.cin, problem.kernel), - dtype=torch.float32, device="cuda", + dtype=torch.float32, + device="cuda", + ) + gw = grad_weight_empty( + problem.cout, problem.cin, problem.kernel, dtype=torch.bfloat16, device="cuda" ) - gw = grad_weight_empty(problem.cout, problem.cin, problem.kernel, - dtype=torch.bfloat16, device="cuda") hoisted = _run(problem, ops, workspace=ws, out=gw) assert hoisted.data_ptr() == gw.data_ptr() assert torch.equal(inline, hoisted) diff --git a/triton_conv3d/tests/test_gather_gemm.py b/triton_conv3d/tests/test_gather_gemm.py index 55da17c..b5e19aa 100644 --- a/triton_conv3d/tests/test_gather_gemm.py +++ b/triton_conv3d/tests/test_gather_gemm.py @@ -62,7 +62,8 @@ #: regime the exactness question is hardest in. See #: :func:`test_corpus_shapes_match_bitwise`. CORPUS_SMALL = [ - p for p in scaffold_corpus() + p + for p in scaffold_corpus() if not p.transposed and math.prod(p.spatial) * max(p.cin, p.cout) <= 1 << 22 ] @@ -73,8 +74,12 @@ def _ids(problems): def _run(problem: ConvProblem, ops: dict, **kwargs) -> torch.Tensor: return conv3d_forward( - ops["input"], ops["weight"], ops["bias"], - problem.stride, problem.padding, **kwargs, + ops["input"], + ops["weight"], + ops["bias"], + problem.stride, + problem.padding, + **kwargs, ) @@ -137,7 +142,7 @@ def test_the_index_width_decision_covers_every_operand_including_the_weight(): from triton_conv3d.gather_gemm import _index_dtype small = torch.empty((1 << 20,), device="meta") - huge = torch.empty((1 << 31,), device="meta") # numel > 2**31 - 1 by one + huge = torch.empty((1 << 31,), device="meta") # numel > 2**31 - 1 by one assert _index_dtype(small, small, small) == tl.int32 assert _index_dtype(huge, small, small) == tl.int64 assert _index_dtype(small, huge, small) == tl.int64 @@ -153,17 +158,26 @@ def test_block_k_constraint_follows_the_intrinsic_not_the_tile(): briefing's claim that ``BLOCK_K=16`` rows get pruned at nonkdim 16 is simply wrong arithmetic (16 % 16 == 0). """ - ok32 = ConvConfig(BLOCK_M=32, BLOCK_N=32, BLOCK_K=8, matrix_instr_nonkdim=32, - num_warps=4, kpack=1) + ok32 = ConvConfig( + BLOCK_M=32, BLOCK_N=32, BLOCK_K=8, matrix_instr_nonkdim=32, num_warps=4, kpack=1 + ) assert ok32.validate(torch.bfloat16) is None - assert ConvConfig(BLOCK_M=32, BLOCK_N=32, BLOCK_K=8, - matrix_instr_nonkdim=16, num_warps=4).validate(torch.bfloat16) - assert ConvConfig(BLOCK_K=16, matrix_instr_nonkdim=16, - kpack=1).validate(torch.bfloat16) is None + assert ConvConfig( + BLOCK_M=32, BLOCK_N=32, BLOCK_K=8, matrix_instr_nonkdim=16, num_warps=4 + ).validate(torch.bfloat16) + assert ( + ConvConfig(BLOCK_K=16, matrix_instr_nonkdim=16, kpack=1).validate( + torch.bfloat16 + ) + is None + ) -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32], - ids=["bf16", "fp16", "fp32"]) +@pytest.mark.parametrize( + "dtype", + [torch.bfloat16, torch.float16, torch.float32], + ids=["bf16", "fp16", "fp32"], +) def test_default_config_fits_in_lds_in_every_dtype(dtype): """The block sizes were chosen against bf16; fp32 operands are twice the bytes. @@ -185,14 +199,11 @@ def test_default_config_fits_in_lds_in_every_dtype(dtype): cfg = default_config(m, cin, cout, dtype) assert cfg.validate(dtype) is None, f"{cin}->{cout} m={m}: {cfg}" assert cfg.lds_bytes(dtype) <= _LDS_BYTES, ( - f"{cin}->{cout} m={m}: {cfg} needs " - f"{cfg.lds_bytes(dtype)} B of LDS" + f"{cin}->{cout} m={m}: {cfg} needs {cfg.lds_bytes(dtype)} B of LDS" ) -@pytest.mark.parametrize( - "problem", EDGE + CORPUS_SMALL, ids=_ids(EDGE + CORPUS_SMALL) -) +@pytest.mark.parametrize("problem", EDGE + CORPUS_SMALL, ids=_ids(EDGE + CORPUS_SMALL)) def test_default_config_is_legal_for_every_shape(problem: ConvProblem): """The heuristic must never hand back a config that loses the matrix core. @@ -203,7 +214,9 @@ def test_default_config_is_legal_for_every_shape(problem: ConvProblem): dtype = reference.torch_dtype(problem) m = problem.n * math.prod(problem.out_spatial) cfg = default_config(m, problem.cin, problem.cout, dtype) - assert cfg.validate(dtype) is None, f"{problem.label}: {cfg} -> {cfg.validate(dtype)}" + assert cfg.validate(dtype) is None, ( + f"{problem.label}: {cfg} -> {cfg.validate(dtype)}" + ) @pytest.mark.parametrize("problem", CORPUS_SMALL[:6], ids=_ids(CORPUS_SMALL[:6])) @@ -246,8 +259,9 @@ def test_is_supported_declines_what_the_kernel_cannot_do(): assert is_supported(x, w, padding=1) assert not is_supported(x, w, padding=1, groups=2) assert not is_supported(x, w.float(), padding=1) - assert not is_supported(x, torch.empty((8, 4, 3, 3, 3), device="cuda", - dtype=torch.bfloat16), padding=1) + assert not is_supported( + x, torch.empty((8, 4, 3, 3, 3), device="cuda", dtype=torch.bfloat16), padding=1 + ) # A kernel wider than the padded input has no output voxels at all, which # the M-unravel cannot express. tiny = torch.empty((1, 8, 1, 4, 4), device="cuda", dtype=torch.bfloat16) @@ -291,9 +305,10 @@ def test_is_supported_declines_a_bias_torch_itself_rejects(): bias = torch.empty(32, device="cuda", dtype=bf16) assert is_supported(x, w, bias, padding=1) - assert not is_supported(x, w, bias[:4], padding=1) # too short - assert not is_supported(x, w, torch.empty(64, device="cuda", - dtype=bf16)[::2], padding=1) + assert not is_supported(x, w, bias[:4], padding=1) # too short + assert not is_supported( + x, w, torch.empty(64, device="cuda", dtype=bf16)[::2], padding=1 + ) assert not is_supported(x, w, bias.float(), padding=1) assert not is_supported(x, w, bias.cpu(), padding=1) assert not is_supported(x, w, bias.view(1, 32), padding=1) @@ -372,13 +387,13 @@ def test_is_supported_all_is_exactly_the_three_gates_conjoined(): bf16 = torch.bfloat16 cases = [ # (x shape, w shape, kwargs) - ((1, 8, 8, 8, 8), (16, 8, 3, 3, 3), dict(padding=1)), # all yes + ((1, 8, 8, 8, 8), (16, 8, 3, 3, 3), dict(padding=1)), # all yes ((1, 8, 8, 8, 8), (16, 8, 3, 3, 3), dict(stride=2, padding=1)), # bwd-data no ((1, 8, 8, 8, 8), (16, 8, 3, 3, 3), dict(padding=1, groups=2)), # fwd no - ((1, 8, 8, 8, 8), (16, 8, 3, 3, 3), dict(padding=0)), # all yes - ((1, 8, 8, 8, 8), (16, 8, 1, 1, 1), dict(padding=0)), # k=1 + ((1, 8, 8, 8, 8), (16, 8, 3, 3, 3), dict(padding=0)), # all yes + ((1, 8, 8, 8, 8), (16, 8, 1, 1, 1), dict(padding=0)), # k=1 ((1, 8, 8, 8, 8), (16, 8, 3, 3, 3), dict(padding=2)), # p > dil*(k-1) - ((1, 8, 1, 8, 8), (16, 8, 3, 3, 3), dict(padding=1)), # thin D + ((1, 8, 1, 8, 8), (16, 8, 3, 3, 3), dict(padding=1)), # thin D ((1, 8, 8, 8, 8), (16, 8, 3, 3, 3), dict(dilation=2, padding=2)), ] for x_shape, w_shape, kwargs in cases: @@ -436,20 +451,25 @@ def test_is_supported_declines_degenerate_extents(): # A zero-length spatial axis: returned a volume of pure padding, where torch # raises "Only zero batch or zero channel inputs are supported". - assert not is_supported(torch.empty((1, 8, 0, 5, 6), device="cuda", - dtype=bf16), w, padding=2) + assert not is_supported( + torch.empty((1, 8, 0, 5, 6), device="cuda", dtype=bf16), w, padding=2 + ) # A zero-size kernel: ``(in + 2p - d(k-1) - 1)//s + 1`` gains one at k=0, so # the returned output was *larger* than the input. - assert not is_supported(x, torch.empty((16, 8, 0, 0, 0), device="cuda", - dtype=bf16), padding=0) + assert not is_supported( + x, torch.empty((16, 8, 0, 0, 0), device="cuda", dtype=bf16), padding=0 + ) # Cin = 0: returned Cout channels of zeros where torch returns a tensor with # no channels at all -- a different shape, not a different value. - assert not is_supported(torch.empty((1, 0, 4, 5, 6), device="cuda", dtype=bf16), - torch.empty((16, 0, 3, 3, 3), device="cuda", dtype=bf16), - padding=1) + assert not is_supported( + torch.empty((1, 0, 4, 5, 6), device="cuda", dtype=bf16), + torch.empty((16, 0, 3, 3, 3), device="cuda", dtype=bf16), + padding=1, + ) - empty_batch = torch.empty((0, 8, 4, 5, 6), device="cuda", - dtype=bf16).contiguous(memory_format=torch.channels_last_3d) + empty_batch = torch.empty((0, 8, 4, 5, 6), device="cuda", dtype=bf16).contiguous( + memory_format=torch.channels_last_3d + ) assert is_supported(empty_batch, w, padding=1) assert tuple(conv3d_forward(empty_batch, w, padding=1).shape) == (0, 16, 4, 5, 6) @@ -515,9 +535,9 @@ def test_bitwise_standard_rejects_a_shifted_gather(): # Same problem, but the reference gathers from one voxel further along W. shifted = torch.roll(ops["input"], shifts=1, dims=-1) - wrong = reference.reference( - problem, {**ops, "input": shifted}, "fwd" - ).to(torch.bfloat16) + wrong = reference.reference(problem, {**ops, "input": shifted}, "fwd").to( + torch.bfloat16 + ) assert not reference.compare(actual, wrong).bitwise, ( "a one-voxel shift of the input produced a bitwise-identical result; " "the comparison is not discriminating" @@ -654,10 +674,10 @@ def test_every_config_gives_the_same_answer(problem: ConvProblem): m = problem.n * math.prod(problem.out_spatial) # Plus the shipped default, which for a shape too small for any seed tile # (``Cout=6``) is the only candidate there is. - cfgs = candidate_configs(m, problem.cin, problem.cout, dtype, - group_ms=(6, 8)) - cfgs = list(dict.fromkeys(cfgs + [default_config(m, problem.cin, - problem.cout, dtype)])) + cfgs = candidate_configs(m, problem.cin, problem.cout, dtype, group_ms=(6, 8)) + cfgs = list( + dict.fromkeys(cfgs + [default_config(m, problem.cin, problem.cout, dtype)]) + ) ran = 0 for cfg in cfgs: try: @@ -679,8 +699,7 @@ def test_every_config_gives_the_same_answer(problem: ConvProblem): @requires_gpu -@pytest.mark.parametrize("problem", EDGE + CORPUS_SMALL, - ids=_ids(EDGE + CORPUS_SMALL)) +@pytest.mark.parametrize("problem", EDGE + CORPUS_SMALL, ids=_ids(EDGE + CORPUS_SMALL)) def test_no_worse_than_miopen(problem: ConvProblem): """The honest bar for a replacement: not better than MIOpen, but not worse. @@ -696,8 +715,9 @@ def test_no_worse_than_miopen(problem: ConvProblem): reference.incumbent(problem, ops, "fwd"), expected ) actual = _run(problem, ops) - reference.assert_close(actual, expected, problem, "fwd", - incumbent_error=incumbent_err) + reference.assert_close( + actual, expected, problem, "fwd", incumbent_error=incumbent_err + ) @requires_gpu @@ -728,13 +748,15 @@ def test_fp32_accumulates_in_fp32(): @requires_gpu def test_bias_is_added_once_and_broadcast_over_channels(): - problem = ConvProblem("bias", 32, 24, (5, 6, 7), (1, 1, 1), - padding=(0, 0, 0), bias=True) + problem = ConvProblem( + "bias", 32, 24, (5, 6, 7), (1, 1, 1), padding=(0, 0, 0), bias=True + ) ops = reference.make_inputs(problem, seed=31, exact=True) with_bias = _run(problem, ops) - without = conv3d_forward(ops["input"], ops["weight"], None, - problem.stride, problem.padding) - delta = (with_bias.float() - without.float()) + without = conv3d_forward( + ops["input"], ops["weight"], None, problem.stride, problem.padding + ) + delta = with_bias.float() - without.float() # The difference must be exactly the bias, in every voxel. expected = ops["bias"].float().view(1, -1, 1, 1, 1).expand_as(delta) assert torch.equal(delta, expected) @@ -792,7 +814,7 @@ def test_out_buffer_is_written_in_place_and_is_validated(): with pytest.raises(ValueError): _run(problem, ops, out=torch.empty_like(expected, dtype=torch.float32)) with pytest.raises(ValueError): - _run(problem, ops, out=torch.empty(shape, dtype=bf16)) # on the CPU + _run(problem, ops, out=torch.empty(shape, dtype=bf16)) # on the CPU @requires_gpu @@ -806,11 +828,12 @@ def test_the_output_is_allocated_directly_in_channels_last(): invisible in the result, so what pins it is the peak allocation: the wrong form needs two output-sized buffers live at once, the right form needs one. """ - x = torch.randn((1, 64, 64, 64, 64), device="cuda", dtype=torch.bfloat16 - ).contiguous(memory_format=torch.channels_last_3d) + x = torch.randn( + (1, 64, 64, 64, 64), device="cuda", dtype=torch.bfloat16 + ).contiguous(memory_format=torch.channels_last_3d) w = torch.randn((64, 64, 3, 3, 3), device="cuda", dtype=torch.bfloat16) wr = to_rsck(w) - conv3d_forward(x, w, padding=1, weight_rsck=wr) # warm the JIT out of the way + conv3d_forward(x, w, padding=1, weight_rsck=wr) # warm the JIT out of the way torch.cuda.synchronize() torch.cuda.reset_peak_memory_stats() @@ -844,8 +867,9 @@ def test_the_layout_conversion_is_a_no_op_only_where_stride_c_is_moot(): ambiguous = 0 for shape in itertools.product((1, 2, 3), repeat=5): t = torch.empty(shape) - if not (t.is_contiguous() - and t.is_contiguous(memory_format=torch.channels_last_3d)): + if not ( + t.is_contiguous() and t.is_contiguous(memory_format=torch.channels_last_3d) + ): continue ambiguous += 1 assert t.stride(1) == 1 or shape[1] == 1, shape @@ -860,10 +884,12 @@ def test_an_ambiguous_layout_still_gives_the_right_answer(): the result still has to be the reference's -- which it is only because the one stride that differs is the one a single-channel input never uses. """ - x = torch.randint(-1, 2, (1, 1, 4, 5, 6), device="cuda", dtype=torch.int8 - ).to(torch.bfloat16) - w = torch.randint(-1, 2, (8, 1, 3, 3, 3), device="cuda", dtype=torch.int8 - ).to(torch.bfloat16) + x = torch.randint(-1, 2, (1, 1, 4, 5, 6), device="cuda", dtype=torch.int8).to( + torch.bfloat16 + ) + w = torch.randint(-1, 2, (8, 1, 3, 3, 3), device="cuda", dtype=torch.int8).to( + torch.bfloat16 + ) assert x.is_contiguous() assert x.contiguous(memory_format=torch.channels_last_3d).data_ptr() == x.data_ptr() assert torch.equal(conv3d_forward(x, w, padding=1), F.conv3d(x, w, padding=1)) @@ -886,13 +912,14 @@ def test_hoisted_weight_transform_is_validated(): other = torch.randn((24, 16, 1, 1, 1), device="cuda", dtype=torch.bfloat16) with pytest.raises(ValueError): - _run(problem, ops, weight_rsck=to_rsck(other)) # a different kernel + _run(problem, ops, weight_rsck=to_rsck(other)) # a different kernel with pytest.raises(ValueError): _run(problem, ops, weight_rsck=good.float()) # Right shape, wrong layout: the B tile load assumes Cout is contiguous. with pytest.raises(ValueError): - _run(problem, ops, - weight_rsck=good.transpose(3, 4).contiguous().transpose(3, 4)) + _run( + problem, ops, weight_rsck=good.transpose(3, 4).contiguous().transpose(3, 4) + ) @requires_gpu @@ -919,19 +946,20 @@ def test_every_weight_layout_gives_the_same_answer(): layouts = { "channels_last": w.contiguous(memory_format=torch.channels_last_3d), "contiguous": w.contiguous(), - "rsck_strided": (w.permute(2, 3, 4, 1, 0).contiguous() - .permute(4, 3, 0, 1, 2)), + "rsck_strided": (w.permute(2, 3, 4, 1, 0).contiguous().permute(4, 3, 0, 1, 2)), } ref = _run(problem, ops) for name, wl in layouts.items(): - assert torch.equal(wl, w), name # same values, different strides + assert torch.equal(wl, w), name # same values, different strides assert torch.equal(ref, _run(problem, {**ops, "weight": wl})), name assert torch.equal(ref, _run(problem, ops, weight_rsck=to_rsck(w))) # ``to_rsck`` of an already-RSCK-strided weight must not copy: that is what # makes the layout free for a caller who chooses it, and ``.contiguous()`` # returning ``self`` is the whole mechanism. - assert to_rsck(layouts["rsck_strided"]).data_ptr() == \ - layouts["rsck_strided"].data_ptr() + assert ( + to_rsck(layouts["rsck_strided"]).data_ptr() + == layouts["rsck_strided"].data_ptr() + ) @requires_gpu @@ -953,8 +981,13 @@ def test_output_is_channels_last_and_matches_torch_shape(): problem = ConvProblem("shape", 16, 40, (3, 11, 5)) ops = reference.make_inputs(problem, seed=61) y = _run(problem, ops) - ref = F.conv3d(ops["input"], ops["weight"], ops["bias"], - stride=problem.stride, padding=problem.padding) + ref = F.conv3d( + ops["input"], + ops["weight"], + ops["bias"], + stride=problem.stride, + padding=problem.padding, + ) assert y.shape == ref.shape assert y.is_contiguous(memory_format=torch.channels_last_3d) @@ -967,8 +1000,9 @@ def test_unsupported_calls_raise_rather_than_return_garbage(): conv3d_forward(x, w, padding=1, groups=2) good_w = torch.randn((8, 8, 3, 3, 3), device="cuda", dtype=torch.bfloat16) with pytest.raises(ValueError): - conv3d_forward(x, good_w, padding=1, - config=ConvConfig(BLOCK_K=8, matrix_instr_nonkdim=16)) + conv3d_forward( + x, good_w, padding=1, config=ConvConfig(BLOCK_K=8, matrix_instr_nonkdim=16) + ) @requires_gpu @@ -1040,8 +1074,12 @@ def test_indices_beyond_int32_are_addressed_correctly_with_taps_and_padding(): if free < 12 << 30: pytest.skip("needs ~12 GiB free") cin, cout, sp = 128, 16, (258, 258, 258) - x = torch.empty((1, cin, *sp), device="cuda", dtype=torch.bfloat16, - memory_format=torch.channels_last_3d).zero_() + x = torch.empty( + (1, cin, *sp), + device="cuda", + dtype=torch.bfloat16, + memory_format=torch.channels_last_3d, + ).zero_() w = torch.zeros((cout, cin, 3, 3, 3), device="cuda", dtype=torch.bfloat16) # Tap (2,2,2) of channel 0 alone. At padding 1 that is y[o] = x[o + 1]. w[0, 0, 2, 2, 2] = 1.0 @@ -1084,14 +1122,15 @@ def test_a_weight_beyond_int32_is_addressed_correctly(): pytest.skip("needs ~8 GiB free") bf16, cin, cout = torch.bfloat16, 16385, 65536 wr = torch.zeros((2, 1, 1, cin, cout), device="cuda", dtype=bf16) - assert (2 * cin - 1) * cout > 2**31 - 1 # the largest row offset + assert (2 * cin - 1) * cout > 2**31 - 1 # the largest row offset w = torch.zeros((), device="cuda", dtype=bf16).expand(cout, cin, 2, 1, 1) x = torch.zeros((1, cin, 2, 1, 1), device="cuda", dtype=bf16).contiguous( - memory_format=torch.channels_last_3d) + memory_format=torch.channels_last_3d + ) - wr[1, 0, 0, cin - 1, cout - 1] = 1.0 # the last element of the weight + wr[1, 0, 0, cin - 1, cout - 1] = 1.0 # the last element of the weight x[0, cin - 1, 1, 0, 0] = 3.0 - wr[0, 0, 0, 0, 0] = 1.0 # and the first + wr[0, 0, 0, 0, 0] = 1.0 # and the first x[0, 0, 0, 0, 0] = 5.0 y = conv3d_forward(x, w, padding=0, weight_rsck=wr) diff --git a/triton_conv3d/tests/test_infra.py b/triton_conv3d/tests/test_infra.py index 825a7b4..1c68d3b 100644 --- a/triton_conv3d/tests/test_infra.py +++ b/triton_conv3d/tests/test_infra.py @@ -33,9 +33,7 @@ scaffold_corpus, ) -requires_gpu = pytest.mark.skipif( - not torch.cuda.is_available(), reason="needs a GPU" -) +requires_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") #: ``include_large=True`` because the shape and cost model tests below are meta #: tensors and integer arithmetic -- a 4 GiB activation costs nothing here, and @@ -101,8 +99,9 @@ def test_transposed_flops_have_no_phantom_tap_factor(): the input volume by the tap count *and* keeping the tap factor would inflate the count eightfold, which is exactly the bug this pins. """ - p = ConvProblem("t", 64, 32, (8, 8, 8), (2, 2, 2), (2, 2, 2), (0, 0, 0), - transposed=True) + p = ConvProblem( + "t", 64, 32, (8, 8, 8), (2, 2, 2), (2, 2, 2), (0, 0, 0), transposed=True + ) assert p.out_spatial == (16, 16, 16) macs = math.prod(p.out_spatial) * p.cin * p.cout assert p.flops("fwd") == 2 * macs @@ -123,8 +122,11 @@ def test_the_int32_edge_cases_bracket_the_element_boundary(): assert not below.index_exceeds_int32 assert above.index_exceeds_int32 # Same channel pair and kernel: the only thing that differs is the volume. - assert (below.cin, below.cout, below.kernel) == (above.cin, above.cout, - above.kernel) + assert (below.cin, below.cout, below.kernel) == ( + above.cin, + above.cout, + above.kernel, + ) # And the boundary is where the largest index -- not the count -- crosses. assert below.max_elements - 1 <= INT32_MAX < above.max_elements - 1 @@ -145,23 +147,32 @@ def test_the_2gib_cliff_is_a_byte_problem_and_not_an_index_problem(): used to be read as though it did. What it exceeds is the *byte* limit on the whole storage, which is what decides buffer-op eligibility. """ - cliff = next(p.halo_variant for p in scaffold_corpus() - if p.halo_variant.label == "conv 128->64 k3x3x3 @ 130x258x258") + cliff = next( + p.halo_variant + for p in scaffold_corpus() + if p.halo_variant.label == "conv 128->64 k3x3x3 @ 130x258x258" + ) assert not cliff.index_exceeds_int32 - assert cliff.max_elements / 2**31 < 0.55 # half int32's range + assert cliff.max_elements / 2**31 < 0.55 # half int32's range assert not cliff.buffer_ops_eligible assert cliff.max_activation_bytes / BUFFER_OP_MAX_BYTES == pytest.approx( - 1.032, abs=0.002 # 3.2% past 2 GiB + 1.032, + abs=0.002, # 3.2% past 2 GiB ) # It is the only corpus shape on either side of that line, in either shape # mode -- and *no* corpus shape needs a 64-bit element index. A test or a # dispatch rule parametrized on the index predicate selects nothing. - over = [p.halo_variant.label for p in scaffold_corpus() - if not p.halo_variant.buffer_ops_eligible] + over = [ + p.halo_variant.label + for p in scaffold_corpus() + if not p.halo_variant.buffer_ops_eligible + ] assert over == ["conv 128->64 k3x3x3 @ 130x258x258"] - assert not any(p.index_exceeds_int32 or p.halo_variant.index_exceeds_int32 - for p in scaffold_corpus()) + assert not any( + p.index_exceeds_int32 or p.halo_variant.index_exceeds_int32 + for p in scaffold_corpus() + ) def test_corpus_covers_the_three_scaffold_configurations(): @@ -252,24 +263,33 @@ def test_the_production_variant_is_what_a_real_step_issues(): from triton_conv3d.shapes import census_corpus, production_corpus def key(p): - return (p.transposed, p.cin, p.cout, tuple(p.kernel), - tuple(p.spatial), tuple(p.padding), tuple(p.stride), p.n) + return ( + p.transposed, + p.cin, + p.cout, + tuple(p.kernel), + tuple(p.spatial), + tuple(p.padding), + tuple(p.stride), + p.n, + ) census = {key(p) for p in census_corpus()} assert len(census) > 60, "the census is missing; nothing is being checked" - missing = [p for p in production_corpus() - if key(p) not in census and p.kernel != (1, 1, 1)] + missing = [ + p for p in production_corpus() if key(p) not in census and p.kernel != (1, 1, 1) + ] assert not missing, ( "the corpus's production form does not match what a real step issued: " + ", ".join(p.qualified_label for p in missing) ) # And the census is in the form it claims: every k>1 convolution padded. - unpadded = [p for p in census_corpus() - if p.kernel == (3, 3, 3) and not any(p.padding)] + unpadded = [ + p for p in census_corpus() if p.kernel == (3, 3, 3) and not any(p.padding) + ] assert not unpadded, ( "a k=3 production convolution arrived unpadded, which would mean the " - "adapter's halo plan changed: " - + ", ".join(p.qualified_label for p in unpadded) + "adapter's halo plan changed: " + ", ".join(p.qualified_label for p in unpadded) ) @@ -284,8 +304,11 @@ def test_the_three_forms_are_told_apart_by_the_qualified_label(): sharded = [p for p in scaffold_corpus() if any(p.shard_halo)] assert sharded, "corpus has no sharded problems; shard_halo_dhw was lost" for p in sharded: - forms = {p.qualified_label, p.production_variant.qualified_label, - p.halo_variant.qualified_label} + forms = { + p.qualified_label, + p.production_variant.qualified_label, + p.halo_variant.qualified_label, + } assert len(forms) == 3, f"{p.label}: forms collide -> {forms}" # The adapter halos D and leaves H and W padded; DistConv does neither. assert p.production_variant.padding == (0, *p.padding[1:]) @@ -342,8 +365,12 @@ def test_stored_efficiency_agrees_with_the_cost_model(): f"stored {m['pct_roofline']}%, cost model {got:.3f}%" ) # And no forward cell exceeds the roof, which is what the artifact implied. - fwd = [m["pct_roofline"] for p in scaffold_corpus() for m in p.measured - if m["direction"] == "fwd"] + fwd = [ + m["pct_roofline"] + for p in scaffold_corpus() + for m in p.measured + if m["direction"] == "fwd" + ] assert fwd and max(fwd) < 100 @@ -389,8 +416,13 @@ def test_reference_agrees_with_miopen_within_tolerance(problem: ConvProblem): # test passed. Only this direction and only the incumbent get it: our # own backward-weight reduces its split-K partials in fp32 and stores # once, so it is held to ``roundings=1`` like everything else. - reference.assert_close(actual, expected, problem, direction, - roundings=2 if direction == "bwd-weight" else 1) + reference.assert_close( + actual, + expected, + problem, + direction, + roundings=2 if direction == "bwd-weight" else 1, + ) @requires_gpu @@ -417,7 +449,9 @@ def test_error_bound_grows_with_reduction_length_and_shrinks_with_precision(): expected = torch.randn(4096, dtype=torch.float64) short = ConvProblem("short", 8, 8, (8, 8, 8)) long_ = ConvProblem("long", 1024, 1024, (8, 8, 8)) - assert reference.error_bound(long_, expected) > reference.error_bound(short, expected) + assert reference.error_bound(long_, expected) > reference.error_bound( + short, expected + ) fp32 = ConvProblem("fp32", 64, 64, (8, 8, 8), dtype="fp32") bf16 = ConvProblem("bf16", 64, 64, (8, 8, 8), dtype="bf16") assert reference.error_bound(fp32, expected) < reference.error_bound(bf16, expected) @@ -434,7 +468,9 @@ def test_error_bound_tracks_peak_not_just_rms(): flat = torch.ones(4096, dtype=torch.float64) spiky = flat.clone() spiky[0] = 100.0 - assert reference.error_bound(problem, spiky) > 10 * reference.error_bound(problem, flat) + assert reference.error_bound(problem, spiky) > 10 * reference.error_bound( + problem, flat + ) def test_the_store_term_is_charged_as_one_rounding_not_four(): @@ -452,13 +488,14 @@ def test_the_store_term_is_charged_as_one_rounding_not_four(): formula rather than on a GPU: at ``K`` short enough that the accumulation term is negligible, the bound must be one ulp of the peak per rounding. """ - problem = ConvProblem("p", 8, 8, (4, 4, 4)) # K = 216 + problem = ConvProblem("p", 8, 8, (4, 4, 4)) # K = 216 peak = torch.zeros(4096, dtype=torch.float64) peak[0] = 64.0 ulp = 2.0 * reference.unit_roundoff(torch.bfloat16) * 64.0 assert reference.error_bound(problem, peak) == pytest.approx(ulp, rel=1e-3) assert reference.error_bound(problem, peak, roundings=2) == pytest.approx( - 2 * ulp, rel=1e-3) + 2 * ulp, rel=1e-3 + ) @requires_gpu @@ -543,8 +580,9 @@ def test_the_incumbent_clause_binds_more_often_than_the_static_bound(): ops = reference.make_inputs(problem, seed=7) expected = reference.reference(problem, ops, "fwd") err = reference.compare(reference.incumbent(problem, ops, "fwd"), expected) - binds.append(4.0 * err.max_abs - > reference.error_bound(problem, expected, "fwd")) + binds.append( + 4.0 * err.max_abs > reference.error_bound(problem, expected, "fwd") + ) assert sum(binds) > len(binds) // 2, ( f"the incumbent clause bound only {sum(binds)}/{len(binds)} cells; the " "static bound has drifted back to swallowing it" @@ -592,6 +630,7 @@ def make(name): def fn(): order.append(name) return a @ a + return fn result = interleaved({k: make(k) for k in seen}, warmup=1, iters=1, rounds=3) @@ -606,8 +645,9 @@ def fn(): # exactly the calls it says: 1 warmup and 3 rounds of 1 iteration each, per # variant, with no calibration probe smuggled in. assert len(order) == 3 * (1 + 3 * 1) - assert all(m.iters == 1 and m.group == 1 and m.stop == "fixed" - for m in result.values()) + assert all( + m.iters == 1 and m.group == 1 and m.stop == "fixed" for m in result.values() + ) def test_the_round_order_is_position_and_adjacency_balanced(): @@ -652,7 +692,7 @@ def test_the_round_order_is_position_and_adjacency_balanced(): f"n={n}: adjacency is not balanced: {adjacency}" ) assert len(adjacency) == n * (n - 1), ( - f"n={n}: only {len(adjacency)} of {n*(n-1)} ordered pairs occur" + f"n={n}: only {len(adjacency)} of {n * (n - 1)} ordered pairs occur" ) @@ -675,6 +715,7 @@ def test_spread_is_a_range_statistic_and_the_interval_is_not(): from triton_conv3d.bench.harness import Measurement rng = random.Random(20260803) + def draw(n): return tuple(1.0 + 0.01 * rng.gauss(0, 1) for _ in range(n)) @@ -760,8 +801,9 @@ def test_a_slow_kernel_stops_on_the_budget_and_says_so(): from triton_conv3d.bench.harness import time_callable t0 = time.perf_counter() - m = time_callable(lambda: torch.cuda._sleep(1_000_000_000), - budget_s=1.0, target_rel=1e-9) + m = time_callable( + lambda: torch.cuda._sleep(1_000_000_000), budget_s=1.0, target_rel=1e-9 + ) elapsed = time.perf_counter() - t0 assert m.stop == "budget", f"stopped for the wrong reason: {m.stop}" assert not m.converged @@ -786,7 +828,6 @@ def test_the_instrument_tax_is_measured_and_grouped_away(): second is the control: if grouping ever stops working, the two agree and this fails, rather than both drifting together unnoticed. """ - import statistics import time from triton_conv3d.bench.harness import time_callable @@ -855,15 +896,14 @@ def test_flush_caches_reuses_one_buffer_and_reaches_only_the_first_sample(): for _ in range(5): H.flush_caches() assert H._flush_buffer.data_ptr() == first, ( - "flush_caches reallocated its buffer; the device comparison is wrong " - "again" + "flush_caches reallocated its buffer; the device comparison is wrong again" ) assert H._flush_buffer.numel() == H._FLUSH_BYTES a = torch.randn(1024, 1024, device="cuda", dtype=torch.bfloat16) cold = H.time_callable(lambda: a @ a, flush=True, rounds=4) assert cold.iters == 1, ( - f"flush=True measured {cold.iters} calls per block, so {cold.iters-1} " + f"flush=True measured {cold.iters} calls per block, so {cold.iters - 1} " "of them are hot and the median reports a hot number" ) assert cold.cold == cold.median # with one sample per block they coincide @@ -1021,8 +1061,10 @@ def test_sporadic_host_stall_is_rejected_from_the_median_and_flagged(): # cannot decide -- failing here would re-enact the original misdiagnosis in # test form, blaming the measurement for observing real contention. clean = min( - (interleaved({"g": lambda: a @ a}, warmup=3, iters=20, rounds=3)["g"] - for _ in range(5)), + ( + interleaved({"g": lambda: a @ a}, warmup=3, iters=20, rounds=3)["g"] + for _ in range(5) + ), key=lambda m: m.stall_ratio, ) if clean.stall_ratio >= 2.0: @@ -1052,7 +1094,9 @@ def sporadic(): # The converse -- that a quiet run is *not* flagged -- is established by the # skip above rather than here, because on a loaded node it is not true and # should not be asserted. - assert clean.stall_ratio < 2.0, f"clean run falsely flagged: {clean.stall_ratio:.2f}" + assert clean.stall_ratio < 2.0, ( + f"clean run falsely flagged: {clean.stall_ratio:.2f}" + ) # --------------------------------------------------------------------------- @@ -1431,9 +1475,18 @@ def test_the_transposed_problems_are_never_haloed_and_the_others_follow_the_form # problem happens to record a zero halo, so a form that called # ``halo_variant`` would be indistinguishable from the right one until the # day one of them did not. Pin the rule instead. - haloed = ConvProblem("would_halo", 32, 16, (4, 4, 4), (2, 2, 2), (2, 2, 2), - (0, 0, 0), transposed=True, halo=(1, 1, 1), - shard_halo=(1, 1, 1)) + haloed = ConvProblem( + "would_halo", + 32, + 16, + (4, 4, 4), + (2, 2, 2), + (2, 2, 2), + (0, 0, 0), + transposed=True, + halo=(1, 1, 1), + shard_halo=(1, 1, 1), + ) for name in _FORMS: assert convt.form(haloed, name) is haloed, ( f"--form {name} gave a transposed problem a halo; at k=2 there is " @@ -1475,27 +1528,33 @@ def test_the_shipped_config_is_the_one_the_entry_point_resolves(): than by re-deriving the answer here -- which would be the same arithmetic twice and would agree with itself while both were wrong. """ - from triton_conv3d.bench.conv_bench import _OPERATORS, _build from triton_conv3d import bwd_data, gather_gemm, reduce_gemm, transposed + from triton_conv3d.bench.conv_bench import _OPERATORS, _build problems = { "conv": ConvProblem("t", 32, 16, (8, 8, 8)), - "convT": ConvProblem("tt", 32, 16, (4, 4, 4), (2, 2, 2), (2, 2, 2), - (0, 0, 0), transposed=True), + "convT": ConvProblem( + "tt", 32, 16, (4, 4, 4), (2, 2, 2), (2, 2, 2), (0, 0, 0), transposed=True + ), } spied = [] def spy(mod, name): real = getattr(mod, name) + def wrapper(*a, **kw): cfg = real(*a, **kw) spied.append(cfg) return cfg + return real, wrapper - patched = [(gather_gemm, "select_config"), (bwd_data, "select_config"), - (reduce_gemm, "bwd_weight_config"), - (transposed, "transposed_config")] + patched = [ + (gather_gemm, "select_config"), + (bwd_data, "select_config"), + (reduce_gemm, "bwd_weight_config"), + (transposed, "transposed_config"), + ] originals = {} for mod, name in patched: real, wrapper = spy(mod, name) @@ -1537,8 +1596,9 @@ def test_the_published_time_is_per_call_and_never_exceeds_the_eager_call(): """ from triton_conv3d.bench.conv_bench import measure_problem - p = ConvProblem("tt", 64, 32, (4, 4, 4), (2, 2, 2), (2, 2, 2), (0, 0, 0), - transposed=True) + p = ConvProblem( + "tt", 64, 32, (4, 4, 4), (2, 2, 2), (2, 2, 2), (0, 0, 0), transposed=True + ) row = measure_problem(p, direction="fwd", shipped=True, budget_s=5.0) assert "error" not in row, row.get("error") assert row["timed_region"] == "kernel", row["timed_region_note"] @@ -1573,8 +1633,9 @@ def test_a_control_free_row_omits_the_control_rather_than_zeroing_it(): """ from triton_conv3d.bench.conv_bench import _build, measure_problem - p = ConvProblem("tt", 64, 32, (4, 4, 4), (2, 2, 2), (2, 2, 2), (0, 0, 0), - transposed=True) + p = ConvProblem( + "tt", 64, 32, (4, 4, 4), (2, 2, 2), (2, 2, 2), (0, 0, 0), transposed=True + ) for direction in DIRECTIONS: case = _build(p, direction, control=False) @@ -1583,8 +1644,9 @@ def test_a_control_free_row_omits_the_control_rather_than_zeroing_it(): del case torch.cuda.empty_cache() - row = measure_problem(p, direction="bwd-weight", shipped=True, - budget_s=5.0, control="none") + row = measure_problem( + p, direction="bwd-weight", shipped=True, budget_s=5.0, control="none" + ) assert "error" not in row, row.get("error") assert row["control"] == "none" # The Triton half is unchanged: same region, same interval, same stop rule. @@ -1597,8 +1659,15 @@ def test_a_control_free_row_omits_the_control_rather_than_zeroing_it(): assert math.isfinite(row["triton_rel_ci"]) and row["triton_rel_ci"] >= 0.0 assert row["measure_stop"] in ("converged", "budget", "max_rounds") # The MIOpen half is gone, not zeroed. - for key in ("miopen_ms", "miopen_rel_ci", "miopen_eager_ms", "speedup", - "speedup_lo", "speedup_hi", "speedup_significant"): + for key in ( + "miopen_ms", + "miopen_rel_ci", + "miopen_eager_ms", + "speedup", + "speedup_lo", + "speedup_hi", + "speedup_significant", + ): assert key not in row, ( f"{key} is present in a --control none row; an absent measurement " "must stay absent, because a zero here reads as a result" diff --git a/triton_conv3d/tests/test_transposed.py b/triton_conv3d/tests/test_transposed.py index 2b99d2c..eb45632 100644 --- a/triton_conv3d/tests/test_transposed.py +++ b/triton_conv3d/tests/test_transposed.py @@ -30,7 +30,6 @@ from __future__ import annotations -import dataclasses import math import pytest @@ -38,7 +37,7 @@ import torch.nn.functional as F from triton_conv3d import reference -from triton_conv3d.gather_gemm import ConvConfig, is_supported +from triton_conv3d.gather_gemm import is_supported from triton_conv3d.shapes import ConvProblem, scaffold_corpus from triton_conv3d.transposed import ( TransposedConfig, @@ -59,11 +58,23 @@ requires_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") -def _problem(name, cin, cout, spatial, k=(2, 2, 2), *, bias=False, - dtype="bf16", n=1) -> ConvProblem: - return ConvProblem(name, cin, cout, spatial, k, k, (0, 0, 0), n=n, - transposed=True, bias=bias, dtype=dtype, - sites=("synthetic",)) +def _problem( + name, cin, cout, spatial, k=(2, 2, 2), *, bias=False, dtype="bf16", n=1 +) -> ConvProblem: + return ConvProblem( + name, + cin, + cout, + spatial, + k, + k, + (0, 0, 0), + n=n, + transposed=True, + bias=bias, + dtype=dtype, + sites=("synthetic",), + ) #: Synthetic problems chosen to break the scatter rather than to be fast. Each @@ -109,10 +120,19 @@ def _problem(name, cin, cout, spatial, k=(2, 2, 2), *, bias=False, #: extents are deliberately not powers of two so ``BLOCK_M`` does not divide #: ``M`` and the store's rows wrap. CORPUS_PAIRS: list[ConvProblem] = [ - _problem(f"{p.cin}to{p.cout}-corpus", p.cin, p.cout, (3, 4, 5), - tuple(p.kernel), bias=p.bias) - for p in {(q.cin, q.cout, tuple(q.kernel), q.bias): q - for q in scaffold_corpus() if q.transposed}.values() + _problem( + f"{p.cin}to{p.cout}-corpus", + p.cin, + p.cout, + (3, 4, 5), + tuple(p.kernel), + bias=p.bias, + ) + for p in { + (q.cin, q.cout, tuple(q.kernel), q.bias): q + for q in scaffold_corpus() + if q.transposed + }.values() ] @@ -132,7 +152,9 @@ def _ops(problem: ConvProblem, seed: int = 0, direction="fwd") -> dict: """ dtype = reference.torch_dtype(problem) return reference.make_inputs( - problem, seed=seed, exact=True, + problem, + seed=seed, + exact=True, density=reference.exact_density(problem, direction, dtype=dtype), ) @@ -149,14 +171,20 @@ def _fwd(problem: ConvProblem, ops: dict, **kw) -> torch.Tensor: def _bwd_data(problem: ConvProblem, ops: dict, **kw) -> torch.Tensor: return conv_transpose3d_backward_data( - ops["grad_output"], ops["weight"], problem.input_shape, problem.stride, + ops["grad_output"], + ops["weight"], + problem.input_shape, + problem.stride, **kw, ) def _bwd_weight(problem: ConvProblem, ops: dict, **kw) -> torch.Tensor: return conv_transpose3d_backward_weight( - ops["input"], problem.weight_shape, ops["grad_output"], problem.stride, + ops["input"], + problem.weight_shape, + ops["grad_output"], + problem.stride, **kw, ) @@ -183,8 +211,7 @@ def test_the_windows_tile_the_output_exactly_once(): for kd in range(k[0]): for kh in range(k[1]): for kw in range(k[2]): - hits[d * k[0] + kd, h * k[1] + kh, - w * k[2] + kw] += 1 + hits[d * k[0] + kd, h * k[1] + kh, w * k[2] + kw] += 1 assert torch.equal(hits, torch.ones_like(hits)), k # ``k=3, s=2`` overlaps: the windows cover some voxels twice. If this ever @@ -211,9 +238,11 @@ def test_transposed_flops_have_no_phantom_tap_factor(): So the count is derived here from first principles: one MAC per (output voxel, output channel, input channel), times two. """ - for cin, cout, spatial, k in [(128, 64, (4, 5, 6), (2, 2, 2)), - (32, 32, (3, 3, 3), (3, 3, 3)), - (16, 8, (2, 3, 4), (1, 2, 4))]: + for cin, cout, spatial, k in [ + (128, 64, (4, 5, 6), (2, 2, 2)), + (32, 32, (3, 3, 3), (3, 3, 3)), + (16, 8, (2, 3, 4), (1, 2, 4)), + ]: p = _problem("f", cin, cout, spatial, k) out_vol = math.prod(p.out_spatial) macs = out_vol * cout * cin @@ -242,11 +271,12 @@ def test_the_gemm_decomposition_matches_the_kernels_grid(): taps = p.tap_count assert n == p.cout * taps assert m == p.n * math.prod(p.spatial) - cfg = transposed_config(m, p.cin, p.cout, p.kernel, - reference.torch_dtype(p)) + cfg = transposed_config(m, p.cin, p.cout, p.kernel, reference.torch_dtype(p)) assert taps % cfg.TAP_BLOCK == 0, (p.label, cfg) - columns = (taps // cfg.TAP_BLOCK) * cfg.TAP_BLOCK * ( - -(-p.cout // cfg.BLOCK_NC) * cfg.BLOCK_NC + columns = ( + (taps // cfg.TAP_BLOCK) + * cfg.TAP_BLOCK + * (-(-p.cout // cfg.BLOCK_NC) * cfg.BLOCK_NC) ) assert columns >= n, (p.label, cfg) @@ -256,8 +286,7 @@ def test_the_gemm_decomposition_matches_the_kernels_grid(): # --------------------------------------------------------------------------- -@pytest.mark.parametrize("problem", EDGE + CORPUS_PAIRS, - ids=_ids(EDGE + CORPUS_PAIRS)) +@pytest.mark.parametrize("problem", EDGE + CORPUS_PAIRS, ids=_ids(EDGE + CORPUS_PAIRS)) def test_selected_config_is_legal_for_every_shape(problem: ConvProblem): """An illegal MFMA config does not fail, it silently emits vector FMA. @@ -285,8 +314,9 @@ def test_every_candidate_config_is_legal(problem: ConvProblem): """ dtype = reference.torch_dtype(problem) m = problem.n * math.prod(problem.spatial) - cands = candidate_transposed_configs(m, problem.cin, problem.cout, - problem.tap_count, dtype) + cands = candidate_transposed_configs( + m, problem.cin, problem.cout, problem.tap_count, dtype + ) assert cands for cfg in cands: assert cfg.validate(dtype) is None, cfg @@ -302,8 +332,14 @@ def test_the_fp32_config_fits_lds(): ScaFFold configuration. This kernel's tile is *wider* than that one's (``TAP_BLOCK`` multiplies the column count), so the same hole is closer. """ - for cin, cout, taps in [(1024, 512, 8), (512, 256, 8), (256, 128, 8), - (128, 64, 8), (64, 64, 27), (2048, 1024, 8)]: + for cin, cout, taps in [ + (1024, 512, 8), + (512, 256, 8), + (256, 128, 8), + (128, 64, 8), + (64, 64, 27), + (2048, 1024, 8), + ]: for dtype in (torch.float32, torch.bfloat16, torch.float16): cfg = default_transposed_config(1 << 16, cin, cout, taps, dtype) assert cfg.validate(dtype) is None, (cin, cout, dtype, cfg) @@ -340,18 +376,23 @@ def test_is_supported_declines_what_the_tiling_argument_does_not_cover(): # is why the positive control below is on the GPU. assert not is_supported_transposed(x, w, None, 2, 0, 0, 1, 1) # not cuda for stride, padding, output_padding, dilation, groups in [ - (1, 0, 0, 1, 1), # k != s: the windows overlap - (3, 0, 0, 1, 1), # k != s: the windows leave gaps - (2, 1, 0, 1, 1), # padding crops the tiled result - (2, 0, 1, 1, 1), # output_padding extends it asymmetrically - (2, 0, 0, 2, 1), # dilation interleaves the window with holes - (2, 0, 0, 1, 2), # groups - ((2, 2, 1), 0, 0, 1, 1), # anisotropic mismatch on one axis only + (1, 0, 0, 1, 1), # k != s: the windows overlap + (3, 0, 0, 1, 1), # k != s: the windows leave gaps + (2, 1, 0, 1, 1), # padding crops the tiled result + (2, 0, 1, 1, 1), # output_padding extends it asymmetrically + (2, 0, 0, 2, 1), # dilation interleaves the window with holes + (2, 0, 0, 1, 2), # groups + ((2, 2, 1), 0, 0, 1, 1), # anisotropic mismatch on one axis only ]: assert not is_supported_transposed( x.cuda() if torch.cuda.is_available() else x, w.cuda() if torch.cuda.is_available() else w, - None, stride, padding, output_padding, dilation, groups + None, + stride, + padding, + output_padding, + dilation, + groups, ), (stride, padding, output_padding, dilation, groups) @@ -370,10 +411,14 @@ def test_the_gates_are_total(): assert is_supported_transposed(x, w, None, 2, 0, bad, 1, 1) is False assert is_supported_transposed(x, w, None, 2, 0, 0, bad, 1) is False assert is_supported_transposed_all(x, w, None, bad, 0, 0, 1, 1) is False - assert is_supported_transposed_bwd_data( - x, w, (1, 8, 4, 4, 4), bad, 0, 0, 1, 1) is False - assert is_supported_transposed_bwd_weight( - x, (8, 4, 2, 2, 2), x, bad, 0, 0, 1, 1) is False + assert ( + is_supported_transposed_bwd_data(x, w, (1, 8, 4, 4, 4), bad, 0, 0, 1, 1) + is False + ) + assert ( + is_supported_transposed_bwd_weight(x, (8, 4, 2, 2, 2), x, bad, 0, 0, 1, 1) + is False + ) # A malformed ``input_shape`` / ``weight_shape`` is the same kind of # question and gets the same kind of answer. for bad in (None, (1, 8, 4, 4), "abcde", 5): @@ -391,14 +436,18 @@ def test_is_supported_declines_degenerate_extents(): """ good_x = torch.zeros(1, 8, 4, 4, 4) good_w = torch.zeros(8, 4, 2, 2, 2) - assert not is_supported_transposed(torch.zeros(1, 8, 0, 4, 4), good_w, - None, 2, 0, 0, 1, 1) - assert not is_supported_transposed(good_x, torch.zeros(0, 4, 2, 2, 2), - None, 2, 0, 0, 1, 1) - assert not is_supported_transposed(good_x, torch.zeros(8, 0, 2, 2, 2), - None, 2, 0, 0, 1, 1) - assert not is_supported_transposed(good_x, torch.zeros(8, 4, 0, 2, 2), - None, (0, 2, 2), 0, 0, 1, 1) + assert not is_supported_transposed( + torch.zeros(1, 8, 0, 4, 4), good_w, None, 2, 0, 0, 1, 1 + ) + assert not is_supported_transposed( + good_x, torch.zeros(0, 4, 2, 2, 2), None, 2, 0, 0, 1, 1 + ) + assert not is_supported_transposed( + good_x, torch.zeros(8, 0, 2, 2, 2), None, 2, 0, 0, 1, 1 + ) + assert not is_supported_transposed( + good_x, torch.zeros(8, 4, 0, 2, 2), None, (0, 2, 2), 0, 0, 1, 1 + ) @requires_gpu @@ -412,19 +461,34 @@ def test_is_supported_reads_the_transposed_weight_convention(): """ x = torch.zeros(1, 8, 4, 4, 4, device="cuda", dtype=torch.bfloat16) assert is_supported_transposed( - x, torch.zeros(8, 4, 2, 2, 2, device="cuda", dtype=torch.bfloat16), - None, 2, 0, 0, 1, 1) + x, + torch.zeros(8, 4, 2, 2, 2, device="cuda", dtype=torch.bfloat16), + None, + 2, + 0, + 0, + 1, + 1, + ) # Same tensor read the other way round: Cin=4 does not match x's 8 channels. assert not is_supported_transposed( - x, torch.zeros(4, 8, 2, 2, 2, device="cuda", dtype=torch.bfloat16), - None, 2, 0, 0, 1, 1) + x, + torch.zeros(4, 8, 2, 2, 2, device="cuda", dtype=torch.bfloat16), + None, + 2, + 0, + 0, + 1, + 1, + ) # A bias is ``Cout`` = w.shape[1] long, not w.shape[0]. w = torch.zeros(8, 4, 2, 2, 2, device="cuda", dtype=torch.bfloat16) assert is_supported_transposed( - x, w, torch.zeros(4, device="cuda", dtype=torch.bfloat16), - None or 2, 0, 0, 1, 1) + x, w, torch.zeros(4, device="cuda", dtype=torch.bfloat16), None or 2, 0, 0, 1, 1 + ) assert not is_supported_transposed( - x, w, torch.zeros(8, device="cuda", dtype=torch.bfloat16), 2, 0, 0, 1, 1) + x, w, torch.zeros(8, device="cuda", dtype=torch.bfloat16), 2, 0, 0, 1, 1 + ) # A stride-2 view of the right length applies every other value; the kernel # indexes the bias with an element stride of 1 and cannot see this. long_bias = torch.zeros(8, device="cuda", dtype=torch.bfloat16) @@ -447,8 +511,7 @@ def test_is_supported_declines_operands_on_different_devices(): @requires_gpu -@pytest.mark.parametrize("problem", EDGE + CORPUS_PAIRS, - ids=_ids(EDGE + CORPUS_PAIRS)) +@pytest.mark.parametrize("problem", EDGE + CORPUS_PAIRS, ids=_ids(EDGE + CORPUS_PAIRS)) def test_all_three_gates_accept_every_problem_this_module_serves(problem): """``is_supported_transposed_all`` must not be narrower than the forward. @@ -463,8 +526,7 @@ def test_all_three_gates_accept_every_problem_this_module_serves(problem): dtype = reference.torch_dtype(problem) x = torch.zeros(problem.input_shape, device="cuda", dtype=dtype) w = torch.zeros(problem.weight_shape, device="cuda", dtype=dtype) - b = (torch.zeros(problem.cout, device="cuda", dtype=dtype) - if problem.bias else None) + b = torch.zeros(problem.cout, device="cuda", dtype=dtype) if problem.bias else None args = (problem.stride, 0, 0, 1, 1) assert is_supported_transposed(x, w, b, *args), problem.label assert is_supported_transposed_all(x, w, b, *args), problem.label @@ -547,8 +609,9 @@ def test_corpus_channel_pairs_match_bitwise(problem: ConvProblem, direction: str expected = _reference(problem, ops, direction) dtype = reference.torch_dtype(problem) assert reference.is_exactly_representable(expected, dtype), problem.label - actual = {"fwd": _fwd, "bwd-data": _bwd_data, - "bwd-weight": _bwd_weight}[direction](problem, ops) + actual = {"fwd": _fwd, "bwd-data": _bwd_data, "bwd-weight": _bwd_weight}[direction]( + problem, ops + ) assert reference.compare(actual, expected.to(dtype)).bitwise, problem.label @@ -576,8 +639,14 @@ def test_a_transposed_tap_permutation_is_detected(): # The same convolution with the weight's three kernel axes permuted, which # is exactly what unpacking the fused index in the wrong order computes. permuted = _reference( - problem, {**ops, "weight": ops["weight"].permute(0, 1, 4, 3, 2) - .contiguous(memory_format=torch.channels_last_3d)}, "fwd", + problem, + { + **ops, + "weight": ops["weight"] + .permute(0, 1, 4, 3, 2) + .contiguous(memory_format=torch.channels_last_3d), + }, + "fwd", ).to(torch.bfloat16) assert permuted.shape == actual.shape assert not reference.compare(actual, permuted).bitwise, ( @@ -634,8 +703,15 @@ def test_backward_weight_operand_swap_is_not_reversible(): expected = _reference(problem, ops, "bwd-weight").to(torch.bfloat16) assert reference.compare(actual, expected).bitwise with pytest.raises(NotImplementedError): - conv3d_backward_weight(ops["input"], problem.weight_shape, - ops["grad_output"], problem.stride, 0, 1, 1) + conv3d_backward_weight( + ops["input"], + problem.weight_shape, + ops["grad_output"], + problem.stride, + 0, + 1, + 1, + ) # ``k=1``: same shapes, so nothing but the argument order decides. flat = _problem("swap1", 32, 32, (3, 4, 5), (1, 1, 1)) @@ -643,8 +719,9 @@ def test_backward_weight_operand_swap_is_not_reversible(): got = _bwd_weight(flat, ops1) want = _reference(flat, ops1, "bwd-weight").to(torch.bfloat16) assert reference.compare(got, want).bitwise - swapped = conv3d_backward_weight(ops1["input"], flat.weight_shape, - ops1["grad_output"], flat.stride, 0, 1, 1) + swapped = conv3d_backward_weight( + ops1["input"], flat.weight_shape, ops1["grad_output"], flat.stride, 0, 1, 1 + ) assert swapped.shape == got.shape assert not reference.compare(got, swapped).bitwise, ( "swapping the two activations gave the same gradient; the operand " @@ -663,8 +740,10 @@ def test_backward_data_is_the_strided_convolution_it_claims_to_be(): problem = _problem("bd", 64, 32, (3, 4, 5)) ops = _ops(problem, seed=23, direction="bwd-data") from triton_conv3d.gather_gemm import conv3d_forward - direct = conv3d_forward(ops["grad_output"], ops["weight"], None, - problem.stride, 0, 1, 1) + + direct = conv3d_forward( + ops["grad_output"], ops["weight"], None, problem.stride, 0, 1, 1 + ) assert torch.equal(direct, _bwd_data(problem, ops)) @@ -689,8 +768,9 @@ def test_every_config_gives_the_same_answer(problem: ConvProblem): expected = _reference(problem, ops, "fwd") assert reference.is_exactly_representable(expected, dtype) m = problem.n * math.prod(problem.spatial) - cands = candidate_transposed_configs(m, problem.cin, problem.cout, - problem.tap_count, dtype) + cands = candidate_transposed_configs( + m, problem.cin, problem.cout, problem.tap_count, dtype + ) by_tb: dict[int, TransposedConfig] = {} for cfg in cands: by_tb.setdefault(cfg.TAP_BLOCK, cfg) @@ -722,8 +802,7 @@ def test_every_weight_layout_gives_the_same_answer(): plain = cl.contiguous() assert not plain.is_contiguous(memory_format=torch.channels_last_3d) for w in (cl, plain, to_tkn(cl).permute(3, 4, 0, 1, 2)): - got = conv_transpose3d_forward(ops["input"], w, ops["bias"], - problem.stride) + got = conv_transpose3d_forward(ops["input"], w, ops["bias"], problem.stride) assert reference.compare(got, expected).bitwise, tuple(w.stride()) @@ -739,23 +818,33 @@ def test_out_buffer_is_written_in_place_and_is_validated(): problem = _problem("outbuf", 32, 16, (3, 4, 5), bias=True) ops = _ops(problem, seed=37) expected = _reference(problem, ops, "fwd").to(torch.bfloat16) - y = torch.empty(problem.output_shape, device="cuda", dtype=torch.bfloat16, - memory_format=torch.channels_last_3d) + y = torch.empty( + problem.output_shape, + device="cuda", + dtype=torch.bfloat16, + memory_format=torch.channels_last_3d, + ) got = _fwd(problem, ops, out=y) assert got.data_ptr() == y.data_ptr() assert reference.compare(y, expected).bitwise - small = torch.empty((1, 16, 2, 2, 2), device="cuda", dtype=torch.bfloat16, - memory_format=torch.channels_last_3d) + small = torch.empty( + (1, 16, 2, 2, 2), + device="cuda", + dtype=torch.bfloat16, + memory_format=torch.channels_last_3d, + ) with pytest.raises(ValueError, match="shape"): _fwd(problem, ops, out=small) - ncdhw = torch.empty(problem.output_shape, device="cuda", - dtype=torch.bfloat16) + ncdhw = torch.empty(problem.output_shape, device="cuda", dtype=torch.bfloat16) with pytest.raises(ValueError, match="channels_last_3d"): _fwd(problem, ops, out=ncdhw) - wrong_dtype = torch.empty(problem.output_shape, device="cuda", - dtype=torch.float32, - memory_format=torch.channels_last_3d) + wrong_dtype = torch.empty( + problem.output_shape, + device="cuda", + dtype=torch.float32, + memory_format=torch.channels_last_3d, + ) with pytest.raises(ValueError, match="dtype"): _fwd(problem, ops, out=wrong_dtype) @@ -777,9 +866,13 @@ def test_an_illegal_tap_block_is_refused_rather_than_run(): _fwd(problem, ops, config=bad) # And an outright illegal MFMA config is refused by the inherited rules. with pytest.raises(ValueError, match="nonkdim"): - _fwd(problem, ops, - config=TransposedConfig(BLOCK_M=64, BLOCK_N=64, BLOCK_K=32, - matrix_instr_nonkdim=64)) + _fwd( + problem, + ops, + config=TransposedConfig( + BLOCK_M=64, BLOCK_N=64, BLOCK_K=32, matrix_instr_nonkdim=64 + ), + ) @requires_gpu @@ -795,8 +888,7 @@ def test_ncdhw_input_is_converted_rather_than_misread(): expected = _reference(problem, ops, "fwd").to(torch.bfloat16) plain = ops["input"].contiguous() assert not plain.is_contiguous(memory_format=torch.channels_last_3d) - got = conv_transpose3d_forward(plain, ops["weight"], ops["bias"], - problem.stride) + got = conv_transpose3d_forward(plain, ops["weight"], ops["bias"], problem.stride) assert got.is_contiguous(memory_format=torch.channels_last_3d) assert reference.compare(got, expected).bitwise @@ -808,8 +900,7 @@ def test_output_matches_torchs_shape_and_layout(): problem = _problem("shape", 32, 16, (3, 4, 5), k, bias=True) ops = _ops(problem, seed=47) got = _fwd(problem, ops) - want = F.conv_transpose3d(ops["input"], ops["weight"], ops["bias"], - stride=k) + want = F.conv_transpose3d(ops["input"], ops["weight"], ops["bias"], stride=k) assert got.shape == want.shape, k assert got.is_contiguous(memory_format=torch.channels_last_3d) assert tuple(got.shape[2:]) == problem.out_spatial @@ -825,17 +916,22 @@ def test_no_worse_than_miopen(): ``roundings`` is 2 for MIOpen's backward-weight because that direction reduces with atomics and disagrees with itself bitwise between two calls. """ - for problem in [_problem("mi", 64, 32, (4, 5, 6), bias=True), - _problem("mi3", 32, 32, (3, 4, 5), (3, 3, 3))]: + for problem in [ + _problem("mi", 64, 32, (4, 5, 6), bias=True), + _problem("mi3", 32, 32, (3, 4, 5), (3, 3, 3)), + ]: for direction in ("fwd", "bwd-data", "bwd-weight"): ops = reference.make_inputs(problem, seed=53) expected = _reference(problem, ops, direction) incumbent = reference.compare( - reference.incumbent(problem, ops, direction), expected) - actual = {"fwd": _fwd, "bwd-data": _bwd_data, - "bwd-weight": _bwd_weight}[direction](problem, ops) - reference.assert_close(actual, expected, problem, direction, - incumbent_error=incumbent) + reference.incumbent(problem, ops, direction), expected + ) + actual = {"fwd": _fwd, "bwd-data": _bwd_data, "bwd-weight": _bwd_weight}[ + direction + ](problem, ops) + reference.assert_close( + actual, expected, problem, direction, incumbent_error=incumbent + ) @requires_gpu @@ -866,22 +962,24 @@ def test_grad_weight_buffer_has_the_transposed_shape(): dimension. """ from triton_conv3d.reduce_gemm import grad_weight_empty - gw = grad_transposed_weight_empty(128, 64, (2, 2, 2), - dtype=torch.bfloat16, device="cuda") + + gw = grad_transposed_weight_empty( + 128, 64, (2, 2, 2), dtype=torch.bfloat16, device="cuda" + ) assert tuple(gw.shape) == (128, 64, 2, 2, 2) assert gw.is_contiguous(memory_format=torch.channels_last_3d) problem = _problem("gw", 32, 16, (3, 4, 5)) ops = _ops(problem, seed=61, direction="bwd-weight") - out = grad_transposed_weight_empty(32, 16, (2, 2, 2), - dtype=torch.bfloat16, device="cuda") + out = grad_transposed_weight_empty( + 32, 16, (2, 2, 2), dtype=torch.bfloat16, device="cuda" + ) got = _bwd_weight(problem, ops, out=out) assert got.data_ptr() == out.data_ptr() expected = _reference(problem, ops, "bwd-weight").to(torch.bfloat16) assert reference.compare(got, expected).bitwise - wrong = grad_weight_empty(32, 16, (2, 2, 2), dtype=torch.bfloat16, - device="cuda") + wrong = grad_weight_empty(32, 16, (2, 2, 2), dtype=torch.bfloat16, device="cuda") assert tuple(wrong.shape) == (32, 16, 2, 2, 2) other = _problem("gw2", 16, 32, (3, 4, 5)) ops2 = _ops(other, seed=61, direction="bwd-weight") @@ -897,15 +995,16 @@ def test_unsupported_calls_raise_rather_than_return_garbage(): with pytest.raises(NotImplementedError): conv_transpose3d_forward(ops["input"], ops["weight"], None, 3) with pytest.raises(NotImplementedError): - conv_transpose3d_backward_data(ops["grad_output"], ops["weight"], - problem.input_shape, 3) + conv_transpose3d_backward_data( + ops["grad_output"], ops["weight"], problem.input_shape, 3 + ) with pytest.raises(NotImplementedError): - conv_transpose3d_backward_weight(ops["input"], problem.weight_shape, - ops["grad_output"], 3) + conv_transpose3d_backward_weight( + ops["input"], problem.weight_shape, ops["grad_output"], 3 + ) # And a padding, which is the one a caller is most likely to pass by habit. with pytest.raises(NotImplementedError): - conv_transpose3d_forward(ops["input"], ops["weight"], None, - problem.stride, 1) + conv_transpose3d_forward(ops["input"], ops["weight"], None, problem.stride, 1) @requires_gpu @@ -939,9 +1038,11 @@ def test_bias_is_per_channel_and_not_per_column(): """ problem = _problem("bias", 64, 48, (3, 4, 5), bias=True) ops = _ops(problem, seed=73) - ops = {**ops, "weight": torch.zeros_like(ops["weight"]), - "bias": torch.arange(1, 49, device="cuda", - dtype=torch.bfloat16)} + ops = { + **ops, + "weight": torch.zeros_like(ops["weight"]), + "bias": torch.arange(1, 49, device="cuda", dtype=torch.bfloat16), + } got = _fwd(problem, ops) want = ops["bias"].view(1, 48, 1, 1, 1).expand(got.shape) assert torch.equal(got, want.to(got.dtype)) diff --git a/triton_conv3d/transposed.py b/triton_conv3d/transposed.py index 1b1681b..2b2647c 100644 --- a/triton_conv3d/transposed.py +++ b/triton_conv3d/transposed.py @@ -96,7 +96,6 @@ from __future__ import annotations import dataclasses -import math from typing import Sequence import torch @@ -104,14 +103,16 @@ import triton.language as tl from .gather_gemm import ( - ConvConfig, _LDS_BYTES, _MFMA_KDIM, + ConvConfig, _check_out, _index_dtype, _pow2_at_most, _triple, conv3d_forward, +) +from .gather_gemm import ( is_supported as _is_supported_fwd, ) from .reduce_gemm import conv3d_backward_weight, is_supported_bwd_weight @@ -142,25 +143,50 @@ @triton.jit def _convT3d_fwd_kernel( - X, W, Y, BIAS, + X, + W, + Y, + BIAS, # Sizes. ``M_TOTAL`` is ``BATCH * IN_D * IN_H * IN_W`` -- the *input* # volume, because that is what a scatter is indexed by. - BATCH, IN_D, IN_H, IN_W, - CIN, COUT, M_TOTAL, + BATCH, + IN_D, + IN_H, + IN_W, + CIN, + COUT, + M_TOTAL, # Element strides. The channel stride of X and Y is 1 by construction -- # that is what NDHWC means -- so it is neither passed nor multiplied by. - stride_xn, stride_xd, stride_xh, stride_xw, + stride_xn, + stride_xd, + stride_xh, + stride_xw, # The weight over the effective GEMM's axes: the fused tap index, the # reduction axis K (Cin), and the output axis N (Cout). Which of the two # channel strides is 1 is a constexpr (``W_ORDER``), as in the forward. - stride_wt, stride_wk, stride_wn, - stride_yn, stride_yd, stride_yh, stride_yw, - KD: tl.constexpr, KH: tl.constexpr, KW: tl.constexpr, - BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_NC: tl.constexpr, - BLOCK_K: tl.constexpr, BLOCK_K_COUNT: tl.constexpr, - TAP_BLOCK: tl.constexpr, GROUP_M: tl.constexpr, - HAS_BIAS: tl.constexpr, EVEN_K: tl.constexpr, EVEN_N: tl.constexpr, - INDEX_DTYPE: tl.constexpr, INPUT_PRECISION: tl.constexpr, + stride_wt, + stride_wk, + stride_wn, + stride_yn, + stride_yd, + stride_yh, + stride_yw, + KD: tl.constexpr, + KH: tl.constexpr, + KW: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_NC: tl.constexpr, + BLOCK_K: tl.constexpr, + BLOCK_K_COUNT: tl.constexpr, + TAP_BLOCK: tl.constexpr, + GROUP_M: tl.constexpr, + HAS_BIAS: tl.constexpr, + EVEN_K: tl.constexpr, + EVEN_N: tl.constexpr, + INDEX_DTYPE: tl.constexpr, + INPUT_PRECISION: tl.constexpr, W_ORDER: tl.constexpr, ): # -- which tile this program owns -------------------------------------- @@ -252,10 +278,7 @@ def _convT3d_fwd_kernel( if W_ORDER == 0: w_col = tap.to(INDEX_DTYPE) * stride_wt + offs_n.to(INDEX_DTYPE) else: - w_col = ( - tap.to(INDEX_DTYPE) * stride_wt - + offs_n.to(INDEX_DTYPE) * stride_wn - ) + w_col = tap.to(INDEX_DTYPE) * stride_wt + offs_n.to(INDEX_DTYPE) * stride_wn acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) @@ -383,8 +406,9 @@ def _fit_transposed( while cfg.lds_bytes(dtype) > _LDS_BYTES: half_k, half_m, half_n = cfg.BLOCK_K // 2, cfg.BLOCK_M // 2, cfg.BLOCK_N // 2 if half_k >= kdim and half_k % kdim == 0: - cfg = dataclasses.replace(cfg, BLOCK_K=half_k, - kpack=1 if half_k <= 16 else cfg.kpack) + cfg = dataclasses.replace( + cfg, BLOCK_K=half_k, kpack=1 if half_k <= 16 else cfg.kpack + ) elif half_m >= nk and half_m % nk == 0: cfg = dataclasses.replace(cfg, BLOCK_M=half_m) elif half_n >= cfg.TAP_BLOCK * nk and half_n % nk == 0: @@ -394,8 +418,8 @@ def _fit_transposed( while ( cfg.BLOCK_M > max(16, nk) and (cfg.BLOCK_M // 2) % nk == 0 - and (-(-m // cfg.BLOCK_M) * -(-cout // cfg.BLOCK_NC) - * (taps // cfg.TAP_BLOCK)) < _MIN_PROGRAMS + and (-(-m // cfg.BLOCK_M) * -(-cout // cfg.BLOCK_NC) * (taps // cfg.TAP_BLOCK)) + < _MIN_PROGRAMS ): cfg = dataclasses.replace(cfg, BLOCK_M=cfg.BLOCK_M // 2) warps = max(1, min(cfg.num_warps, cfg.BLOCK_M * cfg.BLOCK_N // 256)) @@ -432,26 +456,42 @@ def default_transposed_config( block_n = tap_block * block_nc return _fit_transposed( TransposedConfig( - BLOCK_M=block_m, BLOCK_N=block_n, BLOCK_K=block_k, GROUP_M=6, + BLOCK_M=block_m, + BLOCK_N=block_n, + BLOCK_K=block_k, + GROUP_M=6, num_warps=8 if block_n >= 256 or block_k >= 128 else 4, - num_stages=2, matrix_instr_nonkdim=nonkdim, - kpack=1 if block_k <= 16 else 2, TAP_BLOCK=tap_block, + num_stages=2, + matrix_instr_nonkdim=nonkdim, + kpack=1 if block_k <= 16 else 2, + TAP_BLOCK=tap_block, ), - m, cout, taps, dtype, + m, + cout, + taps, + dtype, ) -def _tuned(bm: int, bnc: int, tb: int, bk: int, warps: int, - group_m: int = 6) -> TransposedConfig: +def _tuned( + bm: int, bnc: int, tb: int, bk: int, warps: int, group_m: int = 6 +) -> TransposedConfig: return TransposedConfig( - BLOCK_M=bm, BLOCK_N=tb * bnc, BLOCK_K=bk, GROUP_M=group_m, - num_warps=warps, num_stages=2, matrix_instr_nonkdim=16, - kpack=1 if bk <= 16 else 2, TAP_BLOCK=tb, + BLOCK_M=bm, + BLOCK_N=tb * bnc, + BLOCK_K=bk, + GROUP_M=group_m, + num_warps=warps, + num_stages=2, + matrix_instr_nonkdim=16, + kpack=1 if bk <= 16 else 2, + TAP_BLOCK=tb, ) -def transposed_tune_key(dtype: torch.dtype, cin: int, cout: int, - kernel: tuple[int, ...]) -> tuple: +def transposed_tune_key( + dtype: torch.dtype, cin: int, cout: int, kernel: tuple[int, ...] +) -> tuple: return (str(dtype), cin, cout, tuple(kernel)) @@ -498,7 +538,10 @@ def register_tuned_transposed(dtype, cin, cout, kernel, config) -> None: def transposed_config( - m: int, cin: int, cout: int, kernel: Sequence[int], + m: int, + cin: int, + cout: int, + kernel: Sequence[int], dtype: torch.dtype = torch.bfloat16, ) -> TransposedConfig: """The config :func:`conv_transpose3d_forward` would pick for this problem.""" @@ -530,9 +573,14 @@ def transposed_config( def candidate_transposed_configs( - m: int, cin: int, cout: int, taps: int, + m: int, + cin: int, + cout: int, + taps: int, dtype: torch.dtype = torch.bfloat16, - *, tap_blocks: Sequence[int] = (1, 2, 4, 8), nonkdims: Sequence[int] = (16, 32), + *, + tap_blocks: Sequence[int] = (1, 2, 4, 8), + nonkdims: Sequence[int] = (16, 32), ) -> list[TransposedConfig]: """Configs worth timing for one transposed problem, pruned to legal ones.""" n2 = max(16, triton.next_power_of_2(cout)) @@ -549,14 +597,21 @@ def candidate_transposed_configs( for warps in {4, 8, seed_warps}: for nonkdim in nonkdims: cfg = TransposedConfig( - BLOCK_M=bm, BLOCK_N=tb * bnc, BLOCK_K=bk, GROUP_M=6, - num_warps=warps, num_stages=2, + BLOCK_M=bm, + BLOCK_N=tb * bnc, + BLOCK_K=bk, + GROUP_M=6, + num_warps=warps, + num_stages=2, matrix_instr_nonkdim=nonkdim, - kpack=1 if bk <= 16 else 2, TAP_BLOCK=tb, + kpack=1 if bk <= 16 else 2, + TAP_BLOCK=tb, ) - if (cfg.validate(dtype) is not None - or cfg.lds_bytes(dtype) > _LDS_BYTES - or cfg in seen): + if ( + cfg.validate(dtype) is not None + or cfg.lds_bytes(dtype) > _LDS_BYTES + or cfg in seen + ): continue seen.add(cfg) out.append(cfg) @@ -623,9 +678,11 @@ def _transposed_weight_plan(w: torch.Tensor) -> tuple[int, int, int, int] | None st = s[2] else: st = 0 # one tap: ``tap`` is always 0, so any stride is the right one - if ((kw > 1 and s[4] != st) - or (kh > 1 and s[3] != st * kw) - or (kd > 1 and s[2] != st * kw * kh)): + if ( + (kw > 1 and s[4] != st) + or (kh > 1 and s[3] != st * kw) + or (kd > 1 and s[2] != st * kw * kh) + ): return None if cout == 1 or s[1] == 1: return (_W_N_CONTIG, st, s[0], 1) @@ -649,7 +706,12 @@ def _transposed_out_spatial( def _transposed_shape_ok( - x: torch.Tensor, w: torch.Tensor, stride, padding, output_padding, dilation, + x: torch.Tensor, + w: torch.Tensor, + stride, + padding, + output_padding, + dilation, groups: int, ) -> tuple[int, int, int] | None: """The kernel triple if this is a ``kernel == stride`` upsample, else ``None``. @@ -724,8 +786,7 @@ def is_supported_transposed( of the ordinary convolution were also expected to agree until ``stride > 1`` showed that they do not. """ - k = _transposed_shape_ok(x, w, stride, padding, output_padding, dilation, - groups) + k = _transposed_shape_ok(x, w, stride, padding, output_padding, dilation, groups) if k is None: return False if x.dtype != w.dtype or x.dtype not in _MFMA_KDIM: @@ -745,9 +806,14 @@ def is_supported_transposed( # the right length silently applies every other value. # ``torch.conv_transpose3d`` rejects both; so does this. ``Cout`` is # ``w.shape[1]`` here, not ``w.shape[0]``. - if (bias.dim() != 1 or int(bias.shape[0]) != int(w.shape[1]) - or bias.dtype != x.dtype or not bias.is_cuda - or bias.device != x.device or bias.stride(0) != 1): + if ( + bias.dim() != 1 + or int(bias.shape[0]) != int(w.shape[1]) + or bias.dtype != x.dtype + or not bias.is_cuda + or bias.device != x.device + or bias.stride(0) != 1 + ): return False return True @@ -771,8 +837,9 @@ def is_supported_transposed_bwd_data( ``kernel == stride`` upsample at all, and that ``input_shape`` is the shape this ``grad_output`` came from. """ - k = _transposed_shape_ok(grad_output, w, stride, padding, output_padding, - dilation, groups) + k = _transposed_shape_ok( + grad_output, w, stride, padding, output_padding, dilation, groups + ) if k is None: return False try: @@ -837,9 +904,7 @@ def is_supported_transposed_bwd_weight( return False if any(int(v) < 1 for v in x.shape[2:]): return False - return bool( - is_supported_bwd_weight(grad_output, (ws[0], ws[1], *k), x, k, 0, 1, 1) - ) + return bool(is_supported_bwd_weight(grad_output, (ws[0], ws[1], *k), x, k, 0, 1, 1)) def is_supported_transposed_all( @@ -868,8 +933,9 @@ def is_supported_transposed_all( zeros and answer ``False`` -- a fallback to the caller's other kernel, which is the safe direction. """ - if not is_supported_transposed(x, w, bias, stride, padding, output_padding, - dilation, groups): + if not is_supported_transposed( + x, w, bias, stride, padding, output_padding, dilation, groups + ): return False k = tuple(int(v) for v in w.shape[2:]) grad_shape = (int(x.shape[0]), int(w.shape[1])) + _transposed_out_spatial( @@ -882,7 +948,13 @@ def is_supported_transposed_all( return False return bool( is_supported_transposed_bwd_weight( - x, tuple(w.shape), grad, stride, padding, output_padding, dilation, + x, + tuple(w.shape), + grad, + stride, + padding, + output_padding, + dilation, groups, ) ) @@ -915,8 +987,9 @@ def conv_transpose3d_forward( out-of-bounds device write with no error and an NCDHW one is a full-rate kernel returning a scrambled answer. """ - if not is_supported_transposed(x, w, bias, stride, padding, output_padding, - dilation, groups): + if not is_supported_transposed( + x, w, bias, stride, padding, output_padding, dilation, groups + ): raise NotImplementedError( f"unsupported: x={tuple(x.shape)}/{x.dtype} w={tuple(w.shape)} " f"stride={stride} padding={padding} output_padding={output_padding} " @@ -929,8 +1002,7 @@ def conv_transpose3d_forward( x = x.contiguous(memory_format=torch.channels_last_3d) n, cin, in_d, in_h, in_w = (int(v) for v in x.shape) cout = int(w.shape[1]) - out_d, out_h, out_w = _transposed_out_spatial((in_d, in_h, in_w), - (kd, kh, kw)) + out_d, out_h, out_w = _transposed_out_spatial((in_d, in_h, in_w), (kd, kh, kw)) y_shape = (n, cout, out_d, out_h, out_w) if out is None: @@ -938,8 +1010,12 @@ def conv_transpose3d_forward( # it ``torch.empty(shape).contiguous(memory_format=...)`` allocates NCDHW # and then copies the whole thing -- 235x, measured on the gather # kernel's identically-shaped defect. - y = torch.empty(y_shape, device=x.device, dtype=x.dtype, - memory_format=torch.channels_last_3d) + y = torch.empty( + y_shape, + device=x.device, + dtype=x.dtype, + memory_format=torch.channels_last_3d, + ) else: y = out _check_out(y, y_shape, x) @@ -973,17 +1049,38 @@ def conv_transpose3d_forward( * (taps // config.TAP_BLOCK), ) _convT3d_fwd_kernel[grid]( - x, wt, y, bias, - n, in_d, in_h, in_w, - cin, cout, m_total, - x.stride(0), x.stride(2), x.stride(3), x.stride(4), - plan[1], plan[2], plan[3], - y.stride(0), y.stride(2), y.stride(3), y.stride(4), - KD=kd, KH=kh, KW=kw, - BLOCK_M=config.BLOCK_M, BLOCK_N=config.BLOCK_N, - BLOCK_NC=config.BLOCK_NC, BLOCK_K=config.BLOCK_K, + x, + wt, + y, + bias, + n, + in_d, + in_h, + in_w, + cin, + cout, + m_total, + x.stride(0), + x.stride(2), + x.stride(3), + x.stride(4), + plan[1], + plan[2], + plan[3], + y.stride(0), + y.stride(2), + y.stride(3), + y.stride(4), + KD=kd, + KH=kh, + KW=kw, + BLOCK_M=config.BLOCK_M, + BLOCK_N=config.BLOCK_N, + BLOCK_NC=config.BLOCK_NC, + BLOCK_K=config.BLOCK_K, BLOCK_K_COUNT=triton.cdiv(cin, config.BLOCK_K), - TAP_BLOCK=config.TAP_BLOCK, GROUP_M=config.GROUP_M, + TAP_BLOCK=config.TAP_BLOCK, + GROUP_M=config.GROUP_M, HAS_BIAS=bias is not None, EVEN_K=(cin % config.BLOCK_K == 0), EVEN_N=(cout % config.BLOCK_NC == 0), @@ -1021,8 +1118,7 @@ def conv_transpose3d_backward_data( a reduction of ``grad_output`` and not part of this direction at all. """ if not is_supported_transposed_bwd_data( - grad_output, w, input_shape, stride, padding, output_padding, dilation, - groups + grad_output, w, input_shape, stride, padding, output_padding, dilation, groups ): raise NotImplementedError( f"unsupported: grad_output={tuple(grad_output.shape)}/" @@ -1032,8 +1128,7 @@ def conv_transpose3d_backward_data( f"dilation={dilation} groups={groups}" ) k = tuple(int(v) for v in w.shape[2:]) - return conv3d_forward(grad_output, w, None, k, 0, 1, 1, config=config, - out=out) + return conv3d_forward(grad_output, w, None, k, 0, 1, 1, config=config, out=out) def conv_transpose3d_backward_weight( @@ -1065,8 +1160,7 @@ def conv_transpose3d_backward_weight( in, so the optimizer's elementwise update is contiguous. """ if not is_supported_transposed_bwd_weight( - x, weight_shape, grad_output, stride, padding, output_padding, dilation, - groups + x, weight_shape, grad_output, stride, padding, output_padding, dilation, groups ): raise NotImplementedError( f"unsupported: x={tuple(x.shape)}/{x.dtype} " @@ -1078,8 +1172,16 @@ def conv_transpose3d_backward_weight( ws = tuple(int(v) for v in weight_shape) k = ws[2:] return conv3d_backward_weight( - grad_output, ws, x, k, 0, 1, 1, - config=config, workspace=workspace, out=out, + grad_output, + ws, + x, + k, + 0, + 1, + 1, + config=config, + workspace=workspace, + out=out, deterministic=deterministic, ) @@ -1098,8 +1200,12 @@ def grad_transposed_weight_empty( which ``conv3d_backward_weight``'s ``out=`` check catches. """ k = _triple(kernel, "kernel") - return torch.empty((cin, cout, *k), dtype=dtype, device=device, - memory_format=torch.channels_last_3d) + return torch.empty( + (cin, cout, *k), + dtype=dtype, + device=device, + memory_format=torch.channels_last_3d, + ) # --------------------------------------------------------------------------- @@ -1133,11 +1239,10 @@ def verify_isa_transposed( w = w.contiguous(memory_format=torch.channels_last_3d) elif weight_layout != "tkn": raise ValueError(f"unknown weight_layout {weight_layout!r}") - x = torch.randn((n, cin, d, h, wd), device="cuda", - dtype=torch.bfloat16).contiguous( - memory_format=torch.channels_last_3d) - cfg = config or transposed_config(n * d * h * wd, cin, cout, k, - torch.bfloat16) + x = torch.randn((n, cin, d, h, wd), device="cuda", dtype=torch.bfloat16).contiguous( + memory_format=torch.channels_last_3d + ) + cfg = config or transposed_config(n * d * h * wd, cin, cout, k, torch.bfloat16) y = conv_transpose3d_forward(x, w, None, k, config=cfg) torch.cuda.synchronize() print( From 84e5624b8bec3e69f1f51cddb82e27c6f5903daf Mon Sep 17 00:00:00 2001 From: Nikoli Dryden Date: Thu, 6 Aug 2026 14:47:18 -0700 Subject: [PATCH 8/8] Say which kernel each accelerated module ended up using A run gives no sign of whether the Triton rungs actually served it. Print one line per ladder on rank 0 at startup, naming Triton against everything PyTorch does -- compiled and eager together, since from outside the ladder those are the same answer. Placement is the whole of the problem. ``_triton_ok`` is a latch set when a rung first answers a call, so reporting at construction time would say "Native" about modules that have not run yet. Reporting at the end of warmup is right when there is a warmup, but ``warmup_batches <= 0`` returns before running anything, which is what the benchmark drivers configure -- so a single call site would drop the line on exactly the runs most likely to want it. Both sites call it behind a one-shot flag. The reporter is duck-typed on the latch rather than on isinstance: _rungs is imported by the modules it would otherwise have to import back, and a ladder added later is then reported without touching this code. Labels come from a _rung_label class attribute and fall back to the class name, so a ladder that does not declare one still appears. Deliberately one rank's answer and not a collective: ranks latch independently and a rung can still fall back later, so this is informational rather than a contract, and gathering it would put a barrier on a path with no other reason for one. --- ScaFFold/unet/_rungs.py | 65 ++++++++++++++++++++++++++++ ScaFFold/unet/conv3d.py | 9 ++++ ScaFFold/unet/group_norm.py | 6 +++ ScaFFold/utils/trainer.py | 32 ++++++++++++++ tests/test_platform_guard.py | 82 ++++++++++++++++++++++++++++++++++++ 5 files changed, 194 insertions(+) diff --git a/ScaFFold/unet/_rungs.py b/ScaFFold/unet/_rungs.py index acaaa8e..3dc6e19 100644 --- a/ScaFFold/unet/_rungs.py +++ b/ScaFFold/unet/_rungs.py @@ -422,3 +422,68 @@ def _warn_rung_failure(what, error, fallback, env_var): f"{fallback} for modules that have not already used it. " f"Set {env_var}=0 to skip this attempt entirely." ) + + +#: The per-module "this rung has served me" flag every ladder carries. Read by +#: name rather than by ``isinstance`` so this module does not have to import the +#: modules that import it, and so a ladder added later is reported without +#: touching this function. +_RUNG_FLAG = "_triton_ok" + + +def kernel_selection(model): + """Which kernel each rung-bearing module in ``model`` is currently using. + + Returns ``[(label, triton, total), ...]``, one entry per kind of ladder, + ordered as the modules appear in the model. ``label`` comes from the + class's ``_rung_label`` and falls back to its name, so a new ladder shows up + here whether or not it remembers to declare one. + + **This reads a latch, not a decision.** ``_triton_ok`` is set on the first + call a rung actually answers, so a module that has not run yet reports + ``Native`` -- truthfully, in that it has used no other kernel, but + misleadingly, in that it has used none at all. Callers must therefore run + at least one forward first; :meth:`PyTorchTrainer._log_kernel_selection` + does, and says in its own docstring why it is placed where it is. + + Nothing here is authoritative for the rest of the run either. A rung can + fall back later (a shape it does not serve, an allocator failure), and under + DDP each rank latches independently, so this is one rank's answer at one + moment. It is an informational line, not a contract. + """ + counts = {} + order = [] + for module in model.modules(): + cls = type(module) + if not hasattr(cls, _RUNG_FLAG): + continue + label = getattr(cls, "_rung_label", cls.__name__) + if label not in counts: + counts[label] = [0, 0] + order.append(label) + counts[label][1] += 1 + if getattr(module, _RUNG_FLAG, False): + counts[label][0] += 1 + return [(label, counts[label][0], counts[label][1]) for label in order] + + +def format_kernel_selection(selection): + """:func:`kernel_selection`'s answer as lines fit for a log. + + One line per ladder, naming both kernels only when the ladder is actually + split -- a mixed line is the interesting case (some sites fell back) and + should not be hidden inside a ratio that reads like a uniform one. + """ + if not selection: + return [" (no accelerated modules)"] + width = max(len(label) for label, _, _ in selection) + lines = [] + for label, triton, total in selection: + if triton == total: + used = f"Triton {triton}/{total}" + elif triton == 0: + used = f"Native {total}/{total}" + else: + used = f"Triton {triton}/{total}, Native {total - triton}/{total}" + lines.append(f" {label.ljust(width)} {used}") + return lines diff --git a/ScaFFold/unet/conv3d.py b/ScaFFold/unet/conv3d.py index ec1a4ed..c5d74b6 100644 --- a/ScaFFold/unet/conv3d.py +++ b/ScaFFold/unet/conv3d.py @@ -1249,6 +1249,10 @@ class FastConv3d(nn.Conv3d): #: free. _triton_ok = False + #: How this ladder is named in the startup kernel-selection line; see + #: :func:`ScaFFold.unet._rungs.kernel_selection`. + _rung_label = "Convolution" + def _triton_forward(self, local, plan=None): """Run the Triton rung -- and its fallback -- on an unwrapped tensor. @@ -1778,6 +1782,11 @@ class FastConvTranspose3d(nn.ConvTranspose3d): #: :attr:`FastConv3d._triton_ok`, which this mirrors exactly. _triton_ok = False + #: Reported separately from ``FastConv3d``: it is a different operator with + #: a different kernel behind it, so a mixed run should say which one fell + #: back rather than pooling both into one count. + _rung_label = "Convolution (transposed)" + def _triton_forward(self, local): """Run the Triton rung on an unwrapped tensor. diff --git a/ScaFFold/unet/group_norm.py b/ScaFFold/unet/group_norm.py index fe0f510..6e0729f 100644 --- a/ScaFFold/unet/group_norm.py +++ b/ScaFFold/unet/group_norm.py @@ -630,6 +630,12 @@ class FastGroupNorm(nn.GroupNorm): _triton_ok = False _compiled_ok = False + #: How this ladder is named in the startup kernel-selection line. The line + #: reports Triton against everything else, so ``_compiled_ok`` does not + #: appear there: from the outside the compiled and eager rungs are both + #: "what PyTorch does". + _rung_label = "GroupNorm" + def __init__( self, num_groups, diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py index b284708..201e09f 100644 --- a/ScaFFold/utils/trainer.py +++ b/ScaFFold/utils/trainer.py @@ -29,6 +29,7 @@ from torch.utils.data import DataLoader from tqdm import tqdm +from ScaFFold.unet._rungs import format_kernel_selection, kernel_selection from ScaFFold.utils.checkpointing import CheckpointManager from ScaFFold.utils.data_loading import FractalDataset, SpatialShardSpec from ScaFFold.utils.data_types import AMP_DTYPE, VOLUME_TORCH_DTYPE @@ -90,6 +91,9 @@ class BaseTrainer: def __init__(self, model, config, device, log): self.model = model + # One-shot guard for the startup kernel-selection line; see + # _log_kernel_selection for why it has two call sites. + self._kernel_selection_logged = False self.config = config self.device = device self.log = log @@ -786,6 +790,31 @@ def _warmup_ragged_batches(self, batch): log_prefix=f"warmup ragged ({ragged}): ", ) + def _log_kernel_selection(self): + """Log which kernel each accelerated module is using, once, on rank 0. + + Called from two places -- the end of :meth:`warmup` and after the first + training batch -- because the answer only exists once a forward has run. + ``_triton_ok`` is a latch set when a rung first answers a call, so + reporting at construction time would say "Native" about modules that + have simply not run yet. ``warmup_batches <= 0`` makes :meth:`warmup` + return without running anything, so neither call site alone covers every + run; the flag below makes the pair idempotent rather than making either + one conditional on the other. + + Rank 0 only, and one rank's answer: each rank latches independently, so + under DDP this is representative rather than global. It is an + informational line and is deliberately not a collective -- gathering it + would put a barrier on a path that has no other reason for one. + """ + if self._kernel_selection_logged or self.world_rank != 0: + return + self._kernel_selection_logged = True + model = getattr(self.model, "module", self.model) + self.log.info("Kernel selection (rank 0):") + for line in format_kernel_selection(kernel_selection(model)): + self.log.info(line) + def warmup(self): """Run warmup iterations before the main training loop.""" warmup_batches = self.config.warmup_batches @@ -847,6 +876,7 @@ def warmup(self): torch.distributed.barrier() self.log.info(f"Done warmup. Took {int(time.time() - start_warmup)}s") + self._log_kernel_selection() def train(self, profiler=None): """ @@ -953,6 +983,8 @@ def train(self, profiler=None): # not skew the epoch mean. train_dice_total += batch_dice_score * batch_size end_code_region("run_training_batch") + if first_batch: + self._log_kernel_selection() # Update the loss begin_code_region("update_loss") diff --git a/tests/test_platform_guard.py b/tests/test_platform_guard.py index 36e6a13..f88dc1e 100644 --- a/tests/test_platform_guard.py +++ b/tests/test_platform_guard.py @@ -46,6 +46,7 @@ from ScaFFold.unet import _rungs from ScaFFold.unet import conv3d as conv_mod from ScaFFold.unet import group_norm as gn_mod +from ScaFFold.unet._rungs import format_kernel_selection, kernel_selection from ScaFFold.unet.conv3d import FastConv3d, FastConvTranspose3d from ScaFFold.unet.group_norm import FastGroupNorm @@ -500,3 +501,84 @@ def test_a_whole_unet_step_asks_the_hardware_once(monkeypatch): with torch.no_grad(): model(x) assert asked == [0] + + +# --------------------------------------------------------------------------- +# The startup kernel-selection line +# --------------------------------------------------------------------------- + + +class _Ladder(torch.nn.Module): + """Stands in for a rung-bearing module: the reporter is duck-typed.""" + + _triton_ok = False + _rung_label = "Ladder" + + +class _OtherLadder(torch.nn.Module): + _triton_ok = False + _rung_label = "Other" + + +class _Unlabelled(torch.nn.Module): + _triton_ok = False + + +def test_kernel_selection_counts_each_ladder_separately(): + model = torch.nn.Sequential(_Ladder(), _Ladder(), _OtherLadder(), torch.nn.ReLU()) + model[0]._triton_ok = True + assert kernel_selection(model) == [("Ladder", 1, 2), ("Other", 0, 1)] + + +def test_kernel_selection_ignores_modules_without_a_rung(): + """A plain module must not appear -- the line is about ladders only.""" + model = torch.nn.Sequential(torch.nn.ReLU(), torch.nn.Identity()) + assert kernel_selection(model) == [] + + +def test_a_ladder_without_a_label_still_reports_under_its_class_name(): + """Adding a ladder must not require remembering to declare a label.""" + assert kernel_selection(torch.nn.Sequential(_Unlabelled())) == [ + ("_Unlabelled", 0, 1) + ] + + +def test_a_split_ladder_names_both_kernels(): + """The mixed case is the informative one and must not read as uniform.""" + line = format_kernel_selection([("Convolution", 17, 19)])[0] + assert "Triton 17/19" in line and "Native 2/19" in line + + +@pytest.mark.parametrize( + "selection,expected", + [([("C", 3, 3)], "Triton"), ([("C", 0, 3)], "Native")], +) +def test_an_unsplit_ladder_names_one_kernel(selection, expected): + line = format_kernel_selection(selection)[0] + assert expected in line + assert ("Native" if expected == "Triton" else "Triton") not in line + + +@pytest.mark.gpu +def test_the_real_model_reports_triton_on_every_site_after_a_forward(): + """The shipped configuration is all-Triton, and the line must say so. + + Also pins the placement rule: the same model reports ``Native`` everywhere + *before* a forward, because ``_triton_ok`` is a latch. That is why + ``_log_kernel_selection`` is called after warmup and after the first batch + rather than at construction. + """ + from ScaFFold.unet.unet_model import UNet + + model = UNet(n_channels=3, n_classes=6, trilinear=False, layers=3) + model = model.cuda().to(memory_format=_CHANNELS_LAST) + assert all(triton == 0 for _, triton, _ in kernel_selection(model)) + + x = _gpu_input((1, 3, 32, 32, 32), dtype=torch.float32) + with torch.autocast("cuda", dtype=torch.bfloat16), torch.no_grad(): + model(x) + + selection = kernel_selection(model) + labels = {label for label, _, _ in selection} + assert labels == {"Convolution", "Convolution (transposed)", "GroupNorm"} + assert all(triton == total for _, triton, total in selection), selection