perf: remove the per-sample loops and host stalls from the augmentations

[Hendrik-code]

Every augmentation in smauglab/transforms was read for work that scales with the
batch rather than with the volume. The three patterns that kept turning up:

1. A single-channel F.conv3d is cuDNN's slowest path. Folding the batch into
   the channel axis so groups == B*C picks ATen's depthwise kernel instead. On an
   A40, a 3x3x3 kernel over a 128³ patch: 4.25 ms → 0.12 ms. It is also the more
   accurate path — the implicit-GEMM path it leaves runs in TF32 and lands 6e-3 from
   a float64 reference where the depthwise kernel lands 4e-6 away.
2. Reading a device scalar back to Python blocks the host. float(scales[b]),
   int(mask.sum()), bool(flips[b, axis]), torch.rand(1).item() — each inside a
   per-sample loop. Now one .tolist() before the loop.
3. scatter_add_ into a handful of bins is an atomics-contention test. Region
   sums, k-means cluster sums and the equalisation histogram now go through
   bincount (shared-memory accumulation) or a prefix sum.

Measurements

Single idle NVIDIA L40S, torch 2.6.0+cu124, kornia 0.8.2. Input [2, 1, 128, 128, 128]
with a 3-label segmentation. apply_transform called directly so the probability
gate does not decide what gets timed. Each cell is the median of three runs of
25 calls each, base = origin/main in a separate worktree.

Harness: bench.py / bench_pipeline.py / bench_cpu.py, reproduced in the
"Reproducing" section below.

GPU transforms

Transform	before (ms)	after (ms)	speed-up
PaletteSynthesisGPU (EM partitioner)	43.55	28.27	1.54x
RandomPaletteGPU	17.96	9.46	1.90x
RandomRedistributeSegGPU	5.12	3.42	1.50x
RandomDomainTransferGPU	4.55	4.22	1.08x
RandomRandConvGPU	2.02	0.69	2.91x
RandomUnsharpMaskGPU	1.62	1.18	1.37x
RandomGaussianBlurGPU	1.57	1.13	1.39x
RandomAffineGPU	1.55	1.54	1.01x (untouched)
RandomBiasFieldGPU	1.23	0.57	2.15x
RandomClampGPU	1.18	0.35	3.38x
RandomHistogramEqualizationGPU	0.88	0.64	1.37x
RandomScharrGPU	0.77	0.54	1.42x
RandomInverseGPU	0.39	0.32	1.23x
RandomContrastGPU	0.36	0.26	1.38x
RandomInvGammaGPU	0.36	0.34	1.06x
RandomGammaGPU	0.35	0.32	1.10x
ZscoreNormalizationGPU	0.34	0.31	1.10x
RandomLaplaceGPU	0.33	0.31	1.05x
RandomSigmoidGPU	0.32	0.28	1.11x
RandomGaussianNoiseGPU	0.31	0.25	1.26x
RandomBrightnessGPU	0.29	0.23	1.27x
RandomLog1pGPU	0.28	0.25	1.10x
RandomSqrtGPU	0.27	0.24	1.11x
RandomExpGPU	0.27	0.24	1.11x
RandomSinGPU	0.27	0.25	1.09x
RandomCropTransformGPU	0.24	0.15	1.63x
RandomAcqTransformGPU	0.23	0.18	1.28x
RandomLowResTransformGPU	0.17	0.14	1.15x
RandomFlipTransformGPU	0.10	0.09	1.13x
sum of the column	86.84	56.16	1.55x

RandomAffineGPU is kornia's warp_affine3d and was not touched; it is in the
table as a control.

SynthSeg

Same patch, median of three runs of 10 calls.

	before (ms)	after (ms)	speed-up
RandomSynthSegGPU	3.31	2.51	1.32x
SynthSegGenerator (full pipeline)	12.88	11.52	1.12x
SynthSegGenerator (EM label completion)	66.87	57.50	1.16x

CPU backend

[1, 96, 96, 96], 8 threads, median of three runs of 7 calls.

Transform	before (ms)	after (ms)	speed-up
ScharrConvTransform	110.53	29.20	3.78x
RedistributeTransform	83.05	33.02	2.52x
LaplaceConvTransform	26.30	11.01	2.39x
HistogramEqualTransform	5.92	5.66	1.05x
GaussianBlurTransform	57.36	57.55	1.00x (batchgeneratorsv2)

The rest of the CPU backend is within noise of before; GaussianBlurTransform is
third-party and in the table only to show it was not made worse.

End to end

Full pipeline forward(image, mask), [2, 1, 128, 128, 128], median of five runs
of 40 calls. Which transforms fire is itself random, hence the repeated runs.

Config	mode	before (ms)	after (ms)	speed-up
transform_params_gpu.json	sequential	18.16	17.13	1.06x
transform_params_gpu.json	random_order	13.18	8.26	1.60x
transform_params_gpu.json	random_order_ta	15.68	11.85	1.32x
…one-sequence-to-segment-them-all.json	sequential	16.49	15.34	1.07x
…one-sequence-to-segment-them-all.json	random_order	10.27	7.28	1.41x
…one-sequence-to-segment-them-all.json	random_order_ta	13.87	10.36	1.34x

Why sequential barely moves, and what the next bottleneck is. Run the same
pipelines with every p set to 0 — nothing fires, so what is left is kornia's
AugmentationSequential machinery — and transform_params_gpu.json still takes
10.4 ms before this PR and 10.5 ms after (measured in isolation, three
runs of 60; the two are within noise of each other). A profile of one sequential
call shows ~920 kernel launches for ~3.6 ms of GPU work: the container is
CPU-bound on per-module parameter sampling, not on the augmentations. The
random_order modes gain more because the bucket actually applies the transforms
this PR made faster. Cutting the container overhead is a separate change and is
not attempted here.

Correctness

unit_tests/test_fast_paths.py (new, 28 tests) pins every reformulated primitive
against the straightforward spelling it replaced — the grouped convolution against
a per-plane loop, segment_sum against scatter_add_, _batched_quantiles
against torch.quantile, laplacian_response against the dense kernel,
_region_moments against the explicit masked statistics, the bias-field
contraction against the monomial loop, and the EM quadratic form against the
Gaussian log density written out.

Beyond that, every registered GPU transform was run under 24 seeds on CPU and on
CUDA, at defaults and in 30 explicitly-configured variants (same_on_batch,
in_seg/out_seg/mix_in_out, one-hot and label-map segmentations, non-default
orders and kernel sizes), and compared against origin/main:

• Bit-identical: flips, crop, low-res, acquisition, affine, gamma, inverse
  gamma, brightness, contrast, noise, clamp, inverse, histogram equalisation,
  z-score, the elementwise functions, and domain transfer in all four of its
  modes.
• Float32 round-off only: the convolution paths (≤ 1.4e-6), the bias field
  (≤ 2.7e-4 at order 5, from reordering a 56-term sum), signed Scharr (2.4e-4 on
  responses of magnitude ~700).
• Same distribution, different RNG stream: PALETTE, composed PALETTE,
  redistribute, the flip-axis "at least one" guarantee, the bucketed
  RandomChooseXTransformsGPU, and SynthSeg's background randomisation now draw
  the same independent uniforms in a different order, so a seeded run does not
  reproduce the previous release's output. Each was checked against a control that
  compares origin/main with itself over a disjoint block of seeds; in every
  case the old-vs-new difference in mean, standard deviation and quantiles is
  smaller than that control, i.e. indistinguishable from re-seeding. Example
  (RandomPaletteGPU, 120 seeds): quantile distance control 0.254, old-vs-new
  0.122.

One deliberate behaviour change, with its test updated rather than deleted:
RandomHistogramEqualizationGPU used to raise RuntimeError out of torch.histc
when the caller passed a non-finite voxel, in the middle of a training step. The
batched version derives its bin edges itself, so the NaN now reaches the
transform's own non-finite guard, which returns the channel untouched and prints a
warning — which is what that guard was added for.

Also fixed along the way: RandomCropTransformGPU now moves its sampled
parameters to the input's device. kornia builds the generator's samplers on
whatever device set_rng_device_and_dtype last saw, and
RandomChooseXTransformsGPU calls forward_parameters without going through
forward — the scalar version read the draw out with float(...) and never
noticed.

437 passed, 3 skipped, 1320 subtests, ruff check, ruff format --check and
mypy all clean.

Reproducing

The benchmark scripts are not in the repository (they are throwaway harnesses, not
project code). Each one builds every registered transform with p=1.0, calls
apply_transform directly on a fixed input, and reports the median wall time with
torch.cuda.synchronize() around each call. To repeat the comparison:

git worktree add /tmp/base origin/main
# then run the same harness with PYTHONPATH pointed at each checkout in turn,
# on an idle GPU, interleaving the two so neither sees a different clock state.

Numbers above were taken with CUDA_VISIBLE_DEVICES pinned to one idle L40S.
The convolution speed-up is strongly architecture-dependent: the same
groups == 1 → groups == B*C change is 36x on an A40 (8.63 ms → 0.24 ms for a
per-sample batch of two) and 4.7x on an L40S, because the two pick different
cuDNN heuristics. The table above uses the conservative (L40S) figures.

🤖 Generated with Claude Code