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

[luraess]

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

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

[cwittens]

I observed something similar earlier on CUDA (H200) but significantly less pronounced and it almost goes away when defining wg and ndrange before.

copy 2D dyn WG / dyn ndrange     2706.3 GB/s
copy 2D stat WG / dyn ndrange    3503.2 GB/s
copy 2D stat WG / stat ndrange   3651.0 GB/s
copy 1D dyn WG / dyn ndrange     3263.5 GB/s
copy 1D stat WG / dyn ndrange    3666.0 GB/s
copy 1D stat WG / stat ndrange   3669.5 GB/s
copyto! (vendor memcpy)          4040.8 GB/s

MWE

using KernelAbstractions, CUDA, Adapt, BenchmarkTools, Printf

function bench(f, backend, bytes; reps = 10)
    f(); KernelAbstractions.synchronize(backend)
    t = @belapsed begin
        for _ in 1:$reps
            $f()
        end
        KernelAbstractions.synchronize($backend)
    end
    return bytes * reps / t / 1e9   # GB/s
end

@kernel function kernel_copy!(dst, @Const(src))
    i, j = @index(Global, NTuple)
    @inbounds dst[i, j] = src[i, j]
end

@kernel function kernel_copy_1d!(dst, @Const(src))
    i = @index(Global, Linear)
    @inbounds dst[i] = src[i]
end


backend = CUDABackend()

T = Float64
Nx = Ny = 4096
ndr = (Nx, Ny)
nb  = sizeof(T) * Nx * Ny

u  = adapt(backend, rand(T, Nx, Ny))
du = similar(u)
vu, vdu = vec(u), vec(du)

wgs = (64, 4)

k = kernel_copy!(backend)
r1 = bench(() -> k(du, u, ndrange = ndr), backend, 2nb)

k = kernel_copy!(backend, wgs)
r2 = bench(() -> k(du, u, ndrange = ndr), backend, 2nb)

k = kernel_copy!(backend, wgs, ndr)
r3 = bench(() -> k(du, u), backend, 2nb)

k = kernel_copy_1d!(backend)
r4 = bench(() -> k(vdu, vu, ndrange = length(vu)), backend, 2nb)

k = kernel_copy_1d!(backend, 256)
r5 = bench(() -> k(vdu, vu, ndrange = length(vu)), backend, 2nb)

k = kernel_copy_1d!(backend, 256, length(vu))
r6 = bench(() -> k(vdu, vu), backend, 2nb)

r7 = bench(() -> copyto!(du, u), backend, 2nb)

for (nm, r) in (("copy 2D dyn WG / dyn ndrange",   r1),
                ("copy 2D stat WG / dyn ndrange",  r2),
                ("copy 2D stat WG / stat ndrange", r3),
                ("copy 1D dyn WG / dyn ndrange",   r4),
                ("copy 1D stat WG / dyn ndrange",  r5),
                ("copy 1D stat WG / stat ndrange", r6),
                ("copyto! (vendor memcpy)",        r7))
    @printf("%-32s %6.1f GB/s\n", nm, r)
end

Environment

julia_version = "1.12.6"
CUDA version = "6.2.1"
KernelAbstraction version = "0.9.42"