Skip to content

fix(compile-hlo): lower layerNorm/rmsNorm/batchNorm to real stablehlo.reduce (compile on stock IREE) #769

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
michalharakal merged 1 commit into from
Jun 27, 2026
Merged

fix(compile-hlo): lower layerNorm/rmsNorm/batchNorm to real stablehlo.reduce (compile on stock IREE) #769

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
243 changes: 169 additions & 74 deletions ...hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/NeuralNetOperationsConverter.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<String>,
Expand All @@ -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<?x?x?x?xf32>"

// 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<String>()

// 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
)
}

Expand All @@ -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,
Expand Down Expand Up @@ -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()
Expand All @@ -379,20 +470,26 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter {

val operations = mutableListOf<String>()

// 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>"
Expand All @@ -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
}

Expand Down Expand Up @@ -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()
Expand All @@ -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>"
Expand All @@ -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
Expand Down Expand Up @@ -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.
Expand Down
Loading
Loading