Skip to content

feat(hlo): add GroupNorm StableHLO converter - #752

Merged
michalharakal merged 1 commit into
developfrom
feat/groupnorm-converter
Jun 22, 2026
Merged

michalharakal merged 1 commit into
developfrom
feat/groupnorm-converter

Conversation

@michalharakal

Copy link
Copy Markdown
Contributor

Why

groupNorm had no converter — a groupNorm node fell through NeuralNetOperationsConverter to the "Operation not supported" path and never lowered to StableHLO. It's one of the gaps tracked in #751 (left over after #666 closed).

What

Add convertGroupNorm mirroring the existing LayerNorm/RMSNorm decomposition (no custom_call stub):

xg  = reshape(x, [N, G, M])                  # M = (C/G) * prod(spatial)
out = (xg - mean(xg)) / sqrt(var(xg) + eps)  # reduce over the trailing axis
out = reshape(out, [N, C, *spatial]) * scale + offset   # affine optional
  • Per-group mean/var reuse the @reduce_mean / @reduce_variance custom_calls (exactly as LayerNorm does) by collapsing each group's channels + spatial into one trailing axis.
  • scale/offset (shape C) broadcast over the channel dim (dims=[1]) only.
  • Registered the op names (groupNorm/groupNormalization/GroupNormalization/group_norm) in supportedOperations and the dispatch when.

Emitted MLIR for (N=1,C=4,H=2,W=2), num_groups=2:

%v0  = stablehlo.reshape %arg0 : (tensor<1x4x2x2xf32>) -> tensor<1x2x8xf32>
%v1  = stablehlo.custom_call @reduce_mean(%v0) {dimensions = [2], keepdim = false} : tensor<1x2xf32>
... subtract / reduce_variance / +eps / sqrt / divide ...
%v11 = stablehlo.reshape %v10 : (tensor<1x2x8xf32>) -> tensor<1x4x2x2xf32>
%v12 = stablehlo.broadcast_in_dim %arg1, dims = [1] : (tensor<4xf32>) -> tensor<1x4x2x2xf32>   # scale
%v13 = stablehlo.multiply %v11, %v12 : tensor<1x4x2x2xf32>
%v14 = stablehlo.broadcast_in_dim %arg2, dims = [1] : (tensor<4xf32>) -> tensor<1x4x2x2xf32>   # offset
%v15 = stablehlo.add %v13, %v14 : tensor<1x4x2x2xf32>

Tests

Adds GroupNormConverterTest (structural, mirrors LayerNormConverterTest): asserts the real reshape + reductions + broadcast + sqrt + divide + affine, and that it never hits the stub / unsupported path. The full skainet-compile-hlo jvmTest suite passes locally on develop.

Notes / follow-ups (per #751)

Closes part of #751.

GroupNorm had no converter — a `groupNorm` node fell through to the
"Operation not supported" path. Add `convertGroupNorm` to
NeuralNetOperationsConverter, register the op names in `supportedOperations`
and the dispatch `when`, mirroring the existing LayerNorm/RMSNorm decomposition
(no custom_call stub):

    xg  = reshape(x, [N, G, M])                 # M = (C/G) * prod(spatial)
    out = (xg - mean(xg)) / sqrt(var(xg) + eps) # reduce over the trailing axis
    out = reshape(out, [N, C, *spatial]) * scale + offset   # affine optional

Per-group stats reuse the @reduce_mean / @reduce_variance custom_calls (as
LayerNorm does) by collapsing each group's channels+spatial into one axis;
scale/offset (shape C) broadcast over the channel dim only.

Adds GroupNormConverterTest (structural, mirrors LayerNormConverterTest);
full skainet-compile-hlo jvmTest suite passes.

Refs #751

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@michalharakal
michalharakal merged commit d37303c into develop Jun 22, 2026
4 checks passed
@michalharakal
michalharakal deleted the feat/groupnorm-converter branch June 22, 2026 19:21
MacOS pushed a commit to MacOS/SKaiNET that referenced this pull request Jul 10, 2026
Patch release. Bumps VERSION_NAME 0.31.2 -> 0.32.1 and brings the release
metadata current: develop never received the 0.32.0 release back-merge (PR SKaiNET-developers#753
was tagged + published but blocked from merging), so this consolidates the 0.32.0
AND 0.32.1 CHANGELOG / README "What's New" entries (supersedes SKaiNET-developers#753).

0.32.1 fix: GroupNorm now emits real stablehlo.reduce instead of
@reduce_mean/@reduce_variance custom_calls, so a groupNorm module compiles on
stock iree-compile. Verified end-to-end via skainet-iree-conformance:
PASS max_abs_err=1.2e-7. (PR SKaiNET-developers#754)

0.32.0 (folded in): GroupNorm StableHLO converter (SKaiNET-developers#752), SKEEP proposals docs
(SKaiNET-developers#750), quantization-process doc (SKaiNET-developers#747), dependency bumps.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant