KernelInterface: more sub-group communication (shuffles, votes, constant width)

.buildkite/pipeline.yml

@@ -24,7 +24,9 @@ steps:
       julia -e 'println("--- :julia: Instantiating project")
                 using Pkg
                 Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()),
-                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3
+                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])
+                # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11)
+                Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3
       julia -e 'println("--- :julia: Developing CUDA")
                 using Pkg
                 url="https://github.com/JuliaGPU/CUDA.jl"
@@ -64,7 +66,9 @@ steps:
       julia -e 'println("--- :julia: Instantiating project")
                 using Pkg
                 Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()),
-                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3
+                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])
+                # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11)
+                Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3
       julia -e 'println("--- :julia: Developing Metal")
                 using Pkg
                 Pkg.add([(; url="https://github.com/JuliaGPU/Metal.jl", rev="ka-0.10")])' || exit 3
@@ -100,7 +104,9 @@ steps:
       julia -e 'println("--- :julia: Instantiating project")
                 using Pkg
                 Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()),
-                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3
+                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])
+                # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11)
+                Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3
       julia -e 'println("--- :julia: Developing oneAPI")
                 using Pkg
                 Pkg.add(url="https://github.com/JuliaGPU/AcceleratedKernels.jl", rev="main")
@@ -139,7 +145,9 @@ steps:
                 using Pkg
                 println("--- :julia: Instantiating project")
                 Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()),
-                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3
+                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])
+                # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11)
+                Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3
       julia -e 'println("--- :julia: Developing AMDGPU")
                 using Pkg
                 Pkg.add(url="https://github.com/JuliaGPU/AcceleratedKernels.jl", rev="main")
@@ -178,11 +186,14 @@ steps:
       julia -e 'println("--- :julia: Instantiating project")
                 using Pkg
                 Pkg.develop([PackageSpec(; name="KernelAbstractions", path=pwd()),
-                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])' || exit 3
+                PackageSpec(; name="KernelInterface", path=joinpath(pwd(), "lib", "KernelInterface"))])
+                # TODO: remove once JuliaGPU/OpenCL.jl#526 is released ([sources] needs Julia 1.11)
+                Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3
       julia -e 'println("--- :julia: Developing OpenCL")
                 using Pkg
                 Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="ka-0.10")
-                Pkg.develop(; name="SPIRVIntrinsics")' || exit 3
+                # TODO: develop the ka-0.10 copy again once JuliaGPU/OpenCL.jl#526 is merged
+                Pkg.add(url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")' || exit 3
 
       julia -e 'println("+++ :julia: Running tests")
                 using Pkg

.github/workflows/ci.yml

@@ -56,9 +56,12 @@ jobs:
       # applies to the test run only, not to `Pkg.develop`. Do this before building, which
       # already resolves the environment (and fails if this repo requires an unregistered
       # version of KernelInterface).
+      #
+      # The same applies to the SPIRVIntrinsics branch in `[sources]`.
+      # TODO: remove that once JuliaGPU/OpenCL.jl#526 is released
       - name: Dev KernelInterface
         shell: bash
-        run: julia -e 'using Pkg; Pkg.activate("."); Pkg.develop(; path="lib/KernelInterface")'
+        run: julia -e 'using Pkg; Pkg.activate("."); Pkg.develop(; path="lib/KernelInterface"); Pkg.add(; url="https://github.com/JuliaGPU/OpenCL.jl", rev="vc/subgroup-votes", subdir="lib/intrinsics")'
       # Nightly tracks the next Julia release; don't fail CI on it
       - uses: julia-actions/julia-buildpkg@v1
         continue-on-error: ${{ matrix.version == 'nightly' }}

Project.toml

@@ -31,6 +31,8 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [sources]
 KernelInterface = {path = "lib/KernelInterface"}
+# TODO: remove once JuliaGPU/OpenCL.jl#526 is released
+SPIRVIntrinsics = {url = "https://github.com/JuliaGPU/OpenCL.jl", rev = "vc/subgroup-votes", subdir = "lib/intrinsics"}
 
 [extensions]
 LinearAlgebraExt = "LinearAlgebra"

docs/src/kernelinterface.md

@@ -77,7 +77,7 @@ What a backend implements, at a glance. The docstrings below have the details.
 | **Capabilities** | | [`supports_float64`](@ref), [`supports_atomics`](@ref), [`supports_unified`](@ref), [`supports_subgroups`](@ref), [`supports_shuffle`](@ref) (all `false`) |
 | **Compilation** | [`argconvert`](@ref), [`kernel_function`](@ref), [`launch`](@ref) | |
 | **Device** | [`get_local_id`](@ref), [`get_group_id`](@ref), [`get_local_size`](@ref), [`get_num_groups`](@ref), [`localmemory`](@ref), [`barrier`](@ref) | [`get_global_id`](@ref), [`get_global_size`](@ref) (derived from the primitive queries), [`_print`](@ref KernelInterface._print) (host `print`) |
-| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries, [`sub_group_barrier`](@ref); if `supports_shuffle(backend, T)`: [`shfl_down`](@ref) for `T` | |
+| **Sub-groups** | if `supports_subgroups`: [`sub_group_size`](@ref), the sub-group queries (with a constant [`get_max_sub_group_size`](@ref)), [`sub_group_barrier`](@ref), [`sub_group_any`](@ref), [`sub_group_all`](@ref), and [`sub_group_ballot`](@ref) for widths of at most 64; if `supports_shuffle(backend, T)`: [`shfl`](@ref), [`shfl_down`](@ref), [`shfl_up`](@ref), [`shfl_xor`](@ref) for the primitive `T` supported natively, including `UInt32` | shuffles and votes with a `width`, [`sub_group_match_any`](@ref), [`sub_group_reduce`](@ref), [`sub_group_scan`](@ref) (built on the shuffles and votes); shuffles of other primitive types (as `UInt32` words) and of structs (field by field) |
 
 Everything else, such as [`zeros`](@ref KernelInterface.zeros), [`ones`](@ref KernelInterface.ones),
 the launch-keyword handling of [`Kernel`](@ref) and [`@launch`](@ref KernelInterface.@launch),
@@ -151,17 +151,32 @@ get_global_size
 Sub-groups are optional ([`supports_subgroups`](@ref)). A work-group is divided into
 sub-groups of at most [`sub_group_size(backend)`](@ref sub_group_size) work-items. Which
 work-items form a sub-group, how many sub-groups there are, and which of them are partial
-is unspecified, and differs between devices and work-group shapes. For example, CUDA forms
-warps from consecutive linear work-item indices, while Intel's CPU OpenCL runtime forms
-sub-groups per row of a multi-dimensional work-group, so that a 33×2 work-group consists of
-four sub-groups of 32 and 1 work-items. What KernelInterface guarantees, and backends that
-report sub-group support have to ensure:
+can differ between devices and work-group shapes. For example, CUDA forms warps from
+consecutive linear work-item indices, while Intel's CPU OpenCL runtime forms sub-groups per
+row of a multi-dimensional work-group, so that a 33×2 work-group consists of four sub-groups
+of 32 and 1 work-items. What KernelInterface guarantees, and backends that report sub-group
+support have to ensure:
 
 - every work-item has a unique `(get_sub_group_id(), get_sub_group_local_id())` pair in its
   work-group, which doesn't change during the kernel;
 - the sub-group ids are `1:get_num_sub_groups()`, and the lanes of a sub-group are
   `1:get_sub_group_size()`;
-- a 1-D work-group of at most `sub_group_size(backend)` work-items is a single sub-group.
+- if the work-group is 1-D, or its x extent `get_local_size().x` is a multiple of the
+  sub-group width `W` ([`get_max_sub_group_size`](@ref)), sub-groups are formed from
+  consecutive work-items, x fastest: the work-item with the linear index
+  `lin = x + (y - 1) * size.x + (z - 1) * size.x * size.y` (for `(; x, y, z) =
+  get_local_id()` and `size = get_local_size()`) is in sub-group `(lin - 1) ÷ W + 1`, lane
+  `(lin - 1) % W + 1`. Only the last sub-group can be partial. In particular, a 1-D
+  work-group of at most `W` work-items is a single sub-group.
+
+Other shapes can form sub-groups differently, e.g. per row of the work-group: for portable
+code, make the x extent of multi-dimensional work-groups a multiple of the sub-group width.
+
+A sub-group is partial when it has fewer work-items than the sub-group width: the lanes
+`get_sub_group_size()+1:get_max_sub_group_size()` have no work-item. Shuffles from those lanes
+give unspecified values, and the votes, [`sub_group_match_any`](@ref),
+[`sub_group_reduce`](@ref) and [`sub_group_scan`](@ref) only take the work-items of the
+sub-group into account.
 
 In particular, [`get_num_sub_groups`](@ref) can be larger than
 `cld(prod(get_local_size()), get_max_sub_group_size())`. Storage for a value per sub-group
@@ -191,8 +206,25 @@ localmemory
 
 ### Communication
 
+The shuffles, votes and collectives below exchange values between the work-items of a
+sub-group, not memory: they don't order or make visible the work-items' accesses to local
+or global memory. To communicate through memory within a sub-group, e.g. a work-item reading
+what another one wrote to local memory, use [`sub_group_barrier`](@ref) between the write
+and the read.
+
 ```@docs
+shfl
 shfl_down
+shfl_up
+shfl_xor
+shfl(::Any, ::Integer, ::Integer)
+sub_group_any
+sub_group_all
+sub_group_ballot
+sub_group_match_any
+sub_group_any(::Bool, ::Integer)
+sub_group_reduce
+sub_group_scan
 ```
 
 ### Printing

lib/KernelInterface/src/KernelInterface.jl

@@ -31,7 +31,10 @@ include("host.jl")
             :get_group_id, :get_num_groups,
             :get_sub_group_size, :get_max_sub_group_size, :get_num_sub_groups,
             :get_sub_group_id, :get_sub_group_local_id,
-            :localmemory, :shfl_down, :barrier, :sub_group_barrier, :_print,
+            :localmemory, :shfl, :shfl_down, :shfl_up, :shfl_xor,
+            :sub_group_any, :sub_group_all, :sub_group_ballot, :sub_group_match_any,
+            :sub_group_reduce, :sub_group_scan,
+            :barrier, :sub_group_barrier, :_print,
             # compilation and launch
             :Kernel, :kernel_function, :argconvert, :launch, Symbol("@launch"),
             :launch_configuration, :max_work_group_size, :max_work_group_dims,