diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 6e229c6ea..7e2269723 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -16,7 +16,7 @@ steps: julia -e 'println("--- :julia: Developing CUDA") using Pkg url="https://github.com/christiangnrd/CUDA.jl" - rev="intrinsics" + rev="intrinsicsnew" subdirs = ["CUDACore", "CUDATools", "lib/cusolver", "lib/curand", "lib/cublas", "lib/cufft", "lib/cusparse", "lib/cupti", "lib/nvml", "lib/custatevec", "lib/cudnn", "lib/cutensor", "lib/cutensornet"] Pkg.add([(;url, rev); [(; url, rev, subdir) for subdir in subdirs]])' julia -e 'println("--- :julia: Instantiating project") @@ -81,7 +81,7 @@ steps: command: | julia -e 'println("--- :julia: Developing Metal") using Pkg - Pkg.add(url="https://github.com/JuliaGPU/Metal.jl", rev="kaintr")' + Pkg.add(url="https://github.com/JuliaGPU/Metal.jl", rev="kaintrnew")' julia -e 'println("--- :julia: Instantiating project") using Pkg Pkg.develop(; path=pwd())' || exit 3 @@ -113,8 +113,8 @@ steps: command: | julia -e 'println("--- :julia: Developing oneAPI") using Pkg - Pkg.add(url="https://github.com/christiangnrd/oneAPI.jl", rev="intrinsics") - Pkg.add(url="https://github.com/christiangnrd/AcceleratedKernels.jl", rev="ka0.10simple")' + Pkg.add(url="https://github.com/christiangnrd/oneAPI.jl", rev="intrinsicsnew") + Pkg.add(url="https://github.com/JuliaGPU/AcceleratedKernels.jl", rev="main")' julia -e 'println("--- :julia: Instantiating project") using Pkg Pkg.develop(; path=pwd())' || exit 3 @@ -146,10 +146,10 @@ steps: command: | julia -e 'println("--- :julia: Developing AMDGPU") using Pkg - Pkg.add(url="https://github.com/christiangnrd/AcceleratedKernels.jl", rev="ka0.10simple")' + Pkg.add(url="https://github.com/JuliaGPU/AcceleratedKernels.jl", rev="main")' julia -e ' using Pkg - Pkg.add(url="https://github.com/christiangnrd/AMDGPU.jl", rev="intrinsics") + Pkg.add(url="https://github.com/christiangnrd/AMDGPU.jl", rev="intrinsicsnew") println("--- :julia: Instantiating project") Pkg.develop(; path=pwd())' || exit 3 @@ -181,7 +181,7 @@ steps: command: | julia -e 'println("--- :julia: Developing OpenCL") using Pkg - Pkg.add(url="https://github.com/christiangnrd/OpenCL.jl", rev="intrinsics") + Pkg.add(url="https://github.com/christiangnrd/OpenCL.jl", rev="intrinsicsnew") Pkg.develop(; name="SPIRVIntrinsics")' julia -e 'println("--- :julia: Instantiating project") using Pkg diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 64c294c97..229d78001 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -182,7 +182,7 @@ jobs: run: | julia -e 'println("--- :julia: Developing OpenCL") using Pkg - Pkg.add(url="https://github.com/christiangnrd/OpenCL.jl", rev="intrinsics") + Pkg.add(url="https://github.com/christiangnrd/OpenCL.jl", rev="intrinsicsnew") Pkg.develop(; name="SPIRVIntrinsics")' - name: "Instantiating project" run: | diff --git a/src/interface.jl b/src/interface.jl index f2c21d2cc..6748c227a 100644 --- a/src/interface.jl +++ b/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 diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index 3b0fdb3fa..a72eae309 100644 --- a/src/pocl/backend.jl +++ b/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 diff --git a/src/pocl/compiler/compilation.jl b/src/pocl/compiler/compilation.jl index 8aee85447..e5eefd31c 100644 --- a/src/pocl/compiler/compilation.jl +++ b/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 diff --git a/src/pocl/compiler/execution.jl b/src/pocl/compiler/execution.jl index 595e7ff1a..fd22b57c8 100644 --- a/src/pocl/compiler/execution.jl +++ b/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...) diff --git a/src/pocl/nanoOpenCL.jl b/src/pocl/nanoOpenCL.jl index e7901ad07..eec955990 100644 --- a/src/pocl/nanoOpenCL.jl +++ b/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}() diff --git a/src/pocl/pocl.jl b/src/pocl/pocl.jl index de1bd4b32..681683e70 100644 --- a/src/pocl/pocl.jl +++ b/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 diff --git a/test/interface.jl b/test/interface.jl index a6fdd05bc..416ddcd1e 100644 --- a/test/interface.jl +++ b/test/interface.jl @@ -23,6 +23,65 @@ function test_interface_kernel(results) end return end +struct SubgroupData + sub_group_size::UInt32 + max_sub_group_size::UInt32 + num_sub_groups::UInt32 + sub_group_id::UInt32 + sub_group_local_id::UInt32 +end +function test_subgroup_kernel(results) + i = KI.get_global_id().x + + if i <= length(results) + @inbounds results[i] = SubgroupData( + KI.get_sub_group_size(), + KI.get_max_sub_group_size(), + KI.get_num_sub_groups(), + KI.get_sub_group_id(), + KI.get_sub_group_local_id() + ) + end + return +end + +function shfl_down_test_kernel(a, b, ::Val{N}) where {N} + idx = KI.get_sub_group_local_id() + + val = a[idx] + + offset = 0x00000001 + while offset < N + val += KI.shfl_down(val, offset) + offset <<= 1 + end + + KI.sub_group_barrier() + + if idx == 1 + b[idx] = val + end + return +end + +function shfl_down_test_kernel(a, b, ::Val{N}) where {N} + idx = KI.get_sub_group_local_id() + + val = a[idx] + + offset = 0x00000001 + while offset < N + val += KI.shfl_down(val, offset) + offset <<= 1 + end + + KI.sub_group_barrier() + + if idx == 1 + b[idx] = val + end + return +end function interface_testsuite(backend, AT) @testset "KernelInterface Tests" begin @@ -122,6 +181,60 @@ function interface_testsuite(backend, AT) @test k_data.local_id == expected_local end end - end + + # Used as a proxy for sub-group support + if !isempty(KI.shfl_down_types(backend())) + @testset "Sub-groups" begin + @test KI.sub_group_size(backend()) isa Int + + # Test with small kernel + sg_size = KI.sub_group_size(backend()) + sg_n = 2 + workgroupsize = sg_size * sg_n + numworkgroups = 2 + N = workgroupsize * numworkgroups + + results = AT(Vector{SubgroupData}(undef, N)) + kernel = KI.@kernel backend() launch = false test_subgroup_kernel(results) + + kernel(results; workgroupsize, numworkgroups) + KernelAbstractions.synchronize(backend()) + + host_results = Array(results) + + # Verify results make sense + for (i, sg_data) in enumerate(host_results) + @test sg_data.sub_group_size == sg_size + @test sg_data.max_sub_group_size == sg_size + @test sg_data.num_sub_groups == sg_n + + # Group ID should be 1-based + expected_sub_group = div(((i - 1) % workgroupsize), sg_size) + 1 + @test sg_data.sub_group_id == expected_sub_group + + # Local ID should be 1-based within group + expected_sg_local = ((i - 1) % sg_size) + 1 + @test sg_data.sub_group_local_id == expected_sg_local + end + end + @testset "shfl_down" begin + @test !isempty(KI.shfl_down_types(backend())) + types_to_test = setdiff(KI.shfl_down_types(backend()), [Bool]) + @testset "$T" for T in types_to_test + N = KI.sub_group_size(backend()) + a = zeros(T, N) + rand!(a, (0:1)) + + dev_a = AT(a) + dev_b = AT(zeros(T, N)) + + KI.@kernel backend() workgroupsize = N shfl_down_test_kernel(dev_a, dev_b, Val(N)) + + b = Array(dev_b) + @test sum(a) ≈ b[1] + end + end + end + end # @testset "KernelInterface Tests" begin return nothing end