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 85cd73740..ca6165a06 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 @@ -29,6 +29,7 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter { "batchNorm", "batchNormalization", "BatchNormalization", "layerNorm", "layerNormalization", "LayerNormalization", "rmsNorm", "rms_norm", "RMSNorm", "RmsNorm", + "groupNorm", "groupNormalization", "GroupNormalization", "group_norm", // Attention "scaledDotProductAttention" ) @@ -46,6 +47,7 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter { "batchnorm", "batchnormalization" -> convertBatchNorm(node, operands, context) "layernorm", "layernormalization" -> convertLayerNorm(node, operands, context) "rmsnorm", "rms_norm" -> convertRmsNorm(node, operands, context) + "groupnorm", "groupnormalization", "group_norm" -> convertGroupNorm(node, operands, context) "scaleddotproductattention" -> convertSdpa(node, operands, context) else -> ConversionResult.Unsupported( node.operation.name, @@ -433,6 +435,148 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter { ) } + /** + * Lower GroupNorm to real StableHLO elementwise ops, in the same + * decomposition style as LayerNorm / RMSNorm (no `@group_norm` + * custom_call stub). GroupNorm splits the `C` channels of an + * `(N, C, *spatial)` input into `num_groups` groups and normalizes each + * group over its channels and spatial positions, then applies an + * optional per-channel affine: + * + * xg = reshape(x, [N, G, M]) // M = (C/G) * prod(spatial) + * out = (xg - mean(xg)) / sqrt(var(xg) + eps) // reduce over M + * out = reshape(out, [N, C, *spatial]) + * out = out * scale + offset // scale/offset shape (C,), optional + * + * The per-group reduction reuses the single-axis `@reduce_mean` / + * `@reduce_variance` custom_calls (exactly as LayerNorm does) by collapsing + * each group's channels + spatial into one trailing axis. Scale and offset + * broadcast over the channel dimension only. + */ + private fun convertGroupNorm( + node: GraphNode, + operands: List, + context: ConversionContext + ): ConversionResult { + if (operands.isEmpty()) { + return ConversionResult.Failure( + "GroupNorm operation requires at least 1 operand (input), got ${operands.size}", + "Unsupported groupNorm arity for node ${node.id}" + ) + } + + val outputSpec = node.outputs.firstOrNull() + val outputType = outputSpec?.let { context.getTypeMapper().mapTensorType(it) } + ?: "tensor" + val elementType = outputSpec?.let { context.getTypeMapper().mapDType(it.dtype) } + ?: "f32" + + val inputShape = node.inputs.firstOrNull()?.shape ?: outputSpec?.shape ?: emptyList() + if (inputShape.size < 2) { + return ConversionResult.Failure( + "GroupNorm requires an (N, C, ...) input of rank >= 2, got rank ${inputShape.size}", + "Unsupported groupNorm input rank for node ${node.id}" + ) + } + + val n = inputShape[0] + val c = inputShape[1] + val spatialCount = inputShape.drop(2).fold(1) { acc, d -> acc * d } + + val params = node.operation.parameters + val numGroups = (params["num_groups"] as? Int) + ?: (params["groups"] as? Int) + ?: (params["numGroups"] as? Int) + ?: 1 + val groups = numGroups.coerceIn(1, if (c > 0) c else 1) + if (c % groups != 0) { + return ConversionResult.Failure( + "GroupNorm channels ($c) must be divisible by num_groups ($groups)", + "Unsupported groupNorm grouping for node ${node.id}" + ) + } + val perGroup = (c / groups) * spatialCount // M + + val epsilon = (params["eps"] as? Double) + ?: (params["epsilon"] as? Double) + ?: 1e-5 + + val groupedType = "tensor<${n}x${groups}x${perGroup}x$elementType>" + val reducedType = "tensor<${n}x${groups}x$elementType>" + + val xInput = operands[0] + val scaleOperand: String? = if (operands.size > 1) operands[1] else null + val offsetOperand: String? = if (operands.size > 2) operands[2] else null + + val grouped = context.nextTempValue() + val meanValue = context.nextTempValue() + val meanBroadcast = context.nextTempValue() + val centered = context.nextTempValue() + val varValue = context.nextTempValue() + val epsConst = context.nextTempValue() + val epsBroadcast = context.nextTempValue() + val varPlusEps = context.nextTempValue() + val stdValue = context.nextTempValue() + val stdBroadcast = context.nextTempValue() + val normalized = context.nextTempValue() + val reshapedBack = context.nextTempValue() + + val operations = mutableListOf() + + // Reshape (N, C, *spatial) -> (N, G, M): collapse each group's channels + + // spatial into one trailing axis so a single-axis reduction is per-group. + operations += "$grouped = stablehlo.reshape $xInput : ($outputType) -> $groupedType" + + // mean(xg) over the trailing axis, broadcast back, mean-center. + operations += "$meanValue = stablehlo.custom_call @reduce_mean($grouped) " + + "{dimensions = [2], keepdim = false} : $reducedType" + operations += "$meanBroadcast = stablehlo.broadcast_in_dim $meanValue, " + + "dims = [0, 1] : ($reducedType) -> $groupedType" + operations += "$centered = stablehlo.subtract $grouped, $meanBroadcast : $groupedType" + + // var(xg) over the trailing axis; std = sqrt(var + eps). + operations += "$varValue = stablehlo.custom_call @reduce_variance($grouped) " + + "{dimensions = [2], keepdim = false} : $reducedType" + operations += "$epsConst = stablehlo.constant dense<$epsilon> : tensor<$elementType>" + operations += "$epsBroadcast = stablehlo.broadcast_in_dim $epsConst, " + + "dims = [] : (tensor<$elementType>) -> $reducedType" + operations += "$varPlusEps = stablehlo.add $varValue, $epsBroadcast : $reducedType" + operations += "$stdValue = stablehlo.sqrt $varPlusEps : $reducedType" + operations += "$stdBroadcast = stablehlo.broadcast_in_dim $stdValue, " + + "dims = [0, 1] : ($reducedType) -> $groupedType" + operations += "$normalized = stablehlo.divide $centered, $stdBroadcast : $groupedType" + + // Reshape back to (N, C, *spatial). + operations += "$reshapedBack = stablehlo.reshape $normalized : ($groupedType) -> $outputType" + + // Optional per-channel affine: scale/offset have shape (C,), broadcast over + // the channel dimension (index 1). + var current = reshapedBack + if (scaleOperand != null) { + val scaleBroadcast = context.nextTempValue() + val scaled = context.nextTempValue() + operations += "$scaleBroadcast = stablehlo.broadcast_in_dim $scaleOperand, " + + "dims = [1] : (tensor<${c}x$elementType>) -> $outputType" + operations += "$scaled = stablehlo.multiply $current, $scaleBroadcast : $outputType" + current = scaled + } + if (offsetOperand != null) { + val offsetBroadcast = context.nextTempValue() + val offsetted = context.nextTempValue() + operations += "$offsetBroadcast = stablehlo.broadcast_in_dim $offsetOperand, " + + "dims = [1] : (tensor<${c}x$elementType>) -> $outputType" + operations += "$offsetted = stablehlo.add $current, $offsetBroadcast : $outputType" + current = offsetted + } + + operations.forEach { context.emitOperation(it) } + + return ConversionResult.Success( + outputValueName = current, + emittedOperations = operations + ) + } + /** * Lower RMSNorm to real StableHLO elementwise ops. This is the * normalization every Llama / Mistral / Qwen / Gemma family diff --git a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/GroupNormConverterTest.kt b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/GroupNormConverterTest.kt new file mode 100644 index 000000000..8372f3534 --- /dev/null +++ b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/GroupNormConverterTest.kt @@ -0,0 +1,157 @@ +package sk.ainet.compile.hlo + +import sk.ainet.lang.graph.DefaultComputeGraph +import sk.ainet.lang.graph.GraphEdge +import sk.ainet.lang.graph.GraphNode +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.ops.Operation +import sk.ainet.lang.tensor.ops.TensorSpec +import sk.ainet.lang.tensor.ops.ValidationResult +import sk.ainet.lang.types.DType +import kotlin.test.Test +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +/** + * Covers the new GroupNorm lowering (companion to LayerNorm #480 / RMSNorm). + * + * GroupNorm had no converter at all — `NeuralNetOperationsConverter` did not + * list it in `supportedOperations`, so a `groupNorm` node fell through to the + * "no converter found" path. This test pins the decomposition: + * + * xg = reshape(x, [N, G, M]) // group the channels + * out = (xg - mean) / sqrt(var + eps) // reduce over M + * out = reshape(out, [N, C, *spatial]) * scale + offset + * + * It asserts the real elementwise lowering (reshape + @reduce_mean / + * @reduce_variance + broadcast_in_dim + sqrt + divide), never a stub. + */ +class GroupNormConverterTest { + + @Test + fun groupNorm_does_not_emit_custom_call_stub() { + val graph = buildGroupNormGraph(withScale = true, withOffset = true) + val module = StableHloConverterFactory.createExtended().convert(graph, "test_group_norm") + println("[DEBUG_LOG] GroupNorm lowering:\n${module.content}") + + assertFalse( + module.content.contains("@group_norm"), + "groupNorm must not fall back to a @group_norm custom_call stub" + ) + assertFalse( + module.content.contains("Operation not supported"), + "groupNorm must be routed to a converter, not the unsupported path" + ) + } + + @Test + fun groupNorm_lowers_to_reshape_reductions_and_affine() { + val graph = buildGroupNormGraph(withScale = true, withOffset = true) + val module = StableHloConverterFactory.createExtended().convert(graph, "test_group_norm_full") + val mlir = module.content + + assertTrue(mlir.contains("stablehlo.reshape"), "groupNorm must reshape to group the channels and back") + assertTrue(mlir.contains("@reduce_mean"), "groupNorm must lower mean(x) to a real reduction") + assertTrue(mlir.contains("@reduce_variance"), "groupNorm must lower var(x) to a real reduction") + assertTrue(mlir.contains("stablehlo.subtract"), "groupNorm must mean-center") + assertTrue(mlir.contains("stablehlo.sqrt"), "groupNorm must take sqrt(var + eps)") + assertTrue(mlir.contains("stablehlo.divide"), "groupNorm must divide by the std") + assertTrue(mlir.contains("stablehlo.broadcast_in_dim"), "groupNorm must broadcast reduced stats back") + assertTrue(mlir.contains("stablehlo.multiply"), "groupNorm must apply the per-channel scale") + assertTrue(mlir.contains("stablehlo.add"), "groupNorm must apply the per-channel offset") + } + + @Test + fun groupNorm_without_scale_or_offset_still_lowers() { + val graph = buildGroupNormGraph(withScale = false, withOffset = false) + val module = StableHloConverterFactory.createExtended().convert(graph, "test_group_norm_minimal") + val mlir = module.content + + assertFalse(mlir.contains("@group_norm")) + assertTrue(mlir.contains("stablehlo.reshape")) + assertTrue(mlir.contains("@reduce_mean")) + assertTrue(mlir.contains("@reduce_variance")) + assertTrue(mlir.contains("stablehlo.sqrt")) + assertTrue(mlir.contains("stablehlo.divide")) + } + + // (N=1, C=4, H=2, W=2), num_groups=2 -> M = (4/2)*2*2 = 8; scale/offset shape (4). + private fun buildGroupNormGraph(withScale: Boolean, withOffset: Boolean): DefaultComputeGraph { + val graph = DefaultComputeGraph() + val shape = listOf(1, 4, 2, 2) + + val input = GraphNode( + id = "x", + operation = markerInputOp(), + inputs = emptyList(), + outputs = listOf(TensorSpec("x", shape, "FP32")) + ) + graph.addNode(input) + + val gnInputs = mutableListOf(TensorSpec("x", shape, "FP32")) + val extraEdges = mutableListOf>() + + if (withScale) { + val scaleNode = GraphNode( + id = "scale", + operation = markerInputOp(), + inputs = emptyList(), + outputs = listOf(TensorSpec("scale", listOf(4), "FP32")) + ) + graph.addNode(scaleNode) + gnInputs.add(TensorSpec("scale", listOf(4), "FP32")) + extraEdges.add(scaleNode to (gnInputs.size - 1)) + } + if (withOffset) { + val offsetNode = GraphNode( + id = "offset", + operation = markerInputOp(), + inputs = emptyList(), + outputs = listOf(TensorSpec("offset", listOf(4), "FP32")) + ) + graph.addNode(offsetNode) + gnInputs.add(TensorSpec("offset", listOf(4), "FP32")) + extraEdges.add(offsetNode to (gnInputs.size - 1)) + } + + val groupNorm = GraphNode( + id = "gn1", + operation = groupNormOp(eps = 1e-5, numGroups = 2), + inputs = gnInputs.toList(), + outputs = listOf(TensorSpec("y", shape, "FP32")) + ) + graph.addNode(groupNorm) + graph.addEdge(GraphEdge("e1", input, groupNorm, 0, 0, input.outputs[0])) + extraEdges.forEachIndexed { i, (src, idx) -> + graph.addEdge(GraphEdge("e${i + 2}", src, groupNorm, 0, idx, src.outputs[0])) + } + + return graph + } + + private fun markerInputOp(): Operation = object : Operation { + override val name: String = "input" + override val type: String = "input" + override val parameters: Map = emptyMap() + override fun execute(inputs: List>): List> = + throw UnsupportedOperationException("test fixture only") + override fun validateInputs(inputs: List): ValidationResult = ValidationResult.Valid + override fun inferOutputs(inputs: List): List = emptyList() + override fun clone(newParameters: Map): Operation = this + override fun serialize(): Map = mapOf("name" to name, "type" to type) + } + + private fun groupNormOp(eps: Double, numGroups: Int): Operation = object : Operation { + override val name: String = "groupNorm" + override val type: String = "normalization" + override val parameters: Map = mapOf("eps" to eps, "num_groups" to numGroups) + override fun execute(inputs: List>): List> = + throw UnsupportedOperationException("test fixture only") + override fun validateInputs(inputs: List): ValidationResult = ValidationResult.Valid + override fun inferOutputs(inputs: List): List = inputs.take(1) + override fun clone(newParameters: Map): Operation = this + override fun serialize(): Map = mapOf( + "name" to name, "type" to type, "parameters" to parameters + ) + } +}