diff --git a/docs/src/api.md b/docs/src/api.md index e7c6d6023..646c975b9 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -96,5 +96,11 @@ KernelAbstractions.NDIteration.StaticOffset KernelAbstractions.NDIteration.DynamicOffset KernelAbstractions.NDIteration.extents KernelAbstractions.NDIteration.offsets +KernelAbstractions.NDIteration.expand KernelAbstractions.NDIteration.linear_index +KernelAbstractions.NDIteration.IndexMap +KernelAbstractions.NDIteration.MappedNDRange +KernelAbstractions.NDIteration.MappedIndices +KernelAbstractions.NDIteration.invalid_index +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..cc9c5f56f 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 get an index of `typemin(Int)` along every axis. + 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 c022c7cef..37eceb973 100644 --- a/src/KernelAbstractions.jl +++ b/src/KernelAbstractions.jl @@ -398,9 +398,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) @@ -533,6 +531,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 = """ @@ -588,6 +590,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 @@ -603,7 +638,39 @@ 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 + +# Backends may override `__validindex(ctx)` with this same check, so `expand` yields an +# index `in` `__ndrange(ctx)` exactly for the valid work items of every kind of `ndrange`. +@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..e1730d78f 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -16,12 +16,18 @@ struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace} end end -# `CartesianIndices` covering a launch `ndrange` (any form accepted by `partition`). +# `CartesianIndices` or `MappedIndices` covering a launch `ndrange` (any form accepted by `partition`). cartesian(::Nothing) = nothing cartesian(ci::CartesianIndices) = ci +cartesian(r::MappedIndices) = r cartesian(n::Integer) = CartesianIndices((Int(n),)) cartesian(r::AbstractUnitRange) = CartesianIndices((r,)) cartesian(t::Tuple) = CartesianIndices(t) +cartesian(v::AbstractVector) = MappedIndices(v) +cartesian(m::IndexMap) = MappedIndices(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 1858f3857..1acf07a79 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, MappedIndices, invalid_index 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,84 @@ 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)) + +""" + invalid_index(Val(N)) + +`CartesianIndex{N}` returned by [`expand`](@ref) for a work item past the end of an +[`IndexMap`](@ref): `typemin(Int)` along every axis. +""" +@inline invalid_index(::Val{N}) where {N} = CartesianIndex(ntuple(_ -> typemin(Int), Val(N))) + +# Entry `p` of the map, or `invalid_index` for a position past its end. +@inline function mapped_index(m::IndexMap{N}, p::Integer) where {N} + return p <= length(m) ? (@inbounds m[p]) : invalid_index(Val(N)) +end + +""" + MappedIndices{N} + +The `ndrange` of a launch over an [`IndexMap`](@ref) with `CartesianIndex{N}` entries. +`size` and `length` give the number of listed indices, and a `CartesianIndex{N}` is `in` +it unless it is the [`invalid_index`](@ref), so that backends can check the validity of a +work item as `expand(iterspace, groupidx, idx) in ndrange` for every kind of `ndrange`. +""" +struct MappedIndices{N} + length::Int +end +MappedIndices(m::IndexMap{N}) where {N} = MappedIndices{N}(length(m)) +MappedIndices(v::AbstractVector) = MappedIndices{mapdims(eltype(v))}(length(v)) +Base.length(r::MappedIndices) = r.length +Base.size(r::MappedIndices) = (r.length,) +@inline Base.in(I::CartesianIndex{N}, r::MappedIndices{N}) where {N} = I != invalid_index(Val(N)) + """ 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) @@ -137,8 +198,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 ``` @@ -182,12 +244,32 @@ 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). [`expand`](@ref) +looks up the index listed for a work item, and the `ndrange` of such a launch is a +[`MappedIndices`](@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) Base.length(range::NDRange) = length(blocks(range)) +""" + expand(ndrange::NDRange, groupidx, idx) + +Index of the `ndrange` handled by work item `idx` of workgroup `groupidx`, both given as +`CartesianIndex` or linear positions in the blocked iteration space. For a +[`MappedNDRange`](@ref) this is the entry of the index map for that work item, or the +[`invalid_index`](@ref) for a work item past its end. +""" @inline function expand(ndrange::NDRange{N}, groupidx::CartesianIndex{N}, idx::CartesianIndex{N}) where {N} offset = offsets(ndrange) nI = ntuple(Val(N)) do I @@ -251,6 +333,22 @@ 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]) + +# The listed index of a work item, or `invalid_index` for a work item past the end of the map. +@inline expand(ndrange::MappedNDRange, groupidx::Integer, idx::Integer) = mapped_index(ndrange.mapping, linear_index(ndrange, groupidx, idx)) +@inline expand(ndrange::MappedNDRange, groupidx::CartesianIndex{1}, idx::CartesianIndex{1}) = mapped_index(ndrange.mapping, linear_index(ndrange, groupidx, idx)) +@inline expand(ndrange::MappedNDRange, groupidx::CartesianIndex{1}, idx::Integer) = mapped_index(ndrange.mapping, linear_index(ndrange, groupidx, idx)) +@inline expand(ndrange::MappedNDRange, groupidx::Integer, idx::CartesianIndex{1}) = mapped_index(ndrange.mapping, linear_index(ndrange, groupidx, idx)) + """ partition(ndrange, workgroupsize) diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index fbd5883ef..7bf532a8b 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..c63d746bc 100644 --- a/test/nditeration.jl +++ b/test/nditeration.jl @@ -47,6 +47,35 @@ 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] + @test expand(ndrange, 4, 4) == invalid_index(Val(2)) + @test expand(ndrange, 4, 3) in MappedIndices(m) + @test !(expand(ndrange, 4, 4) in MappedIndices(m)) + @test size(MappedIndices(m)) == (15,) + @test length(MappedIndices(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 c280bf592..7aa23b7b6 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) @@ -267,44 +299,38 @@ function unittest_testsuite(Backend, backend_str, backend_mod, BackendArrayT; sk @conditional_testset "Const" skip_tests begin let kernel = constarg(Backend(), 8, (1024,)) - # this is poking at internals - iterspace = NDRange{1, StaticSize{(128,)}, StaticSize{(8,)}}() - ctx = if Backend == CPU - KernelAbstractions.mkcontext(kernel, 1, nothing, iterspace, Val(NoDynamicCheck())) - else - KernelAbstractions.mkcontext(kernel, nothing, iterspace) - end - AT = if Backend == CPU - Array{Float32, 2} + A = KernelAbstractions.zeros(Backend(), Float32, 1024) + B = adapt(Backend(), rand(Float32, 1024)) + kernel(A, B) + synchronize(Backend()) + @test Array(A) == Array(B) + + if backend_str == "CPU" + # the CPU backend compiles kernels with POCL, whose device arrays have no + # `@Const`-specific lowering to look for in the IR + @test_skip false else - BackendArrayT{Float32, 2, 1} # AS 1 - end - IR = sprint() do io - if backend_str == "CPU" - code_llvm( - io, kernel.f, (typeof(ctx), AT, AT), - optimize = false, raw = true, - ) - else + # this is poking at internals + iterspace = NDRange{1, StaticSize{(128,)}, StaticSize{(8,)}}() + ctx = KernelAbstractions.mkcontext(kernel, nothing, iterspace) + AT = BackendArrayT{Float32, 2, 1} # AS 1 + IR = sprint() do io backend_mod.code_llvm( io, kernel.f, (typeof(ctx), AT, AT), kernel = true, optimize = true, ) end - end - if backend_str == "CPU" - @test occursin("!alias.scope", IR) - @test occursin("!noalias", IR) - elseif backend_str == "CUDA" - if Base.libllvm_version >= v"20" - @test occursin("addrspace(1)", IR) + if backend_str == "CUDA" + if Base.libllvm_version >= v"20" + @test occursin("addrspace(1)", IR) + else + @test occursin("@llvm.nvvm.ldg", IR) + end + elseif backend_str == "ROCM" + @test occursin("addrspace(4)", IR) else - @test occursin("@llvm.nvvm.ldg", IR) + @test_skip false end - elseif backend_str == "ROCM" - @test occursin("addrspace(4)", IR) - else - @test_skip false end end end 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