compile: bf16-native DSL -> StableHLO export path

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
      */

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).

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"

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
         }

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<String>()
+        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<String, GraphNode> = 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<String> = setOf(
+            "F32", "FP32", "FLOAT32", "F16", "FP16", "FLOAT16", "BF16", "BFLOAT16", "F64", "FP64", "FLOAT64",
+        )
+        val DTYPE_CHANGING: Set<String> = setOf(
+            "argmax", "argmin", "greater", "greaterequal", "less", "lessequal",
+            "equal", "notequal", "compare", "logicaland", "logicalor", "logicalnot",
+            "isnan", "isinf",
+        )
+    }
+}

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<T, V>
             }
+            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<T, V>
+            }
             Int32::class -> {
                 val data = IntArray(shape.volume) { 0 }
                 createIntTensorData(shape, data) as TensorData<T, V>
@@ -369,6 +377,7 @@ public class DenseTensorDataFactory: TensorDataFactory {
         return when (dtype) {
             FP32::class -> LazyZeroFloatArrayTensorData<T>(shape) as TensorData<T, V>
             FP16::class -> LazyZeroFloatArrayTensorData<T>(shape) as TensorData<T, V>
+            BF16::class -> LazyZeroFloatArrayTensorData<T>(shape) as TensorData<T, V>
             Int32::class -> LazyZeroIntArrayTensorData<T>(shape) as TensorData<T, V>
             else -> zeros(shape, dtype)
         }