diff --git a/synaptics-torq/src/jvmMain/kotlin/sk/ainet/vendors/torq/TorqAttentionTilingPass.kt b/synaptics-torq/src/jvmMain/kotlin/sk/ainet/vendors/torq/TorqAttentionTilingPass.kt index 6c78680..c35b602 100644 --- a/synaptics-torq/src/jvmMain/kotlin/sk/ainet/vendors/torq/TorqAttentionTilingPass.kt +++ b/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? = + (n.operation.parameters["permutation"] ?: n.operation.parameters["axes"]) + .let { it as? List<*> }?.map { (it as Number).toInt() } + val kSeq: Pair, 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>() // (node, headCount) var s = 0 diff --git a/synaptics-torq/src/jvmMain/kotlin/sk/ainet/vendors/torq/TorqPlugin.kt b/synaptics-torq/src/jvmMain/kotlin/sk/ainet/vendors/torq/TorqPlugin.kt index b5352c4..c64d0e1 100644 --- a/synaptics-torq/src/jvmMain/kotlin/sk/ainet/vendors/torq/TorqPlugin.kt +++ b/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 = 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 = listOf( "--torq-hw=SL2610", "--torq-target-host-triple=native", "--torq-css-qemu", - "--torq-disable-slices", ) } diff --git a/synaptics-torq/src/jvmMain/kotlin/sk/ainet/vendors/torq/TorqPruneOutputsPass.kt b/synaptics-torq/src/jvmMain/kotlin/sk/ainet/vendors/torq/TorqPruneOutputsPass.kt new file mode 100644 index 0000000..5696fca --- /dev/null +++ b/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) + } +} diff --git a/synaptics-torq/src/jvmMain/kotlin/sk/ainet/vendors/torq/TorqRopeSeqMajorPass.kt b/synaptics-torq/src/jvmMain/kotlin/sk/ainet/vendors/torq/TorqRopeSeqMajorPass.kt new file mode 100644 index 0000000..dce99d1 --- /dev/null +++ b/synaptics-torq/src/jvmMain/kotlin/sk/ainet/vendors/torq/TorqRopeSeqMajorPass.kt @@ -0,0 +1,260 @@ +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.DefaultComputeGraph +import sk.ainet.lang.graph.GraphEdge +import sk.ainet.lang.graph.GraphNode +import sk.ainet.lang.tensor.ops.GenericOperation +import sk.ainet.lang.tensor.ops.TensorSpec + +/** + * Torq-**target** pass: rewrite the self-attention **K** path so RoPE is applied SEQ-major. + * + * Why: the Torq global-layout solver trips `MatMulPattern.cpp:57` when the QK^T's K operand + * (`[H,D,Sk]`) is produced by transposing a *head-major* tensor. The moonshine attention reshapes + * the K projection to seq-major `[Sk,H,D]`, then `permute[1,0,2]`→head-major to run RoPE; the + * tiling pass then transposes back to build K^T — and that transpose-sourced seq-major K poisons + * the solver. The hand-written enc6 (no RoPE) keeps K reshape-sourced seq-major and compiles. + * Proven by isolation (scratchpad/min_oproj*.mlir): seq-major RoPE'd K^T compiles; head-major + * K^T does not. + * + * This pass **sinks the head-transpose below RoPE**: it finds the `permute[1,0,2]` on the K path + * (seq-major `[Sk,H,D]` → head-major `[H,Sk,D]`, feeding the RoPE that reaches the SDPA's K input), + * rebuilds every op between that permute and the SDPA in seq-major (swap the leading two dims; + * reshape any `[Sk,D]` cos/sin broadcast operand to `[Sk,1,D]`), and re-inserts a single + * `permute[1,0,2]` at the SDPA boundary. Net effect: RoPE runs seq-major, and the downstream + * TorqAttentionTilingPass builds a reshape-sourced seq-major K^T. HW-agnostic in output (still + * portable StableHLO); registered for the "torq" target only. Runs BEFORE TorqAttentionTilingPass. + * + * Only the K path is rewritten: Q stays head-major (a head-major Q with a seq-sourced K^T + * compiles — v4b/v4f), and V has no RoPE and needs no seq-major layout. + */ +public class TorqRopeSeqMajorPass : GraphOptimizationPass { + override val name: String = "torq-rope-seq-major" + + // Ops on the K RoPE path whose semantics are unchanged by swapping the leading two dims + // (they act on trailing dims / broadcast), so the head-transpose can be sunk past them. All + // rank-3 [H,Sk,D]; the batch `unsqueeze` right below the SDPA is a rank-changing BOUNDARY that + // is NOT rebuilt (it stays after the restore permute). + private val sinkableAlways = setOf( + "convert", "multiply", "add", "subtract", "reshape", "slice", "narrow", "concat", "concatenate", + ) + + /** An op the head-transpose can be sunk past (commutes with a leading-two-dim swap). Internal + * `unsqueeze`/`squeeze` (dim ≥ 2, e.g. the rotate-half pair-axis) qualify; the batch + * `unsqueeze`/`squeeze` (dim 0/1) does not — it is the SDPA boundary, handled separately. */ + private fun isSinkable(n: GraphNode): Boolean { + val op = n.operationName.lowercase() + if (op in sinkableAlways) return true + if (op == "unsqueeze" || op == "squeeze") { + val dim = (n.operation.parameters["dim"] as? Number)?.toInt() ?: return false + return dim >= 2 + } + return false + } + + private fun isHeadTranspose(n: GraphNode): Boolean { + if (n.operationName.lowercase() != "permute" && n.operationName.lowercase() != "transpose") return false + val axes = (n.operation.parameters["axes"] as? List<*>)?.map { (it as Number).toInt() } + val inSh = n.inputs.firstOrNull()?.shape + // [1,0,2] on a rank-3 [Sk,H,D] -> [H,Sk,D]. + return axes == listOf(1, 0, 2) && inSh != null && inSh.size == 3 + } + + /** swap the leading two dims of a shape (head<->seq). */ + private fun swap01(sh: List?): List? = + if (sh == null || sh.size < 2) sh else listOf(sh[1], sh[0]) + sh.drop(2) + + override fun apply(graph: ComputeGraph): GraphOptimizationResult { + val dbg = System.getenv("ROPE_SEQ_DEBUG") == "1" + val sdpas = graph.nodes.filter { it.operationName.lowercase() == "scaleddotproductattention" } + if (dbg) println("[rope-seq] SDPA nodes=${sdpas.size} (all ops: ${graph.nodes.map { it.operationName }.distinct().take(40)})") + if (sdpas.isEmpty()) return GraphOptimizationResult(graph, changed = false) + + val nodeById = graph.nodes.associateBy { it.id }.toMutableMap() + // producer of (nodeId, inputIdx) -> (node, outIdx) + val producerOf = HashMap, Pair>() + for (e in graph.edges) producerOf[e.destination.id to e.destinationInputIndex] = e.source to e.sourceOutputIndex + + val newGraph = DefaultComputeGraph() + graph.nodes.forEach { newGraph.addNode(it) } + graph.edges.forEach { newGraph.addEdge(it) } + var changed = false + var uid = 0 + + for (sdpa in sdpas) { + // 1) The K path (SDPA input 1). The batch `unsqueeze`/`squeeze` right below the SDPA is a + // rank-changing boundary: skip it and rewire ITS rank-3 input; else rewire SDPA.K directly. + val kProd0 = producerOf[sdpa.id to 1] ?: continue + val boundaryIsUnsqueeze = kProd0.first.operationName.lowercase() in setOf("unsqueeze", "squeeze") + val kProd = if (boundaryIsUnsqueeze) (producerOf[kProd0.first.id to 0] ?: continue) else kProd0 + var headPermute: GraphNode? = null + val visited = HashSet() + // BFS up over the (possibly branching, e.g. RoPE add) sinkable region until every branch + // reaches the SAME head-transpose permute. + val frontier = ArrayDeque() + frontier.add(kProd.first) + val region = LinkedHashSet() + var ok = true + // The region top (kProd) MUST be sinkable; otherwise there's nothing to rewrite. + if (!isSinkable(kProd.first)) ok = false + while (ok && frontier.isNotEmpty()) { + val n = frontier.removeFirst() + if (!visited.add(n.id)) continue + if (isHeadTranspose(n)) { + if (headPermute == null) headPermute = n + else if (headPermute!!.id != n.id) { ok = false; break } // two different permutes -> bail + continue + } + if (!isSinkable(n)) continue // external operand (cos/sin/mean/weight): leaf, reuse + region.add(n) + // Follow ONLY sinkable producers or the head-transpose; external operands stay put. + for (i in n.inputs.indices) { + val p = producerOf[n.id to i] ?: continue + if (isHeadTranspose(p.first) || isSinkable(p.first)) frontier.add(p.first) + } + } + val hp = headPermute + if (dbg) println("[rope-seq] sdpa=${sdpa.id} boundaryUnsq=$boundaryIsUnsqueeze ok=$ok headPermute=${hp?.id} regionSize=${region.size} regionOps=${region.map { it.operationName }}") + if (!ok || hp == null || region.isEmpty()) continue + + // 2) The seq-major source feeding the head-transpose. + val seqSrc = producerOf[hp.id to 0] ?: continue // (node, outIdx) producing [Sk,H,D] + + // 3) Rebuild the region in seq-major. Map old node id -> (new node, outIdx) producing seq-major. + val remap = HashMap>() + // The head-transpose's consumers in the region should read the seq-major source directly. + remap[hp.id] = seqSrc + // Rebuild in a stable order: iterate until all region nodes are rebuilt (their inputs available). + val pending = ArrayDeque(region) + var guard = region.size * region.size + 8 + val newEdges = ArrayList() + val newNodes = ArrayList() + while (pending.isNotEmpty() && guard-- > 0) { + val n = pending.removeFirst() + // resolve each input's seq-major producer + val inProds = ArrayList>() + var ready = true + for (i in n.inputs.indices) { + val p = producerOf[n.id to i] + if (p == null) { inProds.add(seqSrc); continue } // constant/no-edge input handled below + when { + p.first.id == hp.id -> inProds.add(seqSrc) + remap.containsKey(p.first.id) -> inProds.add(remap[p.first.id]!!) + p.first.id in region.map { it.id } -> { ready = false } // producer in region not rebuilt yet + else -> inProds.add(p.first to p.second) // outside region (e.g. cos/sin constant): reuse + } + if (!ready) break + } + if (!ready) { pending.addLast(n); continue } + + // Build seq-major specs: swap leading two dims of inputs/outputs. + val newIns = ArrayList() + for (i in n.inputs.indices) { + val orig = n.inputs[i] + val prod = inProds[i] + val prodShape = prod.first.outputs.getOrNull(prod.second)?.shape + // If the producer already yields a seq-major shape (rebuilt), use it; else this is an + // outside-region operand (cos/sin/const) -> keep its shape, but insert a size-1 head + // dim so a [Sk,D] operand broadcasts against seq-major [Sk,H,D]. + val insertHead = prod.first.id !in remap.keys && orig.shape != null && orig.shape!!.size == 2 + val spec = if (insertHead) { + val s = prodShape ?: orig.shape!! + orig.copy(name = "${orig.name}_sq${uid}", shape = listOf(s[0], 1, s[1])) + } else { + orig.copy(name = "${orig.name}_sq${uid}", shape = swap01(orig.shape)) + } + uid++ + newIns.add(spec) + } + val newOuts = n.outputs.map { it.copy(name = "${it.name}_sq${uid++}", shape = swap01(it.shape)) } + // adjust shape-bearing params (reshape newShape, slice indices, concat dim). + val np = adjustParams(n.operation.parameters, n) + val nn = GraphNode("${n.id}_seq", GenericOperation(n.operationName, np), newIns, newOuts) + newNodes.add(nn) + remap[n.id] = nn to 0 + // for insertHead operands, splice a reshape node + for (i in n.inputs.indices) { + val prod = inProds[i] + val needReshape = prod.first.id !in remap.keys && n.inputs[i].shape?.size == 2 + if (needReshape) { + val srcSpec = prod.first.outputs[prod.second] + val rsOut = newIns[i] + val rs = GraphNode( + "${n.id}_bcastfix_${i}_$uid", + GenericOperation("reshape", mapOf("newShape" to (rsOut.shape ?: emptyList()))), + listOf(srcSpec), listOf(rsOut), + ) + uid++ + newNodes.add(rs) + newEdges.add(edge(prod.first, prod.second, rs, 0, srcSpec)) + newEdges.add(edge(rs, 0, nn, i, rsOut)) + } else { + newEdges.add(edge(prod.first, prod.second, nn, i, newIns[i])) + } + } + } + if (dbg) println("[rope-seq] sdpa=${sdpa.id} rebuilt=${remap.size - 1}/${region.size} pendingLeft=${pending.size}") + if (pending.isNotEmpty()) continue // couldn't rebuild cleanly -> skip this SDPA, leave as-is + + // 4) The region's top node (the one whose OLD output fed the SDPA K, i.e. kProd) now has a + // seq-major rebuild. Append a permute[1,0,2] to restore head-major and rewire SDPA.K to it. + val topSeq = remap[kProd.first.id] ?: continue + val seqSpec = topSeq.first.outputs[topSeq.second] + val headSpec = seqSpec.copy(name = "${seqSpec.name}_hd$uid", shape = swap01(seqSpec.shape)) + uid++ + val restore = GraphNode( + "${sdpa.id}_kseq_permute", + GenericOperation("permute", mapOf("axes" to listOf(1, 0, 2))), + listOf(seqSpec), listOf(headSpec), + ) + newNodes.add(restore) + newEdges.add(edge(topSeq.first, topSeq.second, restore, 0, seqSpec)) + + // Commit: add new nodes/edges, then rewire the boundary consumer to the restore permute. + // If an unsqueeze sits below the SDPA, rewire its rank-3 input; else rewire SDPA.K directly. + newNodes.forEach { newGraph.addNode(it) } + newEdges.forEach { newGraph.addEdge(it) } + val (consumerId, consumerIn, consumerSpec) = + if (boundaryIsUnsqueeze) Triple(kProd0.first.id, 0, kProd0.first.inputs[0]) + else Triple(sdpa.id, 1, sdpa.inputs[1]) + val consumerNode = newGraph.nodes.firstOrNull { it.id == consumerId } ?: continue + newGraph.edges.filter { it.destination.id == consumerId && it.destinationInputIndex == consumerIn } + .forEach { newGraph.removeEdge(it) } + newGraph.addEdge(edge(restore, 0, consumerNode, consumerIn, consumerSpec)) + changed = true + } + + return GraphOptimizationResult(newGraph, changed = changed) + } + + private fun edge(src: GraphNode, srcOut: Int, dst: GraphNode, dstIn: Int, spec: TensorSpec) = + GraphEdge("e_${src.id}_${srcOut}__${dst.id}_${dstIn}_rsm", src, dst, srcOut, dstIn, spec) + + /** Adjust shape-bearing params of an op when its leading two dims are swapped. */ + private fun adjustParams(params: Map, n: GraphNode): Map { + val out = params.toMutableMap() + // reshape newShape -> swap leading two + (params["newShape"] as? List<*>)?.let { ns -> + val dims = ns.mapNotNull { (it as? Number)?.toInt() } + if (dims.size >= 2) out["newShape"] = listOf(dims[1], dims[0]) + dims.drop(2) + } + // slice start/limit/strides -> swap leading two entries + for (key in listOf("start_indices", "limit_indices", "strides", "starts", "limits", "start", "limit")) { + (params[key] as? List<*>)?.let { idx -> + val v = idx.mapNotNull { (it as? Number)?.toInt() } + if (v.size >= 2) out[key] = listOf(v[1], v[0]) + v.drop(2) + } + } + // concat/concatenate dim: 0<->1 swap; others unchanged + for (key in listOf("dim", "dimension", "axis")) { + (params[key] as? Number)?.toInt()?.let { d -> + if (d == 0) out[key] = 1 else if (d == 1) out[key] = 0 + } + } + // unsqueeze/squeeze dim in the leading region: 0<->1 + return out + } +}