From 199a9c9ff5ca82af530861689dcea2cecd5be911 Mon Sep 17 00:00:00 2001 From: michal harakal Date: Fri, 3 Jul 2026 09:00:40 +0200 Subject: [PATCH 1/4] compile-opt: pluggable per-phase, per-target optimization mechanism MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the seam that keeps hardware knowledge OUT of the agnostic compiler core: backends register their own passes from outside core and the pipeline asks the registry which passes to run for the selected target at each phase. - CompilePhase { TAPE, DAG, STABLE_HLO } - TargetOptimizer: per-phase pass provider (dagPasses() today; tape / StableHLO hooks follow the same shape) - TargetOptimizers: pluggable registry — register(...) / registerDagPasses(target){} (callback) / forTarget(target) - dagPipelineFor(target, corePasses): builds a GraphOptimizationPipeline of the agnostic core passes plus whatever is plugged in for the target The transforms a target plugs in still emit standard, portable StableHLO — the shared IR emitter and model definitions carry no target-specific code. Co-Authored-By: Claude Opus 4.8 --- .../ainet/compile/opt/TargetOptimization.kt | 78 +++++++++++++++++++ 1 file changed, 78 insertions(+) create mode 100644 skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/TargetOptimization.kt diff --git a/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/TargetOptimization.kt b/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/TargetOptimization.kt new file mode 100644 index 00000000..fbc8232d --- /dev/null +++ b/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/TargetOptimization.kt @@ -0,0 +1,78 @@ +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 = emptyList() + + // Future phases keep the same shape, e.g.: + // public fun tapePasses(): List = emptyList() + // public fun stableHloPasses(): List = 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>() + + /** 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) { + register(object : TargetOptimizer { + override val target: String = target + override fun dagPasses(): List = passes() + }) + } + + /** Optimizers registered for [target] (empty if none / null). */ + public fun forTarget(target: String?): List = + target?.let { registry[it]?.toList() } ?: emptyList() + + /** 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 = emptyList(), + maxIterations: Int = 1, +): GraphOptimizationPipeline { + val targetPasses = TargetOptimizers.forTarget(target).flatMap { it.dagPasses() } + return GraphOptimizationPipeline(corePasses + targetPasses, maxIterations) +} From 6d036a887926fb3d849d3c4e2b5eaa6235ee9eb6 Mon Sep 17 00:00:00 2001 From: michal harakal Date: Fri, 3 Jul 2026 09:00:40 +0200 Subject: [PATCH 2/4] compile-hlo: emit rank-0 tensor type as tensor, not tensor mapTensorType/createTensorType unconditionally inserted the `x` shape separator, so an empty (rank-0 scalar) shape produced the malformed `tensor` that iree-compile rejects. Emit `tensor` when the shape is empty. Co-Authored-By: Claude Opus 4.8 --- .../src/commonMain/kotlin/sk/ainet/compile/hlo/TypeMapper.kt | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 58db6f2e..84f3b268 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 @@ -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`, not `tensor` — 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, 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}>" } From 006c5ab47676a5dd16e0cddc93db92068f17627c Mon Sep 17 00:00:00 2001 From: michal harakal Date: Sat, 4 Jul 2026 23:50:09 +0200 Subject: [PATCH 3/4] nn: compute LayerNorm normalization in f32 regardless of model dtype A bf16 variance (a sum of many bf16 squares) loses enough precision that sqrt(var + eps) can overflow to NaN, and some accelerator backends miscompile the decomposed low-precision reduce/normalize outright. Upcast the mean / variance / std / divide to f32 (standard PyTorch/JAX LayerNorm practice); the affine (gamma/beta) stays in the model dtype. No-op for f32 models. --- .../NeuralNetOperationsConverter.kt | 69 +++++++++++++------ .../nn/normalization/LayerNormalization.kt | 42 ++++++++--- 2 files changed, 81 insertions(+), 30 deletions(-) diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/NeuralNetOperationsConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/NeuralNetOperationsConverter.kt index d87e987d..1de31308 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/NeuralNetOperationsConverter.kt +++ b/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() + // 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 diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/normalization/LayerNormalization.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/normalization/LayerNormalization.kt index 18d27a34..41c79be5 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/normalization/LayerNormalization.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/normalization/LayerNormalization.kt @@ -7,7 +7,10 @@ import sk.ainet.lang.nn.topology.ModuleParameters import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.* +import sk.ainet.lang.types.BF16 import sk.ainet.lang.types.DType +import sk.ainet.lang.types.FP16 +import sk.ainet.lang.types.FP32 import kotlin.reflect.KClass /** @@ -84,13 +87,32 @@ public class LayerNormalization( "Input shape ${normDims.contentToString()} doesn't match normalized shape ${normalizedShape.contentToString()}" } - // Calculate mean and variance across the normalized dimensions - val mean = calculateLayerMean(input) - val variance = calculateLayerVariance(input, mean) - - // Normalize - val normalized = normalize(input, mean, variance) - + // 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 low-precision variance (a sum of many bf16/fp16 + // squares) loses enough precision that `sqrt(var + eps)` can overflow to + // NaN, and some accelerator backends miscompile the decomposed + // low-precision reduce/normalize outright. The affine (gamma/beta) stays in + // the model dtype. No-op when the model already runs f32. + val normalized: Tensor = if (input.dtype != FP32::class) { + val xF32 = input.ops.convert(input, FP32) + val mean = calculateLayerMean(xF32) + val variance = calculateLayerVariance(xF32, mean) + val normF32 = normalize(xF32, mean, variance) + val modelDtype: DType = when (input.dtype) { + BF16::class -> BF16 + FP16::class -> FP16 + else -> FP32 + } + @Suppress("UNCHECKED_CAST") + (input.ops.convert(normF32, modelDtype) as Tensor) + } else { + val mean = calculateLayerMean(input) + val variance = calculateLayerVariance(input, mean) + normalize(input, mean, variance) + } + if (elementwiseAffine) { val gamma = params[0].value // weight parameter val beta = params[1].value // bias parameter @@ -100,7 +122,7 @@ public class LayerNormalization( } } - private fun calculateLayerMean(input: Tensor): Tensor { + private fun calculateLayerMean(input: Tensor): Tensor { var m = input for (i in normalizedShape.indices.reversed()) { m = m.mean(dim = m.rank - 1) @@ -113,7 +135,7 @@ public class LayerNormalization( return result } - private fun calculateLayerVariance(input: Tensor, mean: Tensor): Tensor { + private fun calculateLayerVariance(input: Tensor, mean: Tensor): Tensor { val centered = input - mean var v = centered * centered for (i in normalizedShape.indices.reversed()) { @@ -127,7 +149,7 @@ public class LayerNormalization( return result } - private fun normalize(input: Tensor, mean: Tensor, variance: Tensor): Tensor { + private fun normalize(input: Tensor, mean: Tensor, variance: Tensor): Tensor { return (input - mean) / (variance + eps).sqrt() } } \ No newline at end of file From b2a1e16773c4d3390d674122e1a7bc116ea2bee0 Mon Sep 17 00:00:00 2001 From: michal harakal Date: Sat, 4 Jul 2026 23:50:21 +0200 Subject: [PATCH 4/4] compile: add target + op-granularity plugin seam to StableHLO emitter Thread an optional target id and OpGranularityPolicy through toStableHlo -> StableHloConverterFactory -> StableHloConverter -> ConversionContext, so per-target emission decisions (fused vs decomposed ops = target legalization) are possible without any hardware knowledge in the agnostic core. OpGranularityPolicy (+ FusedOpAllowList) lives in skainet-compile-dag, the module both the emitter (skainet-compile-hlo) and the TargetOptimizer registry (skainet-compile-opt) already depend on, so no new dependency edges and no cycle. TargetOptimizer gains granularity(); TargetOptimizers gains granularityFor(target). The caller resolves the policy and passes it in, keeping the emitter decoupled from the registry. Everything defaults to null (decompose all), so emission is byte-identical today. --- .../compile/target/OpGranularityPolicy.kt | 35 +++++++++++++++++++ .../sk/ainet/compile/hlo/ConversionContext.kt | 14 +++++++- .../ainet/compile/hlo/StableHloConverter.kt | 8 +++-- .../compile/hlo/StableHloConverterFactory.kt | 6 ++-- .../kotlin/sk/ainet/compile/hlo/dag2hlo.kt | 22 ++++++++++-- .../ainet/compile/opt/TargetOptimization.kt | 17 +++++++++ 6 files changed, 94 insertions(+), 8 deletions(-) create mode 100644 skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/compile/target/OpGranularityPolicy.kt diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/compile/target/OpGranularityPolicy.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/compile/target/OpGranularityPolicy.kt new file mode 100644 index 00000000..c9c2ff8a --- /dev/null +++ b/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, +) : OpGranularityPolicy { + override fun keepFused(opName: String): Boolean = opName.lowercase() in fused +} diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/ConversionContext.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/ConversionContext.kt index 3b335be0..0c018a3b 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/ConversionContext.kt +++ b/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() private val valueTypes = mutableMapOf() diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverter.kt index 9f2e428c..2415ad42 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/StableHloConverter.kt +++ b/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() 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 dd8d8c8e..22e04676 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 @@ -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) } /** diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/dag2hlo.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/dag2hlo.kt index a8636e3b..e257b126 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/dag2hlo.kt +++ b/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) } /** diff --git a/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/TargetOptimization.kt b/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/TargetOptimization.kt index fbc8232d..f5c7cdb8 100644 --- a/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/TargetOptimization.kt +++ b/skainet-compile/skainet-compile-opt/src/commonMain/kotlin/sk/ainet/compile/opt/TargetOptimization.kt @@ -26,6 +26,14 @@ public interface TargetOptimizer { /** DAG-phase graph rewrites (produce standard, still-portable ops). */ public fun dagPasses(): List = 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 = emptyList() // public fun stableHloPasses(): List = emptyList() @@ -59,6 +67,15 @@ public object TargetOptimizers { public fun forTarget(target: String?): List = 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() }