Add some logging to how we handle dynamic shapes

Author: narendasan

py/torch_tensorrt/dynamo/backend/backends.py

@@ -10,6 +10,7 @@
 from torch._dynamo.backends.common import aot_autograd
 from torch._dynamo.utils import detect_fake_mode
 from torch._functorch.aot_autograd import aot_export_joint_simple
+
 from torch_tensorrt._utils import is_tegra_platform
 from torch_tensorrt.dynamo import CompilationSettings
 from torch_tensorrt.dynamo._compiler import compile_module
@@ -118,9 +119,10 @@ def _pretraced_backend(
         fake_mode = detect_fake_mode(sample_inputs)
 
         # Place backend tracing within FakeTensor context allowing nonfake Tensors
-        with unittest.mock.patch.object(
-            fake_mode, "allow_non_fake_inputs", True
-        ), fake_mode:
+        with (
+            unittest.mock.patch.object(fake_mode, "allow_non_fake_inputs", True),
+            fake_mode,
+        ):
             repair_input_aliasing(gm, settings)
 
             # Remove sym_int placeholders and inputs
@@ -170,7 +172,7 @@ def _pretraced_backend(
                 engine_cache=engine_cache,
             )
             return trt_compiled
-    except (AssertionError, RuntimeError):
+    except (AssertionError, RuntimeError, TypeError):
         if not settings.pass_through_build_failures:
             logger.warning(
                 "TRT conversion failed on the subgraph. See trace above. "

py/torch_tensorrt/dynamo/partitioning/common.py

@@ -4,6 +4,7 @@
 import torch
 from torch._subclasses.fake_tensor import FakeTensor
 from torch.fx.experimental.proxy_tensor import unset_fake_temporarily
+
 from torch_tensorrt._Input import Input
 from torch_tensorrt.dynamo.utils import contains_sym_int, extract_var_range_info
 
@@ -26,11 +27,21 @@ def construct_dynamic_input(
     min_shape = []
     opt_shape = []
     max_shape = []
-    for dim in input_shape:
+    for d, dim in enumerate(input_shape):
         if isinstance(dim, torch.SymInt):
             min_max_opt = extract_var_range_info(dim)
+            if min_max_opt["max"] is None:
+                logger.warning(
+                    f"Dynamic input {name} (shape: {input_shape}) has no max bound for dim {d}, attempting to use a sane default (max: min({min_max_opt['min']}) * 128). Please set an upper bound using torch._dynamo.mark_dynamic or torch.export.Dim"
+                )
+                min_max_opt["max"] = min_max_opt["min"] * 128
             min_shape.append(min_max_opt["min"])
             # if opt not exist, set it to the mean of min and max
+            if min_max_opt["opt"] is None:
+                logger.info(
+                    f"Dynamic input {name} (shape: {input_shape}) has no opt target i.e. which shape to specialize for, for dim {d}, attempting to use a sane default (opt: min({min_max_opt['min']}) + max({min_max_opt['max']}) / 2). If you want to specialized further, use torch_tensorrt.compile"
+                )
+                min_max_opt["opt"] = int(min_max_opt["min"] + min_max_opt["max"] / 2)
             opt_shape.append(
                 min_max_opt.get("opt", int(min_max_opt["min"] + min_max_opt["max"] / 2))
             )
@@ -61,7 +72,10 @@ def get_input(
     """
     if contains_sym_int(input_shape):
         return construct_dynamic_input(
-            input_shape, dtype, name=name, is_shape_tensor=is_shape_tensor
+            input_shape,
+            dtype,
+            name=name,
+            is_shape_tensor=is_shape_tensor,
         )
     else:
         return Input(

py/torch_tensorrt/dynamo/utils.py

@@ -26,6 +26,9 @@
 import torch
 from torch._subclasses.fake_tensor import FakeTensor
 from torch.fx.experimental.proxy_tensor import unset_fake_temporarily
+from torch.utils._sympy.numbers import int_oo
+
+from packaging import version
 from torch_tensorrt._Device import Device
 from torch_tensorrt._enums import dtype
 from torch_tensorrt._features import ENABLED_FEATURES
@@ -36,8 +39,6 @@
 from torch_tensorrt.dynamo._engine_cache import BaseEngineCache
 from torch_tensorrt.dynamo._settings import CompilationSettings
 
-from packaging import version
-
 from .types import TRTDataType
 
 logger = logging.getLogger(__name__)
@@ -424,7 +425,11 @@ def extract_var_range_info(symbolic_integer: torch.SymInt) -> Dict[str, int]:
         or expr.xreplace(shape_env.var_to_val)
     )
     assert var_range, var_val
-    min_val, max_val = int(var_range.lower), int(var_range.upper)
+    min_val, max_val = (
+        int(var_range.lower),
+        int(var_range.upper) if var_range.upper != int_oo else None,
+    )
+
     # Torchdynamo 0/1 specialization outlier
     min_val = 1 if min_val == 2 else min_val
     min_max_opt = {}
@@ -699,8 +704,9 @@ def check_module_output(
     arg_inputs: Any,
     kwarg_inputs: Any = None,
 ) -> bool:
-    old_outputs, new_outputs = refitted_module(*arg_inputs), new_module(
-        *arg_inputs, **kwarg_inputs
+    old_outputs, new_outputs = (
+        refitted_module(*arg_inputs),
+        new_module(*arg_inputs, **kwarg_inputs),
     )
     if type(old_outputs) != type(new_outputs):
         logger.warning("The output types are different. Output check is skipped.")
@@ -803,9 +809,9 @@ def copy_metadata(match_and_replacements: List[Any]) -> None:
     """
     for match_and_replacement in match_and_replacements:
         anchor_node = match_and_replacement.nodes_map[match_and_replacement.anchor]
-        assert (
-            len(match_and_replacement.replacements) == 1
-        ), "Found more than 1 replacements for the anchor node."
+        assert len(match_and_replacement.replacements) == 1, (
+            "Found more than 1 replacements for the anchor node."
+        )
         replacement_node = match_and_replacement.replacements[0]
         replacement_node.meta = anchor_node.meta