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 diff --git a/docs/src/index.md b/docs/src/index.md index 1f13edeef..e99af31f9 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -203,6 +203,11 @@ end than with 0.9. - `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 4eba70529..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_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), [`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, [`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, [`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,27 +146,74 @@ 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 [`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, [`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 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 + 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 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. +- [`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" + 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 @@ -181,6 +228,7 @@ get_sub_group_local_id ```@docs barrier sub_group_barrier +sub_group_barrier(::Integer) ``` ### Memory @@ -191,8 +239,25 @@ localmemory ### Communication +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 shfl_down +shfl_up +shfl_xor +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 @@ -247,6 +312,9 @@ supports_float64 ```@docs supports_subgroups +supports_linear_subgroups +supports_independent_subgroups +supports_partitionable_subgroups supports_shuffle ``` @@ -301,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 0f0fade1a..0c83b74fa 100644 --- a/lib/KernelInterface/src/KernelInterface.jl +++ b/lib/KernelInterface/src/KernelInterface.jl @@ -31,7 +31,10 @@ 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, + :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"), :launch_configuration, :max_work_group_size, :max_work_group_dims, @@ -42,7 +45,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, :supports_partitionable_subgroups, ) ) end diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl index 5b0090914..29a5a5681 100644 --- a/lib/KernelInterface/src/device.jl +++ b/lib/KernelInterface/src/device.jl @@ -127,19 +127,23 @@ 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, `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 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 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`. @@ -155,12 +159,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. 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`. !!! 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 +186,9 @@ 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 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`. @@ -191,7 +205,9 @@ 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 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`. @@ -209,7 +225,9 @@ 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 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`. @@ -254,26 +272,667 @@ 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. +# +# 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 + +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` 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 +(e.g. `Bool`, `Char` or `Int64`) if the backend supports `UInt32`, and `isbits` structs and +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 + 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{...}} + ``` +""" +@inline shfl(val, lane::Integer) = shfl_fallback(ShuffleOp(shfl, (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` doesn't synchronize memory accesses. 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(ShuffleOp(shfl_down, (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` doesn't synchronize memory accesses. 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(ShuffleOp(shfl_up, (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). For a reduction over the sub-group, use +[`sub_group_reduce`](@ref), which also handles partial sub-groups. + +`shfl_xor` doesn't synchronize memory accesses. 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(ShuffleOp(shfl_xor, (mask,)), val) + +# 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(op, val::T) where {T} + return isprimitivetype(T) ? shfl_words(op, val) : shfl_packed(op, 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] + +# 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(op, val::T) where {T} + shuffle_as_words(T) || return :(shfl_unsupported($T)) + U = word_type(T) + if sizeof(T) <= 4 + return inline_body(:(reinterpret($T, op(reinterpret($U, val) % UInt32) % $U))) + elseif sizeof(T) == 8 && T !== UInt64 + return inline_body(:(reinterpret($T, op(reinterpret(UInt64, val))))) + end + W = sizeof(T) == 8 ? UInt32 : UInt64 + nbits = 8 * sizeof(W) + 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(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 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)) + 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 + +# 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 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) + +""" + 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 `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, so that +the segments partition it. It should be a constant of the generated code. + +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 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 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 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 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 + return ifelse(inside, l0 + offset % Int32, l0) + Int32(1) +end + +@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 + return ifelse(inside, l0 - offset % Int32, l0) + Int32(1) +end + +@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 + +""" + 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). -`shfl_down` exchanges values, not memory: it is not a memory fence. +It doesn't synchronize memory accesses, 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 **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 doesn't synchronize memory accesses, 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. `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 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** + 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 + +""" + 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 + +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 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. + `+` 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 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 + `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 doesn't synchronize memory accesses, 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 + +""" + 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 @@ -306,7 +965,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: @@ -318,6 +978,23 @@ 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 1c3ff0573..9a5cbc41a 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 [Sub-groups](@ref ki_subgroups) +for the layout and participation guarantees. !!! note Backend implementations **must** implement this function if they support sub-groups. @@ -275,15 +279,102 @@ 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 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. 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 +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. -Whether kernels on the active device support [`shfl_down`](@ref) for values of type `T`. +`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 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 + independent. The fallback returns `false`. +""" +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 + +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). 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 +without [`supports_float64`](@ref), as long as they are only stored and loaded, e.g. through +an array of `UInt64`. + +!!! note + 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. +""" +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 + fields = primitive_fields!(Any[], T, :val) + return all(((S, _),) -> sizeof(S) in (1, 2, 4, 8, 16), fields) && supports_shuffle(backend, UInt32) +end """ allocate(::Backend, Type, dims...; unified=false)::AbstractArray diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl index 210b77123..b42625987 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,543 @@ 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 + +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 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) + +# 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 +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)])) +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). +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 + 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 + @test ok + end + return +end - val = a[idx] +# 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 - # 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 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 - offset <<= 1 + @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 + 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) + 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 KI.sub_group_any(val > Int32(1000)) + acc = -acc + end + @inbounds out[i] = acc + KI.shfl(val, 1) + return +end - if idx == 1 - b[idx] = val +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 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} +# 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() - 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 + 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[base + lane] = i + @inbounds scratch[base + lane] = -i + KI.sub_group_barrier() + @inbounds out[i, 1] = lm[other] + @inbounds out[i, 2] = scratch[other] + 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 -# 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} +# 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() - other = mod1(lane + 1, KI.get_sub_group_size()) + 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) - @inbounds lm[lane] = lane - @inbounds scratch[lane] = -lane - KI.sub_group_barrier() - @inbounds out[lane, 1] = lm[other] - @inbounds out[lane, 2] = scratch[other] + 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 + 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 +751,226 @@ 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 + + @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 + + 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)), + # 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) + @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) + @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 + + @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 + 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) + end + end + + if KI.supports_partitionable_subgroups(backend) + @testset "partitionable sub-groups" begin + partitioned_subgroups_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 +1098,13 @@ 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 + @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 @@ -594,6 +1307,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 +1351,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 +1401,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 +1420,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..4978c67c1 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,48 @@ end @test isempty(methods(stub)) end + # 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, Tuple{Bool})) + end + + # 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) + @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(), Int32) + # 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()) + @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. primitives = [ diff --git a/src/KernelAbstractions.jl b/src/KernelAbstractions.jl index c233b8cd2..f7bf3725f 100644 --- a/src/KernelAbstractions.jl +++ b/src/KernelAbstractions.jl @@ -755,6 +755,14 @@ 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, 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 + sub-groups take. Otherwise, results are wrong. """ POCLBackend diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index ca6396700..3b9a0651f 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -233,9 +233,21 @@ 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 +# PoCL limitation: for the same reason, sub-groups can't be partitioned into segments that +# diverge +KI.supports_partitionable_subgroups(::POCLBackend) = false + +# 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} 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 +310,85 @@ 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) + +# 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) + +@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 + +# 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...) 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 1d0b79439..37c3e3b82 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), @@ -115,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/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..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 @@ -306,6 +308,68 @@ 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 + +# 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 + +# 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