Skip to content

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

Merged
michalharakal merged 1 commit into
SKaiNET-developers:developfrom
michalharakal:feature/sdpa-composite
Oct 15, 2026
Merged

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

Conversation

@michalharakal

Copy link
Copy Markdown
Contributor

Stacked on #1348 (part 1). The diff includes its commits until #1348 merges; after that this branch is rebased onto develop and the draft flag dropped. Only the last commit (feat(compile): n-ary composite seam and the skainet.sdpa composite) is this PR.

Summary

Part 2 of #1343, stacked on part 1: the composite seam becomes n-ary and scaledDotProductAttention gets its stablehlo.composite.

  • n-ary seam (sk.ainet.compile.hlo.lowering): HloDecomposition.emit(e, operands, operandTypes, resultType, params), LoweringSpec.lower(context, emitter, operands, operandTypes, resultType, params), ConversionContext.emitComposite with one %argN: T per operand in the private decomposition function and every operand type in the composite line; the symbol covers all operand types. HloDecomposition.unary { e, x, type, params -> … } and a unary lower overload keep GeluLowering/ErfLowering as short as before (byte-identical output).

  • SdpaLowering.spec = LoweringSpec("scaleddotproductattention", "skainet.sdpa", attributes = scale · causal · window · rightContext · maskKind, version 1). When a target's OpGranularityPolicy keeps the op fused the converter emits

    %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>

    with the optional dense mask as a fifth operand; the decomposition function is built by the same SdpaLowering.emit as the inline form, from the operand types alone. scale is the resolved value (1/√headDim when the node left it at 0) so a vendor kernel can match it directly. Attribute keys are the node parameter names (camelCase), as skainet.gelu does, rather than the snake_case sketch in the issue.

Test plan

  • CompositeEmissionTest: sdpa composite with four and five operands, attributes, private func.func with four arguments after the entry function, MlirParser entry = @main, MlirValidator clean, composite body ≡ inline decomposition modulo SSA names, gelu-only policy leaves SDPA inline; gelu/erf cases unchanged
  • SdpaMaskRulesHloTest, SdpaGqaHloExportTest, DynamicShapeHloExportTest and the rest of compile-hlo green
  • apiDump (skainet-compile-hlo.api: HloDecomposition.emit signature, Companion.unary, n-ary lower/emitComposite, SdpaLowering.OP/spec)
  • scripts/pr-gate.sh --quick, full gate
  • stablehlo-opt round trip (--stablehlo-legalize-composite-to-call) not run: the tool is not on the dev/CI machines; the body-equals-inline test is the stand-in.

…KaiNET-developers#1343, part 2)

- `HloDecomposition.emit(e, operands, operandTypes, resultType, params)` and
  `LoweringSpec.lower(context, emitter, operands, operandTypes, resultType,
  params)`: the composite seam of 0.59 takes any number of operands;
  `emitComposite` declares one `%argN` per operand in the private
  decomposition function and lists every operand type in the composite line.
  `HloDecomposition.unary { … }` and a unary `lower` overload keep gelu/erf
  short.
- `SdpaLowering.spec`: `stablehlo.composite "skainet.sdpa" %q, %k, %v[, %mask]
  [, %keyLengths] {composite_attributes = {scale, causal, window, rightContext,
  maskKind}, decomposition = @…, version = 1 : i32}` when the target's
  `OpGranularityPolicy` keeps `scaleddotproductattention` fused; the
  decomposition function is built by the same `SdpaLowering.emit` as the
  inline form. `AttentionOperationsConverter` lowers through the spec.
- Tests: composite with four and five operands, private function after the
  entry function, parser/validator clean, body equals the inline form modulo
  SSA names, gelu-only policy leaves SDPA inline. Tutorial section, changelog,
  api dump.
@michalharakal
michalharakal force-pushed the feature/sdpa-composite branch from b80d638 to 03b9096 Compare October 15, 2026 17:36
@michalharakal
michalharakal marked this pull request as ready for review October 15, 2026 17:36
@michalharakal
michalharakal merged commit 27c5f8d into SKaiNET-developers:develop Oct 15, 2026
17 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant