Update NHWC rewriter for native_group_norm.

PiperOrigin-RevId: 718944411

Author: chunnienc

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_check.py

@@ -201,22 +201,25 @@ def _aten_group_norm_checker(node):
   return NHWCable(can_be=can_be, must_be=must_be)
 
 
-@nhwcable_node_checkers.register(aten.native_group_norm)
+@nhwcable_node_checkers.register(aten.native_group_norm.default)
 def _aten_native_group_norm_checker(node):
+  # aten.group_norm is removed from the decomp table, so aten.native_group_norm
+  # should never exist in the graph. However, torch 2.5.1 could ignore the
+  # decomp table updates, so still add this native_group_norm checker and
+  # rewriter to be safe.
+  # The checker and rewriter are the same as the ones for aten.group_norm.
+
   val = node.meta.get("val")
   if (
       not isinstance(val, (list, tuple))
       or not val
       or not hasattr(val[0], "shape")
   ):
     return NHWCable(can_be=False, must_be=False)
-  if len(node.args) >= 3 and (
-      node.args[1] is not None or node.args[2] is not None
-  ):
-    # Disable NHWC rewriter due to precision issue with weight and bias.
-    # TODO(b/354780253): Re-enable NHWC rewriter with proper lowering.
-    return NHWCable(can_be=False, must_be=False)
-  return NHWCable(can_be=len(val[0].shape) == 4, must_be=False)
+
+  can_be = len(val[0].shape) == 4
+  must_be = can_be and ai_edge_torch.config.enable_group_norm_composite
+  return NHWCable(can_be=can_be, must_be=must_be)
 
 
 # ==== Ops must be NCHW

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py

@@ -391,34 +391,32 @@ def native_group_norm(
       eps: float,
       **kwargs,
   ):
-    input_reshaped = torch.reshape(
-        input,
-        [
-            batch_size,
-            flattened_inner_size,
-            num_groups,
-            num_channels // num_groups,
-        ],
-    )
-    reduction_dims = [1, 3]
-
-    biased_var, mean = torch.var_mean(
-        input_reshaped, dim=reduction_dims, unbiased=False, keepdim=True
+    is_composite_supported = (
+        ai_edge_torch.config.enable_group_norm_composite
+        and weight is not None
+        and bias is not None
     )
-    rstd = torch.rsqrt(biased_var + eps)
-
-    out = (input_reshaped - mean) * rstd
-    out = torch.reshape(out, input.shape)
 
-    if weight is not None:
-      out = out * weight
-    if bias is not None:
-      out = out + bias
+    builder = None
+    if is_composite_supported:
+      builder = StableHLOCompositeBuilder(
+          name="odml.group_norm",
+          attr={
+              "num_groups": num_groups,
+              "epsilon": eps,
+              "reduction_axes": [3],
+              "channel_axis": 3,
+          },
+      )
+      input, weight, bias = builder.mark_inputs(input, weight, bias)
 
-    mean = torch.squeeze(mean, reduction_dims)
-    rstd = torch.squeeze(rstd, reduction_dims)
+    input = utils.tensor_to_nchw(input)
+    output = aten.group_norm.default(input, num_groups, weight, bias, eps=eps)
+    output = utils.tensor_to_nhwc(output)
 
-    return out, mean, rstd
+    if builder is not None:
+      output = builder.mark_outputs(output)
+    return (output, None, None)
 
   node.target = native_group_norm
 

ai_edge_torch/_convert/fx_passes/test/test_optimize_layout_transposes_pass.py

@@ -16,7 +16,9 @@
 
 from typing import Callable, Union
 
+import ai_edge_torch
 from ai_edge_torch import fx_infra
+from ai_edge_torch import lowertools
 from ai_edge_torch._convert import fx_passes
 import torch
 import torch.utils._pytree as pytree
@@ -48,10 +50,7 @@ def forward(self, *args, **kwargs):
   )
   exported_program = fx_infra.run_passes(
       exported_program,
-      [
-          fx_passes.OptimizeLayoutTransposesPass(),
-          fx_passes.CanonicalizePass(),
-      ],
+      [fx_passes.OptimizeLayoutTransposesPass()],
   )
   return exported_program
 
@@ -90,27 +89,19 @@ def test_torchvision_resnet18(self):
         model, exported_program.module(), forward_args()
     )
 
-  def test_native_group_norm_no_weight_bias(self):
-    batch_size = 16
-    num_channels = 640
-    flattened_inner_size = 256
-    num_groups = 32
-    eps = 1e-6
+  def test_group_norm_affine_false(self):
 
     class SampleModel(torch.nn.Module):
 
+      def __init__(self):
+        super().__init__()
+        self.group_norm = torch.nn.GroupNorm(
+            num_groups=32, num_channels=640, affine=False, eps=1e-6
+        )
+
       def forward(self, x):
         x = torch.nn.AvgPool2d(2)(x)
-        x = torch.ops.aten.native_group_norm(
-            x,
-            None,
-            None,
-            batch_size,
-            num_channels,
-            flattened_inner_size,
-            num_groups,
-            eps,
-        )[0]
+        x = self.group_norm(x)
         x = torch.nn.AvgPool2d(2)(x)
         return x
 
@@ -121,41 +112,56 @@ def forward(self, x):
         model, exported_program.module(), forward_args()
     )
 
-  def test_native_group_norm_large_weight_bias(self):
-    batch_size = 16
-    num_channels = 640
-    flattened_inner_size = 256
-    num_groups = 32
-    eps = 1e-6
+  def test_group_norm_large_affine_true(self):
 
     class SampleModel(torch.nn.Module):
 
-      def forward(self, x, weight, bias):
+      def __init__(self):
+        super().__init__()
+        self.group_norm = torch.nn.GroupNorm(
+            num_groups=32, num_channels=640, affine=True, eps=1e-6
+        )
+
+      def forward(self, x):
         x = torch.nn.AvgPool2d(2)(x)
-        x = torch.ops.aten.native_group_norm(
-            x,
-            weight,
-            bias,
-            batch_size,
-            num_channels,
-            flattened_inner_size,
-            num_groups,
-            eps,
-        )[0]
+        x = self.group_norm(x)
         x = torch.nn.AvgPool2d(2)(x)
         return x
 
     model = SampleModel().eval()
-    forward_args = lambda: (
-        torch.rand(16, 640, 32, 32) * 1000,
-        torch.rand([640]) * 1000,
-        torch.rand([640]) * 1000,
+    forward_args = lambda: (torch.rand(16, 640, 32, 32) * 1000,)
+    exported_program = export_with_pass(model, forward_args())
+    self.assert_outputs_allclose(
+        model, exported_program.module(), forward_args()
     )
+
+  def test_group_norm_with_composite_enabled(self):
+    ai_edge_torch.config.enable_group_norm_composite = True
+
+    class SampleModel(torch.nn.Module):
+
+      def __init__(self):
+        super().__init__()
+        self.group_norm = torch.nn.GroupNorm(
+            num_groups=2, num_channels=10, affine=True
+        )
+
+      def forward(self, x):
+        x = torch.nn.AvgPool2d(2)(x)
+        x = self.group_norm(x)
+        x = torch.nn.AvgPool2d(2)(x)
+        return x
+
+    model = SampleModel().eval()
+    forward_args = lambda: (torch.rand(1, 10, 32, 32),)
     exported_program = export_with_pass(model, forward_args())
     self.assert_outputs_allclose(
         model, exported_program.module(), forward_args()
     )
 
+    ir_text = lowertools.exported_program_to_mlir_text(exported_program)
+    self.assertEqual(ir_text.count("stablehlo.custom_call @mark_tensor"), 4)
+
 
-if __name__ == '__main__':
+if __name__ == "__main__":
   googletest.main()