diff --git a/ScaFFold/unet/_rungs.py b/ScaFFold/unet/_rungs.py new file mode 100644 index 00000000..a7a5513b --- /dev/null +++ b/ScaFFold/unet/_rungs.py @@ -0,0 +1,463 @@ +# 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: + +* :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, 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, 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 are shared rather than copied so that the next module to grow a ladder +inherits them instead of rediscovering them. + +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 belongs 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; the part +#: before the first colon is compared instead. +TUNED_ARCH = "gfx942" + +#: ...and the compute-unit count, which is what makes this MI300A rather than +#: ``gfx942`` generally: that arch is the MI300A APU (228 CUs), the MI300X (304) +#: and the MI325X (304). Several tables are *written in terms of* MI300A's +#: 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. Elsewhere +#: those are not bad numbers, just answers about a different machine. +#: +#: The CU count also declines a *partitioned* MI300A, which 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* rather than about +#: either kernel, so someone who has just discovered 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` as the seam the tests replace: + every interesting branch of this guard is one this node cannot take -- an + MI300X, a partitioned MI300A, an NVIDIA device, a driver that will not + answer -- so the query has to be injectable, and keeping it (and nothing + else) behind one function lets a test substitute a tuple and exercise the + real decision, caching and message rather than a copy of them. + + ``gcnArchName`` exists only on a ROCm build, so an NVIDIA device is + described as an empty arch and declines through the same clause an untuned + AMD one does. That is right 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)``, computed once per device index for the life + of the process: it is asked on the routing path of every convolution and + every GroupNorm, and ``get_device_properties`` is a driver query rather than + an attribute read. Computed lazily rather than at import because a CPU-only + run 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. + + Keyed per device index rather than per process because the question is + whether the kernel about to launch will run on the machine its launch + configuration was chosen for, which is a property of the device the tensor + is on. ScaFFold pins one rank per GPU, so in production this dictionary + holds exactly one entry; on a node that exposes two different GPUs a + process-wide answer taken from device 0 would be wrong on one of them. + + 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. The cache above is process-global, so without this the + first test to ask a question fixes the answer for every later one, and 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. + + This protects against silent mistuning, not a crash. Every number that + decides how these kernels launch -- the convolution tile tables and + ``matrix_instr_nonkdim``/``kpack``/``waves_per_eu`` choices in + ``triton_conv3d``, and the entries written in terms of MI300A's geometry + listed at :data:`TUNED_CU_COUNT` -- was raced on one MI300A. Run somewhere + else they are answers to a question about a different machine, and nothing + downstream 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. + * 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, + never 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 asserting a judgement + about their own hardware, the development case this override exists for. + ``False`` never reaches here -- the callers decline on it first. + + Both branches are loud, once per device; see :func:`_warn_platform_decline` + and :func:`_warn_platform_override` for what each owes the reader. + """ + 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 dozens of routed + operations a step would be a log line every few milliseconds, and not + silence, which is the state this guard exists to end. The message names the + device it found, the device it wanted, and both switches + (:data:`_OVERRIDE_SWITCHES`). + + ``is_compiling()`` for the reason :func:`_warn_rung_failure` documents: + Dynamo cannot trace a ``logging.Logger`` method. + """ + 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 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 whose saved tensors are already held. Answering there 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``, or on some shapes a GPU memory fault; 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. 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"). + Dynamo folds ``is_compiling()`` to ``True`` at trace time, leaving ``False`` + here as a constant -- also the right answer, since 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) 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. + + This 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 caller for whom the fallback matters most, since the failure it is + reacting to is usually a compile failure. ``is_compiling()`` folds to + ``True`` at trace time, so the call below is 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. + """ + 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. + + 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 00000000..e0977b26 --- /dev/null +++ b/ScaFFold/unet/conv3d.py @@ -0,0 +1,1727 @@ +# 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. Faster than MIOpen at ScaFFold's sites, and reproducible + on backward-weight at no cost; ``triton_conv3d/bench/conv_bench.py`` + regenerates the comparison. +2. MIOpen, via ``nn.Conv3d.forward``, which every rejection falls back to and + which 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. It shares with :class:`FastConv3d` +what is genuinely one mechanism -- the failure allowlist, the latch, the +autocast reproduction, the ``DCTensor`` unwrap, and the routing conditions that +are not about which operator this is (:func:`_routing_declines`) -- and 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 quietly computes a different operator; +* 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. :func:`_transposed_halo_plan` + *checks* that rather than assuming it. + +Sharding is the whole difficulty +================================ +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 zeroes that +dimension's padding. An adapter that unwraps a ``DCTensor`` and calls a kernel +directly never reaches that code, so it would silently drop the halo and be +wrong at every shard boundary. (GroupNorm's fast path may unwrap and rewrap +freely; its statistics are per-shard at every shard count. A convolution's are +not.) + +So this module performs the exchange itself, *above* autograd, as its own +``autograd.Function`` (:class:`_Halo3d`) stacked over the convolution's +(:class:`_TritonConv3dFn`), with :func:`_exchange_forward` and +:func:`_exchange_backward` reproducing ``distconv``'s two exchanges and the +kernel handed a padding with the exchanged dims zeroed. ``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 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 +bitwise identical through these kernels, so dropping it needs no tolerance +argument. At ``dc_num_shards = (1, 1, 1)`` nothing is split and +:class:`_Halo3d` is not applied at all. + +:func:`_halo_plan` checks every fact that argument rests on *positively* before +the rung is taken, and returns ``None`` for anything it could not read. The +failure mode of getting this wrong is a plausible-looking wrong gradient at +scale, not a crash, so "no evidence of a problem" is not good enough. + +The exchange sends plain-contiguous buffers, which ``forward_halo_exchange`` +does not: it sends ``.contiguous(memory_format=channels_last_3d)``, which for +``C > 1`` is not plain-contiguous, and ``ProcessGroupGloo`` rejects it. That, +not the narrow, is why DistConv's spatial sharding does not run under gloo; +packing NDHWC into a plain-contiguous slab makes the sharded path runnable +without 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 -- happens *in the dispatcher*, so an adapter that calls a kernel +directly bypasses it and would silently run the whole network's convolutions in +fp32, on the operands as the model holds them: a different and much slower +computation, with nothing failing. :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. The +guard is not about capability: the kernels compute the right convolution +wherever Triton lowers them, and ``is_supported*`` is right to say nothing about +the device. It is that every number deciding how they launch was raced on one +machine, and that a launch configuration merely wrong for the hardware raises +nothing at all -- ``gather_gemm``'s docstring records the mechanism, an illegal +MFMA configuration that emits zero MFMA instructions and returns correct results +slowly, which an allowlist of ``triton.errors.TritonError`` cannot catch. So +the check is a routing decision and lives with the others, in +:func:`~ScaFFold.unet._rungs._platform_declines`, which argues it and argues why +the explicit opt-in overrides it while no correctness condition can be +overridden at all. + +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, 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 so hanging the mesh, or cross-pairing two +convolutions' slabs. 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``. Otherwise a Triton compile +failure would be fatal at ``num_shards > 1`` while costing only speed at 1, +which is backwards for a ladder whose reason to exist is degrading rather than +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 the allowlist is +drawn at Triton's boundary instead of at the package's. It is nonetheless +closed and small: every error under that root is raised while compiling a kernel +or sizing its launch, i.e. before any device work, which is what makes the retry +safe, and :func:`_triton_kernel_failures` names them. 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 an entry point whose own + ``is_supported*`` says no, which this module branches on first. If it fires, + predicate and entry point disagree and that has to be seen. +* ``torch.OutOfMemoryError`` and ``torch.AcceleratorError`` -- MIOpen needs + *more* memory than the Triton path at the shapes where the first bites, and + the second means the HIP context is already poisoned. +* bare ``RuntimeError`` -- excluded precisely because the three above are all + subclasses of it. + +No allowlist covers a kernel that stores out of bounds: that takes the process +down with an HSA "Memory access fault" and SIGABRT without raising anything +Python can see. Robustness against it 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``. Once a dim is split the Triton rung +saves the *halo'd* input, whose extent differs, so a flip is caught there as a +plain ``CheckpointError`` -- but the unsharded configuration is the shipped one +and it is the silent case. 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 +both ``backward`` methods type-check slot 0 to turn the ``AttributeError`` into +an actionable message. + +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 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 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``: every configuration the +#: kernel must not serve is refused by a correctness check rather than by this +#: default. The explicit opt-in is *not* merely that default written out -- it +#: additionally overrides the hardware guard, the one routing condition here +#: that is a preference. 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 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``, the same statement spelled + in the environment -- takes the kernels on hardware the default declines, + and says so in the log. That makes an explicit ``1`` stronger than leaving + the variable unset. + + ``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 so a run that never reaches the GPU (the whole CPU unit suite) does + not pay for ``triton``, as ``group_norm`` defers its kernel module. 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 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 is + raised while compiling a kernel or sizing its launch, 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", and the ladder + then re-raises -- right, because with no Triton there is nothing that could + have failed inside one. + """ + 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 forward handlers that catch an allowlisted kernel failure need exactly + this: :meth:`FastConv3d.forward`, which re-runs the whole call on MIOpen, + and :meth:`FastConv3d._triton_forward`, which cannot (its halo is already on + the wire) and 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 + breaks 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 dispatcher applies that cast *below* this module, which is why + calling a kernel directly has to reproduce it -- see the module docstring. + + ``None`` means "no cast": a run with ``torch_amp: 0``, an evaluation outside + the autocast region, or a device autocast does not know about. + """ + 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 runs. + + 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`` -- 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 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. + + ``expand`` gives every dimension a stride of 0, so the result is safe to + read and useless to compute with; nothing computes with it. A predicate + that ever grows a stride or contiguity test 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 while :func:`_policy_declines` is empty. 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. + """ + + __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. + + The strategy-wide facts, each matching a line of ``distconv.py``: + + * ``shard_dim`` has the same length as ``num_shards`` and as ``shard_ind``, + because ``distconv_forward`` indexes both by position in ``shard_dim`` and + a mismatch means they disagree about which axis is which. 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, which is a reason to check rather than 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 and keeps its + padding, for the reason the module docstring gives: no ``P2POp`` is posted + and the ``zeros_like`` receive buffers stay zero, so DistConv's ``cat`` + there is ``cat(zeros, x, zeros)``. + + 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. + + Whether a process group exists is *not* checked here: that is a routing + question rather than a property of the strategy, and 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. + + The hook is empty; the arguments are kept so a future entry has somewhere to + go. Anything added here must be justified per direction. 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 with a + different winner -- so a rule phrased in the forward GEMM's ``M`` decides + three kernels on one direction's evidence. Either measure all three with + ``triton_conv3d/bench/conv_bench.py``, or block per direction. + + An empty block-list also keeps the convolutions out 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`, 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, and 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 -- 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 here 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 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``, as ``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, receives, posting order and result: the local shard with each + neighbour's ``halo``-thick boundary slab concatenated onto the matching + face, and zeros where there is no neighbour. Two deliberate differences: + 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 returned, as ``distconv_backward`` does -- + 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 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 only the :class:`_HaloPlan`, which is + Python objects. 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` read the weight's + channel axes in opposite orders and their gates take different arguments. + + Written once rather than copied because each clause below is a correction + with a failure behind it (the module docstring says which), and a copy would + drift away from them. Every clause is a pure predicate, so their order + costs only the time to reach the answer. + """ + 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``. Identity test, not isinstance: is_supported only asks + # isinstance and would accept any subclass. + 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, what this one + # prevents is a right answer at an unknown speed; see + # :func:`~ScaFFold.unet._rungs._platform_declines`. Cached per device. + 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 that call would otherwise fail + # inside ``dist`` rather than routing to MIOpen. It lives here, not with + # the operator-specific checks, because it is a property of the *plan*; a + # transposed plan never has an exchange in it. + 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. A live branch, not a formality: + # DistConv's narrowed ``_tensor`` stops being channels-last at + # ``local_batch_size > 1``, which is a supported config key. + 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 belongs to :class:`FastConvTranspose3d` and a + # different set of entry points: ``is_supported`` does not take a + # ``transposed`` parameter, so a ``(Cin, Cout, k, k, k)`` weight would be + # *accepted* whenever the two channel counts match and a different operator + # computed. A non-zeros padding mode is an F.pad inside + # nn.Conv3d._conv_forward 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. At (2,1,1) or (4,1,1) only + # D is exchanged, so the kernel sees ``(0,1,1)`` there and ``(1,1,1)`` + # unsharded, 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 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") and has done no observable work, so its failure is a routing miss + # rather than a broken kernel and 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 gate does not imply the + # other two -- ``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: both ladders want one call to one operator, + 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 pays the custom op's dispatcher and autograd-node overhead so + its op composes 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 bypasses -- a real dispatcher op would be intercepted by DistConv's + generic unwrap, which 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, 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, 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 -- a wrong number rather than an error. With it, + # the second differentiation says so. MIOpen's rung *is* twice + # differentiable, so this is a real difference between the rungs. + @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 (see the module docstring's "Latches"): torch's own + # metadata check passes, so a DCTensor lands in this slot instead of + # the plain tensor this forward saved, and without this the next + # line would die inside DistConv with an AttributeError naming + # neither checkpointing nor the rung. + 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 would each call ``.contiguous(channels_last_3d)`` + # on it, so 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: the + # argument that forbids it in the forward is about the saved set, + # and this is not a recompute but the node itself, running once and + # consuming exactly the tensors it was given. Nothing downstream + # can tell which kernel produced the gradients. It matters because + # 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. The segmentation head is ScaFFold's + # only biased ordinary convolution. + 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: 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. Holding it + in RSCK order behind ``state_dict`` hooks instead wins a little in the + kernels and loses considerably more in the optimizer. + + 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; 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. + + The exchange is above the retry and only the kernel call is inside it -- + the split the module docstring's "Latches" argues for. ``_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, and the fallback must consume 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. + + 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, as ``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: + # 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, and MIOpen's result would substitute a DCTensor into a + # slot holding a plain one. 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, since ``_miopen_forward`` goes + # through ``distconv_forward`` and would exchange again. 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. Unconditional, because the + # invariant is that 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 a rule written for the ordinary operator reads the + wrong numbers here: + + * ``w_shape`` is a ``ConvTranspose3d`` parameter, ``(Cin, Cout, kd, kh, + kw)``. Its channel axes are the other way round, so a rule reading + ``w_shape[0]`` as ``Cout`` and ``w_shape[1]`` as ``Cin`` -- which is what + :func:`_policy_declines` may do -- reads each one as the other. + * ``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. + * a small-``M`` cliff 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 an entry here has to be measured on this operator and per direction; + ``triton_conv3d/bench/conv_bench.py`` regenerates that comparison. This + ladder's per-call Python and launch overhead is a property of the adapter + rather than of any shape, so it belongs in a profile and not here. + + 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, and a shard holding + a contiguous block of input voxels 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. + + Hence the parity test below. An odd kernel on a split dim is declined even + though the bijection 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 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. Checked rather than assumed, as 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 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 :func:`_use_triton` 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, ...)``, accept it whenever the two channel counts + # match, and compute 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 + + # Broad catch for the reason _use_triton gives: a predicate that cannot + # answer has a correct answer available ("no") and has done no observable + # work, 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 -- 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 a bias + (``nn.ConvTranspose3d`` defaults to ``bias=True`` and ``unet_parts`` does + not turn it off), 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: here the two rungs save tensors that + # agree on everything ``_default_meta_extractor`` compares at every + # shard count, so this check is the only detector of a rung flip + # across a checkpoint recompute. + 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. 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 expression ATen's convolution + # backward uses. + 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``. + """ + + #: 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 because 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. Unused 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, and + # MIOpen's answer would substitute a DCTensor into a slot + # holding a plain one. + 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 06d7407c..b5bce682 100644 --- a/ScaFFold/unet/group_norm.py +++ b/ScaFFold/unet/group_norm.py @@ -12,44 +12,139 @@ # # 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. - -``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. +"""GroupNorm with a Triton fast path and a ``torch.compile``d one behind it. + +Three kernels, tried in order. They agree to fp32 rounding and all return the +input's memory format (see :func:`_match_memory_format`), so a caller does not +have to know which rung served it: + +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 as a strided gather and returns a + contiguous tensor, breaking the layout chain for the convolutions that + follow. The Triton kernel is NDHWC in and NDHWC out, and can fold the ReLU + into its store (see "Fused activation"). +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 leaves most of the device idle, where + Inductor tiles the reduction across it. +3. Stock eager ``F.group_norm``, which every rejection falls back to and which + defines the semantics the other two must match. + +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 only its speed + is unknown elsewhere, 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: ``triton_group_norm.TritonKernelError`` for +the Triton rung, ``torch._dynamo.exc.TorchDynamoException`` plus +``FailOnRecompileLimitHit`` for the compiled one (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``. A denylist cannot work here: the +set of things torch may raise through a forward is open, while the set of ways a +kernel can be broken is closed at its own boundary. Both allowlisted exceptions +are raised strictly before their rung saves anything for backward (the Triton op +saves in ``_setup_context``, after its launch region; a Dynamo/Inductor failure +happens at compile time, before any execution), so the retry cannot double-fire +saved-tensor hooks. + +A failure latches the rung off only for modules that have never had a call +served by it; a module that has already run on a rung keeps it. That is a +correctness requirement, 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 rungs +save *different tensors* -- Triton ``(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 kill the step with +``CheckpointError: Recomputed values ... have different metadata``. Pinning +each module's choice for the life of the process keeps forward and recompute in +agreement. + +The same reasoning bounds the fallback, 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 forbidden above 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 -- in particular every first one, 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 lives. + +A latch is process-local, so under DDP one rank can end up on a different kernel +from its peers. The kernels agree to fp32 rounding, not bitwise, so a rank that +latches shifts its own gradients and therefore the all-reduced ones, changing +the job's trajectory, and it becomes a straggler besides. 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`` -- transient, and no cheaper on any other +rung -- latches nothing 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, 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 folds it into its forward store -- free, since that store is what bounds +the kernel -- removing a whole streaming pass. +``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`` in the activation slot of its +``nn.Sequential``, which keeps the positional keys -- and therefore every +checkpoint -- unchanged; neither module has parameters or buffers. """ import logging -import os import torch import torch.nn as nn import torch.nn.functional as F +# Shared with ScaFFold.unet.conv3d so the corrections these helpers carry cannot +# drift between the two ladders. Imported by name so they stay module +# attributes here, which is what the tests monkeypatch. +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 +152,29 @@ #: is safe", which is what every production run wants. COMPILE_ENV_VAR = "SCAFFOLD_GROUPNORM_COMPILE" +#: The same spellings and the same default, for the native channels-last Triton +#: kernel, which is tried first. The explicit opt-in is more than that default +#: written out: it additionally overrides the hardware guard, 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. 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 +188,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 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 +217,47 @@ 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. 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 + 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 an untuned GPU gives the right numbers at an unknown speed, where an + unsupported input would give the wrong ones; the opt-in (here or as + ``SCAFFOLD_GROUPNORM_TRITON=1``) is a developer asserting a judgement about + their own hardware, and the log says so. 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 +267,21 @@ 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. + + Not sufficient on its own: ``torch._dynamo.config`` keeps user overrides in + a thread-local ``ContextVar`` (``torch/utils/_config_module.py``), so this + assignment is invisible to every other thread, which goes on reading the + stock default of 8. That matters because ``torch.utils.checkpoint``'s + non-reentrant recompute runs on the autograd engine's device worker thread, + and on a ``DCTensor`` it has to compile (see ``_MIN_RECOMPILE_LIMIT``). + :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 +290,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 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 + 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, since the rung is + # worth 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 +326,177 @@ 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 once the input is known to be a CUDA tensor, which also + keeps ``import triton`` -- deferred again to the first kernel launch -- out + of a CPU-only process. + """ + 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 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 + available even when the thing that failed *is* the module lookup. An empty + tuple (no kernel module at all) means "catch nothing", which is right: with + no kernel module, nothing 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 (an unwritable Inductor cache, a broken + C++/Triton toolchain), and ``InternalTorchDynamoError``. + + ``FailOnRecompileLimitHit`` -- raised when a frame needs more cache entries + than the recompile limit allows, a hard error under ``fullgraph=True`` + rather than a drop to eager -- derives straight from ``Exception`` + (``torch/_dynamo/exc.py``) and not from that root, so catching only + ``TorchDynamoException`` would let it kill the run. It is added only when + it really is outside the root, so a torch that reparents it does not produce + a duplicate entry. + + All of these are raised while *compiling*, 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 + swallowing silently. + """ + 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 tuned on: ``_TUNED``'s + # largest entry names a grid of exactly 228 CUs, so elsewhere its answers + # describe 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 -- see ``_rungs._platform_declines``, + # including why an explicit opt-in overrides this check and not the ones + # around it. Cached per device. + if _platform_declines(input.device, _triton_override): + return False + # is_supported() is cheap and side-effect free: 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") and has done no observable + # work, and its failure is a routing miss rather than a broken kernel, so it + # must not latch the rung off. + 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 - # 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). + # 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 route their ops through __torch_dispatch__, which Dynamo + # cannot trace. DistConv's DCTensor never reaches this check -- forward() + # unwraps to its local shard first -- so anything rejected here is an + # unknown wrapper, whose semantics eager preserves. if type(input) is not torch.Tensor: return False # CPU GroupNorm is not the bottleneck and compiling it would put a @@ -172,33 +509,240 @@ def _use_compiled(input): return True -class FastGroupNorm(nn.GroupNorm): - """``nn.GroupNorm`` that runs its GPU forward through ``torch.compile``. +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. With ``PYTORCH_MIOPEN_SUGGEST_NHWC=1`` the + convolutions on both sides want channels-last, so one fallback would + re-break the layout chain for the rest of the network; relaying out the + GroupNorm output once is cheaper than the transposes those convolutions + would insert, and it is what makes the three rungs agree on everything a + caller can observe rather than only on the values. - Identical state: ``weight``/``bias`` of shape ``(num_channels,)``, no - buffers, so state dicts are interchangeable with plain ``nn.GroupNorm`` - in both directions. + Only a fallback from a channels-last input reaches the copy: the Triton rung + is already channels-last, and a contiguous input has nothing to do. """ + 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) - def forward(self, input): + +class FastGroupNorm(nn.GroupNorm): + """``nn.GroupNorm`` with a Triton GPU kernel and an optional fused ReLU. + + A drop-in replacement with identical state: ``weight``/``bias`` of shape + ``(num_channels,)``, no buffers, so state dicts are interchangeable with + plain ``nn.GroupNorm`` in both directions. ``activation`` is a plain Python + attribute, not a submodule or a buffer, so it adds no 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. Which kernel + runs therefore changes the number of memory passes, never the function the + model computes. + + DistConv's ``DCTensor`` gets the fast kernels by being unwrapped to its + local shard in front of them, rather than by letting the op dispatch through + the wrapper. Dispatch would work -- the Triton kernel is a real dispatcher + op, so ``DCTensor.__torch_dispatch__`` would unwrap, run and rewrap on its + own -- but the explicit unwrap keeps the subclass policy in one place + (``is_supported`` accepts any ``torch.Tensor`` *instance*, so dispatch would + silently extend the fast path to every unknown wrapper), lets the + eligibility predicates examine the tensor the kernel will actually touch + rather than a wrapper's mirrored metadata, and lets both fast rungs share + one unwrap and one fallback ladder. Semantics are unchanged either way: + DistConv's generic ``__torch_dispatch__`` has no GroupNorm-specific + handling, so statistics stay per-shard and nothing communicates at any shard + count. + """ + + #: 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 + #: 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: ``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 and can be assigned after construction, and the Triton + rung would then fuse an activation this method silently skipped -- the + network's function would depend on its input's memory format. It is + also what makes adding a third entry 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 (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 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. + # + # 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 the eager result would make the + # recomputed saved set disagree with the original and torch would + # reject the step with a `CheckpointError`. Degrading is for + # modules with nothing to contradict; here the honest answer is the + # original exception, which 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 00000000..22388b07 --- /dev/null +++ b/ScaFFold/unet/triton_group_norm.py @@ -0,0 +1,1604 @@ +# 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). + +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 (eager or Inductor-compiled) iterates them in logical +NCDHW order -- a strided gather over channels-last memory -- and returns a +contiguous tensor, so the next convolution has to convert back. The kernels +here read and write channels-last directly. + +A channels-last-3d contiguous ``(N, C, D, H, W)`` tensor is physically a dense +``(N, S, C)`` array with ``S = D*H*W``, and group ``g`` owns a contiguous run +of ``C/G`` channels inside every voxel. One program therefore serves all +groups at once for a chunk of voxels: it reads a dense ``(BLOCK_S, C)`` run, +reshapes the inner axis to ``(G, C/G)``, and gets coalesced loads and stores +with the group axis costing nothing. + +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. Accepts anything ``F.group_norm`` accepts; inputs + the Triton kernel cannot serve fall back to ``F.group_norm`` internally. + +``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 with a better fallback than eager -- a + ``torch.compile``d GroupNorm, say -- should test it and route the rejects + themselves. + +Contract +======== +For every input ``is_supported`` accepts, the result matches ``F.group_norm`` +to within fp32 reduction-order noise, with: + +* **dtype** -- exactly ``F.group_norm``'s: the input's dtype outside autocast, + fp32 inside it, since ``at::group_norm`` carries autocast's fp32 cast policy. + ``_autocast_out_dtype`` reproduces that rule without materializing the fp32 + copy of the input that autocast's own cast would make: the kernels read at + the native width and accumulate in fp32, which is the same computation as + upcasting first but reads fewer bytes. Gradients follow suit -- ``d_input`` + has the input's dtype, ``d_weight``/``d_bias`` the parameter's. +* **statistics** -- always accumulated in fp32, never in the input dtype, and + rounded to the output dtype exactly once, on the store. +* **memory format** -- the output has the *input's* memory format. This is the + one deliberate difference from stock GroupNorm, which returns a contiguous + tensor whatever went in; preserving channels-last is the point of the kernel. +* **autograd** -- ``d_input``, ``d_weight``, ``d_bias``, with ``weight`` and/or + ``bias`` ``None`` allowed. First order only: the backward is itself a custom + op with no autograd formula, so a second ``torch.autograd.grad`` raises where + stock ``F.group_norm`` supports double backward. A gradient penalty or a + Hessian-vector product must route around this kernel; it fails loudly rather + than returning garbage. +* **device** -- the kernels run on the *input's* device whatever device is + current, matching ATen's ``DeviceGuard``; see ``_device_guard``. +* **determinism** -- bitwise reproducible run to run and process to process: + no float atomics anywhere, and grid, split count and tile sizes are pure + functions of the shape. The tuning table is frozen in this file for that + reason -- a runtime autotuner would change 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 ``N*(C/G)*D*H*W == 1`` shape, so a caller branching on + ``is_supported`` never gets an answer where the op this replaces would raise. +* **eps** -- one deliberate divergence: for a *subnormal* fp32 ``eps`` (below + the smallest normal, 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. Not clamped, because clamping would perturb every ordinary + call to defend a value far below any plausible one. + +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. +* Every other layout and rank -> ``triton_group_norm`` falls back to + ``F.group_norm``, whose contiguous output again has the input's memory + format. ``is_supported`` returns ``False`` so that callers keep their own, + probably compiled, fallback instead of silently dropping to eager. + +One kernel family rather than two is deliberate. A native NCDHW kernel needs a +different tiling -- with C outermost the fast axis is spatial, so a program owns +one group and streams S instead of owning all groups and streaming voxels -- +and on contiguous input Inductor's compiled GroupNorm is already near this +device's streaming roofline, so the payoff there is a fraction of what the +channels-last path recovers. If a mixed-layout model ever makes it worth +having, the place to add it is :func:`select_strategy`. + +Reduction strategy +================== +Group statistics span ``S * C/G`` elements, far more than one pass can reduce +at the large UNet shapes, so split-K partials 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 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 (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. + +Fusion and tiling are one problem, not two: a fused finalize is recomputed by +every elementwise *program*, so the redundant (L2-resident) traffic is +``nprog_elem * nsplit`` triples, which with an uncapped grid at the largest +shape costs several times the tensor itself. Two knobs in ``GNConfig`` bound +it -- the elementwise grid is capped at ``elem_progs`` programs which then +stride over the tiles, so redundancy scales with the grid and not with the tile +count, and ``nsplit_target`` is tuned per shape against that cap. They were +tuned jointly by coordinate descent; do not change one without re-deriving 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) would collapse each pass to a single +launch, and does win on the smallest, host-dispatch-bound shapes. It is +deliberately not used: a grid barrier needs every workgroup co-resident, which +caps the grid at the CU count (228 here) and with it the achievable bandwidth, +so it loses badly the moment a shape is bandwidth-bound rather than +dispatch-bound, and at the largest UNet shape it does not compile at all. +Serving both regimes would mean two kernel families plus a crossover rule, for +a few percent of the model step -- and nothing at all under CUDA-graph capture, +where launch count is free. + +Numerics: Welford, not ``E[x^2]-E[x]^2`` +======================================== +``E[x^2]-E[x]^2`` is split-friendly and cheap but cancels catastrophically once +``mean >> std``, subtracting two nearly equal large numbers to recover a small +one. Each tile here instead produces ``(count, mean, M2)`` by 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 merge 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. + +The ``corr`` term is what keeps ``rstd`` accurate where a tile mean is formed +from many large values, by orders of magnitude in a single-tile reduction. The +*output* error barely moves, because it is dominated by the fp32 representation +of the mean, which ``corr`` cannot improve -- so do not delete the term on the +strength of an output comparison. Its load-bearing case is pinned by +``test_welford_correction_recovers_rstd_in_a_single_tile_reduction`` in +``tests/test_triton_group_norm_edge.py``. + +Addressing +========== +``[2, 64, 256^3]`` is exactly 2^31 elements, so int32 linear offsets would +block batch>1 at the largest UNet shape and at everything 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, which is why the wide path costs nothing measurable. + +Fused activation +================ +``activation="relu"`` folds the ReLU into the forward store -- one compare and +one select in a store-bound kernel -- and removes an entire 2B streaming pass. + +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)`` +with the expression the forward used: two FLOPs on values already in registers, +and bit-exact, so the sign always agrees with the forward's. Testing ``y > 0`` +on the saved output instead would need the output kept alive in addition to +``x`` (which the backward needs anyway), and in bf16/fp16 would 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 +``tl.maximum(y, 0)`` / ``tl.where(pre > 0, dy, 0)``. The two agree on every +finite value but not on NaN: ``tl.maximum`` returns the non-NaN operand and +``NaN > 0`` is False, so both usual spellings map a NaN to 0.0, while +``F.relu`` propagates it and ``threshold_backward`` -- ReLU's real backward -- +passes its gradient, ``NaN <= 0`` being False too. Matching ``F.relu`` here +matters: a diverging run whose forward comes back finite because the fused +activation ate the NaN sails through ScaFFold's non-finite-loss abort and +checkpoints a broken model. + +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, statistics are per-shard. + +At the smallest shapes the call is dominated by autograd and the dispatcher +rather than by the kernels; that is the price of being an op ``torch.compile`` +and ``DCTensor`` can see. Two ways to claw that host time back are rejected: + +* Bypassing ``JITFunction.run`` for a cached ``CompiledKernel`` handle assumes + Triton's specialization key -- 16-byte pointer alignment included -- 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 + exercised by the test suite, and a stale specialization there is a wrong + answer rather than a crash. +* Caching the scratch buffers across calls makes them shared mutable state + between call sites: correct on one stream, wrong on two, and this module + cannot tell which it is on. + +What is left at those shapes is torch's own plumbing, removable only from +outside this file -- CUDA-graph capture of the training step, or a C++ autograd +node. + +Triton is imported lazily, on the first call that 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 lets a caller with a + fallback (``ScaFFold.unet.group_norm``'s ladder) catch exactly "the kernel + is broken" instead of catching ``Exception`` and trying to enumerate the + framework mechanisms that legitimately raise through a forward -- pack + hooks, ``torch.utils.checkpoint`` recompute, functorch, offloading hooks. + + Two things are deliberately not tagged, and propagate unchanged: + + * ``torch.OutOfMemoryError``, a resource condition rather than a defect -- + every fallback allocates an output of the same size, so retrying one is + just a second OOM at a call site the caller did not ask about; + * the ``ValueError``s ``_validate`` raises, which are caller contract + violations that ``is_supported`` already excludes, so one escaping is a + bug here and must stay loud. + + The tagged region does no autograd-observable work -- allocations and + launches only, with ``save_for_backward`` in ``_setup_context`` strictly + after ``_forward`` returns -- so an exception of this type always predates + anything the op saved, which is what makes retrying 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 number of split-K partials wanted across the whole + 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). + + Not independent knobs: every elementwise program re-reads all ``nsplit`` + partials for the fused finalize, 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() + ) + + +#: Tuning table for the fused kernels, from a coordinate-descent sweep of +#: fwd+bwd time on one MI300A at fp32 with ``num_groups=8``, keyed by the +#: ``(num_channels, cube-root spatial extent)`` of the ScaFFold UNet GroupNorm +#: sites. Frozen -- never autotuned at run time -- because the split count +#: fixes the reduction order and therefore the bits of the result. +#: +#: Keyed on ``(C, edge)`` and not on ``N`` because ``nsplit_target`` targets the +#: split count summed over the batch (per sample it is ``nsplit_target // N``), +#: so one entry serves ``N > 1`` with the same total number of stats programs. +_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, most of that host cost being torch's own +#: autograd and dispatcher plumbing rather than 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: + """The single point where a different kernel strategy is chosen per shape. + + Returns a name from :data:`STRATEGIES`. Today always ``"split_k"``: two + forward and two backward kernels with split-K partial reductions and the + finalize passes fused into their consumers, which is bandwidth-optimal at + the large shapes and close to the floor a Python ``autograd.Function`` can + reach at the small ones. Below roughly ``SMALL_SPATIAL_THRESHOLD`` voxels + the call is host bound, but on autograd and the dispatcher rather than on + launches; see "Why not one launch per pass" in the module docstring for the + strategy that would cut it further and why it is not here. + + A second strategy needs its name in :data:`STRATEGIES`, a rule here that is + a pure function of the shape (determinism depends on it), and a branch in + ``_dispatch`` -- the only caller, sitting in front of the memoized tiling + plan and used by both ``_forward`` and ``_backward``. Nothing else in this + file changes. + """ + 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 that is not + # hypothetical at scale. Only the scalar tile base is widened, in the + # kernels; see "Addressing" in the module docstring. + 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 at the host-bound + sites rebuilding the plan (two power-of-two loops and a dict lookup) on + every call is measurable. 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` must stay a small finite value. `cnt_t` counts only the + # valid lanes, so `mean_t = mean0 + corr` is the true mean whatever + # the masked lanes contributed, and `d`/`dd` are re-masked before + # they reach `m2_t` -- but a huge `other` would swamp `mean0` (and + # the correction with it), and `inf` or `nan` would poison it. + 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 what keeps `rstd` accurate once the mean dominates the + # spread; do not drop it on the strength of an output comparison, + # which barely moves either way -- see the module docstring. + 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, cheaper than the launch it replaces. It is not cheap enough to + pay once per tile at the largest shapes, so the grid is capped at + ``NPROG`` and each program strides over its share of the ``NBLK`` + tiles, making the redundant read ``NPROG * NSPLIT`` rather than + ``NBLK * NSPLIT``. See :class:`GNConfig`: ``nsplit_target``, + ``elem_tile`` and ``elem_progs`` are one joint tuning problem. + + Every program reads the same partials with the same tile shape, so 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 a + # 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 False too, so NaN falls to `y`. Bit-identical + # to `F.relu` on NaN, +-Inf and -0.0 -- 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) so that a NaN pre-activation passes its + # gradient, which is what `threshold_backward` -- ReLU's real + # backward -- does, `NaN <= 0` being False. See the forward. + 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. + + The per-channel dweight/dbias row reduction is done by the first + ``NDW`` programs of ``n == 0``, in one pass over the ``(n*nsplit, C)`` + scratch; the per-``(n, g)`` c1/c2 finalize is recomputed redundantly by + every program 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), which the compiler contracts + # to fma(dy, w, -c1) while c1 was accumulated from the rounded + # product, leaving that rounding error amplified by + # rstd = 1/sqrt(eps). 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 about cost, not correctness: entering + ``torch.cuda.device`` is host time, and ScaFFold pins one device per rank, + so the tensor is already on the current device on every call. It is + charged twice per step (forward and backward) and only shows up at the + host-bound shapes. + """ + 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 -- no autograd-observable + op, so no pack hook, recompute stop or functorch layer inside it -- so the + blanket ``except Exception`` here cannot swallow framework control flow the + way one at the call site would. That is what lets the caller's fallback + ladder key on a single exception type 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``, every group having zero variance), but a caller branching on + :func:`is_supported` would then get a result where the op this replaces + raises -- 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 inside an + enabled autocast region it upcasts its input and returns fp32 whatever came + in. + """ + 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 c9e6cb0f..34035d0b 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,126 @@ _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. 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 ``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`). + """ + 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. The + kernel serves only the ``kernel == stride``, no-padding upsample built + below; anything else falls back to MIOpen. + """ + 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 moving 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. ``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 +172,16 @@ 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, saving both time and peak memory. This rests + on ``self.conv`` beginning with a convolution, which the constructor below + guarantees on either branch. + """ def __init__(self, in_channels, out_channels, group_norm_groups, trilinear=True): super().__init__() @@ -90,7 +196,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 +224,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) at the dtype 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 b284708a..392dcf8b 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,30 @@ 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. + + Rank 0 only, and one rank's answer: each rank latches independently, so + under DDP this is representative rather than global. 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 +875,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 +982,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 d9f2dda8..36fa6a5d 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 00000000..1101e23d --- /dev/null +++ b/scripts/run-tests.sh @@ -0,0 +1,117 @@ +#!/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. +# +# Every export below changes the result when omitted; see the comment at each. + +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. Both parametrizations of +# tests/test_groupnorm.py::test_gpu_triton_dctensor_matches_eager_and_stays_wrapped +# fail deterministically when it is unset; production runs set it too. +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, +# which dominates the convolution suites and 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: such a 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 00000000..c347365e --- /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 00000000..da611995 --- /dev/null +++ b/tests/test_conv3d.py @@ -0,0 +1,1655 @@ +# 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. + +:func:`test_halo_plan_refuses_a_split_dim_whose_arithmetic_it_has_not_checked` +is the test to trust most: every other property here fails loudly, but a +wrong halo plan produces a plausible wrong gradient at every shard boundary, +since the halo DistConv adds below autograd is invisible to a module-level +adapter. The exchange itself needs real ranks and lives in 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. The gate reads exactly ``num_shards``, ``shard_dim`` and + ``shard_ind``, so a stub can pose the sharded question the real strategy + cannot. + + ``shard_ind`` is here so the MIOpen rung -- which takes a DCTensor through + DistConv's own dispatch -- still runs, making "the sharded call went to + the other rung" an assertion about routing rather than about a missing + stub attribute. + """ + + 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 + must agree: an fp32 parameter against a bf16 activation is a call neither + rung serves. The autocast tests pass fp32 on purpose, since that is the + state ``worker.py`` actually leaves the model in, and reproducing the + dispatcher's cast is the point. + """ + 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 bitwise inert; this pins that the plan drops it, and + that 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. + + A test that only checked the refusal would pass against a gate that + refuses everything, and the sharding check sits behind the ``is_cuda`` + test, so on CPU it is never reached. The same module and tensor are asked + twice here, differing only in ``num_shards``. + + The sharded answer is ``False`` only because this process has no process + group to exchange over: the plan is still 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 needs real ranks and lives in + a separate harness. + """ + 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 -- 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; the call count is what is checked, 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 reach ``_miopen_forward`` and + therefore ``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 makes a broken Triton install fatal at ``num_shards > 1``, and + costs only speed at 1. + + The assertion is a count, 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; which of them is called is what is checked. Because + ``cat(zeros, x, zeros)`` at padding 0 is the same arithmetic as the + module's own padding on the unexchanged shard, the result 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): + """None of these three shapes is on the block list; all three route to Triton. + + Parametrized on the 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 the block would be gone in name + only. + + ``cout=512`` also covers the small-``M`` predicate's defect: it read the + forward GEMM's row count and then kept all three directions on MIOpen, + including a backward-data direction that is faster on Triton. 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. + + 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 reach the kernel unchanged, and the whole network's convolutions + would quietly run in fp32 instead of bf16 -- a different, much slower + computation, with nothing raising to say so. + """ + 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. This pins that a module stays on + one rung across the two -- 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 every + convolution otherwise concatenates two zero slabs onto each of the three + sharded dims for no reason. 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 currently empty, so this asserts the routing answer + rather than the shape of a rule -- a future entry in either function has + to re-establish it here. + + The two functions stay separate because 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``. See + :func:`~ScaFFold.unet.conv3d._transposed_policy_declines`. The retired + small-``M`` rule answered ``True`` for ``up1`` using 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 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. + + Checked two ways: the operands reaching the node are bf16 even though the + module's own parameters are fp32 (the state ``worker.py`` leaves the model + in), and the stock op under the same autocast region produces the same + dtype. Without this the four upsamplers would quietly run in fp32 instead + of bf16 -- a different, much slower computation, with nothing raising to + say so. + """ + 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 because MIOpen's transposed + # backward-weight exceeds the static bound at this shape -- the case + # ``error_bound``'s docstring describes. This direction only asks 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 b1e8e6af..73164925 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 once followed 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: since 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, making every + numeric comparison below vacuous. Converting reproduces the pre-fusion + module graph -- ``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 stay + 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 directly guards 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,18 @@ 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 what Dynamo and Inductor + actually 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 +393,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. + + Triton is far faster than the compiled path 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* -- 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. + """ + 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 +562,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 +698,384 @@ 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`` + (``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 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, part of Dynamo's ``GLOBAL_STATE`` guard, so it + misses every entry the forward built), and there the limit would still + read the stock 8 no matter how large the global was set. + ``torch.compile``'s ``recompile_limit=`` is instead applied by Dynamo + around the compile itself, on whichever thread that compile happens on -- + the only spelling that reaches the worker, which is what this pins. + """ + 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 would + let it escape and kill the step. 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 ones + and zeros, 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]``. 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,64 @@ 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. +# +# 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 far off on the bottleneck's parameters, whose gradients +# are tiny relative to the largest in the network. The difference between two +# arms is therefore the difference of two large errors, small only when they +# happen to cancel. Whether they cancel is settled *outside the source tree*: +# some of this model's convolution problems have MIOpen algorithms whose +# benchmark times tie closely enough that which one wins is frozen into the +# machine-local find database, and a different algorithm choice moves the +# per-parameter figure by an order of magnitude with nothing else changed. It +# is not a gate either: injecting a small relative error into the compiled +# rung barely moves the per-parameter figure, because both are already +# saturated by the bf16 floor. +# +# So rung equivalence is asserted where it is measurable -- *without +# autocast*, where cross-rung agreement is tight and an injected error is +# still clearly separable -- 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 against the model's +# own run-to-run floor and still catches a wrong eps or a wrong group count. + +#: Cross-rung agreement without autocast: wide enough to clear normal +#: reduction-order noise between rungs, tight enough to still fail on a +#: genuine relative kernel error. +_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 +1236,439 @@ 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, sitting on the model's own run-to-run floor + under every convolution algorithm. + * compiled vs. eager under bf16 autocast, on the *aggregate* gradient, + against that same floor. + * compiled vs. eager *without autocast*, per parameter -- the sharp one, + held to ``_FP32_RUNG_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 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; without the per-region + ``recompile_limit=`` the ninth compile raises ``FailOnRecompileLimitHit``. + + ``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 fatal three-way combination: checkpointing + compiled rung + DCTensor. + + ``activation_checkpointing: true`` with ``SCAFFOLD_GROUPNORM_TRITON=0`` on + DistConv activations is a supported configuration, but 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 raise + ``FailOnRecompileLimitHit`` or, once the ladder catches it and drops 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, one past the limit the compiled region raises. + 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 every convolution + after it to convert back. 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, against the model's run-to-run floor, and the per-parameter + gradient *without* autocast -- the sharp one, held to the same + ``_FP32_RUNG_TOLERANCE`` the compiled rung is held to. Per-parameter + *under* autocast is not asserted here; see the note for why that + comparison is unbounded regardless of either kernel. + + 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 00000000..6ec1deca --- /dev/null +++ b/tests/test_groupnorm_wiring_edge.py @@ -0,0 +1,1175 @@ +# 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). + +An adversarial pass over the wiring: ``tests/test_groupnorm.py`` covers the +happy paths and the routing predicates, this file covers where the ladder, +the latches and the absorbed ReLU interact with the rest of torch. Properties +pinned here: + +* 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`` checks, but the compiled and eager rungs apply the + activation through ``_activate``, which tests one literal string. Adding a + second activation to both tuples without adding it to ``_activate`` would + fuse it into the Triton store and silently drop it on the other two rungs + -- the network's activation would depend on its input's memory format. So + for every non-``None`` activation, ``_activate`` must *change* an input + 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 instead builds the pre-fusion + ``nn.Sequential`` from scratch -- ``Conv3d, GroupNorm, ReLU, Conv3d, + GroupNorm, ReLU`` -- so the byte comparison does not depend on that shared + construction. + """ + 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 reassigned afterward; ``is_supported`` + would then decline the Triton rung while ``_activate`` silently applied + nothing, quietly turning the module into a bare GroupNorm. The same + failure would follow from adding an activation to ``SUPPORTED_ACTIVATIONS`` + without implementing it in ``_activate``, so validating on every call + closes both holes. + """ + 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's pack hook -- from inside whichever op is + saving a tensor -- exactly like ``_StopRecomputationError``, and like it is + a ``RuntimeError`` subclass, so a handler wide enough to catch "a broken + kernel" by type catches it too. Swallowing it would latch the rung off and + leave the checkpoint frame in a state the machinery never expected, so the + allowlist must be narrow enough to let it propagate untouched. + """ + 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 checkpointing the recompute counts pack-hook firings + and requires both the count and the metadata to match the forward's, so a + rung that packs some tensors before failing -- with the fallback packing + its own set on top -- corrupts the frame. A user's offloading hook has the + same problem in miniature. + + This holds structurally, not by care: the ladder catches only failures + raised *before* their rung saves anything (the Triton op saves in + ``_setup_context``, after which its ``TritonKernelError`` cannot occur; a + Dynamo/Inductor failure is a compile-time error, before execution). + Asserted here: a caught failure packs exactly what a clean fallback packs, + and a failure *after* 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 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 why a broken install + still costs one attempt per module rather than one per call. + + Also tested: because a proven module keeps retrying, the warning must fire + on the latch's edge rather than per call, or a persistently broken kernel + floods the log. + """ + 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: a 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``, or a GPU memory fault). + Neither is a degradation, so this one case re-raises. + + 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. The second half here pins the other side -- 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 backward is marked ``@once_differentiable``. + Which one the traversal reaches first is a routing detail -- the stem + convolution's own backward now stops it one op earlier, before GroupNorm + is ever reached. + + What must not change is that it raises *at all*: without both decorators + in place, ``first`` would silently arrive with no ``grad_fn`` and a second + backward would contribute zero rather than error. + """ + 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. + + ``logger.warning`` cannot be traced by Dynamo, so the fallback handler must + avoid calling it directly or a caller compiling this forward with + ``fullgraph=True`` gets a hard error instead of the fallback -- exactly the + caller for whom the fallback matters most, since what it reacts to is + usually a compile-time failure. Nothing in ScaFFold compiles + ``FastGroupNorm.forward`` today; this pins that it can. + + Both halves of the handler's guard must trace, which is why ``proven`` is + parametrized: with ``_triton_ok`` false, Dynamo folds the ``and`` away + without evaluating ``_replaying_a_forward()``, so only the ``True`` arm + reaches it -- and an untraceable probe there would reintroduce the same + defect. + """ + 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()`` -- e.g. a ``torch.func`` transform turning 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, + since ``is_supported`` inspects an arbitrary tensor and the set of + wrappers whose attribute reads can raise is not closed. A broad catch is + right here and nowhere else in this module: the predicate has done no + observable work and a correct answer ("use the stock kernel") is always + available, so it must fall back without latching -- nothing about the + kernel itself 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. + + Reaching the kernel through ``torch.func.grad`` and failing there must not + latch *both* fast rungs off permanently -- one transform anywhere in a + process must not silently drop 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. + """ + 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 sets -- Triton ``(input, weight, bias, mean, rstd)``, the other + two ``(input, weight, mean, rstd, relu_output)``. A rung change between a + block's forward and its recompute therefore kills the step with a + ``CheckpointError`` -- exactly the failure the ladder exists to prevent -- + and is reachable whenever ``activation_checkpointing`` is on + (``worker.py:230``). Matching the output memory format is not enough on + its own, since the saved sets still differ; 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 fails the same way, 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 naive fallback would re-break the + channels-last chain for every convolution after it, which is the exact + thing this module exists to prevent. It would also make the rungs + distinguishable to anything that inspects metadata + (``torch.utils.checkpoint``, a compiled caller's guards) -- a correctness + problem, not 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. + + ``tl.maximum(y, 0.0)`` returns the non-NaN operand, so a store built on it + would turn a NaN activation into a finite 0.0 on the Triton rung while + ``F.relu`` on the other two rungs keeps it NaN -- worse than a numerics + discrepancy, since the forward would look finite while the backward is + still NaN, letting a diverged run slip past ScaFFold's non-finite-loss + abort. ``activation=None`` is the control: all three rungs already agree + there, which localizes the check 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 values that distinguish + the spellings of ReLU are reached the other way: a zero ``weight`` makes + the pre-activation exactly ``bias``, elementwise, so the bias vector + chooses what the activation sees. 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 first (``+ 0.0`` maps + ``-0.0`` to ``+0.0``, leaving NaN, the infinities and every normal value + alone) before the same bitwise comparison: a ``-0.0`` bias would otherwise + survive to the output, and whether ``xhat * 0 + (-0.0)`` keeps its sign + depends on whether the kernel contracted the multiply-add into an FMA -- + true of Triton and Inductor, false of eager, and unrelated to the + activation. ``torch.equal`` cannot be used 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 00000000..800fc39a --- /dev/null +++ b/tests/test_platform_guard.py @@ -0,0 +1,580 @@ +# 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. A suite that +only exercised the accept path would ship the entire decline path +unexecuted. ``_rungs._device_fingerprint`` exists as a 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 the guard must be evaluated once, which is asserted +as a call count on the seam rather than as a duration, since a timing +assertion on a driver query measures the driver. +""" + +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's 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 too, and that is not tidying: swapping the + hardware out from under a memo that exists to never ask twice would + otherwise leave the previous answer in place -- "yes, MI300A" on this + node -- so every decline test would quietly become another accept test. + """ + + 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 just 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-``; those suffixes describe how + the build was configured, not which silicon is present, so a string + comparison against ``"gfx942"`` would decline the 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 confirms the real device + does carry a suffix. + """ + + 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 checks that the constants have not drifted from the machine the + Triton kernels were tuned on, and that the suffix strip is load-bearing + here rather than defensive: 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 00000000..85e001a7 --- /dev/null +++ b/tests/test_triton_group_norm.py @@ -0,0 +1,857 @@ +# 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``, with the reference computed in +float64 rather than 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 Welford exists to fix +(``test_welford_survives_large_mean``). + +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 ceiling leaves headroom above the reduction noise a much larger +#: production-scale fp32 sum accumulates. The low-precision ceilings sit just +#: above each dtype's own rounding floor: 2^-8 for bf16, 2^-11 for fp16. +_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: on this build, + # ``F.group_norm(x, g, None, bias).backward()`` raises "tensor does + # not have a device" on both CPU and CUDA, 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 variance formula this kernel avoids, reproduced in fp32 torch ops. + + ``var = E[x^2] - E[x]^2`` is split-friendly and cheap, but it is the + formulation 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. + + Past a large enough mean/std ratio, the naive formulation loses the + variance outright: the difference of two similarly-sized fp32 terms falls + below one ulp, so ``rstd`` saturates on ``eps`` and the output is + meaningless. Welford stays correct there, dominated by the fp32 + representation of the mean 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. + + 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 -- confirmed against stock behaviour here rather than + assumed. + """ + 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]`` clears 2^31 elements while staying + non-power-of-two in the fastest spatial dimension, so a truncated offset + cannot accidentally land on the right address. fp32 keeps the comparison + sharp; the reference needs an NCDHW copy on top of the channels-last + input, so the test skips, loudly, if the device doesn't have enough + memory for both. + """ + 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 00000000..533d21e4 --- /dev/null +++ b/tests/test_triton_group_norm_edge.py @@ -0,0 +1,1693 @@ +# 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``: it covers ground that +suite's coverage misses, and pins the divergences an audit of the kernel +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, or 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 tests below also pin six behaviours an audit of the kernel found: 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. The last section adds the coverage a +mutation sweep of the kernels found thinnest: the ``INT64=True`` branch, 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. It leaves +#: headroom above typical fp32 parity at these small shapes, for the +#: reduction noise a production-sized split-K reduction adds. +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, or 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``, so a slot read before it is + written inherits whatever the caching allocator last left there -- + plausible numbers that no parity test would catch. Poisoning every + ``torch.empty``/``empty_like`` with NaN during the call turns that into a + hard failure, and also proves the output buffer 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 forms the variance by + cancellation and does not manage it, which is why this asserts equality + rather than a bound. + """ + 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 these more extreme ratios the naive ``E[x^2]-E[x]^2`` formulation + loses the variance outright; this kernel stays correct, and does at + least as well as ATen's own fp32 GroupNorm on the same input. + """ + 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 anything that needs a + gradient penalty or a Hessian-vector product cannot use this kernel; the + test pins that it raises loudly rather than silently. + """ + 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, out_dtype override, and 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. 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". Without + ``ctx.mark_non_differentiable``, they would come back with + ``requires_grad=True``, and differentiating through them would silently + run the full backward and return a plausible but wrong answer rather + than raising. Marking them non-differentiable 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 true gradient here is genuinely non-zero, so returning zeros would + # never have been a defensible 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``), but a caller branching on ``is_supported`` must get the + same answer as the reference path it stands in for, not a different one + -- silently diverging is worse than being slower. All three of + ``is_supported``, the public wrapper, and the raw op 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``. A naive ``_dx_kernel`` misses + this: the compiler contracts ``dy*w - c1`` to ``fma(dy, w, -c1)`` while + ``c1`` was accumulated from the rounded product, so the rounding error + that survives gets amplified by ``rstd = 1/sqrt(eps)``, which is large at + the default ``eps``. ``_backward`` recognises the degenerate case + instead and returns the exact zero; ATen, on the shapes where it runs at + all, does not. + + 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. 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, since that construction costs real host + time on the smallest, host-dispatch-bound shapes. 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 needs enough device + memory that it is marked ``@pytest.mark.slow``. 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. +_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 diverge enough that no tolerance in this suite could + absorb the difference. + """ + 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 this + mean/std ratio, dropping the correction term measurably degrades + ``rstd``'s relative error; the ceiling below is set well below where the + uncorrected path lands. + + The output is not a witness for this: ``y`` is dominated by the fp32 + representation of the mean whether or not the term is present, so 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 not separate +# 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. + + Every elementwise program re-derives the statistics from the split-K + Welford partials rather than reading them from a separate finalize + launch, 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 all but one dweight block + 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 b61ae21f..4b03c606 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 what must not change: outside autocast the block is +# bitwise identical to before, 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 a block-level + assertion would measure 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, which 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 00000000..1aabedfc --- /dev/null +++ b/triton_conv3d/__init__.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""Triton 3-D convolution kernels for NDHWC (``channels_last_3d``) tensors. + +Self-contained by design: it imports nothing from ScaFFold or DistConv, so it +can be vendored into either unchanged. ScaFFold plugs in through a thin adapter +that lives on the ScaFFold side. + +- :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, + 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. Only its *forward* is a kernel: the windows tile rather than + overlap, so both backward directions are the ordinary strided convolution seen + from the other side, served by the two modules above. +- :mod:`triton_conv3d.shapes` -- the convolution problems that occur in 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 entry points take and return ``channels_last_3d`` tensors, and a caller asks +a gate before calling one. + +The gates say nothing about the GPU, deliberately. They are *capability* +predicates -- "will this call be computed correctly here" -- and 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 it tuned 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". +Whether *this* machine is one whose numbers are trustworthy is therefore the +embedder's routing question, and 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``). + +Which gate to ask 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` 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 modules that import torch and triton, so they are +#: re-exported lazily: ``import triton_conv3d.shapes`` must stay free of both, +#: because the shape and cost model 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 00000000..9cca128f --- /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 00000000..bfdd72fc --- /dev/null +++ b/triton_conv3d/bench/baseline.py @@ -0,0 +1,531 @@ +# 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. 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: same solver + (``ConvHipImplicitGemm3DGroupFwdXdlops``) and same device op, but the + heuristic picks small MFMA tiles with narrow global loads where the search + picks large tiles with wide ones. 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. + + ``--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 tens of seconds per call, and unsharded scale 8 trips an assertion inside +MIOpen that can take the process down with it. + +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 result is not a +#: control. 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 it is sized by the harness. + 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 -- what a step actually costs on a shared + # node. The best round is the diagnostic: 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 chosen per cell, so two cells' + # spreads are not comparable to each other. + 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: 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 -- 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 + 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 00000000..81884b93 --- /dev/null +++ b/triton_conv3d/bench/conv_bench.py @@ -0,0 +1,1606 @@ +# 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 names its own operands, its own control, its own candidate configs +and its own reference. ``_build`` is a lookup, so no builder asks which +operator it has, and what is per-operator rather than per-direction -- which +problems it selects, the shape form, the measurement order -- sits on +:class:`_Op`. Each builder names the config type its own entry point resolves, +so it cannot be handed a config for a different kernel; that is what makes the +transposed backward directions sweepable at all. + +Five things this driver is careful about: + +**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. ``--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 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 far worse times 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. ``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, and that find dwarfs 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 -- so the control would time a solver production never runs. + +**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. No field is + computed by asking the *problem* which operator it is: that 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 it is not free: for a *backward* + #: direction the control is a real forward graph, and running it once costs + #: MIOpen's find, which dwarfs 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, which is all six cells: 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 so that a transform, if one came back, would be charged. + 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 table lookup a shipped call makes. The two are not the + #: same number: at the smallest sites that lookup is a large fraction of the + #: 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. + + 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: the + control would then time a solver production never runs. + + 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 would push a backward cell back onto the eager path -- a + launcher-inclusive number on both arms. + """ + 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, as a ScaFFold parameter is, and the kernel reads + # that layout in place. ``weight_rsck=`` would time a path the + # integration does not take. + 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)``. A ``ConvConfig``, then, not a ``TransposedConfig`` -- + passing the latter here would benchmark a tile nothing selects. + """ + 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. + + Each of these 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. + ``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, so a re-capture lines up row for row with the + # transposed captures already on disk. + 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. No launcher excludes a weight transform, because + 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 -- the expensive + half, because running the control once costs MIOpen's *find* (see the module + docstring). + """ + 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: hoisting Triton's + config lookup out while leaving PyTorch's dispatch inside the MIOpen arm + flatters us, and most on the shortest 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 fair only because the *whole comparison* moves together. + The two arms of one cell can differ in host cost by a large factor, 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. Nor is excluding the + launcher a favour to Triton -- it enlarges some wins and turns others into + losses, depending on which arm was paying the larger host cost. + + 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 the per-call duration + :func:`~triton_conv3d.bench.harness.graph_is_worthwhile` allows, nothing + is captured: the host cost is a negligible fraction of either arm there, + and the eager number is already launcher-exclusive to well inside 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 rounds, no tax probe. A fixed small + # ``iters`` is not neutral between the configs being ranked: the first + # call after a synchronize pays a queue restart whose share of the block + # grows as the kernel shortens, so it is charged hardest to the fastest + # config -- exactly the one the sweep is looking for. 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, which is nearly all of a cell's wall clock; 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 gets; + it is not the same number as the *call* a caller makes, which also pays the + table lookup every time, 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 a fixed grid exactly, for a capture that has to be + byte-comparable with an earlier one -- but a fixed grid issues the same + number of calls whatever the kernel costs, and this corpus holds both + microsecond kernels and multi-second ones, where averaging cannot change a + verdict that is already unambiguous. + """ + 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 if it is asked to run somewhere else -- a cost + # 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 measured result + # rather than 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 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: a ``speedup`` quoted without one is not a + # result. + 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. + # + # An estimate, not a measurement: on kernels long enough for a host-paced + # loop to reach kernel throughput it lands within noise of zero, of + # either sign. Taking it seriously means running both launcher policies + # as full races with intervals, which is a separate experiment. + 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 -- a transposed problem included, which is why + ``--problems`` indices are corpus indices and not positions within an + operator's own list. Every printed and 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, so computing it costs the + # same find the run just declined to pay. Refused rather than silently + # made expensive. 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 00000000..d7972261 --- /dev/null +++ b/triton_conv3d/bench/gemm_probe.py @@ -0,0 +1,1256 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""What ``tl.dot`` can reach on this device, at our shapes. + +Every convolution direction reduces to a GEMM, and no amount of clever +addressing recovers throughput the matrix core never had, so the GEMM's ceiling +is the ceiling of the convolution built on it. + +Three measurements, in increasing specificity: + +``peak`` + A large square GEMM. Calibrates Triton against ``torch.matmul`` + (hipBLASLt) and against the bf16 peak in ``PEAK_FLOPS``, 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 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 what 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 probe measures what an ordinary Triton GEMM + achieves, so that a 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 the buffer-load + # fast path requires: ``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 * stride`` overflows i32 + # at the corpus's real shapes -- 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 ``legacy`` knob set passes none of them, + #: so a distinct sentinel keeps that sweep exactly reproducible. + 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: every config costs a +#: JIT compile, and the product over (BM, BN, BK, warps, stages) would spend +#: hours of compilation on 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 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 its 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 the M0 sweep used, so the + report can say which won rather than assume. + + ``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 multiplies peak memory several times over, + which the largest corpus operands do not have to spare. + """ + 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 out of the picture -- + the ceiling a fused convolution kernel is entitled to aim at. + """ + 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 + # 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 the margin over second place + 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, and repeated passes over one unchanged + grid drift by more than the knob effects being looked 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 stay alongside as ``*_sweep_tflops`` 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 -- see +#: :func:`run_peak_compare` for why a before/after needs that. +_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, so its winner is partly +#: whichever config drew the luckiest sample. Racing the top few and taking +#: each set's best gives both sides 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; otherwise one of them takes + # the post-warmup slot more often than the others, which is worth + # more than the effects being compared. + 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 runs for hours, and one + interrupted near the end must not have to start over. + """ + 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: a materialized im2col does not fit. + 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 what a sweep alone reports; a gap + # between it and the head-to-head 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, and a + # neighbouring job on another die drags whole rounds + # down: the median then measures the neighbour, while + # the minimum is the best available estimate of what the + # kernel can 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 00000000..52956452 --- /dev/null +++ b/triton_conv3d/bench/harness.py @@ -0,0 +1,1061 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""Timing that survives this machine -- and says how well it survived it. + +What threatens a comparison on this node is not the clock, and not the elapsed +time between two measurements in one process; it is a foreign tenant on the +device. That is common-mode within a round -- every arm pays it at once -- so +:func:`interleaved` measures every arm in every round and pairs them there, +which a between-run comparison cannot do. It is not a substitute for an error +bar. + +The instrument is the other hazard, twice over. + +1. For a short kernel the per-iteration event pair is a material part of what + it reports, and that tax is per-arm rather than common-mode, so it does + *not* cancel in a ratio. Iterations are therefore grouped behind one event + interval -- one common width for every arm of a comparison, from a rule that + depends only on the measured duration (:func:`_common_group`) -- and + :attr:`Measurement.tax_ms` reports the residual against an event-free + bracket. +2. ``iters`` is not a neutral knob: the first call after a synchronize pays a + queue restart, so ``iters=1`` over-reports, the more so the shorter the + kernel. The adaptive path picks ``iters`` from the measured duration and + never leaves it at 1 for a small kernel. + +What the timed region contains. For a short kernel the host is the pacer, 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. Host launch costs differ +widely between arms -- 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 -- so *leaving* the launcher in is as much a +per-arm instrument as taking it out asymmetrically would be. 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 :data:`_GRAPH_MAX_MS` +per call the host cost is negligible against either arm and no graph is used. + +Every number 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, so a cell that stopped +on the budget with a wide interval is visibly different from one that +converged. + +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. +#: Set below the smallest per-cell difference this corpus reports. +_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 +#: medians over ``iters`` calls and so close to normal; on heavier tails the +#: factor is conservative (the median's true SE is 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 reproducible without an RNG + seed, and at the round counts in use 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 working set small enough to stay + resident, measured hot, reports bandwidth the same kernel will never see + inside a real step, where everything upstream has already flushed it. + + The trap: only a block's first sample is cold, so a caller that flushes + once per block and then reports the *median* over ``iters`` calls throws + that one sample away. 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 without this a naive ``!=`` is always true and every call + # reallocates the buffer on the critical path of a 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 :attr:`median`, and it comes with + :attr:`half_width` -- a 95% interval -- and :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, so the number is not kernel time. Reported + #: rather than hidden, because folding that idle time into the kernel time + #: hides it in a plausible-looking result. + #: + #: 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, and ``rounds`` is + chosen per cell. 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. + + The gap between an event-per-iteration block and an event-free bracket + of the same kernel, so it is independent of the numbers the median came + from -- 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. A ratio + of two independently-reduced medians throws away the property + :func:`interleaved` exists to buy. + """ + + 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 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. + + So time each iteration separately and return the median, which rejects a + stalled launch instead of averaging it in, alongside the ratio of the 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 host + time that lands in the reported number for a short kernel; grouping divides + that by ``group`` while keeping enough samples per block for the median to + still reject a stall. ``group=1`` is what a large kernel gets, because + there the tax is negligible. + """ + 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 a negligible fraction of +#: the kernel, so the tax probe -- ten extra blocks -- is not worth its wall +#: clock on an expensive cell. +_TAX_PROBE_MAX_MS = 1.0 + +#: Floor on the cost of one event interval, in ms. A measurement chose it; +#: it is 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*, +#: because the host also pays per ``record()`` and a launch path expensive +#: enough to make the host the pacer turns that into device idle -- so this is +#: a floor, which the per-arm probe raises. 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 ``tax_budget``. 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. +_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. + + A *paired* difference -- ev, br, ev, br, ... -- rather than 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. Unpaired, it is noisy enough to pick 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, + Triton's JIT 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: a handful of calls, which is cheap for a small + # kernel and unaffordable for a very expensive one -- hence the hard cap. + 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, which over-states a + # small kernel badly: the sync and the queue restart are most of it, and + # sizing ``iters`` off it leaves the block far too short and the instrument + # tax far above its budget. So re-estimate from an event-free bracket, 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: :func:`_probe` absorbs *one* one-off call before it + estimates anything, and one is not always enough. rocBLAS/hipBLASLt loads + its kernel library lazily, on the call *after* the first launch, and a + decision made from that call puts a tiny kernel on the eager path with + ``iters=1``. So this settles first, bounded by wall clock rather than by a + call count, so that a very expensive kernel is called once and a cheap one + several 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. + + *The group is a function of the duration alone*, not of a per-arm + measurement of the instrument tax: a per-arm probe lets two byte-identical + arms pick different groups, and the difference in residual instrument cost + is then a bias in a ratio whose true value is exactly 1. + ``test_a_paired_ratio_of_two_identical_arms_covers_one`` pins that. + + *And it is common to the whole call*, taken from the shortest arm, so that + arms far apart in duration -- or merely far enough apart to straddle a + power-of-two boundary, which byte-identical arms can -- are measured with + the same ruler. ``iters`` stays per-arm, because that is what the corpus's + range of durations needs; 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``. + + Rotating by one position per round balances *positions* but preserves + *adjacency*: with three or more variants B always runs immediately after A, + so whatever A leaves in the caches is a constant charged to B and averaged + out of nothing -- enough to separate two byte-identical arms. + + 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, from fitting +#: ``per_call(chunk) = kernel + cost / chunk`` over graphs of increasing +#: ``chunk``. It is the *worst* arm measured, not the 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 prevent. 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. Set below both the smallest per-cell difference this +#: corpus reports and the harness's target precision. +_REPLAY_BUDGET = 0.01 + +#: At this many calls the residual replay cost is already below the event +#: instrument's own floor (:data:`_EVENT_INTERVAL_MS`); more would only cost +#: capture time. +_MAX_CHUNK = 128 + +#: Above this per-call time no graph is used: the largest host launch cost on +#: this node (the autograd engine's, on the MIOpen backward control) is then a +#: smaller fraction of either arm than the harness's own target precision, so +#: eager timing is already launcher-exclusive and the capture is not worth its +#: wall clock. Below it the host cost can exceed the kernel and decide 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. 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. That cross-stream tax is per-arm -- the Triton arm has no autograd + graph and does not pay it -- and it decides the answer for a short kernel. + + 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. + """ + + 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 happens on the first + call, and a find inside a capture would synchronize and abort it. + + 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. + """ + 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 replay costs the same device + time whatever is inside it, so at ``chunk = 1`` it is a large part of a + short kernel; 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 worst-case host launch cost is small against + the harness's target precision, so eager timing is already + launcher-exclusive 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``. Per-call times in this + corpus span orders of magnitude, so a fixed ``iters`` and ``rounds`` is a + handful of microseconds for one cell and unaffordable 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`. + * ``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 a deliberately reproducible capture should do. ``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 + # Checked every round, not only on a block boundary, so a cell whose + # single round costs minutes can stop without first completing a + # balanced design it cannot afford. ``balanced`` then 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 00000000..c0438240 --- /dev/null +++ b/triton_conv3d/bwd_data.py @@ -0,0 +1,363 @@ +# 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: it is host code around the + forward's kernel, which is also why the tests here reuse the forward's + correctness standards unchanged. +* There is no weight transform either. The flip is a constexpr index + (``taps - 1 - dij``: flipping all three kernel axes is the complement of a + mixed-radix index), and 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. +* The tuning does not transfer from the forward. The effective GEMM has the + channel widths swapped -- a forward ``Cin`` to ``Cout`` runs here as + ``N = Cin``, ``K = Cout * taps`` -- 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. + +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, + ) + + +#: Backward-data winners, keyed by the *forward* problem's +#: ``(dtype, Cin, Cout, kernel)`` -- 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 the same channel pair mean different things here and in +#: the forward's table, which is a trap not worth setting. +#: +#: A sweep chose the rows and ``triton_conv3d/bench/conv_bench.py`` reproduces +#: it. Only channel pairs that were actually timed appear; a miss falls to the +#: heuristic on the effective widths rather than to an extrapolation. +#: +#: Two rules this table encodes that the forward's does *not*: ``BLOCK_N`` runs +#: up to 256, because the GEMM's N is ``Cin`` and the convolutions whose forward +#: is skinny are the widest ones here; and ``BLOCK_K`` drops to 32 where +#: ``Cout`` is small, because the reduction is ``Cout * taps`` and a deep K-tile +#: then wastes the tail. ``GROUP_M`` is 6 (MI300A's XCD count) and +#: ``matrix_instr_nonkdim`` 16 throughout, as in the forward, both re-measured +#: here rather than inherited. +_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``: the reduction is ``Cout * taps = 6``, one MFMA fragment + # deep, so a second pair of waves has nothing to reduce and only replicates + # the addressing. That holds because this site is narrow; it is not a + # direction-wide rule, and warps below 4 lose badly at the wide pairs. + 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: an empty batch has no output voxels, + # so answering ``True`` would assert something this predicate has not + # checked. Costs only a fallback on a call with 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. + # + # * the parameter itself becomes ``(Cin, Cout, kd, kh, kw)``. Channels-last + # makes its ``Cin`` contiguous, which is this GEMM's N -- the forward's + # ``W_ORDER == 0`` load 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`` here is tap ``flip(t)`` of the weight and its + # matrix is the transpose; both are addressing, neither is a copy. + # + # ``out=`` is validated by the forward rather than here, and that is exact: + # the effective forward's output shape *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 00000000..f7db9df8 --- /dev/null +++ b/triton_conv3d/gather_gemm.py @@ -0,0 +1,1422 @@ +# 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: + +* 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; +* 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, though the +obvious design has one: the B tile wants ``(BLOCK_K, BLOCK_N) = (Cin, Cout)`` +per tap and PyTorch stores neither channel axis in that position. The natural +fix -- materialize ``(kd, kh, kw, Cin, Cout)`` once and reuse it -- costs a copy +of every convolution weight per *optimizer step* rather than per call, because +the optimizer dirties every parameter every step and 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 measure the same, +so the choice between them is not a performance question; the third is much +slower addressed in place, because neither tile axis is dense, and is therefore +copied. + +Two designs that look better and measure worse: loading the tile coalesced and +transposing it in registers, where the transpose costs more than the coalescing +saves; and holding the parameter in RSCK order, which helps the kernels and +costs more in the optimizer, since the gradient this package produces is +channels-last and the elementwise update would then be strided. +``triton_conv3d/bench/conv_bench.py`` regenerates both comparisons. + +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. 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 + # once per tap 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. 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 several times slower, which is why + # :func:`_weight_plan` refuses it rather than compiling it. + # + # The two alternatives were implemented and measured worse: a coalesced + # load plus a register transpose into the ``(BLOCK_K, BLOCK_N)`` the dot + # wants, and holding the parameter in RSCK order, which costs the + # *backward* direction, 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. No weight in + # this project comes near int32, but ``taps * Cin * Cout`` is bounded by + # nothing a caller cannot exceed, and past ``2**31`` a 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 + # its buffer-load eligibility; it loses that only at sizes where the + # *storage* is already over the buffer-op limit (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 -- so such a + config has to be refused, not merely ranked last in a 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. + + There is no double-buffering factor: Triton's gfx942 pipeliner keeps one + LDS buffer at ``num_stages=2``, so ``(BLOCK_M*BLOCK_K + + BLOCK_K*BLOCK_N) * itemsize`` is what the compiler's own + ``metadata.shared`` reports. This is the real number, 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 -- so a *sweep* candidate is +#: left to fail, while the config the entry point picks on its own goes through +#: :func:`_fit_to_lds` first. +_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. + + :func:`default_config`'s block sizes are chosen against bf16, and fp32 + operands are twice the bytes, so a wide ``Cin`` in fp32 asks for more than + ``_LDS_BYTES`` and the *shipped* configuration would raise. + ``more_determinism`` runs the model in fp32, so that is 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, because a sweep found the bottleneck shapes +#: *slower* when pushed to a full wave; ``triton_conv3d/bench/conv_bench.py`` +#: reproduces it. +_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 -- ``triton_conv3d/bench/conv_bench.py`` regenerates + the sweep. The surface is flat and almost entirely determined by the channel + widths, so ``BLOCK_M=128``, ``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``, even though 32 wins on a + plain GEMM at ``N=64``. That result does not transfer: the convolution's + inner loop carries a per-tap boundary predicate and a far longer + reduction, so it is not the instruction mix the GEMM probe measured, and + 16 won every convolution cell. + * 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; the winners in this band are +#: 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 without this grid that site has no +#: tile evidence at all: ``candidate_configs`` prunes on ``bn > 2 * n2``, which +#: at ``Cout = 6`` removes every entry of both grids above and leaves only +#: ``default_config``. +#: +#: 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 mostly padding. +#: One warp wins at every head volume. +#: +#: Gated to ``Cout <= 16`` in :func:`candidate_configs` rather than added to the +#: grids above, because a 16-column tile computes far more padding than work at +#: a wide ``Cout`` and would only lengthen every other site's sweep -- and +#: because ``num_warps=1`` is a catastrophe outside this regime, several times +#: slower on the shipped tile in both directions. +_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 as + Inductor does and AMD's guidance says, because at narrow ``Cout`` that advice + is not the winner; ``triton_conv3d/bench/gemm_probe.py`` reproduces the + comparison. + + ``GROUP_M`` defaults to 6 alone -- MI300A's XCD count, and the value that won + the 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 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; see + :func:`default_config`. 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 -- on any other tile the nonkdim alone + changes 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, the same tile won at each. ``GROUP_M`` is the +#: exception; it flips between 6 and 8 across volumes but is worth little either +#: way, 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. +#: 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, because an extrapolated +#: entry in a table called "measured winners" cannot be told apart from a +#: measured one. Tuning them is not worth much at the step level, but a few of +#: them do lose to MIOpen and belong on the adapter's block-list. +_TUNED: dict[tuple, ConvConfig] = { + # The segmentation head. ``Cout = 6`` prunes every seed tile, so the main + # forward sweep never timed anything but ``default_config`` here; this row + # comes from a follow-up race at this site and is that same ``128x16x64`` + # tile with *one warp instead of four* -- see :data:`_NARROW_N_TILES` for + # why one, and why only here. + 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 makes serving it here rather than + # on MIOpen worthwhile. 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`` a 16-deep dot leaves most of + # the matrix-core work multiplying padding this kernel put there + # itself. + # + # ``BLOCK_K = 8`` is the fix; ``nonkdim=32`` is only how that is + # spelled -- see :func:`_tuned`. It beats both the heuristic and + # MIOpen at every volume the corpus has for this pair, which is what + # a row keyed on channel widths alone has to do. + # + # Bitwise identical to the config it replaces, because at + # ``Cin = 3`` only three products per tap are non-zero whatever + # ``BLOCK_K`` is and the 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 materially worse. ``GROUP_M`` is inert at this site. + (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 measures faster. Cout=256 rather than + # 512 is what distinguishes it, on one data point, so the rule + # stands and this is an exception to it. + (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 they fell to + # the heuristic, under which the forward lost to MIOpen. + # + # ``128x64x128`` wins at every volume both pairs occur at, which is + # what a row keyed on channel widths alone has to do, and it is one + # row rather than three because ``_fit_to_grid`` walks ``BLOCK_M`` + # down as ``M`` falls with the sharded volumes. + # + # One caveat travels with this row. ``(2048, 2048)`` is the one + # site in this project whose time depends on *where its weight + # lands*: the weight nearly fills the 256 MiB MALL, so the heuristic + # is bimodal -- two tight states, each fixed for the life of the + # allocation. This row is stable across rebuilds instead, and at + # the unsharded volume it sits between the heuristic's two states: + # clearly better than the unlucky allocation, slightly worse than + # the lucky one. It is shipped because the expected value and both + # sharded volumes are wins and the variance goes away, not because + # it dominates. + # + # ``(2048, 1024)`` is deliberately absent, and that is a result + # rather than an omission: the tile that wins its unsharded volume + # loses badly at both sharded ones. The discriminator is + # ``BLOCK_K`` against ``M``, which a table keyed per channel pair + # cannot express. ``_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* -- many small strided ``permute().contiguous()`` launches, + latency-bound rather than bandwidth-bound, which no caching could remove + because the optimizer dirties every parameter every step. + + 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 far 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. There is no third value: ``_W_GENERAL`` costs nothing +#: measurable against ``_W_N_CONTIG``, and a "coalesce along K and ``tl.trans``" +#: order was implemented, measured 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 severe 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 tap axis, which is not a tile axis. + # Every element of the B tile is then its own cache line, several times + # slower than materializing RSCK even with the copy charged to every call, + # in both directions. 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 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. + + The three direction gates disagree, and the disagreement is a trap: 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 -- that the gradient cannot be computed. ``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 refused by backward-data, whose + kernel-free formulation turns into a scatter into a sub-lattice once the + stride is not 1. + + The direction gates are deliberately left as they are: narrowing the + forward's to the intersection would take strided convolution away from + inference, which asks only the forward, and there is no single "the + backward" answer to agree with anyway. The asymmetry is a fact about the + three kernels; what was wrong was that a caller had to know it. + + 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 with no error raised, 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 the buffer-op limit. + 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; nothing else enforces that a weight is small. ``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 + would be per optimizer step rather than per call: the optimizer dirties + every parameter every step, so no amount of caching removes 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 in either direction. 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 output, which dwarfs the allocation and + # happens once per convolution. + 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 and still 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`` come from the swapped channel + widths, and its tile comes from a different table -- each of which changes + the code that is emitted. + + ``weight_layout`` selects which of the B loads is compiled, and is 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 00000000..22e9b063 --- /dev/null +++ b/triton_conv3d/reduce_gemm.py @@ -0,0 +1,1416 @@ +# 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. 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. 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. + +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 whole +output volume -- millions of trip counts of a six-compare boundary predicate, +with no split-K anywhere. The reuse is correct and hopelessly 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 -- a few dozen, +on a device with 228 CUs. :func:`split_count` derives the split count from the +shape. + +Several taps per tile. 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). It is by far the slower of the +two, and 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, and with one +tap ``Cout`` caps ``BLOCK_M`` while ``Cin`` caps ``BLOCK_N``. Every operand is +then re-read once per tap, at 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 does not veto the wide tile, and padded is the case production takes: +``ScaFFold/unet/conv3d.py`` performs the halo exchange itself and only on axes +that are *genuinely split*, so every ``k > 1`` convolution reaches this kernel +with a non-zero padding on at least one axis. The two-dimensional predicate is +real and it costs something; it costs far less than the arithmetic intensity the +wide tile buys, at every affected site, and +``triton_conv3d/bench/conv_bench.py`` regenerates that comparison. + +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 a run on the MIOpen rung is not bitwise reproducible +unless ``more_determinism`` -- which disables far more than convolution -- is on. + +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. The claim is the one ``ScaFFold/unet/triton_group_norm.py`` + already 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; it is + # the same device fact that makes ``GROUP_M`` 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 -- 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. + 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 the production case is padded, so + # this is the body that runs. 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 because the gradient + can be *smaller* than one program's tile (the ``k=1`` head is 384 elements): + a one-split-at-a-time loop is then a single workgroup serializing one + dependent memory latency per split, and that is the only shape where the + deterministic path costs anything against the atomic one. + """ + 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 knobs only this direction has. + + A subclass rather than more fields on :class:`ConvConfig`, because the + gather directions would never set them and a config printed in a forward + sweep should not grow suffixes it cannot use. Everything else -- the gfx942 + legality rules, the 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. Six is + #: MI300A's. The swizzle buys little on its own, but every row of this + #: direction's tuning table was measured with it on, so turning it off would + #: make those rows 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. A sweep chose four, but the +#: *whole* wave matters more than the value: every program in this kernel does +#: the same amount of work, so a grid of four and a half waves runs five and +#: idles through half of the last one. Hence the snap in :func:`split_count`, +#: which is worth more than the target itself. +_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 a +#: short reduction into a large gradient asks for splits it cannot pay for. +_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``, which at the widest channel pairs is +#: already large for a single split. Bounded here so that no shape in the +#: corpus can ask for an allocation that fails mid-run. +_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 of 256 columns, or 512 + when ``Cout`` is 64 or less and the tile has no height to trade against. + Where ``Cin`` is 256 or more the channels reach 256 columns unaided and + adding taps on top costs; where ``Cout`` is small the taps are the only way + to get there at all. + + ``padded`` is accepted and *not consulted*: the wide tile wins on padded + problems too (see the module docstring). It stays in the signature because + it describes the problem rather than the policy and every caller already + computes it, but nothing here branches on it. + """ + 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)`` rises with the shorter side. + # ``_pow2_at_most`` keeps ``BLOCK_M <= Cout``, which binds at ``Cout = 128``; + # widening N further instead of raising M there is a dead end. See + # ``_TUNED_BWD_W``. + 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. A 256-row, 256-column + # tile with a 64-deep K falls off a cliff at the long-reduction sites, at an + # identical program count, so it is not parallelism. The LDS budget at + # ``256 + 256`` columns lands on exactly 64 in bf16 -- straight into the + # cliff -- so the cap has to be explicit here. 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: + + * up: enough programs to fill the device :data:`_SPLIT_TARGET_WAVES` 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 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 four and a half waves runs five and idles through half of the + last; the split-count curve is mostly that sawtooth, and reading it as + "the cache prefers short chunks" leads to far 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 large and the volume small this is the bound + # that bites, and without it such a site pays 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; that ordering is + # deliberate. ``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. Re-applying + # the bound after the snap was tried and loses: half an idle device + # costs more than a doubled epilogue whenever the grid is that small, + # and a grid that small is exactly 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, for the ``Cout >= 256`` sites where the M axis has the + # room. Both ``BLOCK_K`` are here because the choice between them is + # shape-dependent rather than a preference: 64 wins at the short-reduction + # sites and collapses 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: 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 tuned row. + + ``nk`` defaults to 32 here and to 16 everywhere else in the package; + :data:`_TUNED_BWD_W` says why that asymmetry is deliberate. + """ + 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, + ) + + +#: 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. +#: +#: A sweep over the corpus chose the rows and +#: ``triton_conv3d/bench/conv_bench.py`` reproduces it. Only channel pairs that +#: were actually timed appear; a miss falls to +#: :func:`default_bwd_weight_config` plus :func:`split_count` rather than to an +#: extrapolation. +#: +#: ``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 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. 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% bank-conflict rate -- which leaves the LDS pipe busy while the +#: matrix core idles. 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 rule is per direction, and within this direction it is per channel pair. +#: The forward has no transposed operand and wants 16; backward-data reuses the +#: forward kernel and wants 16; and here the sign reverses at ``Cout = 1024``, +#: so those two pairs keep 16 as well. :func:`default_bwd_weight_config` -- +#: which serves fp32, the transposed operator and every untuned pair -- stays at +#: 16 too: a heuristic is the path with no measurement behind it, and a knob +#: 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. +_TUNED_BWD_W: dict[tuple, BwdWeightConfig] = { + # The segmentation head. ``BLOCK_M`` is ``Cout = 6`` rounded up to 16, so + # ``nonkdim=32`` is illegal here and the axis above does not apply. What + # does apply is ``num_warps``, which ``candidate_bwd_weight_configs`` never + # takes below 4; one warp wins at every head volume. This site can still + # lose to MIOpen, and the row is here so that the adapter's block-list is + # decided against the best this kernel can do. + 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; the other three have no consistent sign and stay on + # ``default_bwd_weight_config``, i.e. at 16. The split is mechanical, not + # luck: this pair is the only transposed site whose tile is + # ``BLOCK_M = 128``, and the other three reach ``BLOCK_M = 256``, which + # already amortises the transpose -- 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, where the N axis 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, issues 512 columns where only 81 are live. + # ``BLOCK_NC = 4`` (which is what ``bnc=4, tb=16`` spells, keeping + # ``BLOCK_N = 64``) issues 128 columns for the same 81. 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 a dense ``BLOCK_N = 128`` is the same ratio. + # + # ``nk=16`` rather than this table's 32, and one warp rather than + # the 4 ``candidate_bwd_weight_configs`` floors at: both were + # measured for this row rather than inherited. + (3, 64): _tuned(64, 4, 16, 64, 1, nk=16), + # ``Cout = 64`` is where TAP_BLOCK earns its existence: the tile + # can only reach 512 columns through the taps, and reaching them is + # worth a large factor over 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. This table is shared by the + # padded and the unpadded form of the problem, and the row wins in + # both. ``nk=32`` is worth a further margin over the ``nk16`` twin + # of the same tile, which is this table's usual sign; + # ``BLOCK_K=32`` loses and is not taken. + (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*, 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 at every volume the + # corpus has for that channel pair -- per pair, not per site, + # because that is what this table is keyed on -- and the taller tile + # wins at all of them. The two smallest sites in the corpus would + # rather have a 16-deep K-tile, and are not tuned for: their pair's + # other volumes are several times larger and want 64. + # + # ``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 behind at all three volumes, while at the sites + # below it is 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, the only rows here that + # keep the forward's value. ``nonkdim=32`` loses at their larger + # volumes; see the note above the table. + (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 +# every ``BLOCK_K`` and volume the corpus has for them: +# +# * ``(128, 256)``: 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)``: a win at the smallest volume and a loss at the two larger +# ones, so the pair is a net loss, and this table is keyed per channel pair +# and cannot split one by volume. A 256-row, 256-column tile with a 64-deep +# K there is the cliff named in :func:`default_bwd_weight_config`. + + +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`` does not select a different row. Every ScaFFold convolution with + ``k > 1`` reaches here with ``padded=True``, and the wide-tile rows win + there too (see the module docstring), so the argument is threaded through + for :func:`default_bwd_weight_config` alone. + + Determinism is unaffected in the sense the package claims it: a different + config is a different reduction order, so results are *not* bitwise equal + across configs, 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. + + Two configs can therefore differ even where :func:`split_count` returns the + same splits: ``matrix_instr_nonkdim`` changes the MFMA fragment, and a + different fragment sums the same products in a different order within one + ``BLOCK_K``. The split count alone is not enough to assert that two configs + agree. Where they do differ the difference 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 and costs + far more than 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 a small ``Cout`` + is stride-identical to one allocated for a large ``Cout`` with the same + ``Cin`` and kernel. Without this clause the large problem writes past the + end of the small allocation, into whatever the caching allocator handed + out next: no fault and no exception, just some other live tensor 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 -- a plausible wrong gradient instead of a crash. The same clause + # is in ``gather_gemm.is_supported``: the caller picks between the gates by + # direction, so 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 and deliberately does *not* follow + ``torch.are_deterministic_algorithms_enabled()``: the atomic path exists to + price determinism, not to be selected, and tying it to torch's flag would + hand every ScaFFold run the nonreproducible path silently. 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, and the + # arithmetic is repeated in the message so that a traceback + # alone is enough to fix the call. ``workspace_elements`` with + # ``split_count``'s own ``splits`` is the supported way to get + # the number. + 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 belong to 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 the widest channel pairs a round trip through fp32 + # partials moves far more traffic than the kernel itself. + 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. Inspecting the unpadded body would describe 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 00000000..f73aa60d --- /dev/null +++ b/triton_conv3d/reference.py @@ -0,0 +1,397 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""Reference implementations and the tolerance policy. + +Three standards, 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 hides -- 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. 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``. At + a real ScaFFold channel width the dense ``{-1,0,1}`` draw is not exactly + representable in bf16: the forward reduces over ``Cin * taps``, which is + 27 648 terms at ``Cin = 1024`` however 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 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. Thinning the weight instead would multiply whole + ``(tap, Cin)`` rows by an exact zero for a whole output channel -- 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): 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. + + 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. + + ``q*K`` -- the 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. The shape is untouched, which is what + lets the forward's bitwise standard run at the real channel widths rather + than skipping them. + + ``headroom`` guards an order-statistic estimate rather than a bound. + 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. That is asked of the actual result rather than of the + worst-case reduction length: 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). + + Good to the bitwise standard: with operands from :func:`make_inputs` under + ``exact=True`` this is the exact answer, and otherwise it is the fp64 result + the other two standards measure error against. fp64 3-D convolution has no + fast path on any backend, so use a small problem -- every bug this suite is + 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 real time in the test suite. + 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: + """The same convolution in the working dtype, as MIOpen computes it. + + The second standard's baseline: error measured against :func:`reference` is + what a kernel has 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 the working dtype, up to one ulp of each element, + which scales with the *largest* element and not the typical one. + + Conflating the two is a trap: 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. + + The store term takes no safety factor of its own: it is a single + deterministic rounding, bounded by half an ulp of the element and so by + ``u_dtype * peak`` outright. Charging several ulps of the peak instead + inflates the static bound until the "no worse than the incumbent" clause of + :func:`assert_close` never applies. fp32 is the exception and is covered by + the other term: there ``u_dtype`` is 2**16 smaller and the accumulation + dominates. + + ``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 and its error wanders, 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 cannot cover, and no tolerance can: a ``tl.dot`` silently + running at a reduced mantissa width sits *under* one 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. With the store + term at one ulp (see :func:`error_bound`) ``margin * incumbent`` is the + operative arm across most bf16 and fp16 shapes. Which arm wins is closest + to a coin toss in fp32, where ``u_dtype`` is 2**16 smaller: the store term + stops dominating and the bound collapses onto MIOpen's own accumulation + error. + """ + 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 00000000..fbb99391 --- /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 00000000..8b6595c4 --- /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 00000000..a789c6a6 --- /dev/null +++ b/triton_conv3d/shapes.py @@ -0,0 +1,752 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""The convolution problems the kernels have to serve. + +Three sources: + +``scaffold_corpus()`` + The distinct convolutions that occur in real ScaFFold runs, extracted from + ``model-analysis/unet_shapes.py`` dumps for the 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 configurations the 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, and this package's +``bwd_data_config`` derives ``M`` from it. They are easy to confuse, so +:attr:`ConvProblem.form` names which one an instance is and every corpus +accessor says which it returns. + +``"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: a ``k // 2`` halo on *every* + dimension listed in ``dc_shard_dims``, split or not, with the padding + zeroed there. :attr:`ConvProblem.halo_variant`. The MIOpen baseline in + ``measured`` was profiled in this form; it is the incumbent's problem. + +``"adapter"`` + What ``ScaFFold/unet/conv3d.py`` -- the shipped Triton rung -- hands the + kernel: 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 runs, 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, so a kernel that exceeds one is a fact about the +#: measurement and 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 counts *bytes*, so it is a different +#: question from whether an element index overflows int32; 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 describe 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 this is the right form to hold an MIOpen + baseline in and not the form the Triton kernels are handed. + + The distinction is not cosmetic: MIOpen keys its find database on the + whole problem descriptor, padding included, so the padded and unpadded + statements tune independently and can land on different kernels, and + :func:`~triton_conv3d.bwd_data.bwd_data_config` 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 larger input and produces a 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 coincide there, and returning an unequal copy would only + invent 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*. + + ``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; the census of running steps behind + :func:`census_corpus` confirms it rather than leaving it inferred. + + Dropping the zero slabs on the unsplit axes is deliberate and separately + verified: ``cat(zeros, x, zeros)`` at ``padding = 0`` is the same + arithmetic as ``padding = k // 2`` on ``x``, and is 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``. + + This is *not* the buffer-op cliff, which counts bytes: a shape can be + well under the element limit while its storage is over the byte limit, + and that is what :attr:`buffer_ops_eligible` answers. No corpus problem + reaches this one in any of the three forms. + """ + 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 means the kernel loses buffer ops on that operand, which costs + real time; see :data:`BUFFER_OP_MAX_BYTES`. 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 as it is. It does + not name the padding, and the padding is what separates the three forms + of the module docstring. 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 are timings *of*, because the profile that produced them ran through + ``distconv_forward``. So it is 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, which is what the shape dump records + and what the FLOP model is naturally expressed in. + """ + 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* -- 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 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* it does not -- the shape dumps behind + :func:`scaffold_corpus` were taken at ``unet_bottleneck_dim = 4``, a + four-layer model, while every step-level measurement runs the shipped + default of 3, a five-layer one with twice the bottleneck width; + * 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 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 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. + 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`` is the + # unsharded scale-8 activation that makes MIOpen assert and sits just + # above it; ``255^3`` is the largest volume of the same shape family + # that still fits an int32 index. The pair differs only in spatial + # extent, so what it brackets is the boundary and not a change of + # channel width or kernel as well. 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 with a ``--form`` flag wants exactly this, in one place: the + mapping from the word a user typed to the shape a kernel is handed. + """ + 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 00000000..74f2eddb --- /dev/null +++ b/triton_conv3d/tests/test_bwd_data.py @@ -0,0 +1,696 @@ +# 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. That does not work here: backward-data reduces over + ``Cout * taps``, so for the corpus's deepest channels that reduction is + long enough to push the exact sum outside bf16's representable integer + range, and every corpus case would skip. + + 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, so all three are covered here. + """ + 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 (see :func:`_corpus_channel_pairs`). +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 kernel flips a tap by walking the fused tap index backwards rather + than materializing a permuted, flipped copy of the weight. That is only + the same thing because the fused index is a mixed-radix number and + complementing every digit complements the number -- a 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 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. 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 + ``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 an unvalidated buffer fails silently either + way it can be wrong: an undersized one is an out-of-bounds device write, + and an NCDHW one is a full-rate kernel returning a scrambled gradient. + ``reduce_gemm`` validates the same parameter; this direction must too. + + 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 transposed spelling a reader might reach for. + 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 00000000..9ef4cd3b --- /dev/null +++ b/triton_conv3d/tests/test_bwd_weight.py @@ -0,0 +1,1431 @@ +# 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, the only way this GEMM fills the device -- + ``M = Cout`` is one or two tile rows. Correctness is :func:`split_count` + agreeing with what the kernel and the reduction pass assume about each other, + so the arithmetic (every voxel in exactly one split) is pinned as well as the + answer; +* determinism as the package words it: bitwise identical run to run *and* + process to process, for one input, dtype, shape, device and tuning config. + Three tests cover it -- in-process repetition, separate interpreters, and a + negative control on the atomic path that must fail; +* 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 asserts the trip + count that rules it out. + +The bitwise standard carries the two guards it does elsewhere in this suite: a +shifted operand must fail the comparison, and the bug unique to this module -- +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. + + Backward-weight reduces over the whole *output volume*, so at a real + ScaFFold shape a sum of ``{-1,0,1}`` products 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 (see + ``shapes.py``): ``p=(1,1,1)`` is what the adapter hands the kernel + unsharded, ``p=(0,1,1)`` what it hands the kernel at two or four shards, and + ``p=(0,0,0)`` what upstream DistConv hands MIOpen. They compile the same + predicate -- this direction reads X at ``o + t - p``, and ``PADDED`` is on + for any non-zero padding -- but differ in the output extent and therefore in + the reduction length, the split count and whether ``BLOCK_K`` divides a row; + 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 split count rounds the grid to whole waves, and that is the property. + + 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. A sweep chose + ``_SPLIT_TARGET_WAVES`` and the sawtooth it removes 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; + ``triton_conv3d/bench/conv_bench.py`` reproduces the sweep. + """ + 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] + + +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 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 + + +#: The largest fp32 partial workspace any corpus problem asks for, in MiB, on +#: the path :func:`conv3d_backward_weight` actually takes. Pinned by a test +#: rather than recorded in a document because it is the number an integration +#: sizes a hoisted ``workspace=`` from, once and out of the step; a figure taken +#: from a subset of the shapes or from the heuristic config alone understates +#: it, and an undersized workspace raises ``ValueError`` mid-run. +_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 traps for anyone editing this: + + * the config has to come from :func:`bwd_weight_config`, which is what the + entry point uses. :func:`default_bwd_weight_config` skips the tuned table + and so gives a different ``BLOCK_M``/``BLOCK_NC``/``TAP_BLOCK``, a + different tile count and a different split count -- a bound certified on + it is 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. + 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 -- re-applying the epilogue bound after the snap -- was + implemented and raced, and it loses badly, because half an idle device costs + more than a doubled epilogue. 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 mid-run is + not a slow kernel) and is re-applied after the snap, so it is never + exceeded. + """ + + 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. + + The forward's shipped heuristic had exactly this hole: ``128x128x128`` is + 64 KiB in bf16 and 128 KiB in fp32, so it 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, and a hole in either gate is a hole in the fallback 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 bites differently here: ``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. + """ + 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 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 there and nothing else in the suite reaches the combination. +_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 + # predicate at a width the corpus 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 -- that branch exists so 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 a bug planted + in the ``d`` or ``h`` half of it passes everything else in the suite. + + It is a production branch: the 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``: the +#: rows that have to survive the padding every production site carries. +#: 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(): + """A tuned ``TAP_BLOCK > 1`` row must be what a *padded* problem resolves. + + A padded problem is the only kind ScaFFold issues: the adapter halos only + the split axis, so every ``k > 1`` site arrives padded. A resolver that + answered differently under padding would mean the tuned table bought nothing + at any real site, and the tuned row is the faster of the two on the padded + production form of every affected cell -- + ``triton_conv3d/bench/conv_bench.py`` reproduces the comparison. + + :func:`test_the_heuristic_widens_tap_block_under_padding_too` pins the same + property for the pairs the table does not list. + """ + 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` answers the same config padded or not: + 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. + + This is the gradient every ``k = 3`` ScaFFold site computes, so a failure + here is a wrong weight gradient in production. + """ + 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 +#: two entries -- the *quintuple* production launches. ``_PADDED_ROW_ALIGNED`` +#: above reaches the first two but not the third, so nothing else in the suite +#: compiles the combination, which runs at every ``k = 3`` site. ``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 other cases hold ``Cin = BLOCK_NC = 32``, and the +#: ragged-``Cin`` tests elsewhere are not row-aligned. +_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, which is what production launches + at every ``k = 3`` site with a widened row -- see + :func:`test_a_tuned_tap_block_row_survives_the_padding`. The triple is + forced by hand rather than resolved, so a retune cannot quietly stop the + combination being compiled. + + The last two cases add the stem's two raggednesses on top, the combination + the shipped ``3 -> 64`` row launches and 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 weight gradient that passes any tolerance test, and at ``k=3`` with a + symmetric volume nothing else in this file catches 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 identity that would make this file unnecessary, and its 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 ``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 has no split-K. Both numbers + are asserted rather than described. + """ + 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 long 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)``. Pinned so that + "deterministic" is not read as "at some cost in accuracy". + """ + 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 many ways, in fp32. +#: The dtype is the 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 a hundred-odd racing splits a repeat that agrees bitwise every time would + mean the atomic path is not doing what it says. + + The dtype is why this is parametrized. The atomic accumulator is fp32 and + the *result* of a bf16 problem is bf16, so a reordering that perturbs the + sum at the fp32 ulp is invisible after the cast: the splits race the whole + time, but the race is under the resolution of the dtype it is observed in. + So the control certifies nothing on a short-``Cout`` bf16 shape -- a change + that made the deterministic path non-deterministic at the ``k=1`` head would + not show there -- and the ``k=1`` head is covered here in fp32, where the + reordering is resolvable in the output. + + 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``: an unchecked small buffer is written eight times past its end, + with 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: a stride check cannot 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 -- pure waste on a buffer this direction allocates + once per parameter per step. + """ + 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 read elsewhere; if + # that number is short, this exception is the whole of what the caller gets, + # so it has to carry the size to allocate. + 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 00000000..c3f476a8 --- /dev/null +++ b/triton_conv3d/tests/test_gather_gemm.py @@ -0,0 +1,1109 @@ +# 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: 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 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` proves that standard has +teeth rather than being vacuously satisfied. + +The tolerance-based tests are the second line: they 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 non-obvious consequence: 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 where exactness is hardest. 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. + + ``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`` faults the GPU rather than falling back to FMA. + + Its failure mode is not the silent FMA fallback the other fields 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. + + The weight has to be one of the operands it looks at: past ``2**31 - 1`` + elements of ``taps * Cin * Cout`` the int32 offset wraps negative and the + GPU faults, while ``is_supported`` still returns ``True``. Meta tensors + carry a ``numel`` and no storage, so this end of it 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, and ``BLOCK_K = 16`` is legal at both -- the constraint is not a + property of the block size alone. + """ + 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. + + A ``128x128x128`` tile is 64 KiB in bf16 and 128 KiB in fp32, so a heuristic + that picks it for ``Cin >= 512`` raises ``OutOfResources`` on any wide fp32 + convolution. ``more_determinism`` runs the model in fp32, so a real ScaFFold + configuration reaches that. + + 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 ``ValueError`` for a wrong-length one, so both have to be + caught for the predicate to be total. + """ + 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 is a property of the bias, so both failures are silent: a short bias + reads whatever is in memory past its end, 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 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. The trap is that a training caller + who asks :func:`is_supported`, gets ``True`` and builds a graph node finds + 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 remains 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 without failing loudly. The gradient reaches the two + backward predicates as a metadata-only stand-in -- a one-element allocation + expanded to the output shape -- which is sound only while they 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. + """ + 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 returns 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 + raises on all three and the Triton path silently would 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. So the kernel's + output is compared against a reference computed for a *deliberately wrong* + gather, and a mismatch is required. + """ + 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. + + Exactness has to come from the operands here, not from the shape. A dense + ``{-1,0,1}`` draw is exact only while the realized sums stay inside the + mantissa, and bf16 holds integers only to 256, while the forward reduces + over ``Cin * taps`` -- 27 648 terms at ``Cin = 1024`` -- whatever the + volume, so the restatement at ``6x7x8`` the two backward suites use buys + nothing. The activations are thinned to :func:`reference.exact_density` + instead, leaving the shape, the channel widths and the weight as they are: + it is the gather that is under test, not the arithmetic. A thinned draw is + also how one accidentally builds a test that compares zeros against zeros, + so two things are asserted rather than assumed -- that the answer is mostly + nonzero, and that the comparison still rejects a one-voxel shift. + + This test does not skip. If the draw is ever not exact, that is a fact + worth failing on rather than stepping around. + """ + 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 a suite that skips every real-shape case still reports a wall + of passes. The shape of the guard differs from theirs, because + :func:`test_corpus_shapes_match_bitwise` has no skip branch left: 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 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. + + All of ``EDGE``, not a prefix of it: a prefix leaves ``batched`` (the only + ``n > 1`` case), ``kernel_aniso``, ``smaller_than_kernel``, ``unpadded``, + ``pointwise`` and both non-bf16 dtypes at the *default* config alone. The + dtype matters in its own right: it moves the MFMA intrinsic's reduction + depth and therefore which ``BLOCK_K`` values are 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(): + """An unvalidated ``out=`` is an out-of-bounds device write. + + The grid is sized from the problem rather than from the buffer, so an + undersized ``out=`` writes past the end with no error and no fault, + surviving only while the allocator slab happens 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 does, so this is the parameter that most needs the + check. + """ + 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, on a path a training step takes once per convolution. 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 for a reason outside this code: every ambiguous 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 dereferences is 0. It is a property of PyTorch's + predicate rather than of ours, so the enumeration over ``{1,2,3}^5`` pins it + rather than assuming it. + """ + 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. 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 much 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 that most needs a test: PyTorch's shape + over this kernel's storage order, which is what an integration allocates to + make the B tile contiguous, and the case where ``to_rsck`` is a no-op and + ``weight`` and ``weight_rsck`` are the same tensor -- so 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, past the + ``2**31 - 1`` an int32 offset can address; it is also where MIOpen itself + asserts 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. + 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 -- where a truncated + offset aliases back towards the start. + """ + 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``, where an int32 offset faults. + + Not reachable from this model -- its widest weight is far below the limit -- + 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``: ``offs_n`` + is a *second* ``addptr``, sign-extended on its own, so excess put there + passes even against an int32 weight offset. The quantity to push past + ``2**31`` is therefore ``taps*Cin*Cout - Cout``. + + ``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 00000000..781c7ac8 --- /dev/null +++ b/triton_conv3d/tests/test_infra.py @@ -0,0 +1,1562 @@ +# SPDX-License-Identifier: (Apache-2.0) +"""Tests for the measurement infrastructure itself. + +Every performance claim rests on the shape model, the cost model, the reference +and the timing harness, so those are cross-checked against PyTorch rather than +against the same arithmetic 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 run on +#: meta tensors and integer arithmetic, so the largest activations cost nothing +#: here and nothing else instantiates the two int32-boundary cases. 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. + + An off-by-one in the padding arithmetic would otherwise propagate silently + into every FLOP count and every roofline. Run on meta tensors, so the + largest activations cost 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 counts the + taps twice. + """ + 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. + + Asserted on the element count rather than on the predicate, so that it + 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`` reads the predicate under the ``needs_int64`` alias. + 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. + + ``conv 128->64 @ 130x258x258`` holds about half of int32's element range, + so an element-counting predicate says nothing about it -- but DistConv's + halo pushes its activation past the 2 GiB limit on the whole storage, which + is what decides buffer-op eligibility and where MIOpen falls off its solver + database. + """ + 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 + ) + + # The only corpus shape on either side of that line, in either shape mode; + # and no corpus shape needs a 64-bit element index, so 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, so a convolution routed + through it reaches MIOpen two voxels larger per axis and unpadded. + ``halo_variant`` reconstructs that from ``halo`` alone; pinning it against + the independently recorded ``halo_in_shape`` is what keeps every number in + ``measured`` attached to the right problem. + + This is the incumbent's form; ScaFFold's own Triton rung issues + :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 + the padded and the halo'd form tune separately and can land on different + kernels. A baseline that labelled them the same would silently let a cell + measured on one be compared against a profile of the other. + """ + 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 production form is what a real step issues, pinned against a census. + + ScaFFold does not route its convolutions through DistConv, so + :meth:`ConvProblem.halo_variant` is not the form it issues: 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. + + Checked against ``census_corpus()`` -- a recording of the shapes and + paddings real ``FastConv3d`` calls handed the kernels -- rather than + against the same arithmetic twice. The segmentation head is excluded on + ``cout``: it is ``n_categories + 1``, a dataset knob, and the corpus and + the census were taken with different values of it. + """ + 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 forms of one site + can differ in either or in both, so any table that mixes forms 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 adapter form is padded at every ``k = 3`` site; the halo form is not. + + If the padded count ever falls to zero, either the adapter has started + haloing every axis or ``shard_halo`` has been confused with ``halo``, and + the backward-weight kernel silently stops compiling the ``PADDED`` body + production launches -- so every adapter-form table would describe a kernel + that no longer runs. + """ + 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. The failure it catches is a per-step FLOP + count divided by a per-call time, which scales the efficiency by the number + of call sites and pushes the shapes that occur twice past the roof. + """ + 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 that bug's symptom. + 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 error can wander past the one ulp ``error_bound`` + # charges by default -- intermittently, which is the worse failure. + # :func:`test_the_incumbents_extra_roundings_are_the_atomic_ones` pins + # both the nondeterminism and the size of the excess, so this is a + # motivated allowance rather than a tolerance nudged until the test + # passed. Only the incumbent gets 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. + + 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; a bound that tracked the RMS alone would be too tight. + """ + 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. + + A single store is bounded by half an ulp of the element outright and has no + random walk to take a factor against. Inflating it makes the static arm + win :func:`reference.assert_close`'s ``max()`` nearly always, which stops + ``test_no_worse_than_miopen`` in the kernel files holding the kernel to the + standard its name claims. + + Pinned arithmetically 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, + which is what the forward and backward-data do: both are bitwise + reproducible here. MIOpen's backward-weight reduces with atomics instead, + so 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. + + What is asserted is the call-to-call spread, not the error on any one call: + the error wanders because the reduction order does, so a single-call + assertion would be as intermittent as the bound it defends. A single + rounding has a spread of exactly zero, which is what the other two + directions measure. The shape is incidental -- the nondeterminism is a + property of the direction. + """ + 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: if the static bound + swallows the incumbent arm, the "no worse than MIOpen by more than + ``margin``" standard is dead code. So pin that the incumbent arm is the + operative one more often than not. + + A floor over the set rather than a per-cell assertion, because which arm + wins is a measurement and does move: a shape where MIOpen happens to be + unusually accurate is exactly the case the ``max()`` exists to stop from + tightening the test beyond what the numerics justify. Hence not + parametrized -- a per-case fixture cannot state a floor over the set. + """ + 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" + # Pinned ``warmup``/``iters``/``rounds`` must still run exactly the calls + # they say: 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 a cyclic rotation one variant runs immediately after another in every + round, so whatever the first leaves in the caches is a constant charged to + the second and averaged out of nothing -- an effect larger than several of + the per-cell differences this project publishes, so the design property is + asserted rather than trusted. + + 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. A + cyclic :func:`_order` fails the second clause at every ``n >= 3``, where it + makes 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. Since + ``rounds`` is chosen per cell, two cells' spreads are not comparable to + each other at all, and the quotable statistic 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. + + The interval exists because at a small kernel the run-to-run scatter of a + fixed protocol is wider than the differences this project publishes. + """ + 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 in per-call time, and a fixed + ``iters``/``rounds`` issues the same call count either way: too few samples + at one end, an unaffordable wait at 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``, 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(): + """Short kernels, with a negative control for the event instrument. + + An ``hipEventRecord`` costs host time, so at a short kernel a block with an + event between every iteration reports materially more than the same + kernel's wall-clock throughput. That is not noise and it does not cancel + in a ratio, because the cost is per-arm. + + 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 and can only be inflated by contention, never + # deflated. A median here fails inside a busy suite, where the harness's + # own stall rejection holds and the reference's absence of it shows. + 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 gives the ungrouped instrument. + 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 properties of ``flush_caches``, both easy to break silently. + + ``torch.device("cuda")`` carries no index and a tensor made on it does, so + a guard comparing the two is always true and every flush allocates a fresh + buffer while the old one is 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 the one cold sample + is precisely the one the median throws away. 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 hottest problems is far slower -- same solver, same device op, but + 16x16 MFMA tiles with 2-element global loads instead of 32x32 with + 8-element loads -- so leaving it off understates MIOpen and turns every + comparison against it into a fabricated speedup. ScaFFold sets it in + ``worker.py``, and so does the profiler the reference numbers come from, + ``prof_bench.py``. + """ + 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 halo'd input +#: stays 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 what they are for, + and still fails if the harness goes 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 slower + than the profile has lost the find. + + Judged on the best round, not the median: on this shared node a + neighbouring job can inflate every round at once, which is a fact about the + node rather than the harness, and a test that fails on it teaches people to + ignore it. The failures this test is for -- a lost find, the wrong shape, + an inert memory format -- are order-of-magnitude and survive 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. + + Bracketing a whole block of iterations with two events silently adds any + launch gap inside it 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 too. When that + # happens the diagnostic is firing correctly and it is the premise -- "this + # run is clean" -- that is false. So take the quietest of several + # attempts, and if none of them is quiet, say the host was loaded rather + # than blame 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 decision has exactly one way to go wrong -- taking the launcher out of +# one arm and not the other, which at these sizes is a large difference in +# whichever direction flatters us. These tests are the guard on that. + + +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 a fixed amount of device time whatever is inside + it, so at ``chunk = 1`` it is a much larger share of a short kernel than of + a long one and a per-arm chunk would be a per-arm instrument -- the failure + ``_common_group`` already documents one level up. Hence the rule reads + only ``min(durations)``, so two arms of the same cell 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(): + """Past ``graph_is_worthwhile``'s threshold the capture is not worth it. + + The largest host launch cost on this node -- the autograd engine's, on the + MIOpen backward control -- is a negligible share of a call that long, well + inside the harness's own precision target, so both arms stay eager and the + exclusion is negligible for both rather than applied to one. At a short + call the same launch cost exceeds 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. + + An empty capture is a ``UserWarning``, and a caller that ignored it would + publish the cost of ``cudaGraphLaunch`` as a kernel time -- the fastest + wrong answer available, and it looks like a spectacular win. Building a + graph on one stream and capturing it on another is enough to produce one, + so the warning is promoted to a refusal. + """ + 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, so + anything else is the instrument, and this runs them through the shipped + decision path. + + 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 pathologically slow at most sizes and fast only at 512 + # and 1024. That is the ``torch.mm`` bf16 pathology this project owes + # upstream; a test that picked one of the slow sizes would be timing a + # kernel long enough that it is correctly told it needs no 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: one race converges to a + # within-race half-width narrower than the between-race scatter of the same + # pair, so an interval that just misses 1 here is that residual rather than + # a biased instrument. A biased instrument -- per-arm event groups, or a + # one-call graph -- is far larger, and 1% catches both. + 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. + + Two 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. + """ + 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 pathologically slow at most sizes and fast only at 512 + # and 1024. That is the ``torch.mm`` bf16 pathology this project owes + # upstream; a test that picked one of the slow sizes would be timing a + # kernel long enough that it is correctly told it needs no 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``. The injected + # host work is an order of magnitude larger than the kernel, so if it were + # inside the timed region 1% would not come close to hiding it. + 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 a fixed amount of device time whatever is in it, so + at a short kernel a one-call graph is largely instrument; 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 pathologically slow at most sizes and fast only at 512 + # and 1024. That is the ``torch.mm`` bf16 pathology this project owes + # upstream; a test that picked one of the slow sizes would be timing a + # kernel long enough that it is correctly told it needs no 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 a large + unearned difference at the short sites. 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 pathologically slow at most sizes and fast only at 512 + # and 1024. That is the ``torch.mm`` bf16 pathology this project owes + # upstream; a test that picked one of the slow sizes would be timing a + # kernel long enough that it is correctly told it needs no 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. + + A transposed convolution is a different operator, not a fourth value of + ``--direction``, so the table has two axes: 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 operator is resolved once, in ``operator_of``. + + The builder it selects never asks ``problem.transposed`` again; a builder + that did would put back the per-arm switch this factoring exists to remove, + so the rule is asserted on the builders' source rather than trusted. + """ + 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 a fabricated control can make MIOpen decline + its own NDHWC path and hand back an inflated speedup. The rule is a grep + rather than a convention because it is easy to reach for by accident. + """ + 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, giving MIOpen a problem it tunes + independently of the padded one. At ``k = 2`` that arithmetic still yields + a halo of 1, which is why the corpus is the authority instead: every + transposed problem in it records ``halo = (0, 0, 0)``, because ScaFFold's + transposed sites are not sharded convolutions at all. Haloing 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. + """ + 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 some 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 the chunk 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. + + A row carrying ``miopen_ms = 0.0`` and ``speedup = 0.0`` reads to 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", so 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, which dominates + the cell. Dropping the arm from the timing while still building it would + save the timing and none of the cost. + """ + 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 00000000..4f608f38 --- /dev/null +++ b/triton_conv3d/tests/test_transposed.py @@ -0,0 +1,1045 @@ +# 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 applying it +anyway overcounts by a factor of ``taps``. ``shapes.py`` accounts for this +correctly; 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 a factor of ``taps``, 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 describes a kernel that is not the one + that runs. + """ + 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, against a fixed 64 KiB LDS budget. + + ``more_determinism`` runs the model in fp32, a real ScaFFold configuration, + and the sibling gather kernel has exceeded the LDS budget there before. + 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 smaller in every spatial dimension. 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 a corpus + case that stopped being representable would fail loudly instead of + silently dropping out of the suite. + """ + 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 is used unchanged: ``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 00000000..0eb28391 --- /dev/null +++ b/triton_conv3d/transposed.py @@ -0,0 +1,1232 @@ +# 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, and the whole design here follows from one property of that case: + + at ``kernel == stride`` and no padding the scatter windows tile the output + rather than overlapping, so every output voxel receives exactly one + contribution. + +The map ``(d, kd) -> D = d*k + kd`` is then a bijection onto ``[0, k*ID)`` -- +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: 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 ``C(grad_output, w)``, so + :func:`~triton_conv3d.gather_gemm.conv3d_forward` serves it with no permute of + the parameter -- the transposed operator's weight already *is* the shape a + ``Cout -> Cin`` convolution wants. +* backward-weight is ``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. + +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, so applying both factors at once overstates the count +by ``taps``. ``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 design question is which axis carries the taps. + +In M -- one output voxel per row, N plain ``Cout`` -- the weight column a row +needs depends on that row's tap, and ``kw`` alternates between adjacent rows +along W, so the B operand would have to vary down the M axis. That is not a +GEMM. + +In N, B is constant per tile and the tile's tap group is a function of the +program id alone. The store becomes a scatter, but a *structured* one: within +one tap the columns are consecutive output channels at one voxel, 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. + +The tap axis is in the tile and not in the grid because every tap of an input +voxel reads the *same* A row: ``taps`` separate programs would read the input +``taps`` times, where one program spanning ``TAP_BLOCK`` taps loads A once for +all of them -- the way :mod:`~triton_conv3d.reduce_gemm` widens its N across +taps, and for 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, and that is a cache + # decision: the programs sharing 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, each group would sweep the whole volume from HBM before the next + # started. 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 have no power-of-two divisor): ``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: memory order + # ``[Cin][kd][kh][kw][Cout]``, i.e. this GEMM's ``[K][tap][N]`` with N dense, + # so it needs no transform -- the opposite of 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 per-tap + # bounds compares the ordinary forward runs do not exist here -- the + # 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 really 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 and LDS rules + 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 nothing says 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 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`` shrinks last and only down to + ``TAP_BLOCK * nonkdim``, below which ``BLOCK_NC`` drops under the MFMA + granularity and the tile is mostly padding. + + The grid clause differs from the gather kernel's in one term that 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 tap + 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 cuts A traffic rather than filling 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 ``ConvTranspose3d`` channel pairs the model contains, swept +#: over the tile and ``TAP_BLOCK`` grid of +#: :func:`candidate_transposed_configs` and 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: a *speedup ratio* does not transfer across +#: volume even when the winning tile does. So only the *tile* is claimed to +#: transfer, and only for a pair whose entry won at every volume that pair +#: occurs at. +#: +#: 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, +#: 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". +_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. + (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 -- no tap +#: factor -- 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. That 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 call + :func:`~triton_conv3d.gather_gemm._weight_plan`; everything else is the same + computation, including why it is a stride test rather than an + ``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 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 gives 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, a dilation interleaves the + # window with holes -- each breaks the bijection ``(d, kd) -> d*k + kd`` that + # makes every output voxel a single contribution. + 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 gives an empty output 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. 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 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 where the ladder needs a fact. + """ + 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, and with peer access + # between the GPUs of a node that reads another rank's data rather than + # faulting. + 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; a strided view of the + # right length silently applies the wrong values. ``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 cannot 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 a predicate that ever grows a stride test will see those + zeros and answer ``False`` -- falling back 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 returns a + scrambled answer at full speed. + """ + 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=...)`` instead + # allocates NCDHW and then copies the whole volume, which dwarfs the + # kernel -- the gather kernel had exactly that 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. + + An ordinary strided forward convolution, as the module docstring derives: + ``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. + """ + 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. The argument order still matches + ``conv3d_backward_weight``'s and the swap happens inside, once, here -- + expecting ``x`` in the first slot of the inner call is the one way to misuse + this function. + + 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()}" + )