2D `@index(Global, NTuple)` is ~5x slower than linear indexing
- Dominant language
- Julia
- Stars
- 463
- Forks
- 68
- Avg merge
- 1d 19m
- Merged PRs (30d)
- 32
Description
Memcopy at 4096^2, Apple M2. Same access pattern throughout; only the index computation differs.
```
KA 2D @index(Global, NTuple) 7.61 ms 17.6 GB/s
KA 1D @index(Global, Linear) 1.45 ms 92.8 GB/s
KA 2D manual Int32 % and / 2.00 ms 67.1 GB/s
KA 2D manual Int32 & and >> 1.44 ms 92.9 GB/s
Metal.jl native 2D (no KA) 1.47 ms 91.2 GB/s
```
Hand-decomposing the linear index in `Int32` recovers everything, bitwise-identical output (`max|NTuple - shift/mask| = 0.0e+00`). Same kernel with `Int64` arithmetic: 10.3 GB/s.
Already ruled out: `@index(Global, Cartesian)` is identical to `NTuple`; group size is irrelevant (17.5–17.9 across `(16,16)`, `(32,8)`, `(64,4)`, `(128,2)`, `(256,1)`, `(32,32)`, `(8,32)`); `unsafe_indices=true` gives 17.8, so it is not `__validindex`.
## MWE
```julia
using KernelAbstractions, Metal, Printf
const be = MetalBackend()
@kernel inbounds=true function copy2D!(A, B)
ix, iy = @index(Global, NTuple); A[ix,iy] = B[ix,iy]
end
@kernel inbounds=true function copy1D!(A, B)
I = @index(Global, Linear); A[I] = B[I]
end
@kernel inbounds=true function copy2D_i32!(A, B, nx::Int32)
I = Int32(@index(Global, Linear)) - Int32(1)
ix = I % nx + Int32(1); iy = I ÷ nx + Int32(1)
A[ix,iy] = B[ix,iy]
end
@kernel inbounds=true function copy2D_shift!(A, B, mask::Int32, sh::Int32)
I = Int32(@index(Global, Linear)) - Int32(1)
ix = (I & mask) + Int32(1); iy = (I >> sh) + Int32(1)
A[ix,iy] = B[ix,iy]
end
function native2D!(A, B)
ix, iy = thread_position_in_grid_2d()
@inbounds A[ix,iy] = B[ix,iy]
return
end
function timeit(f; nrep=50, ntrial=5)
f(); Metal.synchronize(); best = Inf
for _ in 1:ntrial
Metal.synchronize(); t0 = time()
for _ in 1:nrep; f(); end
Metal.synchronize(); best = min(best, (time()-t0)/nrep)
end
best
end
n = 4096
A = Metal.zeros(Float32,n,n); B = Metal.ones(Float32,n,n)
bw(t) = 2*n*n*sizeof(Float32)/t/1e9
msk, sh = Int32(n-1), Int32(trailing_zeros(n))
k = copy2D!(be,(32,8)); t1 = timeit(()->k(A,B; ndrange=(n,n)))
k = copy1D!(be,256); t2 = timeit(()->k(A,B; ndrange=n*n))
k = copy2D_i32!(be,256); t3 = timeit(()->k(A,B,Int32(n); ndrange=n*n))
k = copy2D_shift!(be,256); t4 = timeit(()->k(A,B,msk,sh; ndrange=n*n))
t5 = timeit(()->Metal.@metal threads=(32,8) groups=(cld(n,32),cld(n,8)) native2D!(A,B))
for (nm,t) in (("2D NTuple",t1), ("1D Linear",t2), ("2D Int32 %,/",t3),
("2D Int32 &,>>",t4), ("native 2D",t5))
@printf("%-16s %6.2f ms %5.1f GB/s\n", nm, t*1e3, bw(t))
end
```
## Cause
`src/MetalKernels.jl:124` collapses the launch to 1D (`groups = length(KA.blocks(iterspace))`), so `:151` rebuilds the index via `LinearIndices(KA.__ndrange(ctx))[I]` in 64-bit arithmetic, and Apple GPUs emulate 64-bit integer division.
`CUDA.jl/CUDACore/src/CUDAKernels.jl:145` does the same collapse, so the 1D dispatch is not the Metal-specific part; the division cost is. Not measured on CUDA hardware.
Would narrowing `__index_Global_*` to `Int32` be viable, dispatching on `length(ndrange)` for the >`typemax(Int32)` case?
This relates possibly to open issue #101 .
## Environment
- Apple M2, macOS 26.6
- Julia 1.12.6
- Metal.jl 1.10.0
- KernelAbstractions.jl 0.9.42
Contributor guide
No contributing guide indexed for this repository
Research direction
Run the provided Metal MWE first and compare the 2D NTuple, linear, and manual-indexing timings. Read src/MetalKernels.jl around lines 124 and 151, then compare the analogous path in CUDA.jl/CUDACore/src/CUDAKernels.jl:145. Done means establishing a safe indexing approach that removes the observed regression without breaking larger ndarrays.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- backend, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100