Torq passes: seq-major RoPE + output prune for the DSL Moonshine encoder

synaptics-torq/src/jvmMain/kotlin/sk/ainet/vendors/torq/TorqAttentionTilingPass.kt

@@ -114,12 +114,33 @@ class TorqAttentionTilingPass(
             // The extra [1,0,2] hop makes this robust to RoPE (which leaves K in [H,Sk,D]) and
             // to the no-RoPE case alike; without RoPE the paired [1,0,2] transposes canonicalize
             // away to the single [1,2,0].
-            val kShdS = spec(sk, h, d)
-            val kShd = op("transpose", mapOf("permutation" to listOf(1, 0, 2)), listOf(k3s), kShdS)
-                .also { edges += Triple(kSrc, it, 0) }
+            // If TorqRopeSeqMajorPass ran, the SDPA's K producer is `permute[1,0,2](seqK)` where seqK
+            // is the RoPE'd K already in seq-major [Sk,H,D] (reshape-sourced). Consume seqK directly and
+            // build K^T with a SINGLE [1,2,0] — the reshape-sourced seq-major K^T the Torq layout solver
+            // accepts (a head-major-sourced K^T trips MatMulPattern.cpp:57). Otherwise fall back to the
+            // route-through-[Sk,H,D] double transpose.
+            fun permAxes(n: GraphNode): List<Int>? =
+                (n.operation.parameters["permutation"] ?: n.operation.parameters["axes"])
+                    .let { it as? List<*> }?.map { (it as Number).toInt() }
+            val kSeq: Pair<Pair<GraphNode, Int>, TensorSpec>? = run {
+                val p = kSrc.first
+                val op = p.operationName.lowercase()
+                if ((op == "permute" || op == "transpose") && permAxes(p) == listOf(1, 0, 2) &&
+                    p.inputs.firstOrNull()?.shape?.size == 3
+                ) producerOf[p.id to 0]?.let { it to p.inputs[0] } else null
+            }
             val kTS = spec(h, d, sk)
-            val kT = op("transpose", mapOf("permutation" to listOf(1, 2, 0)), listOf(kShdS), kTS)
-                .also { edges += Triple(kShd to 0, it, 0) }
+            val kT = if (kSeq != null) {
+                val (seqSrc, seqSpec) = kSeq // seqSpec = [Sk,H,D] seq-major RoPE'd K
+                op("transpose", mapOf("permutation" to listOf(1, 2, 0)), listOf(seqSpec), kTS)
+                    .also { edges += Triple(seqSrc, it, 0) }
+            } else {
+                val kShdS = spec(sk, h, d)
+                val kShd = op("transpose", mapOf("permutation" to listOf(1, 0, 2)), listOf(k3s), kShdS)
+                    .also { edges += Triple(kSrc, it, 0) }
+                op("transpose", mapOf("permutation" to listOf(1, 2, 0)), listOf(kShdS), kTS)
+                    .also { edges += Triple(kShd to 0, it, 0) }
+            }
 
             val groupOuts = mutableListOf<Pair<GraphNode, Int>>() // (node, headCount)
             var s = 0

synaptics-torq/src/jvmMain/kotlin/sk/ainet/vendors/torq/TorqPlugin.kt

@@ -43,16 +43,23 @@ public object TorqPlugin {
         maxHeadsPerTile: Int = 4,
         maxQuerySeqPerTile: Int = 83,
         ffnHiddenTile: Int = 288,
+        modelDim: Int = 288,
     ) {
         TargetOptimizers.register(object : TargetOptimizer {
             override val target: String = TARGET
             override fun dagPasses(): List<GraphOptimizationPass> = listOf(
+                // Apply RoPE seq-major on the K path BEFORE tiling, so the tiling pass builds a
+                // reshape-sourced seq-major K^T (else the Torq layout solver trips MatMulPattern:57).
+                TorqRopeSeqMajorPass(),
                 TorqAttentionTilingPass(
                     maxHeadsPerTile = maxHeadsPerTile,
                     maxQuerySeqPerTile = maxQuerySeqPerTile,
                 ),
                 TorqFfnTilingPass(hiddenTile = ffnHiddenTile),
             )
+            // NOTE: TorqPruneOutputsPass is applied by the app AFTER tiling (the tiling pass
+            // itself leaves the dangling [1,H,S,D] Q/K/V reshapes that must be pruned), so it is
+            // not registered here. See MoonshineEncoderExport.
             // granularity() intentionally left at the default (null / decompose-all): the Torq
             // compiler rejects `stablehlo.composite`, so there is no fused-op emission to select
             // yet. Correctness for Torq is achieved by matching the vendor's op *structure* in
@@ -64,13 +71,16 @@ public object TorqPlugin {
      * Recommended `torq-compile` flags for the SL2610. Applied by the build tool when invoking
      * the vendor compiler — these are compile-time knobs, not DAG passes.
      *
-     * `--torq-disable-slices` is required: Torq's slice-based NSS execution mis-aliases buffers
-     * in large graphs (LayerNorm/attention output collapses to NaN/1e23 without it).
+     * NOTE: `--torq-disable-slices` is deliberately NOT included. It disables the linalg slicing
+     * that tiles the attention softmax to fit the CSS stack, so on the ENCODER it causes
+     * "CSS program allocations of 1328 bytes exceed maximum CSS stack size" (verified: enc6.mlir
+     * compiles WITHOUT it → 277627 bytes, FAILS with it). It was a decoder LayerNorm-NaN
+     * workaround; enable it per-compile via `TORQ_DISABLE_SLICES=1` only if a decoder-on-Torq
+     * path needs it. See the demo's `scripts/.docker/compile-entry.sh`.
      */
     public val compileFlags: List<String> = listOf(
         "--torq-hw=SL2610",
         "--torq-target-host-triple=native",
         "--torq-css-qemu",
-        "--torq-disable-slices",
     )
 }

synaptics-torq/src/jvmMain/kotlin/sk/ainet/vendors/torq/TorqPruneOutputsPass.kt

@@ -0,0 +1,51 @@
+package sk.ainet.vendors.torq
+
+import sk.ainet.compile.opt.GraphOptimizationPass
+import sk.ainet.compile.opt.GraphOptimizationResult
+import sk.ainet.lang.graph.ComputeGraph
+import sk.ainet.lang.graph.GraphNode
+
+/**
+ * Torq-**target** graph pass: prune dangling graph outputs down to the real model output.
+ *
+ * The DSL trace of `moonshineEncoder` records the per-layer RoPE'd Q/K/V tensors as leaf nodes
+ * (no consumers), so `func @main` ends up returning **19** tensors: 18 spurious `[1,H,S,D]`
+ * attention intermediates + the real `[.., S, dim]` encoder memory. The llvm-cpu backend
+ * tolerates the extra results, but the Torq global-layout solver does NOT — the returned
+ * attention intermediates force a K materialization layout that trips `MatMulPattern.cpp:57`
+ * (proven: dropping them to the single real output removes the assertion).
+ *
+ * This pass keeps only the leaves that look like the encoder output — output spec rank ≤ 3 with
+ * a trailing model dimension [modelDim] — and iteratively removes every other leaf together with
+ * any producer left dangling by the removal, until a fixpoint. Purely structural / HW-agnostic in
+ * effect, but only *wanted* for the Torq target, so it is registered by [TorqPlugin.install].
+ */
+public class TorqPruneOutputsPass(
+    private val modelDim: Int,
+) : GraphOptimizationPass {
+    override val name: String = "torq-prune-outputs"
+
+    /** A leaf is a "real" output if it is rank ≤ 3 and its last dim is the model dim. */
+    private fun isKept(n: GraphNode): Boolean {
+        val sh = n.outputs.firstOrNull()?.shape ?: return false
+        return sh.size in 2..3 && sh.last() == modelDim
+    }
+
+    override fun apply(graph: ComputeGraph): GraphOptimizationResult {
+        var changed = false
+        // Peel dangling leaves until every leaf is a kept output (or nothing left to remove).
+        while (true) {
+            val leaves = graph.getOutputNodes() // nodes with no outgoing edges
+            val drop = leaves.filterNot { isKept(it) }
+            if (drop.isEmpty()) break
+            for (n in drop) {
+                // Remove the node's incoming edges first, then the node itself. Producers that
+                // become leaves as a result are revisited on the next iteration.
+                graph.edges.filter { it.destination.id == n.id }.forEach { graph.removeEdge(it) }
+                graph.removeNode(n)
+                changed = true
+            }
+        }
+        return GraphOptimizationResult(graph, changed = changed)
+    }
+}