Fix 2D parallel_for coalescing and remove per-reduce fill! (CUDA/AMDGPU)

[miniskar]

What

Two GPU-backend changes to ext/CUDAExt/CUDAExt.jl and ext/AMDGPUExt/AMDGPUExt.jl:

1. 2D parallel_for coalescing — swap grid axes only when necessary, and keep a
   warp on the stride-1 dimension.
2. Reduction init seeding — pass init as a kernel argument instead of
   pre-filling buffers with fill!.

Why

The M < N swap was a launchability workaround (the y-grid limit is 65535 vs. 2³¹
for x), but it fired for every wide array — including the vast majority that fit the
coalesced path — turning a bandwidth-bound kernel into a strided one. The fill!s
are extra kernel launches on the reduce critical path that only existed to place an
identity value the kernel can hold itself.

---

Issue 1 — 2D parallel_for loses coalescing on wide arrays

A 2D launch maps one array dimension to the grid x-axis and the other to y.
y is limited to 65535 blocks; x to 2³¹. To launch a wide array (M < N) whose
N would overflow grid.y, the backend swaps axes so N lands on x.

The two index mappings differ in which array element each warp touches. Arrays are
column-major (i is stride-1):

# BlockIndexerBasic:   threadIdx().x -> i  (stride 1)  => a warp reads A[i:i+31, j]   → 1 cache line   ✅
# BlockIndexerSwapped: threadIdx().x -> j  (stride M)  => a warp reads A[i, j:j+31]   → 32 cache lines ❌

Before (the bug)

The swap fired for every M < N array, even the majority whose basic grid fits
under the y limit — so a coalesced kernel was needlessly turned into a strided one:

# ext/CUDAExt/CUDAExt.jl (old)
if M < N && maxBlocks.x >= maxBlocks.y            # blanket rule
    _parallel_for(BlockIndexerSwapped(), f, (N, M), (M, N), x...)   # uncoalesced
else
    _parallel_for(BlockIndexerBasic(), f, (M, N), (M, N), x...)
end

Worse, the block-shape heuristic then starved the x (coalesced) dimension for wide
arrays, collapsing the block to (1, threads):

y_thr = clamp(floor(Int, (n / m) * maxThreadsX), 1, config.threads)  # ~threads when n≫m
x_thr = fld(config.threads, y_thr)                                   # -> 1  (no warp on x)

After (the fix)

Swap only when the coalesced grid would actually overflow grid.y, and shape the
block so a full warp walks the stride-1 dimension:

# ext/CUDAExt/CUDAExt.jl (new)
kargs = _kernel_args(BlockIndexerBasic(), (M, N), f, x...)
kernel, shmem = _kernel_maxshmem(_parallel_for_cuda_MN, kargs)
config = CUDA.launch_configuration(kernel.fun; shmem)
x_thr, y_thr = _block_shape_2d(config.threads, M, N)

if cld(N, y_thr) > maxBlocks.y && maxBlocks.x >= maxBlocks.y   # swap only if unavoidable
    _parallel_for(BlockIndexerSwapped(), f, (N, M), (M, N), x...)
else                                                          # coalesced fast path
    kernel(kargs...; threads = (x_thr, y_thr),
        blocks = (cld(M, x_thr), cld(N, y_thr)), shmem = shmem)
    CUDA.synchronize()
end

# keep a warp on the stride-1 (row) dimension; heuristic still handles square/tall
@inline function _block_shape_2d(maxthreads, m, n)
    maxThreadsX = sqrt(maxthreads)
    y_thr = clamp(floor(Int, (n / m) * maxThreadsX), 1, maxthreads)
    x_thr = fld(maxthreads, y_thr)
    if x_thr < 32 && m >= 32
        x_thr = 32; y_thr = fld(maxthreads, x_thr)
    elseif m < 32 && x_thr < m
        x_thr = m;  y_thr = max(1, fld(maxthreads, x_thr))
    end
    return (x_thr, y_thr)
end

Reproduce (NVIDIA A100)

import JACC
JACC.@init_backend           # backend = "cuda"
using CUDA
scale2(i, j, A, B) = (@inbounds B[i, j] = 2f0 * A[i, j]; nothing)

function bw(M, N; iters = 100)
    A = CUDA.rand(Float32, M, N); B = CUDA.zeros(Float32, M, N)
    JACC.parallel_for((M, N), scale2, A, B); CUDA.synchronize()          # warmup
    t = CUDA.@elapsed (for _ in 1:iters; JACC.parallel_for((M, N), scale2, A, B); end;
                       CUDA.synchronize())
    2 * M * N * sizeof(Float32) / (t / iters) / 1e9                      # GB/s
end

bw(256, 262144)   # wide  — before: ~41 GB/s      after: ~770 GB/s
bw(262144, 256)   # tall  — ~770 GB/s (unchanged; reference for the coalesced path)

---

Issue 2 — reductions do a synchronous fill! between kernel launches

The op-identity for out-of-range threads was placed by pre-filling the partial-results
buffer with fill! — two extra launches sitting between the reduction kernels, and the
Unmanaged workspace relied on the caller having pre-filled the buffers.

Before

# host
_init!(wk, blocks, init)                       # fill!(wk.tmp, init); fill!(wk.ret, init)
kernel_1(N, op, wk.tmp, f, x...)
kernel_2(blocks, op, wk.tmp, wk.ret)
# launch sequence:  memset → memset → reduce_1 → reduce_2

# kernel
@inbounds shared_mem[ti] = ret[blockIdx().x]   # identity read back from the filled buffer
...
@inbounds tmp = ret[1]

After

# host — `init` passed as an argument; _init! only sizes the buffer now
_init!(wk, blocks, init)
kernel_1(N, op, wk.tmp, init, f, x...)
kernel_2(blocks, op, wk.tmp, init, wk.ret)
# launch sequence:  reduce_1 → reduce_2   (pure async, no fill in between)

# kernel
@inbounds shared_mem[ti] = init                # identity seeded directly
...
tmp = init

Reproduce (NVIDIA A100)

id(i, x) = (@inbounds x[i])
function us(N; iters = 500)              # microseconds per reduce
    r = JACC.reducer(Float64, N, +); x = JACC.ones(Float64, N)
    r(id, x); CUDA.synchronize()
    t = CUDA.@elapsed (for _ in 1:iters; r(id, x); end; CUDA.synchronize())
    t / iters * 1e6
end
us(1000)     # before: ~44.5 us   after: ~30.8 us   (two memset launches removed)
us(10^7)     # ~unchanged (bandwidth-bound)

---

Performance (NVIDIA A100)

• Wide 2D parallel_for: ~5–19× faster — e.g. 256×262144 went 41 → ~770 GB/s,
  1024×65536 149 → ~770 GB/s — now matching tall/square arrays.
• Reductions: ~25–31% faster for small/medium sizes (the two eliminated memset
  launches), e.g. reused reducer N=1000 44.5 → 30.8 µs. Large bandwidth-bound
  reductions unchanged.
• Tall/square parallel_for unchanged (within noise).

Correctness / tests

Behavior is unchanged for the supported contract (identity init). Validated on A100
(CUDA) and CPU threads — reduce 15, reduce-ND 20, LaunchSpec 14,
parallel_for non-square 16, parallel_reduce non-square 8, all pass. Adds
non-square 2D regression tests to test/unittests.jl; the existing 2D tests only
used square matrices, so they never exercised the axis-swap path:

@testset "parallel_for non-square 2D" begin
    write2d(i, j, A) = (@inbounds A[i, j] = (j - 1) * size(A, 1) + i; nothing)
    for (M, N) in [(2, 5), (5, 2), (7, 13), (128, 1024), (1024, 128), (3, 257)]
        A = JACC.zeros(FloatType, M, N)
        JACC.parallel_for((M, N), write2d, A)
        @test JACC.to_host(A) == FloatType[(j - 1) * M + i for i in 1:M, j in 1:N]
    end
end

Notes

• Out of scope: parallel_reduce on 2D still has no swap path (hard-coded 16×16),
  so it overflows the y-grid for N > ~1.05M — a separate follow-up.

[williamfgc]

Test this please

[miniskar]

The two CI/CD failures are due to unavailability of CUDA runtime for CUDA.jl package and another is some file permission issue.

[williamfgc]

@miniskar thanks! Indeed, it's just CI hiccups. Thanks for the PR contribution!

[williamfgc]

Test this please

[williamfgc]

@miniskar please see the failing test on Intel GPU.