KernelInterface 0.3: a tighter back-end contract

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"

docs/src/api.md

@@ -33,7 +33,6 @@
 
 ```@docs
 Backend
-GPU
 CPU
 POCLBackend
 get_backend

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`
 

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
 

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,29 +119,41 @@ 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.
 
 ### 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.
 

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

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
 

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)

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))