rebase

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -138,12 +138,12 @@ def aten_ops_sigmoid(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.index.Tensor)
+@dynamo_tensorrt_converter(torch.ops.aten.index.Tensor)  # type: ignore[misc]
 @enforce_tensor_types(
     {
         0: (TRTTensor,),
     }
-)
+)  # type: ignore[misc]
 def aten_ops_index(
     ctx: ConversionContext,
     target: Target,
@@ -723,27 +723,8 @@ def one_user_validator(node: Node) -> bool:
 
 
 @dynamo_tensorrt_converter(torch.ops.aten.max.default)  # type: ignore[misc]
-def aten_ops_max(
-    ctx: ConversionContext,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return impl.reduce.max(
-        ctx,
-        target,
-        SourceIR.ATEN,
-        name,
-        args[0],
-        dim=None,
-        keepdim=False,
-        return_indices=False,
-    )
-
-
 @dynamo_tensorrt_converter(torch.ops.aten.max.dim, capability_validator=one_user_validator)  # type: ignore[misc]
-def aten_ops_max_dim(
+def aten_ops_max(
     ctx: ConversionContext,
     target: Target,
     args: Tuple[Argument, ...],
@@ -756,34 +737,15 @@ def aten_ops_max_dim(
         SourceIR.ATEN,
         name,
         args[0],
-        args[1],
-        args_bounds_check(args, 2, replacement=False),
-        return_indices=True,
+        dim=args_bounds_check(args, 1, replacement=None),
+        keepdim=args_bounds_check(args, 2, replacement=False),
+        return_indices=(target == torch.ops.aten.max.dim),
     )
 
 
 @dynamo_tensorrt_converter(torch.ops.aten.min.default)  # type: ignore[misc]
-def aten_ops_min(
-    ctx: ConversionContext,
-    target: Target,
-    args: Tuple[Argument, ...],
-    kwargs: Dict[str, Argument],
-    name: str,
-) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return impl.reduce.min(
-        ctx,
-        target,
-        SourceIR.ATEN,
-        name,
-        args[0],
-        dim=None,
-        keepdim=False,
-        return_indices=False,
-    )
-
-
 @dynamo_tensorrt_converter(torch.ops.aten.min.dim, capability_validator=one_user_validator)  # type: ignore[misc]
-def aten_ops_min_dim(
+def aten_ops_min(
     ctx: ConversionContext,
     target: Target,
     args: Tuple[Argument, ...],
@@ -796,9 +758,9 @@ def aten_ops_min_dim(
         SourceIR.ATEN,
         name,
         args[0],
-        args[1],
-        args_bounds_check(args, 2, replacement=False),
-        return_indices=True,
+        dim=args_bounds_check(args, 1, replacement=None),
+        keepdim=args_bounds_check(args, 2, replacement=False),
+        return_indices=(target == torch.ops.aten.min.dim),
     )
 
 

tests/py/dynamo/conversion/harness.py

@@ -197,35 +197,18 @@ def generate_graph(
         use_dynamo_tracer: bool,
         enable_passes: bool,
     ):
-        # Torchdynamo+aot proxytensor tracer
-        # Below are common passes
-        passes_list = [
-            compose_bmm,
-            compose_chunk,
-            compose_getitem_slice,
-            replace_aten_reshape_alias_with_replace,
-            replace_aten_op_with_indices,
-            replace_transpose_mm_op_with_linear,  # after compose_bmm
-            replace_native_layernorm_with_layernorm,
-            remove_ops,
-            replace_builtin_ops,  # after replace_native_layernorm_with_layernorm
-        ]
-        # Combine with customized passes specific to any model
-        if customized_passes:
-            passes_list.extend(customized_passes)
-
-        if disable_passes:
-            passes_list = []
-
-        fx_module, _ = aten_tracer.trace(mod, original_inputs)
-        for passes in passes_list:
-            pr: PassResult = passes(fx_module)
-            fx_module = pr.graph_module
-        fx_module(*original_inputs)
-
-        fx_module = run_const_fold(fx_module)
-        fx_module.graph.eliminate_dead_code()
-
+        if use_dynamo_tracer:
+            fx_module = torch._dynamo.export(
+                mod,
+                *original_inputs,
+                aten_graph=True,
+                assume_static_by_default=True,
+                tracing_mode="real",
+            ).graph_module
+        else:
+            fx_module = torch.fx.symbolic_trace(mod)
+        if enable_passes:
+            fx_module = apply_lowering_passes(fx_module, original_inputs)
         _LOGGER.info(f"FX graph= {fx_module.graph}")
         return fx_module
 

tests/py/dynamo/conversion/test_max_aten.py

@@ -18,13 +18,12 @@ class TestMaxConverter(DispatchTestCase):
     def test_max_dim_int_default(self, input_shape):
         class Max(nn.Module):
             def forward(self, x):
-                return torch.max(x)
+                return torch.ops.aten.max.default(x)
 
         inputs = [torch.randn(*input_shape)]
         self.run_test(
             Max(),
             inputs,
-            expected_ops={torch.ops.aten.max.default},
         )
 
     @parameterized.expand(
@@ -45,8 +44,6 @@ def forward(self, x):
         self.run_test(
             Max(),
             inputs,
-            expected_ops={torch.ops.aten.max.dim},
-            disable_passes=True,
         )
 
     @parameterized.expand(
@@ -65,9 +62,7 @@ def forward(self, x):
         self.run_test(
             Max(),
             inputs,
-            expected_ops={torch.ops.aten.max.dim},
             check_dtype=False,
-            disable_passes=True,
         )
 
 

tests/py/dynamo/conversion/test_min_aten.py

@@ -24,7 +24,6 @@ def forward(self, x):
         self.run_test(
             Min(),
             inputs,
-            expected_ops={torch.ops.aten.min.default},
         )
 
     @parameterized.expand(
@@ -39,14 +38,12 @@ def forward(self, x):
     def test_min_dim_int(self, input_shape, dim, keep_dims):
         class Min(nn.Module):
             def forward(self, x):
-                return torch.min(x, dim=dim, keepdim=keep_dims)[0]
+                return torch.ops.aten.min.dim(x, dim, keep_dims)[0]
 
         inputs = [torch.randn(*input_shape)]
         self.run_test(
             Min(),
             inputs,
-            expected_ops={torch.ops.aten.min.dim},
-            disable_passes=True,
         )
 
     @parameterized.expand(
@@ -59,15 +56,13 @@ def forward(self, x):
     def test_min_dim_int_int(self, input_shape, dim, keep_dims, dtype, low, high):
         class Min(nn.Module):
             def forward(self, x):
-                return torch.ops.aten.min.dim(x, dim=dim, keepdim=keep_dims)[0]
+                return torch.ops.aten.min.dim(x, dim, keep_dims)[0]
 
         inputs = [torch.randint(low, high, input_shape, dtype=dtype)]
         self.run_test(
             Min(),
             inputs,
-            expected_ops={torch.ops.aten.min.dim},
             check_dtype=False,
-            disable_passes=True,
         )