KI Sub-groups

src/interface.jl

@@ -102,6 +102,78 @@ Returns the unique group ID.
 """
 function get_group_id end
 
+"""
+    get_sub_group_size()::UInt32
+
+Returns the number of work-items in the sub-group.
+
+!!! note
+    Backend implementations **must** implement:
+    ```
+    @device_override get_sub_group_size()::UInt32
+    ```
+"""
+function get_sub_group_size end
+
+"""
+    get_max_sub_group_size()::UInt32
+
+Returns the maximum sub-group size for sub-groups in the current workgroup.
+
+!!! note
+    Backend implementations **must** implement:
+    ```
+    @device_override get_max_sub_group_size()::UInt32
+    ```
+"""
+function get_max_sub_group_size end
+
+"""
+    get_num_sub_groups()::UInt32
+
+Returns the number of sub-groups in the current workgroup.
+
+!!! note
+    Backend implementations **must** implement:
+    ```
+    @device_override get_num_sub_groups()::UInt32
+    ```
+"""
+function get_num_sub_groups end
+
+"""
+    get_sub_group_id()::UInt32
+
+Returns the sub-group ID within the work-group.
+
+!!! note
+    1-based.
+
+!!! note
+    Backend implementations **must** implement:
+    ```
+    @device_override get_sub_group_id()::UInt32
+    ```
+"""
+function get_sub_group_id end
+
+"""
+    get_sub_group_local_id()::UInt32
+
+Returns the work-item ID within the current sub-group.
+
+!!! note
+    1-based.
+
+!!! note
+    Backend implementations **must** implement:
+    ```
+    @device_override get_sub_group_local_id()::UInt32
+    ```
+"""
+function get_sub_group_local_id end
+
+
 """
     localmemory(::Type{T}, dims)
 
@@ -116,6 +188,40 @@ Declare memory that is local to a workgroup.
 """
 localmemory(::Type{T}, dims) where {T} = localmemory(T, Val(dims))
 
+"""
+    shfl_down(val::T, offset::Integer) where T
+
+Read `val` from a lane with higher id given by `offset`.
+
+!!! note
+    `shfl_down` must be encountered by all workitems of a sub-group executing the kernel or by none at all.
+
+!!! note
+    Backend implementations **must** implement:
+    ```
+    @device_override shfl_down(val::T, offset::Integer) where T
+    ```
+    As well as the on-device functionality.
+
+    This implementation **must** be synchronizing.
+    That is, kernels using this function can safely assume that
+    they do **not** need a `sub_group_barrier` before calling
+    this function.
+"""
+function shfl_down end
+
+"""
+    shfl_down_types(::Backend)::Vector{DataType}
+
+Returns a vector of `DataType`s supported on `backend`
+
+!!! note
+    Backend implementations **must** implement this function
+    only if they support `shfl_down` for any types.
+"""
+shfl_down_types(::Backend) = DataType[]
+
+
 """
     barrier()
 
@@ -139,6 +245,29 @@ function barrier()
     error("Group barrier used outside kernel or not captured")
 end
 
+"""
+    sub_group_barrier()
+
+After a `sub_group_barrier()` call, all read and writes to global and local memory
+from each thread in the sub-group are visible in from all other threads in the
+sub-group.
+
+This does **not** guarantee that a write from a thread in a certain sub-group will
+be visible to a thread in a different sub-group.
+
+!!! note
+    `sub_group_barrier()` must be encountered by all workitems of a sub-group executing the kernel or by none at all.
+
+!!! note
+    Backend implementations **must** implement:
+    ```
+    @device_override sub_group_barrier()
+    ```
+"""
+function sub_group_barrier()
+    error("Sub-group barrier used outside kernel or not captured")
+end
+
 """
     _print(args...)
 
@@ -220,6 +349,22 @@ kernel launch with too big a workgroup is attempted.
 """
 function max_work_group_size end
 
+"""
+    sub_group_size(backend)::Int
+
+Returns a reasonable sub-group size supported by the currently
+active device for the specified backend. This would typically
+be 32, or 64 for devices that don't support 32.
+
+!!! note
+    Backend implementations **must** implement:
+    ```
+    sub_group_size(backend::NewBackend)::Int
+    ```
+    As well as the on-device functionality.
+"""
+function sub_group_size end
+
 """
     multiprocessor_count(backend::NewBackend)::Int
 

src/pocl/backend.jl

@@ -9,6 +9,8 @@ using SPIRV_LLVM_Backend_jll, SPIRV_Tools_jll
 import KernelAbstractions as KA
 import KernelAbstractions.KernelInterface as KI
 
+import SPIRVIntrinsics
+
 import StaticArrays
 
 import Adapt
@@ -227,10 +229,39 @@ end
 function KI.max_work_group_size(::POCLBackend)::Int
     return Int(device().max_work_group_size)
 end
+function KI.sub_group_size(::POCLBackend)::Int
+    # POCL can technically support any sub_group size.
+    #  Check for common values used on GPUs then
+    #  return 1 otherwise
+    sg_sizes = cl.device().sub_group_sizes
+    if 32 in sg_sizes
+        return 32
+    elseif 64 in sg_sizes
+        return 64
+    elseif 16 in sg_sizes
+        return 16
+    else
+        return 1
+    end
+end
 function KI.multiprocessor_count(::POCLBackend)::Int
     return Int(device().max_compute_units)
 end
 
+function KI.shfl_down_types(::POCLBackend)
+    res = copy(SPIRVIntrinsics.gentypes)
+
+    backend_extensions = cl.device().extensions
+    if "cl_khr_fp64" ∉ backend_extensions
+        res = setdiff(res, [Float64])
+    end
+    if "cl_khr_fp16" ∉ backend_extensions
+        res = setdiff(res, [Float16])
+    end
+
+    return res
+end
+
 ## Indexing Functions
 
 @device_override @inline function KI.get_local_id()
@@ -257,6 +288,16 @@ end
     return (; x = Int(get_global_size(1)), y = Int(get_global_size(2)), z = Int(get_global_size(3)))
 end
 
+@device_override KI.get_sub_group_size() = get_sub_group_size()
+
+@device_override KI.get_max_sub_group_size() = get_max_sub_group_size()
+
+@device_override KI.get_num_sub_groups() = get_num_sub_groups()
+
+@device_override KI.get_sub_group_id() = get_sub_group_id()
+
+@device_override KI.get_sub_group_local_id() = get_sub_group_local_id()
+
 @device_override @inline function KA.__validindex(ctx)
     if KA.__dynamic_checkbounds(ctx)
         I = @inbounds KA.expand(KA.__iterspace(ctx), get_group_id(1), get_local_id(1))
@@ -285,6 +326,14 @@ end
     work_group_barrier(POCL.LOCAL_MEM_FENCE | POCL.GLOBAL_MEM_FENCE)
 end
 
+@device_override @inline function KI.sub_group_barrier()
+    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)
+end
+
 @device_override @inline function KI._print(args...)
     POCL._print(args...)
 end

src/pocl/compiler/compilation.jl

@@ -1,6 +1,10 @@
 ## gpucompiler interface
 
-struct OpenCLCompilerParams <: AbstractCompilerParams end
+Base.@kwdef struct OpenCLCompilerParams <: AbstractCompilerParams
+    # request a fixed sub-group width via `intel_reqd_sub_group_size`
+    sub_group_size::Union{Nothing, Int} = nothing
+end
+
 const OpenCLCompilerConfig = CompilerConfig{SPIRVCompilerTarget, OpenCLCompilerParams}
 const OpenCLCompilerJob = CompilerJob{SPIRVCompilerTarget, OpenCLCompilerParams}
 
@@ -31,6 +35,11 @@ function GPUCompiler.finish_module!(
         job, mod, entry
     )
 
+    sg_size = job.config.params.sub_group_size
+    if sg_size !== nothing
+        metadata(entry)["intel_reqd_sub_group_size"] = MDNode([ConstantInt(Int32(sg_size))])
+    end
+
     # if this kernel uses our RNG, we should prime the shared state.
     # XXX: these transformations should really happen at the Julia IR level...
     if haskey(functions(mod), "julia.opencl.random_keys") && job.config.kernel
@@ -130,14 +139,17 @@ function compiler_config(dev::cl.Device; kwargs...)
     end
     return config
 end
-@noinline function _compiler_config(dev; kernel = true, name = nothing, always_inline = false, kwargs...)
+@noinline function _compiler_config(dev; kernel = true, name = nothing, always_inline = false, sub_group_size::Union{Nothing, Int} = 32, kwargs...)
     supports_fp16 = "cl_khr_fp16" in dev.extensions
     supports_fp64 = "cl_khr_fp64" in dev.extensions
 
+    if sub_group_size !== nothing && sub_group_size ∉ dev.sub_group_sizes
+        error("$sub_group_size is not a valid sub-group size for this device.")
+    end
 
     # create GPUCompiler objects
     target = SPIRVCompilerTarget(; supports_fp16, supports_fp64, validate = true, kwargs...)
-    params = OpenCLCompilerParams()
+    params = OpenCLCompilerParams(; sub_group_size)
     return CompilerConfig(target, params; kernel, name, always_inline)
 end
 

src/pocl/compiler/execution.jl

@@ -4,7 +4,7 @@ export @opencl, clfunction, clconvert
 ## high-level @opencl interface
 
 const MACRO_KWARGS = [:launch]
-const COMPILER_KWARGS = [:kernel, :name, :always_inline, :validate]
+const COMPILER_KWARGS = [:kernel, :name, :always_inline, :validate, :sub_group_size]
 const LAUNCH_KWARGS = [:global_size, :local_size, :queue]
 
 macro opencl(ex...)

src/pocl/nanoOpenCL.jl

@@ -389,6 +389,8 @@ const CL_KERNEL_EXEC_INFO_SVM_PTRS = 0x11b6
 
 const CL_KERNEL_EXEC_INFO_SVM_FINE_GRAIN_SYSTEM = 0x11b7
 
+const CL_DEVICE_SUB_GROUP_SIZES_INTEL = 0x4108
+
 struct CLError <: Exception
     code::Cint
 end
@@ -951,6 +953,14 @@ devices(p::Platform) = devices(p, CL_DEVICE_TYPE_ALL)
         return tuple([Int(r) for r in result]...)
     end
 
+    if s == :sub_group_sizes
+        res_size = Ref{Csize_t}()
+        clGetDeviceInfo(d, CL_DEVICE_SUB_GROUP_SIZES_INTEL, C_NULL, C_NULL, res_size)
+        result = Vector{Csize_t}(undef, res_size[] ÷ sizeof(Csize_t))
+        clGetDeviceInfo(d, CL_DEVICE_SUB_GROUP_SIZES_INTEL, sizeof(result), result, C_NULL)
+        return tuple([Int(r) for r in result]...)
+    end
+
     if s == :max_image2d_shape
         width = Ref{Csize_t}()
         height = Ref{Csize_t}()

src/pocl/pocl.jl

@@ -42,6 +42,7 @@ end
 
 using GPUCompiler
 using LLVM, LLVM.Interop
+import LLVM: LLVM, MDNode, ConstantInt, metadata
 using SPIRV_LLVM_Backend_jll, SPIRV_Tools_jll
 using Adapt