diff --git a/ScaFFold/unet/_rungs.py b/ScaFFold/unet/_rungs.py new file mode 100644 index 0000000..3dc6e19 --- /dev/null +++ b/ScaFFold/unet/_rungs.py @@ -0,0 +1,489 @@ +# 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." + ) + + +#: 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 new file mode 100644 index 0000000..c5d74b6 --- /dev/null +++ b/ScaFFold/unet/conv3d.py @@ -0,0 +1,1873 @@ +# 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. +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 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 +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 (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 + + #: 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. + + ``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 (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, 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 + + #: 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. + + 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/group_norm.py b/ScaFFold/unet/group_norm.py index 06d7407..6e0729f 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,249 @@ 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 + + #: 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, + 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/triton_group_norm.py b/ScaFFold/unet/triton_group_norm.py new file mode 100644 index 0000000..2dd28e9 --- /dev/null +++ b/ScaFFold/unet/triton_group_norm.py @@ -0,0 +1,1816 @@ +# 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; 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. The two were tuned jointly by coordinate descent, so +do not change one without re-running the other. + +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. + +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 -- 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. + +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/ScaFFold/unet/unet_parts.py b/ScaFFold/unet/unet_parts.py index c9e6cb0..9483b17 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.{}") @@ -28,30 +29,128 @@ _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) + + +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""" + """(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__() 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), - _group_norm(group_norm_groups, mid_channels), - nn.ReLU(inplace=True), - nn.Conv3d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False), - _group_norm(group_norm_groups, out_channels), - nn.ReLU(inplace=True), + _conv3d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False), + _group_norm(group_norm_groups, mid_channels, activation="relu"), + nn.Identity(), + _conv3d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False), + _group_norm(group_norm_groups, out_channels, activation="relu"), + nn.Identity(), ) @_doubleconv_annotate @@ -75,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. + """ def __init__(self, in_channels, out_channels, group_norm_groups, trilinear=True): super().__init__() @@ -90,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) @@ -118,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/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/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/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 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_conv3d.py b/tests/test_conv3d.py new file mode 100644 index 0000000..9a27170 --- /dev/null +++ b/tests/test_conv3d.py @@ -0,0 +1,1657 @@ +# 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 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 +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 a separate harness that needs real ranks. + """ + 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, needs real ranks and + lives in a separate harness. + """ + 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)``. 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_groupnorm.py b/tests/test_groupnorm.py index b1e8e6a..4b1cb05 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). + """ + 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. +# +# **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..86a6597 --- /dev/null +++ b/tests/test_groupnorm_wiring_edge.py @@ -0,0 +1,1191 @@ +# 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 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. + +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..f88dc1e --- /dev/null +++ b/tests/test_platform_guard.py @@ -0,0 +1,584 @@ +# 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. + +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 + +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._rungs import format_kernel_selection, kernel_selection +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 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 + 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] + + +# --------------------------------------------------------------------------- +# 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 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..4cc91ef --- /dev/null +++ b/tests/test_triton_group_norm_edge.py @@ -0,0 +1,1700 @@ +# 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. + +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()) 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 diff --git a/triton_conv3d/__init__.py b/triton_conv3d/__init__.py new file mode 100644 index 0000000..795175a --- /dev/null +++ b/triton_conv3d/__init__.py @@ -0,0 +1,120 @@ +# 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..9cca128 --- /dev/null +++ b/triton_conv3d/bench/__init__.py @@ -0,0 +1,20 @@ +# 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..d7b94ce --- /dev/null +++ b/triton_conv3d/bench/baseline.py @@ -0,0 +1,553 @@ +# 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. 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 + + def fwd(): + return 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,) + + def fn(): + return 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..5a5df07 --- /dev/null +++ b/triton_conv3d/bench/conv_bench.py @@ -0,0 +1,1659 @@ +# 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. + +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. + +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: + +**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, 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 +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 ( + bwd_weight_config, + candidate_bwd_weight_configs, + 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. 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. + #: ``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 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. 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: + + ============================ ======== ======= ============ =============== + 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 -- 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. 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 + 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", + ) + 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} (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..dfc1f6a --- /dev/null +++ b/triton_conv3d/bench/gemm_probe.py @@ -0,0 +1,1267 @@ +# 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 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 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: + # 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 + 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) + # 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")): + 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( + 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), + ) + 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..e5cc8d0 --- /dev/null +++ b/triton_conv3d/bench/harness.py @@ -0,0 +1,1149 @@ +# 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 -- 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. 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``: + +============================= ======= ======== ============= +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``: +#: +#: =========================== ============ +#: 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: 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..6f18e05 --- /dev/null +++ b/triton_conv3d/bwd_data.py @@ -0,0 +1,394 @@ +# 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 ( + _MFMA_KDIM, + ConvConfig, + _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. +#: +#: 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*: +#: +#: * ``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 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 + # 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..fe29f51 --- /dev/null +++ b/triton_conv3d/gather_gemm.py @@ -0,0 +1,1497 @@ +# 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: 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: 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 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 + 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. +#: +#: 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. +#: 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 +#: 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.) +_TUNED: dict[tuple, ConvConfig] = { + # 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 + # 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 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** + # (``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%). + (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 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 + # 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: + # (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. 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 (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. 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..cd4991e --- /dev/null +++ b/triton_conv3d/reduce_gemm.py @@ -0,0 +1,1639 @@ +# 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. 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: 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. 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 ( + _LDS_BYTES, + _MFMA_KDIM, + ConvConfig, + _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 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 + 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; 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: 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. +#: +#: 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. +#: +#: **``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 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 +#: ``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: +#: +#: ============ ========================================= ====== +#: 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. +#: 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): 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 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 + # ``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. 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. 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 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 + # 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. + (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: + # 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 (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, 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. + + 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 (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 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 + 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..120fb12 --- /dev/null +++ b/triton_conv3d/reference.py @@ -0,0 +1,413 @@ +# 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..fbb9939 --- /dev/null +++ b/triton_conv3d/scaffold_census.json @@ -0,0 +1,5065 @@ +{ + "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": [ + { + "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..8b6595c --- /dev/null +++ b/triton_conv3d/scaffold_corpus.json @@ -0,0 +1,4871 @@ +{ + "source": "model-analysis/unet_shapes.py", + "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..50e90cf --- /dev/null +++ b/triton_conv3d/shapes.py @@ -0,0 +1,784 @@ +# 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, 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``. +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, 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. 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 = [] + 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 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 + :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 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 + 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..ed64b25 --- /dev/null +++ b/triton_conv3d/tests/test_bwd_data.py @@ -0,0 +1,705 @@ +# 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..399c3aa --- /dev/null +++ b/triton_conv3d/tests/test_bwd_weight.py @@ -0,0 +1,1493 @@ +# 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**, 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 + 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: 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..b5e19aa --- /dev/null +++ b/triton_conv3d/tests/test_gather_gemm.py @@ -0,0 +1,1140 @@ +# 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 {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..1c68d3b --- /dev/null +++ b/triton_conv3d/tests/test_infra.py @@ -0,0 +1,1674 @@ +# 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. + """ + 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) 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 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. +#: 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 was verified by mutation -- breaking the +# thing it tests and confirming it fails. + + +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, 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: 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 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. + 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%), 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" + ) + 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 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. + 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 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. + 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 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. + 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 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 + ), + } + 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. 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..eb45632 --- /dev/null +++ b/triton_conv3d/tests/test_transposed.py @@ -0,0 +1,1048 @@ +# 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 math + +import pytest +import torch +import torch.nn.functional as F + +from triton_conv3d import reference +from triton_conv3d.gather_gemm import 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..2b2647c --- /dev/null +++ b/triton_conv3d/transposed.py @@ -0,0 +1,1253 @@ +# 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 +from typing import Sequence + +import torch +import triton +import triton.language as tl + +from .gather_gemm import ( + _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 + +__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)``: 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, +#: 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 +#: :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()}" + )