From afd246b25f7373a04551062c3f1950c221a37059 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Mon, 17 Aug 2026 03:37:28 +0200 Subject: [PATCH 01/34] test: add GPU-free correctness tests for the window kernel index math The four kernels in swin_window_process_kernel.cu perform no arithmetic on tensor values -- each is a pure gather whose entire logic lives in the computation of input_offset from blockIdx. Transcribing that offset computation to vectorised PyTorch reproduces the kernels exactly on any device, so their correctness becomes testable without a GPU or a compiled extension. reference.py holds the transcription; test_index_math.py pins each kernel to the composition of PyTorch ops it replaces, checks that the backward kernels invert their forward, and documents two properties of the current CUDA source: - window_merge_and_roll_forward uses `* nH` as the stride between window rows, but windows are laid out row-major as b*nH*nW + wrow*nW + wcol, so the stride is nW. The two agree exactly when nH == nW, which is why no model in this repository has ever hit it. On non-square feature maps it is an out-of-bounds read when nH > nW and silent corruption when nH < nW. - the intra-window modulo omits `% H` / `% W`. This one is genuinely a no-op: the two forms differ by a multiple of H, and H is a multiple of window_h whenever `nH = H / window_h` is exact, which the launcher already requires. --- kernels/window_process/reference.py | 218 ++++++++++++++++++++++ kernels/window_process/test_index_math.py | 203 ++++++++++++++++++++ 2 files changed, 421 insertions(+) create mode 100644 kernels/window_process/reference.py create mode 100644 kernels/window_process/test_index_math.py diff --git a/kernels/window_process/reference.py b/kernels/window_process/reference.py new file mode 100644 index 000000000..ede121028 --- /dev/null +++ b/kernels/window_process/reference.py @@ -0,0 +1,218 @@ +# -------------------------------------------------------- +# Fused kernel for window process for SwinTransformer +# Copyright (c) 2022 Nvidia +# Licensed under The MIT License [see LICENSE for details] +# -------------------------------------------------------- +# Device-independent transcription of the index arithmetic used by the four +# CUDA kernels in swin_window_process_kernel.cu. +# +# The kernels perform no arithmetic on the tensor values: every one of them is a +# pure gather, and all of their logic lives in the computation of `input_offset` +# from `blockIdx`. Transcribing that offset computation to vectorised PyTorch +# therefore reproduces the kernels exactly, on any device, with no compiled +# extension and no GPU. This makes the correctness of the index math testable in +# CI without a GPU runner (see test_index_math.py). +# -------------------------------------------------------- + +import torch + + +def to_pair(value): + """Accept an int (isotropic) or a (h, w) iterable and return a (h, w) tuple.""" + if isinstance(value, int): + return value, value + h, w = value + return int(h), int(w) + + +def window_partition(x, window_size): + """(B, H, W, C) -> (B * nH * nW, window_h, window_w, C).""" + window_h, window_w = to_pair(window_size) + B, H, W, C = x.shape + x = x.view(B, H // window_h, window_h, W // window_w, window_w, C) + return x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_h, window_w, C) + + +def window_reverse(windows, window_size, H, W): + """(B * nH * nW, window_h, window_w, C) -> (B, H, W, C).""" + window_h, window_w = to_pair(window_size) + C = windows.shape[-1] + B = windows.shape[0] // ((H // window_h) * (W // window_w)) + x = windows.view(B, H // window_h, W // window_w, window_h, window_w, C) + return x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, C) + + +# -------------------------------------------------------------------------- +# Index maps +# +# Each function returns the flat index, into the *spatial* grid (B * H * W) or +# into the *window* grid (B * nH * nW * window_h * window_w), that every output +# element reads from. Channels are omitted: the kernels loop over C with a fixed +# stride of 1, so the channel dimension contributes an identical offset to every +# element and can be factored out by gathering over rows of a (-1, C) view. +# +# `shift_h` / `shift_w` follow the kernel convention, which is the *negated* +# torch.roll shift on the partition path — callers pass -shift_size to the +# forward kernels and +shift_size to the reverse ones. See the identities +# asserted in test_index_math.py. +# -------------------------------------------------------------------------- + + +def roll_and_window_partition_forward_index(B, H, W, shift_h, shift_w, window_h, window_w): + """K1: roll + window partition. Reads from the spatial grid. + + Returns: + (B * nH * nW, window_h, window_w) int64 tensor of flat indices into + a (B * H * W, C) view of the input. + """ + nH, nW = H // window_h, W // window_w + flat_win = torch.arange(B * nH * nW).view(-1, 1, 1) # blockIdx.z + ly = torch.arange(window_h).view(1, -1, 1) # blockIdx.y + lx = torch.arange(window_w).view(1, 1, -1) # blockIdx.x + + b = flat_win // (nH * nW) + row = (flat_win % (nH * nW) // nW * window_h + ly - shift_h + H) % H + col = (flat_win % nW * window_w + lx - shift_w + W) % W + return b * (H * W) + row * W + col + + +def roll_and_window_partition_backward_index(B, H, W, shift_h, shift_w, window_h, window_w): + """K2: backward of K1. Reads from the window grid. + + Returns: + (B, H, W) int64 tensor of flat indices into a + (B * nH * nW * window_h * window_w, C) view of the incoming gradient. + """ + nH, nW = H // window_h, W // window_w + b = torch.arange(B).view(-1, 1, 1) # blockIdx.z + y = torch.arange(H).view(1, -1, 1) # blockIdx.y + x = torch.arange(W).view(1, 1, -1) # blockIdx.x + + src_y = (y + shift_h + H) % H + src_x = (x + shift_w + W) % W + win = b * (nH * nW) + src_y // window_h * nW + src_x // window_w + return win * (window_h * window_w) + (src_y % window_h) * window_w + (src_x % window_w) + + +def window_merge_and_roll_forward_index( + B, H, W, shift_h, shift_w, window_h, window_w, + *, + legacy_row_stride=False, + legacy_intra_modulo=False, +): + """K3: window merge + reverse roll. Reads from the window grid. + + Args: + legacy_row_stride: reproduce the upstream ``* nH`` term. The stride + between consecutive window *rows* is nW (there are nW windows per + row in the row-major layout ``b * nH * nW + wrow * nW + wcol``), so + ``* nH`` is only correct when nH == nW, i.e. for square feature + maps. This is the bug fixed by this branch. + legacy_intra_modulo: reproduce the upstream intra-window modulo, which + omits the ``% H`` / ``% W`` reduction before ``% window_h`` / + ``% window_w``. This is equivalent to the explicit form whenever + H % window_h == 0 and W % window_w == 0, which the launcher already + requires; the flag exists to make that equivalence testable rather + than assumed. + + Returns: + (B, H, W) int64 tensor of flat indices into a + (B * nH * nW * window_h * window_w, C) view of the input. + """ + nH, nW = H // window_h, W // window_w + b = torch.arange(B).view(-1, 1, 1) # blockIdx.z + y = torch.arange(H).view(1, -1, 1) # blockIdx.y + x = torch.arange(W).view(1, 1, -1) # blockIdx.x + + src_y = (y - shift_h + H) % H + src_x = (x - shift_w + W) % W + + row_stride = nH if legacy_row_stride else nW + win = b * (nH * nW) + src_y // window_h * row_stride + src_x // window_w + + if legacy_intra_modulo: + intra_y = (y - shift_h + H) % window_h + intra_x = (x - shift_w + W) % window_w + else: + intra_y = src_y % window_h + intra_x = src_x % window_w + + return win * (window_h * window_w) + intra_y * window_w + intra_x + + +def window_merge_and_roll_backward_index(B, H, W, shift_h, shift_w, window_h, window_w): + """K4: backward of K3. Reads from the spatial grid. + + Returns: + (B * nH * nW, window_h, window_w) int64 tensor of flat indices into + a (B * H * W, C) view of the incoming gradient. + """ + nH, nW = H // window_h, W // window_w + flat_win = torch.arange(B * nH * nW).view(-1, 1, 1) # blockIdx.z + ly = torch.arange(window_h).view(1, -1, 1) # blockIdx.y + lx = torch.arange(window_w).view(1, 1, -1) # blockIdx.x + + b = flat_win // (nH * nW) + row = (flat_win % (nH * nW) // nW * window_h + ly + shift_h + H) % H + col = (flat_win % nW * window_w + lx + shift_w + W) % W + return b * (H * W) + row * W + col + + +# -------------------------------------------------------------------------- +# Gathers +# -------------------------------------------------------------------------- + + +def _gather(src, index, out_shape): + """Apply a flat index map to the leading dimensions of ``src``. + + Args: + src: tensor whose last dimension is C. + index: int64 tensor of flat indices into ``src.view(-1, C)``. + out_shape: spatial shape of the result, C appended by this function. + """ + C = src.shape[-1] + flat = src.reshape(-1, C) + return flat[index.reshape(-1)].view(*out_shape, C) + + +def roll_and_window_partition_forward(x, shift_size, window_size): + """Reference implementation of K1. ``x``: (B, H, W, C).""" + shift_h, shift_w = to_pair(shift_size) + window_h, window_w = to_pair(window_size) + B, H, W, _ = x.shape + index = roll_and_window_partition_forward_index( + B, H, W, shift_h, shift_w, window_h, window_w) + nH, nW = H // window_h, W // window_w + return _gather(x, index, (B * nH * nW, window_h, window_w)) + + +def roll_and_window_partition_backward(grad_in, shift_size, window_size, H, W): + """Reference implementation of K2. ``grad_in``: (B*nH*nW, window_h, window_w, C).""" + shift_h, shift_w = to_pair(shift_size) + window_h, window_w = to_pair(window_size) + B = grad_in.shape[0] // ((H // window_h) * (W // window_w)) + index = roll_and_window_partition_backward_index( + B, H, W, shift_h, shift_w, window_h, window_w) + return _gather(grad_in, index, (B, H, W)) + + +def window_merge_and_roll_forward(x, shift_size, window_size, H, W, **legacy): + """Reference implementation of K3. ``x``: (B*nH*nW, window_h, window_w, C).""" + shift_h, shift_w = to_pair(shift_size) + window_h, window_w = to_pair(window_size) + B = x.shape[0] // ((H // window_h) * (W // window_w)) + index = window_merge_and_roll_forward_index( + B, H, W, shift_h, shift_w, window_h, window_w, **legacy) + return _gather(x, index, (B, H, W)) + + +def window_merge_and_roll_backward(grad_in, shift_size, window_size): + """Reference implementation of K4. ``grad_in``: (B, H, W, C).""" + shift_h, shift_w = to_pair(shift_size) + window_h, window_w = to_pair(window_size) + B, H, W, _ = grad_in.shape + index = window_merge_and_roll_backward_index( + B, H, W, shift_h, shift_w, window_h, window_w) + nH, nW = H // window_h, W // window_w + return _gather(grad_in, index, (B * nH * nW, window_h, window_w)) diff --git a/kernels/window_process/test_index_math.py b/kernels/window_process/test_index_math.py new file mode 100644 index 000000000..b5f2b8aff --- /dev/null +++ b/kernels/window_process/test_index_math.py @@ -0,0 +1,203 @@ +# -------------------------------------------------------- +# Fused kernel for window process for SwinTransformer +# Copyright (c) 2022 Nvidia +# Licensed under The MIT License [see LICENSE for details] +# -------------------------------------------------------- +# Correctness tests for the index arithmetic of the fused window kernels. +# +# These run on CPU against reference.py and need neither a GPU nor the compiled +# extension, so they can gate the index math in CI. The GPU parity tests for the +# compiled kernels live in unit_test.py. +# +# python -m unittest kernels.window_process.test_index_math -v +# -------------------------------------------------------- + +import unittest + +import torch + +import reference as ref + + +# (H, W, window_h, window_w) -- covers square and non-square feature maps, and +# square and non-square windows. +SHAPES = [ + (8, 8, 4, 4), # nH == nW == 2 square grid, square window + (16, 8, 4, 4), # nH=4, nW=2 taller than wide + (8, 16, 4, 4), # nH=2, nW=4 wider than tall + (8, 32, 4, 8), # nH=2, nW=4 non-square window + (32, 8, 8, 4), # nH=4, nW=2 non-square window, taller than wide + (8, 8, 8, 8), # nH == nW == 1 single window + (12, 12, 4, 4), # nH == nW == 3 odd window count +] + +BATCH = 2 +CHANNELS = 3 + + +def _shifts(window_h, window_w): + """The two shift regimes used by Swin: W-MSA (no shift) and SW-MSA (half window).""" + return [(0, 0), (window_h // 2, window_w // 2)] + + +def _make_spatial(B, H, W, C): + torch.manual_seed(0) + return torch.arange(B * H * W * C, dtype=torch.float32).view(B, H, W, C) + + +def _make_windows(B, H, W, C, window_h, window_w): + torch.manual_seed(1) + n = B * (H // window_h) * (W // window_w) + return torch.arange(n * window_h * window_w * C, dtype=torch.float32).view( + n, window_h, window_w, C) + + +class TestKernelIndexMath(unittest.TestCase): + """Each kernel must equal the composition of PyTorch ops it replaces. + + The kernel `shift_h`/`shift_w` parameters are the negated torch.roll shift on + the partition path: callers pass -shift_size to the forward kernels and + +shift_size to the reverse ones. These identities pin that convention down. + """ + + def test_k1_equals_roll_then_partition(self): + for H, W, wh, ww in SHAPES: + for sh, sw in _shifts(wh, ww): + with self.subTest(H=H, W=W, wh=wh, ww=ww, shift=(sh, sw)): + x = _make_spatial(BATCH, H, W, CHANNELS) + got = ref.roll_and_window_partition_forward(x, (sh, sw), (wh, ww)) + want = ref.window_partition( + torch.roll(x, shifts=(sh, sw), dims=(1, 2)), (wh, ww)) + self.assertTrue(torch.equal(got, want)) + + def test_k2_equals_reverse_then_unroll(self): + for H, W, wh, ww in SHAPES: + for sh, sw in _shifts(wh, ww): + with self.subTest(H=H, W=W, wh=wh, ww=ww, shift=(sh, sw)): + g = _make_windows(BATCH, H, W, CHANNELS, wh, ww) + got = ref.roll_and_window_partition_backward(g, (sh, sw), (wh, ww), H, W) + want = torch.roll( + ref.window_reverse(g, (wh, ww), H, W), shifts=(-sh, -sw), dims=(1, 2)) + self.assertTrue(torch.equal(got, want)) + + def test_k3_equals_reverse_then_roll(self): + for H, W, wh, ww in SHAPES: + for sh, sw in _shifts(wh, ww): + with self.subTest(H=H, W=W, wh=wh, ww=ww, shift=(sh, sw)): + x = _make_windows(BATCH, H, W, CHANNELS, wh, ww) + got = ref.window_merge_and_roll_forward(x, (sh, sw), (wh, ww), H, W) + want = torch.roll( + ref.window_reverse(x, (wh, ww), H, W), shifts=(sh, sw), dims=(1, 2)) + self.assertTrue(torch.equal(got, want)) + + def test_k4_equals_unroll_then_partition(self): + for H, W, wh, ww in SHAPES: + for sh, sw in _shifts(wh, ww): + with self.subTest(H=H, W=W, wh=wh, ww=ww, shift=(sh, sw)): + g = _make_spatial(BATCH, H, W, CHANNELS) + got = ref.window_merge_and_roll_backward(g, (sh, sw), (wh, ww)) + want = ref.window_partition( + torch.roll(g, shifts=(-sh, -sw), dims=(1, 2)), (wh, ww)) + self.assertTrue(torch.equal(got, want)) + + def test_backward_kernels_invert_their_forward(self): + """K1/K2 and K3/K4 are permutations, so each pair must round-trip.""" + for H, W, wh, ww in SHAPES: + for sh, sw in _shifts(wh, ww): + with self.subTest(H=H, W=W, wh=wh, ww=ww, shift=(sh, sw)): + x = _make_spatial(BATCH, H, W, CHANNELS) + k1 = ref.roll_and_window_partition_forward(x, (sh, sw), (wh, ww)) + self.assertTrue(torch.equal( + ref.roll_and_window_partition_backward(k1, (sh, sw), (wh, ww), H, W), x)) + + w = _make_windows(BATCH, H, W, CHANNELS, wh, ww) + k3 = ref.window_merge_and_roll_forward(w, (sh, sw), (wh, ww), H, W) + self.assertTrue(torch.equal( + ref.window_merge_and_roll_backward(k3, (sh, sw), (wh, ww)), w)) + + +class TestUpstreamRowStrideBug(unittest.TestCase): + """window_merge_and_roll_forward used `* nH` where the row stride is `* nW`. + + Windows are laid out row-major as `b * nH * nW + wrow * nW + wcol`, so the + stride between consecutive window rows is nW. The two agree exactly when + nH == nW, which is why the bug never surfaced: every model in this repository + is trained on square images. + """ + + def _legacy_index(self, B, H, W, sh, sw, wh, ww): + return ref.window_merge_and_roll_forward_index( + B, H, W, sh, sw, wh, ww, legacy_row_stride=True) + + def test_legacy_is_correct_only_on_square_grids(self): + for H, W, wh, ww in SHAPES: + nH, nW = H // wh, W // ww + if nH != nW: + continue + for sh, sw in _shifts(wh, ww): + with self.subTest(H=H, W=W, wh=wh, ww=ww, shift=(sh, sw), nH=nH, nW=nW): + x = _make_windows(BATCH, H, W, CHANNELS, wh, ww) + want = ref.window_merge_and_roll_forward(x, (sh, sw), (wh, ww), H, W) + legacy = ref.window_merge_and_roll_forward( + x, (sh, sw), (wh, ww), H, W, legacy_row_stride=True) + self.assertTrue(torch.equal(legacy, want)) + + def test_legacy_reads_out_of_bounds_when_nH_greater_than_nW(self): + """nH > nW makes the miscomputed offset exceed the input, an illegal read.""" + checked = 0 + for H, W, wh, ww in SHAPES: + nH, nW = H // wh, W // ww + if nH <= nW: + continue + for sh, sw in _shifts(wh, ww): + with self.subTest(H=H, W=W, wh=wh, ww=ww, shift=(sh, sw), nH=nH, nW=nW): + numel = BATCH * nH * nW * wh * ww + legacy_max = int(self._legacy_index(BATCH, H, W, sh, sw, wh, ww).max()) + self.assertGreaterEqual(legacy_max, numel) + checked += 1 + self.assertGreater(checked, 0, "no nH > nW shape exercised") + + def test_legacy_returns_wrong_values_when_nH_less_than_nW(self): + """nH < nW keeps the offset in bounds, so the corruption is silent.""" + checked = 0 + for H, W, wh, ww in SHAPES: + nH, nW = H // wh, W // ww + if nH >= nW: + continue + for sh, sw in _shifts(wh, ww): + with self.subTest(H=H, W=W, wh=wh, ww=ww, shift=(sh, sw), nH=nH, nW=nW): + numel = BATCH * nH * nW * wh * ww + legacy_max = int(self._legacy_index(BATCH, H, W, sh, sw, wh, ww).max()) + self.assertLess(legacy_max, numel) + + x = _make_windows(BATCH, H, W, CHANNELS, wh, ww) + want = ref.window_merge_and_roll_forward(x, (sh, sw), (wh, ww), H, W) + legacy = ref.window_merge_and_roll_forward( + x, (sh, sw), (wh, ww), H, W, legacy_row_stride=True) + self.assertFalse(torch.equal(legacy, want)) + checked += 1 + self.assertGreater(checked, 0, "no nH < nW shape exercised") + + +class TestIntraWindowModuloIsEquivalent(unittest.TestCase): + """The upstream intra-window modulo omits `% H` / `% W`, which is a no-op here. + + `(y - s + H) % window_h` and `((y - s + H) % H) % window_h` differ by a + multiple of H, and H is a multiple of window_h whenever the launcher's + `nH = H / window_h` is exact -- which it must be for the kernel to be valid + at all. The explicit form is kept for readability, not for correctness. + """ + + def test_forms_agree_under_exact_divisibility(self): + for H, W, wh, ww in SHAPES: + for sh, sw in _shifts(wh, ww): + with self.subTest(H=H, W=W, wh=wh, ww=ww, shift=(sh, sw)): + explicit = ref.window_merge_and_roll_forward_index( + BATCH, H, W, sh, sw, wh, ww, legacy_intra_modulo=False) + legacy = ref.window_merge_and_roll_forward_index( + BATCH, H, W, sh, sw, wh, ww, legacy_intra_modulo=True) + self.assertTrue(torch.equal(explicit, legacy)) + + +if __name__ == '__main__': + unittest.main(verbosity=2) From cf5e949d8bbe4c9c0883568587735e56f8457261 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Mon, 17 Aug 2026 03:38:36 +0200 Subject: [PATCH 02/34] fix: use nW as the window-row stride in window_merge_and_roll_forward Windows are laid out row-major as `b * nH * nW + wrow * nW + wcol`, so the stride between consecutive window rows is nW, not nH. The other three kernels already use nW; this one did not. nH and nW are equal for every model in this repository, because they are all trained on square images, which is why this has never been observed. On a non-square feature map the miscomputed offset is: - out of bounds when nH > nW. For B=2, H=16, W=8, window=4 the kernel reads index 351 from a 256-element tensor: an illegal memory access. - in bounds but wrong when nH < nW, silently corrupting half the elements. Covered by TestUpstreamRowStrideBug in test_index_math.py. --- kernels/window_process/swin_window_process_kernel.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kernels/window_process/swin_window_process_kernel.cu b/kernels/window_process/swin_window_process_kernel.cu index 4976949ee..804b5423c 100644 --- a/kernels/window_process/swin_window_process_kernel.cu +++ b/kernels/window_process/swin_window_process_kernel.cu @@ -110,7 +110,7 @@ __global__ void window_merge_and_roll_forward_cuda_kernel( for (int i = index; i < C; i += blockDim.x) { offset = ((blockIdx.z * gridDim.y + blockIdx.y) * gridDim.x + blockIdx.x) * C + i; // C = blocksize int input_offset = - (blockIdx.z * nH * nW + (blockIdx.y - shift_size + H) % H / window_size * nH + (blockIdx.x - shift_size + W) % W / window_size) * window_size * window_size * C + + (blockIdx.z * nH * nW + (blockIdx.y - shift_size + H) % H / window_size * nW + (blockIdx.x - shift_size + W) % W / window_size) * window_size * window_size * C + (blockIdx.y - shift_size + H) % window_size * window_size * C + (blockIdx.x - shift_size + W) % window_size * C + i; From 568ab74014abd60a52b7884cbfb121b3cbab9a05 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Mon, 17 Aug 2026 03:40:50 +0200 Subject: [PATCH 03/34] feat: support non-square windows and per-axis shifts Splits `shift_size` and `window_size` into `shift_h`/`shift_w` and `window_h`/`window_w` throughout the kernels and their launchers, so the fused path is no longer restricted to square windows on square feature maps. Grids become dim3(window_w, window_h, ...) -- x indexes columns, y indexes rows -- and nH/nW are derived per axis. The Python signature is deliberately unchanged: WindowProcess still takes seven arguments, and `shift_size`/`window_size` now accept either an int (isotropic, the previous behaviour) or a (h, w) pair. Both call sites in models/swin_transformer.py keep working untouched. unit_test.py re-declared WindowProcess and WindowProcessReverse instead of importing them, so it called the extension directly and would not have been updated by the change above -- it would have passed seven arguments to a nine-argument function. It now imports the two Function classes, which both fixes that and removes the duplication that allowed the two copies to drift. Two incidental fixes in window_process.py, both visible in the diff: - backward() now calls .contiguous() on the incoming gradient. The C++ side asserts contiguity via CHECK_INPUT, and an upstream op can hand us a non-contiguous gradient. - backward() returned eight values for seven forward inputs. Now seven. window_merge_and_roll_forward also gains the explicit `% H` / `% W` before the intra-window modulo, so it mirrors the structure of its backward counterpart. This is a readability change, not a fix: the two forms are equal whenever H % window_h == 0, which the launcher already requires. Asserted by TestIntraWindowModuloIsEquivalent in test_index_math.py. --- .../window_process/swin_window_process.cpp | 56 ++++-- .../swin_window_process_kernel.cu | 180 +++++++++++------- kernels/window_process/unit_test.py | 54 +----- kernels/window_process/window_process.py | 93 ++++++--- 4 files changed, 214 insertions(+), 169 deletions(-) diff --git a/kernels/window_process/swin_window_process.cpp b/kernels/window_process/swin_window_process.cpp index a7f15b048..24100920c 100644 --- a/kernels/window_process/swin_window_process.cpp +++ b/kernels/window_process/swin_window_process.cpp @@ -25,8 +25,10 @@ at::Tensor roll_and_window_partition_forward_cuda( const int H, const int W, const int C, - const int shift_size, - const int window_size); + const int shift_h, + const int shift_w, + const int window_h, + const int window_w); at::Tensor roll_and_window_partition_backward_cuda( @@ -36,8 +38,10 @@ at::Tensor roll_and_window_partition_backward_cuda( const int H, const int W, const int C, - const int shift_size, - const int window_size); + const int shift_h, + const int shift_w, + const int window_h, + const int window_w); at::Tensor window_merge_and_roll_forward_cuda( @@ -47,8 +51,10 @@ at::Tensor window_merge_and_roll_forward_cuda( const int H, const int W, const int C, - const int shift_size, - const int window_size); + const int shift_h, + const int shift_w, + const int window_h, + const int window_w); at::Tensor window_merge_and_roll_backward_cuda( at::Tensor & grad_in, @@ -57,8 +63,10 @@ at::Tensor window_merge_and_roll_backward_cuda( const int H, const int W, const int C, - const int shift_size, - const int window_size); + const int shift_h, + const int shift_w, + const int window_h, + const int window_w); #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor") @@ -74,10 +82,12 @@ at::Tensor roll_and_window_partition_forward( const int H, const int W, const int C, - const int shift_size, - const int window_size){ + const int shift_h, + const int shift_w, + const int window_h, + const int window_w){ CHECK_INPUT(input); - return roll_and_window_partition_forward_cuda(input, B, H, W, C, shift_size, window_size); + return roll_and_window_partition_forward_cuda(input, B, H, W, C, shift_h, shift_w, window_h, window_w); } @@ -88,10 +98,12 @@ at::Tensor roll_and_window_partition_backward( const int H, const int W, const int C, - const int shift_size, - const int window_size){ + const int shift_h, + const int shift_w, + const int window_h, + const int window_w){ CHECK_INPUT(grad_in); - return roll_and_window_partition_backward_cuda(grad_in, B, H, W, C, shift_size, window_size); + return roll_and_window_partition_backward_cuda(grad_in, B, H, W, C, shift_h, shift_w, window_h, window_w); } @@ -102,10 +114,12 @@ at::Tensor window_merge_and_roll_forward( const int H, const int W, const int C, - const int shift_size, - const int window_size){ + const int shift_h, + const int shift_w, + const int window_h, + const int window_w){ CHECK_INPUT(input); - return window_merge_and_roll_forward_cuda(input, B, H, W, C, shift_size, window_size); + return window_merge_and_roll_forward_cuda(input, B, H, W, C, shift_h, shift_w, window_h, window_w); } @@ -116,10 +130,12 @@ at::Tensor window_merge_and_roll_backward( const int H, const int W, const int C, - const int shift_size, - const int window_size){ + const int shift_h, + const int shift_w, + const int window_h, + const int window_w){ CHECK_INPUT(grad_in); - return window_merge_and_roll_backward_cuda(grad_in, B, H, W, C, shift_size, window_size); + return window_merge_and_roll_backward_cuda(grad_in, B, H, W, C, shift_h, shift_w, window_h, window_w); } diff --git a/kernels/window_process/swin_window_process_kernel.cu b/kernels/window_process/swin_window_process_kernel.cu index 804b5423c..1e3d28444 100644 --- a/kernels/window_process/swin_window_process_kernel.cu +++ b/kernels/window_process/swin_window_process_kernel.cu @@ -38,16 +38,21 @@ int best_block_dim(int feat_dim){ } +// input: [B, H, W, C] +// output: [B*nH*nW, window_h, window_w, C] +// grid: (window_w, window_h, B*nH*nW) -- x indexes columns, y indexes rows template __global__ void roll_and_window_partition_forward_cuda_kernel( - T* input, - T* output, + T* input, + T* output, const int B, const int H, const int W, const int C, - const int shift_size, - const int window_size, + const int shift_h, + const int shift_w, + const int window_h, + const int window_w, const int nH, const int nW){ // start @@ -57,24 +62,29 @@ __global__ void roll_and_window_partition_forward_cuda_kernel( for (int i = index; i < C; i += blockDim.x) { offset = ((blockIdx.z * gridDim.y + blockIdx.y) * gridDim.x + blockIdx.x) * C + i; // C = blocksize int input_offset = blockIdx.z / (nH * nW) * H * W * C + - (blockIdx.z % (nH * nW) / nW * window_size + blockIdx.y - shift_size + H) % H * W * C + - (blockIdx.z % nW * window_size + blockIdx.x - shift_size + W) % W * C + + (blockIdx.z % (nH * nW) / nW * window_h + blockIdx.y - shift_h + H) % H * W * C + + (blockIdx.z % nW * window_w + blockIdx.x - shift_w + W) % W * C + i; output[offset] = (T)(__ldg(input + input_offset)); } } +// grad_in: [B*nH*nW, window_h, window_w, C] +// grad_out: [B, H, W, C] +// grid: (W, H, B) template __global__ void roll_and_window_partition_backward_cuda_kernel( - T* grad_in, - T* grad_out, + T* grad_in, + T* grad_out, const int B, const int H, const int W, const int C, - const int shift_size, - const int window_size, + const int shift_h, + const int shift_w, + const int window_h, + const int window_w, const int nH, const int nW){ // start @@ -82,26 +92,31 @@ __global__ void roll_and_window_partition_backward_cuda_kernel( int offset; for (int i = index; i < C; i += blockDim.x) { offset = ((blockIdx.z * gridDim.y + blockIdx.y) * gridDim.x + blockIdx.x) * C + i; // C = blocksize - int input_offset = - (blockIdx.z * nH * nW + (blockIdx.y + shift_size + H) % H / window_size * nW + (blockIdx.x + shift_size + W) % W / window_size) * window_size * window_size * C + - (blockIdx.y + shift_size + H ) % H % window_size * window_size * C + - (blockIdx.x + shift_size + W ) % W % window_size * C + + int input_offset = + (blockIdx.z * nH * nW + (blockIdx.y + shift_h + H) % H / window_h * nW + (blockIdx.x + shift_w + W) % W / window_w) * window_h * window_w * C + + (blockIdx.y + shift_h + H ) % H % window_h * window_w * C + + (blockIdx.x + shift_w + W ) % W % window_w * C + i; grad_out[offset] = (T)(__ldg(grad_in + input_offset)); } } +// input: [B*nH*nW, window_h, window_w, C] +// output: [B, H, W, C] +// grid: (W, H, B) template __global__ void window_merge_and_roll_forward_cuda_kernel( - T* input, - T* output, + T* input, + T* output, const int B, const int H, const int W, const int C, - const int shift_size, - const int window_size, + const int shift_h, + const int shift_w, + const int window_h, + const int window_w, const int nH, const int nW){ // start @@ -109,10 +124,10 @@ __global__ void window_merge_and_roll_forward_cuda_kernel( int offset; for (int i = index; i < C; i += blockDim.x) { offset = ((blockIdx.z * gridDim.y + blockIdx.y) * gridDim.x + blockIdx.x) * C + i; // C = blocksize - int input_offset = - (blockIdx.z * nH * nW + (blockIdx.y - shift_size + H) % H / window_size * nW + (blockIdx.x - shift_size + W) % W / window_size) * window_size * window_size * C + - (blockIdx.y - shift_size + H) % window_size * window_size * C + - (blockIdx.x - shift_size + W) % window_size * C + + int input_offset = + (blockIdx.z * nH * nW + (blockIdx.y - shift_h + H) % H / window_h * nW + (blockIdx.x - shift_w + W) % W / window_w) * window_h * window_w * C + + (blockIdx.y - shift_h + H) % H % window_h * window_w * C + + (blockIdx.x - shift_w + W) % W % window_w * C + i; output[offset] = (T)(__ldg(input + input_offset)); } @@ -120,16 +135,21 @@ __global__ void window_merge_and_roll_forward_cuda_kernel( +// grad_in: [B, H, W, C] +// grad_out: [B*nH*nW, window_h, window_w, C] +// grid: (window_w, window_h, B*nH*nW) template __global__ void window_merge_and_roll_backward_cuda_kernel( - T* grad_in, - T* grad_out, + T* grad_in, + T* grad_out, const int B, const int H, const int W, const int C, - const int shift_size, - const int window_size, + const int shift_h, + const int shift_w, + const int window_h, + const int window_w, const int nH, const int nW){ // start @@ -137,41 +157,43 @@ __global__ void window_merge_and_roll_backward_cuda_kernel( int offset; for (int i = index; i < C; i += blockDim.x) { offset = ((blockIdx.z * gridDim.y + blockIdx.y) * gridDim.x + blockIdx.x) * C + i; // C = blocksize - int input_offset = + int input_offset = (blockIdx.z / (nH * nW)) * H * W * C + - (blockIdx.z % (nH * nW) / nW * window_size + blockIdx.y + shift_size + H) % H * W * C + - (blockIdx.z % nW * window_size + blockIdx.x + shift_size + W) % W * C + + (blockIdx.z % (nH * nW) / nW * window_h + blockIdx.y + shift_h + H) % H * W * C + + (blockIdx.z % nW * window_w + blockIdx.x + shift_w + W) % W * C + i; grad_out[offset] = (T)(__ldg(grad_in + input_offset)); } } // input: [B, H, W, C] -// output: [B*nH*nW, window_size, window_size, C] +// output: [B*nH*nW, window_h, window_w, C] at::Tensor roll_and_window_partition_forward_cuda( - at::Tensor & input, + at::Tensor & input, //at::Tensor & output, const int B, const int H, const int W, const int C, - const int shift_size, - const int window_size){ - - int nH = H / window_size; - int nW = W / window_size; + const int shift_h, + const int shift_w, + const int window_h, + const int window_w){ - dim3 grid(window_size, window_size, B * nH * nW); + int nH = H / window_h; + int nW = W / window_w; + + dim3 grid(window_w, window_h, B * nH * nW); //dim3 block((C + 31) / 32 * 32); int blocknum = best_block_dim(C); dim3 block(blocknum); at::Tensor output; if (input.scalar_type() == torch::kFloat16){ - output = torch::empty({B*nH*nW, window_size, window_size, C}, torch::dtype(torch::kFloat16).device(torch::kCUDA).requires_grad(true)); + output = torch::empty({B*nH*nW, window_h, window_w, C}, torch::dtype(torch::kFloat16).device(torch::kCUDA).requires_grad(true)); } else{ - output = torch::empty({B*nH*nW, window_size, window_size, C}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(true)); + output = torch::empty({B*nH*nW, window_h, window_w, C}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(true)); } AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.type(), "roll_and_window_partition_forward_cuda_kernel", ([&] { @@ -182,8 +204,10 @@ at::Tensor roll_and_window_partition_forward_cuda( H, W, C, - shift_size, - window_size, + shift_h, + shift_w, + window_h, + window_w, nH, nW); })); @@ -191,19 +215,21 @@ at::Tensor roll_and_window_partition_forward_cuda( } -// grad_in: [B*nH*nW, window_size, window_size, C] +// grad_in: [B*nH*nW, window_h, window_w, C] // grad_out: [B, H, W, C] at::Tensor roll_and_window_partition_backward_cuda( - at::Tensor & grad_in, + at::Tensor & grad_in, const int B, const int H, const int W, const int C, - const int shift_size, - const int window_size){ - - int nH = H / window_size; - int nW = W / window_size; + const int shift_h, + const int shift_w, + const int window_h, + const int window_w){ + + int nH = H / window_h; + int nW = W / window_w; dim3 grid(W, H, B); //dim3 block((C + 31) / 32 * 32); @@ -226,8 +252,10 @@ at::Tensor roll_and_window_partition_backward_cuda( H, W, C, - shift_size, - window_size, + shift_h, + shift_w, + window_h, + window_w, nH, nW); })); @@ -235,20 +263,22 @@ at::Tensor roll_and_window_partition_backward_cuda( } -// input: [B*nH*nW, window_size, window_size, C] +// input: [B*nH*nW, window_h, window_w, C] // output: [B, H, W, C] at::Tensor window_merge_and_roll_forward_cuda( - at::Tensor & input, + at::Tensor & input, //at::Tensor & output, const int B, const int H, const int W, const int C, - const int shift_size, - const int window_size){ - - int nH = H / window_size; - int nW = W / window_size; + const int shift_h, + const int shift_w, + const int window_h, + const int window_w){ + + int nH = H / window_h; + int nW = W / window_w; dim3 grid(W, H, B); //dim3 block((C + 31) / 32 * 32); @@ -272,8 +302,10 @@ at::Tensor window_merge_and_roll_forward_cuda( H, W, C, - shift_size, - window_size, + shift_h, + shift_w, + window_h, + window_w, nH, nW); })); @@ -281,29 +313,33 @@ at::Tensor window_merge_and_roll_forward_cuda( } +// grad_in: [B, H, W, C] +// grad_out: [B*nH*nW, window_h, window_w, C] at::Tensor window_merge_and_roll_backward_cuda( - at::Tensor & grad_in, + at::Tensor & grad_in, const int B, const int H, const int W, const int C, - const int shift_size, - const int window_size){ - - int nH = H / window_size; - int nW = W / window_size; + const int shift_h, + const int shift_w, + const int window_h, + const int window_w){ + + int nH = H / window_h; + int nW = W / window_w; - dim3 grid(window_size, window_size, B * nH * nW); + dim3 grid(window_w, window_h, B * nH * nW); //dim3 block((C + 31) / 32 * 32); int blocknum = best_block_dim(C); dim3 block(blocknum); at::Tensor grad_out; if (grad_in.scalar_type() == torch::kFloat16){ - grad_out = torch::empty({B*nH*nW, window_size, window_size, C}, torch::dtype(torch::kFloat16).device(torch::kCUDA).requires_grad(false)); + grad_out = torch::empty({B*nH*nW, window_h, window_w, C}, torch::dtype(torch::kFloat16).device(torch::kCUDA).requires_grad(false)); } else{ - grad_out = torch::empty({B*nH*nW, window_size, window_size, C}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(false)); + grad_out = torch::empty({B*nH*nW, window_h, window_w, C}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(false)); } AT_DISPATCH_FLOATING_TYPES_AND_HALF(grad_in.type(), "window_merge_and_roll_backward_cuda_kernel", ([&] { @@ -314,10 +350,12 @@ at::Tensor window_merge_and_roll_backward_cuda( H, W, C, - shift_size, - window_size, + shift_h, + shift_w, + window_h, + window_w, nH, nW); })); return grad_out; -} \ No newline at end of file +} diff --git a/kernels/window_process/unit_test.py b/kernels/window_process/unit_test.py index 65dee5661..df35600b5 100644 --- a/kernels/window_process/unit_test.py +++ b/kernels/window_process/unit_test.py @@ -5,63 +5,11 @@ # -------------------------------------------------------- import torch -import swin_window_process import random import time import unittest - -class WindowProcess(torch.autograd.Function): - @staticmethod - def forward(ctx, input, B, H, W, C, shift_size, window_size): - output = swin_window_process.roll_and_window_partition_forward(input, B, H, W, C, shift_size, window_size) - - ctx.B = B - ctx.H = H - ctx.W = W - ctx.C = C - ctx.shift_size = shift_size - ctx.window_size = window_size - return output - - @staticmethod - def backward(ctx, grad_in): - B = ctx.B - H = ctx.H - W = ctx.W - C = ctx.C - shift_size = ctx.shift_size - window_size = ctx.window_size - - grad_out = swin_window_process.roll_and_window_partition_backward(grad_in, B, H, W, C, shift_size, window_size) - return grad_out, None, None, None, None, None, None, None - - -class WindowProcessReverse(torch.autograd.Function): - @staticmethod - def forward(ctx, input, B, H, W, C, shift_size, window_size): - output = swin_window_process.window_merge_and_roll_forward(input, B, H, W, C, shift_size, window_size) - - ctx.B = B - ctx.H = H - ctx.W = W - ctx.C = C - ctx.shift_size = shift_size - ctx.window_size = window_size - - return output - - @staticmethod - def backward(ctx, grad_in): - B = ctx.B - H = ctx.H - W = ctx.W - C = ctx.C - shift_size = ctx.shift_size - window_size = ctx.window_size - - grad_out = swin_window_process.window_merge_and_roll_backward(grad_in, B, H, W, C, shift_size, window_size) - return grad_out, None, None, None, None, None, None, None +from window_process import WindowProcess, WindowProcessReverse def window_partition(x, window_size): diff --git a/kernels/window_process/window_process.py b/kernels/window_process/window_process.py index ee43e9e97..3fbb2ff63 100644 --- a/kernels/window_process/window_process.py +++ b/kernels/window_process/window_process.py @@ -8,43 +8,87 @@ import swin_window_process +def to_pair(value): + """Accept an int (isotropic) or a (h, w) iterable and return a (h, w) tuple. + + Keeps the public signature of WindowProcess unchanged: `window_size=7` + behaves exactly as before, `window_size=(4, 8)` selects a rectangular window. + """ + if isinstance(value, int): + return value, value + h, w = value + return int(h), int(w) + + class WindowProcess(torch.autograd.Function): + """Fused torch.roll + window_partition. + + Args: + input: (B, H, W, C) contiguous float32/float16 CUDA tensor. + B, H, W, C: shape of `input`. + shift_size: int or (shift_h, shift_w). This is the *negated* torch.roll + shift, matching the existing call sites in models/swin_transformer.py. + window_size: int or (window_h, window_w). Must divide H and W per axis. + + Returns: + (B * nH * nW, window_h, window_w, C) + """ + @staticmethod def forward(ctx, input, B, H, W, C, shift_size, window_size): - output = swin_window_process.roll_and_window_partition_forward(input, B, H, W, C, shift_size, window_size) + shift_h, shift_w = to_pair(shift_size) + window_h, window_w = to_pair(window_size) + output = swin_window_process.roll_and_window_partition_forward( + input, B, H, W, C, shift_h, shift_w, window_h, window_w) ctx.B = B ctx.H = H - ctx.W = W - ctx.C = C - ctx.shift_size = shift_size - ctx.window_size = window_size + ctx.W = W + ctx.C = C + ctx.shift_size = (shift_h, shift_w) + ctx.window_size = (window_h, window_w) return output @staticmethod def backward(ctx, grad_in): B = ctx.B H = ctx.H - W = ctx.W - C = ctx.C - shift_size = ctx.shift_size - window_size = ctx.window_size + W = ctx.W + C = ctx.C + shift_h, shift_w = ctx.shift_size + window_h, window_w = ctx.window_size - grad_out = swin_window_process.roll_and_window_partition_backward(grad_in, B, H, W, C, shift_size, window_size) - return grad_out, None, None, None, None, None, None, None + grad_out = swin_window_process.roll_and_window_partition_backward( + grad_in.contiguous(), B, H, W, C, shift_h, shift_w, window_h, window_w) + return grad_out, None, None, None, None, None, None class WindowProcessReverse(torch.autograd.Function): + """Fused window merge + torch.roll, the inverse of WindowProcess. + + Args: + input: (B * nH * nW, window_h, window_w, C) contiguous CUDA tensor. + B, H, W, C: shape of the *output* feature map. + shift_size: int or (shift_h, shift_w), the torch.roll shift. + window_size: int or (window_h, window_w). Must divide H and W per axis. + + Returns: + (B, H, W, C) + """ + @staticmethod def forward(ctx, input, B, H, W, C, shift_size, window_size): - output = swin_window_process.window_merge_and_roll_forward(input, B, H, W, C, shift_size, window_size) + shift_h, shift_w = to_pair(shift_size) + window_h, window_w = to_pair(window_size) + output = swin_window_process.window_merge_and_roll_forward( + input, B, H, W, C, shift_h, shift_w, window_h, window_w) ctx.B = B ctx.H = H - ctx.W = W - ctx.C = C - ctx.shift_size = shift_size - ctx.window_size = window_size + ctx.W = W + ctx.C = C + ctx.shift_size = (shift_h, shift_w) + ctx.window_size = (window_h, window_w) return output @@ -52,12 +96,11 @@ def forward(ctx, input, B, H, W, C, shift_size, window_size): def backward(ctx, grad_in): B = ctx.B H = ctx.H - W = ctx.W - C = ctx.C - shift_size = ctx.shift_size - window_size = ctx.window_size - - #grad_out = ctx.saved_tensors[0] - #grad_out = torch.zeros((B, H, W, C), dtype=dtype).cuda() - grad_out = swin_window_process.window_merge_and_roll_backward(grad_in, B, H, W, C, shift_size, window_size) - return grad_out, None, None, None, None, None, None, None + W = ctx.W + C = ctx.C + shift_h, shift_w = ctx.shift_size + window_h, window_w = ctx.window_size + + grad_out = swin_window_process.window_merge_and_roll_backward( + grad_in.contiguous(), B, H, W, C, shift_h, shift_w, window_h, window_w) + return grad_out, None, None, None, None, None, None From 46d40af953e64b746d26a010c882bd33d2e7fbf2 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Mon, 17 Aug 2026 03:41:48 +0200 Subject: [PATCH 04/34] refactor: replace deprecated ATen APIs All four are marked C10_DEPRECATED_MESSAGE in torch 2.7 and still compile, so this is a warning cleanup rather than a build fix -- except for the allocation change, which is a real one. Tensor::type() -> Tensor::scalar_type() (inside AT_DISPATCH) Tensor::data() -> Tensor::data_ptr() x.type().is_cuda() -> x.is_cuda() (CHECK_CUDA) Output tensors were allocated by branching on kFloat16 and hardcoding torch::kFloat32 otherwise, on torch::kCUDA with no device index, with requires_grad(true). Deriving them from the input's options() instead: - preserves the input dtype rather than silently widening anything that is not fp16 to fp32, which is what unblocks the bfloat16 support added next; - keeps the output on the input's device instead of the ambient current device; - drops requires_grad(true), which has no meaning on a tensor created inside an autograd.Function forward -- autograd builds the graph from the Function itself. torch/extension.h is dropped from the .cu: it was only needed for torch::empty/torch::dtype and pulls pybind11 through nvcc for no reason. ATen/Dispatch.h is now included explicitly rather than relied on transitively through ATen/ATen.h. --- .../window_process/swin_window_process.cpp | 2 +- .../swin_window_process_kernel.cu | 59 ++++++------------- 2 files changed, 18 insertions(+), 43 deletions(-) diff --git a/kernels/window_process/swin_window_process.cpp b/kernels/window_process/swin_window_process.cpp index 24100920c..daf10629c 100644 --- a/kernels/window_process/swin_window_process.cpp +++ b/kernels/window_process/swin_window_process.cpp @@ -69,7 +69,7 @@ at::Tensor window_merge_and_roll_backward_cuda( const int window_w); -#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor") +#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor") #define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous") #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) diff --git a/kernels/window_process/swin_window_process_kernel.cu b/kernels/window_process/swin_window_process_kernel.cu index 1e3d28444..da335f558 100644 --- a/kernels/window_process/swin_window_process_kernel.cu +++ b/kernels/window_process/swin_window_process_kernel.cu @@ -15,10 +15,10 @@ */ #include +#include #include #include #include -#include #include int best_block_dim(int feat_dim){ @@ -188,18 +188,12 @@ at::Tensor roll_and_window_partition_forward_cuda( int blocknum = best_block_dim(C); dim3 block(blocknum); - at::Tensor output; - if (input.scalar_type() == torch::kFloat16){ - output = torch::empty({B*nH*nW, window_h, window_w, C}, torch::dtype(torch::kFloat16).device(torch::kCUDA).requires_grad(true)); - } - else{ - output = torch::empty({B*nH*nW, window_h, window_w, C}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(true)); - } + at::Tensor output = at::empty({B*nH*nW, window_h, window_w, C}, input.options()); - AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.type(), "roll_and_window_partition_forward_cuda_kernel", ([&] { + AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "roll_and_window_partition_forward_cuda_kernel", ([&] { roll_and_window_partition_forward_cuda_kernel<<>>( - input.data(), - output.data(), + input.data_ptr(), + output.data_ptr(), B, H, W, @@ -236,18 +230,12 @@ at::Tensor roll_and_window_partition_backward_cuda( int blocknum = best_block_dim(C); dim3 block(blocknum); - at::Tensor grad_out; - if (grad_in.scalar_type() == torch::kFloat16){ - grad_out = torch::empty({B, H, W, C}, torch::dtype(torch::kFloat16).device(torch::kCUDA).requires_grad(false)); - } - else{ - grad_out = torch::empty({B, H, W, C}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(false)); - } + at::Tensor grad_out = at::empty({B, H, W, C}, grad_in.options()); - AT_DISPATCH_FLOATING_TYPES_AND_HALF(grad_in.type(), "roll_and_window_partition_backward_cuda_kernel", ([&] { + AT_DISPATCH_FLOATING_TYPES_AND_HALF(grad_in.scalar_type(), "roll_and_window_partition_backward_cuda_kernel", ([&] { roll_and_window_partition_backward_cuda_kernel<<>>( - grad_in.data(), - grad_out.data(), + grad_in.data_ptr(), + grad_out.data_ptr(), B, H, W, @@ -285,19 +273,12 @@ at::Tensor window_merge_and_roll_forward_cuda( int blocknum = best_block_dim(C); dim3 block(blocknum); - //generate output tensor inside - at::Tensor output; - if (input.scalar_type() == torch::kFloat16){ - output = torch::empty({B, H, W, C}, torch::dtype(torch::kFloat16).device(torch::kCUDA).requires_grad(true)); - } - else{ - output = torch::empty({B, H, W, C}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(true)); - } + at::Tensor output = at::empty({B, H, W, C}, input.options()); - AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.type(), "window_merge_and_roll_forward_cuda_kernel", ([&] { + AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "window_merge_and_roll_forward_cuda_kernel", ([&] { window_merge_and_roll_forward_cuda_kernel<<>>( - input.data(), - output.data(), + input.data_ptr(), + output.data_ptr(), B, H, W, @@ -334,18 +315,12 @@ at::Tensor window_merge_and_roll_backward_cuda( int blocknum = best_block_dim(C); dim3 block(blocknum); - at::Tensor grad_out; - if (grad_in.scalar_type() == torch::kFloat16){ - grad_out = torch::empty({B*nH*nW, window_h, window_w, C}, torch::dtype(torch::kFloat16).device(torch::kCUDA).requires_grad(false)); - } - else{ - grad_out = torch::empty({B*nH*nW, window_h, window_w, C}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(false)); - } + at::Tensor grad_out = at::empty({B*nH*nW, window_h, window_w, C}, grad_in.options()); - AT_DISPATCH_FLOATING_TYPES_AND_HALF(grad_in.type(), "window_merge_and_roll_backward_cuda_kernel", ([&] { + AT_DISPATCH_FLOATING_TYPES_AND_HALF(grad_in.scalar_type(), "window_merge_and_roll_backward_cuda_kernel", ([&] { window_merge_and_roll_backward_cuda_kernel<<>>( - grad_in.data(), - grad_out.data(), + grad_in.data_ptr(), + grad_out.data_ptr(), B, H, W, From c982fa1c982b41ff738f57522aac17ee788f357c Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Mon, 17 Aug 2026 03:42:21 +0200 Subject: [PATCH 05/34] fix: launch on the current CUDA stream and check for launch errors The kernels were launched with an implicit stream argument of 0, i.e. the default stream, while the surrounding PyTorch ops run on whatever stream is current. These usually coincide, but not always, and when they diverge there is no ordering guarantee between the kernel and the ops around it: - DDP overlaps the backward pass with NCCL all-reduce on side streams; - user code under `with torch.cuda.stream(s)`; - CUDA graph capture, which rejects the default stream outright, so torch.compile(mode="reduce-overhead") cannot capture a model using the fused path. Also adds C10_CUDA_KERNEL_LAUNCH_CHECK() after each launch. Without it a failed launch -- an invalid configuration, or a grid.z above 65535 -- stays silent until the next synchronisation point, which reports it with an unrelated stack trace. --- .../window_process/swin_window_process_kernel.cu | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/kernels/window_process/swin_window_process_kernel.cu b/kernels/window_process/swin_window_process_kernel.cu index da335f558..3c4656610 100644 --- a/kernels/window_process/swin_window_process_kernel.cu +++ b/kernels/window_process/swin_window_process_kernel.cu @@ -16,6 +16,8 @@ #include #include +#include +#include #include #include #include @@ -191,7 +193,7 @@ at::Tensor roll_and_window_partition_forward_cuda( at::Tensor output = at::empty({B*nH*nW, window_h, window_w, C}, input.options()); AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "roll_and_window_partition_forward_cuda_kernel", ([&] { - roll_and_window_partition_forward_cuda_kernel<<>>( + roll_and_window_partition_forward_cuda_kernel<<>>( input.data_ptr(), output.data_ptr(), B, @@ -205,6 +207,7 @@ at::Tensor roll_and_window_partition_forward_cuda( nH, nW); })); + C10_CUDA_KERNEL_LAUNCH_CHECK(); return output; } @@ -233,7 +236,7 @@ at::Tensor roll_and_window_partition_backward_cuda( at::Tensor grad_out = at::empty({B, H, W, C}, grad_in.options()); AT_DISPATCH_FLOATING_TYPES_AND_HALF(grad_in.scalar_type(), "roll_and_window_partition_backward_cuda_kernel", ([&] { - roll_and_window_partition_backward_cuda_kernel<<>>( + roll_and_window_partition_backward_cuda_kernel<<>>( grad_in.data_ptr(), grad_out.data_ptr(), B, @@ -247,6 +250,7 @@ at::Tensor roll_and_window_partition_backward_cuda( nH, nW); })); + C10_CUDA_KERNEL_LAUNCH_CHECK(); return grad_out; } @@ -276,7 +280,7 @@ at::Tensor window_merge_and_roll_forward_cuda( at::Tensor output = at::empty({B, H, W, C}, input.options()); AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "window_merge_and_roll_forward_cuda_kernel", ([&] { - window_merge_and_roll_forward_cuda_kernel<<>>( + window_merge_and_roll_forward_cuda_kernel<<>>( input.data_ptr(), output.data_ptr(), B, @@ -290,6 +294,7 @@ at::Tensor window_merge_and_roll_forward_cuda( nH, nW); })); + C10_CUDA_KERNEL_LAUNCH_CHECK(); return output; } @@ -318,7 +323,7 @@ at::Tensor window_merge_and_roll_backward_cuda( at::Tensor grad_out = at::empty({B*nH*nW, window_h, window_w, C}, grad_in.options()); AT_DISPATCH_FLOATING_TYPES_AND_HALF(grad_in.scalar_type(), "window_merge_and_roll_backward_cuda_kernel", ([&] { - window_merge_and_roll_backward_cuda_kernel<<>>( + window_merge_and_roll_backward_cuda_kernel<<>>( grad_in.data_ptr(), grad_out.data_ptr(), B, @@ -332,5 +337,6 @@ at::Tensor window_merge_and_roll_backward_cuda( nH, nW); })); + C10_CUDA_KERNEL_LAUNCH_CHECK(); return grad_out; } From f51e9fa6bd3bf474a2590b017f9b995a452ed04c Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Mon, 17 Aug 2026 03:42:45 +0200 Subject: [PATCH 06/34] feat: add bfloat16 support AT_DISPATCH_FLOATING_TYPES_AND_HALF covers float64, float32 and float16. bfloat16 is now the default autocast dtype for transformer training on Ampere and later, so the fused path was unreachable for it: a caller had to cast to float32, run the kernel, and cast back. That round trip costs two allocations and two full passes over the tensor -- exactly the two passes the fused kernel exists to eliminate, since it replaces a torch.roll (one copy) plus a window_partition (permute + contiguous, another copy). The workaround therefore cancelled the optimisation. Switching to AT_DISPATCH_FLOATING_TYPES_AND2(Half, BFloat16, ...) makes the kernels dispatch on bfloat16 directly. Nothing in them needs adapting: they perform no arithmetic on the values, only gathers, so any trivially copyable scalar type works. This is only reachable because the previous commit made the output tensor inherit the input dtype instead of widening to float32. --- kernels/window_process/swin_window_process_kernel.cu | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/kernels/window_process/swin_window_process_kernel.cu b/kernels/window_process/swin_window_process_kernel.cu index 3c4656610..ce59d2c14 100644 --- a/kernels/window_process/swin_window_process_kernel.cu +++ b/kernels/window_process/swin_window_process_kernel.cu @@ -192,7 +192,8 @@ at::Tensor roll_and_window_partition_forward_cuda( at::Tensor output = at::empty({B*nH*nW, window_h, window_w, C}, input.options()); - AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "roll_and_window_partition_forward_cuda_kernel", ([&] { + AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, + input.scalar_type(), "roll_and_window_partition_forward_cuda_kernel", ([&] { roll_and_window_partition_forward_cuda_kernel<<>>( input.data_ptr(), output.data_ptr(), @@ -235,7 +236,8 @@ at::Tensor roll_and_window_partition_backward_cuda( at::Tensor grad_out = at::empty({B, H, W, C}, grad_in.options()); - AT_DISPATCH_FLOATING_TYPES_AND_HALF(grad_in.scalar_type(), "roll_and_window_partition_backward_cuda_kernel", ([&] { + AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, + grad_in.scalar_type(), "roll_and_window_partition_backward_cuda_kernel", ([&] { roll_and_window_partition_backward_cuda_kernel<<>>( grad_in.data_ptr(), grad_out.data_ptr(), @@ -279,7 +281,8 @@ at::Tensor window_merge_and_roll_forward_cuda( at::Tensor output = at::empty({B, H, W, C}, input.options()); - AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "window_merge_and_roll_forward_cuda_kernel", ([&] { + AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, + input.scalar_type(), "window_merge_and_roll_forward_cuda_kernel", ([&] { window_merge_and_roll_forward_cuda_kernel<<>>( input.data_ptr(), output.data_ptr(), @@ -322,7 +325,8 @@ at::Tensor window_merge_and_roll_backward_cuda( at::Tensor grad_out = at::empty({B*nH*nW, window_h, window_w, C}, grad_in.options()); - AT_DISPATCH_FLOATING_TYPES_AND_HALF(grad_in.scalar_type(), "window_merge_and_roll_backward_cuda_kernel", ([&] { + AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, + grad_in.scalar_type(), "window_merge_and_roll_backward_cuda_kernel", ([&] { window_merge_and_roll_backward_cuda_kernel<<>>( grad_in.data_ptr(), grad_out.data_ptr(), From 2eca45b71118c8bff79cde0bc0e8f68412066112 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Mon, 17 Aug 2026 03:43:53 +0200 Subject: [PATCH 07/34] fix: validate kernel preconditions instead of failing silently The launchers derive nH = H / window_h and nW = W / window_w with integer division and never verify that the division is exact. Passing a grid that the window does not tile produces no error: the kernels read the wrong elements, or past the end of the tensor. The same holds for a shift larger than the window, for a shape that disagrees with the tensor handed in, and for the two hardware limits the launch configuration can exceed. check_window_args() now enforces, on the host side before the launch: - H % window_h == 0 and W % window_w == 0 - |shift| < window size, per axis - tensor.numel() == B * H * W * C. The window layout and the spatial layout hold the same number of elements, since nH * window_h == H and nW * window_w == W, so a single check covers all four entry points. - B * H * W * C fits in int32. Kernel offsets are computed in int, as they are upstream; this branch keeps that and reports the limit rather than wrapping around silently. - B * nH * nW <= 65535 and H <= 65535, the CUDA grid.z and grid.y limits. --- .../window_process/swin_window_process.cpp | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/kernels/window_process/swin_window_process.cpp b/kernels/window_process/swin_window_process.cpp index daf10629c..3896f44b9 100644 --- a/kernels/window_process/swin_window_process.cpp +++ b/kernels/window_process/swin_window_process.cpp @@ -17,6 +17,9 @@ #include #include +#include +#include + at::Tensor roll_and_window_partition_forward_cuda( at::Tensor & input, @@ -74,6 +77,53 @@ at::Tensor window_merge_and_roll_backward_cuda( #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) +// The kernels derive nH = H / window_h and nW = W / window_w with integer +// division and index with int arithmetic. Violating any of the assumptions below +// does not fail: it silently reads the wrong element, or reads out of bounds. +// All four entry points share the same constraints -- the window layout and the +// spatial layout hold the same number of elements -- so one check covers them. +static void check_window_args( + const at::Tensor & tensor, + const int B, + const int H, + const int W, + const int C, + const int shift_h, + const int shift_w, + const int window_h, + const int window_w){ + + TORCH_CHECK(B > 0 && H > 0 && W > 0 && C > 0, + "B, H, W and C must be positive, got (", B, ", ", H, ", ", W, ", ", C, ")"); + TORCH_CHECK(window_h > 0 && window_w > 0, + "window size must be positive, got (", window_h, ", ", window_w, ")"); + TORCH_CHECK(H % window_h == 0, + "H (", H, ") must be divisible by window_h (", window_h, ")"); + TORCH_CHECK(W % window_w == 0, + "W (", W, ") must be divisible by window_w (", window_w, ")"); + TORCH_CHECK(std::abs(shift_h) < window_h && std::abs(shift_w) < window_w, + "|shift| must be smaller than the window size on each axis, got shift (", + shift_h, ", ", shift_w, ") for window (", window_h, ", ", window_w, ")"); + + const int64_t numel = static_cast(B) * H * W * C; + TORCH_CHECK(tensor.numel() == numel, + "expected ", numel, " elements for B=", B, ", H=", H, ", W=", W, ", C=", C, + ", got ", tensor.numel()); + // Offsets are computed in int inside the kernels, as they are upstream. + TORCH_CHECK(numel <= static_cast(std::numeric_limits::max()), + "tensor holds ", numel, " elements; kernel offsets are computed in int32 " + "and would overflow"); + + // grid.z is B * nH * nW for the partition kernels and B for the merge ones; + // grid.y is window_h or H. Both are capped at 65535 by the CUDA driver. + const int64_t num_windows = static_cast(B) * (H / window_h) * (W / window_w); + TORCH_CHECK(num_windows <= 65535, + "B * nH * nW = ", num_windows, " exceeds the CUDA grid.z limit of 65535"); + TORCH_CHECK(H <= 65535, + "H = ", H, " exceeds the CUDA grid.y limit of 65535"); +} + + at::Tensor roll_and_window_partition_forward( at::Tensor & input, @@ -87,6 +137,7 @@ at::Tensor roll_and_window_partition_forward( const int window_h, const int window_w){ CHECK_INPUT(input); + check_window_args(input, B, H, W, C, shift_h, shift_w, window_h, window_w); return roll_and_window_partition_forward_cuda(input, B, H, W, C, shift_h, shift_w, window_h, window_w); } @@ -103,6 +154,7 @@ at::Tensor roll_and_window_partition_backward( const int window_h, const int window_w){ CHECK_INPUT(grad_in); + check_window_args(grad_in, B, H, W, C, shift_h, shift_w, window_h, window_w); return roll_and_window_partition_backward_cuda(grad_in, B, H, W, C, shift_h, shift_w, window_h, window_w); } @@ -119,6 +171,7 @@ at::Tensor window_merge_and_roll_forward( const int window_h, const int window_w){ CHECK_INPUT(input); + check_window_args(input, B, H, W, C, shift_h, shift_w, window_h, window_w); return window_merge_and_roll_forward_cuda(input, B, H, W, C, shift_h, shift_w, window_h, window_w); } @@ -135,6 +188,7 @@ at::Tensor window_merge_and_roll_backward( const int window_h, const int window_w){ CHECK_INPUT(grad_in); + check_window_args(grad_in, B, H, W, C, shift_h, shift_w, window_h, window_w); return window_merge_and_roll_backward_cuda(grad_in, B, H, W, C, shift_h, shift_w, window_h, window_w); } From b22c53f5b0936c20151b8706ebe2b494c70703c9 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Mon, 17 Aug 2026 03:47:26 +0200 Subject: [PATCH 08/34] test: cover the compiled kernels properly, and split off the benchmark The existing unit_test.py had two problems beyond the duplication already removed in the non-square commit. The two backward tests never checked a gradient. Both call .backward() on the eager and fused outputs and then assert torch.equal(expected, fused_output), which compares the forward results again; input1.grad and input2.grad are never read. The backward kernels were effectively untested. They are now compared directly. test_forward_backward_speed asserted that the fused path is faster than eager, timed with time.time() and no warmup. A wall-clock assertion in a unit test is flaky by construction, so it moves to benchmark.py, which uses CUDA events, warms up, and reports peak memory alongside the timings. Coverage is now the product of: dtype float64, float32, float16, bfloat16 (when supported) shape square; nH > nW; nH < nW; non-square window, both orientations; a single window shift 0 (W-MSA) and window // 2 (SW-MSA) for forward, backward, and a forward/reverse round trip -- plus a check that an int shift/window still matches the (h, w) pair it expands to, which is the path models/swin_transformer.py uses, and that each precondition added earlier raises rather than corrupting memory. Parity is asserted with torch.equal on every dtype, including float16 and bfloat16. That is not too strict: the kernels are exact permutations and perform no arithmetic, so any deviation is an indexing error rather than a rounding one. For the same reason gradcheck adds nothing here, even though float64 does dispatch. Importing the extension is guarded so the file stays collectible on machines without a GPU, where every test skips. --- kernels/window_process/benchmark.py | 157 ++++++++++++ kernels/window_process/unit_test.py | 369 +++++++++++++++------------- 2 files changed, 353 insertions(+), 173 deletions(-) create mode 100644 kernels/window_process/benchmark.py diff --git a/kernels/window_process/benchmark.py b/kernels/window_process/benchmark.py new file mode 100644 index 000000000..98b54cf6c --- /dev/null +++ b/kernels/window_process/benchmark.py @@ -0,0 +1,157 @@ +# -------------------------------------------------------- +# Fused kernel for window process for SwinTransformer +# Copyright (c) 2022 Nvidia +# Licensed under The MIT License [see LICENSE for details] +# -------------------------------------------------------- +# Fused kernels vs the PyTorch ops they replace, in wall time and peak memory. +# +# The eager path materialises the tensor twice -- torch.roll writes a copy, then +# window_partition permutes and calls .contiguous() for another -- where the +# fused kernel writes once. The operation is entirely memory bound, so the +# expected ceiling is roughly 2x. +# +# python benchmark.py # forward + backward, all dtypes +# python benchmark.py --iters 200 --batch 64 +# -------------------------------------------------------- + +import argparse + +import torch + +import reference as ref + +try: + from window_process import WindowProcess, WindowProcessReverse +except ImportError as exc: # extension not built here + WindowProcess = WindowProcessReverse = None + IMPORT_ERROR = exc +else: + IMPORT_ERROR = None + + +# (name, B, H, W, C, window_h, window_w) +CONFIGS = [ + ('swin-tiny stage 1 56x56 w7', 192, 56, 56, 96, 7, 7), + ('swin-tiny stage 2 28x28 w7', 192, 28, 28, 192, 7, 7), + ('non-square 32x16 w8', 192, 32, 16, 96, 8, 8), + ('non-square window 16x64 w4x16', 192, 16, 64, 96, 4, 16), +] + + +def timed(fn, iters, warmup=10): + """Wall time per iteration in ms, and peak allocated memory in MiB.""" + for _ in range(warmup): + fn() + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + + return start.elapsed_time(end) / iters, torch.cuda.max_memory_allocated() / 2 ** 20 + + +def make_runners(B, H, W, C, window_h, window_w, shift_h, shift_w, dtype, backward): + window = (window_h, window_w) + n = B * (H // window_h) * (W // window_w) + + x = torch.randn((B, H, W, C), dtype=dtype, device='cuda', requires_grad=backward) + grad = torch.randn((n, window_h, window_w, C), dtype=dtype, device='cuda') + + def eager(): + shifted = torch.roll(x, shifts=(-shift_h, -shift_w), dims=(1, 2)) + out = ref.window_partition(shifted, window) + if backward: + x.grad = None + out.backward(grad) + + def fused(): + out = WindowProcess.apply(x, B, H, W, C, (-shift_h, -shift_w), window) + if backward: + x.grad = None + out.backward(grad) + + return eager, fused + + +def make_reverse_runners(B, H, W, C, window_h, window_w, shift_h, shift_w, dtype, backward): + window = (window_h, window_w) + n = B * (H // window_h) * (W // window_w) + + x = torch.randn((n, window_h, window_w, C), dtype=dtype, device='cuda', + requires_grad=backward) + grad = torch.randn((B, H, W, C), dtype=dtype, device='cuda') + + def eager(): + merged = ref.window_reverse(x, window, H, W) + out = torch.roll(merged, shifts=(shift_h, shift_w), dims=(1, 2)) + if backward: + x.grad = None + out.backward(grad) + + def fused(): + out = WindowProcessReverse.apply(x, B, H, W, C, (shift_h, shift_w), window) + if backward: + x.grad = None + out.backward(grad) + + return eager, fused + + +def run(args): + dtypes = [torch.float32, torch.float16] + if torch.cuda.is_bf16_supported(): + dtypes.append(torch.bfloat16) + + print(f'device: {torch.cuda.get_device_name()}') + print(f'torch: {torch.__version__} (cuda {torch.version.cuda}, hip {torch.version.hip})') + print(f'mode: {"forward + backward" if not args.forward_only else "forward only"}, ' + f'{args.iters} iterations after 10 warmup\n') + + header = (f'| {"config":<32} | {"op":<7} | {"dtype":<9} | {"eager ms":>9} | ' + f'{"fused ms":>9} | {"speedup":>8} | {"eager MiB":>10} | {"fused MiB":>10} |') + print(header) + print('|' + '|'.join('-' * (len(c) + 2) for c in header.split('|')[1:-1]) + '|') + + for name, B, H, W, C, window_h, window_w in CONFIGS: + B = args.batch or B + shift_h, shift_w = window_h // 2, window_w // 2 + for op, factory in (('partition', make_runners), ('merge', make_reverse_runners)): + for dtype in dtypes: + eager, fused = factory(B, H, W, C, window_h, window_w, + shift_h, shift_w, dtype, + backward=not args.forward_only) + e_ms, e_mem = timed(eager, args.iters) + f_ms, f_mem = timed(fused, args.iters) + print(f'| {name:<32} | {op:<7} | {str(dtype).replace("torch.", ""):<9} | ' + f'{e_ms:>9.3f} | {f_ms:>9.3f} | {e_ms / f_ms:>7.2f}x | ' + f'{e_mem:>10.1f} | {f_mem:>10.1f} |') + del eager, fused + torch.cuda.empty_cache() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--iters', type=int, default=100) + parser.add_argument('--batch', type=int, default=None, + help='override the batch size of every config') + parser.add_argument('--forward-only', action='store_true') + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit('benchmark.py requires a CUDA device') + if IMPORT_ERROR is not None: + raise SystemExit( + f'the swin_window_process extension is not importable ({IMPORT_ERROR}); ' + 'build it with: python setup.py install') + + torch.manual_seed(0) + run(args) + + +if __name__ == '__main__': + main() diff --git a/kernels/window_process/unit_test.py b/kernels/window_process/unit_test.py index df35600b5..73bae42b4 100644 --- a/kernels/window_process/unit_test.py +++ b/kernels/window_process/unit_test.py @@ -3,196 +3,219 @@ # Copyright (c) 2022 Nvidia # Licensed under The MIT License [see LICENSE for details] # -------------------------------------------------------- +# Parity tests for the compiled kernels against the PyTorch ops they replace. +# Requires a CUDA device and the built extension; the index math itself is +# covered without either by test_index_math.py. +# +# python unit_test.py +# -------------------------------------------------------- -import torch -import random -import time import unittest -from window_process import WindowProcess, WindowProcessReverse +import torch +import reference as ref -def window_partition(x, window_size): - """ - Args: - x: (B, H, W, C) - window_size (int): window size - Returns: - windows: (num_windows*B, window_size, window_size, C) - """ - B, H, W, C = x.shape - x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) - windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) - return windows +try: + from window_process import WindowProcess, WindowProcessReverse + EXTENSION_AVAILABLE = True +except ImportError: # extension not built here + WindowProcess = WindowProcessReverse = None + EXTENSION_AVAILABLE = False + + +CUDA_AVAILABLE = torch.cuda.is_available() +requires_cuda = unittest.skipIf( + not (CUDA_AVAILABLE and EXTENSION_AVAILABLE), + 'requires a CUDA device and the compiled swin_window_process extension') -def window_reverse(windows, window_size, H, W): - """ - Args: - windows: (num_windows*B, window_size, window_size, C) - window_size (int): Window size - H (int): Height of image - W (int): Width of image - Returns: - x: (B, H, W, C) - """ - B = int(windows.shape[0] / (H * W / window_size / window_size)) - x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) - x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) - return x +def available_dtypes(): + """float64 and float32 come from AT_DISPATCH_FLOATING_TYPES; the rest are added.""" + dtypes = [torch.float64, torch.float32, torch.float16] + if CUDA_AVAILABLE and torch.cuda.is_bf16_supported(): + dtypes.append(torch.bfloat16) + return dtypes -def pyt_forward(x, shift_size, window_size): - # x in shape(B, H, W, C) - # cyclic shift - if shift_size > 0: - shifted_x = torch.roll(x, shifts=(-shift_size, -shift_size), dims=(1, 2)) - else: - shifted_x = x - # partition windows - x_windows = window_partition(shifted_x, window_size) - return x_windows - - -def reverse_pyt_forward(attn_windows, shift_size, window_size, H, W): - # x in shape(B*nH*nW, window_size, window_size, C) - shifted_x = window_reverse(attn_windows, window_size, H, W) - if shift_size > 0: - x = torch.roll(shifted_x, shifts=(shift_size, shift_size), dims=(1, 2)) - else: - x = shifted_x + +# (B, H, W, C, window_h, window_w) +SHAPES = [ + (2, 56, 56, 96, 7, 7), # nH == nW == 8 the ImageNet configuration + (2, 32, 16, 64, 8, 8), # nH=4, nW=2 out-of-bounds read before the fix + (2, 16, 32, 64, 8, 8), # nH=2, nW=4 silent corruption before the fix + (2, 16, 64, 64, 4, 16), # nH=4, nW=4 non-square window + (2, 64, 16, 64, 16, 4), # nH=4, nW=4 non-square window, transposed + (2, 24, 24, 32, 24, 24), # nH == nW == 1 a single window +] + + +def shifts_for(window_h, window_w): + """W-MSA (no shift) and SW-MSA (half window), the two regimes Swin uses.""" + return [(0, 0), (window_h // 2, window_w // 2)] + + +def pyt_forward(x, shift, window): + """torch.roll followed by window_partition -- what WindowProcess fuses.""" + shift_h, shift_w = ref.to_pair(shift) + if shift_h or shift_w: + x = torch.roll(x, shifts=(-shift_h, -shift_w), dims=(1, 2)) + return ref.window_partition(x, window) + + +def reverse_pyt_forward(windows, shift, window, H, W): + """window_reverse followed by torch.roll -- what WindowProcessReverse fuses.""" + shift_h, shift_w = ref.to_pair(shift) + x = ref.window_reverse(windows, window, H, W) + if shift_h or shift_w: + x = torch.roll(x, shifts=(shift_h, shift_w), dims=(1, 2)) return x -def copy_one_tensor(input, requires_grad=True): - input1 = input.clone().detach().requires_grad_(requires_grad).cuda() - return input1 +def leaf(tensor, requires_grad=True): + return tensor.clone().detach().requires_grad_(requires_grad).cuda() -class Test_WindowProcess(unittest.TestCase): - def setUp(self): - self.B = 192 - self.H = 56 - self.W = 56 - self.C = 96 - self.shift_size = 2 - self.window_size = 7 - self.nH = self.H // self.window_size - self.nW = self.W // self.window_size - - def test_roll_and_window_partition_forward(self, dtype=torch.float32): - input = torch.randn((self.B, self.H, self.W, self.C), dtype=dtype, requires_grad=True).cuda() - - input1 = copy_one_tensor(input, True) - input2 = copy_one_tensor(input, True) - with torch.no_grad(): - # ori - expected = pyt_forward(input1, self.shift_size, self.window_size) - # fused kernel - fused_output = WindowProcess.apply(input2, self.B, self.H, self.W, self.C, -self.shift_size, self.window_size) - - self.assertTrue(torch.equal(expected, fused_output)) - #self.assertTrue(torch.allclose(expected, fused_output, rtol=1e-05, atol=1e-08)) - - def test_roll_and_window_partition_backward(self, dtype=torch.float32): - input = torch.randn((self.B, self.H, self.W, self.C), dtype=dtype, requires_grad=True).cuda() - d_loss_tensor = torch.randn((self.B*self.nW*self.nH, self.window_size, self.window_size, self.C), dtype=dtype).cuda() - - input1 = copy_one_tensor(input, True) - input2 = copy_one_tensor(input, True) - - # ori - expected = pyt_forward(input1, self.shift_size, self.window_size) - expected.backward(d_loss_tensor) - # fused kernel - fused_output = WindowProcess.apply(input2, self.B, self.H, self.W, self.C, -self.shift_size, self.window_size) - fused_output.backward(d_loss_tensor) - - self.assertTrue(torch.equal(expected, fused_output)) - #self.assertTrue(torch.allclose(expected, fused_output, rtol=1e-05, atol=1e-08)) - - def test_window_merge_and_roll_forward(self, dtype=torch.float32): - input = torch.randn((self.B*self.nH*self.nW, self.window_size, self.window_size, self.C), dtype=dtype, requires_grad=True).cuda() - - input1 = copy_one_tensor(input, True) - input2 = copy_one_tensor(input, True) +def spatial(B, H, W, C, dtype): + return torch.randn((B, H, W, C), dtype=dtype) + +def windowed(B, H, W, C, window_h, window_w, dtype): + n = B * (H // window_h) * (W // window_w) + return torch.randn((n, window_h, window_w, C), dtype=dtype) + + +@requires_cuda +class TestWindowProcess(unittest.TestCase): + """The kernels are exact permutations, so parity must be bit-for-bit. + + torch.equal is a stronger check than gradcheck here: no arithmetic is + performed on the values, so any deviation is an indexing error, not a + numerical one. Every dtype is therefore held to exact equality. + """ + + def _cases(self): + for dtype in available_dtypes(): + for B, H, W, C, window_h, window_w in SHAPES: + for shift in shifts_for(window_h, window_w): + yield dtype, B, H, W, C, (window_h, window_w), shift + + def test_partition_forward(self): + for dtype, B, H, W, C, window, shift in self._cases(): + with self.subTest(dtype=dtype, shape=(B, H, W, C), window=window, shift=shift): + x = spatial(B, H, W, C, dtype) + with torch.no_grad(): + expected = pyt_forward(leaf(x), shift, window) + fused = WindowProcess.apply( + leaf(x), B, H, W, C, (-shift[0], -shift[1]), window) + self.assertTrue(torch.equal(expected, fused)) + + def test_partition_backward(self): + for dtype, B, H, W, C, window, shift in self._cases(): + with self.subTest(dtype=dtype, shape=(B, H, W, C), window=window, shift=shift): + x = spatial(B, H, W, C, dtype) + grad = windowed(B, H, W, C, window[0], window[1], dtype).cuda() + + a, b = leaf(x), leaf(x) + pyt_forward(a, shift, window).backward(grad) + WindowProcess.apply( + b, B, H, W, C, (-shift[0], -shift[1]), window).backward(grad) + + self.assertIsNotNone(a.grad) + self.assertTrue(torch.equal(a.grad, b.grad)) + + def test_merge_forward(self): + for dtype, B, H, W, C, window, shift in self._cases(): + with self.subTest(dtype=dtype, shape=(B, H, W, C), window=window, shift=shift): + x = windowed(B, H, W, C, window[0], window[1], dtype) + with torch.no_grad(): + expected = reverse_pyt_forward(leaf(x), shift, window, H, W) + fused = WindowProcessReverse.apply(leaf(x), B, H, W, C, shift, window) + self.assertTrue(torch.equal(expected, fused)) + + def test_merge_backward(self): + for dtype, B, H, W, C, window, shift in self._cases(): + with self.subTest(dtype=dtype, shape=(B, H, W, C), window=window, shift=shift): + x = windowed(B, H, W, C, window[0], window[1], dtype) + grad = spatial(B, H, W, C, dtype).cuda() + + a, b = leaf(x), leaf(x) + reverse_pyt_forward(a, shift, window, H, W).backward(grad) + WindowProcessReverse.apply(b, B, H, W, C, shift, window).backward(grad) + + self.assertIsNotNone(a.grad) + self.assertTrue(torch.equal(a.grad, b.grad)) + + def test_round_trip(self): + """WindowProcessReverse must undo WindowProcess for the same shift.""" + for dtype, B, H, W, C, window, shift in self._cases(): + with self.subTest(dtype=dtype, shape=(B, H, W, C), window=window, shift=shift): + x = leaf(spatial(B, H, W, C, dtype), requires_grad=False) + windows = WindowProcess.apply( + x, B, H, W, C, (-shift[0], -shift[1]), window) + back = WindowProcessReverse.apply( + windows.contiguous(), B, H, W, C, shift, window) + self.assertTrue(torch.equal(x, back)) + + +@requires_cuda +class TestScalarWindowStillWorks(unittest.TestCase): + """An int shift/window must behave exactly as the (h, w) pair it expands to. + + This is the compatibility path used by models/swin_transformer.py, which + passes ints and must keep working unchanged. + """ + + def test_int_and_pair_agree(self): + B, H, W, C, window, shift = 2, 56, 56, 96, 7, 3 + x = spatial(B, H, W, C, torch.float32) with torch.no_grad(): - # ori - expected = reverse_pyt_forward(input1, self.shift_size, self.window_size, self.H, self.W) - # fused kernel - fused_output = WindowProcessReverse.apply(input2, self.B, self.H, self.W, self.C, self.shift_size, self.window_size) - - self.assertTrue(torch.equal(expected, fused_output)) - #self.assertTrue(torch.allclose(expected, fused_output, rtol=1e-05, atol=1e-08)) - - - def test_window_merge_and_roll_backward(self, dtype=torch.float32): - input = torch.randn((self.B*self.nH*self.nW, self.window_size, self.window_size, self.C), dtype=dtype, requires_grad=True).cuda() - d_loss_tensor = torch.randn((self.B, self.H, self.W, self.C), dtype=dtype, requires_grad=True).cuda() - - input1 = copy_one_tensor(input, True) - input2 = copy_one_tensor(input, True) - - # ori - expected = reverse_pyt_forward(input1, self.shift_size, self.window_size, self.H, self.W) - expected.backward(d_loss_tensor) - # fused kernel - fused_output = WindowProcessReverse.apply(input2, self.B, self.H, self.W, self.C, self.shift_size, self.window_size) - fused_output.backward(d_loss_tensor) - - self.assertTrue(torch.equal(expected, fused_output)) - #self.assertTrue(torch.allclose(expected, fused_output, rtol=1e-05, atol=1e-08)) - - def test_forward_backward_speed(self, dtype=torch.float32, times=1000): - input = torch.randn((self.B*self.nH*self.nW, self.window_size, self.window_size, self.C), dtype=dtype, requires_grad=True).cuda() - d_loss_tensor = torch.randn((self.B, self.H, self.W, self.C), dtype=dtype, requires_grad=True).cuda() - - input1 = copy_one_tensor(input, True) - input2 = copy_one_tensor(input, True) - - # SwinTransformer official - def run_pyt(t=1000): - for _ in range(t): - expected = reverse_pyt_forward(input1, self.shift_size, self.window_size, self.H, self.W) - expected.backward(d_loss_tensor) - - # my op - def run_fusedop(t=1000): - for _ in range(t): - fused_output = WindowProcessReverse.apply(input2, self.B, self.H, self.W, self.C, self.shift_size, self.window_size) - fused_output.backward(d_loss_tensor) - - torch.cuda.synchronize() - t1 = time.time() - run_pyt(t=times) - torch.cuda.synchronize() - t2 = time.time() - run_fusedop(t=times) - torch.cuda.synchronize() - t3 = time.time() - self.assertTrue((t3 - t2) < (t2 - t1)) - - print('Run {} times'.format(times)) - print('Original time cost: {}'.format(t2 - t1)) - print('Fused op time cost: {}'.format(t3 - t2)) - - def test_roll_and_window_partition_forward_fp16(self, dtype=torch.float16): - self.test_roll_and_window_partition_forward(dtype=dtype) - - def test_roll_and_window_partition_backward_fp16(self, dtype=torch.float16): - self.test_roll_and_window_partition_backward(dtype=dtype) - - def test_window_merge_and_roll_forward_fp16(self, dtype=torch.float16): - self.test_window_merge_and_roll_forward(dtype=dtype) - - def test_window_merge_and_roll_backward_fp16(self, dtype=torch.float16): - self.test_window_merge_and_roll_backward(dtype=dtype) - - def test_forward_backward_speed_fp16(self, dtype=torch.float16, times=1000): - self.test_forward_backward_speed(dtype=dtype, times=times) + as_int = WindowProcess.apply(leaf(x), B, H, W, C, -shift, window) + as_pair = WindowProcess.apply( + leaf(x), B, H, W, C, (-shift, -shift), (window, window)) + self.assertTrue(torch.equal(as_int, as_pair)) + + +@requires_cuda +class TestPreconditions(unittest.TestCase): + """Invalid arguments must raise instead of reading the wrong memory.""" + + def setUp(self): + self.B, self.H, self.W, self.C = 2, 16, 16, 32 + self.x = leaf(spatial(self.B, self.H, self.W, self.C, torch.float32), False) + + def _apply(self, **overrides): + kwargs = dict(B=self.B, H=self.H, W=self.W, C=self.C, shift=(0, 0), window=(8, 8)) + kwargs.update(overrides) + return WindowProcess.apply( + self.x, kwargs['B'], kwargs['H'], kwargs['W'], kwargs['C'], + kwargs['shift'], kwargs['window']) + + def test_window_must_divide_height(self): + with self.assertRaisesRegex(RuntimeError, 'divisible by window_h'): + self._apply(window=(5, 8)) + + def test_window_must_divide_width(self): + with self.assertRaisesRegex(RuntimeError, 'divisible by window_w'): + self._apply(window=(8, 5)) + + def test_shift_must_be_smaller_than_window(self): + with self.assertRaisesRegex(RuntimeError, 'smaller than the window size'): + self._apply(shift=(8, 0)) + + def test_shape_must_match_tensor(self): + with self.assertRaisesRegex(RuntimeError, 'expected'): + self._apply(C=self.C * 2) + + def test_non_contiguous_input_is_rejected(self): + transposed = self.x.transpose(1, 2) + with self.assertRaises(RuntimeError): + WindowProcess.apply( + transposed, self.B, self.W, self.H, self.C, (0, 0), (8, 8)) if __name__ == '__main__': - print('Pass only two tensors are exactly the same (using torch.equal).\n') + if not (CUDA_AVAILABLE and EXTENSION_AVAILABLE): + print('No CUDA device or extension not built: every test here will be skipped.') + print('Run test_index_math.py to check the index math without a GPU.\n') torch.manual_seed(0) unittest.main(verbosity=2) From ee3ba5c8d51570227c743fbd66dfee6011062774 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Tue, 25 Aug 2026 21:34:03 +0200 Subject: [PATCH 09/34] test: close the remaining coverage gaps Five paths were exercised by nothing. best_block_dim() picks 64, 128 or 256 threads per block from C, but every shape in the suite used C in {1, 32, 64, 96}, all below the first threshold of 384. Only the 64-thread branch ran. Shapes with C = 512 and C = 1024 now reach the other two, and C = 1, 64 and 1024 make the strided `for (i = threadIdx.x; i < C; i += blockDim.x)` loop run in its partial, exact and multi-pass forms. The grid.z guard, B * nH * nW <= 65535, is now tested with a batch that trips it. It costs a 4.5M element tensor because the shape is chosen so the window count comes from B alone. The .contiguous() call added to backward() had no test. An upstream op can hand the backward a non-contiguous gradient, and the C++ side rejects it. The current-stream fix had no test either. A race against the default stream is timing dependent, so this uses CUDA graph capture instead: capture runs on a non-default stream and rejects any launch on the legacy default stream, so the graph fails to capture unless the fix is present. That is deterministic. A second test runs the kernel inside an explicit side stream and checks parity. test_model_parity.py adds the end-to-end statement: a SwinTransformer produces bit-identical logits and bit-identical parameter gradients whether or not the fused kernels are enabled. Both runs flip the flag on the same model instance, so the weights are the same by construction rather than by seeding. --- kernels/window_process/test_model_parity.py | 107 ++++++++++++++++++++ kernels/window_process/unit_test.py | 82 +++++++++++++++ 2 files changed, 189 insertions(+) create mode 100644 kernels/window_process/test_model_parity.py diff --git a/kernels/window_process/test_model_parity.py b/kernels/window_process/test_model_parity.py new file mode 100644 index 000000000..ea7035a45 --- /dev/null +++ b/kernels/window_process/test_model_parity.py @@ -0,0 +1,107 @@ +# -------------------------------------------------------- +# Fused kernel for window process for SwinTransformer +# Copyright (c) 2022 Nvidia +# Licensed under The MIT License [see LICENSE for details] +# -------------------------------------------------------- +# End-to-end check: a SwinTransformer must produce identical output with and +# without the fused window kernels. The kernels replace torch.roll plus +# window_partition, which are exact data movements, so the two paths are not +# merely close -- they are bit-for-bit equal. +# +# python test_model_parity.py +# -------------------------------------------------------- + +import os +import sys +import unittest + +import torch + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))) + +try: + import swin_window_process # noqa: F401 + EXTENSION_AVAILABLE = True +except ImportError: + EXTENSION_AVAILABLE = False + +try: + from models.swin_transformer import SwinTransformer + MODEL_AVAILABLE = True +except ImportError: # timm missing, for instance + SwinTransformer = None + MODEL_AVAILABLE = False + + +requires_everything = unittest.skipIf( + not (torch.cuda.is_available() and EXTENSION_AVAILABLE and MODEL_AVAILABLE), + 'requires a CUDA device, the compiled extension and an importable model') + + +def set_fused(model, enabled): + """Flip the flag on every block, so both runs share exactly the same weights.""" + switched = 0 + for module in model.modules(): + if hasattr(module, 'fused_window_process'): + module.fused_window_process = enabled + switched += 1 + return switched + + +@requires_everything +class TestModelParity(unittest.TestCase): + + def _model(self, img_size=56, window_size=7): + torch.manual_seed(0) + model = SwinTransformer( + img_size=img_size, + patch_size=4, + in_chans=3, + num_classes=10, + embed_dim=48, + depths=[2, 2], + num_heads=[3, 6], + window_size=window_size, + drop_path_rate=0.0, + fused_window_process=False, + ) + return model.cuda().eval() + + def test_forward_is_identical(self): + model = self._model() + x = torch.randn(2, 3, 56, 56, device='cuda') + + with torch.no_grad(): + eager = model(x) + blocks = set_fused(model, True) + fused = model(x) + set_fused(model, False) + + # The shifted blocks are the ones that take the fused path at all. + self.assertGreater(blocks, 0, 'no block exposes fused_window_process') + self.assertTrue(torch.equal(eager, fused)) + + def test_gradients_are_identical(self): + model = self._model() + x = torch.randn(2, 3, 56, 56, device='cuda') + target = torch.randn(2, 10, device='cuda') + + def grads(): + model.zero_grad(set_to_none=True) + torch.nn.functional.mse_loss(model(x), target).backward() + return [p.grad.clone() for p in model.parameters() if p.grad is not None] + + eager = grads() + set_fused(model, True) + fused = grads() + set_fused(model, False) + + self.assertEqual(len(eager), len(fused)) + for a, b in zip(eager, fused): + self.assertTrue(torch.equal(a, b)) + + +if __name__ == '__main__': + if not (torch.cuda.is_available() and EXTENSION_AVAILABLE and MODEL_AVAILABLE): + print('Skipping: needs a CUDA device, the built extension and the model.\n') + unittest.main(verbosity=2) diff --git a/kernels/window_process/unit_test.py b/kernels/window_process/unit_test.py index 73bae42b4..78b4e221d 100644 --- a/kernels/window_process/unit_test.py +++ b/kernels/window_process/unit_test.py @@ -46,6 +46,15 @@ def available_dtypes(): (2, 16, 64, 64, 4, 16), # nH=4, nW=4 non-square window (2, 64, 16, 64, 16, 4), # nH=4, nW=4 non-square window, transposed (2, 24, 24, 32, 24, 24), # nH == nW == 1 a single window + # C selects the block width in best_block_dim(): < 384 -> 64 threads, + # < 1024 -> 128, otherwise 256. All three branches must be exercised, and C + # must be smaller than, equal to and larger than the block width so that the + # strided `for (i = threadIdx.x; i < C; i += blockDim.x)` loop is covered in + # its partial, exact and multi-pass forms. + (2, 16, 16, 1, 8, 8), # C < blockDim 64 threads, 63 idle + (2, 16, 16, 64, 8, 8), # C == blockDim 64 threads, one pass + (2, 16, 16, 512, 8, 8), # 384 <= C < 1024 128 threads + (2, 16, 16, 1024, 8, 8), # C >= 1024 256 threads ] @@ -157,6 +166,72 @@ def test_round_trip(self): self.assertTrue(torch.equal(x, back)) + def test_non_contiguous_gradient(self): + """An upstream op can hand the backward a non-contiguous gradient. + + The C++ side asserts contiguity, so window_process.py normalises it. + """ + B, H, W, C, window, shift = 2, 16, 16, 32, (8, 8), (4, 4) + n = B * (H // window[0]) * (W // window[1]) + x = spatial(B, H, W, C, torch.float32) + grad = torch.randn((n, window[1], window[0], C)).cuda().transpose(1, 2) + self.assertFalse(grad.is_contiguous()) + + a, b = leaf(x), leaf(x) + pyt_forward(a, shift, window).backward(grad) + WindowProcess.apply(b, B, H, W, C, (-shift[0], -shift[1]), window).backward(grad) + self.assertTrue(torch.equal(a.grad, b.grad)) + + +@requires_cuda +class TestStreamSemantics(unittest.TestCase): + """The kernels must run on the current stream, not on the default one. + + A race against the default stream is timing dependent and makes a poor + test. CUDA graph capture is the deterministic form of the same property: + capture runs on a non-default stream and rejects any kernel launched on the + legacy default stream outright, so this fails to capture unless the launch + uses at::cuda::getCurrentCUDAStream(). + """ + + def test_capturable_in_a_cuda_graph(self): + B, H, W, C, window, shift = 2, 16, 16, 32, (8, 8), (4, 4) + static_in = torch.randn((B, H, W, C), device='cuda') + + def call(): + return WindowProcess.apply( + static_in, B, H, W, C, (-shift[0], -shift[1]), window) + + # Warm up on a side stream, as the CUDA graph capture protocol requires. + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(3): + call() + torch.cuda.current_stream().wait_stream(side) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_out = call() + graph.replay() + torch.cuda.synchronize() + + self.assertTrue(torch.equal(static_out, pyt_forward(static_in, shift, window))) + + def test_matches_when_run_on_a_side_stream(self): + B, H, W, C, window, shift = 2, 16, 16, 32, (8, 8), (4, 4) + x = torch.randn((B, H, W, C), device='cuda') + expected = pyt_forward(x, shift, window) + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + fused = WindowProcess.apply(x, B, H, W, C, (-shift[0], -shift[1]), window) + torch.cuda.current_stream().wait_stream(side) + + self.assertTrue(torch.equal(expected, fused)) + + @requires_cuda class TestScalarWindowStillWorks(unittest.TestCase): """An int shift/window must behave exactly as the (h, w) pair it expands to. @@ -206,6 +281,13 @@ def test_shape_must_match_tensor(self): with self.assertRaisesRegex(RuntimeError, 'expected'): self._apply(C=self.C * 2) + def test_window_count_above_the_grid_limit_is_rejected(self): + """B * nH * nW is grid.z, which the CUDA driver caps at 65535.""" + B, H, W, C = 70000, 8, 8, 1 + big = leaf(spatial(B, H, W, C, torch.float32), False) + with self.assertRaisesRegex(RuntimeError, 'grid.z limit'): + WindowProcess.apply(big, B, H, W, C, (0, 0), (8, 8)) + def test_non_contiguous_input_is_rejected(self): transposed = self.x.transpose(1, 2) with self.assertRaises(RuntimeError): From aadeb8d0334ae40f7d518ca4fa19690e801294c5 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Tue, 25 Aug 2026 21:36:33 +0200 Subject: [PATCH 10/34] feat: verify the ROCm/HIP port without AMD hardware setup.py needs no change. torch.utils.cpp_extension.CUDAExtension hipifies its own sources when torch.version.hip is set (cpp_extension.py, the IS_HIP_EXTENSION branch of CUDAExtension), and derives --offload-arch from PYTORCH_ROCM_ARCH or from the arch list the installed torch was built for. Hardcoding an arch in extra_compile_args would be worse than leaving it alone: _get_rocm_arch_flags() skips its own detection entirely once it sees an offload-arch flag from the user, so pinning gfx942 for MI300X would break the build on every other AMD GPU. What ROCm did need is in the current-stream commit earlier in this branch. Without it there is no stream symbol for hipify to map, and the kernels launch on the HIP null stream -- the same class of bug as on CUDA, with the same consequences under DDP and graph capture. What was missing entirely is any way to tell whether the translation is complete. hipify_check.py fills that in, and deliberately requires no GPU: torch.utils.hipify is pure Python, so this runs on a CPU-only CI runner, or on a laptop. It translates the sources out of place, fails if any CUDA-specific symbol survives, lists the symbols left alone because HIP implements them natively, and asserts that every hipLaunchKernelGGL call carries an explicit stream rather than falling back to the null stream. Verified locally: cuda_runtime.h -> hip/hip_runtime.h cuda_fp16.h -> hip/hip_fp16.h ATen/cuda/CUDAContext.h -> ATen/hip/HIPContext.h c10/cuda/CUDAException.h -> c10/hip/HIPException.h at::cuda::getCurrentCUDAStream -> at::hip::getCurrentHIPStreamMasqueradingAsCUDA C10_CUDA_KERNEL_LAUNCH_CHECK -> C10_HIP_KERNEL_LAUNCH_CHECK kernel<<<...>>> -> hipLaunchKernelGGL(...) swin_window_process.cpp translates to itself: it contains nothing device specific. __ldg, __global__, dim3, the blockIdx/threadIdx family and the AT_DISPATCH macros are all native to HIP and pass through unchanged. best_block_dim() now documents why the port is this cheap: the kernels use no shared memory, no __syncthreads() and no warp-level primitives, so CDNA's 64-wide wavefront against NVIDIA's 32-wide warp changes occupancy and nothing else. All three block widths are already multiples of 64, so neither platform schedules a partial wave. The thresholds were tuned on NVIDIA and are not measured on CDNA; -DSWIN_WP_BLOCK_DIM=N overrides them for tuning. Not validated on AMD hardware. The translation is verified, the build path is the standard PyTorch one, and the arithmetic is platform independent -- but no MI300X has run this. --- kernels/window_process/hipify_check.py | 178 ++++++++++++++++++ .../swin_window_process_kernel.cu | 16 ++ 2 files changed, 194 insertions(+) create mode 100644 kernels/window_process/hipify_check.py diff --git a/kernels/window_process/hipify_check.py b/kernels/window_process/hipify_check.py new file mode 100644 index 000000000..72280c751 --- /dev/null +++ b/kernels/window_process/hipify_check.py @@ -0,0 +1,178 @@ +# -------------------------------------------------------- +# Fused kernel for window process for SwinTransformer +# Copyright (c) 2022 Nvidia +# Licensed under The MIT License [see LICENSE for details] +# -------------------------------------------------------- +# Verifies that the CUDA sources translate cleanly to HIP for ROCm. +# +# torch.utils.cpp_extension.CUDAExtension hipifies its sources automatically +# when torch.version.hip is set, so building on ROCm needs no separate source +# tree and no change to setup.py. What it does need is a way to tell whether the +# translation is complete, and that check must not require an AMD GPU -- or any +# GPU. torch.utils.hipify is pure Python, so this runs anywhere, including in CI +# on a CPU runner. +# +# It fails if any CUDA-specific symbol survives translation, and prints the +# symbols that are deliberately left alone because HIP implements them natively. +# +# python hipify_check.py # verify +# python hipify_check.py --diff # verify and show the generated HIP source +# -------------------------------------------------------- + +import argparse +import difflib +import os +import re +import shutil +import sys +import tempfile + +from torch.utils.hipify import hipify_python + + +SOURCES = ['swin_window_process.cpp', 'swin_window_process_kernel.cu'] + +# Present in the CUDA source and required to disappear from the HIP output. +# A survivor here means hipify has no rule for it and the ROCm build would fail +# to compile, or would silently bind to the wrong runtime. +MUST_BE_TRANSLATED = [ + 'cuda_runtime.h', + 'cuda_fp16.h', + 'ATen/cuda/CUDAContext.h', + 'c10/cuda/CUDAException.h', + 'at::cuda::getCurrentCUDAStream', + 'C10_CUDA_KERNEL_LAUNCH_CHECK', +] + +# Deliberately unchanged: HIP implements these with the same spelling and the +# same semantics, so translating them would be wrong. +KNOWN_PORTABLE = [ + '__global__', + '__ldg', + 'blockIdx', + 'threadIdx', + 'blockDim', + 'gridDim', + 'dim3', + 'AT_DISPATCH_FLOATING_TYPES_AND2', + 'at::ScalarType::BFloat16', +] + + +def hipify_into(destination): + """Translate SOURCES out of place and return {source: hipified path}.""" + here = os.path.dirname(os.path.abspath(__file__)) + staging = os.path.join(destination, 'src') + os.makedirs(staging) + for name in SOURCES: + shutil.copy(os.path.join(here, name), staging) + + results = hipify_python.hipify( + project_directory=staging, + output_directory=os.path.join(destination, 'out'), + includes=('*',), + is_pytorch_extension=True, + out_of_place_only=True, + show_progress=False, + ) + + mapping = {} + for source, result in results.items(): + hipified = getattr(result, 'hipified_path', None) or source + mapping[os.path.basename(source)] = hipified + return staging, mapping + + +def report(staging, mapping, show_diff): + failures = [] + + for name in SOURCES: + original = os.path.join(staging, name) + hipified = mapping.get(name) + if hipified is None or not os.path.exists(hipified): + failures.append(f'{name}: hipify produced no output') + continue + + before = open(original).read() + after = open(hipified).read() + + print(f'\n{name} -> {os.path.basename(hipified)}') + + translated, survived = [], [] + for token in MUST_BE_TRANSLATED: + if token not in before: + continue + (survived if token in after else translated).append(token) + + for token in translated: + print(f' translated {token}') + for token in survived: + print(f' SURVIVED {token}') + failures.append(f'{name}: {token} was not translated') + + portable = [t for t in KNOWN_PORTABLE if t in before] + for token in portable: + if token in after: + print(f' portable {token} (native in HIP, left as is)') + else: + failures.append(f'{name}: {token} was rewritten but should be portable') + + if not translated and not portable: + print(' nothing device specific in this file') + + if show_diff: + diff = difflib.unified_diff( + before.splitlines(True), after.splitlines(True), + fromfile=name, tofile=os.path.basename(hipified)) + sys.stdout.writelines(diff) + + return failures + + +def check_launch_form(mapping): + """Every kernel launch must carry an explicit stream argument after translation. + + hipLaunchKernelGGL takes (kernel, grid, block, sharedMem, stream, ...). A + launch that reaches HIP without a stream runs on the null stream, which is + the ROCm form of the default-stream bug. + """ + hipified = mapping.get('swin_window_process_kernel.cu') + if hipified is None: + return ['no hipified kernel source to inspect'] + + text = open(hipified).read() + launches = re.findall(r'hipLaunchKernelGGL\((.*?)\n', text) + if not launches: + return ['no hipLaunchKernelGGL call found in the translated source'] + + failures = [] + for launch in launches: + if 'getCurrentHIPStream' not in launch: + failures.append(f'launch without an explicit stream: {launch.strip()[:70]}') + print(f'\n{len(launches)} kernel launches, all carrying an explicit HIP stream' + if not failures else '') + return failures + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--diff', action='store_true', + help='print the generated HIP source as a unified diff') + args = parser.parse_args() + + with tempfile.TemporaryDirectory() as tmp: + staging, mapping = hipify_into(tmp) + failures = report(staging, mapping, args.diff) + failures += check_launch_form(mapping) + + if failures: + print('\nFAILED') + for failure in failures: + print(f' {failure}') + raise SystemExit(1) + + print('\nOK: the CUDA sources translate to HIP with no unmapped symbols.') + + +if __name__ == '__main__': + main() diff --git a/kernels/window_process/swin_window_process_kernel.cu b/kernels/window_process/swin_window_process_kernel.cu index ce59d2c14..aeee13309 100644 --- a/kernels/window_process/swin_window_process_kernel.cu +++ b/kernels/window_process/swin_window_process_kernel.cu @@ -23,7 +23,22 @@ #include #include +// Width of the block that walks the channel dimension. +// +// These kernels use no shared memory, no __syncthreads() and no warp-level +// primitives, so this width affects occupancy only and never correctness. That +// is also why they port to AMD unchanged: CDNA schedules 64-wide wavefronts +// against NVIDIA's 32-wide warps, and nothing here depends on that width. All +// three values below are multiples of 64, so neither platform is left running a +// partial wave. +// +// The thresholds were tuned on NVIDIA and have not been measured on CDNA. +// Compile with -DSWIN_WP_BLOCK_DIM=N to override them while tuning. int best_block_dim(int feat_dim){ +#ifdef SWIN_WP_BLOCK_DIM + (void)feat_dim; + return SWIN_WP_BLOCK_DIM; +#else int best_dim; if (feat_dim < 384){ best_dim = 64; @@ -37,6 +52,7 @@ int best_block_dim(int feat_dim){ } } return best_dim; +#endif } From 3a038f52c9c2bb8a942c8c944ea130970887a0e5 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Tue, 25 Aug 2026 21:53:44 +0200 Subject: [PATCH 11/34] docs: document the kernels, and run the GPU-free checks in CI README.md covers what the kernels replace, how to build on CUDA and on ROCm, the int/pair form of the shift and window arguments, the constraints now enforced by TORCH_CHECK, and which tests need a GPU and which do not. The CI workflow is the first thing in this repository outside kernels/window_process/. It is a deliberate exception: there is no CI here at all, and the whole point of reference.py and hipify_check.py is that the parts they verify need no GPU. Without a workflow that stays a property nobody observes. The job runs on a CPU runner and covers: - the index arithmetic of all four kernels, via reference.py - the CUDA to HIP translation, via hipify_check.py, including that every launch carries an explicit stream - that the two GPU test files import and skip cleanly rather than erroring when there is no device and no compiled extension The parity tests against the compiled extension are not run here; they need a CUDA device. --- .github/workflows/window-process.yml | 56 +++++++++++++++ kernels/window_process/README.md | 104 +++++++++++++++++++++++++++ 2 files changed, 160 insertions(+) create mode 100644 .github/workflows/window-process.yml create mode 100644 kernels/window_process/README.md diff --git a/.github/workflows/window-process.yml b/.github/workflows/window-process.yml new file mode 100644 index 000000000..94a9f85de --- /dev/null +++ b/.github/workflows/window-process.yml @@ -0,0 +1,56 @@ +# Correctness checks for the fused window kernels that need no GPU. +# +# The kernels perform no arithmetic on tensor values: each is a gather whose +# whole logic is the computation of input_offset from blockIdx. reference.py +# transcribes that to PyTorch, so the index math can be checked on a CPU runner. +# torch.utils.hipify is pure Python for the same reason, so the ROCm/HIP +# translation can be verified here too. +# +# The parity tests against the compiled extension (unit_test.py, +# test_model_parity.py) need a CUDA device and are not run here. + +name: window process + +on: + push: + paths: + - 'kernels/window_process/**' + - '.github/workflows/window-process.yml' + pull_request: + paths: + - 'kernels/window_process/**' + - '.github/workflows/window-process.yml' + workflow_dispatch: + +jobs: + cpu-checks: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ['3.9', '3.11'] + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install PyTorch (CPU build) + run: | + python -m pip install --upgrade pip + python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + + - name: Index math + working-directory: kernels/window_process + run: python -m unittest test_index_math -v + + - name: CUDA to HIP translation + working-directory: kernels/window_process + run: python hipify_check.py + + - name: GPU test files import and skip cleanly + working-directory: kernels/window_process + run: | + python unit_test.py + python test_model_parity.py diff --git a/kernels/window_process/README.md b/kernels/window_process/README.md new file mode 100644 index 000000000..e4b23e8ac --- /dev/null +++ b/kernels/window_process/README.md @@ -0,0 +1,104 @@ +# Fused window process kernels + +`torch.roll` + `window_partition` and its inverse, fused into a single pass. +Both are pure data movement, so the eager path materialises the tensor twice — +`torch.roll` writes a copy, then `window_partition` permutes and calls +`.contiguous()` for another — where the fused kernel writes once. + +Enable with `--fused_window_process` (see `get_started.md`), or by passing +`fused_window_process=True` to the model. + +## Build + +```bash +cd kernels/window_process +python setup.py install +``` + +The same command builds on ROCm. `CUDAExtension` hipifies its sources when +`torch.version.hip` is set, so there is no separate HIP source tree and no +change to `setup.py`. Select the target architecture with the standard +environment variable, for example MI300X: + +```bash +PYTORCH_ROCM_ARCH=gfx942 python setup.py install +``` + +Do not pass `--offload-arch` through `extra_compile_args`: PyTorch's +`_get_rocm_arch_flags()` skips its own detection as soon as it sees one, which +pins the build to a single GPU. + +## Usage + +```python +from window_process import WindowProcess, WindowProcessReverse + +# (B, H, W, C) -> (B * nH * nW, window_h, window_w, C) +windows = WindowProcess.apply(x, B, H, W, C, -shift_size, window_size) + +# and back +x = WindowProcessReverse.apply(windows, B, H, W, C, shift_size, window_size) +``` + +`shift_size` and `window_size` accept an `int` (isotropic) or an `(h, w)` pair: + +```python +windows = WindowProcess.apply(x, B, H, W, C, (-2, -8), (4, 16)) +``` + +`shift_size` is the negated `torch.roll` shift on the partition path, which is +why the call sites pass `-shift_size` forward and `+shift_size` back. + +## Constraints + +| | | +|---|---| +| dtype | float64, float32, float16, bfloat16 | +| layout | contiguous, channels last in memory: `(..., C)` | +| tiling | `H % window_h == 0` and `W % window_w == 0` | +| shift | `abs(shift) < window size`, per axis | +| size | `B * H * W * C < 2^31` — offsets are computed in `int` | +| grid | `B * nH * nW <= 65535` and `H <= 65535` — CUDA `grid.z` and `grid.y` | + +All of these are enforced by `TORCH_CHECK` before the launch. Violating them +used to read the wrong element, or read out of bounds, without any error. + +## Tests + +| file | needs a GPU | what it covers | +|---|---|---| +| `test_index_math.py` | no | the index arithmetic of all four kernels, transcribed to PyTorch in `reference.py` | +| `hipify_check.py` | no | CUDA → HIP translation is complete, and every launch carries a stream | +| `unit_test.py` | yes | parity with the eager path across dtype × shape × shift, forward and backward | +| `test_model_parity.py` | yes | a `SwinTransformer` gives identical logits and gradients either way | +| `benchmark.py` | yes | wall time and peak memory, fused vs eager | + +```bash +python test_index_math.py # runs anywhere +python hipify_check.py # runs anywhere +python unit_test.py # skips without a GPU or the built extension +python test_model_parity.py +python benchmark.py +``` + +The kernels perform no arithmetic on the values, only gathers, so parity is +asserted with `torch.equal` on every dtype including float16 and bfloat16: any +deviation is an indexing error rather than a rounding one. For the same reason +`gradcheck` adds nothing, even though float64 does dispatch. + +`reference.py` transcribes the kernels' `input_offset` computation to vectorised +PyTorch. That is not an approximation of the kernels — the offset computation is +their entire logic — so it makes their correctness testable on a CPU, with no +compiled extension. + +## Portability + +The kernels use no shared memory, no `__syncthreads()` and no warp-level +primitives. Nothing depends on the warp or wavefront width, so CDNA's 64-wide +wavefront against NVIDIA's 32-wide warp affects occupancy and nothing else. The +three block widths in `best_block_dim()` are multiples of 64, so neither +platform schedules a partial wave; the thresholds were tuned on NVIDIA and are +not measured on CDNA. Compile with `-DSWIN_WP_BLOCK_DIM=N` to override them. + +`hipify_check.py` prints the full mapping and fails if anything is left +untranslated. From 226fccd5c46f4e7c8b1fce818ff310e1a0918b12 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Tue, 25 Aug 2026 21:59:06 +0200 Subject: [PATCH 12/34] test: exercise the row stride fix through a stock model configuration The row stride bug needs nH != nW, and that follows from H != W alone: with a single square window_size, nH = H / window_size and nW = W / window_size. A non-square window is not required, and img_size is documented as `int | tuple(int)` and passed through to_2tuple, so a non-square feature map reaches the blocks through a supported public argument. SwinTransformer(img_size=(256, 128), window_size=8, depths=[2, 2, 2]) puts two shifted blocks at 64x32 and 32x16, giving nH/nW of 8/4 and 4/2. Both have nH > nW, which is the regime where the miscomputed offset leaves the tensor rather than merely pointing at the wrong element, so the merge kernel reads out of bounds twice per forward pass. The new test asserts that the configuration really does produce nH != nW in a shifted block before comparing the two paths, so it cannot quietly stop testing what it claims to test. --- kernels/window_process/README.md | 4 +++ kernels/window_process/test_model_parity.py | 37 +++++++++++++++++++-- 2 files changed, 38 insertions(+), 3 deletions(-) diff --git a/kernels/window_process/README.md b/kernels/window_process/README.md index e4b23e8ac..b355563d2 100644 --- a/kernels/window_process/README.md +++ b/kernels/window_process/README.md @@ -63,6 +63,10 @@ why the call sites pass `-shift_size` forward and `+shift_size` back. All of these are enforced by `TORCH_CHECK` before the launch. Violating them used to read the wrong element, or read out of bounds, without any error. +Note that `nH != nW` follows from `H != W` alone: with a square window, +`nH = H / window_size` and `nW = W / window_size`. A non-square window is not +required, and `img_size` is documented as `int | tuple(int)`. + ## Tests | file | needs a GPU | what it covers | diff --git a/kernels/window_process/test_model_parity.py b/kernels/window_process/test_model_parity.py index ea7035a45..b02b91af5 100644 --- a/kernels/window_process/test_model_parity.py +++ b/kernels/window_process/test_model_parity.py @@ -51,7 +51,7 @@ def set_fused(model, enabled): @requires_everything class TestModelParity(unittest.TestCase): - def _model(self, img_size=56, window_size=7): + def _model(self, img_size=56, window_size=7, depths=(2, 2), num_heads=(3, 6)): torch.manual_seed(0) model = SwinTransformer( img_size=img_size, @@ -59,14 +59,45 @@ def _model(self, img_size=56, window_size=7): in_chans=3, num_classes=10, embed_dim=48, - depths=[2, 2], - num_heads=[3, 6], + depths=list(depths), + num_heads=list(num_heads), window_size=window_size, drop_path_rate=0.0, fused_window_process=False, ) return model.cuda().eval() + def test_non_square_image(self): + """img_size is documented as `int | tuple(int)`, and a non-square one + gives nH != nW inside the shifted blocks -- the configuration the row + stride fix is about. + + For img_size=(256, 128) with window_size=8 the shifted blocks run at + 64x32 and 32x16, so nH/nW are 8/4 and 4/2. Before the fix the merge + kernel indexes past the end of its input on both. + """ + model = self._model(img_size=(256, 128), window_size=8, + depths=(2, 2, 2), num_heads=(3, 6, 12)) + + shifted_non_square = [ + m for m in model.modules() + if type(m).__name__ == 'SwinTransformerBlock' + and m.shift_size > 0 + and (m.input_resolution[0] // m.window_size) + != (m.input_resolution[1] // m.window_size) + ] + self.assertGreater(len(shifted_non_square), 0, + 'this configuration must exercise nH != nW') + + x = torch.randn(1, 3, 256, 128, device='cuda') + with torch.no_grad(): + eager = model(x) + set_fused(model, True) + fused = model(x) + set_fused(model, False) + + self.assertTrue(torch.equal(eager, fused)) + def test_forward_is_identical(self): model = self._model() x = torch.randn(2, 3, 56, 56, device='cuda') From 47f9799805de93fcc6194751baf70a27c9c85311 Mon Sep 17 00:00:00 2001 From: Francesco Brigante Date: Thu, 27 Aug 2026 13:19:47 +0200 Subject: [PATCH 13/34] fix: build the leaf tensor on-device before requiring grad leaf() did tensor.clone().detach().requires_grad_(rg).cuda(): moving a tensor that already requires grad produces a non-leaf, so .grad is never populated on it. Every *_backward parity assertion first checks assertIsNotNone(a.grad) on the eager side and so failed on every dtype/shape/shift -- the compiled backward kernels K2/K4 were never actually compared. Order .cuda() before requires_grad_ so the result is a leaf. With this, unit_test.py exercises backward for the first time (the upstream file never checked a gradient at all). --- kernels/window_process/unit_test.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/kernels/window_process/unit_test.py b/kernels/window_process/unit_test.py index 78b4e221d..1f07e6b9d 100644 --- a/kernels/window_process/unit_test.py +++ b/kernels/window_process/unit_test.py @@ -81,7 +81,9 @@ def reverse_pyt_forward(windows, shift, window, H, W): def leaf(tensor, requires_grad=True): - return tensor.clone().detach().requires_grad_(requires_grad).cuda() + # .cuda() must come before requires_grad_: moving a tensor that already + # requires grad makes the result a non-leaf, and .grad never populates on it. + return tensor.clone().detach().cuda().requires_grad_(requires_grad) def spatial(B, H, W, C, dtype): From d18ae537cb274b3c289242fac5f742ef3f36edb9 Mon Sep 17 00:00:00 2001 From: Francesco Brigante Date: Thu, 27 Aug 2026 13:22:11 +0200 Subject: [PATCH 14/34] test: add coprime nH/nW, an asymmetric shift, and a wide non-square image - unit_test.py / test_index_math.py: shapes with nH,nW = 3,2 and 3,5 -- coprime and both > 1, where a wrong row/window stride cannot alias back to a correct offset the way 4,2 and 2,4 can. - both files: a third shift regime (1, 2) with shift_h != shift_w, which only the per-axis signature added on this branch can express. - test_model_parity.py: split the non-square model test into a tall (nH > nW, out of bounds before the fix) and a wide (nH < nW, silent corruption before the fix) case. --- kernels/window_process/test_index_math.py | 7 +++++-- kernels/window_process/test_model_parity.py | 20 +++++++++++++++----- kernels/window_process/unit_test.py | 8 ++++++-- 3 files changed, 26 insertions(+), 9 deletions(-) diff --git a/kernels/window_process/test_index_math.py b/kernels/window_process/test_index_math.py index b5f2b8aff..23727be9b 100644 --- a/kernels/window_process/test_index_math.py +++ b/kernels/window_process/test_index_math.py @@ -27,6 +27,8 @@ (8, 16, 4, 4), # nH=2, nW=4 wider than tall (8, 32, 4, 8), # nH=2, nW=4 non-square window (32, 8, 8, 4), # nH=4, nW=2 non-square window, taller than wide + (12, 8, 4, 4), # nH=3, nW=2 coprime counts, both > 1 + (12, 20, 4, 4), # nH=3, nW=5 coprime counts, nH < nW (8, 8, 8, 8), # nH == nW == 1 single window (12, 12, 4, 4), # nH == nW == 3 odd window count ] @@ -36,8 +38,9 @@ def _shifts(window_h, window_w): - """The two shift regimes used by Swin: W-MSA (no shift) and SW-MSA (half window).""" - return [(0, 0), (window_h // 2, window_w // 2)] + """W-MSA (no shift), SW-MSA (half window), and an asymmetric shift_h != shift_w + that only the per-axis signature on this branch can express.""" + return [(0, 0), (window_h // 2, window_w // 2), (1, 2)] def _make_spatial(B, H, W, C): diff --git a/kernels/window_process/test_model_parity.py b/kernels/window_process/test_model_parity.py index b02b91af5..d24fa83fa 100644 --- a/kernels/window_process/test_model_parity.py +++ b/kernels/window_process/test_model_parity.py @@ -67,16 +67,18 @@ def _model(self, img_size=56, window_size=7, depths=(2, 2), num_heads=(3, 6)): ) return model.cuda().eval() - def test_non_square_image(self): + def _assert_non_square_parity(self, img_size): """img_size is documented as `int | tuple(int)`, and a non-square one gives nH != nW inside the shifted blocks -- the configuration the row stride fix is about. For img_size=(256, 128) with window_size=8 the shifted blocks run at - 64x32 and 32x16, so nH/nW are 8/4 and 4/2. Before the fix the merge - kernel indexes past the end of its input on both. + 64x32 and 32x16, so nH/nW are 8/4 and 4/2; the transpose (128, 256) + gives nW > nH instead. Before the fix the merge kernel mis-indexes on + both -- out of bounds when nH > nW, silently wrong when nH < nW. """ - model = self._model(img_size=(256, 128), window_size=8, + H, W = img_size + model = self._model(img_size=img_size, window_size=8, depths=(2, 2, 2), num_heads=(3, 6, 12)) shifted_non_square = [ @@ -89,7 +91,7 @@ def test_non_square_image(self): self.assertGreater(len(shifted_non_square), 0, 'this configuration must exercise nH != nW') - x = torch.randn(1, 3, 256, 128, device='cuda') + x = torch.randn(1, 3, H, W, device='cuda') with torch.no_grad(): eager = model(x) set_fused(model, True) @@ -98,6 +100,14 @@ def test_non_square_image(self): self.assertTrue(torch.equal(eager, fused)) + def test_non_square_image_tall(self): + """H > W: the shifted blocks hit nH > nW, the out-of-bounds case.""" + self._assert_non_square_parity((256, 128)) + + def test_non_square_image_wide(self): + """W > H: the shifted blocks hit nH < nW, the silent-corruption case.""" + self._assert_non_square_parity((128, 256)) + def test_forward_is_identical(self): model = self._model() x = torch.randn(2, 3, 56, 56, device='cuda') diff --git a/kernels/window_process/unit_test.py b/kernels/window_process/unit_test.py index 1f07e6b9d..924e16cbd 100644 --- a/kernels/window_process/unit_test.py +++ b/kernels/window_process/unit_test.py @@ -45,6 +45,8 @@ def available_dtypes(): (2, 16, 32, 64, 8, 8), # nH=2, nW=4 silent corruption before the fix (2, 16, 64, 64, 4, 16), # nH=4, nW=4 non-square window (2, 64, 16, 64, 16, 4), # nH=4, nW=4 non-square window, transposed + (2, 24, 16, 64, 8, 8), # nH=3, nW=2 coprime counts, both > 1 + (2, 24, 40, 32, 8, 8), # nH=3, nW=5 coprime counts, nH < nW (2, 24, 24, 32, 24, 24), # nH == nW == 1 a single window # C selects the block width in best_block_dim(): < 384 -> 64 threads, # < 1024 -> 128, otherwise 256. All three branches must be exercised, and C @@ -59,8 +61,10 @@ def available_dtypes(): def shifts_for(window_h, window_w): - """W-MSA (no shift) and SW-MSA (half window), the two regimes Swin uses.""" - return [(0, 0), (window_h // 2, window_w // 2)] + """W-MSA (no shift), SW-MSA (half window, the isotropic shift Swin uses), and + an asymmetric shift_h != shift_w -- a regime only the per-axis signature on + this branch can express, so it needs its own coverage.""" + return [(0, 0), (window_h // 2, window_w // 2), (1, 2)] def pyt_forward(x, shift, window): From c6c6b84334afeebe82dbc90044ec66f984f2eb3e Mon Sep 17 00:00:00 2001 From: Francesco Brigante Date: Thu, 27 Aug 2026 13:40:14 +0200 Subject: [PATCH 15/34] test: also cover a square image with a non-square window nH != nW arises whenever H/window_h != W/window_w, so a square feature map with a rectangular window (48x48, window 6x16 -> nH=8, nW=3) reaches the same regime as a non-square map. Add it to both the index-math and the compiled-kernel shape lists so all four image/window combinations are exercised explicitly. --- kernels/window_process/test_index_math.py | 1 + kernels/window_process/unit_test.py | 1 + 2 files changed, 2 insertions(+) diff --git a/kernels/window_process/test_index_math.py b/kernels/window_process/test_index_math.py index 23727be9b..c48b1ae4f 100644 --- a/kernels/window_process/test_index_math.py +++ b/kernels/window_process/test_index_math.py @@ -27,6 +27,7 @@ (8, 16, 4, 4), # nH=2, nW=4 wider than tall (8, 32, 4, 8), # nH=2, nW=4 non-square window (32, 8, 8, 4), # nH=4, nW=2 non-square window, taller than wide + (24, 24, 4, 8), # nH=6, nW=3 square grid, non-square window (12, 8, 4, 4), # nH=3, nW=2 coprime counts, both > 1 (12, 20, 4, 4), # nH=3, nW=5 coprime counts, nH < nW (8, 8, 8, 8), # nH == nW == 1 single window diff --git a/kernels/window_process/unit_test.py b/kernels/window_process/unit_test.py index 924e16cbd..39743702e 100644 --- a/kernels/window_process/unit_test.py +++ b/kernels/window_process/unit_test.py @@ -45,6 +45,7 @@ def available_dtypes(): (2, 16, 32, 64, 8, 8), # nH=2, nW=4 silent corruption before the fix (2, 16, 64, 64, 4, 16), # nH=4, nW=4 non-square window (2, 64, 16, 64, 16, 4), # nH=4, nW=4 non-square window, transposed + (2, 48, 48, 32, 6, 16), # nH=8, nW=3 square image, non-square window (2, 24, 16, 64, 8, 8), # nH=3, nW=2 coprime counts, both > 1 (2, 24, 40, 32, 8, 8), # nH=3, nW=5 coprime counts, nH < nW (2, 24, 24, 32, 24, 24), # nH == nW == 1 a single window From 3e11b66f29ecf78e1019d2cc4bb3dff61832f133 Mon Sep 17 00:00:00 2001 From: Francesco Brigante Date: Thu, 27 Aug 2026 13:40:15 +0200 Subject: [PATCH 16/34] docs: record measured results on RTX 3080 Add the benchmark table (RTX 3080, torch 2.8.0+cu129, CUDA 12.9) and a Validation section: what ran where, the CUDA parity results, and the ROCm path stated honestly as statically verified only. Include the upstream-vs-branch hipified launch diff from torch.utils.hipify. --- kernels/window_process/README.md | 76 ++++++++++++++++++++++++++++++++ 1 file changed, 76 insertions(+) diff --git a/kernels/window_process/README.md b/kernels/window_process/README.md index b355563d2..06f0b281f 100644 --- a/kernels/window_process/README.md +++ b/kernels/window_process/README.md @@ -95,6 +95,82 @@ PyTorch. That is not an approximation of the kernels — the offset computation their entire logic — so it makes their correctness testable on a CPU, with no compiled extension. +## Performance + +`benchmark.py`, the fused kernel against the `torch.roll` + `window_partition` +it replaces, forward and backward, on an **RTX 3080 (10 GB)**, torch +`2.8.0+cu129` / CUDA 12.9, batch 192, 200 iterations after 10 warm-up: + +| config | op | dtype | eager ms | fused ms | speedup | eager MiB | fused MiB | +|---|---|---|---:|---:|---:|---:|---:| +| stage 1 56×56 w7 | partition | float32 | 7.137 | 2.398 | 2.98× | 1543.5 | 882.0 | +| stage 1 56×56 w7 | partition | float16 | 5.223 | 2.083 | 2.51× | 771.8 | 441.0 | +| stage 1 56×56 w7 | partition | bfloat16 | 5.074 | 1.948 | 2.60× | 771.8 | 441.0 | +| stage 1 56×56 w7 | merge | float32 | 7.191 | 2.405 | 2.99× | 1323.0 | 882.0 | +| stage 1 56×56 w7 | merge | float16 | 5.104 | 2.013 | 2.54× | 661.5 | 441.0 | +| stage 1 56×56 w7 | merge | bfloat16 | 5.150 | 1.998 | 2.58× | 661.5 | 441.0 | +| stage 2 28×28 w7 | partition | float32 | 3.711 | 1.156 | 3.21× | 771.8 | 441.0 | +| stage 2 28×28 w7 | partition | float16 | 2.552 | 0.700 | 3.65× | 392.0 | 224.0 | +| stage 2 28×28 w7 | partition | bfloat16 | 2.783 | 0.726 | 3.84× | 392.0 | 224.0 | +| stage 2 28×28 w7 | merge | float32 | 3.907 | 1.222 | 3.20× | 661.5 | 441.0 | +| stage 2 28×28 w7 | merge | float16 | 2.650 | 0.735 | 3.60× | 336.0 | 224.0 | +| stage 2 28×28 w7 | merge | bfloat16 | 2.630 | 0.714 | 3.68× | 336.0 | 224.0 | +| non-square 32×16 w8 | partition | float32 | 1.218 | 0.417 | 2.92× | 252.0 | 144.0 | +| non-square 32×16 w8 | partition | float16 | 0.870 | 0.335 | 2.59× | 126.0 | 72.0 | +| non-square 32×16 w8 | partition | bfloat16 | 0.871 | 0.340 | 2.56× | 126.0 | 72.0 | +| non-square 32×16 w8 | merge | float32 | 1.224 | 0.414 | 2.95× | 216.0 | 144.0 | +| non-square 32×16 w8 | merge | float16 | 0.872 | 0.346 | 2.52× | 108.0 | 72.0 | +| non-square 32×16 w8 | merge | bfloat16 | 0.867 | 0.343 | 2.53× | 108.0 | 72.0 | +| non-square window 16×64 w4×16 | partition | float32 | 2.434 | 0.814 | 2.99× | 504.0 | 288.0 | +| non-square window 16×64 w4×16 | partition | float16 | 1.681 | 0.658 | 2.55× | 252.0 | 144.0 | +| non-square window 16×64 w4×16 | partition | bfloat16 | 1.692 | 0.660 | 2.56× | 252.0 | 144.0 | +| non-square window 16×64 w4×16 | merge | float32 | 2.402 | 0.801 | 3.00× | 432.0 | 288.0 | +| non-square window 16×64 w4×16 | merge | float16 | 1.708 | 0.658 | 2.59× | 216.0 | 144.0 | +| non-square window 16×64 w4×16 | merge | bfloat16 | 1.707 | 0.655 | 2.61× | 216.0 | 144.0 | + +The path is memory bound, so the gain tracks the second materialisation the +eager path pays (`torch.roll` writes a copy, `window_partition` calls +`.contiguous()` for another) and the fused kernel does not; peak memory drops +with it. `bfloat16` is dispatched directly — the upstream kernel raises on a +`bfloat16` input, so the fused path previously needed an fp32 round trip. + +## Validation + +CUDA is validated end to end. The ROCm path is verified statically only — no AMD +GPU was available. + +| check | where | result | +|---|---|---| +| index math (`reference.py` vs the composed PyTorch ops) | CPU (Apple M1), re-run on the RTX 3080 host | 9/9 | +| CUDA → HIP translation (`hipify_check.py`) | CPU | complete, every launch keeps its stream | +| compile the extension | RTX 3080 · torch 2.8.0+cu129 · CUDA 12.9 · MSVC 14.44 | builds, no source change | +| kernel parity (`unit_test.py`) | RTX 3080 | 15/15 — 4 kernels × fwd/bwd × {f64, f32, f16, bf16} × 13 shapes × 3 shifts, bit-exact | +| model parity (`test_model_parity.py`) | RTX 3080 | 4/4 — logits and gradients identical, square and non-square (tall and wide) | +| bug on the compiled kernel | RTX 3080 | pre-fix: a stock `img_size=(256,128)` model gives different logits with the fused path, and `compute-sanitizer` reports an out-of-bounds read; fixed: bit-identical to eager | +| `compute-sanitizer` (memcheck, initcheck, synccheck), fixed kernel | RTX 3080 | 0 errors — all 4 kernels, fwd + bwd, every shape above | +| ROCm runtime (build + parity on an AMD GPU) | — | not run — no AMD hardware | + +The shapes span every image/window combination: square image with a square +window, square image with a non-square window, and non-square images with square +and with non-square windows, including cases where `nH` and `nW` are coprime. + +The one substantive ROCm difference is the launch stream. `torch.utils.hipify` +run on the upstream sources and on this branch: + +```c +// upstream, hipified +hipLaunchKernelGGL((kernel), dim3(grid), dim3(block), 0, 0, ...); // the HIP null stream + +// this branch, hipified +hipLaunchKernelGGL((kernel), dim3(grid), dim3(block), 0, + at::hip::getCurrentHIPStreamMasqueradingAsCUDA(), ...); +``` + +`setup.py` is unchanged: `CUDAExtension` hipifies its own sources when +`torch.version.hip` is set (`cpp_extension.py`, the `IS_HIP_EXTENSION` branch), +substitutes `hipcc`, and derives `--offload-arch` from `PYTORCH_ROCM_ARCH` — +passing an arch through `extra_compile_args` disables that detection. + ## Portability The kernels use no shared memory, no `__syncthreads()` and no warp-level From 40478fcfc641a5e6b6c4bbfd54a64067fb26de55 Mon Sep 17 00:00:00 2001 From: Francesco Brigante Date: Thu, 27 Aug 2026 14:04:02 +0200 Subject: [PATCH 17/34] fix: make hipify_check.py version-agnostic across torch hipify changes torch >= 2.9 hipify no longer rewrites at::cuda::getCurrentCUDAStream or C10_CUDA_KERNEL_LAUNCH_CHECK (they resolve through the CUDA-compat headers) and wraps the hipLaunchKernelGGL call across lines. The old check hard-failed on both. Now: CUDA-only includes must still translate, every launch must still carry a real stream (matched across newlines, either at::cuda:: or at::hip:: spelling), and the compat shims are only reported. The null-stream launch is still caught. Also install numpy in CI so torch stops warning, and note the version split in the README. --- .github/workflows/window-process.yml | 1 + kernels/window_process/README.md | 13 ++++-- kernels/window_process/hipify_check.py | 56 ++++++++++++++++++-------- 3 files changed, 49 insertions(+), 21 deletions(-) diff --git a/.github/workflows/window-process.yml b/.github/workflows/window-process.yml index 94a9f85de..a81d17bd9 100644 --- a/.github/workflows/window-process.yml +++ b/.github/workflows/window-process.yml @@ -39,6 +39,7 @@ jobs: - name: Install PyTorch (CPU build) run: | python -m pip install --upgrade pip + python -m pip install numpy python -m pip install torch --index-url https://download.pytorch.org/whl/cpu - name: Index math diff --git a/kernels/window_process/README.md b/kernels/window_process/README.md index 06f0b281f..c045c42d1 100644 --- a/kernels/window_process/README.md +++ b/kernels/window_process/README.md @@ -158,14 +158,19 @@ The one substantive ROCm difference is the launch stream. `torch.utils.hipify` run on the upstream sources and on this branch: ```c -// upstream, hipified -hipLaunchKernelGGL((kernel), dim3(grid), dim3(block), 0, 0, ...); // the HIP null stream +// upstream, hipified -- 5th arg is 0: the HIP null stream +hipLaunchKernelGGL((kernel), dim3(grid), dim3(block), 0, 0, ...); -// this branch, hipified +// this branch, hipified -- 5th arg is the current stream hipLaunchKernelGGL((kernel), dim3(grid), dim3(block), 0, - at::hip::getCurrentHIPStreamMasqueradingAsCUDA(), ...); + at::cuda::getCurrentCUDAStream(), ...); ``` +Older hipify (torch <= 2.8) rewrites that getter to +`at::hip::getCurrentHIPStreamMasqueradingAsCUDA()`; torch >= 2.9 keeps the +`at::cuda::` spelling and resolves it through the compatibility headers. Both are +the current stream; only the `0` on the upstream side is wrong. + `setup.py` is unchanged: `CUDAExtension` hipifies its own sources when `torch.version.hip` is set (`cpp_extension.py`, the `IS_HIP_EXTENSION` branch), substitutes `hipcc`, and derives `--offload-arch` from `PYTORCH_ROCM_ARCH` — diff --git a/kernels/window_process/hipify_check.py b/kernels/window_process/hipify_check.py index 72280c751..08de6984b 100644 --- a/kernels/window_process/hipify_check.py +++ b/kernels/window_process/hipify_check.py @@ -12,8 +12,11 @@ # GPU. torch.utils.hipify is pure Python, so this runs anywhere, including in CI # on a CPU runner. # -# It fails if any CUDA-specific symbol survives translation, and prints the -# symbols that are deliberately left alone because HIP implements them natively. +# It fails if a CUDA-only include survives translation, or if a kernel launch +# reaches HIP without an explicit stream (the null-stream bug). It only reports, +# without failing, the symbols HIP keeps as is -- either because it implements +# them natively (blockIdx, __ldg, dim3, ...) or because they resolve through the +# CUDA-compatibility headers (at::cuda::getCurrentCUDAStream on torch >= 2.9). # # python hipify_check.py # verify # python hipify_check.py --diff # verify and show the generated HIP source @@ -32,14 +35,21 @@ SOURCES = ['swin_window_process.cpp', 'swin_window_process_kernel.cu'] -# Present in the CUDA source and required to disappear from the HIP output. -# A survivor here means hipify has no rule for it and the ROCm build would fail -# to compile, or would silently bind to the wrong runtime. +# Includes a HIP build genuinely does not provide: these must be rewritten, or +# the ROCm compile fails outright. hipify has rewritten them in every torch +# version. MUST_BE_TRANSLATED = [ 'cuda_runtime.h', 'cuda_fp16.h', 'ATen/cuda/CUDAContext.h', 'c10/cuda/CUDAException.h', +] + +# API symbols older hipify rewrites to an at::hip / C10_HIP spelling, and newer +# hipify (torch >= 2.9) deliberately leaves alone because they resolve to the HIP +# runtime through the CUDA-compatibility headers. Either outcome is correct, so +# this only reports which one happened -- it is never a failure. +CUDA_COMPAT_SHIMS = [ 'at::cuda::getCurrentCUDAStream', 'C10_CUDA_KERNEL_LAUNCH_CHECK', ] @@ -110,6 +120,11 @@ def report(staging, mapping, show_diff): print(f' SURVIVED {token}') failures.append(f'{name}: {token} was not translated') + shims = [t for t in CUDA_COMPAT_SHIMS if t in before] + for token in shims: + how = 'rewritten' if token not in after else 'kept (resolves via the CUDA-compat headers)' + print(f' shim {token} ({how})') + portable = [t for t in KNOWN_PORTABLE if t in before] for token in portable: if token in after: @@ -117,7 +132,7 @@ def report(staging, mapping, show_diff): else: failures.append(f'{name}: {token} was rewritten but should be portable') - if not translated and not portable: + if not translated and not shims and not portable: print(' nothing device specific in this file') if show_diff: @@ -132,25 +147,32 @@ def report(staging, mapping, show_diff): def check_launch_form(mapping): """Every kernel launch must carry an explicit stream argument after translation. - hipLaunchKernelGGL takes (kernel, grid, block, sharedMem, stream, ...). A - launch that reaches HIP without a stream runs on the null stream, which is - the ROCm form of the default-stream bug. + A launch that reaches HIP on the null stream is the ROCm form of the + default-stream bug. hipify may rewrite ``kernel<<<...>>>(...)`` to + ``hipLaunchKernelGGL(...)`` or (newer HIP accepts the triple chevron) leave + it as is, and it may wrap the call across several lines; and it may or may + not rewrite ``getCurrentCUDAStream`` to ``getCurrentHIPStream...``. Accept + every combination -- what matters is that a real stream getter is passed, + not ``0``. """ hipified = mapping.get('swin_window_process_kernel.cu') if hipified is None: return ['no hipified kernel source to inspect'] text = open(hipified).read() - launches = re.findall(r'hipLaunchKernelGGL\((.*?)\n', text) - if not launches: - return ['no hipLaunchKernelGGL call found in the translated source'] + calls = re.findall(r'hipLaunchKernelGGL\s*\((.*?)\)\s*;', text, re.DOTALL) + calls += re.findall(r'<<<(.*?)>>>', text, re.DOTALL) + if not calls: + return ['no kernel launch found in the translated source'] + stream_getters = ('getCurrentHIPStream', 'getCurrentCUDAStream') failures = [] - for launch in launches: - if 'getCurrentHIPStream' not in launch: - failures.append(f'launch without an explicit stream: {launch.strip()[:70]}') - print(f'\n{len(launches)} kernel launches, all carrying an explicit HIP stream' - if not failures else '') + for call in calls: + if not any(g in call for g in stream_getters): + failures.append('launch without an explicit stream: ' + + ' '.join(call.split())[:80]) + if not failures: + print(f'\n{len(calls)} kernel launches, all carrying an explicit stream') return failures From 18950ee7aea897b90ed50d5fdb3c453327eb1a16 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Thu, 27 Aug 2026 14:42:54 +0200 Subject: [PATCH 18/34] docs: list every dispatched dtype in the WindowProcess docstrings The Args block still said "float32/float16", which predates the AT_DISPATCH_FLOATING_TYPES_AND2 change: float64 and bfloat16 dispatch too, and bfloat16 is the reason the fused path no longer needs an fp32 round trip. State all four on both classes. Docstrings only, no behaviour change. --- kernels/window_process/window_process.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/kernels/window_process/window_process.py b/kernels/window_process/window_process.py index 3fbb2ff63..f64c76c68 100644 --- a/kernels/window_process/window_process.py +++ b/kernels/window_process/window_process.py @@ -24,7 +24,8 @@ class WindowProcess(torch.autograd.Function): """Fused torch.roll + window_partition. Args: - input: (B, H, W, C) contiguous float32/float16 CUDA tensor. + input: (B, H, W, C) contiguous CUDA tensor. float64, float32, + float16 and bfloat16 all dispatch. B, H, W, C: shape of `input`. shift_size: int or (shift_h, shift_w). This is the *negated* torch.roll shift, matching the existing call sites in models/swin_transformer.py. @@ -68,6 +69,7 @@ class WindowProcessReverse(torch.autograd.Function): Args: input: (B * nH * nW, window_h, window_w, C) contiguous CUDA tensor. + float64, float32, float16 and bfloat16 all dispatch. B, H, W, C: shape of the *output* feature map. shift_size: int or (shift_h, shift_w), the torch.roll shift. window_size: int or (window_h, window_w). Must divide H and W per axis. From 215e151a362ed9824ff6c2e45b263f623c9fd945 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Thu, 27 Aug 2026 15:24:12 +0200 Subject: [PATCH 19/34] fix: read through a macro so the kernels compile under HIP Building on an MI300X (gfx942, torch 2.12.0+rocm7.14.0) fails on all four kernels with "no matching function for call to '__ldg'", instantiated for c10::Half. Probed one type at a time: float, double and c10::BFloat16 all have a HIP overload, c10::Half does not. Half has been in the dispatch since upstream, so the fused window process has never been buildable on AMD -- not untested, uncompilable. SWIN_WP_LDG is a macro and not an inline function so that the NVIDIA branch expands to exactly the __ldg(ptr) token sequence that was there before: the CUDA translation unit is unchanged, and the parity results measured on an RTX 3080 still describe the code that ships. The hint is advisory on both platforms, so the AMD branch loses nothing but the request. Also drop the "port to AMD unchanged" claim from the block-width comment, which this commit makes untrue. --- .../swin_window_process_kernel.cu | 36 ++++++++++++++----- 1 file changed, 27 insertions(+), 9 deletions(-) diff --git a/kernels/window_process/swin_window_process_kernel.cu b/kernels/window_process/swin_window_process_kernel.cu index aeee13309..9565dd8a3 100644 --- a/kernels/window_process/swin_window_process_kernel.cu +++ b/kernels/window_process/swin_window_process_kernel.cu @@ -23,14 +23,32 @@ #include #include +// Read-only cached load. +// +// CUDA accepts __ldg on every type the dispatch instantiates. HIP does not: +// it has no __ldg overload for c10::Half, so the four kernels below fail to +// compile under ROCm as written (float, double and c10::BFloat16 are fine -- +// c10::Half is the only one). The hint is advisory in both cases: it asks for +// the read-only path and the compiler is free to ignore it, so dropping it on +// AMD changes no result. +// +// This is a macro rather than an inline function on purpose: on NVIDIA it +// expands to the same __ldg(ptr) token sequence as before, so the CUDA +// translation unit is unchanged by this fix. +#if defined(__HIP_PLATFORM_AMD__) +#define SWIN_WP_LDG(ptr) (*(ptr)) +#else +#define SWIN_WP_LDG(ptr) __ldg(ptr) +#endif + // Width of the block that walks the channel dimension. // // These kernels use no shared memory, no __syncthreads() and no warp-level -// primitives, so this width affects occupancy only and never correctness. That -// is also why they port to AMD unchanged: CDNA schedules 64-wide wavefronts -// against NVIDIA's 32-wide warps, and nothing here depends on that width. All -// three values below are multiples of 64, so neither platform is left running a -// partial wave. +// primitives, so this width affects occupancy only and never correctness. It is +// also why the wavefront width is a non-issue on AMD: CDNA schedules 64-wide +// wavefronts against NVIDIA's 32-wide warps, and nothing here depends on that +// width. All three values below are multiples of 64, so neither platform is +// left running a partial wave. // // The thresholds were tuned on NVIDIA and have not been measured on CDNA. // Compile with -DSWIN_WP_BLOCK_DIM=N to override them while tuning. @@ -83,7 +101,7 @@ __global__ void roll_and_window_partition_forward_cuda_kernel( (blockIdx.z % (nH * nW) / nW * window_h + blockIdx.y - shift_h + H) % H * W * C + (blockIdx.z % nW * window_w + blockIdx.x - shift_w + W) % W * C + i; - output[offset] = (T)(__ldg(input + input_offset)); + output[offset] = (T)(SWIN_WP_LDG(input + input_offset)); } } @@ -115,7 +133,7 @@ __global__ void roll_and_window_partition_backward_cuda_kernel( (blockIdx.y + shift_h + H ) % H % window_h * window_w * C + (blockIdx.x + shift_w + W ) % W % window_w * C + i; - grad_out[offset] = (T)(__ldg(grad_in + input_offset)); + grad_out[offset] = (T)(SWIN_WP_LDG(grad_in + input_offset)); } } @@ -147,7 +165,7 @@ __global__ void window_merge_and_roll_forward_cuda_kernel( (blockIdx.y - shift_h + H) % H % window_h * window_w * C + (blockIdx.x - shift_w + W) % W % window_w * C + i; - output[offset] = (T)(__ldg(input + input_offset)); + output[offset] = (T)(SWIN_WP_LDG(input + input_offset)); } } @@ -180,7 +198,7 @@ __global__ void window_merge_and_roll_backward_cuda_kernel( (blockIdx.z % (nH * nW) / nW * window_h + blockIdx.y + shift_h + H) % H * W * C + (blockIdx.z % nW * window_w + blockIdx.x + shift_w + W) % W * C + i; - grad_out[offset] = (T)(__ldg(grad_in + input_offset)); + grad_out[offset] = (T)(SWIN_WP_LDG(grad_in + input_offset)); } } From 9bfbaf76909e374f5773049770dba98979760442 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Thu, 27 Aug 2026 15:24:13 +0200 Subject: [PATCH 20/34] test: hold model gradients to the platform's own reproducibility test_gradients_are_identical compared a whole-model backward with torch.equal and failed on an MI300X. Running the eager path twice, with no fused kernel involved, shows why: 62 of 63 gradients are bit-identical and one moves by 1.9e-9, because the backward of a GEMM may pick a different reduction split per run. Measured on the same machine: eager vs eager identical 62/63 max abs diff 1.863e-09 fused vs fused identical 62/63 max abs diff 2.794e-09 eager vs fused identical 62/63 max abs diff 1.863e-09 The fused path is exactly as far from eager as eager is from itself, on gradients whose scale is 0.25. So measure the noise first and hold each gradient to what that licenses: bit-exactness wherever eager reproduces itself, and no further than eager is from itself elsewhere. Where the backward is deterministic -- CUDA, in every run so far -- every gradient takes the first branch and this is the strict comparison it replaces. A guard keeps the test from going vacuous if a platform reproduced nothing. Bit-exactness of the kernels themselves is unaffected: unit_test.py asserts it with torch.equal per kernel, and passes 15/15 on the MI300X. --- kernels/window_process/test_model_parity.py | 42 ++++++++++++++++++++- 1 file changed, 40 insertions(+), 2 deletions(-) diff --git a/kernels/window_process/test_model_parity.py b/kernels/window_process/test_model_parity.py index d24fa83fa..076a1e322 100644 --- a/kernels/window_process/test_model_parity.py +++ b/kernels/window_process/test_model_parity.py @@ -33,6 +33,11 @@ MODEL_AVAILABLE = False +# The measured noise floor is one sample, so a gradient that is already +# irreproducible is allowed a few multiples of it before the test calls it a +# regression. +NOISE_HEADROOM = 8 + requires_everything = unittest.skipIf( not (torch.cuda.is_available() and EXTENSION_AVAILABLE and MODEL_AVAILABLE), 'requires a CUDA device, the compiled extension and an importable model') @@ -123,6 +128,24 @@ def test_forward_is_identical(self): self.assertTrue(torch.equal(eager, fused)) def test_gradients_are_identical(self): + """The fused path must not perturb a whole-model backward. + + Bit-exactness of the four kernels themselves is asserted in + unit_test.py, where it holds for every dtype and shape. At model level + the bar has to account for the rest of the network: the backward of a + GEMM is not always reproducible run to run, because the library is free + to pick a different reduction split each time. On an MI300X one of the + 63 gradients moves by ~2e-9 between two *identical* eager runs, so + asserting torch.equal against eager would fail without any fused kernel + being involved. + + So the eager path is run twice first, and each gradient is held to what + that measurement licenses: bit-exactness wherever eager reproduces + itself, and no further from eager than eager is from itself elsewhere. + On a platform where the backward is fully deterministic -- CUDA, in + every run of this test so far -- every gradient takes the first branch + and this is exactly the strict comparison it replaces. + """ model = self._model() x = torch.randn(2, 3, 56, 56, device='cuda') target = torch.randn(2, 10, device='cuda') @@ -133,13 +156,28 @@ def grads(): return [p.grad.clone() for p in model.parameters() if p.grad is not None] eager = grads() + eager_again = grads() # the platform's own run-to-run noise set_fused(model, True) fused = grads() set_fused(model, False) self.assertEqual(len(eager), len(fused)) - for a, b in zip(eager, fused): - self.assertTrue(torch.equal(a, b)) + reproducible = 0 + for i, (a, a2, b) in enumerate(zip(eager, eager_again, fused)): + with self.subTest(parameter=i): + if torch.equal(a, a2): + reproducible += 1 + self.assertTrue(torch.equal(a, b)) + else: + noise = (a - a2).abs().max().item() + delta = (a - b).abs().max().item() + self.assertLessEqual( + delta, NOISE_HEADROOM * noise, + f'gradient {i} moves {delta:.3e} with the fused path, ' + f'against {noise:.3e} between two eager runs') + + # A platform that reproduces nothing would make this test vacuous. + self.assertGreater(reproducible, len(eager) // 2) if __name__ == '__main__': From 39c2ec934d64a287b01989de9b0be23d4202d0eb Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Thu, 27 Aug 2026 15:25:22 +0200 Subject: [PATCH 21/34] docs: record the MI300X run The ROCm path is no longer static-only. Build, kernel parity, model parity and benchmark all ran on an AMD Instinct MI300X (gfx942, torch 2.12.0+rocm7.14.0): 15/15 and 4/4, bit-exact, CUDA-graph capture included through hipGraph. Adds the MI300X benchmark table next to the RTX 3080 one. The speedup is 1.03-2.38x against 2.5-3.8x on the 3080, and the honest reading is that the eager baseline moved: stage 1 takes 7.137 ms on the 3080 and 0.911 ms here, so both paths compress toward a floor. The fused rows also stop scaling with dtype, so at these sizes the kernel is not bandwidth bound on CDNA. Replaces the "tuned on NVIDIA, not measured on CDNA" caveat with the measurement: rebuilding with -DSWIN_WP_BLOCK_DIM shows 64 beating 256 beating 1024, and 1024 making the fused path slower than eager. The heuristic transfers. Also records what the static check could not see -- __ldg translates and still does not compile for c10::Half -- and says so in hipify_check.py's own output, since a symbol-level check cannot see an overload set. --- kernels/window_process/README.md | 100 +++++++++++++++++++++++-- kernels/window_process/hipify_check.py | 4 + 2 files changed, 98 insertions(+), 6 deletions(-) diff --git a/kernels/window_process/README.md b/kernels/window_process/README.md index c045c42d1..579056eff 100644 --- a/kernels/window_process/README.md +++ b/kernels/window_process/README.md @@ -134,10 +134,38 @@ eager path pays (`torch.roll` writes a copy, `window_partition` calls with it. `bfloat16` is dispatched directly — the upstream kernel raises on a `bfloat16` input, so the fused path previously needed an fp32 round trip. +The same benchmark on an **AMD Instinct MI300X**, torch `2.12.0+rocm7.14.0`, +same batch and iteration count: + +| config | op | dtype | eager ms | fused ms | speedup | eager MiB | fused MiB | +|---|---|---|---:|---:|---:|---:|---:| +| stage 1 56×56 w7 | partition | float32 | 0.911 | 0.598 | 1.52× | 1543.5 | 882.0 | +| stage 1 56×56 w7 | partition | float16 | 0.652 | 0.588 | 1.11× | 771.8 | 441.0 | +| stage 1 56×56 w7 | partition | bfloat16 | 0.646 | 0.589 | 1.10× | 771.8 | 441.0 | +| stage 1 56×56 w7 | merge | float32 | 0.916 | 0.600 | 1.53× | 1323.0 | 882.0 | +| stage 1 56×56 w7 | merge | bfloat16 | 0.640 | 0.588 | 1.09× | 661.5 | 441.0 | +| stage 2 28×28 w7 | partition | float32 | 0.449 | 0.188 | 2.38× | 771.8 | 441.0 | +| stage 2 28×28 w7 | partition | bfloat16 | 0.320 | 0.152 | 2.10× | 392.0 | 224.0 | +| stage 2 28×28 w7 | merge | float32 | 0.443 | 0.189 | 2.35× | 661.5 | 441.0 | +| stage 2 28×28 w7 | merge | bfloat16 | 0.320 | 0.152 | 2.10× | 336.0 | 224.0 | +| non-square 32×16 w8 | partition | float32 | 0.140 | 0.102 | 1.36× | 252.0 | 144.0 | +| non-square 32×16 w8 | merge | float32 | 0.139 | 0.103 | 1.35× | 216.0 | 144.0 | +| non-square window 16×64 w4×16 | partition | float32 | 0.288 | 0.200 | 1.44× | 504.0 | 288.0 | +| non-square window 16×64 w4×16 | merge | float32 | 0.285 | 0.200 | 1.42× | 432.0 | 288.0 | + +Memory is identical to CUDA, as it must be — the allocations are the same. Time +is not: 1.03–2.38× here against 2.5–3.8× on the RTX 3080. The eager baseline is +what moved. At stage 1 it takes 7.137 ms on the 3080 and 0.911 ms here, so both +paths are compressed toward a floor and the ratio closes with them. + +The fused rows also stop scaling with dtype: 0.598 ms in float32 against 0.588 +in float16, for half the bytes, while the eager path still scales (0.911 → +0.652). At these sizes the fused kernel is no longer bandwidth bound on CDNA. + ## Validation -CUDA is validated end to end. The ROCm path is verified statically only — no AMD -GPU was available. +Validated end to end on both vendors: an RTX 3080 for CUDA and an AMD Instinct +MI300X (gfx942) for ROCm. | check | where | result | |---|---|---| @@ -148,7 +176,10 @@ GPU was available. | model parity (`test_model_parity.py`) | RTX 3080 | 4/4 — logits and gradients identical, square and non-square (tall and wide) | | bug on the compiled kernel | RTX 3080 | pre-fix: a stock `img_size=(256,128)` model gives different logits with the fused path, and `compute-sanitizer` reports an out-of-bounds read; fixed: bit-identical to eager | | `compute-sanitizer` (memcheck, initcheck, synccheck), fixed kernel | RTX 3080 | 0 errors — all 4 kernels, fwd + bwd, every shape above | -| ROCm runtime (build + parity on an AMD GPU) | — | not run — no AMD hardware | +| ROCm build | MI300X · torch 2.12.0+rocm7.14.0 · HIP 7.14 | builds for `gfx942` with `PYTORCH_ROCM_ARCH`, `setup.py` unchanged — after the `__ldg` fix below | +| kernel parity on ROCm (`unit_test.py`) | MI300X | 15/15, bit-exact, including CUDA-graph capture through hipGraph | +| model parity on ROCm (`test_model_parity.py`) | MI300X | 4/4 | +| benchmark on ROCm | MI300X | 1.03–2.38× fwd+bwd; table above | The shapes span every image/window combination: square image with a square window, square image with a non-square window, and non-square images with square @@ -174,7 +205,50 @@ the current stream; only the `0` on the upstream side is wrong. `setup.py` is unchanged: `CUDAExtension` hipifies its own sources when `torch.version.hip` is set (`cpp_extension.py`, the `IS_HIP_EXTENSION` branch), substitutes `hipcc`, and derives `--offload-arch` from `PYTORCH_ROCM_ARCH` — -passing an arch through `extra_compile_args` disables that detection. +passing an arch through `extra_compile_args` disables that detection. Confirmed +on the MI300X: hipify ran on its own and the compile line carried +`--offload-arch=gfx942`, with no change to `setup.py`. + +### What only a real build found + +`__ldg` translates cleanly — hipify leaves it alone, because HIP has it — and +the kernels still failed to compile, on all four, instantiated for `c10::Half`: + +``` +error: no matching function for call to '__ldg' + 87 | output[offset] = (T)(__ldg(input + input_offset)); +note: in instantiation of ... roll_and_window_partition_forward_cuda_kernel +``` + +Probing one type at a time: `float`, `double` and `c10::BFloat16` have a HIP +overload, `c10::Half` does not. Half has been in the dispatch since upstream, so +the fused window process has never been buildable on AMD — not untested, +uncompilable. `SWIN_WP_LDG` is the fix, and it is a macro rather than an inline +function so that the NVIDIA branch expands to the same `__ldg(ptr)` tokens as +before, leaving the CUDA translation unit untouched. + +This is the limit of `hipify_check.py`, which it now prints: the check is +symbol level. A symbol can translate and still not compile, because the two +platforms do not always give it the same overload set. + +The second thing a real build found is not in this repository at all, but is +worth knowing before reproducing the numbers above: AMD's ROCm 7.14 PyTorch +images ship the ROCm SDK as pip wheels with no `/opt/rocm` tree, and those +wheels carry no rocThrust, hipSPARSE, hipBLAS, hipBLASLt or hipSOLVER headers — +which torch's own headers include when compiling device code. In that state *no* +PyTorch HIP extension compiles; a two-line file whose only content is +`#include ` fails the same way. Installing the matching dev +packages and pointing `CPLUS_INCLUDE_PATH` at them is enough, and touches no +source: + +```bash +apt-get install -y amdrocm-blas-dev7.14 amdrocm-hipblas-common-dev7.14 \ + amdrocm-sparse-dev7.14 amdrocm-solver-dev7.14 \ + librocthrust-dev librocprim-dev +export CPLUS_INCLUDE_PATH=/opt/rocm/core-7.14/include:/usr/include +export PYTORCH_ROCM_ARCH=gfx942 +python setup.py install +``` ## Portability @@ -182,8 +256,22 @@ The kernels use no shared memory, no `__syncthreads()` and no warp-level primitives. Nothing depends on the warp or wavefront width, so CDNA's 64-wide wavefront against NVIDIA's 32-wide warp affects occupancy and nothing else. The three block widths in `best_block_dim()` are multiples of 64, so neither -platform schedules a partial wave; the thresholds were tuned on NVIDIA and are -not measured on CDNA. Compile with `-DSWIN_WP_BLOCK_DIM=N` to override them. +platform schedules a partial wave. + +The thresholds were tuned on NVIDIA. They were then measured on CDNA by +rebuilding with `-DSWIN_WP_BLOCK_DIM=N`, stage 1 / stage 2 float32 on the +MI300X, fused time: + +| block width | stage 1 | stage 2 | +|---|---:|---:| +| 64 (what the heuristic picks here) | **0.598 ms** | **0.188 ms** | +| 256 | 0.632 ms | 0.212 ms | +| 1024 | 1.230 ms | 0.326 ms | + +The NVIDIA-tuned choice wins on CDNA too, and widening hurts monotonically — +at 1024 the fused path becomes *slower than eager* (0.74×). The grid is fixed +by the window geometry, so threads past `C` have nothing to do; swin-tiny's +`C` is 96 at stage 1, and a 1024-wide block leaves 928 lanes idle. `hipify_check.py` prints the full mapping and fails if anything is left untranslated. diff --git a/kernels/window_process/hipify_check.py b/kernels/window_process/hipify_check.py index 08de6984b..14b9a6973 100644 --- a/kernels/window_process/hipify_check.py +++ b/kernels/window_process/hipify_check.py @@ -194,6 +194,10 @@ def main(): raise SystemExit(1) print('\nOK: the CUDA sources translate to HIP with no unmapped symbols.') + print('Symbol level only: a symbol can translate and still not compile, ' + 'because\nHIP and CUDA do not always give it the same overload set. ' + '__ldg is the\ncase in point -- it has no HIP overload for c10::Half. ' + 'Only a real build\nfinds that.') if __name__ == '__main__': From 0b1f9d21ad72d7cb55208ff4233c2b74634787ab Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Fri, 28 Aug 2026 14:11:41 +0200 Subject: [PATCH 22/34] docs: attribute the files written for this branch Every file under kernels/window_process/ carries the NVIDIA header, as the whole subtree was contributed by NVIDIA and the repository puts a provenance block on every source file. Three of the new files are original work rather than derived from it, so they now carry a "Written by" line as well, the same way models/swin_transformer.py and the rest of the repository do: test_index_math.py GPU-free tests of the kernels' index arithmetic test_model_parity.py eager vs fused parity through a real SwinTransformer hipify_check.py CUDA -> HIP translation check with no AMD hardware benchmark.py and reference.py deliberately do not get the line: both were split out of the upstream unit_test.py, which held the reference implementations and a timing test alongside the correctness ones, so that code is NVIDIA's. --- kernels/window_process/hipify_check.py | 1 + kernels/window_process/test_index_math.py | 1 + kernels/window_process/test_model_parity.py | 1 + 3 files changed, 3 insertions(+) diff --git a/kernels/window_process/hipify_check.py b/kernels/window_process/hipify_check.py index 14b9a6973..3fbdd108c 100644 --- a/kernels/window_process/hipify_check.py +++ b/kernels/window_process/hipify_check.py @@ -2,6 +2,7 @@ # Fused kernel for window process for SwinTransformer # Copyright (c) 2022 Nvidia # Licensed under The MIT License [see LICENSE for details] +# Written by Francesco Brigante # -------------------------------------------------------- # Verifies that the CUDA sources translate cleanly to HIP for ROCm. # diff --git a/kernels/window_process/test_index_math.py b/kernels/window_process/test_index_math.py index c48b1ae4f..3b0aacd75 100644 --- a/kernels/window_process/test_index_math.py +++ b/kernels/window_process/test_index_math.py @@ -2,6 +2,7 @@ # Fused kernel for window process for SwinTransformer # Copyright (c) 2022 Nvidia # Licensed under The MIT License [see LICENSE for details] +# Written by Francesco Brigante # -------------------------------------------------------- # Correctness tests for the index arithmetic of the fused window kernels. # diff --git a/kernels/window_process/test_model_parity.py b/kernels/window_process/test_model_parity.py index 076a1e322..daa8b7aa0 100644 --- a/kernels/window_process/test_model_parity.py +++ b/kernels/window_process/test_model_parity.py @@ -2,6 +2,7 @@ # Fused kernel for window process for SwinTransformer # Copyright (c) 2022 Nvidia # Licensed under The MIT License [see LICENSE for details] +# Written by Francesco Brigante # -------------------------------------------------------- # End-to-end check: a SwinTransformer must produce identical output with and # without the fused window kernels. The kernels replace torch.roll plus From be1b675371d9eaa1ae49c02f7adc5794d0ff725f Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Fri, 28 Aug 2026 14:11:53 +0200 Subject: [PATCH 23/34] docs: make the README document the kernels, not the change The README had grown into a record of the work that produced it: a validation matrix naming specific machines, the ROCm session's findings written as a narrative, and the full 24-row benchmark matrix for each vendor. That belongs in the pull request, which is where it also lives. A README stays in the tree long after the change is merged, and a reader landing in this directory needs to know what the kernels are, how to build them, what they refuse to do, and how to check them -- not what happened on one afternoon in August. Cut from 277 lines to 207: - The validation matrix becomes one paragraph: both cards, the two test files and their results, and the compute-sanitizer run, with the explicit note that ROCm has no equivalent. - The benchmark tables keep the partition direction and the two swin-tiny stages, with the merge direction bounded (within 6% on the RTX 3080, 1% on the MI300X) and the non-square configurations given as a band. Every number still measured; benchmark.py prints the full matrix for anyone who wants it. - The ROCm material is now "Building on AMD", written as what a builder needs before starting: PYTORCH_ROCM_ARCH, the --offload-arch trap, the __ldg overload gap, and the missing SDK headers in AMD's own images. Adds a line the file did not have: the two GPU-free checks run in CI. --- kernels/window_process/README.md | 269 ++++++++++++------------------- 1 file changed, 100 insertions(+), 169 deletions(-) diff --git a/kernels/window_process/README.md b/kernels/window_process/README.md index 579056eff..07229a4d1 100644 --- a/kernels/window_process/README.md +++ b/kernels/window_process/README.md @@ -1,9 +1,11 @@ # Fused window process kernels -`torch.roll` + `window_partition` and its inverse, fused into a single pass. +`torch.roll` + `window_partition`, and its inverse, fused into a single pass. + Both are pure data movement, so the eager path materialises the tensor twice — `torch.roll` writes a copy, then `window_partition` permutes and calls -`.contiguous()` for another — where the fused kernel writes once. +`.contiguous()` for another — where the fused kernel writes once. Peak memory +drops with the copy that is no longer made. Enable with `--fused_window_process` (see `get_started.md`), or by passing `fused_window_process=True` to the model. @@ -15,18 +17,7 @@ cd kernels/window_process python setup.py install ``` -The same command builds on ROCm. `CUDAExtension` hipifies its sources when -`torch.version.hip` is set, so there is no separate HIP source tree and no -change to `setup.py`. Select the target architecture with the standard -environment variable, for example MI300X: - -```bash -PYTORCH_ROCM_ARCH=gfx942 python setup.py install -``` - -Do not pass `--offload-arch` through `extra_compile_args`: PyTorch's -`_get_rocm_arch_flags()` skips its own detection as soon as it sees one, which -pins the build to a single GPU. +The same command builds on ROCm; see [Building on AMD](#building-on-amd). ## Usage @@ -60,12 +51,11 @@ why the call sites pass `-shift_size` forward and `+shift_size` back. | size | `B * H * W * C < 2^31` — offsets are computed in `int` | | grid | `B * nH * nW <= 65535` and `H <= 65535` — CUDA `grid.z` and `grid.y` | -All of these are enforced by `TORCH_CHECK` before the launch. Violating them -used to read the wrong element, or read out of bounds, without any error. +All of these are enforced by `TORCH_CHECK` before the launch. Note that `nH != nW` follows from `H != W` alone: with a square window, `nH = H / window_size` and `nW = W / window_size`. A non-square window is not -required, and `img_size` is documented as `int | tuple(int)`. +required to reach that case, and `img_size` is documented as `int | tuple(int)`. ## Tests @@ -85,6 +75,8 @@ python test_model_parity.py python benchmark.py ``` +The first two run in CI on every push, on a CPU runner. + The kernels perform no arithmetic on the values, only gathers, so parity is asserted with `torch.equal` on every dtype including float16 and bfloat16: any deviation is an indexing error rather than a rounding one. For the same reason @@ -92,165 +84,102 @@ deviation is an indexing error rather than a rounding one. For the same reason `reference.py` transcribes the kernels' `input_offset` computation to vectorised PyTorch. That is not an approximation of the kernels — the offset computation is -their entire logic — so it makes their correctness testable on a CPU, with no -compiled extension. +their entire logic — which is what makes them testable on a CPU, with no +compiled extension and no GPU of either vendor. + +## Validation + +Parity is asserted against the eager path on an **RTX 3080** (torch +`2.8.0+cu129`, CUDA 12.9) and an **AMD Instinct MI300X** (gfx942, torch +`2.12.0+rocm7.14.0`): `unit_test.py` 15/15 and `test_model_parity.py` 4/4 on +both, bit-exact, over square and non-square images and windows including +coprime `nH`/`nW`. `compute-sanitizer` (memcheck, initcheck, synccheck) is clean +on the RTX 3080 across all four kernels, forward and backward; ROCm has no +equivalent run. ## Performance -`benchmark.py`, the fused kernel against the `torch.roll` + `window_partition` -it replaces, forward and backward, on an **RTX 3080 (10 GB)**, torch -`2.8.0+cu129` / CUDA 12.9, batch 192, 200 iterations after 10 warm-up: - -| config | op | dtype | eager ms | fused ms | speedup | eager MiB | fused MiB | -|---|---|---|---:|---:|---:|---:|---:| -| stage 1 56×56 w7 | partition | float32 | 7.137 | 2.398 | 2.98× | 1543.5 | 882.0 | -| stage 1 56×56 w7 | partition | float16 | 5.223 | 2.083 | 2.51× | 771.8 | 441.0 | -| stage 1 56×56 w7 | partition | bfloat16 | 5.074 | 1.948 | 2.60× | 771.8 | 441.0 | -| stage 1 56×56 w7 | merge | float32 | 7.191 | 2.405 | 2.99× | 1323.0 | 882.0 | -| stage 1 56×56 w7 | merge | float16 | 5.104 | 2.013 | 2.54× | 661.5 | 441.0 | -| stage 1 56×56 w7 | merge | bfloat16 | 5.150 | 1.998 | 2.58× | 661.5 | 441.0 | -| stage 2 28×28 w7 | partition | float32 | 3.711 | 1.156 | 3.21× | 771.8 | 441.0 | -| stage 2 28×28 w7 | partition | float16 | 2.552 | 0.700 | 3.65× | 392.0 | 224.0 | -| stage 2 28×28 w7 | partition | bfloat16 | 2.783 | 0.726 | 3.84× | 392.0 | 224.0 | -| stage 2 28×28 w7 | merge | float32 | 3.907 | 1.222 | 3.20× | 661.5 | 441.0 | -| stage 2 28×28 w7 | merge | float16 | 2.650 | 0.735 | 3.60× | 336.0 | 224.0 | -| stage 2 28×28 w7 | merge | bfloat16 | 2.630 | 0.714 | 3.68× | 336.0 | 224.0 | -| non-square 32×16 w8 | partition | float32 | 1.218 | 0.417 | 2.92× | 252.0 | 144.0 | -| non-square 32×16 w8 | partition | float16 | 0.870 | 0.335 | 2.59× | 126.0 | 72.0 | -| non-square 32×16 w8 | partition | bfloat16 | 0.871 | 0.340 | 2.56× | 126.0 | 72.0 | -| non-square 32×16 w8 | merge | float32 | 1.224 | 0.414 | 2.95× | 216.0 | 144.0 | -| non-square 32×16 w8 | merge | float16 | 0.872 | 0.346 | 2.52× | 108.0 | 72.0 | -| non-square 32×16 w8 | merge | bfloat16 | 0.867 | 0.343 | 2.53× | 108.0 | 72.0 | -| non-square window 16×64 w4×16 | partition | float32 | 2.434 | 0.814 | 2.99× | 504.0 | 288.0 | -| non-square window 16×64 w4×16 | partition | float16 | 1.681 | 0.658 | 2.55× | 252.0 | 144.0 | -| non-square window 16×64 w4×16 | partition | bfloat16 | 1.692 | 0.660 | 2.56× | 252.0 | 144.0 | -| non-square window 16×64 w4×16 | merge | float32 | 2.402 | 0.801 | 3.00× | 432.0 | 288.0 | -| non-square window 16×64 w4×16 | merge | float16 | 1.708 | 0.658 | 2.59× | 216.0 | 144.0 | -| non-square window 16×64 w4×16 | merge | bfloat16 | 1.707 | 0.655 | 2.61× | 216.0 | 144.0 | - -The path is memory bound, so the gain tracks the second materialisation the -eager path pays (`torch.roll` writes a copy, `window_partition` calls -`.contiguous()` for another) and the fused kernel does not; peak memory drops -with it. `bfloat16` is dispatched directly — the upstream kernel raises on a -`bfloat16` input, so the fused path previously needed an fp32 round trip. - -The same benchmark on an **AMD Instinct MI300X**, torch `2.12.0+rocm7.14.0`, -same batch and iteration count: - -| config | op | dtype | eager ms | fused ms | speedup | eager MiB | fused MiB | -|---|---|---|---:|---:|---:|---:|---:| -| stage 1 56×56 w7 | partition | float32 | 0.911 | 0.598 | 1.52× | 1543.5 | 882.0 | -| stage 1 56×56 w7 | partition | float16 | 0.652 | 0.588 | 1.11× | 771.8 | 441.0 | -| stage 1 56×56 w7 | partition | bfloat16 | 0.646 | 0.589 | 1.10× | 771.8 | 441.0 | -| stage 1 56×56 w7 | merge | float32 | 0.916 | 0.600 | 1.53× | 1323.0 | 882.0 | -| stage 1 56×56 w7 | merge | bfloat16 | 0.640 | 0.588 | 1.09× | 661.5 | 441.0 | -| stage 2 28×28 w7 | partition | float32 | 0.449 | 0.188 | 2.38× | 771.8 | 441.0 | -| stage 2 28×28 w7 | partition | bfloat16 | 0.320 | 0.152 | 2.10× | 392.0 | 224.0 | -| stage 2 28×28 w7 | merge | float32 | 0.443 | 0.189 | 2.35× | 661.5 | 441.0 | -| stage 2 28×28 w7 | merge | bfloat16 | 0.320 | 0.152 | 2.10× | 336.0 | 224.0 | -| non-square 32×16 w8 | partition | float32 | 0.140 | 0.102 | 1.36× | 252.0 | 144.0 | -| non-square 32×16 w8 | merge | float32 | 0.139 | 0.103 | 1.35× | 216.0 | 144.0 | -| non-square window 16×64 w4×16 | partition | float32 | 0.288 | 0.200 | 1.44× | 504.0 | 288.0 | -| non-square window 16×64 w4×16 | merge | float32 | 0.285 | 0.200 | 1.42× | 432.0 | 288.0 | - -Memory is identical to CUDA, as it must be — the allocations are the same. Time -is not: 1.03–2.38× here against 2.5–3.8× on the RTX 3080. The eager baseline is -what moved. At stage 1 it takes 7.137 ms on the 3080 and 0.911 ms here, so both -paths are compressed toward a floor and the ratio closes with them. - -The fused rows also stop scaling with dtype: 0.598 ms in float32 against 0.588 -in float16, for half the bytes, while the eager path still scales (0.911 → -0.652). At these sizes the fused kernel is no longer bandwidth bound on CDNA. +`benchmark.py`, forward + backward, batch 192, 200 iterations after 10 warm-up. +The `partition` direction is shown; `merge` tracks it within 6% at every point +on the RTX 3080 and within 1% on the MI300X. `benchmark.py` prints the full +matrix, both directions and all four configurations. -## Validation +**RTX 3080 (10 GB)** -Validated end to end on both vendors: an RTX 3080 for CUDA and an AMD Instinct -MI300X (gfx942) for ROCm. +| config | dtype | eager ms | fused ms | speedup | eager MiB | fused MiB | +|---|---|---:|---:|---:|---:|---:| +| stage 1 56×56 w7 | float32 | 7.137 | 2.398 | 2.98× | 1543.5 | 882.0 | +| stage 1 56×56 w7 | float16 | 5.223 | 2.083 | 2.51× | 771.8 | 441.0 | +| stage 1 56×56 w7 | bfloat16 | 5.074 | 1.948 | 2.60× | 771.8 | 441.0 | +| stage 2 28×28 w7 | float32 | 3.711 | 1.156 | 3.21× | 771.8 | 441.0 | +| stage 2 28×28 w7 | float16 | 2.552 | 0.700 | 3.65× | 392.0 | 224.0 | +| stage 2 28×28 w7 | bfloat16 | 2.783 | 0.726 | 3.84× | 392.0 | 224.0 | -| check | where | result | -|---|---|---| -| index math (`reference.py` vs the composed PyTorch ops) | CPU (Apple M1), re-run on the RTX 3080 host | 9/9 | -| CUDA → HIP translation (`hipify_check.py`) | CPU | complete, every launch keeps its stream | -| compile the extension | RTX 3080 · torch 2.8.0+cu129 · CUDA 12.9 · MSVC 14.44 | builds, no source change | -| kernel parity (`unit_test.py`) | RTX 3080 | 15/15 — 4 kernels × fwd/bwd × {f64, f32, f16, bf16} × 13 shapes × 3 shifts, bit-exact | -| model parity (`test_model_parity.py`) | RTX 3080 | 4/4 — logits and gradients identical, square and non-square (tall and wide) | -| bug on the compiled kernel | RTX 3080 | pre-fix: a stock `img_size=(256,128)` model gives different logits with the fused path, and `compute-sanitizer` reports an out-of-bounds read; fixed: bit-identical to eager | -| `compute-sanitizer` (memcheck, initcheck, synccheck), fixed kernel | RTX 3080 | 0 errors — all 4 kernels, fwd + bwd, every shape above | -| ROCm build | MI300X · torch 2.12.0+rocm7.14.0 · HIP 7.14 | builds for `gfx942` with `PYTORCH_ROCM_ARCH`, `setup.py` unchanged — after the `__ldg` fix below | -| kernel parity on ROCm (`unit_test.py`) | MI300X | 15/15, bit-exact, including CUDA-graph capture through hipGraph | -| model parity on ROCm (`test_model_parity.py`) | MI300X | 4/4 | -| benchmark on ROCm | MI300X | 1.03–2.38× fwd+bwd; table above | - -The shapes span every image/window combination: square image with a square -window, square image with a non-square window, and non-square images with square -and with non-square windows, including cases where `nH` and `nW` are coprime. - -The one substantive ROCm difference is the launch stream. `torch.utils.hipify` -run on the upstream sources and on this branch: - -```c -// upstream, hipified -- 5th arg is 0: the HIP null stream -hipLaunchKernelGGL((kernel), dim3(grid), dim3(block), 0, 0, ...); - -// this branch, hipified -- 5th arg is the current stream -hipLaunchKernelGGL((kernel), dim3(grid), dim3(block), 0, - at::cuda::getCurrentCUDAStream(), ...); -``` +Non-square configurations (32×16 with a square window, 16×64 with a 4×16 window) +land in the same 2.5–3.0× band. -Older hipify (torch <= 2.8) rewrites that getter to -`at::hip::getCurrentHIPStreamMasqueradingAsCUDA()`; torch >= 2.9 keeps the -`at::cuda::` spelling and resolves it through the compatibility headers. Both are -the current stream; only the `0` on the upstream side is wrong. +**AMD Instinct MI300X** -`setup.py` is unchanged: `CUDAExtension` hipifies its own sources when -`torch.version.hip` is set (`cpp_extension.py`, the `IS_HIP_EXTENSION` branch), -substitutes `hipcc`, and derives `--offload-arch` from `PYTORCH_ROCM_ARCH` — -passing an arch through `extra_compile_args` disables that detection. Confirmed -on the MI300X: hipify ran on its own and the compile line carried -`--offload-arch=gfx942`, with no change to `setup.py`. +| config | dtype | eager ms | fused ms | speedup | eager MiB | fused MiB | +|---|---|---:|---:|---:|---:|---:| +| stage 1 56×56 w7 | float32 | 0.911 | 0.598 | 1.52× | 1543.5 | 882.0 | +| stage 1 56×56 w7 | float16 | 0.652 | 0.588 | 1.11× | 771.8 | 441.0 | +| stage 1 56×56 w7 | bfloat16 | 0.646 | 0.589 | 1.10× | 771.8 | 441.0 | +| stage 2 28×28 w7 | float32 | 0.449 | 0.188 | 2.38× | 771.8 | 441.0 | +| stage 2 28×28 w7 | float16 | 0.323 | 0.152 | 2.12× | 392.0 | 224.0 | +| stage 2 28×28 w7 | bfloat16 | 0.320 | 0.152 | 2.10× | 392.0 | 224.0 | -### What only a real build found +Non-square configurations land in a 1.03–1.44× band. -`__ldg` translates cleanly — hipify leaves it alone, because HIP has it — and -the kernels still failed to compile, on all four, instantiated for `c10::Half`: +Memory is identical across vendors, as it must be — the allocations are the +same. Time is not. The eager baseline is what moved: 7.137 ms at stage 1 on the +3080 against 0.911 ms here, so both paths compress toward a floor and the ratio +closes with them. The fused rows also stop scaling with dtype on the MI300X +(0.598 ms float32 against 0.588 float16, for half the bytes) while the eager +path still does, so at these sizes the fused kernel is no longer bandwidth bound +on CDNA. Block width is not the cause — that is measured below. -``` -error: no matching function for call to '__ldg' - 87 | output[offset] = (T)(__ldg(input + input_offset)); -note: in instantiation of ... roll_and_window_partition_forward_cuda_kernel -``` +`bfloat16` is dispatched directly. The upstream kernel raises on a `bfloat16` +input, so the fused path previously needed an fp32 round trip. + +## Building on AMD -Probing one type at a time: `float`, `double` and `c10::BFloat16` have a HIP -overload, `c10::Half` does not. Half has been in the dispatch since upstream, so -the fused window process has never been buildable on AMD — not untested, -uncompilable. `SWIN_WP_LDG` is the fix, and it is a macro rather than an inline -function so that the NVIDIA branch expands to the same `__ldg(ptr)` tokens as -before, leaving the CUDA translation unit untouched. - -This is the limit of `hipify_check.py`, which it now prints: the check is -symbol level. A symbol can translate and still not compile, because the two -platforms do not always give it the same overload set. - -The second thing a real build found is not in this repository at all, but is -worth knowing before reproducing the numbers above: AMD's ROCm 7.14 PyTorch -images ship the ROCm SDK as pip wheels with no `/opt/rocm` tree, and those -wheels carry no rocThrust, hipSPARSE, hipBLAS, hipBLASLt or hipSOLVER headers — -which torch's own headers include when compiling device code. In that state *no* -PyTorch HIP extension compiles; a two-line file whose only content is -`#include ` fails the same way. Installing the matching dev -packages and pointing `CPLUS_INCLUDE_PATH` at them is enough, and touches no -source: +`setup.py` needs no change. `CUDAExtension` hipifies its own sources when +`torch.version.hip` is set, substitutes `hipcc`, and derives `--offload-arch` +from `PYTORCH_ROCM_ARCH`: ```bash -apt-get install -y amdrocm-blas-dev7.14 amdrocm-hipblas-common-dev7.14 \ - amdrocm-sparse-dev7.14 amdrocm-solver-dev7.14 \ - librocthrust-dev librocprim-dev -export CPLUS_INCLUDE_PATH=/opt/rocm/core-7.14/include:/usr/include -export PYTORCH_ROCM_ARCH=gfx942 -python setup.py install +PYTORCH_ROCM_ARCH=gfx942 python setup.py install ``` -## Portability +Do not pass `--offload-arch` through `extra_compile_args`: PyTorch's +`_get_rocm_arch_flags()` skips its own detection as soon as it sees one, which +pins the build to a single GPU. + +Two things are worth knowing before building: + +- **`__ldg` has no HIP overload for `c10::Half`**, and half is in the dispatch. + Reads go through the `SWIN_WP_LDG` macro for that reason: on NVIDIA it expands + to the same `__ldg(ptr)` tokens as before, on AMD to a plain load. The hint is + advisory on both, so no result changes. `hipify_check.py` cannot catch this + and says so — it is a symbol-level check, and a symbol can translate cleanly + and still not compile. + +- **AMD's ROCm 7.14 PyTorch images ship the SDK as pip wheels** with no + `/opt/rocm` tree, and those wheels carry no rocThrust, hipSPARSE, hipBLAS, + hipBLASLt or hipSOLVER headers — which torch's own headers include when + compiling device code. In that state *no* PyTorch HIP extension compiles; a + file whose only content is `#include ` fails the same way. + Installing the matching dev packages is enough, and touches no source: + + ```bash + apt-get install -y amdrocm-blas-dev7.14 amdrocm-hipblas-common-dev7.14 \ + amdrocm-sparse-dev7.14 amdrocm-solver-dev7.14 \ + librocthrust-dev librocprim-dev + export CPLUS_INCLUDE_PATH=/opt/rocm/core-7.14/include:/usr/include + ``` + +### Portability The kernels use no shared memory, no `__syncthreads()` and no warp-level primitives. Nothing depends on the warp or wavefront width, so CDNA's 64-wide @@ -258,9 +187,8 @@ wavefront against NVIDIA's 32-wide warp affects occupancy and nothing else. The three block widths in `best_block_dim()` are multiples of 64, so neither platform schedules a partial wave. -The thresholds were tuned on NVIDIA. They were then measured on CDNA by -rebuilding with `-DSWIN_WP_BLOCK_DIM=N`, stage 1 / stage 2 float32 on the -MI300X, fused time: +Those thresholds were tuned on NVIDIA, then measured on CDNA by rebuilding with +`-DSWIN_WP_BLOCK_DIM=N` (stage 1 / stage 2, float32, MI300X, fused time): | block width | stage 1 | stage 2 | |---|---:|---:| @@ -268,10 +196,13 @@ MI300X, fused time: | 256 | 0.632 ms | 0.212 ms | | 1024 | 1.230 ms | 0.326 ms | -The NVIDIA-tuned choice wins on CDNA too, and widening hurts monotonically — -at 1024 the fused path becomes *slower than eager* (0.74×). The grid is fixed -by the window geometry, so threads past `C` have nothing to do; swin-tiny's -`C` is 96 at stage 1, and a 1024-wide block leaves 928 lanes idle. +The NVIDIA-tuned choice wins on CDNA too, and widening hurts monotonically — at +1024 the fused path becomes *slower than eager* (0.74×). The grid is fixed by +the window geometry, so threads past `C` have nothing to do; swin-tiny's `C` is +96 at stage 1, and a 1024-wide block leaves 928 lanes idle. -`hipify_check.py` prints the full mapping and fails if anything is left -untranslated. +The one substantive difference between the two platforms is the launch stream. +Upstream passes `0`, the null stream; this version passes +`at::cuda::getCurrentCUDAStream()`, which hipify rewrites to the current HIP +stream. `hipify_check.py` verifies that, prints the full symbol mapping, and +fails if anything is left untranslated. From f87cc9c51a0ef8781e9a2a3e80b534d2c36caed8 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Fri, 28 Aug 2026 14:26:05 +0200 Subject: [PATCH 24/34] docs: improve readability of comments and docstrings Comments and docstrings only. The one change to a .cu file is inside //, so nothing is recompiled; test_index_math and hipify_check are green after it. hipify_check.py carried most of the weight and most of the prose: - The explanatory header was a comment block, so __doc__ was None while main() passed description=__doc__ to argparse -- `--help` printed no description at all. It is now a real module docstring, and argparse takes its summary line. - The three constant comments described what the symbols are. They now say what the check does with them: a survivor in MUST_BE_TRANSLATED is a failure, a rewrite in KNOWN_PORTABLE is a failure, and CUDA_COMPAT_SHIMS never fails. Each comment now matches the branch it sits above. - check_launch_form's docstring ran three claims into one sentence. Split: what the bug is, what varies in the translated spelling, what is rejected. - The success caveat was four print continuations with newlines inside the string literals. It is a named constant, printed as one line. - report() had no docstring; open() without a context manager appeared three times and is now a two-line read() helper. swin_window_process_kernel.cu: the note above best_block_dim still said the thresholds "have not been measured on CDNA". They have been, on an MI300X, and the heuristic transferred -- 64 beats 256 beats 1024, and at 1024 the fused path loses to eager at stage 1. The comment said the opposite of the measurement. reference.py: the header said "no GPU" three times in two sentences. --- kernels/window_process/hipify_check.py | 92 ++++++++++--------- kernels/window_process/reference.py | 11 +-- .../swin_window_process_kernel.cu | 6 +- 3 files changed, 58 insertions(+), 51 deletions(-) diff --git a/kernels/window_process/hipify_check.py b/kernels/window_process/hipify_check.py index 3fbdd108c..6802394ea 100644 --- a/kernels/window_process/hipify_check.py +++ b/kernels/window_process/hipify_check.py @@ -4,24 +4,21 @@ # Licensed under The MIT License [see LICENSE for details] # Written by Francesco Brigante # -------------------------------------------------------- -# Verifies that the CUDA sources translate cleanly to HIP for ROCm. -# -# torch.utils.cpp_extension.CUDAExtension hipifies its sources automatically -# when torch.version.hip is set, so building on ROCm needs no separate source -# tree and no change to setup.py. What it does need is a way to tell whether the -# translation is complete, and that check must not require an AMD GPU -- or any -# GPU. torch.utils.hipify is pure Python, so this runs anywhere, including in CI -# on a CPU runner. -# -# It fails if a CUDA-only include survives translation, or if a kernel launch -# reaches HIP without an explicit stream (the null-stream bug). It only reports, -# without failing, the symbols HIP keeps as is -- either because it implements -# them natively (blockIdx, __ldg, dim3, ...) or because they resolve through the -# CUDA-compatibility headers (at::cuda::getCurrentCUDAStream on torch >= 2.9). -# -# python hipify_check.py # verify -# python hipify_check.py --diff # verify and show the generated HIP source -# -------------------------------------------------------- + +"""Check that the CUDA sources translate cleanly to HIP for ROCm. + +CUDAExtension hipifies its own sources when torch.version.hip is set, so a ROCm +build needs no second source tree. What it does need is a way to tell whether +that translation is complete. torch.utils.hipify is pure Python, so this answers +it with no AMD GPU -- with no GPU at all -- and runs on a CPU runner in CI. + +Fails on a CUDA-only include that survives translation, and on a kernel launch +that reaches HIP without an explicit stream. Reports, without failing, the +symbols HIP keeps verbatim. + + python hipify_check.py # verify + python hipify_check.py --diff # verify, and print the generated HIP +""" import argparse import difflib @@ -36,9 +33,8 @@ SOURCES = ['swin_window_process.cpp', 'swin_window_process_kernel.cu'] -# Includes a HIP build genuinely does not provide: these must be rewritten, or -# the ROCm compile fails outright. hipify has rewritten them in every torch -# version. +# Includes a HIP build does not provide. If hipify leaves one of these behind, +# the ROCm compile fails outright, so a survivor here is a failure. MUST_BE_TRANSLATED = [ 'cuda_runtime.h', 'cuda_fp16.h', @@ -46,17 +42,23 @@ 'c10/cuda/CUDAException.h', ] -# API symbols older hipify rewrites to an at::hip / C10_HIP spelling, and newer -# hipify (torch >= 2.9) deliberately leaves alone because they resolve to the HIP -# runtime through the CUDA-compatibility headers. Either outcome is correct, so -# this only reports which one happened -- it is never a failure. +# Symbols older hipify rewrites to an at::hip / C10_HIP spelling, and newer +# hipify (torch >= 2.9) leaves alone because they already reach the HIP runtime +# through the CUDA-compatibility headers. Both outcomes are correct, so this +# reports which one happened and never fails on it. CUDA_COMPAT_SHIMS = [ 'at::cuda::getCurrentCUDAStream', 'C10_CUDA_KERNEL_LAUNCH_CHECK', ] -# Deliberately unchanged: HIP implements these with the same spelling and the -# same semantics, so translating them would be wrong. +# Symbols HIP spells and implements exactly as CUDA does. Rewriting one of these +# would be a bug in the translation, so a rewrite here is a failure. +# Printed on success. hipify works on symbols, and a symbol is not a signature. +SYMBOL_LEVEL_CAVEAT = """\ +Symbol level only: a symbol can translate and still not compile, because HIP and +CUDA do not always give it the same overload set. __ldg is the case in point -- +it has no HIP overload for c10::Half. Only a real build finds that.""" + KNOWN_PORTABLE = [ '__global__', '__ldg', @@ -70,6 +72,11 @@ ] +def read(path): + with open(path) as handle: + return handle.read() + + def hipify_into(destination): """Translate SOURCES out of place and return {source: hipified path}.""" here = os.path.dirname(os.path.abspath(__file__)) @@ -95,6 +102,7 @@ def hipify_into(destination): def report(staging, mapping, show_diff): + """Classify each source's device-specific symbols and return the failures.""" failures = [] for name in SOURCES: @@ -104,8 +112,8 @@ def report(staging, mapping, show_diff): failures.append(f'{name}: hipify produced no output') continue - before = open(original).read() - after = open(hipified).read() + before = read(original) + after = read(hipified) print(f'\n{name} -> {os.path.basename(hipified)}') @@ -123,7 +131,8 @@ def report(staging, mapping, show_diff): shims = [t for t in CUDA_COMPAT_SHIMS if t in before] for token in shims: - how = 'rewritten' if token not in after else 'kept (resolves via the CUDA-compat headers)' + how = ('rewritten' if token not in after + else 'kept (resolves via the CUDA-compat headers)') print(f' shim {token} ({how})') portable = [t for t in KNOWN_PORTABLE if t in before] @@ -146,21 +155,21 @@ def report(staging, mapping, show_diff): def check_launch_form(mapping): - """Every kernel launch must carry an explicit stream argument after translation. + """Every kernel launch must still carry an explicit stream after translation. A launch that reaches HIP on the null stream is the ROCm form of the - default-stream bug. hipify may rewrite ``kernel<<<...>>>(...)`` to - ``hipLaunchKernelGGL(...)`` or (newer HIP accepts the triple chevron) leave - it as is, and it may wrap the call across several lines; and it may or may - not rewrite ``getCurrentCUDAStream`` to ``getCurrentHIPStream...``. Accept - every combination -- what matters is that a real stream getter is passed, - not ``0``. + default-stream bug. + + The spelling of a launch is not fixed, so this accepts all of it: hipify may + turn the triple chevron into hipLaunchKernelGGL or leave it alone, may wrap + the call over several lines, and may or may not rename getCurrentCUDAStream. + Only one thing is rejected -- a literal 0 where the stream belongs. """ hipified = mapping.get('swin_window_process_kernel.cu') if hipified is None: return ['no hipified kernel source to inspect'] - text = open(hipified).read() + text = read(hipified) calls = re.findall(r'hipLaunchKernelGGL\s*\((.*?)\)\s*;', text, re.DOTALL) calls += re.findall(r'<<<(.*?)>>>', text, re.DOTALL) if not calls: @@ -178,7 +187,7 @@ def check_launch_form(mapping): def main(): - parser = argparse.ArgumentParser(description=__doc__) + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument('--diff', action='store_true', help='print the generated HIP source as a unified diff') args = parser.parse_args() @@ -195,10 +204,7 @@ def main(): raise SystemExit(1) print('\nOK: the CUDA sources translate to HIP with no unmapped symbols.') - print('Symbol level only: a symbol can translate and still not compile, ' - 'because\nHIP and CUDA do not always give it the same overload set. ' - '__ldg is the\ncase in point -- it has no HIP overload for c10::Half. ' - 'Only a real build\nfinds that.') + print(SYMBOL_LEVEL_CAVEAT) if __name__ == '__main__': diff --git a/kernels/window_process/reference.py b/kernels/window_process/reference.py index ede121028..d1fa792aa 100644 --- a/kernels/window_process/reference.py +++ b/kernels/window_process/reference.py @@ -6,12 +6,11 @@ # Device-independent transcription of the index arithmetic used by the four # CUDA kernels in swin_window_process_kernel.cu. # -# The kernels perform no arithmetic on the tensor values: every one of them is a -# pure gather, and all of their logic lives in the computation of `input_offset` -# from `blockIdx`. Transcribing that offset computation to vectorised PyTorch -# therefore reproduces the kernels exactly, on any device, with no compiled -# extension and no GPU. This makes the correctness of the index math testable in -# CI without a GPU runner (see test_index_math.py). +# The kernels perform no arithmetic on the tensor values: every one is a pure +# gather, and all of their logic lives in computing `input_offset` from +# `blockIdx`. Transcribing that computation to vectorised PyTorch therefore +# reproduces the kernels exactly -- which is what makes them testable with no +# compiled extension and no GPU, on a CI runner (see test_index_math.py). # -------------------------------------------------------- import torch diff --git a/kernels/window_process/swin_window_process_kernel.cu b/kernels/window_process/swin_window_process_kernel.cu index 9565dd8a3..2091d1e08 100644 --- a/kernels/window_process/swin_window_process_kernel.cu +++ b/kernels/window_process/swin_window_process_kernel.cu @@ -50,8 +50,10 @@ // width. All three values below are multiples of 64, so neither platform is // left running a partial wave. // -// The thresholds were tuned on NVIDIA and have not been measured on CDNA. -// Compile with -DSWIN_WP_BLOCK_DIM=N to override them while tuning. +// The thresholds were tuned on NVIDIA and have since been measured on CDNA, +// where the same choice wins: on an MI300X 64 beats 256 beats 1024, and at 1024 +// the fused path loses to eager at swin-tiny's stage 1. Compile with +// -DSWIN_WP_BLOCK_DIM=N to override them while tuning. int best_block_dim(int feat_dim){ #ifdef SWIN_WP_BLOCK_DIM (void)feat_dim; From ab790896f3d5e2a218e86687019d1e1d70252709 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Fri, 28 Aug 2026 14:37:35 +0200 Subject: [PATCH 25/34] docs: say plainly what the checks do, and fix what did not match A second readability pass, after reading the branch as a reviewer would. hipify_check.py: the comments assumed the reader already knew ROCm. The module docstring now says what hipify is before using its name, and the three constant comments name the concrete thing each list holds -- CUDA header files that do not exist on an AMD machine, PyTorch symbols that work under either spelling, names both languages share -- instead of the word "symbols" three times. The first of them read "Includes a HIP build does not provide", where "Includes" was a noun and read as a verb. SYMBOL_LEVEL_CAVEAT had also been introduced between the KNOWN_PORTABLE comment and the list it describes; it now follows them both. benchmark.py: main() passed description=__doc__ to argparse while the header was a comment block, so __doc__ was None and --help printed no description at all -- the same defect just fixed in hipify_check.py. The header is now a docstring. The claim that "the expected ceiling is roughly 2x" is gone with it: the measured speedups are 2.5-3.8x on an RTX 3080, so the file predicted a ceiling its own README exceeds. It now says what the fused path saves and leaves the number to the measurement. --iters and --forward-only gained help text. test_index_math.py: _make_spatial and _make_windows seeded the RNG and then built their tensors with torch.arange. Nothing in the file is random, so the seeds did nothing. test_model_parity.py: NOISE_HEADROOM = 8 was justified as "a few multiples", which explains the shape of the bound but not the number. It now carries the measurement it came from -- 1.9e-9 of noise against a gradient scale of 0.254, so 8x of it is still seven orders of magnitude below the signal. Two places disagreed with the code and are now aligned: - README said every constraint is enforced by TORCH_CHECK. Device and contiguity go through CHECK_INPUT, and the channel-last layout is the shape contract rather than a separate check. - The CI workflow header said the GPU parity tests "are not run here", while the workflow runs both files. It runs them to check they skip instead of erroring, which is what the step is named and what both files now say. --- .github/workflows/window-process.yml | 4 +- kernels/window_process/README.md | 8 +++- kernels/window_process/benchmark.py | 31 ++++++------ kernels/window_process/hipify_check.py | 52 ++++++++++++--------- kernels/window_process/test_index_math.py | 2 - kernels/window_process/test_model_parity.py | 10 ++-- 6 files changed, 64 insertions(+), 43 deletions(-) diff --git a/.github/workflows/window-process.yml b/.github/workflows/window-process.yml index a81d17bd9..385bf58cf 100644 --- a/.github/workflows/window-process.yml +++ b/.github/workflows/window-process.yml @@ -7,7 +7,9 @@ # translation can be verified here too. # # The parity tests against the compiled extension (unit_test.py, -# test_model_parity.py) need a CUDA device and are not run here. +# test_model_parity.py) need a CUDA device, so they assert nothing on this +# runner. They are still executed, to catch the case where they fail to import +# or error out instead of skipping. name: window process diff --git a/kernels/window_process/README.md b/kernels/window_process/README.md index 07229a4d1..47e8d9d6a 100644 --- a/kernels/window_process/README.md +++ b/kernels/window_process/README.md @@ -51,7 +51,10 @@ why the call sites pass `-shift_size` forward and `+shift_size` back. | size | `B * H * W * C < 2^31` — offsets are computed in `int` | | grid | `B * nH * nW <= 65535` and `H <= 65535` — CUDA `grid.z` and `grid.y` | -All of these are enforced by `TORCH_CHECK` before the launch. +Shape, tiling, shift and size are checked with `TORCH_CHECK` before the launch; +device and contiguity by `CHECK_INPUT`. The channel-last layout is the shape +contract itself and is not separately verified. Violating any of them used to +read the wrong element, or read out of bounds, with no error at all. Note that `nH != nW` follows from `H != W` alone: with a square window, `nH = H / window_size` and `nW = W / window_size`. A non-square window is not @@ -75,7 +78,8 @@ python test_model_parity.py python benchmark.py ``` -The first two run in CI on every push, on a CPU runner. +The first two run in CI on every push, on a CPU runner. The GPU ones are +executed there too, but only to confirm they skip cleanly rather than error. The kernels perform no arithmetic on the values, only gathers, so parity is asserted with `torch.equal` on every dtype including float16 and bfloat16: any diff --git a/kernels/window_process/benchmark.py b/kernels/window_process/benchmark.py index 98b54cf6c..eaa82f18f 100644 --- a/kernels/window_process/benchmark.py +++ b/kernels/window_process/benchmark.py @@ -3,16 +3,19 @@ # Copyright (c) 2022 Nvidia # Licensed under The MIT License [see LICENSE for details] # -------------------------------------------------------- -# Fused kernels vs the PyTorch ops they replace, in wall time and peak memory. -# -# The eager path materialises the tensor twice -- torch.roll writes a copy, then -# window_partition permutes and calls .contiguous() for another -- where the -# fused kernel writes once. The operation is entirely memory bound, so the -# expected ceiling is roughly 2x. -# -# python benchmark.py # forward + backward, all dtypes -# python benchmark.py --iters 200 --batch 64 -# -------------------------------------------------------- + +"""Fused kernels against the PyTorch ops they replace, in time and peak memory. + +The eager path materialises the tensor twice -- torch.roll writes a copy, then +window_partition permutes and calls .contiguous() for another -- where the fused +kernel writes once. The gain is what that saved traffic is worth on the device, +which varies enough between vendors to be worth measuring rather than +predicting: README.md records what these configurations gave on an RTX 3080 and +on an MI300X. + + python benchmark.py # forward + backward, all dtypes + python benchmark.py --iters 200 --batch 64 +""" import argparse @@ -135,11 +138,13 @@ def run(args): def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument('--iters', type=int, default=100) + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument('--iters', type=int, default=100, + help='timed iterations per configuration') parser.add_argument('--batch', type=int, default=None, help='override the batch size of every config') - parser.add_argument('--forward-only', action='store_true') + parser.add_argument('--forward-only', action='store_true', + help='skip the backward pass') args = parser.parse_args() if not torch.cuda.is_available(): diff --git a/kernels/window_process/hipify_check.py b/kernels/window_process/hipify_check.py index 6802394ea..1a575fb58 100644 --- a/kernels/window_process/hipify_check.py +++ b/kernels/window_process/hipify_check.py @@ -7,14 +7,17 @@ """Check that the CUDA sources translate cleanly to HIP for ROCm. -CUDAExtension hipifies its own sources when torch.version.hip is set, so a ROCm -build needs no second source tree. What it does need is a way to tell whether -that translation is complete. torch.utils.hipify is pure Python, so this answers -it with no AMD GPU -- with no GPU at all -- and runs on a CPU runner in CI. - -Fails on a CUDA-only include that survives translation, and on a kernel launch -that reaches HIP without an explicit stream. Reports, without failing, the -symbols HIP keeps verbatim. +AMD GPUs do not compile CUDA. They compile HIP, a nearly identical language, and +PyTorch gets from one to the other with a source-to-source translator called +hipify: CUDA names go in, HIP names come out. CUDAExtension runs it on its own +when torch.version.hip is set, so a ROCm build needs no second copy of these +sources -- but it gives no signal about whether the translation came out +complete. This is that signal, and because hipify is pure Python it can be had +with no AMD GPU (with no GPU at all) on a CPU runner in CI. + +Two things fail the check: a CUDA-only header that survives translation, and a +kernel launch that reaches HIP without an explicit stream. Everything hipify +leaves untouched is reported, not failed. python hipify_check.py # verify python hipify_check.py --diff # verify, and print the generated HIP @@ -33,8 +36,10 @@ SOURCES = ['swin_window_process.cpp', 'swin_window_process_kernel.cu'] -# Includes a HIP build does not provide. If hipify leaves one of these behind, -# the ROCm compile fails outright, so a survivor here is a failure. +# CUDA header files, which simply do not exist on an AMD machine. hipify has to +# rewrite every one of them to its HIP counterpart -- cuda_runtime.h becomes +# hip/hip_runtime.h, and so on. If one of these names is still there afterwards, +# the ROCm build stops at "no such file", so finding one is a failure. MUST_BE_TRANSLATED = [ 'cuda_runtime.h', 'cuda_fp16.h', @@ -42,23 +47,20 @@ 'c10/cuda/CUDAException.h', ] -# Symbols older hipify rewrites to an at::hip / C10_HIP spelling, and newer -# hipify (torch >= 2.9) leaves alone because they already reach the HIP runtime -# through the CUDA-compatibility headers. Both outcomes are correct, so this -# reports which one happened and never fails on it. +# PyTorch functions and macros that work on ROCm under either name. Older hipify +# renamed them to an at::hip / C10_HIP spelling; torch >= 2.9 leaves them alone, +# because ROCm now ships compatibility headers under which the CUDA spelling +# already resolves to the HIP runtime. Both are correct, so this only reports +# which one happened. CUDA_COMPAT_SHIMS = [ 'at::cuda::getCurrentCUDAStream', 'C10_CUDA_KERNEL_LAUNCH_CHECK', ] -# Symbols HIP spells and implements exactly as CUDA does. Rewriting one of these -# would be a bug in the translation, so a rewrite here is a failure. -# Printed on success. hipify works on symbols, and a symbol is not a signature. -SYMBOL_LEVEL_CAVEAT = """\ -Symbol level only: a symbol can translate and still not compile, because HIP and -CUDA do not always give it the same overload set. __ldg is the case in point -- -it has no HIP overload for c10::Half. Only a real build finds that.""" - +# Names that mean the same thing in both languages: HIP spells blockIdx, +# __global__ and the rest exactly as CUDA does, so hipify is meant to leave them +# alone. One coming back renamed would mean the translation went wrong, so that +# is a failure too. KNOWN_PORTABLE = [ '__global__', '__ldg', @@ -71,6 +73,12 @@ 'at::ScalarType::BFloat16', ] +# Printed on success, because passing this check proves less than it sounds like. +SYMBOL_LEVEL_CAVEAT = """\ +Symbol level only: a symbol can translate and still not compile, because HIP and +CUDA do not always give it the same overload set. __ldg is the case in point -- +it has no HIP overload for c10::Half. Only a real build finds that.""" + def read(path): with open(path) as handle: diff --git a/kernels/window_process/test_index_math.py b/kernels/window_process/test_index_math.py index 3b0aacd75..fd4d2f42e 100644 --- a/kernels/window_process/test_index_math.py +++ b/kernels/window_process/test_index_math.py @@ -46,12 +46,10 @@ def _shifts(window_h, window_w): def _make_spatial(B, H, W, C): - torch.manual_seed(0) return torch.arange(B * H * W * C, dtype=torch.float32).view(B, H, W, C) def _make_windows(B, H, W, C, window_h, window_w): - torch.manual_seed(1) n = B * (H // window_h) * (W // window_w) return torch.arange(n * window_h * window_w * C, dtype=torch.float32).view( n, window_h, window_w, C) diff --git a/kernels/window_process/test_model_parity.py b/kernels/window_process/test_model_parity.py index daa8b7aa0..4124b8736 100644 --- a/kernels/window_process/test_model_parity.py +++ b/kernels/window_process/test_model_parity.py @@ -34,9 +34,13 @@ MODEL_AVAILABLE = False -# The measured noise floor is one sample, so a gradient that is already -# irreproducible is allowed a few multiples of it before the test calls it a -# regression. +# How far a gradient may move under the fused path before the test calls it a +# regression, as a multiple of the run-to-run noise measured on the same machine. +# Only gradients that already fail to reproduce against a second eager run are +# held to this; every other one must still be bit-exact. The noise floor is a +# single sample, hence the headroom: on an MI300X it measured 1.9e-9 against a +# gradient scale of 0.254, so even 8x of it stays seven orders of magnitude below +# the signal. NOISE_HEADROOM = 8 requires_everything = unittest.skipIf( From 783624bdbb55f2af978eea2dd25abe0fe2cd6288 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Fri, 28 Aug 2026 15:33:38 +0200 Subject: [PATCH 26/34] docs: correct what the checks claim, and cover the nH == 1 blind spot An external review of the branch found several statements that the code itself contradicts. Each one here is either corrected or dropped, and the two test gaps behind them are closed. Claims that were wrong: - hipify_check.py said torch >= 2.9 stops rewriting the CUDA-compat spellings. It does not follow the version order at all: torch 2.8 and 2.10 rewrite them, 2.13 does not, which the two CI legs show directly. The check never depended on this, so the prediction is gone rather than re-dated, and the "kept" branch no longer explains a mechanism it cannot observe. - Its docstring listed two failure modes where the code has three, and said everything hipify leaves untouched is reported -- the opposite of how KNOWN_PORTABLE is judged three lines below. - check_launch_form claimed to reject only a literal 0. It rejects any launch that does not name a current-stream getter. - The constraint table implied every precondition guarded a silent wrong read. |shift| >= window computes correctly upstream (verified against reference.py); that check is the model's own documented contract made explicit, not a bug guard. The grid bounds caught an async failure, not a bad read, and dtype and layout were already reported. - The README credited CI with running all three GPU test files on every push. It runs two of them, behind a paths filter; benchmark.py exits non-zero without a device. - best_block_dim's thresholds were said to be measured on CDNA. Both swin-tiny stages have C below the first threshold, so only the width the heuristic picks there was measured; 384 and 1024 remain untested on AMD. - test_index_math.py documented an invocation that fails: there is no package, and reference.py is imported top-level. - The __ldg comment credited CUDA with the overloads. PyTorch supplies the Half and BFloat16 ones, and guards them differently -- which is why only Half breaks under hipcc, and is checkable in any torch checkout with no AMD GPU. Coverage: - nH == 1 was in neither shape list. The legacy row stride is src_y / window_h * nH, and src_y / window_h < nH, so at nH == 1 the wrong stride is never multiplied by anything and the bug is invisible -- while its transpose, nW == 1, is the out-of-bounds case. Both orientations are now in SHAPES, the nH < nW test skips the degenerate one explicitly, and test_legacy_is_invisible_when_nH_is_one asserts the equality rather than leaving it unstated. Also adds a square grid with a non-square window, which the file lacked. - unit_test.py promised a partial trailing pass of the strided channel loop that no C produced: 512/128 and 1024/256 are exact. C = 100 on a 64-wide block covers it. The kernel and C++ diffs here are comments only; no compiled file changes behaviour. A note above the four input_offset computations now says that reference.py transcribes them by hand and that no CPU test can bind the two. --- kernels/window_process/README.md | 45 +++++++++++++------ kernels/window_process/hipify_check.py | 33 ++++++++------ .../window_process/swin_window_process.cpp | 16 +++++-- .../swin_window_process_kernel.cu | 30 +++++++++---- kernels/window_process/test_index_math.py | 37 +++++++++++++-- kernels/window_process/unit_test.py | 8 ++-- 6 files changed, 122 insertions(+), 47 deletions(-) diff --git a/kernels/window_process/README.md b/kernels/window_process/README.md index 47e8d9d6a..38aee80a5 100644 --- a/kernels/window_process/README.md +++ b/kernels/window_process/README.md @@ -53,8 +53,15 @@ why the call sites pass `-shift_size` forward and `+shift_size` back. Shape, tiling, shift and size are checked with `TORCH_CHECK` before the launch; device and contiguity by `CHECK_INPUT`. The channel-last layout is the shape -contract itself and is not separately verified. Violating any of them used to -read the wrong element, or read out of bounds, with no error at all. +contract itself and is not separately verified. + +Three of these are new guards over what was previously undetected: violating the +tiling or size rules used to read the wrong element, or read out of bounds, with +no error raised anywhere, and exceeding the grid bounds used to fail the launch +asynchronously somewhere else. The `shift` row is different — the kernels +compute a larger roll correctly — and is enforced because it is the contract the +model already documents, not because it used to break. dtype and layout were +already reported by the dispatch and by `CHECK_CONTIGUOUS`. Note that `nH != nW` follows from `H != W` alone: with a square window, `nH = H / window_size` and `nW = W / window_size`. A non-square window is not @@ -75,11 +82,14 @@ python test_index_math.py # runs anywhere python hipify_check.py # runs anywhere python unit_test.py # skips without a GPU or the built extension python test_model_parity.py -python benchmark.py +python benchmark.py --iters 200 ``` -The first two run in CI on every push, on a CPU runner. The GPU ones are -executed there too, but only to confirm they skip cleanly rather than error. +The first two run in CI on a CPU runner, on every push that touches this +directory. `unit_test.py` and `test_model_parity.py` are executed there too, but +only to confirm they skip cleanly rather than error; `benchmark.py` is not, since +it exits non-zero without a device. The two CI legs install torch 2.8 and 2.13, +so the checks are exercised against both. The kernels perform no arithmetic on the values, only gathers, so parity is asserted with `torch.equal` on every dtype including float16 and bfloat16: any @@ -104,9 +114,10 @@ equivalent run. ## Performance `benchmark.py`, forward + backward, batch 192, 200 iterations after 10 warm-up. -The `partition` direction is shown; `merge` tracks it within 6% at every point -on the RTX 3080 and within 1% on the MI300X. `benchmark.py` prints the full -matrix, both directions and all four configurations. +The `partition` direction is shown. The `merge` direction's fused time tracks it +within 6% at the configurations listed here on the RTX 3080, and within 1% at +every point of the full matrix on the MI300X. `benchmark.py` prints that matrix, +both directions and all four configurations. **RTX 3080 (10 GB)** @@ -191,8 +202,10 @@ wavefront against NVIDIA's 32-wide warp affects occupancy and nothing else. The three block widths in `best_block_dim()` are multiples of 64, so neither platform schedules a partial wave. -Those thresholds were tuned on NVIDIA, then measured on CDNA by rebuilding with -`-DSWIN_WP_BLOCK_DIM=N` (stage 1 / stage 2, float32, MI300X, fused time): +Those thresholds were tuned on NVIDIA. What was measured on CDNA is the width +the heuristic actually picks at swin-tiny's `C` — 64 at both stages, since 96 and +192 are below the first threshold — by rebuilding with `-DSWIN_WP_BLOCK_DIM=N` +(stage 1 / stage 2, float32, MI300X, fused time): | block width | stage 1 | stage 2 | |---|---:|---:| @@ -201,12 +214,16 @@ Those thresholds were tuned on NVIDIA, then measured on CDNA by rebuilding with | 1024 | 1.230 ms | 0.326 ms | The NVIDIA-tuned choice wins on CDNA too, and widening hurts monotonically — at -1024 the fused path becomes *slower than eager* (0.74×). The grid is fixed by +1024 the fused path becomes *slower than eager* (0.74×). Both configurations sit +in the same branch, so the 384 and 1024 thresholds themselves remain untested on +CDNA. The grid is fixed by the window geometry, so threads past `C` have nothing to do; swin-tiny's `C` is 96 at stage 1, and a 1024-wide block leaves 928 lanes idle. The one substantive difference between the two platforms is the launch stream. Upstream passes `0`, the null stream; this version passes -`at::cuda::getCurrentCUDAStream()`, which hipify rewrites to the current HIP -stream. `hipify_check.py` verifies that, prints the full symbol mapping, and -fails if anything is left untranslated. +`at::cuda::getCurrentCUDAStream()`. Some hipify versions rewrite that to the HIP +spelling and some leave it, which resolves to the HIP runtime anyway — both +reach the current stream, and the behaviour does not follow the version order. +`hipify_check.py` reports which one happened rather than requiring either, prints +the full symbol mapping, and fails if anything is left untranslated. diff --git a/kernels/window_process/hipify_check.py b/kernels/window_process/hipify_check.py index 1a575fb58..a0c8f341b 100644 --- a/kernels/window_process/hipify_check.py +++ b/kernels/window_process/hipify_check.py @@ -15,9 +15,10 @@ complete. This is that signal, and because hipify is pure Python it can be had with no AMD GPU (with no GPU at all) on a CPU runner in CI. -Two things fail the check: a CUDA-only header that survives translation, and a -kernel launch that reaches HIP without an explicit stream. Everything hipify -leaves untouched is reported, not failed. +Three things fail the check: a CUDA-only header that survives translation, a +name that is already portable but gets rewritten anyway, and a kernel launch +that reaches HIP without naming a stream. Symbols that translate, and the two +PyTorch spellings that are valid either way, are reported rather than judged. python hipify_check.py # verify python hipify_check.py --diff # verify, and print the generated HIP @@ -41,17 +42,20 @@ # hip/hip_runtime.h, and so on. If one of these names is still there afterwards, # the ROCm build stops at "no such file", so finding one is a failure. MUST_BE_TRANSLATED = [ + 'cuda.h', 'cuda_runtime.h', 'cuda_fp16.h', 'ATen/cuda/CUDAContext.h', 'c10/cuda/CUDAException.h', ] -# PyTorch functions and macros that work on ROCm under either name. Older hipify -# renamed them to an at::hip / C10_HIP spelling; torch >= 2.9 leaves them alone, -# because ROCm now ships compatibility headers under which the CUDA spelling -# already resolves to the HIP runtime. Both are correct, so this only reports -# which one happened. +# PyTorch functions and macros that work on ROCm under either name. Some hipify +# versions rename them to an at::hip / C10_HIP spelling and some leave the CUDA +# spelling in place, which resolves to the HIP runtime anyway; the behaviour does +# not follow the version order, so this makes no prediction about it. Both +# spellings reach the current HIP stream, so the check reports which one happened +# rather than requiring either -- which is why both CI legs pass unchanged while +# installing torch 2.8 and 2.13 respectively. CUDA_COMPAT_SHIMS = [ 'at::cuda::getCurrentCUDAStream', 'C10_CUDA_KERNEL_LAUNCH_CHECK', @@ -139,8 +143,7 @@ def report(staging, mapping, show_diff): shims = [t for t in CUDA_COMPAT_SHIMS if t in before] for token in shims: - how = ('rewritten' if token not in after - else 'kept (resolves via the CUDA-compat headers)') + how = 'rewritten' if token not in after else 'kept' print(f' shim {token} ({how})') portable = [t for t in KNOWN_PORTABLE if t in before] @@ -150,7 +153,7 @@ def report(staging, mapping, show_diff): else: failures.append(f'{name}: {token} was rewritten but should be portable') - if not translated and not shims and not portable: + if not translated and not survived and not shims and not portable: print(' nothing device specific in this file') if show_diff: @@ -169,9 +172,11 @@ def check_launch_form(mapping): default-stream bug. The spelling of a launch is not fixed, so this accepts all of it: hipify may - turn the triple chevron into hipLaunchKernelGGL or leave it alone, may wrap - the call over several lines, and may or may not rename getCurrentCUDAStream. - Only one thing is rejected -- a literal 0 where the stream belongs. + turn the triple chevron into hipLaunchKernelGGL or leave it alone, and may + wrap the call over several lines. What it requires is narrow and textual: the + launch must name a current-stream getter under either spelling. A launch + handed a stream by any other route -- a variable, getStreamFromPool() -- is + rejected too, which is strict rather than wrong for this source. """ hipified = mapping.get('swin_window_process_kernel.cu') if hipified is None: diff --git a/kernels/window_process/swin_window_process.cpp b/kernels/window_process/swin_window_process.cpp index 3896f44b9..1736b7e0b 100644 --- a/kernels/window_process/swin_window_process.cpp +++ b/kernels/window_process/swin_window_process.cpp @@ -78,10 +78,18 @@ at::Tensor window_merge_and_roll_backward_cuda( // The kernels derive nH = H / window_h and nW = W / window_w with integer -// division and index with int arithmetic. Violating any of the assumptions below -// does not fail: it silently reads the wrong element, or reads out of bounds. -// All four entry points share the same constraints -- the window layout and the -// spatial layout hold the same number of elements -- so one check covers them. +// division and index with int arithmetic. All four entry points share the same +// constraints -- the window layout and the spatial layout hold the same number +// of elements -- so one check covers them. +// +// The checks below are not all guarding the same failure. Divisibility, numel +// and the int32 bound catch what used to be a silent wrong read or an +// out-of-bounds one, with no error raised at any point. The grid bounds catch a +// launch that used to fail asynchronously somewhere else, leaving the output +// uninitialised. The shift bound is neither: the formula stays correct for +// |shift| >= window, which is simply a larger roll. It is here because it is the +// contract the model already documents (`0 <= shift_size < window_size` in +// models/swin_transformer.py), made explicit at the boundary. static void check_window_args( const at::Tensor & tensor, const int B, diff --git a/kernels/window_process/swin_window_process_kernel.cu b/kernels/window_process/swin_window_process_kernel.cu index 2091d1e08..b4ae6b527 100644 --- a/kernels/window_process/swin_window_process_kernel.cu +++ b/kernels/window_process/swin_window_process_kernel.cu @@ -25,12 +25,15 @@ // Read-only cached load. // -// CUDA accepts __ldg on every type the dispatch instantiates. HIP does not: -// it has no __ldg overload for c10::Half, so the four kernels below fail to -// compile under ROCm as written (float, double and c10::BFloat16 are fine -- -// c10::Half is the only one). The hint is advisory in both cases: it asks for -// the read-only path and the compiler is free to ignore it, so dropping it on -// AMD changes no result. +// __ldg is defined for every type the dispatch instantiates only under CUDA. +// float and double come from the toolkit; c10::Half and c10::BFloat16 come from +// PyTorch, and the two are not guarded alike. In torch/headeronly/util, the +// BFloat16 overload is behind `__CUDACC__ || __HIPCC__` and degrades to *ptr on +// ROCm, while the Half overload is behind `__CUDA_ARCH__ || __CUDA__` and so +// does not exist under hipcc at all. c10::Half is therefore the one type that +// fails to compile under ROCm as written. The hint is advisory in both cases: +// it asks for the read-only path and the compiler is free to ignore it, so +// dropping it on AMD changes no result. // // This is a macro rather than an inline function on purpose: on NVIDIA it // expands to the same __ldg(ptr) token sequence as before, so the CUDA @@ -50,9 +53,11 @@ // width. All three values below are multiples of 64, so neither platform is // left running a partial wave. // -// The thresholds were tuned on NVIDIA and have since been measured on CDNA, -// where the same choice wins: on an MI300X 64 beats 256 beats 1024, and at 1024 -// the fused path loses to eager at swin-tiny's stage 1. Compile with +// The thresholds were tuned on NVIDIA. On CDNA the width this heuristic picks +// at swin-tiny's C -- 64, for both stage 1 and stage 2 -- was measured to be the +// right one: on an MI300X 64 beats 256 beats 1024, and at 1024 the fused path +// loses to eager at stage 1. Both those configurations fall in the first branch, +// so the 384 and 1024 thresholds themselves remain NVIDIA-tuned. Compile with // -DSWIN_WP_BLOCK_DIM=N to override them while tuning. int best_block_dim(int feat_dim){ #ifdef SWIN_WP_BLOCK_DIM @@ -76,6 +81,13 @@ int best_block_dim(int feat_dim){ } +// The four kernels below are pure gathers: the whole of each one is the +// input_offset it computes. reference.py transcribes those four expressions to +// PyTorch so they can be tested on a CPU, but nothing binds the two files -- +// only a build checks this source. If you change the arithmetic here, change +// reference.py to match, or test_index_math.py will go on passing against the +// old formula. + // input: [B, H, W, C] // output: [B*nH*nW, window_h, window_w, C] // grid: (window_w, window_h, B*nH*nW) -- x indexes columns, y indexes rows diff --git a/kernels/window_process/test_index_math.py b/kernels/window_process/test_index_math.py index fd4d2f42e..7e6466c81 100644 --- a/kernels/window_process/test_index_math.py +++ b/kernels/window_process/test_index_math.py @@ -10,7 +10,7 @@ # extension, so they can gate the index math in CI. The GPU parity tests for the # compiled kernels live in unit_test.py. # -# python -m unittest kernels.window_process.test_index_math -v +# python test_index_math.py # -------------------------------------------------------- import unittest @@ -33,6 +33,9 @@ (12, 20, 4, 4), # nH=3, nW=5 coprime counts, nH < nW (8, 8, 8, 8), # nH == nW == 1 single window (12, 12, 4, 4), # nH == nW == 3 odd window count + (16, 32, 4, 8), # nH == nW == 4 square grid, non-square window + (4, 16, 4, 4), # nH=1, nW=4 one row of windows: the bug hides here + (16, 4, 4, 4), # nH=4, nW=1 one column of windows: it does not ] BATCH = 2 @@ -165,8 +168,8 @@ def test_legacy_returns_wrong_values_when_nH_less_than_nW(self): checked = 0 for H, W, wh, ww in SHAPES: nH, nW = H // wh, W // ww - if nH >= nW: - continue + if nH >= nW or nH == 1: + continue # nH == 1 hides the bug -- covered separately for sh, sw in _shifts(wh, ww): with self.subTest(H=H, W=W, wh=wh, ww=ww, shift=(sh, sw), nH=nH, nW=nW): numel = BATCH * nH * nW * wh * ww @@ -181,6 +184,34 @@ def test_legacy_returns_wrong_values_when_nH_less_than_nW(self): checked += 1 self.assertGreater(checked, 0, "no nH < nW shape exercised") + def test_legacy_is_invisible_when_nH_is_one(self): + """A single row of windows hides the bug entirely; a single column does not. + + The legacy term is `src_y / window_h * nH`, and `src_y / window_h` runs + over [0, nH). At nH == 1 it is always 0, so the wrong stride is never + multiplied by anything and the result is bit-identical to the correct + one -- even though nH != nW and the grid is as non-square as it gets. + The transpose has no such reprieve: nW == 1 with nH > 1 is the + out-of-bounds case, asserted above. + + This is why H != W is not on its own enough to reproduce the bug, and + why the shape list has to contain both orientations. + """ + checked = 0 + for H, W, wh, ww in SHAPES: + nH, nW = H // wh, W // ww + if nH != 1 or nW == 1: + continue + for sh, sw in _shifts(wh, ww): + with self.subTest(H=H, W=W, wh=wh, ww=ww, shift=(sh, sw), nH=nH, nW=nW): + x = _make_windows(BATCH, H, W, CHANNELS, wh, ww) + want = ref.window_merge_and_roll_forward(x, (sh, sw), (wh, ww), H, W) + legacy = ref.window_merge_and_roll_forward( + x, (sh, sw), (wh, ww), H, W, legacy_row_stride=True) + self.assertTrue(torch.equal(legacy, want)) + checked += 1 + self.assertGreater(checked, 0, "no nH == 1 < nW shape exercised") + class TestIntraWindowModuloIsEquivalent(unittest.TestCase): """The upstream intra-window modulo omits `% H` / `% W`, which is a no-op here. diff --git a/kernels/window_process/unit_test.py b/kernels/window_process/unit_test.py index 39743702e..5f409a7f4 100644 --- a/kernels/window_process/unit_test.py +++ b/kernels/window_process/unit_test.py @@ -53,11 +53,13 @@ def available_dtypes(): # < 1024 -> 128, otherwise 256. All three branches must be exercised, and C # must be smaller than, equal to and larger than the block width so that the # strided `for (i = threadIdx.x; i < C; i += blockDim.x)` loop is covered in - # its partial, exact and multi-pass forms. + # all four of its forms: no full pass, exactly one, an exact number, and a + # ragged one where the last pass is partial. (2, 16, 16, 1, 8, 8), # C < blockDim 64 threads, 63 idle (2, 16, 16, 64, 8, 8), # C == blockDim 64 threads, one pass - (2, 16, 16, 512, 8, 8), # 384 <= C < 1024 128 threads - (2, 16, 16, 1024, 8, 8), # C >= 1024 256 threads + (2, 16, 16, 100, 8, 8), # C % blockDim != 0 64 threads, 2 passes, 36 wide + (2, 16, 16, 512, 8, 8), # 384 <= C < 1024 128 threads, 4 exact passes + (2, 16, 16, 1024, 8, 8), # C >= 1024 256 threads, 4 exact passes ] From e60b65b1c82374aafa78165c969eaafa09f0aa86 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Fri, 28 Aug 2026 15:48:06 +0200 Subject: [PATCH 27/34] docs: claim only what the three observations support Replacing "torch >= 2.9 leaves them alone" with "the behaviour does not follow the version order" traded a wrong claim for an unsupported one: 2.8 and 2.10 rewriting while 2.13 does not is a single change of behaviour in version order, and is consistent with a boundary anywhere in (2.10, 2.13]. What the evidence ruled out was the boundary at 2.9, not the existence of one. Both files now state the three observations and stop there. The README also named the two torch versions the CI legs install, but the workflow pins neither: py3.11 will move on, and the README outlives the pin. It now says the legs install different versions, which the Python support window makes true by construction. The PR body keeps the numbers, being a snapshot of one moment. --- kernels/window_process/README.md | 8 ++++---- kernels/window_process/hipify_check.py | 10 +++++----- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/kernels/window_process/README.md b/kernels/window_process/README.md index 38aee80a5..3f2e12b9a 100644 --- a/kernels/window_process/README.md +++ b/kernels/window_process/README.md @@ -88,8 +88,8 @@ python benchmark.py --iters 200 The first two run in CI on a CPU runner, on every push that touches this directory. `unit_test.py` and `test_model_parity.py` are executed there too, but only to confirm they skip cleanly rather than error; `benchmark.py` is not, since -it exits non-zero without a device. The two CI legs install torch 2.8 and 2.13, -so the checks are exercised against both. +it exits non-zero without a device. The two CI legs install different torch +versions, so both hipify behaviours are exercised. The kernels perform no arithmetic on the values, only gathers, so parity is asserted with `torch.equal` on every dtype including float16 and bfloat16: any @@ -224,6 +224,6 @@ The one substantive difference between the two platforms is the launch stream. Upstream passes `0`, the null stream; this version passes `at::cuda::getCurrentCUDAStream()`. Some hipify versions rewrite that to the HIP spelling and some leave it, which resolves to the HIP runtime anyway — both -reach the current stream, and the behaviour does not follow the version order. -`hipify_check.py` reports which one happened rather than requiring either, prints +reach the current stream. `hipify_check.py` makes no prediction about which: +it reports the one that happened rather than requiring either, prints the full symbol mapping, and fails if anything is left untranslated. diff --git a/kernels/window_process/hipify_check.py b/kernels/window_process/hipify_check.py index a0c8f341b..d1ed06ff4 100644 --- a/kernels/window_process/hipify_check.py +++ b/kernels/window_process/hipify_check.py @@ -51,11 +51,11 @@ # PyTorch functions and macros that work on ROCm under either name. Some hipify # versions rename them to an at::hip / C10_HIP spelling and some leave the CUDA -# spelling in place, which resolves to the HIP runtime anyway; the behaviour does -# not follow the version order, so this makes no prediction about it. Both -# spellings reach the current HIP stream, so the check reports which one happened -# rather than requiring either -- which is why both CI legs pass unchanged while -# installing torch 2.8 and 2.13 respectively. +# spelling in place, which resolves to the HIP runtime anyway: torch 2.8 and 2.10 +# rewrite them, 2.13 does not. Both spellings reach the current HIP stream, so +# the check makes no prediction and accepts either, reporting which one happened +# -- which is why it passes unchanged on CI legs that resolve different torch +# versions. CUDA_COMPAT_SHIMS = [ 'at::cuda::getCurrentCUDAStream', 'C10_CUDA_KERNEL_LAUNCH_CHECK', From 79187209efcc82857427bd3417bacbef4344c2c2 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Fri, 28 Aug 2026 20:14:28 +0200 Subject: [PATCH 28/34] test: exercise the sign, the orientation and the window shape the suite missed Three gaps, each one a case where a test could pass while the thing it names is broken. - test_index_math.py only ever passed positive shifts. The model passes -shift_size into the partition path, so the sign that actually reaches the kernels was the one no test used. _shifts() now returns the negative form as well. - unit_test.py had no nH == 1 shape. That is the orientation where the upstream row-stride bug is invisible: the legacy term is src_y / window_h * nH, src_y / window_h runs over [0, nH), so at nH == 1 it is always 0 and the wrong stride multiplies nothing. A suite that reproduces the bug on nH > nW and nH < nW but never on nH == 1 leaves the reader without the reason H != W is not on its own sufficient. test_index_math.py asserts that invisibility; unit_test.py now covers the same shape against the compiled kernels. - test_non_contiguous_gradient built its gradient with an 8x8 window, which makes the transpose shape-invariant: the test would pass with the two window axes swapped anywhere along the path. It now uses a 4x16 window, where the transpose is observable. The C % blockDim row moves to a non-square window for the same reason, rather than adding a sixteenth shape to cover it separately. The comment above the `blocks > 0` assertion in test_model_parity.py said the count was of shifted blocks. set_fused counts every block carrying the flag. The assertion is a guard against the model exposing no flag at all -- in which case the comparison below proves nothing -- and now says so. --- kernels/window_process/test_index_math.py | 11 +++++++---- kernels/window_process/test_model_parity.py | 5 +++-- kernels/window_process/unit_test.py | 10 ++++++++-- 3 files changed, 18 insertions(+), 8 deletions(-) diff --git a/kernels/window_process/test_index_math.py b/kernels/window_process/test_index_math.py index 7e6466c81..22962b9a9 100644 --- a/kernels/window_process/test_index_math.py +++ b/kernels/window_process/test_index_math.py @@ -2,7 +2,6 @@ # Fused kernel for window process for SwinTransformer # Copyright (c) 2022 Nvidia # Licensed under The MIT License [see LICENSE for details] -# Written by Francesco Brigante # -------------------------------------------------------- # Correctness tests for the index arithmetic of the fused window kernels. # @@ -43,9 +42,13 @@ def _shifts(window_h, window_w): - """W-MSA (no shift), SW-MSA (half window), and an asymmetric shift_h != shift_w - that only the per-axis signature on this branch can express.""" - return [(0, 0), (window_h // 2, window_w // 2), (1, 2)] + """W-MSA (no shift), SW-MSA (half window), an asymmetric shift_h != shift_w + that only the per-axis signature on this branch can express, and the + negative form: models/swin_transformer.py calls the partition path with + -shift_size, so a suite that only ever passes positive shifts leaves the + sign the model actually uses untested.""" + return [(0, 0), (window_h // 2, window_w // 2), (1, 2), + (-(window_h // 2), -(window_w // 2))] def _make_spatial(B, H, W, C): diff --git a/kernels/window_process/test_model_parity.py b/kernels/window_process/test_model_parity.py index 4124b8736..081ea3abb 100644 --- a/kernels/window_process/test_model_parity.py +++ b/kernels/window_process/test_model_parity.py @@ -2,7 +2,6 @@ # Fused kernel for window process for SwinTransformer # Copyright (c) 2022 Nvidia # Licensed under The MIT License [see LICENSE for details] -# Written by Francesco Brigante # -------------------------------------------------------- # End-to-end check: a SwinTransformer must produce identical output with and # without the fused window kernels. The kernels replace torch.roll plus @@ -128,7 +127,9 @@ def test_forward_is_identical(self): fused = model(x) set_fused(model, False) - # The shifted blocks are the ones that take the fused path at all. + # set_fused counts every block carrying the flag, not just the shifted + # ones that route through the kernels. Zero would mean the model exposes + # no flag at all, and that the comparison below proved nothing. self.assertGreater(blocks, 0, 'no block exposes fused_window_process') self.assertTrue(torch.equal(eager, fused)) diff --git a/kernels/window_process/unit_test.py b/kernels/window_process/unit_test.py index 5f409a7f4..52880ef5c 100644 --- a/kernels/window_process/unit_test.py +++ b/kernels/window_process/unit_test.py @@ -43,6 +43,7 @@ def available_dtypes(): (2, 56, 56, 96, 7, 7), # nH == nW == 8 the ImageNet configuration (2, 32, 16, 64, 8, 8), # nH=4, nW=2 out-of-bounds read before the fix (2, 16, 32, 64, 8, 8), # nH=2, nW=4 silent corruption before the fix + (2, 8, 32, 64, 8, 8), # nH=1, nW=4 one row of windows: hides the bug (2, 16, 64, 64, 4, 16), # nH=4, nW=4 non-square window (2, 64, 16, 64, 16, 4), # nH=4, nW=4 non-square window, transposed (2, 48, 48, 32, 6, 16), # nH=8, nW=3 square image, non-square window @@ -57,7 +58,8 @@ def available_dtypes(): # ragged one where the last pass is partial. (2, 16, 16, 1, 8, 8), # C < blockDim 64 threads, 63 idle (2, 16, 16, 64, 8, 8), # C == blockDim 64 threads, one pass - (2, 16, 16, 100, 8, 8), # C % blockDim != 0 64 threads, 2 passes, 36 wide + (2, 16, 64, 100, 4, 16), # C % blockDim != 0 64 threads, 2 passes, 36 wide, + # and a non-square window with it (2, 16, 16, 512, 8, 8), # 384 <= C < 1024 128 threads, 4 exact passes (2, 16, 16, 1024, 8, 8), # C >= 1024 256 threads, 4 exact passes ] @@ -179,8 +181,12 @@ def test_non_contiguous_gradient(self): """An upstream op can hand the backward a non-contiguous gradient. The C++ side asserts contiguity, so window_process.py normalises it. + + The window is non-square on purpose: with 8x8 the transpose below is + shape-invariant, so this would pass even if the two window axes were + swapped somewhere along the path. """ - B, H, W, C, window, shift = 2, 16, 16, 32, (8, 8), (4, 4) + B, H, W, C, window, shift = 2, 16, 16, 32, (4, 16), (2, 8) n = B * (H // window[0]) * (W // window[1]) x = spatial(B, H, W, C, torch.float32) grad = torch.randn((n, window[1], window[0], C)).cuda().transpose(1, 2) From 5b63875ecbf3945ee9810b884e3d0e06feda6cc4 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Fri, 28 Aug 2026 20:14:42 +0200 Subject: [PATCH 29/34] fix: stop asserting that __ldg survives translation __ldg was in KNOWN_PORTABLE, so the check required the token to still be there after hipify. It is portable as a symbol -- HIP defines it, hipify leaves it alone -- but not for every type, which is exactly the caveat the script prints on success. And since the kernels now read through SWIN_WP_LDG, the only __ldg left in the source sits in the #else branch of that macro, which hipcc never compiles. The assertion was therefore about dead code: it would have kept passing had the AMD branch been wrong, and it contradicted the caveat three lines below. The macro is reported instead, so the reader sees which arm the build takes rather than a claim about a token that is not in it. Also records a fourth hipify observation. The MI300X container runs torch 2.12, which leaves the at::cuda spelling in place -- so the boundary sits between 2.10 and 2.12, not between 2.10 and 2.13. More to the point, 2.12 is the version that produced the ROCm build, so the kept spelling is not merely accepted by this check: it has compiled under hipcc and run bit-exact on hardware. --- kernels/window_process/hipify_check.py | 27 +++++++++++++++++++------- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/kernels/window_process/hipify_check.py b/kernels/window_process/hipify_check.py index d1ed06ff4..66f31a587 100644 --- a/kernels/window_process/hipify_check.py +++ b/kernels/window_process/hipify_check.py @@ -2,7 +2,6 @@ # Fused kernel for window process for SwinTransformer # Copyright (c) 2022 Nvidia # Licensed under The MIT License [see LICENSE for details] -# Written by Francesco Brigante # -------------------------------------------------------- """Check that the CUDA sources translate cleanly to HIP for ROCm. @@ -52,10 +51,11 @@ # PyTorch functions and macros that work on ROCm under either name. Some hipify # versions rename them to an at::hip / C10_HIP spelling and some leave the CUDA # spelling in place, which resolves to the HIP runtime anyway: torch 2.8 and 2.10 -# rewrite them, 2.13 does not. Both spellings reach the current HIP stream, so -# the check makes no prediction and accepts either, reporting which one happened -# -- which is why it passes unchanged on CI legs that resolve different torch -# versions. +# rewrite them, 2.12 and 2.13 do not. Both spellings reach the current HIP +# stream -- the kept form is the one the MI300X build below actually compiled and +# ran, on torch 2.12 -- so the check makes no prediction and accepts either, +# reporting which one happened. That is why it passes unchanged on CI legs that +# pin different torch versions. CUDA_COMPAT_SHIMS = [ 'at::cuda::getCurrentCUDAStream', 'C10_CUDA_KERNEL_LAUNCH_CHECK', @@ -67,7 +67,6 @@ # is a failure too. KNOWN_PORTABLE = [ '__global__', - '__ldg', 'blockIdx', 'threadIdx', 'blockDim', @@ -77,6 +76,15 @@ 'at::ScalarType::BFloat16', ] +# __ldg is deliberately not in the list above. It is portable as a *symbol* -- +# HIP defines it, and hipify leaves it alone -- but not for every type, which is +# the whole point of the caveat printed on success. Since the kernels read +# through SWIN_WP_LDG, the only __ldg left in the source is inside the #else +# branch of that macro, which hipcc never compiles: asserting that the token +# survives translation would be asserting something about dead code. The macro +# is reported instead, so the reader sees which arm the build will take. +LDG_MACRO = 'SWIN_WP_LDG' + # Printed on success, because passing this check proves less than it sounds like. SYMBOL_LEVEL_CAVEAT = """\ Symbol level only: a symbol can translate and still not compile, because HIP and @@ -146,6 +154,10 @@ def report(staging, mapping, show_diff): how = 'rewritten' if token not in after else 'kept' print(f' shim {token} ({how})') + if LDG_MACRO in before: + print(f' macro {LDG_MACRO} ' + '(__ldg on NVIDIA, plain load on AMD -- see the note below)') + portable = [t for t in KNOWN_PORTABLE if t in before] for token in portable: if token in after: @@ -153,7 +165,8 @@ def report(staging, mapping, show_diff): else: failures.append(f'{name}: {token} was rewritten but should be portable') - if not translated and not survived and not shims and not portable: + if not translated and not survived and not shims and not portable \ + and LDG_MACRO not in before: print(' nothing device specific in this file') if show_diff: From e6a9adc25f6fb3526651cea27d2b3fd2c4a56e50 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Fri, 28 Aug 2026 20:14:42 +0200 Subject: [PATCH 30/34] docs: bound the shift claim where the unsigned arithmetic stops agreeing Both files said the kernels compute a larger roll correctly for |shift| >= window, with no upper bound. There is one. The operand of the modulo is `... + blockIdx.y - shift_h + H`, and blockIdx.y is unsigned, so the whole expression is unsigned. While |shift| <= H the `+ H` keeps it non-negative and the arithmetic is the one reference.py transcribes. Past that it does not go negative -- it wraps modulo 2^32, and `% H` still yields an in-range index. The failure is a silently wrong read, not an out-of-bounds one, which is the opposite of what an unbounded claim invites a reader to assume. reference.py gains the same caveat from its own side: Python's % floors where the kernel wraps, so the transcription is exact inside the supported domain and more forgiving outside it. The TORCH_CHECKs admit a narrower range than either, so no accepted input reaches the divergence. --- kernels/window_process/reference.py | 9 +++++++++ kernels/window_process/swin_window_process.cpp | 6 +++++- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/kernels/window_process/reference.py b/kernels/window_process/reference.py index d1fa792aa..215e0426d 100644 --- a/kernels/window_process/reference.py +++ b/kernels/window_process/reference.py @@ -11,6 +11,15 @@ # `blockIdx`. Transcribing that computation to vectorised PyTorch therefore # reproduces the kernels exactly -- which is what makes them testable with no # compiled extension and no GPU, on a CI runner (see test_index_math.py). +# +# Exactly, but not unconditionally. Python's % and // floor, C++'s truncate, and +# in the kernels the expression is unsigned as well, so it wraps modulo 2^32 +# where this file would carry a negative value through. The two agree for as +# long as the operand of every % stays non-negative, which the `+ H` and `+ W` +# terms guarantee while |shift| <= H and |shift| <= W -- a wider range than the +# launcher's TORCH_CHECKs admit, so inside the supported domain the two are the +# same arithmetic. Outside it they are not, and this file is the more forgiving +# of the two. # -------------------------------------------------------- import torch diff --git a/kernels/window_process/swin_window_process.cpp b/kernels/window_process/swin_window_process.cpp index 1736b7e0b..fec0bc8b6 100644 --- a/kernels/window_process/swin_window_process.cpp +++ b/kernels/window_process/swin_window_process.cpp @@ -87,7 +87,11 @@ at::Tensor window_merge_and_roll_backward_cuda( // out-of-bounds one, with no error raised at any point. The grid bounds catch a // launch that used to fail asynchronously somewhere else, leaving the output // uninitialised. The shift bound is neither: the formula stays correct for -// |shift| >= window, which is simply a larger roll. It is here because it is the +// window <= |shift| <= H (and <= W on the other axis), which is simply a larger +// roll. Past that the `+ H` no longer keeps the operand of `%` non-negative, +// and since the block index is unsigned the expression wraps modulo 2^32 rather +// than going negative -- `% H` still yields an in-range index, so the read is +// silently wrong rather than out of bounds. The bound is here because it is the // contract the model already documents (`0 <= shift_size < window_size` in // models/swin_transformer.py), made explicit at the boundary. static void check_window_args( From 20301397d7561e0f01fcf838945494e04d493564 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Fri, 28 Aug 2026 20:14:53 +0200 Subject: [PATCH 31/34] docs: tighten the three comment blocks in the kernel source Same facts, a third shorter. The __ldg block spent four lines on where float and double come from before reaching the point, and split the "why a macro" rationale into its own paragraph; both fold into the argument they support. The block-width note repeated the wavefront reasoning twice. The reference.py warning loses a clause without losing the warning. One correction rather than a trim: the Half overload guard is now cited in full, `(__CUDA_ARCH__ >= 350) || (__clang__ && __CUDA__)`. The shortened form here was not what the header says. --- .../swin_window_process_kernel.cu | 57 ++++++++----------- 1 file changed, 25 insertions(+), 32 deletions(-) diff --git a/kernels/window_process/swin_window_process_kernel.cu b/kernels/window_process/swin_window_process_kernel.cu index b4ae6b527..f50e667c1 100644 --- a/kernels/window_process/swin_window_process_kernel.cu +++ b/kernels/window_process/swin_window_process_kernel.cu @@ -23,21 +23,17 @@ #include #include -// Read-only cached load. +// Read-only cached load: __ldg on NVIDIA, a plain load on AMD. // -// __ldg is defined for every type the dispatch instantiates only under CUDA. -// float and double come from the toolkit; c10::Half and c10::BFloat16 come from -// PyTorch, and the two are not guarded alike. In torch/headeronly/util, the -// BFloat16 overload is behind `__CUDACC__ || __HIPCC__` and degrades to *ptr on -// ROCm, while the Half overload is behind `__CUDA_ARCH__ || __CUDA__` and so -// does not exist under hipcc at all. c10::Half is therefore the one type that -// fails to compile under ROCm as written. The hint is advisory in both cases: -// it asks for the read-only path and the compiler is free to ignore it, so -// dropping it on AMD changes no result. -// -// This is a macro rather than an inline function on purpose: on NVIDIA it -// expands to the same __ldg(ptr) token sequence as before, so the CUDA -// translation unit is unchanged by this fix. +// __ldg covers every dispatched type only under CUDA. c10::Half and +// c10::BFloat16 come from torch/headeronly/util and are not guarded alike: +// BFloat16 sits behind `__CUDACC__ || __HIPCC__` and degrades to *ptr on ROCm, +// Half behind `(__CUDA_ARCH__ >= 350) || (__clang__ && __CUDA__)`, which no +// hipcc build satisfies. Half is in the dispatch, so this file does not compile +// under ROCm as written. The hint is advisory in both languages -- it asks for +// the read-only path and may be ignored -- so dropping it on AMD changes no +// result. A macro rather than an inline function so that the NVIDIA build still +// sees the same __ldg(ptr) tokens as before. #if defined(__HIP_PLATFORM_AMD__) #define SWIN_WP_LDG(ptr) (*(ptr)) #else @@ -46,19 +42,17 @@ // Width of the block that walks the channel dimension. // -// These kernels use no shared memory, no __syncthreads() and no warp-level -// primitives, so this width affects occupancy only and never correctness. It is -// also why the wavefront width is a non-issue on AMD: CDNA schedules 64-wide -// wavefronts against NVIDIA's 32-wide warps, and nothing here depends on that -// width. All three values below are multiples of 64, so neither platform is -// left running a partial wave. +// No shared memory, no __syncthreads(), no warp-level primitives: this width +// affects occupancy only, never correctness, and nothing here depends on the +// warp or wavefront width -- CDNA's 64 against NVIDIA's 32 is immaterial. All +// three values below are multiples of 64, so neither platform runs a partial +// wave. // -// The thresholds were tuned on NVIDIA. On CDNA the width this heuristic picks -// at swin-tiny's C -- 64, for both stage 1 and stage 2 -- was measured to be the -// right one: on an MI300X 64 beats 256 beats 1024, and at 1024 the fused path -// loses to eager at stage 1. Both those configurations fall in the first branch, -// so the 384 and 1024 thresholds themselves remain NVIDIA-tuned. Compile with -// -DSWIN_WP_BLOCK_DIM=N to override them while tuning. +// The thresholds are NVIDIA-tuned. What CDNA confirms is the width they pick at +// swin-tiny's C, 64 at both stages: on an MI300X 64 beats 256 beats 1024, and +// at 1024 the fused path loses to eager. Both stages land in the first branch, +// so the 384 and 1024 thresholds themselves remain untested there. Override +// with -DSWIN_WP_BLOCK_DIM=N. int best_block_dim(int feat_dim){ #ifdef SWIN_WP_BLOCK_DIM (void)feat_dim; @@ -81,12 +75,11 @@ int best_block_dim(int feat_dim){ } -// The four kernels below are pure gathers: the whole of each one is the -// input_offset it computes. reference.py transcribes those four expressions to -// PyTorch so they can be tested on a CPU, but nothing binds the two files -- -// only a build checks this source. If you change the arithmetic here, change -// reference.py to match, or test_index_math.py will go on passing against the -// old formula. +// The four kernels below are pure gathers: each one *is* the input_offset it +// computes. reference.py transcribes those expressions to PyTorch so they can +// be tested on a CPU. Nothing binds the two files -- only a build checks this +// source -- so if you change the arithmetic here, change reference.py too, or +// test_index_math.py will go on passing against the old formula. // input: [B, H, W, C] // output: [B*nH*nW, window_h, window_w, C] From ff1714bcc6bf954bc2865c89780d5df5721016ff Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Fri, 28 Aug 2026 20:14:53 +0200 Subject: [PATCH 32/34] ci: pin the torch version on each leg The two legs exist to cover both hipify behaviours: one version that rewrites the stream getter and one that leaves it. Neither was pinned, so both legs resolved whatever pip offered for their Python version -- which happens to differ today and will not always. The moment they converge, one of the two behaviours stops being tested and nothing fails, because hipify_check.py accepts either by design. Pinned to 2.8.0 and 2.13.0, the two versions whose behaviour is actually recorded, with the reason in a comment so a future bump keeps the property rather than the numbers. --- .github/workflows/window-process.yml | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/.github/workflows/window-process.yml b/.github/workflows/window-process.yml index 385bf58cf..874811e2f 100644 --- a/.github/workflows/window-process.yml +++ b/.github/workflows/window-process.yml @@ -30,7 +30,11 @@ jobs: strategy: fail-fast: false matrix: - python-version: ['3.9', '3.11'] + include: + - python-version: '3.9' + torch-version: '2.8.0' + - python-version: '3.11' + torch-version: '2.13.0' steps: - uses: actions/checkout@v4 @@ -39,10 +43,15 @@ jobs: python-version: ${{ matrix.python-version }} - name: Install PyTorch (CPU build) + # Pinned, because the two legs are not interchangeable: hipify rewrites + # the stream getter on 2.8 and leaves it on 2.13, and hipify_check.py is + # written to accept either. Resolving whatever is newest would silently + # stop covering one of the two behaviours. run: | python -m pip install --upgrade pip python -m pip install numpy - python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + python -m pip install torch==${{ matrix.torch-version }} \ + --index-url https://download.pytorch.org/whl/cpu - name: Index math working-directory: kernels/window_process From 1d4220dcde478f2a8915f182a767e1f93ac38d22 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Fri, 28 Aug 2026 20:15:08 +0200 Subject: [PATCH 33/34] docs: correct the ROCm build recipe, which did not work The documented recipe was reconstructed from a container that already had the packages. Run against a fresh one it fails three times over, each failure far from its cause: - `export CPLUS_INCLUDE_PATH=...:/usr/include` puts /usr/include ahead of the compiler's own directories, which defeats libstdc++'s #include_next. The build dies on `stdlib.h: No such file or directory` before reaching a single ROCm header. - Ubuntu's librocthrust-dev and librocprim-dev are ROCm 5.7. They install a 5.7 HIP into /usr/include that shadows the image's 7.14 headers, producing a wall of errors inside amd_warp_sync_functions.h. - The amdrocm-*7.14 packages are not on repo.radeon.com but on repo.amd.com/rocm/packages-multi-arch/ubuntu2404, and the four listed are subsumed by one metapackage, amdrocm-core-dev7.14. A fourth step was missing entirely: the image ships libamdhip64.so.7 with no development symlink, so the link resolves -lamdhip64 only with LIBRARY_PATH set. Compilation succeeded without it and the link failed, which is why it went unnoticed. The corrected recipe was validated end to end in a container created from scratch: build clean, unit_test.py 15/15 and test_model_parity.py 4/4 on the MI300X. No signing key is fetched over the network because repo.amd.com publishes none that resolves -- the documented path reuses the keyring the Developer Cloud host already carries, which is what was actually executed. Validation and portability sections updated with that run: all 15 shapes now execute on the MI300X, two of them still only there; and torch 2.12, the version the ROCm build used, keeps the at::cuda stream spelling, so the kept form is the one that compiled under hipcc rather than merely the one the checker tolerates. --- kernels/window_process/README.md | 63 ++++++++++++++++++++++---------- 1 file changed, 44 insertions(+), 19 deletions(-) diff --git a/kernels/window_process/README.md b/kernels/window_process/README.md index 3f2e12b9a..0a72153b2 100644 --- a/kernels/window_process/README.md +++ b/kernels/window_process/README.md @@ -47,6 +47,7 @@ why the call sites pass `-shift_size` forward and `+shift_size` back. | dtype | float64, float32, float16, bfloat16 | | layout | contiguous, channels last in memory: `(..., C)` | | tiling | `H % window_h == 0` and `W % window_w == 0` | +| shape | `tensor.numel() == B * H * W * C` | | shift | `abs(shift) < window size`, per axis | | size | `B * H * W * C < 2^31` — offsets are computed in `int` | | grid | `B * nH * nW <= 65535` and `H <= 65535` — CUDA `grid.z` and `grid.y` | @@ -55,12 +56,15 @@ Shape, tiling, shift and size are checked with `TORCH_CHECK` before the launch; device and contiguity by `CHECK_INPUT`. The channel-last layout is the shape contract itself and is not separately verified. -Three of these are new guards over what was previously undetected: violating the -tiling or size rules used to read the wrong element, or read out of bounds, with -no error raised anywhere, and exceeding the grid bounds used to fail the launch -asynchronously somewhere else. The `shift` row is different — the kernels -compute a larger roll correctly — and is enforced because it is the contract the -model already documents, not because it used to break. dtype and layout were +Four of these are new guards over what was previously undetected: violating the +tiling, shape or size rules used to read the wrong element, or read out of +bounds, with no error raised anywhere, and exceeding the grid bounds used to +fail the launch asynchronously somewhere else. The `shift` row is different — +the kernels compute a larger roll correctly, for `window <= |shift| <= H` per +axis, and past that the operand of the modulo stops being non-negative and wraps +into a silently wrong read rather than an out-of-bounds one. It is enforced +because it is the contract the model already documents, not because it used to +break. dtype and layout were already reported by the dispatch and by `CHECK_CONTIGUOUS`. Note that `nH != nW` follows from `H != W` alone: with a square window, @@ -88,8 +92,9 @@ python benchmark.py --iters 200 The first two run in CI on a CPU runner, on every push that touches this directory. `unit_test.py` and `test_model_parity.py` are executed there too, but only to confirm they skip cleanly rather than error; `benchmark.py` is not, since -it exits non-zero without a device. The two CI legs install different torch -versions, so both hipify behaviours are exercised. +it exits non-zero without a device. The two CI legs pin different torch +versions, one that rewrites the stream getter and one that does not, so both +hipify behaviours are exercised on every run. The kernels perform no arithmetic on the values, only gathers, so parity is asserted with `torch.equal` on every dtype including float16 and bfloat16: any @@ -99,7 +104,10 @@ deviation is an indexing error rather than a rounding one. For the same reason `reference.py` transcribes the kernels' `input_offset` computation to vectorised PyTorch. That is not an approximation of the kernels — the offset computation is their entire logic — which is what makes them testable on a CPU, with no -compiled extension and no GPU of either vendor. +compiled extension and no GPU of either vendor. The two agree wherever the +operand of a modulo stays non-negative, which every input the launcher accepts +guarantees; outside that domain Python floors where the kernel wraps, and the +transcription is the more forgiving of the two. ## Validation @@ -109,7 +117,8 @@ Parity is asserted against the eager path on an **RTX 3080** (torch both, bit-exact, over square and non-square images and windows including coprime `nH`/`nW`. `compute-sanitizer` (memcheck, initcheck, synccheck) is clean on the RTX 3080 across all four kernels, forward and backward; ROCm has no -equivalent run. +equivalent run. The MI300X run covers all 15 shapes in `unit_test.py`; two of +them were added after the RTX 3080 run and have not been executed on NVIDIA. ## Performance @@ -184,16 +193,30 @@ Two things are worth knowing before building: `/opt/rocm` tree, and those wheels carry no rocThrust, hipSPARSE, hipBLAS, hipBLASLt or hipSOLVER headers — which torch's own headers include when compiling device code. In that state *no* PyTorch HIP extension compiles; a - file whose only content is `#include ` fails the same way. - Installing the matching dev packages is enough, and touches no source: + file whose only content is `#include ` fails the same way. One + dev metapackage supplies all of them, and touches no source: ```bash - apt-get install -y amdrocm-blas-dev7.14 amdrocm-hipblas-common-dev7.14 \ - amdrocm-sparse-dev7.14 amdrocm-solver-dev7.14 \ - librocthrust-dev librocprim-dev - export CPLUS_INCLUDE_PATH=/opt/rocm/core-7.14/include:/usr/include + # AMD's Developer Cloud host has this repo configured already; inside the + # container, reuse its keyring and point apt at the same source. + echo 'deb [arch=amd64 signed-by=/etc/apt/keyrings/amdrocm.gpg] https://repo.amd.com/rocm/packages-multi-arch/ubuntu2404 stable main' \ + > /etc/apt/sources.list.d/rocm.list + apt-get update && apt-get install -y amdrocm-core-dev7.14 + export CPLUS_INCLUDE_PATH=/opt/rocm/core-7.14/include + export LIBRARY_PATH=/opt/rocm/core-7.14/lib ``` + `LIBRARY_PATH` is separate from the headers: the image ships + `libamdhip64.so.7` with no development symlink, so the final link cannot + resolve `-lamdhip64` without it. Two near misses are worth naming, because + both fail far from their cause. Do not append `/usr/include` to + `CPLUS_INCLUDE_PATH` — searching it ahead of the compiler's own directories + defeats libstdc++'s `#include_next`, and the build dies on `stdlib.h: No such + file or directory` before it reaches a single ROCm header. And do not reach + for Ubuntu's `librocthrust-dev` / `librocprim-dev`: they are ROCm 5.7, and + they install a 5.7 HIP into `/usr/include` that shadows the image's 7.14 + headers, producing a wall of errors inside `amd_warp_sync_functions.h`. + ### Portability The kernels use no shared memory, no `__syncthreads()` and no warp-level @@ -224,6 +247,8 @@ The one substantive difference between the two platforms is the launch stream. Upstream passes `0`, the null stream; this version passes `at::cuda::getCurrentCUDAStream()`. Some hipify versions rewrite that to the HIP spelling and some leave it, which resolves to the HIP runtime anyway — both -reach the current stream. `hipify_check.py` makes no prediction about which: -it reports the one that happened rather than requiring either, prints -the full symbol mapping, and fails if anything is left untranslated. +reach the current stream. torch 2.8 and 2.10 rewrite it; 2.12 and 2.13 leave it, +and 2.12 is the version the MI300X build used, so the kept spelling is the one +that compiled under hipcc and ran bit-exact above. `hipify_check.py` still makes +no prediction: it reports the one that happened rather than requiring either, +prints the full symbol mapping, and fails if anything is left untranslated. From 63bfb185fe822c2f7072c4be6e754bc374d57512 Mon Sep 17 00:00:00 2001 From: francescobrigante Date: Fri, 28 Aug 2026 20:22:52 +0200 Subject: [PATCH 34/34] docs: say what the .cpp check saw, and what the apt line assumes Three claims that a reader can catch without running anything. hipify_check.py printed "nothing device specific in this file" for swin_window_process.cpp. That file has CHECK_CUDA and is_cuda() in it. The check never looked for those -- it reports on three fixed watch lists and that line is what it prints when none of them matched -- but it is the first line of output for the only file in the PR that anyone executes, and it overstates what was established. It now says so literally: no symbol from the watch lists appears in this file. The corrected apt recipe depends on a keyring it does not install. On the Developer Cloud image the host carries it, which is where the recipe was validated and why the gap was invisible; anywhere else apt-get update dies with NO_PUBKEY -- a fourth failure far from its cause, which is the exact category this recipe exists to remove. repo.amd.com serves no key over HTTP and the repo.radeon.com key signs a different repository (verified: NO_PUBKEY FA296B056C5BB456), so the honest form is to state the assumption rather than invent a fetch. The build command also came before the two exports that make it work; it now points forward to them. The validation line used "15" twice in one sentence for two different things: 15 test methods and 15 entries in the shape list. Both are true and the coincidence is not meaningful, so each is now named. --- kernels/window_process/README.md | 24 ++++++++++++++++-------- kernels/window_process/hipify_check.py | 2 +- 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/kernels/window_process/README.md b/kernels/window_process/README.md index 0a72153b2..f3d16eda9 100644 --- a/kernels/window_process/README.md +++ b/kernels/window_process/README.md @@ -113,12 +113,13 @@ transcription is the more forgiving of the two. Parity is asserted against the eager path on an **RTX 3080** (torch `2.8.0+cu129`, CUDA 12.9) and an **AMD Instinct MI300X** (gfx942, torch -`2.12.0+rocm7.14.0`): `unit_test.py` 15/15 and `test_model_parity.py` 4/4 on -both, bit-exact, over square and non-square images and windows including -coprime `nH`/`nW`. `compute-sanitizer` (memcheck, initcheck, synccheck) is clean -on the RTX 3080 across all four kernels, forward and backward; ROCm has no -equivalent run. The MI300X run covers all 15 shapes in `unit_test.py`; two of -them were added after the RTX 3080 run and have not been executed on NVIDIA. +`2.12.0+rocm7.14.0`): `unit_test.py` 15/15 test methods and +`test_model_parity.py` 4/4 on both, bit-exact, over square and non-square +images and windows including coprime `nH`/`nW`. `compute-sanitizer` (memcheck, +initcheck, synccheck) is clean on the RTX 3080 across all four kernels, forward +and backward; ROCm has no equivalent run. The MI300X run covers all 15 entries +in `unit_test.py`'s shape list; two of them were added after the RTX 3080 run +and have not been executed on NVIDIA. ## Performance @@ -176,6 +177,10 @@ from `PYTORCH_ROCM_ARCH`: PYTORCH_ROCM_ARCH=gfx942 python setup.py install ``` +On AMD's ROCm 7.14 images that command fails until the SDK headers are in +place; the second bullet below has the two `export`s and the one package it +needs, and they have to come first. + Do not pass `--offload-arch` through `extra_compile_args`: PyTorch's `_get_rocm_arch_flags()` skips its own detection as soon as it sees one, which pins the build to a single GPU. @@ -197,8 +202,11 @@ Two things are worth knowing before building: dev metapackage supplies all of them, and touches no source: ```bash - # AMD's Developer Cloud host has this repo configured already; inside the - # container, reuse its keyring and point apt at the same source. + # Assumes AMD's Developer Cloud image, whose host already carries this repo's + # signing key at /etc/apt/keyrings/amdrocm.gpg -- copy it into the container. + # Anywhere else, install AMD's key under that path first: repo.amd.com serves + # none over HTTP, and the repo.radeon.com key signs a different repository, so + # apt-get update fails with NO_PUBKEY without it. echo 'deb [arch=amd64 signed-by=/etc/apt/keyrings/amdrocm.gpg] https://repo.amd.com/rocm/packages-multi-arch/ubuntu2404 stable main' \ > /etc/apt/sources.list.d/rocm.list apt-get update && apt-get install -y amdrocm-core-dev7.14 diff --git a/kernels/window_process/hipify_check.py b/kernels/window_process/hipify_check.py index 66f31a587..f54fcd22d 100644 --- a/kernels/window_process/hipify_check.py +++ b/kernels/window_process/hipify_check.py @@ -167,7 +167,7 @@ def report(staging, mapping, show_diff): if not translated and not survived and not shims and not portable \ and LDG_MACRO not in before: - print(' nothing device specific in this file') + print(' no symbol from the watch lists appears in this file') if show_diff: diff = difflib.unified_diff(