fix type bug

support group norm, and improve batch and layer norms

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -47,6 +47,29 @@ def get_ir(target: Target) -> SourceIR:
 
 
 @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],
+    )
+
+
 @dynamo_tensorrt_converter(torch.ops.aten.batch_norm)  # type: ignore[misc]
 def aten_ops_batch_norm(
     ctx: ConversionContext,
@@ -68,20 +91,19 @@ def aten_ops_batch_norm(
         training=args[5],
         momentum=args[6],
         eps=args[7],
-        cudnn_enabled=args_bounds_check(args, 8, replacement=True),
+        cudnn_enabled=args[8],
     )
 
 
 @dynamo_tensorrt_converter(torch.ops.aten.native_layer_norm.default)  # type: ignore[misc]
-@dynamo_tensorrt_converter(torch.ops.aten.layer_norm.default)  # type: ignore[misc]
-def aten_ops_layer_norm(
+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.layer_norm(
+    return impl.normalization.native_layer_norm(
         ctx,
         target,
         SourceIR.ATEN,
@@ -91,7 +113,74 @@ def aten_ops_layer_norm(
         weight=args[2],
         bias=args[3],
         eps=args[4],
-        cudnn_enable=args_bounds_check(args, 5, replacement=True),
+    )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.layer_norm.default)  # type: ignore[misc]
+def aten_ops_layer_norm(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.normalization.layer_norm(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        input=args[0],
+        normalized_shape=args[1],
+        weight=args_bounds_check(args, 2),
+        bias=args_bounds_check(args, 3),
+        eps=args_bounds_check(args, 4, 1e-05),
+        cudnn_enable=args_bounds_check(args, 5, True),
+    )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.native_group_norm.default)  # type: ignore[misc]
+def aten_ops_native_group_norm(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.normalization.native_group_norm(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        input=args[0],
+        weight=args[1],
+        bias=args[2],
+        N=args[3],
+        C=args[4],
+        HxW=args[5],
+        group=args[6],
+        eps=args[7],
+    )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.group_norm.default)  # type: ignore[misc]
+def aten_ops_group_norm(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.normalization.group_norm(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        input=args[0],
+        num_groups=args[1],
+        weight=args_bounds_check(args, 2, None),
+        bias=args_bounds_check(args, 3, None),
+        eps=args_bounds_check(args, 4, 1e-05),
+        cudnn_enabled=args_bounds_check(args, 5, True),
     )