From e883ad830948b4df055374815aec0c4bdaff2b6d Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 09:33:21 +0200 Subject: [PATCH 1/2] Launch @kernel kernels on N-d grids, indexing in 32 bits `@index` assumes that a kernel was launched on a 1-D grid, so every work-item decomposes its linear group and local id into Cartesian positions, with 64-bit integer divisions for a dynamically-sized ndrange. Backends can now record how they launched a kernel in its context, and `@index` specializes on that: - `NDLaunch{T}`: the grid has the shape of the iteration space (for as many dimensions as the backend's grid has), so the hardware ids are the Cartesian positions and no divisions are needed; - `LinearLaunch{T}`: the 1-D grid used so far. Either way `@index` computes in `T`, `Int32` whenever the padded iteration space fits, and still returns `Int`s. `select_launch` chooses the launch from the iteration space and the backend's limits. When the workgroup size will be tuned, the choice holds for every workgroup size tuning can pick, so the context type (and thus the compiled kernel) doesn't change. Contexts without a launch keep the existing behavior, so backends opt in; `__validindex` is now generic, dispatching on the launch. PoCL launches on N-d grids. --- docs/src/api.md | 4 + docs/src/implementations.md | 70 +++++++++ src/KernelAbstractions.jl | 50 +++--- src/compiler.jl | 19 ++- src/launch.jl | 293 ++++++++++++++++++++++++++++++++++++ src/pocl/backend.jl | 52 +++---- test/codegen_checks.jl | 29 ++++ test/launch.jl | 195 ++++++++++++++++++++++++ test/runtests.jl | 73 +++++++++ test/testsuite.jl | 5 + 10 files changed, 738 insertions(+), 52 deletions(-) create mode 100644 src/launch.jl create mode 100644 test/launch.jl diff --git a/docs/src/api.md b/docs/src/api.md index bb76872fb..1be06929d 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -108,4 +108,8 @@ KernelAbstractions.NDIteration.DynamicOffset KernelAbstractions.NDIteration.extents KernelAbstractions.NDIteration.offsets KernelAbstractions.NDIteration.linear_index +KernelAbstractions.LinearLaunch +KernelAbstractions.NDLaunch +KernelAbstractions.select_launch +KernelAbstractions.launch_workgroupsize ``` diff --git a/docs/src/implementations.md b/docs/src/implementations.md index adc1b132a..91fd68841 100644 --- a/docs/src/implementations.md +++ b/docs/src/implementations.md @@ -75,3 +75,73 @@ to the backend's array type, so that `adapt(backend, x)` and Adapt.adapt_storage(::CUDABackend, x) = adapt(CuArray, x) ``` + +## Launching `@kernel` kernels + +A kernel written with [`@kernel`](@ref) receives a hidden context, a +`KernelAbstractions.CompilerMetadata` built by the backend's `mkcontext`, from which +[`@index`](@ref) computes its indices. By default (a context without a `launch`), +`@index` assumes that the kernel was launched on a 1-D grid of +`length(blocks(iterspace))` groups of `length(workitems(iterspace))` work-items. It then +decomposes the linear hardware ids into Cartesian positions, which takes integer divisions +when the `ndrange` is not known at compile time, and computes in `Int`. + +A backend **may** launch kernels differently, and pass the `launch` keyword to the +`CompilerMetadata` constructor to tell `@index` how: + +- [`NDLaunch{T}`](@ref KernelAbstractions.NDLaunch): the grid has the shape of the + iteration space (for as many dimensions as the backend's grid has, i.e. up to 3), so + `@index` doesn't need any divisions; +- [`LinearLaunch{T}`](@ref KernelAbstractions.LinearLaunch): the default 1-D grid. + +Either way `@index` computes in `T`, e.g. `Int32`, which is faster on GPUs. The backend +has to implement the typed [`KI.get_group_id`](@ref KernelInterface.get_group_id) and +[`KI.get_local_id`](@ref KernelInterface.get_local_id) queries such that they compute in +`T` too, e.g. without checked conversions. + +[`select_launch`](@ref KernelAbstractions.select_launch) chooses the launch from the +iteration space, whether the workgroup size will be tuned, and the limits of the backend +([`KI.max_work_group_size`](@ref KernelInterface.max_work_group_size), +[`KI.max_work_group_dims`](@ref KernelInterface.max_work_group_dims) and +[`KI.max_num_groups`](@ref KernelInterface.max_num_groups)). It doesn't depend on the +workgroup size a backend tunes afterwards, which keeps the context type (and thus the +compiled kernel) the same before and after tuning, as long as the backend tunes with +[`launch_workgroupsize`](@ref KernelAbstractions.launch_workgroupsize). A launch then +looks like this: + +```julia +function (obj::KA.Kernel{MyBackend})(args...; ndrange = nothing, workgroupsize = nothing) + ndrange, workgroupsize, iterspace, dynamic = KA.launch_config(obj, ndrange, workgroupsize) + launch = KA.select_launch(obj, workgroupsize, iterspace) + ctx = KA.CompilerMetadata{KA.ndrange(obj), KA.DynamicCheck}(ndrange, iterspace; launch) + kernel = compile(obj.f, ctx, args...) + + if KA.workgroupsize(obj) <: KA.DynamicSize && workgroupsize === nothing + threads = max_threads(kernel) # at most `KI.max_work_group_size(backend)` + workgroupsize = KA.launch_workgroupsize(backend, launch, threads, ndrange) + iterspace, dynamic = KA.partition(obj, ndrange, workgroupsize) + ctx = KA.CompilerMetadata{KA.ndrange(obj), KA.DynamicCheck}(ndrange, iterspace; launch) + end + + groups, items = size(KA.blocks(iterspace)), size(KA.workitems(iterspace)) + prod(groups) == 0 && return + if launch isa KA.NDLaunch + run(kernel, ctx, args...; groups, items) # padded to 3 dimensions + else + run(kernel, ctx, args...; groups = prod(groups), items = prod(items)) + end +end +``` + +The POCL backend is an example. Backends that launch on an N-d grid **must not** override +`__validindex` or the `__index_*` functions, which dispatch on the launch. + +Packages that customize the iteration space (with a custom `partition` and `expand`) +don't need to do anything for these launches: the index functions only compute the global +index directly for the iteration spaces KernelAbstractions creates itself, and call +`expand`, `in` and `linear_index` otherwise. + +Packages with an `Adapt` rule for `CompilerMetadata` **must** preserve its `launch`, e.g. by +passing `launch = KernelAbstractions.__launch(ctx)` to the constructor. Otherwise the +kernel computes its indices as if it had been launched on a 1-D grid, which gives wrong +results for a kernel launched with an `NDLaunch`. diff --git a/src/KernelAbstractions.jl b/src/KernelAbstractions.jl index 0dacde9a3..6e26bb1fd 100644 --- a/src/KernelAbstractions.jl +++ b/src/KernelAbstractions.jl @@ -389,28 +389,30 @@ end # Internal kernel functions ### -@inline function __index_Local_Linear(ctx) - return KI.get_local_id().x -end - -@inline function __index_Group_Linear(ctx) - return KI.get_group_id().x -end - -@inline function __index_Global_Linear(ctx) +# The index functions dispatch on the launch configuration of the context (see +# `launch.jl`); `nothing` is a 1-D launch, indexed in `Int`. + +@inline __index_Local_Linear(ctx) = local_linear(ctx, __launch(ctx)) +@inline __index_Group_Linear(ctx) = group_linear(ctx, __launch(ctx)) +@inline __index_Global_Linear(ctx) = global_linear(ctx, __launch(ctx)) +@inline __index_Local_Cartesian(ctx) = local_cartesian(ctx, __launch(ctx)) +@inline __index_Group_Cartesian(ctx) = group_cartesian(ctx, __launch(ctx)) +@inline __index_Global_Cartesian(ctx) = global_cartesian(ctx, __launch(ctx)) + +@inline local_linear(ctx, ::Nothing) = KI.get_local_id().x +@inline group_linear(ctx, ::Nothing) = KI.get_group_id().x +@inline function global_linear(ctx, ::Nothing) 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) end - -@inline function __index_Local_Cartesian(ctx) - return @inbounds workitems(__iterspace(ctx))[KI.get_local_id().x] -end -@inline function __index_Group_Cartesian(ctx) - return @inbounds blocks(__iterspace(ctx))[KI.get_group_id().x] -end -@inline function __index_Global_Cartesian(ctx) - return @inbounds expand(__iterspace(ctx), KI.get_group_id().x, KI.get_local_id().x) +@inline local_cartesian(ctx, ::Nothing) = @inbounds workitems(__iterspace(ctx))[KI.get_local_id().x] +@inline group_cartesian(ctx, ::Nothing) = @inbounds blocks(__iterspace(ctx))[KI.get_group_id().x] +@inline global_cartesian(ctx, ::Nothing) = + @inbounds expand(__iterspace(ctx), KI.get_group_id().x, KI.get_local_id().x) +@inline function validindex(ctx, ::Nothing) + I = @inbounds expand(__iterspace(ctx), KI.get_group_id().x, KI.get_local_id().x) + return I in __ndrange(ctx) end @inline __index_Local_NTuple(ctx, I...) = Tuple(__index_Local_Cartesian(ctx, I...)) @@ -606,13 +608,23 @@ end ### include("compiler.jl") +include("launch.jl") ### # Compiler/Frontend ### function __workitems_iterspace end -function __validindex end + +# Whether the current work-item is part of the ndrange, or a padding lane of a partial +# workgroup. Padding lanes still take part in `@synchronize`. +@inline function __validindex(ctx) + if __dynamic_checkbounds(ctx) + return validindex(ctx, __launch(ctx)) + else + return true + end +end # for reflection function mkcontext end diff --git a/src/compiler.jl b/src/compiler.jl index 687442a2d..c61f7d471 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -1,18 +1,26 @@ -struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace} +""" + CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch} + +The hidden context argument of kernels written with [`@kernel`](@ref). The `launch` field +tells the index functions how the backend launched the kernel: `nothing` for a 1-D launch +indexed in `Int`, or a [`LinearLaunch`](@ref) or [`NDLaunch`](@ref). +""" +struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch} groupindex::I ndrange::NDRange iterspace::Iterspace + launch::Launch # CPU variant function CompilerMetadata{NDRange, CB}(idx, ndrange, iterspace) where {NDRange, CB} ndrange = cartesian(ndrange) - return new{NDRange, CB, typeof(idx), typeof(ndrange), typeof(iterspace)}(idx, ndrange, iterspace) + return new{NDRange, CB, typeof(idx), typeof(ndrange), typeof(iterspace), Nothing}(idx, ndrange, iterspace, nothing) end # GPU variante: index is given implicit - function CompilerMetadata{NDRange, CB}(ndrange, iterspace) where {NDRange, CB} + function CompilerMetadata{NDRange, CB}(ndrange, iterspace; launch = nothing) where {NDRange, CB} ndrange = cartesian(ndrange) - return new{NDRange, CB, Nothing, typeof(ndrange), typeof(iterspace)}(nothing, ndrange, iterspace) + return new{NDRange, CB, Nothing, typeof(ndrange), typeof(iterspace), typeof(launch)}(nothing, ndrange, iterspace, launch) end end @@ -25,6 +33,7 @@ cartesian(t::Tuple) = CartesianIndices(t) @inline __iterspace(cm::CompilerMetadata) = cm.iterspace @inline __groupindex(cm::CompilerMetadata) = cm.groupindex +@inline __launch(cm::CompilerMetadata) = cm.launch @inline __groupsize(cm::CompilerMetadata) = size(workitems(__iterspace(cm))) @inline __dynamic_checkbounds(::CompilerMetadata{NDRange, CB}) where {NDRange, CB} = CB <: DynamicCheck @inline __ndrange(::CompilerMetadata{NDRange}) where {NDRange <: StaticSize} = CartesianIndices(get(NDRange)) @@ -39,7 +48,7 @@ cartesian(t::Tuple) = CartesianIndices(t) function Adapt.adapt_structure(to, cm::CompilerMetadata{NDRange, CB, I}) where {NDRange, CB, I} iterspace = Adapt.adapt(to, cm.iterspace) if I === Nothing - return CompilerMetadata{NDRange, CB}(cm.ndrange, iterspace) + return CompilerMetadata{NDRange, CB}(cm.ndrange, iterspace; cm.launch) else return CompilerMetadata{NDRange, CB}(cm.groupindex, cm.ndrange, iterspace) end diff --git a/src/launch.jl b/src/launch.jl new file mode 100644 index 000000000..37667a75b --- /dev/null +++ b/src/launch.jl @@ -0,0 +1,293 @@ +### +# Launch configurations +# +# How a backend maps the hardware groups and work-items of a launch onto the blocked +# iteration space. The launch is stored in the kernel's `CompilerMetadata`, so the index +# functions below can specialize on it. +### + +""" + LinearLaunch{T}() + +Launch configuration for a kernel launched on a 1-D grid: the x-components of the hardware +group and local ids are the column-major positions in `blocks(iterspace)` and +`workitems(iterspace)`, like with the default launch (`nothing`). + +`@index` computes in `T` (`Int32` or `Int`), which has to hold the number of work-items +in the padded iteration space. It still returns `Int`s. +""" +struct LinearLaunch{T <: Integer} end + +""" + NDLaunch{T}() + +Launch configuration for a kernel launched on an N-d grid, where `N = ndims(iterspace)` is +at most the number of grid dimensions of the backend (the length of +[`KI.max_work_group_dims`](@ref KernelInterface.max_work_group_dims), i.e. 3): the grid +consists of `size(blocks(iterspace))` groups of `size(workitems(iterspace))` work-items, so +the x, y and z-components of the hardware group and local ids are the positions along the +first `N` dimensions of the iteration space. This avoids decomposing linear ids into +Cartesian positions, i.e., divisions. + +The linear group and local indices are computed x-fastest, so they agree with the ordering +of a [`LinearLaunch`](@ref) (and with how GPUs typically form sub-groups). + +`@index` computes in `T` (`Int32` or `Int`), which has to hold the number of work-items +in the padded iteration space. It still returns `Int`s. +""" +struct NDLaunch{T <: Integer} end + +const Launch = Union{LinearLaunch, NDLaunch} + +index_type(::LinearLaunch{T}) where {T} = T +index_type(::NDLaunch{T}) where {T} = T + + +## host side + +""" + select_launch(kernel::Kernel, workgroupsize, iterspace)::Union{LinearLaunch, NDLaunch} + +Choose how to launch `kernel` with the (possibly preliminary) iteration space `iterspace` +and the given `workgroupsize` (`nothing` if it will be tuned), as returned by +`launch_config`: an [`NDLaunch`](@ref) if the iteration space fits the grid of +`backend(kernel)`, i.e. its number of dimensions and per-dimension limits, and a +[`LinearLaunch`](@ref) otherwise, both computing indices in `Int32` if possible. + +If the workgroup size will be tuned, the choice holds for any workgroup size the backend +tunes afterwards, so the context and thus the compiled kernel are the same before and +after tuning. That requires tuning with [`launch_workgroupsize`](@ref) and with at most +`KI.max_work_group_size(backend)` work-items, and assumes that `iterspace` covers the +padded iteration space with a single workgroup, as `launch_config` does. + +Throws an `ArgumentError` if the iteration space has more than `typemax(Int)` work-items. + +The result is not inferred concretely: launching with it takes a dynamic dispatch, unless +the caller checks for the common case of an `NDLaunch{Int32}` first. +""" +function select_launch(kernel::Kernel, workgroupsize, iterspace) + b = backend(kernel) + groups = size(blocks(iterspace)) + items = size(workitems(iterspace)) + tuned = KernelAbstractions.workgroupsize(kernel) <: DynamicSize && workgroupsize === nothing + return select_launch( + map(mul_extent, groups, items), tuned ? nothing : items, + KI.max_work_group_size(b), KI.max_work_group_dims(b), KI.max_num_groups(b) + ) +end + +# `groupsize === nothing` means that the workgroup size will be tuned +function select_launch( + extent::Dims{N}, groupsize, max_items::Int, + max_dims::Dims, max_groups::Dims + ) where {N} + nd = N <= length(max_dims) + if groupsize === nothing + if nd + # a tuned workgroup has at least one work-item per dimension + nd = all(ntuple(d -> extent[d] <= max_groups[d], Val(N))) + end + padded = tuned_padded(extent, max_items, nd ? max_dims : ()) + else + blocks, items, _ = NDIteration.partition(extent, groupsize) + if nd + nd = all(ntuple(d -> items[d] <= max_dims[d] && blocks[d] <= max_groups[d], Val(N))) && + prod(items) <= max_items + end + padded = map(mul_extent, blocks, items) + end + if fits(Int32, padded) + return nd ? NDLaunch{Int32}() : LinearLaunch{Int32}() + elseif fits(Int, padded) + return nd ? NDLaunch{Int}() : LinearLaunch{Int}() + else + throw_too_large() + end +end + +# arithmetic on the extents of an iteration space, which has to fit an `Int` +@noinline throw_too_large() = + throw(ArgumentError("Iteration space has more than typemax(Int) work-items")) +function mul_extent(a::Int, b::Int) + n, overflow = Base.mul_with_overflow(a, b) + overflow && throw_too_large() + return n +end +function add_extent(a::Int, b::Int) + n, overflow = Base.add_with_overflow(a, b) + overflow && throw_too_large() + return n +end + +# Upper bound on the padded extents of an iteration space whose workgroup size is chosen +# by `launch_workgroupsize`, i.e., by `KI.threads_to_workgroupsize` with at most `capacity` +# work-items and the per-dimension `limits`. That fills the first dimensions first, so a +# dimension only gets more than one work-item if the previous ones don't use all of them, +# and it pads each dimension by less than its number of work-items. +tuned_padded(::Tuple{}, capacity, limits) = () +function tuned_padded(extent::Tuple, capacity, limits) + limit = isempty(limits) ? typemax(Int) : first(limits) + items = max(1, min(first(extent), capacity, limit)) + padded = iszero(first(extent)) ? 0 : add_extent(first(extent), items - 1) + rest = isempty(limits) ? () : Base.tail(limits) + return (padded, tuned_padded(Base.tail(extent), max(1, capacity รท items), rest)...) +end + +# whether the number of positions in an iteration space with extents `dims` fits `T` +fits(::Type{T}, dims::Dims) where {T} = _fits(T, 1, dims) +_fits(::Type, n, ::Tuple{}) = true +function _fits(::Type{T}, n, dims::Dims) where {T} + n, overflow = Base.mul_with_overflow(n, first(dims)) + return if overflow || n > typemax(T) + any(iszero, dims) + else + _fits(T, n, Base.tail(dims)) + end +end + +""" + launch_workgroupsize(backend, launch, threads, ndrange) + +The workgroup size for launching `threads` work-items per workgroup over `ndrange` with +`launch` (see [`select_launch`](@ref)). Backends that tune the workgroup size of a kernel +launched with a [`LinearLaunch`](@ref) or [`NDLaunch`](@ref) have to use this. +""" +launch_workgroupsize(backend, ::NDLaunch, threads, ndrange) = + KI.threads_to_workgroupsize(threads, extents(ndrange), KI.max_work_group_dims(backend)) +launch_workgroupsize(backend, ::Union{LinearLaunch, Nothing}, threads, ndrange) = + KI.threads_to_workgroupsize(threads, extents(ndrange)) + + +## device side + +# The sizes of the iteration space are stored as `Int`s, but fit the index type `T` +# of the launch, so truncating them is safe. (Written without closures capturing `T`, +# which Julia 1.10 fails to infer.) +@inline narrow(::Type{T}, ::Tuple{}) where {T} = () +@inline narrow(::Type{T}, dims::Tuple) where {T} = (first(dims) % T, narrow(T, Base.tail(dims))...) + +# Widen a (positive) index of type `T` to the `Int` returned by `@index`. +@inline widen_index(i::Int) = i +@inline widen_index(i::T) where {T <: Union{Int8, Int16, Int32}} = (i % unsigned(T)) % Int + +# 0-based linear index into `dims` to 1-based subscripts, using unsigned divisions +@inline ind2sub(::Tuple{}, i) = () +@inline ind2sub(::Tuple{Any}, i) = (i + one(i),) +@inline function ind2sub(dims::Tuple, i) + q = Core.Intrinsics.udiv_int(i, dims[1]) + return (i - q * dims[1] + one(i), ind2sub(Base.tail(dims), q)...) +end + +# 1-based subscripts into `dims` to a 1-based linear index of type `T`, column-major +@inline linearize(::Type{T}, ::Tuple{}, ::Tuple{}) where {T} = one(T) +@inline linearize(::Type{T}, dims::Tuple, I::Tuple) where {T} = + I[1] + dims[1] * (linearize(T, Base.tail(dims), Base.tail(I)) - one(T)) + +@inline hardware_position(id, ::Val{N}) where {N} = ntuple(d -> (id.x, id.y, id.z)[d], Val(N)) + +# The hardware ids are bounded by the launch, which e.g. lets the compiler fold the validity +# check of a statically-sized kernel. +@inline function assume_bounded(pos::Tuple, dims::Tuple) + map(pos, dims) do p, d + NDIteration.assume((p >= one(p)) & (p <= d)) + end + return pos +end +@inline function assume_bounded(i, dims::Tuple) + NDIteration.assume((i >= zero(i)) & (i < prod(dims))) + return i +end + +# 1-based position of the current group in `blocks(iterspace)`, and of the current +# work-item in `workitems(iterspace)`, as tuples of the index type +@inline function group_position(ctx, ::NDLaunch{T}) where {T} + dims = narrow(T, size(blocks(__iterspace(ctx)))) + return assume_bounded(hardware_position(KI.get_group_id(T), Val(ndims(ctx))), dims) +end +@inline function local_position(ctx, ::NDLaunch{T}) where {T} + dims = narrow(T, size(workitems(__iterspace(ctx)))) + return assume_bounded(hardware_position(KI.get_local_id(T), Val(ndims(ctx))), dims) +end +@inline function group_position(ctx, ::LinearLaunch{T}) where {T} + dims = narrow(T, size(blocks(__iterspace(ctx)))) + return ind2sub(dims, assume_bounded(KI.get_group_id(T).x - one(T), dims)) +end +@inline function local_position(ctx, ::LinearLaunch{T}) where {T} + dims = narrow(T, size(workitems(__iterspace(ctx)))) + return ind2sub(dims, assume_bounded(KI.get_local_id(T).x - one(T), dims)) +end + +# 1-based position of the current work-item in the padded iteration space +@inline function blocked_position(ctx, launch::Launch) + T = index_type(launch) + groupsize = narrow(T, size(workitems(__iterspace(ctx)))) + return map( + (g, w, l) -> (g - one(g)) * w + l, + group_position(ctx, launch), groupsize, local_position(ctx, launch) + ) +end + +# The group and work-item indices, i.e. the elements of `blocks(iterspace)` and +# `workitems(iterspace)` at the current positions. These are usually 1-based. +@inline group_index(ctx, launch::Launch) = + CartesianIndex(offset_position(group_position(ctx, launch), blocks(__iterspace(ctx)))) +@inline local_index(ctx, launch::Launch) = + CartesianIndex(offset_position(local_position(ctx, launch), workitems(__iterspace(ctx)))) +@inline offset_position(pos::Tuple, indices::CartesianIndices) = + map((p, f) -> widen_index(p) + f - 1, pos, first(indices).I) + +# For the iteration spaces KernelAbstractions creates itself (1-based blocks and +# work-items, and an identity or offset mapping), the global index is the blocked position +# shifted by the offsets. Other iteration spaces (e.g. with a custom mapping, see #781, +# or a custom `expand` for other contents) have to go through `expand`, `in` and +# `linear_index`. +const OneBasedIndices{N} = CartesianIndices{N, <:NTuple{N, Base.OneTo}} +const BuiltinNDRange = NDRange{ + N, B, W, <:Union{Nothing, OneBasedIndices}, <:Union{Nothing, OneBasedIndices}, + <:Union{Nothing, StaticOffset, DynamicOffset}, +} where {N, B, W} +@inline builtin(iterspace::BuiltinNDRange) = + blocks(iterspace) isa OneBasedIndices && workitems(iterspace) isa OneBasedIndices +@inline builtin(iterspace) = false + +@inline global_cartesian(ctx, launch::Launch, iterspace, ndrange) = +if builtin(iterspace) && ndrange isa CartesianIndices + I = map((i, o) -> widen_index(i) + o, blocked_position(ctx, launch), offsets(iterspace)) + CartesianIndex(I) +else + @inbounds expand(iterspace, group_index(ctx, launch), local_index(ctx, launch)) +end + +@inline global_linear(ctx, launch::Launch, iterspace, ndrange) = +if builtin(iterspace) && ndrange isa CartesianIndices + T = index_type(launch) + widen_index(linearize(T, narrow(T, size(ndrange)), blocked_position(ctx, launch))) +else + linear_index(ndrange, global_cartesian(ctx, launch, iterspace, ndrange)) +end + +@inline validindex(ctx, launch::Launch, iterspace, ndrange) = +if builtin(iterspace) && ndrange isa CartesianIndices + T = index_type(launch) + all(map(<=, blocked_position(ctx, launch), narrow(T, size(ndrange)))) +else + global_cartesian(ctx, launch, iterspace, ndrange) in ndrange +end + +# `@index` entry points, see `__index_*` +@inline local_linear(ctx, ::LinearLaunch{T}) where {T} = widen_index(KI.get_local_id(T).x) +@inline group_linear(ctx, ::LinearLaunch{T}) where {T} = widen_index(KI.get_group_id(T).x) +@inline local_linear(ctx, launch::NDLaunch{T}) where {T} = widen_index( + linearize(T, narrow(T, size(workitems(__iterspace(ctx)))), local_position(ctx, launch)) +) +@inline group_linear(ctx, launch::NDLaunch{T}) where {T} = widen_index( + linearize(T, narrow(T, size(blocks(__iterspace(ctx)))), group_position(ctx, launch)) +) +@inline local_cartesian(ctx, launch::Launch) = local_index(ctx, launch) +@inline group_cartesian(ctx, launch::Launch) = group_index(ctx, launch) +@inline global_cartesian(ctx, launch::Launch) = + global_cartesian(ctx, launch, __iterspace(ctx), __ndrange(ctx)) +@inline global_linear(ctx, launch::Launch) = + global_linear(ctx, launch, __iterspace(ctx), __ndrange(ctx)) +@inline validindex(ctx, launch::Launch) = + validindex(ctx, launch, __iterspace(ctx), __ndrange(ctx)) diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index 88f1d0c64..4bb67491b 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -133,6 +133,9 @@ KI.supports_atomics(::POCLBackend) = true function KA.mkcontext(kernel::KA.Kernel{POCLBackend}, _ndrange, iterspace) return KA.CompilerMetadata{KA.ndrange(kernel), KA.DynamicCheck}(_ndrange, iterspace) end +function KA.mkcontext(kernel::KA.Kernel{POCLBackend}, _ndrange, iterspace, launch) + return KA.CompilerMetadata{KA.ndrange(kernel), KA.DynamicCheck}(_ndrange, iterspace; launch) +end function KA.mkcontext( kernel::KA.Kernel{POCLBackend}, I, _ndrange, iterspace, ::Dynamic @@ -164,47 +167,50 @@ function KA.launch_config(kernel::KA.Kernel{POCLBackend}, ndrange, workgroupsize return ndrange, workgroupsize, iterspace, dynamic end -function threads_to_workgroupsize(threads, ndrange) - total = 1 - return map(ndrange) do n - x = min(div(threads, total), n) - total *= x - return x - end -end - function (obj::KA.Kernel{POCLBackend})(args::Vararg{Any, N}; ndrange = nothing, workgroupsize = nothing) where {N} ndrange, workgroupsize, iterspace, dynamic = KA.launch_config(obj, ndrange, workgroupsize) + # the launch doesn't depend on the tuned workgroup size, so neither does the context + launch = KA.select_launch(obj, workgroupsize, iterspace) + launch_kernel(obj, launch, ndrange, workgroupsize, iterspace, args...) + return nothing +end +function launch_kernel(obj, launch, ndrange, workgroupsize, iterspace, args::Vararg{Any, N}) where {N} # this might not be the final context, since we may tune the workgroupsize - ctx = KA.mkcontext(obj, ndrange, iterspace) + ctx = KA.mkcontext(obj, ndrange, iterspace, launch) kernel = @opencl launch = false obj.f(ctx, args...) # 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, KA.NDIteration.extents(ndrange)) + wg_size_nd = KA.launch_workgroupsize(KA.backend(obj), launch, wg_info.size, ndrange) iterspace, dynamic = KA.partition(obj, ndrange, wg_size_nd) - ctx = KA.mkcontext(obj, ndrange, iterspace) + ctx = KA.mkcontext(obj, ndrange, iterspace, launch) end - groups = length(KA.blocks(iterspace)) - items = length(KA.workitems(iterspace)) - - if groups == 0 + groups = size(KA.blocks(iterspace)) + items = size(KA.workitems(iterspace)) + if prod(groups) == 0 return nothing end # Launch kernel - global_size = groups * items - local_size = items + if launch isa KA.NDLaunch + local_size = pad3(items) + global_size = local_size .* pad3(groups) + else + local_size = prod(items) + global_size = prod(groups) * local_size + end event = kernel(ctx, args...; global_size, local_size) wait(event) cl.clReleaseEvent(event) return nothing end +pad3(t::Tuple) = (t..., ntuple(_ -> 1, 3 - length(t))...) + KI.argconvert(::POCLBackend, arg) = clconvert(arg) function KI.kernel_function(backend::POCLBackend, f::F, tt::TT = Tuple{}; name = nothing, kwargs...) where {F, TT} @@ -303,16 +309,6 @@ end @device_override KI.get_sub_group_local_id(::Type{T}) where {T} = get_sub_group_local_id() % T -@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 @device_override @inline function KI.localmemory(::Type{T}, ::Val{Dims}) where {T, Dims} diff --git a/test/codegen_checks.jl b/test/codegen_checks.jl index 7d4794735..e70881b32 100644 --- a/test/codegen_checks.jl +++ b/test/codegen_checks.jl @@ -52,6 +52,11 @@ end @print("index ", I, "\n") end +@kernel function codegen_global_linear(A) + I = @index(Global, Linear) + @inbounds A[I] = I +end + # `@inbounds` is only honoured under `--check-bounds=auto`; several checks below assert # that it removes code, so refuse to run under anything else rather than fail obscurely. if Base.JLOptions().check_bounds != 0 @@ -152,6 +157,30 @@ end end end + # An N-d launch maps the work-item builtins onto a dynamic 3-D iteration space directly, + # without decomposing linear ids. + @testset "N-d launch" begin + B = KernelAbstractions.zeros(backend, Int, 4, 5, 6) + @test @filecheck implicit_check_not = "{{[us]div i(32|64)}}" begin + @check "define spir_kernel void @{{.*}}gpu_codegen_global_linear" + @check "ret void" + @device_code_llvm debuginfo = :none codegen_global_linear(backend)(B, ndrange = size(B)) + KernelAbstractions.synchronize(backend) + end + + # a linear launch does, so the test above is not vacuous + kernel = codegen_global_linear(backend) + ndrange, workgroupsize, iterspace, _ = KernelAbstractions.launch_config(kernel, size(B), nothing) + @test @filecheck begin + @check "define spir_kernel void @{{.*}}gpu_codegen_global_linear" + @check "udiv i32" + @device_code_llvm debuginfo = :none KernelAbstractions.POCL.POCLKernels.launch_kernel( + kernel, KernelAbstractions.LinearLaunch{Int32}(), ndrange, workgroupsize, iterspace, B + ) + KernelAbstractions.synchronize(backend) + end + end + # `@print` lowers to a single variadic printf call, not to one call per argument. @testset "print" begin @test @filecheck begin diff --git a/test/launch.jl b/test/launch.jl new file mode 100644 index 000000000..429950231 --- /dev/null +++ b/test/launch.jl @@ -0,0 +1,195 @@ +using KernelAbstractions +using KernelAbstractions.NDIteration +import KernelAbstractions.KernelInterface as KI +using Test + +@kernel function launch_indices!(GL, GC, BL, BC, LL, LC, WS, lo) + I = @index(Global, NTuple) + J = I .- lo .+ 1 + @inbounds begin + GL[J...] = @index(Global, Linear) + GC[J...] = @index(Global, Cartesian) + BL[J...] = @index(Group, Linear) + BC[J...] = @index(Group, Cartesian) + LL[J...] = @index(Local, Linear) + LC[J...] = @index(Local, Cartesian) + WS[J...] = CartesianIndex(@groupsize()) + end +end + +# `unsafe_indices` kernels compute global indices from the group and local ones +@kernel unsafe_indices = true function launch_unsafe!(A) + g = @index(Group, NTuple) + l = @index(Local, NTuple) + I = (g .- 1) .* @groupsize() .+ l + if all(I .<= size(A)) + @inbounds A[I...] = LinearIndices(A)[I...] + end +end + +# padding lanes of partial workgroups have to reach the barrier too +@kernel function launch_sync!(A) + I = @index(Global, Linear) + i = @index(Local, Linear) + N = @uniform prod(@groupsize()) + lmem = @localmem Int (N,) + @inbounds lmem[i] = I + @synchronize + @inbounds A[I] = lmem[i] +end + +default_launcher(kernel, args...; ndrange, workgroupsize = nothing) = + kernel(args...; ndrange, workgroupsize) + +# Check every `@index` flavour against the layout KernelAbstractions defines: groups and +# work-items are numbered column-major, and the global index is `(g-1)*groupsize + l`. +function check_indices(launcher, backend, AT, kernel, ndrange; workgroupsize = nothing) + ranges = map(r -> r isa Integer ? (1:r) : r, ndrange) + N = length(ranges) + ext = map(length, ranges) + lo = map(first, ranges) + arrays = map((Int, CartesianIndex{N}, Int, CartesianIndex{N}, Int, CartesianIndex{N}, CartesianIndex{N})) do T + AT(zeros(T, ext)) + end + launcher(kernel, arrays..., lo; ndrange, workgroupsize) + synchronize(backend) + GL, GC, BL, BC, LL, LC, WS = map(Array, arrays) + + wgs = Tuple(first(WS)) + all(==(CartesianIndex(wgs)), WS) || return false + groups = cld.(ext, wgs) + for J in CartesianIndices(ext) + g = cld.(J.I, wgs) + l = J.I .- (g .- 1) .* wgs + GL[J] == LinearIndices(ext)[J] || return false + GC[J] == CartesianIndex(J.I .+ lo .- 1) || return false + BC[J] == CartesianIndex(g) || return false + BL[J] == LinearIndices(groups)[g...] || return false + LC[J] == CartesianIndex(l) || return false + LL[J] == LinearIndices(wgs)[l...] || return false + end + return true +end + +# 0-d ranges only work with a launch configuration +function launch_testsuite(backend, AT; launcher = default_launcher, zerodim = false) + @testset "index layout" begin + shapes = Tuple[(7,), (37,), (5, 7), (33, 3), (3, 5, 7), (2, 3, 4, 5)] + zerodim && pushfirst!(shapes, ()) + @testset "$shape, workgroupsize=$wgs" for shape in shapes, + wgs in (nothing, 4, (2, 3), (4, 1, 2)) + wgs !== nothing && length(wgs) > length(shape) && continue + @test check_indices(launcher, backend(), AT, launch_indices!(backend()), shape; + workgroupsize = wgs) + end + + @testset "static workgroupsize" begin + @test check_indices(launcher, backend(), AT, launch_indices!(backend(), (4, 2)), (9, 5)) + @test check_indices(launcher, backend(), AT, launch_indices!(backend(), 8), (9, 5, 3)) + end + + @testset "static ndrange" begin + @test check_indices(launcher, backend(), AT, launch_indices!(backend(), (4, 2), (9, 5)), (9, 5)) + @test check_indices(launcher, backend(), AT, launch_indices!(backend(), (4, 2), (0:8, -2:2)), (0:8, -2:2)) + end + + @testset "offsets" begin + @test check_indices(launcher, backend(), AT, launch_indices!(backend()), (-3:4,)) + @test check_indices(launcher, backend(), AT, launch_indices!(backend()), (-3:4, 2:11); workgroupsize = (3, 3)) + @test check_indices(launcher, backend(), AT, launch_indices!(backend()), (0:4, 3, 2:3)) + end + end + + @testset "empty ndrange" begin + for shape in ((0,), (0, 5), (5, 0), (3, 4, 0), (2, 0, 2, 2)) + A = AT(zeros(Int, max.(shape, 1))) + launcher(launch_unsafe!(backend()), A; ndrange = shape) + synchronize(backend()) + @test all(iszero, Array(A)) + end + end + + @testset "unsafe_indices" begin + for (shape, wgs) in (((37,), 8), ((33, 7), (8, 4)), ((9, 5, 3), (4, 2, 2)), ((9, 5), nothing)) + A = AT(zeros(Int, shape)) + launcher(launch_unsafe!(backend()), A; ndrange = shape, workgroupsize = wgs) + synchronize(backend()) + @test Array(A) == LinearIndices(A) + end + end + + @testset "synchronize with padding lanes" begin + for (shape, wgs) in (((37,), (8,)), ((7, 6), (4, 4)), ((5, 3, 3), (2, 2, 2))) + A = AT(zeros(Int, shape)) + launcher(launch_sync!(backend(), wgs), A; ndrange = shape) + synchronize(backend()) + @test Array(A) == LinearIndices(A) + end + end + return +end + +function select_launch_testsuite() + # the limits of a CUDA GPU + max_items = 1024 + max_dims = (1024, 1024, 64) + max_groups = (Int(typemax(Int32)), 65535, 65535) + select(extent, groupsize = nothing; max_dims = max_dims, max_groups = max_groups) = + KernelAbstractions.select_launch(extent, groupsize, max_items, max_dims, max_groups) + LinearLaunch = KernelAbstractions.LinearLaunch + NDLaunch = KernelAbstractions.NDLaunch + + @testset "selection" begin + # as many dimensions as the grid has are launched as such + @test select((1000,)) === NDLaunch{Int32}() + @test select((100, 100)) === NDLaunch{Int32}() + @test select((10, 10, 10), (4, 4, 4)) === NDLaunch{Int32}() + @test select(()) === NDLaunch{Int32}() + @test select((2, 3, 4, 5)) === LinearLaunch{Int32}() + @test select((100, 100); max_dims = (1024, 1024), max_groups = (65535, 65535)) === + NDLaunch{Int32}() + @test select((10, 10, 10); max_dims = (1024, 1024), max_groups = (65535, 65535)) === + LinearLaunch{Int32}() + + # unless that exceeds the per-dimension limits + @test select((100, 100_000)) === LinearLaunch{Int32}() + @test select((1, 1, 5000), (1, 1, 128)) === LinearLaunch{Int32}() + @test select((10,), (2048,)) === LinearLaunch{Int32}() + @test select((64, 64), (64, 64)) === LinearLaunch{Int32}() + # tuning respects the per-dimension limits + @test select((1, 1, 5000)) === NDLaunch{Int32}() + + # indices are computed in Int32 if the padded iteration space fits + @test select((1024, 1024, 1024)) === NDLaunch{Int32}() + @test select((1025, 1024, 1024)) === NDLaunch{Int}() + @test select((2^31 - 1024,)) === NDLaunch{Int32}() + @test select((2^31 - 1023,)) === NDLaunch{Int}() + @test select((2^31 - 1,), (1,)) === NDLaunch{Int32}() + @test select((2^31 - 1,), (2,)) === NDLaunch{Int}() + @test select((2^16, 2^16, 2), (256,)) === LinearLaunch{Int}() + # ... and don't have to be representable at all + @test_throws ArgumentError select((2^40, 2^40)) + @test_throws ArgumentError select((typemax(Int),)) + @test_throws ArgumentError select((typemax(Int),), (2,)) + # including when empty + @test select((2^40, 2^40, 0)) === LinearLaunch{Int32}() + @test select((2^20, 2^12, 0), (1024,)) === NDLaunch{Int32}() + end + + # The index type is chosen before the workgroup size is tuned, so the padding that the + # tuned workgroup introduces has to be bounded for every thread count. + @testset "tuned padding bound" begin + for extent in ((5,), (1000,), (3, 7), (33, 1000), (1, 1, 5000), (7, 9, 11), + (1500, 3, 2), (2, 3, 4, 5), (1, 1, 1, 3000), (0, 7)) + for limits in ((), max_dims) + bound = KernelAbstractions.tuned_padded(extent, max_items, limits) + @test all(1:max_items) do threads + wgs = KI.threads_to_workgroupsize(threads, extent, limits) + padded = cld.(extent, wgs) .* wgs + all(padded .<= bound) + end + end + end + end + return +end diff --git a/test/runtests.jl b/test/runtests.jl index 9638ae7f1..13ab5aced 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -145,6 +145,79 @@ end end end +# a mapping KernelAbstractions doesn't know, which the index functions have to go +# through `expand`, `in` and `linear_index` for +struct TransposedMapping end +Base.@propagate_inbounds function KernelAbstractions.expand( + ndrange::KernelAbstractions.NDRange{2, B, W, DB, DW, TransposedMapping}, + groupidx::CartesianIndex{2}, idx::CartesianIndex{2} + ) where {B, W, DB, DW} + I = (groupidx.I .- 1) .* size(KernelAbstractions.workitems(ndrange)) .+ idx.I + return CartesianIndex(reverse(I)) +end + +# like Oceananigans: offsets kept in the dynamic workitems slot, applied by a custom `expand` +struct ItemOffsets{N} + offsets::NTuple{N, Int} +end +Base.@propagate_inbounds function KernelAbstractions.expand( + ndrange::KernelAbstractions.NDRange{N, B, W, Nothing, ItemOffsets{N}}, + groupidx::CartesianIndex{N}, idx::CartesianIndex{N} + ) where {N, B, W} + I = (groupidx.I .- 1) .* size(KernelAbstractions.workitems(ndrange)) .+ idx.I + return CartesianIndex(I .+ ndrange.workitems.offsets) +end + +@kernel function mapped_indices!(A) + I = @index(Global, Cartesian) + @inbounds A[I] = @index(Global, Linear) +end + +# host-side logic, independent of the backend +@testset "select_launch" begin + Testsuite.select_launch_testsuite() +end + +# the shared testsuite only covers the launch configuration POCL selects +@testset "POCL launch configurations" begin + KA = KernelAbstractions + @testset "custom mapping, $launch" for launch in (nothing, KA.LinearLaunch{Int}(), KA.NDLaunch{Int}()) + # a 7x5 iteration space in 4x4 workgroups, mapped onto a 5x7 ndrange + kernel = mapped_indices!(CPU(), (4, 4)) + iterspace = KA.NDRange{2, KA.DynamicSize, KA.StaticSize{(4, 4)}}( + CartesianIndices((2, 2)), nothing, TransposedMapping() + ) + A = zeros(Int, 5, 7) + POCL.POCLKernels.launch_kernel(kernel, launch, CartesianIndices(A), nothing, iterspace, A) + @test A == LinearIndices(A) + end + @testset "custom iteration space, $launch" for launch in (nothing, KA.LinearLaunch{Int}(), KA.NDLaunch{Int}()) + # an 8x8 iteration space in 4x4 workgroups, shifted by (1, 2) onto a 7x5 ndrange + kernel = mapped_indices!(CPU(), (4, 4)) + iterspace = KA.NDRange{2, KA.StaticSize{(2, 2)}, KA.StaticSize{(4, 4)}}(nothing, ItemOffsets((1, 2))) + ndrange = CartesianIndices((2:8, 3:7)) + A = zeros(Int, 9, 8) + POCL.POCLKernels.launch_kernel(kernel, launch, ndrange, nothing, iterspace, A) + @test A[ndrange] == LinearIndices(ndrange) + A[ndrange] .= 0 + @test all(iszero, A) + + # the launch is chosen from the iteration space + @test KA.select_launch(kernel, nothing, iterspace) === KA.NDLaunch{Int32}() + end + for launch in (nothing, KA.LinearLaunch{Int}(), KA.NDLaunch{Int}()) + function launcher(kernel, args...; ndrange, workgroupsize = nothing) + ndrange, workgroupsize, iterspace, _ = KA.launch_config(kernel, ndrange, workgroupsize) + # an N-d launch is limited to three dimensions + l = launch isa KA.NDLaunch && ndims(iterspace) > 3 ? KA.LinearLaunch{Int}() : launch + POCL.POCLKernels.launch_kernel(kernel, l, ndrange, workgroupsize, iterspace, args...) + end + @testset "$launch" begin + Testsuite.launch_testsuite(CPU, Array; launcher, zerodim = launch !== nothing) + end + end +end + @testset "CPU back-end" begin Testsuite.testsuite(CPU, "CPU", POCL, Array, POCL.CLDeviceArray) end diff --git a/test/testsuite.jl b/test/testsuite.jl index c9d650d49..9a0aa0048 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("launch.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 "Launch" skip_tests begin + launch_testsuite(backend, AT) + end + @conditional_testset "copyto!" skip_tests begin copyto_testsuite(backend, AT) end From 7509cf401aab00852f4d5d8967a78276312799ae Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sun, 27 Sep 2026 14:02:39 +0200 Subject: [PATCH 2/2] Index kernels without a launch configuration like a LinearLaunch{Int} The default path decomposed the hardware ids by indexing `blocks(iterspace)` and `workitems(iterspace)`, which throws for 0-d ndranges (`linear_index` reads `I.I[0]`), and which under `--check-bounds=yes` puts bounds checks depending on the local id in front of `@synchronize`. POCL crashes on the resulting barrier in non-uniform control flow, e.g. for a 3-D ndrange with partial workgroups. Computing the positions like a `LinearLaunch{Int}` avoids both, and leaves only one implementation of the index functions. --- src/KernelAbstractions.jl | 35 ++++++++++------------------------- test/launch.jl | 18 ++++++++++-------- test/runtests.jl | 2 +- 3 files changed, 21 insertions(+), 34 deletions(-) diff --git a/src/KernelAbstractions.jl b/src/KernelAbstractions.jl index 6e26bb1fd..c0c7229bc 100644 --- a/src/KernelAbstractions.jl +++ b/src/KernelAbstractions.jl @@ -390,30 +390,15 @@ end ### # The index functions dispatch on the launch configuration of the context (see -# `launch.jl`); `nothing` is a 1-D launch, indexed in `Int`. - -@inline __index_Local_Linear(ctx) = local_linear(ctx, __launch(ctx)) -@inline __index_Group_Linear(ctx) = group_linear(ctx, __launch(ctx)) -@inline __index_Global_Linear(ctx) = global_linear(ctx, __launch(ctx)) -@inline __index_Local_Cartesian(ctx) = local_cartesian(ctx, __launch(ctx)) -@inline __index_Group_Cartesian(ctx) = group_cartesian(ctx, __launch(ctx)) -@inline __index_Global_Cartesian(ctx) = global_cartesian(ctx, __launch(ctx)) - -@inline local_linear(ctx, ::Nothing) = KI.get_local_id().x -@inline group_linear(ctx, ::Nothing) = KI.get_group_id().x -@inline function global_linear(ctx, ::Nothing) - 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) -end -@inline local_cartesian(ctx, ::Nothing) = @inbounds workitems(__iterspace(ctx))[KI.get_local_id().x] -@inline group_cartesian(ctx, ::Nothing) = @inbounds blocks(__iterspace(ctx))[KI.get_group_id().x] -@inline global_cartesian(ctx, ::Nothing) = - @inbounds expand(__iterspace(ctx), KI.get_group_id().x, KI.get_local_id().x) -@inline function validindex(ctx, ::Nothing) - I = @inbounds expand(__iterspace(ctx), KI.get_group_id().x, KI.get_local_id().x) - return I in __ndrange(ctx) -end +# `launch.jl`). A context without one was launched on a 1-D grid, and is indexed in `Int`. +@inline index_launch(ctx) = something(__launch(ctx), LinearLaunch{Int}()) + +@inline __index_Local_Linear(ctx) = local_linear(ctx, index_launch(ctx)) +@inline __index_Group_Linear(ctx) = group_linear(ctx, index_launch(ctx)) +@inline __index_Global_Linear(ctx) = global_linear(ctx, index_launch(ctx)) +@inline __index_Local_Cartesian(ctx) = local_cartesian(ctx, index_launch(ctx)) +@inline __index_Group_Cartesian(ctx) = group_cartesian(ctx, index_launch(ctx)) +@inline __index_Global_Cartesian(ctx) = global_cartesian(ctx, index_launch(ctx)) @inline __index_Local_NTuple(ctx, I...) = Tuple(__index_Local_Cartesian(ctx, I...)) @inline __index_Group_NTuple(ctx, I...) = Tuple(__index_Group_Cartesian(ctx, I...)) @@ -620,7 +605,7 @@ function __workitems_iterspace end # workgroup. Padding lanes still take part in `@synchronize`. @inline function __validindex(ctx) if __dynamic_checkbounds(ctx) - return validindex(ctx, __launch(ctx)) + return validindex(ctx, index_launch(ctx)) else return true end diff --git a/test/launch.jl b/test/launch.jl index 429950231..1c1988fde 100644 --- a/test/launch.jl +++ b/test/launch.jl @@ -71,16 +71,16 @@ function check_indices(launcher, backend, AT, kernel, ndrange; workgroupsize = n return true end -# 0-d ranges only work with a launch configuration -function launch_testsuite(backend, AT; launcher = default_launcher, zerodim = false) +function launch_testsuite(backend, AT; launcher = default_launcher) @testset "index layout" begin - shapes = Tuple[(7,), (37,), (5, 7), (33, 3), (3, 5, 7), (2, 3, 4, 5)] - zerodim && pushfirst!(shapes, ()) + shapes = Tuple[(), (7,), (37,), (5, 7), (33, 3), (3, 5, 7), (2, 3, 4, 5)] @testset "$shape, workgroupsize=$wgs" for shape in shapes, wgs in (nothing, 4, (2, 3), (4, 1, 2)) wgs !== nothing && length(wgs) > length(shape) && continue - @test check_indices(launcher, backend(), AT, launch_indices!(backend()), shape; - workgroupsize = wgs) + @test check_indices( + launcher, backend(), AT, launch_indices!(backend()), shape; + workgroupsize = wgs + ) end @testset "static workgroupsize" begin @@ -179,8 +179,10 @@ function select_launch_testsuite() # The index type is chosen before the workgroup size is tuned, so the padding that the # tuned workgroup introduces has to be bounded for every thread count. @testset "tuned padding bound" begin - for extent in ((5,), (1000,), (3, 7), (33, 1000), (1, 1, 5000), (7, 9, 11), - (1500, 3, 2), (2, 3, 4, 5), (1, 1, 1, 3000), (0, 7)) + for extent in ( + (5,), (1000,), (3, 7), (33, 1000), (1, 1, 5000), (7, 9, 11), + (1500, 3, 2), (2, 3, 4, 5), (1, 1, 1, 3000), (0, 7), + ) for limits in ((), max_dims) bound = KernelAbstractions.tuned_padded(extent, max_items, limits) @test all(1:max_items) do threads diff --git a/test/runtests.jl b/test/runtests.jl index 13ab5aced..07499b704 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -213,7 +213,7 @@ end POCL.POCLKernels.launch_kernel(kernel, l, ndrange, workgroupsize, iterspace, args...) end @testset "$launch" begin - Testsuite.launch_testsuite(CPU, Array; launcher, zerodim = launch !== nothing) + Testsuite.launch_testsuite(CPU, Array; launcher) end end end