Merge pull request #803 from SKaiNET-developers/fix/graph-executor-permute-axes

fix(compile-dag): ComputeGraphExecutor replays permute with its recorded axes

Author: michalharakal

CHANGELOG.md

@@ -2,6 +2,19 @@
 
 ## [Unreleased]
 
+### Fixed
+
+- **`ComputeGraphExecutor` replayed `permute` as a plain transpose, dropping the recorded axes.**
+  The builtin dispatch grouped `permute` with `transpose`/`transpose2d`, so a traced
+  `permute(t, axes)` executed as a last-two-dims swap — only coincidentally correct for rank-2.
+  Any rank-3 permutation (e.g. multi-head attention's heads/sequence swap `[1, 0, 2]` in a
+  full-sequence encoder or prefill trace) silently produced the wrong layout; single-token decode
+  paths never hit it, which is why generation workloads didn't surface the bug. `permute` now
+  replays with its recorded `axes` (the same `List<Int>` convention `permuteBackward` and the
+  StableHLO converter already parse), falling back to the legacy transpose behavior for older
+  traces without axes. Surfaced by the BERT DSL-path encoder work in SKaiNET-transformers, which
+  carries a temporary `LLMFusedOpHandlers` override to be removed once this fix ships.
+
 ## [0.35.0] - 2026-07-08
 
 ### Added

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt

@@ -192,7 +192,25 @@ public class ComputeGraphExecutor(
             }
 
             // Shape ops
-            "transpose", "permute", "transpose2d" -> listOf(ops.transpose(inputs[0]))
+            "transpose", "transpose2d" -> listOf(ops.transpose(inputs[0]))
+            "permute" -> {
+                // The trace records axes as a List<Int> (axes.toList()) — same
+                // convention permuteBackward parses. Replaying permute as a
+                // plain last-two-dims transpose is only correct for rank-2;
+                // e.g. attention's heads/sequence swap [1, 0, 2] on rank-3
+                // silently produced the wrong layout.
+                val axes = when (val a = params["axes"]) {
+                    is IntArray -> a
+                    is List<*> -> IntArray(a.size) { (a[it] as Number).toInt() }
+                    else -> null
+                }
+                if (axes != null && axes.size == inputs[0].rank) {
+                    listOf(ops.permute(inputs[0], axes))
+                } else {
+                    // No usable axes recorded (older traces) — legacy behavior.
+                    listOf(ops.transpose(inputs[0]))
+                }
+            }
             "reshape", "view" -> {
                 // Shape can come from different parameter keys depending on trace source
                 val targetShape = (params["shape"] as? List<Int>)

skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ComputeGraphExecutorTest.kt

@@ -14,6 +14,7 @@ import sk.ainet.lang.tensor.ops.TensorSpec
 import sk.ainet.lang.types.DType
 import sk.ainet.lang.types.FP32
 import kotlin.test.Test
+import kotlin.test.assertContentEquals
 import kotlin.test.assertEquals
 import kotlin.test.assertNotNull
 import kotlin.test.assertTrue
@@ -119,6 +120,62 @@ class ComputeGraphExecutorTest {
         assertNotNull(results["mm"])
     }
 
+    @Test
+    fun permuteReplaysRecordedAxes() {
+        // Regression: permute used to be dispatched as a plain last-two-dims
+        // transpose, silently dropping the recorded axes — wrong for any
+        // rank-3 permutation (e.g. attention's heads/sequence swap [1, 0, 2]).
+        var replayedAxes: IntArray? = null
+        var transposeCalled = false
+        val ops = object : TensorOps by TestTensorOps() {
+            override fun <T : DType, V> permute(tensor: Tensor<T, V>, axes: IntArray): Tensor<T, V> {
+                replayedAxes = axes
+                return tensor
+            }
+
+            override fun <T : DType, V> transpose(tensor: Tensor<T, V>): Tensor<T, V> {
+                transposeCalled = true
+                return tensor
+            }
+        }
+
+        val graph = DefaultComputeGraph()
+        val input = graph.addNode(inputNode("input"))
+        // The tape records axes as a List<Int> (axes.toList())
+        val perm = graph.addNode(opNode("perm", "permute", mapOf("axes" to listOf(1, 0, 2))))
+        graph.addEdge(GraphEdge("e1", input, perm, tensorSpec = spec()))
+
+        val executor = ComputeGraphExecutor(graph, ops)
+        val rank3 = TestTensor(floatArrayOf(1f, 2f, 3f, 4f, 5f, 6f), shape = intArrayOf(2, 1, 3))
+        val results = executor.execute<FP32, Float>(mapOf("input" to rank3))
+
+        assertNotNull(results["perm"])
+        assertContentEquals(intArrayOf(1, 0, 2), replayedAxes, "permute must replay with its recorded axes")
+        assertTrue(!transposeCalled, "permute must not degrade to transpose when axes are recorded")
+    }
+
+    @Test
+    fun permuteWithoutRecordedAxesFallsBackToTranspose() {
+        var transposeCalled = false
+        val ops = object : TensorOps by TestTensorOps() {
+            override fun <T : DType, V> transpose(tensor: Tensor<T, V>): Tensor<T, V> {
+                transposeCalled = true
+                return tensor
+            }
+        }
+
+        val graph = DefaultComputeGraph()
+        val input = graph.addNode(inputNode("input"))
+        val perm = graph.addNode(opNode("perm", "permute"))
+        graph.addEdge(GraphEdge("e1", input, perm, tensorSpec = spec()))
+
+        val executor = ComputeGraphExecutor(graph, ops)
+        val results = executor.execute<FP32, Float>(mapOf("input" to TestTensor(floatArrayOf(1f, 2f))))
+
+        assertNotNull(results["perm"])
+        assertTrue(transposeCalled, "legacy traces without axes keep the transpose behavior")
+    }
+
     @Test
     fun emptyGraphReturnsEmptyResults() {
         val graph = DefaultComputeGraph()
@@ -133,8 +190,11 @@ class ComputeGraphExecutorTest {
  * Used for testing graph execution without needing a full tensor backend.
  */
 @Suppress("UNCHECKED_CAST")
-private class TestTensor(val values: FloatArray) : Tensor<FP32, Float> {
-    private val tensorShape = sk.ainet.lang.tensor.Shape(intArrayOf(1, values.size))
+private class TestTensor(
+    val values: FloatArray,
+    shape: IntArray = intArrayOf(1, values.size),
+) : Tensor<FP32, Float> {
+    private val tensorShape = sk.ainet.lang.tensor.Shape(shape)
     override val data: sk.ainet.lang.tensor.data.TensorData<FP32, Float> = object : sk.ainet.lang.tensor.data.TensorData<FP32, Float> {
         override fun get(vararg indices: Int): Float = values[indices.last()]
         override fun set(vararg indices: Int, value: Float) { values[indices.last()] = value }