[acc_tracer] Remove tensor_meta after NormalizeArgs (#16)

Summary: Pull Request resolved: pytorch/fx2trt#16

Reviewed By: yinghai, wushirong

Differential Revision: D34772842

fbshipit-source-id: fd2577e54b1e5f563df2e1052773ef7b19069abe

Author: jfix71

test/tracer/test_acc_tracer.py

@@ -1989,11 +1989,11 @@ def forward(self, a: List[torch.Tensor]) -> torch.Tensor:
             else:
                 self.fail(f"Unexpected node: {node.format_node()}")
 
-        # Check the tensor metadatas are correct given the input is a list.
-        self.assertTrue(isinstance(ph.meta["tensor_meta"], list))
-        self.assertEqual(len(ph.meta["tensor_meta"]), 2)
-        self.assertEqual(getitem_0.meta["tensor_meta"], ph.meta["tensor_meta"][0])
-        self.assertEqual(getitem_1.meta["tensor_meta"], ph.meta["tensor_meta"][1])
+        # Check the tensor ranks are correct given the input is a list.
+        self.assertTrue(isinstance(ph.meta["tensor_rank"], list))
+        self.assertEqual(len(ph.meta["tensor_rank"]), 2)
+        self.assertEqual(getitem_0.meta["tensor_rank"], ph.meta["tensor_rank"][0])
+        self.assertEqual(getitem_1.meta["tensor_rank"], ph.meta["tensor_rank"][1])
 
         self.assertTrue(torch.equal(m(input), traced(input)))
 

tracer/acc_tracer/acc_ops.py

@@ -320,8 +320,7 @@ def cat(*, tensors, dim):
 )
 def transpose_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
     # Get the dim-permutation/shuffle
-    shape_as_list = node.meta["tensor_meta"].shape
-    ranks = len(shape_as_list)
+    ranks = node.meta["tensor_rank"]
     shuffle = list(range(ranks))
     dim0 = cast(int, node.kwargs["dim0"])
     dim1 = cast(int, node.kwargs["dim1"])
@@ -422,7 +421,7 @@ def addmm_mapper(node: torch.fx.Node, _: nn.Module) -> torch.fx.Node:
     ],
 )
 def t_mapper(node: torch.fx.Node, _: nn.Module):
-    ranks = len(node.meta["tensor_meta"].shape)
+    ranks = node.meta["tensor_rank"]
     shuffle = [1, 0] if (ranks > 1) else [0]
 
     with node.graph.inserting_before(node):
@@ -1842,7 +1841,7 @@ def packed_quantized_linear_mapper(
         }
 
         new_node = node.graph.call_function(quantized_linear, kwargs=kwargs)
-        new_node.meta = node.meta
+        new_node.meta = node.meta.copy()
         return new_node
 
 

tracer/acc_tracer/acc_tracer.py

@@ -10,6 +10,7 @@
 
 import fx2trt_oss.tracer.acc_tracer.acc_normalizer as acc_normalizer
 import fx2trt_oss.tracer.acc_tracer.acc_ops  # noqa: F401
+import fx2trt_oss.tracer.acc_tracer.acc_utils as acc_utils
 import torch
 import torch.jit as jit
 import torch.nn as nn
@@ -249,7 +250,11 @@ def trace(
 }
 
 
-def _rewrite(mod_to_rewrite: nn.Module, allow_list: Optional[Set] = None, leaf_module_list: Optional[Set] = None) -> nn.Module:
+def _rewrite(
+    mod_to_rewrite: nn.Module,
+    allow_list: Optional[Set] = None,
+    leaf_module_list: Optional[Set] = None,
+) -> nn.Module:
     if allow_list is None:
         allow_list = DEFAULT_REWRITE_ALLOW_LIST
     else:
@@ -270,7 +275,7 @@ def rewrite_module(m: nn.Module):
             return m
 
         # If m is an already-rewritten RewrittenModule, then use the original base class.
-        base_class : Type[nn.Module] = getattr(m, "_base_class_origin", type(m))
+        base_class: Type[nn.Module] = getattr(m, "_base_class_origin", type(m))
 
         # Keep track of all the ConditionalExceptionWrappers that the
         # Acc_Rewriter calls into in this module so we can add them in init
@@ -290,7 +295,9 @@ class RewrittenModule(base_class):  # type: ignore[valid-type, misc]
             for method_name in dir(base_class):
                 method = getattr(base_class, method_name, None)
                 if method is None and method_name not in {"__doc__"}:
-                    _LOGGER.warning(f"{__qualname__} does not have attribute {method_name}")
+                    _LOGGER.warning(
+                        f"{__qualname__} does not have attribute {method_name}"
+                    )
 
                 if builtins.type(method) is not FunctionType:
                     continue
@@ -368,6 +375,15 @@ def _remove_exceptions(gm: torch.fx.GraphModule) -> bool:
     return changed
 
 
+def _replace_tensor_meta_with_rank(gm: torch.fx.GraphModule):
+    for node in gm.graph.nodes:
+        if node.op != "output" and "tensor_meta" in node.meta:
+            node.meta["tensor_rank"] = acc_utils.map_tensor_metadata(
+                node.meta["tensor_meta"], lambda x: len(x.shape)
+            )
+            del node.meta["tensor_meta"]
+
+
 def trace(
     mod: nn.Module,
     sample_inputs: Sequence[Any],
@@ -450,9 +466,13 @@ def trace(
     # nodes after removing assertions and exceptions.
     traced.graph.eliminate_dead_code()
 
+    # Run shape prop to add node.meta["type"] to nodes, needed for NormalizeArgs.
+    shape_prop.ShapeProp(traced).propagate(*sample_inputs)
+    # Swap out tensor_meta for tensor_rank, because we don't actually want to rely on
+    # tensor_meta yet for normalization/lowering, though rank shouldn't change.
+    _replace_tensor_meta_with_rank(traced)
     # Now normalize args/kwargs to make default values visible. Leave args/kwargs as
     # they were, since all-kwarg normalization is broken, and we don't need it anyway.
-    shape_prop.ShapeProp(traced).propagate(*sample_inputs)
     traced = NormalizeArgs(traced, normalize_to_only_use_kwargs=False).transform()
 
     # Normalize to acc-specialized wrappers for consistency across op naming and

tracer/acc_tracer/acc_utils.py

@@ -1,17 +1,18 @@
 import inspect
 import json
 import os
-from typing import Any, Tuple, Callable, Union, Dict, List, Optional
 import re
+from typing import Any, Tuple, Callable, Union, Dict, List, Optional
 
 import torch
 import torch.fx
-from torch.fx.passes.graph_manipulation import (
-    serialize_module,
-)
 from torch.fx.graph_module import GraphModule
+from torch.fx.immutable_collections import immutable_list
 from torch.fx.node import _get_qualified_name
 from torch.fx.passes import graph_drawer
+from torch.fx.passes.graph_manipulation import (
+    serialize_module,
+)
 from torch.fx.passes.shape_prop import TensorMetadata
 
 
@@ -173,3 +174,18 @@ def get_unique_attr_name_in_module(mod_traced: torch.fx.GraphModule, name: str)
             name = f"{base}_{int(num) + 1}"
 
     return name
+
+
+def map_tensor_metadata(a: Any, fn: Callable):
+    """
+    Map some `fn` to `a`, where `a` is either a TensorMetadata, or else a tuple/list
+    recursively containing TensorMetadata.
+    """
+    if isinstance(a, TensorMetadata):
+        return fn(a)
+    elif isinstance(a, tuple):
+        return tuple(map_tensor_metadata(elem, fn) for elem in a)
+    assert isinstance(
+        a, list
+    ), f"Only supporting tuple/list/TensorMetadata, but found {type(a)}"
+    return immutable_list(map_tensor_metadata(elem, fn) for elem in a)