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..c0c7229bc 100644 --- a/src/KernelAbstractions.jl +++ b/src/KernelAbstractions.jl @@ -389,29 +389,16 @@ 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 +# The index functions dispatch on the launch configuration of the context (see +# `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 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) -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) -end +@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...)) @@ -606,13 +593,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, index_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..1c1988fde --- /dev/null +++ b/test/launch.jl @@ -0,0 +1,197 @@ +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 + +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)] + @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..07499b704 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) + 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