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/api.md b/docs/src/api.md index 0d1d77696..3ff3630eb 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -15,6 +15,33 @@ @ndrange ``` +### Reductions and scans + +```@docs +@groupreduce +@groupscan +``` + +### Tiles + +```@docs +@tile +Tile +tile_shfl +tile_shfl_up +tile_shfl_down +tile_shfl_xor +tile_any +tile_all +tile_ballot +tile_reduce +tile_scan +tile_exclusive_scan +tile_barrier +tile_width +supports_tiles +``` + ## Host language !!! note diff --git a/docs/src/index.md b/docs/src/index.md index a7644e22c..586be13f8 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -167,6 +167,22 @@ end `GPUCompiler.method_tables` instead. - `foreach_index(f, A)` runs `f` once per index of the array `A` without writing a kernel out, and `foreach_index(f, backend, indices)` once per index in a range or `CartesianIndices`. +- KernelInterface 0.4 adds sub-groups: ids, shuffles, votes, reductions and scans, also + within segments of a sub-group. They are the low-level API for backends, for kernels + written against KernelInterface and for kernel languages built on it; capability queries + say how a backend forms its sub-groups, whether they may diverge, and which types it can + shuffle. See [Sub-groups](@ref ki_subgroups). +- [`@groupreduce`](@ref) and [`@groupscan`](@ref) reduce and scan a value over the + workgroup, on every backend. `@groupreduce` uses sub-groups internally where the backend + can shuffle the values. See [Reductions and scans](@ref group_collectives). +- Tiles let small groups of consecutive work-items exchange values with shuffles, votes, + reductions and scans (`tile_shfl`, `tile_reduce`, ...), while different tiles take + different branches or loop a different number of times. A kernel gets tiles by being + constructed with a width, `f(backend; tile = N)`, without a declaration in the kernel. + [`supports_tiles(backend, N)`](@ref supports_tiles) reports whether a backend supports a + width; the `CPU` backend has no tiles. See [Tiles](@ref tiles). A `@kernel` uses tiles or + the workgroup collectives; KernelInterface's sub-group functions are for kernels written + directly against KernelInterface. ## Semantic differences diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md index 4eba70529..a5c695542 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 + implementations of kernel languages. Don't call them in a KernelAbstractions `@kernel`: + use its [reductions and scans](@ref group_collectives) or [tiles](@ref tiles) instead, + which follow the participation rules of `@kernel`. Host code may use the queries of + both. ```@docs get_sub_group_size @@ -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/docs/src/kernels.md b/docs/src/kernels.md index 5d6df509f..6b9c1fddc 100644 --- a/docs/src/kernels.md +++ b/docs/src/kernels.md @@ -213,6 +213,127 @@ statements, and [`@uniform`](@ref) evaluates an expression outside the work-item can be reused across `@synchronize` statements. For scratch storage that does not need to survive across `@synchronize`, an `MArray` can be used instead. +## [Reductions and scans](@id group_collectives) + +[`@groupreduce`](@ref) combines a value from every work-item of the workgroup and returns the +result to all of them. [`@groupscan`](@ref) gives each work-item the result up to itself, in +the order of `@index(Local, Linear)`; it is inclusive by default, and exclusive with +`inclusive = false`. Both take an associative operator and its neutral element, and +`@groupreduce` also needs the operator to be commutative. + +```julia +@kernel function group_sums!(out, @Const(x)) + i = @index(Global, Linear) + s = @groupreduce(+, x[i], zero(eltype(out)); groupsize = 256) + if @index(Local, Linear) == 1 + out[@index(Group, Linear)] = s + end +end +``` + +Like [`@synchronize`](@ref), they are collectives: every work-item of the workgroup has to +reach them, not in a branch or loop that only some of them take, or after an early `return`. +Use them as a statement of their own, optionally assigning the result, as in +`res = @groupreduce(+, val, zero(T))`. Work-items that pad a partial workgroup take part as +well: they contribute the neutral element, without evaluating `val`. + +Both go through local memory, whose size has to be known when the kernel is compiled: a +kernel without a static workgroup size passes `groupsize`, a constant upper bound on the +number of work-items per workgroup. + +They work on every backend. Where the backend can shuffle the values, `@groupreduce` first +reduces each sub-group, and then combines those results through local memory. +KernelAbstractions picks this automatically. The order in which values are combined can +differ, so floating-point results may vary between backends and launch configurations. + +## [Tiles](@id tiles) + +A tile is a group of `N` consecutive work-items, by `@index(Local, Linear)`. Constructing a +kernel with a tile width, `f(backend; tile = N)`, divides its workgroups into tiles, and +[`@tile()`](@ref @tile) returns the work-item's tile in the kernel. The work-items of a tile +exchange values with operations such as [`tile_shfl`](@ref) and [`tile_reduce`](@ref), without +going through local memory. All work-items of a tile have to execute each tile operation +together, but different tiles may take different branches, loop a different number of times, +or `return` early. + +```julia +@kernel function tile_maxima!(out, @Const(x)) + t = @tile() + m = tile_reduce(t, max, x[@index(Global, Linear)]) + if t.lane == 1 + out[@index(Tile)] = m + end +end + +kernel = tile_maxima!(backend; tile = 16) +kernel(out, x; ndrange = 16 * length(out)) +``` + +The workgroup size has to be a multiple of `N`, and every tile has to lie entirely inside or +entirely outside of the `ndrange`; see [`@tile`](@ref) for the details. The `ndrange` can have +several dimensions, as long as that holds. Here, each tile holds the 4 × 4 nodes of an +element, and each workgroup holds four elements: + +```julia +@kernel function element_sums!(out, @Const(u)) + t = @tile() + I = @index(Global, NTuple) + out[I...] = tile_reduce(t, +, u[I...]) +end +element_sums!(backend, (4, 4, 4); tile = 16)(out, u; ndrange = size(u)) # size(u) == (4, 4, nelements) +``` + +### Choosing the tile width + +Tiles aren't declared in the kernel: pass the width when constructing it, `f(backend; +tile = 16)` or `f(backend, workgroupsize; tile = 16)`. A kernel that uses `@tile()` or +`@index(Tile)` needs one. [`tile_width(kernel)`](@ref tile_width) returns it on the host, and +`tile_width(t)` in the kernel, where it is a constant. Like a static workgroup size, the width +is compiled into the kernel; construct the kernel again for another width. + +Tiles are built on the backend's sub-groups (warps, wavefronts or SIMD-groups). The width +has to be a power of two of at most the sub-group width `W`, and tiles narrower than `W` share +a sub-group, yet still take their own paths. Not every backend can do that: +[`supports_tiles(backend, N)`](@ref supports_tiles) says whether it supports tiles of `N` +work-items, and constructing a kernel with a width it doesn't support is an error. CUDA, +AMDGPU, Metal and oneAPI, and OpenCL.jl on Intel GPUs, support widths of up to 32. For tiles of +the full sub-group width, pass `tile = KernelInterface.sub_group_size(backend)`. + +KernelAbstractions' CPU backend has no tiles: its sub-group operations involve the whole +workgroup, so tiles can't take their own paths. To support it, write a variant of the kernel +that uses local memory or the workgroup collectives instead: + +```julia +kernel = supports_tiles(backend, 16) ? element_sums!(backend; tile = 16) : + element_sums_localmem!(backend) +``` + +## Choosing how work-items communicate + +| To | Use | Available | +|---|---|---| +| share data in any pattern within a workgroup | [`@localmem`](@ref) and [`@synchronize`](@ref) | on every backend | +| reduce or scan a value over the workgroup | [`@groupreduce`](@ref), [`@groupscan`](@ref) | on every backend | +| exchange values within small groups of work-items whose control flow may differ between groups, e.g. the nodes of an element | [tiles](@ref tiles) | where [`supports_tiles(backend, N)`](@ref supports_tiles) is `true` | +| program the hardware's sub-groups directly | a [KernelInterface](@ref kernelinterface) kernel | as its [capability queries](@ref ki_subgroups) report | + +Use the collectives for reductions and scans over the workgroup, and local memory with +`@synchronize` for other patterns; both work everywhere. Tiles fit when values only need to +move within small groups that take their own paths. They avoid local memory, but that doesn't +guarantee a faster kernel: on some GPUs, the local-memory version is as fast or faster. +Compare both on the devices you care about, and keep the local-memory version for the backends +without tiles. + +### Kernel languages + +KernelInterface is the low-level API that backends implement, and `@kernel` is a language +built on it, which takes care of the work-items that pad a partial workgroup and of who takes +part in collectives. A kernel is either a `@kernel` or a kernel written against +[KernelInterface](@ref kernelinterface), never a mix: don't call KernelInterface's device +functions, such as its sub-group shuffles and votes, in a `@kernel`. In a `@kernel`, use the +collectives or tiles instead. Host code can use KernelInterface's queries, e.g. to pick a +kernel or a tile width. + ## Launching kernels Construct a kernel by calling the kernel function on a backend and optional static sizes, then 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..e89889804 100644 --- a/src/KernelAbstractions.jl +++ b/src/KernelAbstractions.jl @@ -42,6 +42,8 @@ and then invoked on the arguments. - [`@uniform`](@ref) - [`@synchronize`](@ref) - [`@print`](@ref) +- [`@groupreduce`](@ref) +- [`@groupscan`](@ref) # Kernel constructor @@ -405,6 +407,8 @@ A cartesian index is a general N-dimensional index that is derived from the iter - `Global`: Used to access global memory. - `Group`: The index of the `workgroup`. - `Local`: The within `workgroup` index. + - `Tile`: The global linear index of the work-item's tile, in kernels with tiles (see + [`@tile`](@ref)). It takes no index kind. # Index kind @@ -426,8 +430,12 @@ If the index kind is not provided it defaults to `Linear`, this is subject to ch ``` """ macro index(locale, args...) + if locale === :Tile + isempty(args) || error("@index(Tile) takes no further arguments") + return :($__index_Tile($(esc(:__ctx__)), $__tile_width($(esc(:__ctx__))))) + end if !(locale === :Global || locale === :Local || locale === :Group) - error("@index requires as first argument either :Global, :Local or :Group") + error("@index requires as first argument either :Global, :Local, :Group or :Tile") end if length(args) >= 1 @@ -523,9 +531,10 @@ import .NDIteration: get ### """ - Kernel{Backend, WorkgroupSize, NDRange, Func} + Kernel{Backend, WorkgroupSize, NDRange, Func, TileWidth} -Host-side handle for a kernel specialized on a backend, workgroup size, and `ndrange`. +Host-side handle for a kernel specialized on a backend, workgroup size, `ndrange`, and for a +kernel with tiles, the tile width (an `Int`, or `nothing`, see [`tile_width`](@ref)). Kernels are created by calling a [`@kernel`](@ref) function on a backend, for example `my_kernel(CUDABackend(), 256)`. The returned object is callable: @@ -542,13 +551,17 @@ and [`backend`](@ref KernelAbstractions.backend) to inspect a kernel's static co Kernels are launched on any backend that implements [KernelInterface](@ref kernelinterface); see the [notes for backend implementations](@ref implementations_notes). """ -struct Kernel{Backend, WorkgroupSize <: _Size, NDRange <: _Size, Fun} +struct Kernel{Backend, WorkgroupSize <: _Size, NDRange <: _Size, Fun, TileWidth} backend::Backend f::Fun end -function Base.similar(kernel::Kernel{D, WS, ND}, f::F) where {D, WS, ND, F} - return Kernel{D, WS, ND, F}(kernel.backend, f) +# a kernel without tiles +Kernel{B, WS, ND, F}(backend, f) where {B, WS, ND, F} = Kernel{B, WS, ND, F, nothing}(backend, f) + +function Base.similar(kernel::Kernel{D, WS, ND, <:Any, T}, f::F) where {D, WS, ND, F, T} + T === nothing && uses_tiles(f) && tile_error(nothing, tile_width_message) + return Kernel{D, WS, ND, F, T}(kernel.backend, f) end """ @@ -652,8 +665,25 @@ last (possibly partial) workgroup. Primarily used by backend implementations and return iterspace, dynamic end -function construct(backend::B, ::S, ::NDRange, xpu_name::XPUName) where {B <: Backend, S <: _Size, NDRange <: _Size, XPUName} - return Kernel{B, S, NDRange, XPUName}(backend, xpu_name) +# The constructors of a `@kernel` function `f`: `f(backend)`, `f(backend, workgroupsize)` and +# `f(backend, workgroupsize, ndrange)` +construct_kernel(backend, f; tile = nothing) = construct(backend, DynamicSize(), DynamicSize(), f; tile) +construct_kernel(backend, f, size; tile = nothing) = construct(backend, StaticSize(size), DynamicSize(), f; tile) +construct_kernel(backend, f, size, range; tile = nothing) = + construct(backend, StaticSize(size), StaticSize(range), f; tile) +construct_kernel(backend, f, size::_Size, range::_Size; tile = nothing) = construct(backend, size, range, f; tile) + +function construct(backend::B, ::S, ::NDRange, xpu_name::XPUName; tile = nothing) where {B <: Backend, S <: _Size, NDRange <: _Size, XPUName} + N = resolve_tile_width(xpu_name, tile) + if N !== nothing + # what is known before the launch, which checks the rest + check_tile_width(backend, N) + NDRange <: StaticSize && check_tile_items(N, extents(get(NDRange))) + S <: StaticSize && check_tile_workgroup(N, extents(get(S))) + NDRange <: StaticSize && S <: StaticSize && + check_tiles(backend, N, extents(get(NDRange)), extents(get(S))) + end + return Kernel{B, S, NDRange, XPUName, N}(backend, xpu_name) end ### @@ -680,6 +710,8 @@ function __workitems_iterspace end end include("macros.jl") +include("groupreduction.jl") +include("tiles.jl") include("spawn.jl") include("foreach_index.jl") @@ -755,6 +787,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/backend_launch.jl b/src/backend_launch.jl index 2b145ef1d..6764f4b9d 100644 --- a/src/backend_launch.jl +++ b/src/backend_launch.jl @@ -8,15 +8,27 @@ ### """ - mkcontext(kernel::Kernel, ndrange, iterspace, [launch]) + mkcontext(kernel::Kernel, ndrange, iterspace, [launch]; subgroups = nothing) The hidden context argument for launching `kernel` over `ndrange`, partitioned as -`iterspace`, with the launch configuration `launch` (see [`select_launch`](@ref)). +`iterspace`, with the launch configuration `launch` (see [`select_launch`](@ref)) and the +backend's sub-group capabilities `subgroups` (see [`kernel_subgroups`](@ref)). """ mkcontext(kernel::Kernel, _ndrange, iterspace) = - CompilerMetadata{ndrange(kernel), DynamicCheck}(_ndrange, iterspace) -mkcontext(kernel::Kernel, _ndrange, iterspace, launch) = - CompilerMetadata{ndrange(kernel), DynamicCheck}(_ndrange, iterspace; launch) + CompilerMetadata{ndrange(kernel), DynamicCheck}(_ndrange, iterspace; tile = tile_val(kernel)) +mkcontext(kernel::Kernel, _ndrange, iterspace, launch; subgroups = nothing) = + CompilerMetadata{ndrange(kernel), DynamicCheck}(_ndrange, iterspace; launch, subgroups, tile = tile_val(kernel)) + +""" + kernel_subgroups(kernel::Kernel) + +The sub-group capabilities of the backend of `kernel` that its work-group collectives can +use, a [`SubgroupCapabilities`](@ref), or `nothing` if the backend has none, or if the kernel +doesn't use collectives (which doesn't depend on the backend, so that launching other +kernels doesn't query it). +""" +kernel_subgroups(kernel::Kernel) = + uses_collectives(kernel.f) ? subgroup_capabilities(backend(kernel)) : nothing mkcontext(kernel::Kernel, I, _ndrange, iterspace, ::Dynamic) where {Dynamic} = CompilerMetadata{ndrange(kernel), Dynamic}(I, _ndrange, iterspace) @@ -84,9 +96,12 @@ Core.kwcall(kwargs::NamedTuple, obj::Kernel{<:KI.Backend}, args::Vararg{Any, N}) launch_tuple(obj, args; kwargs...) function launch_tuple(obj::Kernel, args::Tuple; ndrange = nothing, workgroupsize = nothing) + N = tile_width(obj) ndrange, workgroupsize, iterspace, dynamic = launch_config(obj, ndrange, workgroupsize) # nothing to launch (or compile) for an empty ndrange any(iszero, size(blocks(iterspace))) && return nothing + # reject an invalid launch of a kernel with tiles before compiling it + N === nothing || check_tiles_launch(obj, N, ndrange, workgroupsize, iterspace) # launch on an N-d grid, computing indices in a narrow type, if possible. this doesn't # depend on the tuned workgroup size, so the context (and thus the kernel) doesn't either. @@ -104,10 +119,18 @@ function launch_tuple(obj::Kernel, args::Tuple; ndrange = nothing, workgroupsize end function launch_kernel(obj::Kernel, launch, ndrange, _workgroupsize, iterspace, args::Tuple) + # a constant `nothing` for kernels without collectives; for kernels with, the call below + # is a function barrier that makes the rest of the launch type stable + subgroups = kernel_subgroups(obj) + launch_kernel(obj, launch, subgroups, ndrange, _workgroupsize, iterspace, args) + return nothing +end + +function launch_kernel(obj::Kernel, launch, subgroups, ndrange, _workgroupsize, iterspace, args::Tuple) b = backend(obj) # this might not be the final context, since we may tune the workgroupsize - ctx = mkcontext(obj, ndrange, iterspace, launch) + ctx = mkcontext(obj, ndrange, iterspace, launch; subgroups) args, tt = prepare_arguments(b, obj.f, ctx, args) kernel = KI.kernel_function(b, obj.f, tt; compiler_options(obj)...) @@ -115,13 +138,18 @@ function launch_kernel(obj::Kernel, launch, ndrange, _workgroupsize, iterspace, if workgroupsize(obj) <: DynamicSize && _workgroupsize === nothing range = something(ndrange, static_ndrange(obj)) threads = KI.launch_configuration(kernel; nitems = saturated_prod(extents(range))).workgroupsize - iterspace, _ = partition(obj, ndrange, launch_workgroupsize(b, launch, threads, range)) - ctx = mkcontext(obj, ndrange, iterspace, launch) + N = tile_width(obj) + items = N === nothing ? launch_workgroupsize(b, launch, threads, range) : + tile_launch_workgroupsize(b, launch, N, threads, range) + iterspace, _ = partition(obj, ndrange, items) + ctx = mkcontext(obj, ndrange, iterspace, launch; subgroups) end # the geometry is valid by construction, except for the kernel's limits groups = size(blocks(iterspace)) items = size(workitems(iterspace)) + N = tile_width(obj) + N === nothing || check_tiles(b, N, extents(something(ndrange, static_ndrange(obj))), items) if launch isa NDLaunch groups, items = KI.pad3(groups), KI.pad3(items) else diff --git a/src/compiler.jl b/src/compiler.jl index c61f7d471..7d168b951 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -1,26 +1,31 @@ """ - CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch} + CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch, Subgroups, TileWidth} The hidden context argument of kernels written with [`@kernel`](@ref). The `launch` field tells the index functions how the backend launched the kernel: `nothing` for a 1-D launch -indexed in `Int`, or a [`LinearLaunch`](@ref) or [`NDLaunch`](@ref). +indexed in `Int`, or a [`LinearLaunch`](@ref) or [`NDLaunch`](@ref). The `subgroups` field +holds the backend's [`SubgroupCapabilities`](@ref) for kernels that use work-group +collectives, or `nothing`. The `tile` field holds the tile width `Val(N)` of a kernel with +tiles (see [`@tile`](@ref)), or `nothing`. """ -struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch} +struct CompilerMetadata{StaticNDRange, CheckBounds, I, NDRange, Iterspace, Launch, Subgroups, TileWidth} groupindex::I ndrange::NDRange iterspace::Iterspace launch::Launch + subgroups::Subgroups + tile::TileWidth # CPU variant function CompilerMetadata{NDRange, CB}(idx, ndrange, iterspace) where {NDRange, CB} ndrange = cartesian(ndrange) - return new{NDRange, CB, typeof(idx), typeof(ndrange), typeof(iterspace), Nothing}(idx, ndrange, iterspace, nothing) + return new{NDRange, CB, typeof(idx), typeof(ndrange), typeof(iterspace), Nothing, Nothing, Nothing}(idx, ndrange, iterspace, nothing, nothing, nothing) end # GPU variante: index is given implicit - function CompilerMetadata{NDRange, CB}(ndrange, iterspace; launch = nothing) where {NDRange, CB} + function CompilerMetadata{NDRange, CB}(ndrange, iterspace; launch = nothing, subgroups = nothing, tile = nothing) where {NDRange, CB} ndrange = cartesian(ndrange) - return new{NDRange, CB, Nothing, typeof(ndrange), typeof(iterspace), typeof(launch)}(nothing, ndrange, iterspace, launch) + return new{NDRange, CB, Nothing, typeof(ndrange), typeof(iterspace), typeof(launch), typeof(subgroups), typeof(tile)}(nothing, ndrange, iterspace, launch, subgroups, tile) end end @@ -34,6 +39,8 @@ cartesian(t::Tuple) = CartesianIndices(t) @inline __iterspace(cm::CompilerMetadata) = cm.iterspace @inline __groupindex(cm::CompilerMetadata) = cm.groupindex @inline __launch(cm::CompilerMetadata) = cm.launch +@inline __subgroups(cm::CompilerMetadata) = cm.subgroups +@inline __tile_width(cm::CompilerMetadata) = cm.tile @inline __groupsize(cm::CompilerMetadata) = size(workitems(__iterspace(cm))) @inline __dynamic_checkbounds(::CompilerMetadata{NDRange, CB}) where {NDRange, CB} = CB <: DynamicCheck @inline __ndrange(::CompilerMetadata{NDRange}) where {NDRange <: StaticSize} = CartesianIndices(get(NDRange)) @@ -48,7 +55,7 @@ cartesian(t::Tuple) = CartesianIndices(t) function Adapt.adapt_structure(to, cm::CompilerMetadata{NDRange, CB, I}) where {NDRange, CB, I} iterspace = Adapt.adapt(to, cm.iterspace) if I === Nothing - return CompilerMetadata{NDRange, CB}(cm.ndrange, iterspace; cm.launch) + return CompilerMetadata{NDRange, CB}(cm.ndrange, iterspace; cm.launch, cm.subgroups, cm.tile) else return CompilerMetadata{NDRange, CB}(cm.groupindex, cm.ndrange, iterspace) end diff --git a/src/groupreduction.jl b/src/groupreduction.jl new file mode 100644 index 000000000..5fa103eba --- /dev/null +++ b/src/groupreduction.jl @@ -0,0 +1,351 @@ +### +# Work-group reductions and scans +# - @groupreduce +# - @groupscan +### + +export @groupreduce, @groupscan + +""" + @groupreduce(op, val, neutral; groupsize) + +Reduce `val` over all work-items of the workgroup with the binary operator `op`, and return +the result on every work-item. `op` has to be associative and commutative: the values are +combined in an unspecified order. `neutral` has to be its neutral element +(`op(neutral, x) == op(x, neutral) == x`). For example, to find the smallest value and its +index, reduce `(value, index)` pairs with an operator that breaks ties by the index, such +as `min` on tuples, with the neutral element `(typemax(T), typemax(Int))`. + +The result has the type of `neutral`, which `val` is converted to: `op(x, y)` has to return +that type `T = typeof(neutral)` for arguments of type `T`. + +`@groupreduce` is a collective, like [`@synchronize`](@ref): it has to be used as a +statement on its own (`res = @groupreduce(op, val, neutral)`), and reached by all +work-items of the workgroup, not in a branch or loop that only some of them execute, or +after an early `return`. Work-items that pad a partial workgroup (outside of the `ndrange`) +take part as well: they contribute `neutral`, without evaluating `val`. So `op` and +`neutral` have to be the same for the whole workgroup, and computable on those work-items. + +The reduction goes through local memory, whose size has to be known at compile time: either +the kernel has a static workgroup size, or the keyword `groupsize` gives an upper bound on the +number of work-items per workgroup (a constant, e.g. a literal or a type parameter of the +kernel). `groupsize` only sizes that storage; it doesn't set the launch's workgroup size. + +On backends with sub-groups, KernelAbstractions reduces the values of every sub-group with +[`KernelInterface.sub_group_reduce`](@ref) first, if the backend can shuffle values of the +type of `neutral`. + +```julia +@kernel function sum_kernel!(out, @Const(x)) + i = @index(Global) + res = @groupreduce(+, x[i], zero(eltype(out)); groupsize = 1024) + if @index(Local, Linear) == 1 + out[@index(Group, Linear)] = res + end +end +``` +""" +macro groupreduce(args...) + op, val, neutral, options = parse_collective("@groupreduce", args, (:groupsize,)) + bound = collective_bound(Base.get(options, :groupsize, nothing)) + return __collective_call(neutral, val) do neutral, val + :($__groupreduce($(esc(:__ctx__)), $(esc(op)), $val, $neutral, $bound)) + end +end + +""" + @groupscan(op, val, neutral; groupsize, inclusive = true) + +Scan `val` over the work-items of the workgroup with the binary operator `op`, in the order +of `@index(Local, Linear)`. With `inclusive = true` (the default), the work-item with local +index `i` gets the reduction of the values of work-items `1:i`; with `inclusive = false`, that +of `1:i-1`, and the first work-item gets `neutral`. `op` has to be associative, but needn't be +commutative, and the values may be combined in any grouping. `neutral` has to be its neutral +element (`op(neutral, x) == op(x, neutral) == x`). The result has the type of `neutral`, +which `val` is converted to, and `op` has to return that type. + +Like [`@groupreduce`](@ref), `@groupscan` is a collective that has to be used as a statement +on its own and reached by all work-items of the workgroup; padding work-items contribute +`neutral`. The scan stores two values per work-item in local memory, sized like for +`@groupreduce`: either the kernel has a static workgroup size, or the keyword `groupsize` +gives an upper bound of the workgroup size. `inclusive` has to be a constant as well. + +This example moves the elements of `x` that satisfy `pred` to the start of each workgroup's +part of `out`, and records how many there are in `counts`; the result isn't compacted across +workgroups. The kernel needs a static workgroup size, e.g. `compact!(backend, 256)`: + +```julia +@kernel function compact!(out, counts, @Const(x), pred) + i = @index(Global, Linear) + keep = pred(x[i]) + offset = @groupscan(+, Int32(keep), Int32(0); inclusive = false) + total = @groupreduce(+, Int32(keep), Int32(0)) + base = (@index(Group, Linear) - 1) * prod(@groupsize()) + if keep + out[base + offset + 1] = x[i] + end + if @index(Local, Linear) == 1 + counts[@index(Group, Linear)] = total + end +end +``` +""" +macro groupscan(args...) + op, val, neutral, options = parse_collective("@groupscan", args, (:groupsize, :inclusive)) + bound = collective_bound(Base.get(options, :groupsize, nothing)) + inclusive = Base.get(options, :inclusive, true) + return __collective_call(neutral, val) do neutral, val + :($__groupscan($(esc(:__ctx__)), $(esc(op)), $val, $neutral, $bound, Val($(esc(inclusive))))) + end +end + +collective_bound(::Nothing) = :($__static_groupsize($(esc(:__ctx__)))) +collective_bound(groupsize) = :(Val($(esc(groupsize)))) + +# Convert `val` to the type of `neutral` *before* the call of the collective. The type of +# `val` may differ between the work-items, e.g. `Union{Float32, Float64}` for an accumulator +# that only some work-items added a `Float64` to, or because padding work-items contribute +# `neutral` instead of `val` (see `mask_collective`). Julia union-splits a call with such an +# argument into one call per type, so the work-items would execute different copies of the +# collective, and its barriers and shuffles. Only the `convert` may be split this way. +function __collective_call(f, neutral, val) + n, v = gensym(:neutral), gensym(:val) + return quote + let $v = $(esc(val)), $n = $(esc(neutral)) + $(f(n, :($convert($typeof($n), $v)))) + end + end +end + +# The `op`, `val` and `neutral` of a collective's arguments, and its `key = value` options +# (also after a `;`), which have to be among `keys`. +function parse_collective(name, args, keys) + positional = Any[] + options = Dict{Symbol, Any}() + function option!(ex) + key = ex.args[1] + (key isa Symbol && key in keys) || + error("$name: unknown option `$key`, expected one of $(join(keys, ", "))") + haskey(options, key) && error("$name: option `$key` given more than once") + options[key] = ex.args[2] + return + end + for arg in args + if isexpr(arg, :parameters) + foreach(option!, arg.args) + elseif isexpr(arg, :(=)) || isexpr(arg, :kw) + option!(arg) + else + push!(positional, arg) + end + end + length(positional) == 3 || + error("$name expects `op`, `val` and `neutral`, and the options $(join(keys, ", ")) as keywords") + return positional..., options +end + +const COLLECTIVES = (Symbol("@groupreduce"), Symbol("@groupscan")) + +# Whether `expr` is a collective that all work-items of a workgroup take part in. +is_collective(expr) = any(name -> is_macrocall(expr, name), COLLECTIVES) + +# A collective that is used as a statement: `@groupreduce(...)` or `lhs = @groupreduce(...)`. +function is_collective_stmt(stmt) + is_collective(stmt) && return true + isexpr(stmt, :(=)) && is_collective(stmt.args[2]) || return false + lhs = stmt.args[1] + lhs isa Symbol && return true + isexpr(lhs, :(::)) && lhs.args[1] isa Symbol && return true + isexpr(lhs, :tuple) && all(x -> x isa Symbol, lhs.args) && return true + return false +end + +collective_error(expr) = error( + "`$(expr.args[1])` must be used as a statement of its own, " * + "e.g. `res = $(expr.args[1])(op, val, neutral)`, found `$(expr)`" +) + +# Rewrite a collective statement, so that padding work-items contribute the neutral element +# instead of evaluating the value. `neutral` is evaluated once, before the collective. +function mask_collective(stmt) + if isexpr(stmt, :(=)) + binding, call = mask_collective(stmt.args[2]).args + return Expr(:block, binding, Expr(:(=), stmt.args[1], call)) + end + # `args[2]` is the macro's `LineNumberNode`, options may come first after a `;` + args = copy(stmt.args) + i = 3 + while i <= length(args) && (isexpr(args[i], :parameters) || args[i] isa LineNumberNode) + i += 1 + end + length(args) >= i + 2 || error("`$(args[1])` expects `op`, `val` and `neutral`") + val, neutral = args[i + 1], args[i + 2] + n = gensym(:neutral) + args[i + 1] = :(__active_lane__ ? $val : $n) + args[i + 2] = n + return Expr(:block, :($n = $neutral), Expr(:macrocall, args...)) +end + +# The workgroup size as a `Val`, if it is static. +@inline __static_groupsize(ctx::CompilerMetadata) = __static_groupsize(__iterspace(ctx)) +@inline __static_groupsize(::NDRange{N, B, W}) where {N, B, W <: StaticSize} = Val(prod(get(W))) +@inline __static_groupsize(::NDRange) = throw( + ArgumentError( + "group reductions and scans require a static workgroup size or an upper bound of it" + ) +) + +# The largest power of two smaller than `n`, or 0. +@inline function __prevpow2(n::T) where {T <: Integer} + n <= one(T) && return zero(T) + return one(T) << (8 * sizeof(T) - 1 - leading_zeros(n - one(T))) +end + +@inline function __groupreduce(ctx, op, val, neutral::T, ::Val{N}) where {T, N} + n = prod(groupsize(ctx)) + n <= N || throw(ArgumentError("@groupreduce: the workgroup size exceeds the given upper bound")) + storage = KI.localmemory(T, Val(N)) + if __shuffleable(__subgroups(ctx), T) + res = __groupreduce_subgroups(op, convert(T, val), neutral, storage) + else + res = __groupreduce_tree(op, convert(T, val), storage, __index_Local_Linear(ctx), n) + end + # all work-items have to read the result before the storage can be reused + KI.barrier() + return res +end + +# Tree reduction in local memory, folding the upper half of the values onto the lower half. +@inline function __groupreduce_tree(op, val, storage, lid, n) + @inbounds storage[lid] = val + KI.barrier() + s = __prevpow2(n) + while s > 0 + if lid <= s && lid + s <= n + @inbounds storage[lid] = op(storage[lid], storage[lid + s]) + end + KI.barrier() + s >>= 1 + end + return @inbounds storage[1] +end + +# Reduce every sub-group with `KI.sub_group_reduce`, which backends can implement natively, +# then reduce the results of the sub-groups. This only uses the sub-group ids, not how the +# work-items form sub-groups, and all sub-groups execute the same sub-group operations, so it +# needs neither linear nor independent sub-groups. There are at most as many sub-groups as +# work-items, so `storage` has room for a value per sub-group. +@inline function __groupreduce_subgroups(op, val, neutral, storage) + sg = KI.get_sub_group_id() + lane = KI.get_sub_group_local_id() + val = KI.sub_group_reduce(op, val) + if lane == 1 + @inbounds storage[sg] = val + end + KI.barrier() + + # every sub-group reduces the results of all sub-groups, which avoids running sub-group + # operations in a branch that only some sub-groups take + width = KI.get_sub_group_size() + acc = neutral + i = lane + while i <= KI.get_num_sub_groups() + acc = op(acc, @inbounds storage[i]) + i += width + end + acc = KI.sub_group_reduce(op, acc) + KI.barrier() + # the sub-groups can combine the values in different orders, so broadcast one result + if sg == 1 && lane == 1 + @inbounds storage[1] = acc + end + KI.barrier() + return @inbounds storage[1] +end + +@inline function __groupscan(ctx, op, val, neutral::T, ::Val{N}, ::Val{inclusive}) where {T, N, inclusive} + n = prod(groupsize(ctx)) + n <= N || throw(ArgumentError("@groupscan: the workgroup size exceeds the given upper bound")) + # two buffers, the scan reads from one and writes to the other + storage = KI.localmemory(T, Val(2 * N)) + lid = __index_Local_Linear(ctx) + + # Hillis-Steele: after the step with distance `d`, every work-item holds the scan of the + # (up to) `2d` values ending at its own + src = 0 + @inbounds storage[lid] = convert(T, val) + KI.barrier() + d = 1 + while d < n + x = @inbounds storage[src + lid] + if lid > d + x = op(@inbounds(storage[src + lid - d]), x) + end + @inbounds storage[(N - src) + lid] = x + KI.barrier() + src = N - src + d <<= 1 + end + + if inclusive + res = @inbounds storage[src + lid] + else + res = lid == 1 ? neutral : @inbounds storage[src + lid - 1] + end + # all work-items have to read their result before the storage can be reused + KI.barrier() + return res +end + + +## sub-group capabilities + +""" + SubgroupCapabilities{F16, F32, F64}() + +The sub-group capabilities of a backend that the work-group collectives use, passed to the +kernel in its [`CompilerMetadata`](@ref): the backend supports sub-groups and shuffles of +`UInt32` values (and thus of other integer-like values, as words), and of `Float16`, +`Float32` and `Float64` values if `F16`, `F32` and `F64`. Kernels on backends without these +capabilities get `nothing` instead. +""" +struct SubgroupCapabilities{F16, F32, F64} end + +function subgroup_capabilities(backend::KI.Backend) + (KI.supports_subgroups(backend) && KI.supports_shuffle(backend, UInt32)) || return nothing + return SubgroupCapabilities{ + KI.supports_shuffle(backend, Float16), + KI.supports_shuffle(backend, Float32), + KI.supports_shuffle(backend, Float64), + }() +end + +# Whether values of type `T` can be shuffled. Floating-point types need the backend's support: +# a backend that shuffles a float type natively may do so without regard for the device, e.g. +# `Float64` on a device without it. Other primitive types are shuffled as words, and structs +# packed into `UInt32` words, whatever the types of their fields. +__shuffleable(::Nothing, ::Type) = false +@generated function __shuffleable(::SubgroupCapabilities{F16, F32, F64}, ::Type{T}) where {F16, F32, F64, T} + function shuffleable(S) + if isprimitivetype(S) + S === Float16 && return F16 + S === Float32 && return F32 + S === Float64 && return F64 + return !(S <: AbstractFloat) && sizeof(S) in (1, 2, 4, 8, 16) + end + isbitstype(S) || return false + fields = KI.primitive_fields!(Any[], S, :val) + return all(((F, _),) -> sizeof(F) in (1, 2, 4, 8, 16), fields) + end + return shuffleable(T) +end + +# Whether a kernel's body contains a work-group collective, which needs the backend's +# sub-group capabilities: the `@kernel` macro adds a method taking a `CollectivesQuery` to the +# kernel's function, so that other kernels don't query them at every launch. (`applicable` is +# constant-folded.) It only does so for the first method of a kernel function: other methods +# with collectives use local memory. A callable that forwards its arguments, e.g. a wrapper of +# a kernel function, is taken to use collectives, which only costs the query of the +# capabilities at its launches. +struct CollectivesQuery end +uses_collectives(f) = applicable(f, CollectivesQuery()) diff --git a/src/launch.jl b/src/launch.jl index 7742e749e..eaf88e9f1 100644 --- a/src/launch.jl +++ b/src/launch.jl @@ -77,10 +77,12 @@ function select_launch(kernel::Kernel, workgroupsize, iterspace) groups = size(blocks(iterspace)) items = size(workitems(iterspace)) tuned = KernelAbstractions.workgroupsize(kernel) <: DynamicSize && workgroupsize === nothing + # kernels with tiles are launched on a 1-D grid unless the iteration space is 1-D, so that + # the linear local index is the one the backend forms sub-groups from + dims = tile_width(kernel) !== nothing && length(items) > 1 ? () : KI.max_work_group_dims(b) return select_launch( map(mul_extent, groups, items), tuned ? nothing : items, - KI.max_work_group_size(b), KI.max_work_group_dims(b), KI.max_num_groups(b), - narrow_index_type(b) + KI.max_work_group_size(b), dims, KI.max_num_groups(b), narrow_index_type(b) ) end diff --git a/src/macros.jl b/src/macros.jl index ccbc3bc7e..6f2ea28d3 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -74,17 +74,26 @@ function __kernel(expr, __source__::LineNumberNode, __module__::Module, force_in end gpu_function = combinedef(def_gpu) - # create constructor functions - _name = Symbol(:_, name) + # Properties of the kernel are methods of its function, rather than of a function of + # KernelAbstractions, so that a kernel can also be defined in a local scope. Like the + # constructors, they are defined with the first method of the kernel, since a method can't + # be overwritten during precompilation. + queries = Any[] + if find_tile(def[:body]) + # see `uses_tiles` + push!(queries, :($gpu_name(::$TilesQuery) = true)) + end + if find_collective(def[:body]) + # see `uses_collectives` + push!(queries, :($gpu_name(::$CollectivesQuery) = true)) + end + + # create constructor functions (a single method with keywords: in a local scope, several + # would capture each other in a `Core.Box`) constructors = quote if $(name isa Symbol ? :(!@isdefined($name)) : true) - function $_name(dev::Dev, sz::S, range::NDRange) where {Dev, S <: $_Size, NDRange <: $_Size} - return $construct(dev, sz, range, $gpu_name) - end - Core.@__doc__ $name(dev) = $_name(dev, $DynamicSize(), $DynamicSize()) - $name(dev, size) = $_name(dev, $StaticSize(size), $DynamicSize()) - $name(dev, size, range) = $_name(dev, $StaticSize(size), $StaticSize(range)) - $name(dev, size::$_Size, range::$_Size) = $_name(dev, size, range) + Core.@__doc__ $name(dev, sizes...; tile = nothing) = $construct_kernel(dev, $gpu_name, sizes...; tile) + $(queries...) end end constructors = relocate_lines(constructors, __source__) @@ -207,7 +216,7 @@ function default_forwarders!(def, gpu_name, constargs) default = args[n + 1].args[2] fdef = Dict{Symbol, Any}( :name => gpu_name, - :args => Any[:__ctx__, kept...], + :args => Any[:(__ctx__::Union{$CompilerMetadata, $ConstArgs}), kept...], :kwargs => Any[], :whereparams => filter(in(used), whereparams), :body => quote @@ -226,7 +235,8 @@ end # The easy case, transform the function for GPU execution # - answer which arguments are marked `@Const` (see `ConstArgs`) function transform_gpu!(def, constargs, force_inbounds, unsafe_indices) - pushfirst!(def[:args], :__ctx__) + # typed, so that the queries of `@kernel` (e.g. `uses_collectives`) never match the kernel + def[:args] = Any[:(__ctx__::Union{$CompilerMetadata, $ConstArgs}), def[:args]...] # `Any[]`, since `split` hands back `LineNumberNode`s alongside `Expr`s new_stmts = Any[] push!(new_stmts, :(__ctx__ isa $ConstArgs && return $(constargs_mask(constargs)))) @@ -294,10 +304,32 @@ function is_scope_construct(expr::Expr) # expr.head === :let end +# Whether `stmt` uses the tile of the work-item, `@tile()` or `@index(Tile)` +function find_tile(stmt) + result = Ref(false) + postwalk(stmt) do expr + result[] |= is_macrocall(expr, Symbol("@tile")) || + (is_macrocall(expr, Symbol("@index")) && any(==(:Tile), expr.args)) + expr + end + return result[] +end + +function find_collective(stmt) + result = Ref(false) + postwalk(stmt) do expr + result[] |= is_collective(expr) + expr + end + return result[] +end + +# Whether `stmt` contains a `@synchronize`, or a collective like `@groupreduce` that all +# work-items of the workgroup have to reach as well. function find_sync(stmt) result = Ref(false) postwalk(stmt) do expr - result[] |= is_sync(expr) + result[] |= is_sync(expr) || is_collective(expr) expr end return result[] @@ -329,6 +361,17 @@ function split(stmts) continue end + if is_collective_stmt(stmt) + # executed by all work-items, the padding ones contribute the neutral element + loop = WorkgroupLoop(current, allocations, false, nothing) + push!(new_stmts, emit(loop)) + allocations = Any[] + current = Any[] + take_line!(new_stmts) + push!(new_stmts, mask_collective(stmt)) + continue + end + has_sync = find_sync(stmt) if has_sync loop = WorkgroupLoop(current, allocations, is_sync(stmt), line) @@ -348,6 +391,7 @@ function split(stmts) recurse(x) = x function recurse(expr::Expr) expr = unblock_lines(expr) + is_collective(expr) && collective_error(expr) if expr.head in (:if, :elseif) && find_sync(expr) return split_branches(expr, recurse) elseif is_scope_construct(expr) && any(find_sync, expr.args) diff --git a/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/src/tiles.jl b/src/tiles.jl new file mode 100644 index 000000000..ee2a22765 --- /dev/null +++ b/src/tiles.jl @@ -0,0 +1,419 @@ +### +# Tiles +# - @tile +# - tile_shfl, tile_shfl_up, tile_shfl_down, tile_shfl_xor +# - tile_any, tile_all, tile_ballot, tile_reduce, tile_scan, tile_exclusive_scan, tile_barrier +### + +export @tile, Tile, tile_width, supports_tiles +export tile_shfl, tile_shfl_up, tile_shfl_down, tile_shfl_xor +export tile_any, tile_all, tile_ballot, tile_reduce, tile_scan, tile_exclusive_scan, tile_barrier + +""" + Tile{N} + +The tile of `N` work-items a work-item belongs to, in a kernel with tiles, as returned by +[`@tile`](@ref): `t.index` is the index of the tile within the workgroup, and `t.lane` the +index of the work-item within the tile, both 1-based. [`tile_width(t)`](@ref tile_width) is +`N`. +""" +struct Tile{N} + index::Int + lane::Int +end + +""" + @tile()::Tile{N} + +Return the work-item's tile, in a kernel constructed with a tile width, `f(backend; tile = +N)`; there is no declaration in the kernel. `t.index` is the tile's index within the +workgroup, and `t.lane` the work-item's lane within the tile, both 1-based. + +```julia +@kernel function tile_sums!(out, @Const(x)) + t = @tile() + v = tile_reduce(t, +, x[@index(Global, Linear)]) + if t.lane == 1 + out[@index(Tile)] = v + end +end +kernel = tile_sums!(backend; tile = 32) +kernel(out, x; ndrange = 32 * ntiles) +``` + +The tiles divide the workgroup by its linear local index: the work-item with +`@index(Local, Linear) == lid` is in tile `fld(lid - 1, N) + 1`, at lane `mod(lid - 1, N) + 1`. +[`@index(Tile)`](@ref @index) is the global linear index of the tile, +`(@index(Group, Linear) - 1) * (prod(@groupsize()) ÷ N) + t.index`: the tiles are numbered +workgroup by workgroup. + +The work-items of a tile communicate with the tile operations, which take the tile as their +first argument, so that helper functions can use them: [`tile_shfl`](@ref), +[`tile_shfl_up`](@ref), [`tile_shfl_down`](@ref), [`tile_shfl_xor`](@ref), +[`tile_any`](@ref), [`tile_all`](@ref), [`tile_ballot`](@ref), [`tile_reduce`](@ref), +[`tile_scan`](@ref), [`tile_exclusive_scan`](@ref) and [`tile_barrier`](@ref). All +work-items of a tile have to execute them together, in control flow that is uniform over the +tile, but can differ between tiles: e.g. a tile can `return` early, or loop a different +number of times. The work-items of a tile also have to pass the same `offset` or `mask` to a +shuffle, and the same `op` (and `init`) to a reduction or scan. The shuffles, votes, +reductions and scans don't synchronize memory accesses: use [`tile_barrier`](@ref) when the +work-items of a tile communicate through local or global memory. [`@synchronize`](@ref), +[`@groupreduce`](@ref) and [`@groupscan`](@ref) still need the whole workgroup. + +# Choosing the tile width + +Pass `tile = N` when constructing the kernel, alongside a static workgroup size or `ndrange` +if any: `f(backend; tile = N)`, `f(backend, workgroupsize; tile = N)` or +`f(backend, workgroupsize, ndrange; tile = N)`. A kernel that uses `@tile()` or `@index(Tile)` +needs it. The width is compiled into the kernel; [`tile_width(kernel)`](@ref tile_width) +returns it on the host, and `tile_width(t)` in the kernel, where it is a constant. + +`N` has to be a power of two of at most the backend's sub-group width +([`KernelInterface.sub_group_size`](@ref)). [`supports_tiles(backend, N)`](@ref supports_tiles) +reports whether the backend supports tiles of that width; constructing the kernel fails if it +doesn't. The CPU backend has no tiles. See [Tiles](@ref tiles) in the manual for choosing the +width. + +# Launching + +Kernels with tiles have restrictions on their launch: + +- the workgroup has a multiple of `N` work-items; +- every tile is either entirely within the `ndrange`, or entirely outside of it, so that the + work-items that pad a partial workgroup form whole tiles, which don't execute the kernel. + For a 1-D `ndrange`, that means a multiple of `N`. For an N-D `ndrange`, e.g. one dimension + per index of an element, it holds if the workgroup covers the leading dimensions of the + `ndrange` and the number of work-items is a multiple of `N`: e.g. an `ndrange` of + `(4, 4, nelements)` in workgroups of `(4, 4, k)`, with tiles of 16 work-items or, for an + even `k` and `nelements`, 32. + +Without an explicit workgroup size, KernelAbstractions picks one that satisfies this; an +explicit one that doesn't is an error. Kernels with tiles and an N-D `ndrange` are launched on +a 1-D grid, so that the tiles follow the backend's sub-groups. + +To give each element fewer nodes than a tile has lanes, e.g. 9 nodes of a 3×3 element in a +tile of 16, map the lanes to nodes yourself, and keep the lanes without a node in the tile +operations: clamp their loads, skip their stores, but let them shuffle, vote and reduce (with +a neutral value) like the others. +""" +macro tile() + return :($__tile($(esc(:__ctx__)), $__tile_width($(esc(:__ctx__))))) +end + +""" + tile_width(kernel::Kernel)::Union{Int, Nothing} + tile_width(t::Tile{N}) = N + +The tile width of a kernel with tiles, or `nothing` for a kernel without. In a kernel, +`tile_width(t)` of a tile `t` from [`@tile`](@ref) is a constant. +""" +tile_width(::Kernel{<:Any, <:Any, <:Any, <:Any, T}) where {T} = T +tile_width(::Tile{N}) where {N} = N + +# the tile width for the kernel's metadata, see `CompilerMetadata` +tile_val(kernel::Kernel) = (N = tile_width(kernel); N === nothing ? nothing : Val(N)) + +@inline function __tile(ctx, ::Val{N}) where {N} + lid0 = __index_Local_Linear(ctx) - 1 + return Tile{N}(lid0 ÷ N + 1, lid0 % N + 1) +end + +@inline function __index_Tile(ctx, ::Val{N}) where {N} + lid0 = __index_Local_Linear(ctx) - 1 + return (__index_Group_Linear(ctx) - 1) * (prod(groupsize(ctx)) ÷ N) + lid0 ÷ N + 1 +end + +# a kernel without a tile width, whose use of tiles the `@kernel` macro didn't see (e.g. in the +# expansion of another macro), so that its construction didn't reject it +__tile(ctx, ::Nothing) = throw(ArgumentError(tile_width_message)) +__index_Tile(ctx, ::Nothing) = throw(ArgumentError(tile_width_message)) + +# A tile is a segment of `N` lanes of a sub-group, so the tile operations are KernelInterface's +# sub-group operations with a `width` of `N`: where several tiles share a sub-group, which +# `supports_tiles` only allows with `KI.supports_partitionable_subgroups`, those only need the +# tile to execute them together. + +""" + tile_shfl(t::Tile{N}, val, lane) + +`val` of the work-item at lane `mod(lane - 1, N) + 1` of the tile `t`. The work-items may read +from different lanes. See [`KernelInterface.shfl`](@ref) for the supported types. +""" +@inline tile_shfl(::Tile{N}, val, lane::Integer) where {N} = KI.shfl(val, lane, N) + +""" + tile_shfl_up(t::Tile{N}, val, offset) + +`val` of the work-item `offset` lanes earlier in the tile `t`, or the work-item's own `val` +if that lane is outside of the tile. `offset` has to be at least 0, and the same for the +work-items of the tile. See [`KernelInterface.shfl`](@ref) for the supported types. +""" +@inline tile_shfl_up(::Tile{N}, val, offset::Integer) where {N} = KI.shfl_up(val, offset, N) + +""" + tile_shfl_down(t::Tile{N}, val, offset) + +`val` of the work-item `offset` lanes later in the tile `t`, or the work-item's own `val` if +that lane is outside of the tile. `offset` has to be at least 0, and the same for the +work-items of the tile. See [`KernelInterface.shfl`](@ref) for the supported types. +""" +@inline tile_shfl_down(::Tile{N}, val, offset::Integer) where {N} = KI.shfl_down(val, offset, N) + +""" + tile_shfl_xor(t::Tile{N}, val, mask) + +`val` of the work-item at lane `((t.lane - 1) ⊻ mask) + 1` of the tile `t`. `mask` has to be +between 0 and `N - 1`, and the same for the work-items of the tile. See +[`KernelInterface.shfl`](@ref) for the supported types. +""" +@inline tile_shfl_xor(::Tile{N}, val, mask::Integer) where {N} = KI.shfl_xor(val, mask, N) + +""" + tile_any(t::Tile, pred::Bool)::Bool + +Whether `pred` is `true` for any work-item of the tile `t`. +""" +@inline tile_any(::Tile{N}, pred::Bool) where {N} = KI.sub_group_any(pred, N) + +""" + tile_all(t::Tile, pred::Bool)::Bool + +Whether `pred` is `true` for all work-items of the tile `t`. +""" +@inline tile_all(::Tile{N}, pred::Bool) where {N} = KI.sub_group_all(pred, N) + +""" + tile_ballot(t::Tile, pred::Bool)::UInt64 + +A mask of the work-items of the tile `t` for which `pred` is `true`: bit `i - 1` for the +work-item at lane `i`. Needs a sub-group width of at most 64 +([`KernelInterface.sub_group_size`](@ref)), also for narrower tiles. +""" +@inline tile_ballot(::Tile{N}, pred::Bool) where {N} = KI.sub_group_ballot(pred, N) + +""" + tile_reduce(t::Tile, op, val) + +Reduce `val` over the work-items of the tile `t` with the associative operator `op`, in the +order of the lanes, and return the result on every work-item of the tile. `op` has to be the +same for the work-items of the tile, and needn't be commutative. See [`KernelInterface.sub_group_reduce`](@ref) for the +supported types, and where the order of floating-point operations may differ. +""" +@inline tile_reduce(::Tile{N}, op, val) where {N} = KI.sub_group_reduce(op, val, N) + +""" + tile_scan(t::Tile, op, val) + +The inclusive scan of `val` over the work-items of the tile `t` with the associative +operator `op`, in the order of the lanes. `op` has to be the same for the work-items of the +tile, and needn't be commutative. See [`KernelInterface.sub_group_scan`](@ref) for the +supported types, and where the order of floating-point operations may differ. +""" +@inline tile_scan(::Tile{N}, op, val) where {N} = KI.sub_group_scan(op, val, N) + +""" + tile_exclusive_scan(t::Tile, op, val, init) + +The exclusive scan of `val` over the work-items of the tile `t` with the associative +operator `op`, in the order of the lanes, starting from `init`: lane 1 gets `init`, and lane +`i > 1` gets `op(init, r)` for the reduction `r` of the values of lanes `1:i-1`. `op` and +`init` have to be the same for the work-items of the tile, and `init` of the type of `val`; +it needn't be an identity of `op`. See [`KernelInterface.sub_group_exclusive_scan`](@ref) +for the supported types, and where the order of floating-point operations may differ. +""" +@inline tile_exclusive_scan(::Tile{N}, op, val, init) where {N} = KI.sub_group_exclusive_scan(op, val, init, N) + +""" + tile_barrier(t::Tile) + +Wait until all work-items of the tile `t` have reached the barrier, and make their writes to +global and local memory before it visible to the tile. It doesn't order memory between +tiles. +""" +@inline tile_barrier(::Tile{N}) where {N} = KI.sub_group_barrier(N) + + +## host side + +""" + supports_tiles(backend, N)::Bool + +Whether kernels with tiles of `N` work-items (see [`@tile`](@ref)) can run on `backend`. +Constructing such a kernel on a backend that doesn't support them is an error. + +That needs a power-of-two `N` of at most the sub-group width `W` +([`KernelInterface.sub_group_size`](@ref)), sub-groups that are formed from consecutive +work-items and may diverge from each other +([`KernelInterface.supports_linear_subgroups`](@ref), +[`KernelInterface.supports_independent_subgroups`](@ref)), and for `N < W`, sub-groups that can +be partitioned into segments that diverge as well +([`KernelInterface.supports_partitionable_subgroups`](@ref)). The CPU backend returns `false`. A launch also has to meet the workgroup and `ndrange` restrictions of +[`@tile`](@ref). +""" +supports_tiles(backend, N::Integer) = tile_support_error(backend, N) === nothing + +# Why `backend` doesn't support tiles of `N` work-items, or `nothing` if it does +function tile_support_error(backend, N::Integer) + (N > 0 && ispow2(N)) || return "the width has to be a power of two" + backend isa KI.Backend && KI.supports_linear_subgroups(backend) || + return "the backend doesn't support tiles, as its sub-groups aren't formed from consecutive work-items (see `KernelInterface.supports_linear_subgroups`)" + KI.supports_independent_subgroups(backend) || + return "the backend doesn't support tiles, as its sub-groups can't diverge from each other (see `KernelInterface.supports_independent_subgroups`)" + W = KI.sub_group_size(backend) + N <= W || return "the width is larger than the sub-group width $W of the backend" + N == W || KI.supports_partitionable_subgroups(backend) || + return "the backend only supports tiles of its sub-group width $W, as its sub-groups can't be partitioned into parts that diverge (see `KernelInterface.supports_partitionable_subgroups`)" + return nothing +end + +# Whether the function of a kernel uses its tiles (`@tile()` or `@index(Tile)`), so that it +# needs a tile width. The `@kernel` macro adds a method taking a `TilesQuery` to the function +# of such a kernel, with its first method. (`applicable` is constant-folded, so that other +# kernels don't look up the method.) +struct TilesQuery end +uses_tiles(f) = applicable(f, TilesQuery()) && query_tiles(f) + +# a callable that forwards its arguments, e.g. a wrapper of a kernel function, doesn't use +# tiles itself +@noinline function query_tiles(f) + sig = Base.unwrap_unionall(which(f, Tuple{TilesQuery}).sig) + return sig.parameters[end] === TilesQuery && f(TilesQuery()) +end + +const tile_width_message = "the kernel uses `@tile()` or `@index(Tile)`, so it needs a tile width: construct it with `f(backend; tile = N)`" + +# The tile width of a kernel of `f`, for the `tile` keyword of its constructor +function resolve_tile_width(f, tile) + if tile === nothing + uses_tiles(f) && tile_error(nothing, tile_width_message) + return nothing + end + tile isa Integer || throw(ArgumentError("`tile` has to be an integer, got $tile")) + return Int(tile) +end + +tile_error(N, msg) = throw(ArgumentError(N === nothing ? "Tiles: $msg" : "Tiles of $N work-items: $msg")) + +# The workgroup size for a launch of a kernel with tiles of `N` work-items over `range` with +# `launch` (see `select_launch`), for about `threads` work-items: the largest multiple of `N` +# work-items, up to `threads` (but at least `N`, leaving it to the launch to check the kernel's +# limit), that divides `range` into whole tiles. +function tile_launch_workgroupsize(backend, launch, N, threads, range) + dims = extents(range) + for t in (max(threads, N) ÷ N * N):-N:N + items = length(dims) == 1 ? launch_workgroupsize(backend, launch, t, range) : + tile_shape(N, t, dims, index_type(launch)) + items !== nothing && prod(items) % N == 0 && whole_tiles(N, dims, items) && return items + end + return tile_error(N, "no workgroup size divides the ndrange $dims into whole tiles, pass one") +end + +# The shape of a workgroup of `t` work-items, no larger than `dims` along any dimension, that +# divides `dims` into whole tiles of `N` work-items with the least padding, or `nothing`. The +# padded iteration space has to fit `T`, the index type `select_launch` chose for a workgroup +# from `launch_workgroupsize`, which pads differently. +function tile_shape(N, t, dims::Dims, ::Type{T}) where {T} + best, best_padded = nothing, typemax(Int) + for items in workgroup_shapes(t, dims) + padded = map((d, i) -> Int128(cld(d, i)) * i, dims, items) + all(<=(typemax(Int)), padded) && fits(T, Int.(padded)) || continue + n = prod(Int.(padded)) + if n < best_padded && whole_tiles(N, dims, items) + best, best_padded = items, n + end + end + return best +end + +# The shapes of `t` work-items that are no larger than `dims` along any dimension, larger +# extents of the first dimensions first +function workgroup_shapes(t, dims::Dims) + length(dims) == 1 && return t <= dims[1] ? [(t,)] : Dims{1}[] + shapes = Dims{length(dims)}[] + for a in min(t, dims[1]):-1:1 + t % a == 0 || continue + for rest in workgroup_shapes(t ÷ a, Base.tail(dims)) + push!(shapes, (a, rest...)) + end + end + return shapes +end + +# Check the launch of `kernel` with tiles of `N` work-items with the `ndrange`, `workgroupsize` +# and `iterspace` from `launch_config`, before compiling it: the partitioned workgroup size, or +# that there is one that will be tuned. +function check_tiles_launch(kernel::Kernel, N, ndrange, workgroupsize, iterspace) + b = backend(kernel) + range = something(ndrange, static_ndrange(kernel)) + if workgroupsize === nothing && KernelAbstractions.workgroupsize(kernel) <: DynamicSize + check_tile_width(b, N) + check_tile_items(N, extents(range)) + launch = select_launch(kernel, workgroupsize, iterspace) + tile_launch_workgroupsize(b, launch, N, KI.max_work_group_size(b), range) + return + end + return check_tiles(b, N, extents(range), size(workitems(iterspace))) +end + +function check_tile_width(backend, N) + err = tile_support_error(backend, N) + err === nothing || tile_error(N, err) + return +end + +check_tile_items(N, ndrange::Dims) = saturated_prod(ndrange) % N == 0 || + tile_error(N, "the ndrange $(length(ndrange) == 1 ? only(ndrange) : ndrange) has $(saturated_prod(ndrange)) work-items, which isn't a multiple of $N") + +# Check the launch of a kernel with tiles of `N` work-items over `ndrange`, in workgroups of +# `workgroupsize` (logical sizes, before a launch on a 1-D grid flattens them) +function check_tiles(backend, N, ndrange::Dims, workgroupsize::Dims) + check_tile_width(backend, N) + check_tile_items(N, ndrange) + check_tile_workgroup(N, workgroupsize) + whole_tiles(N, ndrange, workgroupsize) || + tile_error(N, "workgroups of $workgroupsize divide the ndrange $ndrange into partial tiles") + return +end + +function check_tile_workgroup(N, workgroupsize::Dims) + items = prod(workgroupsize) + items % N == 0 || + tile_error(N, "the workgroup size $(length(workgroupsize) == 1 ? items : workgroupsize) isn't a multiple of $N") + return +end + +# Whether workgroups of size `workgroupsize` (a multiple of `N` work-items) divide `ndrange` +# into tiles of `N` consecutive work-items that are each entirely inside or outside of it. The +# part of a workgroup inside the `ndrange` is a box: the whole workgroup, or in the last +# workgroup along a dimension, what's left of the `ndrange`. So it suffices to check one +# workgroup per combination of the dimensions it's the last one along. +function whole_tiles(N, ndrange::Dims, workgroupsize::Dims) + n = max(length(ndrange), length(workgroupsize)) + range = ntuple(d -> Base.get(ndrange, d, 1), n) + items = ntuple(d -> Base.get(workgroupsize, d, 1), n) + # an empty `ndrange` has no workgroups + any(iszero, range) && return true + insides = map(range, items) do r, i + last = r - (cld(r, i) - 1) * i + return r > i && last != i ? (i, last) : (last,) + end + # a 1-D box, a prefix of the workgroup, is whole tiles if its length is a multiple of `N` + n == 1 && return all(inside -> inside % N == 0, only(insides)) + return all(inside -> whole_tiles_box(N, items, inside), Iterators.product(insides...)) +end + +# Whether every tile of `N` consecutive work-items of a workgroup of size `items` is entirely +# inside or outside of the box of size `inside` at its start +function whole_tiles_box(N, items, inside) + inside == items && return true + tile_inside = false + for (l, I) in enumerate(CartesianIndices(items)) + isinside = all(Tuple(I) .<= inside) + if (l - 1) % N == 0 + tile_inside = isinside + elseif isinside != tile_inside + return false + end + end + return true +end diff --git a/test/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/groupreduce.jl b/test/groupreduce.jl new file mode 100644 index 000000000..e18e2ab3f --- /dev/null +++ b/test/groupreduce.jl @@ -0,0 +1,324 @@ +# one result per workgroup, written by every work-item, to check that all of them get it +@kernel function groupreduce_static!(out, @Const(x), op, neutral) + i = @index(Global, Linear) + res = @groupreduce(op, x[i], neutral) + out[i] = res +end + +@kernel function groupreduce_bound!(out, @Const(x), op, neutral, ::Val{N}) where {N} + i = @index(Global, Linear) + res = @groupreduce(op, x[i], neutral; groupsize = N) + out[i] = res +end + +# the same call site, and thus local memory, reused in a loop and in a branch +@kernel function groupreduce_loop!(out, @Const(x)) + i = @index(Global, Linear) + acc = zero(eltype(out)) + for k in 1:3 + res = @groupreduce(+, k * x[i], zero(eltype(out))) + acc += res + end + if true + m = @groupreduce max x[i] typemin(eltype(out)) + end + out[i] = acc + m +end + +@kernel function groupreduce_cartesian!(out, @Const(x)) + I = @index(Global, Cartesian) + res = @groupreduce(+, x[I], zero(eltype(out))) + out[I] = res +end + +@kernel unsafe_indices = true function groupreduce_unsafe!(out, @Const(x)) + i = @index(Global, Linear) + val = i <= length(x) ? x[i] : zero(eltype(out)) + res = @groupreduce(+, val, zero(eltype(out))) + if i <= length(out) + out[i] = res + end +end + +@kernel function groupreduce_noargs!() +end + +# `val` of another type than `neutral`: the padding work-items contribute `neutral`, so the +# value is a `Union` of both types, and the call of the collective must not be union-split +@kernel function groupreduce_mixed!(out, @Const(x)) + i = @index(Global, Linear) + res = @groupreduce(+, x[i], Int64(0)) + out[i] = res +end + +@kernel function groupscan_mixed!(out, @Const(x)) + i = @index(Global, Linear) + res = @groupscan(+, x[i], 0) + out[i] = res +end + +# the composition of affine maps `x -> a * x + b`, first `f` then `g`: associative, but not +# commutative, so that the scans have to combine the values in order +compose(f, g) = (g[1] * f[1], g[1] * f[2] + g[2]) +const affine_identity = (1, 0) + +@kernel function groupscan!(out, @Const(x), op, neutral, ::Val{I}) where {I} + i = @index(Global, Linear) + res = @groupscan(op, x[i], neutral; inclusive = I) + out[i] = res +end + +@kernel function groupscan_bound!(out, @Const(x), op, neutral, ::Val{N}, ::Val{I}) where {N, I} + i = @index(Global, Linear) + res = @groupscan(op, x[i], neutral; groupsize = N, inclusive = I) + out[i] = res +end + +# the scan in the order of the local linear index, with padding in the middle of a workgroup +@kernel function groupscan_cartesian!(out, @Const(x)) + I = @index(Global, Cartesian) + res = @groupscan(compose, x[I], affine_identity) + out[I] = res +end + +# the same call site in a loop, and an exclusive scan of the counts +@kernel function groupscan_loop!(out, @Const(x)) + i = @index(Global, Linear) + acc = 0 + for k in 1:3 + res = @groupscan(+, k * x[i], 0) + acc += res + end + excl = @groupscan (+) x[i] 0 inclusive = false + out[i] = acc + excl +end + +# reference: the reduction of each workgroup of `groupsize` consecutive elements +function groupwise(op, x, groupsize) + return [reduce(op, x[((cld(i, groupsize) - 1) * groupsize + 1):min(cld(i, groupsize) * groupsize, end)]) for i in eachindex(x)] +end + +# reference: the scan of each workgroup of `groupsize` consecutive elements +function groupwise_scan(op, x, groupsize, neutral, inclusive) + out = similar(x, typeof(neutral)) + for first in 1:groupsize:length(x) + group = first:min(first + groupsize - 1, length(x)) + acc = neutral + for i in group + inclusive || (out[i] = acc) + acc = op(acc, x[i]) + inclusive && (out[i] = acc) + end + end + return out +end + +# A 32-bit float type that no backend shuffles, so that `@groupreduce` uses local memory +# only, also on backends with sub-groups. Its `+` adds the bits as integers. +primitive type TreeFloat <: AbstractFloat 32 end +TreeFloat(x::Integer) = reinterpret(TreeFloat, Int32(x)) +Base.:+(a::TreeFloat, b::TreeFloat) = reinterpret(TreeFloat, reinterpret(Int32, a) + reinterpret(Int32, b)) +Base.zero(::Type{TreeFloat}) = TreeFloat(0) +Base.typemin(::Type{TreeFloat}) = TreeFloat(typemin(Int32)) +Base.:(==)(a::TreeFloat, b::TreeFloat) = reinterpret(Int32, a) == reinterpret(Int32, b) + +# Run `f`, and return whether it failed to compile with GPUCompiler.jl#1004 (Metal: a phi of +# a by-reference argument and a device pointer, as in `x[i]` or the `neutral` argument), +# which the local-memory path runs into. +function gpucompiler_1004(f) + try + f() + return false + catch err + occursin("Invalid phi record", sprint(showerror, err)) || rethrow() + return true + end +end + +# `neutral` is evaluated once per work-item, also by the padding work-items +counted_neutral() = 0 +count_symbol(ex::Expr, sym) = sum(arg -> count_symbol(arg, sym), ex.args; init = 0) +count_symbol(ex, sym) = Int(ex === sym) + +function groupreduce_testsuite(backend, AT) + b = backend() + caps = KernelAbstractions.subgroup_capabilities(b) + + @testset "sub-group capabilities" begin + sub_groups = KI.supports_subgroups(b) && KI.supports_shuffle(b, UInt32) + @test (caps !== nothing) == sub_groups + shuffleable(T) = KernelAbstractions.__shuffleable(caps, T) + @test shuffleable(Int32) == sub_groups + @test shuffleable(Tuple{Int64, Bool}) == sub_groups + @test shuffleable(Float32) == (sub_groups && KI.supports_shuffle(b, Float32)) + @test shuffleable(Tuple{Float32, Int32}) == (sub_groups && KI.supports_shuffle(b, Float32)) + @test !shuffleable(TreeFloat) + @test !shuffleable(Ref{Int}) + # kernels without collectives don't query the capabilities + @test KernelAbstractions.kernel_subgroups(groupreduce_static!(b, 64)) == caps + @test KernelAbstractions.kernel_subgroups(groupscan!(b, 64)) == caps + @test KernelAbstractions.uses_collectives(groupreduce_static!(b, 64).f) + # also for a kernel without arguments, whose function takes a single argument + @test !KernelAbstractions.uses_collectives(groupreduce_noargs!(b).f) + end + + types = (Int32, Int64, Float32, TreeFloat) + @testset "$T, $(nameof(typeof(op)))" for T in types, (op, neutral) in ((+, zero(T)), (max, typemin(T))) + T === TreeFloat && op === max && continue + for (groupsize, n) in ((64, 64), (64, 256), (32, 100), (256, 1000), (7, 23), (1, 3)) + x = T.(rand(1:100, n)) + out = AT(fill(zero(T), n)) + if gpucompiler_1004(() -> groupreduce_static!(b, groupsize)(out, AT(x), op, neutral; ndrange = n)) + @test_broken false + continue + end + @test Array(out) == groupwise(op, x, groupsize) + + fill!(out, zero(T)) + if gpucompiler_1004(() -> groupreduce_bound!(b)(out, AT(x), op, neutral, Val(256); ndrange = n, workgroupsize = groupsize)) + @test_broken false + continue + end + @test Array(out) == groupwise(op, x, groupsize) + end + end + + @testset "argmin" begin + # (value, index) pairs, the smallest value with the smallest index + for (groupsize, n) in ((64, 256), (32, 100), (7, 23)) + x = [(Float32(rand(1:20)), Int32(i)) for i in 1:n] + neutral = (Inf32, typemax(Int32)) + out = AT(fill(neutral, n)) + groupreduce_static!(b, groupsize)(out, AT(x), min, neutral; ndrange = n) + @test Array(out) == groupwise(min, x, groupsize) + end + end + + @testset "loop" begin + x = rand(1:100, 100) + out = AT(zeros(Int, 100)) + groupreduce_loop!(b, 64)(out, AT(x); ndrange = 100) + @test Array(out) == 6 .* groupwise(+, x, 64) .+ groupwise(max, x, 64) + end + + @testset "cartesian" begin + x = rand(1:100, 10, 12) + out = AT(zeros(Int, 10, 12)) + groupreduce_cartesian!(b, (4, 8))(out, AT(x); ndrange = size(x)) + ref = similar(x) + for I in CartesianIndices(x) + g = (cld(I[1], 4) - 1) * 4 .+ (1:4), (cld(I[2], 8) - 1) * 8 .+ (1:8) + ref[I] = sum(x[intersect(g[1], axes(x, 1)), intersect(g[2], axes(x, 2))]) + end + @test Array(out) == ref + end + + @testset "unsafe_indices" begin + x = rand(1:100, 100) + out = AT(zeros(Int, 100)) + groupreduce_unsafe!(b, 64)(out, AT(x); ndrange = 128) + @test Array(out) == groupwise(+, x, 64) + end + + @testset "mixed types" begin + x = Int32.(rand(1:100, 100)) + out = AT(zeros(Int64, 100)) + groupreduce_mixed!(b, 64)(out, AT(x); ndrange = 100) + @test Array(out) == groupwise(+, Int64.(x), 64) + end + + @testset "@groupscan" begin + @testset "inclusive = $I" for I in (true, false) + for (groupsize, n) in ((64, 64), (64, 256), (32, 100), (256, 1000), (7, 23), (1, 3)) + x = rand(1:100, n) + out = AT(zeros(Int, n)) + groupscan!(b, groupsize)(out, AT(x), +, 0, Val(I); ndrange = n) + @test Array(out) == groupwise_scan(+, x, groupsize, 0, I) + + y = [(rand((-1, 1, 2)), rand(-5:5)) for _ in 1:n] + out = AT(fill((0, 0), n)) + groupscan_bound!(b)(out, AT(y), compose, affine_identity, Val(256), Val(I); ndrange = n, workgroupsize = groupsize) + @test Array(out) == groupwise_scan(compose, y, groupsize, affine_identity, I) + end + end + + @testset "cartesian" begin + x = [(rand((-1, 1, 2)), rand(-5:5)) for _ in 1:10, _ in 1:12] + out = AT(fill((0, 0), 10, 12)) + groupscan_cartesian!(b, (4, 8))(out, AT(x); ndrange = size(x)) + ref = similar(x) + for gi in 1:4:10, gj in 1:8:12 + acc = affine_identity + # local linear order is column-major within the workgroup + for j in gj:(gj + 7), i in gi:(gi + 3) + (i <= 10 && j <= 12) || continue + acc = compose(acc, x[i, j]) + ref[i, j] = acc + end + end + @test Array(out) == ref + end + + @testset "loop" begin + x = rand(1:100, 100) + out = AT(zeros(Int, 100)) + groupscan_loop!(b, 64)(out, AT(x); ndrange = 100) + @test Array(out) == 6 .* groupwise_scan(+, x, 64, 0, true) .+ groupwise_scan(+, x, 64, 0, false) + end + + @testset "mixed types" begin + x = Int32.(rand(1:100, 100)) + out = AT(zeros(Int, 100)) + groupscan_mixed!(b, 64)(out, AT(x); ndrange = 100) + @test Array(out) == groupwise_scan(+, Int.(x), 64, 0, true) + end + end + + @testset "errors" begin + @test_throws "must be used as a statement" @macroexpand @kernel function f(y, x) + i = @index(Global) + y[i] = @groupreduce(+, x[i], 0) + end + @test_throws "unknown option" @macroexpand @kernel function f(y, x) + res = @groupreduce(+, x[1], 0; foo = 1) + end + @test_throws "given more than once" @macroexpand @kernel function f(y, x) + res = @groupreduce(+, x[1], 0; groupsize = 32, groupsize = 64) + end + # the upper bound of the workgroup size is a keyword + @test_throws "expects `op`, `val` and `neutral`" @macroexpand @kernel function f(y, x) + res = @groupreduce(+, x[1], 0, 64) + end + @test_throws "unknown option" @macroexpand @kernel function f(y, x) + res = @groupreduce(+, x[1], 0; subgroups = true) + end + end + + @testset "kernels in a local scope" begin + # the kernel's properties are methods of its function, which works in a local scope too + function local_groupreduce(b, AT, x) + @kernel function local_groupreduce!(out, @Const(x)) + i = @index(Global, Linear) + res = @groupreduce(+, x[i], zero(eltype(out))) + out[i] = res + end + out = AT(zeros(eltype(x), length(x))) + kernel = local_groupreduce!(b, 32) + @test KernelAbstractions.uses_collectives(kernel.f) + kernel(out, AT(x); ndrange = length(x)) + return Array(out) + end + x = Int32.(rand(1:100, 100)) + @test local_groupreduce(b, AT, x) == groupwise(+, x, 32) + end + + @testset "neutral evaluated once" begin + ex = @macroexpand @kernel function f(y, x) + i = @index(Global, Linear) + res = @groupreduce(+, x[i], counted_neutral()) + y[i] = res + end + @test count_symbol(ex, :counted_neutral) == 1 + end + return +end diff --git a/test/runtests.jl b/test/runtests.jl index d34640414..67b6a8f9d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -28,6 +28,49 @@ end @test KernelAbstractions.precompiled_launch[] skip = !workloads end +# A package with overloaded kernels, which have to define the kernel's properties only once, +# since methods can't be overwritten during precompilation +@testset "Precompiling overloaded kernels" begin + mktempdir() do dir + mkpath(joinpath(dir, "OverloadedKernels", "src")) + write( + joinpath(dir, "OverloadedKernels", "src", "OverloadedKernels.jl"), + """ + module OverloadedKernels + using KernelAbstractions + @kernel function plain!(a) + a[@index(Global)] = 1 + end + @kernel function plain!(a, b) + a[@index(Global)] = b + end + @kernel function tiled!(a) + a[@index(Global)] = tile_width(@tile()) + end + @kernel function tiled!(a, b) + a[@index(Global)] = b + end + @kernel function reduced!(a) + res = @groupreduce(+, a[@index(Global)], zero(eltype(a)); groupsize = 64) + a[@index(Global)] = res + end + @kernel function reduced!(a, b) + res = @groupreduce(+, b, zero(eltype(a)); groupsize = 64) + a[@index(Global)] = res + end + end + """ + ) + pushfirst!(LOAD_PATH, dir) + try + id = Base.PkgId("OverloadedKernels") + @test Base.compilecache(id) isa Union{String, Tuple} + finally + popfirst!(LOAD_PATH) + end + end +end + KernelAbstractions.versioninfo(POCLBackend()) @info "Configuration" pocl = KernelAbstractions.POCL.nanoOpenCL.pocl_standalone_jll.libpocl @@ -75,11 +118,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 +351,84 @@ 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 + +# `@groupreduce` reduces sub-groups with PoCL's native collectives +@kernel function groupreduce_codegen!(out, @Const(x)) + i = @index(Global, Linear) + res = @groupreduce(+, x[i], zero(eltype(out))) + out[i] = res +end +@testset "POCL @groupreduce with sub-groups" begin + x = ones(Int32, 100) + out = zeros(Int32, 100) + ir = sprint() do io + @device_code_llvm io = io debuginfo = :none groupreduce_codegen!(POCLBackend(), 64)(out, x; ndrange = 100) + end + @test out == [fill(64, 64); fill(36, 36)] + @test occursin("sub_group_reduce_add", ir) +end + # Julia doesn't turn a splat of more than 32 elements into a direct call, so a launch with # many arguments allocates unless every layer passes them on as a tuple @testset "POCL launch with many arguments" begin diff --git a/test/testsuite.jl b/test/testsuite.jl index 541028ad8..4ff4a4646 100644 --- a/test/testsuite.jl +++ b/test/testsuite.jl @@ -45,6 +45,8 @@ include("convert.jl") include("specialfunctions.jl") include("random.jl") include("spawn.jl") +include("groupreduce.jl") +include("tiles.jl") function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{String}()) @conditional_testset "Unittests" skip_tests begin @@ -123,6 +125,14 @@ function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{ spawn_testsuite(backend, AT) end + @conditional_testset "Group reductions" skip_tests begin + groupreduce_testsuite(backend, AT) + end + + @conditional_testset "Tiles" skip_tests begin + tile_testsuite(backend, AT) + end + return end diff --git a/test/tiles.jl b/test/tiles.jl new file mode 100644 index 000000000..7d0ecb9a2 --- /dev/null +++ b/test/tiles.jl @@ -0,0 +1,411 @@ +using StaticArrays + +# Kernels with tiles, whose width is chosen when they are constructed, so that the same kernel +# runs with several widths. + +# every tile loops a different number of times, and some return early +@kernel function tile_divergent!(out, @Const(x)) + t = @tile() + i = @index(Global, Linear) + g = @index(Tile) + acc = zero(eltype(out)) + for _ in 1:(g % 3) + acc += tile_reduce(t, +, x[i]) + end + if g % 4 == 2 + return + end + out[i] = acc + tile_shfl(t, x[i], 1) +end + +# the tile, its index and its rotation by a lane, the votes, and the ballot of the odd values +# (which needs a sub-group width of at most 64) +@kernel function tile_ops!(out, ballots, @Const(x), ::Val{ballot}) where {ballot} + t = @tile() + i = @index(Global, Linear) + out[i, 1] = t.index + out[i, 2] = t.lane + out[i, 3] = @index(Tile) + out[i, 4] = tile_shfl(t, x[i], t.lane + 1) + out[i, 5] = tile_shfl_up(t, x[i], 1) + out[i, 6] = tile_shfl_down(t, x[i], 1) + out[i, 7] = tile_shfl_xor(t, x[i], 1) + out[i, 8] = tile_any(t, x[i] == 3) + out[i, 9] = tile_all(t, x[i] > 0) + out[i, 10] = tile_width(t) + ballots[i] = ballot ? tile_ballot(t, isodd(x[i])) : zero(UInt64) +end + +# code that depends on the width: a mask with a bit per lane, of the smallest type +@inline tile_mask_type(t::Tile) = tile_width(t) <= 8 ? UInt8 : tile_width(t) <= 16 ? UInt16 : + tile_width(t) <= 32 ? UInt32 : UInt64 +@kernel function tile_mask!(out, sizes, @Const(x)) + t = @tile() + i = @index(Global, Linear) + mask = tile_ballot(t, isodd(x[i])) % tile_mask_type(t) + out[i] = mask + sizes[i] = sizeof(mask) +end + +# communication through local memory, ordered by `tile_barrier`, with a static workgroup +# size +@kernel function tile_barrier!(out, @Const(x)) + t = @tile() + i = @index(Global, Linear) + lid = @index(Local, Linear) + mem = @localmem eltype(out) (prod(@groupsize()),) + mem[lid] = x[i] + tile_barrier(t) + N = tile_width(t) + out[i] = mem[(t.index - 1) * N + mod1(t.lane + 1, N)] +end + +tile_add(a, b) = a .+ b + +# reductions of tuples and static arrays +@kernel function tile_reduce!(out, out_sv, @Const(x)) + t = @tile() + i = @index(Global, Linear) + out[i] = tile_reduce(t, tile_add, (x[i], 2 * x[i])) + out_sv[i] = tile_reduce(t, +, SVector(x[i], -x[i], Int32(1))) +end + +# the composition of affine maps, first `f` then `g`: associative, but not commutative +tile_compose(f, g) = (g[1] * f[1], g[1] * f[2] + g[2]) + +@kernel function tile_scans!(incl, excl, affine, affine_excl, @Const(x), @Const(maps)) + t = @tile() + i = @index(Global, Linear) + incl[i] = tile_scan(t, +, x[i]) + excl[i] = tile_exclusive_scan(t, +, x[i], Int32(7)) + affine[i] = tile_scan(t, tile_compose, maps[i]) + affine_excl[i] = tile_exclusive_scan(t, tile_compose, maps[i], (Int32(2), Int32(3))) +end + +@kernel unsafe_indices = true function tile_unsafe!(out, @Const(x)) + t = @tile() + i = @index(Global, Linear) + v = i <= length(x) ? x[i] : zero(eltype(x)) + s = tile_reduce(t, +, v) + if i <= length(out) + out[i] = s + end +end + +# a kernel that doesn't compile, to check that invalid launches are rejected before +@kernel function tile_uncompilable!(out) + out[@index(Global, Linear)] = Base.inferencebarrier(identity)(1) +end + +# a kernel without tiles +@kernel function tile_plain!(out) + out[@index(Global, Linear)] = 1 +end + +@kernel function tile_noargs!() +end + +# the nodes of an element per tile, over an N-D ndrange +@kernel function tile_elements!(out, index, @Const(x)) + t = @tile() + I = @index(Global, NTuple) + out[I...] = tile_reduce(t, +, x[I...]) + index[I...] = @index(Tile) +end + +struct TileForward{F} + f::F +end +(w::TileForward)(args...) = w.f(args...) + +# Whether every tile of `N` consecutive work-items of every workgroup is entirely inside or +# outside of the ndrange, by enumerating them +function whole_tiles_reference(N, ndrange, workgroupsize) + for b in CartesianIndices(cld.(ndrange, workgroupsize)) + inside = [all((Tuple(b) .- 1) .* workgroupsize .+ Tuple(I) .<= ndrange) for I in CartesianIndices(workgroupsize)] + all(t -> allequal(inside[t:(t + N - 1)]), 1:N:length(inside)) || return false + end + return true +end + +# the reference results, by tile of `N` consecutive elements +tile_groups(n, N) = [((i - 1) ÷ N * N + 1):((i - 1) ÷ N * N + N) for i in 1:n] + +function tile_testsuite_width(backend, AT, N) + W = KI.sub_group_size(backend) + # several workgroups, and a partial one with a padding tile + groupsize = 2N + ntiles = 5 + n = ntiles * N + x = Int32.(rand(1:100, n)) + groups = tile_groups(n, N) + + @test tile_width(tile_divergent!(backend; tile = N)) == N + @test tile_width(tile_barrier!(backend, groupsize; tile = N)) == N + + divergent = map(1:n) do i + g = (i - 1) ÷ N + 1 + g % 4 == 2 && return Int32(-1) + return Int32(g % 3) * sum(x[groups[i]]) + x[first(groups[i])] + end + @testset "divergent tiles" begin + out = AT(fill(Int32(-1), n)) + tile_divergent!(backend; tile = N)(out, AT(x); ndrange = n) + @test Array(out) == divergent + end + + @testset "tile operations" begin + x3 = copy(x) + x3[2] = 3 + out = AT(zeros(Int32, n, 10)) + ballots = AT(zeros(UInt64, n)) + tile_ops!(backend; tile = N)(out, ballots, AT(x3), Val(W <= 64); ndrange = n) + out, ballots = Array(out), Array(ballots) + for i in 1:n + g = groups[i] + lane = i - first(g) + 1 + @test out[i, 2] == lane + @test out[i, 3] == (i - 1) ÷ N + 1 + @test out[i, 4] == x3[g[mod1(lane + 1, N)]] + @test out[i, 5] == x3[lane > 1 ? i - 1 : i] + @test out[i, 6] == x3[lane < N ? i + 1 : i] + @test out[i, 7] == x3[g[((lane - 1) ⊻ 1) + 1]] + @test out[i, 8] == any(==(3), x3[g]) + @test out[i, 9] == 1 + @test out[i, 10] == N + if W <= 64 + @test ballots[i] == sum(UInt64(1) << (j - 1) for j in 1:N if isodd(x3[g[j]]); init = UInt64(0)) + end + end + # the tile within the workgroup + @test all(i -> out[i, 1] >= 1, 1:n) + end + + if W <= 64 + @testset "width-dependent code" begin + out = AT(zeros(UInt64, n)) + sizes = AT(zeros(Int, n)) + tile_mask!(backend; tile = N)(out, sizes, AT(x); ndrange = n) + @test all(==(cld(N, 8)), Array(sizes)) + @test Array(out) == [sum(UInt64(1) << (j - 1) for j in 1:N if isodd(x[g[j]]); init = UInt64(0)) for g in groups] + end + end + + @testset "tile_barrier" begin + out = AT(zeros(Int32, n)) + tile_barrier!(backend, groupsize; tile = N)(out, AT(x); ndrange = n) + @test Array(out) == [x[g[mod1(i - first(g) + 2, N)]] for (i, g) in enumerate(groups)] + end + + @testset "tuples and static arrays" begin + out = AT(fill((Int32(0), Int32(0)), n)) + out_sv = AT(fill(SVector{3, Int32}(0, 0, 0), n)) + tile_reduce!(backend; tile = N)(out, out_sv, AT(x); ndrange = n) + @test Array(out) == [(sum(x[g]), 2 * sum(x[g])) for g in groups] + @test Array(out_sv) == [SVector(sum(x[g]), -sum(x[g]), Int32(N)) for g in groups] + end + + @testset "scans" begin + maps = [(Int32(rand((-1, 1, 2))), Int32(rand(-5:5))) for _ in 1:n] + incl, excl = AT(zeros(Int32, n)), AT(zeros(Int32, n)) + affine, affine_excl = AT(fill((Int32(0), Int32(0)), n)), AT(fill((Int32(0), Int32(0)), n)) + tile_scans!(backend; tile = N)(incl, excl, affine, affine_excl, AT(x), AT(maps); ndrange = n) + prefix(f, v, g, i) = foldl(f, v[first(g):i]) + @test Array(incl) == [prefix(+, x, g, i) for (i, g) in enumerate(groups)] + @test Array(excl) == [i == first(g) ? Int32(7) : Int32(7) + prefix(+, x, g, i - 1) for (i, g) in enumerate(groups)] + @test Array(affine) == [prefix(tile_compose, maps, g, i) for (i, g) in enumerate(groups)] + init = (Int32(2), Int32(3)) + @test Array(affine_excl) == [i == first(g) ? init : tile_compose(init, prefix(tile_compose, maps, g, i - 1)) for (i, g) in enumerate(groups)] + end + + @testset "unsafe_indices" begin + out = AT(zeros(Int32, n)) + tile_unsafe!(backend; tile = N)(out, AT(x); ndrange = n) + @test Array(out) == [sum(x[g]) for g in groups] + end + + @testset "explicit workgroup sizes" begin + out = AT(fill(Int32(-1), n)) + # a single tile per workgroup + tile_divergent!(backend; tile = N)(out, AT(x); ndrange = n, workgroupsize = N) + @test Array(out) == divergent + # several tiles per workgroup, which diverge, with a padding tile in the last one + for workgroupsize in (2N, 4N) + out = AT(fill(Int32(-1), n)) + tile_divergent!(backend; tile = N)(out, AT(x); ndrange = n, workgroupsize) + @test Array(out) == divergent + out = AT(fill(Int32(-1), n)) + tile_unsafe!(backend; tile = N)(out, AT(x); ndrange = n, workgroupsize) + @test Array(out) == [sum(x[g]) for g in groups] + end + end + + @testset "rejected launches" begin + out = AT(zeros(Int32, n)) + k = tile_divergent!(backend; tile = N) + @test_throws "isn't a multiple of $N" k(out, AT(x); ndrange = n - 1) + @test_throws "isn't a multiple of $N" k(out, AT(x); ndrange = n, workgroupsize = N ÷ 2) + @test_throws "isn't a multiple of $N" tile_barrier!(backend, N + N ÷ 2; tile = N) + @test_throws "isn't a multiple of $N" k(out, AT(x); ndrange = (N - 1, 3)) + @test_throws "isn't a multiple of $N" tile_divergent!(backend, N, n - 1; tile = N) + # before compiling the kernel + u = tile_uncompilable!(backend; tile = N) + @test_throws "isn't a multiple of $N" u(out; ndrange = n - 1) + @test_throws "isn't a multiple of $N" u(out; ndrange = n, workgroupsize = N ÷ 2) + @test_throws "isn't a multiple of $N" u(out; ndrange = (N - 1, 3)) + @test_throws "isn't a multiple of $N" u(out; ndrange = n, workgroupsize = 0) + if N >= 4 + # tiles that would straddle the end of the first dimension + @test_throws "partial tiles" u(out; ndrange = (3, N), workgroupsize = (2, N ÷ 2)) + end + end + + @testset "N-D ndrange" begin + # elements of 2 × N/2 nodes, a tile each, in workgroups of one or several elements + dims = (2, N ÷ 2) + nelem = 5 + x = Int32.(rand(1:100, dims..., nelem)) + sums = [sum(x[:, :, e]) for _ in 1:dims[1], _ in 1:dims[2], e in 1:nelem] + elements = [e for _ in 1:dims[1], _ in 1:dims[2], e in 1:nelem] + k = tile_elements!(backend; tile = N) + for workgroupsize in (nothing, (dims..., 1), (dims..., 2)) + out, index = AT(zeros(Int32, size(x))), AT(zeros(Int, size(x))) + k(out, index, AT(x); ndrange = size(x), workgroupsize) + @test Array(out) == sums + # tiles are numbered workgroup by workgroup, so one per element + @test Array(index) == elements + end + end + return +end + +function tile_testsuite(backend, AT) + b = backend() + @testset "widths" begin + # kernels that use tiles need a width + @test KernelAbstractions.uses_tiles(gpu_tile_divergent!) + @test KernelAbstractions.uses_tiles(gpu_tile_elements!) + @test_throws "needs a tile width" tile_divergent!(b) + @test_throws "needs a tile width" tile_divergent!(b, 64) + @test_throws "has to be an integer" tile_divergent!(b; tile = 8.0) + @test tile_width(tile_plain!(b)) === nothing + @test !KernelAbstractions.uses_tiles(tile_plain!(b).f) + # also for a kernel without arguments, whose function takes a single argument + @test !KernelAbstractions.uses_tiles(tile_noargs!(b).f) + # a kernel with tiles in a local scope + function local_tile_kernel(b) + @kernel function local_tiles!(out) + t = @tile() + out[@index(Global, Linear)] = tile_width(t) + end + return local_tiles!(b; tile = 8) + end + if supports_tiles(b, 8) + @test tile_width(local_tile_kernel(b)) == 8 + end + end + @test !supports_tiles(b, 3) + @test !supports_tiles(b, 0) + @test_throws "power of two" tile_divergent!(b; tile = 3) + + @testset "whole tiles" begin + whole_tiles = KernelAbstractions.whole_tiles + @test whole_tiles(16, (4, 4, 5), (4, 4, 2)) + @test !whole_tiles(32, (4, 4, 5), (4, 4, 2)) + @test whole_tiles(32, (4, 4, 6), (4, 4, 2)) + @test !whole_tiles(16, (6, 4), (4, 4)) + @test !whole_tiles(4, (6, 4), (4, 4)) + @test whole_tiles(2, (6, 4), (4, 4)) + @test whole_tiles(16, (64, 60), (16, 16)) + @test whole_tiles(32, (64, 60), (16, 16)) + @test !whole_tiles(32, (60, 64), (16, 16)) + @test !whole_tiles(64, (64, 62), (16, 16)) + @test whole_tiles(8, (8,), (16,)) + @test !whole_tiles(8, (12,), (16,)) + # tiles that don't follow the dimensions of the workgroup + @test whole_tiles(16, (16, 1), (24, 2)) + ok = true + @test whole_tiles(32, (0, 1), (1, 32)) + for N in (2, 4, 8), w1 in 1:6, w2 in 1:4, r1 in 0:7, r2 in 1:5 + (w1 * w2) % N == 0 || continue + ok &= whole_tiles(N, (r1, r2), (w1, w2)) == whole_tiles_reference(N, (r1, r2), (w1, w2)) + end + for N in (2, 4), w in Iterators.product(1:3, 1:3, 1:2), r in Iterators.product(1:4, 1:3, 1:3) + prod(w) % N == 0 || continue + ok &= whole_tiles(N, r, w) == whole_tiles_reference(N, r, w) + end + @test ok + end + + @testset "padded shapes" begin + # the padded iteration space fits the index type, without overflowing + shape = KernelAbstractions.tile_shape + @test shape(32, 1024, (32_000_016, 66), Int32) === nothing + @test shape(32, 32, (32_000_016, 66), Int32) == (16, 2) + @test shape(32, 1024, (32 * (typemax(Int) ÷ 96 - 100), 3), Int) == (1024, 1) + end + + W = KI.supports_subgroups(b) ? KI.sub_group_size(b) : 0 + if W == 0 || !supports_tiles(b, W) + # a backend without tiles: constructing a kernel with tiles is an error + @test_throws "doesn't support tiles" tile_divergent!(b; tile = 8) + W > 0 && @test_throws "doesn't support tiles" tile_divergent!(b; tile = W) + W > 1 && @test_throws "doesn't support tiles" tile_barrier!(b, W; tile = W ÷ 2) + return + end + @test !supports_tiles(b, 2W) + for N in (W ÷ 2, W ÷ 4) + N >= 1 || continue + @test supports_tiles(b, N) == KI.supports_partitionable_subgroups(b) + end + + @testset "automatic workgroup sizes" begin + size = KernelAbstractions.tile_launch_workgroupsize + launch = KernelAbstractions.LinearLaunch{Int32}() + # a recommendation of less than a tile still gets one + @test size(b, launch, W, max(W ÷ 4, 1), (2W,)) == (W,) + # shapes that the leading dimensions don't divide into whole tiles + for range in ((3, W), (3, 64W), (5, 3, 2W), (6, W ÷ 2)) + items = size(b, launch, W, 4W, range) + @test prod(items) % W == 0 && whole_tiles_reference(W, range, items) + end + @test_throws "no workgroup size" size(b, launch, W, 4W, (3, W + 1)) + end + + @testset "construction" begin + @test tile_width(tile_divergent!(b; tile = W)) == W + @test tile_width(tile_divergent!(b, W; tile = W)) == W + @test tile_width(tile_divergent!(b, W, 4W; tile = W)) == W + # a kernel that doesn't use tiles can still be given a width, for its launch + @test tile_width(tile_plain!(b; tile = W)) == W + @test_throws "larger than the sub-group width" tile_divergent!(b; tile = 2W) + # `similar` keeps the width, and a function that uses tiles needs one + k16 = tile_divergent!(b; tile = W ÷ 2) + @test tile_width(similar(k16, gpu_tile_ops!)) == W ÷ 2 + @test_throws "needs a tile width" similar(tile_plain!(b), k16.f) + @test tile_width(similar(tile_plain!(b), tile_plain!(b).f)) === nothing + # a wrapper that forwards its arguments keeps the width + @test tile_width(similar(k16, TileForward(k16.f))) == W ÷ 2 + @test tile_width(similar(tile_plain!(b), TileForward(tile_plain!(b).f))) === nothing + # the width reaches the device through the kernel's metadata, also when adapted + iterspace, _ = KernelAbstractions.partition(k16, (W,), (W ÷ 2,)) + ctx = KernelAbstractions.mkcontext(k16, (W,), iterspace) + @test KernelAbstractions.__tile_width(ctx) === Val(W ÷ 2) + @test KernelAbstractions.__tile_width(KernelAbstractions.Adapt.adapt(AT, ctx)) === Val(W ÷ 2) + @test KernelAbstractions.__tile_width(KernelAbstractions.mkcontext(tile_plain!(b), (W,), iterspace)) === nothing + end + + # the sub-group width, and narrower widths, several of which share a sub-group, with the + # same kernels + for N in unique((W, W ÷ 2, W ÷ 4)) + N >= 2 || continue + if supports_tiles(b, N) + @testset "tile = $N" begin + tile_testsuite_width(b, AT, N) + end + else + @test_throws "only supports tiles of its sub-group width" tile_divergent!(b; tile = N) + end + end + return +end