feat(compile): n-ary composite seam and the skainet.sdpa composite (#1343, part 2)

CHANGELOG.md

@@ -8,8 +8,11 @@
 - **Mask rules as op metadata**: `RecordingExecution` and the KSP tracer record `parameters["maskKind" | "window" | "rightContext"]` (`maskKind` = `none | causal | slidingWindow | keyPadding | combined`, `AttentionMask.kindOf`) with `keyLengths` as the last operand after the optional dense mask; `ScaledDotProductAttentionOperation.operandLayout` resolves the layout for `ComputeGraphExecutor` and the converter. DIRECT vs OPTIMIZED parity on a padded batch within 1e-6.
 - **StableHLO lowering of the rules** (`SdpaLowering`, the single emitter-based builder behind `AttentionOperationsConverter`): `causal`, the window band and `keyLengths` become one `i1` keep mask from `stablehlo.iota`/`compare`/`and` and a `select 0 / -1e30` added to the scores, so no `[batch, 1, Sq, Sk]` constant or graph input enters the module; `keyLengths` is a plain `[batch] i32` input. The builder reads every shape from the operand MLIR types, which is what the `skainet.sdpa` composite decomposition function (part 2) needs. `TensorOps.scaledDotProductAttention` carries `@Lowering(composite = "skainet.sdpa")`.
 
+- **`skainet.sdpa` composite** (#1343, part 2): the composite seam of 0.59 is n-ary — `HloDecomposition.emit(e, operands, operandTypes, resultType, params)`, `LoweringSpec.lower(context, emitter, operands, operandTypes, resultType, params)`, `ConversionContext.emitComposite` with one `%argN` per operand — and `SdpaLowering.spec` emits `stablehlo.composite "skainet.sdpa" %q, %k, %v[, %mask][, %keyLengths] {composite_attributes = {scale, causal, window, rightContext, maskKind}, decomposition = @…, version = 1 : i32}` when a target's `OpGranularityPolicy` keeps `scaleddotproductattention` fused. The private decomposition function is built by the same `SdpaLowering.emit` as the inline form (tested equal modulo SSA names).
+
 ### Changed
 
+- `HloDecomposition` is n-ary: implementations written against the unary 0.59 signature move to `HloDecomposition.unary { e, x, type, params -> … }`; `LoweringSpec.lower` keeps a unary overload.
 - `TensorOps.scaledDotProductAttention` gained `keyLengths`, `window` and `rightContext`: implementors of the interface must update the override (callers are source-compatible through the defaults). The generated `GraphDsl` (`skainet-lang-dag`) now types optional tensor operands (`mask`, `bias`, `keyLengths`) as `GraphValue<*>? = null` instead of mandatory positionals.
 - StableHLO SDPA export: the causal diagonal now carries the `Sk − Sq` offset the CPU kernel uses (a KV-cache decode step attends to its whole cache; `Sq == Sk` emission is unchanged); a dense mask and `causal = true` now apply together, as on the CPU, instead of the mask silently replacing the rule; positional/padding rules on a dynamic key length return a conversion failure instead of invalid `iota` IR (a dense mask still exports on the dynamic path).
 

docs/modules/ROOT/pages/tutorials/hlo-getting-started.adoc

@@ -347,7 +347,17 @@ func.func private @skainet_gelu_approximate_false_128x768xf32(%arg0: tensor<128x
 }
 ----
 
-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`).
+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`), `ErfLowering.spec` (`skainet.erf`) and `SdpaLowering.spec` (`skainet.sdpa`, version 1). The attention composite takes every operand of the op — `q, k, v`, the optional dense mask and the optional `keyLengths` — and records the mask rules as attributes, so a vendor compiler whose fused attention kernel takes lengths can match it directly:
+
+[source,mlir]
+----
+%v5 = stablehlo.composite "skainet.sdpa" %arg0, %arg1, %arg2, %arg3
+    {composite_attributes = {scale = 0.25 : f32, causal = true, window = 4 : i32, rightContext = 0 : i32, maskKind = "combined"},
+     decomposition = @skainet_scaleddotproductattention_…, version = 1 : i32}
+    : (tensor<2x4x8x16xf32>, tensor<2x2x8x16xf32>, tensor<2x2x8x16xf32>, tensor<2xi32>) -> tensor<2x4x8x16xf32>
+----
+
+The private decomposition function builds the same `iota`/`compare`/`select` keep mask from `%arg3` that the portable export emits inline, so both rows of the policy table have identical semantics by construction. A decomposition builder sees only operand names, their MLIR types, the result type and the node parameters (`HloDecomposition.emit`), which is why every shape it needs must be readable from the types.
 
 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.
 

skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api

@@ -472,8 +472,10 @@ public final class sk/ainet/compile/hlo/converters/ScalarOperationsConverter : s
 
 public final class sk/ainet/compile/hlo/converters/SdpaLowering {
 	public static final field INSTANCE Lsk/ainet/compile/hlo/converters/SdpaLowering;
+	public static final field OP Ljava/lang/String;
 	public final fun effectiveScale (FI)F
 	public final fun emit (Lsk/ainet/compile/hlo/converters/HloEmitter;Ljava/util/List;Ljava/util/List;Ljava/lang/String;Ljava/util/Map;)Ljava/lang/String;
+	public final fun getSpec ()Lsk/ainet/compile/hlo/lowering/LoweringSpec;
 	public final fun parseTensorType (Ljava/lang/String;)Lkotlin/Pair;
 	public final fun renderType (Ljava/util/List;Ljava/lang/String;)Ljava/lang/String;
 }
@@ -587,7 +589,12 @@ public final class sk/ainet/compile/hlo/generate/JvmHloGenerator {
 }
 
 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 static final field Companion Lsk/ainet/compile/hlo/lowering/HloDecomposition$Companion;
+	public abstract fun emit (Lsk/ainet/compile/hlo/converters/HloEmitter;Ljava/util/List;Ljava/util/List;Ljava/lang/String;Ljava/util/Map;)Ljava/lang/String;
+}
+
+public final class sk/ainet/compile/hlo/lowering/HloDecomposition$Companion {
+	public final fun unary (Lkotlin/jvm/functions/Function4;)Lsk/ainet/compile/hlo/lowering/HloDecomposition;
 }
 
 public final class sk/ainet/compile/hlo/lowering/LoweringSpec {
@@ -607,11 +614,13 @@ public final class sk/ainet/compile/hlo/lowering/LoweringSpec {
 	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 final fun lower (Lsk/ainet/compile/hlo/ConversionContext;Lsk/ainet/compile/hlo/converters/ContextEmitter;Ljava/util/List;Ljava/util/List;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 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/util/List;Ljava/lang/String;Ljava/util/Map;)Ljava/lang/String;
 }
 
 public final class sk/ainet/compile/hlo/validation/BenchmarkResults {

skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt

@@ -22,6 +22,10 @@ import sk.ainet.lang.graph.GraphNode
  * SDPA is a core `TensorOps` op (KSP-generated), so its converter lives here in
  * core alongside dot_general/softmax — the transformer modules just decompose to it.
  *
+ * Emission goes through `SdpaLowering.spec`: a `stablehlo.composite "skainet.sdpa"` with its private
+ * decomposition function when the target's `OpGranularityPolicy` keeps `scaleddotproductattention`
+ * fused, the decomposition inline otherwise (#1343 part 2).
+ *
  * Masking: the node's positional rules (`causal`, `window`/`rightContext`) and the `keyLengths`
  * operand become an `iota`/`compare`/`select` keep mask, and a dense `mask` operand is added on
  * top — the same rules the CPU kernel applies, defined once in `AttentionMask`. The rules need a
@@ -126,8 +130,9 @@ public class AttentionOperationsConverter : StableHloOperationConverter {
             "causal" to causal, "window" to window, "rightContext" to rightContext,
         )
         return try {
+            // Composite `skainet.sdpa` when the target keeps the op fused, else the decomposition inline.
             val emitter = ContextEmitter(context)
-            val out = SdpaLowering.emit(emitter, operands, operandTypes, outputType, lowered)
+            val out = SdpaLowering.spec.lower(context, emitter, operands, operandTypes, outputType, lowered)
             ConversionResult.Success(outputValueName = out, emittedOperations = emitter.lines)
         } catch (e: IllegalArgumentException) {
             ConversionResult.Failure(e.message ?: "SDPA lowering failed", "SDPA lowering failed for ${node.id}")

skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ErfLowering.kt

@@ -9,7 +9,8 @@ public object ErfLowering {
         compositeName = "skainet.erf",
         attributes = emptyList(),
         version = 1,
-    ) { e, x, type, _ -> emit(e, x, type) }
+        decomposition = sk.ainet.compile.hlo.lowering.HloDecomposition.unary { e, x, type, _ -> emit(e, x, type) },
+    )
 
     public fun emit(e: HloEmitter, x: String, type: String): String {
         val resultValue = e.newValue()

skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/GeluLowering.kt

@@ -50,7 +50,10 @@ public object GeluLowering {
         compositeName = "skainet.gelu",
         attributes = listOf("approximate"),
         version = 1,
-    ) { e, x, type, params -> emit(e, x, type, (params["approximate"] as? Boolean) ?: true) }
+        decomposition = sk.ainet.compile.hlo.lowering.HloDecomposition.unary { 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)

skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/SdpaLowering.kt

@@ -30,6 +30,25 @@ import sk.ainet.lang.tensor.ops.ScaledDotProductAttentionOperation
  */
 public object SdpaLowering {
 
+    /** Graph op name as [sk.ainet.compile.target.OpGranularityPolicy.keepFused] receives it. */
+    public const val OP: String = "scaleddotproductattention"
+
+    /**
+     * Composite `skainet.sdpa` (version 1): operands `q, k, v[, mask][, keyLengths]`, attributes
+     * `scale` (resolved, `f32`), `causal`, `window`, `rightContext` and `maskKind`, so a vendor
+     * compiler with a fused attention kernel that takes lengths can match it; everyone else
+     * inlines the decomposition or runs `stablehlo-legalize-composite-to-call`.
+     */
+    public val spec: sk.ainet.compile.hlo.lowering.LoweringSpec = sk.ainet.compile.hlo.lowering.LoweringSpec(
+        op = OP,
+        compositeName = "skainet.sdpa",
+        attributes = listOf("scale", "causal", "window", "rightContext", "maskKind"),
+        version = 1,
+        decomposition = sk.ainet.compile.hlo.lowering.HloDecomposition { e, operands, operandTypes, resultType, params ->
+            emit(e, operands, operandTypes, resultType, params)
+        },
+    )
+
     /** `tensor<1x2x?x8xf32>` → dims `[1, 2, -1, 8]` and element type `f32`; `tensor<f32>` → `[]`. */
     public fun parseTensorType(type: String): Pair<List<Int>, String> {
         require(type.startsWith("tensor<") && type.endsWith(">")) { "not a tensor type: $type" }

skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/lowering/LoweringSpec.kt

@@ -9,14 +9,26 @@ import sk.ainet.compile.hlo.converters.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.
+ * The builder sees only operand value names, their MLIR types, the result type and the node
+ * parameters — exactly what a decomposition function has — so every shape it needs must be
+ * readable from the types.
  */
 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, Any>): String
+    /** Emit the op for [operands] of MLIR [operandTypes] producing [resultType]; return the result value name. */
+    public fun emit(e: HloEmitter, operands: List<String>, operandTypes: List<String>, resultType: String, params: Map<String, Any>): String
+
+    public companion object {
+        /** A decomposition of a unary op whose operand and result share one type. */
+        public fun unary(block: (e: HloEmitter, operand: String, type: String, params: Map<String, Any>) -> String): HloDecomposition =
+            HloDecomposition { e, operands, _, resultType, params ->
+                require(operands.size == 1) { "unary decomposition got ${operands.size} operands" }
+                block(e, operands[0], resultType, params)
+            }
+    }
 }
 
 /**
- * Lowering metadata of one op (#1342): how it is named as a `stablehlo.composite`, which
+ * Lowering metadata of one op (#1342, #1343): 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].
@@ -37,67 +49,91 @@ public data class LoweringSpec(
     }
 
     /**
-     * Lower one unary node: a `stablehlo.composite` plus its private decomposition function
-     * when `context.granularity?.keepFused(op) == true`, else the decomposition inline.
+     * Lower one 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,
+        operands: List<String>,
+        operandTypes: List<String>,
+        resultType: String,
         params: Map<String, Any>,
     ): String =
         if (context.granularity?.keepFused(op) == true) {
-            context.emitComposite(this, emitter, listOf(operand), type, params)
+            context.emitComposite(this, emitter, operands, operandTypes, resultType, params)
         } else {
-            decomposition.emit(emitter, operand, type, params)
+            decomposition.emit(emitter, operands, operandTypes, resultType, params)
         }
+
+    /** [lower] for a unary op whose operand and result share [type]. */
+    public fun lower(
+        context: ConversionContext,
+        emitter: ContextEmitter,
+        operand: String,
+        type: String,
+        params: Map<String, Any>,
+    ): String = lower(context, emitter, listOf(operand), listOf(type), type, params)
 }
 
 /**
- * Emit `%r = stablehlo.composite "<name>" %x { composite_attributes = {…}, decomposition = @fn, version = N : i32 } : (T) -> T`
- * and, once per (op, attributes, type), the `func.func private @fn` holding [LoweringSpec.decomposition].
+ * Emit `%r = stablehlo.composite "<name>" %a, %b, … { composite_attributes = {…}, decomposition = @fn, version = N : i32 } : (Ta, Tb, …) -> R`
+ * and, once per (op, attributes, operand types), the `func.func private @fn(%arg0: Ta, %arg1: Tb, …) -> R`
+ * 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.
+ * uniqueness rule holds across the entry function and every helper; `%argN` are the helper's own
+ * arguments. Helpers go to the module trailer, after the entry function.
  */
 public fun ConversionContext.emitComposite(
     spec: LoweringSpec,
     emitter: ContextEmitter,
     operands: List<String>,
-    type: String,
+    operandTypes: List<String>,
+    resultType: String,
     params: Map<String, Any>,
 ): String {
-    require(operands.size == 1) { "composite emission currently supports unary ops, got ${operands.size} operands for ${spec.op}" }
+    require(operands.isNotEmpty() && operands.size == operandTypes.size) { "composite emission needs one type per operand for ${spec.op}" }
     val attrs = spec.attributes.mapNotNull { key -> params[key]?.let { key to it } }
-    val symbol = compositeSymbol(spec, attrs, type)
+    val symbol = compositeSymbol(spec, attrs, operandTypes)
 
     if (claimSymbol(symbol)) {
         val body = BufferEmitter(this)
-        val result = spec.decomposition.emit(body, "%arg0", type, params)
-        emitModuleTrailer("func.func private @$symbol(%arg0: $type) -> $type {")
+        val args = operandTypes.indices.map { "%arg$it" }
+        val result = spec.decomposition.emit(body, args, operandTypes, resultType, params)
+        val signature = operandTypes.indices.joinToString(", ") { "%arg$it: ${operandTypes[it]}" }
+        emitModuleTrailer("func.func private @$symbol($signature) -> $resultType {")
         body.lines.forEach { emitModuleTrailer("  $it") }
-        emitModuleTrailer("  return $result : $type")
+        emitModuleTrailer("  return $result : $resultType")
         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"
+        "$resultValue = stablehlo.composite \"${spec.compositeName}\" ${operands.joinToString(", ")} " +
+            "{${attrText}decomposition = @$symbol, version = ${spec.version} : i32} : (${operandTypes.joinToString(", ")}) -> $resultType"
     )
     return resultValue
 }
 
-/** `skainet_gelu_approximate_false_f32_2x3`: op, attribute values, then the sanitized type. */
-internal fun compositeSymbol(spec: LoweringSpec, attrs: List<Pair<String, Any>>, type: String): String {
+/** Unary form of [emitComposite]: operand and result share [type]. */
+public fun ConversionContext.emitComposite(
+    spec: LoweringSpec,
+    emitter: ContextEmitter,
+    operands: List<String>,
+    type: String,
+    params: Map<String, Any>,
+): String = emitComposite(spec, emitter, operands, operands.map { type }, type, params)
+
+/** `skainet_gelu_approximate_false_2x3xf32`: op, attribute values, then the sanitized operand types. */
+internal fun compositeSymbol(spec: LoweringSpec, attrs: List<Pair<String, Any>>, operandTypes: List<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 typePart = operandTypes.joinToString("_") { type ->
+        type.removePrefix("tensor<").removeSuffix(">").replace('?', 'd').replace('*', 's')
+    }
     val sanitized = (spec.op + attrPart + "_" + typePart).replace(Regex("[^A-Za-z0-9_]"), "_")
     return "skainet_$sanitized"
 }