JuliaGPU / JuliaGPU/Metal.jl

2D `@index(Global, NTuple)` is ~5x slower than linear indexing

Open
#910 1 comment 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.