fix: Add I/O values for nodes

Author: gs-olive

py/torch_tensorrt/dynamo/_compiler.py

@@ -307,7 +307,6 @@ def contains_metadata(gm: torch.fx.GraphModule) -> bool:
     # Partition module into components that can be TRT-accelerated
     fast_partitioner_failed = False
 
-    logger.info("Beginning TensorRT operator Partitioning Phase")
     # If specified, try using the fast partitioner and fall back to the global one on failure
     if settings.use_fast_partitioner:
         try:
@@ -343,11 +342,6 @@ def contains_metadata(gm: torch.fx.GraphModule) -> bool:
     if not settings.use_fast_partitioner:
         dryrun_tracker.to_run_in_torch.extend(parse_non_trt_nodes(partitioned_module))
 
-    logger.info(
-        "Successfully completed graph partitioning phase. "
-        "Beginning the conversion phase."
-    )
-
     # Store TRT replicas of Torch subgraphs
     trt_modules = {}
     # Iterate over all components that can be accelerated

py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py

@@ -1,7 +1,7 @@
 import logging
 import warnings
 from datetime import datetime
-from typing import Any, Callable, Dict, List, NamedTuple, Optional, Sequence, Set
+from typing import Any, Callable, Dict, List, NamedTuple, Optional, Sequence, Set, Tuple
 
 import numpy as np
 import tensorrt as trt
@@ -20,6 +20,7 @@
 )
 from torch_tensorrt.dynamo.conversion._ConverterRegistry import CallingConvention
 from torch_tensorrt.dynamo.conversion.converter_utils import (
+    get_node_io,
     get_node_name,
     get_trt_tensor,
 )
@@ -106,6 +107,9 @@ def __init__(
             [dtype._from(o) for o in output_dtypes] if output_dtypes else None
         )
 
+        # Mapping of constants to shapes and dtypes
+        self.const_mapping: Dict[str, Tuple[Sequence[int], str]] = {}
+
     def validate_conversion(self) -> Set[str]:
         missing_converters: Set[str] = set()
 
@@ -352,6 +356,13 @@ def run_node(self, n: torch.fx.Node) -> torch.fx.Node:
         )
         trt_node: torch.fx.Node = super().run_node(n)
 
+        if n.op == "get_attr":
+            self.const_mapping[str(n)] = (tuple(trt_node.shape), str(trt_node.dtype))
+
+        _LOGGER.debug(
+            f"Ran node {self._cur_node_name} with properties: {get_node_io(n, self.const_mapping)}"
+        )
+
         # remove "_itensor_to_tensor_meta"
         kwargs = dict(n.kwargs)
         del kwargs["_itensor_to_tensor_meta"]

py/torch_tensorrt/dynamo/conversion/converter_utils.py

@@ -8,6 +8,7 @@
 import torch
 import torch_tensorrt.dynamo.conversion.impl as impl
 from torch.fx.node import Argument, Target
+from torch.fx.passes.shape_prop import TensorMetadata
 from torch_tensorrt import _enums
 from torch_tensorrt.dynamo._SourceIR import SourceIR
 from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
@@ -44,6 +45,63 @@ def get_node_name(node: torch.fx.Node) -> str:
     return node_name
 
 
+def get_node_io(
+    node: torch.fx.Node, constant_mapping: Dict[str, Tuple[Sequence[int], str]]
+) -> str:
+    """Gets a string representing the node inputs and outputs including tensor shapes and dtypes"""
+
+    def format_tensor_metadata(
+        metadata: Union[TensorMetadata, Sequence[TensorMetadata]]
+    ) -> str:
+        """Formats the metadata for a single node"""
+        # If the provided data is a simple TensorMetadata object, parse it
+        if isinstance(metadata, TensorMetadata):
+            return f"{tuple(metadata.shape)}@{metadata.dtype}"
+        # If the provided data is a sequence, recursively parse it
+        else:
+            formatted_str = "("
+            for meta in metadata:
+                formatted_str += format_tensor_metadata(meta) + ", "
+
+            return formatted_str[:-2] + ")"
+
+    # Format input tensors
+    metadata_string = "Inputs: ("
+
+    # For each input argument, format it accordingly
+    for arg in node.args:
+        if isinstance(arg, torch.fx.Node):
+            if arg.op == "get_attr":
+                shape, dtype = constant_mapping[str(arg)]
+                arg_repr = f"{shape}@{dtype}"
+            elif arg.meta.get("tensor_meta", False):
+                arg_repr = format_tensor_metadata(arg.meta["tensor_meta"])
+            else:
+                arg_repr = ""
+
+            metadata_string += f"{arg}: {arg_repr}, "
+        else:
+            metadata_string += f"{arg}, "
+
+    metadata_string = (
+        metadata_string[:-2] if metadata_string[-1] != "(" else metadata_string
+    ) + ")"
+
+    # Format output tensors and arguments
+    metadata_string += " | Outputs: ("
+    if node.op == "get_attr":
+        shape, dtype = constant_mapping[str(node)]
+        node_repr = f"{shape}@{dtype}"
+    elif node.meta.get("tensor_meta", False):
+        node_repr = format_tensor_metadata(node.meta["tensor_meta"])
+    else:
+        node_repr = ""
+    metadata_string += f"{node}: {node_repr}, "
+    metadata_string = metadata_string[:-2] + ")"
+
+    return metadata_string
+
+
 def is_only_operator_on_placeholder(node: torch.fx.Node) -> bool:
     """Detects whether a call_function node is the only operator on a placeholder"""
     # Returns true if the node operates on a placeholder and is a direct output

py/torch_tensorrt/dynamo/lowering/passes/propagate_shapes.py

@@ -0,0 +1,24 @@
+import logging
+from typing import Sequence
+
+import torch
+from torch.fx.passes.shape_prop import ShapeProp
+
+logger = logging.getLogger(__name__)
+
+
+def propagate_shapes(
+    gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor]
+) -> torch.fx.GraphModule:
+    """Attempts to propagate shapes through the graph"""
+
+    # Propagate shapes through the graph
+    try:
+        ShapeProp(gm).propagate(*sample_inputs)
+    except (RuntimeError, AssertionError):
+        logger.warning(
+            "Shape Propagation Failed on Graph, skipping propagate_shapes lowering pass",
+            exc_info=True,
+        )
+
+    return gm