Accept index ranges in ndrange

[giordano]

Part 1/2 of upstreaming Oceananigans piracy of KA (part 1 is #765). Ref: #757. I'm open to adapt the implementation if necessary. I also have ready a backport to the v0.9 branch, but will open later after the design here is finalised.

CC @simone-silvestri.

Kernels can now iterate over a region whose indices do not start at 1. Each entry of ndrange may be a range instead of an extent, and the whole ndrange may be a single range or a CartesianIndices, both statically (kernel(backend, workgroupsize, (-2:N+3, 0:M+1))) and at launch (ndrange=(-2:N+3, 0:M+1)). @index(Global, Cartesian) and @index(Global, NTuple) return the shifted indices, @index(Global, Linear) counts the region from 1 in column-major order, @ndrange() returns the extents.

Downstream packages such as Oceananigans implement this today by pirating partition, expand, __ndrange and __groupsize with a custom _Size subtype smuggled into NDRange's dynamic-workitems type parameter, which is fragile and broke with the compiled CPU backend.

Implementation:

• StaticSize stores UnitRange{Int} axes next to Int extents; StaticSize(ranges) and StaticSize(::CartesianIndices) normalise their input, Base.OneTo axes become plain extents.
• NDRange gains a mapping field holding a StaticOffset or DynamicOffset (or nothing), which expand adds to the blocked index. Plain-size launches keep mapping === nothing, so their NDRange types are unchanged.
• partition normalises ndrange/workgroupsize (integer, range, tuple, CartesianIndices) and compares static and launch ndranges by extents and offsets, so (1:128,) matches a static (128,).
• CompilerMetadata keeps an offset CartesianIndices as ndrange (CartesianIndices(::CartesianIndices) would drop the offsets), so the existing I in __ndrange(ctx) bounds check is correct.
• The global linear index is computed by linear_index, since LinearIndices only supports 1-based axes.
• The POCL autotune path uses extents(ndrange).

Offsets need no backend changes as long as the workgroup size is static or given at launch; autotuning a dynamic workgroup size from a range ndrange requires backends to call extents where they use prod(ndrange).

[vchuravy]

Looks generally okay, but I wouldn't backport it to 0.9 at this point.

Does this have any impact on performance if you don't use offset ranges?

[giordano]

Benchmark script (pass output file as argument)

# Launch-overhead and throughput benchmarks for plain size-tuple ndranges on the CPU backend.
# Run once per KernelAbstractions checkout; prints one line per case: name, min and median
# time, allocations per call.
using KernelAbstractions, BenchmarkTools, Statistics, Printf

@kernel function scale_linear!(A, @Const(B))
    i = @index(Global, Linear)
    @inbounds A[i] = 2 * B[i]
end

@kernel function scale_cartesian!(A, @Const(B))
    i, j, k = @index(Global, NTuple)
    @inbounds A[i, j, k] = 2 * B[i, j, k]
end

# 3-D ndrange, linear global index (exercises the linear-index computation)
@kernel function scale_linear3d!(A, @Const(B))
    i = @index(Global, Linear)
    @inbounds A[i] = 2 * B[i]
end

backend = CPU()
results = Dict{String, Tuple{Float64, Float64, Int}}()

function record!(name, trial)
    results[name] = (minimum(trial).time, median(trial).time, trial.allocs)
    return @printf("%-40s  min %10.3f μs   median %10.3f μs   allocs %5d\n", name, minimum(trial).time / 1e3, median(trial).time / 1e3, trial.allocs)
end

BenchmarkTools.DEFAULT_PARAMETERS.seconds = 3

## Host-side partition cost (no launch)
let k_dd = scale_linear!(backend), k_sd = scale_linear!(backend, 16), k_ss = scale_linear!(backend, 16, (1024,))
    nd, wg = Ref((1024,)), Ref((16,))
    record!("partition dyn-wg dyn-nd", @benchmark KernelAbstractions.partition($k_dd, $nd[], $wg[]))
    record!("partition static-wg dyn-nd", @benchmark KernelAbstractions.partition($k_sd, $nd[], nothing))
    record!("partition static-wg static-nd", @benchmark KernelAbstractions.partition($k_ss, nothing, nothing))
end

## Launch overhead: tiny problem, one workgroup
let n = 16, A = zeros(n), B = rand(n)
    k = scale_linear!(backend)
    record!("launch dyn-wg dyn-nd (autotune)", @benchmark $k($A, $B; ndrange = ($n,)))
    record!("launch dyn-wg dyn-nd (wg given)", @benchmark $k($A, $B; ndrange = ($n,), workgroupsize = ($n,)))
    k = scale_linear!(backend, n)
    record!("launch static-wg dyn-nd", @benchmark $k($A, $B; ndrange = ($n,)))
    k = scale_linear!(backend, n, (n,))
    record!("launch static-wg static-nd", @benchmark $k($A, $B))
end
let dims = (4, 4, 4), A = zeros(dims), B = rand(Float64, dims)
    k = scale_cartesian!(backend, dims)
    record!("launch 3d static-wg dyn-nd", @benchmark $k($A, $B; ndrange = $dims))
    k = scale_cartesian!(backend, dims, dims)
    record!("launch 3d static-wg static-nd", @benchmark $k($A, $B))
end

## Throughput: 2^23 Float64 elements (64 MiB per array)
let n = 2^23, A = zeros(n), B = rand(n)
    k = scale_linear!(backend, 256)
    record!("throughput 1d linear", @benchmark $k($A, $B; ndrange = ($n,)))
    @assert A == 2B
end
let dims = (256, 256, 128), A = zeros(dims), B = rand(Float64, dims)
    k = scale_cartesian!(backend, (16, 16, 1))
    record!("throughput 3d cartesian", @benchmark $k($A, $B; ndrange = $dims))
    @assert A == 2B
    k = scale_linear3d!(backend, (16, 16, 1))
    fill!(A, 0)
    record!("throughput 3d linear-index", @benchmark $k($A, $B; ndrange = $dims))
    @assert A == 2B
    k = scale_cartesian!(backend, (16, 16, 1), dims)
    fill!(A, 0)
    record!("throughput 3d cartesian static-nd", @benchmark $k($A, $B))
    @assert A == 2B
end
# Ragged ndrange: every workgroup row is partially filled, so the bounds check is live
let dims = (250, 250, 120), A = zeros(dims), B = rand(Float64, dims)
    k = scale_cartesian!(backend, (16, 16, 1))
    record!("throughput 3d cartesian ragged", @benchmark $k($A, $B; ndrange = $dims))
    @assert A == 2B
    k = scale_linear3d!(backend, (16, 16, 1))
    fill!(A, 0)
    record!("throughput 3d linear-index ragged", @benchmark $k($A, $B; ndrange = $dims))
    @assert A == 2B
end

open(ARGS[1], "w") do io
    for (name, (mn, md, al)) in sort!(collect(results))
        println(io, name, "\t", mn, "\t", md, "\t", al)
    end
end

Summary of results:

Case	main	offsets	ratio
partition dyn wg, dyn ndrange	35 ns	34 ns	0.97
partition static wg, dyn ndrange	32 ns	32 ns	0.99
partition static wg, static ndrange	1 ns	1 ns	0.99
launch, dyn wg autotune	324 μs	301 μs	0.93
launch, dyn wg given	324 μs	299 μs	0.92
launch, static wg, dyn ndrange	325 μs	300 μs	0.92
launch, static wg, static ndrange	324 μs	330 μs	1.02
launch 3-D, static wg, dyn ndrange	327 μs	338 μs	1.03
launch 3-D, static wg, static ndrange	351 μs	337 μs	0.96
throughput 1-D linear	1.595 ms	1.647 ms	1.03
throughput 3-D Cartesian	1.690 ms	1.660 ms	0.98
throughput 3-D linear index	1.685 ms	1.769 ms	1.05
throughput 3-D Cartesian, static ndrange	1.690 ms	1.772 ms	1.05
throughput 3-D Cartesian ragged	1.478 ms	1.482 ms	1.00
throughput 3-D linear index ragged	1.478 ms	1.480 ms	1.00

Bot comments:

The two 5% entries are drift within that run rather than the branch: the "3-D Cartesian, static ndrange" kernel compiles to identical code on both checkouts, and the two adjacent cases moved together. Three further alternating runs of the linear-index kernels gave medians of 1.647 to 1.661 ms (1-D) and 1.655 to 1.664 ms (3-D) on main against 1.641 to 1.669 ms and 1.655 to 1.675 ms on the branch, so both are within about 1%. Launch medians differ by up to 8% in both directions, matching the run-to-run spread of main alone. Allocations per launch are identical everywhere.

[codecov]

Codecov Report

❌ Patch coverage is 92.42424% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 64.70%. Comparing base (f68a281) to head (cc39448).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/nditeration.jl	92.00%	4 Missing ⚠️
src/compiler.jl	85.71%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #764      +/-   ##
==========================================
+ Coverage   63.72%   64.70%   +0.98%     
==========================================
  Files          23       23              
  Lines        1935     1989      +54     
==========================================
+ Hits         1233     1287      +54     
  Misses        702      702              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[vchuravy]

Can you add these as benchmarks to our current benchmark file?

[christiangnrd]

What's the difference that made you say #765 (comment) over there but not here?

Is it that this allows not materializing the index array whereas in #765 the index array already exists?

[vchuravy]

The history goes back to #399 and #403 back then I was hesitant since I was unsure if this is something that folks generally needed.

I think my issue with using an array is that it overloads the semantics of ndrange too much. I think I would be happier with something like Remap(indices) instead of just passing in an array.

[giordano]

Closing in favour of #771.