feat(hlo): add GroupNorm StableHLO converter

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<String>,
+        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<?x?x?x?xf32>"
+        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<String>()
+
+        // 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

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<Pair<GraphNode, Int>>()
+
+        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<String, Any> = emptyMap()
+        override fun <T : DType, V> execute(inputs: List<Tensor<T, V>>): List<Tensor<T, V>> =
+            throw UnsupportedOperationException("test fixture only")
+        override fun validateInputs(inputs: List<TensorSpec>): ValidationResult = ValidationResult.Valid
+        override fun inferOutputs(inputs: List<TensorSpec>): List<TensorSpec> = emptyList()
+        override fun clone(newParameters: Map<String, Any>): Operation = this
+        override fun serialize(): Map<String, Any> = 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<String, Any> = mapOf("eps" to eps, "num_groups" to numGroups)
+        override fun <T : DType, V> execute(inputs: List<Tensor<T, V>>): List<Tensor<T, V>> =
+            throw UnsupportedOperationException("test fixture only")
+        override fun validateInputs(inputs: List<TensorSpec>): ValidationResult = ValidationResult.Valid
+        override fun inferOutputs(inputs: List<TensorSpec>): List<TensorSpec> = inputs.take(1)
+        override fun clone(newParameters: Map<String, Any>): Operation = this
+        override fun serialize(): Map<String, Any> = mapOf(
+            "name" to name, "type" to type, "parameters" to parameters
+        )
+    }
+}