diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 9a7c65d96..189920191 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -24,7 +24,9 @@ steps: julia -e 'println("--- :julia: Instantiating project") using Pkg Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()), - PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3 + PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))]) + # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11) + Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3 julia -e 'println("--- :julia: Developing CUDA") using Pkg url="https://github.com/JuliaGPU/CUDA.jl" @@ -64,7 +66,9 @@ steps: julia -e 'println("--- :julia: Instantiating project") using Pkg Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()), - PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3 + PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))]) + # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11) + Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3 julia -e 'println("--- :julia: Developing Metal") using Pkg Pkg.add([(; url="https://github.com/JuliaGPU/Metal.jl", rev="ka-0.10")])' || exit 3 @@ -100,7 +104,9 @@ steps: julia -e 'println("--- :julia: Instantiating project") using Pkg Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()), - PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3 + PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))]) + # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11) + Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3 julia -e 'println("--- :julia: Developing oneAPI") using Pkg Pkg.add(url="https://github.com/JuliaGPU/AcceleratedKernels.jl", rev="main") @@ -139,7 +145,9 @@ steps: using Pkg println("--- :julia: Instantiating project") Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()), - PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3 + PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))]) + # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11) + Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3 julia -e 'println("--- :julia: Developing AMDGPU") using Pkg Pkg.add(url="https://github.com/JuliaGPU/AcceleratedKernels.jl", rev="main") @@ -178,11 +186,14 @@ steps: julia -e 'println("--- :julia: Instantiating project") using Pkg Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()), - PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3 + PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))]) + # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11) + Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3 julia -e 'println("--- :julia: Developing OpenCL") using Pkg Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="ka-0.10") - Pkg.develop(; name="SPIRVIntrinsics")' || exit 3 + # TODO: develop the ka-0.10 copy again once JuliaGPU/OpenCL.jl#526 is merged + Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3 julia -e 'println("+++ :julia: Running tests") using Pkg diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index bdccba98d..47938dfd7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -56,9 +56,12 @@ jobs: # applies to the test run only, not to `Pkg.develop`. Do this before building, which # already resolves the environment (and fails if this repo requires an unregistered # version of KernelInterface). + # + # The same applies to the SPIRVIntrinsics branch in `[sources]`. + # TODO: remove that once JuliaGPU/OpenCL.jl#526 is released - name: Dev KernelInterface shell: bash - run: julia -e 'using Pkg; Pkg.activate("."); Pkg.develop(; path="lib/KernelInterface")' + run: julia -e 'using Pkg; Pkg.activate("."); Pkg.develop(; path="lib/KernelInterface"); Pkg.add(; url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' # Nightly tracks the next Julia release; don't fail CI on it - uses: julia-actions/julia-buildpkg@v1 continue-on-error: ${{ matrix.version == 'nightly' }} diff --git a/Project.toml b/Project.toml index 4400dd267..6aefc30af 100644 --- a/Project.toml +++ b/Project.toml @@ -31,6 +31,8 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" [sources] KernelInterface = {path = "lib/KernelInterface"} +# TODO: remove once JuliaGPU/OpenCL.jl#526 is released +SPIRVIntrinsics = {url = "https://github.com/JuliaGPU/OpenCL.jl", rev = "vc/subgroup-votes", subdir = "lib/intrinsics"} [extensions] LinearAlgebraExt = "LinearAlgebra" diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md index 4eba70529..9950833f2 100644 --- a/docs/src/kernelinterface.md +++ b/docs/src/kernelinterface.md @@ -77,7 +77,7 @@ What a backend implements, at a glance. The docstrings below have the details. | **Capabilities** | | [`supports_float64`](@ref), [`supports_atomics`](@ref), [`supports_unified`](@ref), [`supports_subgroups`](@ref), [`supports_shuffle`](@ref) (all `false`) | | **Compilation** | [`argconvert`](@ref), [`kernel_function`](@ref), [`launch`](@ref) | | | **Device** | [`get_local_id`](@ref), [`get_group_id`](@ref), [`get_local_size`](@ref), [`get_num_groups`](@ref), [`localmemory`](@ref), [`barrier`](@ref) | [`get_global_id`](@ref), [`get_global_size`](@ref) (derived from the primitive queries), [`_print`](@ref KernelInterface._print) (host `print`) | -| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries, [`sub_group_barrier`](@ref); if `supports_shuffle(backend, T)`: [`shfl_down`](@ref) for `T` | | +| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries (with a constant [`get_max_sub_group_size`](@ref)), [`sub_group_barrier`](@ref), [`sub_group_any`](@ref), [`sub_group_all`](@ref), and [`sub_group_ballot`](@ref) for widths of at most 64; if `supports_shuffle(backend, T)`: [`shfl`](@ref), [`shfl_down`](@ref), [`shfl_up`](@ref), [`shfl_xor`](@ref) for the primitive `T` supported natively, including `UInt32` | shuffles and votes with a `width`, [`sub_group_match_any`](@ref), [`sub_group_reduce`](@ref), [`sub_group_scan`](@ref) (built on the shuffles and votes); shuffles of other primitive types (as `UInt32` words) and of structs (field by field) | Everything else, such as [`zeros`](@ref KernelInterface.zeros), [`ones`](@ref KernelInterface.ones), the launch-keyword handling of [`Kernel`](@ref) and [`@launch`](@ref KernelInterface.@launch), @@ -151,17 +151,32 @@ get_global_size 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: +can differ 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: - 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. +- if the work-group is 1-D, or its x extent `get_local_size().x` is a multiple of the + sub-group width `W` ([`get_max_sub_group_size`](@ref)), sub-groups are formed from + consecutive work-items, x fastest: the work-item with the linear 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 `(lin - 1) ÷ W + 1`, lane + `(lin - 1) % W + 1`. 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. + +Other shapes can form sub-groups differently, e.g. per row of the work-group: for portable +code, make the x extent of multi-dimensional work-groups a multiple of the sub-group width. + +A sub-group is partial when it has fewer work-items than the sub-group width: the lanes +`get_sub_group_size()+1:get_max_sub_group_size()` have no work-item. Shuffles from those lanes +give unspecified values, and the votes, [`sub_group_match_any`](@ref), +[`sub_group_reduce`](@ref) and [`sub_group_scan`](@ref) only take the work-items of the +sub-group into account. 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 @@ -191,8 +206,25 @@ localmemory ### Communication +The shuffles, votes and collectives below exchange values between the work-items of a +sub-group, not memory: they don't order or make visible the work-items' accesses to local +or global memory. To communicate through memory within a sub-group, e.g. a work-item reading +what another one wrote to local memory, use [`sub_group_barrier`](@ref) between the write +and the read. + ```@docs +shfl shfl_down +shfl_up +shfl_xor +shfl(::Any, ::Integer, ::Integer) +sub_group_any +sub_group_all +sub_group_ballot +sub_group_match_any +sub_group_any(::Bool, ::Integer) +sub_group_reduce +sub_group_scan ``` ### Printing diff --git a/lib/KernelInterface/src/KernelInterface.jl b/lib/KernelInterface/src/KernelInterface.jl index 0f0fade1a..61db3a339 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_match_any, + :sub_group_reduce, :sub_group_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, diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl index 5b0090914..405d3caa2 100644 --- a/lib/KernelInterface/src/device.jl +++ b/lib/KernelInterface/src/device.jl @@ -129,9 +129,15 @@ end # 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. +# doesn't change during the kernel's execution, and that sub-groups are formed from +# consecutive work-items, x fastest, if the work-group is 1-D or its x extent is a multiple of +# the sub-group width (see the manual). Backends that can't ensure that don't report sub-group +# support. +# +# 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_match_any`, +# `sub_group_reduce` and `sub_group_scan` only take the work-items of the sub-group into +# account. """ get_sub_group_size([::Type{T}=Int])::T @@ -155,12 +161,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. + +It is a compile-time constant of the generated code, so code that depends on it, e.g. a +loop over the lanes or a shuffle butterfly, is specialized for it. It isn't known during +type inference, though: 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, returning the + width the kernel is compiled for as a constant (not by querying the device at run + time): ``` @device_override get_max_sub_group_size(::Type{T})::T where {T} ``` @@ -254,26 +269,452 @@ localmemory(::Type{T}, ::Val) where {T} = ## communication +# Shuffles exchange values between the work-items of a sub-group. Backends implement them for +# the primitive types for which `supports_shuffle` returns `true`. The fallbacks below shuffle +# other primitive types as `UInt32` words, and other `isbits` types field by field. + +""" + 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. + +`shfl` exchanges values, not memory: it is not a memory fence. + +Types for which [`supports_shuffle`](@ref) returns `true` are supported. Besides the types a +backend supports natively, that includes other primitive types of 1, 2 or a multiple of 4 +bytes (e.g. `Bool`, `Char` or `Int64`) if the backend supports `UInt32`, which are shuffled +as `UInt32` words, and `isbits` structs and tuples of supported types, which are shuffled +field by field. + +!!! note + Backend implementations **must** implement this for the primitive types they support + natively, which have to include `UInt32`, and only for those, so that other types reach + the fallbacks: + ``` + @device_override shfl(val::T, lane::Integer) where {T <: Union{...}} + ``` +""" +@inline shfl(val, lane::Integer) = shfl_fallback(x -> shfl(x, lane), val) + """ shfl_down(val::T, offset::Integer)::T Return `val` of the work-item `offset` lanes further in the sub-group, i.e. with -[`get_sub_group_local_id`](@ref) equal to `get_sub_group_local_id() + offset`. When there is -no such work-item, the result is an unspecified value (of type `T`). +[`get_sub_group_local_id`](@ref) equal to `get_sub_group_local_id() + offset`, for an `offset` +of at least 0. When that lane is past the sub-group width, i.e. +`get_sub_group_local_id() + offset > get_max_sub_group_size()`, the result is `val` of the +work-item itself, like CUDA's `shfl_down_sync`. 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`. -`shfl_down` exchanges values, not memory: it is not a memory fence. +`shfl_down` exchanges values, not memory: it is not a memory fence. See [`shfl`](@ref) for +the supported types. + +!!! note + Backend implementations **must** implement this like [`shfl`](@ref): + ``` + @device_override shfl_down(val::T, offset::Integer) where {T <: Union{...}} + ``` +""" +@inline shfl_down(val, offset::Integer) = shfl_fallback(x -> shfl_down(x, offset), val) + +""" + shfl_up(val::T, offset::Integer)::T + +Return `val` of the work-item `offset` lanes earlier in the sub-group, i.e. with +[`get_sub_group_local_id`](@ref) equal to `get_sub_group_local_id() - offset`, for an `offset` +of at least 0. When there is no such lane, i.e. `get_sub_group_local_id() <= offset`, the +result is `val` of the work-item itself, like CUDA's `shfl_up_sync`. + +All work-items of the sub-group have to execute `shfl_up` together (not in a divergent +branch), with the same `offset`. + +`shfl_up` exchanges values, not memory: it is not a memory fence. See [`shfl`](@ref) for +the supported types. + +!!! note + Backend implementations **must** implement this like [`shfl`](@ref): + ``` + @device_override shfl_up(val::T, offset::Integer) where {T <: Union{...}} + ``` +""" +@inline shfl_up(val, offset::Integer) = shfl_fallback(x -> shfl_up(x, offset), val) + +""" + shfl_xor(val::T, mask::Integer)::T + +Return `val` of the work-item whose 0-based lane id is the 0-based lane id of this work-item +xor `mask`, i.e. with [`get_sub_group_local_id`](@ref) equal to +`((get_sub_group_local_id() - 1) ⊻ mask) + 1`, for a `mask` 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`. A butterfly over the masks `width ÷ 2, …, 2, 1` (for the +sub-group width [`get_max_sub_group_size`](@ref)) reduces a full sub-group such that every +work-item gets the result. + +`shfl_xor` exchanges values, not memory: it is not a memory fence. See [`shfl`](@ref) for +the supported types. + +!!! note + Backend implementations **must** implement this like [`shfl`](@ref): + ``` + @device_override shfl_xor(val::T, mask::Integer) where {T <: Union{...}} + ``` +""" +@inline shfl_xor(val, mask::Integer) = shfl_fallback(x -> shfl_xor(x, mask), val) + +# Shuffle a value of a type that the backend doesn't support natively, with `f` shuffling a +# value of a type it does support. The fallbacks are separate functions for primitive and for +# other types, so that the fallback of a struct with a field the backend doesn't support +# natively, e.g. an `Int64` on Metal, doesn't call itself, which inference gives up on (on +# Julia 1.10). +@inline function shfl_fallback(f, val::T) where {T} + return isprimitivetype(T) ? shfl_words(f, val) : shfl_fields(f, val) +end + +shfl_unsupported(T) = throw( + ArgumentError( + "Shuffling values of type $T is not supported by this backend, see `supports_shuffle`" + ) +) + +# Whether a primitive type that a backend doesn't support natively can be shuffled as `UInt32` +# 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` +const word_types = Dict(1 => UInt8, 2 => UInt16, 4 => UInt32, 8 => UInt64, 16 => UInt128) + +# Shuffle a primitive value as `UInt32` words: smaller values are zero-extended, larger ones +# split into words. +@inline @generated function shfl_words(f, val::T) where {T} + shuffle_as_words(T) || return :(shfl_unsupported($T)) + U = word_types[sizeof(T)] + if sizeof(T) <= 4 + return :(reinterpret($T, f(reinterpret($U, val) % UInt32) % $U)) + end + n = sizeof(T) ÷ 4 + words = (:((f((bits >> $(32 * (i - 1))) % UInt32) % $U) << $(32 * (i - 1))) for i in 1:n) + return quote + bits = reinterpret($U, val) + return reinterpret($T, |($(words...))) + end +end + +# The expression that shuffles `ex::S` field by field, calling `f` on the primitive fields +function shfl_fields_expr(S, ex) + isprimitivetype(S) && return :(f($ex)) + fields = (shfl_fields_expr(fieldtype(S, i), :(getfield($ex, $i))) for i in 1:fieldcount(S)) + return Expr(:new, S, fields...) +end + +# Shuffle a value that the backend doesn't support directly field by field. Nested fields are +# unrolled here, rather than shuffled with a recursive call, which inference gives up on +# (on Julia 1.10), so that `f` is only called on the primitive types. +@inline @generated function shfl_fields(f, val::T) where {T} + isbitstype(T) || return :(shfl_unsupported($T)) + return shfl_fields_expr(T, :val) +end + +# The shuffles within segments of `width` lanes, implemented with `shfl` from a lane. + +""" + shfl(val::T, lane::Integer, width::Integer)::T + shfl_down(val::T, offset::Integer, width::Integer)::T + shfl_up(val::T, offset::Integer, width::Integer)::T + shfl_xor(val::T, mask::Integer, width::Integer)::T + +Shuffles within segments of `width` consecutive lanes of the sub-group, as if each segment +were a sub-group of its own: `lane` is the lane within the segment (between 1 and `width`, and +taken modulo `width` otherwise), and `shfl_down`, `shfl_up` and `shfl_xor` read from lanes of +the same segment. Where these would read from outside of the segment, they return `val` of the +work-item itself, like CUDA's shuffles with a `width` (and like `shfl_down` and `shfl_up` +without a `width` past the sub-group width). 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. + +!!! note + Backends **may** implement these, e.g. if they have native shuffles with a width. The + fallbacks use [`shfl`](@ref) from a lane. +""" +@inline function shfl(val, lane::Integer, width::Integer) + l0 = get_sub_group_local_id(Int32) - Int32(1) + w = width % Int32 + base = l0 & ~(w - Int32(1)) + return shfl(val, base + ((lane % Int32 - Int32(1)) & (w - Int32(1))) + Int32(1)) +end + +@inline function shfl_down(val, offset::Integer, width::Integer) + l0 = get_sub_group_local_id(Int32) - Int32(1) + w = width % Int32 + d = offset % Int32 + src = ifelse((l0 & (w - Int32(1))) + d < w, l0 + d, l0) + return shfl(val, src + Int32(1)) +end + +@inline function shfl_up(val, offset::Integer, width::Integer) + l0 = get_sub_group_local_id(Int32) - Int32(1) + w = width % Int32 + d = offset % Int32 + src = ifelse((l0 & (w - Int32(1))) >= d, l0 - d, l0) + return shfl(val, src + Int32(1)) +end + +@inline function shfl_xor(val, mask::Integer, width::Integer) + l0 = get_sub_group_local_id(Int32) - Int32(1) + w = width % Int32 + x = l0 ⊻ (mask % Int32) + src = ifelse((x & ~(w - Int32(1))) == (l0 & ~(w - Int32(1))), x, l0) + return shfl(val, src + 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). + +It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref). + +!!! note + Backend implementations that support sub-groups **must** implement: + ``` + @device_override sub_group_any(pred::Bool)::Bool + ``` +""" +function sub_group_any end + +""" + sub_group_all(pred::Bool)::Bool + +Whether `pred` is `true` for all work-items of the sub-group. All work-items of the +sub-group get the same result. + +All work-items of the sub-group have to execute `sub_group_all` together (not in a divergent +branch). + +It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref). !!! note - Backend implementations **must** implement this for every `T` for which - [`supports_shuffle`](@ref) returns `true`: + Backend implementations that support sub-groups **must** implement: ``` - @device_override shfl_down(val::T, offset::Integer) where T + @device_override sub_group_all(pred::Bool)::Bool ``` """ -function shfl_down end +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 for the work-item with +[`get_sub_group_local_id`](@ref) equal to `i`. All work-items of the sub-group get the same +result. Use e.g. `count_ones` to count the work-items, or `trailing_zeros` to find the first +one. + +All work-items of the sub-group have to execute `sub_group_ballot` together (not in a +divergent branch). Only sub-groups of at most 64 work-items are supported. + +It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref). + +!!! note + Backend implementations that support sub-groups with a width of at most 64 **must** + implement: + ``` + @device_override sub_group_ballot(pred::Bool)::UInt64 + ``` +""" +function sub_group_ballot end + +""" + sub_group_match_any(val)::UInt64 + +A mask of the work-items of the sub-group whose `val` is the same as the one of this +work-item (compared bitwise, with `===`), in the format of [`sub_group_ballot`](@ref). Work-items +with different values get different masks: e.g. `trailing_zeros(mask) + 1` is the lane of +the first work-item with the same value, and `count_ones(mask)` the number of them. + +All work-items of the sub-group have to execute `sub_group_match_any` together (not in a +divergent branch). Values of the types that the shuffles support are supported, see +[`supports_shuffle`](@ref). + +It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref). + +!!! note + Backends **may** implement this, e.g. with CUDA's `match.any.sync`. The fallback finds + the groups of equal values one by one with [`shfl`](@ref) and [`sub_group_ballot`](@ref), + so it takes as many steps as there are distinct values. +""" +@inline function sub_group_match_any(val) + remaining = sub_group_ballot(true) + mask = zero(UInt64) + # uniform: every step removes the group of the lowest remaining lane from `remaining` + while remaining != zero(UInt64) + leader = trailing_zeros(remaining) + 1 + same = val === shfl(val, leader) + group = sub_group_ballot(same) + mask = ifelse(same, group, mask) + remaining &= ~group + end + return mask +end + +# The votes within segments of `width` lanes, implemented with `sub_group_ballot`. + +# The 0-based lane of the first work-item of this work-item's segment of `width` lanes +@inline segment_base(width) = (get_sub_group_local_id(Int32) - Int32(1)) & ~(width % Int32 - Int32(1)) + +# The bits of the segment of `width` lanes starting at `base` of a mask, moved to the bottom +@inline function segment_bits(mask::UInt64, base, width) + w = width % UInt32 + bits = w >= UInt32(64) ? typemax(UInt64) : (UInt64(1) << w) - UInt64(1) + return (mask >> (base % UInt32)) & bits +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 + sub_group_match_any(val, width::Integer)::UInt64 + +Votes within segments of `width` consecutive lanes of the sub-group, as if each segment were a +sub-group of its own (like the shuffles with a `width`, e.g. [`shfl`](@ref +shfl(::Any, ::Integer, ::Integer))): the work-items of a segment get the result for the +work-items of their segment. The masks of `sub_group_ballot` and `sub_group_match_any` have a +bit per lane of the segment: bit `i - 1` for its `i`-th lane, i.e. the lane +`get_sub_group_local_id()` with `(get_sub_group_local_id() - 1) % width == i - 1`. In a +partial sub-group, only the work-items of the segment take part. + +`width` has to be a power of two of at most the sub-group width +[`get_max_sub_group_size`](@ref) (and of at most 64), and the same for all work-items of the +sub-group, which all have to execute the vote together (not in a divergent branch). + +It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref). + +!!! note + Backends **may** implement these. The fallbacks use [`sub_group_ballot`](@ref) and + [`sub_group_match_any`](@ref) of the whole sub-group. +""" +@inline sub_group_ballot(pred::Bool, width::Integer) = + segment_bits(sub_group_ballot(pred), segment_base(width), width) + +@inline sub_group_any(pred::Bool, width::Integer) = sub_group_ballot(pred, width) != zero(UInt64) + +@inline function sub_group_all(pred::Bool, width::Integer) + # compare with the work-items of the segment, which may be partial + base = segment_base(width) + return segment_bits(sub_group_ballot(pred), base, width) == + segment_bits(sub_group_ballot(true), base, width) +end + +@inline sub_group_match_any(val, width::Integer) = + segment_bits(sub_group_match_any(val), segment_base(width), width) + +""" + sub_group_reduce(op, val::T)::T + +Reduce `val` over the work-items of the sub-group with the associative binary operator `op`, +in the order of the lanes. 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). Values of the types that the shuffles support are supported, see +[`supports_shuffle`](@ref). + +It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref). + +!!! note + Backends **may** implement this for operators and types with a native reduction (e.g. + `+` on `Float32`), dispatching on `typeof(op)`. The 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`. + +All work-items of the sub-group have to execute `sub_group_scan` together (not in a divergent +branch). Values of the types that the shuffles support are supported, see +[`supports_shuffle`](@ref). + +It exchanges values, not memory: it is not a memory fence, see [`sub_group_barrier`](@ref). + +!!! note + Backends **may** implement this for operators and types with a native scan (e.g. `+` on + `Float32`), dispatching on `typeof(op)`. 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 ## synchronization diff --git a/lib/KernelInterface/src/host.jl b/lib/KernelInterface/src/host.jl index b769783cd..2952d1d3c 100644 --- a/lib/KernelInterface/src/host.jl +++ b/lib/KernelInterface/src/host.jl @@ -267,7 +267,9 @@ width [`sub_group_size`](@ref). See the manual for what KernelInterface guarante 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). +Sub-group support includes the votes [`sub_group_any`](@ref), [`sub_group_all`](@ref) and, +for sub-groups of at most 64 work-items, [`sub_group_ballot`](@ref). Which types the +shuffles support is queried separately with [`supports_shuffle`](@ref). !!! note Backend implementations **must** implement this function if they support sub-groups. @@ -278,13 +280,23 @@ supports_subgroups(::Backend) = false """ supports_shuffle(::Backend, ::Type{T})::Bool -Whether kernels on the active device support [`shfl_down`](@ref) for values of type `T`. +Whether kernels on the active device support the shuffles [`shfl`](@ref), +[`shfl_down`](@ref), [`shfl_up`](@ref) and [`shfl_xor`](@ref) for values of type `T`. + +Other primitive types of 1, 2, 4, 8 or 16 bytes are supported if `UInt32` is, and `isbits` +structs and tuples if all of their fields are. !!! note - Backend implementations **must** implement this function for the types they support. - The fallback returns `false`. -""" -supports_shuffle(::Backend, ::Type) = false + Backend implementations **must** implement this function for the primitive types they + support natively, which have to include `UInt32`, with a signature that only matches + those, e.g. `supports_shuffle(::NewBackend, ::Type{<:Union{UInt32, Float32}})`. The + fallback handles other types. +""" +function supports_shuffle(backend::Backend, ::Type{T}) where {T} + isprimitivetype(T) && return shuffle_as_words(T) && supports_shuffle(backend, UInt32) + (isbitstype(T) && !isprimitivetype(T)) || return false + return all(i -> supports_shuffle(backend, fieldtype(T, i)), 1:fieldcount(T)) +end """ allocate(::Backend, Type, dims...; unified=false)::AbstractArray diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl index 210b77123..c6ca72316 100644 --- a/lib/KernelInterface/test/interface.jl +++ b/lib/KernelInterface/test/interface.jl @@ -244,6 +244,166 @@ function shfl_down_lanes_kernel(out, ::Type{T}, offset) where {T} return end +# Every lane reads the lane `shift` further, wrapping around: the rotation of a tile. +function shfl_rotate_kernel(out, a::AbstractArray{T}, shift) where {T} + lane = KI.get_sub_group_local_id() + width = KI.get_sub_group_size() + val = @inbounds a[lane] + @inbounds out[lane] = KI.shfl(val, mod1(lane + shift, width)) + return +end + +function shfl_up_lanes_kernel(out, ::Type{T}, offset) where {T} + lane = KI.get_sub_group_local_id() + shuffled = KI.shfl_up(T(lane), offset) + @inbounds out[lane] = shuffled + return +end + +# An all-reduce with a butterfly: every lane gets the sum of the sub-group. `N` is the +# width, and the work-group is one sub-group. +function shfl_xor_sum_kernel(out, a, ::Val{N}) where {N} + lane = KI.get_sub_group_local_id() + val = @inbounds a[lane] + mask = N >> 1 + while mask > 0 + val += KI.shfl_xor(val, mask) + mask >>= 1 + end + @inbounds out[lane] = val + return +end + +struct ShuffleStruct + a::Float32 + b::Int64 + c::NTuple{3, Int32} +end + +function shfl_struct_kernel(out, a) + lane = KI.get_sub_group_local_id() + width = KI.get_sub_group_size() + val = @inbounds a[lane] + @inbounds out[lane] = KI.shfl(val, mod1(lane + 1, width)) + return +end + +function vote_kernel(out, pred) + lane = KI.get_sub_group_local_id() + p = @inbounds pred[lane] + any = KI.sub_group_any(p) + all = KI.sub_group_all(p) + ballot = KI.sub_group_ballot(p) + @inbounds begin + out[lane, 1] = any + out[lane, 2] = all + out[lane, 3] = ballot + out[lane, 4] = KI.get_max_sub_group_size() + end + return +end + +# Each work-item shuffles its value with `op` (`shfl`, `shfl_down`, `shfl_up`, `shfl_xor`) and +# a `width`, recording its lane. +function shfl_width_kernel(out, lanes, a, op, arg, width) + lane = KI.get_sub_group_local_id() + val = @inbounds a[lane] + @inbounds out[lane] = op(val, arg, width) + @inbounds lanes[lane] = lane + return +end + +# `pred` determines which values are equal, `vals` gives the values +function match_any_kernel(out, vals) + lane = KI.get_sub_group_local_id() + @inbounds out[lane] = KI.sub_group_match_any(vals[lane]) + return +end + +# a work-group of `length(a)` work-items, which is a single (possibly partial) sub-group +function reduce_scan_kernel(red, scan, lanes, a, op) + i = KI.get_local_id().x + val = @inbounds a[i] + r = KI.sub_group_reduce(op, val) + s = KI.sub_group_scan(op, val) + @inbounds begin + red[i] = r + scan[i] = s + lanes[i] = KI.get_sub_group_local_id() + end + return +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) + +# the values come from a divergent branch, like for the padding work-items of a `@kernel` +function reduce_divergent_kernel(red, scan, a, m) + i = KI.get_local_id().x + 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 + +# the votes within segments of `width` lanes; a work-group of `length(pred)` work-items, i.e. a +# single (possibly partial) sub-group +function segmented_vote_kernel(out, lanes, pred, vals, width) + i = KI.get_local_id().x + p = @inbounds pred[i] + any = KI.sub_group_any(p, width) + all = KI.sub_group_all(p, width) + ballot = KI.sub_group_ballot(p, width) + match = KI.sub_group_match_any(@inbounds(vals[i]), width) + @inbounds begin + out[i, 1] = any + out[i, 2] = all + out[i, 3] = ballot + out[i, 4] = match + lanes[i] = KI.get_sub_group_local_id() + end + return +end + +# 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 + +# reductions and scans in a 1-D work-group of several sub-groups, the last one partial +function reduce_scan_multi_kernel(red, scan, sgs, a) + i = KI.get_local_id().x + val = @inbounds a[i] + r = KI.sub_group_reduce(+, val) + s = KI.sub_group_scan(+, val) + @inbounds begin + red[i] = r + scan[i] = s + sgs[i] = KI.get_sub_group_id() + end + return +end + +struct FallbackStruct + flag::Bool + c::Char + x::Bits64 + limbs::NTuple{4, Float64} +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} @@ -265,6 +425,271 @@ function captured_array_kernel(backend, AT, out) return kernel, WeakRef(a) 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 sub-group operations that are built on the shuffles and votes. Separate functions, so +# that `interface_testsuite` doesn't get too large to compile. + +# rotating values of type `T` by one lane, `f(i)` giving the value of lane `i` +function shfl_type_testsuite(backend, AT, sg_size, ::Type{T}, f) where {T} + a = T[f(i) for i in 1:sg_size] + out = AT(fill(f(0), sg_size)) + KI.@launch backend workgroupsize = sg_size shfl_rotate_kernel(out, AT(a), 1) + KI.synchronize(backend) + @test Array(out) == circshift(a, -1) + return +end + +# `ref(lane, arg, width)` is the lane that `op(val, arg, width)` reads from +function shfl_width_testsuite(backend, AT, sg_size, op, ref, width, arg) + a = Int32.(1:sg_size) .* Int32(10) + out = AT(zeros(Int32, sg_size)) + lanes = AT(zeros(Int, sg_size)) + KI.@launch backend workgroupsize = sg_size shfl_width_kernel(out, lanes, AT(a), op, arg, width) + KI.synchronize(backend) + out, lanes = Array(out), Array(lanes) + @test all(i -> out[i] == a[ref(lanes[i], arg, width)], 1:sg_size) + return +end + +shfl_width_ref(l, arg, w) = (l - 1) ÷ w * w + mod1(arg, w) +shfl_down_width_ref(l, arg, w) = (l - 1) % w + arg < w ? l + arg : l +shfl_up_width_ref(l, arg, w) = (l - 1) % w >= arg ? l - arg : l +shfl_xor_width_ref(l, arg, w) = (((l - 1) ⊻ arg) ÷ w == (l - 1) ÷ w) ? ((l - 1) ⊻ arg) + 1 : l + +function match_any_testsuite(backend, AT, sg_size, vals) + out = AT(zeros(UInt64, sg_size)) + KI.@launch backend workgroupsize = sg_size match_any_kernel(out, AT(vals)) + KI.synchronize(backend) + out = Array(out) + for i in 1:sg_size + expected = UInt64(0) + for j in 1:sg_size + vals[j] === vals[i] && (expected |= UInt64(1) << (j - 1)) + end + @test out[i] == expected + end + return +end + +# a work-group of `length(a)` work-items, i.e. a single sub-group, possibly partial +function reduce_scan_testsuite(backend, AT, op, a) + n = length(a) + red = AT(similar(a)) + scan = AT(similar(a)) + lanes = AT(zeros(Int, n)) + KI.@launch backend workgroupsize = n reduce_scan_kernel(red, scan, lanes, AT(a), op) + KI.synchronize(backend) + red, scan, lanes = Array(red), Array(scan), Array(lanes) + # in the order of the lanes + order = sortperm(lanes) + # `isequal`, as the results may be NaN + @test all(isequal(foldl(op, a[order])), red) + @test all(i -> isequal(scan[order[i]], foldl(op, a[order[1:i]])), 1:n) + return +end + +function reduce_divergent_testsuite(backend, AT, sg_size, ::Type{T}) where {T} + m = max(sg_size - 5, 1) + a = T.(rand(1:20, sg_size)) + red = AT(zeros(T, sg_size)) + scan = AT(zeros(T, sg_size)) + KI.@launch backend workgroupsize = sg_size reduce_divergent_kernel(red, scan, AT(a), m) + KI.synchronize(backend) + # the sum doesn't depend on the order of the lanes + @test all(==(sum(a[1:m])), Array(red)) + @test maximum(Array(scan)) == sum(a[1:m]) + return +end + +function segmented_vote_testsuite(backend, AT, n, width) + pred = [rand(Bool) for _ in 1:n] + # make some segments all true and some all false + for i in 1:n + seg = (i - 1) ÷ width + seg % 3 == 1 && (pred[i] = true) + seg % 3 == 2 && (pred[i] = false) + end + vals = Int32.(rand(1:3, n)) + out = AT(zeros(UInt64, n, 4)) + lanes = AT(zeros(Int, n)) + KI.@launch backend workgroupsize = n segmented_vote_kernel(out, lanes, AT(pred), AT(vals), width) + KI.synchronize(backend) + out, lanes = Array(out), Array(lanes) + by_lane = zeros(Int, maximum(lanes)) + for i in 1:n + by_lane[lanes[i]] = i + end + for i in 1:n + base = (lanes[i] - 1) ÷ width * width + # the work-items of the segment, by their position in the segment + seg = [(l - base, by_lane[l]) for l in (base + 1):min(base + width, length(by_lane)) if by_lane[l] != 0] + ballot = UInt64(0) + match = UInt64(0) + for (k, j) in seg + pred[j] && (ballot |= UInt64(1) << (k - 1)) + vals[j] == vals[i] && (match |= UInt64(1) << (k - 1)) + end + @test out[i, 1] == any(pred[j] for (_, j) in seg) + @test out[i, 2] == all(pred[j] for (_, j) in seg) + @test out[i, 3] == ballot + @test out[i, 4] == match + end + return +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 + +function reduce_scan_multi_testsuite(backend, AT, sg_size, ::Type{T}) where {T} + n = 2 * sg_size + 5 + a = T.(rand(1:20, n)) + red, scan, sgs = AT(zeros(T, n)), AT(zeros(T, n)), AT(zeros(Int, n)) + kernel = KI.@launch backend launch = false reduce_scan_multi_kernel(red, scan, sgs, AT(a)) + n <= KI.max_work_group_size(kernel) || return + kernel(red, scan, sgs, AT(a); workgroupsize = n) + KI.synchronize(backend) + red, scan, sgs = Array(red), Array(scan), Array(sgs) + for i in 1:n + group = findall(==(sgs[i]), sgs) + @test red[i] == sum(a[group]) + end + # 1-D work-groups form sub-groups from consecutive work-items + @test scan == [sum(a[((i - 1) ÷ sg_size * sg_size + 1):i]) for i in 1:n] + return +end + +function subgroup_communication_testsuite(backend::KI.Backend, AT, sg_size) + @testset "shuffles of other types" begin + # primitive types that backends need not support natively are shuffled as words + if KI.supports_shuffle(backend, 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) + end + @test !KI.supports_shuffle(backend, Ref{Int}) + @testset "Bool" begin + KI.supports_shuffle(backend, Bool) && shfl_type_testsuite(backend, AT, sg_size, Bool, isodd) + end + @testset "Char" begin + KI.supports_shuffle(backend, Char) && + shfl_type_testsuite(backend, AT, sg_size, Char, i -> Char(0x0001F600 + i)) + end + @testset "Bits64" begin + KI.supports_shuffle(backend, Bits64) && + shfl_type_testsuite(backend, AT, sg_size, Bits64, i -> Bits64((UInt64(i) << 40) - i)) + end + @testset "Bits16" begin + KI.supports_shuffle(backend, Bits16) && + shfl_type_testsuite(backend, AT, sg_size, Bits16, i -> Bits16(0xa000 + i)) + end + @testset "struct" begin + KI.supports_shuffle(backend, FallbackStruct) && shfl_type_testsuite( + backend, AT, sg_size, FallbackStruct, + i -> FallbackStruct(isodd(i), Char(64 + i), Bits64(-i), (i, -i, 1 / i, 2.0^i)) + ) + end + end + + KI.supports_shuffle(backend, Int32) || return + + @testset "shuffles with a width" begin + for w in (1, 2, 4, 8, 16, 32, 64) + (w <= sg_size && sg_size % w == 0) || continue + for arg in unique((1, 3, w - 1, w + 2)) + @testset "width $w, $arg" begin + shfl_width_testsuite(backend, AT, sg_size, KI.shfl, shfl_width_ref, w, arg) + shfl_width_testsuite(backend, AT, sg_size, KI.shfl_down, shfl_down_width_ref, w, arg) + shfl_width_testsuite(backend, AT, sg_size, KI.shfl_up, shfl_up_width_ref, w, arg) + shfl_width_testsuite(backend, AT, sg_size, KI.shfl_xor, shfl_xor_width_ref, w, arg) + end + end + end + end + + if sg_size <= 64 + @testset "sub_group_match_any" begin + match_any_testsuite(backend, AT, sg_size, Int32[i % 3 for i in 1:sg_size]) + # bitwise comparison, so NaN matches NaN, and -0.0 doesn't match 0.0 + match_any_testsuite( + backend, AT, sg_size, + Float32[isodd(i) ? NaN32 : (i % 4 == 0 ? -0.0f0 : 0.0f0) for i in 1:sg_size] + ) + end + end + + if sg_size <= 64 + @testset "votes with a width" begin + for n in unique((sg_size, max(sg_size - 3, 1))), w in (1, 2, 4, 8, 16, 32, 64) + (w <= sg_size && sg_size % w == 0) || continue + @testset "$n work-items, width $w" begin + segmented_vote_testsuite(backend, AT, n, w) + end + end + end + end + + @testset "sub_group_reduce and sub_group_scan, several sub-groups" begin + reduce_scan_multi_testsuite(backend, AT, sg_size, Int32) + KI.supports_shuffle(backend, Float32) && reduce_scan_multi_testsuite(backend, AT, sg_size, Float32) + end + + @testset "sub_group_reduce and sub_group_scan of divergent values" begin + reduce_divergent_testsuite(backend, AT, sg_size, Int32) + KI.supports_shuffle(backend, Float32) && reduce_divergent_testsuite(backend, AT, sg_size, Float32) + end + + for n in unique((sg_size, max(sg_size - 3, 1))) + @testset "sub_group_reduce and sub_group_scan, $n work-items" begin + reduce_scan_testsuite(backend, AT, +, Int32.(rand(1:100, n))) + # operators and types that backends may implement natively + reduce_scan_testsuite(backend, AT, min, Int64.(rand(-100:100, n))) + reduce_scan_testsuite(backend, AT, max, UInt32.(rand(1:100, n))) + KI.supports_shuffle(backend, Float32) && + reduce_scan_testsuite(backend, AT, +, Float32.(rand(1:100, n))) + KI.supports_shuffle(backend, Float32) && + reduce_scan_testsuite(backend, AT, max, Float32[i == 2 ? NaN32 : rand(1:100) for i in 1:n]) + reduce_scan_testsuite( + backend, AT, compose_affine, + [(Int32(rand((-1, 1, 2))), Int32(rand(-5:5))) for _ in 1:n] + ) + reduce_scan_testsuite(backend, AT, argmin_op, [(Float32(rand(1:20)), Int32(i)) for i in 1:n]) + 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 @@ -594,6 +1019,10 @@ 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 + sub_group_layout_testsuite(backend, AT, sg_size, fits) + end + @testset "Sub-group return types" begin @test sg_size isa Int && sg_size >= 1 @@ -696,6 +1125,75 @@ function interface_testsuite(backend::KI.Backend, AT) end end + @testset "shuffles" begin + candidates = (Int32, Int64, UInt32, Float32, Float64) + types = filter(T -> KI.supports_shuffle(backend, T), candidates) + @testset "$T" for T in types + a = T.(rand(1:100, sg_size)) + @testset "shfl, shift $shift" for shift in unique((0, 1, sg_size - 1)) + out = AT(zeros(T, sg_size)) + KI.@launch backend workgroupsize = sg_size shfl_rotate_kernel(out, AT(a), shift) + KI.synchronize(backend) + @test Array(out) == circshift(a, -shift) + end + + @testset "shfl_up, offset $offset" for offset in unique((1, 3, sg_size ÷ 2)) + 1 <= offset < sg_size || continue + out = AT(zeros(T, sg_size)) + KI.@launch backend workgroupsize = sg_size shfl_up_lanes_kernel(out, T, offset) + KI.synchronize(backend) + @test Array(out)[(offset + 1):end] == T.(1:(sg_size - offset)) + # lanes without a lane `offset` earlier get their own value + @test Array(out)[1:offset] == T.(1:offset) + end + + if ispow2(sg_size) + out = AT(zeros(T, sg_size)) + KI.@launch backend workgroupsize = sg_size shfl_xor_sum_kernel(out, AT(a), Val(sg_size)) + KI.synchronize(backend) + @test all(==(sum(a)), Array(out)) + end + end + + @testset "structs" begin + T = ShuffleStruct + @test KI.supports_shuffle(backend, T) == + all(S -> KI.supports_shuffle(backend, S), (Float32, Int64, Int32)) + @test !KI.supports_shuffle(backend, Ref{Int}) + if KI.supports_shuffle(backend, T) + a = [T(i, -i, (i, 2i, 3i)) for i in 1:sg_size] + out = AT(fill(T(0, 0, (0, 0, 0)), sg_size)) + KI.@launch backend workgroupsize = sg_size shfl_struct_kernel(out, AT(a)) + KI.synchronize(backend) + @test Array(out) == circshift(a, -1) + end + end + end + + subgroup_communication_testsuite(backend, AT, sg_size) + + @testset "votes" begin + patterns = ( + "none" => falses(sg_size), + "all" => trues(sg_size), + "some" => [i % 3 == 1 for i in 1:sg_size], + "last" => [i == sg_size for i in 1:sg_size], + ) + @testset "$name" for (name, pred) in patterns + out = AT(zeros(UInt64, sg_size, 4)) + KI.@launch backend workgroupsize = sg_size vote_kernel(out, AT(collect(pred))) + KI.synchronize(backend) + out = Array(out) + @test all(==(any(pred)), out[:, 1]) + @test all(==(all(pred)), out[:, 2]) + @test all(==(sg_size), out[:, 4]) + if sg_size <= 64 + mask = reduce(|, (UInt64(1) << (i - 1) for i in 1:sg_size if pred[i]); init = UInt64(0)) + @test all(==(mask), out[:, 3]) + end + end + end + @testset "shfl_down" begin candidates = ( Int8, Int16, Int32, Int64, UInt8, UInt16, UInt32, UInt64, @@ -730,6 +1228,11 @@ function interface_testsuite(backend::KI.Backend, AT) end @test !isempty(in_range) @test out[in_range, 3] == out[in_range, 1] .+ offset + # a 1-D work-group forms full sub-groups (but the last one), and lanes + # past the sub-group width get their own value + past = findall(i -> out[i, 1] + offset > sg_size, 1:N) + @test !isempty(past) + @test out[past, 3] == out[past, 1] end end end @@ -751,6 +1254,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..3d1325504 100644 --- a/lib/KernelInterface/test/runtests.jl +++ b/lib/KernelInterface/test/runtests.jl @@ -23,13 +23,17 @@ end @test !occursin("[sources]", toml) end +struct ShuffleBackend <: KI.Backend end +KI.supports_shuffle(::ShuffleBackend, ::Type{Int32}) = true +struct WordShuffleBackend <: 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 +43,31 @@ end @test isempty(methods(stub)) end + # The votes only have the fallbacks within segments of a `width`, built on the votes + # of the whole sub-group that backends implement. + for vote in [KI.sub_group_any, KI.sub_group_all, KI.sub_group_ballot] + @test !hasmethod(vote, Tuple{Bool}) + @test hasmethod(vote, Tuple{Bool, Int}) + end + + # The shuffles only have the fallback that shuffles structs field by field, which + # doesn't handle the primitive types a backend has to implement. + for shfl in [KI.shfl, KI.shfl_down, KI.shfl_up, KI.shfl_xor] + @test_throws ArgumentError shfl(1.0f0, 1) + @test_throws ArgumentError shfl((1.0f0, 2), 1) + @test_throws ArgumentError shfl(Ref(1), 1) + end + # primitive types are shuffled as `UInt32` words, if the backend supports `UInt32` + @test KI.supports_shuffle(WordShuffleBackend(), Bool) + @test KI.supports_shuffle(WordShuffleBackend(), Float64) + @test KI.supports_shuffle(WordShuffleBackend(), Tuple{Char, Int128}) + @test !KI.supports_shuffle(WordShuffleBackend(), Ref{Int}) + @test !KI.supports_shuffle(ShuffleBackend(), Float32) + @test !KI.supports_shuffle(ShuffleBackend(), Tuple{Float32, Int}) + @test KI.supports_shuffle(ShuffleBackend(), Int32) + @test KI.supports_shuffle(ShuffleBackend(), Tuple{Int32, NTuple{2, Int32}}) + @test !KI.supports_shuffle(ShuffleBackend(), Tuple{Int32, Float32}) + # The 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 b3200ca30..f549aec38 100644 --- a/src/KernelAbstractions.jl +++ b/src/KernelAbstractions.jl @@ -707,6 +707,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. Beyond KernelInterface's requirement + that all work-items of a sub-group execute them together, 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 4d5df6726..dd430ea9f 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -174,9 +174,11 @@ function KI.multiprocessor_count(::POCLBackend)::Int end KI.supports_subgroups(::POCLBackend) = device_limits().sub_group_size > 0 -function KI.supports_shuffle(backend::POCLBackend, ::Type{T}) where {T} +# the types `sub_group_shuffle` supports; other types are shuffled field by field +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 @@ -247,8 +249,105 @@ 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} = + 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() + src = lane + offset + return sub_group_shuffle(val, ifelse(src <= get_max_sub_group_size(), src, lane)) +end + +@device_override function KI.shfl_up(val::T, offset::Integer) where {T <: ShuffleTypes} + lane = get_sub_group_local_id() + return sub_group_shuffle(val, ifelse(lane > offset, lane - offset, lane)) +end + +@device_override KI.shfl_xor(val::T, mask::Integer) where {T <: ShuffleTypes} = + sub_group_shuffle_xor(val, mask) + +@device_override KI.sub_group_any(pred::Bool) = SPIRVIntrinsics.sub_group_any(pred) + +@device_override KI.sub_group_all(pred::Bool) = SPIRVIntrinsics.sub_group_all(pred) + +@device_override function KI.sub_group_ballot(pred::Bool) + mask = SPIRVIntrinsics.sub_group_ballot(pred) + return UInt64(mask[1].value) | (UInt64(mask[2].value) << 32) +end + +# PoCL 7.2 miscompiles sub-group operations after a branch with an early exit (as bounds checks +# emit), because WorkitemLoops gives the peeled first work-item its own copy of their scratch +# memory: the native collectives below, and the unrolled shuffles of KernelInterface's +# fallbacks of `KI.sub_group_reduce` and `KI.sub_group_scan`. `pocl_standalone_jll` includes +# the fix (pocl/pocl#2239) since 7.2.1+1 (JuliaPackaging/Yggdrasil#15001). +const POCL_REPLICA_FIX = pkgversion(cl.pocl_standalone_jll) >= v"7.2.1+1" + +@static if !POCL_REPLICA_FIX + # Without the fix, loop to a bound that isn't a constant, but is the same for all + # sub-groups of the work-group, which PoCL needs (see `POCLBackend`): the smaller of the + # width and the work-group size. Reduce by combining ranges of doubling length with + # `shfl_down`, skipping the lanes without a work-item, and broadcast the result of the + # first lane. + @inline function uniform_sub_group_bound() + sz = KI.get_local_size(Int32) + return min(KI.get_max_sub_group_size(Int32), sz.x * sz.y * sz.z) + end + + @device_override @inline function KI.sub_group_reduce(op, val) + lane = KI.get_sub_group_local_id(Int32) + sgsize = KI.get_sub_group_size(Int32) + offset = Int32(1) + bound = uniform_sub_group_bound() + while offset < bound + other = KI.shfl_down(val, offset) + if lane + offset <= sgsize + val = op(val, other) + end + offset <<= 1 + end + return KI.shfl(val, 1) + end + + # likewise for `KI.sub_group_scan` + @device_override @inline function KI.sub_group_scan(op, val) + lane = KI.get_sub_group_local_id(Int32) + offset = Int32(1) + bound = uniform_sub_group_bound() + while offset < bound + other = KI.shfl_up(val, offset) + if lane > offset + val = op(other, val) + end + offset <<= 1 + end + return val + end +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), with the fix for PoCL's peeling. +const NATIVE_COLLECTIVES = POCL_REPLICA_FIX + +@static if NATIVE_COLLECTIVES + 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) end @device_override @inline function KI._print(args...) diff --git a/src/pocl/compiler/compilation.jl b/src/pocl/compiler/compilation.jl index 62ff16027..90bb400fc 100644 --- a/src/pocl/compiler/compilation.jl +++ b/src/pocl/compiler/compilation.jl @@ -64,6 +64,20 @@ function GPUCompiler.finish_module!( sg_size = job.config.params.sub_group_size if sg_size !== nothing entry.metadata["intel_reqd_sub_group_size"] = MDNode([ConstantInt(Int32(sg_size))]) + + # the sub-group width is fixed, so make `get_max_sub_group_size` a constant, as + # KernelInterface requires (this runs before optimization) + gvs = LLVM.globals(mod) + if haskey(gvs, "__spirv_BuiltInSubgroupMaxSize") + gv = gvs["__spirv_BuiltInSubgroupMaxSize"] + for use in collect(LLVM.uses(gv)) + load = LLVM.user(use) + load isa LLVM.LoadInst || continue + LLVM.replace_uses!(load, ConstantInt(LLVM.value_type(load), sg_size)) + LLVM.erase!(load) + end + isempty(LLVM.uses(gv)) && LLVM.erase!(gv) + end end # if this kernel uses our RNG, we should prime the shared state. diff --git a/test/runtests.jl b/test/runtests.jl index 0b5d06e91..94e7b9905 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -99,6 +99,39 @@ if "cl_khr_fp16" in POCL.device().extensions end end +# KernelInterface requires the sub-group width to be a constant of the generated code +function max_sub_group_size_kernel(out) + out[1] = KernelAbstractions.KernelInterface.get_max_sub_group_size() + return +end +@testset "POCL constant sub-group width" begin + out = zeros(Int, 1) + ir = sprint() do io + @device_code_llvm io = io debuginfo = :none @opencl max_sub_group_size_kernel(out) + end + width = POCL.device_limits().sub_group_size + @test out[1] == width + @test !occursin("BuiltInSubgroupMaxSize", ir) + @test occursin("store i64 $width", ir) +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 +@testset "POCL native sub-group collectives" begin + for (T, native) in ((Int32, true), (Float32, true)) + x, out = ones(T, 32), zeros(T, 32) + ir = sprint() do io + @device_code_llvm io = io debuginfo = :none @opencl local_size = 32 global_size = 32 sub_group_reduce_kernel(out, x) + end + @test all(==(32), out) + @test occursin("sub_group_reduce_add", ir) == (native && POCL.POCLKernels.NATIVE_COLLECTIVES) + 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