bf16-native DSL → StableHLO export + pluggable per-target optimization

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/compile/target/OpGranularityPolicy.kt

@@ -0,0 +1,35 @@
+package sk.ainet.compile.target
+
+/**
+ * Per-target policy deciding, for a **fused** graph op (by op name, e.g. "layernorm",
+ * "scaleddotproductattention"), whether an emitter should KEEP it as a single target op
+ * (`stablehlo.composite` / native kernel-call) or DECOMPOSE it into primitives.
+ *
+ * This is the seam for **target legalization** — matching the op *granularity* a fragile
+ * vendor compiler expects, without leaking any hardware knowledge into the agnostic core.
+ * The core never calls this directly; the emitters ([toStableHlo] / the C codegen) consult
+ * the policy their caller resolved for the selected target. A `null` policy means "decompose
+ * everything", which is the portable default (llvm-cpu et al.).
+ *
+ * Lives in the shared DAG module so both the StableHLO emitter (`skainet-compile-hlo`) and
+ * the `TargetOptimizer` registry (`skainet-compile-opt`) can reference it without either
+ * gaining a new dependency edge.
+ */
+public interface OpGranularityPolicy {
+    /** Backend/device id this policy applies to (matches the target string / iree device). */
+    public val target: String
+
+    /** `true` = keep [opName] fused (emit composite / kernel-call); `false` = decompose. */
+    public fun keepFused(opName: String): Boolean
+}
+
+/**
+ * Straightforward allow-list policy: an op is kept fused iff its (lower-cased) name is in
+ * [fused]. Everything else decomposes.
+ */
+public class FusedOpAllowList(
+    override val target: String,
+    private val fused: Set<String>,
+) : OpGranularityPolicy {
+    override fun keepFused(opName: String): Boolean = opName.lowercase() in fused
+}

skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/ConversionContext.kt

@@ -23,7 +23,19 @@ public class ConversionContext @kotlin.jvm.JvmOverloads constructor(
      * See issue #523 for the architecture context.
      */
     public val materializationPolicy: ConstantMaterializationPolicy =
-        ConstantMaterializationPolicy.InlineAlways
+        ConstantMaterializationPolicy.InlineAlways,
+    /**
+     * Selected compile target (iree device id, e.g. "torq"); `null` = target-agnostic.
+     * Threaded through so per-target emit decisions (e.g. op granularity) are possible
+     * without any hardware knowledge in the shared emitter.
+     */
+    public val target: String? = null,
+    /**
+     * Per-target op-granularity policy (fused vs decomposed emission). Resolved by the
+     * caller from the [sk.ainet.compile.opt.TargetOptimizers] registry and passed in;
+     * `null` = decompose everything (portable default). The emitter only *reads* it.
+     */
+    public val granularity: sk.ainet.compile.target.OpGranularityPolicy? = null
 ) {
     private val valueNames = mutableMapOf<String, String>()
     private val valueTypes = mutableMapOf<String, String>()

skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverter.kt

@@ -23,14 +23,18 @@ public class StableHloConverter @kotlin.jvm.JvmOverloads constructor(
      * through to every [ConversionContext] this converter creates.
      */
     private val materializationPolicy: ConstantMaterializationPolicy =
-        ConstantMaterializationPolicy.InlineAlways
+        ConstantMaterializationPolicy.InlineAlways,
+    /** Selected compile target (iree device id, e.g. "torq"); handed to every context. */
+    private val target: String? = null,
+    /** Per-target op-granularity policy; handed to every context (null = decompose all). */
+    private val granularity: sk.ainet.compile.target.OpGranularityPolicy? = null
 ) {
 
     /**
      * Convert a ComputeGraph to StableHLO MLIR format
      */
     public fun convert(graph: ComputeGraph, functionName: String = "main"): StableHloModule {
-        val context = ConversionContext(typeMapper, graph, materializationPolicy)
+        val context = ConversionContext(typeMapper, graph, materializationPolicy, target, granularity)
 
         // Pre-conversion validation (allow orphaned nodes for backward compatibility)
         val validationResult = graph.validate()

skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverterFactory.kt

@@ -32,7 +32,9 @@ public object StableHloConverterFactory {
     @JvmStatic
     @kotlin.jvm.JvmOverloads
     public fun createBasic(
-        policy: ConstantMaterializationPolicy = ConstantMaterializationPolicy.InlineAlways
+        policy: ConstantMaterializationPolicy = ConstantMaterializationPolicy.InlineAlways,
+        target: String? = null,
+        granularity: sk.ainet.compile.target.OpGranularityPolicy? = null
     ): StableHloConverter {
         val registry = StableHloOperationRegistry()
         val typeMapper = TypeMapper()
@@ -75,7 +77,7 @@ public object StableHloConverterFactory {
         // LLM front-door op for token-id \u2192 embedding lookups.
         registry.register(GatherOperationsConverter())
 
-        return StableHloConverter(registry, typeMapper, validator, policy)
+        return StableHloConverter(registry, typeMapper, validator, policy, target, granularity)
     }
 
     /**

skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/TypeMapper.kt

@@ -52,7 +52,8 @@ public class TypeMapper {
     public fun mapTensorType(spec: TensorSpec): String {
         val elementType = mapDType(spec.dtype)
         val shapeStr = formatShape(spec.shape)
-        return "tensor<${shapeStr}x${elementType}>"
+        // Rank-0 (scalar) is `tensor<elem>`, not `tensor<xelem>` — no leading `x`.
+        return if (shapeStr.isEmpty()) "tensor<$elementType>" else "tensor<${shapeStr}x${elementType}>"
     }
 
     /**
@@ -131,6 +132,7 @@ public class TypeMapper {
      */
     public fun createTensorType(shape: List<Int>, dtype: String): String {
         val elementType = mapDType(dtype)
+        if (shape.isEmpty()) return "tensor<$elementType>" // rank-0 scalar
         val shapeStr = shape.joinToString("x")
         return "tensor<${shapeStr}x${elementType}>"
     }

skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/NeuralNetOperationsConverter.kt

@@ -474,44 +474,73 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter {
 
         val operations = mutableListOf<String>()
 
+        // Compute the numerically-sensitive normalization (mean / variance / std /
+        // divide) in f32 regardless of the model dtype. This is standard LayerNorm
+        // practice — PyTorch and JAX upcast fp16/bf16 LayerNorm to f32 internally —
+        // because a bf16 variance (a sum of `axisSize` bf16 squares) loses enough
+        // precision that `sqrt(var + eps)` can overflow/NaN, and some accelerator
+        // backends miscompile the decomposed bf16 reduce/normalize outright. Only the
+        // scale/offset affine stays in the model dtype. No-op when the model is f32.
+        val computeElement = "f32"
+        val isMixed = elementType != computeElement
+        val computeType = if (rank > 0) {
+            "tensor<${inputShape.joinToString("x")}x$computeElement>"
+        } else {
+            "tensor<$computeElement>"
+        }
+        val computeReducedType = if (reducedShape.isEmpty()) {
+            "tensor<$computeElement>"
+        } else {
+            "tensor<${reducedShape.joinToString("x")}x$computeElement>"
+        }
+        val xF32 = if (isMixed) context.nextTempValue() else xInput
+        if (isMixed) {
+            operations += "$xF32 = stablehlo.convert $xInput : ($outputType) -> $computeType"
+        }
+
         // mean(x) along the normalization axis, via real `stablehlo.reduce` (sum / count)
         // so the module compiles on stock IREE (no @reduce_* custom_call stubs).
-        operations += "$zeroInit = stablehlo.constant dense<0.0> : tensor<$elementType>"
-        operations += "$countConst = stablehlo.constant dense<${axisSize}.0> : $reducedType"
-        operations += "$sumX = stablehlo.reduce($xInput init: $zeroInit) " +
-            "applies stablehlo.add across dimensions = [$axis] : ($outputType, tensor<$elementType>) -> $reducedType"
-        operations += "$meanValue = stablehlo.divide $sumX, $countConst : $reducedType"
+        operations += "$zeroInit = stablehlo.constant dense<0.0> : tensor<$computeElement>"
+        operations += "$countConst = stablehlo.constant dense<${axisSize}.0> : $computeReducedType"
+        operations += "$sumX = stablehlo.reduce($xF32 init: $zeroInit) " +
+            "applies stablehlo.add across dimensions = [$axis] : ($computeType, tensor<$computeElement>) -> $computeReducedType"
+        operations += "$meanValue = stablehlo.divide $sumX, $countConst : $computeReducedType"
 
         // Broadcast mean back to input shape, then mean-center.
         operations += "$meanBroadcast = stablehlo.broadcast_in_dim $meanValue, " +
-            "dims = [$broadcastDims] : ($reducedType) -> $outputType"
-        operations += "$centered = stablehlo.subtract $xInput, $meanBroadcast : $outputType"
+            "dims = [$broadcastDims] : ($computeReducedType) -> $computeType"
+        operations += "$centered = stablehlo.subtract $xF32, $meanBroadcast : $computeType"
 
         // var(x) = E[x²] - E[x]² (population, ddof=0), again via real reductions.
-        operations += "$squared = stablehlo.multiply $xInput, $xInput : $outputType"
+        operations += "$squared = stablehlo.multiply $xF32, $xF32 : $computeType"
         operations += "$sumSq = stablehlo.reduce($squared init: $zeroInit) " +
-            "applies stablehlo.add across dimensions = [$axis] : ($outputType, tensor<$elementType>) -> $reducedType"
-        operations += "$meanSq = stablehlo.divide $sumSq, $countConst : $reducedType"
-        operations += "$meanSquared = stablehlo.multiply $meanValue, $meanValue : $reducedType"
-        operations += "$varValue = stablehlo.subtract $meanSq, $meanSquared : $reducedType"
+            "applies stablehlo.add across dimensions = [$axis] : ($computeType, tensor<$computeElement>) -> $computeReducedType"
+        operations += "$meanSq = stablehlo.divide $sumSq, $countConst : $computeReducedType"
+        operations += "$meanSquared = stablehlo.multiply $meanValue, $meanValue : $computeReducedType"
+        operations += "$varValue = stablehlo.subtract $meanSq, $meanSquared : $computeReducedType"
 
         // Epsilon constant broadcast into the reduced shape.
-        operations += "$epsConst = stablehlo.constant dense<$epsilon> : tensor<$elementType>"
+        operations += "$epsConst = stablehlo.constant dense<$epsilon> : tensor<$computeElement>"
         operations += "$epsBroadcast = stablehlo.broadcast_in_dim $epsConst, " +
-            "dims = [] : (tensor<$elementType>) -> $reducedType"
+            "dims = [] : (tensor<$computeElement>) -> $computeReducedType"
 
         // variance + eps
-        operations += "$varPlusEps = stablehlo.add $varValue, $epsBroadcast : $reducedType"
+        operations += "$varPlusEps = stablehlo.add $varValue, $epsBroadcast : $computeReducedType"
 
         // std = sqrt(variance + eps)
-        operations += "$stdValue = stablehlo.sqrt $varPlusEps : $reducedType"
+        operations += "$stdValue = stablehlo.sqrt $varPlusEps : $computeReducedType"
 
         // Broadcast std back to the input shape.
         operations += "$stdBroadcast = stablehlo.broadcast_in_dim $stdValue, " +
-            "dims = [$broadcastDims] : ($reducedType) -> $outputType"
-
-        // normalized = (x - mean) / std
-        operations += "$normalized = stablehlo.divide $centered, $stdBroadcast : $outputType"
+            "dims = [$broadcastDims] : ($computeReducedType) -> $computeType"
+
+        // normalized = (x - mean) / std  (in f32), then cast back to the model dtype
+        // before the scale/offset affine.
+        val normalizedF32 = if (isMixed) context.nextTempValue() else normalized
+        operations += "$normalizedF32 = stablehlo.divide $centered, $stdBroadcast : $computeType"
+        if (isMixed) {
+            operations += "$normalized = stablehlo.convert $normalizedF32 : ($computeType) -> $outputType"
+        }
 
         // Apply scale and offset if present. Each has shape [axisSize] and is broadcast over
         // the normalization axis before the elementwise op. Track the running SSA value so

skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/dag2hlo.kt

@@ -49,9 +49,23 @@ public data class StableHloModule(
  * - Currently supports: input, add, matmul, relu. Unsupported ops are emitted as comments.
  * - DType mapping expects TensorSpec.dtype to be strings like "FP32", "F32", "F64", "I32".
  */
-public fun toStableHlo(graph: ComputeGraph, functionName: String = "main"): StableHloModule {
+public fun toStableHlo(
+    graph: ComputeGraph,
+    functionName: String = "main",
+    /**
+     * Selected compile target (iree device id, e.g. "torq", "llvm-cpu"); `null` =
+     * target-agnostic emission (the portable default, unchanged behavior).
+     */
+    target: String? = null,
+    /**
+     * Per-target op-granularity policy (fused vs decomposed). Resolve it at the call
+     * site with `TargetOptimizers.granularityFor(target)` and pass it in, so the emitter
+     * stays decoupled from the optimizer registry. `null` = decompose everything.
+     */
+    granularity: sk.ainet.compile.target.OpGranularityPolicy? = null,
+): StableHloModule {
     // Use the new converter architecture
-    val converter = StableHloConverterFactory.createBasic()
+    val converter = StableHloConverterFactory.createBasic(target = target, granularity = granularity)
     return converter.convert(graph, functionName)
 }
 
@@ -78,6 +92,8 @@ public fun toStableHlo(
     graph: sk.ainet.lang.graph.ResolvedComputeGraph,
     functionName: String = "main",
     validate: Boolean = true,
+    target: String? = null,
+    granularity: sk.ainet.compile.target.OpGranularityPolicy? = null,
 ): StableHloModule {
     if (validate) {
         graph.validate().requireValid()
@@ -86,7 +102,7 @@ public fun toStableHlo(
     // for graphs that pass validation. Future versions can branch here to
     // consume `graph.resolvedLayout(edgeId)` / `graph.backendAssignment(nodeId)`
     // once those passes ship.
-    return toStableHlo(graph.delegate, functionName)
+    return toStableHlo(graph.delegate, functionName, target, granularity)
 }
 
 /**

skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/TargetOptimization.kt

@@ -0,0 +1,95 @@
+package sk.ainet.compile.opt
+
+/**
+ * Phases of the compile pipeline where target-specific optimizations can plug in.
+ * The IR flows TAPE → DAG → STABLE_HLO; each phase has its own pass shape
+ * ([GraphOptimizationPass] for [DAG]; tape/StableHLO pass types are their own).
+ */
+public enum class CompilePhase { TAPE, DAG, STABLE_HLO }
+
+/**
+ * A **pluggable, target-specific optimizer**. This is the seam that keeps
+ * hardware knowledge OUT of the agnostic compiler core: a backend (e.g. `"torq"`,
+ * `"llvm-cpu"`) registers its own passes here from OUTSIDE core, and the core
+ * pipeline merely asks the registry which passes to run for the selected target
+ * at each phase. Nothing target-specific ever lands in the shared IR emitter or in
+ * a model definition — the StableHLO stays portable.
+ *
+ * Contribute passes per phase by overriding the phase you need; the rest default
+ * to empty. Today [dagPasses] (graph rewrites) is wired; tape / StableHLO phase
+ * hooks follow the same shape when their pass types are needed.
+ */
+public interface TargetOptimizer {
+    /** Backend/device id this optimizer contributes to (matches the target name). */
+    public val target: String
+
+    /** DAG-phase graph rewrites (produce standard, still-portable ops). */
+    public fun dagPasses(): List<GraphOptimizationPass> = emptyList()
+
+    /**
+     * Op-granularity policy for the emitters (fused vs decomposed). `null` = decompose
+     * everything (the portable default). A target that wants a fused op kept as a single
+     * `stablehlo.composite` / kernel-call returns a policy here. See
+     * [sk.ainet.compile.target.OpGranularityPolicy].
+     */
+    public fun granularity(): sk.ainet.compile.target.OpGranularityPolicy? = null
+
+    // Future phases keep the same shape, e.g.:
+    //   public fun tapePasses(): List<TapePass> = emptyList()
+    //   public fun stableHloPasses(): List<StableHloPass> = emptyList()
+}
+
+/**
+ * Pluggable registry of [TargetOptimizer]s. Backends register themselves; core
+ * queries by target name. Being a registry (not a hard-coded list) is precisely
+ * what keeps target-specific optimization out of the agnostic core.
+ */
+public object TargetOptimizers {
+    private val registry = mutableMapOf<String, MutableList<TargetOptimizer>>()
+
+    /** Plug an optimizer in. Several may register for one target; all of them apply. */
+    public fun register(optimizer: TargetOptimizer) {
+        registry.getOrPut(optimizer.target) { mutableListOf() }.add(optimizer)
+    }
+
+    /**
+     * Callback-style registration — plug in a target's DAG passes without declaring
+     * a class: `registerDagPasses("torq") { listOf(TorqAttentionTilingPass()) }`.
+     */
+    public fun registerDagPasses(target: String, passes: () -> List<GraphOptimizationPass>) {
+        register(object : TargetOptimizer {
+            override val target: String = target
+            override fun dagPasses(): List<GraphOptimizationPass> = passes()
+        })
+    }
+
+    /** Optimizers registered for [target] (empty if none / null). */
+    public fun forTarget(target: String?): List<TargetOptimizer> =
+        target?.let { registry[it]?.toList() } ?: emptyList()
+
+    /**
+     * The op-granularity policy for [target] — the first non-null [TargetOptimizer.granularity]
+     * among the registered optimizers, or `null` (decompose everything) if none provide one.
+     * Callers (the model-build tool) resolve this and pass it into `toStableHlo(...)`, so the
+     * emitter stays decoupled from this registry.
+     */
+    public fun granularityFor(target: String?): sk.ainet.compile.target.OpGranularityPolicy? =
+        forTarget(target).firstNotNullOfOrNull { it.granularity() }
+
+    /** Test/reset hook. */
+    public fun clear(): Unit = registry.clear()
+}
+
+/**
+ * DAG-phase pipeline for a target: HW-agnostic [corePasses] first, then whatever the
+ * registry has plugged in for [target]. Callers stay agnostic — they never name a
+ * backend's passes, only the target string (which is the iree device name anyway).
+ */
+public fun dagPipelineFor(
+    target: String?,
+    corePasses: List<GraphOptimizationPass> = emptyList(),
+    maxIterations: Int = 1,
+): GraphOptimizationPipeline {
+    val targetPasses = TargetOptimizers.forTarget(target).flatMap { it.dagPasses() }
+    return GraphOptimizationPipeline(corePasses + targetPasses, maxIterations)
+}