Repository navigation
[Feature]: exact (erf) GELU as a first-class op — CPU kernels and HLO lowering #1342
Description
🧠 D: DEFINE — Problem & Opportunity
TensorOps.gelu is the tanh approximation everywhere: DefaultCpuOps computes 0.5·x·(1 + tanh(√(2/π)·(x + 0.044715·x³))) and the HLO converter lowers gelu "using approximation". There is no erf op and no exact GELU.
Encoder checkpoints trained with hidden_act = "gelu" (ModernBERT, mmBERT, BERT-class models from Hugging Face, the Laya decision models) use the exact 0.5·x·(1 + erf(x/√2)). Measured on the Laya multilingual checkpoint while building skainet-transformers-inference-laya (transformers #487): with the engine's tanh GELU the decision-head outputs move by about 7.5e-4 against the reference; the parity gate for that model is 5e-5. The adapter therefore ships a workaround, LayaGelu, which copies the activation to a FloatArray, applies erf on the host and copies back — correct, but a host round trip per FFN per layer (22 layers × 2 per decision) on the eager path, outside every kernel tier and invisible to the compiled path.
🔍 A: ASSESS — Options
gelu(approximate = false)+erfinTensorOps(proposed): a boolean on the existing op plus a standaloneerf, default unchanged (tanh), so no existing model changes numerically.- A separate
geluExactop withouterf: smaller surface, buterfis useful on its own (normal CDF, GELU variants) and is what StableHLO/CHLO expose (chlo.erf). - Keep host-side workarounds in model modules: what Laya does today; every BERT-class consumer re-implements it and the compiled path cannot see it.
✅ R: RESOLVE — Decision
Option 1. Scope:
TensorOps:fun erf(x),fun gelu(x, approximate: Boolean = true);VoidTensorOps,RecordingExecution,ComputeGraphExecutorfollow.- CPU: scalar kernel (
kotlin.mathhas noerf; use the Abramowitz-Stegun 7.1.26 form with max error 1.5e-7, or theerfof the FFM/libm tier where available) and a Vector API variant; JNI/FFM tiers optional in the first PR. - Lowering: the node carries the
approximateflag as op metadata; the StableHLO emitter chooses between a decomposed form (chlo.erf) and astablehlo.compositethroughOpGranularityPolicy— see Lowering metadata below. - NN DSL:
GELU(approximate = …)module parameter; docs operator/kernel matrix row. - Tests:
erfagainst a table (|err| < 2e-7 in [-6, 6]), exact-vs-tanh GELU difference (max ≈ 2e-3 near |x| ≈ 2), parity ofGELU(approximate = false)with the host formula.
🧩 Lowering metadata — one op, exact on every path
The point of a native op is not the CPU kernel alone: gelu(approximate = false) and erf must survive direct execution, the recorded/optimized graph and every export with their semantics intact, and carry enough metadata that a lowering transformation can pick the best legal form for its target without re-deriving what the op means.
What travels with the op
-
Graph node:
Operation.parameters = { "approximate": Boolean }ongelu;erfis its own op name.Operation.parametersis already serialized byBaseOperation.serialize()and the JSON export, so the flag reaches every emitter for free. -
Recording path gap (must be fixed in the same PR):
RecordingExecution.geludelegates tobase.geluwithout recording a node (skainet-compile-core/.../tape/RecordingExecution.kt:443, the "unrecorded/less critical ops" block). A compiled graph therefore never sees GELU today — the activation is computed eagerly and vanishes. The PR records it likePowOperation(parameters = mapOf("scalar_exponent" to n)), andComputeGraphExecutordispatches"gelu" -> ops.gelu(x, approximate = node.parameters["approximate"]),"erf" -> ops.erf(x).DefaultExecutionTape.geluBackwardalready receivesattributesand has both derivative formulas in comments; it picks one from the flag. -
Lowering descriptor (new, declarative): one small record per fused-able op, registered next to its converter:
LoweringSpec( op = "gelu", compositeName = "skainet.gelu", // namespaced, as the StableHLO spec requires attributes = listOf("approximate"), // copied into composite_attributes version = 1, decomposition = GeluDecomposition, // the one "composite constructor" )
GeluDecompositionis a single Kotlin builder that emits the formula as StableHLO primitives. The same builder writes the body of the composite's@decompositionfunction and the inline fallback, so eager semantics, composite and decomposed form cannot drift. Forapproximate = falsethe builder useschlo.erfwhen the target toolchain accepts CHLO, otherwise the engine's own rational approximation with the same coefficients as the CPU kernel (Abramowitz–Stegun 7.1.26), which keeps DIRECT and lowered output within the 2e-7 kernel bound. -
Source annotation (docs): the operator matrix is generated by the docs KSP processor from
@Backend/@InProgress. A@Lowering(composite = "skainet.gelu", decomposes = true)annotation on theTensorOpsfunction (same package as@TensorOp,@DslOp) lets the matrix grow a Lowering column (composite · decompose · native) without hand edits, and gives tooling one place to read what an op can lower to.
Emission rules (StableHLO)
Checked against the StableHLO spec and ChloOps.td on 2026-10-14: the core op set has no erf (the alphabetical list jumps from dynamic_update_slice to exp2); chlo.erf, chlo.erfc, chlo.erf_inv live in CHLO and chlo-legalize-to-stablehlo provides their polynomial expansion; stablehlo.composite takes name, composite_attributes, decomposition, version and by definition "can be replaced with its decomposition without changing program semantics"; stablehlo-legalize-composite-to-call does exactly that.
OpGranularityPolicy.keepFused("gelu") |
emitted form |
|---|---|
null policy / false (portable default, llvm-cpu) |
decomposed: approximate = true → the tanh chain the converter emits today; approximate = false → 0.5·x·(1 + chlo.erf(x·rsqrt(2))) (or the primitive-only rational form with --no-chlo) |
true (vendor target with a GELU/erf kernel, e.g. the torq allow-list) |
%y = stablehlo.composite "skainet.gelu" %x { composite_attributes = { approximate = false }, decomposition = @skainet_gelu_erf_f32_128x768, version = 1 : i32 } : (tensor<128x768xf32>) -> tensor<128x768xf32> plus one func.func private @skainet_gelu_erf_f32_128x768 per (dtype, shape) built by GeluDecomposition |
A vendor compiler pattern-matches the composite name + attributes and maps it to its kernel; everyone else inlines or runs legalize-composite-to-call and gets the exact formula back. erf follows the same two rows with "skainet.erf".
This is also the first real consumer of OpGranularityPolicy: the policy is threaded into ConversionContext.granularity but no converter calls keepFused yet. The GELU PR adds the generic ConversionContext.emitComposite(spec, operands, attrs) helper (composite op + deduplicated decomposition function) that layernorm and scaled-dot-product attention can then reuse, which is what the policy KDoc already promises.
Degrade visibly, never silently
Today the exact→tanh substitution happens implicitly on every path, which is how the 7.5e-4 Laya drift stayed invisible. After this change a target that has no erf anywhere (no CHLO, no composite match) gets an explicit TargetOptimization pass in skainet-compile-opt that rewrites gelu(approximate = false) → gelu(approximate = true) and records it on the node (parameters["legalizedFrom"] = "gelu(approximate=false)", parameters["maxAbsError"] ≈ 2e-3). The export report lists these rewrites, so a numerics regression is a visible legalization decision instead of a surprise.
Other emitters
- C codegen:
erff()from<math.h>;approximate = truekeeps the current tanh form. - Minerva:
MinervaCompatibilityValidator.activationOperationsisrelu, sigmoid, tanh; addgelu/erfor reject with a remediation that names the legalization pass above. - ONNX (if an exporter ever lands):
Gelu(opset 20) has the sameapproximateattribute ("none"/"tanh"),Erfis a core op — the metadata maps 1:1, which is the test that the descriptor is target-neutral.
🚀 C: COMMIT — Acceptance
-
TensorOps.erfandgelu(approximate = false)on the CPU path, scalar and Vector API -
RecordingExecutionrecordsgelu/erfwith theapproximateparameter;ComputeGraphExecutorandgeluBackwarddispatch on it; DIRECT vs OPTIMIZED on a Laya FFN block within 1e-6 - StableHLO: decomposed (
chlo.erf) and composite ("skainet.gelu"/"skainet.erf",version = 1) emission selected byOpGranularityPolicy;ConversionContext.emitCompositehelper; test runs the composite throughstablehlo-legalize-composite-to-calland compares with the decomposed module - One
GeluDecompositionbuilder feeds both the composite body and the inline fallback (asserted by a golden-text test) - Explicit exact→tanh legalization pass with recorded
legalizedFrom/maxAbsError, surfaced in the export report -
@Loweringannotation and the Lowering column in the generated operator matrix -
GELUmodule exposes the flag; operator matrix documents it - transformers:
LayaGelureplaced byctx.ops.gelu(x, approximate = false)with the 5e-5 parity gate still green (follow-up PR there, linked here)
Related: #1343 (SDPA key lengths — reuses the LoweringSpec / emitComposite seam defined here; together the first consumers of OpGranularityPolicy), SKaiNET-transformers#490.
Consumer evidence: transformers llm-inference/laya (LayaModules.kt, LayaGelu), measured 2026-10-12/13.
Activity
michalharakal commentedon Oct 14, 2026
Part 1 is up: #1344 (
erfop,gelu(approximate = false), CPU scalar + JVM Vector API kernels, the flag as op metadata through the recorder / graph executor / backward rule,chlo.erflowering, API dumps and docs).Two corrections to the numbers above, measured on the kernel while writing the tests:
maxAbsError = 4.8e-4.ErfKernel.erfpinsx == 0.0soerf(0) = 0exactly and the function stays odd.One behavioural note for the Lowering-metadata section: the KSP tracing wrapper now records
approximatenext to the shape attributes (gelu/erf joined the unary list inMethodAnalyzer), so the flag reachesnode.operation.parameterswith no extra plumbing, and the legacy recording decorator recordsgelu/erfinstead of delegating them unrecorded. Part 2 (composite emission behindOpGranularityPolicy, the explicit exact→tanh legalization pass,@Loweringcolumn) follows on a separate branch.feat(ops)!: add erf and exact gelu(approximate = false) (#1342, part 1)
feat(ops)!: exact (erf) GELU and erf op (#1344)
michalharakal commentedon Oct 14, 2026
Part 2 is up: #1346 —
LoweringSpec+ConversionContext.emitComposite(stablehlo.composite "skainet.gelu"/"skainet.erf"with a private decomposition function per (op, attributes, type) behindOpGranularityPolicy.keepFused, the policy's first consumer; helpers follow the entry function so the text tooling keeps finding@main),GeluApproximationLegalizationPass(exact→tanh recorded aslegalizedFrom/maxAbsError = 4.8e-4plus aGELU_APPROX_LEGALIZEDdiagnostic, opt-in per target), and@Loweringwith the Lowering column in the coverage matrix.With #1344 merged and #1346 open, every acceptance box above has an implementation except the transformers follow-up (
LayaGelu→ctx.ops.gelu(x, approximate = false)), which waits for the next engine release.emitCompositeis unary-only for now; #1343 extends it for SDPA.feat(compile): composite emission seam, GELU legalization pass, @lowe…
feat(compile): composite emission seam, GELU legalization pass, @lowe…
michalharakal commentedon Oct 14, 2026
Part 2 is merged: #1346 (
58f2e18), on top of #1344 (c1a4d4d). Both parts of the design above are now ondevelop:erfop andgelu(approximate = false)on DIRECT (CPU scalar + JVM Vector API), OPTIMIZED (graph executor, exact backward) and StableHLO (chlo.erfdecomposition).LoweringSpec+ConversionContext.emitComposite:stablehlo.composite "skainet.gelu"/"skainet.erf"when anOpGranularityPolicykeeps the op fused, inline decomposition otherwise, both from one decomposition builder.GeluApproximationLegalizationPass(opt-in per target viaTargetOptimizers.registerDagPasses), recordinglegalizedFrom/maxAbsError = 4.8e-4instead of substituting silently.@LoweringonTensorOps.gelu/erf, flowing intooperators.jsonand the Lowering column of the ops matrix.Follow-ups tracked elsewhere: the transformers Laya adapter replaces
LayaGeluwithctx.ops.gelu(x, approximate = false)after the next engine release; SDPA composite / key lengths is #1343; extendingemitCompositebeyond unary operands comes with it. Closing.