From a7d369d8ce0d15890e1ba635013985359b7be2ee Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sat, 27 Jun 2026 14:31:03 +0200 Subject: [PATCH] fix(compile-hlo): lower layerNorm/rmsNorm/batchNorm to real stablehlo.reduce MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit convertLayerNorm and convertRmsNorm emitted stablehlo.custom_call @reduce_mean / @reduce_variance placeholders that no MLIR toolchain can compile, so the native norm ops were export-only — downstream users got modules that fail in iree-compile. (RMSNorm's only working path was the NN-DSL primitive decomposition, not the converter op.) LayerNorm/RmsNorm also applied scale/offset without a broadcast_in_dim, a shape mismatch for the usual rank-1 affine params. BatchNorm emitted stablehlo.batch_norm_training, whose 3-tuple result the string emitter can't bind to a single SSA value. Lower all three to real, compilable StableHLO, matching convertGroupNorm (already fixed): - mean/variance via stablehlo.reduce (sum) + divide; variance = E[x²] - E[x]² (ddof=0). - broadcast_in_dim the reduced mean/std AND the scale/offset (shape [axisSize]/(C,)) back to the input shape before the elementwise affine. - batchNorm decomposed to the same elementwise form: 5-operand inference (running mean/var) or 3-operand training (batch stats reduced over the non-feature axes); drop the batch_norm_* emission and the now-unused buildBatchNormOperation helper. Validated end-to-end on stock IREE 3.11.0 (via the skainet-iree-conformance harness): native layerNorm/rmsNorm/batchNorm now iree-compile, run, and match numpy (max abs err 0 / 1.2e-7 / 6e-8). Converter unit tests updated to assert real stablehlo.reduce and the absence of @reduce_* custom_calls. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../NeuralNetOperationsConverter.kt | 243 ++++++++++++------ .../compile/hlo/LayerNormConverterTest.kt | 23 +- .../ainet/compile/hlo/RmsNormConverterTest.kt | 11 +- 3 files changed, 189 insertions(+), 88 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 bd87db0c3..551ce27fe 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 @@ -249,6 +249,21 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter { ) } + /** + * Lower BatchNorm to real StableHLO elementwise ops, in the same decomposition style as + * LayerNorm / GroupNorm — instead of `stablehlo.batch_norm_inference` / + * `batch_norm_training` (the training form returns a 3-tuple, which the string emitter + * cannot represent as a single SSA value). Per-channel affine over the `feature_index` + * (channel) axis of an `(N, C, *spatial)` input: + * + * out = (x - mean) / sqrt(var + eps) * scale + offset + * + * `scale` / `offset` (and `mean` / `var` when provided) are shape `(C,)` and broadcast + * over the channel axis. With 5 operands (input, scale, offset, mean, variance) this is + * the **inference / eval** form (running statistics). With 3 (input, scale, offset) the + * batch statistics are computed via real `stablehlo.reduce` over the non-channel axes + * (population variance, ddof=0) — the **training** form's normalized output. + */ private fun convertBatchNorm( node: GraphNode, operands: List, @@ -260,36 +275,102 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter { "Unsupported batchNorm arity for node ${node.id}" ) } - + val outputSpec = node.outputs.firstOrNull() - val outputType = outputSpec?.let { context.getTypeMapper().mapTensorType(it) } + val outputType = outputSpec?.let { context.getTypeMapper().mapTensorType(it) } ?: "tensor" - - // Extract batch norm parameters + val elementType = outputSpec?.let { context.getTypeMapper().mapDType(it.dtype) } + ?: "f32" + + val inputShape = node.inputs.firstOrNull()?.shape ?: outputSpec?.shape ?: emptyList() + val rank = inputShape.size + val params = node.operation.parameters - val epsilon = params["eps"] as? Double ?: 1e-5 - val featureIndex = params["feature_index"] as? Int ?: 1 // Channel dimension - + val epsilon = (params["eps"] as? Double) ?: (params["epsilon"] as? Double) ?: 1e-5 + val rawFeature = params["feature_index"] as? Int ?: 1 // channel dimension + val featureIndex = (if (rawFeature < 0) rank + rawFeature else rawFeature) + .coerceIn(0, (rank - 1).coerceAtLeast(0)) + val channels = if (rank > 0) inputShape[featureIndex] else 1 + val channelType = "tensor<${channels}x$elementType>" + + val xInput = operands[0] + val scaleOperand = operands[1] + val offsetOperand = operands[2] + val meanOperand = if (operands.size > 3) operands[3] else null + val varOperand = if (operands.size > 4) operands[4] else null + + val operations = mutableListOf() + + // Per-channel mean & variance: provided directly (inference) or computed via real + // `stablehlo.reduce` over the non-channel axes (training). Both are shape (C,). + val meanCh: String + val varCh: String + if (meanOperand != null && varOperand != null) { + meanCh = meanOperand + varCh = varOperand + } else { + val reduceDims = (0 until rank).filter { it != featureIndex } + val count = reduceDims.fold(1) { acc, d -> acc * inputShape[d] } + val dimsList = reduceDims.joinToString(", ") + val zeroInit = context.nextTempValue() + val countConst = context.nextTempValue() + val sumX = context.nextTempValue() + val computedMean = context.nextTempValue() + val squared = context.nextTempValue() + val sumSq = context.nextTempValue() + val meanSq = context.nextTempValue() + val meanSquared = context.nextTempValue() + val computedVar = context.nextTempValue() + operations += "$zeroInit = stablehlo.constant dense<0.0> : tensor<$elementType>" + operations += "$countConst = stablehlo.constant dense<${count}.0> : $channelType" + operations += "$sumX = stablehlo.reduce($xInput init: $zeroInit) " + + "applies stablehlo.add across dimensions = [$dimsList] : ($outputType, tensor<$elementType>) -> $channelType" + operations += "$computedMean = stablehlo.divide $sumX, $countConst : $channelType" + operations += "$squared = stablehlo.multiply $xInput, $xInput : $outputType" + operations += "$sumSq = stablehlo.reduce($squared init: $zeroInit) " + + "applies stablehlo.add across dimensions = [$dimsList] : ($outputType, tensor<$elementType>) -> $channelType" + operations += "$meanSq = stablehlo.divide $sumSq, $countConst : $channelType" + operations += "$meanSquared = stablehlo.multiply $computedMean, $computedMean : $channelType" + operations += "$computedVar = stablehlo.subtract $meanSq, $meanSquared : $channelType" + meanCh = computedMean + varCh = computedVar + } + + // std = sqrt(var + eps) per channel. + val epsConst = context.nextTempValue() + val varPlusEps = context.nextTempValue() + val stdCh = context.nextTempValue() + operations += "$epsConst = stablehlo.constant dense<$epsilon> : $channelType" + operations += "$varPlusEps = stablehlo.add $varCh, $epsConst : $channelType" + operations += "$stdCh = stablehlo.sqrt $varPlusEps : $channelType" + + // Broadcast the (C,) tensors over the channel axis and apply the affine. + val meanB = context.nextTempValue() + val centered = context.nextTempValue() + val stdB = context.nextTempValue() + val normalized = context.nextTempValue() + val scaleB = context.nextTempValue() + val scaled = context.nextTempValue() + val offsetB = context.nextTempValue() val resultValue = context.nextTempValue() - - // Build StableHLO batch_norm_inference operation - val batchNormOperation = buildBatchNormOperation( - resultValue = resultValue, - input = operands[0], - scale = operands[1], - offset = operands[2], - mean = if (operands.size > 3) operands[3] else null, - variance = if (operands.size > 4) operands[4] else null, - outputType = outputType, - epsilon = epsilon, - featureIndex = featureIndex - ) - - context.emitOperation(batchNormOperation) - + operations += "$meanB = stablehlo.broadcast_in_dim $meanCh, " + + "dims = [$featureIndex] : ($channelType) -> $outputType" + operations += "$centered = stablehlo.subtract $xInput, $meanB : $outputType" + operations += "$stdB = stablehlo.broadcast_in_dim $stdCh, " + + "dims = [$featureIndex] : ($channelType) -> $outputType" + operations += "$normalized = stablehlo.divide $centered, $stdB : $outputType" + operations += "$scaleB = stablehlo.broadcast_in_dim $scaleOperand, " + + "dims = [$featureIndex] : ($channelType) -> $outputType" + operations += "$scaled = stablehlo.multiply $normalized, $scaleB : $outputType" + operations += "$offsetB = stablehlo.broadcast_in_dim $offsetOperand, " + + "dims = [$featureIndex] : ($channelType) -> $outputType" + operations += "$resultValue = stablehlo.add $scaled, $offsetB : $outputType" + + operations.forEach { context.emitOperation(it) } + return ConversionResult.Success( outputValueName = resultValue, - emittedOperations = listOf(batchNormOperation) + emittedOperations = operations ) } @@ -301,15 +382,14 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter { * * out = scale * (x - mean) / sqrt(var + eps) + offset * - * Emission style matches the softmax fix (#467) and the rest of - * the emitter: reductions go through - * `stablehlo.custom_call @reduce_mean` / `@reduce_variance` (both - * already supported by `ReductionOperationsConverter`), the reduced - * tensors are broadcast back to the input shape via - * `stablehlo.broadcast_in_dim`, and scale / offset are elementwise - * multiplied / added only when their operands are actually present. - * Migrating every reduction to real `stablehlo.reduce` regions is - * a separate, larger refactor. + * Reductions use real `stablehlo.reduce` (sum / count) — not the + * `@reduce_mean` / `@reduce_variance` custom_call stubs — so the + * module compiles on stock IREE (matching `convertGroupNorm`). + * Variance is population (ddof=0) via `E[x²] - E[x]²`. The reduced + * mean / std are broadcast back to the input shape via + * `stablehlo.broadcast_in_dim`; scale / offset (shape `[axisSize]`) + * are broadcast over the normalization axis and applied elementwise + * only when their operands are actually present. */ private fun convertLayerNorm( node: GraphNode, @@ -366,9 +446,20 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter { val scaleOperand: String? = if (operands.size > 1) operands[1] else null val offsetOperand: String? = if (operands.size > 2) operands[2] else null + // Number of elements reduced over the normalization axis (mean/var divisor). + val axisSize = if (rank > 0) inputShape[axis] else 1 + val scaleType = "tensor<${axisSize}x$elementType>" + + val zeroInit = context.nextTempValue() + val countConst = context.nextTempValue() + val sumX = context.nextTempValue() val meanValue = context.nextTempValue() val meanBroadcast = context.nextTempValue() val centered = context.nextTempValue() + val squared = context.nextTempValue() + val sumSq = context.nextTempValue() + val meanSq = context.nextTempValue() + val meanSquared = context.nextTempValue() val varValue = context.nextTempValue() val epsConst = context.nextTempValue() val epsBroadcast = context.nextTempValue() @@ -379,20 +470,26 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter { val operations = mutableListOf() - // mean(x) along the normalization axis. - operations += "$meanValue = stablehlo.custom_call @reduce_mean($xInput) " + - "{dimensions = [$axis], keepdim = false} : $reducedType" + // 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" - // Broadcast mean back to input shape. + // Broadcast mean back to input shape, then mean-center. operations += "$meanBroadcast = stablehlo.broadcast_in_dim $meanValue, " + "dims = [$broadcastDims] : ($reducedType) -> $outputType" - - // Mean-center. operations += "$centered = stablehlo.subtract $xInput, $meanBroadcast : $outputType" - // variance(x) along the normalization axis. - operations += "$varValue = stablehlo.custom_call @reduce_variance($xInput) " + - "{dimensions = [$axis], keepdim = false} : $reducedType" + // var(x) = E[x²] - E[x]² (population, ddof=0), again via real reductions. + operations += "$squared = stablehlo.multiply $xInput, $xInput : $outputType" + 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" // Epsilon constant broadcast into the reduced shape. operations += "$epsConst = stablehlo.constant dense<$epsilon> : tensor<$elementType>" @@ -412,18 +509,24 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter { // normalized = (x - mean) / std operations += "$normalized = stablehlo.divide $centered, $stdBroadcast : $outputType" - // Apply scale and offset if present. Track the current running - // SSA value so omitting either one keeps the emitted MLIR - // faithful to the input graph. + // 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 + // omitting either one keeps the emitted MLIR faithful to the input graph. var current = normalized if (scaleOperand != null) { + val scaleBroadcast = context.nextTempValue() val scaled = context.nextTempValue() - operations += "$scaled = stablehlo.multiply $current, $scaleOperand : $outputType" + operations += "$scaleBroadcast = stablehlo.broadcast_in_dim $scaleOperand, " + + "dims = [$axis] : ($scaleType) -> $outputType" + operations += "$scaled = stablehlo.multiply $current, $scaleBroadcast : $outputType" current = scaled } if (offsetOperand != null) { + val offsetBroadcast = context.nextTempValue() val offsetted = context.nextTempValue() - operations += "$offsetted = stablehlo.add $current, $offsetOperand : $outputType" + operations += "$offsetBroadcast = stablehlo.broadcast_in_dim $offsetOperand, " + + "dims = [$axis] : ($scaleType) -> $outputType" + operations += "$offsetted = stablehlo.add $current, $offsetBroadcast : $outputType" current = offsetted } @@ -665,7 +768,13 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter { val xInput = operands[0] val scaleOperand: String? = if (operands.size >= 2) operands[1] else null + val axisSize = if (rank > 0) inputShape[axis] else 1 + val scaleType = "tensor<${axisSize}x$elementType>" + val xSquared = context.nextTempValue() + val zeroInit = context.nextTempValue() + val countConst = context.nextTempValue() + val sumSq = context.nextTempValue() val meanSquared = context.nextTempValue() val epsConst = context.nextTempValue() val epsBroadcast = context.nextTempValue() @@ -680,9 +789,13 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter { // x^2 operations += "$xSquared = stablehlo.multiply $xInput, $xInput : $outputType" - // reduce_mean(x^2, axis) - operations += "$meanSquared = stablehlo.custom_call @reduce_mean($xSquared) " + - "{dimensions = [$axis], keepdim = false} : $reducedType" + // mean(x^2, axis) via real `stablehlo.reduce` (sum / count) — not the @reduce_mean + // custom_call stub — so the module compiles on stock IREE (matching convertGroupNorm). + operations += "$zeroInit = stablehlo.constant dense<0.0> : tensor<$elementType>" + operations += "$countConst = stablehlo.constant dense<${axisSize}.0> : $reducedType" + operations += "$sumSq = stablehlo.reduce($xSquared init: $zeroInit) " + + "applies stablehlo.add across dimensions = [$axis] : ($outputType, tensor<$elementType>) -> $reducedType" + operations += "$meanSquared = stablehlo.divide $sumSq, $countConst : $reducedType" // eps constant broadcast into the reduced shape operations += "$epsConst = stablehlo.constant dense<$eps> : tensor<$elementType>" @@ -702,11 +815,14 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter { // x / rms operations += "$normalized = stablehlo.divide $xInput, $rmsBroadcast : $outputType" - // Final scale multiply is optional — when the caller did not - // pass a scale operand we return the normalized value directly. + // Final scale multiply is optional. Scale has shape [axisSize] and is broadcast over + // the normalization axis; when the caller passed no scale we return normalized directly. val finalValue: String if (scaleOperand != null) { - operations += "$resultValue = stablehlo.multiply $normalized, $scaleOperand : $outputType" + val scaleBroadcast = context.nextTempValue() + operations += "$scaleBroadcast = stablehlo.broadcast_in_dim $scaleOperand, " + + "dims = [$axis] : ($scaleType) -> $outputType" + operations += "$resultValue = stablehlo.multiply $normalized, $scaleBroadcast : $outputType" finalValue = resultValue } else { finalValue = normalized @@ -955,27 +1071,6 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter { return listOf(initConstant, areaConstant, sumOp, divideOp) } - private fun buildBatchNormOperation( - resultValue: String, - input: String, - scale: String, - offset: String, - mean: String?, - variance: String?, - outputType: String, - epsilon: Double, - featureIndex: Int - ): String { - return if (mean != null && variance != null) { - // Use batch_norm_inference when mean and variance are provided - "$resultValue = stablehlo.batch_norm_inference $input, $scale, $offset, $mean, $variance, " + - "epsilon = $epsilon, feature_index = $featureIndex : $outputType" - } else { - // Use batch_norm_training when mean and variance need to be computed - "$resultValue = stablehlo.batch_norm_training $input, $scale, $offset, " + - "epsilon = $epsilon, feature_index = $featureIndex : $outputType" - } - } /** * Convert scaledDotProductAttention to StableHLO. diff --git a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/LayerNormConverterTest.kt b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/LayerNormConverterTest.kt index 9746861b2..f1c1f4a1f 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/LayerNormConverterTest.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/LayerNormConverterTest.kt @@ -18,9 +18,9 @@ import kotlin.test.assertTrue * Before this fix, `NeuralNetOperationsConverter.convertLayerNorm` * emitted `stablehlo.custom_call @layer_norm(...)`, which no MLIR * tool in the repo understands. This test pins the new lowering: - * a real elementwise decomposition using @reduce_mean / @reduce_variance - * / broadcast_in_dim / sqrt / divide — matching softmax #467 and the - * codebase's existing reduction-via-custom-call style. + * a real elementwise decomposition using `stablehlo.reduce` (sum / count) + * / broadcast_in_dim / sqrt / divide — compilable on stock IREE, matching + * convertGroupNorm. No @reduce_* custom_call stubs. * * layer_norm(x) = scale * (x - mean) / sqrt(var + eps) + offset */ @@ -45,14 +45,14 @@ class LayerNormConverterTest { val converter = StableHloConverterFactory.createExtended() val module = converter.convert(graph, "test_layer_norm_full") - // Core elementwise decomposition. + // Core elementwise decomposition — real reductions, no custom_call stubs. assertTrue( - module.content.contains("@reduce_mean"), - "layerNorm must lower mean(x) to a real reduction" + module.content.contains("stablehlo.reduce"), + "layerNorm must lower mean/var to real stablehlo.reduce (compiles on stock IREE)" ) - assertTrue( - module.content.contains("@reduce_variance"), - "layerNorm must lower var(x) to a real reduction" + assertFalse( + module.content.contains("@reduce_mean") || module.content.contains("@reduce_variance"), + "layerNorm must not emit @reduce_* custom_call stubs (uncompilable)" ) assertTrue( module.content.contains("stablehlo.subtract"), @@ -87,8 +87,9 @@ class LayerNormConverterTest { val module = converter.convert(graph, "test_layer_norm_minimal") assertFalse(module.content.contains("@layer_norm")) - assertTrue(module.content.contains("@reduce_mean")) - assertTrue(module.content.contains("@reduce_variance")) + assertTrue(module.content.contains("stablehlo.reduce")) + assertFalse(module.content.contains("@reduce_mean")) + assertFalse(module.content.contains("@reduce_variance")) assertTrue(module.content.contains("stablehlo.subtract")) assertTrue(module.content.contains("stablehlo.sqrt")) assertTrue(module.content.contains("stablehlo.divide")) diff --git a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/RmsNormConverterTest.kt b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/RmsNormConverterTest.kt index f09e4ede4..9e562571d 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/RmsNormConverterTest.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/RmsNormConverterTest.kt @@ -68,8 +68,12 @@ class RmsNormConverterTest { "RMSNorm must emit at least one multiply (x*x and/or scale*x/rms)" ) assertTrue( - module.content.contains("@reduce_mean"), - "RMSNorm must lower mean(x^2) to a real reduction (custom_call style matches the rest of the emitter)" + module.content.contains("stablehlo.reduce"), + "RMSNorm must lower mean(x^2) to a real stablehlo.reduce (compiles on stock IREE)" + ) + assertFalse( + module.content.contains("@reduce_mean") || module.content.contains("@reduce_variance"), + "RMSNorm must not emit @reduce_* custom_call stubs (uncompilable)" ) assertTrue( module.content.contains("stablehlo.sqrt"), @@ -104,7 +108,8 @@ class RmsNormConverterTest { "RMSNorm without a scale operand must still be claimed by the converter" ) // The core norm still happens — we just skip the final scale multiply. - assertTrue(module.content.contains("@reduce_mean")) + assertTrue(module.content.contains("stablehlo.reduce")) + assertFalse(module.content.contains("@reduce_mean")) assertTrue(module.content.contains("stablehlo.sqrt")) assertTrue(module.content.contains("stablehlo.divide")) }