diff --git a/CHANGELOG.md b/CHANGELOG.md index 14712174..9d420f5b 100644 --- a/CHANGELOG.md +++ b/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). diff --git a/docs/modules/ROOT/pages/tutorials/hlo-getting-started.adoc b/docs/modules/ROOT/pages/tutorials/hlo-getting-started.adoc index 49065b4e..598667ce 100644 --- a/docs/modules/ROOT/pages/tutorials/hlo-getting-started.adoc +++ b/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. 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 8f228dfd..73303d53 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 @@ -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 { diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt index 08b789c8..21bee0c0 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt +++ b/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}") 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 index ca2156f1..06d8fe14 100644 --- 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 @@ -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() 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 1d9d21d9..ee3932f7 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 @@ -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) diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/SdpaLowering.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/SdpaLowering.kt index 21537589..a3d4f07e 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/SdpaLowering.kt +++ b/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` → `[]`. */ public fun parseTensorType(type: String): Pair, String> { require(type.startsWith("tensor<") && type.endsWith(">")) { "not a tensor type: $type" } 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 index 185b2f9a..7c5936fd 100644 --- 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 @@ -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 + /** Emit the op for [operands] of MLIR [operandTypes] producing [resultType]; return the result value name. */ + public fun emit(e: HloEmitter, operands: List, operandTypes: List, resultType: String, params: Map): 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): 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,49 +49,63 @@ 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, + operandTypes: List, + resultType: String, params: Map, ): 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 = lower(context, emitter, listOf(operand), listOf(type), 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]. + * Emit `%r = stablehlo.composite "" %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, - type: String, + operandTypes: List, + resultType: String, params: Map, ): 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("}") } @@ -87,17 +113,27 @@ public fun ConversionContext.emitComposite( 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>, type: String): String { +/** Unary form of [emitComposite]: operand and result share [type]. */ +public fun ConversionContext.emitComposite( + spec: LoweringSpec, + emitter: ContextEmitter, + operands: List, + type: String, + params: Map, +): 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>, operandTypes: List): 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" } 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 index 6db3fc4b..ddd48b82 100644 --- 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 @@ -110,6 +110,57 @@ class CompositeEmissionTest { assertEquals(normalize(inlineOps), normalize(body), "one builder must produce both forms") } + // ---- #1343 part 2: the n-ary seam, `skainet.sdpa` ---- + + private fun sdpaGraph(lengths: Boolean, mask: Boolean): DefaultComputeGraph { + val g = DefaultComputeGraph() + fun input(id: String, shape: List, dtype: String = "FP32") = + GraphNode(id, InputOperation(), emptyList(), listOf(TensorSpec(id, shape, dtype))).also { g.addNode(it) } + val q = input("q", listOf(2, 4, 8, 16)); val k = input("k", listOf(2, 2, 8, 16)); val v = input("v", listOf(2, 2, 8, 16)) + val m = if (mask) input("m", listOf(2, 1, 8, 8)) else null + val l = if (lengths) input("len", listOf(2), "Int32") else null + val operands = listOfNotNull(q, k, v, m, l) + val params = sk.ainet.lang.tensor.ops.ScaledDotProductAttentionOperation.parameters(0f, causal = true, window = 4, rightContext = 0, hasLengths = lengths) + val sdpa = GraphNode("sdpa", sk.ainet.lang.tensor.ops.ScaledDotProductAttentionOperation(params), operands.map { it.outputs[0] }, listOf(TensorSpec("o", listOf(2, 4, 8, 16), "FP32"))) + g.addNode(sdpa) + operands.forEachIndexed { i, n -> g.addEdge(GraphEdge("e$i", n, sdpa, 0, i, n.outputs[0])) } + return g + } + + @Test + fun sdpa_composite_carries_every_operand_and_the_mask_rules_as_attributes() { + val module = StableHloConverterFactory.createBasic(granularity = fused("scaleddotproductattention")).convert(sdpaGraph(lengths = true, mask = false), "sdpa_fused") + val c = module.content + val line = c.lines().single { "stablehlo.composite \"skainet.sdpa\"" in it }.trim() + assertTrue(line.contains("stablehlo.composite \"skainet.sdpa\" %arg0, %arg1, %arg2, %arg3 {"), line) + assertTrue(line.contains("composite_attributes = {scale = 0.25 : f32, causal = true, window = 4 : i32, rightContext = 0 : i32, maskKind = \"combined\"}"), line) + assertTrue(line.endsWith(": (tensor<2x4x8x16xf32>, tensor<2x2x8x16xf32>, tensor<2x2x8x16xf32>, tensor<2xi32>) -> tensor<2x4x8x16xf32>"), line) + val decl = Regex("decomposition = @(\\S+), version = 1 : i32").find(line)!!.groupValues[1] + assertTrue(c.contains("func.func private @$decl(%arg0: tensor<2x4x8x16xf32>, %arg1: tensor<2x2x8x16xf32>, %arg2: tensor<2x2x8x16xf32>, %arg3: tensor<2xi32>) -> tensor<2x4x8x16xf32> {"), c) + assertTrue(c.indexOf("func.func @sdpa_fused") < c.indexOf("func.func private"), "entry function must stay first") + assertTrue(c.substring(c.indexOf("func.func private")).contains("stablehlo.broadcast_in_dim %arg3, dims = [0]"), "the body lowers the lengths operand") + assertEquals("sdpa_fused", MlirParser().parse(c).getOrThrow().functionName) + assertTrue(MlirValidator().validate(c).isEmpty(), MlirValidator().validate(c).toString()) + + // Five operands when a dense mask is present as well. + val withMask = StableHloConverterFactory.createBasic(granularity = fused("scaleddotproductattention")).convert(sdpaGraph(lengths = true, mask = true), "sdpa_mask").content + assertTrue(withMask.contains("stablehlo.composite \"skainet.sdpa\" %arg0, %arg1, %arg2, %arg3, %arg4 {"), withMask) + assertTrue(withMask.contains("%arg3: tensor<2x1x8x8xf32>, %arg4: tensor<2xi32>)"), withMask) + } + + @Test + fun sdpa_composite_body_equals_the_inline_decomposition_modulo_ssa_names() { + fun normalize(lines: List) = lines.map { it.trim().replace(Regex("%v\\d+"), "%v") } + val inline = StableHloConverterFactory.createBasic().convert(sdpaGraph(lengths = true, mask = true), "a").content + val fusedModule = StableHloConverterFactory.createBasic(granularity = fused("scaleddotproductattention")).convert(sdpaGraph(lengths = true, mask = true), "b").content + assertFalse(inline.contains("stablehlo.composite")) + val inlineOps = inline.lines().filter { it.trim().startsWith("%v") } + val body = fusedModule.substring(fusedModule.indexOf("func.func private")).lines().filter { it.trim().startsWith("%v") } + assertEquals(normalize(inlineOps), normalize(body), "one builder must produce both forms") + // A policy that fuses only gelu leaves SDPA inline. + assertFalse(StableHloConverterFactory.createBasic(granularity = fused("gelu")).convert(sdpaGraph(lengths = false, mask = false), "c").content.contains("stablehlo.composite")) + } + @Test fun text_optimizer_passes_tolerate_a_module_with_a_helper_function() { val converter = StableHloConverterFactory.createBasic(granularity = fused("gelu"))