rebase and update three norms

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -46,31 +46,22 @@ def get_ir(target: Target) -> SourceIR:
     return SourceIR.UNKNOWN
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.native_batch_norm.default)  # type: ignore[misc]
-def aten_ops_native_batch_norm(
-    ctx: ConversionContext,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return impl.normalization.native_batch_norm(
-        ctx,
-        target,
-        SourceIR.ATEN,
-        name,
-        input=args[0],
-        weight=args[1],
-        bias=args[2],
-        running_mean=args[3],
-        running_var=args[4],
-        training=args[5],
-        momentum=args[6],
-        eps=args[7],
+def one_user_validator(node: Node) -> bool:
+    # Validate only one user, which is a getitem node that accesses the first element in the list
+    return (
+        len(node.users) == 1
+        and list(node.users)[0].target == operator.getitem
+        and list(node.users)[0].args[1] == 0
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.batch_norm)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.native_batch_norm.default, capability_validator=one_user_validator)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.batch_norm.default)  # type: ignore[misc]
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+    }
+)  # type: ignore[misc]
 def aten_ops_batch_norm(
     ctx: ConversionContext,
     target: Target,
@@ -91,32 +82,18 @@ def aten_ops_batch_norm(
         training=args[5],
         momentum=args[6],
         eps=args[7],
-        cudnn_enabled=args[8],
-    )
-
-
-@dynamo_tensorrt_converter(torch.ops.aten.native_layer_norm.default)  # type: ignore[misc]
-def aten_ops_native_layer_norm(
-    ctx: ConversionContext,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return impl.normalization.native_layer_norm(
-        ctx,
-        target,
-        SourceIR.ATEN,
-        name,
-        input=args[0],
-        normalized_shape=args[1],
-        weight=args[2],
-        bias=args[3],
-        eps=args[4],
+        cudnn_enabled=args_bounds_check(args, 8, True),
+        return_mean_rstd=(target == torch.ops.aten.native_batch_norm.default),
     )
 
 
+@dynamo_tensorrt_converter(torch.ops.aten.native_layer_norm.default, capability_validator=one_user_validator)  # type: ignore[misc]
 @dynamo_tensorrt_converter(torch.ops.aten.layer_norm.default)  # type: ignore[misc]
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+    }
+)  # type: ignore[misc]
 def aten_ops_layer_norm(
     ctx: ConversionContext,
     target: Target,
@@ -135,10 +112,16 @@ def aten_ops_layer_norm(
         bias=args_bounds_check(args, 3),
         eps=args_bounds_check(args, 4, 1e-05),
         cudnn_enable=args_bounds_check(args, 5, True),
+        return_mean_rstd=(target == torch.ops.aten.native_layer_norm.default),
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.native_group_norm.default)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.native_group_norm.default, capability_validator=one_user_validator)  # type: ignore[misc]
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+    }
+)  # type: ignore[misc]
 def aten_ops_native_group_norm(
     ctx: ConversionContext,
     target: Target,
@@ -163,6 +146,11 @@ def aten_ops_native_group_norm(
 
 
 @dynamo_tensorrt_converter(torch.ops.aten.group_norm.default)  # type: ignore[misc]
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+    }
+)  # type: ignore[misc]
 def aten_ops_group_norm(
     ctx: ConversionContext,
     target: Target,
@@ -838,15 +826,6 @@ def aten_ops_prod(
     )
 
 
-def one_user_validator(node: Node) -> bool:
-    # Validate only one user, which is a getitem node that accesses the first element in the list
-    return (
-        len(node.users) == 1
-        and list(node.users)[0].target == operator.getitem
-        and list(node.users)[0].args[1] == 0
-    )
-
-
 @dynamo_tensorrt_converter(torch.ops.aten.max.default)  # type: ignore[misc]
 @dynamo_tensorrt_converter(torch.ops.aten.max.dim, capability_validator=one_user_validator)  # type: ignore[misc]
 def aten_ops_max(