[Relax][PyTorch] Add decomposed operator support for normalization (#18460)

## Related Issue

- #18401

## How

This PR

- added `_batch_norm_legit_no_stats`
- added `_native_group_norm`
- added `any.dims`
- refctored `_reshape`

Author: guan404ming

python/tvm/relax/frontend/torch/base_fx_graph_translator.py

@@ -1848,6 +1848,12 @@ def _reshape(self, node: fx.Node) -> relax.Var:
         args = self.retrieve_args(node)
         x = args[0]
         dims = args[1] if isinstance(args[1], (torch.Size, tuple, list)) else args[1:]
+
+        # Skip identity reshape
+        current_shape = self.shape_of(x)
+        if list(current_shape) == list(dims):
+            return x
+
         return self.block_builder.emit(relax.op.reshape(x, dims))
 
     def _reshape_as(self, node: fx.Node) -> relax.Var:

python/tvm/relax/frontend/torch/exported_program_translator.py

@@ -113,6 +113,31 @@ def _batch_norm_legit_no_training(self, node: fx.Node) -> relax.Var:
         training = False
         return self._batch_norm(node, training)
 
+    def _batch_norm_legit_no_stats(self, node: fx.Node) -> relax.Var:
+        import numpy as np
+
+        x = self.env[node.args[0]]
+        channel = int(self.shape_of(x)[1])
+        dtype = x.struct_info.dtype
+        weight = self.env.get(node.args[1], relax.const(np.ones(channel), dtype=dtype))
+        bias = self.env.get(node.args[2], relax.const(np.zeros(channel), dtype=dtype))
+        eps = node.args[5] if len(node.args) > 5 else node.kwargs.get("eps", 1e-05)
+
+        # Determine axes for instance norm (all spatial dimensions after channel)
+        dim = len(self.shape_of(x))
+        axes = list(range(2, dim))
+
+        return self.block_builder.emit(
+            relax.op.nn.instance_norm(
+                x,
+                weight,
+                bias,
+                channel_axis=1,
+                axes=axes,
+                epsilon=eps,
+            )
+        )
+
     def _cross_entropy_default(self, node: fx.Node) -> relax.Expr:
         preds = self.env[node.args[0]]
         targets = self.env[node.args[1]]
@@ -141,6 +166,28 @@ def _group_norm(self, node: fx.Node) -> relax.Var:
             )
         )
 
+    def _native_group_norm(self, node: fx.Node) -> relax.Var:
+        # native_group_norm signature: (input, weight, bias, N, C, HxW, group, eps)
+        x = self.env[node.args[0]]
+        gamma = self.env.get(node.args[1], None) if len(node.args) > 1 else None
+        beta = self.env.get(node.args[2], None) if len(node.args) > 2 else None
+        # args[3] = N (batch size), args[4] = C (channels), args[5] = HxW (spatial size)
+        num_groups = node.args[6] if len(node.args) > 6 else 1
+        eps = node.args[7] if len(node.args) > 7 else 1e-05
+
+        dim = len(self.shape_of(x))
+        return self.block_builder.emit(
+            relax.op.nn.group_norm(
+                x,
+                gamma,
+                beta,
+                num_groups=num_groups,
+                channel_axis=1,
+                axes=list(range(2, dim)),
+                epsilon=eps,
+            )
+        )
+
     def _upsample_impl(
         self,
         x: relax.Expr,
@@ -963,6 +1010,7 @@ def create_convert_map(
             "_adaptive_avg_pool3d.default": self._adaptive_avg_pool3d,
             "_native_batch_norm_legit_functional.default": self._batch_norm_legit_functional,
             "_native_batch_norm_legit_no_training.default": self._batch_norm_legit_no_training,
+            "_native_batch_norm_legit.no_stats": self._batch_norm_legit_no_stats,
             "batch_norm.default": self._batch_norm_legit_no_training,
             "adaptive_avg_pool1d.default": self._adaptive_avg_pool1d,
             "adaptive_avg_pool2d.default": self._adaptive_avg_pool2d,
@@ -988,6 +1036,7 @@ def create_convert_map(
             ),
             "group_norm.default": self._group_norm,
             "instance_norm.default": self._instance_norm,
+            "native_group_norm.default": self._native_group_norm,
             "layer_norm.default": self._layer_norm,
             "linear.default": self._linear,
             "lstm.input": self._lstm,
@@ -1004,6 +1053,7 @@ def create_convert_map(
             "upsample_bicubic2d.vec": self._upsample_bicubic2d,
             # statistical
             "any.dim": self._any,
+            "any.dims": self._any,
             "mean.dim": self._mean,
             "prod.default": self._prod,
             "std.correction": self._std,

tests/python/relax/test_frontend_from_exported_program.py

@@ -1514,20 +1514,18 @@ def main(
             x: R.Tensor((10, 10), dtype="float32"), test_elements: R.Tensor((8,), dtype="float32")
         ) -> R.Tuple(R.Tensor((10, 10), dtype="bool")):
             with R.dataflow():
-                lv: R.Tensor((10, 10, 1), dtype="float32") = R.expand_dims(x, axis=[-1])
-                lv1: R.Tensor((8,), dtype="float32") = R.reshape(test_elements, R.shape([8]))
-                lv2: R.Tensor((10, 10, 8), dtype="bool") = R.equal(lv, lv1)
-                lv3: R.Tensor((10, 10), dtype="bool") = R.sum(lv2, axis=[-1], keepdims=False)
-                lv4: R.Tensor((10, 10), dtype="bool") = R.greater(lv3, R.const(0.0, "float32"))
-                gv: R.Tuple(R.Tensor((10, 10), dtype="bool")) = (lv4,)
+                lv: R.Tensor((10, 10, 1), dtype="float32") = R.reshape(x, R.shape([10, 10, 1]))
+                lv1: R.Tensor((10, 10, 8), dtype="bool") = R.equal(lv, test_elements)
+                lv2: R.Tensor((10, 10), dtype="bool") = R.max(lv1, axis=[-1], keepdims=False)
+                gv: R.Tuple(R.Tensor((10, 10), dtype="bool")) = (lv2,)
                 R.output(gv)
             return gv
 
     example_args = (
         torch.randn(10, 10, dtype=torch.float32),
         torch.randn(8, dtype=torch.float32),
     )
-    verify_model(IsInModel(), example_args, {}, expected)
+    verify_model(IsInModel(), example_args, {}, expected, run_ep_decomposition=True)
 
 
 def test_div_mode():
@@ -3155,7 +3153,7 @@ def main(
         "w1": model.gn.weight.detach().numpy(),
         "w2": model.gn.bias.detach().numpy(),
     }
-    verify_model(model, example_args, binding, expected1)
+    verify_model(model, example_args, binding, expected1, run_ep_decomposition=True)
 
 
 def test_instancenorm2d():
@@ -3200,7 +3198,7 @@ def main(
         "w1": torch.ones(3).detach().numpy(),
         "w2": torch.zeros(3).detach().numpy(),
     }
-    verify_model(model, example_args, binding, expected1)
+    verify_model(model, example_args, binding, expected1, run_ep_decomposition=True)
 
 
 def test_layernorm():
@@ -5556,7 +5554,9 @@ def main(
 
     example_args = (torch.randn(256, 256, dtype=torch.float32),)
     exported_program = export(Identity(), args=example_args)
-    mod = from_exported_program(exported_program, unwrap_unit_return_tuple=True)
+    mod = from_exported_program(
+        exported_program, unwrap_unit_return_tuple=True, run_ep_decomposition=True
+    )
     tvm.ir.assert_structural_equal(mod, Expected)
 
 
@@ -5586,7 +5586,9 @@ def main(
         torch.randn(256, 256, dtype=torch.float32),
     )
     exported_program = export(Identity(), args=example_args)
-    mod = from_exported_program(exported_program, no_bind_return_tuple=True)
+    mod = from_exported_program(
+        exported_program, no_bind_return_tuple=True, run_ep_decomposition=True
+    )
     tvm.ir.assert_structural_equal(mod, Expected)