diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/TypeMapper.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/TypeMapper.kt index 6de465a8..58db6f2e 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/TypeMapper.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/TypeMapper.kt @@ -33,6 +33,19 @@ public class TypeMapper { } } + /** + * Bit-pattern literal for -inf in the given MLIR float element type. The + * width MUST match the type — a 32-bit `0xFF800000` in a `bf16` constant is + * "out of range" to iree-compile. Used as the identity for `stablehlo.maximum` + * (softmax / attention max-reduce). + */ + public fun negInfBits(mlirElementType: String): String = when (mlirElementType) { + "f64" -> "0xFFF0000000000000" + "f16" -> "0xFC00" + "bf16" -> "0xFF80" + else -> "0xFF800000" // f32 and fallback + } + /** * Map a TensorSpec to MLIR tensor type string */ diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ActivationOperationsConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ActivationOperationsConverter.kt index dc9512cb..4fd5c221 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ActivationOperationsConverter.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/ActivationOperationsConverter.kt @@ -157,8 +157,9 @@ public class ActivationOperationsConverter : StableHloOperationConverter { val resultValue = context.nextTempValue() // Identity for stablehlo.maximum on floats: -inf. Spell it via the bit - // pattern so MLIR parses it regardless of how the element type prints. - val maxIdentity = "0xFF800000" + // pattern (width-matched to the element type — a 32-bit pattern in a bf16 + // constant is out of range). + val maxIdentity = context.getTypeMapper().negInfBits(elementType) val operations = listOf( // Reduce-max along the softmax axis (for numerical stability). diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt index b221e676..283b4f44 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/AttentionOperationsConverter.kt @@ -144,7 +144,7 @@ public class AttentionOperationsConverter : StableHloOperationConverter { } // softmax(softmaxIn) over the key-length axis - ops += "$maxInit = stablehlo.constant dense<0xFF800000> : tensor<$elem>" + ops += "$maxInit = stablehlo.constant dense<${mapper.negInfBits(elem)}> : tensor<$elem>" ops += "$maxV = stablehlo.reduce($softmaxIn init: $maxInit) applies stablehlo.maximum across dimensions = [$sdAxis] : ($scoresType, tensor<$elem>) -> $reducedType" ops += "$maxB = stablehlo.broadcast_in_dim $maxV, dims = [$bcastDims] : ($reducedType) -> $scoresType" ops += "$shifted = stablehlo.subtract $softmaxIn, $maxB : $scoresType" diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/BasicMathConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/BasicMathConverter.kt index 6c66b37b..a5fc34bb 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/BasicMathConverter.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/BasicMathConverter.kt @@ -131,7 +131,9 @@ public class BasicMathConverter : StableHloOperationConverter { dtype = targetSpec.dtype ) val toType = typeMapper.mapTensorType(promotedSpec) - context.emitOperation("$converted = stablehlo.convert $current : $fromType to $toType") + // Functional-type form `: (A) -> B`; the `A to B` short form is not + // valid stablehlo.convert assembly and is rejected by iree-compile. + context.emitOperation("$converted = stablehlo.convert $current : ($fromType) -> $toType") current = converted currentSpec = promotedSpec } diff --git a/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/DtypeForwardPropagationPass.kt b/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/DtypeForwardPropagationPass.kt new file mode 100644 index 00000000..82c33e4b --- /dev/null +++ b/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/passes/DtypeForwardPropagationPass.kt @@ -0,0 +1,131 @@ +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.GraphNode + +/** + * Forward dtype propagation — makes graph edges dtype-consistent so the + * StableHLO emitter produces well-typed MLIR. + * + * Walks the graph in topological order and, for every node: + * 1. sets each **input** spec's dtype to its producer's actual **output** dtype + * (an edge whose producer emits `bf16` but whose consumer input spec says + * `f32` renders the same SSA value with two types — malformed MLIR); and + * 2. for dtype-**preserving** ops, sets the **output** spec dtype to the (now + * unified) input dtype, so the whole f32/bf16 island collapses to one dtype. + * + * This fixes bf16-native traces where reductions / normalizations were recorded + * with a stale FP32 dtype while their producers emit bf16 (e.g. the Moonshine + * encoder's LayerNorm reduce). It is a **no-op for uniformly-typed graphs** + * (all-FP32 models are already edge-consistent, so nothing changes). + * + * When [targetFloatDtype] is set (e.g. `"BF16"`), every **float** source node's + * output is coerced to it first, so a bf16-native model unifies to bf16 *end to + * end* — including matmul activations, which the Torq NPU requires (bf16 weights + * alone leave `f32 activation × bf16 weight` matmuls). Weights already at the + * target dtype are a no-op; integer/bool sources (indices, masks) are left alone. + * With [targetFloatDtype] null the pass only enforces edge-consistency. + * + * Left untouched: + * - **source** nodes' non-float dtypes (indices/masks): authoritative. + * - explicit **dtype-changing** ops (`convert`/`cast`/`quantize`/`dequantize`, + * `argmax`/`argmin`, comparisons): their output dtype is intentional. + */ +public class DtypeForwardPropagationPass( + private val targetFloatDtype: String? = null, +) : GraphOptimizationPass { + + override val name: String = "dtype-forward-propagation" + + override fun apply(graph: ComputeGraph): GraphOptimizationResult { + val diagnostics = mutableListOf() + val topo = graph.getTopologicalOrder() + val incoming = graph.edges.groupBy { it.destination.id } + + // Updated node per id; seeded with the originals, overwritten in topo order + // so a consumer always reads its producer's already-updated output dtype. + val updated: MutableMap = graph.nodes.associateBy { it.id }.toMutableMap() + var changed = false + + // Coerce float source nodes to the target dtype so the whole model unifies + // to it (not just edges downstream of the weights). + if (targetFloatDtype != null) { + for (node in graph.nodes) { + if (incoming[node.id]?.isNotEmpty() == true) continue // not a source + val newOutputs = node.outputs.map { + if (isFloat(it.dtype) && it.dtype != targetFloatDtype) it.copy(dtype = targetFloatDtype) else it + } + if (newOutputs != node.outputs) { + updated[node.id] = node.copy(outputs = newOutputs) + changed = true + } + } + } + + for (orig in topo) { + val node = updated[orig.id] ?: continue + val edges = incoming[orig.id].orEmpty() + if (edges.isEmpty()) continue // source node — keep its declared dtype + + // (1) input spec dtype := producer output dtype + val newInputs = node.inputs.toMutableList() + for (e in edges) { + val producer = updated[e.source.id] ?: continue + val prodDtype = producer.outputs.getOrNull(e.sourceOutputIndex)?.dtype ?: continue + val i = e.destinationInputIndex + val spec = newInputs.getOrNull(i) ?: continue + if (spec.dtype != prodDtype) newInputs[i] = spec.copy(dtype = prodDtype) + } + + // (2) dtype-preserving ops inherit the (data) input dtype on outputs + val inheritDtype = newInputs.firstOrNull()?.dtype + val newOutputs = if (inheritDtype != null && !isDtypeChanging(node.operationName)) { + node.outputs.map { if (it.dtype != inheritDtype) it.copy(dtype = inheritDtype) else it } + } else { + node.outputs + } + + if (newInputs != node.inputs || newOutputs != node.outputs) { + updated[orig.id] = node.copy(inputs = newInputs, outputs = newOutputs) + changed = true + } + } + + if (!changed) return GraphOptimizationResult(graph, changed = false, diagnostics = diagnostics) + + // Rebuild the graph so edges reference the updated node instances (the + // topo sort keys on node identity, so edges must point at the new nodes). + val newGraph = DefaultComputeGraph() + for (node in graph.nodes) newGraph.addNode(updated[node.id] ?: node) + for (edge in graph.edges) { + val src = updated[edge.source.id] ?: continue + val dst = updated[edge.destination.id] ?: continue + val spec = src.outputs.getOrNull(edge.sourceOutputIndex) ?: edge.tensorSpec + newGraph.addEdge(edge.copy(source = src, destination = dst, tensorSpec = spec)) + } + return GraphOptimizationResult(newGraph, changed = true, diagnostics = diagnostics) + } + + private fun isFloat(dtype: String): Boolean = dtype.uppercase() in FLOAT_DTYPES + + private fun isDtypeChanging(op: String): Boolean { + val n = op.lowercase() + return n in DTYPE_CHANGING || + n.startsWith("convert") || n.startsWith("cast") || + n.startsWith("quant") || n.startsWith("dequant") + } + + private companion object { + val FLOAT_DTYPES: Set = setOf( + "F32", "FP32", "FLOAT32", "F16", "FP16", "FLOAT16", "BF16", "BFLOAT16", "F64", "FP64", "FLOAT64", + ) + val DTYPE_CHANGING: Set = setOf( + "argmax", "argmin", "greater", "greaterequal", "less", "lessequal", + "equal", "notequal", "compare", "logicaland", "logicalor", "logicalnot", + "isnan", "isinf", + ) + } +} diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt index f95e0599..66e1f577 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/data/DenseTensorDataFactory.kt @@ -3,6 +3,7 @@ package sk.ainet.lang.tensor.data import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.data.dense.DenseByteTensorArray import sk.ainet.lang.tensor.storage.ActiveMemoryTracker +import sk.ainet.lang.types.BF16 import sk.ainet.lang.types.DType import sk.ainet.lang.types.FP16 import sk.ainet.lang.types.FP32 @@ -346,6 +347,13 @@ public class DenseTensorDataFactory: TensorDataFactory { val data = FloatArray(shape.volume) { 0.0f } createFloatTensorData(shape, data, FP16 as T) as TensorData } + BF16::class -> { + // Float-backed, BF16-tagged (mirrors FP16). The dtype tag is what + // the DSL trace / StableHLO export reads to emit `bf16` element + // types — required for the Torq NPU (bf16-native weights). + val data = FloatArray(shape.volume) { 0.0f } + createFloatTensorData(shape, data, BF16 as T) as TensorData + } Int32::class -> { val data = IntArray(shape.volume) { 0 } createIntTensorData(shape, data) as TensorData @@ -369,6 +377,7 @@ public class DenseTensorDataFactory: TensorDataFactory { return when (dtype) { FP32::class -> LazyZeroFloatArrayTensorData(shape) as TensorData FP16::class -> LazyZeroFloatArrayTensorData(shape) as TensorData + BF16::class -> LazyZeroFloatArrayTensorData(shape) as TensorData Int32::class -> LazyZeroIntArrayTensorData(shape) as TensorData else -> zeros(shape, dtype) }