update batch_norm and layer_norm

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -31,14 +31,37 @@ def aten_ops_batch_norm(
         target,
         SourceIR.ATEN,
         name,
-        args[0],
-        args[1],
-        args[2],
-        args[3],
-        args[4],
-        args[5],
-        args[6],
-        args[7],
+        input=args[0],
+        weight=args_bounds_check(args, 1, replacement=1),
+        bias=args_bounds_check(args, 2, replacement=0),
+        running_mean=args_bounds_check(args, 3),
+        running_var=args_bounds_check(args, 4),
+        training=args_bounds_check(args, 5),
+        momentum=args_bounds_check(args, 6, replacement=0.1),
+        eps=args_bounds_check(args, 7, replacement=1e-05),
+        cudnn_enabled=args_bounds_check(args, 8, replacement=False),
+    )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.layer_norm.default)  # type: ignore[misc]
+def aten_ops_layer_norm(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.normalization.layer_norm(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        input=args[0],
+        normalized_shape=args[1],
+        weight=args_bounds_check(args, 2, replacement=1),
+        bias=args_bounds_check(args, 3, replacement=0),
+        eps=args_bounds_check(args, 4, replacement=1e-05),
+        cudnn_enable=args_bounds_check(args, 5, replacement=True),
     )
 
 
@@ -258,27 +281,6 @@ def aten_ops_matmul(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.layer_norm.default)  # type: ignore[misc]
-def aten_ops_layernorm(
-    network: TRTNetwork,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return impl.normalization.layer_norm(
-        network,
-        target,
-        SourceIR.ATEN,
-        name,
-        args[0],
-        args[1],
-        args[2],
-        args[3],
-        args[4],
-    )
-
-
 @dynamo_tensorrt_converter(torch.ops.aten.rsqrt.default)  # type: ignore[misc]
 def aten_ops_rsqrt(
     network: TRTNetwork,

py/torch_tensorrt/dynamo/conversion/impl/normalization/ops.py

@@ -36,6 +36,7 @@ def batch_norm(
     training: torch.Tensor,
     momentum: torch.Tensor,
     eps: List[float],
+    cudnn_enabled: bool,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
     if not isinstance(input, TRTTensor):
         raise RuntimeError(
@@ -69,16 +70,16 @@ def batch_norm(
                 input.shape[2],
                 1,
             )
-        set_layer_name(reshape_layer, target, f"{name}_reshape_2d")
+        set_layer_name(reshape_layer, target, f"{name}_reshape_2d", source_ir)
         input = reshape_layer.get_output(0)
     layer = network.add_scale(input, trt.ScaleMode.CHANNEL, bias, scale, power)
-    set_layer_name(layer, target, name)
+    set_layer_name(layer, target, name, source_ir)
 
     # For BatchNorm1d, reshape output back to 1d
     if not network.has_implicit_batch_dimension and len(output_shape) < 4:
         reshape_output_layer = network.add_shuffle(layer.get_output(0))
         reshape_output_layer.reshape_dims = tuple(output_shape)
-        set_layer_name(reshape_output_layer, target, f"{name}_reshape_1d")
+        set_layer_name(reshape_output_layer, target, f"{name}_reshape_1d", source_ir)
         layer = reshape_output_layer
     return layer.get_output(0)
 
@@ -93,6 +94,7 @@ def layer_norm(
     weight: torch.Tensor,
     bias: torch.Tensor,
     eps: List[float],
+    cudnn_enable: bool,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
     if not isinstance(input, trt.tensorrt.ITensor):
         raise RuntimeError(
@@ -173,7 +175,7 @@ def layer_norm_no_plugin(
     mean_expected_layer = network.add_reduce(
         input, trt.ReduceOperation.AVG, axes, keep_dims=True
     )
-    set_layer_name(mean_expected_layer, target, f"{name}_mean_expected")
+    set_layer_name(mean_expected_layer, target, f"{name}_mean_expected", source_ir)
 
     # X-E[x]
     sub_trt = convert_binary_elementwise(
@@ -203,7 +205,7 @@ def layer_norm_no_plugin(
     mean_trt_layer = network.add_reduce(
         pow_var, trt.ReduceOperation.AVG, axes, keep_dims=True
     )
-    set_layer_name(mean_trt_layer, target, f"{name}_mean")
+    set_layer_name(mean_trt_layer, target, f"{name}_mean", source_ir)
     # Variance + eps
     eps_tensor = network.add_constant(
         (1,) * len(input.shape),