From e2e2cf8d3a3eb4bd129362deab0d12f96f1a5342 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 09:12:21 +0200 Subject: [PATCH 1/9] KernelInterface: remove the GPU backend supertype `GPU` didn't mean GPU hardware: the CPU backend is PoCL and subtyped it, so `CPU <: GPU`. Backends now subtype `Backend` directly. This is the first breaking change of KernelInterface 0.3. --- Project.toml | 2 +- docs/src/api.md | 1 - docs/src/implementations.md | 2 +- docs/src/index.md | 4 ++++ docs/src/kernelinterface.md | 12 +++++------- docs/src/quickstart.md | 2 +- ext/EnzymeCore07Ext.jl | 12 ++++++------ ext/EnzymeCore08Ext.jl | 12 ++++++------ ext/EnzymeExt.jl | 1 - lib/KernelInterface/Project.toml | 2 +- lib/KernelInterface/src/backend.jl | 13 +------------ lib/KernelInterface/test/runtests.jl | 2 -- src/KernelAbstractions.jl | 8 ++++---- src/pocl/backend.jl | 2 +- test/runtests.jl | 2 +- 15 files changed, 32 insertions(+), 45 deletions(-) diff --git a/Project.toml b/Project.toml index 4d02e6782..0cf5fc004 100644 --- a/Project.toml +++ b/Project.toml @@ -42,7 +42,7 @@ Adapt = "0.4, 1.0, 2.0, 3.0, 4" Atomix = "1.2.1" EnzymeCore = "0.7, 0.8.1" GPUCompiler = "2.7" -KernelInterface = "0.2.3" +KernelInterface = "0.3" LLVM = "9.9" LinearAlgebra = "1.6" MacroTools = "0.5" diff --git a/docs/src/api.md b/docs/src/api.md index eba5e05db..bb76872fb 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -33,7 +33,6 @@ ```@docs Backend -GPU CPU POCLBackend get_backend diff --git a/docs/src/implementations.md b/docs/src/implementations.md index b0002a3ed..adc1b132a 100644 --- a/docs/src/implementations.md +++ b/docs/src/implementations.md @@ -1,6 +1,6 @@ # [Notes for backend implementations](@id implementations_notes) -The [KernelInterface](@ref kernelinterface) sibling package defines the core interface a backend must implement. A backend must implement a backend type that subtypes `KernelInterface.GPU`, or `KernelInterface.Backend` for non-gpu backends. This documentation contains the host and devices side functions that backends can define, as well as whether they are mandatory or not. +The [KernelInterface](@ref kernelinterface) sibling package defines the core interface a backend must implement. A backend must implement a backend type that subtypes `KernelInterface.Backend`. This documentation contains the host and devices side functions that backends can define, as well as whether they are mandatory or not. ## Semantics of `KernelAbstractions.synchronize` diff --git a/docs/src/index.md b/docs/src/index.md index 9a9dc067c..f485a76a7 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -141,6 +141,10 @@ end the CPU backend became an OpenCL backend meant every backend. Use `KernelAbstractions.@device_code_llvm` and `KernelAbstractions.@device_code_typed` instead, which report on the code a backend actually generates; see [Reflection](@ref). +- `KernelAbstractions.GPU` has been removed: it didn't mean GPU hardware (the `CPU` backend + was a subtype), and backends now subtype `KernelAbstractions.Backend` directly. Code + dispatching on `::GPU` should dispatch on `::Backend`, on concrete backend types, or on a + capability such as `KernelAbstractions.supports_float64`. ## Semantic differences diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md index 0c09aa088..c44a2e903 100644 --- a/docs/src/kernelinterface.md +++ b/docs/src/kernelinterface.md @@ -46,15 +46,13 @@ KernelInterface ## Backend hierarchy -A backend package subtypes [`GPU`](@ref) (or [`Backend`](@ref) directly for -non-GPU backends), and everything else in the interface dispatches on that -type. These types and the host-side management functions below are re-exported -by `KernelAbstractions`, so their canonical docstrings are on the +Backends subtype [`Backend`](@ref), and everything else in the interface dispatches on +that type. It and the host-side management functions below are re-exported by +`KernelAbstractions`, so their canonical docstrings are on the [API page](@ref api_backends_arrays). ```@docs; canonical=false Backend -GPU get_backend ``` @@ -198,8 +196,8 @@ KernelInterface.@kernel A backend must, at minimum: -1. Define a backend type subtyping [`GPU`](@ref) (or [`Backend`](@ref) for - non-GPU backends), and implement [`get_backend`](@ref) for its array type. +1. Define a backend type subtyping [`Backend`](@ref), and implement [`get_backend`](@ref) + for its array type. 2. Implement the host-side management functions for that type: [`allocate`](@ref), [`copyto!`](@ref), [`synchronize`](@ref) and [`unsafe_free!`](@ref) are required; the remaining functions under diff --git a/docs/src/quickstart.md b/docs/src/quickstart.md index 759f74a0a..f6458eebe 100644 --- a/docs/src/quickstart.md +++ b/docs/src/quickstart.md @@ -63,7 +63,7 @@ See also [Memcopy with static NDRange](@ref memcopy_static). ## Launching kernel on the backend -To launch the kernel on a backend-supported backend `isa(backend, KA.GPU)` (e.g., `CUDABackend()`, `ROCBackend()`, `oneAPIBackend()`, `MetalBackend()`), we generate the kernel +To launch the kernel on a backend (e.g., `CUDABackend()`, `ROCBackend()`, `oneAPIBackend()`, `MetalBackend()`), we generate the kernel for this backend. First, we initialize the array using the Array constructor of the chosen backend with diff --git a/ext/EnzymeCore07Ext.jl b/ext/EnzymeCore07Ext.jl index 93159886c..e286cf5e1 100644 --- a/ext/EnzymeCore07Ext.jl +++ b/ext/EnzymeCore07Ext.jl @@ -26,7 +26,7 @@ function EnzymeRules.forward( end function EnzymeRules.forward( - func::Const{<:Kernel{<:GPU}}, + func::Const{<:Kernel{<:Backend}}, ::Type{Const{Nothing}}, args...; ndrange = nothing, @@ -41,7 +41,7 @@ end _enzyme_mkcontext(kernel::Kernel{CPU}, ndrange, iterspace, dynamic) = mkcontext(kernel, first(blocks(iterspace)), ndrange, iterspace, dynamic) -_enzyme_mkcontext(kernel::Kernel{<:GPU}, ndrange, iterspace, dynamic) = +_enzyme_mkcontext(kernel::Kernel{<:Backend}, ndrange, iterspace, dynamic) = mkcontext(kernel, ndrange, iterspace) _augmented_return(::Kernel{CPU}, subtape, arg_refs, tape_type) = @@ -50,7 +50,7 @@ _augmented_return(::Kernel{CPU}, subtape, arg_refs, tape_type) = nothing, (subtape, arg_refs, tape_type), ) -_augmented_return(::Kernel{<:GPU}, subtape, arg_refs, tape_type) = +_augmented_return(::Kernel{<:Backend}, subtape, arg_refs, tape_type) = AugmentedReturn{Nothing, Nothing, Any}(nothing, nothing, (subtape, arg_refs, tape_type)) function _create_tape_kernel( @@ -75,7 +75,7 @@ function _create_tape_kernel( end function _create_tape_kernel( - kernel::Kernel{<:GPU}, + kernel::Kernel{<:Backend}, ModifiedBetween, FT, ctxTy, @@ -109,7 +109,7 @@ function _create_tape_kernel( end _create_rev_kernel(kernel::Kernel{CPU}) = similar(kernel, cpu_rev) -_create_rev_kernel(kernel::Kernel{<:GPU}) = similar(kernel, gpu_rev) +_create_rev_kernel(kernel::Kernel{<:Backend}) = similar(kernel, gpu_rev) function cpu_aug_fwd( ctx, @@ -232,7 +232,7 @@ function EnzymeRules.augmented_primal( arg_refs = ntuple(Val(N)) do i Base.@_inline_meta if args[i] isa Active - if func.val isa Kernel{<:GPU} + if func.val isa Kernel{<:Backend} error("Active kernel arguments not supported on GPU") else Ref(EnzymeCore.make_zero(args[i].val)) diff --git a/ext/EnzymeCore08Ext.jl b/ext/EnzymeCore08Ext.jl index 1fda85120..467f073cc 100644 --- a/ext/EnzymeCore08Ext.jl +++ b/ext/EnzymeCore08Ext.jl @@ -28,7 +28,7 @@ end function EnzymeRules.forward( config, - func::Const{<:Kernel{<:GPU}}, + func::Const{<:Kernel{<:Backend}}, ::Type{Const{Nothing}}, args...; ndrange = nothing, @@ -43,7 +43,7 @@ end _enzyme_mkcontext(kernel::Kernel{CPU}, ndrange, iterspace, dynamic) = mkcontext(kernel, first(blocks(iterspace)), ndrange, iterspace, dynamic) -_enzyme_mkcontext(kernel::Kernel{<:GPU}, ndrange, iterspace, dynamic) = +_enzyme_mkcontext(kernel::Kernel{<:Backend}, ndrange, iterspace, dynamic) = mkcontext(kernel, ndrange, iterspace) _augmented_return(::Kernel{CPU}, subtape, arg_refs, tape_type) = @@ -52,7 +52,7 @@ _augmented_return(::Kernel{CPU}, subtape, arg_refs, tape_type) = nothing, (subtape, arg_refs, tape_type), ) -_augmented_return(::Kernel{<:GPU}, subtape, arg_refs, tape_type) = +_augmented_return(::Kernel{<:Backend}, subtape, arg_refs, tape_type) = AugmentedReturn{Nothing, Nothing, Any}(nothing, nothing, (subtape, arg_refs, tape_type)) function _create_tape_kernel( @@ -77,7 +77,7 @@ function _create_tape_kernel( end function _create_tape_kernel( - kernel::Kernel{<:GPU}, + kernel::Kernel{<:Backend}, Mode, FT, ctxTy, @@ -111,7 +111,7 @@ function _create_tape_kernel( end _create_rev_kernel(kernel::Kernel{CPU}) = similar(kernel, cpu_rev) -_create_rev_kernel(kernel::Kernel{<:GPU}) = similar(kernel, gpu_rev) +_create_rev_kernel(kernel::Kernel{<:Backend}) = similar(kernel, gpu_rev) function cpu_aug_fwd( ctx, @@ -234,7 +234,7 @@ function EnzymeRules.augmented_primal( arg_refs = ntuple(Val(N)) do i Base.@_inline_meta if args[i] isa Active - if func.val isa Kernel{<:GPU} + if func.val isa Kernel{<:Backend} error("Active kernel arguments not supported on GPU") else Ref(EnzymeCore.make_zero(args[i].val)) diff --git a/ext/EnzymeExt.jl b/ext/EnzymeExt.jl index 1fda4e051..c707d18e8 100644 --- a/ext/EnzymeExt.jl +++ b/ext/EnzymeExt.jl @@ -16,7 +16,6 @@ import KernelAbstractions: mkcontext, CompilerMetadata, CPU, - GPU, argconvert, supports_enzyme, __fake_compiler_job, diff --git a/lib/KernelInterface/Project.toml b/lib/KernelInterface/Project.toml index e1660a954..e34a91b0e 100644 --- a/lib/KernelInterface/Project.toml +++ b/lib/KernelInterface/Project.toml @@ -1,7 +1,7 @@ name = "KernelInterface" uuid = "4ee993da-d684-4d17-a7dd-4e58e78d92bf" authors = ["Valentin Churavy and contributors"] -version = "0.2.3" +version = "0.3.0" [compat] julia = "1.10" diff --git a/lib/KernelInterface/src/backend.jl b/lib/KernelInterface/src/backend.jl index 409fd0522..ad1433b46 100644 --- a/lib/KernelInterface/src/backend.jl +++ b/lib/KernelInterface/src/backend.jl @@ -6,7 +6,7 @@ """ Backend -Abstract supertype for all KernelAbstractions backends. +Abstract supertype for all KernelInterface backends. Backends subtype it directly. Concrete backends (for example `CUDABackend` from CUDA.jl or `CPU` from KernelAbstractions) determine where arrays are allocated and where kernels execute. Use [`get_backend`](@ref) to @@ -23,17 +23,6 @@ synchronize(backend) """ abstract type Backend end -""" -Abstract type for all GPU based KernelAbstractions backends. - -!!! note - New backend implementations **must** sub-type this abstract type. - -!!! note - `GPU` will be removed in KernelAbstractions v1.0 -""" -abstract type GPU <: Backend end - """ get_backend(A::AbstractArray)::Backend diff --git a/lib/KernelInterface/test/runtests.jl b/lib/KernelInterface/test/runtests.jl index 944e5648c..1da2b4da4 100644 --- a/lib/KernelInterface/test/runtests.jl +++ b/lib/KernelInterface/test/runtests.jl @@ -92,8 +92,6 @@ end end @testset "get_backend" begin - @test KI.GPU <: KI.Backend - # The fallback finds the backend of wrapper arrays by walking `parent`. arr = BackedArray([1, 2, 3]) @test KI.get_backend(arr) === StubBackend() diff --git a/src/KernelAbstractions.jl b/src/KernelAbstractions.jl index e5186cfc0..0dacde9a3 100644 --- a/src/KernelAbstractions.jl +++ b/src/KernelAbstractions.jl @@ -4,7 +4,7 @@ export @kernel export @Const, @localmem, @private, @uniform, @synchronize export @index, @groupsize, @ndrange export @print -export Backend, GPU, CPU +export Backend, CPU export synchronize, get_backend, allocate import PrecompileTools @@ -14,7 +14,7 @@ import Atomix: @atomic, @atomicswap, @atomicreplace using MacroTools using Adapt -using KernelInterface: KernelInterface, Backend, GPU, get_backend, functional, synchronize, versioninfo, supports_unified, supports_float64, supports_atomics, copyto!, allocate, zeros, ones, device, device!, ndevices, priority!, pagelock!, unsafe_free!, record_event, wait_event +using KernelInterface: KernelInterface, Backend, get_backend, functional, synchronize, versioninfo, supports_unified, supports_float64, supports_atomics, copyto!, allocate, zeros, ones, device, device!, ndevices, priority!, pagelock!, unsafe_free!, record_event, wait_event import KernelInterface as KI export KernelInterface @@ -597,8 +597,8 @@ last (possibly partial) workgroup. Primarily used by backend implementations and return iterspace, dynamic end -function construct(backend::Backend, ::S, ::NDRange, xpu_name::XPUName) where {Backend <: GPU, S <: _Size, NDRange <: _Size, XPUName} - return Kernel{Backend, S, NDRange, XPUName}(backend, xpu_name) +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) end ### diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index 9e9c4177f..e51a0795d 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -18,7 +18,7 @@ import Adapt export POCLBackend -struct POCLBackend <: KA.GPU +struct POCLBackend <: KI.Backend end function KI.versioninfo(io::IO, ::POCLBackend) diff --git a/test/runtests.jl b/test/runtests.jl index 728feff1c..70c8a3778 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -152,7 +152,7 @@ end Testsuite.testsuite(CPU, "CPU", POCL, Array, POCL.CLDeviceArray) end -struct NewBackend <: KernelAbstractions.GPU end +struct NewBackend <: KernelAbstractions.Backend end @testset "Default host implementation" begin backend = NewBackend() From d39f7cb53e7b34e7942936c4de9f5645a306af1b Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 09:16:07 +0200 Subject: [PATCH 2/9] KernelInterface testsuite: take a backend value and its array type `testsuite(backend, AT)` replaces `testsuite(Backend, name, mod, AT, DAT)`, whose name, module and device array type were unused, and tests the backend value it is given, e.g. one with non-default options. The events tests now skip backends without events of their own (whose `record_event` returns `nothing`), instead of them having to pass `skip_tests`. --- lib/KernelInterface/test/events.jl | 9 ++- lib/KernelInterface/test/interface.jl | 86 +++++++++++++-------------- lib/KernelInterface/test/testsuite.jl | 8 ++- test/runtests.jl | 5 +- 4 files changed, 58 insertions(+), 50 deletions(-) diff --git a/lib/KernelInterface/test/events.jl b/lib/KernelInterface/test/events.jl index 6ff22e089..5f3f8da05 100644 --- a/lib/KernelInterface/test/events.jl +++ b/lib/KernelInterface/test/events.jl @@ -19,8 +19,13 @@ function slow_fill_kernel(A, v, iters::UInt32) return end -function events_testsuite(backend) - b = backend() +function events_testsuite(b::KI.Backend) + # A backend without events of its own records by synchronizing, so there is nothing + # for these tests to order. + if KI.record_event(b) === nothing + @test KI.wait_event(b, nothing) === nothing + return + end dev = KI.device(b) N = 64 diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl index 30172bf72..b1888bb3f 100644 --- a/lib/KernelInterface/test/interface.jl +++ b/lib/KernelInterface/test/interface.jl @@ -131,7 +131,7 @@ function shfl_down_test_kernel(a, b, ::Val{N}) where {N} return end -function interface_testsuite(backend, AT) +function interface_testsuite(backend::KI.Backend, AT) @testset "Launch parameters" begin # 1d function launch_kernel1d(arr) @@ -143,14 +143,14 @@ function interface_testsuite(backend, AT) return end arr1d = AT(zeros(Float32, 4)) - KI.@kernel backend() numworkgroups = 2 workgroupsize = 2 launch_kernel1d(arr1d) - KI.synchronize(backend()) + KI.@kernel backend numworkgroups = 2 workgroupsize = 2 launch_kernel1d(arr1d) + KI.synchronize(backend) @test all(Array(arr1d) .== 1) # 1d tuple arr1dt = AT(zeros(Float32, 4)) - KI.@kernel backend() numworkgroups = (2,) workgroupsize = (2,) launch_kernel1d(arr1dt) - KI.synchronize(backend()) + KI.@kernel backend numworkgroups = (2,) workgroupsize = (2,) launch_kernel1d(arr1dt) + KI.synchronize(backend) @test all(Array(arr1dt) .== 1) # 2d @@ -163,8 +163,8 @@ function interface_testsuite(backend, AT) return end arr2d = AT(zeros(Float32, 4, 4)) - KI.@kernel backend() numworkgroups = (2, 2) workgroupsize = (2, 2) launch_kernel2d(arr2d) - KI.synchronize(backend()) + KI.@kernel backend numworkgroups = (2, 2) workgroupsize = (2, 2) launch_kernel2d(arr2d) + KI.synchronize(backend) @test all(Array(arr2d) .== 1) # 3d @@ -177,18 +177,18 @@ function interface_testsuite(backend, AT) return end arr3d = AT(zeros(Float32, 4, 4, 4)) - KI.@kernel backend() numworkgroups = (2, 2, 2) workgroupsize = (2, 2, 2) launch_kernel3d(arr3d) - KI.synchronize(backend()) + KI.@kernel backend numworkgroups = (2, 2, 2) workgroupsize = (2, 2, 2) launch_kernel3d(arr3d) + KI.synchronize(backend) @test all(Array(arr3d) .== 1) # 4d (Errors) - @test_throws ArgumentError (KI.@kernel backend() numworkgroups = (2, 2, 2, 2) workgroupsize = (2, 2, 2) launch_kernel3d(arr3d)) - @test_throws ArgumentError (KI.@kernel backend() numworkgroups = (2, 2, 2) workgroupsize = (2, 2, 2, 2) launch_kernel3d(arr3d)) + @test_throws ArgumentError (KI.@kernel backend numworkgroups = (2, 2, 2, 2) workgroupsize = (2, 2, 2) launch_kernel3d(arr3d)) + @test_throws ArgumentError (KI.@kernel backend numworkgroups = (2, 2, 2) workgroupsize = (2, 2, 2, 2) launch_kernel3d(arr3d)) end @testset "Launch limits" begin - max_dims = KI.max_work_group_dims(backend()) - max_groups = KI.max_num_groups(backend()) + max_dims = KI.max_work_group_dims(backend) + max_groups = KI.max_num_groups(backend) @test max_dims isa NTuple{3, Int} && all(>=(1), max_dims) @test max_groups isa NTuple{3, Int} && all(>=(1), max_groups) @@ -199,11 +199,11 @@ function interface_testsuite(backend, AT) end return end - kernel = KI.@kernel backend() launch = false fill_kernel(AT(zeros(Float32, 1, 1, 1))) + kernel = KI.@kernel backend launch = false fill_kernel(AT(zeros(Float32, 1, 1, 1))) function fill_test(dims; kwargs...) arr = AT(zeros(Float32, dims)) kernel(arr; kwargs...) - KI.synchronize(backend()) + KI.synchronize(backend) return all(Array(arr) .== 1) end @@ -227,7 +227,7 @@ function interface_testsuite(backend, AT) end @testset "Host return types" begin - b = backend() + b = backend @test KI.supports_unified(b) isa Bool @test KI.supports_atomics(b) isa Bool @@ -250,15 +250,15 @@ function interface_testsuite(backend, AT) end @testset "Device return types" begin - results = KI.zeros(backend(), Bool, 6) - KI.@kernel backend() typecheck_kernel(results) - KI.synchronize(backend()) + results = KI.zeros(backend, Bool, 6) + KI.@kernel backend typecheck_kernel(results) + KI.synchronize(backend) @test all(Array(results)) @testset "$T" for T in (Int32, Int64, UInt32, UInt64) - typed_results = KI.zeros(backend(), Bool, 6) - KI.@kernel backend() typed_typecheck_kernel(typed_results, T) - KI.synchronize(backend()) + typed_results = KI.zeros(backend, Bool, 6) + KI.@kernel backend typed_typecheck_kernel(typed_results, T) + KI.synchronize(backend) @test all(Array(typed_results)) end end @@ -270,9 +270,9 @@ function interface_testsuite(backend, AT) # `Int` is the reference: it is what the zero-argument form returns. function run_typed(::Type{T}) where {T} - results = KI.zeros(backend(), T, N, 18) - KI.@kernel backend() workgroupsize = workgroupsize numworkgroups = numworkgroups typed_index_kernel(results, T) - KI.synchronize(backend()) + results = KI.zeros(backend, T, N, 18) + KI.@kernel backend workgroupsize = workgroupsize numworkgroups = numworkgroups typed_index_kernel(results, T) + KI.synchronize(backend) return Array(results) end reference = run_typed(Int) @@ -293,21 +293,21 @@ function interface_testsuite(backend, AT) @testset "Basic interface functionality" begin - @test KI.max_work_group_size(backend()) isa Int - @test KI.multiprocessor_count(backend()) isa Int + @test KI.max_work_group_size(backend) isa Int + @test KI.multiprocessor_count(backend) isa Int # Test with small kernel workgroupsize = 4 numworkgroups = 4 N = workgroupsize * numworkgroups results = AT(Vector{KernelData}(undef, N)) - kernel = KI.@kernel backend() launch = false test_interface_kernel(results) + kernel = KI.@kernel backend launch = false test_interface_kernel(results) @test KI.kernel_max_work_group_size(kernel) isa Int @test KI.kernel_max_work_group_size(kernel; max_work_items = 1) == 1 kernel(results; workgroupsize, numworkgroups) - KI.synchronize(backend()) + KI.synchronize(backend) host_results = Array(results) @@ -335,32 +335,32 @@ function interface_testsuite(backend, AT) end # Used as a proxy for sub-group support - if !isempty(KI.shfl_down_types(backend())) + if !isempty(KI.shfl_down_types(backend)) @testset "Sub-group return types" begin - @test KI.sub_group_size(backend()) isa Int + @test KI.sub_group_size(backend) isa Int - T = first(setdiff(KI.shfl_down_types(backend()), [Bool])) - results = KI.zeros(backend(), Bool, 6) - KI.@kernel backend() workgroupsize = KI.sub_group_size(backend()) subgroup_typecheck_kernel(results, one(T)) - KI.synchronize(backend()) + T = first(setdiff(KI.shfl_down_types(backend), [Bool])) + results = KI.zeros(backend, Bool, 6) + KI.@kernel backend workgroupsize = KI.sub_group_size(backend) subgroup_typecheck_kernel(results, one(T)) + KI.synchronize(backend) @test all(Array(results)) end @testset "Sub-groups" begin - @test KI.sub_group_size(backend()) isa Int + @test KI.sub_group_size(backend) isa Int # Test with small kernel - sg_size = KI.sub_group_size(backend()) + 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 = KI.@kernel backend launch = false test_subgroup_kernel(results) kernel(results; workgroupsize, numworkgroups) - KI.synchronize(backend()) + KI.synchronize(backend) host_results = Array(results) @@ -380,17 +380,17 @@ function interface_testsuite(backend, AT) end end @testset "shfl_down" begin - @test !isempty(KI.shfl_down_types(backend())) - types_to_test = setdiff(KI.shfl_down_types(backend()), [Bool]) + @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()) + 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)) + KI.@kernel backend workgroupsize = N shfl_down_test_kernel(dev_a, dev_b, Val(N)) b = Array(dev_b) @test sum(a) ≈ b[1] diff --git a/lib/KernelInterface/test/testsuite.jl b/lib/KernelInterface/test/testsuite.jl index 4457ef3df..fa0b598af 100644 --- a/lib/KernelInterface/test/testsuite.jl +++ b/lib/KernelInterface/test/testsuite.jl @@ -29,7 +29,13 @@ end include("interface.jl") include("events.jl") -function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{String}()) +""" + testsuite(backend::KI.Backend, AT; skip_tests=Set{String}()) + +Run the KernelInterface tests for `backend`, whose array type is `AT`. Test sets named in +`skip_tests` are skipped. +""" +function testsuite(backend::KI.Backend, AT; skip_tests = Set{String}()) @conditional_testset "Interface" skip_tests begin interface_testsuite(backend, AT) end diff --git a/test/runtests.jl b/test/runtests.jl index 70c8a3778..b0d499553 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -31,10 +31,7 @@ module KernelInterfaceTests include(joinpath(pkgdir(KernelInterface), "test", "testsuite.jl")) end @testset "POCL KernelInterface" begin - # POCL launches synchronously, so there's nothing for the events tests to order - KernelInterfaceTests.Testsuite.testsuite( - POCLBackend, "POCL", POCL, Array, POCL.CLDeviceArray; skip_tests = Set(["Events"]) - ) + KernelInterfaceTests.Testsuite.testsuite(POCLBackend(), Array) end @testset "POCL float atomics" begin From 18d33cb280085294ce49b0f3321f3aba0fc93c40 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 09:18:00 +0200 Subject: [PATCH 3/9] KernelInterface: rename KI.@kernel to KI.@launch, numworkgroups to numgroups `KI.@kernel` prefixes a call and launches it, like `@cuda` or `@metal`, while `KA.@kernel` defines a kernel; sharing the name was a source of confusion. `numgroups` matches `get_num_groups` and `max_num_groups`. `KI.@launch` now also evaluates its backend expression once (it was evaluated once per argument), and passes the keywords it doesn't know to `kernel_function` as compiler options instead of rejecting them. --- docs/src/kernelinterface.md | 11 +-- examples/histogram.jl | 2 +- examples/performant_matmul.jl | 4 +- lib/KernelInterface/src/launch.jl | 130 +++++++++++++------------- lib/KernelInterface/test/events.jl | 2 +- lib/KernelInterface/test/interface.jl | 58 ++++++------ lib/KernelInterface/test/runtests.jl | 62 +++++++----- src/pocl/backend.jl | 10 +- test/hostinterface.jl | 2 +- 9 files changed, 144 insertions(+), 137 deletions(-) diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md index c44a2e903..446b34ceb 100644 --- a/docs/src/kernelinterface.md +++ b/docs/src/kernelinterface.md @@ -182,16 +182,9 @@ Kernel kernel_function kernel_max_work_group_size argconvert -KernelInterface.@kernel +KernelInterface.@launch ``` -!!! note - `KI.@kernel` is **not** `KernelAbstractions.@kernel`. `KI.@kernel` wraps a - backend's own compile-and-launch path — the equivalent of `@cuda` or - `@metal` — and prefixes a *call*. [`KernelAbstractions.@kernel`](@ref) - prefixes a *definition* and produces a kernel written in the higher-level - KernelAbstractions language. - ## Implementing a backend A backend must, at minimum: @@ -212,7 +205,7 @@ A backend must, at minimum: [`shfl_down`](@ref) support is optional. 5. Implement [`argconvert`](@ref) and [`kernel_function`](@ref) for its backend type, returning a [`Kernel`](@ref). -6. Make that `Kernel` callable, accepting `numworkgroups`, `workgroupsize` and +6. Make that `Kernel` callable, accepting `numgroups`, `workgroupsize` and `ndrange` as a scalar `Integer` or a 1-, 2- or 3-element tuple. Use `KI.check_launch_args` to validate them, or check them directly. A zero-sized `ndrange` — launching over an empty array is not uncommon — must be a no-op diff --git a/examples/histogram.jl b/examples/histogram.jl index 15c241cab..3b7572e29 100644 --- a/examples/histogram.jl +++ b/examples/histogram.jl @@ -60,7 +60,7 @@ end function histogram!(histogram_output, input, groupsize = 256) backend = get_backend(histogram_output) # Need static block size - KI.@kernel backend workgroupsize = groupsize numworkgroups = cld(length(input), groupsize) histogram_kernel!(histogram_output, input, Val(groupsize)) + KI.@launch backend workgroupsize = groupsize numgroups = cld(length(input), groupsize) histogram_kernel!(histogram_output, input, Val(groupsize)) return end diff --git a/examples/performant_matmul.jl b/examples/performant_matmul.jl index f732c2abb..8ea431f17 100644 --- a/examples/performant_matmul.jl +++ b/examples/performant_matmul.jl @@ -84,9 +84,9 @@ B = copyto!(allocate(backend, Float32, R, M), rand(Float32, R, M)) C = KernelAbstractions.zeros(backend, Float32, N, M) workgroupsize = (TILE_DIM, TILE_DIM) -numworkgroups = (cld(size(C, 1), TILE_DIM), cld(size(C, 2), TILE_DIM)) +numgroups = (cld(size(C, 1), TILE_DIM), cld(size(C, 2), TILE_DIM)) -KI.@kernel backend workgroupsize numworkgroups coalesced_matmul_kernel!(C, A, B, N, R, M, Val(TILE_DIM)) +KI.@launch backend workgroupsize numgroups coalesced_matmul_kernel!(C, A, B, N, R, M, Val(TILE_DIM)) KernelAbstractions.synchronize(backend) @test isapprox(A * B, C) diff --git a/lib/KernelInterface/src/launch.jl b/lib/KernelInterface/src/launch.jl index a5126af13..925491a39 100644 --- a/lib/KernelInterface/src/launch.jl +++ b/lib/KernelInterface/src/launch.jl @@ -9,12 +9,12 @@ kernel on the host. !!! note Backend implementations **must** implement: ``` - (kernel::Kernel{<:NewBackend})(args...; numworkgroups=(), workgroupsize=(), ndrange=(), max_work_group_size=typemax(Int)) + (kernel::Kernel{<:NewBackend})(args...; numgroups=(), workgroupsize=(), ndrange=(), max_work_group_size=typemax(Int)) ``` - `numworkgroups`, `workgroupsize`, and `ndrange` must accept a scalar Integer, a 1, 2, + `numgroups`, `workgroupsize`, and `ndrange` must accept a scalar Integer, a 1, 2, or 3 Integer tuple, or an empty tuple. Otherwise, it must throw an `ArgumentError`. An - `ArgumentError` must also be thrown if `ndrange` and `numworkgroups` are both specified. - The helper function `KI.check_launch_args(numworkgroups, workgroupsize, ndrange)` can be + `ArgumentError` must also be thrown if `ndrange` and `numgroups` are both specified. + The helper function `KI.check_launch_args(numgroups, workgroupsize, ndrange)` can be used by the backend or a custom check can be implemented. `max_work_group_size` is to allow algorithms to request a max workgroupsize with `ndrange`. @@ -34,20 +34,20 @@ struct Kernel{B, Kern} end """ - check_launch_args(numworkgroups, workgroupsize, ndrange) + check_launch_args(numgroups, workgroupsize, ndrange) Validate the launch configuration passed to a [`Kernel`](@ref), throwing an `ArgumentError` if either argument has more than 3 dimensions, or if `ndrange` -and `numworkgroups` are both defined. +and `numgroups` are both defined. Backends may call this from their kernel-launch method instead of writing their own check. """ -function check_launch_args(numworkgroups, workgroupsize, ndrange) - length(ndrange) > 0 && length(numworkgroups) > 0 && - throw(ArgumentError("Only one of `numworkgroups` and `ndrange` can be used")) - length(numworkgroups) <= 3 || - throw(ArgumentError("`numworkgroups` only accepts up to 3 dimensions")) +function check_launch_args(numgroups, workgroupsize, ndrange) + length(ndrange) > 0 && length(numgroups) > 0 && + throw(ArgumentError("Only one of `numgroups` and `ndrange` can be used")) + length(numgroups) <= 3 || + throw(ArgumentError("`numgroups` only accepts up to 3 dimensions")) length(workgroupsize) <= 3 || throw(ArgumentError("`workgroupsize` only accepts up to 3 dimensions")) length(ndrange) <= 3 || @@ -78,13 +78,13 @@ function _threads_to_workgroupsize(threads, total, ndrange::Tuple, limits) end """ - auto_launch_sizes(kernel::KI.Kernel, numworkgroups, workgroupsize, ndrange, [max_work_items]) + auto_launch_sizes(kernel::KI.Kernel, numgroups, workgroupsize, ndrange, [max_work_items]) -Returns a suggested `numworkgroups` and `workgroupsize` based on +Returns a suggested `numgroups` and `workgroupsize` based on the input arguments. This function assumes arguments have been validated by `check_launch_args`. -If any `ndrange` dimension is zero, the returned `numworkgroups` is zero in +If any `ndrange` dimension is zero, the returned `numgroups` is zero in that dimension; backends should skip the launch in that case. Note that very large `ndrange`s can produce total grid sizes >= 2^32, which is problematic on some backends. @@ -92,9 +92,9 @@ on some backends. Backends may call this from their kernel-launch method instead of writing their own heuristic for calculating launch size. """ -@inline function auto_launch_sizes(kernel::Kernel, numworkgroups, workgroupsize, ndrange, max_work_items = typemax(Int)) - numworkgroups, workgroupsize = if ndrange == () - numworkgroups == () ? 1 : numworkgroups, workgroupsize == () ? 1 : workgroupsize +@inline function auto_launch_sizes(kernel::Kernel, numgroups, workgroupsize, ndrange, max_work_items = typemax(Int)) + numgroups, workgroupsize = if ndrange == () + numgroups == () ? 1 : numgroups, workgroupsize == () ? 1 : workgroupsize else workgroupsize = if workgroupsize == () max_wgs = kernel_max_work_group_size(kernel; max_work_items = min(prod(ndrange), max_work_items)) @@ -102,11 +102,11 @@ writing their own heuristic for calculating launch size. else workgroupsize end - numworkgroups = cld.(ndrange, workgroupsize) - Int.(numworkgroups), Int.(workgroupsize) + numgroups = cld.(ndrange, workgroupsize) + Int.(numgroups), Int.(workgroupsize) end - return numworkgroups, workgroupsize + return numgroups, workgroupsize end """ @@ -222,13 +222,13 @@ function argconvert end Low-level interface to compile a function invocation for the currently-active GPU, returning a callable kernel object. For a higher-level interface, use -[`KernelInterface.@kernel`](@ref). - -Currently, `kernel_function` only supports the `name` keyword argument as it is the only one -by all backends. +[`KernelInterface.@launch`](@ref). Keyword arguments: -- `name`: override the name that the kernel will have in the generated code +- `name`: override the name that the kernel will have in the generated code. + +Other keyword arguments are backend-specific compiler options (e.g. `maxthreads` for +CUDA.jl); backends throw an error for options they don't support. !!! note Backend implementations **must** implement: @@ -239,34 +239,43 @@ Keyword arguments: function kernel_function end const MACRO_KWARGS = [:launch] -const COMPILER_KWARGS = [:name] -const LAUNCH_KWARGS = [:numworkgroups, :workgroupsize, :ndrange, :max_work_group_size] +const LAUNCH_KWARGS = [:numgroups, :workgroupsize, :ndrange, :max_work_group_size] """ - KI.@kernel backend [workgroupsize=... numworkgroups=... ndrange=...] [kwargs...] func(args...) + KI.@launch backend [launch=true] [numgroups=...] [workgroupsize=...] [ndrange=...] [max_work_group_size=...] [kwargs...] f(args...) -High-level interface for executing code on a GPU. +Compile `f(args...)` for `backend` and launch it, like `@cuda` or `@metal` do. -The `KI.@kernel` macro should prefix a call, with `func` a callable function or object that -should return nothing. It will be compiled to a function native to the specified `backend` -upon first use, and to a certain extent arguments will be converted and managed automatically -using `argconvert`. Finally, if `launch=true`, the newly created callable kernel object is -called and launched according to the specified `backend`. +`f` and the arguments are converted with [`argconvert`](@ref) and compiled with +[`kernel_function`](@ref), and the resulting [`Kernel`](@ref) is called with the launch +keywords `numgroups`, `workgroupsize`, `ndrange` and `max_work_group_size`, whose meaning +is documented there. The arguments are kept alive while the launch is being queued. -There are a few keyword arguments that influence the behavior of `KI.@kernel`: +Other keyword arguments: +- `launch`: whether to launch the kernel, defaults to `true`. With `launch=false`, the + kernel is only compiled and returned, and the launch keywords can't be used: launch it by + calling it with the arguments and the launch keywords. +- `name` and any other keyword are passed to [`kernel_function`](@ref) as compiler options. -- `launch`: whether to launch this kernel, defaults to `true`. If `false`, the returned - kernel object should be launched by calling it and passing arguments again. -- `name`: the name of the kernel in the generated code. Defaults to an automatically- - generated name. +Launch options specific to a backend (such as a CUDA stream) can't be passed to +`@launch`; use `launch=false` and pass them when calling the kernel. -!!! note - `KI.@kernel` differs from the `KernelAbstractions` macro in that this macro acts - a wrapper around backend kernel compilation/launching (such as `@cuda`, `@metal`, etc.). It is - used when calling a function to be run on a specific backend, while `KernelAbstractions.@kernel` - is used kernel definition for use with the original higher-level `KernelAbstractions` API. +`backend` is evaluated once. Returns the `Kernel`. + +```julia +function vadd(c, a, b) + i = KI.get_global_id().x + if i <= length(c) + @inbounds c[i] = a[i] + b[i] + end + return +end + +KI.@launch backend ndrange=length(c) vadd(c, a, b) +``` """ -macro kernel(backend, ex...) +macro launch(backend, ex...) + isempty(ex) && throw(ArgumentError("KI.@launch needs a function call to launch")) call = ex[end] kwargs = map(ex[1:(end - 1)]) do kwarg if kwarg isa Symbol @@ -279,51 +288,44 @@ macro kernel(backend, ex...) end # destructure the kernel call - Meta.isexpr(call, :call) || throw(ArgumentError("final argument to KI.@kernel should be a function call")) + Meta.isexpr(call, :call) || throw(ArgumentError("final argument to KI.@launch should be a function call")) f = call.args[1] args = call.args[2:end] code = quote end vars, var_exprs = assign_args!(code, args) - # group keyword argument - macro_kwargs, compiler_kwargs, call_kwargs, other_kwargs = - split_kwargs(kwargs, MACRO_KWARGS, COMPILER_KWARGS, LAUNCH_KWARGS) - if !isempty(other_kwargs) - key, val = first(other_kwargs).args - throw(ArgumentError("Unsupported keyword argument '$key'")) - end + # group keyword argument; everything we don't know is a compiler option + macro_kwargs, call_kwargs, compiler_kwargs = + split_kwargs(kwargs, MACRO_KWARGS, LAUNCH_KWARGS) # handle keyword arguments that influence the macro's behavior launch = true for kwarg in macro_kwargs key, val = kwarg.args - if key === :launch - isa(val, Bool) || throw(ArgumentError("`launch` keyword argument to KI.@kernel should be a Bool")) - launch = val::Bool - else - throw(ArgumentError("Unsupported keyword argument '$key'")) - end + isa(val, Bool) || throw(ArgumentError("`launch` keyword argument to KI.@launch should be a Bool")) + launch = val::Bool end if !launch && !isempty(call_kwargs) - error("KI.@kernel with launch=false does not support launch-time keyword arguments; use them when calling the kernel") + throw(ArgumentError("KI.@launch with launch=false does not support launch keyword arguments; use them when calling the kernel")) end # FIXME: macro hygiene wrt. escaping kwarg values (this broke with 1.5) # we esc() the whole thing now, necessitating gensyms... - @gensym f_var kernel_f kernel_args kernel_tt kernel + @gensym backend_var f_var kernel_f kernel_args kernel_tt kernel # convert the arguments, call the compiler and launch the kernel # while keeping the original arguments alive push!( code.args, quote + $backend_var = $backend $f_var = $f GC.@preserve $(vars...) $f_var begin - $kernel_f = $argconvert($backend, $f_var) - $kernel_args = Base.map(x -> $argconvert($backend, x), ($(var_exprs...),)) + $kernel_f = $argconvert($backend_var, $f_var) + $kernel_args = Base.map(x -> $argconvert($backend_var, x), ($(var_exprs...),)) $kernel_tt = Tuple{Base.map(Core.Typeof, $kernel_args)...} - $kernel = $kernel_function($backend, $kernel_f, $kernel_tt; $(compiler_kwargs...)) + $kernel = $kernel_function($backend_var, $kernel_f, $kernel_tt; $(compiler_kwargs...)) if $launch $kernel($(var_exprs...); $(call_kwargs...)) end diff --git a/lib/KernelInterface/test/events.jl b/lib/KernelInterface/test/events.jl index 5f3f8da05..f78179086 100644 --- a/lib/KernelInterface/test/events.jl +++ b/lib/KernelInterface/test/events.jl @@ -30,7 +30,7 @@ function events_testsuite(b::KI.Backend) N = 64 A = KI.zeros(b, Float32, N) - slow_fill(v, iters) = KI.@kernel b numworkgroups = 1 workgroupsize = N slow_fill_kernel(A, v, UInt32(iters)) + slow_fill(v, iters) = KI.@launch b numgroups = 1 workgroupsize = N slow_fill_kernel(A, v, UInt32(iters)) # Time a launch as the minimum of a few runs: a backend's `synchronize` may run a # GC or otherwise stall once in a while, and the minimum discards that. diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl index b1888bb3f..92395b0bf 100644 --- a/lib/KernelInterface/test/interface.jl +++ b/lib/KernelInterface/test/interface.jl @@ -143,13 +143,13 @@ function interface_testsuite(backend::KI.Backend, AT) return end arr1d = AT(zeros(Float32, 4)) - KI.@kernel backend numworkgroups = 2 workgroupsize = 2 launch_kernel1d(arr1d) + KI.@launch backend numgroups = 2 workgroupsize = 2 launch_kernel1d(arr1d) KI.synchronize(backend) @test all(Array(arr1d) .== 1) # 1d tuple arr1dt = AT(zeros(Float32, 4)) - KI.@kernel backend numworkgroups = (2,) workgroupsize = (2,) launch_kernel1d(arr1dt) + KI.@launch backend numgroups = (2,) workgroupsize = (2,) launch_kernel1d(arr1dt) KI.synchronize(backend) @test all(Array(arr1dt) .== 1) @@ -163,7 +163,7 @@ function interface_testsuite(backend::KI.Backend, AT) return end arr2d = AT(zeros(Float32, 4, 4)) - KI.@kernel backend numworkgroups = (2, 2) workgroupsize = (2, 2) launch_kernel2d(arr2d) + KI.@launch backend numgroups = (2, 2) workgroupsize = (2, 2) launch_kernel2d(arr2d) KI.synchronize(backend) @test all(Array(arr2d) .== 1) @@ -177,13 +177,13 @@ function interface_testsuite(backend::KI.Backend, AT) return end arr3d = AT(zeros(Float32, 4, 4, 4)) - KI.@kernel backend numworkgroups = (2, 2, 2) workgroupsize = (2, 2, 2) launch_kernel3d(arr3d) + KI.@launch backend numgroups = (2, 2, 2) workgroupsize = (2, 2, 2) launch_kernel3d(arr3d) KI.synchronize(backend) @test all(Array(arr3d) .== 1) # 4d (Errors) - @test_throws ArgumentError (KI.@kernel backend numworkgroups = (2, 2, 2, 2) workgroupsize = (2, 2, 2) launch_kernel3d(arr3d)) - @test_throws ArgumentError (KI.@kernel backend numworkgroups = (2, 2, 2) workgroupsize = (2, 2, 2, 2) launch_kernel3d(arr3d)) + @test_throws ArgumentError (KI.@launch backend numgroups = (2, 2, 2, 2) workgroupsize = (2, 2, 2) launch_kernel3d(arr3d)) + @test_throws ArgumentError (KI.@launch backend numgroups = (2, 2, 2) workgroupsize = (2, 2, 2, 2) launch_kernel3d(arr3d)) end @testset "Launch limits" begin @@ -199,7 +199,7 @@ function interface_testsuite(backend::KI.Backend, AT) end return end - kernel = KI.@kernel backend launch = false fill_kernel(AT(zeros(Float32, 1, 1, 1))) + kernel = KI.@launch backend launch = false fill_kernel(AT(zeros(Float32, 1, 1, 1))) function fill_test(dims; kwargs...) arr = AT(zeros(Float32, dims)) kernel(arr; kwargs...) @@ -217,12 +217,12 @@ function interface_testsuite(backend::KI.Backend, AT) @testset "dimension $d" for d in 1:3 items = min(max_dims[d], max_items) workgroupsize = ntuple(i -> i == d ? items : 1, 3) - @test fill_test(workgroupsize; workgroupsize, numworkgroups = (1, 1, 1)) + @test fill_test(workgroupsize; workgroupsize, numgroups = (1, 1, 1)) # don't launch (practically) unlimited grids groups = min(max_groups[d], 2^16) - numworkgroups = ntuple(i -> i == d ? groups : 1, 3) - @test fill_test(numworkgroups; workgroupsize = (1, 1, 1), numworkgroups) + numgroups = ntuple(i -> i == d ? groups : 1, 3) + @test fill_test(numgroups; workgroupsize = (1, 1, 1), numgroups) end end @@ -251,13 +251,13 @@ function interface_testsuite(backend::KI.Backend, AT) @testset "Device return types" begin results = KI.zeros(backend, Bool, 6) - KI.@kernel backend typecheck_kernel(results) + KI.@launch backend typecheck_kernel(results) KI.synchronize(backend) @test all(Array(results)) @testset "$T" for T in (Int32, Int64, UInt32, UInt64) typed_results = KI.zeros(backend, Bool, 6) - KI.@kernel backend typed_typecheck_kernel(typed_results, T) + KI.@launch backend typed_typecheck_kernel(typed_results, T) KI.synchronize(backend) @test all(Array(typed_results)) end @@ -265,22 +265,22 @@ function interface_testsuite(backend::KI.Backend, AT) @testset "Typed indexing" begin workgroupsize = (2, 2, 2) - numworkgroups = (3, 2, 1) - N = prod(workgroupsize) * prod(numworkgroups) + numgroups = (3, 2, 1) + N = prod(workgroupsize) * prod(numgroups) # `Int` is the reference: it is what the zero-argument form returns. function run_typed(::Type{T}) where {T} results = KI.zeros(backend, T, N, 18) - KI.@kernel backend workgroupsize = workgroupsize numworkgroups = numworkgroups typed_index_kernel(results, T) + KI.@launch backend workgroupsize = workgroupsize numgroups = numgroups typed_index_kernel(results, T) KI.synchronize(backend) return Array(results) end reference = run_typed(Int) - global_size = workgroupsize .* numworkgroups + global_size = workgroupsize .* numgroups @test all(eachrow(reference[:, 1:3]) .== Ref(collect(global_size))) @test all(eachrow(reference[:, 7:9]) .== Ref(collect(workgroupsize))) - @test all(eachrow(reference[:, 13:15]) .== Ref(collect(numworkgroups))) + @test all(eachrow(reference[:, 13:15]) .== Ref(collect(numgroups))) # every global id is seen exactly once @test sort(Tuple.(eachrow(reference[:, 4:6]))) == sort(vec(Tuple.(CartesianIndices(global_size)))) @@ -298,15 +298,15 @@ function interface_testsuite(backend::KI.Backend, AT) # Test with small kernel workgroupsize = 4 - numworkgroups = 4 - N = workgroupsize * numworkgroups + numgroups = 4 + N = workgroupsize * numgroups results = AT(Vector{KernelData}(undef, N)) - kernel = KI.@kernel backend launch = false test_interface_kernel(results) + kernel = KI.@launch backend launch = false test_interface_kernel(results) @test KI.kernel_max_work_group_size(kernel) isa Int @test KI.kernel_max_work_group_size(kernel; max_work_items = 1) == 1 - kernel(results; workgroupsize, numworkgroups) + kernel(results; workgroupsize, numgroups) KI.synchronize(backend) host_results = Array(results) @@ -322,10 +322,10 @@ function interface_testsuite(backend::KI.Backend, AT) @test k_data.local_size == workgroupsize - @test k_data.num_groups == numworkgroups + @test k_data.num_groups == numgroups # Group ID should be 1-based - expected_group = div(i - 1, numworkgroups) + 1 + expected_group = div(i - 1, numgroups) + 1 @test k_data.group_id == expected_group # Local ID should be 1-based within group @@ -341,7 +341,7 @@ function interface_testsuite(backend::KI.Backend, AT) T = first(setdiff(KI.shfl_down_types(backend), [Bool])) results = KI.zeros(backend, Bool, 6) - KI.@kernel backend workgroupsize = KI.sub_group_size(backend) subgroup_typecheck_kernel(results, one(T)) + KI.@launch backend workgroupsize = KI.sub_group_size(backend) subgroup_typecheck_kernel(results, one(T)) KI.synchronize(backend) @test all(Array(results)) end @@ -353,13 +353,13 @@ function interface_testsuite(backend::KI.Backend, AT) sg_size = KI.sub_group_size(backend) sg_n = 2 workgroupsize = sg_size * sg_n - numworkgroups = 2 - N = workgroupsize * numworkgroups + numgroups = 2 + N = workgroupsize * numgroups results = AT(Vector{SubgroupData}(undef, N)) - kernel = KI.@kernel backend launch = false test_subgroup_kernel(results) + kernel = KI.@launch backend launch = false test_subgroup_kernel(results) - kernel(results; workgroupsize, numworkgroups) + kernel(results; workgroupsize, numgroups) KI.synchronize(backend) host_results = Array(results) @@ -390,7 +390,7 @@ function interface_testsuite(backend::KI.Backend, AT) 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)) + KI.@launch backend workgroupsize = N shfl_down_test_kernel(dev_a, dev_b, Val(N)) b = Array(dev_b) @test sum(a) ≈ b[1] diff --git a/lib/KernelInterface/test/runtests.jl b/lib/KernelInterface/test/runtests.jl index 1da2b4da4..bb7affff7 100644 --- a/lib/KernelInterface/test/runtests.jl +++ b/lib/KernelInterface/test/runtests.jl @@ -200,8 +200,8 @@ end @test_throws ArgumentError KI.check_launch_args((1, 2, 3, 4), 1, ()) @test_throws ArgumentError KI.check_launch_args(1, (1, 2, 3, 4), ()) @test_throws ArgumentError KI.check_launch_args((), 1, (1, 2, 3, 4)) - @test_throws ArgumentError KI.check_launch_args(2, 4, 2) # both numworkgroupsize and ndrange defined - @test_throws ArgumentError KI.check_launch_args(2, (), 2) # both numworkgroupsize and ndrange defined + @test_throws ArgumentError KI.check_launch_args(2, 4, 2) # both numgroupsize and ndrange defined + @test_throws ArgumentError KI.check_launch_args(2, (), 2) # both numgroupsize and ndrange defined end @testset "threads_to_workgroupsize" begin @@ -282,14 +282,11 @@ KI.max_work_group_dims(::DimsBackend) = (1024, 1024, 64) end @testset "split_kwargs" begin - kwargs = [:(launch = false), :(name = "foo"), :(numworkgroups = 2)] - macro_kw, compiler_kw, launch_kw, other = KI.split_kwargs( - kwargs, KI.MACRO_KWARGS, KI.COMPILER_KWARGS, KI.LAUNCH_KWARGS - ) + kwargs = [:(launch = false), :(name = "foo"), :(numgroups = 2)] + macro_kw, launch_kw, other = KI.split_kwargs(kwargs, KI.MACRO_KWARGS, KI.LAUNCH_KWARGS) @test macro_kw == [:(launch = false)] - @test compiler_kw == [:(name = "foo")] - @test launch_kw == [:(numworkgroups = 2)] - @test isempty(other) + @test launch_kw == [:(numgroups = 2)] + @test other == [:(name = "foo")] # Unmatched keywords land in the trailing group rather than erroring. _, unmatched = KI.split_kwargs([:(bogus = 1)], [:launch]) @@ -313,19 +310,20 @@ end @test var_exprs[2] == Expr(:..., vars[2]) end -# A minimal backend, exercising the contract `KI.@kernel` expects of one. +# A minimal backend, exercising the contract `KI.@launch` expects of one. struct MockBackend end struct MockKernel f::Any tt::Any name::Any + options::Any launches::Vector{Any} end KI.argconvert(::MockBackend, arg) = arg function KI.kernel_function(::MockBackend, f, tt = Tuple{}; name = nothing, kwargs...) - return MockKernel(f, tt, name, []) + return MockKernel(f, tt, name, Dict(kwargs), []) end function (kernel::MockKernel)(args...; kwargs...) push!(kernel.launches, (args, Dict(kwargs))) @@ -334,27 +332,41 @@ end dummy(a, b) = nothing -@testset "@kernel" begin +const backend_evaluations = Ref(0) +function counted_backend() + backend_evaluations[] += 1 + return MockBackend() +end + +@testset "@launch" begin backend = MockBackend() - kernel = KI.@kernel backend numworkgroups = 2 workgroupsize = 4 dummy(1, 2.0) + kernel = KI.@launch backend numgroups = 2 workgroupsize = 4 dummy(1, 2.0) @test kernel isa MockKernel @test kernel.f === dummy @test kernel.tt == Tuple{Int, Float64} args, launch_kwargs = only(kernel.launches) @test args == (1, 2.0) - @test launch_kwargs == Dict(:numworkgroups => 2, :workgroupsize => 4) + @test launch_kwargs == Dict(:numgroups => 2, :workgroupsize => 4) + + # the backend expression is evaluated once + backend_evaluations[] = 0 + KI.@launch counted_backend() ndrange = 4 dummy(1, 2.0) + @test backend_evaluations[] == 1 # `launch=false` compiles only; the caller launches later. - deferred = KI.@kernel backend launch = false dummy(1, 2.0) + deferred = KI.@launch backend launch = false dummy(1, 2.0) @test isempty(deferred.launches) - # Compiler kwargs reach `kernel_function` instead of the launch. - named = KI.@kernel backend launch = false name = "mykernel" dummy(1, 2.0) + # Other keywords are compiler options for `kernel_function`. + named = KI.@launch backend launch = false name = "mykernel" maxthreads = 32 dummy(1, 2.0) @test named.name == "mykernel" + @test named.options == Dict(:maxthreads => 32) + optioned = KI.@launch backend ndrange = 4 maxthreads = 32 dummy(1, 2.0) + @test last(only(optioned.launches)) == Dict(:ndrange => 4) # Splatted arguments are supported. - splatted = KI.@kernel backend launch = false dummy((1, 2.0)...) + splatted = KI.@launch backend launch = false dummy((1, 2.0)...) @test splatted.tt == Tuple{Int, Float64} @testset "errors" begin @@ -369,14 +381,14 @@ dummy(a, b) = nothing return nothing end - @test expansion_error(:(KI.@kernel backend bogus = 1 dummy(1))) isa ArgumentError - @test expansion_error(:(KI.@kernel backend dummy)) isa ArgumentError - @test expansion_error(:(KI.@kernel backend launch = 1 dummy(1))) isa ArgumentError - @test expansion_error(:(KI.@kernel backend "notakwarg" dummy(1))) isa ArgumentError - # launch-time kwargs are meaningless when we are not launching + @test expansion_error(:(KI.@launch backend)) isa ArgumentError + @test expansion_error(:(KI.@launch backend dummy)) isa ArgumentError + @test expansion_error(:(KI.@launch backend launch = 1 dummy(1))) isa ArgumentError + @test expansion_error(:(KI.@launch backend "notakwarg" dummy(1))) isa ArgumentError + # launch keywords are meaningless when we are not launching @test expansion_error( - :(KI.@kernel backend launch = false numworkgroups = 2 dummy(1)) - ) isa ErrorException + :(KI.@launch backend launch = false numgroups = 2 dummy(1)) + ) isa ArgumentError end end diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index e51a0795d..66597624d 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -211,18 +211,18 @@ function KI.kernel_function(::POCLBackend, f::F, tt::TT = Tuple{}; name = nothin return KI.Kernel{POCLBackend, typeof(kern)}(POCLBackend(), kern) end -function (obj::KI.Kernel{POCLBackend})(args...; numworkgroups = (), workgroupsize = (), ndrange = (), max_work_group_size = typemax(Int)) - KI.check_launch_args(numworkgroups, workgroupsize, ndrange) +function (obj::KI.Kernel{POCLBackend})(args...; numgroups = (), workgroupsize = (), ndrange = (), max_work_group_size = typemax(Int)) + KI.check_launch_args(numgroups, workgroupsize, ndrange) # zero-sized ndrange: nothing to launch prod(ndrange) == 0 && return nothing - numworkgroups, workgroupsize = KI.auto_launch_sizes(obj, numworkgroups, workgroupsize, ndrange, max_work_group_size) + numgroups, workgroupsize = KI.auto_launch_sizes(obj, numgroups, workgroupsize, ndrange, max_work_group_size) local_size = (workgroupsize..., ntuple(_ -> 1, 3 - length(workgroupsize))...) - numworkgroups = (numworkgroups..., ntuple(_ -> 1, 3 - length(numworkgroups))...) - global_size = local_size .* numworkgroups + numgroups = (numgroups..., ntuple(_ -> 1, 3 - length(numgroups))...) + global_size = local_size .* numgroups event = obj.kern(args...; local_size, global_size) wait(event) diff --git a/test/hostinterface.jl b/test/hostinterface.jl index 17728150e..26aca3ed5 100644 --- a/test/hostinterface.jl +++ b/test/hostinterface.jl @@ -68,7 +68,7 @@ function hostinterface_testsuite(_backend, AT) end x = AT(zeros(Float32, 4)) - kernel = KI.@kernel _backend() launch = false ki_hostinterface_kernel(x) + kernel = KI.@launch _backend() launch = false ki_hostinterface_kernel(x) @test kernel isa KI.Kernel @test KI.kernel_max_work_group_size(kernel) isa Int @test KI.kernel_max_work_group_size(kernel; max_work_items = 1) == 1 From 381c17b3246cad2064b2b51566d82a211d56da98 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 09:20:24 +0200 Subject: [PATCH 4/9] KernelInterface: separate the work-group size limit from the recommendation `kernel_max_work_group_size` returned an occupancy recommendation on CUDA, AMDGPU and oneAPI, but the hard limit on Metal, OpenCL and PoCL: on an RTX 5080, 768 for a trivial kernel that can be launched with 1024 threads. It is replaced by two queries: - `max_work_group_size(kernel)`, the largest work-group the compiled kernel can be launched with; - `launch_configuration(kernel; nitems, max_work_group_size)`, the recommended work-group size for a launch of `nitems` work-items, which `ndrange` launches without a `workgroupsize` use. It falls back to the limit. The problem size is passed separately from the bound, so that heuristics like CUDA's `prefer_blocks` don't have to guess it. `max_work_group_dims` and `max_num_groups` are now required: their `typemax(Int)` fallbacks meant both "unknown" and "unlimited". --- docs/src/kernelinterface.md | 11 ++-- lib/KernelInterface/src/launch.jl | 78 +++++++++++++++++---------- lib/KernelInterface/test/interface.jl | 19 +++++-- lib/KernelInterface/test/runtests.jl | 57 ++++++++++++++------ src/pocl/backend.jl | 6 ++- test/hostinterface.jl | 5 +- 6 files changed, 117 insertions(+), 59 deletions(-) diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md index 446b34ceb..293c6a75f 100644 --- a/docs/src/kernelinterface.md +++ b/docs/src/kernelinterface.md @@ -165,10 +165,11 @@ supports_atomics supports_float64 ``` -### Backend queries +### Limits ```@docs max_work_group_size +launch_configuration max_work_group_dims max_num_groups sub_group_size @@ -180,7 +181,6 @@ multiprocessor_count ```@docs Kernel kernel_function -kernel_max_work_group_size argconvert KernelInterface.@launch ``` @@ -210,9 +210,10 @@ A backend must, at minimum: `KI.check_launch_args` to validate them, or check them directly. A zero-sized `ndrange` — launching over an empty array is not uncommon — must be a no-op returning `nothing`, not an error. -7. Report its limits through [`kernel_max_work_group_size`](@ref) and, where - applicable, [`max_work_group_size`](@ref), [`sub_group_size`](@ref) and - [`multiprocessor_count`](@ref). +7. Report its limits through [`max_work_group_size`](@ref) (for the backend and for a + kernel), [`max_work_group_dims`](@ref) and [`max_num_groups`](@ref), and where + applicable [`sub_group_size`](@ref) and [`multiprocessor_count`](@ref). It may + recommend work-group sizes with [`launch_configuration`](@ref). The PoCL backend in `src/pocl/backend.jl` is a complete worked example. diff --git a/lib/KernelInterface/src/launch.jl b/lib/KernelInterface/src/launch.jl index 925491a39..ab30237bf 100644 --- a/lib/KernelInterface/src/launch.jl +++ b/lib/KernelInterface/src/launch.jl @@ -97,8 +97,8 @@ writing their own heuristic for calculating launch size. numgroups == () ? 1 : numgroups, workgroupsize == () ? 1 : workgroupsize else workgroupsize = if workgroupsize == () - max_wgs = kernel_max_work_group_size(kernel; max_work_items = min(prod(ndrange), max_work_items)) - threads_to_workgroupsize(max_wgs, ndrange, max_work_group_dims(kernel.backend)) + config = launch_configuration(kernel; nitems = prod(ndrange), max_work_group_size = max_work_items) + threads_to_workgroupsize(config.workgroupsize, ndrange, max_work_group_dims(kernel.backend)) else workgroupsize end @@ -109,68 +109,88 @@ writing their own heuristic for calculating launch size. return numgroups, workgroupsize end + +## limits and advice + """ - kernel_max_work_group_size(kern; [max_work_items::Int])::Int + max_work_group_size(backend)::Int + max_work_group_size(kernel::Kernel)::Int + +The largest number of work-items a work-group can have: on the active device of `backend`, +or for launches of the compiled `kernel` (which may be lower, e.g. because of the kernel's +register use). Launching a larger work-group is an error. -The maximum workgroup size limit for a kernel as reported by the backend. -This function should always be used to determine the workgroup size before -launching a kernel. +The work-group size that performs best is often smaller; see [`launch_configuration`](@ref). !!! note - Backend implementations **must** implement: + Backend implementations **must** implement both: ``` - kernel_max_work_group_size(kern::Kernel{<:NewBackend}; max_work_items::Int=typemax(Int))::Int + max_work_group_size(backend::NewBackend)::Int + max_work_group_size(kernel::Kernel{<:NewBackend})::Int ``` - As well as the on-device functionality. + The kernel form answers for the device the kernel was compiled for. """ -function kernel_max_work_group_size end +function max_work_group_size end """ - max_work_group_size(backend, kern; [max_work_items::Int])::Int + launch_configuration(kernel::Kernel; nitems=nothing, max_work_group_size=typemax(Int))::@NamedTuple{workgroupsize::Int} -The maximum workgroup size limit for a kernel as reported by the backend. -This function represents a theoretical maximum; `kernel_max_work_group_size` -should be used before launching a kernel as some backends may error if -kernel launch with too big a workgroup is attempted. +The recommended number of work-items per work-group for launching `kernel` over `nitems` +work-items in total (`nothing` if unknown), at most `max_work_group_size`. This is what an +`ndrange` launch without a `workgroupsize` uses, passing the number of work-items in the +`ndrange`. `nitems` and `max_work_group_size` are positive. + +Unlike [`max_work_group_size`](@ref), this is advice: backends may base it on occupancy or +on the size of the launch, e.g. to prefer more work-groups over larger ones. !!! note - Backend implementations **must** implement: + Backend implementations **may** implement: ``` - max_work_group_size(backend::NewBackend)::Int + launch_configuration(kernel::Kernel{<:NewBackend}; nitems::Union{Int, Nothing}=nothing, + max_work_group_size::Int=typemax(Int))::@NamedTuple{workgroupsize::Int} ``` - As well as the on-device functionality. + The result has to be positive and at most `max_work_group_size` and + `max_work_group_size(kernel)`. The fallback recommends the largest legal work-group + size. """ -function max_work_group_size end +function launch_configuration( + kernel::Kernel; nitems::Union{Integer, Nothing} = nothing, + max_work_group_size::Integer = typemax(Int) + ) + return (; workgroupsize = Int(min(KernelInterface.max_work_group_size(kernel), max_work_group_size))) +end """ max_work_group_dims(backend)::NTuple{3, Int} -The maximum number of work-items along each dimension of a workgroup, for the currently -active device of `backend`. [`max_work_group_size`](@ref) bounds their product. +The maximum number of work-items along each dimension of a work-group, for the active +device of `backend`. [`max_work_group_size`](@ref) bounds their product. !!! note - Backend implementations **should** implement: + Backend implementations **must** implement: ``` max_work_group_dims(backend::NewBackend)::NTuple{3, Int} ``` - The fallback does not limit individual dimensions. """ -max_work_group_dims(::Backend) = (typemax(Int), typemax(Int), typemax(Int)) +function max_work_group_dims end """ max_num_groups(backend)::NTuple{3, Int} -The maximum number of workgroups along each dimension of a launch, for the currently -active device of `backend`. +The maximum number of work-groups along each dimension of a launch, for the active device +of `backend`. + +This is conservative: a launch within these limits works for any work-group size, but some +backends accept more work-groups for smaller work-groups (e.g. HIP bounds the number of +work-items per dimension). The backend's validation at launch time is authoritative. !!! note - Backend implementations **should** implement: + Backend implementations **must** implement: ``` max_num_groups(backend::NewBackend)::NTuple{3, Int} ``` - The fallback does not limit the number of workgroups. """ -max_num_groups(::Backend) = (typemax(Int), typemax(Int), typemax(Int)) +function max_num_groups end """ sub_group_size(backend)::Int diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl index 92395b0bf..6d24983c2 100644 --- a/lib/KernelInterface/test/interface.jl +++ b/lib/KernelInterface/test/interface.jl @@ -191,6 +191,7 @@ function interface_testsuite(backend::KI.Backend, AT) max_groups = KI.max_num_groups(backend) @test max_dims isa NTuple{3, Int} && all(>=(1), max_dims) @test max_groups isa NTuple{3, Int} && all(>=(1), max_groups) + @test KI.max_work_group_size(backend) isa Int function fill_kernel(arr) i, j, k = KI.get_global_id() @@ -207,13 +208,26 @@ function interface_testsuite(backend::KI.Backend, AT) return all(Array(arr) .== 1) end + max_items = KI.max_work_group_size(kernel) + @test max_items isa Int && 1 <= max_items <= KI.max_work_group_size(backend) + config = KI.launch_configuration(kernel) + @test config isa @NamedTuple{workgroupsize::Int} + @test 1 <= config.workgroupsize <= max_items + @test KI.launch_configuration(kernel; max_work_group_size = 1).workgroupsize == 1 + # the recommendation is legal, whatever the size of the launch + @testset "nitems = $nitems, max_work_group_size = $cap" for nitems in (nothing, 1, 1000, typemax(Int)), + cap in (1, 7, typemax(Int)) + config = KI.launch_configuration(kernel; nitems, max_work_group_size = cap) + @test 1 <= config.workgroupsize <= min(cap, max_items) + end + # automatically chosen workgroup sizes respect the per-dimension limits @testset "ndrange = $ndrange" for ndrange in ((1, 1, 5000), (1, 5000, 1), (1, 3, 2000)) @test fill_test(ndrange; ndrange) + @test fill_test(ndrange; ndrange, max_work_group_size = 7) end # the reported limits can be launched - max_items = KI.kernel_max_work_group_size(kernel) @testset "dimension $d" for d in 1:3 items = min(max_dims[d], max_items) workgroupsize = ntuple(i -> i == d ? items : 1, 3) @@ -303,9 +317,6 @@ function interface_testsuite(backend::KI.Backend, AT) results = AT(Vector{KernelData}(undef, N)) kernel = KI.@launch backend launch = false test_interface_kernel(results) - @test KI.kernel_max_work_group_size(kernel) isa Int - @test KI.kernel_max_work_group_size(kernel; max_work_items = 1) == 1 - kernel(results; workgroupsize, numgroups) KI.synchronize(backend) diff --git a/lib/KernelInterface/test/runtests.jl b/lib/KernelInterface/test/runtests.jl index bb7affff7..234f3751f 100644 --- a/lib/KernelInterface/test/runtests.jl +++ b/lib/KernelInterface/test/runtests.jl @@ -33,7 +33,7 @@ end KI.get_num_sub_groups, KI.get_sub_group_id, KI.get_sub_group_local_id, KI.shfl_down, - KI.kernel_max_work_group_size, KI.max_work_group_size, KI.sub_group_size, + KI.max_work_group_size, KI.max_work_group_dims, KI.max_num_groups, KI.sub_group_size, KI.argconvert, KI.kernel_function, # Host-side stubs: required backend methods with no sensible fallback. KI.synchronize, KI.copyto!, @@ -225,12 +225,6 @@ end @test KI.threads_to_workgroupsize(64, (1, 1, 1, 100), (1024, 1024, 64)) == (1, 1, 1, 64) end -@testset "per-dimension limits" begin - # Unlimited unless the backend says otherwise. - @test KI.max_work_group_dims(StubBackend()) == (typemax(Int), typemax(Int), typemax(Int)) - @test KI.max_num_groups(StubBackend()) == (typemax(Int), typemax(Int), typemax(Int)) -end - @testset "Kernel" begin kernel = KI.Kernel(:backend, :kern) @test kernel.backend === :backend @@ -242,19 +236,27 @@ end struct SizedBackend <: KI.Backend maxThreads::Int end -function KI.kernel_max_work_group_size(k::KI.Kernel{SizedBackend}; max_work_items::Int = typemax(Int)) - return min(k.backend.maxThreads, max_work_items) -end +KI.max_work_group_size(k::KI.Kernel{SizedBackend}) = k.backend.maxThreads +KI.max_work_group_dims(::SizedBackend) = (1024, 1024, 64) -# ... and a fixed limit per workgroup dimension -struct DimsBackend <: KI.Backend end -KI.kernel_max_work_group_size(k::KI.Kernel{DimsBackend}; max_work_items::Int = typemax(Int)) = - min(1024, max_work_items) -KI.max_work_group_dims(::DimsBackend) = (1024, 1024, 64) +# ... and one recommending smaller work-groups than it can launch, like CUDA's occupancy API, +# recording what it was asked +struct OccupancyBackend <: KI.Backend + queries::Vector{Any} +end +OccupancyBackend() = OccupancyBackend([]) +KI.max_work_group_size(::KI.Kernel{OccupancyBackend}) = 1024 +KI.max_work_group_dims(::OccupancyBackend) = (1024, 1024, 64) +function KI.launch_configuration( + kernel::KI.Kernel{OccupancyBackend}; nitems = nothing, max_work_group_size = typemax(Int) + ) + push!(kernel.backend.queries, (; nitems, max_work_group_size)) + return (; workgroupsize = min(96, max_work_group_size)) +end @testset "auto_launch_sizes" begin # the per-dimension limit is respected - @test KI.auto_launch_sizes(KI.Kernel(DimsBackend(), nothing), (), (), (1, 1, 5000)) === + @test KI.auto_launch_sizes(KI.Kernel(SizedBackend(1024), nothing), (), (), (1, 1, 5000)) === ((1, 1, 79), (1, 1, 64)) kernel = KI.Kernel(SizedBackend(256), nothing) @@ -277,8 +279,29 @@ KI.max_work_group_dims(::DimsBackend) = (1024, 1024, 64) # A zero-sized ndrange yields zero workgroups; backends skip the launch. @test KI.auto_launch_sizes(kernel, (), (), (0,)) === ((0,), (1,)) - @test KI.auto_launch_sizes(kernel, (), (), (0, 4)) === ((0, 4), (1, 1)) + @test KI.auto_launch_sizes(kernel, (), (), (0, 4)) === ((0, 1), (1, 4)) @test KI.auto_launch_sizes(kernel, (), (), 0) === (0, 1) + + # The backend's recommendation is used, not the limit, and it's told both the size of + # the launch and the cap. + occupancy = KI.Kernel(OccupancyBackend(), nothing) + @test KI.auto_launch_sizes(occupancy, (), (), (1000,)) === ((11,), (96,)) + @test only(occupancy.backend.queries) == (; nitems = 1000, max_work_group_size = typemax(Int)) + @test KI.auto_launch_sizes(occupancy, (), (), (1000, 3), 64) === ((16, 3), (64, 1)) + @test last(occupancy.backend.queries) == (; nitems = 3000, max_work_group_size = 64) +end + +@testset "launch_configuration" begin + kernel = KI.Kernel(SizedBackend(256), nothing) + # the fallback recommends the limit + @test KI.launch_configuration(kernel) === (; workgroupsize = 256) + @test KI.launch_configuration(kernel; max_work_group_size = 100) === (; workgroupsize = 100) + @test KI.launch_configuration(kernel; nitems = 10) === (; workgroupsize = 256) + + # backends can recommend less than the limit + occupancy = KI.Kernel(OccupancyBackend(), nothing) + @test KI.launch_configuration(occupancy) === (; workgroupsize = 96) + @test KI.max_work_group_size(occupancy) == 1024 end @testset "split_kwargs" begin diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index 66597624d..3c545de36 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -230,9 +230,9 @@ function (obj::KI.Kernel{POCLBackend})(args...; numgroups = (), workgroupsize = return nothing end -function KI.kernel_max_work_group_size(kernel::KI.Kernel{<:POCLBackend}; max_work_items::Int = typemax(Int))::Int +function KI.max_work_group_size(kernel::KI.Kernel{<:POCLBackend})::Int wginfo = cl.work_group_info(kernel.kern.fun, device()) - return Int(min(wginfo.size, max_work_items)) + return Int(wginfo.size) end # querying the device allocates, so cache the limits that every launch needs function device_limits() @@ -247,6 +247,8 @@ function device_limits() end KI.max_work_group_size(::POCLBackend)::Int = device_limits().max_work_group_size KI.max_work_group_dims(::POCLBackend)::NTuple{3, Int} = device_limits().max_work_group_dims +# the grid is only limited by the size of `size_t` +KI.max_num_groups(::POCLBackend)::NTuple{3, Int} = (typemax(Int), typemax(Int), typemax(Int)) function KI.sub_group_size(::POCLBackend)::Int # POCL can technically support any sub_group size. # Check for common values used on GPUs then diff --git a/test/hostinterface.jl b/test/hostinterface.jl index 26aca3ed5..68cce2f12 100644 --- a/test/hostinterface.jl +++ b/test/hostinterface.jl @@ -70,8 +70,9 @@ function hostinterface_testsuite(_backend, AT) x = AT(zeros(Float32, 4)) kernel = KI.@launch _backend() launch = false ki_hostinterface_kernel(x) @test kernel isa KI.Kernel - @test KI.kernel_max_work_group_size(kernel) isa Int - @test KI.kernel_max_work_group_size(kernel; max_work_items = 1) == 1 + @test KI.max_work_group_size(kernel) isa Int + @test KI.launch_configuration(kernel) isa @NamedTuple{workgroupsize::Int} + @test KI.launch_configuration(kernel; max_work_group_size = 1).workgroupsize == 1 end return nothing From 960cb6b7f8160c12eca076ba67b918b9e327ac28 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 09:24:47 +0200 Subject: [PATCH 5/9] KernelInterface: validate launches in KernelInterface, backends implement KI.launch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every backend repeated the launch validation and sizing in front of its native launch, and the copies had drifted apart (return values, a `prod(ndrange) == 0` check that can overflow). KernelInterface now implements the `Kernel` call and passes validated 3-D sizes to one required method: KI.launch(kernel, groups::Dims{3}, items::Dims{3}, args...; kwargs...) The launch semantics are documented on `Kernel`: `ndrange` is rounded up to whole work-groups and not masked, a zero in `ndrange` or `numgroups` launches nothing, work-group sizes are checked against `max_work_group_dims` and `max_work_group_size(kernel)`, and the number of work-items per dimension has to fit an `Int`. Invalid launches throw an `ArgumentError` before the backend sees them. Unknown keywords are passed on to `KI.launch`, so options like CUDA's `stream` still reach the driver. The `Kernel` call and PoCL's `launch` declare their arguments as `Vararg{Any, N}`: Julia doesn't specialize a method on `args...` that it only passes through, which made every launch dispatch dynamically (on CUDA, 1.2 µs and 1 kB more per `KI.@launch`). The `launch` docstring recommends the same to backends. `check_launch_args` and `auto_launch_sizes` are removed. The launch tests now use different group counts and sizes per dimension, which exposes two tests whose formulas only worked because they were equal. --- docs/src/kernelinterface.md | 12 +- lib/KernelInterface/src/launch.jl | 248 ++++++++++++++++++-------- lib/KernelInterface/test/interface.jl | 140 ++++++++------- lib/KernelInterface/test/runtests.jl | 197 +++++++++++--------- src/pocl/backend.jl | 17 +- 5 files changed, 373 insertions(+), 241 deletions(-) diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md index 293c6a75f..dc6456ace 100644 --- a/docs/src/kernelinterface.md +++ b/docs/src/kernelinterface.md @@ -182,6 +182,7 @@ multiprocessor_count Kernel kernel_function argconvert +launch KernelInterface.@launch ``` @@ -205,11 +206,12 @@ A backend must, at minimum: [`shfl_down`](@ref) support is optional. 5. Implement [`argconvert`](@ref) and [`kernel_function`](@ref) for its backend type, returning a [`Kernel`](@ref). -6. Make that `Kernel` callable, accepting `numgroups`, `workgroupsize` and - `ndrange` as a scalar `Integer` or a 1-, 2- or 3-element tuple. Use - `KI.check_launch_args` to validate them, or check them directly. A zero-sized - `ndrange` — launching over an empty array is not uncommon — must be a no-op - returning `nothing`, not an error. +6. Implement [`launch`](@ref), which receives an already validated `NTuple{3, Int}` of + work-groups and of work-items. For CUDA.jl, that is + ```julia + KI.launch(k::KI.Kernel{CUDABackend}, groups::Dims{3}, items::Dims{3}, args::Vararg{Any, N}; kwargs...) where {N} = + k.kern(args...; threads = items, blocks = groups, kwargs...) + ``` 7. Report its limits through [`max_work_group_size`](@ref) (for the backend and for a kernel), [`max_work_group_dims`](@ref) and [`max_num_groups`](@ref), and where applicable [`sub_group_size`](@ref) and [`multiprocessor_count`](@ref). It may diff --git a/lib/KernelInterface/src/launch.jl b/lib/KernelInterface/src/launch.jl index ab30237bf..42869a567 100644 --- a/lib/KernelInterface/src/launch.jl +++ b/lib/KernelInterface/src/launch.jl @@ -3,55 +3,185 @@ """ Kernel{Backend, Kern} -Kernel closure struct that is used to represent the backend -kernel on the host. +A kernel compiled by [`kernel_function`](@ref) for `backend`, wrapping the backend's own +kernel object `kern`. `kernel.backend` is the backend value that was passed to +`kernel_function`. -!!! note - Backend implementations **must** implement: - ``` - (kernel::Kernel{<:NewBackend})(args...; numgroups=(), workgroupsize=(), ndrange=(), max_work_group_size=typemax(Int)) - ``` - `numgroups`, `workgroupsize`, and `ndrange` must accept a scalar Integer, a 1, 2, - or 3 Integer tuple, or an empty tuple. Otherwise, it must throw an `ArgumentError`. An - `ArgumentError` must also be thrown if `ndrange` and `numgroups` are both specified. - The helper function `KI.check_launch_args(numgroups, workgroupsize, ndrange)` can be - used by the backend or a custom check can be implemented. +Calling a `Kernel` launches it: + + (kernel::Kernel)(args...; numgroups=(), workgroupsize=(), ndrange=(), + max_work_group_size=typemax(Int), kwargs...) + +`args` are the host-side arguments, e.g. a `CuArray` rather than a `CuDeviceArray`. The +backend converts them with [`argconvert`](@ref), and the converted types have to match the +argument types the kernel was compiled for. + +The launch geometry is given in one of three ways: + +- `numgroups` and `workgroupsize`: launch exactly that many work-groups of that many + work-items. Either defaults to 1. +- `ndrange` and `workgroupsize`: launch `cld.(ndrange, workgroupsize)` work-groups. +- `ndrange` alone: the work-group size is chosen with [`launch_configuration`](@ref), bounded + by `max_work_group_size` and by [`max_work_group_dims`](@ref), and filled first dimension + first. - `max_work_group_size` is to allow algorithms to request a max workgroupsize with `ndrange`. - This is a maximum value because a kernel's maximum workitems per workgroup may be lower than - requested. +Each is an `Integer` or a tuple of up to 3 `Integer`s; missing dimensions are 1. `ndrange` +and `numgroups` are mutually exclusive. - An `ndrange` with a zero-sized dimension, as when launching over an empty array, is - not an error: the call must be a no-op and return `nothing` instead of launching. +`ndrange` is rounded up to whole work-groups and is not masked: the kernel runs for every +work-item of every launched group, and [`get_global_size`](@ref) returns the padded size. +Kernels have to check their own bounds. - By default, kernels must launch with 1 workgroup containing 1 workitem. +A zero anywhere in `ndrange` or `numgroups` launches nothing. Work-group sizes must be +positive and fit [`max_work_group_dims`](@ref) and [`max_work_group_size`](@ref)`(kernel)`, +and the number of work-items in each dimension must fit an `Int`; anything else throws an +`ArgumentError` before the backend sees it. The number of work-groups is validated by the +backend. - Backends must also implement the on-device kernel launch functionality. +Other keyword arguments are passed to [`launch`](@ref) unchanged. They are +backend-specific: a backend throws an error for keywords it doesn't support. + +A launch is queued on the calling task's queue of the active device, and returns `nothing`. """ struct Kernel{B, Kern} backend::B kern::Kern end +# `Vararg{Any, N}` makes Julia specialize on the arguments, which are only passed through +function (kernel::Kernel)( + args::Vararg{Any, N}; numgroups = (), workgroupsize = (), ndrange = (), + max_work_group_size::Integer = typemax(Int), kwargs... + ) where {N} + groups, items = launch_geometry(kernel, numgroups, workgroupsize, ndrange, max_work_group_size) + any(iszero, groups) && return nothing + launch(kernel, groups, items, args...; kwargs...) + return nothing +end + """ - check_launch_args(numgroups, workgroupsize, ndrange) + launch(kernel::Kernel, groups::Dims{3}, items::Dims{3}, args...; kwargs...) + +Launch `kernel` with `groups` work-groups of `items` work-items each, passing the host-side +arguments `args`. This is what calling a [`Kernel`](@ref) does after validating and +normalizing the launch geometry; users call the kernel instead. + +`groups` and `items` are positive, `items` fits [`max_work_group_dims`](@ref) and +[`max_work_group_size`](@ref)`(kernel)`, and `groups .* items` doesn't overflow `Int`. +`kwargs` are the keyword arguments of the call that KernelInterface doesn't know. + +!!! note + Backend implementations **must** implement: + ``` + launch(kernel::Kernel{<:NewBackend}, groups::Dims{3}, items::Dims{3}, args...; kwargs...) + ``` + It converts `args` with [`argconvert`](@ref) (or lets its native launcher do so), and + queues the launch on the calling task's queue; it doesn't have to wait for the kernel to + complete. Declare the arguments as `args::Vararg{Any, N}` (with `where {N}`): Julia + doesn't specialize a method on `args...` that it only passes through, which makes every + launch dispatch dynamically. It must throw for keywords it does not support, and may + throw for a number of work-groups the device cannot launch, or for a geometry that + backend-specific compiler options of the kernel don't allow. +""" +function launch end + +# the launch keywords are either scalars or tuples of up to 3 integers +const LaunchDims = Union{Integer, Tuple{}, NTuple{1, Integer}, NTuple{2, Integer}, NTuple{3, Integer}} + +@inline pad3(x::Integer) = (Int(x), 1, 1) +@inline pad3(x::Tuple) = (map(Int, x)..., ntuple(_ -> 1, Val(3 - length(x)))...) + +@noinline function throw_launch_error(name, value) + throw(ArgumentError("`$name` must be an integer or a tuple of up to 3 integers, got $(repr(value))")) +end + +@noinline function throw_range_error(name, value) + throw(ArgumentError("`$name` must be between 0 and typemax(Int), got $(repr(value))")) +end + +@inline function check_dims(name, x) + x isa LaunchDims || throw_launch_error(name, x) + any(d -> d < 0 || d > typemax(Int), x) && throw_range_error(name, x) + return +end + +# the product of positive `dims`, or `cap` if it is larger, without overflowing +@inline function capped_prod(dims::Dims, cap::Int) + p = 1 + for d in dims + d > cap ÷ p && return cap + p *= d + end + return p +end + +# whether the product of positive `dims` exceeds `limit`, without overflowing +@inline function prod_exceeds(dims::Dims, limit::Int) + p = 1 + for d in dims + d > limit ÷ p && return true + p *= d + end + return false +end -Validate the launch configuration passed to a [`Kernel`](@ref), throwing an -`ArgumentError` if either argument has more than 3 dimensions, or if `ndrange` -and `numgroups` are both defined. +""" + launch_geometry(kernel, numgroups, workgroupsize, ndrange, max_work_group_size)::Tuple{Dims{3}, Dims{3}} -Backends may call this from their kernel-launch method instead of writing their -own check. +Validate the launch keywords of a [`Kernel`](@ref) call and turn them into the number of +work-groups and the work-group size. A zero number of work-groups means nothing is launched. """ -function check_launch_args(numgroups, workgroupsize, ndrange) - length(ndrange) > 0 && length(numgroups) > 0 && +@inline function launch_geometry(kernel::Kernel, numgroups, workgroupsize, ndrange, max_work_group_size) + check_dims("numgroups", numgroups) + check_dims("workgroupsize", workgroupsize) + check_dims("ndrange", ndrange) + if ndrange != () && numgroups != () throw(ArgumentError("Only one of `numgroups` and `ndrange` can be used")) - length(numgroups) <= 3 || - throw(ArgumentError("`numgroups` only accepts up to 3 dimensions")) - length(workgroupsize) <= 3 || - throw(ArgumentError("`workgroupsize` only accepts up to 3 dimensions")) - length(ndrange) <= 3 || - throw(ArgumentError("`ndrange` only accepts up to 3 dimensions")) + end + max_work_group_size > 0 || + throw(ArgumentError("`max_work_group_size` must be positive, got $max_work_group_size")) + + items = if workgroupsize != () + wgsize = pad3(workgroupsize) + any(iszero, wgsize) && + throw(ArgumentError("`workgroupsize` must be positive, got $(repr(workgroupsize))")) + check_work_group_size(kernel, wgsize) + wgsize + elseif ndrange == () + (1, 1, 1) + elseif any(iszero, ndrange) + return (0, 0, 0), (1, 1, 1) + else + wanted = pad3(ndrange) + config = launch_configuration( + kernel; nitems = capped_prod(wanted, typemax(Int)), + max_work_group_size = min(max_work_group_size, typemax(Int)) % Int + ) + threads_to_workgroupsize(config.workgroupsize, wanted, max_work_group_dims(kernel.backend)) + end + + groups = if ndrange != () + cld.(pad3(ndrange), items) + elseif numgroups != () + pad3(numgroups) + else + (1, 1, 1) + end + + # the global size has to fit an `Int`, as `get_global_size()` returns it + if !any(iszero, groups) && any(map((g, i) -> g > typemax(Int) ÷ i, groups, items)) + throw(ArgumentError("Launch of $groups work-groups of $items work-items has more than typemax(Int) work-items in a dimension")) + end + return groups, items +end + +function check_work_group_size(kernel::Kernel, items::Dims{3}) + max_dims = max_work_group_dims(kernel.backend) + all(items .<= max_dims) || + throw(ArgumentError("Work-group size $items exceeds the maximum of $max_dims per dimension")) + max_items = max_work_group_size(kernel) + prod_exceeds(items, max_items) && + throw(ArgumentError("Work-group size $items has more than $max_items work-items, the maximum for this kernel")) return end @@ -63,6 +193,9 @@ dimension first. Dimension `d` gets at most `limits[d]` work-items; dimensions p of `limits` are only bounded by `threads`. Every dimension gets at least one work-item, even for a zero-sized `ndrange`. + +Not part of the public interface; used by KernelInterface's and KernelAbstractions' launch +code. """ threads_to_workgroupsize(threads, ndrange::Tuple, limits = ()) = _threads_to_workgroupsize(threads, 1, ndrange, limits) @@ -77,38 +210,6 @@ function _threads_to_workgroupsize(threads, total, ndrange::Tuple, limits) return (x, _threads_to_workgroupsize(threads, total * x, Base.tail(ndrange), rest)...) end -""" - auto_launch_sizes(kernel::KI.Kernel, numgroups, workgroupsize, ndrange, [max_work_items]) - -Returns a suggested `numgroups` and `workgroupsize` based on -the input arguments. This function assumes arguments have been -validated by `check_launch_args`. - -If any `ndrange` dimension is zero, the returned `numgroups` is zero in -that dimension; backends should skip the launch in that case. Note that very -large `ndrange`s can produce total grid sizes >= 2^32, which is problematic -on some backends. - -Backends may call this from their kernel-launch method instead of -writing their own heuristic for calculating launch size. -""" -@inline function auto_launch_sizes(kernel::Kernel, numgroups, workgroupsize, ndrange, max_work_items = typemax(Int)) - numgroups, workgroupsize = if ndrange == () - numgroups == () ? 1 : numgroups, workgroupsize == () ? 1 : workgroupsize - else - workgroupsize = if workgroupsize == () - config = launch_configuration(kernel; nitems = prod(ndrange), max_work_group_size = max_work_items) - threads_to_workgroupsize(config.workgroupsize, ndrange, max_work_group_dims(kernel.backend)) - else - workgroupsize - end - numgroups = cld.(ndrange, workgroupsize) - Int.(numgroups), Int.(workgroupsize) - end - - return numgroups, workgroupsize -end - ## limits and advice @@ -138,7 +239,7 @@ function max_work_group_size end The recommended number of work-items per work-group for launching `kernel` over `nitems` work-items in total (`nothing` if unknown), at most `max_work_group_size`. This is what an `ndrange` launch without a `workgroupsize` uses, passing the number of work-items in the -`ndrange`. `nitems` and `max_work_group_size` are positive. +`ndrange` (saturated at `typemax(Int)`). `nitems` and `max_work_group_size` are positive. Unlike [`max_work_group_size`](@ref), this is advice: backends may base it on occupancy or on the size of the launch, e.g. to prefer more work-groups over larger ones. @@ -180,9 +281,10 @@ function max_work_group_dims end The maximum number of work-groups along each dimension of a launch, for the active device of `backend`. -This is conservative: a launch within these limits works for any work-group size, but some -backends accept more work-groups for smaller work-groups (e.g. HIP bounds the number of -work-items per dimension). The backend's validation at launch time is authoritative. +This is conservative: a launch within these limits works for any work-group size (as long +as the number of work-items in each dimension fits an `Int`), but some backends accept more +work-groups for smaller work-groups (e.g. HIP bounds the number of work-items per +dimension). The backend's validation at launch time is authoritative. !!! note Backend implementations **must** implement: @@ -238,10 +340,10 @@ converting them to their device side representation. function argconvert end """ - KI.kernel_function(::NewBackend, f::F, tt::TT=Tuple{}; name=nothing, kwargs...) where {F,TT} + kernel_function(backend, f::F, tt::TT=Tuple{}; name=nothing, kwargs...)::Kernel -Low-level interface to compile a function invocation for the currently-active GPU, returning -a callable kernel object. For a higher-level interface, use +Compile the function `f` for arguments of the (device-side) types `tt`, for the active +device of `backend`, returning a [`Kernel`](@ref). For a higher-level interface, use [`KernelInterface.@launch`](@ref). Keyword arguments: @@ -253,7 +355,7 @@ CUDA.jl); backends throw an error for options they don't support. !!! note Backend implementations **must** implement: ``` - kernel_function(::NewBackend, f::F, tt::TT=Tuple{}; name=nothing, kwargs...) where {F,TT} + kernel_function(backend::NewBackend, f::F, tt::TT=Tuple{}; name=nothing, kwargs...) where {F,TT} ``` """ function kernel_function end diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl index 6d24983c2..96a91e6c3 100644 --- a/lib/KernelInterface/test/interface.jl +++ b/lib/KernelInterface/test/interface.jl @@ -1,6 +1,21 @@ import KernelInterface as KI using Random +# Counts every work-item at the element it identifies, computed from the group and local +# ids, so that a mix-up between group counts and group sizes shows. +function launch_kernel(arr) + l = KI.get_local_id() + g = KI.get_group_id() + s = KI.get_local_size() + i = (g.x - 1) * s.x + l.x + j = (g.y - 1) * s.y + l.y + k = (g.z - 1) * s.z + l.z + if i <= size(arr, 1) && j <= size(arr, 2) && k <= size(arr, 3) + @inbounds arr[i, j, k] += 1 + end + return +end + struct KernelData global_size::Int global_id::Int @@ -133,57 +148,68 @@ end function interface_testsuite(backend::KI.Backend, AT) @testset "Launch parameters" begin - # 1d - function launch_kernel1d(arr) - i, _, _ = KI.get_local_id() - gi, _, _ = KI.get_group_id() - ngi, _, _ = KI.get_num_groups() - - arr[(gi - 1) * ngi + i] = 1.0f0 - return - end - arr1d = AT(zeros(Float32, 4)) - KI.@launch backend numgroups = 2 workgroupsize = 2 launch_kernel1d(arr1d) - KI.synchronize(backend) - @test all(Array(arr1d) .== 1) - - # 1d tuple - arr1dt = AT(zeros(Float32, 4)) - KI.@launch backend numgroups = (2,) workgroupsize = (2,) launch_kernel1d(arr1dt) - KI.synchronize(backend) - @test all(Array(arr1dt) .== 1) - - # 2d - function launch_kernel2d(arr) - i, j, _ = KI.get_local_id() - gi, gj, _ = KI.get_group_id() - ngi, ngj, _ = KI.get_num_groups() - - arr[(gi - 1) * ngi + i, (gj - 1) * ngj + j] = 1.0f0 - return + # unequal group counts and sizes in every dimension, so that confusing them shows + function run(dims; kwargs...) + arr = KI.zeros(backend, Int32, dims) + kernel = KI.@launch backend launch = false launch_kernel(arr) + kernel(arr; kwargs...) + KI.synchronize(backend) + return Array(arr) end - arr2d = AT(zeros(Float32, 4, 4)) - KI.@launch backend numgroups = (2, 2) workgroupsize = (2, 2) launch_kernel2d(arr2d) + @test all(==(1), run((6, 1, 1); numgroups = 3, workgroupsize = 2)) + @test all(==(1), run((6, 1, 1); numgroups = (2,), workgroupsize = (3,))) + @test all(==(1), run((6, 10, 1); numgroups = (2, 5), workgroupsize = (3, 2))) + @test all(==(1), run((6, 10, 12); numgroups = (2, 5, 3), workgroupsize = (3, 2, 4))) + + # `ndrange` rounds up to whole groups, and doesn't mask the padding + @test all(==(1), run((7, 5, 3); ndrange = (7, 5, 3), workgroupsize = (2, 3, 2))) + @test all(==(1), run((7, 5, 3); ndrange = (7, 5, 3))) + @test all(==(1), run((1000, 1, 1); ndrange = 1000)) + + # defaults: one work-group of one work-item + @test run((2, 2, 1)) == reshape(Int32[1, 0, 0, 0], 2, 2, 1) + @test run((4, 1, 1); numgroups = 2) == reshape(Int32[1, 1, 0, 0], 4, 1, 1) + @test run((4, 1, 1); workgroupsize = 2) == reshape(Int32[1, 1, 0, 0], 4, 1, 1) + + # nothing to launch + @test all(==(0), run((2, 2, 2); ndrange = 0)) + @test all(==(0), run((2, 2, 2); ndrange = (2, 0))) + @test all(==(0), run((2, 2, 2); numgroups = (2, 0, 2), workgroupsize = 2)) + + # the global size is the padded ndrange + results = AT(Vector{KernelData}(undef, 12)) + kernel = KI.@launch backend launch = false test_interface_kernel(results) + kernel(results; ndrange = 10, workgroupsize = 4) KI.synchronize(backend) - @test all(Array(arr2d) .== 1) + @test all(d -> d.global_size == 12 && d.num_groups == 3, Array(results)) + end - # 3d - function launch_kernel3d(arr) - i, j, k = KI.get_local_id() - gi, gj, gk = KI.get_group_id() - ngi, ngj, ngk = KI.get_num_groups() + @testset "Launch validation" begin + arr = KI.zeros(backend, Int32, (1, 1, 1)) + kernel = KI.@launch backend launch = false launch_kernel(arr) + max_items = KI.max_work_group_size(kernel) + max_dims = KI.max_work_group_dims(backend) - arr[(gi - 1) * ngi + i, (gj - 1) * ngj + j, (gk - 1) * ngk + k] = 1.0f0 - return + @test_throws ArgumentError kernel(arr; numgroups = (2, 2, 2, 2), workgroupsize = (2, 2, 2)) + @test_throws ArgumentError kernel(arr; numgroups = (2, 2, 2), workgroupsize = (2, 2, 2, 2)) + @test_throws ArgumentError kernel(arr; ndrange = (2, 2, 2, 2)) + @test_throws ArgumentError kernel(arr; ndrange = 4, numgroups = 2) + @test_throws ArgumentError kernel(arr; workgroupsize = 0) + @test_throws ArgumentError kernel(arr; workgroupsize = (2, 0)) + @test_throws ArgumentError kernel(arr; workgroupsize = -1) + @test_throws ArgumentError kernel(arr; numgroups = -1) + @test_throws ArgumentError kernel(arr; ndrange = (4, -1)) + @test_throws ArgumentError kernel(arr; ndrange = 4.0) + @test_throws ArgumentError kernel(arr; ndrange = 4, max_work_group_size = 0) + @test_throws ArgumentError kernel(arr; workgroupsize = max_items + 1) + if max_dims[3] < max_items + @test_throws ArgumentError kernel(arr; workgroupsize = (1, 1, max_dims[3] + 1)) end - arr3d = AT(zeros(Float32, 4, 4, 4)) - KI.@launch backend numgroups = (2, 2, 2) workgroupsize = (2, 2, 2) launch_kernel3d(arr3d) KI.synchronize(backend) - @test all(Array(arr3d) .== 1) + @test Array(arr) == zeros(Int32, 1, 1, 1) - # 4d (Errors) - @test_throws ArgumentError (KI.@launch backend numgroups = (2, 2, 2, 2) workgroupsize = (2, 2, 2) launch_kernel3d(arr3d)) - @test_throws ArgumentError (KI.@launch backend numgroups = (2, 2, 2) workgroupsize = (2, 2, 2, 2) launch_kernel3d(arr3d)) + # other keywords are passed to the backend, which rejects the ones it doesn't know + @test_throws Exception kernel(arr; this_is_not_a_launch_option = 1) end @testset "Launch limits" begin @@ -247,6 +273,7 @@ function interface_testsuite(backend::KI.Backend, AT) @test KI.supports_atomics(b) isa Bool @test KI.supports_float64(b) isa Bool @test KI.functional(b) isa Union{Missing, Bool} + @test KI.multiprocessor_count(b) isa Int @test KI.device(b) isa Int @test KI.ndevices(b) isa Int @@ -306,13 +333,8 @@ function interface_testsuite(backend::KI.Backend, AT) end @testset "Basic interface functionality" begin - - @test KI.max_work_group_size(backend) isa Int - @test KI.multiprocessor_count(backend) isa Int - - # Test with small kernel workgroupsize = 4 - numgroups = 4 + numgroups = 3 N = workgroupsize * numgroups results = AT(Vector{KernelData}(undef, N)) kernel = KI.@launch backend launch = false test_interface_kernel(results) @@ -321,27 +343,13 @@ function interface_testsuite(backend::KI.Backend, AT) KI.synchronize(backend) host_results = Array(results) - - # Verify results make sense for (i, k_data) in enumerate(host_results) - - # Global IDs should be 1-based and sequential @test k_data.global_id == i - - # Global size should match our ndrange @test k_data.global_size == N - @test k_data.local_size == workgroupsize - @test k_data.num_groups == numgroups - - # Group ID should be 1-based - expected_group = div(i - 1, numgroups) + 1 - @test k_data.group_id == expected_group - - # Local ID should be 1-based within group - expected_local = ((i - 1) % workgroupsize) + 1 - @test k_data.local_id == expected_local + @test k_data.group_id == div(i - 1, workgroupsize) + 1 + @test k_data.local_id == ((i - 1) % workgroupsize) + 1 end end diff --git a/lib/KernelInterface/test/runtests.jl b/lib/KernelInterface/test/runtests.jl index 234f3751f..cb776ede4 100644 --- a/lib/KernelInterface/test/runtests.jl +++ b/lib/KernelInterface/test/runtests.jl @@ -34,7 +34,7 @@ end KI.get_sub_group_local_id, 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.argconvert, KI.kernel_function, KI.launch, # Host-side stubs: required backend methods with no sensible fallback. KI.synchronize, KI.copyto!, ] @@ -190,20 +190,6 @@ end @test capture_stdout(() -> KI._print(Val(3), " ", Val(:sym))) == "3 sym" end -@testset "check_launch_args" begin - # Validation only: valid configurations pass through without normalization. - @test KI.check_launch_args(1, 1, ()) === nothing - @test KI.check_launch_args((1, 2, 3), (1, 2, 3), ()) === nothing - @test KI.check_launch_args((), 1, 1) === nothing - @test KI.check_launch_args((), (), ()) === nothing - - @test_throws ArgumentError KI.check_launch_args((1, 2, 3, 4), 1, ()) - @test_throws ArgumentError KI.check_launch_args(1, (1, 2, 3, 4), ()) - @test_throws ArgumentError KI.check_launch_args((), 1, (1, 2, 3, 4)) - @test_throws ArgumentError KI.check_launch_args(2, 4, 2) # both numgroupsize and ndrange defined - @test_throws ArgumentError KI.check_launch_args(2, (), 2) # both numgroupsize and ndrange defined -end - @testset "threads_to_workgroupsize" begin # Fills dimensions left to right without exceeding the thread budget. @test KI.threads_to_workgroupsize(256, (1000,)) == (256,) @@ -231,13 +217,30 @@ end @test kernel.kern === :kern end -# A backend reporting a fixed workgroup-size limit, for exercising the -# auto-sizing helper. -struct SizedBackend <: KI.Backend - maxThreads::Int +# A minimal backend, recording the compilations and launches that KernelInterface asks for. +struct MockBackend <: KI.Backend + max_items::Int +end +MockBackend() = MockBackend(256) + +struct MockKernel + f::Any + tt::Any + name::Any + options::Any + launches::Vector{Any} +end + +KI.argconvert(::MockBackend, arg) = arg +function KI.kernel_function(backend::MockBackend, f, tt = Tuple{}; name = nothing, kwargs...) + return KI.Kernel(backend, MockKernel(f, tt, name, Dict(kwargs), [])) +end +function KI.launch(kernel::KI.Kernel{MockBackend}, groups::Dims{3}, items::Dims{3}, args...; kwargs...) + push!(kernel.kern.launches, (; groups, items, args, kwargs = Dict(kwargs))) + return :ignored end -KI.max_work_group_size(k::KI.Kernel{SizedBackend}) = k.backend.maxThreads -KI.max_work_group_dims(::SizedBackend) = (1024, 1024, 64) +KI.max_work_group_size(kernel::KI.Kernel{MockBackend}) = kernel.backend.max_items +KI.max_work_group_dims(::MockBackend) = (1024, 1024, 64) # ... and one recommending smaller work-groups than it can launch, like CUDA's occupancy API, # recording what it was asked @@ -253,53 +256,100 @@ function KI.launch_configuration( push!(kernel.backend.queries, (; nitems, max_work_group_size)) return (; workgroupsize = min(96, max_work_group_size)) end - -@testset "auto_launch_sizes" begin - # the per-dimension limit is respected - @test KI.auto_launch_sizes(KI.Kernel(SizedBackend(1024), nothing), (), (), (1, 1, 5000)) === - ((1, 1, 79), (1, 1, 64)) - - kernel = KI.Kernel(SizedBackend(256), nothing) +KI.launch(kernel::KI.Kernel{OccupancyBackend}, groups::Dims{3}, items::Dims{3}, args...) = + push!(kernel.kern, (groups, items)) + +@testset "launch geometry" begin + kernel = KI.kernel_function(MockBackend(), identity, Tuple{Int}) + function geometry(; kwargs...) + empty!(kernel.kern.launches) + @test kernel(1; kwargs...) === nothing + isempty(kernel.kern.launches) && return nothing + launch = only(kernel.kern.launches) + return launch.groups, launch.items + end # Without an ndrange the sizes pass through, defaulting to 1. - @test KI.auto_launch_sizes(kernel, (), (), ()) === (1, 1) - @test KI.auto_launch_sizes(kernel, 4, (), ()) === (4, 1) - @test KI.auto_launch_sizes(kernel, (), (2, 2), ()) === (1, (2, 2)) - @test KI.auto_launch_sizes(kernel, (4, 4), (2, 2), ()) === ((4, 4), (2, 2)) + @test geometry() == ((1, 1, 1), (1, 1, 1)) + @test geometry(numgroups = 4) == ((4, 1, 1), (1, 1, 1)) + @test geometry(workgroupsize = (2, 2)) == ((1, 1, 1), (2, 2, 1)) + @test geometry(numgroups = (4, 3), workgroupsize = (2, 5)) == ((4, 3, 1), (2, 5, 1)) + @test geometry(numgroups = (4, 3, 2), workgroupsize = (2, 5, 3)) == ((4, 3, 2), (2, 5, 3)) # With an ndrange and no workgroupsize, the workgroupsize is derived from # the kernel's limit and the workgroup count covers the ndrange. - @test KI.auto_launch_sizes(kernel, (), (), (1000,)) === ((4,), (256,)) - @test KI.auto_launch_sizes(kernel, (), (), (10,)) === ((1,), (10,)) - @test KI.auto_launch_sizes(kernel, (), (), (100, 50)) === ((1, 25), (100, 2)) - @test KI.auto_launch_sizes(kernel, (), (), 1000) === (4, 256) - - # An explicit workgroupsize is kept as-is. - @test KI.auto_launch_sizes(kernel, (), (16,), (100,)) === ((7,), (16,)) - - # A zero-sized ndrange yields zero workgroups; backends skip the launch. - @test KI.auto_launch_sizes(kernel, (), (), (0,)) === ((0,), (1,)) - @test KI.auto_launch_sizes(kernel, (), (), (0, 4)) === ((0, 1), (1, 4)) - @test KI.auto_launch_sizes(kernel, (), (), 0) === (0, 1) - - # The backend's recommendation is used, not the limit, and it's told both the size of - # the launch and the cap. - occupancy = KI.Kernel(OccupancyBackend(), nothing) - @test KI.auto_launch_sizes(occupancy, (), (), (1000,)) === ((11,), (96,)) + @test geometry(ndrange = 1000) == ((4, 1, 1), (256, 1, 1)) + @test geometry(ndrange = (1000,)) == ((4, 1, 1), (256, 1, 1)) + @test geometry(ndrange = 10) == ((1, 1, 1), (10, 1, 1)) + @test geometry(ndrange = (100, 50)) == ((1, 25, 1), (100, 2, 1)) + @test geometry(ndrange = 1000, max_work_group_size = 100) == ((10, 1, 1), (100, 1, 1)) + # ... also for ndranges with more elements than an `Int` can count + @test geometry(ndrange = (2^40, 2^40)) == ((2^32, 2^40, 1), (256, 1, 1)) + # ... respecting the per-dimension limit + let k = KI.kernel_function(MockBackend(1024), identity) + k(; ndrange = (1, 1, 5000)) + launch = only(k.kern.launches) + @test (launch.groups, launch.items) == ((1, 1, 79), (1, 1, 64)) + end + + # An explicit workgroupsize is kept as-is, and the ndrange rounded up to it. + @test geometry(ndrange = 100, workgroupsize = 16) == ((7, 1, 1), (16, 1, 1)) + @test geometry(ndrange = (7, 5), workgroupsize = (2, 3)) == ((4, 2, 1), (2, 3, 1)) + + # Zero anywhere in the ndrange or the number of groups launches nothing. + @test geometry(ndrange = 0) === nothing + @test geometry(ndrange = (0, 4)) === nothing + @test geometry(ndrange = (4, 0), workgroupsize = 2) === nothing + @test geometry(numgroups = (2, 0)) === nothing + @test geometry(ndrange = (0, typemax(Int)), workgroupsize = (1, 2)) === nothing + + # Invalid launches are rejected before the backend sees them. + for kwargs in [ + (; numgroups = (1, 1, 1, 1)), (; workgroupsize = (1, 1, 1, 1)), + (; ndrange = (1, 1, 1, 1)), (; ndrange = 2, numgroups = 2), + (; workgroupsize = 0), (; workgroupsize = (1, 0)), (; workgroupsize = -1), + (; numgroups = -1), (; ndrange = -1), (; ndrange = 2.0), (; numgroups = [1]), + (; ndrange = 4, max_work_group_size = 0), (; ndrange = typemax(UInt)), + # the kernel's limit, and the per-dimension limit + (; workgroupsize = 257), (; workgroupsize = (1, 1, 65)), + # more work-items in a dimension than an `Int` can count + (; ndrange = typemax(Int), workgroupsize = 2), + (; numgroups = typemax(Int), workgroupsize = 2), + (; numgroups = (1, typemax(Int) ÷ 2 + 1), workgroupsize = (1, 2)), + ] + @test_throws ArgumentError kernel(1; kwargs...) + end + @test isempty(kernel.kern.launches) + + # Other keywords are for the backend. + kernel(1; ndrange = 4, stream = :mine) + @test last(kernel.kern.launches).kwargs == Dict(:stream => :mine) + + # Auto-sizing uses the backend's recommendation, not the limit, and tells it both the + # size of the launch and the cap. + occupancy = KI.Kernel(OccupancyBackend(), []) + occupancy(; ndrange = 1000) + @test only(occupancy.kern) == ((11, 1, 1), (96, 1, 1)) @test only(occupancy.backend.queries) == (; nitems = 1000, max_work_group_size = typemax(Int)) - @test KI.auto_launch_sizes(occupancy, (), (), (1000, 3), 64) === ((16, 3), (64, 1)) + occupancy(; ndrange = (1000, 3), max_work_group_size = 64) + @test last(occupancy.kern) == ((16, 3, 1), (64, 1, 1)) @test last(occupancy.backend.queries) == (; nitems = 3000, max_work_group_size = 64) + occupancy(; ndrange = (2^40, 2^40)) + @test last(occupancy.backend.queries).nitems == typemax(Int) + # ... while explicit sizes can go up to the limit + occupancy(; workgroupsize = 1024) + @test last(occupancy.kern) == ((1, 1, 1), (1024, 1, 1)) end @testset "launch_configuration" begin - kernel = KI.Kernel(SizedBackend(256), nothing) + kernel = KI.kernel_function(MockBackend(), identity) # the fallback recommends the limit @test KI.launch_configuration(kernel) === (; workgroupsize = 256) @test KI.launch_configuration(kernel; max_work_group_size = 100) === (; workgroupsize = 100) @test KI.launch_configuration(kernel; nitems = 10) === (; workgroupsize = 256) # backends can recommend less than the limit - occupancy = KI.Kernel(OccupancyBackend(), nothing) + occupancy = KI.Kernel(OccupancyBackend(), []) @test KI.launch_configuration(occupancy) === (; workgroupsize = 96) @test KI.max_work_group_size(occupancy) == 1024 end @@ -333,26 +383,6 @@ end @test var_exprs[2] == Expr(:..., vars[2]) end -# A minimal backend, exercising the contract `KI.@launch` expects of one. -struct MockBackend end - -struct MockKernel - f::Any - tt::Any - name::Any - options::Any - launches::Vector{Any} -end - -KI.argconvert(::MockBackend, arg) = arg -function KI.kernel_function(::MockBackend, f, tt = Tuple{}; name = nothing, kwargs...) - return MockKernel(f, tt, name, Dict(kwargs), []) -end -function (kernel::MockKernel)(args...; kwargs...) - push!(kernel.launches, (args, Dict(kwargs))) - return nothing -end - dummy(a, b) = nothing const backend_evaluations = Ref(0) @@ -365,12 +395,13 @@ end backend = MockBackend() kernel = KI.@launch backend numgroups = 2 workgroupsize = 4 dummy(1, 2.0) - @test kernel isa MockKernel - @test kernel.f === dummy - @test kernel.tt == Tuple{Int, Float64} - args, launch_kwargs = only(kernel.launches) - @test args == (1, 2.0) - @test launch_kwargs == Dict(:numgroups => 2, :workgroupsize => 4) + @test kernel isa KI.Kernel{MockBackend} + @test kernel.backend === backend + @test kernel.kern.f === dummy + @test kernel.kern.tt == Tuple{Int, Float64} + launch = only(kernel.kern.launches) + @test launch.args == (1, 2.0) + @test (launch.groups, launch.items) == ((2, 1, 1), (4, 1, 1)) # the backend expression is evaluated once backend_evaluations[] = 0 @@ -379,18 +410,18 @@ end # `launch=false` compiles only; the caller launches later. deferred = KI.@launch backend launch = false dummy(1, 2.0) - @test isempty(deferred.launches) + @test isempty(deferred.kern.launches) # Other keywords are compiler options for `kernel_function`. named = KI.@launch backend launch = false name = "mykernel" maxthreads = 32 dummy(1, 2.0) - @test named.name == "mykernel" - @test named.options == Dict(:maxthreads => 32) + @test named.kern.name == "mykernel" + @test named.kern.options == Dict(:maxthreads => 32) optioned = KI.@launch backend ndrange = 4 maxthreads = 32 dummy(1, 2.0) - @test last(only(optioned.launches)) == Dict(:ndrange => 4) + @test isempty(only(optioned.kern.launches).kwargs) # Splatted arguments are supported. splatted = KI.@launch backend launch = false dummy((1, 2.0)...) - @test splatted.tt == Tuple{Int, Float64} + @test splatted.kern.tt == Tuple{Int, Float64} @testset "errors" begin # These throw during macro expansion, so they cannot be written as a plain diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index 3c545de36..4162f6b9c 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -211,20 +211,9 @@ function KI.kernel_function(::POCLBackend, f::F, tt::TT = Tuple{}; name = nothin return KI.Kernel{POCLBackend, typeof(kern)}(POCLBackend(), kern) end -function (obj::KI.Kernel{POCLBackend})(args...; numgroups = (), workgroupsize = (), ndrange = (), max_work_group_size = typemax(Int)) - KI.check_launch_args(numgroups, workgroupsize, ndrange) - - # zero-sized ndrange: nothing to launch - prod(ndrange) == 0 && return nothing - - numgroups, workgroupsize = KI.auto_launch_sizes(obj, numgroups, workgroupsize, ndrange, max_work_group_size) - - local_size = (workgroupsize..., ntuple(_ -> 1, 3 - length(workgroupsize))...) - - numgroups = (numgroups..., ntuple(_ -> 1, 3 - length(numgroups))...) - global_size = local_size .* numgroups - - event = obj.kern(args...; local_size, global_size) +function KI.launch(obj::KI.Kernel{POCLBackend}, groups::Dims{3}, items::Dims{3}, args::Vararg{Any, N}) where {N} + # POCL launches synchronously, see the implementation note on `synchronize` + event = obj.kern(args...; local_size = items, global_size = groups .* items) wait(event) cl.clReleaseEvent(event) return nothing From 37ec17ac574b7392ec72ca5a3d43780af451154e Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 09:25:54 +0200 Subject: [PATCH 6/9] KernelInterface: specify the typed index queries as `% T` The 0.2 docs said that typed queries are "computed in `T`" and undefined when the value doesn't fit, but backends implemented them as `T(x)`, a checked conversion that leaves a `throw_inexacterror` branch in every kernel. The result is now the exact value modulo `T`, as with `x % T`, which is cheap and testable: a `UInt8` query over 384 work-items has to wrap, where `T(x)` throws. Backends implement the four primitive queries (`get_local_id`, `get_group_id`, `get_local_size`, `get_num_groups`); `get_global_id` and `get_global_size` have fallbacks derived from them for backends without a builtin (CUDA, HIP), which backends with one (SPIR-V, Metal) should override. --- docs/src/kernelinterface.md | 30 ++++--- lib/KernelInterface/src/device.jl | 110 +++++++++++++++----------- lib/KernelInterface/test/interface.jl | 46 ++++++++++- lib/KernelInterface/test/runtests.jl | 14 +++- src/pocl/backend.jl | 8 +- 5 files changed, 140 insertions(+), 68 deletions(-) diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md index dc6456ace..abb178577 100644 --- a/docs/src/kernelinterface.md +++ b/docs/src/kernelinterface.md @@ -61,7 +61,7 @@ get_backend These are called from inside a kernel. A backend provides each one with ```julia -@device_override KI.get_global_id(::Type{T}) where {T} = ... +@device_override KI.get_local_id(::Type{T}) where {T} = ... ``` along with the corresponding on-device functionality. @@ -69,17 +69,23 @@ along with the corresponding on-device functionality. ### Indexing All index queries are **1-based** and return a named tuple of `x`, `y` and `z` -components. They take an optional element type `T` for the components, defaulting +components. They take an optional integer type `T` for the components, defaulting to `Int`, so a kernel can request e.g. `Int32` indices with -`KI.get_global_id(Int32)`. +`KI.get_global_id(Int32)`. The operands are converted to `T` before any arithmetic, +and the result is the exact value modulo `T`: a query never throws, and a value that +doesn't fit wraps around, as with `x % T`. + +Backends implement the four primitive queries. [`get_global_id`](@ref) and +[`get_global_size`](@ref) have fallbacks derived from them, which backends with a native +builtin (e.g. SPIR-V and Metal) should override. ```@docs -get_global_size -get_global_id -get_local_size get_local_id -get_num_groups get_group_id +get_local_size +get_num_groups +get_global_id +get_global_size ``` ### Sub-groups @@ -201,9 +207,9 @@ A backend must, at minimum: [`adapt(backend, x)`](@ref Adapt.adapt_storage(::Backend, ::Any)) moves data to the backend, preferably by delegating to its array type: `Adapt.adapt_storage(::NewBackend, x) = adapt(NewArray, x)`. -4. `@device_override` the device-side functions it supports. The indexing - queries and [`barrier`](@ref) are required; sub-group and - [`shfl_down`](@ref) support is optional. +4. `@device_override` the device-side functions it supports. The four primitive index + queries and [`barrier`](@ref) are required; sub-group and [`shfl_down`](@ref) support + is optional. 5. Implement [`argconvert`](@ref) and [`kernel_function`](@ref) for its backend type, returning a [`Kernel`](@ref). 6. Implement [`launch`](@ref), which receives an already validated `NTuple{3, Int}` of @@ -212,7 +218,9 @@ A backend must, at minimum: KI.launch(k::KI.Kernel{CUDABackend}, groups::Dims{3}, items::Dims{3}, args::Vararg{Any, N}; kwargs...) where {N} = k.kern(args...; threads = items, blocks = groups, kwargs...) ``` -7. Report its limits through [`max_work_group_size`](@ref) (for the backend and for a +7. Compute the typed index queries with `% T`, not `T(x)`: a checked conversion leaves + an error branch in every kernel. +8. Report its limits through [`max_work_group_size`](@ref) (for the backend and for a kernel), [`max_work_group_dims`](@ref) and [`max_num_groups`](@ref), and where applicable [`sub_group_size`](@ref) and [`multiprocessor_count`](@ref). It may recommend work-group sizes with [`launch_configuration`](@ref). diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl index 8814544f1..53202c2cd 100644 --- a/lib/KernelInterface/src/device.jl +++ b/lib/KernelInterface/src/device.jl @@ -1,46 +1,51 @@ +## indexing + +# The index queries are 1-based, and take the integer type `T` of their result. Backends +# implement the four primitive ones; the global ones have fallbacks derived from them. +# +# Supported `T` are the fixed-width integer types up to 64 bits. The operands are converted +# to `T` before any arithmetic, and the result is the exact value modulo `T` (as with `x % T`): +# a query never throws, and a value that doesn't fit wraps. + """ - get_global_size([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T} + get_local_id([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T} -Return the number of global work-items specified as a tuple of type `T`. -`T` defaults to `Int`. +The 1-based index of the work-item within its work-group, as integers of type `T`. -The value is computed in `T`, and is undefined when it does not fit in `T`. +`T` is a fixed-width integer type of at most 64 bits (e.g. `Int32` or `UInt64`); the result +is the exact value modulo `T`, as if computed with `x % T`. !!! note Backend implementations **must** implement: ``` - @device_override get_global_size(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T} + @device_override get_local_id(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T} ``` - The zero-argument form forwards to `get_global_size(Int)`. + The zero-argument form forwards to `get_local_id(Int)`. """ -@inline get_global_size() = get_global_size(Int) +@inline get_local_id() = get_local_id(Int) """ - get_global_id([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T} - -Returns the unique global work-item ID as a tuple of type `T`. `T` defaults to `Int`. + get_group_id([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T} -The value is computed in `T`, and is undefined when it does not fit in `T`. +The 1-based index of the work-group within the launch, as integers of type `T`. -!!! note - 1-based. +See [`get_local_id`](@ref) for the supported types `T`. !!! note Backend implementations **must** implement: ``` - @device_override get_global_id(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T} + @device_override get_group_id(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T} ``` - The zero-argument form forwards to `get_global_id(Int)`. + The zero-argument form forwards to `get_group_id(Int)`. """ -@inline get_global_id() = get_global_id(Int) +@inline get_group_id() = get_group_id(Int) """ get_local_size([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T} -Return the number of local work-items specified as a tuple of type `T`. -`T` defaults to `Int`. +The number of work-items in a work-group, as integers of type `T`. -The value is computed in `T`, and is undefined when it does not fit in `T`. +See [`get_local_id`](@ref) for the supported types `T`. !!! note Backend implementations **must** implement: @@ -52,58 +57,73 @@ The value is computed in `T`, and is undefined when it does not fit in `T`. @inline get_local_size() = get_local_size(Int) """ - get_local_id([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T} - -Returns the unique local work-item ID as a tuple of type `T`. `T` defaults to `Int`. + get_num_groups([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T} -The value is computed in `T`, and is undefined when it does not fit in `T`. +The number of work-groups in the launch, as integers of type `T`. -!!! note - 1-based. +See [`get_local_id`](@ref) for the supported types `T`. !!! note Backend implementations **must** implement: ``` - @device_override get_local_id(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T} + @device_override get_num_groups(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T} ``` - The zero-argument form forwards to `get_local_id(Int)`. + The zero-argument form forwards to `get_num_groups(Int)`. """ -@inline get_local_id() = get_local_id(Int) +@inline get_num_groups() = get_num_groups(Int) """ - get_num_groups([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T} + get_global_id([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T} -Returns the number of groups as a tuple of type `T`. `T` defaults to `Int`. +The 1-based index of the work-item within the launch, as integers of type `T`: +`(get_group_id(T) - 1) * get_local_size(T) + get_local_id(T)` per dimension. -The value is computed in `T`, and is undefined when it does not fit in `T`. +See [`get_local_id`](@ref) for the supported types `T`. !!! note - Backend implementations **must** implement: + The fallback derives this from the primitive queries. Backend implementations with a + native builtin **should** override it, returning the same values: ``` - @device_override get_num_groups(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T} + @device_override get_global_id(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T} ``` - The zero-argument form forwards to `get_num_groups(Int)`. """ -@inline get_num_groups() = get_num_groups(Int) +@inline function get_global_id(::Type{T}) where {T} + group = get_group_id(T) + size = get_local_size(T) + local_id = get_local_id(T) + return (; + x = (group.x - one(T)) * size.x + local_id.x, + y = (group.y - one(T)) * size.y + local_id.y, + z = (group.z - one(T)) * size.z + local_id.z, + ) +end +@inline get_global_id() = get_global_id(Int) """ - get_group_id([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T} - -Returns the unique group ID as a tuple of type `T`. `T` defaults to `Int`. + get_global_size([::Type{T}=Int])::@NamedTuple{x::T, y::T, z::T} -The value is computed in `T`, and is undefined when it does not fit in `T`. +The number of work-items in the launch, as integers of type `T`: +`get_local_size(T) * get_num_groups(T)` per dimension. For an `ndrange` launch, this is the +`ndrange` padded to whole work-groups. -!!! note - 1-based. +See [`get_local_id`](@ref) for the supported types `T`. !!! note - Backend implementations **must** implement: + The fallback derives this from the primitive queries. Backend implementations with a + native builtin **should** override it, returning the same values: ``` - @device_override get_group_id(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T} + @device_override get_global_size(::Type{T})::@NamedTuple{x::T, y::T, z::T} where {T} ``` - The zero-argument form forwards to `get_group_id(Int)`. """ -@inline get_group_id() = get_group_id(Int) +@inline function get_global_size(::Type{T}) where {T} + size = get_local_size(T) + groups = get_num_groups(T) + return (; x = size.x * groups.x, y = size.y * groups.y, z = size.z * groups.z) +end +@inline get_global_size() = get_global_size(Int) + + +## sub-groups """ get_sub_group_size()::UInt32 diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl index 96a91e6c3..110047805 100644 --- a/lib/KernelInterface/test/interface.jl +++ b/lib/KernelInterface/test/interface.jl @@ -111,6 +111,23 @@ function typed_index_kernel(results, ::Type{T}) where {T} return end +# Records the `x` components of the typed queries in a type too small for the launch, +# indexed by the (untyped) global id. +function wrapping_index_kernel(results, ::Type{T}) where {T} + i = KI.get_global_id().x + if i <= size(results, 1) + @inbounds begin + results[i, 1] = KI.get_global_id(T).x + results[i, 2] = KI.get_global_size(T).x + results[i, 3] = KI.get_local_id(T).x + results[i, 4] = KI.get_local_size(T).x + results[i, 5] = KI.get_group_id(T).x + results[i, 6] = KI.get_num_groups(T).x + end + end + return +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, 0x00000001) @@ -305,8 +322,8 @@ function interface_testsuite(backend::KI.Backend, AT) end @testset "Typed indexing" begin - workgroupsize = (2, 2, 2) - numgroups = (3, 2, 1) + workgroupsize = (2, 3, 2) + numgroups = (3, 2, 4) N = prod(workgroupsize) * prod(numgroups) # `Int` is the reference: it is what the zero-argument form returns. @@ -322,16 +339,37 @@ function interface_testsuite(backend::KI.Backend, AT) @test all(eachrow(reference[:, 1:3]) .== Ref(collect(global_size))) @test all(eachrow(reference[:, 7:9]) .== Ref(collect(workgroupsize))) @test all(eachrow(reference[:, 13:15]) .== Ref(collect(numgroups))) - # every global id is seen exactly once + # every global id is seen exactly once, and agrees with the group and local ids @test sort(Tuple.(eachrow(reference[:, 4:6]))) == sort(vec(Tuple.(CartesianIndices(global_size)))) + @test reference[:, 4:6] == (reference[:, 16:18] .- 1) .* reference[:, 7:9] .+ reference[:, 10:12] - @testset "$T" for T in (Int32, UInt32, UInt64) + @testset "$T" for T in (Int32, UInt32, Int64, UInt64) typed = run_typed(T) @test typed isa AbstractMatrix{T} @test typed == reference end end + @testset "Typed indexing wraps" begin + # a type too small for the launch wraps around (like `x % T`) instead of throwing + @testset "$T" for (T, workgroupsize, numgroups) in ( + (UInt8, 128, 3), (Int8, 128, 3), (Int16, 256, 160), + ) + N = workgroupsize * numgroups + results = KI.zeros(backend, T, N, 6) + KI.@launch backend workgroupsize = workgroupsize numgroups = numgroups wrapping_index_kernel(results, T) + KI.synchronize(backend) + results = Array(results) + ids = 1:N + @test results[:, 1] == ids .% T + @test all(==(N % T), results[:, 2]) + @test results[:, 3] == mod1.(ids, workgroupsize) .% T + @test all(==(workgroupsize % T), results[:, 4]) + @test results[:, 5] == cld.(ids, workgroupsize) .% T + @test all(==(numgroups % T), results[:, 6]) + end + end + @testset "Basic interface functionality" begin workgroupsize = 4 numgroups = 3 diff --git a/lib/KernelInterface/test/runtests.jl b/lib/KernelInterface/test/runtests.jl index cb776ede4..55fecc092 100644 --- a/lib/KernelInterface/test/runtests.jl +++ b/lib/KernelInterface/test/runtests.jl @@ -42,20 +42,26 @@ end @test isempty(methods(stub)) end - # The indexing queries take an element type; only the zero-argument form has a + # 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. - indexing = [ - KI.get_global_size, KI.get_global_id, + primitives = [ KI.get_local_size, KI.get_local_id, KI.get_num_groups, KI.get_group_id, ] - for f in indexing + for f in primitives @test length(methods(f)) == 1 @test hasmethod(f, Tuple{}) @test !hasmethod(f, Tuple{Type{Int}}) @test_throws MethodError f() @test_throws MethodError f(Int32) end + + # The global queries are derived from the primitive ones. + for f in [KI.get_global_size, KI.get_global_id] + @test hasmethod(f, Tuple{Type{Int}}) + @test_throws MethodError f() + @test_throws MethodError f(Int32) + end end struct StubBackend <: KI.Backend end diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index 4162f6b9c..cc6f60ba1 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -284,10 +284,6 @@ end return (; x = get_group_id(1) % T, y = get_group_id(2) % T, z = get_group_id(3) % T) end -@device_override @inline function KI.get_global_id(::Type{T}) where {T} - return (; x = get_global_id(1) % T, y = get_global_id(2) % T, z = get_global_id(3) % T) -end - @device_override @inline function KI.get_local_size(::Type{T}) where {T} return (; x = get_local_size(1) % T, y = get_local_size(2) % T, z = get_local_size(3) % T) end @@ -296,6 +292,10 @@ end return (; x = get_num_groups(1) % T, y = get_num_groups(2) % T, z = get_num_groups(3) % T) end +@device_override @inline function KI.get_global_id(::Type{T}) where {T} + return (; x = get_global_id(1) % T, y = get_global_id(2) % T, z = get_global_id(3) % T) +end + @device_override @inline function KI.get_global_size(::Type{T}) where {T} return (; x = get_global_size(1) % T, y = get_global_size(2) % T, z = get_global_size(3) % T) end From a63762a0bfb309dd31b8f9843e4c97fd5f5ca237 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 09:28:52 +0200 Subject: [PATCH 7/9] KernelInterface: specify sub-groups Sub-groups were underspecified exactly where backends disagree: - Support is a capability: `supports_subgroups(backend)` and `supports_shuffle(backend, T)` replace `shfl_down_types`, whose non-empty result the testsuite used as the flag for sub-group support. - `sub_group_size(backend)` is a guarantee instead of "a reasonable size": kernels from `kernel_function` run with exactly that width, so host code can pick a `Val(N)` for a warp-level reduction. A backend that can't guarantee it reports no sub-group support. PoCL fixes the width when compiling. - Partial sub-groups follow OpenCL and SYCL: `get_sub_group_size()` counts the work-items that are present, and `get_num_sub_groups()` is `cld(items, width)`. - How work-items are assigned to sub-groups is unspecified, but each has a unique `(sub-group id, lane)` pair that doesn't change during the kernel. - A `shfl_down` from a lane that doesn't exist returns an unspecified value, and shuffles are not memory fences. - The sub-group queries take a result type like the index queries, defaulting to `Int`; they returned `UInt32` on most backends but `Int32` on CUDA. The testsuite checks partial sub-groups, lane uniqueness, `shfl_down` on every lane, and that `sub_group_barrier` makes memory writes visible. --- docs/src/kernelinterface.md | 12 +- lib/KernelInterface/src/device.jl | 121 ++++++------ lib/KernelInterface/src/host.jl | 26 +++ lib/KernelInterface/src/launch.jl | 16 +- lib/KernelInterface/test/interface.jl | 256 ++++++++++++++++++-------- lib/KernelInterface/test/runtests.jl | 15 +- src/pocl/backend.jl | 60 +++--- test/hostinterface.jl | 7 +- 8 files changed, 332 insertions(+), 181 deletions(-) diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md index abb178577..a1fc1debf 100644 --- a/docs/src/kernelinterface.md +++ b/docs/src/kernelinterface.md @@ -90,6 +90,12 @@ get_global_size ### Sub-groups +Sub-groups are optional ([`supports_subgroups`](@ref)). A work-group is divided into +sub-groups of [`sub_group_size(backend)`](@ref sub_group_size) work-items, the last of which +can be partial. How work-items are assigned to sub-groups is unspecified, but 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. + ```@docs get_sub_group_size get_max_sub_group_size @@ -115,7 +121,6 @@ localmemory ```@docs shfl_down -shfl_down_types ``` ### Printing @@ -171,6 +176,11 @@ supports_atomics supports_float64 ``` +```@docs +supports_subgroups +supports_shuffle +``` + ### Limits ```@docs diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl index 53202c2cd..3d7cb0b71 100644 --- a/lib/KernelInterface/src/device.jl +++ b/lib/KernelInterface/src/device.jl @@ -125,76 +125,93 @@ end ## sub-groups +# Sub-group support is optional, see `supports_subgroups`. A work-group is divided into +# sub-groups of `sub_group_size(backend)` work-items. How work-items are assigned to +# sub-groups 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. + """ - get_sub_group_size()::UInt32 + get_sub_group_size([::Type{T}=Int])::T + +The number of work-items in the sub-group: the sub-group width +([`get_max_sub_group_size`](@ref)), or fewer for the last sub-group of a work-group whose +size isn't a multiple of the width. -Returns the number of work-items in the sub-group. +See [`get_local_id`](@ref) for the supported types `T`. !!! note - Backend implementations **must** implement: + Backend implementations that support sub-groups **must** implement: ``` - @device_override get_sub_group_size()::UInt32 + @device_override get_sub_group_size(::Type{T})::T where {T} ``` + The zero-argument form forwards to `get_sub_group_size(Int)`. """ -function get_sub_group_size end +@inline get_sub_group_size() = get_sub_group_size(Int) """ - get_max_sub_group_size()::UInt32 + get_max_sub_group_size([::Type{T}=Int])::T -Returns the maximum sub-group size for sub-groups in the current workgroup. +The sub-group width, [`sub_group_size(backend)`](@ref sub_group_size) on the host. + +See [`get_local_id`](@ref) for the supported types `T`. !!! note - Backend implementations **must** implement: + Backend implementations that support sub-groups **must** implement: ``` - @device_override get_max_sub_group_size()::UInt32 + @device_override get_max_sub_group_size(::Type{T})::T where {T} ``` + The zero-argument form forwards to `get_max_sub_group_size(Int)`. """ -function get_max_sub_group_size end +@inline get_max_sub_group_size() = get_max_sub_group_size(Int) """ - get_num_sub_groups()::UInt32 + get_num_sub_groups([::Type{T}=Int])::T -Returns the number of sub-groups in the current workgroup. +The number of sub-groups in the work-group: `cld(prod(get_local_size()), get_max_sub_group_size())`. + +See [`get_local_id`](@ref) for the supported types `T`. !!! note - Backend implementations **must** implement: + Backend implementations that support sub-groups **must** implement: ``` - @device_override get_num_sub_groups()::UInt32 + @device_override get_num_sub_groups(::Type{T})::T where {T} ``` + The zero-argument form forwards to `get_num_sub_groups(Int)`. """ -function get_num_sub_groups end +@inline get_num_sub_groups() = get_num_sub_groups(Int) """ - get_sub_group_id()::UInt32 + get_sub_group_id([::Type{T}=Int])::T -Returns the sub-group ID within the work-group. +The 1-based index of the sub-group within the work-group. -!!! note - 1-based. +See [`get_local_id`](@ref) for the supported types `T`. !!! note - Backend implementations **must** implement: + Backend implementations that support sub-groups **must** implement: ``` - @device_override get_sub_group_id()::UInt32 + @device_override get_sub_group_id(::Type{T})::T where {T} ``` + The zero-argument form forwards to `get_sub_group_id(Int)`. """ -function get_sub_group_id end +@inline get_sub_group_id() = get_sub_group_id(Int) """ - get_sub_group_local_id()::UInt32 + get_sub_group_local_id([::Type{T}=Int])::T -Returns the work-item ID within the current sub-group. +The 1-based index of the work-item within its sub-group (its lane). It doesn't depend on +which work-items of the sub-group are active, e.g. in a divergent branch. -!!! note - 1-based. +See [`get_local_id`](@ref) for the supported types `T`. !!! note - Backend implementations **must** implement: + Backend implementations that support sub-groups **must** implement: ``` - @device_override get_sub_group_local_id()::UInt32 + @device_override get_sub_group_local_id(::Type{T})::T where {T} ``` + The zero-argument form forwards to `get_sub_group_local_id(Int)`. """ -function get_sub_group_local_id end +@inline get_sub_group_local_id() = get_sub_group_local_id(Int) """ @@ -218,38 +235,26 @@ localmemory(::Type{T}, ::Val) where {T} = """ - shfl_down(val::T, offset::Integer) where T + shfl_down(val::T, offset::Integer)::T -Read `val` from a lane with higher id given by `offset`. +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`). -!!! note - `shfl_down` must be encountered by all workitems of a sub-group executing the kernel or by none at all. +All work-items of the sub-group have to execute `shfl_down` together (not in a divergent +branch), with the same `offset`. + +`shfl_down` exchanges values, not memory: it is not a memory fence. !!! note - Backend implementations **must** implement: + Backend implementations **must** implement this for every `T` for which + [`supports_shuffle`](@ref) returns `true`: ``` @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() @@ -277,18 +282,14 @@ 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. +Like [`barrier`](@ref), for the work-items of a sub-group: wait until all work-items of the +sub-group have reached the barrier, and make their writes to global and local memory +before it visible to 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. +All work-items of a sub-group have to reach the same `sub_group_barrier()`. !!! 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: + Backend implementations that support sub-groups **must** implement: ``` @device_override sub_group_barrier() ``` diff --git a/lib/KernelInterface/src/host.jl b/lib/KernelInterface/src/host.jl index 419698eb2..ee105e3e0 100644 --- a/lib/KernelInterface/src/host.jl +++ b/lib/KernelInterface/src/host.jl @@ -236,6 +236,32 @@ Returns whether `Float64` values are supported by the backend. """ supports_float64(::Backend) = true +""" + 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). + +Which types [`shfl_down`](@ref) supports is queried separately with [`supports_shuffle`](@ref). + +!!! note + Backend implementations **must** implement this function if they support sub-groups. + The fallback returns `false`. +""" +supports_subgroups(::Backend) = false + +""" + supports_shuffle(::Backend, ::Type{T})::Bool + +Whether kernels on the active device support [`shfl_down`](@ref) for values of type `T`. + +!!! note + Backend implementations **must** implement this function for the types they support. + The fallback returns `false`. +""" +supports_shuffle(::Backend, ::Type) = false + """ allocate(::Backend, Type, dims...; unified=false)::AbstractArray diff --git a/lib/KernelInterface/src/launch.jl b/lib/KernelInterface/src/launch.jl index 42869a567..3c023d065 100644 --- a/lib/KernelInterface/src/launch.jl +++ b/lib/KernelInterface/src/launch.jl @@ -297,16 +297,20 @@ function max_num_groups 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. +The sub-group width of kernels compiled for the active device of `backend`. + +Kernels compiled by [`kernel_function`](@ref) execute with exactly this width: on the device, +[`get_max_sub_group_size`](@ref) returns it, and full sub-groups have this many work-items. +Host code can rely on it, e.g. to pick a `Val(N)` for a warp-level reduction. !!! note - Backend implementations **must** implement: + Backend implementations **must** implement this if [`supports_subgroups`](@ref) returns + `true`: ``` sub_group_size(backend::NewBackend)::Int ``` - As well as the on-device functionality. + A backend that cannot guarantee the width for every kernel has to report + `supports_subgroups(backend) = false`. """ function sub_group_size end @@ -357,6 +361,8 @@ CUDA.jl); backends throw an error for options they don't support. ``` kernel_function(backend::NewBackend, f::F, tt::TT=Tuple{}; name=nothing, kwargs...) where {F,TT} ``` + Kernels must execute with sub-group width [`sub_group_size(backend)`](@ref sub_group_size) + if the backend supports sub-groups. """ function kernel_function end diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl index 110047805..d659377a7 100644 --- a/lib/KernelInterface/test/interface.jl +++ b/lib/KernelInterface/test/interface.jl @@ -39,28 +39,6 @@ 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 - # The interface documents a concrete return type for each device-side function; # these kernels record whether the backend honors them. const WorkItemNT{T} = @NamedTuple{x::T, y::T, z::T} @@ -128,17 +106,46 @@ function wrapping_index_kernel(results, ::Type{T}) where {T} return end +struct SubgroupData + sub_group_size::Int + max_sub_group_size::Int + num_sub_groups::Int + sub_group_id::Int + sub_group_local_id::Int +end +function test_subgroup_kernel(results) + l = KI.get_local_id() + s = KI.get_local_size() + i = (l.y - 1) * s.x + l.x + (KI.get_group_id().x - 1) * s.x * s.y + + 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 subgroup_typecheck_kernel(results, val::T) where {T} # uniformly executed by the whole sub-group, as `shfl_down` requires - shuffled = KI.shfl_down(val, 0x00000001) + shuffled = KI.shfl_down(val, 1) if KI.get_sub_group_local_id() == 1 @inbounds begin - results[1] = KI.get_sub_group_size() isa UInt32 - results[2] = KI.get_max_sub_group_size() isa UInt32 - results[3] = KI.get_num_sub_groups() isa UInt32 - results[4] = KI.get_sub_group_id() isa UInt32 - results[5] = KI.get_sub_group_local_id() isa UInt32 + results[1] = KI.get_sub_group_size() isa Int + results[2] = KI.get_max_sub_group_size() isa Int + results[3] = KI.get_num_sub_groups() isa Int + results[4] = KI.get_sub_group_id() isa Int + results[5] = KI.get_sub_group_local_id() isa Int results[6] = shuffled isa T + results[7] = KI.get_sub_group_size(UInt32) isa UInt32 + results[8] = KI.get_max_sub_group_size(Int32) isa Int32 + results[9] = KI.get_num_sub_groups(UInt32) isa UInt32 + results[10] = KI.get_sub_group_id(Int32) isa Int32 + results[11] = KI.get_sub_group_local_id(UInt32) isa UInt32 end end return @@ -149,9 +156,13 @@ function shfl_down_test_kernel(a, b, ::Val{N}) where {N} val = a[idx] - offset = 0x00000001 + # the result of shuffling from a lane past the end is unspecified, so don't add it + offset = 1 while offset < N - val += KI.shfl_down(val, offset) + shuffled = KI.shfl_down(val, offset) + if idx + offset <= N + val += shuffled + end offset <<= 1 end @@ -163,6 +174,34 @@ function shfl_down_test_kernel(a, b, ::Val{N}) where {N} 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} + 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 + 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} + lane = KI.get_sub_group_local_id() + other = mod1(lane + 1, KI.get_sub_group_size()) + 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] + 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 @@ -297,7 +336,8 @@ function interface_testsuite(backend::KI.Backend, AT) # @test KI.device!(b, KI.device(b)) isa Nothing # @test KI.priority!(b, :normal) isa Nothing - @test KI.shfl_down_types(b) isa Vector{DataType} + @test KI.supports_subgroups(b) isa Bool + @test KI.supports_shuffle(b, Float32) isa Bool arr = KI.allocate(b, Float32, 2) @test arr isa AT{Float32, 1} @@ -391,23 +431,47 @@ function interface_testsuite(backend::KI.Backend, AT) end end - # Used as a proxy for sub-group support - if !isempty(KI.shfl_down_types(backend)) + if KI.supports_subgroups(backend) + sg_size = KI.sub_group_size(backend) + max_dims = KI.max_work_group_dims(backend) + # 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 return types" begin - @test KI.sub_group_size(backend) isa Int + @test sg_size isa Int && sg_size >= 1 + + types = filter(T -> KI.supports_shuffle(backend, T), [Int32, Float32, Int64, UInt32]) + if !isempty(types) + results = KI.zeros(backend, Bool, 11) + val = one(first(types)) + kernel = KI.@launch backend launch = false subgroup_typecheck_kernel(results, val) + if fits(kernel, (sg_size,)) + kernel(results, val; workgroupsize = sg_size) + KI.synchronize(backend) + @test all(Array(results)) + else + @test_skip "work-groups of $sg_size work-items" + end + end + end - T = first(setdiff(KI.shfl_down_types(backend), [Bool])) - results = KI.zeros(backend, Bool, 6) - KI.@launch backend workgroupsize = KI.sub_group_size(backend) subgroup_typecheck_kernel(results, one(T)) - KI.synchronize(backend) - @test all(Array(results)) + # checks the sub-groups of a work-group of `items` work-items + function check_subgroups(data, items) + @test all(d -> d.max_sub_group_size == sg_size, data) + @test all(d -> d.num_sub_groups == cld(items, sg_size), data) + @test all(d -> 1 <= d.sub_group_id <= cld(items, sg_size), data) + @test all(d -> 1 <= d.sub_group_local_id <= d.sub_group_size, data) + # every work-item has its own (sub-group, lane) pair + @test allunique(map(d -> (d.sub_group_id, d.sub_group_local_id), data)) + # each sub-group has as many members as its size says + for id in unique(map(d -> d.sub_group_id, data)) + members = filter(d -> d.sub_group_id == id, data) + @test all(d -> d.sub_group_size == length(members), members) + end + return end @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 numgroups = 2 @@ -415,42 +479,90 @@ function interface_testsuite(backend::KI.Backend, AT) results = AT(Vector{SubgroupData}(undef, N)) kernel = KI.@launch backend launch = false test_subgroup_kernel(results) + if fits(kernel, (workgroupsize,)) + kernel(results; workgroupsize, numgroups) + KI.synchronize(backend) + + host_results = Array(results) + @test all(d -> d.sub_group_size == sg_size, host_results) + for group in Iterators.partition(host_results, workgroupsize) + check_subgroups(collect(group), workgroupsize) + end + else + @test_skip "work-groups of $workgroupsize work-items" + end + end - kernel(results; workgroupsize, numgroups) - KI.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 + @testset "Partial sub-groups" begin + # a 2-D work-group whose size isn't a multiple of the sub-group size, or else a + # 1-D one with one work-item more than a sub-group + numgroups = 2 + results = AT(Vector{SubgroupData}(undef, max(66, sg_size + 1) * numgroups)) + kernel = KI.@launch backend launch = false test_subgroup_kernel(results) + workgroupsize = fits(kernel, (33, 2)) ? (33, 2) : (sg_size + 1,) + items = prod(workgroupsize) + if fits(kernel, workgroupsize) && items % sg_size != 0 + kernel(results; workgroupsize, numgroups) + KI.synchronize(backend) + + host_results = Array(results)[1:(items * numgroups)] + for group in Iterators.partition(host_results, items) + group = collect(group) + check_subgroups(group, items) + # the sizes of the sub-groups add up to the work-group, with one partial one + sizes = Dict(d.sub_group_id => d.sub_group_size for d in group) + @test sum(values(sizes)) == items + @test count(<(sg_size), values(sizes)) == (items % sg_size == 0 ? 0 : 1) + end + else + @test_skip "work-groups of $workgroupsize work-items" + end + end - # 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 + @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 - @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) + 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, N)) - - KI.@launch backend workgroupsize = N shfl_down_test_kernel(dev_a, dev_b, Val(N)) - - b = Array(dev_b) - @test sum(a) ≈ b[1] + 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) + @test length(in_range) == N - (N ÷ sg_size) * offset + @test out[in_range, 3] == out[in_range, 1] .+ offset + end end end end diff --git a/lib/KernelInterface/test/runtests.jl b/lib/KernelInterface/test/runtests.jl index 55fecc092..804473dfb 100644 --- a/lib/KernelInterface/test/runtests.jl +++ b/lib/KernelInterface/test/runtests.jl @@ -29,12 +29,9 @@ end # 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.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, 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, + 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. KI.synchronize, KI.copyto!, ] @@ -47,6 +44,9 @@ end primitives = [ KI.get_local_size, KI.get_local_id, KI.get_num_groups, KI.get_group_id, + 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, ] for f in primitives @test length(methods(f)) == 1 @@ -86,9 +86,10 @@ end @test_throws "used outside kernel" KI.barrier() @test_throws "used outside kernel" KI.sub_group_barrier() - # Permissive defaults: a backend only implements these if it can do better. - @test KI.shfl_down_types(StubBackend()) == DataType[] + # Conservative defaults: a backend only implements these if it can do better. @test KI.multiprocessor_count(StubBackend()) == 0 + @test KI.supports_subgroups(StubBackend()) === false + @test KI.supports_shuffle(StubBackend(), Float32) === false # `localmemory` forwards the untyped `dims` to the `Val` form backends override. # Off-device that form is unimplemented, and must error rather than recurse diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index cc6f60ba1..acd1a5cdf 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -207,7 +207,13 @@ end KI.argconvert(::POCLBackend, arg) = clconvert(arg) function KI.kernel_function(::POCLBackend, f::F, tt::TT = Tuple{}; name = nothing, kwargs...) where {F, TT} - kern = clfunction(f, tt; name, kwargs...) + # fix the sub-group width, as `KI.sub_group_size` promises + sub_group_size = device_limits().sub_group_size + kern = if sub_group_size > 0 + clfunction(f, tt; name, sub_group_size, kwargs...) + else + clfunction(f, tt; name, kwargs...) + end return KI.Kernel{POCLBackend, typeof(kern)}(POCLBackend(), kern) end @@ -228,47 +234,33 @@ function device_limits() return get!(task_local_storage(), :POCLLimits) do dev = device() sizes = dev.max_work_item_size + # POCL can technically support any sub-group size; prefer the common GPU ones + sg_sizes = dev.sub_group_sizes + common = filter(in(sg_sizes), [32, 64, 16, sg_sizes...]) (; max_work_group_size = Int(dev.max_work_group_size), max_work_group_dims = ntuple(d -> d <= length(sizes) ? sizes[d] : 1, 3), + # 0 if the device has no sub-groups + sub_group_size = isempty(common) ? 0 : first(common), ) - end::@NamedTuple{max_work_group_size::Int, max_work_group_dims::NTuple{3, Int}} + end::@NamedTuple{max_work_group_size::Int, max_work_group_dims::NTuple{3, Int}, sub_group_size::Int} end KI.max_work_group_size(::POCLBackend)::Int = device_limits().max_work_group_size KI.max_work_group_dims(::POCLBackend)::NTuple{3, Int} = device_limits().max_work_group_dims # the grid is only limited by the size of `size_t` KI.max_num_groups(::POCLBackend)::NTuple{3, Int} = (typemax(Int), typemax(Int), typemax(Int)) -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 +KI.sub_group_size(::POCLBackend)::Int = device_limits().sub_group_size 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 +KI.supports_subgroups(::POCLBackend) = device_limits().sub_group_size > 0 +function KI.supports_shuffle(backend::POCLBackend, ::Type{T}) where {T} + 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 end ## Indexing Functions @@ -300,15 +292,15 @@ end return (; x = get_global_size(1) % T, y = get_global_size(2) % T, z = get_global_size(3) % T) end -@device_override KI.get_sub_group_size() = get_sub_group_size() % UInt32 +@device_override KI.get_sub_group_size(::Type{T}) where {T} = get_sub_group_size() % T -@device_override KI.get_max_sub_group_size() = get_max_sub_group_size() % UInt32 +@device_override KI.get_max_sub_group_size(::Type{T}) where {T} = get_max_sub_group_size() % T -@device_override KI.get_num_sub_groups() = get_num_sub_groups() % UInt32 +@device_override KI.get_num_sub_groups(::Type{T}) where {T} = get_num_sub_groups() % T -@device_override KI.get_sub_group_id() = get_sub_group_id() % UInt32 +@device_override KI.get_sub_group_id(::Type{T}) where {T} = get_sub_group_id() % T -@device_override KI.get_sub_group_local_id() = get_sub_group_local_id() % UInt32 +@device_override KI.get_sub_group_local_id(::Type{T}) where {T} = get_sub_group_local_id() % T @device_override @inline function KA.__validindex(ctx) if KA.__dynamic_checkbounds(ctx) diff --git a/test/hostinterface.jl b/test/hostinterface.jl index 68cce2f12..5693d6e4e 100644 --- a/test/hostinterface.jl +++ b/test/hostinterface.jl @@ -56,8 +56,11 @@ function hostinterface_testsuite(_backend, AT) @testset "KernelInterface host functions" begin @test KI.max_work_group_size(backend) isa Int @test KI.multiprocessor_count(backend) isa Int - @test KI.sub_group_size(backend) isa Int - @test KI.shfl_down_types(backend) isa Vector{DataType} + @test KI.supports_subgroups(backend) isa Bool + if KI.supports_subgroups(backend) + @test KI.sub_group_size(backend) isa Int + end + @test KI.supports_shuffle(backend, Float32) isa Bool function ki_hostinterface_kernel(x) i = KI.get_global_id().x From 8cbd50e013742e4c89b79401a04e65339996dd94 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 09:30:56 +0200 Subject: [PATCH 8/9] KernelInterface: specify execution, device and memory semantics - Execution is task-local: each task has an active device per backend and a queue on it, which host queries, compilation, allocations, copies and launches use. `kernel_function` has to store the backend value it was given (oneAPI dropped its compiler options, and OpenCL its platform, by constructing a new default), and a compiled kernel launched after switching devices either works or throws, but never runs on the wrong device. - New: `device(backend, A)`, the device that owns `A`. `device`, `device!`, `ndevices` and `device(backend, A)` are required for backends with more than one device; the single-device fallbacks now throw when `ndevices` reports more instead of answering for the wrong device. - Capability defaults are conservative: `supports_float64` and `supports_atomics` default to `false`, so that a missing method never claims support. `supports_atomics` means Atomix add and CAS on 32-bit integers and floats in global memory. - `copyto!(backend, dst, src)` copies in queue order and returns `dst`, and throws an `ArgumentError` for arrays of different lengths. CUDA's implementation copied `length(dst)` elements from `pointer(src)` without a check. - `localmemory` and `barrier` say what is shared and what becomes visible, and `unsafe_free!` is an optional hint with a legal no-op fallback. The testsuite covers devices, `copyto!`, local memory (two allocations don't alias; contents are visible after a barrier) and global-memory barriers. --- docs/src/kernelinterface.md | 26 ++++- lib/KernelInterface/src/backend.jl | 8 +- lib/KernelInterface/src/device.jl | 41 +++++--- lib/KernelInterface/src/host.jl | 144 +++++++++++++++----------- lib/KernelInterface/src/launch.jl | 30 ++++-- lib/KernelInterface/test/interface.jl | 127 ++++++++++++++++++++++- lib/KernelInterface/test/runtests.jl | 19 +++- src/pocl/backend.jl | 15 +-- test/runtests.jl | 5 +- 9 files changed, 310 insertions(+), 105 deletions(-) diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md index a1fc1debf..4ecdf144b 100644 --- a/docs/src/kernelinterface.md +++ b/docs/src/kernelinterface.md @@ -44,6 +44,25 @@ unchanged. KernelInterface ``` +## Semantics + +A few rules hold throughout the interface: + +- **Execution is task-local.** A backend value (e.g. `CUDABackend()`) identifies a + backend and its configuration, such as compiler options. Each Julia task has an active + device per backend (selected with [`device!`](@ref)) and a queue on it. Host-side + queries and compilation use the active device, allocations go to it, and copies and + launches go to the calling task's queue. [`synchronize`](@ref) waits for that queue, + and [`record_event`](@ref)/[`wait_event`](@ref) order work across queues. Switching + devices doesn't synchronize. +- **Compiled kernels belong to a device.** Queries on a [`Kernel`](@ref) + ([`max_work_group_size`](@ref), [`launch_configuration`](@ref)) answer for the device it + was compiled for. Launching it after switching to another device either works or + throws, but never runs on the wrong device. +- **Indices are 1-based**, and `x` is the fastest-varying dimension. +- **Capabilities default to "unsupported".** A backend that doesn't implement a + `supports_*` query never claims support. + ## Backend hierarchy Backends subtype [`Backend`](@ref), and everything else in the interface dispatches on @@ -209,10 +228,9 @@ A backend must, at minimum: 1. Define a backend type subtyping [`Backend`](@ref), and implement [`get_backend`](@ref) for its array type. 2. Implement the host-side management functions for that type: - [`allocate`](@ref), [`copyto!`](@ref), [`synchronize`](@ref) and - [`unsafe_free!`](@ref) are required; the remaining functions under - [Host-side API](@ref) have fallbacks that only need overriding when the - defaults don't apply. + [`allocate`](@ref), [`copyto!`](@ref) and [`synchronize`](@ref) are required, and so + are the device functions for backends with more than one device; the remaining + functions under [Host-side API](@ref) have conservative fallbacks. 3. Extend `Adapt.adapt_storage(::NewBackend, x)` so that [`adapt(backend, x)`](@ref Adapt.adapt_storage(::Backend, ::Any)) moves data to the backend, preferably by delegating to its array type: diff --git a/lib/KernelInterface/src/backend.jl b/lib/KernelInterface/src/backend.jl index ad1433b46..0a8539072 100644 --- a/lib/KernelInterface/src/backend.jl +++ b/lib/KernelInterface/src/backend.jl @@ -8,9 +8,11 @@ Abstract supertype for all KernelInterface backends. Backends subtype it directly. -Concrete backends (for example `CUDABackend` from CUDA.jl or `CPU` from KernelAbstractions) -determine where arrays are allocated and where kernels execute. Use [`get_backend`](@ref) to -obtain the backend for an array and [`allocate`](@ref) to create storage on a backend. +A backend value identifies a backend and its configuration (e.g. compiler options). The +device and the queue that operations use are task-local: each task selects its active device +with [`device!`](@ref). Host-side queries answer for the calling task's active device, and +work is queued on the calling task's queue of that device. Use [`get_backend`](@ref) to obtain the +backend of an array and [`allocate`](@ref) to create storage on a backend. # Example diff --git a/lib/KernelInterface/src/device.jl b/lib/KernelInterface/src/device.jl index 3d7cb0b71..40081923f 100644 --- a/lib/KernelInterface/src/device.jl +++ b/lib/KernelInterface/src/device.jl @@ -214,17 +214,26 @@ See [`get_local_id`](@ref) for the supported types `T`. @inline get_sub_group_local_id() = get_sub_group_local_id(Int) +## memory + """ localmemory(::Type{T}, dims) -Declare memory that is local to a workgroup. +Declare an array of element type `T` and size `dims` in memory that is local to a +work-group. `dims` has to be known at compile time. + +Every call site of `localmemory` in a kernel has its own memory, shared by all work-items +of a work-group. It is uninitialized, and lives until the work-group finishes. Executing the +same call site again, e.g. in a loop, returns the same memory. A function containing a call +that is itself called from several places may get the same memory at each of them, or +different memory, depending on whether it is inlined: don't rely on either. Use +[`barrier`](@ref) to make writes visible to the other work-items. !!! note Backend implementations **must** implement: ``` @device_override localmemory(::Type{T}, ::Val{Dims}) where {T, Dims} ``` - As well as the on-device functionality. """ localmemory(::Type{T}, dims) where {T} = localmemory(T, Val(dims)) @@ -234,6 +243,8 @@ localmemory(::Type{T}, ::Val) where {T} = error("Local memory used outside kernel or not captured") +## communication + """ shfl_down(val::T, offset::Integer)::T @@ -256,18 +267,18 @@ branch), with the same `offset`. function shfl_down end +## synchronization + """ barrier() -After a `barrier()` call, all read and writes to global and local memory -from each thread in the workgroup are visible in from all other threads in the -workgroup. +Wait until all work-items of the work-group have reached the barrier. Afterwards, the +writes to global and local memory that each work-item made before the barrier are visible +to all work-items of the work-group. -This does **not** guarantee that a write from a thread in a certain workgroup will -be visible to a thread in a different workgroup. +This does **not** order memory between work-groups. -!!! note - `barrier()` must be encountered by all workitems of a work-group executing the kernel or by none at all. +All work-items of a work-group have to reach the same `barrier()` (not in a divergent branch). !!! note Backend implementations **must** implement: @@ -298,19 +309,21 @@ function sub_group_barrier() error("Sub-group barrier used outside kernel or not captured") end + +## printing + """ _print(args...) - Overloaded by backends to enable `KernelAbstractions.@print` - functionality. +Print `args` from a kernel; the backend hook behind `KernelAbstractions.@print`. !!! note - Backend implementations **must** implement: + Backend implementations **should** implement: ``` @device_override _print(args...) ``` - If the backend does not support printing, - define it to return `nothing`. + A backend that can't print from a kernel defines it to return `nothing`, and + documents that. The generic fallback prints on the host, which keeps CPU backends working. `Val` arguments are unwrapped, since `KernelAbstractions.@print` uses them to diff --git a/lib/KernelInterface/src/host.jl b/lib/KernelInterface/src/host.jl index ee105e3e0..09f55238e 100644 --- a/lib/KernelInterface/src/host.jl +++ b/lib/KernelInterface/src/host.jl @@ -32,8 +32,8 @@ end """ synchronize(::Backend) -Synchronize the current backend: block the calling task until all work it has queued on -`backend` has completed. +Block the calling task until all work it has queued on the active device of `backend` has +completed. !!! note Backend implementations **must** implement this function, and it **must** be @@ -117,11 +117,28 @@ end Return the 1-based index of the currently active device for `backend`. !!! note - The default implementation assumes a single device. Backends supporting multiple devices - **must** implement `device(backend::Backend)::Int`, [`ndevices`](@ref), - and [`device!`](@ref). + Backends supporting multiple devices **must** implement `device(backend::Backend)::Int`, + along with [`ndevices`](@ref), [`device!`](@ref) and `device(backend, A)`. The fallback + only works for a single device, and throws if [`ndevices`](@ref) reports more. """ -function device(::Backend) +function device(backend::Backend) + ndevices(backend) == 1 || throw_multi_device(device, backend) + return 1 +end + +""" + device(backend::Backend, A::AbstractArray)::Int + +Return the 1-based index of the device that owns the memory of `A`, independently of the +currently active device. + +!!! note + Backends supporting multiple devices **must** implement this for their array type. + The fallback only works for a single device, and throws if [`ndevices`](@ref) reports + more. +""" +function device(backend::Backend, ::AbstractArray) + ndevices(backend) == 1 || throw_multi_device(device, backend) return 1 end @@ -131,9 +148,8 @@ end Return the number of devices available to `backend`. !!! note - The default implementation assumes a single device. Backends supporting multiple devices - **must** implement `ndevices(backend::Backend)::Int`, [`device`](@ref), - and [`device!`](@ref). + Backends supporting multiple devices **must** implement `ndevices(backend::Backend)::Int`, + along with [`device`](@ref) and [`device!`](@ref). The fallback returns 1. """ function ndevices(::Backend) return 1 @@ -142,8 +158,8 @@ end """ device!(backend::Backend, id::Int)::Nothing -Select the active device for `backend`. `id` is a 1-based device index and must satisfy -`1 <= id <= ndevices(backend)`. +Select the active device for `backend`. `id` is a 1-based device index; an `id` outside +`1:ndevices(backend)` throws an `ArgumentError`. `device!` is not a synchronization point: work queued before the switch is not ordered with respect to work queued after it. To order across a switch, either [`synchronize`](@ref) @@ -156,24 +172,31 @@ device!(CUDABackend(), 2) # use the second CUDA device ``` !!! note - The default implementation assumes a single device. Backends supporting multiple devices - **must** implement `device!(backend::Backend, id::Int)`, [`ndevices`](@ref), - and [`device`](@ref). + Backends supporting multiple devices **must** implement `device!(backend::Backend, id::Int)`, + along with [`ndevices`](@ref) and [`device`](@ref). The fallback only works for a single + device, and throws if [`ndevices`](@ref) reports more. """ function device!(backend::Backend, id::Int) - if !(0 < id <= ndevices(backend)) + n = ndevices(backend) + if !(0 < id <= n) throw(ArgumentError("Device id $id out of bounds.")) end + n == 1 || throw_multi_device(device!, backend) return nothing end +# a backend with several devices has to implement the device functions itself: the +# single-device fallbacks would silently answer for the wrong device +@noinline throw_multi_device(f, backend) = + error("`$(typeof(backend))` has multiple devices, so it must implement `KernelInterface.$(nameof(f))`") + """ pagelock!(::Backend, dest::AbstractArray)::Union{Nothing, Missing} Pagelock (pin) a host memory buffer for a backend device. This may be necessary for [`copyto!`](@ref) -to perform asynchronously w.r.t to the host/ +to perform asynchronously with respect to the host. -This function should return `nothing`; or `missing` if not implemented. +This function returns `nothing`, or `missing` if not implemented. !!! note @@ -186,17 +209,14 @@ end """ unsafe_free!(x::AbstractArray) -Release the memory of an array for reuse by future allocations -and reduce pressure on the allocator. -After releasing the memory of an array, it should no longer be accessed. +Release the memory of an array for reuse by future allocations, reducing pressure on the +allocator. The array may not be used afterwards. -!!! note - On CPU backend this is always a no-op. +This is a hint: releasing the memory is allowed to do nothing. !!! note - Backend implementations **may** implement this function. - If not implemented for a particular backend, default action is a no-op. - Otherwise, it should be defined for backend's array type. + Backend implementations **may** implement this function for their array type, and + should forward it to their own `unsafe_free!` if they have one. The fallback is a no-op. """ function unsafe_free! end @@ -206,35 +226,37 @@ unsafe_free!(::AbstractArray) = return """ supports_unified(::Backend)::Bool -Returns whether unified memory arrays are supported by the backend. +Whether [`allocate`](@ref) supports `unified=true` on the active device: memory that can be +accessed from both the host and the device without explicit copies. !!! note - Backend implementations **should** implement this function - only if they **do** support unified memory. + Backend implementations **must** implement this function if they support unified + memory. The fallback returns `false`. """ supports_unified(::Backend) = false """ supports_atomics(::Backend)::Bool -Returns whether `@atomic` operations are supported by the backend. +Whether kernels on the active device support Atomix.jl's atomic operations: at least `add` +and compare-and-swap on 32-bit integers and floats in global memory. !!! note - Backend implementations **must** implement this function - only if they **do not** support atomic operations with Atomix. + Backend implementations **must** implement this function if they support atomics. + The fallback returns `false`. """ -supports_atomics(::Backend) = true +supports_atomics(::Backend) = false """ supports_float64(::Backend)::Bool -Returns whether `Float64` values are supported by the backend. +Whether kernels on the active device support `Float64` values. !!! note - Backend implementations **must** implement this function - only if they **do not** support `Float64`. + Backend implementations **must** implement this function if they support `Float64`. + The fallback returns `false`. """ -supports_float64(::Backend) = true +supports_float64(::Backend) = false """ supports_subgroups(::Backend)::Bool @@ -265,9 +287,10 @@ supports_shuffle(::Backend, ::Type) = false """ allocate(::Backend, Type, dims...; unified=false)::AbstractArray -Allocate a storage array appropriate for the computational backend. `unified=true` -allocates an array using unified memory if the backend supports it and throws otherwise. -Use [`supports_unified`](@ref) to determine whether it is supported by a backend. +Allocate an uninitialized array on the active device of the backend. `unified=true` +allocates unified memory, accessible from the host and the device without explicit copies, +if the backend supports it and throws otherwise. Use [`supports_unified`](@ref) to +determine whether it is supported by a backend. !!! note Backend implementations **must** implement `allocate(::NewBackend, T, dims::Tuple)` @@ -288,13 +311,13 @@ end """ zeros(::Backend, Type, dims...; unified=false)::AbstractArray -Allocate a storage array appropriate for the computational backend filled with zeros. -`unified=true` allocates an array using unified memory if the backend supports it and -throws otherwise. +Allocate an array with [`allocate`](@ref) and fill it with zeros. + +This is generic: backends implement `allocate` (and `fill!` for their array type). """ zeros(backend::Backend, T::Type, dims...; kwargs...) = zeros(backend, T, dims; kwargs...) function zeros(backend::Backend, ::Type{T}, dims::Tuple; kwargs...) where {T} - data = allocate(backend, T, dims...; kwargs...) + data = allocate(backend, T, dims; kwargs...) fill!(data, zero(T)) return data end @@ -302,9 +325,9 @@ end """ ones(::Backend, Type, dims...; unified=false)::AbstractArray -Allocate a storage array appropriate for the computational backend filled with ones. -`unified=true` allocates an array using unified memory if the backend supports it and -throws otherwise. +Allocate an array with [`allocate`](@ref) and fill it with ones. + +This is generic: backends implement `allocate` (and `fill!` for their array type). """ ones(backend::Backend, T::Type, dims...; kwargs...) = ones(backend, T, dims; kwargs...) function ones(backend::Backend, ::Type{T}, dims::Tuple; kwargs...) where {T} @@ -315,21 +338,25 @@ end """ - copyto!(::Backend, dest::AbstractArray, src::AbstractArray) + copyto!(::Backend, dest::AbstractArray, src::AbstractArray)::typeof(dest) + +Copy the elements of `src` to `dest`, ordered with respect to the other work on the calling +task's queue: after work queued before the copy, and before work queued after it. Returns +`dest`. -Perform an asynchronous `copyto!` operation that is execution ordered with respect to the back-end. +Either array can be a host array or an array of `backend`. `dest` and `src` must have the +same length, otherwise an `ArgumentError` is thrown. Backends only have to support dense +(contiguous) arrays with the same element type. -For most users, `Base.copyto!` should suffice, performance a simple, synchronous copy. -Only when you know you need asynchronicity w.r.t. the host, you should consider using -this asynchronous version, which requires additional lifetime guarantees as documented below. +The copy may be asynchronous with respect to the host, but doesn't have to be: it can also +block until it has completed. For a simple, synchronous copy, use `Base.copyto!`. !!! warning - Because of the asynchronous nature of this operation, the user is required to guarantee that the lifetime - of the source extends past the *completion* of the copy operation as to avoid a use-after-free. It is not - sufficient to simply use `GC.@preserve` around the call to `copyto!`, because that only extends the - lifetime past the operation getting queued. Instead, it may be required to `synchronize()`, - or otherwise guarantee that the source will still be around when the copy is executed: + Because the copy may be asynchronous, the caller has to keep both arrays alive, and not + access them from the host, until the copy has *completed*, e.g. by calling + [`synchronize`](@ref) before using them. A `GC.@preserve` around `copyto!` only keeps + them alive until the copy is queued: ```julia arr = zeros(64) @@ -342,10 +369,11 @@ this asynchronous version, which requires additional lifetime guarantees as docu !!! note - On some back-ends it may be necessary to first call [`pagelock!`](@ref) on host memory + On some backends it may be necessary to first call [`pagelock!`](@ref) on host memory to enable fully asynchronous behavior w.r.t to the host. !!! note - Backends **must** implement this function. + Backends **must** implement this function, for host-to-device, device-to-host and + device-to-device copies. """ function copyto! end diff --git a/lib/KernelInterface/src/launch.jl b/lib/KernelInterface/src/launch.jl index 3c023d065..97809596f 100644 --- a/lib/KernelInterface/src/launch.jl +++ b/lib/KernelInterface/src/launch.jl @@ -315,25 +315,30 @@ Host code can rely on it, e.g. to pick a `Val(N)` for a warp-level reduction. function sub_group_size end """ - multiprocessor_count(backend::NewBackend)::Int + multiprocessor_count(backend)::Int -The multiprocessor count for the current device used by `backend`. -Used for certain algorithm optimizations. +The number of multiprocessors (CUDA SMs, AMD CUs, Intel Xe cores, ...) of the active device +of `backend`, or 0 if unknown. The unit differs between backends, so this is only useful for +heuristics, e.g. to choose how many work-groups a grid-stride loop launches. !!! note Backend implementations **may** implement: ``` multiprocessor_count(backend::NewBackend)::Int ``` - As well as the on-device functionality. """ multiprocessor_count(::Backend) = 0 + +## compilation + """ - argconvert(::NewBackend, arg) + argconvert(backend, arg) -This function is called for every argument to be passed to a kernel, -converting them to their device side representation. +Convert `arg` to its device-side representation, e.g. a `CuArray` to a `CuDeviceArray`. +Called for every kernel argument, and for the kernel function itself. + +It has to be pure: it may be called more than once for the same launch. !!! note Backend implementations **must** implement: @@ -356,13 +361,20 @@ Keyword arguments: Other keyword arguments are backend-specific compiler options (e.g. `maxthreads` for CUDA.jl); backends throw an error for options they don't support. +The returned kernel doesn't keep any arguments alive: they are passed again at launch. + !!! note Backend implementations **must** implement: ``` kernel_function(backend::NewBackend, f::F, tt::TT=Tuple{}; name=nothing, kwargs...) where {F,TT} ``` - Kernels must execute with sub-group width [`sub_group_size(backend)`](@ref sub_group_size) - if the backend supports sub-groups. + The returned `Kernel` stores `backend` itself (not a new default backend), so that + options it carries apply to the launch. Kernels must execute with sub-group width + [`sub_group_size(backend)`](@ref sub_group_size) if the backend supports sub-groups. + + Launching a kernel after [`device!`](@ref) switched to a device other than the one it + was compiled for must either work, or throw an error: it may never run on the wrong + device. """ function kernel_function end diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl index d659377a7..c575d92c0 100644 --- a/lib/KernelInterface/test/interface.jl +++ b/lib/KernelInterface/test/interface.jl @@ -39,6 +39,42 @@ function test_interface_kernel(results) end return end +# Local arrays, written by one work-item and read back by another one after a barrier. +# `c` has the same type and size as `a`, but is a different call site, so different memory. +function localmem_kernel(out, ::Val{N}) where {N} + i = KI.get_local_id().x + a = KI.localmemory(Int32, N) + b = KI.localmemory(Int32, (2, N)) + c = KI.localmemory(Int32, N) + @inbounds begin + a[i] = i + b[1, i] = -i + b[2, i] = 2i + c[i] = 3i + end + KI.barrier() + j = N - i + 1 + gid = KI.get_global_id().x + @inbounds begin + out[gid, 1] = a[j] + out[gid, 2] = b[1, j] + out[gid, 3] = b[2, j] + out[gid, 4] = c[j] + end + return +end + +# Every work-item writes global memory, and reads another work-item's write after a barrier. +function global_barrier_kernel(scratch, out) + n = KI.get_local_size().x + base = (KI.get_group_id().x - 1) * n + i = KI.get_local_id().x + @inbounds scratch[base + i] = base + i + KI.barrier() + @inbounds out[base + i] = scratch[base + n - i + 1] + return +end + # The interface documents a concrete return type for each device-side function; # these kernels record whether the backend honors them. const WorkItemNT{T} = @NamedTuple{x::T, y::T, z::T} @@ -328,23 +364,83 @@ function interface_testsuite(backend::KI.Backend, AT) @test KI.supports_unified(b) isa Bool @test KI.supports_atomics(b) isa Bool @test KI.supports_float64(b) isa Bool + @test KI.supports_subgroups(b) isa Bool + @test KI.supports_shuffle(b, Float32) isa Bool @test KI.functional(b) isa Union{Missing, Bool} @test KI.multiprocessor_count(b) isa Int @test KI.device(b) isa Int @test KI.ndevices(b) isa Int - # @test KI.device!(b, KI.device(b)) isa Nothing - # @test KI.priority!(b, :normal) isa Nothing - - @test KI.supports_subgroups(b) isa Bool - @test KI.supports_shuffle(b, Float32) isa Bool arr = KI.allocate(b, Float32, 2) @test arr isa AT{Float32, 1} @test KI.zeros(b, Float32, 2) isa AT{Float32, 1} @test KI.ones(b, Float32, 2) isa AT{Float32, 1} @test KI.get_backend(arr) isa KI.Backend + @test KI.device(b, arr) == KI.device(b) + kernel = KI.@launch b launch = false test_interface_kernel(AT(Vector{KernelData}(undef, 1))) + @test kernel isa KI.Kernel + @test kernel.backend === b + end + + @testset "Devices" begin + b = backend + dev = KI.device(b) + @test 1 <= dev <= KI.ndevices(b) + KI.device!(b, dev) + @test KI.device(b) == dev + @test_throws ArgumentError KI.device!(b, 0) + @test_throws ArgumentError KI.device!(b, KI.ndevices(b) + 1) + + if KI.ndevices(b) > 1 + other = mod1(dev + 1, KI.ndevices(b)) + arr_dev = KI.ones(b, Float32, 4) + try + KI.device!(b, other) + # the owner of an array doesn't change with the active device + @test KI.device(b, arr_dev) == dev + @test KI.device(b) == other + arr = KI.ones(b, Float32, 4) + @test KI.device(b, arr) == other + @test Array(arr) == ones(Float32, 4) + finally + KI.device!(b, dev) + end + @test KI.device(b) == dev + end + end + + @testset "copyto!" begin + b = backend + host = rand(Float32, 16) + dev = KI.allocate(b, Float32, 16) + @test KI.copyto!(b, dev, host) === dev + dev2 = KI.allocate(b, Float32, 16) + @test KI.copyto!(b, dev2, dev) === dev2 + back = zeros(Float32, 16) + @test KI.copyto!(b, back, dev2) === back + KI.synchronize(b) + @test back == host + + # arrays of different shapes copy by linear index + dev3 = KI.allocate(b, Float32, (4, 4)) + KI.copyto!(b, dev3, host) + KI.synchronize(b) + @test Array(dev3) == reshape(host, 4, 4) + + # ordered with respect to kernels on the same queue + arr = KI.zeros(b, Int32, (4, 1, 1)) + KI.@launch b numgroups = 1 workgroupsize = 4 launch_kernel(arr) + result = zeros(Int32, 4) + KI.copyto!(b, result, arr) + KI.synchronize(b) + @test result == ones(Int32, 4) + + @test_throws ArgumentError KI.copyto!(b, KI.allocate(b, Float32, 8), host) + @test_throws ArgumentError KI.copyto!(b, zeros(Float32, 8), dev) + @test_throws ArgumentError KI.copyto!(b, KI.allocate(b, Float32, 32), dev) + KI.synchronize(b) end @testset "Device return types" begin @@ -431,6 +527,27 @@ function interface_testsuite(backend::KI.Backend, AT) end end + @testset "Local memory and barriers" begin + N = 32 + groups = 3 + out = KI.zeros(backend, Int32, N * groups, 4) + KI.@launch backend numgroups = groups workgroupsize = N localmem_kernel(out, Val(N)) + KI.synchronize(backend) + out = Array(out) + # each work-item sees what its mirror image wrote before the barrier, in both arrays + expected = repeat(N:-1:1, groups) + @test out[:, 1] == expected + @test out[:, 2] == -expected + @test out[:, 3] == 2 .* expected + @test out[:, 4] == 3 .* expected + + scratch = KI.zeros(backend, Int, N * groups) + out = KI.zeros(backend, Int, N * groups) + KI.@launch backend numgroups = groups workgroupsize = N global_barrier_kernel(scratch, out) + KI.synchronize(backend) + @test Array(out) == vcat([(g * N) .+ (N:-1:1) for g in 0:(groups - 1)]...) + end + if KI.supports_subgroups(backend) sg_size = KI.sub_group_size(backend) max_dims = KI.max_work_group_dims(backend) diff --git a/lib/KernelInterface/test/runtests.jl b/lib/KernelInterface/test/runtests.jl index 804473dfb..956417597 100644 --- a/lib/KernelInterface/test/runtests.jl +++ b/lib/KernelInterface/test/runtests.jl @@ -66,6 +66,10 @@ end struct StubBackend <: KI.Backend end +# A backend with two devices that forgot the other device functions. +struct MultiDeviceBackend <: KI.Backend end +KI.ndevices(::MultiDeviceBackend) = 2 + # An array type with a known backend, for exercising the `get_backend` fallback # that unwraps wrapper arrays. struct BackedArray{T, N} <: AbstractArray{T, N} @@ -125,19 +129,28 @@ end @test KI.device(b) == 1 @test KI.ndevices(b) == 1 @test KI.device!(b, 1) === nothing + @test KI.device(b, zeros(2)) == 1 @test_throws ArgumentError KI.device!(b, 0) @test_throws ArgumentError KI.device!(b, 2) + # A backend with several devices that only implements `ndevices` gets errors from + # the single-device fallbacks, not answers for the wrong device. + mb = MultiDeviceBackend() + @test_throws "must implement `KernelInterface.device`" KI.device(mb) + @test_throws "must implement `KernelInterface.device`" KI.device(mb, zeros(2)) + @test_throws "must implement `KernelInterface.device!`" KI.device!(mb, 2) + @test_throws ArgumentError KI.device!(mb, 3) + # `priority!` validates the symbol even when the backend ignores it. for prio in (:high, :normal, :low) @test KI.priority!(b, prio) === nothing end @test_throws "priority must be one of" KI.priority!(b, :bogus) - # Capability defaults: pessimistic for unified memory, optimistic otherwise. + # Capability defaults are conservative: a missing method never claims support. @test KI.supports_unified(b) === false - @test KI.supports_atomics(b) === true - @test KI.supports_float64(b) === true + @test KI.supports_atomics(b) === false + @test KI.supports_float64(b) === false # Pinning is optional and freeing is a no-op unless a backend does better. @test KI.pagelock!(b, zeros(2)) === missing diff --git a/src/pocl/backend.jl b/src/pocl/backend.jl index acd1a5cdf..88f1d0c64 100644 --- a/src/pocl/backend.jl +++ b/src/pocl/backend.jl @@ -96,10 +96,9 @@ end function KI.copyto!(backend::POCLBackend, A, B) + length(A) == length(B) || + throw(ArgumentError("Arrays must match in length, got $(length(A)) and $(length(B))")) if KI.get_backend(A) == KI.get_backend(B) && KI.get_backend(A) isa POCLBackend - if length(A) != length(B) - error("Arrays must match in length") - end if Base.mightalias(A, B) error("Arrays may not alias") end @@ -107,7 +106,8 @@ function KI.copyto!(backend::POCLBackend, A, B) kernel(A, B, ndrange = length(A)) return A else - return Base.copyto!(A, B) + Base.copyto!(A, B) + return A end end @@ -123,8 +123,9 @@ KI.get_backend(::Array) = POCLBackend() ## must synchronize upon kernel launch and can't rely on synchronization upon ## array access. Therefore, `synchronize` is a no-op. KI.synchronize(::POCLBackend) = nothing -KI.supports_float64(::POCLBackend) = true +KI.supports_float64(::POCLBackend) = "cl_khr_fp64" in device().extensions KI.supports_unified(::POCLBackend) = true +KI.supports_atomics(::POCLBackend) = true ## Kernel Launch @@ -206,7 +207,7 @@ end KI.argconvert(::POCLBackend, arg) = clconvert(arg) -function KI.kernel_function(::POCLBackend, f::F, tt::TT = Tuple{}; name = nothing, kwargs...) where {F, TT} +function KI.kernel_function(backend::POCLBackend, f::F, tt::TT = Tuple{}; name = nothing, kwargs...) where {F, TT} # fix the sub-group width, as `KI.sub_group_size` promises sub_group_size = device_limits().sub_group_size kern = if sub_group_size > 0 @@ -214,7 +215,7 @@ function KI.kernel_function(::POCLBackend, f::F, tt::TT = Tuple{}; name = nothin else clfunction(f, tt; name, kwargs...) end - return KI.Kernel{POCLBackend, typeof(kern)}(POCLBackend(), kern) + return KI.Kernel{POCLBackend, typeof(kern)}(backend, kern) end function KI.launch(obj::KI.Kernel{POCLBackend}, groups::Dims{3}, items::Dims{3}, args::Vararg{Any, N}) where {N} diff --git a/test/runtests.jl b/test/runtests.jl index b0d499553..9638ae7f1 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -162,8 +162,9 @@ struct NewBackend <: KernelAbstractions.Backend end @test_throws MethodError KernelAbstractions.zeros(backend, Float32, 1) @test_throws MethodError KernelAbstractions.ones(backend, Float32, 1) - @test KernelAbstractions.supports_atomics(backend) == true - @test KernelAbstractions.supports_float64(backend) == true + # conservative capability defaults + @test KernelAbstractions.supports_atomics(backend) == false + @test KernelAbstractions.supports_float64(backend) == false @test KernelAbstractions.priority!(backend, :high) === nothing @test KernelAbstractions.priority!(backend, :normal) === nothing From 2e505eaceb81e8a7b3acb2c575440fd86489ef4a Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 09:31:16 +0200 Subject: [PATCH 9/9] KernelInterface: declare the public API, document the contract - The public API is declared with `public` (Julia 1.11+); nothing is exported, so `KI.` prefixes the interface everywhere. - The docs get a contract table listing what a backend implements, which methods are optional and what their fallbacks do, and a versioning rule: required methods only change in breaking releases, optional methods can be added in any release if their fallback is conservative, and patch releases only add tests for behavior that was already specified. (0.2.3 broke backend CI in a patch release by adding tests for new obligations.) - The testsuite checks that a backend implements the methods without a fallback. --- docs/src/kernelinterface.md | 67 +++++++++++++++------- lib/KernelInterface/README.md | 24 ++++++++ lib/KernelInterface/src/KernelInterface.jl | 28 +++++++++ lib/KernelInterface/test/interface.jl | 18 ++++++ lib/KernelInterface/test/testsuite.jl | 4 ++ 5 files changed, 120 insertions(+), 21 deletions(-) diff --git a/docs/src/kernelinterface.md b/docs/src/kernelinterface.md index 4ecdf144b..1f6fa0f5e 100644 --- a/docs/src/kernelinterface.md +++ b/docs/src/kernelinterface.md @@ -63,6 +63,45 @@ A few rules hold throughout the interface: - **Capabilities default to "unsupported".** A backend that doesn't implement a `supports_*` query never claims support. +## Contract + +What a backend implements, at a glance. The docstrings below have the details. + +| | Required | Optional (fallback) | +|---|---|---| +| **Backend** | subtype [`Backend`](@ref); [`get_backend`](@ref) for its array type | | +| **Memory** | [`allocate`](@ref), [`copyto!`](@ref) | `allocate(...; unified=true)` (throws), [`pagelock!`](@ref) (`missing`), [`unsafe_free!`](@ref) (no-op) | +| **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`) | +| **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` | | + +Everything else, such as [`zeros`](@ref KernelInterface.zeros), [`ones`](@ref KernelInterface.ones), +the launch-keyword handling of [`Kernel`](@ref) and [`@launch`](@ref KernelInterface.@launch), +is generic and not meant to be overridden. + +### Versioning + +- Required methods only change in breaking releases (0.x → 0.x+1). +- Optional methods can be added in any release, with a fallback that is conservative: + never claiming support, never wrong. Tests for them pass on the fallback, or are gated + on a capability query. +- A patch release may add tests of behavior that was already specified; tests for newly + specified behavior are new obligations and wait for a breaking release. + +Backends test themselves against the contract with the testsuite in +`lib/KernelInterface/test`: + +```julia +import KernelInterface +using Test +include(joinpath(pkgdir(KernelInterface), "test", "testsuite.jl")) +Testsuite.testsuite(MyBackend(), MyArray) +``` + ## Backend hierarchy Backends subtype [`Backend`](@ref), and everything else in the interface dispatches on @@ -154,9 +193,6 @@ what makes [`KernelAbstractions.@print`](@ref) usable outside of a kernel. ## Host-side API -Several of these have generic fallbacks. Each docstring notes which methods -a backend **must** implement and which ones are optional. - ### Memory ```@docs; canonical=false @@ -223,35 +259,24 @@ KernelInterface.@launch ## Implementing a backend -A backend must, at minimum: +A backend implements the required methods from the [contract](@ref Contract), and those +optional methods where it can do better than the fallback. In particular: 1. Define a backend type subtyping [`Backend`](@ref), and implement [`get_backend`](@ref) for its array type. -2. Implement the host-side management functions for that type: - [`allocate`](@ref), [`copyto!`](@ref) and [`synchronize`](@ref) are required, and so - are the device functions for backends with more than one device; the remaining - functions under [Host-side API](@ref) have conservative fallbacks. -3. Extend `Adapt.adapt_storage(::NewBackend, x)` so that +2. Extend `Adapt.adapt_storage(::NewBackend, x)` so that [`adapt(backend, x)`](@ref Adapt.adapt_storage(::Backend, ::Any)) moves data to the backend, preferably by delegating to its array type: `Adapt.adapt_storage(::NewBackend, x) = adapt(NewArray, x)`. -4. `@device_override` the device-side functions it supports. The four primitive index - queries and [`barrier`](@ref) are required; sub-group and [`shfl_down`](@ref) support - is optional. -5. Implement [`argconvert`](@ref) and [`kernel_function`](@ref) for its backend - type, returning a [`Kernel`](@ref). -6. Implement [`launch`](@ref), which receives an already validated `NTuple{3, Int}` of - work-groups and of work-items. For CUDA.jl, that is +3. Implement [`kernel_function`](@ref), returning a [`Kernel`](@ref) that holds the + backend value it was given, and [`launch`](@ref), which receives an already validated + `NTuple{3, Int}` of work-groups and of work-items. For CUDA.jl, the latter is ```julia KI.launch(k::KI.Kernel{CUDABackend}, groups::Dims{3}, items::Dims{3}, args::Vararg{Any, N}; kwargs...) where {N} = k.kern(args...; threads = items, blocks = groups, kwargs...) ``` -7. Compute the typed index queries with `% T`, not `T(x)`: a checked conversion leaves +4. Compute the typed index queries with `% T`, not `T(x)`: a checked conversion leaves an error branch in every kernel. -8. Report its limits through [`max_work_group_size`](@ref) (for the backend and for a - kernel), [`max_work_group_dims`](@ref) and [`max_num_groups`](@ref), and where - applicable [`sub_group_size`](@ref) and [`multiprocessor_count`](@ref). It may - recommend work-group sizes with [`launch_configuration`](@ref). The PoCL backend in `src/pocl/backend.jl` is a complete worked example. diff --git a/lib/KernelInterface/README.md b/lib/KernelInterface/README.md index 6994f2321..fd55a7052 100644 --- a/lib/KernelInterface/README.md +++ b/lib/KernelInterface/README.md @@ -10,6 +10,30 @@ KernelInterface focuses on the lower-level functionality shared amongst backends such as kernel launching, device intrinsics, and host-side operations such as allocation and synchronization. +Backends implement a small set of required methods (allocation, copies, +synchronization, compilation, a `launch` method and the primitive device +queries) and may override optional ones whose fallbacks are conservative; see +the [contract table](https://juliagpu.github.io/KernelAbstractions.jl/dev/kernelinterface/#Contract) +in the documentation. The testsuite in `test/testsuite.jl` checks it: + +```julia +import KernelInterface +using Test +include(joinpath(pkgdir(KernelInterface), "test", "testsuite.jl")) +Testsuite.testsuite(MyBackend(), MyArray) +``` + + +## Versioning + +- Required methods only change in breaking releases (0.x → 0.x+1). +- Optional methods can be added in any release, with a conservative fallback + (never claiming support, never wrong); tests for them pass on the fallback or + are gated on a capability query. +- A patch release may add tests of behavior that was already specified; tests + for newly specified behavior are new obligations and wait for a breaking + release. + ## License diff --git a/lib/KernelInterface/src/KernelInterface.jl b/lib/KernelInterface/src/KernelInterface.jl index 083f024ba..0f0fade1a 100644 --- a/lib/KernelInterface/src/KernelInterface.jl +++ b/lib/KernelInterface/src/KernelInterface.jl @@ -19,4 +19,32 @@ include("device.jl") include("launch.jl") include("host.jl") +# the public API; nothing is exported, so that `KI.` prefixes the interface everywhere +@static if VERSION >= v"1.11" + eval( + Expr( + :public, + # backends + :Backend, :get_backend, + # device side + :get_global_id, :get_global_size, :get_local_id, :get_local_size, + :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, + # compilation and launch + :Kernel, :kernel_function, :argconvert, :launch, Symbol("@launch"), + :launch_configuration, :max_work_group_size, :max_work_group_dims, + :max_num_groups, :sub_group_size, :multiprocessor_count, + # host side + :allocate, :zeros, :ones, :copyto!, :pagelock!, :unsafe_free!, + :synchronize, :record_event, :wait_event, :priority!, + :device, :ndevices, :device!, + :functional, :versioninfo, + :supports_unified, :supports_atomics, :supports_float64, + :supports_subgroups, :supports_shuffle, + ) + ) +end + end diff --git a/lib/KernelInterface/test/interface.jl b/lib/KernelInterface/test/interface.jl index c575d92c0..3cd9a11e6 100644 --- a/lib/KernelInterface/test/interface.jl +++ b/lib/KernelInterface/test/interface.jl @@ -685,3 +685,21 @@ function interface_testsuite(backend::KI.Backend, AT) end return nothing end + +# Checks that a backend implements the methods that have no fallback. +function contract_testsuite(backend::KI.Backend, AT) + B = typeof(backend) + @test hasmethod(KI.synchronize, Tuple{B}) + @test hasmethod(KI.copyto!, Tuple{B, AT, Array}) + @test hasmethod(KI.argconvert, Tuple{B, Any}) + @test hasmethod(KI.kernel_function, Tuple{B, Any, Type}) + @test hasmethod(KI.launch, Tuple{KI.Kernel{B}, Dims{3}, Dims{3}}) + @test hasmethod(KI.max_work_group_size, Tuple{B}) + @test hasmethod(KI.max_work_group_size, Tuple{KI.Kernel{B}}) + @test hasmethod(KI.max_work_group_dims, Tuple{B}) + @test hasmethod(KI.max_num_groups, Tuple{B}) + if KI.supports_subgroups(backend) + @test hasmethod(KI.sub_group_size, Tuple{B}) + end + return +end diff --git a/lib/KernelInterface/test/testsuite.jl b/lib/KernelInterface/test/testsuite.jl index fa0b598af..d175bc2ad 100644 --- a/lib/KernelInterface/test/testsuite.jl +++ b/lib/KernelInterface/test/testsuite.jl @@ -36,6 +36,10 @@ Run the KernelInterface tests for `backend`, whose array type is `AT`. Test sets `skip_tests` are skipped. """ function testsuite(backend::KI.Backend, AT; skip_tests = Set{String}()) + @conditional_testset "Contract" skip_tests begin + contract_testsuite(backend, AT) + end + @conditional_testset "Interface" skip_tests begin interface_testsuite(backend, AT) end