Proposal: portable compiler options for kernels (fastmath, always_inline, maxthreads, …) #839

Draft for discussion. Line references are to the vc/groupreduce branch of #830 (stacked on #831); backend references are to the vc/ki-subgroup-ops branches of CUDA.jl, AMDGPU.jl and Metal.jl (JuliaGPU/CUDA.jl#3331, JuliaGPU/AMDGPU.jl#1133, JuliaGPU/Metal.jl#998), OpenCL.jl and oneAPI.jl master, and GPUCompiler 2.11.1.

1. Motivation

Molly's pairwise tile kernels used to be compiled with

@cuda launch=false always_inline=true fastmath=(T == Float32) maxregs=force_maxregs force_kernel!(...)

where maxregs comes from autotuning or MOLLY_CUDA_FORCE_MAXREGS. KernelAbstractions
has no portable way to say any of this, so the port (JuliaMolSim/Molly.jl#297) does two things
around KA:

1. It builds a new backend that carries options. tile_kernel_config(::CUDABackend, T, maxregs)
   (ext/MollyCUDAExt.jl:47-52) returns
   CUDABackend(; prefer_blocks=backend.prefer_blocks, always_inline=true, fastmath=pairwise_fastmath(T)).
   This only works because CUDA.jl happens to store always_inline/fastmath in
   CUDABackend (CUDA.jl #3282). It has to copy the other fields by hand, and it doesn't
   generalize: ROCBackend and MetalBackend have no fields.
2. It skips KA's launch and calls KernelInterface directly. launch_ki_kernel!
   (src/gpu_tiles.jl:170-175) does
   KI.kernel_function(backend, f, KI.argument_types(backend, args); options...).
   KI.argument_types is internal (it isn't in KI's public list,
   lib/KernelInterface/src/KernelInterface.jl:39). KI.@launch can't be used here,
   because the options are a backend-dependent NamedTuple, (;) or (; maxregs), and
   @launch only takes literal key=value keywords (launch.jl:446-460).

Two more issues show up along the way:

• Precedence is accidental. KI.kernel_function(::CUDABackend, ...) calls
  cufunction(...; name, backend.always_inline, backend.fastmath, kwargs...)
  (CUDA.jl CUDACore/src/CUDAKernels.jl:108-111). In Julia the later keyword in a splat
  wins, so an explicit fastmath= overrides the backend field without an error. That is the
  right behavior, but nothing documents it, and backends that don't put the user's
  kwargs... last behave differently.
• @kernel users can't pass options at all. The only route is the backend hook
  KA.compiler_options(::Kernel{B}) (src/backend_launch.jl:53-61). CUDA.jl and AMDGPU.jl
  override it to return (; maxthreads = prod(static workgroupsize)). It is called from
  compile (src/backend_launch.jl:128-132), and users have no input into it.

Prior discussion: #628 (open; "options for more aggressive inlining". @maleadt suggested
having the kernel object carry the option and the backends act on it), #550 (open PR;
@kernel fastmath=true as a source-level @fastmath transform, modeled on
inbounds=true from #429), #643 (@fastmath + Float64 PTX failure, i.e. source-level
fastmath has sharp edges), #740 (querying a backend's recommended group size). No issue
covers maxregs or launch bounds.

2. What backends accept today

Options accepted by kernel_function / the native *function. A dash means the backend
doesn't accept the option. Unless noted otherwise, an unknown keyword ends up as a
MethodError from the GPUCompiler target's @kwdef constructor.

Notes:

• fastmath means different things. Julia's @fastmath rewrites the source. CUDA/PTX
  sets FTZ and afn per instruction. Metal picks air.fast_*. OpenCL's equivalent is
  -cl-fast-relaxed-math. GPUCompiler.apply_fastmath! (src/optim.jl:19) doesn't depend
  on the target, so SPIR-V and GCN could support fastmath by adding a fastmath field to
  their targets.

• Names diverge: minthreads is an exact requirement on CUDA but a lower bound on AMDGPU.

• In other ecosystems:

  • CUDA C++: __launch_bounds__(maxThreads, minBlocksPerSM, maxThreadsPerCluster) and
    __maxnreg__(N). -use_fast_math and --maxrregcount are per translation unit.
  • HIP: __launch_bounds__(maxThreads, minWarpsPerEU). The second argument means
    something different from CUDA's. HIP also has amdgpu_flat_work_group_size,
    amdgpu_waves_per_eu, -munsafe-fp-atomics and -ffast-math.
  • OpenCL C: reqd_work_group_size, work_group_size_hint,
    intel_reqd_sub_group_size, and the build option -cl-fast-relaxed-math.
  • SYCL: [[sycl::reqd_work_group_size]], [[sycl::reqd_sub_group_size]], and oneAPI
    compile-time kernel properties (properties{work_group_size<...>, sub_group_size<N>},
    Intel grf_size<256>, the analog of maxregs), attached at the launch site but
    type-level.
  • Metal: [[max_total_threads_per_threadgroup(N)]], or
    MTLComputePipelineDescriptor.maxTotalThreadsPerThreadgroup.
  • Kokkos: LaunchBounds<maxT, minB> as an execution-policy template parameter.
  • Triton: num_warps, num_stages and maxnreg are launch-site meta-parameters that are
    part of the compile cache key, and triton.autotune searches over them.

  The common pattern: a handful of portable properties (launch bounds, sub-group size), plus
  vendor extensions. They are attached at or near the launch site, and their values are
  part of the compilation key.

3. Taxonomy and semantics

(a) Portable options. Every KI backend must accept these. Each one has a defined
meaning, and says whether it is a hint or a requirement.

maxthreads keeps the CUDA/AMDGPU spelling, which KA's own compiler_options already
emits, so no migration is needed. Open question 1 covers renaming it.

Options left out of the portable set on purpose:

• sub_group_size. KI already fixes the sub-group width to KI.sub_group_size(backend)
  (launch.jl:393-395). A per-kernel request would break that promise. Backends keep
  rejecting a conflicting value, as AMDGPU does for wavefrontsize64.
• minthreads / required size. The two existing meanings conflict, see §2. A future
  portable reqd_work_group_size could be added, but KA's static workgroup size already
  covers most uses.

(b) Backend-specific options. maxregs, blocks_per_sm, unsafe_fp_atomics, arch,
ptx, Metal versions, extensions, debug_level and similar are passed through
unchanged. An unknown key errors, as it does today. Portable code selects them by
dispatching on the backend type, which is the Molly pattern:
opts(::CUDABackend, T, r) = (; maxregs = r), opts(_, T, r) = (;). Silently ignoring
unknown keys would hide typos like maxreg=64.

Precedence, documented in KI.kernel_function: explicit keyword > option carried by
the backend value (CUDABackend(fastmath=...)) > global default (--math-mode=fast).
Backends implement this by splatting the user's kwargs... last. CUDA.jl already does.

4. API design options

Type stability and caching. Options don't need to be compile-time constants. Every
backend's compiler_config(dev; kwargs...) hashes the keyword values at run time
(CUDA.jl compilation.jl:214-221, POCL src/pocl/compiler/compilation.jl:195-208), and
GPUCompiler's target hash covers fastmath, maxregs and the launch bounds
(ptx.jl:46-50), as its cache token covers always_inline. A concrete NamedTuple
type is enough for inference: (; maxregs::Int) with a runtime value is fine. Values
only need to be type parameters if KA itself rewrote the kernel source based on them
(option A). So B stores options as a field, and the field's type is a type parameter.

Recommendation: D + B now, C not at all, A later at most as sugar.

D. KernelInterface

# lib/KernelInterface/src/launch.jl, kernel_function docstring (replacing lines 376-377):
# Compiler options. Backends **must** accept the portable options below, with the
# documented meaning; they **may** ignore hints:
# - `always_inline::Bool` (hint), `fastmath::Bool` (hint)
# - `maxthreads::Union{Integer,Dims}` (requirement: `max_work_group_size(kernel)` must not
#   exceed `prod(maxthreads)`)
# Other options are backend-specific: backends throw for options they don't support.
# Explicit options take precedence over options carried by `backend`.

const PORTABLE_COMPILER_OPTIONS = (:always_inline, :fastmath, :maxthreads)  # public
@public argument_types  # make it public; document it next to kernel_function

KI.@launch also accepts a splatted NamedTuple of compiler options. In split_kwargs, a
:... expression goes to compiler_kwargs (launch.jl:449-458 currently throws
"Invalid keyword argument"):

opts = tile_options(backend, T, maxregs)       # (;) or (; maxregs = 64)
KI.@launch backend numgroups=n workgroupsize=w always_inline=true opts... force_kernel!(args...)

A conformance test in KI's test suite (and in each backend's KI tests) compiles a trivial
kernel with each portable option and checks that max_work_group_size(k) <= maxthreads.

B. KernelAbstractions

# src/KernelAbstractions.jl:502, new trailing type parameter with a default
struct Kernel{Backend, WorkgroupSize <: _Size, NDRange <: _Size, Fun, Opts <: NamedTuple}
    backend::Backend
    f::Fun
    options::Opts
end
Kernel{B, WS, ND, F}(backend, f) where {B, WS, ND, F} = Kernel{B, WS, ND, F, @NamedTuple{}}(backend, f, (;))
options(k::Kernel) = k.options          # public accessor
Base.similar(k::Kernel{D, WS, ND, F0, O}, f::F) where {D, WS, ND, F0, O, F} =
    Kernel{D, WS, ND, F, O}(k.backend, f, k.options)

# src/macros.jl:53-62, constructors gain keywords
$name(dev; kw...)              = $_name(dev, DynamicSize(), DynamicSize(); kw...)
$name(dev, size; kw...)        = $_name(dev, StaticSize(size), DynamicSize(); kw...)
$name(dev, size, range; kw...) = $_name(dev, StaticSize(size), StaticSize(range); kw...)
# construct(backend, sz, range, f; kw...) = Kernel{...}(backend, f, values(kw))

# src/backend_launch.jl:128-132: user options win over the backend hook
@inline function compile(obj::Kernel, ctx, args::Tuple)
    b = backend(obj)
    tt = argument_types(b, ctx, args)
    return KI.kernel_function(b, obj.f, tt; merge(compiler_options(obj), options(obj))...)
end

Usage (Molly, after the change):

fm(::Type{Float32}) = true; fm(::Type) = false
backend_opts(::CUDABackend, maxregs) = maxregs === nothing ? (;) : (; maxregs)
backend_opts(_, _) = (;)

k = force_kernel!(backend, n_items; always_inline=true, fastmath=fm(T), backend_opts(backend, maxregs)...)
k(args...; ndrange = n_items * ngroups)

Details:

• Backend hook unchanged. compiler_options(::Kernel{B}) stays the backend hook
  (implementations.md:93-94). KA merges the user's options over it, so CUDA's derived
  maxthreads gives way to an explicit one. No backend has to change for B.
• Static workgroup size vs maxthreads. If a static workgroup size exceeds a
  user-given maxthreads, KA throws ArgumentError at construction, not at the first
  launch.
• Autotuning. launch_kernel compiles before it tunes (backend_launch.jl:107-115).
  KI.launch_configuration(kernel) and max_work_group_size(kernel) are queried on the
  compiled kernel, so they already reflect maxregs/maxthreads (CUDA's occupancy API
  sees the register count; HIP's MAX_THREADS_PER_BLOCK reflects the flat-work-group
  bound). Searching over maxregs means constructing one Kernel per candidate. Each
  candidate is its own cache entry.
• Backend-carried options (CUDABackend(always_inline, fastmath), oneAPIBackend(always_inline))
  stay as backend-wide defaults with the precedence above. They don't need deprecating.
  The docs should point to the per-kernel options, and new backends shouldn't add such
  fields.

Why not A or C (yet)

C adds nothing that B can't do. f(backend, ws; opts...)(args...; ndrange) is already
a one-liner, and putting it in the launch keywords would make launch_tuple
(backend_launch.jl:85) carry compile-time state.

A is only worth adding as sugar that sets default construction options:
@kernel always_inline=true function ... would store defaults that B's keywords override.
That would be a separate mechanism from #550's source-level @fastmath. Defer it.

5. Backend contract changes and migration

1. KI docs: the kernel_function text above, the precedence rule, and argument_types
   made public. implementations.md §"Launching @kernel kernels" (lines 79-99) gets one
   paragraph: "KA merges KA.options(kernel) over compiler_options(kernel); backends
   must not drop keys they don't know from kwargs; they must forward them so that
   unsupported options error."

2. Per backend:

   backend	always_inline	fastmath	maxthreads
   CUDA.jl	done	done	done
   AMDGPU.jl	accept the keyword and ignore it, or honor false	accept; implement via GCNCompilerTarget.fastmath + apply_fastmath! (GPUCompiler)	done
   Metal.jl	done	done (pass-through); add it to COMPILER_KWARGS	accept; clamp max_work_group_size(kernel) (or set the pipeline descriptor's maxTotalThreadsPerThreadgroup)
   POCL / OpenCL.jl / oneAPI.jl	done	accept; implement via SPIRVCompilerTarget.fastmath + apply_fastmath!	accept; clamp max_work_group_size(kernel)

   OpenCL.jl and oneAPI.jl still subtype KA.GPU and use their own launch. They get this
   when they move to KI.

3. Migration: everything is additive. Kernel's new type parameter is trailing and
   has a default constructor. Code that writes Kernel{B,WS,ND,F} as a type (dispatch)
   still matches. Code that constructs Kernel{B,WS,ND,F}(backend, f) keeps working
   through the compatibility constructor.

Minimal first PR (KI only, about 30 lines): make KI.argument_types public, and let
KI.@launch accept opts.... That unblocks Molly's raw-KI path immediately. Second
PR: the portable-option wording plus a KI conformance test, and POCL accepting
fastmath/maxthreads. Third PR: the Kernel options field and constructor
keywords (B).

6. Open questions

1. Naming. maxthreads (existing, zero migration) vs max_work_group_size, which
   matches KI's vocabulary but collides with the launch keyword of the same name
   (launch.jl:12-13, LAUNCH_KWARGS at line 411). Or work_group_size_bound?
2. fastmath semantics. Should portable fastmath promise FTZ, or only afn-style
   approximations? Should fastmath=false override a global --math-mode=fast?
3. Occupancy hint. Is there a portable form of blocks_per_sm / waves_per_eu
   (e.g. min_groups_per_compute_unit)? Only CUDA implements one today.
4. Hints on backends that ignore them. Should those warn once (debug-level log), so
   that users see fastmath doing nothing on AMDGPU?
5. Capability query. Should backends expose
   KI.supported_compiler_options(backend)::Tuple{Vararg{Symbol}}, so that portable code
   can filter backend-specific options instead of dispatching on the backend type?
6. KA.options(kernel) and similar. Should they be public for wrappers that rebuild
   kernels (Enzyme once it's re-added, AcceleratedKernels)?
7. Kernel-level @kernel defaults (A). Worth it, and should fastmath there mean
   compiler fastmath or attempt to force fastmath at the kernel level #550's source transform?

🤖 Generated with Claude Code