Fix an out-of-bounds read in the fused window kernels, generalise to non-square shapes, port to ROCm - #389
Open
Francesco Brigante (francescobrigante) wants to merge 34 commits into
Conversation
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.
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.
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.
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<T>() -> Tensor::data_ptr<T>()
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.
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.
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.
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.
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.
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.
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.
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.
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.
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).
…mage - 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.
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.
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.
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.
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.
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.
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.
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.
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.
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.
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.
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.
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.
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.
…te 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.
__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.
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.
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.
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.
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.
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.
Author
|
@microsoft-github-policy-service agree |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
window_merge_and_roll_forward_cuda_kernelusesnHas the stride betweenwindow rows. The stride is
nW. On any feature map withH != Wthe kernelreads the wrong element, and when
nH > nWit reads past the end of thetensor — an out-of-bounds read on every forward pass, silent, with
fused_window_process=Trueand a stockimg_size=(256, 128).This PR fixes that one-token bug, then generalises the four kernels to
non-square windows and per-axis shifts, adds bfloat16 support, replaces
deprecated ATen APIs, launches on the current stream, validates preconditions,
and adds a test suite — including one that checks the index arithmetic with no
GPU at all, so it runs in CI.
It also ports the kernels to ROCm. The port uncovered a second defect:
__ldghas no HIP overload for
c10::Half, and half has been in the dispatch from thestart, so the fused window process has never compiled on AMD. With that
fixed, everything is validated on an MI300X as well as on an RTX 3080.
Everything is under
kernels/window_process/. The single exception is a CIworkflow. The Python API,
models/swin_transformer.pyandsetup.pyareunchanged.
1. The bug
Upstream,
swin_window_process_kernel.cu:Windows are laid out row-major: window
(wrow, wcol)of batch elementbsitsat flat index
b * nH * nW + wrow * nW + wcol. Advancing one window rowskips
nWwindows, notnH. The other three kernels already usenW; this oneis the odd one out.
No non-square window is needed to trigger it. With a single scalar
window_size,nH = H / window_sizeandnW = W / window_size, sonH != nWwhenever
H != W.img_sizeis documented asint | tuple(int)and goesthrough
to_2tuple, so a non-square feature map reaches the kernels through asupported public argument. Every released Swin checkpoint is trained on square
images, which is the only reason this has never surfaced.
Square 8×8 window, shift 4,
B=2,C=1, onlyH×Wvarying — reproducible ona CPU with
reference.py:nH > nWwalks off the end of the tensor.nH < nWstays inside it andcorrupts silently — the worse of the two.
A stock configuration reaches it:
Measured on an RTX 3080, reverting only the one-token fix on top of this branch:
torch.equal(eager, fused)→ False, peak logit error3.2e-2;compute-sanitizer→Invalid __global__ read, 41 008 errors. With the fix:torch.equalTrue, error exactly 0, memcheck + initcheck + synccheck report0 errors.
The pre-fix kernel is bit-exact on every square configuration, including the
56×56 w7 case in the existing
unit_test.py— which is why the existing suitepasses over it.
One thing that is not a bug. Upstream computes the intra-window offset as
(blockIdx.y - shift + H) % window_sizewithout reducing% Hfirst. Thisbranch writes the explicit form, but the two are equivalent: they differ by
a multiple of
H, andHis a multiple ofwindow_hwhenevernHis exact,which the launcher already requires.
TestIntraWindowModuloIsEquivalentassertsit. The change is for readability and is not claimed as a fix.
2. Generalisation and cleanup
Beyond the fix, the kernels now accept non-square windows (
window_h,window_w) and per-axis shifts (shift_h,shift_w) instead of twoscalars. The Python arity is unchanged —
intworks exactly as before, an(h, w)pair is new:TestScalarWindowStillWorksasserts the int path and the pair it expands to arebit-identical, so
models/swin_transformer.pykeeps working untouched.Other changes in the range:
half, so a bf16 tensor needed an fp32 round trip to use the fused path.
input.type()→scalar_type(),.data<T>()→.data_ptr<T>(), output frominput.options()instead ofhardcoded
kCUDA+kFloat32). The pre-fix sources fail to compile ontorch 2.8+: the
AT_DISPATCH_*overload takingDeprecatedTypePropertiesis removed, and
T* data()no longer exists onTensorBase.at::cuda::getCurrentCUDAStream()) instead of the legacydefault stream, which serialises against any user stream and cannot be captured
in a CUDA graph.
C10_CUDA_KERNEL_LAUNCH_CHECK()added.TORCH_CHECKed before launch: tiling divisibility, elementcount,
int32offset overflow, and thegrid.z/grid.ylimit of 65 535.backward()returned 8 values for 7 forwardinputs;
unit_test.pynever checked a gradient (both "backward" testsre-compared the forward outputs);
test_forward_backward_speedassertedfused < eagerontime.time()with no warm-up or synchronisation — moved tobenchmark.py, which reports rather than asserts.3. Correctness
The kernels perform no arithmetic on tensor values. Each is a pure gather
whose entire logic is the computation of
input_offsetfromblockIdx. Twoconsequences:
Parity is asserted with
torch.equal, exact, on every dtype includingfloat16 and bfloat16 — any deviation is an indexing error, not a rounding one.
The index math is testable with no GPU.
reference.pytranscribes the fourinput_offsetcomputations to vectorised PyTorch — the same arithmetic insidethe domain the launcher enforces, not an approximation — so
test_index_math.pyruns on a CPU runner in CI. It asserts that each kernel equals the composition
of the eager ops it fuses, that K1/K2 and K3/K4 are inverse permutations, and —
with a
legacy_row_stride=Trueflag — that the upstream map is correct onsquare grids, wrong on
nH < nW, out of bounds onnH > nW, and invisible whennH == 1. Every case runs at four shifts, including the negative one: the modelcalls the partition path with
-shift_size, and this suite used to pass onlypositive shifts, leaving that sign untested.
unit_test.pynever had the gap —it negates at the call site — so the fix was needed here and only here.
The regression guarantee for the compiled kernel rests on
unit_test.py, on aGPU: 4 kernels × 4 dtypes × 15 shapes × 3 shifts, plus CUDA-graph capture, a
side-stream run, a non-contiguous incoming gradient, the int-vs-pair
compatibility path and six precondition tests — 15 test methods total. The full
matrix has run on the MI300X; two of the 15 shapes postdate the RTX 3080 run,
and §6 records exactly what each device swept.
test_model_parity.pyruns a wholeSwinTransformer— square, tall(256,128), wide(128,256)— and requires the logits and the gradients tobe identical with and without
fused_window_process.4. Performance
benchmark.py, fused againsttorch.roll+window_partition, forward andbackward, batch 192, 200 iterations after 10 warm-up. Selected rows from a
24-row matrix — 4 configurations × 2 directions × 3 dtypes — which
benchmark.pyprints in full and the README samples.RTX 3080, torch
2.8.0+cu129/ CUDA 12.9:2.51–3.84× on fwd+bwd: that is the min and max over all 24 measured rows, of
which the six above are a sample —
with peak memory dropping proportionally. The path is
memory bound: the gain comes from eliminating the second materialisation that the
eager path pays (
torch.rollwrites a copy,window_partitioncalls.contiguous()for another).AMD Instinct MI300X, torch
2.12.0+rocm7.14.0: 1.03–2.38×, memory identicalto CUDA. The smaller ratio reflects the faster eager baseline on this hardware
(stage 1 float32: 7.137 ms on the 3080 vs 0.911 ms on the MI300X), which
compresses both paths toward a floor. The block-width thresholds in
best_block_dim()were measured on CDNA at swin-tiny'sCand confirmed tomatch the NVIDIA-tuned choice; full numbers are in the README.
5. ROCm / HIP port
Why no new kernel code was written
The four kernels use no NVIDIA-specific primitives — no shared memory, no
__syncthreads(), no warp-level intrinsics, no atomics. Each kernel is a puregather: read one element, write it to a permuted location. PyTorch's
CUDAExtensionruns hipify automatically whentorch.version.hipis set, and_get_rocm_arch_flags()derives--offload-archfromPYTORCH_ROCM_ARCH,so the build command is:
PYTORCH_ROCM_ARCH=gfx942 python setup.py install # gfx90a for MI210setup.pyis untouched — and deliberately so: passing an arch throughextra_compile_argswould makeCUDAExtensionskip its own detection, pinningthe build to one GPU.
The
__ldgdefectDespite the clean translation, the kernels did not compile on AMD — for all
four kernels, instantiated for
c10::Half:The symbol
__ldghipifies correctly (HIP defines it), but not for everytype.
float,doubleandc10::BFloat16have HIP overloads;c10::Halfdoes not. The cause is in PyTorch's own headers, which do not guard the two
alike:
Half.hguards the overload on(__CUDA_ARCH__ >= 350) || (__clang__ && __CUDA__), which no hipcc buildsatisfies, while
BFloat16.hguards on__CUDACC__ || __HIPCC__and degradesto
*ptrunder ROCm. Since half has beenin the dispatch from the start, the fused window process has never been
buildable on AMD.
The fix is a macro:
On NVIDIA this expands to
__ldg(ptr), identical to the original code, so theCUDA translation unit is unchanged. On AMD, it degrades to a plain dereference —
__ldgis an advisory cache hint on both platforms, so the AMD branch loses thehint and nothing else.
Hipify verification in CI
hipify_check.pyruns the real hipify with no GPU and verifies that everylaunch names a stream getter (upstream's
0hipifies to the HIP null stream;this branch's
getCurrentCUDAStream()hipifies to the current one). Differenttorch versions rewrite the getter differently — 2.8 and 2.10 rewrite it to
at::hip::getCurrentHIPStreamMasqueradingAsCUDA(), 2.12 and 2.13 leave theat::cuda::spelling — so the check accepts either, and the two CI legs pin2.8 and 2.13 precisely so that both behaviours stay covered. The MI300X build
ran on 2.12, so the kept spelling is not only accepted by the checker but has
compiled under hipcc and produced bit-exact results on hardware.
6. Validation status
test_index_math.py)hipify_check.py)unit_test.py)test_model_parity.py)img_size=(256,128)gives different logits,compute-sanitizerreports OOB reads; post-fix: bit-identical to eagercompute-sanitizermemcheck + initcheck + synccheckPYTORCH_ROCM_ARCH,setup.pyunchangedunit_test.py)test_model_parity.py)CI on the branch head, both legs green:
https://github.com/francescobrigante/Swin-Transformer/actions/runs/33199033099
(full history:
https://github.com/francescobrigante/Swin-Transformer/actions/workflows/window-process.yml).
The GPU rows are reported rather than linkable: both cards are rented, and the
RTX 3080 was released once its rows were recorded, which is why the two shapes
added afterwards carry an MI300X result and no NVIDIA one.
A note on gradient reproducibility.
test_gradients_are_identicalinitiallyused
torch.equaland failed on the MI300X. Investigation showed this isplatform non-determinism, not a kernel defect: running the eager path (no fused
kernel) twice produces
max abs diff = 1.863e-09on 1 of 63 gradient tensors,and the fused-vs-eager diff is the same
1.863e-09. The test now measures theplatform's own reproducibility first, and holds each gradient to that bound.
On CUDA — deterministic in every run so far — every gradient takes the strict
torch.equalbranch. Bit-exactness of the kernels themselves is unaffected andstill asserted per kernel in
unit_test.py, 15/15 on both vendors.7. Review guide
The commits are ordered so each one is separately readable, and
git bisectismeaningful across the whole range.
Start with two.
afd246baddsreference.pyandtest_index_math.py—the bug provable without a GPU;
cf5e949is the fix itself, one token, isolated.Then
568ab74(per-axis window and shift),2eca45b(TORCH_CHECKpreconditions),
aadeb8d(hipify_check.pyand the ROCm findings) and215e151(the__ldgmacro, so the kernels compile under HIP).The rest of the range covers the deprecated ATen APIs, the current stream,
bfloat16, the test suite, the MI300X run, the README and CI workflow, and
documentation corrections.
setup.pyis deliberately untouched.