From 486e8f6c1aea6c49d24c6a75fa651e9f61750e3b Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 2 Nov 2026 11:25:25 +0000 Subject: [PATCH 1/8] learning: 3D neighbourhood attention, with CUDA passes in the library NeighborhoodAttention3d lets each voxel of a 3D feature map attend to the window^3 voxels around it (the window shifted inwards at an edge), with a learned bias per relative position. Its cost is linear in the volume, so it fits a denoiser at 64^3 where full attention would need over 100 GB. On a card the operation and its gradient run as three passes compiled into the library (src/csrc/ops/attention.cu): gather_dot, gather_sum and a scatter_sum written as a gather, so no atomics and repeatable sums. Float32, float16 and bfloat16, 8 to 64 channels per head, windows up to 9, sums in float32. A build without CUDA exports the entry points and reports them absent; anywhere else a plain-torch path serves. scripts/check_attention.py compares the passes and their gradients with the plain-torch path and times them; --standalone builds attention.cu alone with nvcc, so the kernels can be checked without building BART. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MFB6QPF5rgL8sKizNYzB8D --- CMakeLists.txt | 3 +- scripts/check_attention.py | 164 ++++++++++++++ src/bartorch/_abi.py | 51 +++++ src/bartorch/learning/__init__.py | 5 + src/bartorch/learning/_neighborhood.py | 240 ++++++++++++++++++++ src/csrc/include/bartorch.h | 27 +++ src/csrc/ops/attention.c | 37 ++++ src/csrc/ops/attention.cu | 292 +++++++++++++++++++++++++ tests/test_learning_attention.py | 119 ++++++++++ 9 files changed, 937 insertions(+), 1 deletion(-) create mode 100644 scripts/check_attention.py create mode 100644 src/bartorch/learning/_neighborhood.py create mode 100644 src/csrc/ops/attention.c create mode 100644 src/csrc/ops/attention.cu create mode 100644 tests/test_learning_attention.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 2bc5dad58..8ea87ad6f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -326,6 +326,7 @@ set(BARTORCH_SOURCES src/csrc/ops/sense.c src/csrc/ops/grid.c src/csrc/ops/iter.c + src/csrc/ops/attention.c src/csrc/substitute/backend.c src/csrc/substitute/ref_blas.c @@ -341,7 +342,7 @@ if(OpenMP_C_FOUND AND (WIN32 OR APPLE)) endif() if(BARTORCH_CUDA) - list(APPEND BARTORCH_SOURCES src/csrc/ops/kernels.cu src/csrc/ops/fft_callbacks.cu) + list(APPEND BARTORCH_SOURCES src/csrc/ops/kernels.cu src/csrc/ops/fft_callbacks.cu src/csrc/ops/attention.cu) # The passes around a volume's transform, which cuFFT links into the # transform when a plan is made (src/csrc/ops/fft_callbacks.cu). It takes them as diff --git a/scripts/check_attention.py b/scripts/check_attention.py new file mode 100644 index 000000000..1981032d5 --- /dev/null +++ b/scripts/check_attention.py @@ -0,0 +1,164 @@ +"""Check and time the neighbourhood-attention CUDA passes on a card. + + python scripts/check_attention.py # the library's own passes + python scripts/check_attention.py --standalone # attention.cu alone, built by nvcc + +With ``--standalone`` the one translation unit is compiled into a library of +its own and used in place of ``libbartorch``'s passes, so a change to the +kernels can be checked on a card without building BART: it needs ``nvcc``. + +The passes and their gradients are compared with the plain-torch path in +float64, in each element type, and then timed at the sizes a denoiser meets: +forward and forward-backward, against the plain-torch path, with peak memory. +""" + +from __future__ import annotations + +import argparse +import ctypes +import json +import subprocess +import tempfile +import time +from pathlib import Path + +import torch + +from bartorch.learning import _neighborhood as na + +ROOT = Path(__file__).resolve().parent.parent + + +def standalone(arch: str | None) -> ctypes.CDLL: + """``attention.cu`` compiled on its own, with the ctypes signatures of the library's ABI.""" + out = Path(tempfile.mkdtemp()) / "libna3d.so" + if arch is None: + major, minor = torch.cuda.get_device_capability() + arch = f"{major}{minor}" + command = [ + "nvcc", "-O3", "-std=c++17", "-shared", "-Xcompiler", "-fPIC", + f"-gencode=arch=compute_{arch},code=sm_{arch}", + "-I", str(ROOT / "src" / "csrc"), + "-o", str(out), str(ROOT / "src" / "csrc" / "ops" / "attention.cu"), + ] # fmt: skip + subprocess.run(command, check=True) + lib = ctypes.CDLL(str(out)) + floats = ctypes.POINTER(ctypes.c_float) + ints = [ + ctypes.c_int64, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ] + lib.bartorch_na3d_built.restype = ctypes.c_int + for name, first, third in ( + ("gather_dot", ctypes.c_void_p, floats), + ("gather_sum", floats, ctypes.c_void_p), + ("scatter_sum", floats, ctypes.c_void_p), + ): + fn = getattr(lib, f"bartorch_na3d_{name}") + fn.restype = ctypes.c_int + fn.argtypes = [first, ctypes.c_void_p, third, *ints, ctypes.c_void_p] + return lib + + +def check(dtype, channels, window, shape=(2, 2, 11, 9, 13)): + """Largest error of the output and of each gradient against the float64 plain-torch path.""" + torch.manual_seed(0) + full = (*shape, channels) + base = [torch.randn(full, device="cuda", dtype=torch.float64) for _ in range(3)] + bias64 = 0.1 * torch.randn(shape[1], window**3, device="cuda", dtype=torch.float64) + grad = torch.randn(full, device="cuda", dtype=torch.float64) + + q, k, v = (t.to(dtype).requires_grad_() for t in base) + bias = bias64.to(dtype).requires_grad_() + made = na.neighborhood_attention3d(q, k, v, window, bias) + made.backward(grad.to(dtype)) + + rq, rk, rv = (t.clone().requires_grad_() for t in base) + rb = bias64.clone().requires_grad_() + groups = shape[0] * shape[1] + flat = [t.reshape(groups, -1, channels) for t in (rq, rk, rv)] + want = na._attend_reference(*flat, shape[2:], window, rb, shape[1]).reshape(full) + want.backward(grad) + pairs = {"out": (made, want), "dq": (q.grad, rq.grad), "dk": (k.grad, rk.grad)} + pairs.update({"dv": (v.grad, rv.grad), "dbias": (bias.grad, rb.grad)}) + return {k: float((a.double() - b).abs().max() / b.abs().max()) for k, (a, b) in pairs.items()} + + +def timed(fn, repeat=5): + fn() + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + start = time.perf_counter() + for _ in range(repeat): + fn() + torch.cuda.synchronize() + return (time.perf_counter() - start) / repeat, (torch.cuda.max_memory_allocated() - base) / 1e9 + + +def bench(extent, heads, channels, window, dtype=torch.float16): + block = na.NeighborhoodAttention3d(heads * channels, heads, window=window).cuda() + x = torch.randn(1, heads * channels, *(extent,) * 3, device="cuda", requires_grad=True) + row = {"extent": extent, "heads": heads, "channels/head": channels, "window": window} + + def forward(): + with torch.no_grad(), torch.autocast("cuda", dtype=dtype): + block(x) + + def train(): + with torch.autocast("cuda", dtype=dtype): + block(x).float().square().mean().backward() + + row["forward s"], row["forward GB"] = timed(forward) + row["train s"], row["train GB"] = timed(train) + saved = na.kernels_available + na.kernels_available = lambda: False + try: + row["torch forward s"], row["torch forward GB"] = timed(forward) + except torch.cuda.OutOfMemoryError: + row["torch forward s"] = row["torch forward GB"] = "out of memory" + finally: + na.kernels_available = saved + torch.cuda.empty_cache() + return {k: round(v, 4) if isinstance(v, float) else v for k, v in row.items()} + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--standalone", action="store_true", help="build attention.cu alone with nvcc") + ap.add_argument( + "--arch", help="compute capability without the dot, e.g. 89; the card's by default" + ) + args = ap.parse_args() + if args.standalone: + lib = standalone(args.arch) + na._library = lambda: lib + if not na.kernels_available(): + raise SystemExit("no card, or no CUDA passes in this library") + print(torch.cuda.get_device_name(), flush=True) + + worst = {} + for dtype in (torch.float32, torch.float16, torch.bfloat16): + for channels in (8, 16, 32, 64): + for window in (3, 5, 7): + errors = check(dtype, channels, window) + key = str(dtype).removeprefix("torch.") + worst[key] = max(worst.get(key, 0.0), *errors.values()) + print("largest relative error, output and gradients:", json.dumps(worst), flush=True) + + for extent, heads, channels, window in ( + (64, 1, 16, 5), (64, 2, 16, 5), (32, 2, 32, 5), + (32, 4, 32, 7), (16, 4, 32, 7), (16, 8, 16, 9), + ): # fmt: skip + print(json.dumps(bench(extent, heads, channels, window)), flush=True) + + +if __name__ == "__main__": + main() diff --git a/src/bartorch/_abi.py b/src/bartorch/_abi.py index 4c5e4400f..c04f3f03d 100644 --- a/src/bartorch/_abi.py +++ b/src/bartorch/_abi.py @@ -38,6 +38,9 @@ BARTORCH_ENCODING_SEGMENTED = 6 BARTORCH_ENCODING_STACKED = 7 BARTORCH_ENCODING_ITEMS = 8 +BARTORCH_NA_FLOAT32 = 0 +BARTORCH_NA_FLOAT16 = 1 +BARTORCH_NA_BFLOAT16 = 2 class Encoding(ctypes.Structure): @@ -346,6 +349,10 @@ class Encoding(ctypes.Structure): "bartorch_nufft_reset_counters", "bartorch_toeplitz_counter", "bartorch_toeplitz_reset_counters", + "bartorch_na3d_built", + "bartorch_na3d_gather_dot", + "bartorch_na3d_gather_sum", + "bartorch_na3d_scatter_sum", "bartorch_on_device", ) @@ -1255,6 +1262,50 @@ def bind(lib: ctypes.CDLL) -> ctypes.CDLL: lib.bartorch_toeplitz_counter.argtypes = [ctypes.c_int] lib.bartorch_toeplitz_reset_counters.restype = None lib.bartorch_toeplitz_reset_counters.argtypes = [] + lib.bartorch_na3d_built.restype = ctypes.c_int + lib.bartorch_na3d_built.argtypes = [] + lib.bartorch_na3d_gather_dot.restype = ctypes.c_int + lib.bartorch_na3d_gather_dot.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_float), + ctypes.c_int64, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_void_p, + ] + lib.bartorch_na3d_gather_sum.restype = ctypes.c_int + lib.bartorch_na3d_gather_sum.argtypes = [ + ctypes.POINTER(ctypes.c_float), + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int64, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_void_p, + ] + lib.bartorch_na3d_scatter_sum.restype = ctypes.c_int + lib.bartorch_na3d_scatter_sum.argtypes = [ + ctypes.POINTER(ctypes.c_float), + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int64, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_void_p, + ] lib.bartorch_on_device.restype = ctypes.c_int lib.bartorch_on_device.argtypes = [ctypes.c_void_p] return lib diff --git a/src/bartorch/learning/__init__.py b/src/bartorch/learning/__init__.py index 5cc9249a3..0603ab1d4 100644 --- a/src/bartorch/learning/__init__.py +++ b/src/bartorch/learning/__init__.py @@ -9,6 +9,8 @@ :class:`bartorch.priors.ImplicitPrior`. :func:`split` partitions acquired samples for self-supervised training, and :func:`moments` and :func:`calibrate` give a voxel-wise uncertainty. +:class:`NeighborhoodAttention3d` lets each voxel of a 3D feature map attend to +a window around it, through CUDA kernels of this library on a card. :class:`Reconstruction` trains a network in one of the three stages of an unrolled network, on ``lightning``, and :class:`RandomGain` is a ``torchio`` @@ -19,6 +21,7 @@ from __future__ import annotations from bartorch.learning._channels import ComplexNet, as_complex, as_real +from bartorch.learning._neighborhood import NeighborhoodAttention3d, neighborhood_attention3d from bartorch.learning._nets import UNet from bartorch.learning._patches import Patchwise from bartorch.learning._splitting import split @@ -30,6 +33,7 @@ __all__ = [ "ComplexNet", + "NeighborhoodAttention3d", "Patchwise", "RandomGain", "Reconstruction", @@ -39,6 +43,7 @@ "as_real", "calibrate", "moments", + "neighborhood_attention3d", "split", ] diff --git a/src/bartorch/learning/_neighborhood.py b/src/bartorch/learning/_neighborhood.py new file mode 100644 index 000000000..bebf6f62e --- /dev/null +++ b/src/bartorch/learning/_neighborhood.py @@ -0,0 +1,240 @@ +"""Neighbourhood attention over 3D volumes: each voxel attends to a window around it.""" + +from __future__ import annotations + +import ctypes + +import torch +from torch import nn + +__all__ = ["NeighborhoodAttention3d", "neighborhood_attention3d"] + +_DTYPES = {torch.float32: 0, torch.float16: 1, torch.bfloat16: 2} +_CHANNELS = (8, 16, 32, 64) +_MAX_WINDOW = 9 + + +def _starts(n: int, window: int, device) -> torch.Tensor: + """First voxel of each voxel's window along an axis of ``n``, shifted inwards at the edges.""" + i = torch.arange(n, device=device) + return (i - window // 2).clamp(0, n - window) + + +def _neighbours(extent, window: int, device) -> torch.Tensor: + """``(voxels, window^3)`` index of each neighbour of each voxel, z-major.""" + nz, ny, nx = extent + offsets = torch.arange(window, device=device) + z = (_starts(nz, window, device)[:, None] + offsets)[:, None, None, :, None, None] + y = (_starts(ny, window, device)[:, None] + offsets)[None, :, None, None, :, None] + x = (_starts(nx, window, device)[:, None] + offsets)[None, None, :, None, None, :] + return ((z * ny + y) * nx + x).reshape(nz * ny * nx, window**3) + + +def _attend_reference(q, k, v, extent, window, bias, heads): + """Neighbourhood attention in plain torch on ``(groups, voxels, channels)``.""" + index = _neighbours(extent, window, q.device) + scale = q.shape[-1] ** -0.5 + kn = k[:, index] # (groups, voxels, window^3, channels) + logits = torch.einsum("gvc,gvmc->gvm", q.float() * scale, kn.float()) + if bias is not None: + logits = logits + bias.float().repeat(q.shape[0] // heads, 1)[:, None, :] + attention = logits.softmax(-1) + return torch.einsum("gvm,gvmc->gvc", attention, v[:, index].float()).to(q.dtype) + + +def _library(): + from bartorch._lib import library + + return library() + + +def kernels_available() -> bool: + """Whether this library carries the CUDA kernels and a card is present.""" + try: + return bool(_library().bartorch_na3d_built()) and torch.cuda.is_available() + except (OSError, AttributeError): + return False + + +def _floats(t: torch.Tensor): + return ctypes.cast(ctypes.c_void_p(t.data_ptr()), ctypes.POINTER(ctypes.c_float)) + + +def _call(name: str, x, y, out, extent, window: int): + """One CUDA pass; ``x`` and ``out`` are float32 weights or elements, as the pass takes them.""" + groups, _, channels = y.shape + nz, ny, nx = extent + if name == "gather_dot": # elements, elements -> weights + first, third = ctypes.c_void_p(x.data_ptr()), _floats(out) + else: # weights, elements -> elements + first, third = _floats(x), ctypes.c_void_p(out.data_ptr()) + with torch.cuda.device(y.device): + stream = torch.cuda.current_stream(y.device).cuda_stream + error = getattr(_library(), f"bartorch_na3d_{name}")( + first, + ctypes.c_void_p(y.data_ptr()), + third, + groups, + nz, + ny, + nx, + channels, + window, + _DTYPES[y.dtype], + ctypes.c_void_p(stream), + ) + if error: + raise RuntimeError(f"bartorch_na3d_{name} failed with CUDA error {error}") + return out + + +def _gather_dot(a, b, extent, window): + out = torch.empty(*a.shape[:2], window**3, device=a.device, dtype=torch.float32) + return _call("gather_dot", a, b, out, extent, window) + + +def _gather_sum(s, b, extent, window): + return _call("gather_sum", s, b, torch.empty_like(b), extent, window) + + +def _scatter_sum(s, a, extent, window): + return _call("scatter_sum", s, a, torch.empty_like(a), extent, window) + + +class _Attend(torch.autograd.Function): + """Neighbourhood attention through the CUDA passes, with its gradient through them too.""" + + @staticmethod + def forward(ctx, q, k, v, bias, extent, window, heads): + scale = q.shape[-1] ** -0.5 + qs = (q * scale).contiguous() + k, v = k.to(q.dtype).contiguous(), v.to(q.dtype).contiguous() + logits = _gather_dot(qs, k, extent, window) + if bias is not None: + logits = logits.view(-1, heads, *logits.shape[1:]).add_(bias.float()[None, :, None, :]) + logits = logits.view(-1, *logits.shape[2:]) + attention = logits.softmax(-1) + out = _gather_sum(attention, v, extent, window) + ctx.save_for_backward(qs, k, v, attention) + ctx.extent, ctx.window, ctx.heads, ctx.scale = extent, window, heads, scale + ctx.has_bias = bias is not None + return out + + @staticmethod + def backward(ctx, grad): + qs, k, v, attention = ctx.saved_tensors + extent, window = ctx.extent, ctx.window + grad = grad.to(qs.dtype).contiguous() + d_attention = _gather_dot(grad, v, extent, window) + d_logits = attention * (d_attention - (d_attention * attention).sum(-1, keepdim=True)) + dq = _gather_sum(d_logits, k, extent, window) * ctx.scale + dk = _scatter_sum(d_logits, qs, extent, window) + dv = _scatter_sum(attention, grad, extent, window) + d_bias = None + if ctx.has_bias: + d_bias = d_logits.view(-1, ctx.heads, *d_logits.shape[1:]).sum((0, 2)) + return dq, dk, dv, d_bias, None, None, None + + +def neighborhood_attention3d(q, k, v, window: int = 5, bias=None) -> torch.Tensor: + """Scaled dot-product attention of each voxel over the ``window^3`` voxels around it. + + The window is shifted inwards at an edge rather than cut, so each voxel + sees the same number of neighbours. On a card with the library's CUDA + kernels, and for 8, 16, 32 or 64 channels per head, the passes and their + gradients run there, with sums in float32 and memory linear in the + volume; anywhere else a plain-torch version serves, which gathers every + neighbour and needs ``window^3`` times the memory of ``k``. + + Parameters + ---------- + q, k, v : torch.Tensor + ``(n, heads, z, y, x, channels)``: queries, keys and values. + window : int, default=5 + Odd edge of the window, at most 9 and at most each extent. + bias : torch.Tensor, default=None + ``(heads, window^3)`` added to the logits: a bias for each relative position. + + Returns + ------- + torch.Tensor + ``(n, heads, z, y, x, channels)``, of ``q``'s dtype. + """ + window = int(window) + if window < 1 or window % 2 == 0 or window > _MAX_WINDOW: + raise ValueError(f"the window is an odd edge up to {_MAX_WINDOW}, not {window}") + if q.ndim != 6 or k.shape != q.shape or v.shape != q.shape: + raise ValueError("q, k and v are (n, heads, z, y, x, channels), all of one shape") + n, heads, nz, ny, nx, channels = q.shape + extent = (nz, ny, nx) + if min(extent) < window: + raise ValueError(f"every extent is at least the window ({window}), and {extent} is not") + if bias is not None and tuple(bias.shape) != (heads, window**3): + raise ValueError( + f"bias is (heads, window^3) = ({heads}, {window**3}), not {tuple(bias.shape)}" + ) + + flat = [t.reshape(n * heads, nz * ny * nx, channels) for t in (q, k, v)] + fast = ( + q.is_cuda + and channels in _CHANNELS + and q.dtype in _DTYPES + and k.device == q.device + and v.device == q.device + and kernels_available() + ) + if fast: + out = _Attend.apply(*flat, bias, extent, window, heads) + else: + out = _attend_reference(*flat, extent, window, bias, heads) + return out.reshape(q.shape) + + +class NeighborhoodAttention3d(nn.Module): + """Multi-head neighbourhood attention over ``(n, channels, z, y, x)`` feature maps. + + Queries, keys and values are 1-wide convolutions of the input, each head + attends over the ``window^3`` voxels around each voxel with a learned bias + for each relative position, and a 1-wide convolution mixes the heads. The + cost is linear in the volume, ``window^3`` times that of a 1-wide + convolution, against the square of the volume for full attention. + + Parameters + ---------- + channels : int + Input and output channels. + heads : int, default=None + Heads; ``channels`` split into heads of 16 when omitted, or one head + for fewer than 16 channels. + window : int, default=5 + Odd edge of each voxel's window, at most 9. + + Examples + -------- + >>> block = learning.NeighborhoodAttention3d(32, window=5) + >>> features = features + block(features) + """ + + def __init__(self, channels: int, heads: int | None = None, *, window: int = 5): + super().__init__() + channels = int(channels) + heads = max(channels // 16, 1) if heads is None else int(heads) + if channels % heads: + raise ValueError(f"{channels} channels do not split into {heads} heads") + self.heads, self.window = heads, int(window) + self.qkv = nn.Conv3d(channels, 3 * channels, 1) + self.project = nn.Conv3d(channels, channels, 1) + self.bias = nn.Parameter(torch.zeros(heads, self.window**3)) + nn.init.trunc_normal_(self.bias, std=0.02) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + n, c, nz, ny, nx = x.shape + q, k, v = ( + self.qkv(x) + .reshape(n, 3, self.heads, c // self.heads, nz, ny, nx) + .permute(1, 0, 2, 4, 5, 6, 3) + ) + made = neighborhood_attention3d( + q.contiguous(), k.contiguous(), v.contiguous(), self.window, self.bias.to(q.dtype) + ) + return self.project(made.permute(0, 1, 5, 2, 3, 4).reshape(n, c, nz, ny, nx)) diff --git a/src/csrc/include/bartorch.h b/src/csrc/include/bartorch.h index 335158487..dee835568 100644 --- a/src/csrc/include/bartorch.h +++ b/src/csrc/include/bartorch.h @@ -939,6 +939,33 @@ BARTORCH_API void bartorch_nufft_reset_counters(void); BARTORCH_API int64_t bartorch_toeplitz_counter(int which); BARTORCH_API void bartorch_toeplitz_reset_counters(void); +/* + * Neighbourhood attention over a 3D grid, for bartorch.learning. Each voxel + * attends to the window x window x window voxels around it, the window shifted + * inwards at an edge. Arrays are (groups, voxels, channels), channels fastest, + * voxels z-major; weights are (groups, voxels, window^3) float32. Channels are + * 8, 16, 32 or 64, the window odd and at most 9, and the element type one of + * the BARTORCH_NA_* below. Each runs on the cudaStream_t passed as a void* and + * returns 0 or the CUDA error it met; without CUDA, built() is 0 and the rest + * return -1. + * + * gather_dot s[i, m] = sum_c a[i, c] b[n(i, m), c] + * gather_sum o[i, c] = sum_m s[i, m] b[n(i, m), c] + * scatter_sum o[j, c] = sum over (i, m) with n(i, m) = j of s[i, m] a[i, c] + */ +enum bartorch_na_dtype { + BARTORCH_NA_FLOAT32 = 0, + BARTORCH_NA_FLOAT16 = 1, + BARTORCH_NA_BFLOAT16 = 2, +}; +BARTORCH_API int bartorch_na3d_built(void); +BARTORCH_API int bartorch_na3d_gather_dot(const void* a, const void* b, float* s, int64_t groups, + int nz, int ny, int nx, int channels, int window, int dtype, void* stream); +BARTORCH_API int bartorch_na3d_gather_sum(const float* s, const void* b, void* o, int64_t groups, + int nz, int ny, int nx, int channels, int window, int dtype, void* stream); +BARTORCH_API int bartorch_na3d_scatter_sum(const float* s, const void* a, void* o, int64_t groups, + int nz, int ny, int nx, int channels, int window, int dtype, void* stream); + /* Whether a pointer is device memory; always false without CUDA. */ BARTORCH_API int bartorch_on_device(const void* ptr); diff --git a/src/csrc/ops/attention.c b/src/csrc/ops/attention.c new file mode 100644 index 000000000..557e65301 --- /dev/null +++ b/src/csrc/ops/attention.c @@ -0,0 +1,37 @@ +/* + * The neighbourhood-attention entry points of a build without CUDA: the + * passes are in attention.cu, and here they report that they are absent. + */ +#include + +#include "include/bartorch.h" + +#ifndef USE_CUDA + +int bartorch_na3d_built(void) { return 0; } + +int bartorch_na3d_gather_dot(const void* a, const void* b, float* s, int64_t groups, + int nz, int ny, int nx, int channels, int window, int dtype, void* stream) +{ + (void)a; (void)b; (void)s; (void)groups; (void)nz; (void)ny; (void)nx; + (void)channels; (void)window; (void)dtype; (void)stream; + return -1; +} + +int bartorch_na3d_gather_sum(const float* s, const void* b, void* o, int64_t groups, + int nz, int ny, int nx, int channels, int window, int dtype, void* stream) +{ + (void)s; (void)b; (void)o; (void)groups; (void)nz; (void)ny; (void)nx; + (void)channels; (void)window; (void)dtype; (void)stream; + return -1; +} + +int bartorch_na3d_scatter_sum(const float* s, const void* a, void* o, int64_t groups, + int nz, int ny, int nx, int channels, int window, int dtype, void* stream) +{ + (void)s; (void)a; (void)o; (void)groups; (void)nz; (void)ny; (void)nx; + (void)channels; (void)window; (void)dtype; (void)stream; + return -1; +} + +#endif diff --git a/src/csrc/ops/attention.cu b/src/csrc/ops/attention.cu new file mode 100644 index 000000000..df581945b --- /dev/null +++ b/src/csrc/ops/attention.cu @@ -0,0 +1,292 @@ +/* + * Neighbourhood attention over a 3D grid, on a card. + * + * Each voxel attends to the w x w x w voxels around it. Near an edge the + * window is shifted inwards rather than cut, so every voxel sees w^3 + * neighbours: along an axis of n voxels the window of voxel i starts at + * min(max(i - w/2, 0), n - w), which needs n >= w. + * + * Three passes make up the operation and its gradient. The arrays are + * (groups, voxels, channels), channels fastest, where a group is one head + * of one item; the voxels are z-major (x fastest) and the w^3 neighbours of + * a voxel are ordered the same way. Weights between a voxel and its + * neighbours are (groups, voxels, w^3) in float32 throughout. + * + * gather_dot s[i, m] = sum_c a[i, c] b[n(i, m), c] logits; dL/dattn + * gather_sum o[i, c] = sum_m s[i, m] b[n(i, m), c] output; dL/dq + * scatter_sum o[j, c] = sum_{i, m : n(i, m) = j} s[i, m] a[i, c] dL/dv; dL/dk + * + * scatter_sum is written as a gather over the voxels whose window holds j, + * so it needs no atomics and its sums are the same from run to run. + * + * Channels are a template parameter, and the element type is float, half or + * bfloat16 with every sum in float. Each pass runs on the caller's stream + * and returns 0, or the CUDA error it met. No torch here: the host owns the + * memory and hands over its pointers. + */ +#include +#include +#include + +#include + +#include "include/bartorch.h" + +namespace { + +struct grid3 { + int nz, ny, nx; /* voxels per axis */ + int w; /* window edge */ + int64_t groups; /* heads times items */ +}; + +__device__ __forceinline__ float to_float(float v) { return v; } +__device__ __forceinline__ float to_float(__half v) { return __half2float(v); } +__device__ __forceinline__ float to_float(__nv_bfloat16 v) { return __bfloat162float(v); } + +template __device__ __forceinline__ T from_float(float v); +template <> __device__ __forceinline__ float from_float(float v) { return v; } +template <> __device__ __forceinline__ __half from_float<__half>(float v) { return __float2half(v); } +template <> __device__ __forceinline__ __nv_bfloat16 from_float<__nv_bfloat16>(float v) { return __float2bfloat16(v); } + +/* First voxel of the window of voxel i along an axis of n voxels. */ +__device__ __forceinline__ int window_start(int i, int n, int w) +{ + int s = i - w / 2; + s = (s < 0) ? 0 : s; + return (s > n - w) ? n - w : s; +} + +template +__global__ void gather_dot(grid3 g, const T* __restrict__ a, const T* __restrict__ b, float* __restrict__ s) +{ + const int64_t voxels = (int64_t)g.nz * g.ny * g.nx; + const int w3 = g.w * g.w * g.w; + + for (int64_t t = blockIdx.x * (int64_t)blockDim.x + threadIdx.x; t < g.groups * voxels; + t += (int64_t)blockDim.x * gridDim.x) { + + const int64_t group = t / voxels; + const int64_t i = t % voxels; + const int x = (int)(i % g.nx), y = (int)((i / g.nx) % g.ny), z = (int)(i / ((int64_t)g.nx * g.ny)); + const int sz = window_start(z, g.nz, g.w), sy = window_start(y, g.ny, g.w), sx = window_start(x, g.nx, g.w); + + float q[C]; +#pragma unroll + for (int c = 0; c < C; c++) + q[c] = to_float(a[t * C + c]); + + const T* base = b + group * voxels * C; + float* out = s + t * w3; + int m = 0; + + for (int dz = 0; dz < g.w; dz++) + for (int dy = 0; dy < g.w; dy++) + for (int dx = 0; dx < g.w; dx++, m++) { + + const T* kv = base + (((int64_t)(sz + dz) * g.ny + (sy + dy)) * g.nx + (sx + dx)) * C; + float acc = 0.f; +#pragma unroll + for (int c = 0; c < C; c++) + acc += q[c] * to_float(kv[c]); + out[m] = acc; + } + } +} + +template +__global__ void gather_sum(grid3 g, const float* __restrict__ s, const T* __restrict__ b, T* __restrict__ o) +{ + const int64_t voxels = (int64_t)g.nz * g.ny * g.nx; + const int w3 = g.w * g.w * g.w; + + for (int64_t t = blockIdx.x * (int64_t)blockDim.x + threadIdx.x; t < g.groups * voxels; + t += (int64_t)blockDim.x * gridDim.x) { + + const int64_t group = t / voxels; + const int64_t i = t % voxels; + const int x = (int)(i % g.nx), y = (int)((i / g.nx) % g.ny), z = (int)(i / ((int64_t)g.nx * g.ny)); + const int sz = window_start(z, g.nz, g.w), sy = window_start(y, g.ny, g.w), sx = window_start(x, g.nx, g.w); + + float acc[C]; +#pragma unroll + for (int c = 0; c < C; c++) + acc[c] = 0.f; + + const T* base = b + group * voxels * C; + const float* weights = s + t * w3; + int m = 0; + + for (int dz = 0; dz < g.w; dz++) + for (int dy = 0; dy < g.w; dy++) + for (int dx = 0; dx < g.w; dx++, m++) { + + const float wm = weights[m]; + const T* kv = base + (((int64_t)(sz + dz) * g.ny + (sy + dy)) * g.nx + (sx + dx)) * C; +#pragma unroll + for (int c = 0; c < C; c++) + acc[c] += wm * to_float(kv[c]); + } + +#pragma unroll + for (int c = 0; c < C; c++) + o[t * C + c] = from_float(acc[c]); + } +} + +/* The voxels along an axis whose window holds voxel j, and j's place in each window. */ +__device__ __forceinline__ int holders(int j, int n, int w, int* who, int* where) +{ + int count = 0; + const int lo = (j - w + 1 < 0) ? 0 : j - w + 1; + const int hi = (j + w - 1 > n - 1) ? n - 1 : j + w - 1; + + for (int i = lo; i <= hi; i++) { + + const int s = window_start(i, n, w); + + if (s <= j && j < s + w) { + + who[count] = i; + where[count] = j - s; + count++; + } + } + + return count; +} + +#define BARTORCH_NA_MAX_WINDOW 9 + +template +__global__ void scatter_sum(grid3 g, const float* __restrict__ s, const T* __restrict__ a, T* __restrict__ o) +{ + const int64_t voxels = (int64_t)g.nz * g.ny * g.nx; + const int w3 = g.w * g.w * g.w; + + for (int64_t t = blockIdx.x * (int64_t)blockDim.x + threadIdx.x; t < g.groups * voxels; + t += (int64_t)blockDim.x * gridDim.x) { + + const int64_t group = t / voxels; + const int64_t j = t % voxels; + const int x = (int)(j % g.nx), y = (int)((j / g.nx) % g.ny), z = (int)(j / ((int64_t)g.nx * g.ny)); + + int iz[2 * BARTORCH_NA_MAX_WINDOW], mz[2 * BARTORCH_NA_MAX_WINDOW]; + int iy[2 * BARTORCH_NA_MAX_WINDOW], my[2 * BARTORCH_NA_MAX_WINDOW]; + int ix[2 * BARTORCH_NA_MAX_WINDOW], mx[2 * BARTORCH_NA_MAX_WINDOW]; + const int cz = holders(z, g.nz, g.w, iz, mz); + const int cy = holders(y, g.ny, g.w, iy, my); + const int cx = holders(x, g.nx, g.w, ix, mx); + + float acc[C]; +#pragma unroll + for (int c = 0; c < C; c++) + acc[c] = 0.f; + + const float* sbase = s + group * voxels * w3; + const T* abase = a + group * voxels * C; + + for (int pz = 0; pz < cz; pz++) + for (int py = 0; py < cy; py++) + for (int px = 0; px < cx; px++) { + + const int64_t i = ((int64_t)iz[pz] * g.ny + iy[py]) * g.nx + ix[px]; + const int m = (mz[pz] * g.w + my[py]) * g.w + mx[px]; + const float wm = sbase[i * w3 + m]; + const T* av = abase + i * C; +#pragma unroll + for (int c = 0; c < C; c++) + acc[c] += wm * to_float(av[c]); + } + +#pragma unroll + for (int c = 0; c < C; c++) + o[t * C + c] = from_float(acc[c]); + } +} + +enum pass { GATHER_DOT, GATHER_SUM, SCATTER_SUM }; + +template +cudaError_t launch(enum pass p, grid3 g, const void* x, const void* y, void* out, cudaStream_t stream) +{ + const int64_t work = g.groups * (int64_t)g.nz * g.ny * g.nx; + const int threads = 128; + int64_t blocks = (work + threads - 1) / threads; + blocks = (blocks > 65535 * 16) ? 65535 * 16 : blocks; + + switch (p) { + case GATHER_DOT: + gather_dot<<<(unsigned)blocks, threads, 0, stream>>>(g, (const T*)x, (const T*)y, (float*)out); + break; + case GATHER_SUM: + gather_sum<<<(unsigned)blocks, threads, 0, stream>>>(g, (const float*)x, (const T*)y, (T*)out); + break; + case SCATTER_SUM: + scatter_sum<<<(unsigned)blocks, threads, 0, stream>>>(g, (const float*)x, (const T*)y, (T*)out); + break; + } + + return cudaGetLastError(); +} + +template +cudaError_t by_channels(enum pass p, grid3 g, int channels, const void* x, const void* y, void* out, cudaStream_t stream) +{ + switch (channels) { + case 8: return launch(p, g, x, y, out, stream); + case 16: return launch(p, g, x, y, out, stream); + case 32: return launch(p, g, x, y, out, stream); + case 64: return launch(p, g, x, y, out, stream); + default: return cudaErrorInvalidValue; + } +} + +int run(enum pass p, const void* x, const void* y, void* out, int64_t groups, int nz, int ny, int nx, + int channels, int window, int dtype, void* stream) +{ + if (window < 1 || window > BARTORCH_NA_MAX_WINDOW || nz < window || ny < window || nx < window || groups < 1) + return (int)cudaErrorInvalidValue; + + grid3 g = { nz, ny, nx, window, groups }; + cudaStream_t s = (cudaStream_t)stream; + cudaError_t e; + + switch (dtype) { + case BARTORCH_NA_FLOAT32: e = by_channels(p, g, channels, x, y, out, s); break; + case BARTORCH_NA_FLOAT16: e = by_channels<__half>(p, g, channels, x, y, out, s); break; + case BARTORCH_NA_BFLOAT16: e = by_channels<__nv_bfloat16>(p, g, channels, x, y, out, s); break; + default: e = cudaErrorInvalidValue; + } + + return (cudaSuccess == e) ? 0 : (int)e; +} + +} /* namespace */ + +extern "C" { + +int bartorch_na3d_gather_dot(const void* a, const void* b, float* s, int64_t groups, int nz, int ny, int nx, + int channels, int window, int dtype, void* stream) +{ + return run(GATHER_DOT, a, b, s, groups, nz, ny, nx, channels, window, dtype, stream); +} + +int bartorch_na3d_gather_sum(const float* s, const void* b, void* o, int64_t groups, int nz, int ny, int nx, + int channels, int window, int dtype, void* stream) +{ + return run(GATHER_SUM, s, b, o, groups, nz, ny, nx, channels, window, dtype, stream); +} + +int bartorch_na3d_scatter_sum(const float* s, const void* a, void* o, int64_t groups, int nz, int ny, int nx, + int channels, int window, int dtype, void* stream) +{ + return run(SCATTER_SUM, s, a, o, groups, nz, ny, nx, channels, window, dtype, stream); +} + +int bartorch_na3d_built(void) +{ + return 1; +} + +} diff --git a/tests/test_learning_attention.py b/tests/test_learning_attention.py new file mode 100644 index 000000000..31e90fcfe --- /dev/null +++ b/tests/test_learning_attention.py @@ -0,0 +1,119 @@ +"""Neighbourhood attention: the plain-torch path against attention written out voxel by voxel, +and the library's CUDA passes, with their gradients, against the plain-torch path. +""" + +import itertools + +import pytest +import torch + +from bartorch import learning +from bartorch.learning._neighborhood import _attend_reference, kernels_available + + +def _brute(q, k, v, window, bias): + """Softmax over each voxel's shifted window, one voxel at a time.""" + n, heads, nz, ny, nx, c = q.shape + out = torch.zeros_like(q) + + def start(i, size): + return min(max(i - window // 2, 0), size - window) + + for b, h, z, y, x in itertools.product(range(n), range(heads), range(nz), range(ny), range(nx)): + sz, sy, sx = start(z, nz), start(y, ny), start(x, nx) + keys = k[b, h, sz : sz + window, sy : sy + window, sx : sx + window].reshape(-1, c) + values = v[b, h, sz : sz + window, sy : sy + window, sx : sx + window].reshape(-1, c) + logits = keys @ q[b, h, z, y, x] / c**0.5 + if bias is not None: + logits = logits + bias[h] + out[b, h, z, y, x] = logits.softmax(0) @ values + return out + + +@pytest.mark.parametrize("window", [1, 3, 5]) +def test_reference_is_attention_over_the_shifted_window(window): + torch.manual_seed(0) + shape = (2, 2, 5, 6, 7, 4) + q, k, v = (torch.randn(shape, dtype=torch.float64) for _ in range(3)) + bias = torch.randn(2, window**3, dtype=torch.float64) + made = learning.neighborhood_attention3d(q, k, v, window, bias) + torch.testing.assert_close(made, _brute(q, k, v, window, bias)) + + +def test_full_window_is_full_attention(): + torch.manual_seed(1) + q, k, v = (torch.randn(1, 1, 3, 3, 3, 8, dtype=torch.float64) for _ in range(3)) + made = learning.neighborhood_attention3d(q, k, v, 3) + flat = [t.reshape(27, 8) for t in (q, k, v)] + full = torch.softmax(flat[0] @ flat[1].T / 8**0.5, -1) @ flat[2] + torch.testing.assert_close(made.reshape(27, 8), full) + + +def test_rejects_what_it_cannot_do(): + q = torch.zeros(1, 1, 4, 4, 4, 8) + with pytest.raises(ValueError, match="odd"): + learning.neighborhood_attention3d(q, q, q, 4) + with pytest.raises(ValueError, match="at least the window"): + learning.neighborhood_attention3d(q, q, q, 5) + with pytest.raises(ValueError, match="bias"): + learning.neighborhood_attention3d(q, q, q, 3, torch.zeros(2, 27)) + + +def test_module_keeps_the_shape_and_trains(): + torch.manual_seed(2) + block = learning.NeighborhoodAttention3d(16, window=3) + x = torch.randn(2, 16, 5, 4, 6, requires_grad=True) + out = block(x) + assert out.shape == x.shape + out.square().mean().backward() + assert x.grad is not None and block.bias.grad is not None and block.qkv.weight.grad is not None + + +cuda = pytest.mark.skipif( + not kernels_available(), reason="needs a card and the library's CUDA kernels" +) + + +@cuda +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("channels,window", [(8, 3), (16, 5), (32, 3), (64, 5)]) +def test_cuda_passes_match_the_reference_with_gradients(dtype, channels, window): + torch.manual_seed(3) + shape = (2, 2, 7, 9, 6, channels) + tensors = [torch.randn(shape, device="cuda", dtype=torch.float64) for _ in range(3)] + bias64 = 0.1 * torch.randn(2, window**3, device="cuda", dtype=torch.float64) + grad = torch.randn(shape, device="cuda", dtype=torch.float64) + + q, k, v = (t.to(dtype).requires_grad_() for t in tensors) + bias = bias64.to(dtype).requires_grad_() + made = learning.neighborhood_attention3d(q, k, v, window, bias) + made.backward(grad.to(dtype)) + + rq, rk, rv = (t.clone().requires_grad_() for t in tensors) + rbias = bias64.clone().requires_grad_() + flat = [t.reshape(4, -1, channels) for t in (rq, rk, rv)] + expected = _attend_reference(*flat, shape[2:5], window, rbias, 2).reshape(shape) + expected.backward(grad) + + tolerance = {torch.float32: 1e-4, torch.float16: 2e-2, torch.bfloat16: 1e-1}[dtype] + pairs = [ + (made, expected), + (q.grad, rq.grad), + (k.grad, rk.grad), + (v.grad, rv.grad), + (bias.grad, rbias.grad), + ] + for got, want in pairs: + torch.testing.assert_close(got.double(), want, atol=tolerance, rtol=tolerance) + + +@cuda +def test_cuda_uses_memory_linear_in_the_volume(): + q = torch.randn(1, 2, 64, 64, 64, 16, device="cuda", dtype=torch.float16) + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + before = torch.cuda.memory_allocated() + learning.neighborhood_attention3d(q, q, q, 5) + torch.cuda.synchronize() + # float32 weights (2 x 64^3 x 125 x 4 bytes = 262 MB) and the output; no gathered keys. + assert torch.cuda.max_memory_allocated() - before < 600e6 From c0daa890f078bbcc4d6f484abf41463dbc2fd614 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 2 Nov 2026 11:25:57 +0000 Subject: [PATCH 2/8] check_attention: --standalone reads the module from the checkout, without importing bartorch Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MFB6QPF5rgL8sKizNYzB8D --- scripts/check_attention.py | 27 ++++++++++++++++++++++---- src/bartorch/learning/_neighborhood.py | 2 +- 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/scripts/check_attention.py b/scripts/check_attention.py index 1981032d5..ef6309fe0 100644 --- a/scripts/check_attention.py +++ b/scripts/check_attention.py @@ -4,8 +4,9 @@ python scripts/check_attention.py --standalone # attention.cu alone, built by nvcc With ``--standalone`` the one translation unit is compiled into a library of -its own and used in place of ``libbartorch``'s passes, so a change to the -kernels can be checked on a card without building BART: it needs ``nvcc``. +its own and used in place of ``libbartorch``'s passes, and the Python side is +read from this checkout without importing bartorch, so a change to the kernels +can be checked on a card without building BART: it needs ``nvcc`` and torch. The passes and their gradients are compared with the plain-torch path in float64, in each element type, and then timed at the sizes a denoiser meets: @@ -24,10 +25,26 @@ import torch -from bartorch.learning import _neighborhood as na - ROOT = Path(__file__).resolve().parent.parent +#: The module under test, set by main(): the package's, or with --standalone the file itself, +#: loaded without importing bartorch, whose library may predate these passes. +na = None + + +def load_module(standalone_file: bool): + if not standalone_file: + from bartorch.learning import _neighborhood + + return _neighborhood + import importlib.util + + path = ROOT / "src" / "bartorch" / "learning" / "_neighborhood.py" + spec = importlib.util.spec_from_file_location("na3d_under_test", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + def standalone(arch: str | None) -> ctypes.CDLL: """``attention.cu`` compiled on its own, with the ctypes signatures of the library's ABI.""" @@ -137,6 +154,8 @@ def main(): "--arch", help="compute capability without the dot, e.g. 89; the card's by default" ) args = ap.parse_args() + global na + na = load_module(args.standalone) if args.standalone: lib = standalone(args.arch) na._library = lambda: lib diff --git a/src/bartorch/learning/_neighborhood.py b/src/bartorch/learning/_neighborhood.py index bebf6f62e..75cbfdb18 100644 --- a/src/bartorch/learning/_neighborhood.py +++ b/src/bartorch/learning/_neighborhood.py @@ -52,7 +52,7 @@ def kernels_available() -> bool: """Whether this library carries the CUDA kernels and a card is present.""" try: return bool(_library().bartorch_na3d_built()) and torch.cuda.is_available() - except (OSError, AttributeError): + except (ImportError, OSError, AttributeError): return False From 620c542073f1d1103e2058a131aa5c10ae9f8bd4 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 2 Nov 2026 11:29:00 +0000 Subject: [PATCH 3/8] learning: neighbourhood attention reaches its pointers through _marshal; documented; exact reference The plain-torch path computes in float64 for float64 inputs, so the tests hold it to the voxel-by-voxel attention to round-off. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MFB6QPF5rgL8sKizNYzB8D --- docs/api/learning.md | 2 ++ scripts/check_attention.py | 3 +++ src/bartorch/_marshal.py | 6 ++++++ src/bartorch/learning/_neighborhood.py | 27 +++++++++++++------------- 4 files changed, 25 insertions(+), 13 deletions(-) diff --git a/docs/api/learning.md b/docs/api/learning.md index c63e02a14..c85ce2e02 100644 --- a/docs/api/learning.md +++ b/docs/api/learning.md @@ -22,6 +22,8 @@ backward strategy uses. | {obj}`~bartorch.learning.UNet` | Residual UNet on 2D or 3D real channels, optionally over a frame axis, conditioned on the iteration index, the noise level and a class | | {obj}`~bartorch.learning.ComplexNet` | A real network applied to complex images, their leading axes and real and imaginary parts as its channels, normalized by the peak or whitened | | {obj}`~bartorch.learning.Patchwise` | A network applied patch by patch on a device, with a randomly shifted or rolled grid, empty patches skipped and mixed precision, to an image held elsewhere | +| {obj}`~bartorch.learning.NeighborhoodAttention3d` | Multi-head attention of each voxel of a 3D feature map over a window around it, with a bias per relative position; linear in the volume | +| {obj}`~bartorch.learning.neighborhood_attention3d` | The operation itself on queries, keys and values: through this library's CUDA kernels on a card, in plain torch elsewhere | | Precision on the device | `Patchwise(dtype="auto")` | | --- | --- | diff --git a/scripts/check_attention.py b/scripts/check_attention.py index ef6309fe0..c09c3d7d4 100644 --- a/scripts/check_attention.py +++ b/scripts/check_attention.py @@ -159,6 +159,9 @@ def main(): if args.standalone: lib = standalone(args.arch) na._library = lambda: lib + na._float_pointer = lambda address: ctypes.cast( + ctypes.c_void_p(address), ctypes.POINTER(ctypes.c_float) + ) if not na.kernels_available(): raise SystemExit("no card, or no CUDA passes in this library") print(torch.cuda.get_device_name(), flush=True) diff --git a/src/bartorch/_marshal.py b/src/bartorch/_marshal.py index a6c17c034..3d8e57efa 100644 --- a/src/bartorch/_marshal.py +++ b/src/bartorch/_marshal.py @@ -19,6 +19,7 @@ "dim_vector", "dims", "float_buffer", + "float_pointer", "floats", "ints", "int64s", @@ -131,6 +132,11 @@ def float_buffer(ptr: int, count: int) -> ctypes.Array: return (ctypes.c_float * count).from_address(ptr) +def float_pointer(ptr: int): + """A ``float*`` at the address ``ptr``, for an entry point that takes one.""" + return ctypes.cast(ctypes.c_void_p(ptr), ctypes.POINTER(ctypes.c_float)) + + def null_apply() -> ctypes.Array: """The NULL an apply callback the caller did not give is passed as. diff --git a/src/bartorch/learning/_neighborhood.py b/src/bartorch/learning/_neighborhood.py index 75cbfdb18..e128065c5 100644 --- a/src/bartorch/learning/_neighborhood.py +++ b/src/bartorch/learning/_neighborhood.py @@ -2,8 +2,6 @@ from __future__ import annotations -import ctypes - import torch from torch import nn @@ -34,12 +32,13 @@ def _attend_reference(q, k, v, extent, window, bias, heads): """Neighbourhood attention in plain torch on ``(groups, voxels, channels)``.""" index = _neighbours(extent, window, q.device) scale = q.shape[-1] ** -0.5 + work = torch.promote_types(q.dtype, torch.float32) kn = k[:, index] # (groups, voxels, window^3, channels) - logits = torch.einsum("gvc,gvmc->gvm", q.float() * scale, kn.float()) + logits = torch.einsum("gvc,gvmc->gvm", q.to(work) * scale, kn.to(work)) if bias is not None: - logits = logits + bias.float().repeat(q.shape[0] // heads, 1)[:, None, :] + logits = logits + bias.to(work).repeat(q.shape[0] // heads, 1)[:, None, :] attention = logits.softmax(-1) - return torch.einsum("gvm,gvmc->gvc", attention, v[:, index].float()).to(q.dtype) + return torch.einsum("gvm,gvmc->gvc", attention, v[:, index].to(work)).to(q.dtype) def _library(): @@ -48,6 +47,12 @@ def _library(): return library() +def _float_pointer(address: int): + from bartorch._marshal import float_pointer + + return float_pointer(address) + + def kernels_available() -> bool: """Whether this library carries the CUDA kernels and a card is present.""" try: @@ -56,23 +61,19 @@ def kernels_available() -> bool: return False -def _floats(t: torch.Tensor): - return ctypes.cast(ctypes.c_void_p(t.data_ptr()), ctypes.POINTER(ctypes.c_float)) - - def _call(name: str, x, y, out, extent, window: int): """One CUDA pass; ``x`` and ``out`` are float32 weights or elements, as the pass takes them.""" groups, _, channels = y.shape nz, ny, nx = extent if name == "gather_dot": # elements, elements -> weights - first, third = ctypes.c_void_p(x.data_ptr()), _floats(out) + first, third = x.data_ptr(), _float_pointer(out.data_ptr()) else: # weights, elements -> elements - first, third = _floats(x), ctypes.c_void_p(out.data_ptr()) + first, third = _float_pointer(x.data_ptr()), out.data_ptr() with torch.cuda.device(y.device): stream = torch.cuda.current_stream(y.device).cuda_stream error = getattr(_library(), f"bartorch_na3d_{name}")( first, - ctypes.c_void_p(y.data_ptr()), + y.data_ptr(), third, groups, nz, @@ -81,7 +82,7 @@ def _call(name: str, x, y, out, extent, window: int): channels, window, _DTYPES[y.dtype], - ctypes.c_void_p(stream), + stream, ) if error: raise RuntimeError(f"bartorch_na3d_{name} failed with CUDA error {error}") From 7594b3ff530279daa922415b6c7357edb52ffd0d Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 3 Nov 2026 12:11:17 +0000 Subject: [PATCH 4/8] learning: group_norm_silu, a diffusion UNet's group norm, scale and shift and SiLU in one CUDA pass Two short passes take a group's float32 statistics (Welford, a group spread over several blocks and merged), and one more reads the activations once more and writes them once, in their own element type, with the norm's affine, the per-item scale and shift and SiLU folded into one multiply-add per row. Eager torch reads and writes them four or five times, with casts around a float32 norm, which in MRI2Qmap's teacher costs more than its convolutions. Inference only: under autograd, without a card or without CUDA in the library, plain torch computes the same. scripts/check_groupnorm.py --standalone builds groupnorm.cu alone with nvcc to check and time it. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MFB6QPoLKTK5gZEQeb1W8D --- CMakeLists.txt | 3 +- docs/api/learning.md | 1 + scripts/check_groupnorm.py | 181 +++++++++++++++++++++ src/bartorch/_abi.py | 30 ++++ src/bartorch/learning/__init__.py | 6 +- src/bartorch/learning/_groupnorm.py | 133 ++++++++++++++++ src/csrc/include/bartorch.h | 20 +++ src/csrc/ops/groupnorm.c | 28 ++++ src/csrc/ops/groupnorm.cu | 238 ++++++++++++++++++++++++++++ tests/test_learning_groupnorm.py | 66 ++++++++ 10 files changed, 704 insertions(+), 2 deletions(-) create mode 100644 scripts/check_groupnorm.py create mode 100644 src/bartorch/learning/_groupnorm.py create mode 100644 src/csrc/ops/groupnorm.c create mode 100644 src/csrc/ops/groupnorm.cu create mode 100644 tests/test_learning_groupnorm.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 8ea87ad6f..9239fffb9 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -327,6 +327,7 @@ set(BARTORCH_SOURCES src/csrc/ops/grid.c src/csrc/ops/iter.c src/csrc/ops/attention.c + src/csrc/ops/groupnorm.c src/csrc/substitute/backend.c src/csrc/substitute/ref_blas.c @@ -342,7 +343,7 @@ if(OpenMP_C_FOUND AND (WIN32 OR APPLE)) endif() if(BARTORCH_CUDA) - list(APPEND BARTORCH_SOURCES src/csrc/ops/kernels.cu src/csrc/ops/fft_callbacks.cu src/csrc/ops/attention.cu) + list(APPEND BARTORCH_SOURCES src/csrc/ops/kernels.cu src/csrc/ops/fft_callbacks.cu src/csrc/ops/attention.cu src/csrc/ops/groupnorm.cu) # The passes around a volume's transform, which cuFFT links into the # transform when a plan is made (src/csrc/ops/fft_callbacks.cu). It takes them as diff --git a/docs/api/learning.md b/docs/api/learning.md index c85ce2e02..c144c7794 100644 --- a/docs/api/learning.md +++ b/docs/api/learning.md @@ -24,6 +24,7 @@ backward strategy uses. | {obj}`~bartorch.learning.Patchwise` | A network applied patch by patch on a device, with a randomly shifted or rolled grid, empty patches skipped and mixed precision, to an image held elsewhere | | {obj}`~bartorch.learning.NeighborhoodAttention3d` | Multi-head attention of each voxel of a 3D feature map over a window around it, with a bias per relative position; linear in the volume | | {obj}`~bartorch.learning.neighborhood_attention3d` | The operation itself on queries, keys and values: through this library's CUDA kernels on a card, in plain torch elsewhere | +| {obj}`~bartorch.learning.group_norm_silu` | A diffusion UNet's group norm, noise-level scale and shift and SiLU in one pass over memory on a card (float32 statistics), plain torch elsewhere or under autograd | | Precision on the device | `Patchwise(dtype="auto")` | | --- | --- | diff --git a/scripts/check_groupnorm.py b/scripts/check_groupnorm.py new file mode 100644 index 000000000..875a612d0 --- /dev/null +++ b/scripts/check_groupnorm.py @@ -0,0 +1,181 @@ +"""Check and time the fused group norm (group norm, scale and shift, SiLU) on a card. + + python scripts/check_groupnorm.py # the library's own pass + python scripts/check_groupnorm.py --standalone # groupnorm.cu alone, built by nvcc + +With ``--standalone`` the one translation unit is compiled into a library of +its own and used in place of ``libbartorch``'s pass, and the Python side is +read from this checkout without importing bartorch, so the kernel can be +checked on a card without building BART: it needs ``nvcc`` and torch. + +The pass is compared with float64 arithmetic in each element type, then timed +at the activations of a diffusion UNet's levels against the way such a network +runs it in eager torch (a float32 group norm, a cast back, the scale and shift, +SiLU) and against a float16 group norm alone, with the memory floor (one read +for the statistics, one read and one write for the rest) for reference. +""" + +from __future__ import annotations + +import argparse +import ctypes +import json +import subprocess +import tempfile +import time +from pathlib import Path + +import torch +import torch.nn.functional as F + +ROOT = Path(__file__).resolve().parent.parent + +gn = None + + +def load_module(standalone_file: bool): + if not standalone_file: + from bartorch.learning import _groupnorm + + return _groupnorm + import importlib.util + + path = ROOT / "src" / "bartorch" / "learning" / "_groupnorm.py" + spec = importlib.util.spec_from_file_location("gn_under_test", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def standalone(arch: str | None) -> ctypes.CDLL: + """``groupnorm.cu`` compiled on its own, with the ctypes signatures of the library's ABI.""" + out = Path(tempfile.mkdtemp()) / "libgn.so" + if arch is None: + major, minor = torch.cuda.get_device_capability() + arch = f"{major}{minor}" + command = [ + "nvcc", "-O3", "-std=c++17", "-shared", "-Xcompiler", "-fPIC", + f"-gencode=arch=compute_{arch},code=sm_{arch}", + "-I", str(ROOT / "src" / "csrc"), + "-o", str(out), str(ROOT / "src" / "csrc" / "ops" / "groupnorm.cu"), + ] # fmt: skip + subprocess.run(command, check=True) + lib = ctypes.CDLL(str(out)) + lib.bartorch_gn_built.restype = ctypes.c_int + lib.bartorch_gn_workspace.restype = ctypes.c_int64 + lib.bartorch_gn_workspace.argtypes = [ + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_int, + ] + v = ctypes.c_void_p + lib.bartorch_gn_silu.restype = ctypes.c_int + lib.bartorch_gn_silu.argtypes = [ + v, v, ctypes.c_int64, ctypes.c_int64, ctypes.c_int64, ctypes.c_int, + v, v, v, v, ctypes.c_float, ctypes.c_int, ctypes.c_int, v, v, + ] # fmt: skip + return lib + + +def exact(x, groups, weight, bias, scale, shift): + n, c = x.shape[:2] + g = x.double().reshape(n, groups, -1) + g = (g - g.mean(-1, keepdim=True)) / (g.var(-1, unbiased=False, keepdim=True) + 1e-5).sqrt() + shape = (1, c, *([1] * (x.ndim - 2))) + h = g.reshape(x.shape) * weight.double().reshape(shape) + bias.double().reshape(shape) + shape = (n, c, *([1] * (x.ndim - 2))) + return F.silu(h * (1 + scale.double().reshape(shape)) + shift.double().reshape(shape)) + + +def check(dtype, shape, groups=32): + torch.manual_seed(0) + n, c = shape[:2] + x = (2 + 3 * torch.randn(shape, device="cuda")).to(dtype) + weight, bias = torch.randn(c, device="cuda"), torch.randn(c, device="cuda") + scale, shift = 0.3 * torch.randn(n, c, device="cuda"), torch.randn(n, c, device="cuda") + with torch.no_grad(): + made = gn.group_norm_silu(x, groups, weight, bias, scale, shift) + want = exact(x, groups, weight, bias, scale, shift) + return float((made.double() - want).abs().max() / want.abs().max()) + + +def timed(fn, repeat=20): + fn() + torch.cuda.synchronize() + start = time.perf_counter() + for _ in range(repeat): + fn() + torch.cuda.synchronize() + return (time.perf_counter() - start) / repeat + + +def bench(channels, extent, dtype=torch.float16, groups=32): + x = torch.randn(1, channels, *(extent,) * 3, device="cuda", dtype=dtype) + weight, bias = torch.ones(channels, device="cuda"), torch.zeros(channels, device="cuda") + scale, shift = torch.zeros(1, channels, device="cuda"), torch.zeros(1, channels, device="cuda") + s5, t5 = scale.reshape(1, -1, 1, 1, 1).to(dtype), shift.reshape(1, -1, 1, 1, 1).to(dtype) + + def eager(): # guided diffusion's GroupNorm32, then FiLM, then SiLU + h = F.group_norm(x.float(), groups, weight, bias).type(x.dtype) + return F.silu(h * (1 + s5) + t5) + + def half_norm(): + return F.silu(F.group_norm(x, groups, weight.to(dtype), bias.to(dtype)) * (1 + s5) + t5) + + def fused(): + return gn.group_norm_silu(x, groups, weight, bias, scale, shift) + + with torch.no_grad(): + row = {"channels": channels, "extent": extent} + row["eager fp32 norm ms"] = 1e3 * timed(eager) + row["eager fp16 norm ms"] = 1e3 * timed(half_norm) + row["fused ms"] = 1e3 * timed(fused) + row["bytes moved MB (floor)"] = 3 * x.numel() * x.element_size() / 1e6 + return {k: round(v, 3) if isinstance(v, float) else v for k, v in row.items()} + + +def main(): + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument("--standalone", action="store_true", help="build groupnorm.cu alone with nvcc") + ap.add_argument( + "--arch", help="compute capability without the dot, e.g. 86; the card's by default" + ) + args = ap.parse_args() + global gn + gn = load_module(args.standalone) + if args.standalone: + lib = standalone(args.arch) + gn._library = lambda: lib + if not gn.kernels_available(): + raise SystemExit("no card, or no CUDA pass in this library") + print(torch.cuda.get_device_name(), flush=True) + + worst = {} + for dtype in (torch.float32, torch.float16, torch.bfloat16): + for shape in ( + (2, 64, 9, 10, 11), + (1, 128, 32, 32, 32), + (1, 64, 64, 64, 64), + (3, 512, 8, 8, 8), + ): + key = str(dtype).removeprefix("torch.") + worst[key] = max(worst.get(key, 0.0), check(dtype, shape)) + print("largest relative error against float64:", json.dumps(worst), flush=True) + + for channels, extent in ( + (64, 64), + (128, 64), + (128, 32), + (256, 32), + (256, 16), + (512, 16), + (512, 8), + ): + print(json.dumps(bench(channels, extent)), flush=True) + + +if __name__ == "__main__": + main() diff --git a/src/bartorch/_abi.py b/src/bartorch/_abi.py index c04f3f03d..18545c77e 100644 --- a/src/bartorch/_abi.py +++ b/src/bartorch/_abi.py @@ -353,6 +353,9 @@ class Encoding(ctypes.Structure): "bartorch_na3d_gather_dot", "bartorch_na3d_gather_sum", "bartorch_na3d_scatter_sum", + "bartorch_gn_built", + "bartorch_gn_workspace", + "bartorch_gn_silu", "bartorch_on_device", ) @@ -1306,6 +1309,33 @@ def bind(lib: ctypes.CDLL) -> ctypes.CDLL: ctypes.c_int, ctypes.c_void_p, ] + lib.bartorch_gn_built.restype = ctypes.c_int + lib.bartorch_gn_built.argtypes = [] + lib.bartorch_gn_workspace.restype = ctypes.c_int64 + lib.bartorch_gn_workspace.argtypes = [ + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_int, + ] + lib.bartorch_gn_silu.restype = ctypes.c_int + lib.bartorch_gn_silu.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_float, + ctypes.c_int, + ctypes.c_int, + ctypes.c_void_p, + ctypes.c_void_p, + ] lib.bartorch_on_device.restype = ctypes.c_int lib.bartorch_on_device.argtypes = [ctypes.c_void_p] return lib diff --git a/src/bartorch/learning/__init__.py b/src/bartorch/learning/__init__.py index 0603ab1d4..06a4ecca2 100644 --- a/src/bartorch/learning/__init__.py +++ b/src/bartorch/learning/__init__.py @@ -10,7 +10,9 @@ samples for self-supervised training, and :func:`moments` and :func:`calibrate` give a voxel-wise uncertainty. :class:`NeighborhoodAttention3d` lets each voxel of a 3D feature map attend to -a window around it, through CUDA kernels of this library on a card. +a window around it, through CUDA kernels of this library on a card, and +:func:`group_norm_silu` fuses a diffusion UNet's group norm, noise-level scale +and shift and SiLU into one pass there. :class:`Reconstruction` trains a network in one of the three stages of an unrolled network, on ``lightning``, and :class:`RandomGain` is a ``torchio`` @@ -21,6 +23,7 @@ from __future__ import annotations from bartorch.learning._channels import ComplexNet, as_complex, as_real +from bartorch.learning._groupnorm import group_norm_silu from bartorch.learning._neighborhood import NeighborhoodAttention3d, neighborhood_attention3d from bartorch.learning._nets import UNet from bartorch.learning._patches import Patchwise @@ -42,6 +45,7 @@ "as_complex", "as_real", "calibrate", + "group_norm_silu", "moments", "neighborhood_attention3d", "split", diff --git a/src/bartorch/learning/_groupnorm.py b/src/bartorch/learning/_groupnorm.py new file mode 100644 index 000000000..cfeede75e --- /dev/null +++ b/src/bartorch/learning/_groupnorm.py @@ -0,0 +1,133 @@ +"""Group normalisation fused with a per-item scale and shift and SiLU, in one pass on a card.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +__all__ = ["group_norm_silu"] + +_DTYPES = {torch.float32: 0, torch.float16: 1, torch.bfloat16: 2} + + +def _library(): + from bartorch._lib import library + + return library() + + +def kernels_available() -> bool: + """Whether this library carries the CUDA pass and a card is present.""" + try: + return bool(_library().bartorch_gn_built()) and torch.cuda.is_available() + except (ImportError, OSError, AttributeError): + return False + + +def _reference(x, groups, weight, bias, scale, shift, eps, silu): + """The same arithmetic in plain torch, statistics in float32.""" + h = F.group_norm(x.float(), groups, weight, bias, eps) + shape = (*x.shape[:2], *([1] * (x.ndim - 2))) + if scale is not None: + h = h * (1 + scale.float().reshape(shape)) + if shift is not None: + h = h + shift.float().reshape(shape) + if silu: + h = F.silu(h) + return h.to(x.dtype) + + +def _pointer(t): + return None if t is None else t.data_ptr() + + +def _as_float(t, device, shape): + if t is None: + return None + return t.detach().to(device=device, dtype=torch.float32).reshape(shape).contiguous() + + +def group_norm_silu( + x: torch.Tensor, + groups: int, + weight: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + scale: torch.Tensor | None = None, + shift: torch.Tensor | None = None, + *, + eps: float = 1e-5, + silu: bool = True, +) -> torch.Tensor: + """``silu((group_norm(x) * weight + bias) * (1 + scale) + shift)``: a diffusion UNet's block. + + On a card with this library's CUDA pass, and when no gradient is asked + for, the statistics take one read of ``x`` and the rest one more read and + one write, in ``x``'s own element type (float32, float16 or bfloat16) with + float32 statistics -- where eager torch reads and writes the activations + four or five times, with casts around the norm. Anywhere else plain torch + computes the same, also with float32 statistics. + + Parameters + ---------- + x : torch.Tensor + ``(n, channels, *spatial)``. + groups : int + Groups the channels are normalised in; it divides ``channels``. + weight, bias : torch.Tensor, default=None + ``(channels,)`` affine of the norm. + scale, shift : torch.Tensor, default=None + ``(n, channels)`` per-item scale (applied as ``1 + scale``) and shift, + e.g. from a noise-level embedding. + eps : float, default=1e-5 + Added to the variance. + silu : bool, default=True + Apply SiLU last. + + Returns + ------- + torch.Tensor + Of ``x``'s shape and dtype. + """ + n, channels = x.shape[:2] + if channels % groups: + raise ValueError(f"{channels} channels do not split into {groups} groups") + grad = torch.is_grad_enabled() and any( + t is not None and t.requires_grad for t in (x, weight, bias, scale, shift) + ) + fast = x.is_cuda and x.dtype in _DTYPES and not grad and x.numel() > 0 and kernels_available() + if not fast: + return _reference(x, groups, weight, bias, scale, shift, eps, silu) + + lib = _library() + voxels = x[0, 0].numel() + x = x.contiguous() + device = x.device + w = _as_float(weight, device, (channels,)) + b = _as_float(bias, device, (channels,)) + s = _as_float(scale, device, (n, channels)) + t = _as_float(shift, device, (n, channels)) + size = lib.bartorch_gn_workspace(n, channels, voxels, groups) + workspace = torch.empty(size, device=device, dtype=torch.float32) + out = torch.empty_like(x) + with torch.cuda.device(device): + stream = torch.cuda.current_stream(device).cuda_stream + error = lib.bartorch_gn_silu( + x.data_ptr(), + out.data_ptr(), + n, + channels, + voxels, + groups, + _pointer(w), + _pointer(b), + _pointer(s), + _pointer(t), + float(eps), + int(bool(silu)), + _DTYPES[x.dtype], + workspace.data_ptr(), + stream, + ) + if error: + raise RuntimeError(f"bartorch_gn_silu failed with CUDA error {error}") + return out diff --git a/src/csrc/include/bartorch.h b/src/csrc/include/bartorch.h index dee835568..67a743c3e 100644 --- a/src/csrc/include/bartorch.h +++ b/src/csrc/include/bartorch.h @@ -966,6 +966,26 @@ BARTORCH_API int bartorch_na3d_gather_sum(const float* s, const void* b, void* o BARTORCH_API int bartorch_na3d_scatter_sum(const float* s, const void* a, void* o, int64_t groups, int nz, int ny, int nx, int channels, int window, int dtype, void* stream); +/* + * Group normalisation fused with a per-item scale and shift and SiLU, for + * bartorch.learning: + * + * y[n, c, v] = act(((x[n, c, v] - mean[n, g]) * rstd[n, g] * w[c] + b[c]) * (1 + s[n, c]) + t[n, c]) + * + * x and y are (items, channels, voxels), voxels fastest, of a BARTORCH_NA_* + * element type, and y may be x; weight and bias are (channels) float32, scale + * and shift (items, channels) float32, each NULL when absent; act is SiLU + * when silu is nonzero. Statistics are float32. workspace holds + * bartorch_gn_workspace(...) floats on the card (-1 for a bad shape). Runs on + * the cudaStream_t passed as a void* and returns 0 or the CUDA error it met; + * without CUDA, built() is 0 and gn_silu returns -1. + */ +BARTORCH_API int bartorch_gn_built(void); +BARTORCH_API int64_t bartorch_gn_workspace(int64_t items, int64_t channels, int64_t voxels, int groups); +BARTORCH_API int bartorch_gn_silu(const void* x, void* y, int64_t items, int64_t channels, int64_t voxels, + int groups, const void* weight, const void* bias, const void* scale, const void* shift, float eps, + int silu, int dtype, void* workspace, void* stream); + /* Whether a pointer is device memory; always false without CUDA. */ BARTORCH_API int bartorch_on_device(const void* ptr); diff --git a/src/csrc/ops/groupnorm.c b/src/csrc/ops/groupnorm.c new file mode 100644 index 000000000..295507310 --- /dev/null +++ b/src/csrc/ops/groupnorm.c @@ -0,0 +1,28 @@ +/* + * The fused group-norm entry points of a build without CUDA: the passes are + * in groupnorm.cu, and here they report that they are absent. + */ +#include + +#include "include/bartorch.h" + +#ifndef USE_CUDA + +int bartorch_gn_built(void) { return 0; } + +int64_t bartorch_gn_workspace(int64_t items, int64_t channels, int64_t voxels, int groups) +{ + (void)items; (void)channels; (void)voxels; (void)groups; + return -1; +} + +int bartorch_gn_silu(const void* x, void* y, int64_t items, int64_t channels, int64_t voxels, int groups, + const void* weight, const void* bias, const void* scale, const void* shift, float eps, int silu, + int dtype, void* workspace, void* stream) +{ + (void)x; (void)y; (void)items; (void)channels; (void)voxels; (void)groups; (void)weight; (void)bias; + (void)scale; (void)shift; (void)eps; (void)silu; (void)dtype; (void)workspace; (void)stream; + return -1; +} + +#endif diff --git a/src/csrc/ops/groupnorm.cu b/src/csrc/ops/groupnorm.cu new file mode 100644 index 000000000..955b5997e --- /dev/null +++ b/src/csrc/ops/groupnorm.cu @@ -0,0 +1,238 @@ +/* + * Group normalisation, a per-item scale and shift, and SiLU, in one pass over memory. + * + * The block every diffusion-style UNet repeats -- GroupNorm, then the noise + * level's scale and shift (FiLM), then SiLU -- as eager PyTorch runs it reads + * and writes the activations four or five times, and in float16 networks it + * adds two casts around the norm. Here the statistics take one read, and the + * normalisation, affine, scale, shift and SiLU one more read and one write: + * + * y[n, c, v] = act( ((x[n, c, v] - mean[n, g]) * rstd[n, g] * w[c] + b[c]) * (1 + s[n, c]) + t[n, c] ) + * + * with the norm's affine w, b and the FiLM s, t each optional (NULL), and act + * SiLU when asked for, the identity otherwise. x and y are + * (items, channels, voxels), voxels fastest -- a contiguous NC(D)HW tensor -- + * in float, half or bfloat16; w, b are (channels) float32, s, t (items, + * channels) float32. Statistics are in float32 by Welford's method; a group + * is spread over several blocks, whose partial sums a second pass merges. + * y may be x. + * + * The caller hands over a float32 workspace of bartorch_gn_workspace(...) + * floats. Each pass runs on the caller's stream; the entry point returns 0 + * or the CUDA error it met. + */ +#include +#include +#include + +#include + +#include "include/bartorch.h" + +namespace { + +__device__ __forceinline__ float to_float(float v) { return v; } +__device__ __forceinline__ float to_float(__half v) { return __half2float(v); } +__device__ __forceinline__ float to_float(__nv_bfloat16 v) { return __bfloat162float(v); } + +template __device__ __forceinline__ T from_float(float v); +template <> __device__ __forceinline__ float from_float(float v) { return v; } +template <> __device__ __forceinline__ __half from_float<__half>(float v) { return __float2half(v); } +template <> __device__ __forceinline__ __nv_bfloat16 from_float<__nv_bfloat16>(float v) { return __float2bfloat16(v); } + +const int THREADS = 256; +/* Elements of a group each statistics block takes, at least. */ +const int64_t CHUNK = 16384; +const int MAX_PARTS = 64; + +int64_t parts_of(int64_t group_size) +{ + int64_t p = (group_size + CHUNK - 1) / CHUNK; + return (p < 1) ? 1 : (p > MAX_PARTS) ? MAX_PARTS : p; +} + +struct welford { + float count, mean, m2; +}; + +__device__ __forceinline__ welford merge(welford a, welford b) +{ + const float count = a.count + b.count; + if (count == 0.f) + return a; + const float delta = b.mean - a.mean; + const float fb = b.count / count; + welford r; + r.count = count; + r.mean = a.mean + delta * fb; + r.m2 = a.m2 + b.m2 + delta * delta * a.count * fb; + return r; +} + +__device__ __forceinline__ welford shuffle_down(welford w, int offset) +{ + welford r; + r.count = __shfl_down_sync(0xffffffffu, w.count, offset); + r.mean = __shfl_down_sync(0xffffffffu, w.mean, offset); + r.m2 = __shfl_down_sync(0xffffffffu, w.m2, offset); + return r; +} + +__device__ welford block_merge(welford w) +{ + __shared__ welford warps[THREADS / 32]; + + for (int offset = 16; offset > 0; offset /= 2) + w = merge(w, shuffle_down(w, offset)); + + const int lane = threadIdx.x % 32, warp = threadIdx.x / 32; + if (lane == 0) + warps[warp] = w; + __syncthreads(); + + if (warp == 0) { + w = (lane < THREADS / 32) ? warps[lane] : welford{ 0.f, 0.f, 0.f }; + for (int offset = 16; offset > 0; offset /= 2) + w = merge(w, shuffle_down(w, offset)); + } + return w; /* valid in thread 0 */ +} + +/* Partial statistics of part blockIdx.y of group blockIdx.x (an item's group). */ +template +__global__ void partial_stats(const T* __restrict__ x, int64_t group_size, int parts, float* __restrict__ partial) +{ + const int64_t group = blockIdx.x; + const int part = blockIdx.y; + const int64_t per = (group_size + parts - 1) / parts; + const int64_t lo = part * per; + const int64_t hi = (lo + per < group_size) ? lo + per : group_size; + const T* base = x + group * group_size; + + welford w = { 0.f, 0.f, 0.f }; + for (int64_t i = lo + threadIdx.x; i < hi; i += THREADS) { + const float v = to_float(base[i]); + w.count += 1.f; + const float delta = v - w.mean; + w.mean += delta / w.count; + w.m2 += delta * (v - w.mean); + } + w = block_merge(w); + + if (threadIdx.x == 0) { + float* out = partial + (group * parts + part) * 3; + out[0] = w.count; + out[1] = w.mean; + out[2] = w.m2; + } +} + +/* Mean and reciprocal standard deviation of each group from its parts. */ +__global__ void finish_stats(const float* __restrict__ partial, int64_t groups, int parts, float eps, float* __restrict__ stats) +{ + for (int64_t g = blockIdx.x * (int64_t)blockDim.x + threadIdx.x; g < groups; g += (int64_t)blockDim.x * gridDim.x) { + welford w = { 0.f, 0.f, 0.f }; + for (int p = 0; p < parts; p++) { + const float* in = partial + (g * parts + p) * 3; + w = merge(w, welford{ in[0], in[1], in[2] }); + } + stats[2 * g] = w.mean; + stats[2 * g + 1] = rsqrtf(w.m2 / ((w.count > 0.f) ? w.count : 1.f) + eps); + } +} + +/* y = act(x * a + b), with a and b worked out once for each row (an item's channel, blockIdx.x). */ +template +__global__ void apply(const T* x, T* y, int64_t channels, int64_t voxels, int per_group, + const float* __restrict__ stats, const float* __restrict__ weight, const float* __restrict__ bias, + const float* __restrict__ scale, const float* __restrict__ shift, int silu) +{ + const int64_t row = blockIdx.x; /* item * channels + channel */ + const int64_t c = row % channels; + const int64_t item = row / channels; + const int64_t group = item * (channels / per_group) + c / per_group; + + const float mean = stats[2 * group], rstd = stats[2 * group + 1]; + const float w = weight ? weight[c] : 1.f; + const float b = bias ? bias[c] : 0.f; + const float s = scale ? 1.f + scale[row] : 1.f; + const float t = shift ? shift[row] : 0.f; + const float a = rstd * w * s; + const float d = (b - mean * rstd * w) * s + t; + + const T* in = x + row * voxels; + T* out = y + row * voxels; + for (int64_t i = blockIdx.y * (int64_t)THREADS + threadIdx.x; i < voxels; i += (int64_t)THREADS * gridDim.y) { + float v = to_float(in[i]) * a + d; + if (silu) + v = v / (1.f + __expf(-v)); + out[i] = from_float(v); + } +} + +template +cudaError_t run(const void* x, void* y, int64_t items, int64_t channels, int64_t voxels, int groups, + const float* weight, const float* bias, const float* scale, const float* shift, float eps, int silu, + float* workspace, cudaStream_t stream) +{ + const int per_group = (int)(channels / groups); + const int64_t group_size = per_group * voxels; + const int parts = (int)parts_of(group_size); + const int64_t n_groups = items * groups; + float* partial = workspace; + float* stats = workspace + n_groups * parts * 3; + + partial_stats<<>>((const T*)x, group_size, parts, partial); + finish_stats<<<(unsigned)((n_groups + 127) / 128), 128, 0, stream>>>(partial, n_groups, parts, eps, stats); + + int64_t ychunks = (voxels + THREADS * 8 - 1) / (THREADS * 8); + ychunks = (ychunks < 1) ? 1 : (ychunks > 65535) ? 65535 : ychunks; + apply<<>>( + (const T*)x, (T*)y, channels, voxels, per_group, stats, weight, bias, scale, shift, silu); + + return cudaGetLastError(); +} + +} /* namespace */ + +extern "C" { + +int64_t bartorch_gn_workspace(int64_t items, int64_t channels, int64_t voxels, int groups) +{ + if (groups < 1 || channels % groups) + return -1; + const int64_t n_groups = items * groups; + return n_groups * parts_of(channels / groups * voxels) * 3 + 2 * n_groups; +} + +int bartorch_gn_silu(const void* x, void* y, int64_t items, int64_t channels, int64_t voxels, int groups, + const void* weight, const void* bias, const void* scale, const void* shift, float eps, int silu, + int dtype, void* workspace, void* stream) +{ + if (items < 1 || channels < 1 || voxels < 1 || groups < 1 || channels % groups || items * channels > 2147483647LL) + return (int)cudaErrorInvalidValue; + + const float* w = (const float*)weight; + const float* b = (const float*)bias; + const float* s = (const float*)scale; + const float* t = (const float*)shift; + float* ws = (float*)workspace; + cudaStream_t st = (cudaStream_t)stream; + cudaError_t e; + + switch (dtype) { + case BARTORCH_NA_FLOAT32: e = run(x, y, items, channels, voxels, groups, w, b, s, t, eps, silu, ws, st); break; + case BARTORCH_NA_FLOAT16: e = run<__half>(x, y, items, channels, voxels, groups, w, b, s, t, eps, silu, ws, st); break; + case BARTORCH_NA_BFLOAT16: e = run<__nv_bfloat16>(x, y, items, channels, voxels, groups, w, b, s, t, eps, silu, ws, st); break; + default: e = cudaErrorInvalidValue; + } + + return (cudaSuccess == e) ? 0 : (int)e; +} + +int bartorch_gn_built(void) +{ + return 1; +} + +} diff --git a/tests/test_learning_groupnorm.py b/tests/test_learning_groupnorm.py new file mode 100644 index 000000000..60bf6c08a --- /dev/null +++ b/tests/test_learning_groupnorm.py @@ -0,0 +1,66 @@ +"""The fused group norm: the plain-torch path against the steps written out, and the +library's CUDA pass against the plain-torch path. +""" + +import pytest +import torch +import torch.nn.functional as F + +from bartorch import learning +from bartorch.learning._groupnorm import kernels_available + + +def _steps(x, groups, weight, bias, scale, shift, silu): + n, c = x.shape[:2] + g = x.double().reshape(n, groups, -1) + g = (g - g.mean(-1, keepdim=True)) / (g.var(-1, unbiased=False, keepdim=True) + 1e-5).sqrt() + h = g.reshape(x.shape) + shape = (n, c, *([1] * (x.ndim - 2))) + if weight is not None: + h = h * weight.double().reshape(1, c, *([1] * (x.ndim - 2))) + bias.double().reshape( + 1, c, *([1] * (x.ndim - 2)) + ) + if scale is not None: + h = h * (1 + scale.double().reshape(shape)) + shift.double().reshape(shape) + return F.silu(h) if silu else h + + +@pytest.mark.parametrize("affine", [False, True]) +@pytest.mark.parametrize("film", [False, True]) +@pytest.mark.parametrize("silu", [False, True]) +def test_reference_is_the_steps(affine, film, silu): + torch.manual_seed(0) + x = 3 + 2 * torch.randn(2, 8, 5, 6, 7) + weight, bias = (torch.randn(8), torch.randn(8)) if affine else (None, None) + scale, shift = (torch.randn(2, 8), torch.randn(2, 8)) if film else (None, None) + made = learning.group_norm_silu(x, 4, weight, bias, scale, shift, silu=silu) + torch.testing.assert_close( + made.double(), _steps(x, 4, weight, bias, scale, shift, silu), atol=1e-5, rtol=1e-5 + ) + + +def test_channels_must_split_into_groups(): + with pytest.raises(ValueError, match="groups"): + learning.group_norm_silu(torch.zeros(1, 6, 4, 4, 4), 4) + + +def test_gradients_go_through_plain_torch(): + x = torch.randn(1, 4, 3, 3, 3, requires_grad=True) + learning.group_norm_silu(x, 2).sum().backward() + assert x.grad is not None + + +@pytest.mark.skipif(not kernels_available(), reason="no card, or no CUDA pass in this library") +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("shape", [(2, 64, 9, 10, 11), (1, 128, 32, 32, 32), (3, 32, 7)]) +def test_pass_matches_plain_torch(dtype, shape): + torch.manual_seed(0) + x = (1 + 3 * torch.randn(shape, device="cuda")).to(dtype) + n, c = shape[:2] + weight, bias = torch.randn(c, device="cuda"), torch.randn(c, device="cuda") + scale, shift = torch.randn(n, c, device="cuda"), torch.randn(n, c, device="cuda") + with torch.no_grad(): + made = learning.group_norm_silu(x, 32 if c >= 32 else 8, weight, bias, scale, shift) + want = _steps(x.float(), 32 if c >= 32 else 8, weight, bias, scale, shift, True) + tolerance = {torch.float32: 1e-4, torch.float16: 1e-2, torch.bfloat16: 5e-2}[dtype] + assert float((made.double() - want).abs().max() / want.abs().max()) < tolerance From 371f622e0cacbd42de667cc8e50fa938c5c8a9c1 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 3 Nov 2026 14:39:25 +0000 Subject: [PATCH 5/8] learning: the CUDA passes' calls run on any device, and are tested on the host _neighborhood._on gives the card's stream with that card current, and no stream off a card; the group norm's call moves into _launch. A stand-in library doing the passes' arithmetic in torch on the host memory it is handed then checks everything on the Python side of the ABI without a card: pointers, shapes, workspace, the error report, and the attention's autograd function with its gradients against the plain-torch path. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MFB6QPoLKTK5gZEQeb1W8D --- src/bartorch/learning/_groupnorm.py | 11 +- src/bartorch/learning/_neighborhood.py | 15 ++- tests/test_learning_kernel_calls.py | 168 +++++++++++++++++++++++++ 3 files changed, 190 insertions(+), 4 deletions(-) create mode 100644 tests/test_learning_kernel_calls.py diff --git a/src/bartorch/learning/_groupnorm.py b/src/bartorch/learning/_groupnorm.py index cfeede75e..b8da5da1b 100644 --- a/src/bartorch/learning/_groupnorm.py +++ b/src/bartorch/learning/_groupnorm.py @@ -98,7 +98,15 @@ def group_norm_silu( if not fast: return _reference(x, groups, weight, bias, scale, shift, eps, silu) + return _launch(x, groups, weight, bias, scale, shift, eps, silu) + + +def _launch(x, groups, weight, bias, scale, shift, eps, silu): + """The library's pass on ``x`` wherever its memory is; the caller has checked it can run.""" + from bartorch.learning._neighborhood import _on + lib = _library() + n, channels = x.shape[:2] voxels = x[0, 0].numel() x = x.contiguous() device = x.device @@ -109,8 +117,7 @@ def group_norm_silu( size = lib.bartorch_gn_workspace(n, channels, voxels, groups) workspace = torch.empty(size, device=device, dtype=torch.float32) out = torch.empty_like(x) - with torch.cuda.device(device): - stream = torch.cuda.current_stream(device).cuda_stream + with _on(device) as stream: error = lib.bartorch_gn_silu( x.data_ptr(), out.data_ptr(), diff --git a/src/bartorch/learning/_neighborhood.py b/src/bartorch/learning/_neighborhood.py index e128065c5..5facb48b4 100644 --- a/src/bartorch/learning/_neighborhood.py +++ b/src/bartorch/learning/_neighborhood.py @@ -2,6 +2,8 @@ from __future__ import annotations +import contextlib + import torch from torch import nn @@ -61,6 +63,16 @@ def kernels_available() -> bool: return False +@contextlib.contextmanager +def _on(device): + """The current stream of ``device``'s card, with that card current; no stream off a card.""" + if device.type != "cuda": + yield None + return + with torch.cuda.device(device): + yield torch.cuda.current_stream(device).cuda_stream + + def _call(name: str, x, y, out, extent, window: int): """One CUDA pass; ``x`` and ``out`` are float32 weights or elements, as the pass takes them.""" groups, _, channels = y.shape @@ -69,8 +81,7 @@ def _call(name: str, x, y, out, extent, window: int): first, third = x.data_ptr(), _float_pointer(out.data_ptr()) else: # weights, elements -> elements first, third = _float_pointer(x.data_ptr()), out.data_ptr() - with torch.cuda.device(y.device): - stream = torch.cuda.current_stream(y.device).cuda_stream + with _on(y.device) as stream: error = getattr(_library(), f"bartorch_na3d_{name}")( first, y.data_ptr(), diff --git a/tests/test_learning_kernel_calls.py b/tests/test_learning_kernel_calls.py new file mode 100644 index 000000000..c253eb1d5 --- /dev/null +++ b/tests/test_learning_kernel_calls.py @@ -0,0 +1,168 @@ +"""How the learning module's CUDA passes are called, checked on the host. + +The passes themselves need a card (``test_learning_attention.py`` and +``test_learning_groupnorm.py`` check them there). Here a stand-in library does +the same arithmetic in torch on host memory reached through the pointers it is +handed, so everything on the Python side of the ABI -- shapes, pointers, +workspace, the autograd function around the attention passes -- runs anywhere. +""" + +import ctypes + +import pytest +import torch + +from bartorch.learning import _groupnorm, _neighborhood + + +def _address(pointer) -> int: + return pointer if isinstance(pointer, int) else ctypes.cast(pointer, ctypes.c_void_p).value + + +def _view(pointer, *shape) -> torch.Tensor: + """The float32 host memory at ``pointer`` as a tensor of ``shape``, without a copy.""" + count = 1 + for size in shape: + count *= size + array = (ctypes.c_float * count).from_address(_address(pointer)) + return torch.frombuffer(array, dtype=torch.float32).reshape(shape) + + +class _HostAttention: + """``bartorch_na3d_*`` on host float32 memory.""" + + def bartorch_na3d_built(self): + return 1 + + def _index(self, nz, ny, nx, window): + return _neighborhood._neighbours((nz, ny, nx), window, "cpu") + + def bartorch_na3d_gather_dot( + self, a, b, s, groups, nz, ny, nx, channels, window, dtype, stream + ): + voxels, index = nz * ny * nx, self._index(nz, ny, nx, window) + a, b = _view(a, groups, voxels, channels), _view(b, groups, voxels, channels) + _view(s, groups, voxels, window**3)[:] = torch.einsum("gvc,gvmc->gvm", a, b[:, index]) + return 0 + + def bartorch_na3d_gather_sum( + self, s, b, o, groups, nz, ny, nx, channels, window, dtype, stream + ): + voxels, index = nz * ny * nx, self._index(nz, ny, nx, window) + s, b = _view(s, groups, voxels, window**3), _view(b, groups, voxels, channels) + _view(o, groups, voxels, channels)[:] = torch.einsum("gvm,gvmc->gvc", s, b[:, index]) + return 0 + + def bartorch_na3d_scatter_sum( + self, s, a, o, groups, nz, ny, nx, channels, window, dtype, stream + ): + voxels, index = nz * ny * nx, self._index(nz, ny, nx, window) + s, a = _view(s, groups, voxels, window**3), _view(a, groups, voxels, channels) + out = torch.zeros(groups, voxels, channels) + contributions = s[..., None] * a[:, :, None, :] # (groups, voxels, window^3, channels) + out.index_add_(1, index.reshape(-1), contributions.reshape(groups, -1, channels)) + _view(o, groups, voxels, channels)[:] = out + return 0 + + +class _HostGroupNorm: + """``bartorch_gn_*`` on host float32 memory; ``fail`` makes the pass report a CUDA error.""" + + def __init__(self, fail=0): + self.fail = fail + + def bartorch_gn_built(self): + return 1 + + def bartorch_gn_workspace(self, items, channels, voxels, groups): + return 2 * items * groups + + def bartorch_gn_silu( + self, x, y, n, c, voxels, groups, w, b, s, t, eps, silu, dtype, ws, stream + ): + if self.fail: + return self.fail + assert dtype == 0 and stream is None + + def optional(pointer, *shape): + return None if pointer is None else _view(pointer, *shape) + + made = _groupnorm._reference( + _view(x, n, c, voxels), + groups, + optional(w, c), + optional(b, c), + optional(s, n, c), + optional(t, n, c), + eps, + bool(silu), + ) + _view(y, n, c, voxels)[:] = made + return 0 + + +def test_attention_passes_and_their_gradients_through_the_abi(monkeypatch): + monkeypatch.setattr(_neighborhood, "_library", _HostAttention) + torch.manual_seed(0) + n, heads, extent, channels, window = 2, 2, (4, 5, 3), 8, 3 + shape = (n * heads, extent[0] * extent[1] * extent[2], channels) + q, k, v = (torch.randn(shape, requires_grad=True) for _ in range(3)) + bias = (0.1 * torch.randn(heads, window**3)).requires_grad_() + made = _neighborhood._Attend.apply(q, k, v, bias, extent, window, heads) + grad = torch.randn(shape) + made.backward(grad) + + rq, rk, rv, rbias = (t.detach().clone().requires_grad_() for t in (q, k, v, bias)) + expected = _neighborhood._attend_reference(rq, rk, rv, extent, window, rbias, heads) + expected.backward(grad) + torch.testing.assert_close(made, expected, atol=1e-5, rtol=1e-5) + for got, want in ( + (q.grad, rq.grad), + (k.grad, rk.grad), + (v.grad, rv.grad), + (bias.grad, rbias.grad), + ): + torch.testing.assert_close(got, want, atol=1e-5, rtol=1e-4) + + +def test_attention_without_a_bias_and_a_failing_pass(monkeypatch): + monkeypatch.setattr(_neighborhood, "_library", _HostAttention) + q = torch.randn(1, 27, 8) + made = _neighborhood._Attend.apply(q, q, q, None, (3, 3, 3), 3, 1) + want = _neighborhood._attend_reference(q, q, q, (3, 3, 3), 3, None, 1) + torch.testing.assert_close(made, want, atol=1e-5, rtol=1e-5) + + class Failing(_HostAttention): + def bartorch_na3d_gather_dot(self, *args): + return 700 + + monkeypatch.setattr(_neighborhood, "_library", Failing) + with pytest.raises(RuntimeError, match="CUDA error 700"): + _neighborhood._Attend.apply(q, q, q, None, (3, 3, 3), 3, 1) + + +def test_kernels_are_unavailable_without_a_card(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(_neighborhood, "_library", _HostAttention) + monkeypatch.setattr(_groupnorm, "_library", _HostGroupNorm) + assert not _neighborhood.kernels_available() + assert not _groupnorm.kernels_available() + + +@pytest.mark.parametrize("affine", [False, True]) +@pytest.mark.parametrize("film", [False, True]) +def test_group_norm_pass_through_the_abi(monkeypatch, affine, film): + monkeypatch.setattr(_groupnorm, "_library", _HostGroupNorm) + torch.manual_seed(1) + x = 2 + torch.randn(2, 8, 3, 4, 5) + weight, bias = (torch.randn(8), torch.randn(8)) if affine else (None, None) + scale, shift = (torch.randn(2, 8), torch.randn(2, 8)) if film else (None, None) + made = _groupnorm._launch(x, 4, weight, bias, scale, shift, 1e-5, True) + want = _groupnorm._reference(x, 4, weight, bias, scale, shift, 1e-5, True) + torch.testing.assert_close(made, want) + + +def test_group_norm_pass_reports_its_error(monkeypatch): + monkeypatch.setattr(_groupnorm, "_library", lambda: _HostGroupNorm(fail=700)) + with pytest.raises(RuntimeError, match="CUDA error 700"): + _groupnorm._launch(torch.randn(1, 4, 8), 2, None, None, None, None, 1e-5, False) From c14fd3deab7155cb4a34aab420fec82a2e3194ba Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 3 Nov 2026 14:40:19 +0000 Subject: [PATCH 6/8] check scripts: mark the nvcc call as reviewed for bandit (B404, B603) The standalone checks run nvcc with a fixed argument list on a file of this checkout; nothing reaches it from outside. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MFB6QPoLKTK5gZEQeb1W8D --- scripts/check_attention.py | 4 ++-- scripts/check_groupnorm.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/scripts/check_attention.py b/scripts/check_attention.py index c09c3d7d4..e2a6004b9 100644 --- a/scripts/check_attention.py +++ b/scripts/check_attention.py @@ -18,7 +18,7 @@ import argparse import ctypes import json -import subprocess +import subprocess # nosec B404 - runs nvcc, a fixed command import tempfile import time from pathlib import Path @@ -58,7 +58,7 @@ def standalone(arch: str | None) -> ctypes.CDLL: "-I", str(ROOT / "src" / "csrc"), "-o", str(out), str(ROOT / "src" / "csrc" / "ops" / "attention.cu"), ] # fmt: skip - subprocess.run(command, check=True) + subprocess.run(command, check=True) # nosec B603 - nvcc on this checkout's source lib = ctypes.CDLL(str(out)) floats = ctypes.POINTER(ctypes.c_float) ints = [ diff --git a/scripts/check_groupnorm.py b/scripts/check_groupnorm.py index 875a612d0..9b4ee1b19 100644 --- a/scripts/check_groupnorm.py +++ b/scripts/check_groupnorm.py @@ -20,7 +20,7 @@ import argparse import ctypes import json -import subprocess +import subprocess # nosec B404 - runs nvcc, a fixed command import tempfile import time from pathlib import Path @@ -59,7 +59,7 @@ def standalone(arch: str | None) -> ctypes.CDLL: "-I", str(ROOT / "src" / "csrc"), "-o", str(out), str(ROOT / "src" / "csrc" / "ops" / "groupnorm.cu"), ] # fmt: skip - subprocess.run(command, check=True) + subprocess.run(command, check=True) # nosec B603 - nvcc on this checkout's source lib = ctypes.CDLL(str(out)) lib.bartorch_gn_built.restype = ctypes.c_int lib.bartorch_gn_workspace.restype = ctypes.c_int64 From c2d3c2c6dfc4badd591079ad7cca0324c431ee04 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 3 Nov 2026 15:48:04 +0000 Subject: [PATCH 7/8] Fused group norm: a channels-last (NDHWC) pass that keeps the layout cuDNN's fastest 3-D convolutions run channels last; a norm that only reads NCDHW forces a transpose on either side. The pass now takes channels-last input as it lies (statistics per channel strip, groups merged after) and writes channels-last output. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MFB6QPF5rgL8sKizNYzB8D --- scripts/check_groupnorm.py | 24 ++++- src/bartorch/_abi.py | 1 + src/bartorch/learning/_groupnorm.py | 17 +++- src/csrc/include/bartorch.h | 7 +- src/csrc/ops/groupnorm.c | 4 +- src/csrc/ops/groupnorm.cu | 141 ++++++++++++++++++++++++++-- tests/test_learning_kernel_calls.py | 24 ++++- 7 files changed, 198 insertions(+), 20 deletions(-) diff --git a/scripts/check_groupnorm.py b/scripts/check_groupnorm.py index 9b4ee1b19..5b832807f 100644 --- a/scripts/check_groupnorm.py +++ b/scripts/check_groupnorm.py @@ -73,7 +73,7 @@ def standalone(arch: str | None) -> ctypes.CDLL: lib.bartorch_gn_silu.restype = ctypes.c_int lib.bartorch_gn_silu.argtypes = [ v, v, ctypes.c_int64, ctypes.c_int64, ctypes.c_int64, ctypes.c_int, - v, v, v, v, ctypes.c_float, ctypes.c_int, ctypes.c_int, v, v, + v, v, v, v, ctypes.c_float, ctypes.c_int, ctypes.c_int, ctypes.c_int, v, v, ] # fmt: skip return lib @@ -88,15 +88,19 @@ def exact(x, groups, weight, bias, scale, shift): return F.silu(h * (1 + scale.double().reshape(shape)) + shift.double().reshape(shape)) -def check(dtype, shape, groups=32): +def check(dtype, shape, groups=32, last=False): torch.manual_seed(0) n, c = shape[:2] x = (2 + 3 * torch.randn(shape, device="cuda")).to(dtype) + if last: + x = x.to(memory_format=torch.channels_last_3d) weight, bias = torch.randn(c, device="cuda"), torch.randn(c, device="cuda") scale, shift = 0.3 * torch.randn(n, c, device="cuda"), torch.randn(n, c, device="cuda") with torch.no_grad(): made = gn.group_norm_silu(x, groups, weight, bias, scale, shift) want = exact(x, groups, weight, bias, scale, shift) + if last and not made.is_contiguous(memory_format=torch.channels_last_3d): + raise AssertionError("the channels-last pass changed the layout") return float((made.double() - want).abs().max() / want.abs().max()) @@ -126,11 +130,22 @@ def half_norm(): def fused(): return gn.group_norm_silu(x, groups, weight, bias, scale, shift) + xl = x.to(memory_format=torch.channels_last_3d) + + def eager_last(): + h = F.group_norm(xl.float(), groups, weight, bias).type(xl.dtype) + return F.silu(h * (1 + s5) + t5) + + def fused_last(): + return gn.group_norm_silu(xl, groups, weight, bias, scale, shift) + with torch.no_grad(): row = {"channels": channels, "extent": extent} row["eager fp32 norm ms"] = 1e3 * timed(eager) row["eager fp16 norm ms"] = 1e3 * timed(half_norm) row["fused ms"] = 1e3 * timed(fused) + row["eager channels-last ms"] = 1e3 * timed(eager_last) + row["fused channels-last ms"] = 1e3 * timed(fused_last) row["bytes moved MB (floor)"] = 3 * x.numel() * x.element_size() / 1e6 return {k: round(v, 3) if isinstance(v, float) else v for k, v in row.items()} @@ -161,8 +176,9 @@ def main(): (1, 64, 64, 64, 64), (3, 512, 8, 8, 8), ): - key = str(dtype).removeprefix("torch.") - worst[key] = max(worst.get(key, 0.0), check(dtype, shape)) + for last in (False, True): + key = str(dtype).removeprefix("torch.") + (" channels-last" if last else "") + worst[key] = max(worst.get(key, 0.0), check(dtype, shape, last=last)) print("largest relative error against float64:", json.dumps(worst), flush=True) for channels, extent in ( diff --git a/src/bartorch/_abi.py b/src/bartorch/_abi.py index 18545c77e..a07e224dd 100644 --- a/src/bartorch/_abi.py +++ b/src/bartorch/_abi.py @@ -1333,6 +1333,7 @@ def bind(lib: ctypes.CDLL) -> ctypes.CDLL: ctypes.c_float, ctypes.c_int, ctypes.c_int, + ctypes.c_int, ctypes.c_void_p, ctypes.c_void_p, ] diff --git a/src/bartorch/learning/_groupnorm.py b/src/bartorch/learning/_groupnorm.py index b8da5da1b..144250874 100644 --- a/src/bartorch/learning/_groupnorm.py +++ b/src/bartorch/learning/_groupnorm.py @@ -37,6 +37,14 @@ def _reference(x, groups, weight, bias, scale, shift, eps, silu): return h.to(x.dtype) +def _channels_last(x) -> bool: + """Whether ``x`` is channels last (N(D)HWC, as cuDNN prefers) and not also contiguous.""" + if x.ndim not in (4, 5) or x.is_contiguous(): + return False + layout = torch.channels_last_3d if x.ndim == 5 else torch.channels_last + return x.is_contiguous(memory_format=layout) and x.shape[1] <= 4096 + + def _pointer(t): return None if t is None else t.data_ptr() @@ -65,7 +73,9 @@ def group_norm_silu( one write, in ``x``'s own element type (float32, float16 or bfloat16) with float32 statistics -- where eager torch reads and writes the activations four or five times, with casts around the norm. Anywhere else plain torch - computes the same, also with float32 statistics. + computes the same, also with float32 statistics. ``x`` may be contiguous or + channels last (``torch.channels_last_3d``, which cuDNN's convolutions + prefer), and the output keeps its layout. Parameters ---------- @@ -108,7 +118,9 @@ def _launch(x, groups, weight, bias, scale, shift, eps, silu): lib = _library() n, channels = x.shape[:2] voxels = x[0, 0].numel() - x = x.contiguous() + last = _channels_last(x) + if not last: + x = x.contiguous() device = x.device w = _as_float(weight, device, (channels,)) b = _as_float(bias, device, (channels,)) @@ -131,6 +143,7 @@ def _launch(x, groups, weight, bias, scale, shift, eps, silu): _pointer(t), float(eps), int(bool(silu)), + int(last), _DTYPES[x.dtype], workspace.data_ptr(), stream, diff --git a/src/csrc/include/bartorch.h b/src/csrc/include/bartorch.h index 67a743c3e..59e4a1dac 100644 --- a/src/csrc/include/bartorch.h +++ b/src/csrc/include/bartorch.h @@ -972,8 +972,9 @@ BARTORCH_API int bartorch_na3d_scatter_sum(const float* s, const void* a, void* * * y[n, c, v] = act(((x[n, c, v] - mean[n, g]) * rstd[n, g] * w[c] + b[c]) * (1 + s[n, c]) + t[n, c]) * - * x and y are (items, channels, voxels), voxels fastest, of a BARTORCH_NA_* - * element type, and y may be x; weight and bias are (channels) float32, scale + * x and y are (items, channels, voxels), voxels fastest, or with channels_last + * nonzero (items, voxels, channels), channels fastest (at most 4096), of a + * BARTORCH_NA_* element type, and y may be x; weight and bias are (channels) float32, scale * and shift (items, channels) float32, each NULL when absent; act is SiLU * when silu is nonzero. Statistics are float32. workspace holds * bartorch_gn_workspace(...) floats on the card (-1 for a bad shape). Runs on @@ -984,7 +985,7 @@ BARTORCH_API int bartorch_gn_built(void); BARTORCH_API int64_t bartorch_gn_workspace(int64_t items, int64_t channels, int64_t voxels, int groups); BARTORCH_API int bartorch_gn_silu(const void* x, void* y, int64_t items, int64_t channels, int64_t voxels, int groups, const void* weight, const void* bias, const void* scale, const void* shift, float eps, - int silu, int dtype, void* workspace, void* stream); + int silu, int channels_last, int dtype, void* workspace, void* stream); /* Whether a pointer is device memory; always false without CUDA. */ BARTORCH_API int bartorch_on_device(const void* ptr); diff --git a/src/csrc/ops/groupnorm.c b/src/csrc/ops/groupnorm.c index 295507310..b9ccb4850 100644 --- a/src/csrc/ops/groupnorm.c +++ b/src/csrc/ops/groupnorm.c @@ -18,10 +18,10 @@ int64_t bartorch_gn_workspace(int64_t items, int64_t channels, int64_t voxels, i int bartorch_gn_silu(const void* x, void* y, int64_t items, int64_t channels, int64_t voxels, int groups, const void* weight, const void* bias, const void* scale, const void* shift, float eps, int silu, - int dtype, void* workspace, void* stream) + int channels_last, int dtype, void* workspace, void* stream) { (void)x; (void)y; (void)items; (void)channels; (void)voxels; (void)groups; (void)weight; (void)bias; - (void)scale; (void)shift; (void)eps; (void)silu; (void)dtype; (void)workspace; (void)stream; + (void)scale; (void)shift; (void)eps; (void)silu; (void)channels_last; (void)dtype; (void)workspace; (void)stream; return -1; } diff --git a/src/csrc/ops/groupnorm.cu b/src/csrc/ops/groupnorm.cu index 955b5997e..0b5ad92d7 100644 --- a/src/csrc/ops/groupnorm.cu +++ b/src/csrc/ops/groupnorm.cu @@ -12,7 +12,9 @@ * with the norm's affine w, b and the FiLM s, t each optional (NULL), and act * SiLU when asked for, the identity otherwise. x and y are * (items, channels, voxels), voxels fastest -- a contiguous NC(D)HW tensor -- - * in float, half or bfloat16; w, b are (channels) float32, s, t (items, + * or with channels_last (items, voxels, channels), channels fastest -- a + * channels-last N(D)HWC one, which cuDNN's convolutions prefer -- in float, + * half or bfloat16; w, b are (channels) float32, s, t (items, * channels) float32. Statistics are in float32 by Welford's method; a group * is spread over several blocks, whose partial sums a second pass merges. * y may be x. @@ -170,6 +172,127 @@ __global__ void apply(const T* x, T* y, int64_t channels, int64_t voxels, int pe } } +/* + * Channels-last (items, voxels, channels): the partial statistics of every group of an + * item over part blockIdx.y of its voxels. Each thread follows one channel down the + * voxels (or several, when there are more channels than threads), so a row of channels + * is read whole by neighbouring threads; the channels of each group are merged at the end. + */ +template +__global__ void partial_stats_last(const T* __restrict__ x, int64_t voxels, int channels, int groups, int parts, + float* __restrict__ partial) +{ + extern __shared__ welford entries[]; + + const int64_t item = blockIdx.x; + const int part = blockIdx.y; + const int64_t per = (voxels + parts - 1) / parts; + const int64_t lo = part * per; + const int64_t hi = (lo + per < voxels) ? lo + per : voxels; + const T* base = x + item * voxels * channels; + const int rows = (channels <= THREADS) ? THREADS / channels : 1; + + if (channels <= THREADS) { + const int t = threadIdx.x; + if (t < rows * channels) { + const int r = t / channels, c = t % channels; + welford w = { 0.f, 0.f, 0.f }; + for (int64_t v = lo + r; v < hi; v += rows) { + const float value = to_float(base[v * channels + c]); + w.count += 1.f; + const float delta = value - w.mean; + w.mean += delta / w.count; + w.m2 += delta * (value - w.mean); + } + entries[t] = w; + } + } else { + for (int c = threadIdx.x; c < channels; c += THREADS) { + welford w = { 0.f, 0.f, 0.f }; + for (int64_t v = lo; v < hi; v++) { + const float value = to_float(base[v * channels + c]); + w.count += 1.f; + const float delta = value - w.mean; + w.mean += delta / w.count; + w.m2 += delta * (value - w.mean); + } + entries[c] = w; + } + } + __syncthreads(); + + const int per_group = channels / groups; + for (int g = threadIdx.x; g < groups; g += THREADS) { + welford w = { 0.f, 0.f, 0.f }; + for (int r = 0; r < rows; r++) + for (int j = 0; j < per_group; j++) + w = merge(w, entries[r * channels + g * per_group + j]); + float* out = partial + ((item * groups + g) * parts + part) * 3; + out[0] = w.count; + out[1] = w.mean; + out[2] = w.m2; + } +} + +/* Each item's channel as one multiply-add: y = act(x * ab[0] + ab[1]). */ +__global__ void coefficients(int64_t items, int channels, int per_group, const float* __restrict__ stats, + const float* __restrict__ weight, const float* __restrict__ bias, const float* __restrict__ scale, + const float* __restrict__ shift, float* __restrict__ ab) +{ + for (int64_t row = blockIdx.x * (int64_t)blockDim.x + threadIdx.x; row < items * channels; + row += (int64_t)blockDim.x * gridDim.x) { + const int64_t c = row % channels, item = row / channels; + const int64_t group = item * (channels / per_group) + c / per_group; + const float mean = stats[2 * group], rstd = stats[2 * group + 1]; + const float w = weight ? weight[c] : 1.f; + const float b = bias ? bias[c] : 0.f; + const float s = scale ? 1.f + scale[row] : 1.f; + const float t = shift ? shift[row] : 0.f; + ab[2 * row] = rstd * w * s; + ab[2 * row + 1] = (b - mean * rstd * w) * s + t; + } +} + +template +__global__ void apply_last(const T* x, T* y, int64_t items, int64_t voxels, int channels, const float* __restrict__ ab, int silu) +{ + const int64_t total = items * voxels * channels; + for (int64_t i = blockIdx.x * (int64_t)blockDim.x + threadIdx.x; i < total; i += (int64_t)blockDim.x * gridDim.x) { + const int64_t c = i % channels, item = i / (voxels * channels); + const int64_t row = item * channels + c; + float v = to_float(x[i]) * ab[2 * row] + ab[2 * row + 1]; + if (silu) + v = v / (1.f + __expf(-v)); + y[i] = from_float(v); + } +} + +template +cudaError_t run_last(const void* x, void* y, int64_t items, int64_t channels, int64_t voxels, int groups, + const float* weight, const float* bias, const float* scale, const float* shift, float eps, int silu, + float* workspace, cudaStream_t stream) +{ + const int per_group = (int)(channels / groups); + const int parts = (int)parts_of(per_group * voxels); + const int64_t n_groups = items * groups; + float* partial = workspace; + float* stats = partial + n_groups * parts * 3; + float* ab = stats + 2 * n_groups; + const int entries = (channels <= THREADS) ? (int)(THREADS / channels * channels) : (int)channels; + + partial_stats_last<<>>( + (const T*)x, voxels, (int)channels, groups, parts, partial); + finish_stats<<<(unsigned)((n_groups + 127) / 128), 128, 0, stream>>>(partial, n_groups, parts, eps, stats); + coefficients<<<(unsigned)((items * channels + 127) / 128), 128, 0, stream>>>( + items, (int)channels, per_group, stats, weight, bias, scale, shift, ab); + + int64_t blocks = (items * voxels * channels + THREADS * 8 - 1) / (THREADS * 8); + blocks = (blocks < 1) ? 1 : (blocks > 65535 * 16) ? 65535 * 16 : blocks; + apply_last<<<(unsigned)blocks, THREADS, 0, stream>>>((const T*)x, (T*)y, items, voxels, (int)channels, ab, silu); + + return cudaGetLastError(); +} + template cudaError_t run(const void* x, void* y, int64_t items, int64_t channels, int64_t voxels, int groups, const float* weight, const float* bias, const float* scale, const float* shift, float eps, int silu, @@ -202,15 +325,17 @@ int64_t bartorch_gn_workspace(int64_t items, int64_t channels, int64_t voxels, i if (groups < 1 || channels % groups) return -1; const int64_t n_groups = items * groups; - return n_groups * parts_of(channels / groups * voxels) * 3 + 2 * n_groups; + return n_groups * parts_of(channels / groups * voxels) * 3 + 2 * n_groups + 2 * items * channels; } int bartorch_gn_silu(const void* x, void* y, int64_t items, int64_t channels, int64_t voxels, int groups, const void* weight, const void* bias, const void* scale, const void* shift, float eps, int silu, - int dtype, void* workspace, void* stream) + int channels_last, int dtype, void* workspace, void* stream) { if (items < 1 || channels < 1 || voxels < 1 || groups < 1 || channels % groups || items * channels > 2147483647LL) return (int)cudaErrorInvalidValue; + if (channels_last && channels > 4096) /* the statistics' shared memory holds a Welford sum per channel */ + return (int)cudaErrorInvalidValue; const float* w = (const float*)weight; const float* b = (const float*)bias; @@ -220,12 +345,16 @@ int bartorch_gn_silu(const void* x, void* y, int64_t items, int64_t channels, in cudaStream_t st = (cudaStream_t)stream; cudaError_t e; +#define BARTORCH_GN_RUN(T) (channels_last \ + ? run_last(x, y, items, channels, voxels, groups, w, b, s, t, eps, silu, ws, st) \ + : run(x, y, items, channels, voxels, groups, w, b, s, t, eps, silu, ws, st)) switch (dtype) { - case BARTORCH_NA_FLOAT32: e = run(x, y, items, channels, voxels, groups, w, b, s, t, eps, silu, ws, st); break; - case BARTORCH_NA_FLOAT16: e = run<__half>(x, y, items, channels, voxels, groups, w, b, s, t, eps, silu, ws, st); break; - case BARTORCH_NA_BFLOAT16: e = run<__nv_bfloat16>(x, y, items, channels, voxels, groups, w, b, s, t, eps, silu, ws, st); break; + case BARTORCH_NA_FLOAT32: e = BARTORCH_GN_RUN(float); break; + case BARTORCH_NA_FLOAT16: e = BARTORCH_GN_RUN(__half); break; + case BARTORCH_NA_BFLOAT16: e = BARTORCH_GN_RUN(__nv_bfloat16); break; default: e = cudaErrorInvalidValue; } +#undef BARTORCH_GN_RUN return (cudaSuccess == e) ? 0 : (int)e; } diff --git a/tests/test_learning_kernel_calls.py b/tests/test_learning_kernel_calls.py index c253eb1d5..12d5910ab 100644 --- a/tests/test_learning_kernel_calls.py +++ b/tests/test_learning_kernel_calls.py @@ -78,7 +78,7 @@ def bartorch_gn_workspace(self, items, channels, voxels, groups): return 2 * items * groups def bartorch_gn_silu( - self, x, y, n, c, voxels, groups, w, b, s, t, eps, silu, dtype, ws, stream + self, x, y, n, c, voxels, groups, w, b, s, t, eps, silu, last, dtype, ws, stream ): if self.fail: return self.fail @@ -87,8 +87,13 @@ def bartorch_gn_silu( def optional(pointer, *shape): return None if pointer is None else _view(pointer, *shape) + def volume(pointer): # (n, c, voxels), whichever way it lies in memory + if last: + return _view(pointer, n, voxels, c).transpose(1, 2) + return _view(pointer, n, c, voxels) + made = _groupnorm._reference( - _view(x, n, c, voxels), + volume(x), groups, optional(w, c), optional(b, c), @@ -97,7 +102,7 @@ def optional(pointer, *shape): eps, bool(silu), ) - _view(y, n, c, voxels)[:] = made + volume(y)[:] = made return 0 @@ -166,3 +171,16 @@ def test_group_norm_pass_reports_its_error(monkeypatch): monkeypatch.setattr(_groupnorm, "_library", lambda: _HostGroupNorm(fail=700)) with pytest.raises(RuntimeError, match="CUDA error 700"): _groupnorm._launch(torch.randn(1, 4, 8), 2, None, None, None, None, 1e-5, False) + + +def test_group_norm_pass_keeps_a_channels_last_layout(monkeypatch): + monkeypatch.setattr(_groupnorm, "_library", _HostGroupNorm) + torch.manual_seed(2) + x = (1 + torch.randn(2, 8, 3, 4, 5)).to(memory_format=torch.channels_last_3d) + assert _groupnorm._channels_last(x) + scale, shift = torch.randn(2, 8), torch.randn(2, 8) + made = _groupnorm._launch(x, 4, None, None, scale, shift, 1e-5, True) + assert made.is_contiguous(memory_format=torch.channels_last_3d) + want = _groupnorm._reference(x, 4, None, None, scale, shift, 1e-5, True) + torch.testing.assert_close(made, want) + assert not _groupnorm._channels_last(torch.zeros(1, 4, 1, 1, 1)) # contiguous either way: not From e63b8cd4c65bacc60ba0990c6a77cf31cf23fac1 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 3 Nov 2026 16:00:22 +0000 Subject: [PATCH 8/8] Fused group norm, channels last: split an item by its whole size The statistics pass split each item into as many runs as a single group's size called for, so a deep level (512 channels at 8^3, three items) ran on three blocks; the A40 showed it 3-12x slower than the NCDHW pass below 64^3 and 2.5-4x at 64^3. Runs are now sized by the item's whole extent, at least 4096 elements each and at most 256 per item. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MFB6QPF5rgL8sKizNYzB8D --- src/csrc/ops/groupnorm.cu | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/src/csrc/ops/groupnorm.cu b/src/csrc/ops/groupnorm.cu index 0b5ad92d7..7d108fa0e 100644 --- a/src/csrc/ops/groupnorm.cu +++ b/src/csrc/ops/groupnorm.cu @@ -53,6 +53,22 @@ int64_t parts_of(int64_t group_size) return (p < 1) ? 1 : (p > MAX_PARTS) ? MAX_PARTS : p; } +/* + * Channels last, a block reads all of an item's channels over a run of voxels, so + * an item is split by its whole size, not a group's, and into smaller runs: with + * few items and many channels (a UNet's deep levels) a split by group would leave + * the card with a handful of blocks. + */ +const int64_t CHUNK_LAST = 4096; +const int MAX_PARTS_LAST = 256; + +int64_t parts_last(int64_t channels, int64_t voxels) +{ + int64_t p = (channels * voxels + CHUNK_LAST - 1) / CHUNK_LAST; + p = (p > voxels) ? voxels : p; + return (p < 1) ? 1 : (p > MAX_PARTS_LAST) ? MAX_PARTS_LAST : p; +} + struct welford { float count, mean, m2; }; @@ -273,7 +289,7 @@ cudaError_t run_last(const void* x, void* y, int64_t items, int64_t channels, in float* workspace, cudaStream_t stream) { const int per_group = (int)(channels / groups); - const int parts = (int)parts_of(per_group * voxels); + const int parts = (int)parts_last(channels, voxels); const int64_t n_groups = items * groups; float* partial = workspace; float* stats = partial + n_groups * parts * 3; @@ -325,7 +341,8 @@ int64_t bartorch_gn_workspace(int64_t items, int64_t channels, int64_t voxels, i if (groups < 1 || channels % groups) return -1; const int64_t n_groups = items * groups; - return n_groups * parts_of(channels / groups * voxels) * 3 + 2 * n_groups + 2 * items * channels; + const int64_t a = parts_of(channels / groups * voxels), b = parts_last(channels, voxels); + return n_groups * ((a > b) ? a : b) * 3 + 2 * n_groups + 2 * items * channels; } int bartorch_gn_silu(const void* x, void* y, int64_t items, int64_t channels, int64_t voxels, int groups,