Add enable_group_norm_composite global flag

PiperOrigin-RevId: 712602061

Author: chunnienc

ai_edge_torch/_config.py

@@ -22,6 +22,18 @@
 __all__ = ["config"]
 
 
+def _get_bool_env_var(name: str, default: bool) -> bool:
+  var = os.environ.get(name, "false")
+  var = var.lower().strip()
+  if var in ("y", "yes", "t", "true", "on", "1"):
+    return True
+  elif var in ("n", "no", "f", "false", "off", "0"):
+    return False
+  else:
+    logging.warning("Invalid %s value is ignored: %s.", name, var)
+    return default
+
+
 class _Config:
   """ai-edge-torch global configs."""
 
@@ -33,20 +45,25 @@ def use_torch_xla(self) -> bool:
     To use torch_xla as the lowering backend, set environment variable
     `USE_TORCH_XLA` to "true".
     """
-    var = os.environ.get("USE_TORCH_XLA", "false")
-    var = var.lower().strip()
-    if var in ("y", "yes", "t", "true", "on", "1"):
-      return True
-    elif var in ("n", "no", "f", "false", "off", "0"):
-      return False
-    else:
-      logging.warning("Invalid USE_TORCH_XLA value is ignored: %s.", var)
-      return False
+    return _get_bool_env_var("USE_TORCH_XLA", default=False)
 
   @property
   def in_oss(self) -> bool:
     """True if the code is not running in google internal environment."""
     return True
 
+  @property
+  def enable_group_norm_composite(self) -> bool:
+    """True if lowering group norm in StableHLO composite.
+
+    Currently only supports NHWC group norm generated by
+    OptimizeLayoutTransposesPass.
+    """
+    return _get_bool_env_var("ENABLE_GROUP_NORM_COMPOSITE", default=False)
+
+  @enable_group_norm_composite.setter
+  def enable_group_norm_composite(self, value: bool):
+    os.environ["ENABLE_GROUP_NORM_COMPOSITE"] = "y" if value else "n"
+
 
 config = _Config()

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_check.py

@@ -17,6 +17,7 @@
 import dataclasses
 import operator
 
+import ai_edge_torch
 from ai_edge_torch import lowertools
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_rewrite
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import utils
@@ -194,7 +195,10 @@ def _aten_group_norm_checker(node):
   val = node.meta.get("val")
   if not hasattr(val, "shape"):
     return NHWCable(can_be=False, must_be=False)
-  return NHWCable(can_be=len(val.shape) == 4, must_be=False)
+
+  can_be = len(val.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)
 
 
 @nhwcable_node_checkers.register(aten.native_group_norm)

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_rewrite.py

@@ -16,13 +16,15 @@
 
 import operator
 
+import ai_edge_torch
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_mark
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import op_func_registry
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import utils
 import torch
 import torch.utils._pytree as pytree
 
 aten = torch.ops.aten
+StableHLOCompositeBuilder = ai_edge_torch.hlfb.StableHLOCompositeBuilder
 
 __all__ = ["rewrite_nhwc_node", "has_nhwc_rewriter"]
 
@@ -345,11 +347,32 @@ def batch_norm(input, weight, bias, running_mean, running_var, momentum, eps):
 @rewriters.register(aten.group_norm.default)
 def _aten_group_norm(node):
   def group_norm(input, num_groups: int, weight=None, bias=None, eps=1e-5):
-    # Disable NHWC rewriter with native decomposied ops due to precision issue.
-    # TODO(b/354780253): Re-enable NHWC rewriter with proper lowering.
+    is_composite_supported = (
+        ai_edge_torch.config.enable_group_norm_composite
+        and weight is not None
+        and bias is not None
+    )
+
+    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)
+
     input = utils.tensor_to_nchw(input)
-    res = aten.group_norm.default(input, num_groups, weight, bias, eps=eps)
-    return utils.tensor_to_nhwc(res)
+    output = aten.group_norm.default(input, num_groups, weight, bias, eps=eps)
+    output = utils.tensor_to_nhwc(output)
+
+    if builder is not None:
+      output = builder.mark_outputs(output)
+    return output
 
   node.target = group_norm
 

ai_edge_torch/generative/layers/normalization.py

@@ -80,6 +80,7 @@ def forward(self, x):
       output = self._norm(x.float()).type_as(x)
       return output * w
 
+
 class GroupNorm(torch.nn.Module):
 
   def __init__(
@@ -115,16 +116,7 @@ def forward(self, x):
     Returns:
       torch.Tensor: output tensor after applying GroupNorm.
     """
-    if self.enable_hlfb:
-      return group_norm_with_hlfb(
-          x,
-          self.weight,
-          self.bias,
-          self.group_num,
-          self.eps,
-      )
-    else:
-      return F.group_norm(x, self.group_num, self.weight, self.bias, self.eps)
+    return F.group_norm(x, self.group_num, self.weight, self.bias, self.eps)
 
 
 class LayerNorm(torch.nn.Module):
@@ -169,46 +161,6 @@ def forward(self, x):
     )
 
 
-def group_norm_with_hlfb(
-    x: torch.Tensor,
-    w: torch.Tensor,
-    b: torch.Tensor,
-    num_groups: int,
-    eps: float,
-):
-  """Group Normalization with high-level function boundary enabled.
-
-  Args:
-    x (torch.Tensor): Input tensor for Group Normalization, with BCHW shape.
-    w (torch.Tensor): The weight tensor for the normalization.
-    b (torch.Tensor): The bias tensor for the normalization.
-    num_groups (int): Number of groups to separate the channels into.
-    eps (float): A small float value to ensure numerical stability.
-
-  Returns:
-    The output tensor of Group Normalization.
-  """
-  x = torch.permute(x, (0, 2, 3, 1))
-
-  builder = StableHLOCompositeBuilder(
-      name="odml.group_norm",
-      attr={
-          "num_groups": num_groups,
-          "epsilon": eps,
-          "reduction_axes": [3],
-          "channel_axis": 3,
-      },
-  )
-  x, w, b = builder.mark_inputs(x, w, b)
-  x = torch.permute(x, (0, 3, 1, 2))
-  y = F.group_norm(x, num_groups, weight=w, bias=b, eps=eps)
-  y = torch.permute(y, (0, 2, 3, 1))
-  y = builder.mark_outputs(y)
-
-  y = torch.permute(y, (0, 3, 1, 2))
-  return y
-
-
 def rms_norm_with_hlfb(
     x: torch.Tensor,
     w: torch.Tensor,