Make Whisper encoder compile end-to-end via SKaiNET → StableHLO → IREE

Prior state: `iree-compile` rejected the Whisper-encoder MLIR emitted by
`StableHloConverterFactory.createExtended()` — 141 `// Unsupported` comments
over 296 graph nodes and 12 `stablehlo.custom_call @reduce_*` ops. Cascade
analysis in `docs/whisper-iree-issues/` identified a mix of missing converters
and invalid-MLIR emissions; a single upstream miss fanned out via
`StableHloConverter.processNode`'s `mapNotNull` operand lookup and appeared
downstream as bogus "wrong arity" failures.

Result: 0 Unsupported, 0 custom_call reductions, 35 real `stablehlo.reduce`,
32 `stablehlo.dot_general` (including the 8 rank-4 attention matmuls), and a
40 MB `.vmfb` compiled with `iree-compile --iree-hal-target-backends=llvm-cpu`.

Individual fixes, each guarded by an end-to-end assertion in the new
`Conv1dTapeToHloTest`:

* `TraceToGraphBuilder.buildOutputSpecs`: fall back to
  `trace.outputs[i].shape.dimensions` when the KSP-generated tracer omits
  `outputShapes` from attributes. Fixes `tensor<?xf32>` on every conv/gelu/etc.
  whose wrapper didn't go through `OpAttributeFactory.shapesAndDTypes`.

* `NeuralNetOperationsConverter.buildConv1d/ConvolutionOperation`: emit
  `stablehlo.broadcast_in_dim %bias, dims = [1]` before the bias add so the
  add's operands share a type; previously produced
  `stablehlo.add %conv, %bias : tensor<N,C,L>` with a rank-mismatched bias.

* `BasicMathConverter`: rewrite `add`/`subtract`/`multiply`/`divide` lowering.
  Old code short-circuited broadcasting whenever dtypes matched, and the
  broadcast path itself emitted split-line syntax MLIR rejects. New path
  takes `node.outputs[0]` as the target spec, adapts each operand via
  `stablehlo.convert` then `stablehlo.broadcast_in_dim` with NumPy-style
  right-aligned dim mapping.

* New `UnaryMathConverter`: `sqrt`, `rsqrt`, `exp`, `expm1`, `log`, `log1p`,
  `abs`, `sign`, `negate`, `ceil`, `floor`, `round`, `cos`, `sin` → 1:1
  StableHLO primitives.

* New `ScalarOperationsConverter`: `addScalar`/`subScalar`/`mulScalar`/
  `divScalar`/`rsubScalar`/`rdivScalar` materialize the scalar as a splat
  `stablehlo.constant` then apply the matching binary op (reversed for
  `r*Scalar`). Reads the scalar from `operation.parameters["b"]` (the KSP
  tracer stores it there via `OpAttributeFactory.scalarOp`).

* `ReductionOperationsConverter`: emit real `stablehlo.reduce` instead of
  `stablehlo.custom_call @reduce_*` (the previous form was both non-standard
  and syntactically malformed). `mean` = reduce-sum + divide-by-count;
  `variance` decomposes to `E[X²] − E[X]²` via two reduces.

* `ActivationOperationsConverter.convertSoftmax`: softmax's two reductions
  now use `stablehlo.reduce` applying `stablehlo.maximum` and `stablehlo.add`.
  Same motivation as the reductions above.

* `LinalgOperationsConverter`: route every matmul variant
  (`matmul`/`dot`/`mm`/`bmm`/`batch_matmul`) through the rank-generic batched
  lowering. Prior `matmul` path hardcoded `contracting_dims = [1] x [0]` and
  produced invalid MLIR for Whisper attention's rank-4 matmul
  `[1,6,1500,64] × [1,6,64,1500]`; now emits
  `batching_dims = [0, 1] x [0, 1], contracting_dims = [3] x [2]`. Rank-2
  continues to emit the compact form (empty batching clause).

* `StableHloConverter.processNodes`: enrich the fallback "Unsupported op"
  comment with the quoted name and the registry's full key list. Makes the
  MLIR self-diagnostic so a future missing-converter case surfaces without
  needing a local reproducer.

Registry entries added for the two new converters in all three factory
entry points (`createBasic`, `createExtended`, `createFast`).

`docs/whisper-iree-issues/` is included for context — it captured the
failure surface that these changes close.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: michalharakal

docs/whisper-iree-issues/ISSUE-A-conv1d-infer-outputs.md

@@ -0,0 +1,72 @@
+# Conv1d/2d/3dOperation.inferOutputs() echoes input shape instead of computing output shape
+
+## Problem
+
+`Conv1dOperation.inferOutputs()` in `TensorOperations.kt` (line ~439) returns the
+input tensor's shape as the output shape, ignoring weight shape, stride, padding,
+and dilation:
+
+```kotlin
+override fun inferOutputs(inputs: List<TensorSpec>): List<TensorSpec> {
+    require(inputs.size >= 2) { "Conv1d operation requires at least 2 inputs" }
+    val outputShape = inputs[0].shape   // <-- BUG: just copies input shape
+    return listOf(TensorSpec("conv1d_output", outputShape, inputs[0].dtype, ...))
+}
+```
+
+Conv2dOperation (line ~471) and Conv3dOperation (line ~503) have the identical bug.
+
+## Expected
+
+```kotlin
+override fun inferOutputs(inputs: List<TensorSpec>): List<TensorSpec> {
+    val inShape = inputs[0].shape        // [N, Cin, L]
+    val wShape  = inputs[1].shape        // [Cout, Cin/g, K]
+    val stride   = (parameters["stride"]   as? Int) ?: 1
+    val padding  = (parameters["padding"]  as? Int) ?: 0
+    val dilation = (parameters["dilation"] as? Int) ?: 1
+    val outShape = if (inShape != null && wShape != null && inShape.size == 3 && wShape.size == 3)
+        listOf(inShape[0], wShape[0],
+               (inShape[2] + 2*padding - dilation*(wShape[2]-1) - 1)/stride + 1)
+    else null
+    return listOf(TensorSpec("conv1d_output", outShape, inputs[0].dtype, ...))
+}
+```
+
+The formula already exists in `VoidTensorOps.calculateConv1dShape()` (line ~747).
+`ConvShapeUtils` was added to the JAR but `inferOutputs()` does not call it yet.
+
+## Impact
+
+When the StableHLO converter calls `inferOutputs()` to determine the MLIR output
+type, it gets the wrong shape. For Whisper's first conv1d:
+
+```
+Input:  [1, 80, 3000]  Weight: [384, 80, 3]  stride=1 padding=1
+Actual: [1, 80, 3000]  ← wrong (echoed input)
+Expect: [1, 384, 3000] ← correct
+```
+
+This produces `tensor<?xf32>` in the MLIR (12 occurrences), which `iree-compile`
+rejects.
+
+## Parameters are available
+
+PR #532 stores stride/padding/dilation in `operation.parameters`:
+
+```kotlin
+// RecordingExecution.kt:238-261
+val params = mapOf("stride" to stride, "padding" to padding, "dilation" to dilation, "groups" to groups)
+record(Conv1dOperation<T, V>(params), ...)
+```
+
+Verified by test: `assertEquals(1, recorded.operation.parameters["stride"])`
+
+## Suggested fix
+
+Extract `ConvShapeUtils` calls into all three `inferOutputs()` methods.
+Single PR covering conv1d/2d/3d since the bug and fix are identical.
+
+## Test
+
+See `Conv1dTapeToHloTest.kt` — asserts `tensor<?` does not appear in output MLIR.

docs/whisper-iree-issues/ISSUE-B-graph-arity-and-op-type.md

@@ -0,0 +1,73 @@
+# toComputeGraph() loses edge wiring and produces wrong op types
+
+## Problem
+
+After `tape.toComputeGraph(synthesizeExternalInputs = true)`, many graph nodes have:
+
+1. **Wrong input edge count** — binary ops (add, matmul, subtract) don't get 2 input
+   edges wired; unary ops (gelu, softmax, reshape) don't get 1 input edge wired.
+   The StableHLO converter checks arity and emits "Unsupported X arity" comments.
+
+2. **Wrong operation type** — some ops have `operation.type = "trace"` instead of a
+   recognized category. The converter uses `operation.name` for dispatch but some
+   converters also check `type`, and "trace" doesn't match any registered converter.
+
+## Scope
+
+Whisper encoder tape produces 296 graph nodes. After StableHLO conversion:
+- 166 nodes emit valid `stablehlo.*` ops
+- 157 nodes emit `// Unsupported ...` comments (some nodes emit both)
+
+Breakdown of unsupported:
+
+```
+ 32  add             — wrong arity (expected 2 inputs)
+ 24  matmul          — wrong arity (expected 2 inputs)
+ 16  unsqueeze       — wrong arity (expected 1 input)
+ 12  reshape         — wrong arity (expected 1 input)
+  9  subtract        — wrong arity (expected 2 inputs)
+  9  sqrt            — type "trace" not recognized
+  9  addScalar       — type "trace" not recognized
+  9  multiply        — wrong arity (expected 2 inputs)
+  9  divide          — wrong arity (expected 2 inputs)
+  7  variance        — wrong arity (expected 1 input)
+  7  mean            — wrong arity (expected 1 input)
+  4  softmax         — wrong arity (expected 1 input)
+  4  mulScalar       — type "trace" not recognized
+  4  gelu            — wrong arity (expected 1 input)
+  2  mean            — type "trace" not recognized
+```
+
+## Root cause hypothesis
+
+`DefaultExecutionTape.toComputeGraph()` builds graph edges by matching tensor ref
+IDs between operation outputs and subsequent operation inputs. If:
+
+- The ref ID scheme changed between recording and graph construction, edges don't
+  connect and binary ops appear to have 0 or 1 inputs.
+- Weight tensors created before `startRecording()` may not have their ref IDs in
+  the tape's scope, so edges from weights to consumers are missing.
+
+The `type = "trace"` issue: `KspTensorOps` (the auto-generated tracing wrapper)
+may record operations with a generic "trace" type string for ops that don't have
+an explicit `Operation` subclass (sqrt, addScalar, mulScalar, etc.).
+
+## Impact
+
+The generated MLIR is structurally incomplete — most ops are comments instead of
+valid StableHLO operations. `iree-compile` cannot process it.
+
+## Suggested investigation
+
+1. In `DefaultExecutionTape.toComputeGraph()`, check how `GraphEdge` source/target
+   are resolved from tape trace inputs/outputs. Are ref IDs stable?
+2. For the `type = "trace"` ops: check what `KspTensorOps` records as the operation
+   type for `sqrt`, `addScalar`, `mulScalar`. The HLO converter's operation
+   registry should recognize these names regardless of type.
+3. The `synthesizeExternalInputs = true` flag should create input/weight nodes for
+   external tensors — verify these get edges to their consumers.
+
+## Test
+
+See `Conv1dTapeToHloTest.kt` — asserts that `Unsupported` does not appear in
+output MLIR for a simple conv1d → gelu → add pipeline.

docs/whisper-iree-issues/README.md

@@ -0,0 +1,37 @@
+# SKaiNET Upstream Issues — Whisper IREE Pipeline
+
+Two issues block the native SKaiNET DSL → StableHLO → IREE compilation path.
+
+## Issue A: Conv1dOperation.inferOutputs echoes input shape
+
+`Conv1dOperation.inferOutputs()` returns `inputs[0].shape` instead of
+computing `[batch, outChannels, outLength]`. Same bug in Conv2d/Conv3d.
+
+**File:** `skainet-lang/skainet-lang-core/.../tensor/ops/TensorOperations.kt`
+**Fix:** Use `ConvShapeUtils` (already in JAR) from `inferOutputs()`.
+
+## Issue B: toComputeGraph loses edge wiring and op types
+
+`tape.toComputeGraph()` produces nodes where:
+- Binary ops (add, matmul, subtract, ...) have wrong input edge count
+- Some ops have `operation.type = "trace"` instead of recognized names
+
+157 of 296 Whisper encoder nodes emit "Unsupported ... arity" in MLIR.
+
+**File:** `skainet-compile/skainet-compile-dag/.../tape/extensions.kt` or
+`DefaultExecutionTape.toComputeGraph()`
+
+## Test
+
+`Conv1dTapeToHloTest.kt` is a KMP commonTest that:
+1. Builds a tape-recording context
+2. Runs conv1d → gelu → add through `ctx.ops`
+3. Converts tape to ComputeGraph
+4. Exports to StableHLO MLIR
+5. Asserts: no `tensor<?`, no `Unsupported`, valid `stablehlo.convolution`
+
+Place in: `skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/`
+
+Run: `./gradlew :skainet-compile:skainet-compile-hlo:allTests --tests "*Conv1dTapeToHloTest*"`
+
+Currently **fails** on both issues. Will **pass** when both are fixed.

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/TraceToGraphBuilder.kt

@@ -350,8 +350,8 @@ public class TraceToGraphBuilder(
         val count = trace.outputs.size
         return List(count) { i ->
             val name = trace.outputs[i].id
-            val shape = shapes?.getOrNull(i)
-            val dtype = dtypes?.getOrNull(i) ?: "unknown"
+            val shape = shapes?.getOrNull(i) ?: trace.outputs[i].shape.dimensions.toList()
+            val dtype = dtypes?.getOrNull(i) ?: trace.outputs[i].dtype::class.simpleName ?: "unknown"
             TensorSpec(name = name, shape = shape, dtype = dtype)
         }
     }

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

@@ -178,7 +178,14 @@ public class StableHloConverter @kotlin.jvm.JvmOverloads constructor(
                 processNode(node, context)
             } catch (e: Exception) {
                 context.emitComment("Error processing node ${node.id}: ${e.message}")
-                context.emitComment("Unsupported op ${node.operation.name} (type=${node.operation.type}) for node ${node.id}")
+                // Quote the name so trailing whitespace / casing surprises are visible,
+                // and include the registry's full key set so "no converter found"
+                // failures are self-diagnostic (is the name missing, or mis-matched?).
+                val known = registry.getSupportedOperations().sorted().joinToString(", ")
+                context.emitComment(
+                    "Unsupported op '${node.operation.name}' (type=${node.operation.type}) " +
+                        "for node ${node.id}. Known names: [$known]"
+                )
             }
         }
     }

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

@@ -8,7 +8,9 @@ import sk.ainet.compile.hlo.converters.LinalgOperationsConverter
 import sk.ainet.compile.hlo.converters.MathOperationsConverter
 import sk.ainet.compile.hlo.converters.NeuralNetOperationsConverter
 import sk.ainet.compile.hlo.converters.ReductionOperationsConverter
+import sk.ainet.compile.hlo.converters.ScalarOperationsConverter
 import sk.ainet.compile.hlo.converters.ShapeOperationsConverter
+import sk.ainet.compile.hlo.converters.UnaryMathConverter
 import kotlin.jvm.JvmStatic
 
 /**
@@ -53,6 +55,15 @@ public object StableHloConverterFactory {
         // Register reduction operations converter
         registry.register(ReductionOperationsConverter())
 
+        // Register elementwise unary math converter (sqrt, exp, log, abs, …).
+        // Must be present so downstream consumers don't cascade-fail with
+        // "wrong arity" when an upstream op is silently dropped.
+        registry.register(UnaryMathConverter())
+
+        // Register tensor+scalar ops (addScalar / mulScalar / …) emitted by the
+        // KSP-generated tracing wrapper for `tensor op Number` expressions.
+        registry.register(ScalarOperationsConverter())
+
         // Register constant operations converter
         registry.register(ConstantOperationsConverter())
 
@@ -98,6 +109,15 @@ public object StableHloConverterFactory {
         // Register reduction operations converter
         registry.register(ReductionOperationsConverter())
 
+        // Register elementwise unary math converter (sqrt, exp, log, abs, …).
+        // Must be present so downstream consumers don't cascade-fail with
+        // "wrong arity" when an upstream op is silently dropped.
+        registry.register(UnaryMathConverter())
+
+        // Register tensor+scalar ops (addScalar / mulScalar / …) emitted by the
+        // KSP-generated tracing wrapper for `tensor op Number` expressions.
+        registry.register(ScalarOperationsConverter())
+
         // Register constant operations converter
         registry.register(ConstantOperationsConverter())
 
@@ -128,6 +148,8 @@ public object StableHloConverterFactory {
         registry.register(ActivationOperationsConverter())
         registry.register(ShapeOperationsConverter())
         registry.register(ReductionOperationsConverter())
+        registry.register(UnaryMathConverter())
+        registry.register(ScalarOperationsConverter())
         registry.register(ConstantOperationsConverter())
 
         return StableHloConverter(registry, typeMapper, null, policy)

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

@@ -146,18 +146,26 @@ public class ActivationOperationsConverter : StableHloOperationConverter {
         // mapped to its position in the reduced tensor.
         val broadcastDims = (0 until rank).filter { it != axis }.joinToString(", ")
 
+        val maxInit = context.nextTempValue()
         val maxValue = context.nextTempValue()
         val maxBroadcast = context.nextTempValue()
         val shiftedValue = context.nextTempValue()
         val expValue = context.nextTempValue()
+        val sumInit = context.nextTempValue()
         val sumValue = context.nextTempValue()
         val sumBroadcast = context.nextTempValue()
         val resultValue = context.nextTempValue()
 
+        // Identity for stablehlo.maximum on floats: -inf. Spell it via the bit
+        // pattern so MLIR parses it regardless of how the element type prints.
+        val maxIdentity = "0xFF800000"
+
         val operations = listOf(
             // Reduce-max along the softmax axis (for numerical stability).
-            "$maxValue = stablehlo.custom_call @reduce_max(${operands[0]}) " +
-                "{dimensions = [$axis], keepdim = false} : $reducedType",
+            "$maxInit = stablehlo.constant dense<$maxIdentity> : tensor<$elementType>",
+            "$maxValue = stablehlo.reduce(${operands[0]} init: $maxInit) " +
+                "applies stablehlo.maximum across dimensions = [$axis] : " +
+                "($outputType, tensor<$elementType>) -> $reducedType",
 
             // Broadcast reduced max back to the input shape.
             "$maxBroadcast = stablehlo.broadcast_in_dim $maxValue, " +
@@ -170,8 +178,10 @@ public class ActivationOperationsConverter : StableHloOperationConverter {
             "$expValue = stablehlo.exponential $shiftedValue : $outputType",
 
             // Reduce-sum along the softmax axis.
-            "$sumValue = stablehlo.custom_call @reduce_sum($expValue) " +
-                "{dimensions = [$axis], keepdim = false} : $reducedType",
+            "$sumInit = stablehlo.constant dense<0.0> : tensor<$elementType>",
+            "$sumValue = stablehlo.reduce($expValue init: $sumInit) " +
+                "applies stablehlo.add across dimensions = [$axis] : " +
+                "($outputType, tensor<$elementType>) -> $reducedType",
 
             // Broadcast the sum back to the input shape.
             "$sumBroadcast = stablehlo.broadcast_in_dim $sumValue, " +