Rework KernelInterface into a sibling package

.github/workflows/ci.yml

@@ -164,6 +164,30 @@ jobs:
         with:
           files: lcov.info
 
+  KernelInterface:
+    name: KernelInterface
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        version: ['1.10', '1.11', '1.12', '1.13-nightly']
+        os: [ubuntu-24.04, macOS-15, windows-2022]
+    steps:
+      - uses: actions/checkout@v7
+      - uses: julia-actions/install-juliaup@v3
+        with:
+          channel: ${{ matrix.version }}
+      - uses: julia-actions/cache@v3
+      # Tested on its own, without KernelAbstractions, to keep the sibling
+      # package standalone and dependency-free.
+      - uses: julia-actions/julia-buildpkg@v1
+        with:
+          project: lib/KernelInterface
+      - uses: julia-actions/julia-runtest@v1
+        with:
+          project: lib/KernelInterface
+          annotate: true
+
   OpenCL:
     name: OpenCL (POCL)
     runs-on: ubuntu-latest

Project.toml

@@ -8,6 +8,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Atomix = "a9b6321e-bd34-4604-b9c9-b65b8de01458"
 GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
@@ -21,6 +22,9 @@ SPIRV_Tools_jll = "6ac6d60f-d740-5983-97d7-a4482c0689f4"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 pocl_standalone_jll = "54f56a70-6062-5590-a942-1226658f6c83"
 
+[sources]
+KernelInterface = {path = "lib/KernelInterface"}
+
 [weakdeps]
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -38,6 +42,7 @@ Atomix = "0.1, 1"
 EnzymeCore = "0.7, 0.8.1"
 GPUCompiler = "2"
 InteractiveUtils = "1.6"
+KernelInterface = "0.1"
 LLVM = "9.9"
 LinearAlgebra = "1.6"
 MacroTools = "0.5"

docs/Project.toml

@@ -1,9 +1,11 @@
 [deps]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf"
 
 [compat]
 Documenter = "1"
 
 [sources]
 KernelAbstractions = {path = ".."}
+KernelInterface = {path = "../lib/KernelInterface"}

docs/make.jl

@@ -1,11 +1,12 @@
 using KernelAbstractions
+using KernelInterface
 using Documenter
 
 function main()
     ci = get(ENV, "CI", "") == "true"
 
     makedocs(;
-        modules = [KernelAbstractions],
+        modules = [KernelAbstractions, KernelInterface],
         authors = "JuliaGPU and contributors",
         repo = "https://github.com/JuliaGPU/KernelAbstractions.jl/blob/{commit}{path}#L{line}",
         sitename = "KernelAbstractions.jl",
@@ -35,6 +36,7 @@ function main()
                 "examples/atomix.md",
             ], # Examples
             "API" => "api.md",
+            "KernelInterface" => "kernelinterface.md",
             "Extras" => [
                 "extras/unrolling.md",
                 "extras/pocl_debugging.md",

docs/src/kernelinterface.md

@@ -0,0 +1,143 @@
+# [KernelInterface](@id kernelinterface)
+
+```@meta
+CurrentModule = KernelInterface
+```
+
+`KernelInterface` (conventionally imported as `KI`) is the low-level API that
+backends implement, and that `KernelAbstractions` builds its higher-level kernel
+language on top of.
+
+It ships as a standalone package under `lib/KernelInterface` with **no
+dependencies outside the standard library**, so a backend can implement the
+interface without taking on `KernelAbstractions` or its compiler stack:
+
+```julia
+using KernelInterface
+const KI = KernelInterface
+```
+
+`KernelAbstractions` re-exports it, so `KernelAbstractions.KernelInterface` and
+`KernelAbstractions.KI` refer to the same module.
+
+!!! note
+    Most of the functions below are stubs with no methods. They exist so that
+    backends can add device-side implementations with
+    `GPUCompiler.@device_override`, and so kernels can call them generically.
+    Calling one without a backend that implements it is a `MethodError`.
+
+```@docs
+KernelInterface
+```
+
+## Device-side API
+
+These are called from inside a kernel. A backend provides each one with
+
+```julia
+@device_override KI.get_global_id() = ...
+```
+
+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.
+
+```@docs
+get_global_size
+get_global_id
+get_local_size
+get_local_id
+get_num_groups
+get_group_id
+```
+
+### Sub-groups
+
+```@docs
+get_sub_group_size
+get_max_sub_group_size
+get_num_sub_groups
+get_sub_group_id
+get_sub_group_local_id
+```
+
+### Barriers
+
+```@docs
+barrier
+sub_group_barrier
+```
+
+### Memory
+
+```@docs
+localmemory
+```
+
+### Communication
+
+```@docs
+shfl_down
+shfl_down_types
+```
+
+### Printing
+
+```@docs
+KernelInterface._print
+```
+
+`_print` is the one device-side function with a working host fallback: it prints
+its arguments with `Base.print`, unwrapping any `Val`-wrapped literals. That is
+what makes [`KernelAbstractions.@print`](@ref) usable outside of a kernel.
+
+## Host-side API
+
+### Backend queries
+
+```@docs
+max_work_group_size
+sub_group_size
+multiprocessor_count
+```
+
+### Compilation and launching
+
+```@docs
+Kernel
+kernel_function
+kernel_max_work_group_size
+check_launch_args
+argconvert
+KernelInterface.@kernel
+```
+
+!!! 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:
+
+1. `@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.
+2. Implement [`argconvert`](@ref) and [`kernel_function`](@ref) for its backend
+   type, returning a [`Kernel`](@ref).
+3. Make that `Kernel` callable, accepting `numworkgroups` and `workgroupsize` as
+   a scalar `Integer` or a 1-, 2- or 3-element tuple. Use
+   [`check_launch_args`](@ref) to validate them, or check them directly.
+4. 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).
+
+The PoCL backend in `src/pocl/backend.jl` is a complete worked example.
+
+See also the [notes for backend implementations](@ref implementations_notes).

lib/KernelInterface/Project.toml

@@ -0,0 +1,7 @@
+name = "KernelInterface"
+uuid = "4ee993da-d684-4d17-a7dd-4e58e78d92bf"
+authors = ["Valentin Churavy <v.churavy@gmail.com> and contributors"]
+version = "0.1.0"
+
+[compat]
+julia = "1.10"

src/interface.jl → lib/KernelInterface/src/KernelInterface.jl

@@ -12,8 +12,70 @@ like allocating arrays on a backend.
 """
 module KernelInterface
 
-import ..KernelAbstractions: Backend
-import GPUCompiler: split_kwargs, assign_args!
+## macro tools
+# Vendored from GPUCompiler to keep KernelInterface free of heavy dependencies.
+
+# split keyword arguments expressions into groups. returns vectors of keyword argument
+# values, one more than the number of groups (unmatched keywords in the last vector).
+# intended for use in macros; the resulting groups can be used in expressions.
+# can be used at run time, but not in performance critical code.
+function split_kwargs(kwargs, kw_groups...)
+    kwarg_groups = ntuple(_ -> [], length(kw_groups) + 1)
+    for kwarg in kwargs
+        # decode
+        if Meta.isexpr(kwarg, :(=))
+            # use in macros
+            key, val = kwarg.args
+        elseif kwarg isa Pair{Symbol, <:Any}
+            # use in functions
+            key, val = kwarg
+        else
+            throw(ArgumentError("non-keyword argument like option '$kwarg'"))
+        end
+        isa(key, Symbol) || throw(ArgumentError("non-symbolic keyword '$key'"))
+
+        # find a matching group
+        group = length(kwarg_groups)
+        for (i, kws) in enumerate(kw_groups)
+            if key in kws
+                group = i
+                break
+            end
+        end
+        push!(kwarg_groups[group], kwarg)
+    end
+
+    return kwarg_groups
+end
+
+# assign arguments to variables, handle splatting
+function assign_args!(code, _args)
+    nargs = length(_args)
+
+    # handle splatting
+    splats = Vector{Bool}(undef, nargs)
+    args = Vector{Any}(undef, nargs)
+    for i in 1:nargs
+        splats[i] = Meta.isexpr(_args[i], :(...))
+        args[i] = splats[i] ? _args[i].args[1] : _args[i]
+    end
+
+    # assign arguments to variables
+    vars = Vector{Symbol}(undef, nargs)
+    for i in 1:nargs
+        vars[i] = gensym()
+        push!(code.args, :($(vars[i]) = $(args[i])))
+    end
+
+    # convert the arguments, compile the function and call the kernel
+    # while keeping the original arguments alive
+    var_exprs = Vector{Any}(undef, nargs)
+    for i in 1:nargs
+        var_exprs[i] = splats[i] ? Expr(:(...), vars[i]) : vars[i]
+    end
+
+    return vars, var_exprs
+end
 
 """
     get_global_size()::@NamedTuple{x::Int, y::Int, z::Int}
@@ -188,6 +250,11 @@ Declare memory that is local to a workgroup.
 """
 localmemory(::Type{T}, dims) where {T} = localmemory(T, Val(dims))
 
+# The `Val` form only exists in a backend's overlay method table, so off-device it
+# would otherwise fall back to the forwarding method above and recurse forever.
+localmemory(::Type{T}, ::Val) where {T} =
+    error("Local memory used outside kernel or not captured")
+
 """
     shfl_down(val::T, offset::Integer) where T
 
@@ -219,7 +286,7 @@ Returns a vector of `DataType`s supported on `backend`
     Backend implementations **must** implement this function
     only if they support `shfl_down` for any types.
 """
-shfl_down_types(::Backend) = DataType[]
+shfl_down_types(_) = DataType[]
 
 
 """
@@ -281,8 +348,27 @@ end
     ```
     If the backend does not support printing,
     define it to return `nothing`.
+
+The generic fallback prints on the host, which keeps CPU backends working.
+`Val` arguments are unwrapped, since `KernelAbstractions.@print` uses them to
+pass literal strings through to backends that require compile-time format strings.
 """
-function _print end
+@generated function _print(items...)
+    args = []
+
+    for i in 1:length(items)
+        item = :(items[$i])
+        T = items[i]
+        if T <: Val
+            item = QuoteNode(T.parameters[1])
+        end
+        push!(args, item)
+    end
+
+    return quote
+        print($(args...))
+    end
+end
 
 
 """
@@ -308,6 +394,15 @@ struct Kernel{B, Kern}
     kern::Kern
 end
 
+"""
+    check_launch_args(numworkgroups, workgroupsize)
+
+Validate the launch configuration passed to a [`Kernel`](@ref), throwing an
+`ArgumentError` if either argument has more than 3 dimensions.
+
+Backends may call this from their kernel-launch method instead of writing their
+own check.
+"""
 function check_launch_args(numworkgroups, workgroupsize)
     length(numworkgroups) <= 3 ||
         throw(ArgumentError("`numworkgroups` only accepts up to 3 dimensions"))
@@ -398,7 +493,8 @@ function argconvert end
     KI.kernel_function(::NewBackend, f::F, tt::TT=Tuple{}; name=nothing, kwargs...) where {F,TT}
 
 Low-level interface to compile a function invocation for the currently-active GPU, returning
-a callable kernel object. For a higher-level interface, use [`KI.@kernel`](@ref).
+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.

lib/KernelInterface/test/Project.toml

@@ -0,0 +1,7 @@
+[deps]
+Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[compat]
+Aqua = "0.8"