From aa64cc59b480d1d0145bae8e318f58886d3ef214 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Sat, 10 Oct 2026 18:52:48 +0200 Subject: [PATCH 1/9] KernelInterface: sub-group shuffles and votes Add `shfl`, `shfl_up` and `shfl_xor` next to `shfl_down`, their variants within segments of a `width`, and the votes `sub_group_any`, `sub_group_all` and `sub_group_ballot`. Back ends implement the shuffles for their native primitive types; other primitive types are shuffled as `UInt32` or `UInt64` words and `isbits` structs field by field, with `supports_shuffle` derived accordingly. Which work-items form a sub-group stays unspecified. Back ends that form them from consecutive work-items report `supports_linear_subgroups`, and those whose sub-groups can communicate in control flow that differs between the sub-groups of a work-group report `supports_independent_subgroups`. The testsuite computes the expected results from the sub-groups that formed, so that it runs without either guarantee. The POCL back end implements the shuffles and votes with SPIRVIntrinsics, for widths of up to 64 so that the ballot covers every lane. PoCL forms sub-groups from consecutive work-items, but executes sub-group operations as work-group barriers, so it reports linear but not independent sub-groups. Co-authored-by: Tim Besard --- docs/src/kernelinterface.md | 86 +++- lib/KernelInterface/src/KernelInterface.jl | 7 +- lib/KernelInterface/src/device.jl | 334 ++++++++++++++- lib/KernelInterface/src/host.jl | 90 +++- lib/KernelInterface/test/interface.jl | 468 +++++++++++++++++---- lib/KernelInterface/test/runtests.jl | 35 +- src/KernelAbstractions.jl | 7 + src/pocl/backend.jl | 43 +- src/pocl/pocl.jl | 5 +- test/codegen_checks.jl | 24 ++ test/runtests.jl | 12 + 11 files changed, 968 insertions(+), 143 deletions(-) diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md index 4eba70529..c613c5a98 100644 --- a/docs/src/kernelinterface.md +++ b/docs/src/kernelinterface.md @@ -74,10 +74,10 @@ What a backend implements, at a glance. The docstrings below have the details. | **Execution** | [`synchronize`](@ref) (cooperative) | [`record_event`](@ref)/[`wait_event`](@ref) (synchronize), [`priority!`](@ref) (no-op) | | **Devices** | with more than one device: [`ndevices`](@ref), [`device`](@ref), [`device!`](@ref), `device(backend, A)` | all four (a single device) | | **Queries** | [`max_work_group_size`](@ref) (for the backend and for a kernel), [`max_work_group_dims`](@ref), [`max_num_groups`](@ref) | [`launch_configuration`](@ref) (the limit), [`multiprocessor_count`](@ref) (0), [`functional`](@ref) (`missing`), [`versioninfo`](@ref) | -| **Capabilities** | | [`supports_float64`](@ref), [`supports_atomics`](@ref), [`supports_unified`](@ref), [`supports_subgroups`](@ref), [`supports_shuffle`](@ref) (all `false`) | +| **Capabilities** | | [`supports_float64`](@ref), [`supports_atomics`](@ref), [`supports_unified`](@ref), [`supports_subgroups`](@ref), [`supports_linear_subgroups`](@ref), [`supports_independent_subgroups`](@ref) (all `false`), [`supports_shuffle`](@ref) (`false` unless derived from the natively supported types, see its docstring) | | **Compilation** | [`argconvert`](@ref), [`kernel_function`](@ref), [`launch`](@ref) | | | **Device** | [`get_local_id`](@ref), [`get_group_id`](@ref), [`get_local_size`](@ref), [`get_num_groups`](@ref), [`localmemory`](@ref), [`barrier`](@ref) | [`get_global_id`](@ref), [`get_global_size`](@ref) (derived from the primitive queries), [`_print`](@ref KernelInterface._print) (host `print`) | -| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries, [`sub_group_barrier`](@ref); if `supports_shuffle(backend, T)`: [`shfl_down`](@ref) for `T` | | +| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries (with a constant [`get_max_sub_group_size`](@ref)), [`sub_group_barrier`](@ref), [`sub_group_any`](@ref), [`sub_group_all`](@ref), and [`sub_group_ballot`](@ref) for widths of at most 64; if `supports_shuffle(backend, T)`: [`shfl`](@ref), [`shfl_down`](@ref), [`shfl_up`](@ref), [`shfl_xor`](@ref) for the primitive `T` supported natively, including `UInt32` | shuffles with a `width` (built on [`shfl`](@ref)); shuffles of other primitive types (as `UInt32` or `UInt64` words) and of structs (field by field) | Everything else, such as [`zeros`](@ref KernelInterface.zeros), [`ones`](@ref KernelInterface.ones), the launch-keyword handling of [`Kernel`](@ref) and [`@launch`](@ref KernelInterface.@launch), @@ -149,24 +149,63 @@ get_global_size ### Sub-groups Sub-groups are optional ([`supports_subgroups`](@ref)). A work-group is divided into -sub-groups of at most [`sub_group_size(backend)`](@ref sub_group_size) work-items. Which -work-items form a sub-group, how many sub-groups there are, and which of them are partial -is unspecified, and differs between devices and work-group shapes. For example, CUDA forms -warps from consecutive linear work-item indices, while Intel's CPU OpenCL runtime forms -sub-groups per row of a multi-dimensional work-group, so that a 33×2 work-group consists of -four sub-groups of 32 and 1 work-items. What KernelInterface guarantees, and backends that -report sub-group support have to ensure: +sub-groups of at most `W = `[`sub_group_size(backend)`](@ref sub_group_size) work-items, the +sub-group width. What KernelInterface guarantees, and backends that report sub-group support +have to ensure: - every work-item has a unique `(get_sub_group_id(), get_sub_group_local_id())` pair in its work-group, which doesn't change during the kernel; - the sub-group ids are `1:get_num_sub_groups()`, and the lanes of a sub-group are - `1:get_sub_group_size()`; -- a 1-D work-group of at most `sub_group_size(backend)` work-items is a single sub-group. - -In particular, [`get_num_sub_groups`](@ref) can be larger than -`cld(prod(get_local_size()), get_max_sub_group_size())`. Storage for a value per sub-group -has to be sized for up to one sub-group per work-item, and code combining those values has -to use `get_num_sub_groups()` rather than compute the count. + `1:get_sub_group_size()`. + +Which work-items form a sub-group is **not specified**: it differs between APIs, devices and +work-group shapes. For example, Intel's CPU OpenCL runtime forms sub-groups per row of a +multi-dimensional work-group, so that a 33×2 work-group consists of four sub-groups of 32, 1, +32 and 1 work-items. In particular, [`get_num_sub_groups`](@ref) can be larger than +`cld(prod(get_local_size()), W)`. Storage for a value per sub-group has to be sized for up to +one sub-group per work-item, and code combining those values has to use +`get_num_sub_groups()` rather than compute the count. + +A sub-group is partial when it has fewer work-items than the width: the lanes +`get_sub_group_size()+1:W` have no work-item. Shuffles from those lanes give unspecified +values, and the votes only take the work-items of the sub-group into account. + +Backends that guarantee more say so with two queries: + +- [`supports_linear_subgroups`](@ref): if the work-group is 1-D, or its x extent is a + multiple of `W`, sub-groups are formed from consecutive work-items, x fastest: the + work-item with the linear local index + `lin = x + (y - 1) * size.x + (z - 1) * size.x * size.y` (for `(; x, y, z) = + get_local_id()` and `size = get_local_size()`) is in sub-group `fld(lin - 1, W) + 1`, lane + `mod(lin - 1, W) + 1`, and only the last sub-group can be partial. Code that relates local + ids to lanes needs this, e.g. stencils exchanging neighbouring values with shuffles. Code + that doesn't, like a work-group reduction combining a value per sub-group, can index by the + sub-group slot `(get_sub_group_id() - 1) * W + get_sub_group_local_id()` instead. +- [`supports_independent_subgroups`](@ref): the communication functions (shuffles, votes + and [`sub_group_barrier`](@ref)) can be executed in control flow that is uniform + over each sub-group, but differs between the sub-groups of a work-group, e.g. a loop whose + trip count depends on the sub-group, or after an early `return` of whole sub-groups. + Otherwise, all sub-groups of a work-group have to execute the same communication + functions in the same order. + +| Backend | `supports_linear_subgroups` | `supports_independent_subgroups` | +|---|---|---| +| CUDA, AMDGPU | `true` | `true` | +| Metal | `true` (not specified by Metal, but tested) | `true` | +| oneAPI | once the layout is tested on Intel's GPU compiler | `true` | +| OpenCL.jl | on runtimes known to form them linearly | `true`, except on PoCL | +| KernelAbstractions' POCL backend | `true` | `false`: sub-group operations are work-group barriers | + +In all cases, all work-items of a sub-group have to execute a communication function +together, not in a branch that only some of its lanes take. The id and size queries can be +used anywhere. + +!!! note "Kernel languages" + The sub-group functions are meant for kernels written against KernelInterface, and for + kernel languages built on it. A kernel is either a KernelInterface kernel or a + KernelAbstractions `@kernel`, never a mix: `@kernel` adds padding work-items to partial + work-groups, which skip the kernel's body, so KernelInterface's sub-group functions called + in a `@kernel` aren't reached by all work-items. Host code may use the queries of both. ```@docs get_sub_group_size @@ -191,8 +230,21 @@ localmemory ### Communication +The shuffles, votes and collectives below exchange values between the work-items of a +sub-group, not memory: they don't order or make visible the work-items' accesses to local +or global memory. To communicate through memory within a sub-group, e.g. a work-item reading +what another one wrote to local memory, use [`sub_group_barrier`](@ref) between the write +and the read. + ```@docs +shfl shfl_down +shfl_up +shfl_xor +shfl(::Any, ::Integer, ::Integer) +sub_group_any +sub_group_all +sub_group_ballot ``` ### Printing @@ -247,6 +299,8 @@ supports_float64 ```@docs supports_subgroups +supports_linear_subgroups +supports_independent_subgroups supports_shuffle ``` diff --git a/lib/KernelInterface/src/KernelInterface.jl b/lib/KernelInterface/src/KernelInterface.jl index 0f0fade1a..6932ed8ba 100644 --- a/lib/KernelInterface/src/KernelInterface.jl +++ b/lib/KernelInterface/src/KernelInterface.jl @@ -31,7 +31,9 @@ include("host.jl") :get_group_id, :get_num_groups, :get_sub_group_size, :get_max_sub_group_size, :get_num_sub_groups, :get_sub_group_id, :get_sub_group_local_id, - :localmemory, :shfl_down, :barrier, :sub_group_barrier, :_print, + :localmemory, :shfl, :shfl_down, :shfl_up, :shfl_xor, + :sub_group_any, :sub_group_all, :sub_group_ballot, + :barrier, :sub_group_barrier, :_print, # compilation and launch :Kernel, :kernel_function, :argconvert, :launch, Symbol("@launch"), :launch_configuration, :max_work_group_size, :max_work_group_dims, @@ -42,7 +44,8 @@ include("host.jl") :device, :ndevices, :device!, :functional, :versioninfo, :supports_unified, :supports_atomics, :supports_float64, - :supports_subgroups, :supports_shuffle, + :supports_subgroups, :supports_shuffle, :supports_linear_subgroups, + :supports_independent_subgroups, ) ) end diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl index 5b0090914..09af74ca3 100644 --- a/lib/KernelInterface/src/device.jl +++ b/lib/KernelInterface/src/device.jl @@ -127,19 +127,22 @@ end # Sub-group support is optional, see `supports_subgroups`. A work-group is divided into # sub-groups of at most `sub_group_size(backend)` work-items. Which work-items form a -# sub-group, how many sub-groups there are and which are partial is unspecified, except -# that `(get_sub_group_id(), get_sub_group_local_id())` is unique within a work-group and -# doesn't change during the kernel's execution, and that a 1-D work-group of at most -# `sub_group_size(backend)` work-items is a single sub-group. Backends that can't ensure -# that don't report sub-group support. See the manual. +# sub-group, how many sub-groups there are and which are partial is unspecified (unless +# `supports_linear_subgroups`), except that `(get_sub_group_id(), get_sub_group_local_id())` +# is unique within a work-group and doesn't change during the kernel's execution (see the +# manual). +# +# In a partial sub-group, the lanes `get_sub_group_size()+1:get_max_sub_group_size()` have no +# work-item: shuffles from them give unspecified values, and the votes only take the +# work-items of the sub-group into account. """ get_sub_group_size([::Type{T}=Int])::T The number of work-items in the sub-group, at most the sub-group width ([`get_max_sub_group_size`](@ref)). Which sub-groups have fewer work-items than the width -is unspecified: when the work-group size isn't a multiple of the width, there can be more -than one, e.g. one per row of a multi-dimensional work-group. +is unspecified (unless [`supports_linear_subgroups`](@ref)): there can be more than one, +e.g. one per row of a multi-dimensional work-group. See [`get_local_id`](@ref) for the supported types `T`. @@ -155,12 +158,21 @@ See [`get_local_id`](@ref) for the supported types `T`. """ get_max_sub_group_size([::Type{T}=Int])::T -The sub-group width, [`sub_group_size(backend)`](@ref sub_group_size) on the host. +The sub-group width (the warp or wavefront size), [`sub_group_size(backend)`](@ref +sub_group_size) on the host. + +Backends should make it a constant of the generated code, so that code depending on it, e.g. +a loop over the lanes or a shuffle butterfly, is specialized for it; on some, it is only +folded late in compilation, or not at all. It is never known during type inference: to pick +types or `Val` parameters from the width, e.g. the integer type of a mask with a bit per +lane, pass [`sub_group_size(backend)`](@ref sub_group_size) from the host. See [`get_local_id`](@ref) for the supported types `T`. !!! note - Backend implementations that support sub-groups **must** implement: + Backend implementations that support sub-groups **must** implement this, and **should** + return the width the kernel is compiled for as a constant rather than query the device + at run time, so that code depending on it is specialized for the width: ``` @device_override get_max_sub_group_size(::Type{T})::T where {T} ``` @@ -173,8 +185,8 @@ See [`get_local_id`](@ref) for the supported types `T`. The number of sub-groups in the work-group. It is at least `cld(prod(get_local_size()), get_max_sub_group_size())`, but can be larger, since more than -one sub-group can be partial. Size storage for a value per sub-group for up to one -sub-group per work-item. +one sub-group can be partial (unless [`supports_linear_subgroups`](@ref)). Size storage for a +value per sub-group for up to one sub-group per work-item. See [`get_local_id`](@ref) for the supported types `T`. @@ -191,7 +203,8 @@ See [`get_local_id`](@ref) for the supported types `T`. get_sub_group_id([::Type{T}=Int])::T The 1-based index of the sub-group within the work-group, between 1 and -[`get_num_sub_groups`](@ref). How it relates to [`get_local_id`](@ref) is unspecified. +[`get_num_sub_groups`](@ref). How it relates to [`get_local_id`](@ref) is unspecified, +unless [`supports_linear_subgroups`](@ref). See [`get_local_id`](@ref) for the supported types `T`. @@ -209,7 +222,8 @@ See [`get_local_id`](@ref) for the supported types `T`. The 1-based index of the work-item within its sub-group (its lane), between 1 and [`get_sub_group_size`](@ref). It doesn't depend on which work-items of the sub-group are -active, e.g. in a divergent branch. +active, e.g. in a divergent branch. How it relates to [`get_local_id`](@ref) is unspecified, +unless [`supports_linear_subgroups`](@ref). See [`get_local_id`](@ref) for the supported types `T`. @@ -254,26 +268,301 @@ localmemory(::Type{T}, ::Val) where {T} = ## communication +# The shuffles, votes, collectives and `sub_group_barrier` have to be executed by all +# work-items of a sub-group together, in control flow that is uniform over the work-group +# unless `supports_independent_subgroups`. The id and size queries above can be used anywhere. + +""" + shfl(val::T, lane::Integer)::T + +Return `val` of the work-item with [`get_sub_group_local_id`](@ref) equal to `lane` in the +sub-group. When there is no such work-item, the result is an unspecified value (of type `T`). + +All work-items of the sub-group have to execute `shfl` together (not in a divergent branch), +but they may read from different lanes. Unless [`supports_independent_subgroups`](@ref), all +sub-groups of the work-group have to execute it. + +`shfl` exchanges values, not memory: it is not a memory fence. + +Types for which [`supports_shuffle`](@ref) returns `true` are supported. Besides the types a +backend supports natively, that includes other primitive types of 1, 2, 4, 8 or 16 bytes +(e.g. `Bool`, `Char` or `Int64`) if the backend supports `UInt32`, and `isbits` structs and +tuples of supported types, which are shuffled field by field. Primitive types of up to 4 +bytes are shuffled as a `UInt32`, and larger ones as `UInt64` words, each of which is shuffled +natively if the backend supports `UInt64`, and as two `UInt32` words otherwise. + +!!! note + Backend implementations **must** implement this for the primitive types they support + natively, which have to include `UInt32`, and only for those, so that other types reach + the fallbacks: + ``` + @device_override shfl(val::T, lane::Integer) where {T <: Union{...}} + ``` +""" +@inline shfl(val, lane::Integer) = shfl_fallback(x -> shfl(x, lane), val) + """ shfl_down(val::T, offset::Integer)::T Return `val` of the work-item `offset` lanes further in the sub-group, i.e. with -[`get_sub_group_local_id`](@ref) equal to `get_sub_group_local_id() + offset`. When there is -no such work-item, the result is an unspecified value (of type `T`). +[`get_sub_group_local_id`](@ref) equal to `get_sub_group_local_id() + offset`, for an `offset` +of at least 0. When that lane is past the sub-group width, i.e. +`get_sub_group_local_id() + offset > get_max_sub_group_size()`, the result is `val` of the +work-item itself, as on CUDA, HIP and Metal. When the lane is within the width but has no +work-item, in a partial sub-group, the result is an unspecified value (of type `T`). All work-items of the sub-group have to execute `shfl_down` together (not in a divergent -branch), with the same `offset`. +branch), with the same `offset`; see [`shfl`](@ref). + +`shfl_down` exchanges values, not memory: it is not a memory fence. See [`shfl`](@ref) for +the supported types. + +!!! note + Backend implementations **must** implement this like [`shfl`](@ref): + ``` + @device_override shfl_down(val::T, offset::Integer) where {T <: Union{...}} + ``` +""" +@inline shfl_down(val, offset::Integer) = shfl_fallback(x -> shfl_down(x, offset), val) + +""" + shfl_up(val::T, offset::Integer)::T + +Return `val` of the work-item `offset` lanes earlier in the sub-group, i.e. with +[`get_sub_group_local_id`](@ref) equal to `get_sub_group_local_id() - offset`, for an `offset` +of at least 0. When there is no such lane, i.e. `get_sub_group_local_id() <= offset`, the +result is `val` of the work-item itself, as on CUDA, HIP and Metal. + +All work-items of the sub-group have to execute `shfl_up` together (not in a divergent +branch), with the same `offset`; see [`shfl`](@ref). + +`shfl_up` exchanges values, not memory: it is not a memory fence. See [`shfl`](@ref) for +the supported types. + +!!! note + Backend implementations **must** implement this like [`shfl`](@ref): + ``` + @device_override shfl_up(val::T, offset::Integer) where {T <: Union{...}} + ``` +""" +@inline shfl_up(val, offset::Integer) = shfl_fallback(x -> shfl_up(x, offset), val) + +""" + shfl_xor(val::T, mask::Integer)::T + +Return `val` of the work-item whose 0-based lane id is the 0-based lane id of this work-item +xor `mask`, i.e. with [`get_sub_group_local_id`](@ref) equal to +`((get_sub_group_local_id() - 1) ⊻ mask) + 1`. `mask` has to be between 0 and +`get_max_sub_group_size() - 1`. When that lane has no work-item, in a partial sub-group, the +result is an unspecified value (of type `T`). + +All work-items of the sub-group have to execute `shfl_xor` together (not in a divergent +branch), with the same `mask`; see [`shfl`](@ref). A butterfly over the masks `width ÷ 2, …, 2, 1` (for the +sub-group width [`get_max_sub_group_size`](@ref)) reduces a full sub-group such that every +work-item gets the result. + +`shfl_xor` exchanges values, not memory: it is not a memory fence. See [`shfl`](@ref) for +the supported types. + +!!! note + Backend implementations **must** implement this like [`shfl`](@ref): + ``` + @device_override shfl_xor(val::T, mask::Integer) where {T <: Union{...}} + ``` +""" +@inline shfl_xor(val, mask::Integer) = shfl_fallback(x -> shfl_xor(x, mask), val) + +# Shuffle a value of a type that the backend doesn't support natively, with `f` shuffling a +# value of a type it does support. The fallbacks are separate functions for primitive and for +# other types, so that the fallback of a struct with a field the backend doesn't support +# natively, e.g. an `Int64` on Metal, doesn't call itself, which inference gives up on (on +# Julia 1.10). +@inline function shfl_fallback(f, val::T) where {T} + return isprimitivetype(T) ? shfl_words(f, val) : shfl_fields(f, val) +end + +shfl_unsupported(T) = throw( + ArgumentError( + "Shuffling values of type $T is not supported by this backend, see `supports_shuffle`" + ) +) + +# Whether a primitive type that a backend doesn't support natively can be shuffled as unsigned +# words +shuffle_as_words(T) = T !== UInt32 && sizeof(T) in (1, 2, 4, 8, 16) + +# The unsigned integer type of the size of a primitive type `T` (of 1, 2, 4, 8 or 16 bytes) +const word_types = (UInt8, UInt16, UInt32, UInt64, UInt128) +word_type(T) = word_types[trailing_zeros(sizeof(T)) + 1] + +# Shuffle a primitive value as unsigned words: values of up to 4 bytes are zero-extended to a +# `UInt32`. Other 8-byte values are shuffled as a `UInt64`, which the backend may support +# natively, and a `UInt64` it doesn't is split into two `UInt32` words. 16-byte values are +# split into two `UInt64` words. +@inline @generated function shfl_words(f, val::T) where {T} + shuffle_as_words(T) || return :(shfl_unsupported($T)) + U = word_type(T) + if sizeof(T) <= 4 + return :(reinterpret($T, f(reinterpret($U, val) % UInt32) % $U)) + elseif sizeof(T) == 8 && T !== UInt64 + return :(reinterpret($T, f(reinterpret(UInt64, val)))) + end + W = sizeof(T) == 8 ? UInt32 : UInt64 + nbits = 8 * sizeof(W) + words = (:((f((bits >> $(nbits * (i - 1))) % $W) % $U) << $(nbits * (i - 1))) for i in 1:2) + return quote + bits = reinterpret($U, val) + return reinterpret($T, |($(words...))) + end +end + +# A `UInt64` that the backend doesn't shuffle natively is split into two `UInt32` words by a +# method of its own, so that another 8-byte type, shuffled as a `UInt64`, doesn't reach the +# fallback through itself, which inference gives up on (on Julia 1.10). Backends that shuffle +# `UInt64` natively override these. +@inline shfl(val::UInt64, lane::Integer) = shfl_words(x -> shfl(x, lane), val) +@inline shfl_down(val::UInt64, offset::Integer) = shfl_words(x -> shfl_down(x, offset), val) +@inline shfl_up(val::UInt64, offset::Integer) = shfl_words(x -> shfl_up(x, offset), val) +@inline shfl_xor(val::UInt64, mask::Integer) = shfl_words(x -> shfl_xor(x, mask), val) + +# The expression that shuffles `ex::S` field by field, calling `f` on the primitive fields +function shfl_fields_expr(S, ex) + isprimitivetype(S) && return :(f($ex)) + fields = (shfl_fields_expr(fieldtype(S, i), :(getfield($ex, $i))) for i in 1:fieldcount(S)) + return Expr(:new, S, fields...) +end + +# Shuffle a value that the backend doesn't support directly field by field. Nested fields are +# unrolled here, rather than shuffled with a recursive call, which inference gives up on +# (on Julia 1.10), so that `f` is only called on the primitive types. +@inline @generated function shfl_fields(f, val::T) where {T} + isbitstype(T) || return :(shfl_unsupported($T)) + return shfl_fields_expr(T, :val) +end + +# The shuffles within segments of `width` lanes, implemented with `shfl` from a lane. + +""" + shfl(val::T, lane::Integer, width::Integer)::T + shfl_down(val::T, offset::Integer, width::Integer)::T + shfl_up(val::T, offset::Integer, width::Integer)::T + shfl_xor(val::T, mask::Integer, width::Integer)::T + +Shuffles within segments of `width` consecutive lanes of the sub-group, with CUDA's semantics +for a `width`: + +- `shfl` reads from lane `lane` of the caller's segment, with `lane` taken modulo `width` + (lane `width + 1` is the segment's first lane); +- `shfl_down` and `shfl_up` read from the lane `offset` lanes further or earlier, and return + `val` of the work-item itself where that lane is outside of the caller's segment; +- `shfl_xor` reads from the lane whose 0-based position in the segment is the caller's xor + `mask`, which has to be between 0 and `width - 1`. + +Reading from a lane of the segment that has no work-item (in a partial sub-group) gives an +unspecified value. `width` has to be a power of two of at most the sub-group width +[`get_max_sub_group_size`](@ref), and the same for all work-items of the sub-group. + +The `width` only changes which lanes are read from: as for the shuffles without a `width`, all +work-items of the sub-group have to execute them together, not just those of a segment. + +!!! note + Backends **may** implement these, e.g. if they have native shuffles with a width. The + fallbacks use [`shfl`](@ref) from a lane. +""" +@inline function shfl(val, lane::Integer, width::Integer) + l0 = get_sub_group_local_id(Int32) - Int32(1) + w = width % Int32 + base = l0 & ~(w - Int32(1)) + # `lane - 1` modulo the power of two `width`, before narrowing `lane` + return shfl(val, base + ((lane - 1) & (width - 1)) % Int32 + Int32(1)) +end + +# The offsets are compared with the width before they are narrowed to `Int32`, so that large +# offsets read from outside of the segment. +@inline function shfl_down(val, offset::Integer, width::Integer) + l0 = get_sub_group_local_id(Int32) - Int32(1) + w = width % Int32 + pos = l0 & (w - Int32(1)) + inside = offset < width - pos + src = ifelse(inside, l0 + offset % Int32, l0) + return shfl(val, src + Int32(1)) +end + +@inline function shfl_up(val, offset::Integer, width::Integer) + l0 = get_sub_group_local_id(Int32) - Int32(1) + w = width % Int32 + pos = l0 & (w - Int32(1)) + inside = offset <= pos + src = ifelse(inside, l0 - offset % Int32, l0) + return shfl(val, src + Int32(1)) +end + +@inline function shfl_xor(val, mask::Integer, width::Integer) + l0 = get_sub_group_local_id(Int32) - Int32(1) + return shfl(val, (l0 ⊻ (mask % Int32)) + Int32(1)) +end + +""" + sub_group_any(pred::Bool)::Bool + +Whether `pred` is `true` for any work-item of the sub-group. All work-items of the sub-group +get the same result. + +All work-items of the sub-group have to execute `sub_group_any` together (not in a divergent +branch); see [`shfl`](@ref). + +It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref). + +!!! note + Backend implementations that support sub-groups **must** implement: + ``` + @device_override sub_group_any(pred::Bool)::Bool + ``` +""" +function sub_group_any end + +""" + sub_group_all(pred::Bool)::Bool + +Whether `pred` is `true` for all work-items of the sub-group. All work-items of the +sub-group get the same result. + +All work-items of the sub-group have to execute `sub_group_all` together (not in a divergent +branch); see [`shfl`](@ref). + +It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref). + +!!! note + Backend implementations that support sub-groups **must** implement: + ``` + @device_override sub_group_all(pred::Bool)::Bool + ``` +""" +function sub_group_all end + +""" + sub_group_ballot(pred::Bool)::UInt64 + +A mask of the work-items of the sub-group for which `pred` is `true`: bit `i - 1` (counting +from the least significant bit) is set if and only if the sub-group has a work-item with +[`get_sub_group_local_id`](@ref) equal to `i` and its `pred` is `true`. All other bits, e.g. +those of lanes without a work-item in a partial sub-group, are zero. All work-items of the +sub-group get the same result. Use e.g. `count_ones` to count the work-items, or +`trailing_zeros` to find the first one. + +All work-items of the sub-group have to execute `sub_group_ballot` together (not in a +divergent branch); see [`shfl`](@ref). Only sub-groups of at most 64 work-items are supported. -`shfl_down` exchanges values, not memory: it is not a memory fence. +It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref). !!! note - Backend implementations **must** implement this for every `T` for which - [`supports_shuffle`](@ref) returns `true`: + Backend implementations that support sub-groups with a width of at most 64 **must** + implement: ``` - @device_override shfl_down(val::T, offset::Integer) where T + @device_override sub_group_ballot(pred::Bool)::UInt64 ``` """ -function shfl_down end +function sub_group_ballot end ## synchronization @@ -306,7 +595,8 @@ Like [`barrier`](@ref), for the work-items of a sub-group: wait until all work-i sub-group have reached the barrier, and make their writes to global and local memory before it visible to the sub-group. -All work-items of a sub-group have to reach the same `sub_group_barrier()`. +All work-items of a sub-group have to reach the same `sub_group_barrier()`; see +[`shfl`](@ref). !!! note Backend implementations that support sub-groups **must** implement: diff --git a/lib/KernelInterface/src/host.jl b/lib/KernelInterface/src/host.jl index 1c3ff0573..848d1b55b 100644 --- a/lib/KernelInterface/src/host.jl +++ b/lib/KernelInterface/src/host.jl @@ -260,13 +260,17 @@ supports_float64(::Backend) = false """ supports_subgroups(::Backend)::Bool -Whether kernels on the active device support sub-groups: the sub-group queries -([`get_sub_group_size`](@ref) etc.), [`sub_group_barrier`](@ref), and a fixed sub-group -width [`sub_group_size`](@ref). See the manual for what KernelInterface guarantees about -how work-groups are divided into sub-groups; a backend that can't ensure that reports -`false`. - -Which types [`shfl_down`](@ref) supports is queried separately with [`supports_shuffle`](@ref). +Whether kernels on the active device support sub-groups with a fixed width +[`sub_group_size`](@ref): the sub-group queries ([`get_sub_group_size`](@ref) etc.), +[`sub_group_barrier`](@ref), and the votes [`sub_group_any`](@ref), [`sub_group_all`](@ref) +and, for widths of at most 64, [`sub_group_ballot`](@ref). Which types the shuffles support +is queried separately with [`supports_shuffle`](@ref). + +Which work-items form a sub-group is unspecified, unless +[`supports_linear_subgroups`](@ref) returns `true`, and the communication functions have to +be executed in control flow that is uniform over the work-group, unless +[`supports_independent_subgroups`](@ref) returns `true`. See the manual for what +KernelInterface guarantees about sub-groups. !!! note Backend implementations **must** implement this function if they support sub-groups. @@ -275,15 +279,79 @@ Which types [`shfl_down`](@ref) supports is queried separately with [`supports_s supports_subgroups(::Backend) = false """ - supports_shuffle(::Backend, ::Type{T})::Bool + supports_linear_subgroups(::Backend)::Bool + +Whether sub-groups are formed from consecutive work-items on the active device, for every +kernel and launch: if the work-group is 1-D, or its x extent is a multiple of the sub-group +width `W` ([`sub_group_size`](@ref)), the work-item with the linear local index `lin` (x +fastest, see [`get_local_id`](@ref)) is in sub-group `fld(lin - 1, W) + 1`, at lane +`mod(lin - 1, W) + 1`, the sub-group ids are `1:cld(n, W)` for a work-group of `n` +work-items, and only the last sub-group can be partial. In particular, a 1-D work-group of at +most `W` work-items is a single sub-group. Implies [`supports_subgroups`](@ref). -Whether kernels on the active device support [`shfl_down`](@ref) for values of type `T`. +Code needs this when it relates local ids to lanes: e.g. stencils that shuffle values between +neighbouring work-items, transposes through shuffles, or work-group scans built from +sub-group scans. Code that doesn't can index storage by the sub-group slot +`(get_sub_group_id() - 1) * W + get_sub_group_local_id()`, which is unique within the +work-group, but not dense: it can exceed the work-group size, so storage has to have room for +`get_num_sub_groups() * W` values. + +`false` means that the layout isn't guaranteed, not that it is different: OpenCL, Vulkan and +other APIs leave it to the implementation. !!! note - Backend implementations **must** implement this function for the types they support. + Backend implementations **should** implement this function if they guarantee the layout. The fallback returns `false`. """ -supports_shuffle(::Backend, ::Type) = false +supports_linear_subgroups(::Backend) = false + +""" + supports_independent_subgroups(::Backend)::Bool + +Whether the sub-group communication functions (shuffles, votes, collectives and +[`sub_group_barrier`](@ref)) can be executed in control flow that is uniform over the +sub-group, but not over the work-group: e.g. a loop whose trip count differs between +sub-groups, or after an early `return` of whole sub-groups. Otherwise, all sub-groups of a +work-group have to execute the same communication functions, in the same order. + +This covers physical sub-groups, not the segments of the shuffles with a `width`, and doesn't +promise forward progress: a sub-group spin-waiting for another one can still hang. Implies +[`supports_subgroups`](@ref). + +!!! note + Backend implementations **should** implement this function if sub-groups are + independent. The fallback returns `false`. +""" +supports_independent_subgroups(::Backend) = false + +""" + supports_shuffle(::Backend, ::Type{T})::Bool + +Whether kernels on the active device support the shuffles [`shfl`](@ref), +[`shfl_down`](@ref), [`shfl_up`](@ref) and [`shfl_xor`](@ref) for values of type `T`. + +Other primitive types of 1, 2, 4, 8 or 16 bytes are supported if `UInt32` is, and `isbits` +structs and tuples if all of their fields are (and, for those without fields, if `UInt32` +is). + +This reports whether values of type `T` can be moved between work-items, not whether the +device can compute with them: e.g. `Float64` values can be shuffled (as words) on a device +without [`supports_float64`](@ref), as long as they are only stored and loaded, e.g. through +an array of `UInt64`. + +!!! note + Backend implementations **must** implement this function for the primitive types they + support natively, which have to include `UInt32`, with a signature that only matches + those, e.g. `supports_shuffle(::NewBackend, ::Type{<:Union{UInt32, Float32}})`. The + fallback handles other types. +""" +function supports_shuffle(backend::Backend, ::Type{T}) where {T} + isprimitivetype(T) && return shuffle_as_words(T) && supports_shuffle(backend, UInt32) + (isbitstype(T) && !isprimitivetype(T)) || return false + # not `true` for a type without fields on a backend without shuffles + fieldcount(T) == 0 && return supports_shuffle(backend, UInt32) + return all(i -> supports_shuffle(backend, fieldtype(T, i)), 1:fieldcount(T)) +end """ allocate(::Backend, Type, dims...; unified=false)::AbstractArray diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl index 210b77123..9b5b1537d 100644 --- a/lib/KernelInterface/test/interface.jl +++ b/lib/KernelInterface/test/interface.jl @@ -189,7 +189,7 @@ end function subgroup_typecheck_kernel(results, val::T) where {T} # uniformly executed by the whole sub-group, as `shfl_down` requires shuffled = KI.shfl_down(val, 1) - if KI.get_sub_group_local_id() == 1 + if KI.get_local_id().x == 1 @inbounds begin results[1] = KI.get_sub_group_size() isa Int results[2] = KI.get_max_sub_group_size() isa Int @@ -207,57 +207,274 @@ function subgroup_typecheck_kernel(results, val::T) where {T} return end -function shfl_down_test_kernel(a, b, ::Val{N}) where {N} - idx = KI.get_sub_group_local_id() +# the sub-group and lane of every work-item, by its linear index (x fastest) +function sub_group_layout_kernel(out) + l = KI.get_local_id() + sz = KI.get_local_size() + lin = l.x + (l.y - 1) * sz.x + (l.z - 1) * sz.x * sz.y + @inbounds begin + out[1, lin] = KI.get_sub_group_id() + out[2, lin] = KI.get_sub_group_local_id() + out[3, lin] = KI.get_sub_group_size() + end + return +end + +# The tests of the communication functions below don't assume which work-items form a +# sub-group (unless they are about `supports_linear_subgroups`): the kernels run in 1-D +# work-groups, every work-item records its sub-group, lane and sub-group size at its local +# index, and the expected results are computed on the host for the sub-groups that formed. + +@inline function record_sub_group!(ids, i) + @inbounds begin + ids[1, i] = KI.get_sub_group_id() + ids[2, i] = KI.get_sub_group_local_id() + ids[3, i] = KI.get_sub_group_size() + end + return +end + +# every work-item applies `f`, which communicates within the sub-group, to its value of `a` +function sub_group_apply_kernel(ids, out, a, f) + i = KI.get_local_id().x + record_sub_group!(ids, i) + @inbounds out[i] = f(a[i]) + return +end + +# the work-items of every sub-group by lane, from the ids `record_sub_group!` wrote +function observed_sub_groups(ids) + groups = [Int[] for _ in 1:maximum(ids[1, :])] + for i in axes(ids, 2) + push!(groups[ids[1, i]], i) + end + for g in groups + sort!(g; by = i -> ids[2, i]) + @test [ids[2, i] for i in g] == 1:length(g) + @test all(i -> ids[3, i] == length(g), g) + end + return groups +end + +# Run `f` on the values `a` in a 1-D work-group of `length(a)` work-items, with results of +# type `T`. Returns the results by local index and the work-items of each sub-group, or +# `nothing` if the work-group is too large for the kernel. +function sub_group_apply(backend, AT, f, a, T = eltype(a)) + n = length(a) + ids = AT(zeros(Int, 3, n)) + out = AT(Vector{T}(undef, n)) + dev_a = AT(a) + kernel = KI.@launch backend launch = false sub_group_apply_kernel(ids, out, dev_a, f) + n <= KI.max_work_group_size(kernel) || return nothing + kernel(ids, out, dev_a, f; workgroupsize = n) + KI.synchronize(backend) + return Array(out), observed_sub_groups(Array(ids)) +end + +# work-group sizes that form full and partial sub-groups, several of them, and a single +# work-item +sub_group_test_sizes(sg_size) = unique((sg_size, max(sg_size - 3, 1), 2 * sg_size + 5, 1)) + +# The shuffles, as callables that are the same type for every offset, so that a kernel +# is compiled once for all of them. +struct Shuffle{F, A} + f::F + arg::A +end +(s::Shuffle)(x) = s.f(x, s.arg) + +struct WidthShuffle{F, A} + f::F + arg::A + width::Int +end +(s::WidthShuffle)(x) = s.f(x, s.arg, s.width) + +# read from the lane `shift` further, wrapping around within the sub-group +struct Rotate + shift::Int +end +(r::Rotate)(x) = KI.shfl(x, mod1(KI.get_sub_group_local_id() + r.shift, KI.get_sub_group_size())) + +# apply `f` to values of type `T`, stored as values of another type of the same size, so +# that types the device can't compute with (e.g. `Float64`) are only moved +struct AsType{T, F} + f::F +end +AsType{T}(f) where {T} = AsType{T, typeof(f)}(f) +(s::AsType{T})(x) where {T} = reinterpret(typeof(x), s.f(reinterpret(T, x))) + +# The lane a shuffle reads from, given the lane and size of the sub-group of the work-item +# and the width `W`, or `nothing` where the result is unspecified. +in_sub_group(src, size) = src <= size ? src : nothing +function shuffle_source(s::Shuffle, l, size, W) + arg = s.arg + src = s.f === KI.shfl ? arg : + s.f === KI.shfl_down ? (arg > W - l ? l : l + arg) : + s.f === KI.shfl_up ? (l > arg ? l - arg : l) : + ((l - 1) ⊻ arg) + 1 + return in_sub_group(src, size) +end +function shuffle_source(s::WidthShuffle, l, size, W) + arg, w = s.arg, s.width + pos = (l - 1) % w + src = s.f === KI.shfl ? l - 1 - pos + mod1(arg, w) : + s.f === KI.shfl_down ? (arg < w - pos ? l + arg : l) : + s.f === KI.shfl_up ? (pos >= arg ? l - arg : l) : + ((l - 1) ⊻ arg) + 1 + return in_sub_group(src, size) +end +shuffle_source(r::Rotate, l, size, W) = mod1(l + r.shift, size) +shuffle_source(s::AsType, l, size, W) = shuffle_source(s.f, l, size, W) + +function shuffle_testsuite(backend, AT, sg_size, f, a) + res = sub_group_apply(backend, AT, f, a) + res === nothing && return + out, groups = res + ok = true + for g in groups, (l, i) in enumerate(g) + src = shuffle_source(f, l, length(g), sg_size) + src === nothing && continue + ok &= out[i] === a[g[src]] + end + @test ok + return +end + +# the shuffles that are tested for every supported type, `wrap` turning them into the +# shuffle that is run +function basic_shuffles(sg_size) + return [ + Rotate(1), Rotate(sg_size - 1), Shuffle(KI.shfl, 1), + Shuffle(KI.shfl_down, 1), Shuffle(KI.shfl_up, 1), Shuffle(KI.shfl_xor, 1), + ] +end - val = a[idx] +function all_shuffles(sg_size) + shuffles = Any[Rotate(0), Rotate(1), Rotate(sg_size - 1)] + for lane in unique((1, 2, sg_size, sg_size + 1)) + push!(shuffles, Shuffle(KI.shfl, lane)) + end + # offsets past the width, also ones that don't fit in 32 bits + for d in unique((0, 1, 3, sg_size ÷ 2, sg_size - 1, sg_size, sg_size + 1, Int64(2)^32 + 1, typemax(Int64))) + push!(shuffles, Shuffle(KI.shfl_down, d), Shuffle(KI.shfl_up, d)) + end + for mask in unique((0, 1, 3, sg_size ÷ 2, sg_size - 1)) + 0 <= mask < sg_size && push!(shuffles, Shuffle(KI.shfl_xor, mask)) + end + for w in (1, 2, 4, 8, 16, 32, 64) + w <= sg_size && ispow2(sg_size) || continue + for arg in unique((1, 3, w - 1, w + 2, Int64(2)^32 + 1, typemax(Int64))) + push!(shuffles, WidthShuffle(KI.shfl, arg, w)) + push!(shuffles, WidthShuffle(KI.shfl_down, arg, w)) + push!(shuffles, WidthShuffle(KI.shfl_up, arg, w)) + end + for mask in unique((0, 1, w - 1)) + 0 <= mask < w && push!(shuffles, WidthShuffle(KI.shfl_xor, mask, w)) + end + end + return shuffles +end - # the result of shuffling from a lane past the end is unspecified, so don't add it - offset = 1 - while offset < N - shuffled = KI.shfl_down(val, offset) - if idx + offset <= N - val += shuffled +function vote_testsuite(backend, AT, sg_size, pred) + for (vote, T) in ((KI.sub_group_any, Bool), (KI.sub_group_all, Bool), (KI.sub_group_ballot, UInt64)) + vote === KI.sub_group_ballot && sg_size > 64 && continue + res = sub_group_apply(backend, AT, vote, pred, T) + res === nothing && continue + out, groups = res + ok = true + for g in groups + p = pred[g] + expected = vote === KI.sub_group_any ? any(p) : + vote === KI.sub_group_all ? all(p) : + reduce(|, (UInt64(1) << (l - 1) for l in eachindex(p) if p[l]); init = UInt64(0)) + ok &= all(i -> out[i] == expected, g) end - offset <<= 1 + @test ok end + return +end +# Every sub-group executes the communication functions a different number of times, and +# some return early, which needs `supports_independent_subgroups`. +function independent_sub_groups_kernel(ids, out, a) + i = KI.get_local_id().x + record_sub_group!(ids, i) + sg = KI.get_sub_group_id() + next = mod1(KI.get_sub_group_local_id() + 1, KI.get_sub_group_size()) + val = @inbounds a[i] + acc = zero(val) + for _ in 1:sg + acc += KI.shfl(val, next) + end + @inbounds out[i] = acc + iseven(sg) && return KI.sub_group_barrier() - - if idx == 1 - b[idx] = val + if KI.sub_group_any(val > Int32(1000)) + acc = -acc end + @inbounds out[i] = acc + KI.shfl(val, 1) return end -# Every lane shuffles its lane id down by `offset`; `out` gets the lane id, the sub-group -# size and the result. -function shfl_down_lanes_kernel(out, ::Type{T}, offset) where {T} - lane = KI.get_sub_group_local_id() - shuffled = KI.shfl_down(T(lane), offset) - i = KI.get_global_id().x - @inbounds begin - out[i, 1] = lane - out[i, 2] = KI.get_sub_group_size() - out[i, 3] = shuffled +function independent_sub_groups_testsuite(backend, AT, sg_size) + n = 4 * sg_size + a = Int32.(rand(1:2000, n)) + ids, out = AT(zeros(Int, 3, n)), AT(zeros(Int32, n)) + kernel = KI.@launch backend launch = false independent_sub_groups_kernel(ids, out, AT(a)) + n = min(n, KI.max_work_group_size(kernel)) + kernel(ids, out, AT(a); workgroupsize = n) + KI.synchronize(backend) + out = Array(out) + groups = observed_sub_groups(Array(ids)[:, 1:n]) + for (sg, g) in enumerate(groups), (l, i) in enumerate(g) + expected = Int32(sg) * a[g[mod1(l + 1, length(g))]] + if isodd(sg) + any(>(1000), a[g]) && (expected = -expected) + expected += a[g[1]] + end + @test out[i] == expected end return end -# Every lane writes local and global memory, and reads another lane's write after a -# sub-group barrier. `N` is the sub-group width; the work-group is one sub-group. -function sub_group_barrier_kernel(scratch, out, ::Val{N}) where {N} +# Every work-item writes local and global memory at its sub-group slot, and reads the next +# lane's write after a sub-group barrier. `N` bounds the slots: the work-group size times the +# width. +function sub_group_barrier_kernel(ids, scratch, out, ::Val{N}) where {N} + i = KI.get_local_id().x + record_sub_group!(ids, i) lane = KI.get_sub_group_local_id() - other = mod1(lane + 1, KI.get_sub_group_size()) + base = (KI.get_sub_group_id() - 1) * KI.get_max_sub_group_size() + other = base + mod1(lane + 1, KI.get_sub_group_size()) lm = KI.localmemory(Int32, N) - @inbounds lm[lane] = lane - @inbounds scratch[lane] = -lane + @inbounds lm[base + lane] = i + @inbounds scratch[base + lane] = -i KI.sub_group_barrier() - @inbounds out[lane, 1] = lm[other] - @inbounds out[lane, 2] = scratch[other] + @inbounds out[i, 1] = lm[other] + @inbounds out[i, 2] = scratch[other] return end +struct ShuffleStruct + a::Float32 + b::Int64 + c::NTuple{3, Int32} +end + +# primitive types that no backend supports natively, shuffled as `UInt32` words +primitive type Bits64 64 end +primitive type Bits16 16 end +Bits64(x::Integer) = reinterpret(Bits64, x % UInt64) +Bits16(x::Integer) = reinterpret(Bits16, x % UInt16) + +struct FallbackStruct + flag::Bool + c::Char + x::Bits64 + limbs::NTuple{4, Int64} +end + # a kernel whose callable captures an array, compiled but not launched yet function captured_array_kernel(backend, AT, out) a = AT(Int32[42]) @@ -265,6 +482,124 @@ function captured_array_kernel(backend, AT, out) return kernel, WeakRef(a) end +# 1-D work-groups and ones whose x extent is a multiple of the sub-group width form +# sub-groups from consecutive work-items, x fastest +function sub_group_layout_testsuite(backend, AT, sg_size, fits) + # including one with a partial last sub-group past 256 work-items, which catches 8-bit + # arithmetic in the index computations + shapes = ( + (sg_size + 5,), (3 * sg_size,), (9 * sg_size + 5,), + (sg_size, 4), (2 * sg_size, 2), (sg_size, 2, 2), + ) + for dims in shapes + n = prod(dims) + out = AT(zeros(Int, 3, n)) + kernel = KI.@launch backend launch = false sub_group_layout_kernel(out) + if !fits(kernel, dims) + @test_skip "work-groups of $dims work-items" + continue + end + kernel(out; workgroupsize = dims) + KI.synchronize(backend) + out = Array(out) + @testset "$dims" begin + @test out[1, :] == [(lin - 1) ÷ sg_size + 1 for lin in 1:n] + @test out[2, :] == [(lin - 1) % sg_size + 1 for lin in 1:n] + @test out[3, :] == [min(sg_size, n - (lin - 1) ÷ sg_size * sg_size) for lin in 1:n] + end + end + return +end + +# The sub-group communication functions. A separate function, so that `interface_testsuite` +# doesn't get too large to compile. +function subgroup_communication_testsuite(backend::KI.Backend, AT, sg_size) + sizes = sub_group_test_sizes(sg_size) + fp64 = KI.supports_float64(backend) + + @testset "sub_group_barrier" begin + n = sg_size + ids = AT(zeros(Int, 3, n)) + scratch = AT(zeros(Int32, n * sg_size)) + out = AT(zeros(Int32, n, 2)) + KI.@launch backend workgroupsize = n sub_group_barrier_kernel(ids, scratch, out, Val(n * sg_size)) + KI.synchronize(backend) + out = Array(out) + for g in observed_sub_groups(Array(ids)), (l, i) in enumerate(g) + other = g[mod1(l + 1, length(g))] + @test out[i, :] == [other, -other] + end + end + + @testset "votes" begin + @testset "$n work-items" for n in sizes + for pred in (falses(n), trues(n), [i % 3 == 1 for i in 1:n], [i == n for i in 1:n], rand(Bool, n)) + vote_testsuite(backend, AT, sg_size, collect(pred)) + end + end + end + + KI.supports_shuffle(backend, UInt32) || return + + @testset "shuffles" begin + @testset "$n work-items" for n in sizes + a = Int32.(1:n) .* Int32(10) + for f in all_shuffles(sg_size) + shuffle_testsuite(backend, AT, sg_size, f, a) + end + end + + # other types, natively or as words or fields + values = Any[ + Int8 => i -> Int8(i % 100), UInt8 => i -> UInt8(i), Int16 => i -> Int16(-i), + UInt16 => i -> (1000i) % UInt16, UInt32 => i -> UInt32(i) << 20, Float16 => Float16, + Int64 => i -> (Int64(i) << 40) - i, UInt64 => i -> typemax(UInt64) - i, + Float32 => i -> Float32(i) / 3, Bool => isodd, Char => i -> Char(0x0001F600 + i), + Bits64 => i -> Bits64((UInt64(i) << 40) - i), Bits16 => i -> Bits16(0xa000 + i), + ShuffleStruct => i -> ShuffleStruct(i, -i, (i, 2i, 3i)), + FallbackStruct => i -> FallbackStruct(isodd(i), Char(64 + i), Bits64(-i), (i, -i, 2i, typemin(Int64) + i)), + ] + fp64 && push!(values, Float64 => i -> 1 / i) + @testset "$T" for (T, f) in values + KI.supports_shuffle(backend, T) || continue + @testset "$n work-items" for n in sizes + a = T[f(i) for i in 1:n] + for s in basic_shuffles(sg_size) + shuffle_testsuite(backend, AT, sg_size, s, a) + end + end + end + + # `Float64` values are moved through `UInt64` storage, so that this also works where + # the device can't compute with them + @testset "Float64 as UInt64" begin + if KI.supports_shuffle(backend, Float64) + a = [reinterpret(UInt64, 1 / i) for i in 1:sg_size] + for s in basic_shuffles(sg_size) + shuffle_testsuite(backend, AT, sg_size, AsType{Float64}(s), a) + end + end + end + + # word and field fallbacks are derived from `UInt32` + @test KI.supports_shuffle(backend, Bool) + @test KI.supports_shuffle(backend, Char) + @test KI.supports_shuffle(backend, Bits64) + @test KI.supports_shuffle(backend, Bits16) + @test KI.supports_shuffle(backend, FallbackStruct) + @test KI.supports_shuffle(backend, ShuffleStruct) == + all(S -> KI.supports_shuffle(backend, S), (Float32, Int64, Int32)) + @test !KI.supports_shuffle(backend, Ref{Int}) + end + + if KI.supports_independent_subgroups(backend) + @testset "independent sub-groups" begin + independent_sub_groups_testsuite(backend, AT, sg_size) + end + end + return +end + function interface_testsuite(backend::KI.Backend, AT) @testset "Launch parameters" begin # unequal group counts and sizes in every dimension, so that confusing them shows @@ -392,6 +727,11 @@ function interface_testsuite(backend::KI.Backend, AT) @test KI.supports_atomics(b) isa Bool @test KI.supports_float64(b) isa Bool @test KI.supports_subgroups(b) isa Bool + @test KI.supports_linear_subgroups(b) isa Bool + @test KI.supports_independent_subgroups(b) isa Bool + # both imply sub-group support + KI.supports_linear_subgroups(b) && @test KI.supports_subgroups(b) + KI.supports_independent_subgroups(b) && @test KI.supports_subgroups(b) @test KI.supports_shuffle(b, Float32) isa Bool @test KI.functional(b) isa Union{Missing, Bool} @test KI.multiprocessor_count(b) isa Int @@ -594,6 +934,12 @@ function interface_testsuite(backend::KI.Backend, AT) # whether `kernel` can be launched with work-groups of size `dims` fits(kernel, dims) = all(dims .<= max_dims[1:length(dims)]) && prod(dims) <= KI.max_work_group_size(kernel) + @testset "Sub-group layout" begin + if KI.supports_linear_subgroups(backend) + sub_group_layout_testsuite(backend, AT, sg_size, fits) + end + end + @testset "Sub-group return types" begin @test sg_size isa Int && sg_size >= 1 @@ -632,8 +978,8 @@ function interface_testsuite(backend::KI.Backend, AT) @test all(d -> d.sub_group_size == size, members) @test sort(map(d -> d.sub_group_local_id, members)) == 1:size end - # a 1-D work-group that fits a sub-group is one - if length(dims) == 1 && items <= sg_size + # with the linear layout, a 1-D work-group that fits a sub-group is one + if KI.supports_linear_subgroups(backend) && length(dims) == 1 && items <= sg_size @test n == 1 end return @@ -682,57 +1028,7 @@ function interface_testsuite(backend::KI.Backend, AT) end end - @testset "sub_group_barrier" begin - out = KI.zeros(backend, Int32, sg_size, 2) - scratch = KI.zeros(backend, Int32, sg_size) - kernel = KI.@launch backend launch = false sub_group_barrier_kernel(scratch, out, Val(sg_size)) - if fits(kernel, (sg_size,)) - kernel(scratch, out, Val(sg_size); workgroupsize = sg_size) - KI.synchronize(backend) - other = mod1.(2:(sg_size + 1), sg_size) - @test Array(out) == hcat(other, -other) - else - @test_skip "work-groups of $sg_size work-items" - end - end - - @testset "shfl_down" begin - candidates = ( - Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64, - Float16, Float32, Float64, - ) - types = filter(T -> KI.supports_shuffle(backend, T), candidates) - @testset "$T" for T in types - a = zeros(T, sg_size) - rand!(a, (0:1)) - dev_a = AT(a) - dev_b = AT(zeros(T, sg_size)) - KI.@launch backend workgroupsize = sg_size shfl_down_test_kernel(dev_a, dev_b, Val(sg_size)) - KI.synchronize(backend) - @test sum(a) ≈ Array(dev_b)[1] - - # every lane whose source is in range gets the source's value - @testset "offset $offset" for offset in unique((1, 3, sg_size ÷ 2)) - 1 <= offset < sg_size || continue - N = 2 * sg_size - out = KI.zeros(backend, T, N, 3) - kernel = KI.@launch backend launch = false shfl_down_lanes_kernel(out, T, offset) - # one sub-group if two don't fit in a work-group - fits(kernel, (N,)) || (N = sg_size) - kernel(out, T, offset; workgroupsize = N) - KI.synchronize(backend) - out = Array(out) - in_range = findall(i -> out[i, 1] + offset <= out[i, 2], 1:N) - # a work-group of one sub-group's width is a single sub-group; with more - # work-items, how many are in range depends on the sub-groups' sizes - if N == sg_size - @test length(in_range) == N - offset - end - @test !isempty(in_range) - @test out[in_range, 3] == out[in_range, 1] .+ offset - end - end - end + subgroup_communication_testsuite(backend, AT, sg_size) end return nothing end @@ -751,6 +1047,8 @@ function contract_testsuite(backend::KI.Backend, AT) @test hasmethod(KI.max_num_groups, Tuple{B}) if KI.supports_subgroups(backend) @test hasmethod(KI.sub_group_size, Tuple{B}) + # the device functions are overlays, so they can't be checked here; the sub-group + # testsuite runs them end return end diff --git a/lib/KernelInterface/test/runtests.jl b/lib/KernelInterface/test/runtests.jl index a8e5b8def..1f13a3214 100644 --- a/lib/KernelInterface/test/runtests.jl +++ b/lib/KernelInterface/test/runtests.jl @@ -23,13 +23,18 @@ end @test !occursin("[sources]", toml) end +struct ShuffleBackend <: KI.Backend end +KI.supports_shuffle(::ShuffleBackend, ::Type{Int32}) = true +struct WordShuffleBackend <: KI.Backend end +struct PlainBackend <: KI.Backend end +KI.supports_shuffle(::WordShuffleBackend, ::Type{UInt32}) = true + # NOTE: this runs before the mock backend below defines methods on `argconvert` # and `kernel_function`. @testset "interface stubs" begin # These have no fallback on purpose: a backend that forgets to `@device_override` # them should get a MethodError rather than silently wrong behaviour. stubs = [ - KI.shfl_down, KI.max_work_group_size, KI.max_work_group_dims, KI.max_num_groups, KI.sub_group_size, KI.argconvert, KI.kernel_function, KI.launch, # Host-side stubs: required backend methods with no sensible fallback. @@ -39,6 +44,34 @@ end @test isempty(methods(stub)) end + # The votes have no fallbacks: backends that support sub-groups implement them. + for vote in [KI.sub_group_any, KI.sub_group_all, KI.sub_group_ballot] + @test isempty(methods(vote)) + end + + # The shuffles only have the fallback that shuffles structs field by field, which + # doesn't handle the primitive types a backend has to implement. + for shfl in [KI.shfl, KI.shfl_down, KI.shfl_up, KI.shfl_xor] + @test_throws ArgumentError shfl(1.0f0, 1) + @test_throws ArgumentError shfl((1.0f0, 2), 1) + @test_throws ArgumentError shfl(Ref(1), 1) + end + # primitive types are shuffled as `UInt32` words, if the backend supports `UInt32` + @test KI.supports_shuffle(WordShuffleBackend(), Bool) + @test KI.supports_shuffle(WordShuffleBackend(), Float64) + @test KI.supports_shuffle(WordShuffleBackend(), Tuple{Char, Int128}) + @test !KI.supports_shuffle(WordShuffleBackend(), Ref{Int}) + @test !KI.supports_shuffle(ShuffleBackend(), Float32) + @test !KI.supports_shuffle(ShuffleBackend(), Tuple{Float32, Int}) + @test KI.supports_shuffle(ShuffleBackend(), Int32) + @test KI.supports_shuffle(ShuffleBackend(), Tuple{Int32, NTuple{2, Int32}}) + @test !KI.supports_shuffle(ShuffleBackend(), Tuple{Int32, Float32}) + + # the sub-group capabilities default to unsupported + @test !KI.supports_subgroups(PlainBackend()) + @test !KI.supports_linear_subgroups(PlainBackend()) + @test !KI.supports_independent_subgroups(PlainBackend()) + # The primitive queries take an element type; only the zero-argument form has a # (forwarding) method, and it must reach the typed stub rather than recurse. primitives = [ diff --git a/src/KernelAbstractions.jl b/src/KernelAbstractions.jl index c233b8cd2..4dc013822 100644 --- a/src/KernelAbstractions.jl +++ b/src/KernelAbstractions.jl @@ -755,6 +755,13 @@ export POCLBackend CPU backend that compiles kernels to OpenCL via [POCL](https://portablecl.org/) and executes them on the host. This is the concrete type behind the [`CPU`](@ref) alias. + +!!! warning "Sub-group operations" + PoCL implements sub-group operations (shuffles and votes) with barriers across the whole work-group, so it doesn't support independent + sub-groups (`KernelInterface.supports_independent_subgroups` is `false`): all sub-groups + of a work-group have to execute the same sub-group operations, in the same order, e.g. + not in a loop whose trip count depends on the sub-group, or in a branch that only some + sub-groups take. Otherwise, results are wrong. """ POCLBackend diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index ca6396700..d2f6fa449 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -233,9 +233,19 @@ function cpu_workgroupsize(nitems, threads, max_items, multiple) end KI.supports_subgroups(::POCLBackend) = device_limits().sub_group_size > 0 -function KI.supports_shuffle(backend::POCLBackend, ::Type{T}) where {T} +# PoCL forms sub-groups from consecutive work-items, but executes sub-group operations as +# work-group barriers, so they need work-group-uniform control flow (see `POCLBackend`) +KI.supports_linear_subgroups(backend::POCLBackend) = KI.supports_subgroups(backend) +KI.supports_independent_subgroups(::POCLBackend) = false + +# The types `sub_group_shuffle` supports; other types are shuffled as words or field by +# field. `Float16` and `Float64` are reported as unsupported on devices without `cl_khr_fp16` +# or `cl_khr_fp64`, rather than shuffled as words: the overrides below are selected by type, +# not by device. +const ShuffleTypes = Union{SPIRVIntrinsics.gentypes...} + +function KI.supports_shuffle(backend::POCLBackend, ::Type{T}) where {T <: ShuffleTypes} KI.supports_subgroups(backend) || return false - T in SPIRVIntrinsics.gentypes || return false T === Float64 && return "cl_khr_fp64" in device().extensions T === Float16 && return "cl_khr_fp16" in device().extensions return true @@ -298,8 +308,33 @@ end sub_group_barrier(POCL.LOCAL_MEM_FENCE | POCL.GLOBAL_MEM_FENCE) end -@device_override function KI.shfl_down(val::T, offset::Integer) where {T} - sub_group_shuffle(val, get_sub_group_local_id() + offset) +@device_override KI.shfl(val::T, lane::Integer) where {T <: ShuffleTypes} = + SPIRVIntrinsics.sub_group_shuffle(val, lane) + +# past the sub-group width, `shfl_down` and `shfl_up` return the work-item's own value, which +# `sub_group_shuffle` (like SPIR-V's `OpGroupNonUniformShuffleDown`) leaves undefined +@device_override function KI.shfl_down(val::T, offset::Integer) where {T <: ShuffleTypes} + lane = get_sub_group_local_id() + # compared before adding, so that large offsets don't overflow + inside = offset <= get_max_sub_group_size() - lane + return SPIRVIntrinsics.sub_group_shuffle(val, ifelse(inside, lane + offset, lane)) +end + +@device_override function KI.shfl_up(val::T, offset::Integer) where {T <: ShuffleTypes} + lane = get_sub_group_local_id() + return SPIRVIntrinsics.sub_group_shuffle(val, ifelse(lane > offset, lane - offset, lane)) +end + +@device_override KI.shfl_xor(val::T, mask::Integer) where {T <: ShuffleTypes} = + SPIRVIntrinsics.sub_group_shuffle_xor(val, mask) + +@device_override KI.sub_group_any(pred::Bool) = SPIRVIntrinsics.sub_group_any(pred) + +@device_override KI.sub_group_all(pred::Bool) = SPIRVIntrinsics.sub_group_all(pred) + +@device_override function KI.sub_group_ballot(pred::Bool) + mask = SPIRVIntrinsics.sub_group_ballot(pred) + return UInt64(mask[1].value) | (UInt64(mask[2].value) << 32) end @device_override @inline function KI._print(args...) diff --git a/src/pocl/pocl.jl b/src/pocl/pocl.jl index 1d0b79439..cd0b3b055 100644 --- a/src/pocl/pocl.jl +++ b/src/pocl/pocl.jl @@ -37,8 +37,9 @@ function Session() context = cl.Context(device) sizes = device.max_work_item_size - # POCL can technically support any sub-group size; prefer the common GPU ones - sg_sizes = device.sub_group_sizes + # POCL can technically support any sub-group size; prefer the common GPU ones, and only + # widths of up to 64, so that `KI.sub_group_ballot` (a `UInt64` mask) covers every lane + sg_sizes = filter(<=(64), device.sub_group_sizes) common = filter(in(sg_sizes), [32, 64, 16, sg_sizes...]) limits = (; max_work_group_size = Int(device.max_work_group_size), diff --git a/test/codegen_checks.jl b/test/codegen_checks.jl index 786cdaa71..536881ea6 100644 --- a/test/codegen_checks.jl +++ b/test/codegen_checks.jl @@ -17,6 +17,7 @@ using FileCheck using KernelAbstractions using KernelAbstractions: @atomic using StaticArrays +import KernelAbstractions.KernelInterface as KI using Test import KernelAbstractions.POCL: @device_code_llvm @@ -83,6 +84,14 @@ end @inbounds A[I] = 2 * A[I] end +# a 64-bit primitive type that no backend shuffles natively +primitive type CodegenBits64 64 end + +@kernel function codegen_shfl_bits64(A) + I = @index(Global, Linear) + @inbounds A[I] = reinterpret(UInt64, KI.shfl(reinterpret(CodegenBits64, A[I]), 1)) +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 @@ -269,6 +278,21 @@ end end end + # A 64-bit primitive type that the backend doesn't shuffle natively uses its native 64-bit + # shuffle (`__spirv_GroupNonUniformShuffle` mangled for an `i64` value is `...jmj`, for an + # `i32` value `...jjj`), rather than being split into two 32-bit words. + @testset "shuffles of other primitive types" begin + A = KernelAbstractions.zeros(backend, UInt64, 32) + @test @filecheck implicit_check_not = "GroupNonUniformShufflejjj" begin + @check "define spir_kernel void @{{.*}}gpu_codegen_shfl_bits64" + @check "GroupNonUniformShufflejmj" + @check_not "GroupNonUniformShuffle" + @check "ret void" + @device_code_llvm debuginfo = :none codegen_shfl_bits64(backend, 32)(A, ndrange = 32) + 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/runtests.jl b/test/runtests.jl index d34640414..9590e483f 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -306,6 +306,18 @@ if "cl_khr_fp16" in POCL.device().extensions end end +# the sub-group width kernels are compiled for is the one `KI.sub_group_size` reports +# (PoCL folds `get_max_sub_group_size` to that constant itself, see pocl/pocl#2375) +function max_sub_group_size_kernel(out) + out[1] = KernelAbstractions.KernelInterface.get_max_sub_group_size() + return +end +@testset "POCL sub-group width" begin + out = zeros(Int, 1) + @opencl max_sub_group_size_kernel(out) + @test out[1] == POCL.device_limits().sub_group_size +end + # Julia doesn't turn a splat of more than 32 elements into a direct call, so a launch with # many arguments allocates unless every layer passes them on as a tuple @testset "POCL launch with many arguments" begin From b6903d20bae58c886993a6ea3d0fa8c79bed4ff1 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Sat, 10 Oct 2026 18:53:53 +0200 Subject: [PATCH 2/9] KernelInterface: sub-group reductions and scans Add `sub_group_reduce`, `sub_group_scan` and `sub_group_exclusive_scan`, in the order of the lanes, for associative operators that needn't be commutative. The fallbacks are built on the shuffles. Back ends may implement them natively per operator and type, as long as the result is the documented one, except that floating-point `+` and `*` may be reassociated. The POCL back end uses PoCL's native reductions and scans where they have Julia's semantics: `+` on integers and floats, and `min`/`max` on integers. Co-authored-by: Tim Besard --- docs/src/kernelinterface.md | 12 +- lib/KernelInterface/src/KernelInterface.jl | 1 + lib/KernelInterface/src/device.jl | 132 ++++++++++++++++++- lib/KernelInterface/test/interface.jl | 145 +++++++++++++++++++++ src/KernelAbstractions.jl | 3 +- src/pocl/backend.jl | 27 ++++ test/runtests.jl | 29 +++++ 7 files changed, 342 insertions(+), 7 deletions(-) diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md index c613c5a98..8e3351f8d 100644 --- a/docs/src/kernelinterface.md +++ b/docs/src/kernelinterface.md @@ -77,7 +77,7 @@ What a backend implements, at a glance. The docstrings below have the details. | **Capabilities** | | [`supports_float64`](@ref), [`supports_atomics`](@ref), [`supports_unified`](@ref), [`supports_subgroups`](@ref), [`supports_linear_subgroups`](@ref), [`supports_independent_subgroups`](@ref) (all `false`), [`supports_shuffle`](@ref) (`false` unless derived from the natively supported types, see its docstring) | | **Compilation** | [`argconvert`](@ref), [`kernel_function`](@ref), [`launch`](@ref) | | | **Device** | [`get_local_id`](@ref), [`get_group_id`](@ref), [`get_local_size`](@ref), [`get_num_groups`](@ref), [`localmemory`](@ref), [`barrier`](@ref) | [`get_global_id`](@ref), [`get_global_size`](@ref) (derived from the primitive queries), [`_print`](@ref KernelInterface._print) (host `print`) | -| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries (with a constant [`get_max_sub_group_size`](@ref)), [`sub_group_barrier`](@ref), [`sub_group_any`](@ref), [`sub_group_all`](@ref), and [`sub_group_ballot`](@ref) for widths of at most 64; if `supports_shuffle(backend, T)`: [`shfl`](@ref), [`shfl_down`](@ref), [`shfl_up`](@ref), [`shfl_xor`](@ref) for the primitive `T` supported natively, including `UInt32` | shuffles with a `width` (built on [`shfl`](@ref)); shuffles of other primitive types (as `UInt32` or `UInt64` words) and of structs (field by field) | +| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries (with a constant [`get_max_sub_group_size`](@ref)), [`sub_group_barrier`](@ref), [`sub_group_any`](@ref), [`sub_group_all`](@ref), and [`sub_group_ballot`](@ref) for widths of at most 64; if `supports_shuffle(backend, T)`: [`shfl`](@ref), [`shfl_down`](@ref), [`shfl_up`](@ref), [`shfl_xor`](@ref) for the primitive `T` supported natively, including `UInt32` | shuffles with a `width`, [`sub_group_reduce`](@ref), [`sub_group_scan`](@ref), [`sub_group_exclusive_scan`](@ref) (built on the shuffles); shuffles of other primitive types (as `UInt32` or `UInt64` words) and of structs (field by field) | Everything else, such as [`zeros`](@ref KernelInterface.zeros), [`ones`](@ref KernelInterface.ones), the launch-keyword handling of [`Kernel`](@ref) and [`@launch`](@ref KernelInterface.@launch), @@ -168,7 +168,8 @@ one sub-group per work-item, and code combining those values has to use A sub-group is partial when it has fewer work-items than the width: the lanes `get_sub_group_size()+1:W` have no work-item. Shuffles from those lanes give unspecified -values, and the votes only take the work-items of the sub-group into account. +values, and the votes, [`sub_group_reduce`](@ref) and the scans only take the work-items of +the sub-group into account. Backends that guarantee more say so with two queries: @@ -181,8 +182,8 @@ Backends that guarantee more say so with two queries: ids to lanes needs this, e.g. stencils exchanging neighbouring values with shuffles. Code that doesn't, like a work-group reduction combining a value per sub-group, can index by the sub-group slot `(get_sub_group_id() - 1) * W + get_sub_group_local_id()` instead. -- [`supports_independent_subgroups`](@ref): the communication functions (shuffles, votes - and [`sub_group_barrier`](@ref)) can be executed in control flow that is uniform +- [`supports_independent_subgroups`](@ref): the communication functions (shuffles, votes, + collectives and [`sub_group_barrier`](@ref)) can be executed in control flow that is uniform over each sub-group, but differs between the sub-groups of a work-group, e.g. a loop whose trip count depends on the sub-group, or after an early `return` of whole sub-groups. Otherwise, all sub-groups of a work-group have to execute the same communication @@ -245,6 +246,9 @@ shfl(::Any, ::Integer, ::Integer) sub_group_any sub_group_all sub_group_ballot +sub_group_reduce +sub_group_scan +sub_group_exclusive_scan ``` ### Printing diff --git a/lib/KernelInterface/src/KernelInterface.jl b/lib/KernelInterface/src/KernelInterface.jl index 6932ed8ba..8eb260b4b 100644 --- a/lib/KernelInterface/src/KernelInterface.jl +++ b/lib/KernelInterface/src/KernelInterface.jl @@ -33,6 +33,7 @@ include("host.jl") :get_sub_group_id, :get_sub_group_local_id, :localmemory, :shfl, :shfl_down, :shfl_up, :shfl_xor, :sub_group_any, :sub_group_all, :sub_group_ballot, + :sub_group_reduce, :sub_group_scan, :sub_group_exclusive_scan, :barrier, :sub_group_barrier, :_print, # compilation and launch :Kernel, :kernel_function, :argconvert, :launch, Symbol("@launch"), diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl index 09af74ca3..478fb455f 100644 --- a/lib/KernelInterface/src/device.jl +++ b/lib/KernelInterface/src/device.jl @@ -133,8 +133,8 @@ end # manual). # # In a partial sub-group, the lanes `get_sub_group_size()+1:get_max_sub_group_size()` have no -# work-item: shuffles from them give unspecified values, and the votes only take the -# work-items of the sub-group into account. +# work-item: shuffles from them give unspecified values, and the votes, `sub_group_reduce` and +# the scans only take the work-items of the sub-group into account. """ get_sub_group_size([::Type{T}=Int])::T @@ -564,6 +564,134 @@ It exchanges values, not memory: it is not a memory fence, see [`sub_group_barri """ function sub_group_ballot end +""" + sub_group_reduce(op, val::T)::T + +Reduce `val` over the work-items of the sub-group with the binary operator `op`, in the order +of the lanes: `op(…op(op(v₁, v₂), v₃)…, vₙ)` for the values `vᵢ` of lanes `1:n`, with +`n = get_sub_group_size()`, up to associativity. `op` has to be associative, but needn't be +commutative. All work-items of the sub-group get the result. + +All work-items of the sub-group have to execute `sub_group_reduce` together (not in a +divergent branch), with the same `op`; see [`shfl`](@ref). Values of the types that the +shuffles support are supported, see [`supports_shuffle`](@ref), and others for which the +backend has a native reduction. + +It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref). + +!!! note + Backends **may** implement this for operators and types with a native reduction (e.g. + `+` on `Float32`), dispatching on `typeof(op)`. The result has to be the one documented + above, with Julia's semantics of `op` (e.g. NaN propagation and the sign of zero for `min` + and `max`, wrap-around for integers), with one exception: for `+` and `*` on floating-point + values, it may differ as if the values were combined in another order and grouping. The + fallback combines ranges of doubling length with a butterfly of [`shfl_xor`](@ref), and + in a partial sub-group broadcasts the result of the first lane with [`shfl`](@ref). +""" +@inline function sub_group_reduce(op, val) + width = get_max_sub_group_size(Int32) + sgsize = get_sub_group_size(Int32) + lane0 = get_sub_group_local_id(Int32) - Int32(1) + if ispow2(width) + # A butterfly with `shfl_xor`, which unrolls for the constant width. Each step combines + # the block of a work-item with the neighboring block, the lower one first, so that + # only associativity is needed. In a full sub-group, every work-item ends up with the + # reduction; in a partial one, blocks without work-items are skipped, which keeps the + # result of the first lane correct, and it is broadcast. All sub-groups run the same + # shuffles, since some backends (PoCL) need that across the sub-groups of a work-group. + mask = Int32(1) + while mask < width + other = shfl_xor(val, mask) + if (lane0 ⊻ mask) < sgsize + if lane0 & mask == Int32(0) + val = op(val, other) + else + val = op(other, val) + end + end + mask <<= 1 + end + first = shfl(val, 1) + return ifelse(sgsize == width, val, first) + else + # combine ranges of doubling length with `shfl_down`, skipping the lanes without a + # work-item, and broadcast the result of the first lane + lane = lane0 + Int32(1) + offset = Int32(1) + while offset < width + other = shfl_down(val, offset) + if lane + offset <= sgsize + val = op(val, other) + end + offset <<= 1 + end + return shfl(val, 1) + end +end + +""" + sub_group_scan(op, val::T)::T + +The inclusive scan of `val` over the work-items of the sub-group with the associative binary +operator `op`, in the order of the lanes: the work-item in lane `i` gets the reduction of the +values of lanes `1` to `i`, as with [`sub_group_reduce`](@ref). + +All work-items of the sub-group have to execute `sub_group_scan` together (not in a divergent +branch), with the same `op`; see [`shfl`](@ref). Values of the types that the shuffles +support are supported, see [`supports_shuffle`](@ref), and others for which the backend has a +native scan. + +It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref). + +!!! note + Backends **may** implement this for operators and types with a native scan (e.g. `+` on + `Float32`), dispatching on `typeof(op)`, with the results documented above, as for + [`sub_group_reduce`](@ref): for floating-point `+` and `*`, the values of each lane's + prefix may be combined in another order and grouping. The fallback is a Hillis-Steele scan + with [`shfl_up`](@ref). +""" +@inline function sub_group_scan(op, val) + lane = get_sub_group_local_id(Int32) + # loop to the constant width, so that the loop unrolls: the lanes `shfl_up` reads from + # always have a work-item, also in a partial sub-group + width = get_max_sub_group_size(Int32) + offset = Int32(1) + while offset < width + other = shfl_up(val, offset) + if lane > offset + val = op(other, val) + end + offset <<= 1 + end + return val +end + +""" + sub_group_exclusive_scan(op, val::T, init::T)::T + +The exclusive scan of `val` over the work-items of the sub-group with the associative binary +operator `op`, in the order of the lanes, starting from `init`: the work-item in lane 1 gets +`init` itself, and the one in lane `i > 1` gets `op(init, r)` for the reduction `r` of the +values of lanes `1` to `i - 1`, as with [`sub_group_scan`](@ref). `init` has to be the same +for all work-items of the sub-group, and of the type of `val`; it needn't be an identity of +`op`. + +All work-items of the sub-group have to execute `sub_group_exclusive_scan` together (not in a +divergent branch), with the same `op`; see [`shfl`](@ref). Values of the types that the +shuffles support are supported, see [`supports_shuffle`](@ref), and others for which the +backend has a native scan. + +It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref). + +!!! note + Backends **may** implement this like [`sub_group_scan`](@ref). The fallback shifts the + result of [`sub_group_scan`](@ref) up by a lane with [`shfl_up`](@ref). +""" +@inline function sub_group_exclusive_scan(op, val, init) + prefix = shfl_up(sub_group_scan(op, val), 1) + return get_sub_group_local_id(Int32) == Int32(1) ? init : op(init, prefix) +end + ## synchronization diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl index 9b5b1537d..a87172edd 100644 --- a/lib/KernelInterface/test/interface.jl +++ b/lib/KernelInterface/test/interface.jl @@ -376,6 +376,60 @@ function all_shuffles(sg_size) return shuffles end +# The collectives, as callables +struct Reduce{O} + op::O +end +(r::Reduce)(x) = KI.sub_group_reduce(r.op, x) + +struct Scan{O} + op::O +end +(s::Scan)(x) = KI.sub_group_scan(s.op, x) + +struct ExclusiveScan{O, T} + op::O + init::T +end +(s::ExclusiveScan)(x) = KI.sub_group_exclusive_scan(s.op, x, s.init) + +# an associative but not commutative operator: the composition of affine maps +compose_affine(f, g) = (g[1] * f[1], g[1] * f[2] + g[2]) +# the (value, index) of the smallest value, the first one of equal values +argmin_op(x, y) = ifelse(y[1] < x[1], y, x) + +# The expected result of a collective for lane `l` of the sub-group with the values `vals` +collective_result(r::Reduce, vals, l) = foldl(r.op, vals) +collective_result(s::Scan, vals, l) = foldl(s.op, vals[1:l]) +collective_result(s::ExclusiveScan, vals, l) = l == 1 ? s.init : s.op(s.init, foldl(s.op, vals[1:(l - 1)])) + +# Floating-point `+` and `*` may be reassociated, so they are compared with a tolerance +# (other than for the `init` of an exclusive scan, which lane 1 gets unchanged). +collective_op(f) = f.op +reassociable(f, ::Type{T}) where {T} = T <: AbstractFloat && collective_op(f) in (+, *) + +function collective_testsuite(backend, AT, f, a) + res = sub_group_apply(backend, AT, f, a) + res === nothing && return + out, groups = res + T = eltype(a) + ok = true + for g in groups + vals = a[g] + for (l, i) in enumerate(g) + expected = collective_result(f, vals, l) + if reassociable(f, T) && !(f isa ExclusiveScan && l == 1) + ok &= isapprox(out[i], expected; nans = true, rtol = sqrt(eps(T))) + else + # `isequal`, as the results may be NaN or a signed zero + ok &= isequal(out[i], expected) + end + end + end + @test ok + return +end + function vote_testsuite(backend, AT, sg_size, pred) for (vote, T) in ((KI.sub_group_any, Bool), (KI.sub_group_all, Bool), (KI.sub_group_ballot, UInt64)) vote === KI.sub_group_ballot && sg_size > 64 && continue @@ -395,6 +449,36 @@ function vote_testsuite(backend, AT, sg_size, pred) return end +# the values come from a divergent branch, like for the padding work-items of a `@kernel` +function reduce_divergent_kernel(ids, red, scan, a, m) + i = KI.get_local_id().x + record_sub_group!(ids, i) + val = i <= m ? (@inbounds a[i]) : zero(eltype(a)) + r = KI.sub_group_reduce(+, val) + s = KI.sub_group_scan(+, val) + @inbounds red[i] = r + @inbounds scan[i] = s + return +end + +function reduce_divergent_testsuite(backend, AT, sg_size, ::Type{T}) where {T} + n = sg_size + m = max(n - 5, 1) + a = T.(rand(1:20, n)) + ids = AT(zeros(Int, 3, n)) + red = AT(zeros(T, n)) + scan = AT(zeros(T, n)) + KI.@launch backend workgroupsize = n reduce_divergent_kernel(ids, red, scan, AT(a), m) + KI.synchronize(backend) + red, scan = Array(red), Array(scan) + vals = [i <= m ? a[i] : zero(T) for i in 1:n] + for g in observed_sub_groups(Array(ids)) + @test all(i -> red[i] == sum(vals[g]), g) + @test all(l -> scan[g[l]] == sum(vals[g[1:l]]), eachindex(g)) + end + return +end + # Every sub-group executes the communication functions a different number of times, and # some return early, which needs `supports_independent_subgroups`. function independent_sub_groups_kernel(ids, out, a) @@ -592,6 +676,67 @@ function subgroup_communication_testsuite(backend::KI.Backend, AT, sg_size) @test !KI.supports_shuffle(backend, Ref{Int}) end + @testset "sub_group_reduce and scans" begin + @testset "$n work-items" for n in sizes + ops = Any[ + (+, Int32.(rand(1:100, n))), + # wrap-around + (+, Int32[typemax(Int32) - rand(Int32(0):Int32(10)) for _ in 1:n]), + (*, Int32.(rand(-3:3, n))), + # operators and types that backends may implement natively + (+, UInt64.(rand(1:100, n))), (+, Int64.(rand(-100:100, n))), + (+, UInt32.(rand(1:100, n))), + (min, Int64.(rand(-100:100, n))), (min, Int32.(rand(-100:100, n))), + (max, UInt32.(rand(1:100, n))), (max, Int64.(rand(-100:100, n))), + (|, UInt32.(rand(UInt32, n))), (&, UInt32.(rand(UInt32, n))), + (xor, UInt64.(rand(UInt64, n))), + # not commutative + (compose_affine, [(Int32(rand((-1, 1, 2))), Int32(rand(-5:5))) for _ in 1:n]), + (argmin_op, [(Float32(rand(1:20)), Int32(i)) for i in 1:n]), + ] + float_types = Any[Float32] + KI.supports_shuffle(backend, Float16) && push!(float_types, Float16) + fp64 && KI.supports_shuffle(backend, Float64) && push!(float_types, Float64) + for T in float_types + # exactly representable sums for `Float16`, rounding for the others + push!(ops, (+, T === Float16 ? T.(rand(1:8, n)) ./ 4 : T.(rand(n)))) + T === Float16 || push!(ops, (*, T.(rand(0.9:0.01:1.1, n)))) + # NaN and infinities propagate through `+` whatever the order + push!(ops, (+, T[i == 2 ? -T(Inf) : i == 3 ? T(Inf) : rand(1:100) for i in 1:n])) + push!(ops, (+, T[i == 2 ? T(NaN) : rand(1:100) for i in 1:n])) + # Julia's `max` and `min` propagate NaN, unlike OpenCL's + push!(ops, (max, T[i == 2 ? T(NaN) : rand(1:100) for i in 1:n])) + push!(ops, (min, T[i == n ? T(NaN) : rand(1:100) for i in 1:n])) + # ... and tell the sign of zero apart + push!(ops, (min, T[isodd(i) ? -zero(T) : zero(T) for i in 1:n])) + push!(ops, (max, T[isodd(i) ? zero(T) : -zero(T) for i in 1:n])) + end + for (op, a) in ops + KI.supports_shuffle(backend, eltype(a)) || continue + collective_testsuite(backend, AT, Reduce(op), a) + collective_testsuite(backend, AT, Scan(op), a) + end + + # exclusive scans with an `init` that isn't the identity + collective_testsuite(backend, AT, ExclusiveScan(+, Int32(7)), Int32.(rand(1:100, n))) + collective_testsuite(backend, AT, ExclusiveScan(+, UInt64(7)), UInt64.(rand(1:100, n))) + collective_testsuite(backend, AT, ExclusiveScan(max, Int32(50)), Int32.(rand(1:100, n))) + collective_testsuite(backend, AT, ExclusiveScan(+, 7.0f0), Float32.(rand(1:100, n))) + collective_testsuite( + backend, AT, ExclusiveScan(compose_affine, (Int32(2), Int32(3))), + [(Int32(rand((-1, 1, 2))), Int32(rand(-5:5))) for _ in 1:n] + ) + # lane 1 gets `init` itself, not `init + 0.0` + collective_testsuite(backend, AT, ExclusiveScan(+, -0.0f0), Float32.(rand(1:100, n))) + collective_testsuite(backend, AT, ExclusiveScan(+, -0.0f0), fill(-0.0f0, n)) + end + end + + @testset "sub_group_reduce and sub_group_scan of divergent values" begin + reduce_divergent_testsuite(backend, AT, sg_size, Int32) + reduce_divergent_testsuite(backend, AT, sg_size, Float32) + end + if KI.supports_independent_subgroups(backend) @testset "independent sub-groups" begin independent_sub_groups_testsuite(backend, AT, sg_size) diff --git a/src/KernelAbstractions.jl b/src/KernelAbstractions.jl index 4dc013822..f7bf3725f 100644 --- a/src/KernelAbstractions.jl +++ b/src/KernelAbstractions.jl @@ -757,7 +757,8 @@ CPU backend that compiles kernels to OpenCL via [POCL](https://portablecl.org/) them on the host. This is the concrete type behind the [`CPU`](@ref) alias. !!! warning "Sub-group operations" - PoCL implements sub-group operations (shuffles and votes) with barriers across the whole work-group, so it doesn't support independent + PoCL implements sub-group operations (shuffles, votes, and the collectives built on + them) with barriers across the whole work-group, so it doesn't support independent sub-groups (`KernelInterface.supports_independent_subgroups` is `false`): all sub-groups of a work-group have to execute the same sub-group operations, in the same order, e.g. not in a loop whose trip count depends on the sub-group, or in a branch that only some diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index d2f6fa449..93c96565e 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -337,6 +337,33 @@ end return UInt64(mask[1].value) | (UInt64(mask[2].value) << 32) end +# Native reductions and scans of `cl_khr_subgroups`, for `+` on 32- and 64-bit integers and +# floats, and `min`/`max` on integers (OpenCL's `min` and `max` treat NaN and the sign of zero +# differently from Julia's). They need the fix of PoCL 7.2's peeling of the first work-item +# (pocl/pocl#2239), which `pocl_standalone_jll` includes since 7.2.1+1. +const CollectiveIntTypes = Union{Int32, UInt32, Int64, UInt64} +const CollectiveTypes = Union{CollectiveIntTypes, Float16, Float32, Float64} + +@device_override KI.sub_group_reduce(::typeof(+), val::CollectiveTypes) = + SPIRVIntrinsics.sub_group_reduce_add(val) +@device_override KI.sub_group_reduce(::typeof(min), val::CollectiveIntTypes) = + SPIRVIntrinsics.sub_group_reduce_min(val) +@device_override KI.sub_group_reduce(::typeof(max), val::CollectiveIntTypes) = + SPIRVIntrinsics.sub_group_reduce_max(val) + +@device_override KI.sub_group_scan(::typeof(+), val::CollectiveTypes) = + SPIRVIntrinsics.sub_group_scan_inclusive_add(val) +@device_override KI.sub_group_scan(::typeof(min), val::CollectiveIntTypes) = + SPIRVIntrinsics.sub_group_scan_inclusive_min(val) +@device_override KI.sub_group_scan(::typeof(max), val::CollectiveIntTypes) = + SPIRVIntrinsics.sub_group_scan_inclusive_max(val) + +# the exclusive scans of `cl_khr_subgroups` start from the identity, not from `init` +@device_override function KI.sub_group_exclusive_scan(::typeof(+), val::T, init::T) where {T <: CollectiveTypes} + prefix = SPIRVIntrinsics.sub_group_scan_exclusive_add(val) + return ifelse(get_sub_group_local_id() == 1, init, init + prefix) +end + @device_override @inline function KI._print(args...) POCL._print(args...) end diff --git a/test/runtests.jl b/test/runtests.jl index 9590e483f..d96538c0f 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -318,6 +318,35 @@ end @test out[1] == POCL.device_limits().sub_group_size end +# the native sub-group collectives are used where they have Julia's semantics +function sub_group_reduce_kernel(out, x) + i = KernelAbstractions.KernelInterface.get_global_id().x + out[i] = KernelAbstractions.KernelInterface.sub_group_reduce(+, x[i]) + return +end +function sub_group_exclusive_scan_kernel(out, x) + i = KernelAbstractions.KernelInterface.get_global_id().x + out[i] = KernelAbstractions.KernelInterface.sub_group_exclusive_scan(+, x[i], eltype(x)(7)) + return +end +@testset "POCL native sub-group collectives" begin + W = POCL.device_limits().sub_group_size + for T in (Int32, Float32) + x, out = ones(T, W), zeros(T, W) + ir = sprint() do io + @device_code_llvm io = io debuginfo = :none @opencl local_size = W global_size = W sub_group_reduce_kernel(out, x) + end + @test all(==(W), out) + @test occursin("sub_group_reduce_add", ir) + + ir = sprint() do io + @device_code_llvm io = io debuginfo = :none @opencl local_size = W global_size = W sub_group_exclusive_scan_kernel(out, x) + end + @test out == 7 .+ (0:(W - 1)) + @test occursin("sub_group_scan_exclusive_add", ir) + end +end + # Julia doesn't turn a splat of more than 32 elements into a direct call, so a launch with # many arguments allocates unless every layer passes them on as a tuple @testset "POCL launch with many arguments" begin From c497ccd95845ba55f0721bdc9bb387a947e8b840 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sun, 11 Oct 2026 11:07:57 +0200 Subject: [PATCH 3/9] KernelInterface: shuffle structs as packed words Structs and tuples were shuffled field by field, a shuffle per primitive field, so that e.g. an `NTuple{2, Float16}` took two. Pack their fields into `UInt32` words instead, with fields of less than 4 bytes sharing a word, and shuffle the words through `shfl_word_tuple`. A backend can implement that to exchange all words at once: on PoCL, where every shuffle is a work-group barrier, that would help, once its CPU device accepts vector shuffles from SPIR-V. The generated fallbacks also weren't inlined: `@inline` on a generated function doesn't apply to a body that the generator returns with an explicit `return`, so shuffling an `SVector` in a larger kernel left a call, with the value passed through memory. They now mark their bodies for inlining. --- docs/src/kernelinterface.md | 2 +- lib/KernelInterface/src/device.jl | 158 ++++++++++++++++++++------ lib/KernelInterface/src/host.jl | 11 +- lib/KernelInterface/test/interface.jl | 5 +- lib/KernelInterface/test/runtests.jl | 22 +++- src/pocl/backend.jl | 10 +- 6 files changed, 156 insertions(+), 52 deletions(-) diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md index 8e3351f8d..e68ce8cc8 100644 --- a/docs/src/kernelinterface.md +++ b/docs/src/kernelinterface.md @@ -77,7 +77,7 @@ What a backend implements, at a glance. The docstrings below have the details. | **Capabilities** | | [`supports_float64`](@ref), [`supports_atomics`](@ref), [`supports_unified`](@ref), [`supports_subgroups`](@ref), [`supports_linear_subgroups`](@ref), [`supports_independent_subgroups`](@ref) (all `false`), [`supports_shuffle`](@ref) (`false` unless derived from the natively supported types, see its docstring) | | **Compilation** | [`argconvert`](@ref), [`kernel_function`](@ref), [`launch`](@ref) | | | **Device** | [`get_local_id`](@ref), [`get_group_id`](@ref), [`get_local_size`](@ref), [`get_num_groups`](@ref), [`localmemory`](@ref), [`barrier`](@ref) | [`get_global_id`](@ref), [`get_global_size`](@ref) (derived from the primitive queries), [`_print`](@ref KernelInterface._print) (host `print`) | -| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries (with a constant [`get_max_sub_group_size`](@ref)), [`sub_group_barrier`](@ref), [`sub_group_any`](@ref), [`sub_group_all`](@ref), and [`sub_group_ballot`](@ref) for widths of at most 64; if `supports_shuffle(backend, T)`: [`shfl`](@ref), [`shfl_down`](@ref), [`shfl_up`](@ref), [`shfl_xor`](@ref) for the primitive `T` supported natively, including `UInt32` | shuffles with a `width`, [`sub_group_reduce`](@ref), [`sub_group_scan`](@ref), [`sub_group_exclusive_scan`](@ref) (built on the shuffles); shuffles of other primitive types (as `UInt32` or `UInt64` words) and of structs (field by field) | +| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries (with a constant [`get_max_sub_group_size`](@ref)), [`sub_group_barrier`](@ref), [`sub_group_any`](@ref), [`sub_group_all`](@ref), and [`sub_group_ballot`](@ref) for widths of at most 64; if `supports_shuffle(backend, T)`: [`shfl`](@ref), [`shfl_down`](@ref), [`shfl_up`](@ref), [`shfl_xor`](@ref) for the primitive `T` supported natively, including `UInt32` | shuffles with a `width`, [`sub_group_reduce`](@ref), [`sub_group_scan`](@ref), [`sub_group_exclusive_scan`](@ref) (built on the shuffles); shuffles of other primitive types (as `UInt32` or `UInt64` words) and of structs (packed into `UInt32` words) | Everything else, such as [`zeros`](@ref KernelInterface.zeros), [`ones`](@ref KernelInterface.ones), the launch-keyword handling of [`Kernel`](@ref) and [`@launch`](@ref KernelInterface.@launch), diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl index 478fb455f..0afef9f6b 100644 --- a/lib/KernelInterface/src/device.jl +++ b/lib/KernelInterface/src/device.jl @@ -287,9 +287,10 @@ sub-groups of the work-group have to execute it. Types for which [`supports_shuffle`](@ref) returns `true` are supported. Besides the types a backend supports natively, that includes other primitive types of 1, 2, 4, 8 or 16 bytes (e.g. `Bool`, `Char` or `Int64`) if the backend supports `UInt32`, and `isbits` structs and -tuples of supported types, which are shuffled field by field. Primitive types of up to 4 -bytes are shuffled as a `UInt32`, and larger ones as `UInt64` words, each of which is shuffled -natively if the backend supports `UInt64`, and as two `UInt32` words otherwise. +tuples of such types. Primitive types of up to 4 bytes are shuffled as a `UInt32`, and larger +ones as `UInt64` words, each of which is shuffled natively if the backend supports `UInt64`, +and as two `UInt32` words otherwise. Structs and tuples are packed into `UInt32` words, with +fields of less than 4 bytes sharing words, and shuffled word by word. !!! note Backend implementations **must** implement this for the primitive types they support @@ -299,7 +300,7 @@ natively if the backend supports `UInt64`, and as two `UInt32` words otherwise. @device_override shfl(val::T, lane::Integer) where {T <: Union{...}} ``` """ -@inline shfl(val, lane::Integer) = shfl_fallback(x -> shfl(x, lane), val) +@inline shfl(val, lane::Integer) = shfl_fallback(ShuffleOp(shfl, (lane,)), val) """ shfl_down(val::T, offset::Integer)::T @@ -323,7 +324,7 @@ the supported types. @device_override shfl_down(val::T, offset::Integer) where {T <: Union{...}} ``` """ -@inline shfl_down(val, offset::Integer) = shfl_fallback(x -> shfl_down(x, offset), val) +@inline shfl_down(val, offset::Integer) = shfl_fallback(ShuffleOp(shfl_down, (offset,)), val) """ shfl_up(val::T, offset::Integer)::T @@ -345,7 +346,7 @@ the supported types. @device_override shfl_up(val::T, offset::Integer) where {T <: Union{...}} ``` """ -@inline shfl_up(val, offset::Integer) = shfl_fallback(x -> shfl_up(x, offset), val) +@inline shfl_up(val, offset::Integer) = shfl_fallback(ShuffleOp(shfl_up, (offset,)), val) """ shfl_xor(val::T, mask::Integer)::T @@ -370,15 +371,24 @@ the supported types. @device_override shfl_xor(val::T, mask::Integer) where {T <: Union{...}} ``` """ -@inline shfl_xor(val, mask::Integer) = shfl_fallback(x -> shfl_xor(x, mask), val) +@inline shfl_xor(val, mask::Integer) = shfl_fallback(ShuffleOp(shfl_xor, (mask,)), val) -# Shuffle a value of a type that the backend doesn't support natively, with `f` shuffling a +# A shuffle of values of the types that the backend supports natively: `op(x)` is +# `op.f(x, op.args...)`. A struct rather than a closure, so that a backend can dispatch on it in +# `shfl_word_tuple`. +struct ShuffleOp{F, A <: Tuple} + f::F + args::A +end +@inline (op::ShuffleOp)(x) = op.f(x, op.args...) + +# Shuffle a value of a type that the backend doesn't support natively, with `op` shuffling a # value of a type it does support. The fallbacks are separate functions for primitive and for # other types, so that the fallback of a struct with a field the backend doesn't support # natively, e.g. an `Int64` on Metal, doesn't call itself, which inference gives up on (on # Julia 1.10). -@inline function shfl_fallback(f, val::T) where {T} - return isprimitivetype(T) ? shfl_words(f, val) : shfl_fields(f, val) +@inline function shfl_fallback(op, val::T) where {T} + return isprimitivetype(T) ? shfl_words(op, val) : shfl_packed(op, val) end shfl_unsupported(T) = throw( @@ -395,52 +405,130 @@ shuffle_as_words(T) = T !== UInt32 && sizeof(T) in (1, 2, 4, 8, 16) const word_types = (UInt8, UInt16, UInt32, UInt64, UInt128) word_type(T) = word_types[trailing_zeros(sizeof(T)) + 1] +# The generated functions mark their bodies for inlining: `@inline` on a generated function +# doesn't apply to a body that the generator returns with an explicit `return`. +inline_body(ex) = Expr(:block, Expr(:meta, :inline), ex) + # Shuffle a primitive value as unsigned words: values of up to 4 bytes are zero-extended to a # `UInt32`. Other 8-byte values are shuffled as a `UInt64`, which the backend may support # natively, and a `UInt64` it doesn't is split into two `UInt32` words. 16-byte values are # split into two `UInt64` words. -@inline @generated function shfl_words(f, val::T) where {T} +@inline @generated function shfl_words(op, val::T) where {T} shuffle_as_words(T) || return :(shfl_unsupported($T)) U = word_type(T) if sizeof(T) <= 4 - return :(reinterpret($T, f(reinterpret($U, val) % UInt32) % $U)) + return inline_body(:(reinterpret($T, op(reinterpret($U, val) % UInt32) % $U))) elseif sizeof(T) == 8 && T !== UInt64 - return :(reinterpret($T, f(reinterpret(UInt64, val)))) + return inline_body(:(reinterpret($T, op(reinterpret(UInt64, val))))) end W = sizeof(T) == 8 ? UInt32 : UInt64 nbits = 8 * sizeof(W) - words = (:((f((bits >> $(nbits * (i - 1))) % $W) % $U) << $(nbits * (i - 1))) for i in 1:2) - return quote - bits = reinterpret($U, val) - return reinterpret($T, |($(words...))) - end + words = (:((op((bits >> $(nbits * (i - 1))) % $W) % $U) << $(nbits * (i - 1))) for i in 1:2) + return inline_body( + quote + bits = reinterpret($U, val) + reinterpret($T, |($(words...))) + end + ) end # A `UInt64` that the backend doesn't shuffle natively is split into two `UInt32` words by a # method of its own, so that another 8-byte type, shuffled as a `UInt64`, doesn't reach the # fallback through itself, which inference gives up on (on Julia 1.10). Backends that shuffle # `UInt64` natively override these. -@inline shfl(val::UInt64, lane::Integer) = shfl_words(x -> shfl(x, lane), val) -@inline shfl_down(val::UInt64, offset::Integer) = shfl_words(x -> shfl_down(x, offset), val) -@inline shfl_up(val::UInt64, offset::Integer) = shfl_words(x -> shfl_up(x, offset), val) -@inline shfl_xor(val::UInt64, mask::Integer) = shfl_words(x -> shfl_xor(x, mask), val) - -# The expression that shuffles `ex::S` field by field, calling `f` on the primitive fields -function shfl_fields_expr(S, ex) - isprimitivetype(S) && return :(f($ex)) - fields = (shfl_fields_expr(fieldtype(S, i), :(getfield($ex, $i))) for i in 1:fieldcount(S)) - return Expr(:new, S, fields...) +@inline shfl(val::UInt64, lane::Integer) = shfl_words(ShuffleOp(shfl, (lane,)), val) +@inline shfl_down(val::UInt64, offset::Integer) = shfl_words(ShuffleOp(shfl_down, (offset,)), val) +@inline shfl_up(val::UInt64, offset::Integer) = shfl_words(ShuffleOp(shfl_up, (offset,)), val) +@inline shfl_xor(val::UInt64, mask::Integer) = shfl_words(ShuffleOp(shfl_xor, (mask,)), val) + +# The primitive fields of a value `ex::S`, nested ones included, in order, as +# `(type, expression)` pairs +function primitive_fields!(fields, S, ex) + if isprimitivetype(S) + push!(fields, (S, ex)) + else + for i in 1:fieldcount(S) + primitive_fields!(fields, fieldtype(S, i), :(getfield($ex, $i))) + end + end + return fields +end + +# The bit offset of each of the primitive types `types` when packed into `UInt32` words, and +# the number of words: types of less than 4 bytes are aligned to their size, so that they don't +# straddle words, and larger ones to a word. +function packed_layout(types) + offsets = Int[] + bit = 0 + for T in types + align = min(8 * sizeof(T), 32) + bit = cld(bit, align) * align + push!(offsets, bit) + bit += 8 * sizeof(T) + end + return offsets, cld(bit, 32) +end + +# The expression that builds a value of type `S` from the expressions `leaves` of its primitive +# fields, in the order of `primitive_fields!`, starting at `leaves[i]`; and the next index +function rebuild_expr(S, leaves, i) + isprimitivetype(S) && return leaves[i], i + 1 + fields = Any[] + for j in 1:fieldcount(S) + ex, i = rebuild_expr(fieldtype(S, j), leaves, i) + push!(fields, ex) + end + return Expr(:new, S, fields...), i end -# Shuffle a value that the backend doesn't support directly field by field. Nested fields are -# unrolled here, rather than shuffled with a recursive call, which inference gives up on -# (on Julia 1.10), so that `f` is only called on the primitive types. -@inline @generated function shfl_fields(f, val::T) where {T} +# Shuffle a struct or tuple: its primitive fields are packed into `UInt32` words, which are +# shuffled with `shfl_word_tuple`. Nested fields are unrolled here, rather than shuffled with a +# recursive call, which inference gives up on (on Julia 1.10). +@inline @generated function shfl_packed(op, val::T) where {T} isbitstype(T) || return :(shfl_unsupported($T)) - return shfl_fields_expr(T, :val) + fields = primitive_fields!(Any[], T, :val) + all(((S, _),) -> sizeof(S) in (1, 2, 4, 8, 16), fields) || return :(shfl_unsupported($T)) + offsets, nwords = packed_layout(first.(fields)) + bits = [Symbol(:bits, i) for i in eachindex(fields)] + body = Any[] + terms = [Any[] for _ in 1:nwords] + for (i, (S, ex)) in enumerate(fields) + U = word_type(S) + push!(body, :($(bits[i]) = reinterpret($U, $ex))) + word, shift = fld(offsets[i], 32) + 1, offsets[i] % 32 + if sizeof(S) <= 4 + push!(terms[word], :(($(bits[i]) % UInt32) << $shift)) + else + for j in 0:(sizeof(S) ÷ 4 - 1) + push!(terms[word + j], :(($(bits[i]) >> $(32j)) % UInt32)) + end + end + end + words = [isempty(t) ? :(UInt32(0)) : :(|($(t...))) for t in terms] + push!(body, :(words = shfl_word_tuple(op, ($(words...),)))) + leaves = Any[] + for (i, (S, _)) in enumerate(fields) + U = word_type(S) + word, shift = fld(offsets[i], 32) + 1, offsets[i] % 32 + if sizeof(S) <= 4 + push!(leaves, :(reinterpret($S, (words[$word] >> $shift) % $U))) + else + parts = (:($U(words[$(word + j)]) << $(32j)) for j in 0:(sizeof(S) ÷ 4 - 1)) + push!(leaves, :(reinterpret($S, |($(parts...))))) + end + end + push!(body, first(rebuild_expr(T, leaves, 1))) + return inline_body(Expr(:block, body...)) end -# The shuffles within segments of `width` lanes, implemented with `shfl` from a lane. +# Shuffle the words a struct or tuple is packed into, with `op` shuffling a `UInt32`, one at a +# time. Not public: a backend may exchange all words at once instead, e.g. with a vector shuffle. +# +# PoCL limitation: every shuffle on PoCL is a work-group barrier, so one exchange of all words +# would make these shuffles several times faster there. PoCL has vector shuffles +# (`intel_sub_group_shuffle` on `uint2` to `uint16`), but its CPU device doesn't accept them +# from SPIR-V (`cl_intel_spirv_subgroups`), so no backend implements this yet. +@inline shfl_word_tuple(op, words::NTuple{N, UInt32}) where {N} = map(op, words) """ shfl(val::T, lane::Integer, width::Integer)::T @@ -693,6 +781,7 @@ It exchanges values, not memory: it is not a memory fence, see [`sub_group_barri end + ## synchronization """ @@ -737,6 +826,7 @@ function sub_group_barrier() end + ## printing """ diff --git a/lib/KernelInterface/src/host.jl b/lib/KernelInterface/src/host.jl index 848d1b55b..f28d2601e 100644 --- a/lib/KernelInterface/src/host.jl +++ b/lib/KernelInterface/src/host.jl @@ -330,9 +330,9 @@ supports_independent_subgroups(::Backend) = false Whether kernels on the active device support the shuffles [`shfl`](@ref), [`shfl_down`](@ref), [`shfl_up`](@ref) and [`shfl_xor`](@ref) for values of type `T`. -Other primitive types of 1, 2, 4, 8 or 16 bytes are supported if `UInt32` is, and `isbits` -structs and tuples if all of their fields are (and, for those without fields, if `UInt32` -is). +Other primitive types of 1, 2, 4, 8 or 16 bytes are supported if `UInt32` is, and so are +`isbits` structs and tuples of such types (nested ones included), which are shuffled as +`UInt32` words. This reports whether values of type `T` can be moved between work-items, not whether the device can compute with them: e.g. `Float64` values can be shuffled (as words) on a device @@ -348,9 +348,8 @@ an array of `UInt64`. function supports_shuffle(backend::Backend, ::Type{T}) where {T} isprimitivetype(T) && return shuffle_as_words(T) && supports_shuffle(backend, UInt32) (isbitstype(T) && !isprimitivetype(T)) || return false - # not `true` for a type without fields on a backend without shuffles - fieldcount(T) == 0 && return supports_shuffle(backend, UInt32) - return all(i -> supports_shuffle(backend, fieldtype(T, i)), 1:fieldcount(T)) + fields = primitive_fields!(Any[], T, :val) + return all(((S, _),) -> sizeof(S) in (1, 2, 4, 8, 16), fields) && supports_shuffle(backend, UInt32) end """ diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl index a87172edd..711582fdb 100644 --- a/lib/KernelInterface/test/interface.jl +++ b/lib/KernelInterface/test/interface.jl @@ -641,6 +641,8 @@ function subgroup_communication_testsuite(backend::KI.Backend, AT, sg_size) Float32 => i -> Float32(i) / 3, Bool => isodd, Char => i -> Char(0x0001F600 + i), Bits64 => i -> Bits64((UInt64(i) << 40) - i), Bits16 => i -> Bits16(0xa000 + i), ShuffleStruct => i -> ShuffleStruct(i, -i, (i, 2i, 3i)), + # fields of less than 4 bytes share words + Tuple{Bool, Int8, Float16, UInt16} => i -> (isodd(i), Int8(-i), Float16(i), UInt16(1000i % 65536)), FallbackStruct => i -> FallbackStruct(isodd(i), Char(64 + i), Bits64(-i), (i, -i, 2i, typemin(Int64) + i)), ] fp64 && push!(values, Float64 => i -> 1 / i) @@ -671,8 +673,7 @@ function subgroup_communication_testsuite(backend::KI.Backend, AT, sg_size) @test KI.supports_shuffle(backend, Bits64) @test KI.supports_shuffle(backend, Bits16) @test KI.supports_shuffle(backend, FallbackStruct) - @test KI.supports_shuffle(backend, ShuffleStruct) == - all(S -> KI.supports_shuffle(backend, S), (Float32, Int64, Int32)) + @test KI.supports_shuffle(backend, ShuffleStruct) @test !KI.supports_shuffle(backend, Ref{Int}) end diff --git a/lib/KernelInterface/test/runtests.jl b/lib/KernelInterface/test/runtests.jl index 1f13a3214..c590b1197 100644 --- a/lib/KernelInterface/test/runtests.jl +++ b/lib/KernelInterface/test/runtests.jl @@ -49,8 +49,8 @@ KI.supports_shuffle(::WordShuffleBackend, ::Type{UInt32}) = true @test isempty(methods(vote)) end - # The shuffles only have the fallback that shuffles structs field by field, which - # doesn't handle the primitive types a backend has to implement. + # The shuffles only have the fallbacks that shuffle other types as `UInt32` words, which + # don't handle the primitive types a backend has to implement. for shfl in [KI.shfl, KI.shfl_down, KI.shfl_up, KI.shfl_xor] @test_throws ArgumentError shfl(1.0f0, 1) @test_throws ArgumentError shfl((1.0f0, 2), 1) @@ -62,10 +62,22 @@ KI.supports_shuffle(::WordShuffleBackend, ::Type{UInt32}) = true @test KI.supports_shuffle(WordShuffleBackend(), Tuple{Char, Int128}) @test !KI.supports_shuffle(WordShuffleBackend(), Ref{Int}) @test !KI.supports_shuffle(ShuffleBackend(), Float32) - @test !KI.supports_shuffle(ShuffleBackend(), Tuple{Float32, Int}) @test KI.supports_shuffle(ShuffleBackend(), Int32) - @test KI.supports_shuffle(ShuffleBackend(), Tuple{Int32, NTuple{2, Int32}}) - @test !KI.supports_shuffle(ShuffleBackend(), Tuple{Int32, Float32}) + # structs and tuples are packed into `UInt32` words, whatever their fields + @test KI.supports_shuffle(WordShuffleBackend(), Tuple{Float32, Int}) + @test KI.supports_shuffle(WordShuffleBackend(), Tuple{Int32, NTuple{2, Float16}}) + @test !KI.supports_shuffle(ShuffleBackend(), Tuple{Int32, NTuple{2, Int32}}) + + # fields of less than 4 bytes share words, and every word is shuffled once + for (val, nwords) in ( + (Float16(1), Float16(2)) => 1, (true, Int8(3), 2.0f0) => 2, + (1.0, 0x01, Int128(-5)) => 7, () => 0, + ) + count = Ref(0) + @test KI.shfl_packed(w -> (count[] += 1; ~w), KI.shfl_packed(w -> ~w, val)) === val + KI.shfl_packed(w -> (count[] += 1; w), val) + @test count[] == 2 * nwords + end # the sub-group capabilities default to unsupported @test !KI.supports_subgroups(PlainBackend()) diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index 93c96565e..dd9033277 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -238,10 +238,12 @@ KI.supports_subgroups(::POCLBackend) = device_limits().sub_group_size > 0 KI.supports_linear_subgroups(backend::POCLBackend) = KI.supports_subgroups(backend) KI.supports_independent_subgroups(::POCLBackend) = false -# The types `sub_group_shuffle` supports; other types are shuffled as words or field by -# field. `Float16` and `Float64` are reported as unsupported on devices without `cl_khr_fp16` -# or `cl_khr_fp64`, rather than shuffled as words: the overrides below are selected by type, -# not by device. +# PoCL limitation: every shuffle is a work-group barrier, also for each word of a struct; see +# `KI.shfl_word_tuple`, which could exchange them at once with `cl_intel_spirv_subgroups`. + +# The types `sub_group_shuffle` supports; other types are shuffled as words. `Float16` and +# `Float64` are reported as unsupported on devices without `cl_khr_fp16` or `cl_khr_fp64`, +# rather than shuffled as words: the overrides below are selected by type, not by device. const ShuffleTypes = Union{SPIRVIntrinsics.gentypes...} function KI.supports_shuffle(backend::POCLBackend, ::Type{T}) where {T <: ShuffleTypes} From 9e6a0cad917ca57588c77a6ed5b24c3dca17bcb0 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sun, 11 Oct 2026 11:07:57 +0200 Subject: [PATCH 4/9] KernelInterface: partitionable sub-groups Tiles narrower than a sub-group have to share it, and to be useful they have to be able to diverge from each other, like CUDA's `tiled_partition`. Add the votes, collectives and barrier within segments of `width` lanes (the shuffles already had a `width`), and `supports_partitionable_subgroups`: whether a sub-group can be partitioned into segments that diverge independently, so that only the work-items of a segment have to execute these together. It extends `supports_independent_subgroups` to segments, and implies it. The fallbacks build on the functions without a `width`, so that they are correct wherever the whole sub-group executes them. Backends that report partitionable sub-groups have to make them work for a segment alone: CUDA needs a mask of the segment's lanes, other backends a ballot that only reports active lanes and a barrier that a segment can execute. OpenCL and Metal have no barrier for part of a sub-group, so there it takes hardware that executes a sub-group in lock-step, like Intel's and Apple's GPUs, and a fence. PoCL never reports it, as its sub-group functions are work-group barriers. The shuffles with a `width` of other types now shuffle `UInt32` words with a `width`, so that a backend's shuffles with a `width` handle them. --- docs/src/index.md | 5 + docs/src/kernelinterface.md | 70 +++++-- lib/KernelInterface/src/KernelInterface.jl | 2 +- lib/KernelInterface/src/device.jl | 231 +++++++++++++++++---- lib/KernelInterface/src/host.jl | 50 +++-- lib/KernelInterface/test/interface.jl | 229 +++++++++++++++++++- lib/KernelInterface/test/runtests.jl | 6 +- src/pocl/backend.jl | 3 + 8 files changed, 515 insertions(+), 81 deletions(-) diff --git a/docs/src/index.md b/docs/src/index.md index a7644e22c..cdb8669da 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -167,6 +167,11 @@ end `GPUCompiler.method_tables` instead. - `foreach_index(f, A)` runs `f` once per index of the array `A` without writing a kernel out, and `foreach_index(f, backend, indices)` once per index in a range or `CartesianIndices`. +- KernelInterface 0.4 adds sub-groups: ids, shuffles, votes, reductions and scans, also + within segments of a sub-group. They are the low-level API for backends, for kernels + written against KernelInterface and for kernel languages built on it; capability queries + say how a backend forms its sub-groups, whether they may diverge, and which types it can + shuffle. See [Sub-groups](@ref ki_subgroups). ## Semantic differences diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md index e68ce8cc8..31015fde9 100644 --- a/docs/src/kernelinterface.md +++ b/docs/src/kernelinterface.md @@ -74,10 +74,10 @@ What a backend implements, at a glance. The docstrings below have the details. | **Execution** | [`synchronize`](@ref) (cooperative) | [`record_event`](@ref)/[`wait_event`](@ref) (synchronize), [`priority!`](@ref) (no-op) | | **Devices** | with more than one device: [`ndevices`](@ref), [`device`](@ref), [`device!`](@ref), `device(backend, A)` | all four (a single device) | | **Queries** | [`max_work_group_size`](@ref) (for the backend and for a kernel), [`max_work_group_dims`](@ref), [`max_num_groups`](@ref) | [`launch_configuration`](@ref) (the limit), [`multiprocessor_count`](@ref) (0), [`functional`](@ref) (`missing`), [`versioninfo`](@ref) | -| **Capabilities** | | [`supports_float64`](@ref), [`supports_atomics`](@ref), [`supports_unified`](@ref), [`supports_subgroups`](@ref), [`supports_linear_subgroups`](@ref), [`supports_independent_subgroups`](@ref) (all `false`), [`supports_shuffle`](@ref) (`false` unless derived from the natively supported types, see its docstring) | +| **Capabilities** | | [`supports_float64`](@ref), [`supports_atomics`](@ref), [`supports_unified`](@ref), [`supports_subgroups`](@ref), [`supports_linear_subgroups`](@ref), [`supports_independent_subgroups`](@ref), [`supports_partitionable_subgroups`](@ref) (all `false`), [`supports_shuffle`](@ref) (`false` unless derived from the natively supported types, see its docstring) | | **Compilation** | [`argconvert`](@ref), [`kernel_function`](@ref), [`launch`](@ref) | | | **Device** | [`get_local_id`](@ref), [`get_group_id`](@ref), [`get_local_size`](@ref), [`get_num_groups`](@ref), [`localmemory`](@ref), [`barrier`](@ref) | [`get_global_id`](@ref), [`get_global_size`](@ref) (derived from the primitive queries), [`_print`](@ref KernelInterface._print) (host `print`) | -| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries (with a constant [`get_max_sub_group_size`](@ref)), [`sub_group_barrier`](@ref), [`sub_group_any`](@ref), [`sub_group_all`](@ref), and [`sub_group_ballot`](@ref) for widths of at most 64; if `supports_shuffle(backend, T)`: [`shfl`](@ref), [`shfl_down`](@ref), [`shfl_up`](@ref), [`shfl_xor`](@ref) for the primitive `T` supported natively, including `UInt32` | shuffles with a `width`, [`sub_group_reduce`](@ref), [`sub_group_scan`](@ref), [`sub_group_exclusive_scan`](@ref) (built on the shuffles); shuffles of other primitive types (as `UInt32` or `UInt64` words) and of structs (packed into `UInt32` words) | +| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries, [`sub_group_barrier`](@ref), [`sub_group_any`](@ref), [`sub_group_all`](@ref), and [`sub_group_ballot`](@ref) for widths of at most 64; if `supports_shuffle(backend, T)`: [`shfl`](@ref), [`shfl_down`](@ref), [`shfl_up`](@ref), [`shfl_xor`](@ref) for the primitive `T` supported natively, including `UInt32` | [`sub_group_reduce`](@ref), [`sub_group_scan`](@ref), [`sub_group_exclusive_scan`](@ref) (built on the shuffles); the functions with a `width` (built on those without, or on the ballot); shuffles of other primitive types (as `UInt32` or `UInt64` words) and of structs (packed into `UInt32` words) | Everything else, such as [`zeros`](@ref KernelInterface.zeros), [`ones`](@ref KernelInterface.ones), the launch-keyword handling of [`Kernel`](@ref) and [`@launch`](@ref KernelInterface.@launch), @@ -146,7 +146,7 @@ get_global_id get_global_size ``` -### Sub-groups +### [Sub-groups](@id ki_subgroups) Sub-groups are optional ([`supports_subgroups`](@ref)). A work-group is divided into sub-groups of at most `W = `[`sub_group_size(backend)`](@ref sub_group_size) work-items, the @@ -171,7 +171,7 @@ A sub-group is partial when it has fewer work-items than the width: the lanes values, and the votes, [`sub_group_reduce`](@ref) and the scans only take the work-items of the sub-group into account. -Backends that guarantee more say so with two queries: +Backends that guarantee more say so with these queries: - [`supports_linear_subgroups`](@ref): if the work-group is 1-D, or its x extent is a multiple of `W`, sub-groups are formed from consecutive work-items, x fastest: the @@ -180,25 +180,32 @@ Backends that guarantee more say so with two queries: get_local_id()` and `size = get_local_size()`) is in sub-group `fld(lin - 1, W) + 1`, lane `mod(lin - 1, W) + 1`, and only the last sub-group can be partial. Code that relates local ids to lanes needs this, e.g. stencils exchanging neighbouring values with shuffles. Code - that doesn't, like a work-group reduction combining a value per sub-group, can index by the - sub-group slot `(get_sub_group_id() - 1) * W + get_sub_group_local_id()` instead. + that doesn't, like a work-group reduction that stores a result per sub-group, can index + those results by `get_sub_group_id()` and combine them using `get_num_sub_groups()`. - [`supports_independent_subgroups`](@ref): the communication functions (shuffles, votes, collectives and [`sub_group_barrier`](@ref)) can be executed in control flow that is uniform over each sub-group, but differs between the sub-groups of a work-group, e.g. a loop whose trip count depends on the sub-group, or after an early `return` of whole sub-groups. Otherwise, all sub-groups of a work-group have to execute the same communication functions in the same order. - -| Backend | `supports_linear_subgroups` | `supports_independent_subgroups` | -|---|---|---| -| CUDA, AMDGPU | `true` | `true` | -| Metal | `true` (not specified by Metal, but tested) | `true` | -| oneAPI | once the layout is tested on Intel's GPU compiler | `true` | -| OpenCL.jl | on runtimes known to form them linearly | `true`, except on PoCL | -| KernelAbstractions' POCL backend | `true` | `false`: sub-group operations are work-group barriers | - -In all cases, all work-items of a sub-group have to execute a communication function -together, not in a branch that only some of its lanes take. The id and size queries can be +- [`supports_partitionable_subgroups`](@ref): a sub-group can be partitioned into segments that + diverge independently. The communication functions with a `width` work on segments of + `width` consecutive lanes; all work-items of a segment have to execute them together, but + different segments of a sub-group may take different branches or loop a different number of + times. Implies `supports_independent_subgroups`. + +| Backend | `supports_linear_subgroups` | `supports_independent_subgroups` | `supports_partitionable_subgroups` | +|---|---|---|---| +| CUDA, AMDGPU | `true` | `true` | `true` | +| Metal | `true` (not specified by Metal, but tested) | `true` | `true`: Apple GPUs execute a SIMD-group in lock-step, so the barrier for a segment is a SIMD-group-scope fence (not specified by Metal, but validated on M1 and M3) | +| oneAPI | `true` | `true` | `true`: Intel GPUs execute a sub-group as the SIMD lanes of one hardware thread, so the barrier for a segment is a sub-group-scope fence | +| OpenCL.jl | on PoCL and Intel's runtimes, which form them linearly | `true`, except on PoCL and Intel's CPU runtime | on Intel's GPU runtime only, as with oneAPI | +| KernelAbstractions' POCL backend | `true` | `false`: sub-group operations are work-group barriers | `false` | + +Unless sub-groups are partitionable, the functions with a `width` need the whole sub-group, like +the others: all its work-items have to execute them together, not in a branch that only some +of its lanes take. Without independent sub-groups, all sub-groups of the work-group have to +execute the same communication functions in the same order. The id and size queries can be used anywhere. !!! note "Kernel languages" @@ -221,6 +228,7 @@ get_sub_group_local_id ```@docs barrier sub_group_barrier +sub_group_barrier(::Integer) ``` ### Memory @@ -231,11 +239,10 @@ localmemory ### Communication -The shuffles, votes and collectives below exchange values between the work-items of a -sub-group, not memory: they don't order or make visible the work-items' accesses to local -or global memory. To communicate through memory within a sub-group, e.g. a work-item reading -what another one wrote to local memory, use [`sub_group_barrier`](@ref) between the write -and the read. +The shuffles, votes and collectives below don't synchronize accesses to local or global +memory. When a work-item reads memory that another one of its sub-group wrote, put a +[`sub_group_barrier`](@ref) between the write and the read, or +[`sub_group_barrier(width)`](@ref sub_group_barrier(::Integer)) within a segment. ```@docs shfl @@ -246,9 +253,11 @@ shfl(::Any, ::Integer, ::Integer) sub_group_any sub_group_all sub_group_ballot +sub_group_ballot(::Bool, ::Integer) sub_group_reduce sub_group_scan sub_group_exclusive_scan +sub_group_reduce(::Any, ::Any, ::Integer) ``` ### Printing @@ -305,6 +314,7 @@ supports_float64 supports_subgroups supports_linear_subgroups supports_independent_subgroups +supports_partitionable_subgroups supports_shuffle ``` @@ -359,6 +369,22 @@ optional methods where it can do better than the fallback. In particular: ``` 4. Compute the typed index queries with `% T`, not `T(x)`: a checked conversion leaves an error branch in every kernel. +5. Sub-groups are optional. To support them, report [`supports_subgroups`](@ref) and + [`sub_group_size`](@ref), and implement the device functions the + [contract](@ref Contract) requires. If the backend has shuffles, implement them only for the + primitive types the hardware shuffles natively, including `UInt32`, and report those with + [`supports_shuffle`](@ref); KernelInterface shuffles other types as words. It also + provides the reductions, scans and functions with a `width`; override those where a + native implementation helps, keeping their documented results. + + Report linear, independent and partitionable sub-groups only when the backend guarantees + them. Partitionable sub-groups need every function with a `width`, + including [`sub_group_barrier(width)`](@ref sub_group_barrier(::Integer)), to work when + only that segment executes it. Without a barrier for part of a sub-group, as in OpenCL + and Metal, that can rest on a hardware execution model, such as sub-groups that run in + lock-step, combined with a fence. Base these guarantees on the API or the hardware, not + on the testsuite passing: it checks the layout and divergent control flow, but can't + prove them. The PoCL backend in `src/pocl/backend.jl` is a complete worked example. diff --git a/lib/KernelInterface/src/KernelInterface.jl b/lib/KernelInterface/src/KernelInterface.jl index 8eb260b4b..0c83b74fa 100644 --- a/lib/KernelInterface/src/KernelInterface.jl +++ b/lib/KernelInterface/src/KernelInterface.jl @@ -46,7 +46,7 @@ include("host.jl") :functional, :versioninfo, :supports_unified, :supports_atomics, :supports_float64, :supports_subgroups, :supports_shuffle, :supports_linear_subgroups, - :supports_independent_subgroups, + :supports_independent_subgroups, :supports_partitionable_subgroups, ) ) end diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl index 0afef9f6b..3df9e7606 100644 --- a/lib/KernelInterface/src/device.jl +++ b/lib/KernelInterface/src/device.jl @@ -141,7 +141,8 @@ end The number of work-items in the sub-group, at most the sub-group width ([`get_max_sub_group_size`](@ref)). Which sub-groups have fewer work-items than the width -is unspecified (unless [`supports_linear_subgroups`](@ref)): there can be more than one, +is unspecified (unless the layout is linear, see [`supports_linear_subgroups`](@ref)): there +can be more than one, e.g. one per row of a multi-dimensional work-group. See [`get_local_id`](@ref) for the supported types `T`. @@ -163,9 +164,9 @@ sub_group_size) on the host. Backends should make it a constant of the generated code, so that code depending on it, e.g. a loop over the lanes or a shuffle butterfly, is specialized for it; on some, it is only -folded late in compilation, or not at all. It is never known during type inference: to pick -types or `Val` parameters from the width, e.g. the integer type of a mask with a bit per -lane, pass [`sub_group_size(backend)`](@ref sub_group_size) from the host. +folded late in compilation, or not at all. Don't rely on it being known during type +inference: to pick types or `Val` parameters from the width, e.g. the integer type of a mask +with a bit per lane, pass [`sub_group_size(backend)`](@ref sub_group_size) from the host. See [`get_local_id`](@ref) for the supported types `T`. @@ -185,7 +186,8 @@ See [`get_local_id`](@ref) for the supported types `T`. The number of sub-groups in the work-group. It is at least `cld(prod(get_local_size()), get_max_sub_group_size())`, but can be larger, since more than -one sub-group can be partial (unless [`supports_linear_subgroups`](@ref)). Size storage for a +one sub-group can be partial (unless the layout is linear, see +[`supports_linear_subgroups`](@ref)). Size storage for a value per sub-group for up to one sub-group per work-item. See [`get_local_id`](@ref) for the supported types `T`. @@ -204,7 +206,8 @@ See [`get_local_id`](@ref) for the supported types `T`. The 1-based index of the sub-group within the work-group, between 1 and [`get_num_sub_groups`](@ref). How it relates to [`get_local_id`](@ref) is unspecified, -unless [`supports_linear_subgroups`](@ref). +unless the layout is linear (see [`supports_linear_subgroups`](@ref), which only covers some +work-group shapes). See [`get_local_id`](@ref) for the supported types `T`. @@ -223,7 +226,8 @@ See [`get_local_id`](@ref) for the supported types `T`. The 1-based index of the work-item within its sub-group (its lane), between 1 and [`get_sub_group_size`](@ref). It doesn't depend on which work-items of the sub-group are active, e.g. in a divergent branch. How it relates to [`get_local_id`](@ref) is unspecified, -unless [`supports_linear_subgroups`](@ref). +unless the layout is linear (see [`supports_linear_subgroups`](@ref), which only covers some +work-group shapes). See [`get_local_id`](@ref) for the supported types `T`. @@ -271,6 +275,11 @@ localmemory(::Type{T}, ::Val) where {T} = # The shuffles, votes, collectives and `sub_group_barrier` have to be executed by all # work-items of a sub-group together, in control flow that is uniform over the work-group # unless `supports_independent_subgroups`. The id and size queries above can be used anywhere. +# +# Their variants with a `width` work within segments of `width` consecutive lanes. If +# `supports_partitionable_subgroups`, only the work-items of a segment have to execute those +# together, so that the segments of a sub-group can diverge; otherwise, they have to be +# executed like the functions without a `width`. """ shfl(val::T, lane::Integer)::T @@ -282,7 +291,8 @@ All work-items of the sub-group have to execute `shfl` together (not in a diverg but they may read from different lanes. Unless [`supports_independent_subgroups`](@ref), all sub-groups of the work-group have to execute it. -`shfl` exchanges values, not memory: it is not a memory fence. +`shfl` doesn't synchronize memory accesses: use [`sub_group_barrier`](@ref) to communicate +through memory. Types for which [`supports_shuffle`](@ref) returns `true` are supported. Besides the types a backend supports natively, that includes other primitive types of 1, 2, 4, 8 or 16 bytes @@ -293,9 +303,9 @@ and as two `UInt32` words otherwise. Structs and tuples are packed into `UInt32` fields of less than 4 bytes sharing words, and shuffled word by word. !!! note - Backend implementations **must** implement this for the primitive types they support - natively, which have to include `UInt32`, and only for those, so that other types reach - the fallbacks: + Backends that support shuffles **must** implement this for the primitive types they + support natively, which have to include `UInt32`, and only for those, so that other types + reach the fallbacks: ``` @device_override shfl(val::T, lane::Integer) where {T <: Union{...}} ``` @@ -315,8 +325,7 @@ work-item, in a partial sub-group, the result is an unspecified value (of type ` All work-items of the sub-group have to execute `shfl_down` together (not in a divergent branch), with the same `offset`; see [`shfl`](@ref). -`shfl_down` exchanges values, not memory: it is not a memory fence. See [`shfl`](@ref) for -the supported types. +`shfl_down` doesn't synchronize memory accesses. See [`shfl`](@ref) for the supported types. !!! note Backend implementations **must** implement this like [`shfl`](@ref): @@ -337,8 +346,7 @@ result is `val` of the work-item itself, as on CUDA, HIP and Metal. All work-items of the sub-group have to execute `shfl_up` together (not in a divergent branch), with the same `offset`; see [`shfl`](@ref). -`shfl_up` exchanges values, not memory: it is not a memory fence. See [`shfl`](@ref) for -the supported types. +`shfl_up` doesn't synchronize memory accesses. See [`shfl`](@ref) for the supported types. !!! note Backend implementations **must** implement this like [`shfl`](@ref): @@ -358,12 +366,10 @@ xor `mask`, i.e. with [`get_sub_group_local_id`](@ref) equal to result is an unspecified value (of type `T`). All work-items of the sub-group have to execute `shfl_xor` together (not in a divergent -branch), with the same `mask`; see [`shfl`](@ref). A butterfly over the masks `width ÷ 2, …, 2, 1` (for the -sub-group width [`get_max_sub_group_size`](@ref)) reduces a full sub-group such that every -work-item gets the result. +branch), with the same `mask`; see [`shfl`](@ref). For a reduction over the sub-group, use +[`sub_group_reduce`](@ref), which also handles partial sub-groups. -`shfl_xor` exchanges values, not memory: it is not a memory fence. See [`shfl`](@ref) for -the supported types. +`shfl_xor` doesn't synchronize memory accesses. See [`shfl`](@ref) for the supported types. !!! note Backend implementations **must** implement this like [`shfl`](@ref): @@ -539,25 +545,50 @@ end Shuffles within segments of `width` consecutive lanes of the sub-group, with CUDA's semantics for a `width`: -- `shfl` reads from lane `lane` of the caller's segment, with `lane` taken modulo `width` - (lane `width + 1` is the segment's first lane); -- `shfl_down` and `shfl_up` read from the lane `offset` lanes further or earlier, and return - `val` of the work-item itself where that lane is outside of the caller's segment; +- `shfl` reads from lane `mod(lane - 1, width) + 1` of the caller's segment; +- `shfl_down` and `shfl_up` read from the lane `offset` lanes further or earlier, for an + `offset` of at least 0, and return `val` of the work-item itself where that lane is outside + of the caller's segment; - `shfl_xor` reads from the lane whose 0-based position in the segment is the caller's xor `mask`, which has to be between 0 and `width - 1`. Reading from a lane of the segment that has no work-item (in a partial sub-group) gives an unspecified value. `width` has to be a power of two of at most the sub-group width -[`get_max_sub_group_size`](@ref), and the same for all work-items of the sub-group. +[`get_max_sub_group_size`](@ref), and the same for all work-items of the sub-group, so that +the segments partition it. It should be a constant of the generated code. -The `width` only changes which lanes are read from: as for the shuffles without a `width`, all -work-items of the sub-group have to execute them together, not just those of a segment. +If [`supports_partitionable_subgroups`](@ref), only the work-items of the caller's segment have +to execute them together: the segments of a sub-group can diverge, e.g. loop a different +number of times. Otherwise, all work-items of the sub-group have to execute them together, +as for the shuffles without a `width`. Either way, `offset` and `mask` have to be the same +for the work-items that execute them together, while `lane` may differ. !!! note - Backends **may** implement these, e.g. if they have native shuffles with a width. The - fallbacks use [`shfl`](@ref) from a lane. -""" -@inline function shfl(val, lane::Integer, width::Integer) + Backends **may** implement these for the types they support natively, e.g. if they have + native shuffles with a width. A backend that implements + [`supports_partitionable_subgroups`](@ref) **must** implement them unless its `shfl` from a + lane can be executed by a segment alone. The fallbacks pack other types into `UInt32` + words, and shuffle those with `shfl` from a lane. +""" +@inline shfl(val, lane::Integer, width::Integer) = shfl_fallback(ShuffleOp(shfl, (lane, width)), val) +@inline shfl_down(val, offset::Integer, width::Integer) = + shfl_fallback(ShuffleOp(shfl_down, (offset, width)), val) +@inline shfl_up(val, offset::Integer, width::Integer) = + shfl_fallback(ShuffleOp(shfl_up, (offset, width)), val) +@inline shfl_xor(val, mask::Integer, width::Integer) = + shfl_fallback(ShuffleOp(shfl_xor, (mask, width)), val) + +@inline shfl(val::UInt64, lane::Integer, width::Integer) = + shfl_words(ShuffleOp(shfl, (lane, width)), val) +@inline shfl_down(val::UInt64, offset::Integer, width::Integer) = + shfl_words(ShuffleOp(shfl_down, (offset, width)), val) +@inline shfl_up(val::UInt64, offset::Integer, width::Integer) = + shfl_words(ShuffleOp(shfl_up, (offset, width)), val) +@inline shfl_xor(val::UInt64, mask::Integer, width::Integer) = + shfl_words(ShuffleOp(shfl_xor, (mask, width)), val) + +# A `UInt32` within a segment is read with `shfl` from the computed lane +@inline function shfl(val::UInt32, lane::Integer, width::Integer) l0 = get_sub_group_local_id(Int32) - Int32(1) w = width % Int32 base = l0 & ~(w - Int32(1)) @@ -567,7 +598,7 @@ end # The offsets are compared with the width before they are narrowed to `Int32`, so that large # offsets read from outside of the segment. -@inline function shfl_down(val, offset::Integer, width::Integer) +@inline function shfl_down(val::UInt32, offset::Integer, width::Integer) l0 = get_sub_group_local_id(Int32) - Int32(1) w = width % Int32 pos = l0 & (w - Int32(1)) @@ -576,7 +607,7 @@ end return shfl(val, src + Int32(1)) end -@inline function shfl_up(val, offset::Integer, width::Integer) +@inline function shfl_up(val::UInt32, offset::Integer, width::Integer) l0 = get_sub_group_local_id(Int32) - Int32(1) w = width % Int32 pos = l0 & (w - Int32(1)) @@ -585,7 +616,7 @@ end return shfl(val, src + Int32(1)) end -@inline function shfl_xor(val, mask::Integer, width::Integer) +@inline function shfl_xor(val::UInt32, mask::Integer, width::Integer) l0 = get_sub_group_local_id(Int32) - Int32(1) return shfl(val, (l0 ⊻ (mask % Int32)) + Int32(1)) end @@ -599,7 +630,7 @@ get the same result. All work-items of the sub-group have to execute `sub_group_any` together (not in a divergent branch); see [`shfl`](@ref). -It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref). +It doesn't synchronize memory accesses, see [`sub_group_barrier`](@ref). !!! note Backend implementations that support sub-groups **must** implement: @@ -618,7 +649,7 @@ sub-group get the same result. All work-items of the sub-group have to execute `sub_group_all` together (not in a divergent branch); see [`shfl`](@ref). -It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref). +It doesn't synchronize memory accesses, see [`sub_group_barrier`](@ref). !!! note Backend implementations that support sub-groups **must** implement: @@ -635,13 +666,13 @@ A mask of the work-items of the sub-group for which `pred` is `true`: bit `i - 1 from the least significant bit) is set if and only if the sub-group has a work-item with [`get_sub_group_local_id`](@ref) equal to `i` and its `pred` is `true`. All other bits, e.g. those of lanes without a work-item in a partial sub-group, are zero. All work-items of the -sub-group get the same result. Use e.g. `count_ones` to count the work-items, or -`trailing_zeros` to find the first one. +sub-group get the same result. `count_ones(mask)` counts the work-items, and for a nonzero +mask, `trailing_zeros(mask) + 1` is the lane of the first one. All work-items of the sub-group have to execute `sub_group_ballot` together (not in a divergent branch); see [`shfl`](@ref). Only sub-groups of at most 64 work-items are supported. -It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref). +It doesn't synchronize memory accesses, see [`sub_group_barrier`](@ref). !!! note Backend implementations that support sub-groups with a width of at most 64 **must** @@ -652,6 +683,46 @@ It exchanges values, not memory: it is not a memory fence, see [`sub_group_barri """ function sub_group_ballot end +""" + sub_group_any(pred::Bool, width::Integer)::Bool + sub_group_all(pred::Bool, width::Integer)::Bool + sub_group_ballot(pred::Bool, width::Integer)::UInt64 + +The votes within segments of `width` consecutive lanes of the sub-group, as for the shuffles +with a `width` (see [`shfl`](@ref shfl(::Any, ::Integer, ::Integer))): whether `pred` is +`true` for any or all work-items of the caller's segment, or a mask of those for which it is, +relative to the segment: bit `i - 1` for the `i`-th lane of the segment. Lanes without a +work-item don't count. + +For a `width` smaller than the sub-group width, they need a sub-group width of at most 64. + +!!! note + Backends **may** implement these. A backend that implements + [`supports_partitionable_subgroups`](@ref) **must** implement them unless its + [`sub_group_ballot`](@ref) can be executed by a segment alone and only reports the + work-items that execute it. The fallbacks use the votes without a `width` if `width` is + the sub-group width, and [`sub_group_ballot`](@ref) otherwise. +""" +@inline function sub_group_ballot(pred::Bool, width::Integer) + width == get_max_sub_group_size(Int32) && return sub_group_ballot(pred) + l0 = get_sub_group_local_id(Int32) - Int32(1) + base = l0 & ~(width % Int32 - Int32(1)) + return (sub_group_ballot(pred) >> base) & segment_mask(width) +end + +@inline function sub_group_any(pred::Bool, width::Integer) + width == get_max_sub_group_size(Int32) && return sub_group_any(pred) + return sub_group_ballot(pred, width) != 0 +end + +@inline function sub_group_all(pred::Bool, width::Integer) + width == get_max_sub_group_size(Int32) && return sub_group_all(pred) + return sub_group_ballot(!pred, width) == 0 +end + +# The bits of the lanes of a segment of `width` lanes in a ballot relative to the segment +@inline segment_mask(width) = ifelse(width >= 64, typemax(UInt64), (UInt64(1) << (width % UInt32)) - UInt64(1)) + """ sub_group_reduce(op, val::T)::T @@ -665,7 +736,7 @@ divergent branch), with the same `op`; see [`shfl`](@ref). Values of the types t shuffles support are supported, see [`supports_shuffle`](@ref), and others for which the backend has a native reduction. -It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref). +It doesn't synchronize memory accesses, see [`sub_group_barrier`](@ref). !!! note Backends **may** implement this for operators and types with a native reduction (e.g. @@ -729,7 +800,7 @@ branch), with the same `op`; see [`shfl`](@ref). Values of the types that the sh support are supported, see [`supports_shuffle`](@ref), and others for which the backend has a native scan. -It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref). +It doesn't synchronize memory accesses, see [`sub_group_barrier`](@ref). !!! note Backends **may** implement this for operators and types with a native scan (e.g. `+` on @@ -769,7 +840,7 @@ divergent branch), with the same `op`; see [`shfl`](@ref). Values of the types t shuffles support are supported, see [`supports_shuffle`](@ref), and others for which the backend has a native scan. -It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref). +It doesn't synchronize memory accesses, see [`sub_group_barrier`](@ref). !!! note Backends **may** implement this like [`sub_group_scan`](@ref). The fallback shifts the @@ -780,6 +851,66 @@ It exchanges values, not memory: it is not a memory fence, see [`sub_group_barri return get_sub_group_local_id(Int32) == Int32(1) ? init : op(init, prefix) end +""" + sub_group_reduce(op, val::T, width::Integer)::T + sub_group_scan(op, val::T, width::Integer)::T + sub_group_exclusive_scan(op, val::T, init::T, width::Integer)::T + +The collectives within segments of `width` consecutive lanes of the sub-group, as for the +shuffles with a `width` (see [`shfl`](@ref shfl(::Any, ::Integer, ::Integer))): the +reduction and scans of [`sub_group_reduce`](@ref), [`sub_group_scan`](@ref) and +[`sub_group_exclusive_scan`](@ref), over the work-items of the caller's segment, in the order +of its lanes. `init` has to be the same for the work-items of a segment. + +!!! note + Backends **may** implement these, like the collectives without a `width`. The fallbacks + use the collectives without a `width` if `width` is the sub-group width, and the shuffles + with a `width` otherwise, like the fallbacks of the collectives without one. +""" +@inline function sub_group_reduce(op, val, width::Integer) + width == get_max_sub_group_size(Int32) && return sub_group_reduce(op, val) + w = width % Int32 + sgsize = get_sub_group_size(Int32) + lane0 = get_sub_group_local_id(Int32) - Int32(1) + # the butterfly of `sub_group_reduce`, within the segment; a partial segment (in a partial + # sub-group) gets the result of its first lane + mask = Int32(1) + while mask < w + other = shfl_xor(val, mask, width) + if (lane0 ⊻ mask) < sgsize + if lane0 & mask == Int32(0) + val = op(val, other) + else + val = op(other, val) + end + end + mask <<= 1 + end + first = shfl(val, 1, width) + return ifelse((lane0 & ~(w - Int32(1))) + w <= sgsize, val, first) +end + +@inline function sub_group_scan(op, val, width::Integer) + width == get_max_sub_group_size(Int32) && return sub_group_scan(op, val) + w = width % Int32 + pos = (get_sub_group_local_id(Int32) - Int32(1)) & (w - Int32(1)) + offset = Int32(1) + while offset < w + other = shfl_up(val, offset, width) + if pos >= offset + val = op(other, val) + end + offset <<= 1 + end + return val +end + +@inline function sub_group_exclusive_scan(op, val, init, width::Integer) + width == get_max_sub_group_size(Int32) && return sub_group_exclusive_scan(op, val, init) + prefix = shfl_up(sub_group_scan(op, val, width), 1, width) + first = (get_sub_group_local_id(Int32) - Int32(1)) & (width % Int32 - Int32(1)) == Int32(0) + return first ? init : op(init, prefix) +end ## synchronization @@ -825,6 +956,22 @@ function sub_group_barrier() error("Sub-group barrier used outside kernel or not captured") end +""" + sub_group_barrier(width::Integer) + +Like [`sub_group_barrier`](@ref), for the work-items of the caller's segment of `width` +consecutive lanes of the sub-group, as for the shuffles with a `width` (see +[`shfl`](@ref shfl(::Any, ::Integer, ::Integer))): wait until they have reached the barrier, +and make their writes to global and local memory before it visible to the segment. It +doesn't order memory with other segments. + +!!! note + Backends **may** implement this. A backend that implements + [`supports_partitionable_subgroups`](@ref) **must** implement it unless its + [`sub_group_barrier`](@ref) can be executed by a segment alone. The fallback is + [`sub_group_barrier`](@ref). +""" +@inline sub_group_barrier(width::Integer) = sub_group_barrier() ## printing diff --git a/lib/KernelInterface/src/host.jl b/lib/KernelInterface/src/host.jl index f28d2601e..9a5cbc41a 100644 --- a/lib/KernelInterface/src/host.jl +++ b/lib/KernelInterface/src/host.jl @@ -269,8 +269,8 @@ is queried separately with [`supports_shuffle`](@ref). Which work-items form a sub-group is unspecified, unless [`supports_linear_subgroups`](@ref) returns `true`, and the communication functions have to be executed in control flow that is uniform over the work-group, unless -[`supports_independent_subgroups`](@ref) returns `true`. See the manual for what -KernelInterface guarantees about sub-groups. +[`supports_independent_subgroups`](@ref) returns `true`. See [Sub-groups](@ref ki_subgroups) +for the layout and participation guarantees. !!! note Backend implementations **must** implement this function if they support sub-groups. @@ -281,13 +281,14 @@ supports_subgroups(::Backend) = false """ supports_linear_subgroups(::Backend)::Bool -Whether sub-groups are formed from consecutive work-items on the active device, for every -kernel and launch: if the work-group is 1-D, or its x extent is a multiple of the sub-group -width `W` ([`sub_group_size`](@ref)), the work-item with the linear local index `lin` (x +Whether the active device forms sub-groups from consecutive work-items, for every kernel and +launch with a 1-D work-group or one whose x extent is a multiple of the sub-group width `W` +([`sub_group_size`](@ref)): then the work-item with the linear local index `lin` (x fastest, see [`get_local_id`](@ref)) is in sub-group `fld(lin - 1, W) + 1`, at lane `mod(lin - 1, W) + 1`, the sub-group ids are `1:cld(n, W)` for a work-group of `n` work-items, and only the last sub-group can be partial. In particular, a 1-D work-group of at -most `W` work-items is a single sub-group. Implies [`supports_subgroups`](@ref). +most `W` work-items is a single sub-group. It promises nothing about other work-group shapes. +Implies [`supports_subgroups`](@ref). Code needs this when it relates local ids to lanes: e.g. stencils that shuffle values between neighbouring work-items, transposes through shuffles, or work-group scans built from @@ -314,9 +315,9 @@ sub-group, but not over the work-group: e.g. a loop whose trip count differs bet sub-groups, or after an early `return` of whole sub-groups. Otherwise, all sub-groups of a work-group have to execute the same communication functions, in the same order. -This covers physical sub-groups, not the segments of the shuffles with a `width`, and doesn't -promise forward progress: a sub-group spin-waiting for another one can still hang. Implies -[`supports_subgroups`](@ref). +This covers whole sub-groups, not the segments of the functions with a `width` (see +[`supports_partitionable_subgroups`](@ref)), and doesn't promise forward progress: a sub-group +spin-waiting for another one can still hang. Implies [`supports_subgroups`](@ref). !!! note Backend implementations **should** implement this function if sub-groups are @@ -324,6 +325,29 @@ promise forward progress: a sub-group spin-waiting for another one can still han """ supports_independent_subgroups(::Backend) = false +""" + supports_partitionable_subgroups(::Backend)::Bool + +Whether a sub-group can be partitioned into segments of `width` consecutive lanes that diverge +independently, for the sub-group communication functions with a `width` (the shuffles, +votes, collectives and [`sub_group_barrier`](@ref) with a `width`, e.g. +[`shfl`](@ref shfl(::Any, ::Integer, ::Integer))): those only have to be executed by the +work-items of the caller's segment, rather than by the whole sub-group, so the segments of a +sub-group can take different branches, loop a different number of times, or skip the +functions altogether. + +This extends [`supports_independent_subgroups`](@ref) to segments, and implies it. Like that +query, it doesn't promise forward progress between segments. + +!!! note + Backend implementations **should** return `true` only if every function with a + `width`, including the barrier, works when only one segment executes it. That can rest + on the API, or on an execution model of the hardware that the backend relies on and + validates, such as the lock-step execution of SIMD-groups on Apple GPUs; seeing + lock-step execution in a test isn't enough. The fallback returns `false`. +""" +supports_partitionable_subgroups(::Backend) = false + """ supports_shuffle(::Backend, ::Type{T})::Bool @@ -331,8 +355,8 @@ Whether kernels on the active device support the shuffles [`shfl`](@ref), [`shfl_down`](@ref), [`shfl_up`](@ref) and [`shfl_xor`](@ref) for values of type `T`. Other primitive types of 1, 2, 4, 8 or 16 bytes are supported if `UInt32` is, and so are -`isbits` structs and tuples of such types (nested ones included), which are shuffled as -`UInt32` words. +`isbits` structs and tuples of such types (nested ones included). The fallbacks shuffle +these as words; see [`shfl`](@ref) for how. This reports whether values of type `T` can be moved between work-items, not whether the device can compute with them: e.g. `Float64` values can be shuffled (as words) on a device @@ -340,8 +364,8 @@ without [`supports_float64`](@ref), as long as they are only stored and loaded, an array of `UInt64`. !!! note - Backend implementations **must** implement this function for the primitive types they - support natively, which have to include `UInt32`, with a signature that only matches + Backends that support shuffles **must** implement this function for the primitive types + they support natively, which have to include `UInt32`, with a signature that only matches those, e.g. `supports_shuffle(::NewBackend, ::Type{<:Union{UInt32, Float32}})`. The fallback handles other types. """ diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl index 711582fdb..b42625987 100644 --- a/lib/KernelInterface/test/interface.jl +++ b/lib/KernelInterface/test/interface.jl @@ -393,6 +393,43 @@ struct ExclusiveScan{O, T} end (s::ExclusiveScan)(x) = KI.sub_group_exclusive_scan(s.op, x, s.init) +# The votes and collectives within segments of `width` lanes. The width is a field, rather +# than a type parameter, so that a kernel is compiled once for all widths. +struct WidthVote{F} + f::F + width::Int +end +(v::WidthVote)(p) = v.f(p, v.width) + +struct WidthReduce{O} + op::O + width::Int +end +(r::WidthReduce)(x) = KI.sub_group_reduce(r.op, x, r.width) + +struct WidthScan{O} + op::O + width::Int +end +(s::WidthScan)(x) = KI.sub_group_scan(s.op, x, s.width) + +struct WidthExclusiveScan{O, T} + op::O + init::T + width::Int +end +(s::WidthExclusiveScan)(x) = KI.sub_group_exclusive_scan(s.op, x, s.init, s.width) + +# the segment widths to test: the powers of two up to the sub-group width +segment_widths(sg_size) = ispow2(sg_size) ? filter(<=(sg_size), [1, 2, 4, 8, 16, 32, 64]) : Int[] + +# the lanes of the segment of `width` lanes that lane `l` of a sub-group of `size` work-items +# is in +function segment_lanes(l, width, size) + first = fld(l - 1, width) * width + 1 + return first:min(first + width - 1, size) +end + # an associative but not commutative operator: the composition of affine maps compose_affine(f, g) = (g[1] * f[1], g[1] * f[2] + g[2]) # the (value, index) of the smallest value, the first one of equal values @@ -402,6 +439,12 @@ argmin_op(x, y) = ifelse(y[1] < x[1], y, x) collective_result(r::Reduce, vals, l) = foldl(r.op, vals) collective_result(s::Scan, vals, l) = foldl(s.op, vals[1:l]) collective_result(s::ExclusiveScan, vals, l) = l == 1 ? s.init : s.op(s.init, foldl(s.op, vals[1:(l - 1)])) +collective_result(r::WidthReduce, vals, l) = foldl(r.op, vals[segment_lanes(l, r.width, length(vals))]) +collective_result(s::WidthScan, vals, l) = foldl(s.op, vals[first(segment_lanes(l, s.width, length(vals))):l]) +function collective_result(s::WidthExclusiveScan, vals, l) + f = first(segment_lanes(l, s.width, length(vals))) + return l == f ? s.init : s.op(s.init, foldl(s.op, vals[f:(l - 1)])) +end # Floating-point `+` and `*` may be reassociated, so they are compared with a tolerance # (other than for the `init` of an exclusive scan, which lane 1 gets unchanged). @@ -449,6 +492,29 @@ function vote_testsuite(backend, AT, sg_size, pred) return end +# the expected result of a vote over the predicates `p` of a sub-group or segment +function vote_result(vote, p) + vote === KI.sub_group_any && return any(p) + vote === KI.sub_group_all && return all(p) + return reduce(|, (UInt64(1) << (l - 1) for l in eachindex(p) if p[l]); init = UInt64(0)) +end + +function width_vote_testsuite(backend, AT, sg_size, pred) + for w in segment_widths(sg_size), (vote, T) in ((KI.sub_group_any, Bool), (KI.sub_group_all, Bool), (KI.sub_group_ballot, UInt64)) + # the votes within segments narrower than the sub-group use its ballot + (vote === KI.sub_group_ballot || w < sg_size) && sg_size > 64 && continue + res = sub_group_apply(backend, AT, WidthVote(vote, w), pred, T) + res === nothing && continue + out, groups = res + ok = true + for g in groups, (l, i) in enumerate(g) + ok &= out[i] == vote_result(vote, pred[g[segment_lanes(l, w, length(g))]]) + end + @test ok + end + return +end + # the values come from a divergent branch, like for the padding work-items of a `@kernel` function reduce_divergent_kernel(ids, red, scan, a, m) i = KI.get_local_id().x @@ -540,6 +606,125 @@ function sub_group_barrier_kernel(ids, scratch, out, ::Val{N}) where {N} return end +# Like `sub_group_barrier_kernel`, reading the next lane's write within segments of `width` +# lanes after a barrier of the segment. +function segment_barrier_kernel(ids, scratch, out, width, ::Val{N}) where {N} + i = KI.get_local_id().x + record_sub_group!(ids, i) + lane = KI.get_sub_group_local_id() + base = (KI.get_sub_group_id() - 1) * KI.get_max_sub_group_size() + first = fld(lane - 1, width) * width + other = base + first + mod1(lane - first + 1, min(width, KI.get_sub_group_size() - first)) + lm = KI.localmemory(Int32, N) + @inbounds lm[base + lane] = i + @inbounds scratch[base + lane] = -i + KI.sub_group_barrier(width) + @inbounds out[i, 1] = lm[other] + @inbounds out[i, 2] = scratch[other] + return +end + +# Every segment of `width` lanes executes the communication functions with that width a +# different number of times, and some return early, which needs +# `supports_partitionable_subgroups`. Each segment exchanges values through local memory after a +# barrier of the segment. `N` is the size of the work-group, of full sub-groups. +function partitioned_subgroups_kernel(ids, out, a, width, ::Val{N}) where {N} + i = KI.get_local_id().x + record_sub_group!(ids, i) + lane = KI.get_sub_group_local_id() + W = KI.get_max_sub_group_size() + pos = (lane - 1) % width + seg = (KI.get_sub_group_id() - 1) * (W ÷ width) + (lane - 1) ÷ width + val = @inbounds a[i] + acc = zero(val) + for _ in 1:(seg % 3 + 1) + acc += KI.shfl(val, pos + 2, width) + end + @inbounds out[1, i] = acc + seg % 4 == 1 && return + lm = KI.localmemory(Int32, N) + slot = (KI.get_sub_group_id() - 1) * W + lane + @inbounds lm[slot] = val + KI.sub_group_barrier(width) + @inbounds out[2, i] = lm[slot - pos + (pos + 1) % width] + @inbounds out[3, i] = KI.sub_group_reduce(+, val, width) + @inbounds out[4, i] = KI.sub_group_scan(+, val, width) + @inbounds out[5, i] = KI.sub_group_any(val > Int32(1900), width) + if KI.sub_group_all(val > Int32(500), width) + @inbounds out[6, i] = 1 + end + return +end + +# Segments of `width` lanes repeatedly exchange values through local and global memory, a +# different number of times each, with a barrier of the segment after writing and after +# reading: counts the values that weren't the neighbour's write of the same iteration. +function segment_exchange_kernel(errors, scratch, width, ::Val{N}) where {N} + i = KI.get_local_id().x + lane = KI.get_sub_group_local_id() + W = KI.get_max_sub_group_size() + pos = (lane - 1) % width + seg = (KI.get_sub_group_id() - 1) * (W ÷ width) + (lane - 1) ÷ width + slot = (KI.get_sub_group_id() - 1) * W + lane + other = slot - pos + (pos + 1) % width + other_lane = lane - pos + (pos + 1) % width + offset = (KI.get_group_id().x - 1) * N + lm = KI.localmemory(Int32, N) + err = Int32(0) + for it in Int32(1):Int32(seg % 7 + 1) + @inbounds lm[slot] = it * Int32(4096) + Int32(lane) + @inbounds scratch[offset + slot] = -(it * Int32(4096) + Int32(lane)) + KI.sub_group_barrier(width) + expected = it * Int32(4096) + Int32(other_lane) + err += (@inbounds(lm[other]) != expected) + (@inbounds(scratch[offset + other]) != -expected) + KI.sub_group_barrier(width) + end + @inbounds errors[offset + i] = err + return +end + +function partitioned_subgroups_testsuite(backend, AT, sg_size) + for w in segment_widths(sg_size) + w < sg_size || continue + n = 4 * sg_size + a = Int32[rand(Bool) ? rand(501:2000) : rand(1:2000) for _ in 1:n] + ids, out = AT(zeros(Int, 3, n)), AT(zeros(Int32, 6, n)) + kernel = KI.@launch backend launch = false partitioned_subgroups_kernel(ids, out, AT(a), w, Val(n)) + m = min(n, fld(KI.max_work_group_size(kernel), sg_size) * sg_size) + kernel(ids, out, AT(a), w, Val(n); workgroupsize = m) + KI.synchronize(backend) + out = Array(out) + ok = true + for (sg, g) in enumerate(observed_sub_groups(Array(ids)[:, 1:m])), (l, i) in enumerate(g) + lanes = segment_lanes(l, w, length(g)) + vals = a[g[lanes]] + pos = l - first(lanes) + 1 + seg = (sg - 1) * (sg_size ÷ w) + (l - 1) ÷ w + next = vals[mod1(pos + 1, length(vals))] + expected = zeros(Int32, 6) + expected[1] = (seg % 3 + 1) * next + if seg % 4 != 1 + expected[2:6] = [next, sum(vals), sum(vals[1:pos]), any(>(1900), vals), all(>(500), vals)] + end + ok &= out[:, i] == expected + end + @test ok + + groups = 64 + errors, scratch = AT(fill(Int32(-1), n * groups)), AT(zeros(Int32, n * groups)) + kernel = KI.@launch backend launch = false segment_exchange_kernel(errors, scratch, w, Val(n)) + m = min(n, fld(KI.max_work_group_size(kernel), sg_size) * sg_size) + for _ in 1:10 + kernel(errors, scratch, w, Val(n); workgroupsize = m, numgroups = groups) + KI.synchronize(backend) + errs = reshape(Array(errors), n, groups) + ok &= all(iszero, errs[1:m, :]) + end + @test ok + end + return +end + struct ShuffleStruct a::Float32 b::Int64 @@ -613,12 +798,27 @@ function subgroup_communication_testsuite(backend::KI.Backend, AT, sg_size) other = g[mod1(l + 1, length(g))] @test out[i, :] == [other, -other] end + + @testset "with a width of $w" for w in segment_widths(sg_size) + ids = AT(zeros(Int, 3, n)) + scratch = AT(zeros(Int32, n * sg_size)) + out = AT(zeros(Int32, n, 2)) + KI.@launch backend workgroupsize = n segment_barrier_kernel(ids, scratch, out, w, Val(n * sg_size)) + KI.synchronize(backend) + out = Array(out) + for g in observed_sub_groups(Array(ids)), (l, i) in enumerate(g) + lanes = segment_lanes(l, w, length(g)) + other = g[first(lanes) + mod1(l - first(lanes) + 2, length(lanes)) - 1] + @test out[i, :] == [other, -other] + end + end end @testset "votes" begin @testset "$n work-items" for n in sizes for pred in (falses(n), trues(n), [i % 3 == 1 for i in 1:n], [i == n for i in 1:n], rand(Bool, n)) vote_testsuite(backend, AT, sg_size, collect(pred)) + width_vote_testsuite(backend, AT, sg_size, collect(pred)) end end end @@ -731,6 +931,25 @@ function subgroup_communication_testsuite(backend::KI.Backend, AT, sg_size) collective_testsuite(backend, AT, ExclusiveScan(+, -0.0f0), Float32.(rand(1:100, n))) collective_testsuite(backend, AT, ExclusiveScan(+, -0.0f0), fill(-0.0f0, n)) end + + @testset "with a width, $n work-items" for n in sizes + for w in segment_widths(sg_size) + ops = Any[ + (+, Int32.(rand(1:100, n))), (max, Int64.(rand(-100:100, n))), + (+, Float32.(rand(n))), + (compose_affine, [(Int32(rand((-1, 1, 2))), Int32(rand(-5:5))) for _ in 1:n]), + ] + for (op, a) in ops + collective_testsuite(backend, AT, WidthReduce(op, w), a) + collective_testsuite(backend, AT, WidthScan(op, w), a) + end + collective_testsuite(backend, AT, WidthExclusiveScan(+, Int32(7), w), Int32.(rand(1:100, n))) + collective_testsuite( + backend, AT, WidthExclusiveScan(compose_affine, (Int32(2), Int32(3)), w), + [(Int32(rand((-1, 1, 2))), Int32(rand(-5:5))) for _ in 1:n] + ) + end + end end @testset "sub_group_reduce and sub_group_scan of divergent values" begin @@ -743,6 +962,12 @@ function subgroup_communication_testsuite(backend::KI.Backend, AT, sg_size) independent_sub_groups_testsuite(backend, AT, sg_size) end end + + if KI.supports_partitionable_subgroups(backend) + @testset "partitionable sub-groups" begin + partitioned_subgroups_testsuite(backend, AT, sg_size) + end + end return end @@ -875,9 +1100,11 @@ function interface_testsuite(backend::KI.Backend, AT) @test KI.supports_subgroups(b) isa Bool @test KI.supports_linear_subgroups(b) isa Bool @test KI.supports_independent_subgroups(b) isa Bool - # both imply sub-group support + @test KI.supports_partitionable_subgroups(b) isa Bool + # all imply sub-group support KI.supports_linear_subgroups(b) && @test KI.supports_subgroups(b) KI.supports_independent_subgroups(b) && @test KI.supports_subgroups(b) + KI.supports_partitionable_subgroups(b) && @test KI.supports_independent_subgroups(b) @test KI.supports_shuffle(b, Float32) isa Bool @test KI.functional(b) isa Union{Missing, Bool} @test KI.multiprocessor_count(b) isa Int diff --git a/lib/KernelInterface/test/runtests.jl b/lib/KernelInterface/test/runtests.jl index c590b1197..4978c67c1 100644 --- a/lib/KernelInterface/test/runtests.jl +++ b/lib/KernelInterface/test/runtests.jl @@ -44,9 +44,10 @@ KI.supports_shuffle(::WordShuffleBackend, ::Type{UInt32}) = true @test isempty(methods(stub)) end - # The votes have no fallbacks: backends that support sub-groups implement them. + # The votes have no fallbacks: backends that support sub-groups implement them. (Their + # variants with a `width` are built on them.) for vote in [KI.sub_group_any, KI.sub_group_all, KI.sub_group_ballot] - @test isempty(methods(vote)) + @test isempty(methods(vote, Tuple{Bool})) end # The shuffles only have the fallbacks that shuffle other types as `UInt32` words, which @@ -83,6 +84,7 @@ KI.supports_shuffle(::WordShuffleBackend, ::Type{UInt32}) = true @test !KI.supports_subgroups(PlainBackend()) @test !KI.supports_linear_subgroups(PlainBackend()) @test !KI.supports_independent_subgroups(PlainBackend()) + @test !KI.supports_partitionable_subgroups(PlainBackend()) # The primitive queries take an element type; only the zero-argument form has a # (forwarding) method, and it must reach the typed stub rather than recurse. diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index dd9033277..a4dbaa8a8 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -237,6 +237,9 @@ KI.supports_subgroups(::POCLBackend) = device_limits().sub_group_size > 0 # work-group barriers, so they need work-group-uniform control flow (see `POCLBackend`) KI.supports_linear_subgroups(backend::POCLBackend) = KI.supports_subgroups(backend) KI.supports_independent_subgroups(::POCLBackend) = false +# PoCL limitation: for the same reason, sub-groups can't be partitioned into segments that +# diverge +KI.supports_partitionable_subgroups(::POCLBackend) = false # PoCL limitation: every shuffle is a work-group barrier, also for each word of a struct; see # `KI.shfl_word_tuple`, which could exchange them at once with `cl_intel_spirv_subgroups`. From 8d7ce1fea3b81f92508be4591099457f21bea371 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sun, 11 Oct 2026 11:32:12 +0200 Subject: [PATCH 5/9] POCL: shuffle packed structs with vector shuffles Every shuffle on PoCL is a work-group barrier, so a struct packed into words cost a barrier per word: tile kernels shuffling an `SVector` were several times slower than ones using local memory. Exchange the words with the vector shuffles of `SPV_INTEL_subgroups`, a barrier per up to 16 words, which PoCL's CPU device accepts with `cl_intel_subgroups` since pocl/pocl#2394. Four shuffles of an `SVector{4, Float32}` over 4M work-items take 1.7 ms instead of 4.6 ms. KernelInterface gets `shuffle_source`, the lane a shuffle reads from, for backends that exchange the words with a shuffle from a lane. --- lib/KernelInterface/src/device.jl | 56 +++++++++++++++++++++---------- src/pocl/backend.jl | 28 ++++++++++++++-- src/pocl/compiler/compilation.jl | 12 +++++-- src/pocl/device/shuffle.jl | 10 ++++++ src/pocl/pocl.jl | 1 + test/runtests.jl | 29 ++++++++++++++-- 6 files changed, 111 insertions(+), 25 deletions(-) create mode 100644 src/pocl/device/shuffle.jl diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl index 3df9e7606..29a5a5681 100644 --- a/lib/KernelInterface/src/device.jl +++ b/lib/KernelInterface/src/device.jl @@ -528,12 +528,8 @@ end end # Shuffle the words a struct or tuple is packed into, with `op` shuffling a `UInt32`, one at a -# time. Not public: a backend may exchange all words at once instead, e.g. with a vector shuffle. -# -# PoCL limitation: every shuffle on PoCL is a work-group barrier, so one exchange of all words -# would make these shuffles several times faster there. PoCL has vector shuffles -# (`intel_sub_group_shuffle` on `uint2` to `uint16`), but its CPU device doesn't accept them -# from SPIR-V (`cl_intel_spirv_subgroups`), so no backend implements this yet. +# time. Not public: a backend may exchange several words at once instead, e.g. with vector +# shuffles from the lane `shuffle_source(op)`, as KernelAbstractions' POCL backend does. @inline shfl_word_tuple(op, words::NTuple{N, UInt32}) where {N} = map(op, words) """ @@ -588,37 +584,63 @@ for the work-items that execute them together, while `lane` may differ. shfl_words(ShuffleOp(shfl_xor, (mask, width)), val) # A `UInt32` within a segment is read with `shfl` from the computed lane -@inline function shfl(val::UInt32, lane::Integer, width::Integer) +@inline shfl(val::UInt32, lane::Integer, width::Integer) = shfl(val, shuffle_source(shfl, lane, width)) +@inline shfl_down(val::UInt32, offset::Integer, width::Integer) = + shfl(val, shuffle_source(shfl_down, offset, width)) +@inline shfl_up(val::UInt32, offset::Integer, width::Integer) = + shfl(val, shuffle_source(shfl_up, offset, width)) +@inline shfl_xor(val::UInt32, mask::Integer, width::Integer) = + shfl(val, shuffle_source(shfl_xor, mask, width)) + +# The lane of the sub-group (1-based, as an `Int32`) that the shuffle `op` reads from, e.g. for +# a backend's `shfl_word_tuple` that exchanges the words with `shfl` from a lane. Without a +# `width`, the shuffles read like those with the sub-group width. +@inline shuffle_source(op::ShuffleOp) = shuffle_source(op.f, op.args...) + +@inline function shuffle_source(::typeof(shfl), lane::Integer) + return lane % Int32 +end +@inline function shuffle_source(::typeof(shfl_down), offset::Integer) + l = get_sub_group_local_id(Int32) + # compared before adding, so that large offsets don't overflow + return ifelse(offset <= get_max_sub_group_size(Int32) - l, l + offset % Int32, l) +end +@inline function shuffle_source(::typeof(shfl_up), offset::Integer) + l = get_sub_group_local_id(Int32) + return ifelse(l > offset, l - offset % Int32, l) +end +@inline function shuffle_source(::typeof(shfl_xor), mask::Integer) + return ((get_sub_group_local_id(Int32) - Int32(1)) ⊻ (mask % Int32)) + Int32(1) +end + +@inline function shuffle_source(::typeof(shfl), lane::Integer, width::Integer) l0 = get_sub_group_local_id(Int32) - Int32(1) w = width % Int32 base = l0 & ~(w - Int32(1)) # `lane - 1` modulo the power of two `width`, before narrowing `lane` - return shfl(val, base + ((lane - 1) & (width - 1)) % Int32 + Int32(1)) + return base + ((lane - 1) & (width - 1)) % Int32 + Int32(1) end # The offsets are compared with the width before they are narrowed to `Int32`, so that large # offsets read from outside of the segment. -@inline function shfl_down(val::UInt32, offset::Integer, width::Integer) +@inline function shuffle_source(::typeof(shfl_down), offset::Integer, width::Integer) l0 = get_sub_group_local_id(Int32) - Int32(1) w = width % Int32 pos = l0 & (w - Int32(1)) inside = offset < width - pos - src = ifelse(inside, l0 + offset % Int32, l0) - return shfl(val, src + Int32(1)) + return ifelse(inside, l0 + offset % Int32, l0) + Int32(1) end -@inline function shfl_up(val::UInt32, offset::Integer, width::Integer) +@inline function shuffle_source(::typeof(shfl_up), offset::Integer, width::Integer) l0 = get_sub_group_local_id(Int32) - Int32(1) w = width % Int32 pos = l0 & (w - Int32(1)) inside = offset <= pos - src = ifelse(inside, l0 - offset % Int32, l0) - return shfl(val, src + Int32(1)) + return ifelse(inside, l0 - offset % Int32, l0) + Int32(1) end -@inline function shfl_xor(val::UInt32, mask::Integer, width::Integer) - l0 = get_sub_group_local_id(Int32) - Int32(1) - return shfl(val, (l0 ⊻ (mask % Int32)) + Int32(1)) +@inline function shuffle_source(::typeof(shfl_xor), mask::Integer, width::Integer) + return ((get_sub_group_local_id(Int32) - Int32(1)) ⊻ (mask % Int32)) + Int32(1) end """ diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index a4dbaa8a8..3b9a0651f 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -241,9 +241,6 @@ KI.supports_independent_subgroups(::POCLBackend) = false # diverge KI.supports_partitionable_subgroups(::POCLBackend) = false -# PoCL limitation: every shuffle is a work-group barrier, also for each word of a struct; see -# `KI.shfl_word_tuple`, which could exchange them at once with `cl_intel_spirv_subgroups`. - # The types `sub_group_shuffle` supports; other types are shuffled as words. `Float16` and # `Float64` are reported as unsupported on devices without `cl_khr_fp16` or `cl_khr_fp64`, # rather than shuffled as words: the overrides below are selected by type, not by device. @@ -333,6 +330,31 @@ end @device_override KI.shfl_xor(val::T, mask::Integer) where {T <: ShuffleTypes} = SPIRVIntrinsics.sub_group_shuffle_xor(val, mask) +# Every shuffle is a work-group barrier, so the words of a packed struct are exchanged with +# vector shuffles (`POCL.vector_shuffle`), a barrier per up to 16 words rather than per word. +# The last vector is padded to a size that exists (2, 3, 4, 8 or 16 words). +@generated function shuffle_word_vectors(words::NTuple{N, UInt32}, lane::Int32) where {N} + body = Any[Expr(:meta, :inline)] + results = Any[] + for first in 1:16:N + k = min(N - first + 1, 16) + if k == 1 + push!(results, :(KI.shfl(words[$first], lane))) + continue + end + n = k <= 4 ? k : k <= 8 ? 8 : 16 + elements = ((j < k ? :(VecElement(words[$(first + j)])) : :(VecElement(UInt32(0)))) for j in 0:(n - 1)) + v = Symbol(:v, first) + push!(body, :($v = POCL.vector_shuffle(($(elements...),), lane))) + append!(results, (:($v[$j].value) for j in 1:k)) + end + push!(body, :(($(results...),))) + return Expr(:block, body...) +end + +@device_override @inline KI.shfl_word_tuple(op::KI.ShuffleOp, words::NTuple{N, UInt32}) where {N} = + shuffle_word_vectors(words, KI.shuffle_source(op)) + @device_override KI.sub_group_any(pred::Bool) = SPIRVIntrinsics.sub_group_any(pred) @device_override KI.sub_group_all(pred::Bool) = SPIRVIntrinsics.sub_group_all(pred) diff --git a/src/pocl/compiler/compilation.jl b/src/pocl/compiler/compilation.jl index 919077cc0..60844c17d 100644 --- a/src/pocl/compiler/compilation.jl +++ b/src/pocl/compiler/compilation.jl @@ -173,8 +173,10 @@ keyword arguments. Besides those of `CompilerConfig` (`kernel`, `name`, `always_ compilation fail or crash the driver's compiler, while disabling them relies on integer compare-and-swap for the fallback. - `extensions`: SPIR-V extensions to enable, as a `--spirv-ext` specifier (e.g. - `"+SPV_KHR_expect_assume"`), in addition to those the atomics need and to - `SPV_KHR_no_integer_wrap_decoration`. `nothing`, the default, enables no others. + `"+SPV_KHR_expect_assume"`), in addition to those the atomics need, to + `SPV_KHR_no_integer_wrap_decoration`, and to `SPV_INTEL_subgroups` on devices with + `cl_intel_subgroups` (for the vector shuffles of packed structs). `nothing`, the default, + enables no others. Other keyword arguments are passed on to `SPIRVCompilerTarget`. """ @@ -213,6 +215,12 @@ end error("$sub_group_size is not a valid sub-group size for this device.") end + # the vector shuffles of `KI.shfl_word_tuple` (see `POCLKernels`); `SPV_INTEL_subgroups` + # is only declared by modules that use it + if "cl_intel_subgroups" in dev.extensions + extensions = extensions === nothing ? "+SPV_INTEL_subgroups" : "+SPV_INTEL_subgroups,$extensions" + end + # create GPUCompiler objects target = SPIRVCompilerTarget(; supports_fp16, supports_fp64, atomics, diff --git a/src/pocl/device/shuffle.jl b/src/pocl/device/shuffle.jl new file mode 100644 index 000000000..5881eca1c --- /dev/null +++ b/src/pocl/device/shuffle.jl @@ -0,0 +1,10 @@ +# The vector shuffles of `SPV_INTEL_subgroups` (`intel_sub_group_shuffle` on `uint2` to +# `uint16`), which PoCL's CPU device implements with a single work-group barrier, like a +# shuffle of a scalar. The standard shuffles only take scalars. +for N in (2, 3, 4, 8, 16) + V = NTuple{N, VecElement{UInt32}} + @eval @device_function @inline vector_shuffle(x::$V, lane::Int32) = @builtin_ccall( + "__spirv_SubgroupShuffleINTEL", $V, ($V, UInt32), x, (lane - Int32(1)) % UInt32, + convergent = true + ) +end diff --git a/src/pocl/pocl.jl b/src/pocl/pocl.jl index cd0b3b055..37c3e3b82 100644 --- a/src/pocl/pocl.jl +++ b/src/pocl/pocl.jl @@ -116,6 +116,7 @@ include("device/array.jl") include("device/quirks.jl") include("device/runtime.jl") include("device/random.jl") +include("device/shuffle.jl") include("compiler/compilation.jl") include("compiler/execution.jl") diff --git a/test/runtests.jl b/test/runtests.jl index d96538c0f..81c8e8fa0 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -75,11 +75,13 @@ end fadd_f64_global = true, fadd_f64_local = true ) @test POCL.compiler_config(dev).target.atomics == atomics - # explicitly requested extensions are kept, in addition to the default ones - @test POCL.compiler_config(dev).target.extensions == "+SPV_KHR_no_integer_wrap_decoration" + # explicitly requested extensions are kept, in addition to the default ones and the one + # for vector shuffles + @test POCL.compiler_config(dev).target.extensions == + "+SPV_KHR_no_integer_wrap_decoration,+SPV_INTEL_subgroups" config = POCL.compiler_config(dev; extensions = "+SPV_KHR_expect_assume") @test config.target.extensions == - "+SPV_KHR_no_integer_wrap_decoration,+SPV_KHR_expect_assume" + "+SPV_KHR_no_integer_wrap_decoration,+SPV_INTEL_subgroups,+SPV_KHR_expect_assume" @test config.target.atomics == atomics @test POCL.compiler_config(dev; extensions = nothing).target == POCL.compiler_config(dev).target @@ -347,6 +349,27 @@ end end end +# structs are packed into words, which are exchanged with as few vector shuffles as possible +function packed_shuffle_kernel(out, x) + KI = KernelAbstractions.KernelInterface + i = KI.get_global_id().x + out[i] = KI.shfl(x[i], KI.get_sub_group_local_id() % KI.get_sub_group_size() + 1) + return +end +@testset "POCL vector shuffles of packed structs" begin + W = POCL.device_limits().sub_group_size + for n in (1, 2, 5, 16, 17, 18, 33) + x = [ntuple(k -> UInt32(1000i + k), n) for i in 1:2W] + out = similar(x) + ir = sprint() do io + @device_code_llvm io = io debuginfo = :none @opencl local_size = 2W global_size = 2W packed_shuffle_kernel(out, x) + end + @test out == [x[fld(i - 1, W) * W + mod1(i - fld(i - 1, W) * W + 1, W)] for i in 1:2W] + # vectors of up to 16 words, and a single word as a scalar + @test count(r"call [^\n]*SubgroupShuffleINTEL", ir) == fld(n, 16) + (n % 16 >= 2) + end +end + # Julia doesn't turn a splat of more than 32 elements into a direct call, so a launch with # many arguments allocates unless every layer passes them on as a tuple @testset "POCL launch with many arguments" begin From c46beaea3df63cd83df6639cb445d6fac18f8f28 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sat, 10 Oct 2026 20:29:09 +0200 Subject: [PATCH 6/9] [TEMP] Test the back ends' sub-group branches on Buildkite --- .buildkite/pipeline.yml | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 808d8b9d3..ac6a29b58 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -26,7 +26,7 @@ steps: Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()), PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3 echo "--- :julia: Developing CUDA" - julia .buildkite/develop_backend.jl https://github.com/JuliaGPU/CUDA.jl ka-0.10 \ + julia .buildkite/develop_backend.jl https://github.com/JuliaGPU/CUDA.jl vc/ki-subgroup-ops \ CUDACore CUDATools lib/cusolver lib/curand lib/cublas lib/cufft lib/cusparse lib/cupti lib/nvml lib/custatevec lib/cudnn lib/cutensor lib/cutensornet || exit 3 julia -e 'println("+++ :julia: Running tests") @@ -64,7 +64,7 @@ steps: Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()), PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3 echo "--- :julia: Developing Metal" - julia .buildkite/develop_backend.jl https://github.com/JuliaGPU/Metal.jl ka-0.10 || exit 3 + julia .buildkite/develop_backend.jl https://github.com/JuliaGPU/Metal.jl vc/ki-subgroup-ops || exit 3 julia -e 'println("+++ :julia: Running tests") using Pkg check = joinpath(pwd(), ".buildkite", "check_sources.jl") @@ -102,7 +102,7 @@ steps: julia -e 'println("--- :julia: Developing oneAPI") using Pkg Pkg.add(url="https://github.com/JuliaGPU/AcceleratedKernels.jl", rev="main")' || exit 3 - julia .buildkite/develop_backend.jl https://github.com/JuliaGPU/oneAPI.jl ka-0.10 || exit 3 + julia .buildkite/develop_backend.jl https://github.com/JuliaGPU/oneAPI.jl tb/ki-subgroup-ops || exit 3 julia -e 'println("+++ :julia: Running tests") using Pkg @@ -142,7 +142,7 @@ steps: julia -e 'println("--- :julia: Developing AMDGPU") using Pkg Pkg.add(url="https://github.com/JuliaGPU/AcceleratedKernels.jl", rev="main")' || exit 3 - julia .buildkite/develop_backend.jl https://github.com/JuliaGPU/AMDGPU.jl ka-0.10 || exit 3 + julia .buildkite/develop_backend.jl https://github.com/JuliaGPU/AMDGPU.jl vc/ki-subgroup-ops || exit 3 julia -e 'println("+++ :julia: Running tests") using Pkg @@ -180,7 +180,7 @@ steps: Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()), PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3 echo "--- :julia: Developing OpenCL" - julia .buildkite/develop_backend.jl https://github.com/JuliaGPU/OpenCL.jl ka-0.10 lib/intrinsics || exit 3 + julia .buildkite/develop_backend.jl https://github.com/JuliaGPU/OpenCL.jl tb/ki-subgroup-ops lib/intrinsics || exit 3 julia -e 'println("+++ :julia: Running tests") using Pkg From ae616559c5b37c0335666a35465691b4377a9fb4 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Sat, 10 Oct 2026 18:58:27 +0200 Subject: [PATCH 7/9] Add `@groupreduce` and `@groupscan` Work-group reductions and scans for `@kernel`, as collectives like `@synchronize`: they are statements of their own that every work-item of the workgroup reaches, and the split of the kernel body makes the work-items that pad a partial workgroup take part, contributing the neutral element without evaluating the value. The value is converted to the type of `neutral` before the call, so that a union-typed value doesn't split the collective, and `neutral` is evaluated once. The upper bound of a dynamic workgroup size is the keyword `groupsize`. `@groupreduce` needs a commutative operator, `@groupscan` only an associative one, in the order of `@index(Local, Linear)`. Both use local memory. Co-authored-by: Tim Besard --- docs/src/api.md | 7 + docs/src/kernelinterface.md | 4 +- docs/src/kernels.md | 15 +++ src/KernelAbstractions.jl | 3 + src/groupreduction.jl | 255 ++++++++++++++++++++++++++++++++++++ src/macros.jl | 16 ++- test/groupreduce.jl | 254 +++++++++++++++++++++++++++++++++++ test/testsuite.jl | 5 + 8 files changed, 557 insertions(+), 2 deletions(-) create mode 100644 src/groupreduction.jl create mode 100644 test/groupreduce.jl diff --git a/docs/src/api.md b/docs/src/api.md index 0d1d77696..f5f3e3845 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -15,6 +15,13 @@ @ndrange ``` +### Reductions and scans + +```@docs +@groupreduce +@groupscan +``` + ## Host language !!! note diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md index 31015fde9..f42467194 100644 --- a/docs/src/kernelinterface.md +++ b/docs/src/kernelinterface.md @@ -213,7 +213,9 @@ used anywhere. kernel languages built on it. A kernel is either a KernelInterface kernel or a KernelAbstractions `@kernel`, never a mix: `@kernel` adds padding work-items to partial work-groups, which skip the kernel's body, so KernelInterface's sub-group functions called - in a `@kernel` aren't reached by all work-items. Host code may use the queries of both. + in a `@kernel` aren't reached by all work-items. A `@kernel` uses KernelAbstractions' + collectives instead, e.g. [`@groupreduce`](@ref KernelAbstractions.@groupreduce). Host + code may use the queries of both. ```@docs get_sub_group_size diff --git a/docs/src/kernels.md b/docs/src/kernels.md index 5d6df509f..75ea30e9d 100644 --- a/docs/src/kernels.md +++ b/docs/src/kernels.md @@ -213,6 +213,21 @@ statements, and [`@uniform`](@ref) evaluates an expression outside the work-item can be reused across `@synchronize` statements. For scratch storage that does not need to survive across `@synchronize`, an `MArray` can be used instead. +## Reductions and scans + +[`@groupreduce`](@ref) and [`@groupscan`](@ref) reduce and scan values over the workgroup. +Like [`@synchronize`](@ref), they are collectives: all work-items of the workgroup have to +reach them, not in a branch or loop that only some of them take, and they have to be used as +statements of their own, e.g. `res = @groupreduce(+, val, zero(T))`. Work-items that pad a +partial workgroup take part as well, contributing the neutral element. + +A kernel is either a `@kernel` or a kernel written against +[KernelInterface](@ref kernelinterface), never a mix: don't call KernelInterface's device +functions, such as its sub-group shuffles and votes, in a `@kernel`. They have to be executed +by all work-items of a sub-group, but in a `@kernel` with the default bounds checking, every +statement other than the collectives above only runs on the work-items that are part of the +`ndrange`. Host code can use KernelInterface's queries, e.g. to pick a kernel. + ## Launching kernels Construct a kernel by calling the kernel function on a backend and optional static sizes, then diff --git a/src/KernelAbstractions.jl b/src/KernelAbstractions.jl index f7bf3725f..b7cb5c36d 100644 --- a/src/KernelAbstractions.jl +++ b/src/KernelAbstractions.jl @@ -42,6 +42,8 @@ and then invoked on the arguments. - [`@uniform`](@ref) - [`@synchronize`](@ref) - [`@print`](@ref) +- [`@groupreduce`](@ref) +- [`@groupscan`](@ref) # Kernel constructor @@ -680,6 +682,7 @@ function __workitems_iterspace end end include("macros.jl") +include("groupreduction.jl") include("spawn.jl") include("foreach_index.jl") diff --git a/src/groupreduction.jl b/src/groupreduction.jl new file mode 100644 index 000000000..1b92d9958 --- /dev/null +++ b/src/groupreduction.jl @@ -0,0 +1,255 @@ +### +# Work-group reductions and scans +# - @groupreduce +# - @groupscan +### + +export @groupreduce, @groupscan + +""" + @groupreduce(op, val, neutral; groupsize) + +Reduce `val` over all work-items of the workgroup with the binary operator `op`, and return +the result on every work-item. `op` has to be associative and commutative: the values are +combined in an unspecified order. `neutral` has to be its neutral element +(`op(neutral, x) == op(x, neutral) == x`). For example, to find the smallest value and its +index, reduce `(value, index)` pairs with an operator that breaks ties by the index, such +as `min` on tuples, with the neutral element `(typemax(T), typemax(Int))`. + +The result has the type of `neutral`, which `val` is converted to: `op(x, y)` has to return +that type `T = typeof(neutral)` for arguments of type `T`. + +`@groupreduce` is a collective, like [`@synchronize`](@ref): it has to be used as a +statement on its own (`res = @groupreduce(op, val, neutral)`), and reached by all +work-items of the workgroup, not in a branch or loop that only some of them execute, or +after an early `return`. Work-items that pad a partial workgroup (outside of the `ndrange`) +take part as well: they contribute `neutral`, without evaluating `val`. So `op` and +`neutral` have to be the same for the whole workgroup, and computable on those work-items. + +The reduction stores values in local memory, so its size has to be known at compile time: +either the kernel has a static workgroup size, or the keyword `groupsize` gives an upper +bound of the workgroup size (a constant, e.g. a literal or a type parameter of the kernel). + +```julia +@kernel function sum_kernel!(out, @Const(x)) + i = @index(Global) + res = @groupreduce(+, x[i], zero(eltype(out)); groupsize = 1024) + if @index(Local, Linear) == 1 + out[@index(Group, Linear)] = res + end +end +``` +""" +macro groupreduce(args...) + op, val, neutral, options = parse_collective("@groupreduce", args, (:groupsize,)) + bound = collective_bound(Base.get(options, :groupsize, nothing)) + return __collective_call(neutral, val) do neutral, val + :($__groupreduce($(esc(:__ctx__)), $(esc(op)), $val, $neutral, $bound)) + end +end + +""" + @groupscan(op, val, neutral; groupsize, inclusive = true) + +Scan `val` over the work-items of the workgroup with the binary operator `op`, in the order +of `@index(Local, Linear)`: the work-item with local index `i` gets +`op(...op(op(val₁, val₂), val₃)..., valᵢ)` (inclusive), or the same up to `valᵢ₋₁` and +`neutral` for the first work-item (with `inclusive = false`). `op` has to be associative, +but needn't be commutative, and `neutral` its neutral element +(`op(neutral, x) == op(x, neutral) == x`). The result has the type of `neutral`, which `val` +is converted to, and `op` has to return that type. + +Like [`@groupreduce`](@ref), `@groupscan` is a collective that has to be used as a statement +on its own and reached by all work-items of the workgroup; padding work-items contribute +`neutral`. The scan stores two values per work-item in local memory, sized like for +`@groupreduce`: either the kernel has a static workgroup size, or the keyword `groupsize` +gives an upper bound of the workgroup size. `inclusive` has to be a constant as well. + +For example, to compact the elements of `x` that satisfy `pred` within each workgroup: + +```julia +@kernel function compact!(out, counts, @Const(x), pred) + i = @index(Global, Linear) + keep = pred(x[i]) + offset = @groupscan(+, Int32(keep), Int32(0); inclusive = false) + total = @groupreduce(+, Int32(keep), Int32(0)) + base = (@index(Group, Linear) - 1) * prod(@groupsize()) + if keep + out[base + offset + 1] = x[i] + end + if @index(Local, Linear) == 1 + counts[@index(Group, Linear)] = total + end +end +``` +""" +macro groupscan(args...) + op, val, neutral, options = parse_collective("@groupscan", args, (:groupsize, :inclusive)) + bound = collective_bound(Base.get(options, :groupsize, nothing)) + inclusive = Base.get(options, :inclusive, true) + return __collective_call(neutral, val) do neutral, val + :($__groupscan($(esc(:__ctx__)), $(esc(op)), $val, $neutral, $bound, Val($(esc(inclusive))))) + end +end + +collective_bound(::Nothing) = :($__static_groupsize($(esc(:__ctx__)))) +collective_bound(groupsize) = :(Val($(esc(groupsize)))) + +# Convert `val` to the type of `neutral` *before* the call of the collective. The type of +# `val` may differ between the work-items, e.g. `Union{Float32, Float64}` for an accumulator +# that only some work-items added a `Float64` to, or because padding work-items contribute +# `neutral` instead of `val` (see `mask_collective`). Julia union-splits a call with such an +# argument into one call per type, so the work-items would execute different copies of the +# collective, and its barriers and shuffles. Only the `convert` may be split this way. +function __collective_call(f, neutral, val) + n, v = gensym(:neutral), gensym(:val) + return quote + let $v = $(esc(val)), $n = $(esc(neutral)) + $(f(n, :($convert($typeof($n), $v)))) + end + end +end + +# The `op`, `val` and `neutral` of a collective's arguments, and its `key = value` options +# (also after a `;`), which have to be among `keys`. +function parse_collective(name, args, keys) + positional = Any[] + options = Dict{Symbol, Any}() + function option!(ex) + key = ex.args[1] + (key isa Symbol && key in keys) || + error("$name: unknown option `$key`, expected one of $(join(keys, ", "))") + haskey(options, key) && error("$name: option `$key` given more than once") + options[key] = ex.args[2] + return + end + for arg in args + if isexpr(arg, :parameters) + foreach(option!, arg.args) + elseif isexpr(arg, :(=)) || isexpr(arg, :kw) + option!(arg) + else + push!(positional, arg) + end + end + length(positional) == 3 || + error("$name expects `op`, `val` and `neutral`, and the options $(join(keys, ", ")) as keywords") + return positional..., options +end + +const COLLECTIVES = (Symbol("@groupreduce"), Symbol("@groupscan")) + +# Whether `expr` is a collective that all work-items of a workgroup take part in. +is_collective(expr) = any(name -> is_macrocall(expr, name), COLLECTIVES) + +# A collective that is used as a statement: `@groupreduce(...)` or `lhs = @groupreduce(...)`. +function is_collective_stmt(stmt) + is_collective(stmt) && return true + isexpr(stmt, :(=)) && is_collective(stmt.args[2]) || return false + lhs = stmt.args[1] + lhs isa Symbol && return true + isexpr(lhs, :(::)) && lhs.args[1] isa Symbol && return true + isexpr(lhs, :tuple) && all(x -> x isa Symbol, lhs.args) && return true + return false +end + +collective_error(expr) = error( + "`$(expr.args[1])` must be used as a statement of its own, " * + "e.g. `res = $(expr.args[1])(op, val, neutral)`, found `$(expr)`" +) + +# Rewrite a collective statement, so that padding work-items contribute the neutral element +# instead of evaluating the value. `neutral` is evaluated once, before the collective. +function mask_collective(stmt) + if isexpr(stmt, :(=)) + binding, call = mask_collective(stmt.args[2]).args + return Expr(:block, binding, Expr(:(=), stmt.args[1], call)) + end + # `args[2]` is the macro's `LineNumberNode`, options may come first after a `;` + args = copy(stmt.args) + i = 3 + while i <= length(args) && (isexpr(args[i], :parameters) || args[i] isa LineNumberNode) + i += 1 + end + length(args) >= i + 2 || error("`$(args[1])` expects `op`, `val` and `neutral`") + val, neutral = args[i + 1], args[i + 2] + n = gensym(:neutral) + args[i + 1] = :(__active_lane__ ? $val : $n) + args[i + 2] = n + return Expr(:block, :($n = $neutral), Expr(:macrocall, args...)) +end + +# The workgroup size as a `Val`, if it is static. +@inline __static_groupsize(ctx::CompilerMetadata) = __static_groupsize(__iterspace(ctx)) +@inline __static_groupsize(::NDRange{N, B, W}) where {N, B, W <: StaticSize} = Val(prod(get(W))) +@inline __static_groupsize(::NDRange) = throw( + ArgumentError( + "group reductions and scans require a static workgroup size or an upper bound of it" + ) +) + +# The largest power of two smaller than `n`, or 0. +@inline function __prevpow2(n::T) where {T <: Integer} + n <= one(T) && return zero(T) + return one(T) << (8 * sizeof(T) - 1 - leading_zeros(n - one(T))) +end + +@inline function __groupreduce(ctx, op, val, neutral::T, ::Val{N}) where {T, N} + n = prod(groupsize(ctx)) + n <= N || throw(ArgumentError("@groupreduce: the workgroup size exceeds the given upper bound")) + storage = KI.localmemory(T, Val(N)) + res = __groupreduce_tree(op, convert(T, val), storage, __index_Local_Linear(ctx), n) + # all work-items have to read the result before the storage can be reused + KI.barrier() + return res +end + +# Tree reduction in local memory, folding the upper half of the values onto the lower half. +@inline function __groupreduce_tree(op, val, storage, lid, n) + @inbounds storage[lid] = val + KI.barrier() + s = __prevpow2(n) + while s > 0 + if lid <= s && lid + s <= n + @inbounds storage[lid] = op(storage[lid], storage[lid + s]) + end + KI.barrier() + s >>= 1 + end + return @inbounds storage[1] +end + +@inline function __groupscan(ctx, op, val, neutral::T, ::Val{N}, ::Val{inclusive}) where {T, N, inclusive} + n = prod(groupsize(ctx)) + n <= N || throw(ArgumentError("@groupscan: the workgroup size exceeds the given upper bound")) + # two buffers, the scan reads from one and writes to the other + storage = KI.localmemory(T, Val(2 * N)) + lid = __index_Local_Linear(ctx) + + # Hillis-Steele: after the step with distance `d`, every work-item holds the scan of the + # (up to) `2d` values ending at its own + src = 0 + @inbounds storage[lid] = convert(T, val) + KI.barrier() + d = 1 + while d < n + x = @inbounds storage[src + lid] + if lid > d + x = op(@inbounds(storage[src + lid - d]), x) + end + @inbounds storage[(N - src) + lid] = x + KI.barrier() + src = N - src + d <<= 1 + end + + if inclusive + res = @inbounds storage[src + lid] + else + res = lid == 1 ? neutral : @inbounds storage[src + lid - 1] + end + # all work-items have to read their result before the storage can be reused + KI.barrier() + return res +end + diff --git a/src/macros.jl b/src/macros.jl index ccbc3bc7e..48774b0b6 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -294,10 +294,12 @@ function is_scope_construct(expr::Expr) # expr.head === :let end +# Whether `stmt` contains a `@synchronize`, or a collective like `@groupreduce` that all +# work-items of the workgroup have to reach as well. function find_sync(stmt) result = Ref(false) postwalk(stmt) do expr - result[] |= is_sync(expr) + result[] |= is_sync(expr) || is_collective(expr) expr end return result[] @@ -329,6 +331,17 @@ function split(stmts) continue end + if is_collective_stmt(stmt) + # executed by all work-items, the padding ones contribute the neutral element + loop = WorkgroupLoop(current, allocations, false, nothing) + push!(new_stmts, emit(loop)) + allocations = Any[] + current = Any[] + take_line!(new_stmts) + push!(new_stmts, mask_collective(stmt)) + continue + end + has_sync = find_sync(stmt) if has_sync loop = WorkgroupLoop(current, allocations, is_sync(stmt), line) @@ -348,6 +361,7 @@ function split(stmts) recurse(x) = x function recurse(expr::Expr) expr = unblock_lines(expr) + is_collective(expr) && collective_error(expr) if expr.head in (:if, :elseif) && find_sync(expr) return split_branches(expr, recurse) elseif is_scope_construct(expr) && any(find_sync, expr.args) diff --git a/test/groupreduce.jl b/test/groupreduce.jl new file mode 100644 index 000000000..3989578b0 --- /dev/null +++ b/test/groupreduce.jl @@ -0,0 +1,254 @@ +# one result per workgroup, written by every work-item, to check that all of them get it +@kernel function groupreduce_static!(out, @Const(x), op, neutral) + i = @index(Global, Linear) + res = @groupreduce(op, x[i], neutral) + out[i] = res +end + +@kernel function groupreduce_bound!(out, @Const(x), op, neutral, ::Val{N}) where {N} + i = @index(Global, Linear) + res = @groupreduce(op, x[i], neutral; groupsize = N) + out[i] = res +end + +# the same call site, and thus local memory, reused in a loop and in a branch +@kernel function groupreduce_loop!(out, @Const(x)) + i = @index(Global, Linear) + acc = zero(eltype(out)) + for k in 1:3 + res = @groupreduce(+, k * x[i], zero(eltype(out))) + acc += res + end + if true + m = @groupreduce max x[i] typemin(eltype(out)) + end + out[i] = acc + m +end + +@kernel function groupreduce_cartesian!(out, @Const(x)) + I = @index(Global, Cartesian) + res = @groupreduce(+, x[I], zero(eltype(out))) + out[I] = res +end + +@kernel unsafe_indices = true function groupreduce_unsafe!(out, @Const(x)) + i = @index(Global, Linear) + val = i <= length(x) ? x[i] : zero(eltype(out)) + res = @groupreduce(+, val, zero(eltype(out))) + if i <= length(out) + out[i] = res + end +end + +# `val` of another type than `neutral`: the padding work-items contribute `neutral`, so the +# value is a `Union` of both types, and the call of the collective must not be union-split +@kernel function groupreduce_mixed!(out, @Const(x)) + i = @index(Global, Linear) + res = @groupreduce(+, x[i], Int64(0)) + out[i] = res +end + +@kernel function groupscan_mixed!(out, @Const(x)) + i = @index(Global, Linear) + res = @groupscan(+, x[i], 0) + out[i] = res +end + +# the composition of affine maps `x -> a * x + b`, first `f` then `g`: associative, but not +# commutative, so that the scans have to combine the values in order +compose(f, g) = (g[1] * f[1], g[1] * f[2] + g[2]) +const affine_identity = (1, 0) + +@kernel function groupscan!(out, @Const(x), op, neutral, ::Val{I}) where {I} + i = @index(Global, Linear) + res = @groupscan(op, x[i], neutral; inclusive = I) + out[i] = res +end + +@kernel function groupscan_bound!(out, @Const(x), op, neutral, ::Val{N}, ::Val{I}) where {N, I} + i = @index(Global, Linear) + res = @groupscan(op, x[i], neutral; groupsize = N, inclusive = I) + out[i] = res +end + +# the scan in the order of the local linear index, with padding in the middle of a workgroup +@kernel function groupscan_cartesian!(out, @Const(x)) + I = @index(Global, Cartesian) + res = @groupscan(compose, x[I], affine_identity) + out[I] = res +end + +# the same call site in a loop, and an exclusive scan of the counts +@kernel function groupscan_loop!(out, @Const(x)) + i = @index(Global, Linear) + acc = 0 + for k in 1:3 + res = @groupscan(+, k * x[i], 0) + acc += res + end + excl = @groupscan (+) x[i] 0 inclusive = false + out[i] = acc + excl +end + +# reference: the reduction of each workgroup of `groupsize` consecutive elements +function groupwise(op, x, groupsize) + return [reduce(op, x[((cld(i, groupsize) - 1) * groupsize + 1):min(cld(i, groupsize) * groupsize, end)]) for i in eachindex(x)] +end + +# reference: the scan of each workgroup of `groupsize` consecutive elements +function groupwise_scan(op, x, groupsize, neutral, inclusive) + out = similar(x, typeof(neutral)) + for first in 1:groupsize:length(x) + group = first:min(first + groupsize - 1, length(x)) + acc = neutral + for i in group + inclusive || (out[i] = acc) + acc = op(acc, x[i]) + inclusive && (out[i] = acc) + end + end + return out +end + +# `neutral` is evaluated once per work-item, also by the padding work-items +counted_neutral() = 0 +count_symbol(ex::Expr, sym) = sum(arg -> count_symbol(arg, sym), ex.args; init = 0) +count_symbol(ex, sym) = Int(ex === sym) + +function groupreduce_testsuite(backend, AT) + b = backend() + types = (Int32, Int64, Float32) + @testset "$T, $(nameof(typeof(op)))" for T in types, (op, neutral) in ((+, zero(T)), (max, typemin(T))) + for (groupsize, n) in ((64, 64), (64, 256), (32, 100), (256, 1000), (7, 23), (1, 3)) + x = T.(rand(1:100, n)) + out = AT(fill(zero(T), n)) + groupreduce_static!(b, groupsize)(out, AT(x), op, neutral; ndrange = n) + @test Array(out) == groupwise(op, x, groupsize) + + fill!(out, zero(T)) + groupreduce_bound!(b)(out, AT(x), op, neutral, Val(256); ndrange = n, workgroupsize = groupsize) + @test Array(out) == groupwise(op, x, groupsize) + end + end + + @testset "argmin" begin + # (value, index) pairs, the smallest value with the smallest index + for (groupsize, n) in ((64, 256), (32, 100), (7, 23)) + x = [(Float32(rand(1:20)), Int32(i)) for i in 1:n] + neutral = (Inf32, typemax(Int32)) + out = AT(fill(neutral, n)) + groupreduce_static!(b, groupsize)(out, AT(x), min, neutral; ndrange = n) + @test Array(out) == groupwise(min, x, groupsize) + end + end + + @testset "loop" begin + x = rand(1:100, 100) + out = AT(zeros(Int, 100)) + groupreduce_loop!(b, 64)(out, AT(x); ndrange = 100) + @test Array(out) == 6 .* groupwise(+, x, 64) .+ groupwise(max, x, 64) + end + + @testset "cartesian" begin + x = rand(1:100, 10, 12) + out = AT(zeros(Int, 10, 12)) + groupreduce_cartesian!(b, (4, 8))(out, AT(x); ndrange = size(x)) + ref = similar(x) + for I in CartesianIndices(x) + g = (cld(I[1], 4) - 1) * 4 .+ (1:4), (cld(I[2], 8) - 1) * 8 .+ (1:8) + ref[I] = sum(x[intersect(g[1], axes(x, 1)), intersect(g[2], axes(x, 2))]) + end + @test Array(out) == ref + end + + @testset "unsafe_indices" begin + x = rand(1:100, 100) + out = AT(zeros(Int, 100)) + groupreduce_unsafe!(b, 64)(out, AT(x); ndrange = 128) + @test Array(out) == groupwise(+, x, 64) + end + + @testset "mixed types" begin + x = Int32.(rand(1:100, 100)) + out = AT(zeros(Int64, 100)) + groupreduce_mixed!(b, 64)(out, AT(x); ndrange = 100) + @test Array(out) == groupwise(+, Int64.(x), 64) + end + + @testset "@groupscan" begin + @testset "inclusive = $I" for I in (true, false) + for (groupsize, n) in ((64, 64), (64, 256), (32, 100), (256, 1000), (7, 23), (1, 3)) + x = rand(1:100, n) + out = AT(zeros(Int, n)) + groupscan!(b, groupsize)(out, AT(x), +, 0, Val(I); ndrange = n) + @test Array(out) == groupwise_scan(+, x, groupsize, 0, I) + + y = [(rand((-1, 1, 2)), rand(-5:5)) for _ in 1:n] + out = AT(fill((0, 0), n)) + groupscan_bound!(b)(out, AT(y), compose, affine_identity, Val(256), Val(I); ndrange = n, workgroupsize = groupsize) + @test Array(out) == groupwise_scan(compose, y, groupsize, affine_identity, I) + end + end + + @testset "cartesian" begin + x = [(rand((-1, 1, 2)), rand(-5:5)) for _ in 1:10, _ in 1:12] + out = AT(fill((0, 0), 10, 12)) + groupscan_cartesian!(b, (4, 8))(out, AT(x); ndrange = size(x)) + ref = similar(x) + for gi in 1:4:10, gj in 1:8:12 + acc = affine_identity + # local linear order is column-major within the workgroup + for j in gj:(gj + 7), i in gi:(gi + 3) + (i <= 10 && j <= 12) || continue + acc = compose(acc, x[i, j]) + ref[i, j] = acc + end + end + @test Array(out) == ref + end + + @testset "loop" begin + x = rand(1:100, 100) + out = AT(zeros(Int, 100)) + groupscan_loop!(b, 64)(out, AT(x); ndrange = 100) + @test Array(out) == 6 .* groupwise_scan(+, x, 64, 0, true) .+ groupwise_scan(+, x, 64, 0, false) + end + + @testset "mixed types" begin + x = Int32.(rand(1:100, 100)) + out = AT(zeros(Int, 100)) + groupscan_mixed!(b, 64)(out, AT(x); ndrange = 100) + @test Array(out) == groupwise_scan(+, Int.(x), 64, 0, true) + end + end + + @testset "errors" begin + @test_throws "must be used as a statement" @macroexpand @kernel function f(y, x) + i = @index(Global) + y[i] = @groupreduce(+, x[i], 0) + end + @test_throws "unknown option" @macroexpand @kernel function f(y, x) + res = @groupreduce(+, x[1], 0; foo = 1) + end + @test_throws "given more than once" @macroexpand @kernel function f(y, x) + res = @groupreduce(+, x[1], 0; groupsize = 32, groupsize = 64) + end + # the upper bound of the workgroup size is a keyword + @test_throws "expects `op`, `val` and `neutral`" @macroexpand @kernel function f(y, x) + res = @groupreduce(+, x[1], 0, 64) + end + @test_throws "unknown option" @macroexpand @kernel function f(y, x) + res = @groupreduce(+, x[1], 0; subgroups = true) + end + end + + @testset "neutral evaluated once" begin + ex = @macroexpand @kernel function f(y, x) + i = @index(Global, Linear) + res = @groupreduce(+, x[i], counted_neutral()) + y[i] = res + end + @test count_symbol(ex, :counted_neutral) == 1 + end + return +end diff --git a/test/testsuite.jl b/test/testsuite.jl index 541028ad8..375407df2 100644 --- a/test/testsuite.jl +++ b/test/testsuite.jl @@ -45,6 +45,7 @@ include("convert.jl") include("specialfunctions.jl") include("random.jl") include("spawn.jl") +include("groupreduce.jl") function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{String}()) @conditional_testset "Unittests" skip_tests begin @@ -123,6 +124,10 @@ function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{ spawn_testsuite(backend, AT) end + @conditional_testset "Group reductions" skip_tests begin + groupreduce_testsuite(backend, AT) + end + return end From 8fe0876227f0ea6d1f98ae10bdf12d11c273c1ea Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Sat, 10 Oct 2026 18:58:32 +0200 Subject: [PATCH 8/9] Reduce sub-groups first in `@groupreduce` where the backend can On backends with sub-groups, `@groupreduce` reduces every sub-group with `KernelInterface.sub_group_reduce`, which backends can implement natively, and then the results of the sub-groups. That only relies on the sub-group ids, so it needs neither linear nor independent sub-groups. KernelAbstractions picks the algorithm instead of a user-facing option: kernels that use collectives get the backend's sub-group capabilities in their `CompilerMetadata`, and use the sub-groups for types the backend can shuffle. Floating-point types need the backend's support for that type, since a native shuffle may be selected by type alone; other primitive types and structs are shuffled as packed words. Kernels without collectives don't query the capabilities. Co-authored-by: Tim Besard --- docs/src/index.md | 3 ++ docs/src/kernels.md | 37 ++++++++++--- src/backend_launch.jl | 32 +++++++++--- src/compiler.jl | 18 ++++--- src/groupreduction.jl | 118 ++++++++++++++++++++++++++++++++++++++---- src/macros.jl | 25 ++++++++- test/groupreduce.jl | 76 +++++++++++++++++++++++++-- test/runtests.jl | 16 ++++++ 8 files changed, 290 insertions(+), 35 deletions(-) diff --git a/docs/src/index.md b/docs/src/index.md index cdb8669da..d9c40dd8c 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -172,6 +172,9 @@ end written against KernelInterface and for kernel languages built on it; capability queries say how a backend forms its sub-groups, whether they may diverge, and which types it can shuffle. See [Sub-groups](@ref ki_subgroups). +- [`@groupreduce`](@ref) and [`@groupscan`](@ref) reduce and scan a value over the + workgroup, on every backend. `@groupreduce` uses sub-groups internally where the backend + can shuffle the values. See [Reductions and scans](@ref group_collectives). ## Semantic differences diff --git a/docs/src/kernels.md b/docs/src/kernels.md index 75ea30e9d..6ead9aa21 100644 --- a/docs/src/kernels.md +++ b/docs/src/kernels.md @@ -213,13 +213,38 @@ statements, and [`@uniform`](@ref) evaluates an expression outside the work-item can be reused across `@synchronize` statements. For scratch storage that does not need to survive across `@synchronize`, an `MArray` can be used instead. -## Reductions and scans +## [Reductions and scans](@id group_collectives) -[`@groupreduce`](@ref) and [`@groupscan`](@ref) reduce and scan values over the workgroup. -Like [`@synchronize`](@ref), they are collectives: all work-items of the workgroup have to -reach them, not in a branch or loop that only some of them take, and they have to be used as -statements of their own, e.g. `res = @groupreduce(+, val, zero(T))`. Work-items that pad a -partial workgroup take part as well, contributing the neutral element. +[`@groupreduce`](@ref) combines a value from every work-item of the workgroup and returns the +result to all of them. [`@groupscan`](@ref) gives each work-item the result up to itself, in +the order of `@index(Local, Linear)`; it is inclusive by default, and exclusive with +`inclusive = false`. Both take an associative operator and its neutral element, and +`@groupreduce` also needs the operator to be commutative. + +```julia +@kernel function group_sums!(out, @Const(x)) + i = @index(Global, Linear) + s = @groupreduce(+, x[i], zero(eltype(out)); groupsize = 256) + if @index(Local, Linear) == 1 + out[@index(Group, Linear)] = s + end +end +``` + +Like [`@synchronize`](@ref), they are collectives: every work-item of the workgroup has to +reach them, not in a branch or loop that only some of them take, or after an early `return`. +Use them as a statement of their own, optionally assigning the result, as in +`res = @groupreduce(+, val, zero(T))`. Work-items that pad a partial workgroup take part as +well: they contribute the neutral element, without evaluating `val`. + +Both go through local memory, whose size has to be known when the kernel is compiled: a +kernel without a static workgroup size passes `groupsize`, a constant upper bound on the +number of work-items per workgroup. + +They work on every backend. Where the backend can shuffle the values, `@groupreduce` first +reduces each sub-group, and then combines those results through local memory. +KernelAbstractions picks this automatically. The order in which values are combined can +differ, so floating-point results may vary between backends and launch configurations. A kernel is either a `@kernel` or a kernel written against [KernelInterface](@ref kernelinterface), never a mix: don't call KernelInterface's device diff --git a/src/backend_launch.jl b/src/backend_launch.jl index 2b145ef1d..991db0a8c 100644 --- a/src/backend_launch.jl +++ b/src/backend_launch.jl @@ -8,15 +8,27 @@ ### """ - mkcontext(kernel::Kernel, ndrange, iterspace, [launch]) + mkcontext(kernel::Kernel, ndrange, iterspace, [launch]; subgroups = nothing) The hidden context argument for launching `kernel` over `ndrange`, partitioned as -`iterspace`, with the launch configuration `launch` (see [`select_launch`](@ref)). +`iterspace`, with the launch configuration `launch` (see [`select_launch`](@ref)) and the +backend's sub-group capabilities `subgroups` (see [`kernel_subgroups`](@ref)). """ mkcontext(kernel::Kernel, _ndrange, iterspace) = CompilerMetadata{ndrange(kernel), DynamicCheck}(_ndrange, iterspace) -mkcontext(kernel::Kernel, _ndrange, iterspace, launch) = - CompilerMetadata{ndrange(kernel), DynamicCheck}(_ndrange, iterspace; launch) +mkcontext(kernel::Kernel, _ndrange, iterspace, launch; subgroups = nothing) = + CompilerMetadata{ndrange(kernel), DynamicCheck}(_ndrange, iterspace; launch, subgroups) + +""" + kernel_subgroups(kernel::Kernel) + +The sub-group capabilities of the backend of `kernel` that its work-group collectives can +use, a [`SubgroupCapabilities`](@ref), or `nothing` if the backend has none, or if the kernel +doesn't use collectives (which doesn't depend on the backend, so that launching other +kernels doesn't query it). +""" +kernel_subgroups(kernel::Kernel) = + uses_collectives(kernel.f) ? subgroup_capabilities(backend(kernel)) : nothing mkcontext(kernel::Kernel, I, _ndrange, iterspace, ::Dynamic) where {Dynamic} = CompilerMetadata{ndrange(kernel), Dynamic}(I, _ndrange, iterspace) @@ -104,10 +116,18 @@ function launch_tuple(obj::Kernel, args::Tuple; ndrange = nothing, workgroupsize end function launch_kernel(obj::Kernel, launch, ndrange, _workgroupsize, iterspace, args::Tuple) + # a constant `nothing` for kernels without collectives; for kernels with, the call below + # is a function barrier that makes the rest of the launch type stable + subgroups = kernel_subgroups(obj) + launch_kernel(obj, launch, subgroups, ndrange, _workgroupsize, iterspace, args) + return nothing +end + +function launch_kernel(obj::Kernel, launch, subgroups, ndrange, _workgroupsize, iterspace, args::Tuple) b = backend(obj) # this might not be the final context, since we may tune the workgroupsize - ctx = mkcontext(obj, ndrange, iterspace, launch) + ctx = mkcontext(obj, ndrange, iterspace, launch; subgroups) args, tt = prepare_arguments(b, obj.f, ctx, args) kernel = KI.kernel_function(b, obj.f, tt; compiler_options(obj)...) @@ -116,7 +136,7 @@ function launch_kernel(obj::Kernel, launch, ndrange, _workgroupsize, iterspace, range = something(ndrange, static_ndrange(obj)) threads = KI.launch_configuration(kernel; nitems = saturated_prod(extents(range))).workgroupsize iterspace, _ = partition(obj, ndrange, launch_workgroupsize(b, launch, threads, range)) - ctx = mkcontext(obj, ndrange, iterspace, launch) + ctx = mkcontext(obj, ndrange, iterspace, launch; subgroups) end # the geometry is valid by construction, except for the kernel's limits diff --git a/src/compiler.jl b/src/compiler.jl index c61f7d471..cbc22721b 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -1,26 +1,29 @@ """ - CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch} + CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch, Subgroups} 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). +indexed in `Int`, or a [`LinearLaunch`](@ref) or [`NDLaunch`](@ref). The `subgroups` field +holds the backend's [`SubgroupCapabilities`](@ref) for kernels that use work-group +collectives, or `nothing`. """ -struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch} +struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch, Subgroups} groupindex::I ndrange::NDRange iterspace::Iterspace launch::Launch + subgroups::Subgroups # 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), Nothing}(idx, ndrange, iterspace, nothing) + return new{NDRange, CB, typeof(idx), typeof(ndrange), typeof(iterspace), Nothing, Nothing}(idx, ndrange, iterspace, nothing, nothing) end # GPU variante: index is given implicit - function CompilerMetadata{NDRange, CB}(ndrange, iterspace; launch = nothing) where {NDRange, CB} + function CompilerMetadata{NDRange, CB}(ndrange, iterspace; launch = nothing, subgroups = nothing) where {NDRange, CB} ndrange = cartesian(ndrange) - return new{NDRange, CB, Nothing, typeof(ndrange), typeof(iterspace), typeof(launch)}(nothing, ndrange, iterspace, launch) + return new{NDRange, CB, Nothing, typeof(ndrange), typeof(iterspace), typeof(launch), typeof(subgroups)}(nothing, ndrange, iterspace, launch, subgroups) end end @@ -34,6 +37,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 __subgroups(cm::CompilerMetadata) = cm.subgroups @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)) @@ -48,7 +52,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; cm.launch) + return CompilerMetadata{NDRange, CB}(cm.ndrange, iterspace; cm.launch, cm.subgroups) else return CompilerMetadata{NDRange, CB}(cm.groupindex, cm.ndrange, iterspace) end diff --git a/src/groupreduction.jl b/src/groupreduction.jl index 1b92d9958..5fa103eba 100644 --- a/src/groupreduction.jl +++ b/src/groupreduction.jl @@ -26,9 +26,14 @@ after an early `return`. Work-items that pad a partial workgroup (outside of the take part as well: they contribute `neutral`, without evaluating `val`. So `op` and `neutral` have to be the same for the whole workgroup, and computable on those work-items. -The reduction stores values in local memory, so its size has to be known at compile time: -either the kernel has a static workgroup size, or the keyword `groupsize` gives an upper -bound of the workgroup size (a constant, e.g. a literal or a type parameter of the kernel). +The reduction goes through local memory, whose size has to be known at compile time: either +the kernel has a static workgroup size, or the keyword `groupsize` gives an upper bound on the +number of work-items per workgroup (a constant, e.g. a literal or a type parameter of the +kernel). `groupsize` only sizes that storage; it doesn't set the launch's workgroup size. + +On backends with sub-groups, KernelAbstractions reduces the values of every sub-group with +[`KernelInterface.sub_group_reduce`](@ref) first, if the backend can shuffle values of the +type of `neutral`. ```julia @kernel function sum_kernel!(out, @Const(x)) @@ -52,12 +57,12 @@ end @groupscan(op, val, neutral; groupsize, inclusive = true) Scan `val` over the work-items of the workgroup with the binary operator `op`, in the order -of `@index(Local, Linear)`: the work-item with local index `i` gets -`op(...op(op(val₁, val₂), val₃)..., valᵢ)` (inclusive), or the same up to `valᵢ₋₁` and -`neutral` for the first work-item (with `inclusive = false`). `op` has to be associative, -but needn't be commutative, and `neutral` its neutral element -(`op(neutral, x) == op(x, neutral) == x`). The result has the type of `neutral`, which `val` -is converted to, and `op` has to return that type. +of `@index(Local, Linear)`. With `inclusive = true` (the default), the work-item with local +index `i` gets the reduction of the values of work-items `1:i`; with `inclusive = false`, that +of `1:i-1`, and the first work-item gets `neutral`. `op` has to be associative, but needn't be +commutative, and the values may be combined in any grouping. `neutral` has to be its neutral +element (`op(neutral, x) == op(x, neutral) == x`). The result has the type of `neutral`, +which `val` is converted to, and `op` has to return that type. Like [`@groupreduce`](@ref), `@groupscan` is a collective that has to be used as a statement on its own and reached by all work-items of the workgroup; padding work-items contribute @@ -65,7 +70,9 @@ on its own and reached by all work-items of the workgroup; padding work-items co `@groupreduce`: either the kernel has a static workgroup size, or the keyword `groupsize` gives an upper bound of the workgroup size. `inclusive` has to be a constant as well. -For example, to compact the elements of `x` that satisfy `pred` within each workgroup: +This example moves the elements of `x` that satisfy `pred` to the start of each workgroup's +part of `out`, and records how many there are in `counts`; the result isn't compacted across +workgroups. The kernel needs a static workgroup size, e.g. `compact!(backend, 256)`: ```julia @kernel function compact!(out, counts, @Const(x), pred) @@ -198,7 +205,11 @@ end n = prod(groupsize(ctx)) n <= N || throw(ArgumentError("@groupreduce: the workgroup size exceeds the given upper bound")) storage = KI.localmemory(T, Val(N)) - res = __groupreduce_tree(op, convert(T, val), storage, __index_Local_Linear(ctx), n) + if __shuffleable(__subgroups(ctx), T) + res = __groupreduce_subgroups(op, convert(T, val), neutral, storage) + else + res = __groupreduce_tree(op, convert(T, val), storage, __index_Local_Linear(ctx), n) + end # all work-items have to read the result before the storage can be reused KI.barrier() return res @@ -219,6 +230,39 @@ end return @inbounds storage[1] end +# Reduce every sub-group with `KI.sub_group_reduce`, which backends can implement natively, +# then reduce the results of the sub-groups. This only uses the sub-group ids, not how the +# work-items form sub-groups, and all sub-groups execute the same sub-group operations, so it +# needs neither linear nor independent sub-groups. There are at most as many sub-groups as +# work-items, so `storage` has room for a value per sub-group. +@inline function __groupreduce_subgroups(op, val, neutral, storage) + sg = KI.get_sub_group_id() + lane = KI.get_sub_group_local_id() + val = KI.sub_group_reduce(op, val) + if lane == 1 + @inbounds storage[sg] = val + end + KI.barrier() + + # every sub-group reduces the results of all sub-groups, which avoids running sub-group + # operations in a branch that only some sub-groups take + width = KI.get_sub_group_size() + acc = neutral + i = lane + while i <= KI.get_num_sub_groups() + acc = op(acc, @inbounds storage[i]) + i += width + end + acc = KI.sub_group_reduce(op, acc) + KI.barrier() + # the sub-groups can combine the values in different orders, so broadcast one result + if sg == 1 && lane == 1 + @inbounds storage[1] = acc + end + KI.barrier() + return @inbounds storage[1] +end + @inline function __groupscan(ctx, op, val, neutral::T, ::Val{N}, ::Val{inclusive}) where {T, N, inclusive} n = prod(groupsize(ctx)) n <= N || throw(ArgumentError("@groupscan: the workgroup size exceeds the given upper bound")) @@ -253,3 +297,55 @@ end return res end + +## sub-group capabilities + +""" + SubgroupCapabilities{F16, F32, F64}() + +The sub-group capabilities of a backend that the work-group collectives use, passed to the +kernel in its [`CompilerMetadata`](@ref): the backend supports sub-groups and shuffles of +`UInt32` values (and thus of other integer-like values, as words), and of `Float16`, +`Float32` and `Float64` values if `F16`, `F32` and `F64`. Kernels on backends without these +capabilities get `nothing` instead. +""" +struct SubgroupCapabilities{F16, F32, F64} end + +function subgroup_capabilities(backend::KI.Backend) + (KI.supports_subgroups(backend) && KI.supports_shuffle(backend, UInt32)) || return nothing + return SubgroupCapabilities{ + KI.supports_shuffle(backend, Float16), + KI.supports_shuffle(backend, Float32), + KI.supports_shuffle(backend, Float64), + }() +end + +# Whether values of type `T` can be shuffled. Floating-point types need the backend's support: +# a backend that shuffles a float type natively may do so without regard for the device, e.g. +# `Float64` on a device without it. Other primitive types are shuffled as words, and structs +# packed into `UInt32` words, whatever the types of their fields. +__shuffleable(::Nothing, ::Type) = false +@generated function __shuffleable(::SubgroupCapabilities{F16, F32, F64}, ::Type{T}) where {F16, F32, F64, T} + function shuffleable(S) + if isprimitivetype(S) + S === Float16 && return F16 + S === Float32 && return F32 + S === Float64 && return F64 + return !(S <: AbstractFloat) && sizeof(S) in (1, 2, 4, 8, 16) + end + isbitstype(S) || return false + fields = KI.primitive_fields!(Any[], S, :val) + return all(((F, _),) -> sizeof(F) in (1, 2, 4, 8, 16), fields) + end + return shuffleable(T) +end + +# Whether a kernel's body contains a work-group collective, which needs the backend's +# sub-group capabilities: the `@kernel` macro adds a method taking a `CollectivesQuery` to the +# kernel's function, so that other kernels don't query them at every launch. (`applicable` is +# constant-folded.) It only does so for the first method of a kernel function: other methods +# with collectives use local memory. A callable that forwards its arguments, e.g. a wrapper of +# a kernel function, is taken to use collectives, which only costs the query of the +# capabilities at its launches. +struct CollectivesQuery end +uses_collectives(f) = applicable(f, CollectivesQuery()) diff --git a/src/macros.jl b/src/macros.jl index 48774b0b6..a7e5dfb3b 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -74,6 +74,16 @@ function __kernel(expr, __source__::LineNumberNode, __module__::Module, force_in end gpu_function = combinedef(def_gpu) + # Properties of the kernel are methods of its function, rather than of a function of + # KernelAbstractions, so that a kernel can also be defined in a local scope. Like the + # constructors, they are defined with the first method of the kernel, since a method can't + # be overwritten during precompilation. + queries = Any[] + if find_collective(def[:body]) + # see `uses_collectives` + push!(queries, :($gpu_name(::$CollectivesQuery) = true)) + end + # create constructor functions _name = Symbol(:_, name) constructors = quote @@ -85,6 +95,7 @@ function __kernel(expr, __source__::LineNumberNode, __module__::Module, force_in $name(dev, size) = $_name(dev, $StaticSize(size), $DynamicSize()) $name(dev, size, range) = $_name(dev, $StaticSize(size), $StaticSize(range)) $name(dev, size::$_Size, range::$_Size) = $_name(dev, size, range) + $(queries...) end end constructors = relocate_lines(constructors, __source__) @@ -207,7 +218,7 @@ function default_forwarders!(def, gpu_name, constargs) default = args[n + 1].args[2] fdef = Dict{Symbol, Any}( :name => gpu_name, - :args => Any[:__ctx__, kept...], + :args => Any[:(__ctx__::Union{$CompilerMetadata, $ConstArgs}), kept...], :kwargs => Any[], :whereparams => filter(in(used), whereparams), :body => quote @@ -226,7 +237,8 @@ end # The easy case, transform the function for GPU execution # - answer which arguments are marked `@Const` (see `ConstArgs`) function transform_gpu!(def, constargs, force_inbounds, unsafe_indices) - pushfirst!(def[:args], :__ctx__) + # typed, so that the queries of `@kernel` (e.g. `uses_collectives`) never match the kernel + def[:args] = Any[:(__ctx__::Union{$CompilerMetadata, $ConstArgs}), def[:args]...] # `Any[]`, since `split` hands back `LineNumberNode`s alongside `Expr`s new_stmts = Any[] push!(new_stmts, :(__ctx__ isa $ConstArgs && return $(constargs_mask(constargs)))) @@ -294,6 +306,15 @@ function is_scope_construct(expr::Expr) # expr.head === :let end +function find_collective(stmt) + result = Ref(false) + postwalk(stmt) do expr + result[] |= is_collective(expr) + expr + end + return result[] +end + # Whether `stmt` contains a `@synchronize`, or a collective like `@groupreduce` that all # work-items of the workgroup have to reach as well. function find_sync(stmt) diff --git a/test/groupreduce.jl b/test/groupreduce.jl index 3989578b0..e18e2ab3f 100644 --- a/test/groupreduce.jl +++ b/test/groupreduce.jl @@ -40,6 +40,9 @@ end end end +@kernel function groupreduce_noargs!() +end + # `val` of another type than `neutral`: the padding work-items contribute `neutral`, so the # value is a `Union` of both types, and the call of the collective must not be union-split @kernel function groupreduce_mixed!(out, @Const(x)) @@ -110,6 +113,28 @@ function groupwise_scan(op, x, groupsize, neutral, inclusive) return out end +# A 32-bit float type that no backend shuffles, so that `@groupreduce` uses local memory +# only, also on backends with sub-groups. Its `+` adds the bits as integers. +primitive type TreeFloat <: AbstractFloat 32 end +TreeFloat(x::Integer) = reinterpret(TreeFloat, Int32(x)) +Base.:+(a::TreeFloat, b::TreeFloat) = reinterpret(TreeFloat, reinterpret(Int32, a) + reinterpret(Int32, b)) +Base.zero(::Type{TreeFloat}) = TreeFloat(0) +Base.typemin(::Type{TreeFloat}) = TreeFloat(typemin(Int32)) +Base.:(==)(a::TreeFloat, b::TreeFloat) = reinterpret(Int32, a) == reinterpret(Int32, b) + +# Run `f`, and return whether it failed to compile with GPUCompiler.jl#1004 (Metal: a phi of +# a by-reference argument and a device pointer, as in `x[i]` or the `neutral` argument), +# which the local-memory path runs into. +function gpucompiler_1004(f) + try + f() + return false + catch err + occursin("Invalid phi record", sprint(showerror, err)) || rethrow() + return true + end +end + # `neutral` is evaluated once per work-item, also by the padding work-items counted_neutral() = 0 count_symbol(ex::Expr, sym) = sum(arg -> count_symbol(arg, sym), ex.args; init = 0) @@ -117,16 +142,43 @@ count_symbol(ex, sym) = Int(ex === sym) function groupreduce_testsuite(backend, AT) b = backend() - types = (Int32, Int64, Float32) + caps = KernelAbstractions.subgroup_capabilities(b) + + @testset "sub-group capabilities" begin + sub_groups = KI.supports_subgroups(b) && KI.supports_shuffle(b, UInt32) + @test (caps !== nothing) == sub_groups + shuffleable(T) = KernelAbstractions.__shuffleable(caps, T) + @test shuffleable(Int32) == sub_groups + @test shuffleable(Tuple{Int64, Bool}) == sub_groups + @test shuffleable(Float32) == (sub_groups && KI.supports_shuffle(b, Float32)) + @test shuffleable(Tuple{Float32, Int32}) == (sub_groups && KI.supports_shuffle(b, Float32)) + @test !shuffleable(TreeFloat) + @test !shuffleable(Ref{Int}) + # kernels without collectives don't query the capabilities + @test KernelAbstractions.kernel_subgroups(groupreduce_static!(b, 64)) == caps + @test KernelAbstractions.kernel_subgroups(groupscan!(b, 64)) == caps + @test KernelAbstractions.uses_collectives(groupreduce_static!(b, 64).f) + # also for a kernel without arguments, whose function takes a single argument + @test !KernelAbstractions.uses_collectives(groupreduce_noargs!(b).f) + end + + types = (Int32, Int64, Float32, TreeFloat) @testset "$T, $(nameof(typeof(op)))" for T in types, (op, neutral) in ((+, zero(T)), (max, typemin(T))) + T === TreeFloat && op === max && continue for (groupsize, n) in ((64, 64), (64, 256), (32, 100), (256, 1000), (7, 23), (1, 3)) x = T.(rand(1:100, n)) out = AT(fill(zero(T), n)) - groupreduce_static!(b, groupsize)(out, AT(x), op, neutral; ndrange = n) + if gpucompiler_1004(() -> groupreduce_static!(b, groupsize)(out, AT(x), op, neutral; ndrange = n)) + @test_broken false + continue + end @test Array(out) == groupwise(op, x, groupsize) fill!(out, zero(T)) - groupreduce_bound!(b)(out, AT(x), op, neutral, Val(256); ndrange = n, workgroupsize = groupsize) + if gpucompiler_1004(() -> groupreduce_bound!(b)(out, AT(x), op, neutral, Val(256); ndrange = n, workgroupsize = groupsize)) + @test_broken false + continue + end @test Array(out) == groupwise(op, x, groupsize) end end @@ -242,6 +294,24 @@ function groupreduce_testsuite(backend, AT) end end + @testset "kernels in a local scope" begin + # the kernel's properties are methods of its function, which works in a local scope too + function local_groupreduce(b, AT, x) + @kernel function local_groupreduce!(out, @Const(x)) + i = @index(Global, Linear) + res = @groupreduce(+, x[i], zero(eltype(out))) + out[i] = res + end + out = AT(zeros(eltype(x), length(x))) + kernel = local_groupreduce!(b, 32) + @test KernelAbstractions.uses_collectives(kernel.f) + kernel(out, AT(x); ndrange = length(x)) + return Array(out) + end + x = Int32.(rand(1:100, 100)) + @test local_groupreduce(b, AT, x) == groupwise(+, x, 32) + end + @testset "neutral evaluated once" begin ex = @macroexpand @kernel function f(y, x) i = @index(Global, Linear) diff --git a/test/runtests.jl b/test/runtests.jl index 81c8e8fa0..088b010a2 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -370,6 +370,22 @@ end end end +# `@groupreduce` reduces sub-groups with PoCL's native collectives +@kernel function groupreduce_codegen!(out, @Const(x)) + i = @index(Global, Linear) + res = @groupreduce(+, x[i], zero(eltype(out))) + out[i] = res +end +@testset "POCL @groupreduce with sub-groups" begin + x = ones(Int32, 100) + out = zeros(Int32, 100) + ir = sprint() do io + @device_code_llvm io = io debuginfo = :none groupreduce_codegen!(POCLBackend(), 64)(out, x; ndrange = 100) + end + @test out == [fill(64, 64); fill(36, 36)] + @test occursin("sub_group_reduce_add", ir) +end + # Julia doesn't turn a splat of more than 32 elements into a direct call, so a launch with # many arguments allocates unless every layer passes them on as a tuple @testset "POCL launch with many arguments" begin From 23d70e6f9ca298ae3e8ba68b09f43639b6b40985 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Mon, 12 Oct 2026 10:08:59 +0200 Subject: [PATCH 9/9] Add tiles A kernel constructed with a tile width, `f(backend; tile = N)`, divides its workgroups into tiles of `N` consecutive work-items (by `@index(Local, Linear)`), like CUDA's `tiled_partition`. The work-items of a tile communicate with operations that take the tile (`tile_shfl` and its variants, `tile_any`, `tile_all`, `tile_ballot`, `tile_reduce`, `tile_scan`, `tile_exclusive_scan` and `tile_barrier`), in control flow that only has to be uniform over the tile. `@tile()` returns the work-item's tile and lane, and `@index(Tile)` the global index of its tile. Tiles aren't declared in the kernel: what sets a kernel with tiles apart is its width, which the host needs for the launch geometry, and which is passed when the kernel is constructed. A kernel that uses `@tile()` or `@index(Tile)` has to get one; the `@kernel` macro notes that use, so that construction rejects a missing width. There is no default: the sub-group width differs between backends (32 or 64 on AMD GPUs), so kernels name the width they're written for. The width is a type parameter of `Kernel` (`tile_width(kernel)`) and reaches the device through the kernel's `CompilerMetadata`, where `tile_width(t)` is a constant. Tiles are a capability: a power-of-two `N` of at most the sub-group width `W`, any number of tiles per workgroup, and tiles that diverge independently, also where several share a sub-group (KernelInterface's segments). `supports_tiles(backend, N)` reports it, and constructing a kernel with tiles a backend doesn't support is an error. All GPU backends support widths up to 32. The CPU backend has no tiles, as its sub-group functions are work-group barriers, which can't diverge. Falling back to a tile per workgroup there would leak into how kernels choose their workgroup size, and be slower than a kernel using local memory, which is what such backends need instead. The ndrange can have several dimensions, e.g. one per index of the nodes of an element, if the workgroups divide it into whole tiles. Such kernels are launched on a 1-D grid, so that the tiles stay on the backend's sub-groups. Without an explicit workgroup size, the launch picks one that does; construction and launch check the geometry before compiling. --- docs/src/api.md | 20 ++ docs/src/index.md | 8 + docs/src/kernelinterface.md | 12 +- docs/src/kernels.md | 91 +++++++- src/KernelAbstractions.jl | 45 +++- src/backend_launch.jl | 14 +- src/compiler.jl | 17 +- src/launch.jl | 6 +- src/macros.jl | 27 ++- src/tiles.jl | 419 ++++++++++++++++++++++++++++++++++++ test/runtests.jl | 43 ++++ test/testsuite.jl | 5 + test/tiles.jl | 411 +++++++++++++++++++++++++++++++++++ 13 files changed, 1077 insertions(+), 41 deletions(-) create mode 100644 src/tiles.jl create mode 100644 test/tiles.jl diff --git a/docs/src/api.md b/docs/src/api.md index f5f3e3845..3ff3630eb 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -22,6 +22,26 @@ @groupscan ``` +### Tiles + +```@docs +@tile +Tile +tile_shfl +tile_shfl_up +tile_shfl_down +tile_shfl_xor +tile_any +tile_all +tile_ballot +tile_reduce +tile_scan +tile_exclusive_scan +tile_barrier +tile_width +supports_tiles +``` + ## Host language !!! note diff --git a/docs/src/index.md b/docs/src/index.md index d9c40dd8c..586be13f8 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -175,6 +175,14 @@ end - [`@groupreduce`](@ref) and [`@groupscan`](@ref) reduce and scan a value over the workgroup, on every backend. `@groupreduce` uses sub-groups internally where the backend can shuffle the values. See [Reductions and scans](@ref group_collectives). +- Tiles let small groups of consecutive work-items exchange values with shuffles, votes, + reductions and scans (`tile_shfl`, `tile_reduce`, ...), while different tiles take + different branches or loop a different number of times. A kernel gets tiles by being + constructed with a width, `f(backend; tile = N)`, without a declaration in the kernel. + [`supports_tiles(backend, N)`](@ref supports_tiles) reports whether a backend supports a + width; the `CPU` backend has no tiles. See [Tiles](@ref tiles). A `@kernel` uses tiles or + the workgroup collectives; KernelInterface's sub-group functions are for kernels written + directly against KernelInterface. ## Semantic differences diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md index f42467194..a5c695542 100644 --- a/docs/src/kernelinterface.md +++ b/docs/src/kernelinterface.md @@ -209,13 +209,11 @@ execute the same communication functions in the same order. The id and size quer used anywhere. !!! note "Kernel languages" - The sub-group functions are meant for kernels written against KernelInterface, and for - kernel languages built on it. A kernel is either a KernelInterface kernel or a - KernelAbstractions `@kernel`, never a mix: `@kernel` adds padding work-items to partial - work-groups, which skip the kernel's body, so KernelInterface's sub-group functions called - in a `@kernel` aren't reached by all work-items. A `@kernel` uses KernelAbstractions' - collectives instead, e.g. [`@groupreduce`](@ref KernelAbstractions.@groupreduce). Host - code may use the queries of both. + The sub-group functions are meant for kernels written against KernelInterface and for + implementations of kernel languages. Don't call them in a KernelAbstractions `@kernel`: + use its [reductions and scans](@ref group_collectives) or [tiles](@ref tiles) instead, + which follow the participation rules of `@kernel`. Host code may use the queries of + both. ```@docs get_sub_group_size diff --git a/docs/src/kernels.md b/docs/src/kernels.md index 6ead9aa21..6b9c1fddc 100644 --- a/docs/src/kernels.md +++ b/docs/src/kernels.md @@ -246,12 +246,93 @@ reduces each sub-group, and then combines those results through local memory. KernelAbstractions picks this automatically. The order in which values are combined can differ, so floating-point results may vary between backends and launch configurations. -A kernel is either a `@kernel` or a kernel written against +## [Tiles](@id tiles) + +A tile is a group of `N` consecutive work-items, by `@index(Local, Linear)`. Constructing a +kernel with a tile width, `f(backend; tile = N)`, divides its workgroups into tiles, and +[`@tile()`](@ref @tile) returns the work-item's tile in the kernel. The work-items of a tile +exchange values with operations such as [`tile_shfl`](@ref) and [`tile_reduce`](@ref), without +going through local memory. All work-items of a tile have to execute each tile operation +together, but different tiles may take different branches, loop a different number of times, +or `return` early. + +```julia +@kernel function tile_maxima!(out, @Const(x)) + t = @tile() + m = tile_reduce(t, max, x[@index(Global, Linear)]) + if t.lane == 1 + out[@index(Tile)] = m + end +end + +kernel = tile_maxima!(backend; tile = 16) +kernel(out, x; ndrange = 16 * length(out)) +``` + +The workgroup size has to be a multiple of `N`, and every tile has to lie entirely inside or +entirely outside of the `ndrange`; see [`@tile`](@ref) for the details. The `ndrange` can have +several dimensions, as long as that holds. Here, each tile holds the 4 × 4 nodes of an +element, and each workgroup holds four elements: + +```julia +@kernel function element_sums!(out, @Const(u)) + t = @tile() + I = @index(Global, NTuple) + out[I...] = tile_reduce(t, +, u[I...]) +end +element_sums!(backend, (4, 4, 4); tile = 16)(out, u; ndrange = size(u)) # size(u) == (4, 4, nelements) +``` + +### Choosing the tile width + +Tiles aren't declared in the kernel: pass the width when constructing it, `f(backend; +tile = 16)` or `f(backend, workgroupsize; tile = 16)`. A kernel that uses `@tile()` or +`@index(Tile)` needs one. [`tile_width(kernel)`](@ref tile_width) returns it on the host, and +`tile_width(t)` in the kernel, where it is a constant. Like a static workgroup size, the width +is compiled into the kernel; construct the kernel again for another width. + +Tiles are built on the backend's sub-groups (warps, wavefronts or SIMD-groups). The width +has to be a power of two of at most the sub-group width `W`, and tiles narrower than `W` share +a sub-group, yet still take their own paths. Not every backend can do that: +[`supports_tiles(backend, N)`](@ref supports_tiles) says whether it supports tiles of `N` +work-items, and constructing a kernel with a width it doesn't support is an error. CUDA, +AMDGPU, Metal and oneAPI, and OpenCL.jl on Intel GPUs, support widths of up to 32. For tiles of +the full sub-group width, pass `tile = KernelInterface.sub_group_size(backend)`. + +KernelAbstractions' CPU backend has no tiles: its sub-group operations involve the whole +workgroup, so tiles can't take their own paths. To support it, write a variant of the kernel +that uses local memory or the workgroup collectives instead: + +```julia +kernel = supports_tiles(backend, 16) ? element_sums!(backend; tile = 16) : + element_sums_localmem!(backend) +``` + +## Choosing how work-items communicate + +| To | Use | Available | +|---|---|---| +| share data in any pattern within a workgroup | [`@localmem`](@ref) and [`@synchronize`](@ref) | on every backend | +| reduce or scan a value over the workgroup | [`@groupreduce`](@ref), [`@groupscan`](@ref) | on every backend | +| exchange values within small groups of work-items whose control flow may differ between groups, e.g. the nodes of an element | [tiles](@ref tiles) | where [`supports_tiles(backend, N)`](@ref supports_tiles) is `true` | +| program the hardware's sub-groups directly | a [KernelInterface](@ref kernelinterface) kernel | as its [capability queries](@ref ki_subgroups) report | + +Use the collectives for reductions and scans over the workgroup, and local memory with +`@synchronize` for other patterns; both work everywhere. Tiles fit when values only need to +move within small groups that take their own paths. They avoid local memory, but that doesn't +guarantee a faster kernel: on some GPUs, the local-memory version is as fast or faster. +Compare both on the devices you care about, and keep the local-memory version for the backends +without tiles. + +### Kernel languages + +KernelInterface is the low-level API that backends implement, and `@kernel` is a language +built on it, which takes care of the work-items that pad a partial workgroup and of who takes +part in collectives. A kernel is either a `@kernel` or a kernel written against [KernelInterface](@ref kernelinterface), never a mix: don't call KernelInterface's device -functions, such as its sub-group shuffles and votes, in a `@kernel`. They have to be executed -by all work-items of a sub-group, but in a `@kernel` with the default bounds checking, every -statement other than the collectives above only runs on the work-items that are part of the -`ndrange`. Host code can use KernelInterface's queries, e.g. to pick a kernel. +functions, such as its sub-group shuffles and votes, in a `@kernel`. In a `@kernel`, use the +collectives or tiles instead. Host code can use KernelInterface's queries, e.g. to pick a +kernel or a tile width. ## Launching kernels diff --git a/src/KernelAbstractions.jl b/src/KernelAbstractions.jl index b7cb5c36d..e89889804 100644 --- a/src/KernelAbstractions.jl +++ b/src/KernelAbstractions.jl @@ -407,6 +407,8 @@ A cartesian index is a general N-dimensional index that is derived from the iter - `Global`: Used to access global memory. - `Group`: The index of the `workgroup`. - `Local`: The within `workgroup` index. + - `Tile`: The global linear index of the work-item's tile, in kernels with tiles (see + [`@tile`](@ref)). It takes no index kind. # Index kind @@ -428,8 +430,12 @@ If the index kind is not provided it defaults to `Linear`, this is subject to ch ``` """ macro index(locale, args...) + if locale === :Tile + isempty(args) || error("@index(Tile) takes no further arguments") + return :($__index_Tile($(esc(:__ctx__)), $__tile_width($(esc(:__ctx__))))) + end if !(locale === :Global || locale === :Local || locale === :Group) - error("@index requires as first argument either :Global, :Local or :Group") + error("@index requires as first argument either :Global, :Local, :Group or :Tile") end if length(args) >= 1 @@ -525,9 +531,10 @@ import .NDIteration: get ### """ - Kernel{Backend, WorkgroupSize, NDRange, Func} + Kernel{Backend, WorkgroupSize, NDRange, Func, TileWidth} -Host-side handle for a kernel specialized on a backend, workgroup size, and `ndrange`. +Host-side handle for a kernel specialized on a backend, workgroup size, `ndrange`, and for a +kernel with tiles, the tile width (an `Int`, or `nothing`, see [`tile_width`](@ref)). Kernels are created by calling a [`@kernel`](@ref) function on a backend, for example `my_kernel(CUDABackend(), 256)`. The returned object is callable: @@ -544,13 +551,17 @@ and [`backend`](@ref KernelAbstractions.backend) to inspect a kernel's static co Kernels are launched on any backend that implements [KernelInterface](@ref kernelinterface); see the [notes for backend implementations](@ref implementations_notes). """ -struct Kernel{Backend, WorkgroupSize <: _Size, NDRange <: _Size, Fun} +struct Kernel{Backend, WorkgroupSize <: _Size, NDRange <: _Size, Fun, TileWidth} backend::Backend f::Fun end -function Base.similar(kernel::Kernel{D, WS, ND}, f::F) where {D, WS, ND, F} - return Kernel{D, WS, ND, F}(kernel.backend, f) +# a kernel without tiles +Kernel{B, WS, ND, F}(backend, f) where {B, WS, ND, F} = Kernel{B, WS, ND, F, nothing}(backend, f) + +function Base.similar(kernel::Kernel{D, WS, ND, <:Any, T}, f::F) where {D, WS, ND, F, T} + T === nothing && uses_tiles(f) && tile_error(nothing, tile_width_message) + return Kernel{D, WS, ND, F, T}(kernel.backend, f) end """ @@ -654,8 +665,25 @@ last (possibly partial) workgroup. Primarily used by backend implementations and return iterspace, dynamic end -function construct(backend::B, ::S, ::NDRange, xpu_name::XPUName) where {B <: Backend, S <: _Size, NDRange <: _Size, XPUName} - return Kernel{B, S, NDRange, XPUName}(backend, xpu_name) +# The constructors of a `@kernel` function `f`: `f(backend)`, `f(backend, workgroupsize)` and +# `f(backend, workgroupsize, ndrange)` +construct_kernel(backend, f; tile = nothing) = construct(backend, DynamicSize(), DynamicSize(), f; tile) +construct_kernel(backend, f, size; tile = nothing) = construct(backend, StaticSize(size), DynamicSize(), f; tile) +construct_kernel(backend, f, size, range; tile = nothing) = + construct(backend, StaticSize(size), StaticSize(range), f; tile) +construct_kernel(backend, f, size::_Size, range::_Size; tile = nothing) = construct(backend, size, range, f; tile) + +function construct(backend::B, ::S, ::NDRange, xpu_name::XPUName; tile = nothing) where {B <: Backend, S <: _Size, NDRange <: _Size, XPUName} + N = resolve_tile_width(xpu_name, tile) + if N !== nothing + # what is known before the launch, which checks the rest + check_tile_width(backend, N) + NDRange <: StaticSize && check_tile_items(N, extents(get(NDRange))) + S <: StaticSize && check_tile_workgroup(N, extents(get(S))) + NDRange <: StaticSize && S <: StaticSize && + check_tiles(backend, N, extents(get(NDRange)), extents(get(S))) + end + return Kernel{B, S, NDRange, XPUName, N}(backend, xpu_name) end ### @@ -683,6 +711,7 @@ end include("macros.jl") include("groupreduction.jl") +include("tiles.jl") include("spawn.jl") include("foreach_index.jl") diff --git a/src/backend_launch.jl b/src/backend_launch.jl index 991db0a8c..6764f4b9d 100644 --- a/src/backend_launch.jl +++ b/src/backend_launch.jl @@ -15,9 +15,9 @@ The hidden context argument for launching `kernel` over `ndrange`, partitioned a backend's sub-group capabilities `subgroups` (see [`kernel_subgroups`](@ref)). """ mkcontext(kernel::Kernel, _ndrange, iterspace) = - CompilerMetadata{ndrange(kernel), DynamicCheck}(_ndrange, iterspace) + CompilerMetadata{ndrange(kernel), DynamicCheck}(_ndrange, iterspace; tile = tile_val(kernel)) mkcontext(kernel::Kernel, _ndrange, iterspace, launch; subgroups = nothing) = - CompilerMetadata{ndrange(kernel), DynamicCheck}(_ndrange, iterspace; launch, subgroups) + CompilerMetadata{ndrange(kernel), DynamicCheck}(_ndrange, iterspace; launch, subgroups, tile = tile_val(kernel)) """ kernel_subgroups(kernel::Kernel) @@ -96,9 +96,12 @@ Core.kwcall(kwargs::NamedTuple, obj::Kernel{<:KI.Backend}, args::Vararg{Any, N}) launch_tuple(obj, args; kwargs...) function launch_tuple(obj::Kernel, args::Tuple; ndrange = nothing, workgroupsize = nothing) + N = tile_width(obj) ndrange, workgroupsize, iterspace, dynamic = launch_config(obj, ndrange, workgroupsize) # nothing to launch (or compile) for an empty ndrange any(iszero, size(blocks(iterspace))) && return nothing + # reject an invalid launch of a kernel with tiles before compiling it + N === nothing || check_tiles_launch(obj, N, ndrange, workgroupsize, iterspace) # launch on an N-d grid, computing indices in a narrow type, if possible. this doesn't # depend on the tuned workgroup size, so the context (and thus the kernel) doesn't either. @@ -135,13 +138,18 @@ function launch_kernel(obj::Kernel, launch, subgroups, ndrange, _workgroupsize, if workgroupsize(obj) <: DynamicSize && _workgroupsize === nothing range = something(ndrange, static_ndrange(obj)) threads = KI.launch_configuration(kernel; nitems = saturated_prod(extents(range))).workgroupsize - iterspace, _ = partition(obj, ndrange, launch_workgroupsize(b, launch, threads, range)) + N = tile_width(obj) + items = N === nothing ? launch_workgroupsize(b, launch, threads, range) : + tile_launch_workgroupsize(b, launch, N, threads, range) + iterspace, _ = partition(obj, ndrange, items) ctx = mkcontext(obj, ndrange, iterspace, launch; subgroups) end # the geometry is valid by construction, except for the kernel's limits groups = size(blocks(iterspace)) items = size(workitems(iterspace)) + N = tile_width(obj) + N === nothing || check_tiles(b, N, extents(something(ndrange, static_ndrange(obj))), items) if launch isa NDLaunch groups, items = KI.pad3(groups), KI.pad3(items) else diff --git a/src/compiler.jl b/src/compiler.jl index cbc22721b..7d168b951 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -1,29 +1,31 @@ """ - CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch, Subgroups} + CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch, Subgroups, TileWidth} 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). The `subgroups` field holds the backend's [`SubgroupCapabilities`](@ref) for kernels that use work-group -collectives, or `nothing`. +collectives, or `nothing`. The `tile` field holds the tile width `Val(N)` of a kernel with +tiles (see [`@tile`](@ref)), or `nothing`. """ -struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch, Subgroups} +struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch, Subgroups, TileWidth} groupindex::I ndrange::NDRange iterspace::Iterspace launch::Launch subgroups::Subgroups + tile::TileWidth # 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), Nothing, Nothing}(idx, ndrange, iterspace, nothing, nothing) + return new{NDRange, CB, typeof(idx), typeof(ndrange), typeof(iterspace), Nothing, Nothing, Nothing}(idx, ndrange, iterspace, nothing, nothing, nothing) end # GPU variante: index is given implicit - function CompilerMetadata{NDRange, CB}(ndrange, iterspace; launch = nothing, subgroups = nothing) where {NDRange, CB} + function CompilerMetadata{NDRange, CB}(ndrange, iterspace; launch = nothing, subgroups = nothing, tile = nothing) where {NDRange, CB} ndrange = cartesian(ndrange) - return new{NDRange, CB, Nothing, typeof(ndrange), typeof(iterspace), typeof(launch), typeof(subgroups)}(nothing, ndrange, iterspace, launch, subgroups) + return new{NDRange, CB, Nothing, typeof(ndrange), typeof(iterspace), typeof(launch), typeof(subgroups), typeof(tile)}(nothing, ndrange, iterspace, launch, subgroups, tile) end end @@ -38,6 +40,7 @@ cartesian(t::Tuple) = CartesianIndices(t) @inline __groupindex(cm::CompilerMetadata) = cm.groupindex @inline __launch(cm::CompilerMetadata) = cm.launch @inline __subgroups(cm::CompilerMetadata) = cm.subgroups +@inline __tile_width(cm::CompilerMetadata) = cm.tile @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)) @@ -52,7 +55,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; cm.launch, cm.subgroups) + return CompilerMetadata{NDRange, CB}(cm.ndrange, iterspace; cm.launch, cm.subgroups, cm.tile) else return CompilerMetadata{NDRange, CB}(cm.groupindex, cm.ndrange, iterspace) end diff --git a/src/launch.jl b/src/launch.jl index 7742e749e..eaf88e9f1 100644 --- a/src/launch.jl +++ b/src/launch.jl @@ -77,10 +77,12 @@ function select_launch(kernel::Kernel, workgroupsize, iterspace) groups = size(blocks(iterspace)) items = size(workitems(iterspace)) tuned = KernelAbstractions.workgroupsize(kernel) <: DynamicSize && workgroupsize === nothing + # kernels with tiles are launched on a 1-D grid unless the iteration space is 1-D, so that + # the linear local index is the one the backend forms sub-groups from + dims = tile_width(kernel) !== nothing && length(items) > 1 ? () : KI.max_work_group_dims(b) 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), - narrow_index_type(b) + KI.max_work_group_size(b), dims, KI.max_num_groups(b), narrow_index_type(b) ) end diff --git a/src/macros.jl b/src/macros.jl index a7e5dfb3b..6f2ea28d3 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -79,22 +79,20 @@ function __kernel(expr, __source__::LineNumberNode, __module__::Module, force_in # constructors, they are defined with the first method of the kernel, since a method can't # be overwritten during precompilation. queries = Any[] + if find_tile(def[:body]) + # see `uses_tiles` + push!(queries, :($gpu_name(::$TilesQuery) = true)) + end if find_collective(def[:body]) # see `uses_collectives` push!(queries, :($gpu_name(::$CollectivesQuery) = true)) end - # create constructor functions - _name = Symbol(:_, name) + # create constructor functions (a single method with keywords: in a local scope, several + # would capture each other in a `Core.Box`) constructors = quote if $(name isa Symbol ? :(!@isdefined($name)) : true) - function $_name(dev::Dev, sz::S, range::NDRange) where {Dev, S <: $_Size, NDRange <: $_Size} - return $construct(dev, sz, range, $gpu_name) - end - Core.@__doc__ $name(dev) = $_name(dev, $DynamicSize(), $DynamicSize()) - $name(dev, size) = $_name(dev, $StaticSize(size), $DynamicSize()) - $name(dev, size, range) = $_name(dev, $StaticSize(size), $StaticSize(range)) - $name(dev, size::$_Size, range::$_Size) = $_name(dev, size, range) + Core.@__doc__ $name(dev, sizes...; tile = nothing) = $construct_kernel(dev, $gpu_name, sizes...; tile) $(queries...) end end @@ -306,6 +304,17 @@ function is_scope_construct(expr::Expr) # expr.head === :let end +# Whether `stmt` uses the tile of the work-item, `@tile()` or `@index(Tile)` +function find_tile(stmt) + result = Ref(false) + postwalk(stmt) do expr + result[] |= is_macrocall(expr, Symbol("@tile")) || + (is_macrocall(expr, Symbol("@index")) && any(==(:Tile), expr.args)) + expr + end + return result[] +end + function find_collective(stmt) result = Ref(false) postwalk(stmt) do expr diff --git a/src/tiles.jl b/src/tiles.jl new file mode 100644 index 000000000..ee2a22765 --- /dev/null +++ b/src/tiles.jl @@ -0,0 +1,419 @@ +### +# Tiles +# - @tile +# - tile_shfl, tile_shfl_up, tile_shfl_down, tile_shfl_xor +# - tile_any, tile_all, tile_ballot, tile_reduce, tile_scan, tile_exclusive_scan, tile_barrier +### + +export @tile, Tile, tile_width, supports_tiles +export tile_shfl, tile_shfl_up, tile_shfl_down, tile_shfl_xor +export tile_any, tile_all, tile_ballot, tile_reduce, tile_scan, tile_exclusive_scan, tile_barrier + +""" + Tile{N} + +The tile of `N` work-items a work-item belongs to, in a kernel with tiles, as returned by +[`@tile`](@ref): `t.index` is the index of the tile within the workgroup, and `t.lane` the +index of the work-item within the tile, both 1-based. [`tile_width(t)`](@ref tile_width) is +`N`. +""" +struct Tile{N} + index::Int + lane::Int +end + +""" + @tile()::Tile{N} + +Return the work-item's tile, in a kernel constructed with a tile width, `f(backend; tile = +N)`; there is no declaration in the kernel. `t.index` is the tile's index within the +workgroup, and `t.lane` the work-item's lane within the tile, both 1-based. + +```julia +@kernel function tile_sums!(out, @Const(x)) + t = @tile() + v = tile_reduce(t, +, x[@index(Global, Linear)]) + if t.lane == 1 + out[@index(Tile)] = v + end +end +kernel = tile_sums!(backend; tile = 32) +kernel(out, x; ndrange = 32 * ntiles) +``` + +The tiles divide the workgroup by its linear local index: the work-item with +`@index(Local, Linear) == lid` is in tile `fld(lid - 1, N) + 1`, at lane `mod(lid - 1, N) + 1`. +[`@index(Tile)`](@ref @index) is the global linear index of the tile, +`(@index(Group, Linear) - 1) * (prod(@groupsize()) ÷ N) + t.index`: the tiles are numbered +workgroup by workgroup. + +The work-items of a tile communicate with the tile operations, which take the tile as their +first argument, so that helper functions can use them: [`tile_shfl`](@ref), +[`tile_shfl_up`](@ref), [`tile_shfl_down`](@ref), [`tile_shfl_xor`](@ref), +[`tile_any`](@ref), [`tile_all`](@ref), [`tile_ballot`](@ref), [`tile_reduce`](@ref), +[`tile_scan`](@ref), [`tile_exclusive_scan`](@ref) and [`tile_barrier`](@ref). All +work-items of a tile have to execute them together, in control flow that is uniform over the +tile, but can differ between tiles: e.g. a tile can `return` early, or loop a different +number of times. The work-items of a tile also have to pass the same `offset` or `mask` to a +shuffle, and the same `op` (and `init`) to a reduction or scan. The shuffles, votes, +reductions and scans don't synchronize memory accesses: use [`tile_barrier`](@ref) when the +work-items of a tile communicate through local or global memory. [`@synchronize`](@ref), +[`@groupreduce`](@ref) and [`@groupscan`](@ref) still need the whole workgroup. + +# Choosing the tile width + +Pass `tile = N` when constructing the kernel, alongside a static workgroup size or `ndrange` +if any: `f(backend; tile = N)`, `f(backend, workgroupsize; tile = N)` or +`f(backend, workgroupsize, ndrange; tile = N)`. A kernel that uses `@tile()` or `@index(Tile)` +needs it. The width is compiled into the kernel; [`tile_width(kernel)`](@ref tile_width) +returns it on the host, and `tile_width(t)` in the kernel, where it is a constant. + +`N` has to be a power of two of at most the backend's sub-group width +([`KernelInterface.sub_group_size`](@ref)). [`supports_tiles(backend, N)`](@ref supports_tiles) +reports whether the backend supports tiles of that width; constructing the kernel fails if it +doesn't. The CPU backend has no tiles. See [Tiles](@ref tiles) in the manual for choosing the +width. + +# Launching + +Kernels with tiles have restrictions on their launch: + +- the workgroup has a multiple of `N` work-items; +- every tile is either entirely within the `ndrange`, or entirely outside of it, so that the + work-items that pad a partial workgroup form whole tiles, which don't execute the kernel. + For a 1-D `ndrange`, that means a multiple of `N`. For an N-D `ndrange`, e.g. one dimension + per index of an element, it holds if the workgroup covers the leading dimensions of the + `ndrange` and the number of work-items is a multiple of `N`: e.g. an `ndrange` of + `(4, 4, nelements)` in workgroups of `(4, 4, k)`, with tiles of 16 work-items or, for an + even `k` and `nelements`, 32. + +Without an explicit workgroup size, KernelAbstractions picks one that satisfies this; an +explicit one that doesn't is an error. Kernels with tiles and an N-D `ndrange` are launched on +a 1-D grid, so that the tiles follow the backend's sub-groups. + +To give each element fewer nodes than a tile has lanes, e.g. 9 nodes of a 3×3 element in a +tile of 16, map the lanes to nodes yourself, and keep the lanes without a node in the tile +operations: clamp their loads, skip their stores, but let them shuffle, vote and reduce (with +a neutral value) like the others. +""" +macro tile() + return :($__tile($(esc(:__ctx__)), $__tile_width($(esc(:__ctx__))))) +end + +""" + tile_width(kernel::Kernel)::Union{Int, Nothing} + tile_width(t::Tile{N}) = N + +The tile width of a kernel with tiles, or `nothing` for a kernel without. In a kernel, +`tile_width(t)` of a tile `t` from [`@tile`](@ref) is a constant. +""" +tile_width(::Kernel{<:Any, <:Any, <:Any, <:Any, T}) where {T} = T +tile_width(::Tile{N}) where {N} = N + +# the tile width for the kernel's metadata, see `CompilerMetadata` +tile_val(kernel::Kernel) = (N = tile_width(kernel); N === nothing ? nothing : Val(N)) + +@inline function __tile(ctx, ::Val{N}) where {N} + lid0 = __index_Local_Linear(ctx) - 1 + return Tile{N}(lid0 ÷ N + 1, lid0 % N + 1) +end + +@inline function __index_Tile(ctx, ::Val{N}) where {N} + lid0 = __index_Local_Linear(ctx) - 1 + return (__index_Group_Linear(ctx) - 1) * (prod(groupsize(ctx)) ÷ N) + lid0 ÷ N + 1 +end + +# a kernel without a tile width, whose use of tiles the `@kernel` macro didn't see (e.g. in the +# expansion of another macro), so that its construction didn't reject it +__tile(ctx, ::Nothing) = throw(ArgumentError(tile_width_message)) +__index_Tile(ctx, ::Nothing) = throw(ArgumentError(tile_width_message)) + +# A tile is a segment of `N` lanes of a sub-group, so the tile operations are KernelInterface's +# sub-group operations with a `width` of `N`: where several tiles share a sub-group, which +# `supports_tiles` only allows with `KI.supports_partitionable_subgroups`, those only need the +# tile to execute them together. + +""" + tile_shfl(t::Tile{N}, val, lane) + +`val` of the work-item at lane `mod(lane - 1, N) + 1` of the tile `t`. The work-items may read +from different lanes. See [`KernelInterface.shfl`](@ref) for the supported types. +""" +@inline tile_shfl(::Tile{N}, val, lane::Integer) where {N} = KI.shfl(val, lane, N) + +""" + tile_shfl_up(t::Tile{N}, val, offset) + +`val` of the work-item `offset` lanes earlier in the tile `t`, or the work-item's own `val` +if that lane is outside of the tile. `offset` has to be at least 0, and the same for the +work-items of the tile. See [`KernelInterface.shfl`](@ref) for the supported types. +""" +@inline tile_shfl_up(::Tile{N}, val, offset::Integer) where {N} = KI.shfl_up(val, offset, N) + +""" + tile_shfl_down(t::Tile{N}, val, offset) + +`val` of the work-item `offset` lanes later in the tile `t`, or the work-item's own `val` if +that lane is outside of the tile. `offset` has to be at least 0, and the same for the +work-items of the tile. See [`KernelInterface.shfl`](@ref) for the supported types. +""" +@inline tile_shfl_down(::Tile{N}, val, offset::Integer) where {N} = KI.shfl_down(val, offset, N) + +""" + tile_shfl_xor(t::Tile{N}, val, mask) + +`val` of the work-item at lane `((t.lane - 1) ⊻ mask) + 1` of the tile `t`. `mask` has to be +between 0 and `N - 1`, and the same for the work-items of the tile. See +[`KernelInterface.shfl`](@ref) for the supported types. +""" +@inline tile_shfl_xor(::Tile{N}, val, mask::Integer) where {N} = KI.shfl_xor(val, mask, N) + +""" + tile_any(t::Tile, pred::Bool)::Bool + +Whether `pred` is `true` for any work-item of the tile `t`. +""" +@inline tile_any(::Tile{N}, pred::Bool) where {N} = KI.sub_group_any(pred, N) + +""" + tile_all(t::Tile, pred::Bool)::Bool + +Whether `pred` is `true` for all work-items of the tile `t`. +""" +@inline tile_all(::Tile{N}, pred::Bool) where {N} = KI.sub_group_all(pred, N) + +""" + tile_ballot(t::Tile, pred::Bool)::UInt64 + +A mask of the work-items of the tile `t` for which `pred` is `true`: bit `i - 1` for the +work-item at lane `i`. Needs a sub-group width of at most 64 +([`KernelInterface.sub_group_size`](@ref)), also for narrower tiles. +""" +@inline tile_ballot(::Tile{N}, pred::Bool) where {N} = KI.sub_group_ballot(pred, N) + +""" + tile_reduce(t::Tile, op, val) + +Reduce `val` over the work-items of the tile `t` with the associative operator `op`, in the +order of the lanes, and return the result on every work-item of the tile. `op` has to be the +same for the work-items of the tile, and needn't be commutative. See [`KernelInterface.sub_group_reduce`](@ref) for the +supported types, and where the order of floating-point operations may differ. +""" +@inline tile_reduce(::Tile{N}, op, val) where {N} = KI.sub_group_reduce(op, val, N) + +""" + tile_scan(t::Tile, op, val) + +The inclusive scan of `val` over the work-items of the tile `t` with the associative +operator `op`, in the order of the lanes. `op` has to be the same for the work-items of the +tile, and needn't be commutative. See [`KernelInterface.sub_group_scan`](@ref) for the +supported types, and where the order of floating-point operations may differ. +""" +@inline tile_scan(::Tile{N}, op, val) where {N} = KI.sub_group_scan(op, val, N) + +""" + tile_exclusive_scan(t::Tile, op, val, init) + +The exclusive scan of `val` over the work-items of the tile `t` with the associative +operator `op`, in the order of the lanes, starting from `init`: lane 1 gets `init`, and lane +`i > 1` gets `op(init, r)` for the reduction `r` of the values of lanes `1:i-1`. `op` and +`init` have to be the same for the work-items of the tile, and `init` of the type of `val`; +it needn't be an identity of `op`. See [`KernelInterface.sub_group_exclusive_scan`](@ref) +for the supported types, and where the order of floating-point operations may differ. +""" +@inline tile_exclusive_scan(::Tile{N}, op, val, init) where {N} = KI.sub_group_exclusive_scan(op, val, init, N) + +""" + tile_barrier(t::Tile) + +Wait until all work-items of the tile `t` have reached the barrier, and make their writes to +global and local memory before it visible to the tile. It doesn't order memory between +tiles. +""" +@inline tile_barrier(::Tile{N}) where {N} = KI.sub_group_barrier(N) + + +## host side + +""" + supports_tiles(backend, N)::Bool + +Whether kernels with tiles of `N` work-items (see [`@tile`](@ref)) can run on `backend`. +Constructing such a kernel on a backend that doesn't support them is an error. + +That needs a power-of-two `N` of at most the sub-group width `W` +([`KernelInterface.sub_group_size`](@ref)), sub-groups that are formed from consecutive +work-items and may diverge from each other +([`KernelInterface.supports_linear_subgroups`](@ref), +[`KernelInterface.supports_independent_subgroups`](@ref)), and for `N < W`, sub-groups that can +be partitioned into segments that diverge as well +([`KernelInterface.supports_partitionable_subgroups`](@ref)). The CPU backend returns `false`. A launch also has to meet the workgroup and `ndrange` restrictions of +[`@tile`](@ref). +""" +supports_tiles(backend, N::Integer) = tile_support_error(backend, N) === nothing + +# Why `backend` doesn't support tiles of `N` work-items, or `nothing` if it does +function tile_support_error(backend, N::Integer) + (N > 0 && ispow2(N)) || return "the width has to be a power of two" + backend isa KI.Backend && KI.supports_linear_subgroups(backend) || + return "the backend doesn't support tiles, as its sub-groups aren't formed from consecutive work-items (see `KernelInterface.supports_linear_subgroups`)" + KI.supports_independent_subgroups(backend) || + return "the backend doesn't support tiles, as its sub-groups can't diverge from each other (see `KernelInterface.supports_independent_subgroups`)" + W = KI.sub_group_size(backend) + N <= W || return "the width is larger than the sub-group width $W of the backend" + N == W || KI.supports_partitionable_subgroups(backend) || + return "the backend only supports tiles of its sub-group width $W, as its sub-groups can't be partitioned into parts that diverge (see `KernelInterface.supports_partitionable_subgroups`)" + return nothing +end + +# Whether the function of a kernel uses its tiles (`@tile()` or `@index(Tile)`), so that it +# needs a tile width. The `@kernel` macro adds a method taking a `TilesQuery` to the function +# of such a kernel, with its first method. (`applicable` is constant-folded, so that other +# kernels don't look up the method.) +struct TilesQuery end +uses_tiles(f) = applicable(f, TilesQuery()) && query_tiles(f) + +# a callable that forwards its arguments, e.g. a wrapper of a kernel function, doesn't use +# tiles itself +@noinline function query_tiles(f) + sig = Base.unwrap_unionall(which(f, Tuple{TilesQuery}).sig) + return sig.parameters[end] === TilesQuery && f(TilesQuery()) +end + +const tile_width_message = "the kernel uses `@tile()` or `@index(Tile)`, so it needs a tile width: construct it with `f(backend; tile = N)`" + +# The tile width of a kernel of `f`, for the `tile` keyword of its constructor +function resolve_tile_width(f, tile) + if tile === nothing + uses_tiles(f) && tile_error(nothing, tile_width_message) + return nothing + end + tile isa Integer || throw(ArgumentError("`tile` has to be an integer, got $tile")) + return Int(tile) +end + +tile_error(N, msg) = throw(ArgumentError(N === nothing ? "Tiles: $msg" : "Tiles of $N work-items: $msg")) + +# The workgroup size for a launch of a kernel with tiles of `N` work-items over `range` with +# `launch` (see `select_launch`), for about `threads` work-items: the largest multiple of `N` +# work-items, up to `threads` (but at least `N`, leaving it to the launch to check the kernel's +# limit), that divides `range` into whole tiles. +function tile_launch_workgroupsize(backend, launch, N, threads, range) + dims = extents(range) + for t in (max(threads, N) ÷ N * N):-N:N + items = length(dims) == 1 ? launch_workgroupsize(backend, launch, t, range) : + tile_shape(N, t, dims, index_type(launch)) + items !== nothing && prod(items) % N == 0 && whole_tiles(N, dims, items) && return items + end + return tile_error(N, "no workgroup size divides the ndrange $dims into whole tiles, pass one") +end + +# The shape of a workgroup of `t` work-items, no larger than `dims` along any dimension, that +# divides `dims` into whole tiles of `N` work-items with the least padding, or `nothing`. The +# padded iteration space has to fit `T`, the index type `select_launch` chose for a workgroup +# from `launch_workgroupsize`, which pads differently. +function tile_shape(N, t, dims::Dims, ::Type{T}) where {T} + best, best_padded = nothing, typemax(Int) + for items in workgroup_shapes(t, dims) + padded = map((d, i) -> Int128(cld(d, i)) * i, dims, items) + all(<=(typemax(Int)), padded) && fits(T, Int.(padded)) || continue + n = prod(Int.(padded)) + if n < best_padded && whole_tiles(N, dims, items) + best, best_padded = items, n + end + end + return best +end + +# The shapes of `t` work-items that are no larger than `dims` along any dimension, larger +# extents of the first dimensions first +function workgroup_shapes(t, dims::Dims) + length(dims) == 1 && return t <= dims[1] ? [(t,)] : Dims{1}[] + shapes = Dims{length(dims)}[] + for a in min(t, dims[1]):-1:1 + t % a == 0 || continue + for rest in workgroup_shapes(t ÷ a, Base.tail(dims)) + push!(shapes, (a, rest...)) + end + end + return shapes +end + +# Check the launch of `kernel` with tiles of `N` work-items with the `ndrange`, `workgroupsize` +# and `iterspace` from `launch_config`, before compiling it: the partitioned workgroup size, or +# that there is one that will be tuned. +function check_tiles_launch(kernel::Kernel, N, ndrange, workgroupsize, iterspace) + b = backend(kernel) + range = something(ndrange, static_ndrange(kernel)) + if workgroupsize === nothing && KernelAbstractions.workgroupsize(kernel) <: DynamicSize + check_tile_width(b, N) + check_tile_items(N, extents(range)) + launch = select_launch(kernel, workgroupsize, iterspace) + tile_launch_workgroupsize(b, launch, N, KI.max_work_group_size(b), range) + return + end + return check_tiles(b, N, extents(range), size(workitems(iterspace))) +end + +function check_tile_width(backend, N) + err = tile_support_error(backend, N) + err === nothing || tile_error(N, err) + return +end + +check_tile_items(N, ndrange::Dims) = saturated_prod(ndrange) % N == 0 || + tile_error(N, "the ndrange $(length(ndrange) == 1 ? only(ndrange) : ndrange) has $(saturated_prod(ndrange)) work-items, which isn't a multiple of $N") + +# Check the launch of a kernel with tiles of `N` work-items over `ndrange`, in workgroups of +# `workgroupsize` (logical sizes, before a launch on a 1-D grid flattens them) +function check_tiles(backend, N, ndrange::Dims, workgroupsize::Dims) + check_tile_width(backend, N) + check_tile_items(N, ndrange) + check_tile_workgroup(N, workgroupsize) + whole_tiles(N, ndrange, workgroupsize) || + tile_error(N, "workgroups of $workgroupsize divide the ndrange $ndrange into partial tiles") + return +end + +function check_tile_workgroup(N, workgroupsize::Dims) + items = prod(workgroupsize) + items % N == 0 || + tile_error(N, "the workgroup size $(length(workgroupsize) == 1 ? items : workgroupsize) isn't a multiple of $N") + return +end + +# Whether workgroups of size `workgroupsize` (a multiple of `N` work-items) divide `ndrange` +# into tiles of `N` consecutive work-items that are each entirely inside or outside of it. The +# part of a workgroup inside the `ndrange` is a box: the whole workgroup, or in the last +# workgroup along a dimension, what's left of the `ndrange`. So it suffices to check one +# workgroup per combination of the dimensions it's the last one along. +function whole_tiles(N, ndrange::Dims, workgroupsize::Dims) + n = max(length(ndrange), length(workgroupsize)) + range = ntuple(d -> Base.get(ndrange, d, 1), n) + items = ntuple(d -> Base.get(workgroupsize, d, 1), n) + # an empty `ndrange` has no workgroups + any(iszero, range) && return true + insides = map(range, items) do r, i + last = r - (cld(r, i) - 1) * i + return r > i && last != i ? (i, last) : (last,) + end + # a 1-D box, a prefix of the workgroup, is whole tiles if its length is a multiple of `N` + n == 1 && return all(inside -> inside % N == 0, only(insides)) + return all(inside -> whole_tiles_box(N, items, inside), Iterators.product(insides...)) +end + +# Whether every tile of `N` consecutive work-items of a workgroup of size `items` is entirely +# inside or outside of the box of size `inside` at its start +function whole_tiles_box(N, items, inside) + inside == items && return true + tile_inside = false + for (l, I) in enumerate(CartesianIndices(items)) + isinside = all(Tuple(I) .<= inside) + if (l - 1) % N == 0 + tile_inside = isinside + elseif isinside != tile_inside + return false + end + end + return true +end diff --git a/test/runtests.jl b/test/runtests.jl index 088b010a2..67b6a8f9d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -28,6 +28,49 @@ end @test KernelAbstractions.precompiled_launch[] skip = !workloads end +# A package with overloaded kernels, which have to define the kernel's properties only once, +# since methods can't be overwritten during precompilation +@testset "Precompiling overloaded kernels" begin + mktempdir() do dir + mkpath(joinpath(dir, "OverloadedKernels", "src")) + write( + joinpath(dir, "OverloadedKernels", "src", "OverloadedKernels.jl"), + """ + module OverloadedKernels + using KernelAbstractions + @kernel function plain!(a) + a[@index(Global)] = 1 + end + @kernel function plain!(a, b) + a[@index(Global)] = b + end + @kernel function tiled!(a) + a[@index(Global)] = tile_width(@tile()) + end + @kernel function tiled!(a, b) + a[@index(Global)] = b + end + @kernel function reduced!(a) + res = @groupreduce(+, a[@index(Global)], zero(eltype(a)); groupsize = 64) + a[@index(Global)] = res + end + @kernel function reduced!(a, b) + res = @groupreduce(+, b, zero(eltype(a)); groupsize = 64) + a[@index(Global)] = res + end + end + """ + ) + pushfirst!(LOAD_PATH, dir) + try + id = Base.PkgId("OverloadedKernels") + @test Base.compilecache(id) isa Union{String, Tuple} + finally + popfirst!(LOAD_PATH) + end + end +end + KernelAbstractions.versioninfo(POCLBackend()) @info "Configuration" pocl = KernelAbstractions.POCL.nanoOpenCL.pocl_standalone_jll.libpocl diff --git a/test/testsuite.jl b/test/testsuite.jl index 375407df2..4ff4a4646 100644 --- a/test/testsuite.jl +++ b/test/testsuite.jl @@ -46,6 +46,7 @@ include("specialfunctions.jl") include("random.jl") include("spawn.jl") include("groupreduce.jl") +include("tiles.jl") function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{String}()) @conditional_testset "Unittests" skip_tests begin @@ -128,6 +129,10 @@ function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{ groupreduce_testsuite(backend, AT) end + @conditional_testset "Tiles" skip_tests begin + tile_testsuite(backend, AT) + end + return end diff --git a/test/tiles.jl b/test/tiles.jl new file mode 100644 index 000000000..7d0ecb9a2 --- /dev/null +++ b/test/tiles.jl @@ -0,0 +1,411 @@ +using StaticArrays + +# Kernels with tiles, whose width is chosen when they are constructed, so that the same kernel +# runs with several widths. + +# every tile loops a different number of times, and some return early +@kernel function tile_divergent!(out, @Const(x)) + t = @tile() + i = @index(Global, Linear) + g = @index(Tile) + acc = zero(eltype(out)) + for _ in 1:(g % 3) + acc += tile_reduce(t, +, x[i]) + end + if g % 4 == 2 + return + end + out[i] = acc + tile_shfl(t, x[i], 1) +end + +# the tile, its index and its rotation by a lane, the votes, and the ballot of the odd values +# (which needs a sub-group width of at most 64) +@kernel function tile_ops!(out, ballots, @Const(x), ::Val{ballot}) where {ballot} + t = @tile() + i = @index(Global, Linear) + out[i, 1] = t.index + out[i, 2] = t.lane + out[i, 3] = @index(Tile) + out[i, 4] = tile_shfl(t, x[i], t.lane + 1) + out[i, 5] = tile_shfl_up(t, x[i], 1) + out[i, 6] = tile_shfl_down(t, x[i], 1) + out[i, 7] = tile_shfl_xor(t, x[i], 1) + out[i, 8] = tile_any(t, x[i] == 3) + out[i, 9] = tile_all(t, x[i] > 0) + out[i, 10] = tile_width(t) + ballots[i] = ballot ? tile_ballot(t, isodd(x[i])) : zero(UInt64) +end + +# code that depends on the width: a mask with a bit per lane, of the smallest type +@inline tile_mask_type(t::Tile) = tile_width(t) <= 8 ? UInt8 : tile_width(t) <= 16 ? UInt16 : + tile_width(t) <= 32 ? UInt32 : UInt64 +@kernel function tile_mask!(out, sizes, @Const(x)) + t = @tile() + i = @index(Global, Linear) + mask = tile_ballot(t, isodd(x[i])) % tile_mask_type(t) + out[i] = mask + sizes[i] = sizeof(mask) +end + +# communication through local memory, ordered by `tile_barrier`, with a static workgroup +# size +@kernel function tile_barrier!(out, @Const(x)) + t = @tile() + i = @index(Global, Linear) + lid = @index(Local, Linear) + mem = @localmem eltype(out) (prod(@groupsize()),) + mem[lid] = x[i] + tile_barrier(t) + N = tile_width(t) + out[i] = mem[(t.index - 1) * N + mod1(t.lane + 1, N)] +end + +tile_add(a, b) = a .+ b + +# reductions of tuples and static arrays +@kernel function tile_reduce!(out, out_sv, @Const(x)) + t = @tile() + i = @index(Global, Linear) + out[i] = tile_reduce(t, tile_add, (x[i], 2 * x[i])) + out_sv[i] = tile_reduce(t, +, SVector(x[i], -x[i], Int32(1))) +end + +# the composition of affine maps, first `f` then `g`: associative, but not commutative +tile_compose(f, g) = (g[1] * f[1], g[1] * f[2] + g[2]) + +@kernel function tile_scans!(incl, excl, affine, affine_excl, @Const(x), @Const(maps)) + t = @tile() + i = @index(Global, Linear) + incl[i] = tile_scan(t, +, x[i]) + excl[i] = tile_exclusive_scan(t, +, x[i], Int32(7)) + affine[i] = tile_scan(t, tile_compose, maps[i]) + affine_excl[i] = tile_exclusive_scan(t, tile_compose, maps[i], (Int32(2), Int32(3))) +end + +@kernel unsafe_indices = true function tile_unsafe!(out, @Const(x)) + t = @tile() + i = @index(Global, Linear) + v = i <= length(x) ? x[i] : zero(eltype(x)) + s = tile_reduce(t, +, v) + if i <= length(out) + out[i] = s + end +end + +# a kernel that doesn't compile, to check that invalid launches are rejected before +@kernel function tile_uncompilable!(out) + out[@index(Global, Linear)] = Base.inferencebarrier(identity)(1) +end + +# a kernel without tiles +@kernel function tile_plain!(out) + out[@index(Global, Linear)] = 1 +end + +@kernel function tile_noargs!() +end + +# the nodes of an element per tile, over an N-D ndrange +@kernel function tile_elements!(out, index, @Const(x)) + t = @tile() + I = @index(Global, NTuple) + out[I...] = tile_reduce(t, +, x[I...]) + index[I...] = @index(Tile) +end + +struct TileForward{F} + f::F +end +(w::TileForward)(args...) = w.f(args...) + +# Whether every tile of `N` consecutive work-items of every workgroup is entirely inside or +# outside of the ndrange, by enumerating them +function whole_tiles_reference(N, ndrange, workgroupsize) + for b in CartesianIndices(cld.(ndrange, workgroupsize)) + inside = [all((Tuple(b) .- 1) .* workgroupsize .+ Tuple(I) .<= ndrange) for I in CartesianIndices(workgroupsize)] + all(t -> allequal(inside[t:(t + N - 1)]), 1:N:length(inside)) || return false + end + return true +end + +# the reference results, by tile of `N` consecutive elements +tile_groups(n, N) = [((i - 1) ÷ N * N + 1):((i - 1) ÷ N * N + N) for i in 1:n] + +function tile_testsuite_width(backend, AT, N) + W = KI.sub_group_size(backend) + # several workgroups, and a partial one with a padding tile + groupsize = 2N + ntiles = 5 + n = ntiles * N + x = Int32.(rand(1:100, n)) + groups = tile_groups(n, N) + + @test tile_width(tile_divergent!(backend; tile = N)) == N + @test tile_width(tile_barrier!(backend, groupsize; tile = N)) == N + + divergent = map(1:n) do i + g = (i - 1) ÷ N + 1 + g % 4 == 2 && return Int32(-1) + return Int32(g % 3) * sum(x[groups[i]]) + x[first(groups[i])] + end + @testset "divergent tiles" begin + out = AT(fill(Int32(-1), n)) + tile_divergent!(backend; tile = N)(out, AT(x); ndrange = n) + @test Array(out) == divergent + end + + @testset "tile operations" begin + x3 = copy(x) + x3[2] = 3 + out = AT(zeros(Int32, n, 10)) + ballots = AT(zeros(UInt64, n)) + tile_ops!(backend; tile = N)(out, ballots, AT(x3), Val(W <= 64); ndrange = n) + out, ballots = Array(out), Array(ballots) + for i in 1:n + g = groups[i] + lane = i - first(g) + 1 + @test out[i, 2] == lane + @test out[i, 3] == (i - 1) ÷ N + 1 + @test out[i, 4] == x3[g[mod1(lane + 1, N)]] + @test out[i, 5] == x3[lane > 1 ? i - 1 : i] + @test out[i, 6] == x3[lane < N ? i + 1 : i] + @test out[i, 7] == x3[g[((lane - 1) ⊻ 1) + 1]] + @test out[i, 8] == any(==(3), x3[g]) + @test out[i, 9] == 1 + @test out[i, 10] == N + if W <= 64 + @test ballots[i] == sum(UInt64(1) << (j - 1) for j in 1:N if isodd(x3[g[j]]); init = UInt64(0)) + end + end + # the tile within the workgroup + @test all(i -> out[i, 1] >= 1, 1:n) + end + + if W <= 64 + @testset "width-dependent code" begin + out = AT(zeros(UInt64, n)) + sizes = AT(zeros(Int, n)) + tile_mask!(backend; tile = N)(out, sizes, AT(x); ndrange = n) + @test all(==(cld(N, 8)), Array(sizes)) + @test Array(out) == [sum(UInt64(1) << (j - 1) for j in 1:N if isodd(x[g[j]]); init = UInt64(0)) for g in groups] + end + end + + @testset "tile_barrier" begin + out = AT(zeros(Int32, n)) + tile_barrier!(backend, groupsize; tile = N)(out, AT(x); ndrange = n) + @test Array(out) == [x[g[mod1(i - first(g) + 2, N)]] for (i, g) in enumerate(groups)] + end + + @testset "tuples and static arrays" begin + out = AT(fill((Int32(0), Int32(0)), n)) + out_sv = AT(fill(SVector{3, Int32}(0, 0, 0), n)) + tile_reduce!(backend; tile = N)(out, out_sv, AT(x); ndrange = n) + @test Array(out) == [(sum(x[g]), 2 * sum(x[g])) for g in groups] + @test Array(out_sv) == [SVector(sum(x[g]), -sum(x[g]), Int32(N)) for g in groups] + end + + @testset "scans" begin + maps = [(Int32(rand((-1, 1, 2))), Int32(rand(-5:5))) for _ in 1:n] + incl, excl = AT(zeros(Int32, n)), AT(zeros(Int32, n)) + affine, affine_excl = AT(fill((Int32(0), Int32(0)), n)), AT(fill((Int32(0), Int32(0)), n)) + tile_scans!(backend; tile = N)(incl, excl, affine, affine_excl, AT(x), AT(maps); ndrange = n) + prefix(f, v, g, i) = foldl(f, v[first(g):i]) + @test Array(incl) == [prefix(+, x, g, i) for (i, g) in enumerate(groups)] + @test Array(excl) == [i == first(g) ? Int32(7) : Int32(7) + prefix(+, x, g, i - 1) for (i, g) in enumerate(groups)] + @test Array(affine) == [prefix(tile_compose, maps, g, i) for (i, g) in enumerate(groups)] + init = (Int32(2), Int32(3)) + @test Array(affine_excl) == [i == first(g) ? init : tile_compose(init, prefix(tile_compose, maps, g, i - 1)) for (i, g) in enumerate(groups)] + end + + @testset "unsafe_indices" begin + out = AT(zeros(Int32, n)) + tile_unsafe!(backend; tile = N)(out, AT(x); ndrange = n) + @test Array(out) == [sum(x[g]) for g in groups] + end + + @testset "explicit workgroup sizes" begin + out = AT(fill(Int32(-1), n)) + # a single tile per workgroup + tile_divergent!(backend; tile = N)(out, AT(x); ndrange = n, workgroupsize = N) + @test Array(out) == divergent + # several tiles per workgroup, which diverge, with a padding tile in the last one + for workgroupsize in (2N, 4N) + out = AT(fill(Int32(-1), n)) + tile_divergent!(backend; tile = N)(out, AT(x); ndrange = n, workgroupsize) + @test Array(out) == divergent + out = AT(fill(Int32(-1), n)) + tile_unsafe!(backend; tile = N)(out, AT(x); ndrange = n, workgroupsize) + @test Array(out) == [sum(x[g]) for g in groups] + end + end + + @testset "rejected launches" begin + out = AT(zeros(Int32, n)) + k = tile_divergent!(backend; tile = N) + @test_throws "isn't a multiple of $N" k(out, AT(x); ndrange = n - 1) + @test_throws "isn't a multiple of $N" k(out, AT(x); ndrange = n, workgroupsize = N ÷ 2) + @test_throws "isn't a multiple of $N" tile_barrier!(backend, N + N ÷ 2; tile = N) + @test_throws "isn't a multiple of $N" k(out, AT(x); ndrange = (N - 1, 3)) + @test_throws "isn't a multiple of $N" tile_divergent!(backend, N, n - 1; tile = N) + # before compiling the kernel + u = tile_uncompilable!(backend; tile = N) + @test_throws "isn't a multiple of $N" u(out; ndrange = n - 1) + @test_throws "isn't a multiple of $N" u(out; ndrange = n, workgroupsize = N ÷ 2) + @test_throws "isn't a multiple of $N" u(out; ndrange = (N - 1, 3)) + @test_throws "isn't a multiple of $N" u(out; ndrange = n, workgroupsize = 0) + if N >= 4 + # tiles that would straddle the end of the first dimension + @test_throws "partial tiles" u(out; ndrange = (3, N), workgroupsize = (2, N ÷ 2)) + end + end + + @testset "N-D ndrange" begin + # elements of 2 × N/2 nodes, a tile each, in workgroups of one or several elements + dims = (2, N ÷ 2) + nelem = 5 + x = Int32.(rand(1:100, dims..., nelem)) + sums = [sum(x[:, :, e]) for _ in 1:dims[1], _ in 1:dims[2], e in 1:nelem] + elements = [e for _ in 1:dims[1], _ in 1:dims[2], e in 1:nelem] + k = tile_elements!(backend; tile = N) + for workgroupsize in (nothing, (dims..., 1), (dims..., 2)) + out, index = AT(zeros(Int32, size(x))), AT(zeros(Int, size(x))) + k(out, index, AT(x); ndrange = size(x), workgroupsize) + @test Array(out) == sums + # tiles are numbered workgroup by workgroup, so one per element + @test Array(index) == elements + end + end + return +end + +function tile_testsuite(backend, AT) + b = backend() + @testset "widths" begin + # kernels that use tiles need a width + @test KernelAbstractions.uses_tiles(gpu_tile_divergent!) + @test KernelAbstractions.uses_tiles(gpu_tile_elements!) + @test_throws "needs a tile width" tile_divergent!(b) + @test_throws "needs a tile width" tile_divergent!(b, 64) + @test_throws "has to be an integer" tile_divergent!(b; tile = 8.0) + @test tile_width(tile_plain!(b)) === nothing + @test !KernelAbstractions.uses_tiles(tile_plain!(b).f) + # also for a kernel without arguments, whose function takes a single argument + @test !KernelAbstractions.uses_tiles(tile_noargs!(b).f) + # a kernel with tiles in a local scope + function local_tile_kernel(b) + @kernel function local_tiles!(out) + t = @tile() + out[@index(Global, Linear)] = tile_width(t) + end + return local_tiles!(b; tile = 8) + end + if supports_tiles(b, 8) + @test tile_width(local_tile_kernel(b)) == 8 + end + end + @test !supports_tiles(b, 3) + @test !supports_tiles(b, 0) + @test_throws "power of two" tile_divergent!(b; tile = 3) + + @testset "whole tiles" begin + whole_tiles = KernelAbstractions.whole_tiles + @test whole_tiles(16, (4, 4, 5), (4, 4, 2)) + @test !whole_tiles(32, (4, 4, 5), (4, 4, 2)) + @test whole_tiles(32, (4, 4, 6), (4, 4, 2)) + @test !whole_tiles(16, (6, 4), (4, 4)) + @test !whole_tiles(4, (6, 4), (4, 4)) + @test whole_tiles(2, (6, 4), (4, 4)) + @test whole_tiles(16, (64, 60), (16, 16)) + @test whole_tiles(32, (64, 60), (16, 16)) + @test !whole_tiles(32, (60, 64), (16, 16)) + @test !whole_tiles(64, (64, 62), (16, 16)) + @test whole_tiles(8, (8,), (16,)) + @test !whole_tiles(8, (12,), (16,)) + # tiles that don't follow the dimensions of the workgroup + @test whole_tiles(16, (16, 1), (24, 2)) + ok = true + @test whole_tiles(32, (0, 1), (1, 32)) + for N in (2, 4, 8), w1 in 1:6, w2 in 1:4, r1 in 0:7, r2 in 1:5 + (w1 * w2) % N == 0 || continue + ok &= whole_tiles(N, (r1, r2), (w1, w2)) == whole_tiles_reference(N, (r1, r2), (w1, w2)) + end + for N in (2, 4), w in Iterators.product(1:3, 1:3, 1:2), r in Iterators.product(1:4, 1:3, 1:3) + prod(w) % N == 0 || continue + ok &= whole_tiles(N, r, w) == whole_tiles_reference(N, r, w) + end + @test ok + end + + @testset "padded shapes" begin + # the padded iteration space fits the index type, without overflowing + shape = KernelAbstractions.tile_shape + @test shape(32, 1024, (32_000_016, 66), Int32) === nothing + @test shape(32, 32, (32_000_016, 66), Int32) == (16, 2) + @test shape(32, 1024, (32 * (typemax(Int) ÷ 96 - 100), 3), Int) == (1024, 1) + end + + W = KI.supports_subgroups(b) ? KI.sub_group_size(b) : 0 + if W == 0 || !supports_tiles(b, W) + # a backend without tiles: constructing a kernel with tiles is an error + @test_throws "doesn't support tiles" tile_divergent!(b; tile = 8) + W > 0 && @test_throws "doesn't support tiles" tile_divergent!(b; tile = W) + W > 1 && @test_throws "doesn't support tiles" tile_barrier!(b, W; tile = W ÷ 2) + return + end + @test !supports_tiles(b, 2W) + for N in (W ÷ 2, W ÷ 4) + N >= 1 || continue + @test supports_tiles(b, N) == KI.supports_partitionable_subgroups(b) + end + + @testset "automatic workgroup sizes" begin + size = KernelAbstractions.tile_launch_workgroupsize + launch = KernelAbstractions.LinearLaunch{Int32}() + # a recommendation of less than a tile still gets one + @test size(b, launch, W, max(W ÷ 4, 1), (2W,)) == (W,) + # shapes that the leading dimensions don't divide into whole tiles + for range in ((3, W), (3, 64W), (5, 3, 2W), (6, W ÷ 2)) + items = size(b, launch, W, 4W, range) + @test prod(items) % W == 0 && whole_tiles_reference(W, range, items) + end + @test_throws "no workgroup size" size(b, launch, W, 4W, (3, W + 1)) + end + + @testset "construction" begin + @test tile_width(tile_divergent!(b; tile = W)) == W + @test tile_width(tile_divergent!(b, W; tile = W)) == W + @test tile_width(tile_divergent!(b, W, 4W; tile = W)) == W + # a kernel that doesn't use tiles can still be given a width, for its launch + @test tile_width(tile_plain!(b; tile = W)) == W + @test_throws "larger than the sub-group width" tile_divergent!(b; tile = 2W) + # `similar` keeps the width, and a function that uses tiles needs one + k16 = tile_divergent!(b; tile = W ÷ 2) + @test tile_width(similar(k16, gpu_tile_ops!)) == W ÷ 2 + @test_throws "needs a tile width" similar(tile_plain!(b), k16.f) + @test tile_width(similar(tile_plain!(b), tile_plain!(b).f)) === nothing + # a wrapper that forwards its arguments keeps the width + @test tile_width(similar(k16, TileForward(k16.f))) == W ÷ 2 + @test tile_width(similar(tile_plain!(b), TileForward(tile_plain!(b).f))) === nothing + # the width reaches the device through the kernel's metadata, also when adapted + iterspace, _ = KernelAbstractions.partition(k16, (W,), (W ÷ 2,)) + ctx = KernelAbstractions.mkcontext(k16, (W,), iterspace) + @test KernelAbstractions.__tile_width(ctx) === Val(W ÷ 2) + @test KernelAbstractions.__tile_width(KernelAbstractions.Adapt.adapt(AT, ctx)) === Val(W ÷ 2) + @test KernelAbstractions.__tile_width(KernelAbstractions.mkcontext(tile_plain!(b), (W,), iterspace)) === nothing + end + + # the sub-group width, and narrower widths, several of which share a sub-group, with the + # same kernels + for N in unique((W, W ÷ 2, W ÷ 4)) + N >= 2 || continue + if supports_tiles(b, N) + @testset "tile = $N" begin + tile_testsuite_width(b, AT, N) + end + else + @test_throws "only supports tiles of its sub-group width" tile_divergent!(b; tile = N) + end + end + return +end