Skip to content

feat(compile): composite emission seam, GELU legalization pass, @Lowering (#1342, part 2) - #1346

Merged
michalharakal merged 1 commit into
SKaiNET-developers:developfrom
michalharakal:feature/gelu-composite-lowering
Oct 14, 2026
Merged

michalharakal merged 1 commit into
SKaiNET-developers:developfrom
michalharakal:feature/gelu-composite-lowering

Conversation

@michalharakal

Copy link
Copy Markdown
Contributor

Summary

Part 2 of #1342, on top of #1344: the op metadata becomes a per-target lowering decision.

  • LoweringSpec + ConversionContext.emitComposite (sk.ainet.compile.hlo.lowering): one record per fused-able op — composite name, the node parameters that become composite_attributes, a version, and a single decomposition builder. spec.lower(...) emits stablehlo.composite "skainet.gelu" %x {composite_attributes = {approximate = false}, decomposition = @skainet_gelu_approximate_false_2x3xf32, version = 1 : i32} : (T) -> T plus one func.func private per (op, attributes, type) when context.granularity?.keepFused(op) is true, and the same builder's lines inline otherwise. This is the first consumer of OpGranularityPolicy, which was threaded through the converter but never read.
  • Module trailer: helper functions are emitted after the entry function (ConversionContext.emitModuleTrailer, claimSymbol for dedup), so MlirParser still finds @main first and MlirValidator's module-wide SSA rule holds because the helper bodies draw names from the shared counter. The text-level StableHloOptimizer passes tolerate the two-function module (tested).
  • Specs: GeluLowering.spec (skainet.gelu, attribute approximate) and ErfLowering.spec (skainet.erf); convertGelu and the erf case of UnaryMathConverter go through spec.lower. Portable emission is byte-identical to feat(ops)!: exact (erf) GELU and erf op (#1342, part 1) #1344.
  • GeluApproximationLegalizationPass (skainet-compile-opt): a target with no erf lowering rewrites gelu(approximate = false) → gelu(approximate = true) and records legalizedFrom / maxAbsError = 4.8e-4 on the node plus a GELU_APPROX_LEGALIZED diagnostic; erf nodes are left to fail loudly. Opt-in per target via TargetOptimizers.registerDagPasses.
  • @Lowering annotation (sk.ainet.lang.ops) on TensorOps.gelu / erf; OperatorDocProcessor writes a lowering field into operators.json (both FunctionDoc model copies default it to empty, so old JSON still deserializes); the coverage matrix gains a Lowering column (composite · decompose · native).
  • Docs: a "Fused ops and stablehlo.composite" section in the HLO tutorial; CHANGELOG.

Nothing registers a fused policy or the legalization pass in-repo yet — both are seams a target plugs into (FusedOpAllowList("my-npu", setOf("gelu")), TargetOptimizers.registerDagPasses("my-npu") { … }), as the #1342 lowering-metadata section describes. Composite emission is currently unary-only (emitComposite requires one operand); SDPA (#1343) will extend it.

Test plan

  • CompositeEmissionTest (hlo): composite + private fn after @main, MlirParser entry = @main, MlirValidator clean, inline when policy is null or decomposes gelu, dedup per (attributes, shape), skainet.erf, composite body ≡ inline decomposition modulo SSA names, convertWithOptimization on a two-function module
  • GeluApproximationLegalizationPassTest (opt): rewrite + recorded parameters + diagnostics + edges, no-op on tanh nodes, idempotent, registry opt-in through dagPipelineFor vs a portable target
  • OperatorDocProcessorTest.testLoweringFlowsIntoJson
  • ActivationOperationsConverterTest / UnaryMath erf tests from feat(ops)!: exact (erf) GELU and erf op (#1342, part 1) #1344 unchanged
  • apiDump, generateDocs (matrix column), scripts/pr-gate.sh --quick and full gate

…ring (SKaiNET-developers#1342, part 2)

The exact-GELU flag from part 1 becomes a per-target lowering decision.

- LoweringSpec + ConversionContext.emitComposite (sk.ainet.compile.hlo.lowering):
  one record per fused-able op — composite name, the node parameters that
  become composite_attributes, a version, one decomposition builder.
  spec.lower(...) emits `stablehlo.composite "skainet.gelu" %x {...,
  decomposition = @fn, version = 1 : i32}` plus one `func.func private` per
  (op, attributes, type) when the target's OpGranularityPolicy keeps the op
  fused, and the same builder's lines inline otherwise. First consumer of
  OpGranularityPolicy.keepFused.
- Module trailer: helper functions follow the entry function
  (emitModuleTrailer / claimSymbol), so MlirParser still finds @main first
  and the module-wide SSA rule holds (helper bodies use the shared counter).
- Specs: GeluLowering.spec (skainet.gelu, attribute approximate) and
  ErfLowering.spec (skainet.erf); convertGelu and UnaryMathConverter's erf
  case lower through them. Portable emission is unchanged.
- GeluApproximationLegalizationPass (skainet-compile-opt): targets without an
  erf lowering rewrite gelu(approximate = false) to the tanh form and record
  legalizedFrom / maxAbsError = 4.8e-4 on the node plus a
  GELU_APPROX_LEGALIZED diagnostic; erf nodes are left to fail loudly.
  Opt-in via TargetOptimizers.registerDagPasses.
- @lowering annotation on TensorOps.gelu/erf; OperatorDocProcessor writes a
  `lowering` field into operators.json; the coverage matrix gains a Lowering
  column (composite · decompose · native).
- Docs: "Fused ops and stablehlo.composite" in the HLO tutorial; CHANGELOG.

Tests: CompositeEmissionTest (7), GeluApproximationLegalizationPassTest (3),
OperatorDocProcessorTest.testLoweringFlowsIntoJson. API dumps regenerated.
@michalharakal
michalharakal merged commit 58f2e18 into SKaiNET-developers:develop Oct 14, 2026
18 checks passed
@michalharakal michalharakal mentioned this pull request Oct 14, 2026
1 of 4 tasks
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