feat(Fields): add soa_broadcast! fast path for point-wise kernels#2543
feat(Fields): add soa_broadcast! fast path for point-wise kernels#2543haakon-e wants to merge 1 commit into
Conversation
|
cc: @imreddyTeja here's the hacky fix that helps my P3 quadrature kernels compiling in reasonable time. |
Add `Fields.soa_broadcast!(f, out, args...)`, a restricted point-wise fill that indexes the underlying arrays linearly rather than through `CartesianIndex`. This avoids the `CartesianIndex` code generation that makes many-call-site point-wise kernels (e.g. quadrature integrals) compile slowly on the GPU. Restricted to uniform same-space, unmasked, scalar-field arguments; provides CPU and CUDA methods.
haakon-e
force-pushed
the
he/soa-broadcast
branch
from
27279a3 to
51f7bba
Compare
| n > 0 || return nothing | ||
| # A plain launch (no forced `always_inline`): forcing inlining of the per-point | ||
| # function balloons the IR and the LLVM compile for many-call-site functions. | ||
| kernel = CUDA.@cuda launch = false soa_broadcast_kernel!(f, outs, ins, n) |
There was a problem hiding this comment.
I think the main improvement in compilation time might come from not having always_inline. I previously found that the Cartesian indexing does slow down compilation, but not severely.
|
The compile-time decomposition promised above. The standalone benchmark scripts are at On a many-field point-wise kernel, the compile-time advantage of Compile time, full Rosenbrock 2M+P3 tendency, A100:
Method, variant definitions, full tables, per-component decomposition, and caveatsMethod. Compile time is the first kernel launch in a fresh Julia process (CUDA compiles lazily on first launch), with a warm launch subtracted to remove device runtime. Each configuration runs in its own fresh process, so the CUDA compile cache and Julia codegen are cold. Cheap variants are repeated and the median reported with the min-max spread; the intractable Variants. Each reverts exactly one choice of
Full Rosenbrock tendency.
On the exact kernel the warm runtime separates: Instantaneous tendency (lighter reference). The instantaneous 2M+P3 tendency the Rosenbrock step wraps compiles for every variant, so it gives the full decomposition the intractable variants cannot.
Per-component, relative to Synthetic case (ClimaCore only). A tunable point-wise function of Caveats.
|
|
I'll close this since this isn't actually being merged. |
2 participants
Purpose
Add
Fields.soa_broadcast!(f, out, args...), a restricted, fast-compiling stand-in forBase.copyto!of a point-wise broadcast.What
soa_broadcast!applies a point-wise functionfto the per-point scalar values ofargsand writes the result into the pre-existing fieldoutin a single device pass, indexing the underlying arrays linearly rather than throughCartesianIndex.freturns a scalar (for a scalarout) or aNamedTuplematchingeltype(out).Why
For many-call-site point-wise functions, for example the quadrature-based 2M+P3 microphysics tendencies in ClimaAtmos, the standard broadcast makes the GPU kernel compile very slowly: minutes for the instantaneous tendency, and more than 30 minutes without completion for the full Rosenbrock 2M+P3 tendency. A compile-time decomposition (see the benchmark comment below) identifies the cause: forcing
always_inlineon the broadcast kernel balloons the IR and the LLVM compile.soa_broadcast!recovers the fast compile chiefly by not forcing inlining. Two further choices are secondary for compile time and kept for other reasons: linear indexing (instead ofCartesianIndex), which does not change compile time but gives a roughly 1.6x faster warm kernel on the heavy tendency and reaches the speed of a hand-writtenCuArraykernel without leaving the Fields interface; and extracting inputs withmaprather thanntuple(_, Val(N)), the safe choice for heterogeneous inputs (theN > 10boxing fallback cited earlier did not reproduce on the benchmarked Julia/CUDA).Notes
dr/data_style_broadcast_getindextargets a related goal (a specializedBase.getindexforBroadcasted{<:DataStyle}) via the generic broadcast path; this is an orthogonal, restricted fast path and the two do not overlap at the file level.