Skip to content

feat(ops)!: exact (erf) GELU and erf op (#1342, part 1) - #1344

Merged
michalharakal merged 2 commits into
SKaiNET-developers:developfrom
michalharakal:feature/exact-gelu-erf
Oct 14, 2026
Merged

michalharakal merged 2 commits into
SKaiNET-developers:developfrom
michalharakal:feature/exact-gelu-erf

Conversation

@michalharakal

Copy link
Copy Markdown
Contributor

Summary

Part 1 of #1342: erf as a first-class op and an exact GELU, carried as op metadata so DIRECT, OPTIMIZED and the StableHLO export all see the same choice.

  • TensorOps.erf(tensor) and gelu(tensor, approximate: Boolean = true). The default stays the tanh form, byte-for-byte, so no existing model changes numerically; approximate = false is 0.5·x·(1 + erf(x/√2)), the form BERT-class encoders and the Laya decision models were trained with. Measured on the kernel: the two differ by at most 4.73e-4 at |x| ≈ 2.70 (the issue's "≈ 2e-3" was an overestimate; the test pins the real band).
  • CPU: ErfKernel (Abramowitz–Stegun 7.1.26 in double, |err| < 1.5e-7, the same expression order as the Laya host workaround so the scalar path reproduces it bit-for-bit) plus a JVM Vector API lane kernel (VectorOperators.EXP + Horner) with a single-precision scalar tail of identical operation order; scalar-vs-vector parity ≤ 1e-6 across SPECIES±1 tails. Drive-by: DefaultCpuOps.gelu printed a literal ${tensor.dtype} in its error message; the stale @InProgress(... task-ops.md) reference is gone.
  • Metadata: GeluOperation(parameters["approximate"]) / ErfOperation in TensorOperations.kt; the KSP tracing wrapper records approximate next to the shape attributes (gelu/erf joined the unary list in MethodAnalyzer), so TraceToGraphBuilder puts it on the node; ComputeGraphExecutor dispatches on it; geluBackward picks the exact derivative Φ(x) + x·φ(x) from it; new erfBackward. The legacy RecordingTensorOpsDecorator now records gelu/erf instead of delegating them unrecorded (before this PR a compiled graph never saw GELU at all).
  • StableHLO: erf → chlo.erf %x : T -> T (CHLO unary assembly); exact GELU → 0.5·x·(1 + chlo.erf(x·√½)). One GeluLowering builder emits both the tanh chain (unchanged text) and the exact chain through an HloEmitter seam, so part 2 can wrap the same builder into a stablehlo.composite decomposition without a second copy of the formula.
  • DSLs / API: network {} gets gelu(approximate: Boolean = true, id) and erf(id) (the KSP default tables are name-keyed — approximate was added to both generators); GELU(name, approximate) module; Java TensorJavaOps.gelu(a, approximate) / erf(a), SequentialModelBuilder.gelu(approximate); ground-truth OperationParams.approximate.

Breaking

TensorOps.gelu gained a parameter: implementors of the interface must update the override (callers are source-compatible through the default). In the generated network DSL the activation convention puts parameters before id, so a positional gelu("name") becomes gelu(id = "name") (same as leakyRelu(slope, id)).

Test plan

  • skainet-backend-cpu: erf reference table (2e-7), oddness/saturation, exact gelu bit-equal to the host formula, default byte-identical to the old tanh path, exact-vs-tanh band 4.5e-4..5.0e-4 at |x| ∈ [2.5, 2.9], Int32 rejection; JVM erf_and_exact_gelu_match_scalar_reference_with_tails
  • skainet-compile-dag: GeluExactRecordingTest (flag on trace and node, DIRECT vs OPTIMIZED replay ≤ 1e-6), exact-gelu and erf finite-difference gradients, AutodiffCoverageTest (erf rule registered)
  • skainet-compile-core: legacy recorder records GeluOperation(approximate) / ErfOperation
  • skainet-compile-hlo: testGeluOperation unchanged, testGeluExactOperation, testErfOperationLowersToChlo, approximate=true equals default
  • skainet-lang-core: generated DSL gelu(approximate = false, id = …) compiles and builds; KSP processor tests
  • apiDump, generateDocs, scripts/pr-gate.sh --quick

Part 2 (composite emission behind OpGranularityPolicy, explicit exact→tanh legalization pass, @Lowering matrix column) follows in a separate PR, as laid out in #1342.

…lopers#1342, part 1)

Encoder checkpoints trained with hidden_act = "gelu" (BERT-class, Laya) use
the exact 0.5·x·(1 + erf(x/√2)); the engine only had the tanh approximation,
and no erf at all. The two differ by up to 4.7e-4 near |x| ≈ 2.7, which the
Laya decision head turns into a 7.5e-4 drift against a 5e-5 parity gate.

Op surface
- TensorOps.erf(tensor); TensorOps.gelu(tensor, approximate: Boolean = true).
  The default is the tanh form, byte-for-byte unchanged; approximate = false
  selects the exact form. Tensor extensions, VoidTensorOps, GELU(approximate)
  module, Java TensorJavaOps.gelu(a, approximate)/erf(a) and
  SequentialModelBuilder.gelu(approximate), ground-truth OperationParams.
- KSP: gelu and erf join the unary list in MethodAnalyzer so the tracing
  wrapper records the approximate flag next to the shape attributes; the
  name-keyed default tables of the network and graph DSL generators learn
  `approximate -> true` (KSP cannot read declared defaults).

CPU
- ErfKernel: Abramowitz–Stegun 7.1.26 in double (|err| < 1.5e-7), the same
  expression order as the Laya host workaround, origin pinned to signed zero;
  single-precision variants for the SIMD tail.
- DefaultCpuOps.gelu(approximate) / erf with the FP32 fast path; DefaultCpuOpsJvm
  lane kernels over VectorOperators.EXP with a scalar tail of identical
  operation order (parity ≤ 1e-6 across SPECIES±1 tails). Fixes the literal
  "${tensor.dtype}" in gelu's error message; drops the stale @inprogress.

Metadata through the compiled path
- GeluOperation(parameters["approximate"]) and ErfOperation; the legacy
  RecordingTensorOpsDecorator records both (gelu used to be delegated
  unrecorded, so compiled graphs never saw it); ComputeGraphExecutor
  dispatches on the flag; geluBackward picks the exact derivative
  Φ(x) + x·φ(x) from it; new erfBackward registered for autodiff coverage.
- StableHLO: erf lowers to `chlo.erf %x : T -> T`, exact GELU to
  0.5·x·(1 + chlo.erf(x·√½)). GeluLowering emits both chains through an
  HloEmitter seam (ContextEmitter today) so part 2 can reuse the builder for
  a stablehlo.composite decomposition.

Tests: CPU erf table / oddness / exact-vs-host bit equality / default
unchanged / exact-vs-tanh band 4.5e-4..5.0e-4; JVM vector parity;
GeluExactRecordingTest (flag on trace and node, DIRECT vs OPTIMIZED ≤ 1e-6);
exact-gelu and erf finite-difference gradients; recorder test; HLO exact and
erf lowering tests; generated DSL gelu(approximate = false) in the network
builder coverage test. API dumps and operator docs regenerated.

BREAKING CHANGE: TensorOps.gelu gained the approximate parameter — interface
implementors must update the override (callers keep compiling through the
default). In the generated network DSL a positional gelu("name") becomes
gelu(id = "name"), matching leakyRelu(slope, id).
…/gelu tests

JS and Wasm evaluate Float in double precision, so the exact-GELU host
comparison allows one ulp instead of bit equality, and the unsupported-dtype
check matches the dtype name rather than the JVM KClass rendering.
@michalharakal
michalharakal merged commit c1a4d4d into SKaiNET-developers:develop Oct 14, 2026
27 of 29 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant