chore: address review comments

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

py/torch_tensorrt/_compile.py

@@ -73,7 +73,7 @@ def _get_target_ir(module_type: _ModuleType, ir: str) -> _IRType:
                 )
                 return _IRType.ts
             else:
-                raise ValueError("Module was provided with in an unsupported format")
+                raise ValueError("Module was provided in an unsupported format")
         else:
             raise ValueError("Unknown ir was requested")
 
@@ -157,37 +157,38 @@ def compile(
             **kwargs,
         )
     elif target_ir == _IRType.dynamo:
+        from torch_tensorrt import Device
+        from torch_tensorrt.dynamo.utils import prepare_inputs, prepare_device
+        import collections.abc
+
+        if not isinstance(inputs, collections.abc.Sequence):
+            inputs = [inputs]
+        device = kwargs.get("device", Device._current_device())
+        torchtrt_inputs, torch_inputs = prepare_inputs(inputs, prepare_device(device))
+        module = torch_tensorrt.dynamo.trace(module, torch_inputs, **kwargs)
         return torch_tensorrt.dynamo.compile(
             module,
             inputs=inputs,
             enabled_precisions=enabled_precisions,
             **kwargs,
         )
     elif target_ir == _IRType.torch_compile:
-        return torch_compile(
-            module, inputs, enabled_precisions=enabled_precisions, **kwargs
-        )
+        return torch_compile(module, enabled_precisions=enabled_precisions, **kwargs)
     else:
         raise RuntimeError("Module is an unknown format or the ir requested is unknown")
 
 
-def torch_compile(module, inputs, **kwargs):
-
-    from torch_tensorrt.dynamo.utils import prepare_inputs, prepare_device
+def torch_compile(module, **kwargs):
+    """
+    Returns a boxed model which is the output of torch.compile.
+    This does not compile the model to TRT. Execute this model on
+    sample inputs to compile the model to TRT.
+    """
     from torch_tensorrt.dynamo.backend import torch_tensorrt_backend
-    from torch_tensorrt import Device
-    import collections.abc
-
-    if not isinstance(inputs, collections.abc.Sequence):
-        inputs = [inputs]
 
-    device = kwargs.get("device", Device._current_device())
-    torchtrt_inputs, torch_inputs = prepare_inputs(inputs, prepare_device(device))
-    model = torch.compile(module, backend=torch_tensorrt_backend, options={**kwargs})
-    # Ensure compilation occurs by calling the function with provided inputs
-    model(*torch_inputs)
+    boxed_fn = torch.compile(module, backend=torch_tensorrt_backend, options={**kwargs})
 
-    return model
+    return boxed_fn
 
 
 def convert_method_to_trt_engine(
@@ -246,6 +247,16 @@ def convert_method_to_trt_engine(
             **kwargs,
         )
     elif target_ir == _IRType.fx:
-        raise RuntimeError("fx is currently not supported")
+        raise RuntimeError(
+            "convert_method_to_trt_engine call is not supported for ir=fx"
+        )
+    elif target_ir == _IRType.dynamo:
+        raise RuntimeError(
+            "convert_method_to_trt_engine call is not supported for ir=dynamo."
+        )
+    elif target_ir == _IRType.torch_compile:
+        raise RuntimeError(
+            "convert_method_to_trt_engine call is not supported for ir=torch_compile"
+        )
     else:
         raise RuntimeError("Module is an unknown format or the ir requested is unknown")

py/torch_tensorrt/dynamo/compile.py

@@ -9,7 +9,6 @@
 from torch.fx.passes.pass_manager import PassManager
 from torch.fx.passes.shape_prop import ShapeProp
 from torch.fx.passes.splitter_base import SplitResult
-from torch_tensorrt.dynamo.aten_tracer import trace
 from torch_tensorrt.fx.tools.trt_splitter import TRTSplitter, TRTSplitterSetting
 from torch_tensorrt.dynamo.lowering import (
     fuse_permute_linear,
@@ -38,7 +37,7 @@
 
 
 def compile(
-    gm: torch.nn.Module,
+    gm: Any,
     inputs: Any,
     *,
     device=Device._current_device(),
@@ -113,13 +112,11 @@ def compile(
     }
 
     settings = CompilationSettings(**compilation_options)
-    model = trace(gm, torch_inputs, **kwargs)
-
     if kwargs.get("use_capability_partitioner", None):
-        model = lower_model(model, torch_inputs)
+        model = lower_model(gm, torch_inputs)
         return _compile_module(model, torch_inputs, settings)
     else:
-        split_result = lower_model_using_trt_splitter(model, torch_inputs)
+        split_result = lower_model_using_trt_splitter(gm, torch_inputs)
         trt_module = _compile_graph(split_result, torch_inputs, settings)
 
         return trt_module