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..1f6fa0f5e 100644 --- a/docs/src/kernelinterface.md +++ b/docs/src/kernelinterface.md @@ -44,17 +44,73 @@ 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. + +## 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 -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 ``` @@ -63,7 +119,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. @@ -71,21 +127,33 @@ 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 +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 @@ -111,7 +179,6 @@ localmemory ```@docs shfl_down -shfl_down_types ``` ### Printing @@ -126,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 @@ -167,10 +231,16 @@ supports_atomics supports_float64 ``` -### Backend queries +```@docs +supports_subgroups +supports_shuffle +``` + +### Limits ```@docs max_work_group_size +launch_configuration max_work_group_dims max_num_groups sub_group_size @@ -182,46 +252,31 @@ multiprocessor_count ```@docs Kernel kernel_function -kernel_max_work_group_size argconvert -KernelInterface.@kernel +launch +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: +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 [`GPU`](@ref) (or [`Backend`](@ref) for - non-GPU backends), 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. -3. Extend `Adapt.adapt_storage(::NewBackend, x)` so that +1. Define a backend type subtyping [`Backend`](@ref), and implement [`get_backend`](@ref) + for its array type. +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 indexing - 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. Make that `Kernel` callable, accepting `numworkgroups`, `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. -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). +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...) + ``` +4. Compute the typed index queries with `% T`, not `T(x)`: a checked conversion leaves + an error branch in every kernel. The PoCL backend in `src/pocl/backend.jl` is a complete worked example. 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/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/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/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/src/backend.jl b/lib/KernelInterface/src/backend.jl index 409fd0522..0a8539072 100644 --- a/lib/KernelInterface/src/backend.jl +++ b/lib/KernelInterface/src/backend.jl @@ -6,11 +6,13 @@ """ 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 -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 @@ -23,17 +25,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/src/device.jl b/lib/KernelInterface/src/device.jl index 8814544f1..40081923f 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,142 +57,183 @@ 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 + +# 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) + +## 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)) @@ -197,52 +243,42 @@ localmemory(::Type{T}, ::Val) where {T} = error("Local memory used outside kernel or not captured") +## communication + """ - 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[] +## 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: @@ -257,18 +293,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. - -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. +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. -!!! note - `sub_group_barrier()` must be encountered by all workitems of a sub-group executing the kernel or by none at all. +All work-items of a sub-group have to reach the same `sub_group_barrier()`. !!! note - Backend implementations **must** implement: + Backend implementations that support sub-groups **must** implement: ``` @device_override sub_group_barrier() ``` @@ -277,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 419698eb2..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,42 +226,71 @@ 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 if they support `Float64`. + The fallback returns `false`. +""" +supports_float64(::Backend) = false + +""" + supports_subgroups(::Backend)::Bool + +Whether kernels on the active device support sub-groups: the sub-group queries +([`get_sub_group_size`](@ref) etc.), [`sub_group_barrier`](@ref), and a fixed sub-group +width [`sub_group_size`](@ref). + +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 - only if they **do not** support `Float64`. + Backend implementations **must** implement this function for the types they support. + The fallback returns `false`. """ -supports_float64(::Backend) = true +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)` @@ -262,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 @@ -276,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} @@ -289,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) @@ -316,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 a5126af13..97809596f 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...; numworkgroups=(), workgroupsize=(), ndrange=(), max_work_group_size=typemax(Int)) - ``` - `numworkgroups`, `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 - 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. - `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. +The launch geometry is given in one of three ways: - 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. +- `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. - By default, kernels must launch with 1 workgroup containing 1 workitem. +Each is an `Integer` or a tuple of up to 3 `Integer`s; missing dimensions are 1. `ndrange` +and `numgroups` are mutually exclusive. - Backends must also implement the on-device kernel launch functionality. +`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. + +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. + +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(numworkgroups, 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 `numworkgroups` 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(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")) - length(workgroupsize) <= 3 || - throw(ArgumentError("`workgroupsize` only accepts up to 3 dimensions")) - length(ndrange) <= 3 || - throw(ArgumentError("`ndrange` only accepts up to 3 dimensions")) +@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")) + 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,137 +210,135 @@ 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, numworkgroups, workgroupsize, ndrange, [max_work_items]) - -Returns a suggested `numworkgroups` 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 -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. +## limits and advice -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 - 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)) - else - workgroupsize - end - numworkgroups = cld.(ndrange, workgroupsize) - Int.(numworkgroups), Int.(workgroupsize) - end + max_work_group_size(backend)::Int + max_work_group_size(kernel::Kernel)::Int - return numworkgroups, workgroupsize -end - -""" - kernel_max_work_group_size(kern; [max_work_items::Int])::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` (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. !!! 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 (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 **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 -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 """ - 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) + +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. -This function is called for every argument to be passed to a kernel, -converting them to their device side representation. +It has to be pure: it may be called more than once for the same launch. !!! note Backend implementations **must** implement: @@ -218,55 +349,73 @@ 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 -[`KernelInterface.@kernel`](@ref). - -Currently, `kernel_function` only supports the `name` keyword argument as it is the only one -by all backends. +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: -- `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. + +The returned kernel doesn't keep any arguments alive: they are passed again at launch. !!! 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} ``` + 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 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 +428,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 6ff22e089..f78179086 100644 --- a/lib/KernelInterface/test/events.jl +++ b/lib/KernelInterface/test/events.jl @@ -19,13 +19,18 @@ 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 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 30172bf72..3cd9a11e6 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 @@ -24,25 +39,39 @@ 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 +# 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 -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 +# 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 @@ -96,17 +125,63 @@ 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 + +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 @@ -117,9 +192,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 @@ -131,66 +210,106 @@ function shfl_down_test_kernel(a, b, ::Val{N}) where {N} return end -function interface_testsuite(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() +# 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 - arr[(gi - 1) * ngi + i] = 1.0f0 - return - end - arr1d = AT(zeros(Float32, 4)) - 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()) - @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 +# 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 + 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.@kernel backend() numworkgroups = (2, 2) workgroupsize = (2, 2) launch_kernel2d(arr2d) - KI.synchronize(backend()) - @test all(Array(arr2d) .== 1) - - # 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() - - arr[(gi - 1) * ngi + i, (gj - 1) * ngj + j, (gk - 1) * ngk + k] = 1.0f0 - return + @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(d -> d.global_size == 12 && d.num_groups == 3, Array(results)) + end + + @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) + + @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.@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)) + KI.synchronize(backend) + @test Array(arr) == zeros(Int32, 1, 1, 1) + + # 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 - 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) + @test KI.max_work_group_size(backend) isa Int function fill_kernel(arr) i, j, k = KI.get_global_id() @@ -199,203 +318,388 @@ function interface_testsuite(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...) - KI.synchronize(backend()) + KI.synchronize(backend) 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) - @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 @testset "Host return types" begin - b = backend() + b = backend @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.shfl_down_types(b) isa Vector{DataType} 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 - results = KI.zeros(backend(), Bool, 6) - KI.@kernel backend() typecheck_kernel(results) - KI.synchronize(backend()) + results = KI.zeros(backend, Bool, 6) + 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.synchronize(backend()) + typed_results = KI.zeros(backend, Bool, 6) + KI.@launch backend typed_typecheck_kernel(typed_results, T) + KI.synchronize(backend) @test all(Array(typed_results)) end end @testset "Typed indexing" begin - workgroupsize = (2, 2, 2) - numworkgroups = (3, 2, 1) - N = prod(workgroupsize) * prod(numworkgroups) + 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. 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.@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))) - # every global id is seen exactly once + @test all(eachrow(reference[:, 13:15]) .== Ref(collect(numgroups))) + # 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 "Basic interface functionality" begin - - @test KI.max_work_group_size(backend()) isa Int - @test KI.multiprocessor_count(backend()) isa Int + @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 - # Test with small kernel + @testset "Basic interface functionality" begin workgroupsize = 4 - numworkgroups = 4 - N = workgroupsize * numworkgroups + numgroups = 3 + N = workgroupsize * numgroups results = AT(Vector{KernelData}(undef, N)) - 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 = KI.@launch backend launch = false test_interface_kernel(results) - kernel(results; workgroupsize, numworkgroups) - KI.synchronize(backend()) + kernel(results; workgroupsize, numgroups) + 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 == numworkgroups - - # Group ID should be 1-based - expected_group = div(i - 1, numworkgroups) + 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.num_groups == numgroups + @test k_data.group_id == div(i - 1, workgroupsize) + 1 + @test k_data.local_id == ((i - 1) % workgroupsize) + 1 end end - # Used as a proxy for sub-group support - if !isempty(KI.shfl_down_types(backend())) + @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) + # 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.@kernel 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 - 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(results; workgroupsize, numworkgroups) - 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 + 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 - # 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.@kernel 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 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/runtests.jl b/lib/KernelInterface/test/runtests.jl index 944e5648c..956417597 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.kernel_max_work_group_size, KI.max_work_group_size, KI.sub_group_size, - KI.argconvert, KI.kernel_function, + 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!, ] @@ -42,24 +39,37 @@ 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, + 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 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 +# 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} @@ -80,9 +90,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 @@ -92,8 +103,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() @@ -120,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 @@ -192,20 +210,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 numworkgroupsize and ndrange defined - @test_throws ArgumentError KI.check_launch_args(2, (), 2) # both numworkgroupsize 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,) @@ -227,71 +231,155 @@ 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 @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 -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) +# A minimal backend, recording the compilations and launches that KernelInterface asks for. +struct MockBackend <: KI.Backend + max_items::Int end +MockBackend() = MockBackend(256) -# ... 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) +struct MockKernel + f::Any + tt::Any + name::Any + options::Any + launches::Vector{Any} +end -@testset "auto_launch_sizes" begin - # the per-dimension limit is respected - @test KI.auto_launch_sizes(KI.Kernel(DimsBackend(), nothing), (), (), (1, 1, 5000)) === - ((1, 1, 79), (1, 1, 64)) +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(kernel::KI.Kernel{MockBackend}) = kernel.backend.max_items +KI.max_work_group_dims(::MockBackend) = (1024, 1024, 64) - kernel = KI.Kernel(SizedBackend(256), nothing) +# ... 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 +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, 4), (1, 1)) - @test KI.auto_launch_sizes(kernel, (), (), 0) === (0, 1) + @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)) + 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_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(), []) + @test KI.launch_configuration(occupancy) === (; workgroupsize = 96) + @test KI.max_work_group_size(occupancy) == 1024 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]) @@ -315,49 +403,45 @@ end @test var_exprs[2] == Expr(:..., vars[2]) end -# A minimal backend, exercising the contract `KI.@kernel` expects of one. -struct MockBackend end - -struct MockKernel - f::Any - tt::Any - name::Any - launches::Vector{Any} -end +dummy(a, b) = nothing -KI.argconvert(::MockBackend, arg) = arg -function KI.kernel_function(::MockBackend, f, tt = Tuple{}; name = nothing, kwargs...) - return MockKernel(f, tt, name, []) -end -function (kernel::MockKernel)(args...; kwargs...) - push!(kernel.launches, (args, Dict(kwargs))) - return nothing +const backend_evaluations = Ref(0) +function counted_backend() + backend_evaluations[] += 1 + return MockBackend() end -dummy(a, b) = nothing - -@testset "@kernel" begin +@testset "@launch" begin backend = MockBackend() - kernel = KI.@kernel backend numworkgroups = 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) + kernel = KI.@launch backend numgroups = 2 workgroupsize = 4 dummy(1, 2.0) + @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 + 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) - @test isempty(deferred.launches) + deferred = KI.@launch backend launch = false dummy(1, 2.0) + @test isempty(deferred.kern.launches) - # Compiler kwargs reach `kernel_function` instead of the launch. - named = KI.@kernel backend launch = false name = "mykernel" dummy(1, 2.0) - @test named.name == "mykernel" + # Other keywords are compiler options for `kernel_function`. + named = KI.@launch backend launch = false name = "mykernel" maxthreads = 32 dummy(1, 2.0) + @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 isempty(only(optioned.kern.launches).kwargs) # Splatted arguments are supported. - splatted = KI.@kernel backend launch = false dummy((1, 2.0)...) - @test splatted.tt == Tuple{Int, Float64} + splatted = KI.@launch backend launch = false dummy((1, 2.0)...) + @test splatted.kern.tt == Tuple{Int, Float64} @testset "errors" begin # These throw during macro expansion, so they cannot be written as a plain @@ -371,14 +455,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/lib/KernelInterface/test/testsuite.jl b/lib/KernelInterface/test/testsuite.jl index 4457ef3df..d175bc2ad 100644 --- a/lib/KernelInterface/test/testsuite.jl +++ b/lib/KernelInterface/test/testsuite.jl @@ -29,7 +29,17 @@ 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 "Contract" skip_tests begin + contract_testsuite(backend, AT) + end + @conditional_testset "Interface" skip_tests begin interface_testsuite(backend, AT) end 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..88f1d0c64 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) @@ -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,78 +207,61 @@ 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...) - return KI.Kernel{POCLBackend, typeof(kern)}(POCLBackend(), kern) +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 + clfunction(f, tt; name, sub_group_size, kwargs...) + else + clfunction(f, tt; name, kwargs...) + end + return KI.Kernel{POCLBackend, typeof(kern)}(backend, kern) end -function (obj::KI.Kernel{POCLBackend})(args...; numworkgroups = (), workgroupsize = (), ndrange = (), max_work_group_size = typemax(Int)) - KI.check_launch_args(numworkgroups, 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) - - local_size = (workgroupsize..., ntuple(_ -> 1, 3 - length(workgroupsize))...) - - numworkgroups = (numworkgroups..., ntuple(_ -> 1, 3 - length(numworkgroups))...) - global_size = local_size .* numworkgroups - - 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 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() 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 -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 +# 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)) +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 @@ -293,10 +277,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 @@ -305,19 +285,23 @@ 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 -@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 17728150e..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 @@ -68,10 +71,11 @@ 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 + @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 diff --git a/test/runtests.jl b/test/runtests.jl index 728feff1c..9638ae7f1 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 @@ -152,7 +149,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() @@ -165,8 +162,9 @@ struct NewBackend <: KernelAbstractions.GPU 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