From dfb9fcc33cd248b2e79f03f5a8dbe5332d1be297 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mos=C3=A8=20Giordano?= Date: Thu, 10 Sep 2026 11:48:13 -0500 Subject: [PATCH 1/2] Accept index ranges in `ndrange` 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)`. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01VHciC8x39gm97sABrSvBkt --- docs/src/api.md | 6 ++ docs/src/index.md | 5 ++ docs/src/kernels.md | 21 +++++++ src/KernelAbstractions.jl | 12 ++-- src/compiler.jl | 15 +++-- src/nditeration.jl | 119 +++++++++++++++++++++++++++++++++++--- src/pocl/backend.jl | 2 +- test/nditeration.jl | 31 ++++++++++ test/offsets.jl | 95 ++++++++++++++++++++++++++++++ test/test.jl | 41 +++++++++++++ test/testsuite.jl | 5 ++ 11 files changed, 333 insertions(+), 19 deletions(-) create mode 100644 test/offsets.jl diff --git a/docs/src/api.md b/docs/src/api.md index 91f1981a0..e7c6d6023 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -91,4 +91,10 @@ KernelAbstractions.@context KernelAbstractions.argconvert KernelAbstractions.NDIteration.DynamicSize KernelAbstractions.NDIteration.StaticSize +KernelAbstractions.NDIteration.NDRange +KernelAbstractions.NDIteration.StaticOffset +KernelAbstractions.NDIteration.DynamicOffset +KernelAbstractions.NDIteration.extents +KernelAbstractions.NDIteration.offsets +KernelAbstractions.NDIteration.linear_index ``` diff --git a/docs/src/index.md b/docs/src/index.md index 164f6c637..1593c2609 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -132,6 +132,11 @@ but users must avoid the use of `@index(Global)` and instead use their own deriv end ``` +### 0.10 +- `ndrange` entries may be index ranges, given statically (`kernel(backend, workgroupsize, (-2:N+3, 0:M+1))`) + or at launch (`ndrange=(-2:N+3, 0:M+1)`, a single range, or a `CartesianIndices`). + `@index(Global, Cartesian)` and `@index(Global, NTuple)` return the shifted indices. + ## Semantic differences ### To CUDA.jl/AMDGPU.jl diff --git a/docs/src/kernels.md b/docs/src/kernels.md index a01504358..e1bf57a8b 100644 --- a/docs/src/kernels.md +++ b/docs/src/kernels.md @@ -232,5 +232,26 @@ kernel = my_kernel(backend, 32, size(A)) kernel(A) ``` +### Index ranges + +Each entry of `ndrange` is either an extent (indices `1:n`) or a range of indices, so a kernel +can iterate over a region whose indices do not start at 1. `ndrange` can also be given as a +single range or as a `CartesianIndices`: + +```julia +# static ndrange over the indices -2:N+3 along x and 0:M+1 along y +kernel = my_kernel(backend, (16, 16), (-2:N+3, 0:M+1)) +kernel(A) + +# dynamic ndrange +kernel = my_kernel(backend, (16, 16)) +kernel(A, ndrange=(-2:N+3, 0:M+1)) +kernel(A, ndrange=CartesianIndices(A)) # e.g. for an OffsetArray +``` + +Inside the kernel `@index(Global, Cartesian)` and `@index(Global, NTuple)` return the shifted +indices, `@index(Global, Linear)` counts the indices of the region from 1 in column-major order, +and `@ndrange()` returns the extents. + Obtain the backend from an array with [`get_backend`](@ref) and always call [`synchronize`](@ref) before reading results on the host. See the [Quickstart](@ref) for a full walkthrough and the Examples section of the manual for larger patterns. diff --git a/src/KernelAbstractions.jl b/src/KernelAbstractions.jl index 17a725427..464b55d98 100644 --- a/src/KernelAbstractions.jl +++ b/src/KernelAbstractions.jl @@ -414,7 +414,7 @@ end @inline function __index_Global_Linear(ctx) I = @inbounds expand(__iterspace(ctx), KI.get_group_id().x, KI.get_local_id().x) # TODO: This is unfortunate, can we get the linear index cheaper - return @inbounds LinearIndices(__ndrange(ctx))[I] + return linear_index(__ndrange(ctx), I) end @inline function __index_Local_Cartesian(ctx) @@ -544,6 +544,8 @@ last (possibly partial) workgroup. Primarily used by backend implementations and @inline function partition(kernel, ndrange, workgroupsize) static_ndrange = KernelAbstractions.ndrange(kernel) static_workgroupsize = KernelAbstractions.workgroupsize(kernel) + ndrange = NDIteration.normalize_ndrange(ndrange) + workgroupsize = NDIteration.normalize_workgroupsize(workgroupsize) if ndrange === nothing && static_ndrange <: DynamicSize || workgroupsize === nothing && static_workgroupsize <: DynamicSize @@ -562,7 +564,7 @@ last (possibly partial) workgroup. Primarily used by backend implementations and end if static_ndrange <: StaticSize - if ndrange !== nothing && ndrange != get(static_ndrange) + if ndrange !== nothing && !NDIteration.same_axes(ndrange, get(static_ndrange)) error("Static NDRange ($static_ndrange) and launch NDRange ($ndrange) differ") end ndrange = get(static_ndrange) @@ -577,14 +579,16 @@ last (possibly partial) workgroup. Primarily used by backend implementations and @assert workgroupsize !== nothing @assert ndrange !== nothing - blocks, workgroupsize, dynamic = NDIteration.partition(ndrange, workgroupsize) + blocks, workgroupsize, dynamic = NDIteration.partition(extents(ndrange), workgroupsize) if static_ndrange <: StaticSize static_blocks = StaticSize{blocks} blocks = nothing + mapping = NDIteration.static_mapping(ndrange) else static_blocks = DynamicSize blocks = CartesianIndices(blocks) + mapping = NDIteration.dynamic_mapping(ndrange) end if static_workgroupsize <: StaticSize @@ -594,7 +598,7 @@ last (possibly partial) workgroup. Primarily used by backend implementations and workgroupsize = CartesianIndices(workgroupsize) end - iterspace = NDRange{length(ndrange), static_blocks, static_workgroupsize}(blocks, workgroupsize) + iterspace = NDRange{length(ndrange), static_blocks, static_workgroupsize}(blocks, workgroupsize, mapping) return iterspace, dynamic end diff --git a/src/compiler.jl b/src/compiler.jl index 2950ea275..b7d388d62 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -5,21 +5,24 @@ struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace} # CPU variant function CompilerMetadata{NDRange, CB}(idx, ndrange, iterspace) where {NDRange, CB} - if ndrange !== nothing - ndrange = CartesianIndices(ndrange) - end + ndrange = cartesian(ndrange) return new{NDRange, CB, typeof(idx), typeof(ndrange), typeof(iterspace)}(idx, ndrange, iterspace) end # GPU variante: index is given implicit function CompilerMetadata{NDRange, CB}(ndrange, iterspace) where {NDRange, CB} - if ndrange !== nothing - ndrange = CartesianIndices(ndrange) - end + ndrange = cartesian(ndrange) return new{NDRange, CB, Nothing, typeof(ndrange), typeof(iterspace)}(nothing, ndrange, iterspace) end end +# `CartesianIndices` covering a launch `ndrange` (any form accepted by `partition`). +cartesian(::Nothing) = nothing +cartesian(ci::CartesianIndices) = ci +cartesian(n::Integer) = CartesianIndices((Int(n),)) +cartesian(r::AbstractUnitRange) = CartesianIndices((r,)) +cartesian(t::Tuple) = CartesianIndices(t) + @inline __iterspace(cm::CompilerMetadata) = cm.iterspace @inline __groupindex(cm::CompilerMetadata) = cm.groupindex @inline __groupsize(cm::CompilerMetadata) = size(workitems(__iterspace(cm))) diff --git a/src/nditeration.jl b/src/nditeration.jl index aacaa8bff..9b1274d00 100644 --- a/src/nditeration.jl +++ b/src/nditeration.jl @@ -2,6 +2,7 @@ module NDIteration export _Size, StaticSize, DynamicSize, get export NDRange, blocks, workitems, expand +export StaticOffset, DynamicOffset, offsets, extents, linear_index export DynamicCheck, NoDynamicCheck import Base.@pure @@ -9,6 +10,69 @@ import Base.@pure struct DynamicCheck end struct NoDynamicCheck end +# An axis of an `ndrange` is either an extent (`Int`) or a range of indices (`UnitRange{Int}`). +axis(n::Integer) = Int(n) +axis(r::Base.OneTo) = Int(length(r)) +axis(r::AbstractUnitRange) = UnitRange{Int}(r) + +extent(n::Integer) = Int(n) +extent(r::AbstractUnitRange) = length(r) + +axis_offset(::Integer) = 0 +axis_offset(r::AbstractUnitRange) = first(r) - 1 + +""" + extents(ndrange) + +Number of indices along each axis of `ndrange`, given as a tuple of extents and/or ranges, +a `CartesianIndices`, a single range, or an integer. +""" +extents(t::Tuple) = map(extent, t) +extents(ci::CartesianIndices) = size(ci) +extents(r::AbstractUnitRange) = (length(r),) +extents(n::Integer) = (Int(n),) + +""" + offsets(ndrange) + +Offset of the first index along each axis of `ndrange` relative to 1. +""" +offsets(t::Tuple) = map(axis_offset, t) + +""" + normalize_ndrange(ndrange) + +Canonical form of a launch `ndrange`: `nothing`, or a tuple of `Int` extents and +`UnitRange{Int}` axes. +""" +normalize_ndrange(::Nothing) = nothing +normalize_ndrange(n::Integer) = (Int(n),) +normalize_ndrange(r::AbstractUnitRange) = (axis(r),) +normalize_ndrange(ci::CartesianIndices) = map(axis, ci.indices) +normalize_ndrange(t::Tuple) = map(axis, t) + +""" + normalize_workgroupsize(workgroupsize) + +Canonical form of a launch `workgroupsize`: `nothing`, or a tuple of `Int` extents. +""" +normalize_workgroupsize(::Nothing) = nothing +normalize_workgroupsize(n::Integer) = (Int(n),) +normalize_workgroupsize(t::Tuple) = extents(t) + +# Two ndranges denote the same indices. +same_axes(a::Tuple, b::Tuple) = extents(a) == extents(b) && offsets(a) == offsets(b) + +""" + linear_index(ndrange::CartesianIndices, I::CartesianIndex) + +Column-major position of `I` within `ndrange`, counted from 1. +""" +@inline function linear_index(ndrange::CartesianIndices{N}, I::CartesianIndex{N}) where {N} + lo = map(first, ndrange.indices) + return @inbounds LinearIndices(size(ndrange))[CartesianIndex(I.I .- lo .+ 1)] +end + abstract type _Size end """ @@ -22,29 +86,54 @@ struct DynamicSize <: _Size end StaticSize{S} Marker type encoding a compile-time workgroup size or `ndrange` as a tuple `S`. +Each entry of `S` is an `Int` extent or, for an `ndrange` axis whose indices do not start +at 1, a `UnitRange{Int}`. """ struct StaticSize{S} <: _Size function StaticSize{S}() where {S} - return new{S::Tuple{Vararg{Int}}}() + return new{S::Tuple{Vararg{Union{Int, UnitRange{Int}}}}}() end end @pure StaticSize(s::Tuple{Vararg{Int}}) = StaticSize{s}() @pure StaticSize(s::Int...) = StaticSize{s}() @pure StaticSize(s::Type{<:Tuple}) = StaticSize{tuple(s.parameters...)}() +StaticSize(s::Tuple{Vararg{Union{Integer, AbstractUnitRange{<:Integer}}}}) = StaticSize{map(axis, s)}() +StaticSize(ci::CartesianIndices) = StaticSize(ci.indices) # Some @pure convenience functions for `StaticSize` @pure get(::Type{StaticSize{S}}) where {S} = S @pure get(::StaticSize{S}) where {S} = S @pure Base.getindex(::StaticSize{S}, i::Int) where {S} = i <= length(S) ? S[i] : 1 @pure Base.ndims(::StaticSize{S}) where {S} = length(S) -@pure Base.length(::StaticSize{S}) where {S} = prod(S) +@pure Base.length(::StaticSize{S}) where {S} = prod(extents(S)) + +""" + StaticOffset{O} + +Compile-time offset `O::NTuple{N, Int}` added to the indices produced by an [`NDRange`](@ref). +""" +struct StaticOffset{O} + function StaticOffset{O}() where {O} + return new{O::Tuple{Vararg{Int}}}() + end +end +""" + DynamicOffset{N} + +Runtime offset added to the indices produced by an [`NDRange`](@ref). +""" +struct DynamicOffset{N} + offset::NTuple{N, Int} +end """ NDRange -Encodes a blocked iteration space. +Encodes a blocked iteration space. The `mapping` field relates blocked indices to +`ndrange` indices: `nothing` for the identity, or a [`StaticOffset`](@ref)/[`DynamicOffset`](@ref) +for an `ndrange` whose indices do not start at 1. # Example ``` @@ -58,16 +147,17 @@ for block in ndrange end ``` """ -struct NDRange{N, StaticBlocks, StaticWorkitems, DynamicBlock, DynamicWorkitems} +struct NDRange{N, StaticBlocks, StaticWorkitems, DynamicBlock, DynamicWorkitems, Mapping} blocks::DynamicBlock workitems::DynamicWorkitems + mapping::Mapping function NDRange{N, B, W}() where {N, B, W} - return new{N, B, W, Nothing, Nothing}(nothing, nothing) + return new{N, B, W, Nothing, Nothing, Nothing}(nothing, nothing, nothing) end - function NDRange{N, B, W}(blocks, workitems) where {N, B, W} - return new{N, B, W, typeof(blocks), typeof(workitems)}(blocks, workitems) + function NDRange{N, B, W}(blocks, workitems, mapping = nothing) where {N, B, W} + return new{N, B, W, typeof(blocks), typeof(workitems), typeof(mapping)}(blocks, workitems, mapping) end end @@ -77,6 +167,16 @@ end @inline blocks(range::NDRange{N, B}) where {N, B <: StaticSize} = CartesianIndices(get(B))::CartesianIndices{N} @inline Base.ndims(::NDRange{N}) where {N} = N +@inline offsets(::NDRange{N, B, W, DB, DW, Nothing}) where {N, B, W, DB, DW} = ntuple(_ -> 0, Val(N)) +@inline offsets(::NDRange{N, B, W, DB, DW, StaticOffset{O}}) where {N, B, W, DB, DW, O} = O +@inline offsets(range::NDRange{N, B, W, DB, DW, DynamicOffset{N}}) where {N, B, W, DB, DW} = range.mapping.offset + +# Mapping of a partitioned `ndrange` (in canonical form); a plain size tuple has no mapping. +static_mapping(::Tuple{Vararg{Int}}) = nothing +static_mapping(t::Tuple) = StaticOffset{offsets(t)}() +dynamic_mapping(::Tuple{Vararg{Int}}) = nothing +dynamic_mapping(t::Tuple) = DynamicOffset(offsets(t)) + import Base.iterate @inline iterate(range::NDRange) = iterate(blocks(range)) @inline iterate(range::NDRange, state) = iterate(blocks(range), state) @@ -84,11 +184,12 @@ import Base.iterate Base.length(range::NDRange) = length(blocks(range)) @inline function expand(ndrange::NDRange{N}, groupidx::CartesianIndex{N}, idx::CartesianIndex{N}) where {N} + offset = offsets(ndrange) nI = ntuple(Val(N)) do I Base.@_inline_meta stride = size(workitems(ndrange), I) gidx = groupidx.I[I] - (gidx - 1) * stride + idx.I[I] + (gidx - 1) * stride + idx.I[I] + offset[I] end return CartesianIndex(nI) end @@ -153,6 +254,8 @@ Returns the number of workgroups necessary and whether the last workgroup needs to perform dynamic bounds-checking. """ @inline function partition(ndrange, __workgroupsize) + ndrange = extents(ndrange) + __workgroupsize = extents(__workgroupsize) @assert length(__workgroupsize) <= length(ndrange) # pad workgroupsize with ones workgroupsize = ntuple(Val(length(ndrange))) do I diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index 98a606afe..344e44791 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -196,7 +196,7 @@ function (obj::KA.Kernel{POCLBackend})(args...; ndrange = nothing, workgroupsize # figure out the optimal workgroupsize automatically if KA.workgroupsize(obj) <: KA.DynamicSize && workgroupsize === nothing wg_info = cl.work_group_info(kernel.fun, device()) - wg_size_nd = threads_to_workgroupsize(wg_info.size, ndrange) + wg_size_nd = threads_to_workgroupsize(wg_info.size, KA.NDIteration.extents(ndrange)) iterspace, dynamic = KA.partition(obj, ndrange, wg_size_nd) ctx = KA.mkcontext(obj, ndrange, iterspace) end diff --git a/test/nditeration.jl b/test/nditeration.jl index c0a8c08c9..5b0a762f3 100644 --- a/test/nditeration.jl +++ b/test/nditeration.jl @@ -16,6 +16,37 @@ function nditeration_testsuite() end end + @testset "offsets" begin + @test NDIteration.get(StaticSize((1:4, 0:9))) == (1:4, 0:9) + @test NDIteration.get(StaticSize(CartesianIndices((3, 0:9)))) == (3, 0:9) + @test length(StaticSize((1:4, 0:9))) == 40 + @test extents((1:4, 0:9, 7)) == (4, 10, 7) + @test extents(CartesianIndices((3, 0:9))) == (3, 10) + @test extents(0:9) == (10,) + @test offsets((1:4, 0:9, 7)) == (0, -1, 0) + + let ndrange = NDRange{2, StaticSize{(4, 4)}, StaticSize{(8, 8)}}(nothing, nothing, StaticOffset{(-8, 3)}()) + @test offsets(ndrange) == (-8, 3) + @test expand(ndrange, CartesianIndex(1, 1), CartesianIndex(1, 1)) == CartesianIndex(-7, 4) + @test expand(ndrange, CartesianIndex(4, 4), CartesianIndex(8, 8)) == CartesianIndex(24, 35) + end + let ndrange = NDRange{2, DynamicSize, DynamicSize}(CartesianIndices((4, 4)), CartesianIndices((8, 8)), DynamicOffset((-8, 3))) + @test offsets(ndrange) == (-8, 3) + @test expand(ndrange, 1, 1) == CartesianIndex(-7, 4) + @test expand(ndrange, 16, 64) == CartesianIndex(24, 35) + end + let ndrange = NDRange{2, DynamicSize, DynamicSize}(CartesianIndices((4, 4)), CartesianIndices((8, 8))) + @test offsets(ndrange) == (0, 0) + @test ndrange.mapping === nothing + end + + let ci = CartesianIndices((-3:4, 2:11)) + @test linear_index(ci, CartesianIndex(-3, 2)) == 1 + @test linear_index(ci, CartesianIndex(4, 2)) == 8 + @test linear_index(ci, CartesianIndex(4, 11)) == 80 + end + end + # GPU scenario where we get a linear index into workitems/blocks function linear_iteration(ndrange) idx = Array{CartesianIndex{2}}(undef, length(blocks(ndrange)) * length(workitems(ndrange))) diff --git a/test/offsets.jl b/test/offsets.jl new file mode 100644 index 000000000..0367f6438 --- /dev/null +++ b/test/offsets.jl @@ -0,0 +1,95 @@ +using KernelAbstractions +using KernelAbstractions.NDIteration +using Test + +@kernel function offsets_fill_indices!(out, lo) + I = @index(Global, NTuple) + i = @index(Global, Linear) + @inbounds out[(I .- lo .+ 1)...] = i +end + +@kernel function offsets_fill_ndrange!(out, lo) + I = @index(Global, NTuple) + sz = @ndrange() + @inbounds out[(I .- lo .+ 1)...] = sz[1] +end + +@kernel function offsets_fill_cartesian!(out, lo) + I = @index(Global, Cartesian) + @inbounds out[I - lo + oneunit(I)] = 1 +end + +function offsets_testsuite(Backend, AT) + backend = Backend() + ranges = (-3:4, 2:11) + lo = map(first, ranges) + ref = reshape(1:80, 8, 10) + fresh() = AT(zeros(Int, 8, 10)) + + @testset "static ndrange" begin + out = fresh() + offsets_fill_indices!(backend, (4, 4), ranges)(out, lo) + synchronize(backend) + @test Array(out) == ref + end + + @testset "dynamic ndrange" begin + out = fresh() + offsets_fill_indices!(backend, (4, 4))(out, lo; ndrange = ranges) + synchronize(backend) + @test Array(out) == ref + + out = fresh() + offsets_fill_indices!(backend, (4, 4))(out, lo; ndrange = CartesianIndices(ranges)) + synchronize(backend) + @test Array(out) == ref + + out = fresh() + offsets_fill_indices!(backend)(out, lo; ndrange = ranges, workgroupsize = (4, 4)) + synchronize(backend) + @test Array(out) == ref + end + + @testset "mixed extents and ranges" begin + out = fresh() + offsets_fill_indices!(backend, (4, 4))(out, (1, 2); ndrange = (8, 2:11)) + synchronize(backend) + @test Array(out) == ref + end + + @testset "ragged workgroups" begin + out = fresh() + offsets_fill_indices!(backend, (3, 3))(out, lo; ndrange = ranges) + synchronize(backend) + @test Array(out) == ref + end + + @testset "bare range" begin + out = AT(zeros(Int, 16)) + offsets_fill_indices!(backend, 4)(out, (5,); ndrange = 5:20) + synchronize(backend) + @test Array(out) == 1:16 + end + + @testset "cartesian index" begin + out = fresh() + offsets_fill_cartesian!(backend, (4, 4))(out, CartesianIndex(lo); ndrange = ranges) + synchronize(backend) + @test all(==(1), Array(out)) + end + + @testset "@ndrange returns extents" begin + out = fresh() + offsets_fill_ndrange!(backend, (4, 4))(out, lo; ndrange = ranges) + synchronize(backend) + @test all(==(8), Array(out)) + end + + @testset "empty range" begin + out = fresh() + offsets_fill_indices!(backend, (4, 4))(out, lo; ndrange = (5:4, 1:3)) + synchronize(backend) + @test all(iszero, Array(out)) + end + return +end diff --git a/test/test.jl b/test/test.jl index 53ef1fa41..d202b51a8 100644 --- a/test/test.jl +++ b/test/test.jl @@ -48,6 +48,47 @@ function unittest_testsuite(Backend, backend_str, backend_mod, BackendArrayT; sk @test_throws ErrorException KernelAbstractions.partition(kernel, (129,), nothing) @test KernelAbstractions.backend(kernel) == backend end + let kernel = KernelAbstractions.Kernel{typeof(backend), StaticSize{(64,)}, DynamicSize, typeof(identity)}(backend, identity) + iterspace, dynamic = KernelAbstractions.partition(kernel, (-63:64,), nothing) + @test length(blocks(iterspace)) == 2 + @test dynamic isa NoDynamicCheck + @test offsets(iterspace) == (-64,) + @test iterspace.mapping isa DynamicOffset + + iterspace, dynamic = KernelAbstractions.partition(kernel, CartesianIndices((0:128,)), (64,)) + @test length(blocks(iterspace)) == 3 + @test dynamic isa DynamicCheck + @test offsets(iterspace) == (-1,) + + iterspace, dynamic = KernelAbstractions.partition(kernel, 0:127, nothing) + @test length(blocks(iterspace)) == 2 + @test offsets(iterspace) == (-1,) + + iterspace, dynamic = KernelAbstractions.partition(kernel, (128,), nothing) + @test iterspace.mapping === nothing + + # a range in place of the workgroup size is taken as its length + iterspace, dynamic = KernelAbstractions.partition(kernel, (-63:64,), (-63:0,)) + @test length(blocks(iterspace)) == 2 + end + let kernel = KernelAbstractions.Kernel{typeof(backend), StaticSize{(64,)}, StaticSize{(-63:64,)}, typeof(identity)}(backend, identity) + iterspace, dynamic = KernelAbstractions.partition(kernel, nothing, nothing) + @test length(blocks(iterspace)) == 2 + @test dynamic isa NoDynamicCheck + @test offsets(iterspace) == (-64,) + @test iterspace.mapping isa StaticOffset + + iterspace, dynamic = KernelAbstractions.partition(kernel, (-63:64,), nothing) + @test length(blocks(iterspace)) == 2 + + @test_throws ErrorException KernelAbstractions.partition(kernel, (128,), nothing) + @test_throws ErrorException KernelAbstractions.partition(kernel, (-62:65,), nothing) + end + let kernel = KernelAbstractions.Kernel{typeof(backend), StaticSize{(64,)}, StaticSize{(128,)}, typeof(identity)}(backend, identity) + iterspace, dynamic = KernelAbstractions.partition(kernel, (1:128,), nothing) + @test length(blocks(iterspace)) == 2 + @test iterspace.mapping === nothing + end end @kernel function index_linear_global(A) diff --git a/test/testsuite.jl b/test/testsuite.jl index c647b14e6..9c9db2c8e 100644 --- a/test/testsuite.jl +++ b/test/testsuite.jl @@ -33,6 +33,7 @@ include("localmem.jl") include("private.jl") include("unroll.jl") include("nditeration.jl") +include("offsets.jl") include("copyto.jl") include("devices.jl") include("print_test.jl") @@ -75,6 +76,10 @@ function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{ nditeration_testsuite() end + @conditional_testset "Offsets" skip_tests begin + offsets_testsuite(backend, AT) + end + @conditional_testset "copyto!" skip_tests begin copyto_testsuite(backend, AT) end From c0b180e35a2e27dfe5c5cd649a48943716662f5e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mos=C3=A8=20Giordano?= Date: Thu, 10 Sep 2026 11:56:52 -0500 Subject: [PATCH 2/2] Launch kernels over a list of indices `ndrange` may now be a device vector of `CartesianIndex{N}` or `NTuple{N, <:Integer}` elements. Each work item handles one listed index: `@index(Global, Cartesian)` and `@index(Global, NTuple)` return it, `@index(Global, Linear)` its position in the vector, `@ndrange()` the vector's length. The kernel must have a dynamic `ndrange` and a 1-D workgroup size, static or given at launch; the partial last workgroup is bounds-checked against the vector's length. This is how Oceananigans iterates over the active cells of an immersed boundary grid. It currently wraps the kernel function, pirates `Base.getproperty` on `Kernel`, `Adapt.adapt_structure` on `NDRange` and `CompilerMetadata`, `partition`, `expand`, and adds one `@device_override __validindex` per GPU backend, which the compiled CPU backend's own `::Any` override shadows. Implementation: - `IndexMap{N}` wraps the vector and yields `CartesianIndex{N}` on indexing; it is the `mapping` of a 1-D `NDRange` (`MappedNDRange`). `expand` and `linear_index` on a `MappedNDRange` look the index up. - `partition` dispatches vectors to `mapped_partition`, which validates the kernel and workgroup size and partitions `(length(map),)`; the dynamic-check flag is computed as for any other ndrange. - Index validity moves into the generic `__validindex(ctx, groupidx, idx)`, which dispatches on the iteration space (`I in ndrange` for regular spaces, position `<= length(map)` for index maps). The 1-arg `__validindex(ctx)` is now a plain method built on `KernelInterface.get_group_id`/`get_local_id`, so the POCL override is removed; `__index_Global_Linear` goes through `__global_linear` in the same way. Backends that still override `__validindex(ctx)` keep working for regular ndranges and should forward to the 3-arg method to support index maps. - `Adapt.adapt_structure` for `NDRange`, `IndexMap` and `CompilerMetadata` moves the vector to the device with the kernel arguments; `CompilerMetadata` treats a vector `ndrange` as `CartesianIndices((length(v),))`. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01VHciC8x39gm97sABrSvBkt --- docs/src/api.md | 3 ++ docs/src/index.md | 4 ++ docs/src/kernels.md | 18 +++++++++ src/KernelAbstractions.jl | 73 ++++++++++++++++++++++++++++++++-- src/compiler.jl | 5 +++ src/nditeration.jl | 65 +++++++++++++++++++++++++++++-- src/pocl/backend.jl | 9 ----- test/indexmap.jl | 82 +++++++++++++++++++++++++++++++++++++++ test/nditeration.jl | 24 ++++++++++++ test/test.jl | 32 +++++++++++++++ test/testsuite.jl | 5 +++ 11 files changed, 303 insertions(+), 17 deletions(-) create mode 100644 test/indexmap.jl diff --git a/docs/src/api.md b/docs/src/api.md index e7c6d6023..65a6a9863 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -97,4 +97,7 @@ KernelAbstractions.NDIteration.DynamicOffset KernelAbstractions.NDIteration.extents KernelAbstractions.NDIteration.offsets KernelAbstractions.NDIteration.linear_index +KernelAbstractions.NDIteration.IndexMap +KernelAbstractions.NDIteration.MappedNDRange +KernelAbstractions.__validindex ``` diff --git a/docs/src/index.md b/docs/src/index.md index 1593c2609..a1a385f62 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -136,6 +136,10 @@ end - `ndrange` entries may be index ranges, given statically (`kernel(backend, workgroupsize, (-2:N+3, 0:M+1))`) or at launch (`ndrange=(-2:N+3, 0:M+1)`, a single range, or a `CartesianIndices`). `@index(Global, Cartesian)` and `@index(Global, NTuple)` return the shifted indices. +- `ndrange` may be a device vector of `CartesianIndex`/`NTuple` indices, running one work item per + listed index (`kernel(A, ndrange=active_cells)`). +- Index validity is decided by the generic `__validindex(ctx, groupidx, idx)`, so backends only + supply the hardware indices. ## Semantic differences diff --git a/docs/src/kernels.md b/docs/src/kernels.md index e1bf57a8b..cc3b0a8d2 100644 --- a/docs/src/kernels.md +++ b/docs/src/kernels.md @@ -253,5 +253,23 @@ Inside the kernel `@index(Global, Cartesian)` and `@index(Global, NTuple)` retur indices, `@index(Global, Linear)` counts the indices of the region from 1 in column-major order, and `@ndrange()` returns the extents. +### Index maps + +`ndrange` can be a vector of indices (`CartesianIndex{N}` or `NTuple{N, <:Integer}` elements, +stored on the backend's device) to run one work item per listed index, for example over the +active cells of a masked domain: + +```julia +active = CuArray([CartesianIndex(i, j, k) for (i, j, k) in cells if mask[i, j, k]]) +kernel = my_kernel(backend, 256) # 1-D workgroup size is required +kernel(A, ndrange=active) +``` + +Inside the kernel `@index(Global, Cartesian)` and `@index(Global, NTuple)` return the listed +index, `@index(Global, Linear)` its position in the vector, and `@ndrange()` the length of +the vector. The kernel must be constructed with a dynamic `ndrange`, and the workgroup size +must be static or given with `workgroupsize`. With `@kernel unsafe_indices=true` the work +items of a partial last workgroup have no valid index. + Obtain the backend from an array with [`get_backend`](@ref) and always call [`synchronize`](@ref) before reading results on the host. See the [Quickstart](@ref) for a full walkthrough and the Examples section of the manual for larger patterns. diff --git a/src/KernelAbstractions.jl b/src/KernelAbstractions.jl index 464b55d98..3b41f7546 100644 --- a/src/KernelAbstractions.jl +++ b/src/KernelAbstractions.jl @@ -412,9 +412,7 @@ end end @inline function __index_Global_Linear(ctx) - I = @inbounds expand(__iterspace(ctx), KI.get_group_id().x, KI.get_local_id().x) - # TODO: This is unfortunate, can we get the linear index cheaper - return linear_index(__ndrange(ctx), I) + return __global_linear(__iterspace(ctx), __ndrange(ctx), KI.get_group_id().x, KI.get_local_id().x) end @inline function __index_Local_Cartesian(ctx) @@ -547,6 +545,10 @@ last (possibly partial) workgroup. Primarily used by backend implementations and ndrange = NDIteration.normalize_ndrange(ndrange) workgroupsize = NDIteration.normalize_workgroupsize(workgroupsize) + if ndrange isa IndexMap + return mapped_partition(kernel, ndrange, workgroupsize) + end + if ndrange === nothing && static_ndrange <: DynamicSize || workgroupsize === nothing && static_workgroupsize <: DynamicSize errmsg = """ @@ -602,6 +604,39 @@ last (possibly partial) workgroup. Primarily used by backend implementations and return iterspace, dynamic end +# Partition of an index map: a 1-D blocked space over the positions in the map. +@inline function mapped_partition(kernel, map::IndexMap, workgroupsize) + static_ndrange = KernelAbstractions.ndrange(kernel) + static_workgroupsize = KernelAbstractions.workgroupsize(kernel) + + if static_ndrange <: StaticSize + error("An index map is a runtime iteration space; construct the kernel with a dynamic ndrange") + end + if static_workgroupsize <: StaticSize + if workgroupsize !== nothing && workgroupsize != get(static_workgroupsize) + error("Static WorkgroupSize ($static_workgroupsize) and launch WorkgroupSize $(workgroupsize) differ") + end + workgroupsize = get(static_workgroupsize) + elseif !(workgroupsize isa Tuple) + error("An index map requires a workgroup size, either static or given with `workgroupsize`") + end + if length(workgroupsize) != 1 + error("An index map requires a 1-D workgroup size, got $(workgroupsize)") + end + + blocks, workgroupsize, dynamic = NDIteration.partition((length(map),), workgroupsize) + + if static_workgroupsize <: StaticSize + static_workgroupsize = StaticSize{workgroupsize} + workgroupsize = nothing + else + workgroupsize = CartesianIndices(workgroupsize) + end + + iterspace = NDRange{1, DynamicSize, static_workgroupsize}(CartesianIndices(blocks), workgroupsize, map) + return iterspace, dynamic +end + function construct(backend::Backend, ::S, ::NDRange, xpu_name::XPUName) where {Backend <: GPU, S <: _Size, NDRange <: _Size, XPUName} return Kernel{Backend, S, NDRange, XPUName}(backend, xpu_name) end @@ -617,7 +652,37 @@ include("compiler.jl") ### function __workitems_iterspace end -function __validindex end + +""" + __validindex(ctx, groupidx, idx) + +Whether work item `idx` of workgroup `groupidx` has an index within the `ndrange`. +Both indices are linear or `CartesianIndex` positions within the blocked iteration space. +""" +@inline function __validindex(ctx, groupidx, idx) + if __dynamic_checkbounds(ctx) + return __inrange(__iterspace(ctx), __ndrange(ctx), groupidx, idx) + else + return true + end +end + +@inline function __validindex(ctx) + return __validindex(ctx, KI.get_group_id().x, KI.get_local_id().x) +end + +@inline function __inrange(iterspace::NDRange, ndrange, groupidx, idx) + I = @inbounds expand(iterspace, groupidx, idx) + return I in ndrange +end +@inline __inrange(iterspace::MappedNDRange, ndrange, groupidx, idx) = linear_index(iterspace, groupidx, idx) <= length(iterspace.mapping) + +# Global linear index of work item `idx` of workgroup `groupidx`. +@inline function __global_linear(iterspace::NDRange, ndrange, groupidx, idx) + I = @inbounds expand(iterspace, groupidx, idx) + return linear_index(ndrange, I) +end +@inline __global_linear(iterspace::MappedNDRange, ndrange, groupidx, idx) = linear_index(iterspace, groupidx, idx) # for reflection function mkcontext end diff --git a/src/compiler.jl b/src/compiler.jl index b7d388d62..f0ade67cc 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -22,6 +22,11 @@ cartesian(ci::CartesianIndices) = ci cartesian(n::Integer) = CartesianIndices((Int(n),)) cartesian(r::AbstractUnitRange) = CartesianIndices((r,)) cartesian(t::Tuple) = CartesianIndices(t) +cartesian(v::AbstractVector) = CartesianIndices((length(v),)) +cartesian(m::IndexMap) = CartesianIndices((length(m),)) + +Adapt.adapt_structure(to, cm::CompilerMetadata{NDRange, CB}) where {NDRange, CB} = + CompilerMetadata{NDRange, CB}(cm.groupindex, cm.ndrange, Adapt.adapt(to, cm.iterspace)) @inline __iterspace(cm::CompilerMetadata) = cm.iterspace @inline __groupindex(cm::CompilerMetadata) = cm.groupindex diff --git a/src/nditeration.jl b/src/nditeration.jl index 9b1274d00..f899c768b 100644 --- a/src/nditeration.jl +++ b/src/nditeration.jl @@ -3,9 +3,11 @@ module NDIteration export _Size, StaticSize, DynamicSize, get export NDRange, blocks, workitems, expand export StaticOffset, DynamicOffset, offsets, extents, linear_index +export IndexMap, MappedNDRange export DynamicCheck, NoDynamicCheck import Base.@pure +import Adapt struct DynamicCheck end struct NoDynamicCheck end @@ -31,6 +33,7 @@ extents(t::Tuple) = map(extent, t) extents(ci::CartesianIndices) = size(ci) extents(r::AbstractUnitRange) = (length(r),) extents(n::Integer) = (Int(n),) +extents(v::AbstractVector) = (length(v),) """ offsets(ndrange) @@ -39,26 +42,54 @@ Offset of the first index along each axis of `ndrange` relative to 1. """ offsets(t::Tuple) = map(axis_offset, t) +""" + IndexMap{N}(map::AbstractVector) + +Iteration space given by the indices listed in `map`, whose elements are `CartesianIndex{N}` +or `NTuple{N, <:Integer}`. Work item `p` handles the index `map[p]`. +""" +struct IndexMap{N, A <: AbstractVector} + map::A + IndexMap{N}(map::AbstractVector) where {N} = new{N, typeof(map)}(map) +end +IndexMap(map::AbstractVector) = IndexMap{mapdims(eltype(map))}(map) + +mapdims(::Type{CartesianIndex{N}}) where {N} = N +mapdims(::Type{<:NTuple{N, Integer}}) where {N} = N +mapdims(::Type{T}) where {T} = throw(ArgumentError("an index map must have elements of type `CartesianIndex{N}` or `NTuple{N, Integer}`, got `$T`")) + +Base.length(m::IndexMap) = length(m.map) +extents(m::IndexMap) = (length(m),) +Base.@propagate_inbounds Base.getindex(m::IndexMap{N}, i::Integer) where {N} = mapindex(Val(N), m.map[i]) +mapindex(::Val{N}, I::CartesianIndex{N}) where {N} = I +mapindex(::Val{N}, I::Tuple) where {N} = CartesianIndex{N}(I) + +Adapt.adapt_structure(to, m::IndexMap{N}) where {N} = IndexMap{N}(Adapt.adapt(to, m.map)) + """ normalize_ndrange(ndrange) -Canonical form of a launch `ndrange`: `nothing`, or a tuple of `Int` extents and -`UnitRange{Int}` axes. +Canonical form of a launch `ndrange`: `nothing`, a tuple of `Int` extents and +`UnitRange{Int}` axes, or an [`IndexMap`](@ref). """ normalize_ndrange(::Nothing) = nothing normalize_ndrange(n::Integer) = (Int(n),) normalize_ndrange(r::AbstractUnitRange) = (axis(r),) normalize_ndrange(ci::CartesianIndices) = map(axis, ci.indices) normalize_ndrange(t::Tuple) = map(axis, t) +normalize_ndrange(m::IndexMap) = m +normalize_ndrange(v::AbstractVector) = IndexMap(v) """ normalize_workgroupsize(workgroupsize) Canonical form of a launch `workgroupsize`: `nothing`, or a tuple of `Int` extents. +Anything else is passed through to be rejected by `partition`. """ normalize_workgroupsize(::Nothing) = nothing normalize_workgroupsize(n::Integer) = (Int(n),) normalize_workgroupsize(t::Tuple) = extents(t) +normalize_workgroupsize(x) = x # Two ndranges denote the same indices. same_axes(a::Tuple, b::Tuple) = extents(a) == extents(b) && offsets(a) == offsets(b) @@ -132,8 +163,9 @@ end NDRange Encodes a blocked iteration space. The `mapping` field relates blocked indices to -`ndrange` indices: `nothing` for the identity, or a [`StaticOffset`](@ref)/[`DynamicOffset`](@ref) -for an `ndrange` whose indices do not start at 1. +`ndrange` indices: `nothing` for the identity, a [`StaticOffset`](@ref)/[`DynamicOffset`](@ref) +for an `ndrange` whose indices do not start at 1, or an [`IndexMap`](@ref) for a 1-D +blocked space whose work items look up their index in a list. # Example ``` @@ -177,6 +209,16 @@ static_mapping(t::Tuple) = StaticOffset{offsets(t)}() dynamic_mapping(::Tuple{Vararg{Int}}) = nothing dynamic_mapping(t::Tuple) = DynamicOffset(offsets(t)) +""" + MappedNDRange + +A 1-D blocked iteration space whose `mapping` is an [`IndexMap`](@ref). +""" +const MappedNDRange = NDRange{1, <:Any, <:Any, <:Any, <:Any, <:IndexMap} + +Adapt.adapt_structure(to, range::NDRange{N, B, W}) where {N, B, W} = + NDRange{N, B, W}(Adapt.adapt(to, range.blocks), Adapt.adapt(to, range.workitems), Adapt.adapt(to, range.mapping)) + import Base.iterate @inline iterate(range::NDRange) = iterate(blocks(range)) @inline iterate(range::NDRange, state) = iterate(blocks(range), state) @@ -246,6 +288,21 @@ Base.@propagate_inbounds function expand(ndrange::NDRange{N}, groupidx::Integer, return expand(ndrange, blocks(ndrange)[groupidx], idx) end +""" + linear_index(ndrange::MappedNDRange, groupidx, idx) + +Position in the index map of work item `idx` of workgroup `groupidx`. +""" +@inline linear_index(ndrange::MappedNDRange, groupidx::Integer, idx::Integer) = (groupidx - 1) * length(workitems(ndrange)) + idx +@inline linear_index(ndrange::MappedNDRange, groupidx::CartesianIndex{1}, idx::CartesianIndex{1}) = linear_index(ndrange, groupidx.I[1], idx.I[1]) +@inline linear_index(ndrange::MappedNDRange, groupidx::CartesianIndex{1}, idx::Integer) = linear_index(ndrange, groupidx.I[1], idx) +@inline linear_index(ndrange::MappedNDRange, groupidx::Integer, idx::CartesianIndex{1}) = linear_index(ndrange, groupidx, idx.I[1]) + +Base.@propagate_inbounds expand(ndrange::MappedNDRange, groupidx::Integer, idx::Integer) = ndrange.mapping[linear_index(ndrange, groupidx, idx)] +Base.@propagate_inbounds expand(ndrange::MappedNDRange, groupidx::CartesianIndex{1}, idx::CartesianIndex{1}) = ndrange.mapping[linear_index(ndrange, groupidx, idx)] +Base.@propagate_inbounds expand(ndrange::MappedNDRange, groupidx::CartesianIndex{1}, idx::Integer) = ndrange.mapping[linear_index(ndrange, groupidx, idx)] +Base.@propagate_inbounds expand(ndrange::MappedNDRange, groupidx::Integer, idx::CartesianIndex{1}) = ndrange.mapping[linear_index(ndrange, groupidx, idx)] + """ partition(ndrange, workgroupsize) diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index 344e44791..a8ebcaebf 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -319,15 +319,6 @@ end @device_override KI.get_sub_group_local_id() = get_sub_group_local_id() % UInt32 -@device_override @inline function KA.__validindex(ctx) - if KA.__dynamic_checkbounds(ctx) - I = @inbounds KA.expand(KA.__iterspace(ctx), get_group_id(1), get_local_id(1)) - return I in KA.__ndrange(ctx) - else - return true - end -end - ## Shared and Scratch Memory diff --git a/test/indexmap.jl b/test/indexmap.jl new file mode 100644 index 000000000..6d05b284e --- /dev/null +++ b/test/indexmap.jl @@ -0,0 +1,82 @@ +using KernelAbstractions +using KernelAbstractions.NDIteration +using Test + +@kernel function indexmap_mark!(A, count) + I = @index(Global, Cartesian) + p = @index(Global, Linear) + @inbounds A[I] += 1 + @inbounds count[p] = p +end + +@kernel function indexmap_mark_ntuple!(A) + i, j, k = @index(Global, NTuple) + @inbounds A[i, j, k] += 1 +end + +@kernel function indexmap_positions!(count) + p = @index(Global, Linear) + g = @index(Group, Linear) + l = @index(Local, Linear) + nd = @ndrange() + @inbounds count[p] = p + 1000 * g + 100000 * l + 10^7 * nd[1] +end + +function indexmap_testsuite(Backend, AT) + backend = Backend() + dims = (4, 5, 6) + indices = [CartesianIndex(i, j, k) for i in 1:4, j in 1:5, k in 1:6 if (i + j + k) % 3 == 0] + n = length(indices) + ref = zeros(Int, dims) + for I in indices + ref[I] = 1 + end + + @testset "$(eltype(map))" for map in (indices, Tuple.(indices), [Int32.(Tuple(I)) for I in indices]) + A = AT(zeros(Int, dims)) + count = AT(zeros(Int, n)) + indexmap_mark!(backend, 4)(A, count; ndrange = AT(map)) + synchronize(backend) + @test Array(A) == ref + @test Array(count) == 1:n + + A = AT(zeros(Int, dims)) + indexmap_mark_ntuple!(backend)(A; ndrange = AT(map), workgroupsize = 8) + synchronize(backend) + @test Array(A) == ref + end + + @testset "group and local indices" begin + count = AT(zeros(Int, n)) + indexmap_positions!(backend, 4)(count; ndrange = AT(indices)) + synchronize(backend) + @test Array(count) == [p + 1000 * ((p - 1) รท 4 + 1) + 100000 * ((p - 1) % 4 + 1) + 10^7 * n for p in 1:n] + end + + @testset "exact multiple of the workgroup size" begin + A = AT(zeros(Int, dims)) + count = AT(zeros(Int, 8)) + indexmap_mark!(backend, 4)(A, count; ndrange = AT(indices[1:8])) + synchronize(backend) + @test sum(Array(A)) == 8 + @test Array(count) == 1:8 + end + + @testset "empty map" begin + A = AT(zeros(Int, dims)) + indexmap_mark!(backend, 4)(A, AT(Int[]); ndrange = AT(CartesianIndex{3}[])) + synchronize(backend) + @test all(iszero, Array(A)) + end + + @testset "errors" begin + A = AT(zeros(Int, dims)) + count = AT(zeros(Int, n)) + map = AT(indices) + @test_throws ErrorException indexmap_mark!(backend)(A, count; ndrange = map) + @test_throws ErrorException indexmap_mark!(backend, (2, 2))(A, count; ndrange = map) + @test_throws ErrorException KernelAbstractions.partition(indexmap_mark!(backend, 4, (n,)), map, nothing) + @test_throws ArgumentError indexmap_mark!(backend, 4)(A, count; ndrange = AT([1, 2, 3])) + end + return +end diff --git a/test/nditeration.jl b/test/nditeration.jl index 5b0a762f3..7e4f14679 100644 --- a/test/nditeration.jl +++ b/test/nditeration.jl @@ -47,6 +47,30 @@ function nditeration_testsuite() end end + @testset "index map" begin + indices = [CartesianIndex(i, j) for i in 1:3 for j in 1:5] + m = IndexMap(indices) + @test m isa IndexMap{2} + @test length(m) == 15 + @test m[7] == indices[7] + @test IndexMap(Tuple.(indices))[7] == indices[7] + @test IndexMap([Int32.(Tuple(I)) for I in indices])[7] == indices[7] + @test_throws ArgumentError IndexMap([1, 2, 3]) + + let ndrange = NDRange{1, DynamicSize, StaticSize{(4,)}}(CartesianIndices((4,)), nothing, m) + @test ndrange isa MappedNDRange + @test length(ndrange) == 4 + @test linear_index(ndrange, 2, 3) == 7 + @test linear_index(ndrange, CartesianIndex(2), CartesianIndex(3)) == 7 + @test expand(ndrange, 2, 3) == indices[7] + @test expand(ndrange, CartesianIndex(2), CartesianIndex(3)) == indices[7] + @test expand(ndrange, 4, 3) == indices[15] + end + let ndrange = NDRange{1, DynamicSize, DynamicSize}(CartesianIndices((4,)), CartesianIndices((4,)), m) + @test expand(ndrange, 2, 3) == indices[7] + end + end + # GPU scenario where we get a linear index into workitems/blocks function linear_iteration(ndrange) idx = Array{CartesianIndex{2}}(undef, length(blocks(ndrange)) * length(workitems(ndrange))) diff --git a/test/test.jl b/test/test.jl index d202b51a8..6d079a6d5 100644 --- a/test/test.jl +++ b/test/test.jl @@ -89,6 +89,38 @@ function unittest_testsuite(Backend, backend_str, backend_mod, BackendArrayT; sk @test length(blocks(iterspace)) == 2 @test iterspace.mapping === nothing end + let kernel = KernelAbstractions.Kernel{typeof(backend), StaticSize{(4,)}, DynamicSize, typeof(identity)}(backend, identity) + map = [CartesianIndex(i, j) for i in 1:3 for j in 1:5] + iterspace, dynamic = KernelAbstractions.partition(kernel, map, nothing) + @test iterspace isa MappedNDRange + @test length(blocks(iterspace)) == 4 + @test dynamic isa DynamicCheck + @test ndims(iterspace) == 1 + + iterspace, dynamic = KernelAbstractions.partition(kernel, map[1:8], (4,)) + @test length(blocks(iterspace)) == 2 + @test dynamic isa NoDynamicCheck + + iterspace, dynamic = KernelAbstractions.partition(kernel, CartesianIndex{2}[], nothing) + @test length(blocks(iterspace)) == 0 + + @test_throws ErrorException KernelAbstractions.partition(kernel, map, (8,)) + @test_throws ArgumentError KernelAbstractions.partition(kernel, [1, 2, 3], nothing) + end + let kernel = KernelAbstractions.Kernel{typeof(backend), DynamicSize, DynamicSize, typeof(identity)}(backend, identity) + map = [CartesianIndex(i, j) for i in 1:3 for j in 1:5] + iterspace, dynamic = KernelAbstractions.partition(kernel, map, (4,)) + @test length(blocks(iterspace)) == 4 + @test length(workitems(iterspace)) == 4 + + @test_throws ErrorException KernelAbstractions.partition(kernel, map, nothing) + @test_throws ErrorException KernelAbstractions.partition(kernel, map, map) + @test_throws ErrorException KernelAbstractions.partition(kernel, map, (2, 2)) + end + let kernel = KernelAbstractions.Kernel{typeof(backend), StaticSize{(4,)}, StaticSize{(15,)}, typeof(identity)}(backend, identity) + map = [CartesianIndex(i, j) for i in 1:3 for j in 1:5] + @test_throws ErrorException KernelAbstractions.partition(kernel, map, nothing) + end end @kernel function index_linear_global(A) diff --git a/test/testsuite.jl b/test/testsuite.jl index 9c9db2c8e..273fa73e1 100644 --- a/test/testsuite.jl +++ b/test/testsuite.jl @@ -34,6 +34,7 @@ include("private.jl") include("unroll.jl") include("nditeration.jl") include("offsets.jl") +include("indexmap.jl") include("copyto.jl") include("devices.jl") include("print_test.jl") @@ -80,6 +81,10 @@ function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{ offsets_testsuite(backend, AT) end + @conditional_testset "IndexMap" skip_tests begin + indexmap_testsuite(backend, AT) + end + @conditional_testset "copyto!" skip_tests begin copyto_testsuite(backend, AT) end