feat(backend): bridge the packed SPI kernels through the ordered key (#973.2)

skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelDispatch.kt

@@ -1,5 +1,6 @@
 package sk.ainet.backend.api.kernel
 
+import sk.ainet.lang.memory.BlockOrder
 import sk.ainet.lang.memory.ExperimentalMemoryApi
 import sk.ainet.lang.memory.Format
 import sk.ainet.lang.memory.Scope
@@ -70,6 +71,18 @@ public object KernelDispatch {
         out: TensorView,
         scope: Scope = Scope.Ambient,
         sink: TraceSink = NoopTraceSink,
+        /**
+         * Relayout a canonical packed weight into kernel order when that is what unlocks a packed
+         * kernel (#973/#1095).
+         *
+         * **Off by default, deliberately.** The relayout is O(bytes), so doing it inside a decode
+         * step would copy the whole weight on every token — the per-forward copy #973 objects to,
+         * merely moved. A weight is prepacked *once*, at load
+         * ([sk.ainet.lang.memory.TensorView.prepack]); a canonical weight handed straight to the
+         * dispatcher gets the decoding reference kernel, which is correct and slower. Pass `true`
+         * for a one-shot call where the copy is cheaper than the reference path.
+         */
+        prepackWeights: Boolean = false,
     ) {
         val key = KernelKey.matmul(a, b)
         val exact = find(key)
@@ -92,6 +105,19 @@ public object KernelDispatch {
                 }
             }
         }
+        // A packed kernel reads its weight input-block-major; a weight loaded from a file is
+        // canonical. Now that the order is in the key (#973/#1094) the two can be bridged — but the
+        // relayout is O(bytes), so it happens only when the caller asks, and the right shape for a
+        // hot loop is a weight prepacked once at load, which hits the exact key above and copies
+        // nothing.
+        if (prepackWeights && b.layout.blocked && b.layout.blockOrder == BlockOrder.ROW_MAJOR) {
+            val prepacked = b.prepack(BlockOrder.INPUT_BLOCK_MAJOR, scope, sink)
+            val packedKernel = find(KernelKey.matmul(a, prepacked))
+            if (packedKernel != null) {
+                runTraced(packedKernel, listOf(a, prepacked), out, sink)
+                return
+            }
+        }
         // No exact kernel: adapt the operands a kernel would accept, then fall back to the reference,
         // which reads any format through decoding get().
         val adaptedA = adapt(a, scope, sink, "gather")

skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelPacks.kt

@@ -4,6 +4,7 @@ import sk.ainet.lang.memory.ExperimentalMemoryApi
 import sk.ainet.lang.memory.Format
 import sk.ainet.lang.memory.Storage
 import sk.ainet.lang.memory.TensorView
+import sk.ainet.lang.tensor.storage.TensorEncoding
 import sk.ainet.lang.types.FP32
 
 /**
@@ -14,15 +15,11 @@ import sk.ainet.lang.types.FP32
  * under the keys the dispatcher looks up, so the generic path gets the fast kernel instead of the
  * decoding reference whenever the operands' formats and layouts match what the pack declares.
  *
- * **Scope note (SKEEP-003 migration, #1029).** Only the dense FP32 kernel is bridged here. The
- * packed (Q4_0…Q6_K) SPI kernels take their weight bytes in **block-major** order — the layout
- * `DefaultCpuOpsBase.transposePackedBlocks` produces — while a packed `TensorView` describes the
- * canonical row-major block order of the file. That contract is exactly what #973 reports as
- * unwritten and contradictory across the engine and the converters, and getting it wrong is the
- * silent-wrong-numbers class of #968/#971. Bridging the packed kernels therefore waits until #973
- * pins the byte order down; until then the packed fast paths stay on their existing (working)
- * ladder in `DefaultCpuOps`/`DefaultCpuOpsJvm`, and the registry serves packed operands with the
- * decoding reference kernel, which is correct for any layout.
+ * **The packed kernels are bridged too, now that the order is in the key (#973.2, #1095).** They
+ * read their weight input-block-major; a packed `TensorView` from a file is canonical row-major;
+ * [sk.ainet.lang.memory.BlockOrder] says which is which, so the kernel declares what it takes and
+ * the dispatcher relayouts. That distinction is what #1029 was missing and what made mixing the two
+ * a silent-wrong-numbers bug rather than a crash (#968, #971).
  */
 @ExperimentalMemoryApi
 public object KernelPacks {
@@ -46,6 +43,30 @@ public object KernelPacks {
         val strided = OperandKey(Format.dense(FP32), LayoutClass.STRIDED)
         KernelDispatch.register(Fp32ViewMatmulKernel(p.name, fp32, KernelKey("matmul", listOf(dense, strided))))
         KernelDispatch.register(Fp32ViewMatmulKernel(p.name, fp32, KernelKey("matmul", listOf(dense, dense))))
+        installPacked(p)
+    }
+
+    /**
+     * Install [provider]'s packed matmul kernels as [PackedViewMatmulKernel]s, keyed on
+     * `BLOCKED_INPUT_MAJOR` — the order they actually read.
+     *
+     * A provider that offers no kernel for a format simply does not get one registered, and the
+     * decoding reference kernel keeps serving that format, as it does today.
+     */
+    public fun installPacked(provider: KernelProvider) {
+        fun register(encoding: TensorEncoding, kernel: ((FloatArray, Int, ByteArray, Int, Int, Int, FloatArray, Int) -> Unit)?) {
+            if (kernel == null) return
+            KernelDispatch.register(
+                PackedViewMatmulKernel(provider.name, encoding.name, PackedViewMatmulKernel.keyFor(encoding), kernel),
+            )
+        }
+        register(TensorEncoding.Q4_0, provider.matmulQ4_0()?.let { k -> { i, io, w, wo, id, od, o, oo -> k.matmul(i, io, w, wo, id, od, o, oo) } })
+        register(TensorEncoding.Q5_0, provider.matmulQ5_0()?.let { k -> { i, io, w, wo, id, od, o, oo -> k.matmul(i, io, w, wo, id, od, o, oo) } })
+        register(TensorEncoding.Q5_1, provider.matmulQ5_1()?.let { k -> { i, io, w, wo, id, od, o, oo -> k.matmul(i, io, w, wo, id, od, o, oo) } })
+        register(TensorEncoding.Q8_0, provider.matmulQ8_0()?.let { k -> { i, io, w, wo, id, od, o, oo -> k.matmul(i, io, w, wo, id, od, o, oo) } })
+        register(TensorEncoding.Q4_K, provider.matmulQ4K()?.let { k -> { i, io, w, wo, id, od, o, oo -> k.matmul(i, io, w, wo, id, od, o, oo) } })
+        register(TensorEncoding.Q5_K, provider.matmulQ5K()?.let { k -> { i, io, w, wo, id, od, o, oo -> k.matmul(i, io, w, wo, id, od, o, oo) } })
+        register(TensorEncoding.Q6_K, provider.matmulQ6K()?.let { k -> { i, io, w, wo, id, od, o, oo -> k.matmul(i, io, w, wo, id, od, o, oo) } })
     }
 
     /** The reference matmul for dense FP32 — always available, so a key is never unserved. */

skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/PackedViewMatmulKernel.kt

@@ -0,0 +1,87 @@
+package sk.ainet.backend.api.kernel
+
+import sk.ainet.lang.memory.BlockOrder
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+import sk.ainet.lang.memory.Format
+import sk.ainet.lang.memory.Storage
+import sk.ainet.lang.memory.TensorView
+import sk.ainet.lang.tensor.storage.TensorEncoding
+import sk.ainet.lang.types.FP32
+
+/**
+ * A [ViewKernel] over one of the packed SPI matmul kernels (SKEEP-003 §5.2; #973.2, #1095).
+ *
+ * These kernels read their weight **input-block-major** — `(blockIdx * outputDim + o)`, every
+ * output row's block for one input block contiguous — which is why bridging them was deferred in
+ * #1029: a packed `TensorView` loaded from a file is canonical row-major, and nothing in the key
+ * said so. Now it does ([BlockOrder]), so the kernel declares the order it reads and the dispatcher
+ * relayouts when the operand disagrees.
+ *
+ * One row at a time: the SPI is a matrix-vector kernel, so an activation of `m` rows is `m` calls,
+ * which is what the decode path does anyway (`m == 1`).
+ */
+@ExperimentalMemoryApi
+public class PackedViewMatmulKernel(
+    providerName: String,
+    private val encodingName: String,
+    override val key: KernelKey,
+    private val matmul: (
+        input: FloatArray, inputOffset: Int,
+        weight: ByteArray, weightByteOffset: Int,
+        inputDim: Int, outputDim: Int,
+        output: FloatArray, outputOffset: Int,
+    ) -> Unit,
+) : ViewKernel {
+
+    override val name: String = "$providerName-$encodingName"
+
+    override fun run(inputs: List<TensorView>, out: TensorView) {
+        require(inputs.size == 2) { "matmul takes two operands" }
+        val a = inputs[0]
+        val w = inputs[1]
+        val rows = a.shape[0]
+        val k = a.shape[1]
+        val n = w.shape[0]
+        require(w.shape[1] == k) { "inner dimensions disagree: [$rows, $k] × [$n, ${w.shape[1]}]" }
+        check(w.layout.blockOrder == BlockOrder.INPUT_BLOCK_MAJOR) {
+            "$name reads input-block-major weights; this one is ${w.layout.blockOrder} — the dispatcher " +
+                "should have prepacked it (#973)"
+        }
+
+        val aHeap = a.storage as? Storage.Heap ?: return fallback(inputs, out)
+        val wHeap = w.storage as? Storage.Heap ?: return fallback(inputs, out)
+        val oHeap = out.storage as? Storage.Heap ?: return fallback(inputs, out)
+        val activation = aHeap.floats ?: return fallback(inputs, out)
+        val weight = wHeap.bytes ?: return fallback(inputs, out)
+        val output = oHeap.floats ?: return fallback(inputs, out)
+        // The SPI takes a contiguous activation row; a strided one would be mis-indexed, so it goes
+        // to the reference kernel rather than silently reading the wrong floats.
+        if (!a.isContiguous) return fallback(inputs, out)
+
+        val weightOffset = wHeap.arrayOffset + (w.layout.offsetElements * w.layout.elementBytes).toInt()
+        for (r in 0 until rows) {
+            matmul(
+                activation, aHeap.arrayOffset + (a.layout.offsetElements + r.toLong() * k).toInt(),
+                weight, weightOffset,
+                k, n,
+                output, oHeap.arrayOffset + (out.layout.offsetElements + r.toLong() * n).toInt(),
+            )
+        }
+    }
+
+    private fun fallback(inputs: List<TensorView>, out: TensorView) {
+        ReferenceMatmulKernel(key).run(inputs, out)
+    }
+
+    public companion object {
+        /** The key a packed kernel serves: a dense contiguous activation × an input-major weight. */
+        public fun keyFor(encoding: TensorEncoding, capabilities: Set<String> = emptySet()): KernelKey = KernelKey(
+            op = "matmul",
+            operands = listOf(
+                OperandKey.contiguous(Format.dense(FP32)),
+                OperandKey(Format(FP32, encoding), LayoutClass.BLOCKED_INPUT_MAJOR),
+            ),
+            capabilities = capabilities,
+        )
+    }
+}

skainet-backends/skainet-backend-cpu/src/goldenTest/kotlin/sk/ainet/exec/golden/Goldens.kt

@@ -9,6 +9,17 @@ package sk.ainet.exec.golden
  */
 internal object Goldens {
     val expected: Map<String, String> = mapOf(
+        // #1095 — the packed kernels reached through the registry: a canonical weight relayouted
+        // into kernel order by the dispatcher. The `head` values are the first row of the existing
+        // scalar-matmul goldens, which is the point: the bridge changes nothing about the numbers.
+        "dispatch-bridge/Q4_0" to "n=4 fnv=5338c3fcbe1b41a0 head=3f15be82,3f0f4b4e,3e49fb8e,be220b31",
+        "dispatch-bridge/Q4_K" to "n=4 fnv=45ad612547bdb769 head=c29e9232,c3116957,c3a2cb4f,c2da54e5",
+        "dispatch-bridge/Q5_0" to "n=4 fnv=3504943a5a9d6a37 head=3ffafe21,4073853c,c0a97fbe,3e9fe3a0",
+        "dispatch-bridge/Q5_1" to "n=4 fnv=5ab4974269167da2 head=3fe5eab0,bfe35b0c,40821cc5,40ca2f5a",
+        "dispatch-bridge/Q5_K" to "n=4 fnv=f6d101a0033fc64e head=4317c800,c3350201,4380baa3,4286212d",
+        "dispatch-bridge/Q6_K" to "n=4 fnv=de59f18bcd4472af head=c41704aa,425d6ccb,4423cdf0,43bffcfd",
+        "dispatch-bridge/Q8_0" to "n=4 fnv=f1f93362776a955c head=bf462edd,3f92c634,c151b4cb,c0e14735",
+
         // #1034 — the zero-copy packed transpose: a TensorView whose block axis moved, decoding
         // the same matrix the block-grid permutation in DefaultCpuOps.transpose produces.
         "transpose/Q4_0" to "n=384 fnv=272715de48d49929 head=3c442000,3df90000,be8ed800,be2e0000",

skainet-backends/skainet-backend-cpu/src/goldenTest/kotlin/sk/ainet/exec/golden/PackedDispatchBridgeGoldenTest.kt

@@ -0,0 +1,159 @@
+package sk.ainet.exec.golden
+
+import sk.ainet.backend.api.kernel.KernelDispatch
+import sk.ainet.backend.api.kernel.KernelPacks
+import sk.ainet.backend.api.kernel.PackedViewMatmulKernel
+import sk.ainet.backend.api.kernel.KernelRegistry
+import sk.ainet.exec.golden.GoldenSupport.Packed
+import sk.ainet.exec.kernel.ScalarKernelProvider
+import sk.ainet.lang.memory.BlockOrder
+import sk.ainet.lang.memory.ExperimentalMemoryApi
+import sk.ainet.lang.memory.PackedBlockDecoder
+import sk.ainet.lang.memory.Scope
+import sk.ainet.lang.memory.Storage
+import sk.ainet.lang.memory.TensorView
+import sk.ainet.lang.memory.trace.RecordingTraceSink
+import sk.ainet.lang.memory.trace.TraceEvent
+import sk.ainet.lang.tensor.Shape
+import sk.ainet.lang.tensor.data.Q4_0BlockTensorData
+import sk.ainet.lang.tensor.data.Q4_KBlockTensorData
+import sk.ainet.lang.tensor.data.Q5_0BlockTensorData
+import sk.ainet.lang.tensor.data.Q5_1BlockTensorData
+import sk.ainet.lang.tensor.data.Q5_KBlockTensorData
+import sk.ainet.lang.tensor.data.Q6_KBlockTensorData
+import sk.ainet.lang.tensor.data.Q8_0BlockTensorData
+import sk.ainet.lang.tensor.storage.PackedBlockStorage
+import sk.ainet.lang.tensor.storage.TensorEncoding
+import sk.ainet.lang.types.FP32
+import kotlin.test.AfterTest
+import kotlin.test.BeforeTest
+import kotlin.test.Test
+import kotlin.test.assertContentEquals
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+
+/**
+ * #1095 (#973.2): the packed SPI kernels reached through the registry produce **exactly** what
+ * calling them directly produces.
+ *
+ * This is the parity that #1029 could not assert, because nothing said which block order a view's
+ * bytes were in. The dispatcher now relayouts a canonical weight into kernel order and selects the
+ * packed kernel; the result must be bit-identical to feeding that kernel block-major bytes by hand,
+ * which is what `ScalarKernelGoldenTest` already pins. Every weight here is **three blocks wide**,
+ * the case where the two orders differ.
+ */
+@OptIn(ExperimentalMemoryApi::class)
+class PackedDispatchBridgeGoldenTest {
+
+    private companion object {
+        const val OUT = 4
+        const val BLOCKS = 3
+        const val SEED = 0x5EED_0002L      // the seed ScalarKernelGoldenTest uses
+    }
+
+    @BeforeTest fun setUp() {
+        KernelDispatch.clearForTesting()
+        KernelPacks.installReference()
+        KernelPacks.installPacked(ScalarKernelProvider)
+    }
+
+    @AfterTest fun tearDown() = KernelDispatch.clearForTesting()
+
+    private fun encodingOf(p: Packed): TensorEncoding = when (p) {
+        Packed.Q4_0 -> TensorEncoding.Q4_0
+        Packed.Q5_0 -> TensorEncoding.Q5_0
+        Packed.Q5_1 -> TensorEncoding.Q5_1
+        Packed.Q8_0 -> TensorEncoding.Q8_0
+        Packed.Q4_K -> TensorEncoding.Q4_K
+        Packed.Q5_K -> TensorEncoding.Q5_K
+        Packed.Q6_K -> TensorEncoding.Q6_K
+    }
+
+    private fun dataOf(p: Packed, shape: Shape, bytes: ByteArray): PackedBlockStorage = when (p) {
+        Packed.Q4_0 -> Q4_0BlockTensorData(shape, bytes)
+        Packed.Q5_0 -> Q5_0BlockTensorData(shape, bytes)
+        Packed.Q5_1 -> Q5_1BlockTensorData(shape, bytes)
+        Packed.Q8_0 -> Q8_0BlockTensorData(shape, bytes)
+        Packed.Q4_K -> Q4_KBlockTensorData(shape, bytes)
+        Packed.Q5_K -> Q5_KBlockTensorData(shape, bytes)
+        Packed.Q6_K -> Q6_KBlockTensorData(shape, bytes)
+    }
+
+    private fun bridge(p: Packed) {
+        val blocks = GoldenSupport.weightBlocks(p, OUT, BLOCKS, SEED)
+        val inputDim = BLOCKS * p.blockSize
+        val shape = Shape(OUT, inputDim)
+        val input = GoldenSupport.floats(inputDim, SEED + 100)
+
+        // what the kernel is fed by hand today: block-major bytes
+        val direct = FloatArray(OUT)
+        val spi = requireNotNull(spiKernel(p)) { "${p.name}: the scalar provider must offer this kernel" }
+        spi(input, 0, GoldenSupport.blockMajor(blocks), 0, inputDim, OUT, direct, 0)
+
+        // what the registry produces from a canonical view of the same weight
+        val canonicalBytes = GoldenSupport.rowMajor(blocks)
+        val weight = TensorView.packed(
+            Storage.Heap.wrap(canonicalBytes, mutable = false), shape, encodingOf(p),
+            PackedBlockDecoder(dataOf(p, shape, canonicalBytes)),
+        )
+        val activation = TensorView.dense(Storage.Heap.wrap(input), Shape(1, inputDim), FP32)
+        val out = TensorView.dense(Storage.Heap.floats(OUT), Shape(1, OUT), FP32)
+        val sink = RecordingTraceSink()
+        KernelDispatch.matmul(activation, weight, out, Scope.Ambient, sink, prepackWeights = true)
+
+        val kernelRun = sink.eventsOf<TraceEvent.KernelRun>().single()
+        assertTrue(
+            kernelRun.kernel.endsWith(p.name),
+            "${p.name}: the registry should select the packed kernel, ran '${kernelRun.kernel}'",
+        )
+        val adapter = sink.eventsOf<TraceEvent.AdapterInserted>().single()
+        assertEquals("prepack-input_block_major", adapter.kind, "${p.name}: the relayout must be visible")
+
+        val fromRegistry = FloatArray(OUT) { out.get(0, it) }
+        assertContentEquals(direct, fromRegistry, "${p.name}: registry vs direct call must be bit-identical")
+        GoldenSupport.check("dispatch-bridge/${p.name}", GoldenSupport.digest(fromRegistry))
+    }
+
+    private fun spiKernel(p: Packed): ((FloatArray, Int, ByteArray, Int, Int, Int, FloatArray, Int) -> Unit)? = when (p) {
+        Packed.Q4_0 -> ScalarKernelProvider.matmulQ4_0()?.let { k -> { i, io, w, wo, id, od, o, oo -> k.matmul(i, io, w, wo, id, od, o, oo) } }
+        Packed.Q5_0 -> ScalarKernelProvider.matmulQ5_0()?.let { k -> { i, io, w, wo, id, od, o, oo -> k.matmul(i, io, w, wo, id, od, o, oo) } }
+        Packed.Q5_1 -> ScalarKernelProvider.matmulQ5_1()?.let { k -> { i, io, w, wo, id, od, o, oo -> k.matmul(i, io, w, wo, id, od, o, oo) } }
+        Packed.Q8_0 -> ScalarKernelProvider.matmulQ8_0()?.let { k -> { i, io, w, wo, id, od, o, oo -> k.matmul(i, io, w, wo, id, od, o, oo) } }
+        Packed.Q4_K -> ScalarKernelProvider.matmulQ4K()?.let { k -> { i, io, w, wo, id, od, o, oo -> k.matmul(i, io, w, wo, id, od, o, oo) } }
+        Packed.Q5_K -> ScalarKernelProvider.matmulQ5K()?.let { k -> { i, io, w, wo, id, od, o, oo -> k.matmul(i, io, w, wo, id, od, o, oo) } }
+        Packed.Q6_K -> ScalarKernelProvider.matmulQ6K()?.let { k -> { i, io, w, wo, id, od, o, oo -> k.matmul(i, io, w, wo, id, od, o, oo) } }
+    }
+
+    @Test fun q4_0() = bridge(Packed.Q4_0)
+    @Test fun q5_0() = bridge(Packed.Q5_0)
+    @Test fun q5_1() = bridge(Packed.Q5_1)
+    @Test fun q8_0() = bridge(Packed.Q8_0)
+    @Test fun q4_K() = bridge(Packed.Q4_K)
+    @Test fun q5_K() = bridge(Packed.Q5_K)
+    @Test fun q6_K() = bridge(Packed.Q6_K)
+
+    @Test
+    fun aWeightAlreadyInKernelOrderCostsNoAdapter() {
+        // the shape a caller should hand over in a hot loop: prepacked once, at load (#1097)
+        val p = Packed.Q8_0
+        val blocks = GoldenSupport.weightBlocks(p, OUT, BLOCKS, SEED)
+        val inputDim = BLOCKS * p.blockSize
+        val shape = Shape(OUT, inputDim)
+        val canonicalBytes = GoldenSupport.rowMajor(blocks)
+        val weight = TensorView.packed(
+            Storage.Heap.wrap(canonicalBytes, mutable = false), shape, encodingOf(p),
+            PackedBlockDecoder(dataOf(p, shape, canonicalBytes)),
+        ).prepack(BlockOrder.INPUT_BLOCK_MAJOR)
+
+        val input = GoldenSupport.floats(inputDim, SEED + 100)
+        val activation = TensorView.dense(Storage.Heap.wrap(input), Shape(1, inputDim), FP32)
+        val out = TensorView.dense(Storage.Heap.floats(OUT), Shape(1, OUT), FP32)
+        val sink = RecordingTraceSink()
+        KernelDispatch.matmul(activation, weight, out, Scope.Ambient, sink)
+
+        assertTrue(sink.eventsOf<TraceEvent.AdapterInserted>().isEmpty(), "nothing to relayout: it is already in kernel order")
+        val direct = FloatArray(OUT)
+        spiKernel(p)!!(input, 0, GoldenSupport.blockMajor(blocks), 0, inputDim, OUT, direct, 0)
+        assertContentEquals(direct, FloatArray(OUT) { out.get(0, it) })
+    }
+}