From 2af6251542ceaaad84a8898669555da8f89dfdb8 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Mon, 24 Aug 2026 18:44:56 +0200 Subject: [PATCH] feat(backend): bridge the packed SPI kernels through the ordered key MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #1095 (#973.2). Unblocks what #1029 deferred. The packed matmul kernels read their weight input-block-major; a packed `TensorView` loaded from a file is canonical. Until #1094 nothing said which was which, so #1029 left them unbridged rather than risk the silent-wrong-numbers class of #968/#971. Now the order is in the key. - `PackedViewMatmulKernel`: a `ViewKernel` over any of the packed SPI kernels, keyed on `BLOCKED_INPUT_MAJOR` — the order it actually reads, checked at run time rather than assumed. Falls back to the reference for a strided activation or non-heap storage instead of mis-indexing. - `KernelPacks.installPacked(provider)` registers Q4_0, Q5_0, Q5_1, Q8_0, Q4_K, Q5_K and Q6_K when the provider offers them; a format it does not offer keeps being served by the decoding reference. - `KernelDispatch.matmul(..., prepackWeights = false)` can relayout a canonical weight to reach those kernels — **off by default, on purpose**. The relayout is O(bytes), so doing it inside a decode step would copy the whole weight per token, which is the per-forward copy #973 objects to, merely moved. Wiring it on by default broke M1-A3 (zero forward-scope allocations per step) in exactly that way, which is how the default got chosen. A weight is prepacked once at load; a canonical weight handed straight to the dispatcher still gets the correct, slower reference path. `PackedDispatchBridgeGoldenTest`: for all seven encodings, the registry path is **bit-identical** to calling the SPI kernel by hand with block-major bytes — and the recorded digests' leading values are the existing `scalar-matmul/*` goldens, which is the point: the bridge changes nothing about the numbers. Every weight is three blocks wide, the case where the two orders differ. A weight already in kernel order costs no adapter at all. Gate: scripts/pr-gate.sh — all legs passed; --golden passed (46 golden tests, the 39 existing digests unchanged). Co-Authored-By: Claude Fable 5 --- .../backend/api/kernel/KernelDispatch.kt | 26 +++ .../ainet/backend/api/kernel/KernelPacks.kt | 39 ++++- .../api/kernel/PackedViewMatmulKernel.kt | 87 ++++++++++ .../kotlin/sk/ainet/exec/golden/Goldens.kt | 11 ++ .../golden/PackedDispatchBridgeGoldenTest.kt | 159 ++++++++++++++++++ 5 files changed, 313 insertions(+), 9 deletions(-) create mode 100644 skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/PackedViewMatmulKernel.kt create mode 100644 skainet-backends/skainet-backend-cpu/src/goldenTest/kotlin/sk/ainet/exec/golden/PackedDispatchBridgeGoldenTest.kt diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelDispatch.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelDispatch.kt index f7cd6cef6..8d5f8a9ef 100644 --- a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelDispatch.kt +++ b/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") diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelPacks.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelPacks.kt index 173df2ce5..60a748558 100644 --- a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/KernelPacks.kt +++ b/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. */ diff --git a/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/PackedViewMatmulKernel.kt b/skainet-backends/skainet-backend-api/src/commonMain/kotlin/sk/ainet/backend/api/kernel/PackedViewMatmulKernel.kt new file mode 100644 index 000000000..c4b968072 --- /dev/null +++ b/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, 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, 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 = emptySet()): KernelKey = KernelKey( + op = "matmul", + operands = listOf( + OperandKey.contiguous(Format.dense(FP32)), + OperandKey(Format(FP32, encoding), LayoutClass.BLOCKED_INPUT_MAJOR), + ), + capabilities = capabilities, + ) + } +} diff --git a/skainet-backends/skainet-backend-cpu/src/goldenTest/kotlin/sk/ainet/exec/golden/Goldens.kt b/skainet-backends/skainet-backend-cpu/src/goldenTest/kotlin/sk/ainet/exec/golden/Goldens.kt index 30e6ebdb9..4f1939209 100644 --- a/skainet-backends/skainet-backend-cpu/src/goldenTest/kotlin/sk/ainet/exec/golden/Goldens.kt +++ b/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 = 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", diff --git a/skainet-backends/skainet-backend-cpu/src/goldenTest/kotlin/sk/ainet/exec/golden/PackedDispatchBridgeGoldenTest.kt b/skainet-backends/skainet-backend-cpu/src/goldenTest/kotlin/sk/ainet/exec/golden/PackedDispatchBridgeGoldenTest.kt new file mode 100644 index 000000000..246a47f91 --- /dev/null +++ b/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().single() + assertTrue( + kernelRun.kernel.endsWith(p.name), + "${p.name}: the registry should select the packed kernel, ran '${kernelRun.kernel}'", + ) + val adapter = sink.eventsOf().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().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) }) + } +}