diff --git a/CHANGELOG.md b/CHANGELOG.md index b641c1113..1c644f479 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,11 @@ ## [Unreleased] +### Added +- **Composite emission behind `OpGranularityPolicy`** (#1342 part 2): a `LoweringSpec` per fused-able op (`skainet.gelu`, `skainet.erf`) and `ConversionContext.emitComposite` emit `stablehlo.composite "" {composite_attributes, decomposition = @fn, version}` plus one private decomposition function per (op, attributes, type) when a target's policy keeps the op fused; otherwise the same builder emits the decomposition inline. Helper functions go to a module trailer after the entry function so existing text tooling keeps finding `@main` first. First consumer of `OpGranularityPolicy.keepFused`. +- **`GeluApproximationLegalizationPass`** (`skainet-compile-opt`): targets without any erf lowering rewrite `gelu(approximate = false)` to the tanh form explicitly, recording `legalizedFrom` and `maxAbsError = 4.8e-4` on the node and a `GELU_APPROX_LEGALIZED` diagnostic, instead of a silent default. Registered per target through `TargetOptimizers.registerDagPasses`. +- **`@Lowering` annotation** and a *Lowering* column in the generated operator coverage matrix (`composite · decompose · native`); `operators.json` carries the field. + ### Added - **`TensorOps.erf` and exact GELU** (#1342): `gelu(tensor, approximate: Boolean = true)` keeps the tanh form as the default and selects the exact `0.5·x·(1 + erf(x/√2))` with `approximate = false`, the form BERT-class encoders and the Laya decision models were trained with (the two differ by up to 4.7e-4 near |x| ≈ 2.7). CPU: Abramowitz–Stegun 7.1.26 erf evaluated in double (|err| < 1.5e-7) with a JVM Vector API lane kernel; `GELU(approximate)` module, `network {}` `gelu(approximate)`, Java `TensorJavaOps.gelu(a, approximate)` / `erf(a)` and `SequentialModelBuilder.gelu(approximate)`; exact backward rule and `erfBackward`. - **GELU carries its form as op metadata**: the flag is recorded as `parameters["approximate"]` on the graph node (`GeluOperation`, `ErfOperation`), read by `ComputeGraphExecutor`, the backward rule and the StableHLO converter, so the compiled path cannot silently diverge from eager execution. StableHLO lowers `erf` to `chlo.erf` and exact GELU to `0.5·x·(1 + chlo.erf(x·√½))`; one `GeluLowering` builder emits both forms. diff --git a/build-logic/convention/src/main/kotlin/GenerateDocumentationTask.kt b/build-logic/convention/src/main/kotlin/GenerateDocumentationTask.kt index dabbae093..e6f540819 100644 --- a/build-logic/convention/src/main/kotlin/GenerateDocumentationTask.kt +++ b/build-logic/convention/src/main/kotlin/GenerateDocumentationTask.kt @@ -130,7 +130,7 @@ abstract class GenerateDocumentationTask : DefaultTask() { appendLine("") appendLine("Generated from `operators.json` version `${module.version}` on ${formatTimestamp(module.timestamp)}.") appendLine("") - appendLine("Rows are `Operator.function` pairs. The `Validated` column shows whether the function's documentation has been DARC-validated by a reviewer (see xref:contributing/darc-workflow.adoc[DARC workflow]). Remaining columns are backends that appear in any function's `statusByBackend` map — a missing entry means the backend makes no claim about the function (treat it as \"unknown\", not \"not supported\").") + appendLine("Rows are `Operator.function` pairs. The `Validated` column shows whether the function's documentation has been DARC-validated by a reviewer (see xref:contributing/darc-workflow.adoc[DARC workflow]). The `Lowering` column lists the compile-path forms the function declares with `@Lowering` (`composite` = a `stablehlo.composite` when a target keeps it fused, `decompose` = StableHLO/CHLO primitives, `native` = a 1:1 target op); `—` means undeclared. Remaining columns are backends that appear in any function's `statusByBackend` map — a missing entry means the backend makes no claim about the function (treat it as \"unknown\", not \"not supported\").") appendLine("") if (rows.isEmpty() || allBackends.isEmpty()) { appendLine("NOTE: No backend status information found in the source data.") @@ -139,10 +139,10 @@ abstract class GenerateDocumentationTask : DefaultTask() { } // Header: row label, validation column, then one column per backend. - val colSpec = (listOf("2", "1") + List(allBackends.size) { "1" }).joinToString(",") + val colSpec = (listOf("2", "1", "1") + List(allBackends.size) { "1" }).joinToString(",") appendLine("[cols=\"$colSpec\", options=\"header\"]") appendLine("|===") - append("| Operator.function | Validated ") + append("| Operator.function | Validated | Lowering ") allBackends.forEach { append("| $it ") } appendLine("") appendLine("") @@ -150,6 +150,7 @@ abstract class GenerateDocumentationTask : DefaultTask() { rows.forEach { row -> append("| `${row.operator}.${row.function.name}` ") append("| ${darcCell(row.function, row.hasPartial)} ") + append("| ${row.function.lowering.ifEmpty { "—" }} ") allBackends.forEach { backend -> val raw = row.function.statusByBackend[backend] val cell = if (raw == null) "—" else shortStatus(raw) @@ -158,10 +159,11 @@ abstract class GenerateDocumentationTask : DefaultTask() { appendLine("") } - // Totals footer: validated count followed by per-backend done count. + // Totals footer: validated count, declared-lowering count, then per-backend done count. appendLine("") val validatedCount = rows.count { it.function.validated } - append("| *Done* | *$validatedCount / ${rows.size}* ") + val loweringCount = rows.count { it.function.lowering.isNotEmpty() } + append("| *Done* | *$validatedCount / ${rows.size}* | *$loweringCount / ${rows.size}* ") allBackends.forEach { backend -> val n = rows.count { isDone(it.function.statusByBackend[backend]) } append("| *$n / ${rows.size}* ") diff --git a/build-logic/convention/src/main/kotlin/models/DocumentationModels.kt b/build-logic/convention/src/main/kotlin/models/DocumentationModels.kt index 7b8717d9f..758995073 100644 --- a/build-logic/convention/src/main/kotlin/models/DocumentationModels.kt +++ b/build-logic/convention/src/main/kotlin/models/DocumentationModels.kt @@ -35,6 +35,8 @@ data class FunctionDoc( val validatedOn: String = "", val validatedCommit: String = "", val referencesChecked: Boolean = true, + /** Lowering forms from `@Lowering`: `composite`, `decompose`, `native` joined by ` · `; empty = undeclared. */ + val lowering: String = "", ) @Serializable diff --git a/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc b/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc index 8ec87a6ca..05805175f 100644 --- a/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc +++ b/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc @@ -3,80 +3,80 @@ Generated from `operators.json` version `0.58.0` on 2026-10-14. -Rows are `Operator.function` pairs. The `Validated` column shows whether the function's documentation has been DARC-validated by a reviewer (see xref:contributing/darc-workflow.adoc[DARC workflow]). Remaining columns are backends that appear in any function's `statusByBackend` map — a missing entry means the backend makes no claim about the function (treat it as "unknown", not "not supported"). +Rows are `Operator.function` pairs. The `Validated` column shows whether the function's documentation has been DARC-validated by a reviewer (see xref:contributing/darc-workflow.adoc[DARC workflow]). The `Lowering` column lists the compile-path forms the function declares with `@Lowering` (`composite` = a `stablehlo.composite` when a target keeps it fused, `decompose` = StableHLO/CHLO primitives, `native` = a 1:1 target op); `—` means undeclared. Remaining columns are backends that appear in any function's `statusByBackend` map — a missing entry means the backend makes no claim about the function (treat it as "unknown", not "not supported"). -[cols="2,1,1,1,1", options="header"] +[cols="2,1,1,1,1,1", options="header"] |=== -| Operator.function | Validated | apple | cpu | wasm +| Operator.function | Validated | Lowering | apple | cpu | wasm -| `TensorOps.add` | ✖ | — | — | — -| `TensorOps.subtract` | ✖ | — | — | — -| `TensorOps.multiply` | ✖ | — | — | — -| `TensorOps.divide` | ✖ | — | — | — -| `TensorOps.addScalar` | ✖ | — | — | — -| `TensorOps.subScalar` | ✖ | — | — | — -| `TensorOps.mulScalar` | ✖ | — | — | — -| `TensorOps.divScalar` | ✖ | — | — | — -| `TensorOps.rsubScalar` | ✖ | — | — | — -| `TensorOps.rdivScalar` | ✖ | — | — | — -| `TensorOps.matmul` | ✅ | — | — | — -| `TensorOps.matmulWeightTransposed` | ✖ | — | — | — -| `TensorOps.relayoutPackedWeightForKernels` | ✖ | — | — | — -| `TensorOps.transpose` | ✖ | — | — | — -| `TensorOps.permute` | ✖ | — | — | — -| `TensorOps.conv1d` | ✖ | — | — | — -| `TensorOps.conv2d` | ✖ | — | — | — -| `TensorOps.conv3d` | ✖ | — | — | — -| `TensorOps.convTranspose1d` | ✖ | — | — | — -| `TensorOps.maxPool2d` | ✖ | — | — | — -| `TensorOps.avgPool2d` | ✖ | — | — | — -| `TensorOps.upsample2d` | ✖ | — | — | — -| `TensorOps.reshape` | ✖ | — | — | — -| `TensorOps.flatten` | ✖ | — | — | — -| `TensorOps.concat` | ✖ | — | — | — -| `TensorOps.split` | ✖ | — | — | — -| `TensorOps.squeeze` | ✖ | — | — | — -| `TensorOps.unsqueeze` | ✖ | — | — | — -| `TensorOps.relu` | ✖ | — | — | — -| `TensorOps.leakyRelu` | ✖ | — | — | — -| `TensorOps.elu` | ✖ | — | — | — -| `TensorOps.softmax` | ✖ | — | — | — -| `TensorOps.logSoftmax` | ✖ | — | — | — -| `TensorOps.sigmoid` | ✖ | — | — | — -| `TensorOps.tanh` | ✖ | — | — | — -| `TensorOps.silu` | ✖ | — | — | — -| `TensorOps.gelu` | ✖ | — | — | — -| `TensorOps.sum` | ✖ | — | — | — -| `TensorOps.mean` | ✖ | — | — | — -| `TensorOps.variance` | ✖ | — | — | — -| `TensorOps.argMax` | ⚠ | — | — | — -| `TensorOps.sqrt` | ✖ | — | — | — -| `TensorOps.pow` | ✖ | — | — | — -| `TensorOps.powScalar` | ✖ | — | — | — -| `TensorOps.log` | ✖ | — | — | — -| `TensorOps.log2` | ✖ | — | — | — -| `TensorOps.log10` | ✖ | — | — | — -| `TensorOps.abs` | ✖ | — | — | — -| `TensorOps.sign` | ✖ | — | — | — -| `TensorOps.clamp` | ✖ | — | — | — -| `TensorOps.narrow` | ✖ | — | — | — -| `TensorOps.pad2d` | ✖ | — | — | — -| `TensorOps.unfold` | ✖ | — | — | — -| `TensorOps.lt` | ✖ | — | — | — -| `TensorOps.ge` | ✖ | — | — | — -| `TensorOps.tril` | ✖ | — | — | — -| `TensorOps.convert` | ✖ | — | — | — -| `TensorOps.gather` | ✖ | — | — | — -| `TensorOps.indexSelect` | ✖ | — | — | — -| `TensorOps.exp` | ✖ | — | — | — -| `TensorOps.expm1` | ✖ | — | — | — -| `TensorOps.erf` | ✖ | — | — | — -| `TensorOps.sin` | ✖ | — | — | — -| `TensorOps.cos` | ✖ | — | — | — -| `TensorOps.scaledDotProductAttention` | ✖ | — | — | — -| `Similarity.cosineDistance` | ✖ | ✅ | ✅ | ✅ +| `TensorOps.add` | ✖ | — | — | — | — +| `TensorOps.subtract` | ✖ | — | — | — | — +| `TensorOps.multiply` | ✖ | — | — | — | — +| `TensorOps.divide` | ✖ | — | — | — | — +| `TensorOps.addScalar` | ✖ | — | — | — | — +| `TensorOps.subScalar` | ✖ | — | — | — | — +| `TensorOps.mulScalar` | ✖ | — | — | — | — +| `TensorOps.divScalar` | ✖ | — | — | — | — +| `TensorOps.rsubScalar` | ✖ | — | — | — | — +| `TensorOps.rdivScalar` | ✖ | — | — | — | — +| `TensorOps.matmul` | ✅ | — | — | — | — +| `TensorOps.matmulWeightTransposed` | ✖ | — | — | — | — +| `TensorOps.relayoutPackedWeightForKernels` | ✖ | — | — | — | — +| `TensorOps.transpose` | ✖ | — | — | — | — +| `TensorOps.permute` | ✖ | — | — | — | — +| `TensorOps.conv1d` | ✖ | — | — | — | — +| `TensorOps.conv2d` | ✖ | — | — | — | — +| `TensorOps.conv3d` | ✖ | — | — | — | — +| `TensorOps.convTranspose1d` | ✖ | — | — | — | — +| `TensorOps.maxPool2d` | ✖ | — | — | — | — +| `TensorOps.avgPool2d` | ✖ | — | — | — | — +| `TensorOps.upsample2d` | ✖ | — | — | — | — +| `TensorOps.reshape` | ✖ | — | — | — | — +| `TensorOps.flatten` | ✖ | — | — | — | — +| `TensorOps.concat` | ✖ | — | — | — | — +| `TensorOps.split` | ✖ | — | — | — | — +| `TensorOps.squeeze` | ✖ | — | — | — | — +| `TensorOps.unsqueeze` | ✖ | — | — | — | — +| `TensorOps.relu` | ✖ | — | — | — | — +| `TensorOps.leakyRelu` | ✖ | — | — | — | — +| `TensorOps.elu` | ✖ | — | — | — | — +| `TensorOps.softmax` | ✖ | — | — | — | — +| `TensorOps.logSoftmax` | ✖ | — | — | — | — +| `TensorOps.sigmoid` | ✖ | — | — | — | — +| `TensorOps.tanh` | ✖ | — | — | — | — +| `TensorOps.silu` | ✖ | — | — | — | — +| `TensorOps.gelu` | ✖ | composite · decompose | — | — | — +| `TensorOps.sum` | ✖ | — | — | — | — +| `TensorOps.mean` | ✖ | — | — | — | — +| `TensorOps.variance` | ✖ | — | — | — | — +| `TensorOps.argMax` | ⚠ | — | — | — | — +| `TensorOps.sqrt` | ✖ | — | — | — | — +| `TensorOps.pow` | ✖ | — | — | — | — +| `TensorOps.powScalar` | ✖ | — | — | — | — +| `TensorOps.log` | ✖ | — | — | — | — +| `TensorOps.log2` | ✖ | — | — | — | — +| `TensorOps.log10` | ✖ | — | — | — | — +| `TensorOps.abs` | ✖ | — | — | — | — +| `TensorOps.sign` | ✖ | — | — | — | — +| `TensorOps.clamp` | ✖ | — | — | — | — +| `TensorOps.narrow` | ✖ | — | — | — | — +| `TensorOps.pad2d` | ✖ | — | — | — | — +| `TensorOps.unfold` | ✖ | — | — | — | — +| `TensorOps.lt` | ✖ | — | — | — | — +| `TensorOps.ge` | ✖ | — | — | — | — +| `TensorOps.tril` | ✖ | — | — | — | — +| `TensorOps.convert` | ✖ | — | — | — | — +| `TensorOps.gather` | ✖ | — | — | — | — +| `TensorOps.indexSelect` | ✖ | — | — | — | — +| `TensorOps.exp` | ✖ | — | — | — | — +| `TensorOps.expm1` | ✖ | — | — | — | — +| `TensorOps.erf` | ✖ | composite · decompose | — | — | — +| `TensorOps.sin` | ✖ | — | — | — | — +| `TensorOps.cos` | ✖ | — | — | — | — +| `TensorOps.scaledDotProductAttention` | ✖ | — | — | — | — +| `Similarity.cosineDistance` | ✖ | — | ✅ | ✅ | ✅ -| *Done* | *1 / 66* | *1 / 66* | *1 / 66* | *1 / 66* +| *Done* | *1 / 66* | *2 / 66* | *1 / 66* | *1 / 66* | *1 / 66* |=== Per-function detail including notes lives in xref:reference/operators/generated/index.adoc[Operator reference]. diff --git a/docs/modules/ROOT/pages/tutorials/hlo-getting-started.adoc b/docs/modules/ROOT/pages/tutorials/hlo-getting-started.adoc index d73e3ee2c..49065b4ee 100644 --- a/docs/modules/ROOT/pages/tutorials/hlo-getting-started.adoc +++ b/docs/modules/ROOT/pages/tutorials/hlo-getting-started.adoc @@ -325,6 +325,32 @@ class RgbEnhanceOp : HloConverter { } ---- +==== Fused ops and `stablehlo.composite` + +Some ops have a `LoweringSpec` (`sk.ainet.compile.hlo.lowering`): a namespaced composite name, the node parameters that become `composite_attributes`, a version, and one decomposition builder. Whether a target sees the composite or the primitives is decided per op by the target's `OpGranularityPolicy`: + +[source,kotlin] +---- +// A target that has a fused GELU kernel keeps the op; everything else decomposes. +val policy = FusedOpAllowList("my-npu", setOf("gelu")) +val module = toStableHlo(graph, "main", target = "my-npu", granularity = policy) +---- + +[source,mlir] +---- +%v8 = stablehlo.composite "skainet.gelu" %v7 {composite_attributes = {approximate = false}, + decomposition = @skainet_gelu_approximate_false_128x768xf32, version = 1 : i32} + : (tensor<128x768xf32>) -> tensor<128x768xf32> +... +func.func private @skainet_gelu_approximate_false_128x768xf32(%arg0: tensor<128x768xf32>) -> tensor<128x768xf32> { + // the same lines the portable emission inlines +} +---- + +The private decomposition functions follow the entry function in the module. A composite is replaceable by its decomposition by definition, so a vendor compiler may pattern-match `skainet.gelu` while `stablehlo-legalize-composite-to-call` (or plain inlining) restores the exact formula everywhere else. Specs today: `GeluLowering.spec` (`skainet.gelu`, attribute `approximate`) and `ErfLowering.spec` (`skainet.erf`). + +Targets that cannot lower `erf` at all register `GeluApproximationLegalizationPass` through `TargetOptimizers.registerDagPasses`; it rewrites `gelu(approximate = false)` to the tanh form and records `legalizedFrom` / `maxAbsError` on the node, so the substitution shows up in the pass diagnostics instead of happening silently. + ==== Debugging HLO Use SKaiNET's built-in debugging tools: diff --git a/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api b/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api index 12a5e408e..81d59d78e 100644 --- a/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api +++ b/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api @@ -53,11 +53,13 @@ public final class sk/ainet/compile/hlo/ConversionContext { public fun (Lsk/ainet/compile/hlo/TypeMapper;Lsk/ainet/lang/graph/ComputeGraph;Lsk/ainet/compile/hlo/ConstantMaterializationPolicy;Ljava/lang/String;Lsk/ainet/compile/target/OpGranularityPolicy;)V public fun (Lsk/ainet/compile/hlo/TypeMapper;Lsk/ainet/lang/graph/ComputeGraph;Lsk/ainet/compile/hlo/ConstantMaterializationPolicy;Ljava/lang/String;Lsk/ainet/compile/target/OpGranularityPolicy;Lsk/ainet/compile/hlo/ConversionErrorPolicy;)V public synthetic fun (Lsk/ainet/compile/hlo/TypeMapper;Lsk/ainet/lang/graph/ComputeGraph;Lsk/ainet/compile/hlo/ConstantMaterializationPolicy;Ljava/lang/String;Lsk/ainet/compile/target/OpGranularityPolicy;Lsk/ainet/compile/hlo/ConversionErrorPolicy;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public final fun claimSymbol (Ljava/lang/String;)Z public final fun clear ()V public final fun emitComment (Ljava/lang/String;)V public final fun emitEncodingAnnotation (Ljava/lang/String;ILsk/ainet/lang/tensor/ops/TensorSpec;)V public final fun emitLine (Ljava/lang/String;)V public final fun emitModuleDeclaration (Ljava/lang/String;)V + public final fun emitModuleTrailer (Ljava/lang/String;)V public final fun emitOperation (Ljava/lang/String;)V public final fun getContent ()Ljava/lang/String; public final fun getErrorPolicy ()Lsk/ainet/compile/hlo/ConversionErrorPolicy; @@ -66,6 +68,7 @@ public final class sk/ainet/compile/hlo/ConversionContext { public final fun getInputNodes (Lsk/ainet/lang/graph/GraphNode;)Ljava/util/List; public final fun getMaterializationPolicy ()Lsk/ainet/compile/hlo/ConstantMaterializationPolicy; public final fun getModuleDeclarations ()Ljava/lang/String; + public final fun getModuleTrailer ()Ljava/lang/String; public final fun getTarget ()Ljava/lang/String; public final fun getTypeMapper ()Lsk/ainet/compile/hlo/TypeMapper; public final fun getValueName (Ljava/lang/String;)Ljava/lang/String; @@ -405,6 +408,12 @@ public final class sk/ainet/compile/hlo/converters/ContextEmitter : sk/ainet/com public fun op (Ljava/lang/String;)V } +public final class sk/ainet/compile/hlo/converters/ErfLowering { + public static final field INSTANCE Lsk/ainet/compile/hlo/converters/ErfLowering; + public final fun emit (Lsk/ainet/compile/hlo/converters/HloEmitter;Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String; + public final fun getSpec ()Lsk/ainet/compile/hlo/lowering/LoweringSpec; +} + public final class sk/ainet/compile/hlo/converters/GatherOperationsConverter : sk/ainet/compile/hlo/StableHloOperationConverter { public fun ()V public fun convert (Lsk/ainet/lang/graph/GraphNode;Ljava/util/List;Lsk/ainet/compile/hlo/ConversionContext;)Lsk/ainet/compile/hlo/ConversionResult; @@ -417,6 +426,7 @@ public final class sk/ainet/compile/hlo/converters/GeluLowering { public final fun emit (Lsk/ainet/compile/hlo/converters/HloEmitter;Ljava/lang/String;Ljava/lang/String;Z)Ljava/lang/String; public final fun emitExact (Lsk/ainet/compile/hlo/converters/HloEmitter;Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String; public final fun emitTanh (Lsk/ainet/compile/hlo/converters/HloEmitter;Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String; + public final fun getSpec ()Lsk/ainet/compile/hlo/lowering/LoweringSpec; } public abstract interface class sk/ainet/compile/hlo/converters/HloEmitter { @@ -568,6 +578,34 @@ public final class sk/ainet/compile/hlo/generate/JvmHloGenerator { public static synthetic fun generateBlocking$default (Lsk/ainet/lang/model/Model;Lsk/ainet/lang/tensor/Tensor;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/compile/hlo/StableHloModule; } +public abstract interface class sk/ainet/compile/hlo/lowering/HloDecomposition { + public abstract fun emit (Lsk/ainet/compile/hlo/converters/HloEmitter;Ljava/lang/String;Ljava/lang/String;Ljava/util/Map;)Ljava/lang/String; +} + +public final class sk/ainet/compile/hlo/lowering/LoweringSpec { + public fun (Ljava/lang/String;Ljava/lang/String;Ljava/util/List;ILsk/ainet/compile/hlo/lowering/HloDecomposition;)V + public final fun component1 ()Ljava/lang/String; + public final fun component2 ()Ljava/lang/String; + public final fun component3 ()Ljava/util/List; + public final fun component4 ()I + public final fun component5 ()Lsk/ainet/compile/hlo/lowering/HloDecomposition; + public final fun copy (Ljava/lang/String;Ljava/lang/String;Ljava/util/List;ILsk/ainet/compile/hlo/lowering/HloDecomposition;)Lsk/ainet/compile/hlo/lowering/LoweringSpec; + public static synthetic fun copy$default (Lsk/ainet/compile/hlo/lowering/LoweringSpec;Ljava/lang/String;Ljava/lang/String;Ljava/util/List;ILsk/ainet/compile/hlo/lowering/HloDecomposition;ILjava/lang/Object;)Lsk/ainet/compile/hlo/lowering/LoweringSpec; + public fun equals (Ljava/lang/Object;)Z + public final fun getAttributes ()Ljava/util/List; + public final fun getCompositeName ()Ljava/lang/String; + public final fun getDecomposition ()Lsk/ainet/compile/hlo/lowering/HloDecomposition; + public final fun getOp ()Ljava/lang/String; + public final fun getVersion ()I + public fun hashCode ()I + public final fun lower (Lsk/ainet/compile/hlo/ConversionContext;Lsk/ainet/compile/hlo/converters/ContextEmitter;Ljava/lang/String;Ljava/lang/String;Ljava/util/Map;)Ljava/lang/String; + public fun toString ()Ljava/lang/String; +} + +public final class sk/ainet/compile/hlo/lowering/LoweringSpecKt { + public static final fun emitComposite (Lsk/ainet/compile/hlo/ConversionContext;Lsk/ainet/compile/hlo/lowering/LoweringSpec;Lsk/ainet/compile/hlo/converters/ContextEmitter;Ljava/util/List;Ljava/lang/String;Ljava/util/Map;)Ljava/lang/String; +} + public final class sk/ainet/compile/hlo/validation/BenchmarkResults { public fun (Ljava/lang/String;Ljava/util/List;Lsk/ainet/compile/hlo/validation/ConversionMetrics;Lsk/ainet/compile/hlo/validation/ConversionMetrics;Lsk/ainet/compile/hlo/validation/ConversionMetrics;)V public final fun component1 ()Ljava/lang/String; diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/ConversionContext.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/ConversionContext.kt index 260e7932e..8eb21ba53 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/ConversionContext.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/ConversionContext.kt @@ -47,6 +47,8 @@ public class ConversionContext @kotlin.jvm.JvmOverloads constructor( private val valueTypes = mutableMapOf() private val stringBuilder = StringBuilder() private val moduleDeclarationsBuilder = StringBuilder() + private val moduleTrailerBuilder = StringBuilder() + private val claimedSymbols = mutableSetOf() private val externalParams = mutableListOf() private var tempCounter = 0 @@ -168,6 +170,25 @@ public class ConversionContext @kotlin.jvm.JvmOverloads constructor( */ public fun getModuleDeclarations(): String = moduleDeclarationsBuilder.toString() + /** + * Emit a module-scope line that must follow the entry function — the private + * `func.func` decompositions a `stablehlo.composite` points at. They go *after* + * `@main` so text tooling that takes the first `func.func` as the entry + * ([sk.ainet.compile.hlo.validation.MlirParser]) keeps working. + */ + public fun emitModuleTrailer(line: String) { + moduleTrailerBuilder.appendLine(" $line") + } + + /** Every trailer line emitted so far; appended by [StableHloConverter] after the entry function. */ + public fun getModuleTrailer(): String = moduleTrailerBuilder.toString() + + /** + * Claim a module-level symbol (e.g. a decomposition function name). Returns `true` the + * first time, `false` afterwards, so a helper emitted once per (op, type) is not duplicated. + */ + public fun claimSymbol(symbol: String): Boolean = claimedSymbols.add(symbol) + /** * Register an externalized constant tensor. The converter records * these alongside MLIR emission so a downstream packager (see PR C @@ -247,6 +268,8 @@ public class ConversionContext @kotlin.jvm.JvmOverloads constructor( valueTypes.clear() stringBuilder.clear() moduleDeclarationsBuilder.clear() + moduleTrailerBuilder.clear() + claimedSymbols.clear() externalParams.clear() tempCounter = 0 } diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverter.kt index 00e2a86be..80e527a67 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverter.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverter.kt @@ -109,6 +109,10 @@ public class StableHloConverter @kotlin.jvm.JvmOverloads constructor( assembled.appendLine(" func.func $functionSignature {") assembled.append(context.getContent()) assembled.appendLine(" }") + val trailer = context.getModuleTrailer() + if (trailer.isNotEmpty()) { + assembled.append(trailer) + } assembled.appendLine("}") val content = assembled.toString() diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ActivationOperationsConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ActivationOperationsConverter.kt index 11a3a16cf..f2018d623 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ActivationOperationsConverter.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ActivationOperationsConverter.kt @@ -277,7 +277,9 @@ public class ActivationOperationsConverter : StableHloOperationConverter { val approximate = (node.operation.parameters["approximate"] as? Boolean) ?: true val emitter = ContextEmitter(context) - val resultValue = GeluLowering.emit(emitter, operands[0], outputType, approximate) + val resultValue = GeluLowering.spec.lower( + context, emitter, operands[0], outputType, mapOf("approximate" to approximate) + ) return ConversionResult.Success( outputValueName = resultValue, diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ErfLowering.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ErfLowering.kt new file mode 100644 index 000000000..ca2156f10 --- /dev/null +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ErfLowering.kt @@ -0,0 +1,19 @@ +package sk.ainet.compile.hlo.converters + +import sk.ainet.compile.hlo.lowering.LoweringSpec + +/** `erf` lowering: `chlo.erf` inline, or composite `skainet.erf` (version 1) when kept fused. */ +public object ErfLowering { + public val spec: LoweringSpec = LoweringSpec( + op = "erf", + compositeName = "skainet.erf", + attributes = emptyList(), + version = 1, + ) { e, x, type, _ -> emit(e, x, type) } + + public fun emit(e: HloEmitter, x: String, type: String): String { + val resultValue = e.newValue() + e.op("$resultValue = chlo.erf $x : $type -> $type") + return resultValue + } +} diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/GeluLowering.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/GeluLowering.kt index 5317b3284..1d9d21d93 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/GeluLowering.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/GeluLowering.kt @@ -44,6 +44,14 @@ public object GeluLowering { /** 1/√2 as the converter prints it. */ public const val SQRT1_2: String = "0.70710678118654752" + /** Composite `skainet.gelu` with the `approximate` flag as its attribute (version 1). */ + public val spec: sk.ainet.compile.hlo.lowering.LoweringSpec = sk.ainet.compile.hlo.lowering.LoweringSpec( + op = "gelu", + compositeName = "skainet.gelu", + attributes = listOf("approximate"), + version = 1, + ) { e, x, type, params -> emit(e, x, type, (params["approximate"] as? Boolean) ?: true) } + public fun emit(e: HloEmitter, x: String, type: String, approximate: Boolean): String = if (approximate) emitTanh(e, x, type) else emitExact(e, x, type) diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/UnaryMathConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/UnaryMathConverter.kt index c71e7e773..e79df7bd0 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/UnaryMathConverter.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/UnaryMathConverter.kt @@ -64,6 +64,14 @@ public class UnaryMathConverter : StableHloOperationConverter { val outputType = outputSpec?.let { context.getTypeMapper().mapTensorType(it) } ?: "tensor" + if (name == "erf") { + // erf has a lowering spec: chlo.erf inline, or a `skainet.erf` composite when the + // target's OpGranularityPolicy keeps it fused. + val emitter = ContextEmitter(context) + val resultValue = ErfLowering.spec.lower(context, emitter, operands[0], outputType, emptyMap()) + return ConversionResult.Success(outputValueName = resultValue, emittedOperations = emitter.lines) + } + val resultValue = context.nextTempValue() val op = if (hloOp.startsWith("chlo.")) { "$resultValue = $hloOp ${operands[0]} : $outputType -> $outputType" diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/lowering/LoweringSpec.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/lowering/LoweringSpec.kt new file mode 100644 index 000000000..185b2f9af --- /dev/null +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/lowering/LoweringSpec.kt @@ -0,0 +1,119 @@ +package sk.ainet.compile.hlo.lowering + +import sk.ainet.compile.hlo.ConversionContext +import sk.ainet.compile.hlo.converters.ContextEmitter +import sk.ainet.compile.hlo.converters.HloEmitter + +/** + * The formula of one fused-able op, written once through an [HloEmitter]. + * + * The same builder produces the inline (decomposed) form and the body of the + * `func.func private` a `stablehlo.composite` points at, so the two can never drift. + */ +public fun interface HloDecomposition { + /** Emit the op for [operand] of MLIR [type] and return the result value name. */ + public fun emit(e: HloEmitter, operand: String, type: String, params: Map): String +} + +/** + * Lowering metadata of one op (#1342): how it is named as a `stablehlo.composite`, which + * node parameters become `composite_attributes`, the composite `version`, and its + * [decomposition]. Whether a target sees the composite or the decomposition is decided + * per op by [sk.ainet.compile.target.OpGranularityPolicy] in [lower]. + */ +public data class LoweringSpec( + /** Graph op name (lower-case), the key [sk.ainet.compile.target.OpGranularityPolicy.keepFused] receives. */ + val op: String, + /** Namespaced composite name, e.g. `skainet.gelu` (the StableHLO spec requires a namespace). */ + val compositeName: String, + /** Node parameter keys copied into `composite_attributes`, in this order. */ + val attributes: List, + /** Composite version; bump when the decomposition's semantics change. */ + val version: Int, + val decomposition: HloDecomposition, +) { + init { + require('.' in compositeName) { "composite name must be namespaced: $compositeName" } + } + + /** + * Lower one unary node: a `stablehlo.composite` plus its private decomposition function + * when `context.granularity?.keepFused(op) == true`, else the decomposition inline. + * Returns the result value name; the emitted lines are in [emitter]. + */ + public fun lower( + context: ConversionContext, + emitter: ContextEmitter, + operand: String, + type: String, + params: Map, + ): String = + if (context.granularity?.keepFused(op) == true) { + context.emitComposite(this, emitter, listOf(operand), type, params) + } else { + decomposition.emit(emitter, operand, type, params) + } +} + +/** + * Emit `%r = stablehlo.composite "" %x { composite_attributes = {…}, decomposition = @fn, version = N : i32 } : (T) -> T` + * and, once per (op, attributes, type), the `func.func private @fn` holding [LoweringSpec.decomposition]. + * + * The decomposition body takes its SSA names from this context's counter, so the module-wide + * uniqueness rule holds across the entry function and every helper; `%arg0` is the helper's own + * argument. Helpers go to the module trailer, after the entry function. + */ +public fun ConversionContext.emitComposite( + spec: LoweringSpec, + emitter: ContextEmitter, + operands: List, + type: String, + params: Map, +): String { + require(operands.size == 1) { "composite emission currently supports unary ops, got ${operands.size} operands for ${spec.op}" } + val attrs = spec.attributes.mapNotNull { key -> params[key]?.let { key to it } } + val symbol = compositeSymbol(spec, attrs, type) + + if (claimSymbol(symbol)) { + val body = BufferEmitter(this) + val result = spec.decomposition.emit(body, "%arg0", type, params) + emitModuleTrailer("func.func private @$symbol(%arg0: $type) -> $type {") + body.lines.forEach { emitModuleTrailer(" $it") } + emitModuleTrailer(" return $result : $type") + emitModuleTrailer("}") + } + + val resultValue = nextTempValue() + val attrText = if (attrs.isEmpty()) "" else "composite_attributes = {" + + attrs.joinToString(", ") { (k, v) -> "$k = ${mlirAttribute(v)}" } + "}, " + emitter.op( + "$resultValue = stablehlo.composite \"${spec.compositeName}\" ${operands[0]} " + + "{${attrText}decomposition = @$symbol, version = ${spec.version} : i32} : ($type) -> $type" + ) + return resultValue +} + +/** `skainet_gelu_approximate_false_f32_2x3`: op, attribute values, then the sanitized type. */ +internal fun compositeSymbol(spec: LoweringSpec, attrs: List>, type: String): String { + val attrPart = attrs.joinToString("") { (k, v) -> "_${k}_${v.toString().lowercase()}" } + val typePart = type.removePrefix("tensor<").removeSuffix(">") + .replace('?', 'd').replace('*', 's').replace("x", "x") + val sanitized = (spec.op + attrPart + "_" + typePart).replace(Regex("[^A-Za-z0-9_]"), "_") + return "skainet_$sanitized" +} + +/** Render a node parameter as an MLIR attribute value. */ +internal fun mlirAttribute(value: Any): String = when (value) { + is Boolean -> value.toString() + is Int, is Long, is Short, is Byte -> "$value : i32" + is Float -> "${value.toDouble()} : f32" + is Double -> "$value : f64" + else -> "\"$value\"" +} + +/** [HloEmitter] that buffers lines but still draws SSA names from the shared context counter. */ +private class BufferEmitter(private val context: ConversionContext) : HloEmitter { + val lines = mutableListOf() + override fun newValue(): String = context.nextTempValue() + override fun op(line: String) { lines += line } +} diff --git a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/CompositeEmissionTest.kt b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/CompositeEmissionTest.kt new file mode 100644 index 000000000..6db3fc4b4 --- /dev/null +++ b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/CompositeEmissionTest.kt @@ -0,0 +1,121 @@ +package sk.ainet.compile.hlo + +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.test.assertTrue +import sk.ainet.compile.hlo.validation.MlirParser +import sk.ainet.compile.target.FusedOpAllowList +import sk.ainet.lang.graph.DefaultComputeGraph +import sk.ainet.lang.graph.GraphEdge +import sk.ainet.lang.graph.GraphNode +import sk.ainet.lang.tensor.ops.ErfOperation +import sk.ainet.lang.tensor.ops.GeluOperation +import sk.ainet.lang.tensor.ops.InputOperation +import sk.ainet.lang.tensor.ops.Operation +import sk.ainet.lang.tensor.ops.TensorSpec +import sk.ainet.lang.types.DType + +/** + * #1342 part 2: `stablehlo.composite` emission behind [sk.ainet.compile.target.OpGranularityPolicy] + * — the first consumer of the policy. A target that keeps `gelu` fused gets the composite plus a + * private decomposition function; every other target gets the decomposition inline, and both + * come from the same builder. + */ +class CompositeEmissionTest { + + private val shape = listOf(2, 3) + private val type = "tensor<2x3xf32>" + + private fun chain(vararg ops: Operation): DefaultComputeGraph { + val graph = DefaultComputeGraph() + var prev = GraphNode("input", InputOperation(), emptyList(), listOf(TensorSpec("input", shape, "FP32"))) + graph.addNode(prev) + ops.forEachIndexed { i, op -> + val node = GraphNode("n$i", op, listOf(TensorSpec("in$i", shape, "FP32")), listOf(TensorSpec("out$i", shape, "FP32"))) + graph.addNode(node) + graph.addEdge(GraphEdge("e$i", prev, node, 0, 0, prev.outputs[0])) + prev = node + } + return graph + } + + private fun exactGelu() = GeluOperation(mapOf("approximate" to false)) + private fun fused(vararg ops: String) = FusedOpAllowList("test-npu", ops.toSet()) + + @Test + fun fused_target_gets_composite_and_private_decomposition_after_main() { + val module = StableHloConverterFactory.createBasic(granularity = fused("gelu")) + .convert(chain(exactGelu()), "test_fused") + val c = module.content + + assertTrue(c.contains("stablehlo.composite \"skainet.gelu\" %arg0 {composite_attributes = {approximate = false}, decomposition = @skainet_gelu_approximate_false_2x3xf32, version = 1 : i32} : (tensor<2x3xf32>) -> tensor<2x3xf32>"), c) + assertTrue(c.contains("func.func private @skainet_gelu_approximate_false_2x3xf32(%arg0: $type) -> $type {"), c) + assertTrue(c.contains("chlo.erf"), "the private body must hold the exact decomposition") + assertTrue(c.indexOf("func.func @test_fused") < c.indexOf("func.func private"), "entry function must stay first") + assertTrue(c.trimEnd().endsWith("}"), c) + + // The text tooling still treats @test_fused as the entry function and the module validates. + assertEquals("test_fused", MlirParser().parse(c).getOrThrow().functionName) + assertTrue(MlirValidator().validate(c).isEmpty(), MlirValidator().validate(c).toString()) + } + + @Test + fun portable_target_gets_the_decomposition_inline() { + val module = StableHloConverterFactory.createBasic().convert(chain(exactGelu()), "test_inline") + assertFalse(module.content.contains("stablehlo.composite")) + assertFalse(module.content.contains("func.func private")) + assertTrue(module.content.contains("chlo.erf")) + } + + @Test + fun policy_that_decomposes_gelu_emits_inline_even_when_set() { + val module = StableHloConverterFactory.createBasic(granularity = fused("layernorm")) + .convert(chain(exactGelu()), "test_other") + assertFalse(module.content.contains("stablehlo.composite")) + } + + @Test + fun decomposition_function_is_emitted_once_per_shape_and_attributes() { + val module = StableHloConverterFactory.createBasic(granularity = fused("gelu")) + .convert(chain(exactGelu(), exactGelu(), GeluOperation(mapOf("approximate" to true))), "test_dedup") + val c = module.content + assertEquals(3, Regex("stablehlo\\.composite \"skainet\\.gelu\"").findAll(c).count()) + assertEquals(1, Regex("func\\.func private @skainet_gelu_approximate_false_2x3xf32\\(").findAll(c).count()) + assertEquals(1, Regex("func\\.func private @skainet_gelu_approximate_true_2x3xf32\\(").findAll(c).count()) + assertTrue(c.contains("composite_attributes = {approximate = true}")) + assertTrue(MlirValidator().validate(c).isEmpty(), "SSA names must stay unique across main and helpers") + } + + @Test + fun erf_has_its_own_composite() { + val module = StableHloConverterFactory.createBasic(granularity = fused("erf")) + .convert(chain(ErfOperation()), "test_erf") + val c = module.content + assertTrue(c.contains("stablehlo.composite \"skainet.erf\" %arg0 {decomposition = @skainet_erf_2x3xf32, version = 1 : i32} : ($type) -> $type"), c) + assertTrue(c.contains("func.func private @skainet_erf_2x3xf32(%arg0: $type) -> $type {"), c) + assertTrue(c.contains("chlo.erf %arg0 : $type -> $type"), c) + assertTrue(MlirValidator().validate(c).isEmpty()) + } + + @Test + fun composite_body_equals_the_inline_decomposition_modulo_ssa_names() { + fun normalize(lines: List) = lines.map { it.trim().replace(Regex("%v\\d+"), "%v").replace("%arg0", "%x") } + val inline = StableHloConverterFactory.createBasic().convert(chain(exactGelu()), "a").content + val fusedModule = StableHloConverterFactory.createBasic(granularity = fused("gelu")).convert(chain(exactGelu()), "b").content + + val inlineOps = inline.lines().filter { it.trim().startsWith("%v") } + val bodyStart = fusedModule.indexOf("func.func private") + val body = fusedModule.substring(bodyStart).lines().filter { it.trim().startsWith("%v") } + assertEquals(normalize(inlineOps), normalize(body), "one builder must produce both forms") + } + + @Test + fun text_optimizer_passes_tolerate_a_module_with_a_helper_function() { + val converter = StableHloConverterFactory.createBasic(granularity = fused("gelu")) + val optimized = converter.convertWithOptimization(chain(exactGelu()), StableHloOptimizer.createDefault()) + assertTrue(optimized.content.contains("stablehlo.composite \"skainet.gelu\"")) + assertTrue(optimized.content.contains("func.func private @skainet_gelu_approximate_false_2x3xf32")) + assertTrue(optimized.content.contains("chlo.erf"), "DCE must not drop the helper body") + } +} diff --git a/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api b/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api index eb14b967b..ee739822f 100644 --- a/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api +++ b/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api @@ -121,6 +121,25 @@ public final class sk/ainet/compile/opt/passes/DtypeForwardPropagationPass : sk/ public fun getName ()Ljava/lang/String; } +public final class sk/ainet/compile/opt/passes/GeluApproximationLegalizationPass : sk/ainet/compile/opt/GraphOptimizationPass { + public static final field APPROXIMATE Ljava/lang/String; + public static final field Companion Lsk/ainet/compile/opt/passes/GeluApproximationLegalizationPass$Companion; + public static final field LEGALIZED_FROM Ljava/lang/String; + public static final field MAX_ABS_ERROR Ljava/lang/String; + public fun ()V + public fun (Ljava/lang/String;)V + public synthetic fun (Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V + public fun apply (Lsk/ainet/lang/graph/ComputeGraph;)Lsk/ainet/compile/opt/GraphOptimizationResult; + public fun getName ()Ljava/lang/String; +} + +public final class sk/ainet/compile/opt/passes/GeluApproximationLegalizationPass$Companion { +} + +public final class sk/ainet/compile/opt/passes/GeluApproximationLegalizationPassKt { + public static final field GELU_TANH_MAX_ABS_ERROR D +} + public final class sk/ainet/compile/opt/passes/LLMFusionPass : sk/ainet/compile/opt/GraphOptimizationPass { public fun ()V public fun apply (Lsk/ainet/lang/graph/ComputeGraph;)Lsk/ainet/compile/opt/GraphOptimizationResult; diff --git a/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/GeluApproximationLegalizationPass.kt b/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/GeluApproximationLegalizationPass.kt new file mode 100644 index 000000000..4139bee1c --- /dev/null +++ b/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/GeluApproximationLegalizationPass.kt @@ -0,0 +1,84 @@ +package sk.ainet.compile.opt.passes + +import sk.ainet.compile.opt.GraphOptimizationPass +import sk.ainet.compile.opt.GraphOptimizationResult +import sk.ainet.lang.graph.ComputeGraph +import sk.ainet.lang.graph.GraphEdge +import sk.ainet.lang.graph.GraphNode + +/** + * Maximum absolute difference between the exact and the tanh-approximated GELU, measured on + * the CPU kernel over `[-6, 6]` (4.73e-4 at |x| ≈ 2.70), rounded up. + */ +public const val GELU_TANH_MAX_ABS_ERROR: Double = 4.8e-4 + +/** + * Target legalization for targets without any `erf` lowering (#1342): rewrites + * `gelu(approximate = false)` nodes to `gelu(approximate = true)` and **records** the + * substitution on the node (`legalizedFrom`, `maxAbsError`) and in the pass diagnostics, + * so a numerics change is a visible decision instead of a silent default. + * + * `erf` nodes are left alone on purpose: there is no cheap substitute, and a loud failure in + * the emitter beats a quiet approximation. + * + * A target opts in by registering the pass: + * ```kotlin + * TargetOptimizers.registerDagPasses("my-npu") { listOf(GeluApproximationLegalizationPass("my-npu has no erf kernel")) } + * ``` + */ +public class GeluApproximationLegalizationPass( + private val reason: String = "target has no erf lowering", +) : GraphOptimizationPass { + + override val name: String = "gelu-approximation-legalization" + + override fun apply(graph: ComputeGraph): GraphOptimizationResult { + val diagnostics = mutableListOf() + var changed = false + + val candidates = graph.nodes.filter { node -> + node.operation.name.equals("gelu", ignoreCase = true) && + node.operation.parameters[APPROXIMATE] == false + } + + for (node in candidates) { + val legalized = GraphNode( + id = node.id, + operation = node.operation.clone( + node.operation.parameters + mapOf( + APPROXIMATE to true, + LEGALIZED_FROM to "gelu(approximate=false)", + MAX_ABS_ERROR to GELU_TANH_MAX_ABS_ERROR, + ) + ), + inputs = node.inputs, + outputs = node.outputs, + metadata = node.metadata, + ) + + val incoming = graph.edges.filter { it.destination.id == node.id } + val outgoing = graph.edges.filter { it.source.id == node.id } + graph.removeNode(node) + graph.addNode(legalized) + for (edge in incoming) { + graph.addEdge(GraphEdge(edge.id, edge.source, legalized, edge.sourceOutputIndex, edge.destinationInputIndex, edge.tensorSpec)) + } + for (edge in outgoing) { + graph.addEdge(GraphEdge(edge.id, legalized, edge.destination, edge.sourceOutputIndex, edge.destinationInputIndex, edge.tensorSpec)) + } + + diagnostics += "GELU_APPROX_LEGALIZED node=${node.id}: gelu(approximate=false) -> gelu(approximate=true) [$reason] maxAbsError=$GELU_TANH_MAX_ABS_ERROR" + changed = true + } + + return GraphOptimizationResult(graph, changed = changed, diagnostics = diagnostics) + } + + public companion object { + public const val APPROXIMATE: String = "approximate" + /** Node parameter naming the original op the node was legalized from. */ + public const val LEGALIZED_FROM: String = "legalizedFrom" + /** Node parameter carrying the worst-case absolute error introduced by the rewrite. */ + public const val MAX_ABS_ERROR: String = "maxAbsError" + } +} diff --git a/skainet-compile/skainet-compile-opt/src/commonTest/kotlin/sk/ainet/compile/opt/GeluApproximationLegalizationPassTest.kt b/skainet-compile/skainet-compile-opt/src/commonTest/kotlin/sk/ainet/compile/opt/GeluApproximationLegalizationPassTest.kt new file mode 100644 index 000000000..fa12fce28 --- /dev/null +++ b/skainet-compile/skainet-compile-opt/src/commonTest/kotlin/sk/ainet/compile/opt/GeluApproximationLegalizationPassTest.kt @@ -0,0 +1,80 @@ +package sk.ainet.compile.opt + +import kotlin.test.AfterTest +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFalse +import kotlin.test.assertTrue +import sk.ainet.compile.opt.passes.GELU_TANH_MAX_ABS_ERROR +import sk.ainet.compile.opt.passes.GeluApproximationLegalizationPass +import sk.ainet.lang.graph.DefaultComputeGraph +import sk.ainet.lang.graph.GraphEdge +import sk.ainet.lang.graph.GraphNode +import sk.ainet.lang.tensor.ops.GeluOperation +import sk.ainet.lang.tensor.ops.GenericOperation +import sk.ainet.lang.tensor.ops.TensorSpec +import sk.ainet.lang.types.DType + +class GeluApproximationLegalizationPassTest { + + @AfterTest + fun reset() = TargetOptimizers.clear() + + private fun spec(name: String = "t") = TensorSpec(name = name, shape = listOf(4), dtype = "Float32") + + /** input -> gelu(approximate) -> consumer */ + private fun graphWith(approximate: Boolean): DefaultComputeGraph { + val g = DefaultComputeGraph() + val input = g.addNode(GraphNode("x", GenericOperation("input"), emptyList(), listOf(spec("x")))) + val gelu = g.addNode( + GraphNode("g", GeluOperation(mapOf("approximate" to approximate)), listOf(spec("x")), listOf(spec("g"))) + ) + val relu = g.addNode(GraphNode("r", GenericOperation("relu"), listOf(spec("g")), listOf(spec("r")))) + g.addEdge(GraphEdge("e0", input, gelu, tensorSpec = spec())) + g.addEdge(GraphEdge("e1", gelu, relu, tensorSpec = spec())) + return g + } + + @Test + fun rewrites_exact_gelu_to_tanh_and_records_the_substitution() { + val result = GeluApproximationLegalizationPass("no erf kernel").apply(graphWith(approximate = false)) + assertTrue(result.changed) + val node = result.graph.nodes.first { it.id == "g" } + assertTrue(node.operation is GeluOperation<*, *>) + assertEquals(true, node.operation.parameters["approximate"]) + assertEquals("gelu(approximate=false)", node.operation.parameters[GeluApproximationLegalizationPass.LEGALIZED_FROM]) + assertEquals(GELU_TANH_MAX_ABS_ERROR, node.operation.parameters[GeluApproximationLegalizationPass.MAX_ABS_ERROR]) + assertEquals(1, result.diagnostics.size) + assertTrue(result.diagnostics[0].startsWith("GELU_APPROX_LEGALIZED node=g:"), result.diagnostics[0]) + assertTrue(result.diagnostics[0].contains("[no erf kernel]")) + + // Edges survive the node swap. + assertEquals(listOf("x"), result.graph.edges.filter { it.destination.id == "g" }.map { it.source.id }) + assertEquals(listOf("r"), result.graph.edges.filter { it.source.id == "g" }.map { it.destination.id }) + } + + @Test + fun leaves_tanh_gelu_untouched_and_is_idempotent() { + val untouched = GeluApproximationLegalizationPass().apply(graphWith(approximate = true)) + assertFalse(untouched.changed) + assertTrue(untouched.diagnostics.isEmpty()) + + val first = GeluApproximationLegalizationPass().apply(graphWith(approximate = false)) + val second = GeluApproximationLegalizationPass().apply(first.graph) + assertFalse(second.changed, "a legalized node must not be rewritten again") + } + + @Test + fun a_target_opts_in_through_the_registry() { + TargetOptimizers.registerDagPasses("test-no-erf") { listOf(GeluApproximationLegalizationPass("test-no-erf has no erf")) } + + val result = dagPipelineFor("test-no-erf").optimize(graphWith(approximate = false)) + val node = result.graph.nodes.first { it.id == "g" } + assertEquals(true, node.operation.parameters["approximate"]) + assertTrue(result.passResults.flatMap { it.diagnostics }.any { it.startsWith("GELU_APPROX_LEGALIZED") }) + + // A target that did not register keeps the exact form. + val portable = dagPipelineFor("llvm-cpu").optimize(graphWith(approximate = false)) + assertEquals(false, portable.graph.nodes.first { it.id == "g" }.operation.parameters["approximate"]) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt index 75bcd0977..0e82b8750 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt @@ -8,6 +8,7 @@ import sk.ainet.lang.trace.Diff import sk.ainet.lang.nn.dsl.GenerateNetworkDsl import sk.ainet.lang.nn.dsl.ActivationDsl import sk.ainet.lang.ops.DarcValidated +import sk.ainet.lang.ops.Lowering @GenerateTracingWrapper @GenerateNetworkDsl @@ -230,6 +231,7 @@ public interface TensorOps { */ @Diff @ActivationDsl + @Lowering(composite = "skainet.gelu") public fun gelu(tensor: Tensor, approximate: Boolean = true): Tensor // Reduction operations @@ -369,6 +371,7 @@ public interface TensorOps { /** Element-wise Gauss error function `erf(x) = 2/√π · ∫₀ˣ e^(−t²) dt`, odd, saturating at ±1. */ @Diff @ActivationDsl + @Lowering(composite = "skainet.erf") public fun erf(tensor: Tensor): Tensor // Trigonometric operations diff --git a/skainet-lang/skainet-lang-ksp-annotations/src/commonMain/kotlin/sk/ainet/lang/ops/Lowering.kt b/skainet-lang/skainet-lang-ksp-annotations/src/commonMain/kotlin/sk/ainet/lang/ops/Lowering.kt new file mode 100644 index 000000000..d309cdade --- /dev/null +++ b/skainet-lang/skainet-lang-ksp-annotations/src/commonMain/kotlin/sk/ainet/lang/ops/Lowering.kt @@ -0,0 +1,21 @@ +package sk.ainet.lang.ops + +/** + * Declares how an op lowers on the compile path, for the generated operator documentation + * (the *Lowering* column of the coverage matrix) and for tooling that wants one place to read + * what an op can become (#1342). + * + * @param composite Namespaced `stablehlo.composite` name emitted when a target's + * `OpGranularityPolicy` keeps the op fused (e.g. `"skainet.gelu"`); empty when the op has no + * composite form. + * @param decomposes `true` when the op has a decomposition into StableHLO/CHLO primitives + * (the portable default emission). + * @param native `true` when the op maps 1:1 onto a target primitive (e.g. `stablehlo.tanh`). + */ +@Target(AnnotationTarget.FUNCTION) +@Retention(AnnotationRetention.SOURCE) +public annotation class Lowering( + val composite: String = "", + val decomposes: Boolean = true, + val native: Boolean = false, +) diff --git a/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ops/ksp/OperatorDocProcessor.kt b/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ops/ksp/OperatorDocProcessor.kt index 332067602..930c27a2b 100644 --- a/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ops/ksp/OperatorDocProcessor.kt +++ b/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ops/ksp/OperatorDocProcessor.kt @@ -43,6 +43,8 @@ data class FunctionDoc( val validatedOn: String = "", val validatedCommit: String = "", val referencesChecked: Boolean = true, + // Lowering forms declared with @Lowering (composite · decompose · native); empty = undeclared. + val lowering: String = "", ) data class ParameterDoc( @@ -207,6 +209,7 @@ class OperatorDocProcessor( validatedOn = validation.on, validatedCommit = validation.commit, referencesChecked = validation.referencesChecked, + lowering = extractLowering(fn), ) } @@ -295,9 +298,24 @@ class OperatorDocProcessor( validatedOn = validation.on, validatedCommit = validation.commit, referencesChecked = validation.referencesChecked, + lowering = extractLowering(function), ) } + /** + * Render the `@Lowering` annotation as the matrix cell text: the declared forms joined by + * ` · ` (`composite`, `decompose`, `native`), or empty when the op declares nothing. + */ + private fun extractLowering(function: KSFunctionDeclaration): String { + val annotation = function.annotations.find { it.shortName.asString() == "Lowering" } ?: return "" + fun arg(name: String): Any? = annotation.arguments.find { it.name?.asString() == name }?.value + val forms = mutableListOf() + if ((arg("composite") as? String).orEmpty().isNotEmpty()) forms += "composite" + if ((arg("decomposes") as? Boolean) ?: true) forms += "decompose" + if ((arg("native") as? Boolean) == true) forms += "native" + return forms.joinToString(" · ") + } + /** * Returns the adjoint rule name when the function carries `@Diff` (its `ruleName` argument, or * the empty string for bare `@Diff`), or `null` when the op is not differentiable. Mirrors the @@ -593,6 +611,9 @@ class OperatorDocProcessor( // Autodiff coverage. Always emitted so the manifest is the single source of // differentiability truth; diffRuleName only when a custom @Diff(ruleName=...) is set. append(",\n \"isDifferentiable\": ${function.isDifferentiable}") + if (function.lowering.isNotEmpty()) { + append(",\n \"lowering\": \"${escapeJson(function.lowering)}\"") + } if (function.diffRuleName.isNotEmpty()) { append(",\n \"diffRuleName\": \"${escapeJson(function.diffRuleName)}\"") } diff --git a/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ops/metadata/DocumentationModels.kt b/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ops/metadata/DocumentationModels.kt index e92079779..d61ff2d67 100644 --- a/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ops/metadata/DocumentationModels.kt +++ b/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ops/metadata/DocumentationModels.kt @@ -42,6 +42,8 @@ data class FunctionDoc( val validatedOn: String = "", val validatedCommit: String = "", val referencesChecked: Boolean = true, + /** Lowering forms from `@Lowering`: `composite`, `decompose`, `native` joined by ` · `; empty = undeclared. */ + val lowering: String = "", ) /** diff --git a/skainet-lang/skainet-lang-ksp-processor/src/test/kotlin/sk/ainet/lang/ops/ksp/OperatorDocProcessorTest.kt b/skainet-lang/skainet-lang-ksp-processor/src/test/kotlin/sk/ainet/lang/ops/ksp/OperatorDocProcessorTest.kt index 07ce9c556..81623efa1 100644 --- a/skainet-lang/skainet-lang-ksp-processor/src/test/kotlin/sk/ainet/lang/ops/ksp/OperatorDocProcessorTest.kt +++ b/skainet-lang/skainet-lang-ksp-processor/src/test/kotlin/sk/ainet/lang/ops/ksp/OperatorDocProcessorTest.kt @@ -67,6 +67,58 @@ class OperatorDocProcessorTest { // but the key requirement is that the test compiles and runs } + @Test + fun testLoweringFlowsIntoJson() { + val sourceCode = """ + package sk.ainet.lang.ops + + @Target(AnnotationTarget.CLASS, AnnotationTarget.FUNCTION) + @Retention(AnnotationRetention.SOURCE) + annotation class InProgress( + vararg val backends: String, + val owner: String = "", + val issue: String = "" + ) + + @Target(AnnotationTarget.FUNCTION) + @Retention(AnnotationRetention.SOURCE) + annotation class Lowering( + val composite: String = "", + val decomposes: Boolean = true, + val native: Boolean = false, + ) + + @InProgress("cpu", owner = "ops-team", issue = "GH-1") + @Lowering(composite = "skainet.gelu") + fun geluFn(): String = "ok" + + @InProgress("cpu", owner = "ops-team", issue = "GH-2") + @Lowering(decomposes = false, native = true) + fun tanhFn(): String = "ok" + + @InProgress("cpu", owner = "ops-team", issue = "GH-3") + fun plainFn(): String = "ok" + """.trimIndent() + + val source = SourceFile.kotlin("sk/ainet/lang/ops/TestLoweringOps.kt", sourceCode) + val compilation = KotlinCompilation().apply { + sources = listOf(source) + configureKsp {} + symbolProcessorProviders = mutableListOf(OperatorDocProcessorProvider()) + inheritClassPath = true + messageOutputStream = System.out + } + val result = compilation.compile() + assertTrue(result.messages.contains("Generated operators.json"), "Processor should generate operators.json") + + val operatorsJson = compilation.workingDir.walkTopDown().firstOrNull { it.isFile && it.name == "operators.json" } + assertTrue(operatorsJson != null && operatorsJson.exists(), "operators.json should be present") + val text = operatorsJson!!.readText() + assertTrue(text.contains("\"lowering\": \"composite · decompose\""), "composite + default decompose expected: $text") + assertTrue(text.contains("\"lowering\": \"native\""), "native-only expected: $text") + assertTrue(Regex("\"lowering\"").findAll(text).count() == 2, "plainFn must not carry a lowering field: $text") + } + @Test fun testDarcValidatedFlowsIntoJson() { // Two functions annotated with @InProgress so the processor picks