chore: add update_metadata utility

py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py

@@ -11,6 +11,7 @@
 from .remove_input_alias_fixing_clones import remove_input_alias_fixing_clones
 from .repair_input_as_output import repair_input_as_output
 from .replace_max_pool_with_indices import replace_max_pool_with_indices
+from .view_to_reshape import view_to_reshape
 
 ATEN_LOWERING_PASSES = DynamoPassManager.build_from_passlist(
     [
@@ -21,6 +22,7 @@
         lower_linear,
         fuse_prims_broadcast,
         replace_max_pool_with_indices,
+        view_to_reshape,
     ]
 )
 

py/torch_tensorrt/dynamo/lowering/passes/pass_utils.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import Any, Dict, List
 
 import torch
 
@@ -29,3 +29,20 @@ def get_tensor_placeholders(
     ]
 
     return placeholders
+
+
+def update_metadata(
+    gm: torch.fx.GraphModule, target_op: Any, metadata: Dict[int, torch._ops.OpOverload]
+) -> None:
+    """
+    Given a graph and a node which has target_op in the graph,
+    a) If the node has metadata, store it in the map
+    b) If the node does not have metadata, retrieve it from the map
+       and assign to the node.
+    """
+    for idx, node in enumerate(gm.graph.nodes):
+        if node.target == target_op:
+            if idx not in metadata and node.meta:
+                metadata[idx] = node.meta
+            elif idx in metadata and not node.meta:
+                node.meta = metadata[idx]

py/torch_tensorrt/dynamo/lowering/passes/view_to_reshape.py

@@ -1,9 +1,10 @@
 import logging
-from typing import Callable, List, Sequence, Tuple
+from typing import Dict, List, Sequence
 
 import torch
 from torch_tensorrt.dynamo.lowering.passes.pass_utils import (
     clean_up_graph_after_modifications,
+    update_metadata,
 )
 
 logger = logging.getLogger(__name__)
@@ -13,29 +14,25 @@ def view_to_reshape(
     gm: torch.fx.GraphModule, sample_inputs: Sequence[torch.Tensor]
 ) -> torch.fx.GraphModule:
     """Replace aten.view with an equivalent implementation which avoids Tensor memory issues"""
-    orig, replacement = view_replacement()
-
-    if torch.fx.subgraph_rewriter.replace_pattern(gm, orig, replacement):
-        gm = clean_up_graph_after_modifications(gm)
-        logger.debug(f"Graph after replacing view with reshape:\n{gm.graph}")
-
-    return gm
-
-
-def view_replacement() -> (
-    Tuple[
-        torch.fx.GraphModule,
-        Callable[[torch.Tensor, List[torch.SymInt]], torch.Tensor],
-    ]
-):
-    """Constructs the original and replacement functions for view"""
+    orig_op = torch.ops.aten.view.default
+    replacement_op = torch.ops.aten.reshape.default
 
     # Original graph
     def orig(input: torch.Tensor, shape: List[torch.SymInt]) -> torch.Tensor:
-        return torch.ops.aten.view.default(input, shape)
+        return orig_op(input, shape)
 
     # Replacement graph
     def replacement(input: torch.Tensor, shape: List[torch.SymInt]) -> torch.Tensor:
-        return torch.ops.aten.reshape.default(input, shape)
+        return replacement_op(input, shape)
 
-    return orig, replacement
+    # Store metadata of the orig_op and copy it to the replacement op
+    meta_map: Dict[int, torch._ops.OpOverload] = {}
+    update_metadata(gm, orig_op, meta_map)
+
+    if torch.fx.subgraph_rewriter.replace_pattern(gm, orig, replacement):
+        gm = clean_up_graph_after_modifications(gm)
+        logger.debug(f"Graph after replacing view with reshape:\n{gm.graph}")
+
+    update_metadata(gm, replacement_op, meta_map)
+
+    return gm

py/torch_tensorrt/dynamo/partitioning/common.py

@@ -43,9 +43,9 @@ def construct_dynamic_input(input: Any) -> Input:
             if var_range.lower == 2:
                 min_shape.append(1)
             else:
-                min_shape.append(var_range.lower)
-            opt_shape.append(var_val)
-            max_shape.append(var_range.upper)
+                min_shape.append(int(var_range.lower))
+            opt_shape.append(int(var_val))
+            max_shape.append(int(var_range.upper))
         else:
             min_shape.append(dim)
             opt_shape.append(dim)

tests/py/dynamo/models/test_dyn_models.py

@@ -3,9 +3,8 @@
 import pytest
 import timm
 import torch
-from torch_tensorrt.dynamo.utils import COSINE_THRESHOLD, cosine_similarity
-
 import torch_tensorrt as torchtrt
+from torch_tensorrt.dynamo.utils import COSINE_THRESHOLD, cosine_similarity
 
 assertions = unittest.TestCase()