From 7c8e25c8394b1df41bb1c465ef9fb8da798e19ed Mon Sep 17 00:00:00 2001 From: Michal Harakal <471011+michalharakal@users.noreply.github.com> Date: Thu, 5 Nov 2026 16:38:05 +0100 Subject: [PATCH] feat(compile): export Q8_0 / Q4_0 Linear weights as a block-quant linalg chain Keep packed weights packed through export so an NPU compiler can match them to a block-quantized matmul kernel instead of a dense FP32 copy. - PackedConstantHandling.KEEP_PACKED: the weight node carries its PackedBlockStorage; HloGenerator.generate takes packedConstants. - BlockQuantMatmulFusionPass: folds weight -> transpose -> matmul (or a transposeB matmul) into a block_quant_matmul node; skips are reported. - BlockQuantMatmulConverter: emits (sitofp(codes) - zp) * scale over [nB, 32, N] i8/i4 codes with bf16 scales, then linalg.batch_matmul, constants inline. A negative block scale d is stored exactly as codes -1 - c, zero point -1, scale -d. A packed weight left unfolded fails the export. On an SL2610, Torq 2.3.0 with --iree-input-type=stablehlo-torq compiles the emitted Q8_0 and Q4_0 projections to mixed_quant_matmul with no generic fallback; outputs match the dequantized weights at cosine 0.999996 / 0.999997. TorqBlockQuantGoldenTest checks the kernel match when SKAINET_TORQ_COMPILE is set. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 1 + .../explanation/compile-flexibility.adoc | 27 ++ .../api/jvm/skainet-compile-dag.api | 6 + .../lang/trace/PackedConstantHandling.kt | 15 ++ .../ainet/lang/trace/TraceToGraphBuilder.kt | 14 +- .../api/jvm/skainet-compile-hlo.api | 10 +- .../compile/hlo/StableHloConverterFactory.kt | 7 + .../converters/BlockQuantMatmulConverter.kt | 252 ++++++++++++++++++ .../converters/ConstantOperationsConverter.kt | 14 + .../compile/hlo/generate/HloGenerator.kt | 9 +- .../compile/hlo/BlockQuantMatmulExportTest.kt | 149 +++++++++++ .../compile/hlo/TorqBlockQuantGoldenTest.kt | 99 +++++++ .../api/jvm/skainet-compile-opt.api | 13 + .../opt/passes/BlockQuantMatmulFusionPass.kt | 147 ++++++++++ 14 files changed, 756 insertions(+), 7 deletions(-) create mode 100644 skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/BlockQuantMatmulConverter.kt create mode 100644 skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/BlockQuantMatmulExportTest.kt create mode 100644 skainet-compile/skainet-compile-hlo/src/jvmTest/kotlin/sk/ainet/compile/hlo/TorqBlockQuantGoldenTest.kt create mode 100644 skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/BlockQuantMatmulFusionPass.kt diff --git a/CHANGELOG.md b/CHANGELOG.md index 145a3135..f85744a7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ ### Added +- **Block-quantized weights through StableHLO export** (`skainet-compile-dag`, `skainet-compile-opt`, `skainet-compile-hlo`): `PackedConstantHandling.KEEP_PACKED` keeps a packed weight as its storage instead of failing or dequantizing; `BlockQuantMatmulFusionPass` folds a Q8_0 / Q4_0 weight used as `x · Wᵀ` into a `block_quant_matmul` node; `BlockQuantMatmulConverter` (registered in every converter set) emits it as `(sitofp(codes) - zp) * scale` over `[nB, 32, N]` `i8` / `i4` codes with `bf16` scales, then `linalg.batch_matmul`, all constants inline. A packed weight left unfolded fails the export. On an SL2610, Torq 2.3.0 with `--iree-input-type=stablehlo-torq` compiles the emitted Q8_0 and Q4_0 projections to its `mixed_quant_matmul` kernel and matches the dequantized weight (cosine 0.999996); `TorqBlockQuantGoldenTest` checks the match when `SKAINET_TORQ_COMPILE` names the compiler. `HloGenerator.generate` takes `packedConstants`. - **StableHLO-phase target passes** (`skainet-compile-opt`, `skainet-compile-hlo`): `TargetOptimizer.stableHloPasses()` and `TargetOptimizers.registerStableHloPasses(target) { … }` / `stableHloPassesFor(target)` wire the `STABLE_HLO` phase. A `StableHloPass` rewrites the emitted module's text into valid StableHLO; `toStableHlo(graph, …, passes = …)` and `StableHloModule.applyPasses(passes)` apply them in order. The existing `toStableHlo` signatures are unchanged. - **TOSA preparation passes** (`sk.ainet.compile.hlo.tosa`): `DivideByBroadcastPass` (a divide by a broadcast becomes a multiply by the broadcast reciprocal, or by `rsqrt` for a `sqrt`) and `RankTwoWeightMatmulPass` (`[1, M, K] · [K, N]` as a rank-2 dot) shape a module for StableHLO's own `--stablehlo-prepare-for-tosa --stablehlo-legalize-to-tosa`; `TosaFinish` translates what that legalization leaves (`broadcast_in_dim`, `rsqrt`, batched attention `dot_general`s) and throws `TosaFinishException` on anything else. Moonshine v2 tiny streaming encoder (64 frames) on an SL2610 with the Torq NPU's `tosa-torq` input: one dispatch per call, golden transcripts identical to the recipe-side rewrite they replace. diff --git a/docs/modules/ROOT/pages/explanation/compile-flexibility.adoc b/docs/modules/ROOT/pages/explanation/compile-flexibility.adoc index 5bd56c86..1b1a502a 100644 --- a/docs/modules/ROOT/pages/explanation/compile-flexibility.adoc +++ b/docs/modules/ROOT/pages/explanation/compile-flexibility.adoc @@ -100,6 +100,33 @@ carries what goes around it: The legalization runs outside SKaiNET (`stablehlo-opt`, built from the same LLVM as the compiler that reads its output). +=== Block-quantized weights for an NPU kernel + +A Q8_0 or Q4_0 `Linear` can reach the compiler as integer codes and per-block scales instead of +a dense FP32 copy. Three pieces carry it: + +* the trace keeps the weight packed: `toComputeGraph(…, packedConstants = KEEP_PACKED)` (or + `HloGenerator.generate(…, packedConstants = …)`); +* `BlockQuantMatmulFusionPass` (DAG phase) folds `weight → transpose → matmul` into one + `block_quant_matmul` node that holds the packed storage; +* `BlockQuantMatmulConverter` emits that node as a `linalg` chain: a `linalg.generic` computing + `(sitofp(codes) - zp) * scale` over `[nB, 32, N]` (`i8` or `i4` codes, `bf16` scales broadcast + over each 32-element block), then `collapse_shape`, `expand_shape` and `linalg.batch_matmul` + with a `bf16` activation. All constants are inline. + +[source,kotlin] +---- +TargetOptimizers.registerDagPasses("torq") { listOf(BlockQuantMatmulFusionPass()) } +---- + +This is the form the Torq compiler (2.3.0, `--iree-input-type=stablehlo-torq`) matches to its +`mixed_quant_matmul` kernel when the activation has one row. That kernel takes only positive +scales, so a block with a negative scale `d` is stored exactly as codes `-1 - c`, zero point +`-1` and scale `-d` (ggml Q4_0 sets `d = max / -8`, about half its blocks are negative). A +packed weight the pass did not fold, such as a tied embedding, makes the export fail; trace +with `DEQUANTIZE` for those. Stock IREE `llvm-cpu` has given wrong results for `i4` weights, so +keep Q4_0 to targets that were checked. + == Why exact GELU needs care [cols="1,3"] diff --git a/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api b/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api index 0cee8f6a..2e522281 100644 --- a/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api +++ b/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api @@ -530,13 +530,19 @@ public final class sk/ainet/lang/trace/PackedConstantException : java/lang/Illeg } public final class sk/ainet/lang/trace/PackedConstantHandling : java/lang/Enum { + public static final field Companion Lsk/ainet/lang/trace/PackedConstantHandling$Companion; public static final field DEQUANTIZE Lsk/ainet/lang/trace/PackedConstantHandling; public static final field FAIL Lsk/ainet/lang/trace/PackedConstantHandling; + public static final field KEEP_PACKED Lsk/ainet/lang/trace/PackedConstantHandling; + public static final field PACKED_WEIGHT_PARAM Ljava/lang/String; public static fun getEntries ()Lkotlin/enums/EnumEntries; public static fun valueOf (Ljava/lang/String;)Lsk/ainet/lang/trace/PackedConstantHandling; public static fun values ()[Lsk/ainet/lang/trace/PackedConstantHandling; } +public final class sk/ainet/lang/trace/PackedConstantHandling$Companion { +} + public final class sk/ainet/lang/trace/TapeSink : sk/ainet/lang/trace/OpSink { public fun (Lsk/ainet/tape/ExecutionTape;)V public fun onOpExecuted (Lsk/ainet/lang/trace/OpTrace;)V diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/PackedConstantHandling.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/PackedConstantHandling.kt index dfcf345e..0a02ba46 100644 --- a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/PackedConstantHandling.kt +++ b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/PackedConstantHandling.kt @@ -25,6 +25,21 @@ public enum class PackedConstantHandling { * use only when the export target genuinely needs dense constants. */ DEQUANTIZE, + + /** + * Keep the packed data as is: the weight node carries the + * [sk.ainet.lang.tensor.storage.PackedBlockStorage] under [PACKED_WEIGHT_PARAM] + * and keeps its encoding, with no float copy. Something downstream must + * consume it: `BlockQuantMatmulFusionPass` folds a Q8_0 / Q4_0 weight into + * the matmul that reads it, and the StableHLO emitter refuses a packed + * weight that is still a plain constant. + */ + KEEP_PACKED; + + public companion object { + /** Weight-node parameter holding the packed storage under [KEEP_PACKED]. */ + public const val PACKED_WEIGHT_PARAM: String = "packed_weight" + } } /** diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/TraceToGraphBuilder.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/TraceToGraphBuilder.kt index fad7331d..230f1018 100644 --- a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/TraceToGraphBuilder.kt +++ b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/trace/TraceToGraphBuilder.kt @@ -292,26 +292,32 @@ public class TraceToGraphBuilder( // packed encoding forward would misdescribe it. encoding = null } + PackedConstantHandling.KEEP_PACKED -> Unit } } + val keptPacked = if (constantValues == null && packedConstants == PackedConstantHandling.KEEP_PACKED) packedData else null val syntheticNode: GraphNode val producedSpec: TensorSpec - if (constantValues != null) { - // Create a constant/weight node with embedded values + if (constantValues != null || keptPacked != null) { + // Create a constant/weight node with embedded values (or, under + // KEEP_PACKED, the packed storage itself) val weightShape = tensor!!.shape.dimensions.toList() val weightDtype = tensor.dtype.simpleName ?: "FP32" val nodeId = "n${nextNodeId++}_weight" val op = TraceBackedOperation( name = "weight", type = "constant", - parameters = mapOf( + parameters = if (keptPacked != null) mapOf( + PackedConstantHandling.PACKED_WEIGHT_PARAM to keptPacked, + "trainable" to false + ) else mapOf( // Store the primitive FloatArray, NOT .toList(): boxing a // real LLM weight (e.g. 262153x640 embedding) into a // List is ~2.7GB and OOMs the trace. The HLO // converter handles FloatArray for both inline and // external (.irpa) materialization. - "initial_value" to constantValues, + "initial_value" to constantValues!!, "trainable" to false ) ) diff --git a/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api b/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api index ca768bc4..5aa9bdf9 100644 --- a/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api +++ b/skainet-compile/skainet-compile-hlo/api/jvm/skainet-compile-hlo.api @@ -398,6 +398,12 @@ public final class sk/ainet/compile/hlo/converters/BasicMathConverter : sk/ainet public fun getSupportedOperations ()Ljava/util/Set; } +public final class sk/ainet/compile/hlo/converters/BlockQuantMatmulConverter : sk/ainet/compile/hlo/StableHloOperationConverter { + public fun ()V + public fun convert (Lsk/ainet/lang/graph/GraphNode;Ljava/util/List;Lsk/ainet/compile/hlo/ConversionContext;)Lsk/ainet/compile/hlo/ConversionResult; + public fun getSupportedOperations ()Ljava/util/Set; +} + public final class sk/ainet/compile/hlo/converters/ConstantOperationsConverter : sk/ainet/compile/hlo/StableHloOperationConverter { public fun ()V public fun convert (Lsk/ainet/lang/graph/GraphNode;Ljava/util/List;Lsk/ainet/compile/hlo/ConversionContext;)Lsk/ainet/compile/hlo/ConversionResult; @@ -570,8 +576,8 @@ public final class sk/ainet/compile/hlo/examples/ValidationDemonstrationResult { public final class sk/ainet/compile/hlo/generate/HloGenerator { public static final field INSTANCE Lsk/ainet/compile/hlo/generate/HloGenerator; public final fun corePasses (Ljava/lang/String;)Ljava/util/List; - public final fun generate (Lsk/ainet/lang/model/Model;Lsk/ainet/lang/tensor/Tensor;Ljava/lang/String;Ljava/lang/String;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; - public static synthetic fun generate$default (Lsk/ainet/compile/hlo/generate/HloGenerator;Lsk/ainet/lang/model/Model;Lsk/ainet/lang/tensor/Tensor;Ljava/lang/String;Ljava/lang/String;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object; + public final fun generate (Lsk/ainet/lang/model/Model;Lsk/ainet/lang/tensor/Tensor;Ljava/lang/String;Ljava/lang/String;Lsk/ainet/lang/trace/PackedConstantHandling;Lkotlin/coroutines/Continuation;)Ljava/lang/Object; + public static synthetic fun generate$default (Lsk/ainet/compile/hlo/generate/HloGenerator;Lsk/ainet/lang/model/Model;Lsk/ainet/lang/tensor/Tensor;Ljava/lang/String;Ljava/lang/String;Lsk/ainet/lang/trace/PackedConstantHandling;Lkotlin/coroutines/Continuation;ILjava/lang/Object;)Ljava/lang/Object; } public final class sk/ainet/compile/hlo/generate/HloGeneratorJvmKt { diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverterFactory.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverterFactory.kt index 2d768826..2b4530a8 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverterFactory.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverterFactory.kt @@ -2,6 +2,7 @@ package sk.ainet.compile.hlo import sk.ainet.compile.hlo.converters.ActivationOperationsConverter import sk.ainet.compile.hlo.converters.ArgMaxOperationsConverter +import sk.ainet.compile.hlo.converters.BlockQuantMatmulConverter import sk.ainet.compile.hlo.converters.AttentionOperationsConverter import sk.ainet.compile.hlo.converters.ConstantOperationsConverter import sk.ainet.compile.hlo.converters.GatherOperationsConverter @@ -81,6 +82,8 @@ public object StableHloConverterFactory { // Register constant operations converter registry.register(ConstantOperationsConverter()) + // block_quant_matmul exists only where BlockQuantMatmulFusionPass ran; packed weights stay packed + registry.register(BlockQuantMatmulConverter()) // Register attention (scaledDotProductAttention) converter registry.register(AttentionOperationsConverter()) @@ -141,6 +144,8 @@ public object StableHloConverterFactory { // Register constant operations converter registry.register(ConstantOperationsConverter()) + // block_quant_matmul exists only where BlockQuantMatmulFusionPass ran; packed weights stay packed + registry.register(BlockQuantMatmulConverter()) // Register attention (scaledDotProductAttention) converter registry.register(AttentionOperationsConverter()) @@ -178,6 +183,8 @@ public object StableHloConverterFactory { registry.register(UnaryMathConverter()) registry.register(ScalarOperationsConverter()) registry.register(ConstantOperationsConverter()) + // block_quant_matmul exists only where BlockQuantMatmulFusionPass ran; packed weights stay packed + registry.register(BlockQuantMatmulConverter()) return StableHloConverter(registry, typeMapper, null, policy, errorPolicy = errorPolicy) } diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/BlockQuantMatmulConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/BlockQuantMatmulConverter.kt new file mode 100644 index 00000000..19d6e0c7 --- /dev/null +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/BlockQuantMatmulConverter.kt @@ -0,0 +1,252 @@ +package sk.ainet.compile.hlo.converters + +import sk.ainet.compile.hlo.ConversionContext +import sk.ainet.compile.hlo.ConversionResult +import sk.ainet.compile.hlo.StableHloOperationConverter +import sk.ainet.compile.opt.passes.BlockQuantMatmulFusionPass +import sk.ainet.lang.graph.GraphNode +import sk.ainet.lang.tensor.data.Q4_0TensorData +import sk.ainet.lang.tensor.data.Q8_0TensorData +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.tensor.storage.TensorEncoding + +/** + * Lowers `block_quant_matmul` (made by [BlockQuantMatmulFusionPass]) to a linalg chain that keeps + * the weight as integer codes plus per-block scales: + * ``` + * w[nB, 32, N] = (sitofp(codes) - zp) * scale linalg.generic, scale/zp broadcast over the block + * y[1, M, N] = batch_matmul(x[1, M, K], expand(collapse(w))) + * ``` + * Codes are `i8` for Q8_0 and `i4` for Q4_0; scales and the optional zero point are `bf16` + * `[nB, 1, N]`; all of them are inline constants. The activation is converted to `bf16` and + * reshaped to `[1, M, K]` around the chain, and the result is brought back to the node's type. + * + * This is the form the Torq compiler (2.3.0, `--iree-input-type=stablehlo-torq`) matches to its + * `mixed_quant_matmul` kernel when `M = 1`. The kernel accepts only positive scales, so a block + * with a negative scale `d` is stored exactly as codes `-1 - c`, zero point `-1` and scale `-d`; + * the zero-point tensor is emitted only when such a block exists. The chain is portable linalg, + * but stock IREE `llvm-cpu` has given wrong results for `i4` weights, so keep Q4_0 for targets + * that were checked. + */ +public class BlockQuantMatmulConverter : StableHloOperationConverter { + + override val supportedOperations: Set = setOf(BlockQuantMatmulFusionPass.OP_NAME) + + override fun convert(node: GraphNode, operands: List, context: ConversionContext): ConversionResult { + val storage = node.operation.parameters[BlockQuantMatmulFusionPass.WEIGHT] as? PackedBlockStorage + ?: return ConversionResult.Failure("block_quant_matmul ${node.id} has no packed weight", "missing weight on ${node.id}") + if (operands.size != 1) { + return ConversionResult.Failure("block_quant_matmul takes one operand, got ${operands.size}", "bad arity on ${node.id}") + } + val inSpec = node.inputs.firstOrNull() + val outSpec = node.outputs.firstOrNull() + val inShape = inSpec?.shape + val outShape = outSpec?.shape + if (inShape.isNullOrEmpty() || outShape.isNullOrEmpty()) { + return ConversionResult.Failure("block_quant_matmul ${node.id} needs static input and output shapes", "unknown shape on ${node.id}") + } + val weight = try { + BlockQuantWeight.of(storage) + } catch (e: IllegalArgumentException) { + return ConversionResult.Failure(e.message ?: "unsupported weight", "weight of ${node.id}") + } + val k = weight.k + val n = weight.n + val m = inShape.dropLast(1).fold(1) { acc, d -> acc * d } + if (inShape.last() != k || outShape.last() != n || outShape.dropLast(1).fold(1) { acc, d -> acc * d } != m) { + return ConversionResult.Failure( + "block_quant_matmul ${node.id}: x $inShape · W[$n, $k]ᵀ does not give $outShape", + "shape mismatch on ${node.id}" + ) + } + + val types = context.getTypeMapper() + val inElem = types.mapDType(inSpec.dtype) + val outElem = types.mapDType(outSpec.dtype) + val nB = weight.blockCount + val b = BlockQuantWeight.BLOCK + val codeType = "i${weight.bits}" + fun t(vararg dims: Int, elem: String = "bf16") = "tensor<${dims.joinToString("x")}x$elem>" + fun tt(dims: List, elem: String) = "tensor<${dims.joinToString("x")}x$elem>" + val ops = mutableListOf() + fun emit(op: String) { context.emitOperation(op); ops += op } + + val weightName = node.operation.parameters[BlockQuantMatmulFusionPass.WEIGHT_NAME] ?: node.id + context.emitComment("block_quant_matmul ${node.id}: $weightName ${storage.encoding.name} [$n, $k] as $codeType codes + bf16 scales") + + var x = operands[0] + if (inElem != "bf16") { + val v = context.nextTempValue() + emit("$v = stablehlo.convert $x : (${tt(inShape, inElem)}) -> ${tt(inShape, "bf16")}") + x = v + } + if (inShape != listOf(1, m, k)) { + val v = context.nextTempValue() + emit("$v = stablehlo.reshape $x : (${tt(inShape, "bf16")}) -> ${t(1, m, k)}") + x = v + } + + val codes = context.nextTempValue() + emit("$codes = arith.constant dense<\"${hex(weight.codes)}\"> : ${t(nB, b, n, elem = codeType)}") + val scales = context.nextTempValue() + emit("$scales = arith.constant dense<\"${hex(weight.scales)}\"> : ${t(nB, 1, n)}") + val zp = weight.zeroPoints?.let { + val v = context.nextTempValue() + emit("$v = arith.constant dense<\"${hex(it)}\"> : ${t(nB, 1, n)}") + v + } + + val init = context.nextTempValue() + emit("$init = tensor.empty() : ${t(nB, b, n)}") + val identity = "affine_map<(d0, d1, d2) -> (d0, d1, d2)>" + val perBlock = "affine_map<(d0, d1, d2) -> (d0, 0, d2)>" + val maps = if (zp == null) listOf(identity, perBlock, identity) else listOf(identity, perBlock, perBlock, identity) + val ins = if (zp == null) "$codes, $scales : ${t(nB, b, n, elem = codeType)}, ${t(nB, 1, n)}" + else "$codes, $scales, $zp : ${t(nB, b, n, elem = codeType)}, ${t(nB, 1, n)}, ${t(nB, 1, n)}" + val (fl, df, r) = Triple(context.nextTempValue(), context.nextTempValue(), context.nextTempValue()) + val body = if (zp == null) { + "^bb0(%c: $codeType, %s: bf16, %o: bf16):\n" + + " $fl = arith.sitofp %c : $codeType to bf16\n" + + " $r = arith.mulf $fl, %s : bf16\n" + } else { + "^bb0(%c: $codeType, %s: bf16, %z: bf16, %o: bf16):\n" + + " $fl = arith.sitofp %c : $codeType to bf16\n" + + " $df = arith.subf $fl, %z : bf16\n" + + " $r = arith.mulf $df, %s : bf16\n" + } + val dq = context.nextTempValue() + emit( + "$dq = linalg.generic {indexing_maps = [${maps.joinToString(", ")}], " + + "iterator_types = [\"parallel\", \"parallel\", \"parallel\"]} ins($ins) outs($init : ${t(nB, b, n)}) {\n" + + " $body" + + " linalg.yield $r : bf16\n" + + " } -> ${t(nB, b, n)}" + ) + val collapsed = context.nextTempValue() + emit("$collapsed = tensor.collapse_shape $dq [[0, 1], [2]] : ${t(nB, b, n)} into ${t(k, n)}") + val expanded = context.nextTempValue() + emit("$expanded = tensor.expand_shape $collapsed [[0, 1], [2]] output_shape [1, $k, $n] : ${t(k, n)} into ${t(1, k, n)}") + val zero = context.nextTempValue() + emit("$zero = arith.constant 0.000000e+00 : bf16") + val acc0 = context.nextTempValue() + emit("$acc0 = tensor.empty() : ${t(1, m, n)}") + val acc = context.nextTempValue() + emit("$acc = linalg.fill ins($zero : bf16) outs($acc0 : ${t(1, m, n)}) -> ${t(1, m, n)}") + var y = context.nextTempValue() + emit("$y = linalg.batch_matmul ins($x, $expanded : ${t(1, m, k)}, ${t(1, k, n)}) outs($acc : ${t(1, m, n)}) -> ${t(1, m, n)}") + + if (outShape != listOf(1, m, n)) { + val v = context.nextTempValue() + emit("$v = stablehlo.reshape $y : (${t(1, m, n)}) -> ${tt(outShape, "bf16")}") + y = v + } + if (outElem != "bf16") { + val v = context.nextTempValue() + emit("$v = stablehlo.convert $y : (${tt(outShape, "bf16")}) -> ${tt(outShape, outElem)}") + y = v + } + return ConversionResult.Success(outputValueName = y, emittedOperations = ops) + } + + private companion object { + private val HEX = "0123456789ABCDEF".toCharArray() + + fun hex(bytes: ByteArray): String { + val out = CharArray(2 + bytes.size * 2) + out[0] = '0'; out[1] = 'x' + for (i in bytes.indices) { + val v = bytes[i].toInt() and 0xFF + out[2 + 2 * i] = HEX[v ushr 4] + out[3 + 2 * i] = HEX[v and 0xF] + } + return out.concatToString() + } + + /** bf16 values as little-endian bytes, the order MLIR reads a hex dense attribute in. */ + fun hex(bf16: ShortArray): String { + val bytes = ByteArray(bf16.size * 2) + for (i in bf16.indices) { + bytes[2 * i] = bf16[i].toByte() + bytes[2 * i + 1] = (bf16[i].toInt() ushr 8).toByte() + } + return hex(bytes) + } + } +} + +/** + * A Q8_0 or Q4_0 weight `[N, K]` regrouped for the block-quant chain: `codes[b][j][o]` is element + * `(o, 32b + j)`, `scales[b][o]` and `zeroPoints[b][o]` belong to block `b` of output row `o`. + * The values it stands for are `(codes - zeroPoints) * scales`. + */ +internal class BlockQuantWeight( + val k: Int, + val n: Int, + /** Code width: 8 for Q8_0, 4 for Q4_0 (one code per byte, sign-extended). */ + val bits: Int, + /** `[nB, 32, N]` signed codes. */ + val codes: ByteArray, + /** `[nB, N]` bf16 bits, all positive. */ + val scales: ShortArray, + /** `[nB, N]` bf16 bits, 0 or -1; `null` when every block's scale was non-negative. */ + val zeroPoints: ShortArray?, +) { + val blockCount: Int get() = k / BLOCK + + companion object { + const val BLOCK: Int = 32 + private const val BF16_ONE: Short = 0x3F80 + private const val BF16_MINUS_ONE: Short = 0xBF80.toShort() + + fun of(storage: PackedBlockStorage): BlockQuantWeight { + require(storage.shape.rank == 2) { "block-quant weight must be [out, in], got ${storage.shape.dimensions.toList()}" } + val n = storage.shape[0] + val k = storage.shape[1] + require(k % BLOCK == 0) { "block-quant weight in = $k is not a multiple of $BLOCK" } + val (bits, scaleOf, codeOf) = when (storage.encoding) { + TensorEncoding.Q8_0 -> { + val q = storage as? Q8_0TensorData ?: throw IllegalArgumentException("Q8_0 storage ${storage::class.simpleName} has no block accessors") + Triple(8, q::getBlockScale, { blk: Int, j: Int -> q.getCode(blk, j).toInt() }) + } + TensorEncoding.Q4_0 -> { + val q = storage as? Q4_0TensorData ?: throw IllegalArgumentException("Q4_0 storage ${storage::class.simpleName} has no block accessors") + Triple(4, q::getBlockScale, { blk: Int, j: Int -> q.getCode(blk, j).toInt() - 8 }) + } + else -> throw IllegalArgumentException("block-quant matmul supports Q8_0 and Q4_0, not ${storage.encoding.name}") + } + val nB = k / BLOCK + val codes = ByteArray(nB * BLOCK * n) + val scales = ShortArray(nB * n) + var zeroPoints: ShortArray? = null + for (o in 0 until n) { + for (b in 0 until nB) { + val blk = o * nB + b + val d = scaleOf(blk) + val at = b * n + o + when { + d > 0f -> { + scales[at] = bf16(d) + for (j in 0 until BLOCK) codes[(b * BLOCK + j) * n + o] = codeOf(blk, j).toByte() + } + d < 0f -> { + // (c - 0) * d == (-1 - c - (-1)) * -d, and -1 - c stays in range for i8 and i4. + scales[at] = bf16(-d) + val zp = zeroPoints ?: ShortArray(nB * n).also { zeroPoints = it } + zp[at] = BF16_MINUS_ONE + for (j in 0 until BLOCK) codes[(b * BLOCK + j) * n + o] = (-1 - codeOf(blk, j)).toByte() + } + // An all-zero block: zero codes, and a scale the kernel accepts (it rejects 0). + else -> scales[at] = BF16_ONE + } + } + } + return BlockQuantWeight(k, n, bits, codes, scales, zeroPoints) + } + + /** FP32 → bf16 bits, rounded to nearest even. */ + fun bf16(value: Float): Short { + val bits = value.toRawBits() + return ((bits + 0x7FFF + ((bits ushr 16) and 1)) ushr 16).toShort() + } + } +} diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ConstantOperationsConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ConstantOperationsConverter.kt index 01c30d6b..0b880581 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ConstantOperationsConverter.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ConstantOperationsConverter.kt @@ -10,10 +10,12 @@ import sk.ainet.compile.hlo.elementCountFromShape import sk.ainet.compile.hlo.floatArrayToLittleEndianBytes import sk.ainet.compile.hlo.numberListToLittleEndianBytes import sk.ainet.lang.graph.GraphNode +import sk.ainet.lang.trace.PackedConstantHandling import sk.ainet.lang.tensor.ops.blockOrder import sk.ainet.lang.tensor.ops.TensorSpec import sk.ainet.lang.tensor.ops.tensorEncoding import sk.ainet.lang.tensor.storage.BufferHandle +import sk.ainet.lang.tensor.storage.PackedBlockStorage import sk.ainet.lang.tensor.storage.TensorEncoding /** @@ -222,6 +224,18 @@ public class ConstantOperationsConverter : StableHloOperationConverter { val initialValue = node.operation.parameters["initial_value"] val isTrainable = node.operation.parameters["trainable"] as? Boolean ?: true + // A weight kept packed (PackedConstantHandling.KEEP_PACKED) has no dense form to emit; only + // a consumer that reads it packed may take it, and none did. + val packed = node.operation.parameters[PackedConstantHandling.PACKED_WEIGHT_PARAM] as? PackedBlockStorage + if (packed != null) { + return ConversionResult.Failure( + "${node.operation.name} ${node.id} (${node.outputs.firstOrNull()?.name}) is packed ${packed.encoding.name} " + + "and no op consumes it packed. Only a Q8_0 / Q4_0 weight used as x · Wᵀ by a single matmul is " + + "folded (BlockQuantMatmulFusionPass); trace with PackedConstantHandling.DEQUANTIZE for the rest.", + "packed weight ${node.id} reached constant emission" + ) + } + val outputSpec = node.outputs.firstOrNull() val outputType = outputSpec?.let { context.getTypeMapper().mapTensorType(it) } ?: "tensor" diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/generate/HloGenerator.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/generate/HloGenerator.kt index 9fb33ff3..7ca8a379 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/generate/HloGenerator.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/generate/HloGenerator.kt @@ -49,6 +49,12 @@ public object HloGenerator { * [sk.ainet.compile.opt.TargetOptimizers]. */ target: String? = null, + /** + * How packed (quantized) weights are traced. [sk.ainet.lang.trace.PackedConstantHandling.KEEP_PACKED] + * keeps Q8_0 / Q4_0 weights packed for a target that registered + * [sk.ainet.compile.opt.passes.BlockQuantMatmulFusionPass]. + */ + packedConstants: sk.ainet.lang.trace.PackedConstantHandling = sk.ainet.lang.trace.PackedConstantHandling.FAIL, ): StableHloModule { val ctx = DefaultGraphExecutionContext.tape(baseOps = VoidTensorOps()) val traceInput = sampleInput.bind(ctx) @@ -68,7 +74,8 @@ public object HloGenerator { val computeGraph = tape?.toComputeGraph( synthesizeExternalInputs = true, - inputTensorIds = setOf(inputRefId) + inputTensorIds = setOf(inputRefId), + packedConstants = packedConstants ) ?: error("Failed to create compute graph: no execution tape was recorded") // The optimizer pipeline joins the production path here (#1180): with no target named, diff --git a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/BlockQuantMatmulExportTest.kt b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/BlockQuantMatmulExportTest.kt new file mode 100644 index 00000000..69210576 --- /dev/null +++ b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/BlockQuantMatmulExportTest.kt @@ -0,0 +1,149 @@ +package sk.ainet.compile.hlo + +import sk.ainet.compile.hlo.converters.BlockQuantWeight +import sk.ainet.compile.opt.passes.BlockQuantMatmulFusionPass +import sk.ainet.lang.graph.ComputeGraph +import sk.ainet.lang.graph.DefaultGraphExecutionContext +import sk.ainet.lang.tape.toComputeGraph +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Q4_0Quantizer +import sk.ainet.lang.tensor.data.Q8_0Quantizer +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.tensor.ops.VoidTensorOps +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.trace.PackedConstantHandling +import sk.ainet.lang.types.FP32 +import kotlin.math.abs +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertContains +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertFalse +import kotlin.test.assertNotNull +import kotlin.test.assertNull +import kotlin.test.assertTrue + +/** + * A Q8_0 / Q4_0 `Linear` traced with `KEEP_PACKED`, folded by [BlockQuantMatmulFusionPass] and + * emitted as the dequantize-then-`batch_matmul` linalg chain — the form the Torq compiler matches + * to its block-quantized matmul kernel (checked against `torq-compile` by `TorqBlockQuantGoldenTest`). + */ +class BlockQuantMatmulExportTest { + + private val k = 64 + private val n = 48 + + private fun weights(seed: Int, size: Int = n * k): FloatArray { + val r = Random(seed) + return FloatArray(size) { r.nextFloat() * 2f - 1f } + } + + /** Trace `x · Wᵀ` with [packed] as W, then build the graph keeping packed weights. */ + private fun traceLinear(packed: TensorData<*, *>, xShape: Shape = Shape(1, k), op: String = "linear"): ComputeGraph { + val ctx = DefaultGraphExecutionContext.tape(baseOps = VoidTensorOps()) + val x = ctx.fromFloatArray(xShape, FP32::class, FloatArray(xShape.volume)) + @Suppress("UNCHECKED_CAST") + val w = ctx.fromData(packed as TensorData, FP32::class) + val xRef = ctx.session.refOf(x as Tensor<*, *>).id + val (tape, _) = ctx.record { + if (op == "linear") ctx.ops.matmulWeightTransposed(x, w) else ctx.ops.add(x, w) + } + return assertNotNull(tape).toComputeGraph( + synthesizeExternalInputs = true, + inputTensorIds = setOf(xRef), + packedConstants = PackedConstantHandling.KEEP_PACKED, + ) + } + + private fun fuseAndExport(graph: ComputeGraph): String { + val result = BlockQuantMatmulFusionPass().apply(graph) + assertTrue(result.changed, "the pass must fold the packed Linear: ${result.diagnostics}") + return toStableHlo(result.graph).content + } + + @Test + fun q8_0LinearBecomesTheI8Chain() { + val mlir = fuseAndExport(traceLinear(Q8_0Quantizer.quantize(weights(1), Shape(n, k)))) + + assertContains(mlir, "tensor<2x32x48xi8>") + assertContains(mlir, "tensor<2x1x48xbf16>") + assertContains(mlir, "affine_map<(d0, d1, d2) -> (d0, 0, d2)>") + assertContains(mlir, "arith.sitofp") + assertContains(mlir, "tensor.collapse_shape") + assertContains(mlir, "output_shape [1, 64, 48]") + assertContains(mlir, "linalg.batch_matmul") + // f32 graph around a bf16 kernel + assertContains(mlir, "stablehlo.convert") + assertContains(mlir, "-> tensor<1x1x64xbf16>") + assertFalse("stablehlo.dot_general" in mlir, "the matmul must not also be emitted dense") + assertFalse("arith.subf" in mlir, "Q8_0 scales are positive: no zero point") + assertFalse("dense<0.0>" in mlir, "no zero-filled stand-in for the weight") + } + + @Test + fun q4_0LinearWithNegativeScalesGetsAZeroPoint() { + // ggml Q4_0 sets d = max / -8, so about half the blocks have a negative scale. + val mlir = fuseAndExport(traceLinear(Q4_0Quantizer.quantize(weights(2), Shape(n, k)))) + + assertContains(mlir, "tensor<2x32x48xi4>") + assertContains(mlir, "arith.subf") + assertEquals(2, Regex("""affine_map<\(d0, d1, d2\) -> \(d0, 0, d2\)>""").findAll(mlir).count(), "scale and zero point are both per block") + } + + @Test + fun batchedActivationsKeepTheirRows() { + val mlir = fuseAndExport(traceLinear(Q8_0Quantizer.quantize(weights(3), Shape(n, k)), xShape = Shape(4, k))) + assertContains(mlir, "tensor<1x4x64xbf16>") + assertContains(mlir, "-> tensor<4x48xf32>") + } + + @Test + fun regroupedQ8_0WeightDecodesToTheSameMatrix() = assertRegroupedMatches(Q8_0Quantizer.quantize(weights(4), Shape(n, k)), bits = 8) + + @Test + fun regroupedQ4_0WeightDecodesToTheSameMatrix() = assertRegroupedMatches(Q4_0Quantizer.quantize(weights(5), Shape(n, k)), bits = 4) + + private fun assertRegroupedMatches(storage: PackedBlockStorage, bits: Int) { + val w = BlockQuantWeight.of(storage) + assertEquals(bits, w.bits) + val expected = storage.toFloatArray() + val nB = k / 32 + val lo = -(1 shl (bits - 1)) + val hi = (1 shl (bits - 1)) - 1 + for (o in 0 until n) for (b in 0 until nB) { + val scale = bf16ToFloat(w.scales[b * n + o]) + assertTrue(scale > 0f, "scale ($o, $b) must be positive for the kernel, was $scale") + val zp = w.zeroPoints?.let { bf16ToFloat(it[b * n + o]) } ?: 0f + for (j in 0 until 32) { + val c = w.codes[(b * 32 + j) * n + o].toInt() + assertTrue(c in lo..hi, "code $c out of i$bits range") + val want = expected[o * k + b * 32 + j] + val got = (c - zp) * scale + // only the scale is rounded (f16 → bf16, 8 significant bits) + assertTrue(abs(got - want) <= abs(want) / 128f + 1e-7f, "($o, ${b * 32 + j}): $got vs $want") + } + } + } + + @Test + fun allZeroBlockGetsAScaleTheKernelAccepts() { + val w = BlockQuantWeight.of(Q8_0Quantizer.quantize(FloatArray(2 * 32), Shape(2, 32))) + assertTrue(w.scales.all { bf16ToFloat(it) == 1f }) + assertTrue(w.codes.all { it == 0.toByte() }) + assertNull(w.zeroPoints) + } + + @Test + fun packedWeightNobodyFoldedFailsTheExport() { + val graph = traceLinear(Q8_0Quantizer.quantize(weights(6, 64), Shape(1, 64)), xShape = Shape(1, 64), op = "add") + val result = BlockQuantMatmulFusionPass().apply(graph) + assertFalse(result.changed) + assertTrue(result.diagnostics.any { it.startsWith("BLOCK_QUANT_SKIPPED") }, "${result.diagnostics}") + val e = assertFailsWith { toStableHlo(result.graph) } + assertContains(e.message.orEmpty(), "BlockQuantMatmulFusionPass") + } + + private fun bf16ToFloat(bits: Short): Float = Float.fromBits((bits.toInt() and 0xFFFF) shl 16) +} diff --git a/skainet-compile/skainet-compile-hlo/src/jvmTest/kotlin/sk/ainet/compile/hlo/TorqBlockQuantGoldenTest.kt b/skainet-compile/skainet-compile-hlo/src/jvmTest/kotlin/sk/ainet/compile/hlo/TorqBlockQuantGoldenTest.kt new file mode 100644 index 00000000..44ef49d8 --- /dev/null +++ b/skainet-compile/skainet-compile-hlo/src/jvmTest/kotlin/sk/ainet/compile/hlo/TorqBlockQuantGoldenTest.kt @@ -0,0 +1,99 @@ +package sk.ainet.compile.hlo + +import sk.ainet.compile.opt.passes.BlockQuantMatmulFusionPass +import sk.ainet.lang.graph.DefaultGraphExecutionContext +import sk.ainet.lang.tape.toComputeGraph +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Q4_0Quantizer +import sk.ainet.lang.tensor.data.Q8_0Quantizer +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.tensor.ops.VoidTensorOps +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.trace.PackedConstantHandling +import sk.ainet.lang.types.FP32 +import java.io.File +import java.nio.ByteBuffer +import java.nio.ByteOrder +import java.util.concurrent.TimeUnit +import kotlin.random.Random +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNotNull +import kotlin.test.assertTrue + +/** + * Golden check against the Torq compiler: a Q8_0 and a Q4_0 `Linear` exported through + * [BlockQuantMatmulFusionPass] must compile to `torq_hl.mixed_quant_matmul`, never to the generic + * `torq_hl.matmul` fallback. + * + * Runs only when `SKAINET_TORQ_COMPILE` names a command that takes `torq-compile`'s arguments + * (Torq 2.3.0 or later, with a linker for the aarch64 host code); otherwise it passes without + * checking. Each case leaves `in.mlir`, `out.vmfb` and the dequantized weight `w.f32` (`[out, in]`) + * in `build/torq-golden//` for a run on the board. + */ +class TorqBlockQuantGoldenTest { + + private val compiler: List? = System.getenv("SKAINET_TORQ_COMPILE")?.takeIf { it.isNotBlank() }?.split(' ') + + @Test + fun q8_0LinearCompilesToTheBlockQuantKernel() = check("q8_0", Q8_0Quantizer.quantize(weights(11), Shape(N, K))) + + @Test + fun q4_0LinearWithNegativeScalesCompilesToTheBlockQuantKernel() = check("q4_0", Q4_0Quantizer.quantize(weights(12), Shape(N, K))) + + private fun check(name: String, weight: PackedBlockStorage) { + val command = compiler ?: return + val dir = File("build/torq-golden/$name").absoluteFile.apply { deleteRecursively(); mkdirs() } + val input = File(dir, "in.mlir").apply { writeText(export(weight)) } + File(dir, "w.f32").writeBytes(littleEndian(weight.toFloatArray())) + val phases = File(dir, "phases") + + val process = ProcessBuilder( + command + listOf( + input.path, + "--iree-input-type=stablehlo-torq", + "--iree-hal-target-device=torq", + "--torq-hw=SL2610", + "--torq-fallback-f32-to-host", + "--mlir-elide-elementsattrs-if-larger=8", + "--dump-compilation-phases-to=${phases.path}", + "-o", File(dir, "out.vmfb").path, + ) + ).redirectErrorStream(true).redirectOutput(File(dir, "log.txt")).start() + assertTrue(process.waitFor(10, TimeUnit.MINUTES), "torq-compile timed out") + assertEquals(0, process.exitValue(), "torq-compile failed, see ${dir.path}/log.txt") + + val targets = assertNotNull(phases.listFiles()?.singleOrNull { it.name.endsWith(".executable-targets.mlir") }).readText() + assertTrue("torq_hl.mixed_quant_matmul" in targets, "the block-quant kernel did not match, see ${phases.path}") + assertTrue("torq_hl.matmul " !in targets, "part of the matmul fell back to the generic kernel") + } + + private fun export(weight: PackedBlockStorage): String { + val ctx = DefaultGraphExecutionContext.tape(baseOps = VoidTensorOps()) + val x = ctx.fromFloatArray(Shape(1, K), FP32::class, FloatArray(K)) + @Suppress("UNCHECKED_CAST") + val w = ctx.fromData(weight as TensorData, FP32::class) + val xRef = ctx.session.refOf(x as Tensor<*, *>).id + val (tape, _) = ctx.record { ctx.ops.matmulWeightTransposed(x, w) } + val graph = assertNotNull(tape).toComputeGraph( + synthesizeExternalInputs = true, + inputTensorIds = setOf(xRef), + packedConstants = PackedConstantHandling.KEEP_PACKED, + ) + return toStableHlo(BlockQuantMatmulFusionPass().apply(graph).graph).content + } + + private fun weights(seed: Int): FloatArray { + val r = Random(seed) + return FloatArray(N * K) { r.nextFloat() * 2f - 1f } + } + + private fun littleEndian(values: FloatArray): ByteArray = + ByteBuffer.allocate(values.size * 4).order(ByteOrder.LITTLE_ENDIAN).apply { asFloatBuffer().put(values) }.array() + + private companion object { + const val K = 256 + const val N = 96 + } +} diff --git a/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api b/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api index 3b4ddbb5..8b77e65a 100644 --- a/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api +++ b/skainet-compile/skainet-compile-opt/api/jvm/skainet-compile-opt.api @@ -94,6 +94,19 @@ public final class sk/ainet/compile/opt/TargetOptimizers { public final fun stableHloPassesFor (Ljava/lang/String;)Ljava/util/List; } +public final class sk/ainet/compile/opt/passes/BlockQuantMatmulFusionPass : sk/ainet/compile/opt/GraphOptimizationPass { + public static final field Companion Lsk/ainet/compile/opt/passes/BlockQuantMatmulFusionPass$Companion; + public static final field OP_NAME Ljava/lang/String; + public static final field WEIGHT Ljava/lang/String; + public static final field WEIGHT_NAME Ljava/lang/String; + public fun ()V + public fun apply (Lsk/ainet/lang/graph/ComputeGraph;)Lsk/ainet/compile/opt/GraphOptimizationResult; + public fun getName ()Ljava/lang/String; +} + +public final class sk/ainet/compile/opt/passes/BlockQuantMatmulFusionPass$Companion { +} + public final class sk/ainet/compile/opt/passes/ConstantFoldingPass : sk/ainet/compile/opt/GraphOptimizationPass { public fun ()V public fun apply (Lsk/ainet/lang/graph/ComputeGraph;)Lsk/ainet/compile/opt/GraphOptimizationResult; diff --git a/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/BlockQuantMatmulFusionPass.kt b/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/BlockQuantMatmulFusionPass.kt new file mode 100644 index 00000000..b6b9884b --- /dev/null +++ b/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/BlockQuantMatmulFusionPass.kt @@ -0,0 +1,147 @@ +package sk.ainet.compile.opt.passes + +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.memory.BlockOrder +import sk.ainet.lang.tensor.ops.GenericOperation +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.tensor.storage.TensorEncoding +import sk.ainet.lang.trace.PackedConstantHandling + +/** + * Folds a block-quantized `Linear` into one `block_quant_matmul` node that keeps the weight packed. + * + * Matches a weight node carrying packed Q8_0 or Q4_0 storage (traced with + * [PackedConstantHandling.KEEP_PACKED]) that feeds a matmul as `x · Wᵀ`, in either traced form: + * ``` + * weight[N, K] → transpose → matmul(x, ·) + * weight[N, K] → matmul(x, ·, transposeB = true) (after TransposeEliminationPass) + * ``` + * and replaces the three (or two) nodes with + * ``` + * block_quant_matmul(x) parameters: weight = the PackedBlockStorage, weight_name + * ``` + * The StableHLO emitter lowers that node to a dequantize-then-`batch_matmul` linalg chain with the + * codes and per-block scales inline, the form an NPU compiler can match to a block-quantized kernel. + * + * Only an exclusive weight is folded: one that has other consumers (a tied embedding, say) stays as + * it is, and so does one in another encoding, in kernel block order, or whose `K` is not a multiple + * of the block size. Each skip is reported in the diagnostics, because a packed weight left behind + * makes the emitter fail. + * + * A target opts in by registering the pass: + * ```kotlin + * TargetOptimizers.registerDagPasses("torq") { listOf(BlockQuantMatmulFusionPass()) } + * ``` + */ +public class BlockQuantMatmulFusionPass : GraphOptimizationPass { + + override val name: String = "block-quant-matmul-fusion" + + private data class Match(val weight: GraphNode, val transpose: GraphNode?, val matmul: GraphNode, val storage: PackedBlockStorage) + + override fun apply(graph: ComputeGraph): GraphOptimizationResult { + val diagnostics = mutableListOf() + val consumers = graph.edges.groupBy { it.source.id } + + val matches = mutableListOf() + for (weight in graph.nodes) { + val storage = weight.operation.parameters[PackedConstantHandling.PACKED_WEIGHT_PARAM] as? PackedBlockStorage ?: continue + val reason = ineligibility(storage) + if (reason != null) { + diagnostics += "BLOCK_QUANT_SKIPPED weight=${weight.id}: $reason" + continue + } + val match = matchConsumer(weight, storage, consumers) + if (match == null) { + diagnostics += "BLOCK_QUANT_SKIPPED weight=${weight.id}: not the only input of a matmul as x · Wᵀ" + continue + } + matches += match + } + if (matches.isEmpty()) return GraphOptimizationResult(graph, changed = false, diagnostics = diagnostics) + + val removed = matches.flatMap { listOfNotNull(it.weight.id, it.transpose?.id) }.toSet() + val fused = matches.associate { m -> + val spec = m.weight.outputs.firstOrNull() + m.matmul.id to GraphNode( + id = m.matmul.id, + operation = GenericOperation( + name = OP_NAME, + parameters = mapOf(WEIGHT to m.storage, WEIGHT_NAME to (spec?.name ?: m.weight.id)), + type = "linalg", + ), + inputs = m.matmul.inputs.take(1), + outputs = m.matmul.outputs, + metadata = m.matmul.metadata, + ) + } + + val rebuilt = DefaultComputeGraph() + val nodeMap = mutableMapOf() + for (node in graph.nodes) { + if (node.id in removed) continue + val toAdd = fused[node.id] ?: node + rebuilt.addNode(toAdd) + nodeMap[toAdd.id] = toAdd + } + for (edge in graph.edges) { + if (edge.source.id in removed) continue + // The weight operand of a fused matmul is gone with its producer. + if (edge.destination.id in fused && edge.destinationInputIndex != 0) continue + val src = nodeMap[edge.source.id] ?: continue + val dst = nodeMap[edge.destination.id] ?: continue + rebuilt.addEdge(GraphEdge(edge.id, src, dst, edge.sourceOutputIndex, edge.destinationInputIndex, edge.tensorSpec)) + } + + for (m in matches) { + diagnostics += "BLOCK_QUANT_FUSED node=${m.matmul.id} weight=${m.weight.id} encoding=${m.storage.encoding.name} shape=${m.storage.shape.dimensions.toList()}" + } + return GraphOptimizationResult(rebuilt, changed = true, diagnostics = diagnostics) + } + + private fun ineligibility(storage: PackedBlockStorage): String? = when { + storage.encoding != TensorEncoding.Q8_0 && storage.encoding != TensorEncoding.Q4_0 -> + "encoding ${storage.encoding.name} is not Q8_0 or Q4_0" + storage.blockOrder != BlockOrder.ROW_MAJOR -> "blocks are in ${storage.blockOrder} order, not ROW_MAJOR" + storage.shape.rank != 2 -> "shape ${storage.shape.dimensions.toList()} is not [out, in]" + storage.shape[1] % storage.blockSize != 0 -> "in = ${storage.shape[1]} is not a multiple of ${storage.blockSize}" + else -> null + } + + private fun matchConsumer(weight: GraphNode, storage: PackedBlockStorage, consumers: Map>): Match? { + val out = consumers[weight.id]?.singleOrNull() ?: return null + val next = out.destination + if (isMatmul(next)) { + val transposed = next.operation.parameters["transposeB"] == true + return if (transposed && out.destinationInputIndex == 1) Match(weight, null, next, storage) else null + } + if (!isPlainTranspose(next)) return null + val mm = consumers[next.id]?.singleOrNull() ?: return null + if (!isMatmul(mm.destination) || mm.destinationInputIndex != 1) return null + if (mm.destination.operation.parameters["transposeA"] == true || mm.destination.operation.parameters["transposeB"] == true) return null + return Match(weight, next, mm.destination, storage) + } + + private fun isMatmul(node: GraphNode): Boolean = node.operation.name.lowercase() == "matmul" + + /** A transpose that swaps the two axes of a rank-2 tensor — the only one `x · Wᵀ` produces. */ + private fun isPlainTranspose(node: GraphNode): Boolean { + if (node.operation.name.lowercase() != "transpose") return false + val perm = (node.operation.parameters["permutation"] ?: node.operation.parameters["perm"] ?: node.operation.parameters["axes"]) as? List<*> + return perm == null || perm.map { it.toString() } == listOf("1", "0") + } + + public companion object { + /** Op name of the fused node; the StableHLO emitter's block-quant converter lowers it. */ + public const val OP_NAME: String = "block_quant_matmul" + /** Parameter holding the weight's [PackedBlockStorage], logical shape `[out, in]`. */ + public const val WEIGHT: String = "weight" + /** Parameter holding the weight tensor's name, for comments and diagnostics. */ + public const val WEIGHT_NAME: String = "weight_name" + } +}