[ts_converter] Fix prim::If buffer names (#136648)

angelayi · AnantGulati · commit fe007e4b9fc7 · 2024-10-02T14:29:03.000+03:00
Summary: We previously incorrectly handled the following graph, specifically for the node `w.3` in `block0`: ``` graph(%x.1 : Float(3, strides=[1], requires_grad=0, device=cpu), %y.1 : int): %2 : __torch__.___torch_mangle_1.M = prim::CreateObject() %3 : int = prim::Constant[value=20](), scope: M:: # /data/users/angelayi/pytorch/test/export/test_converter.py:747:34 %4 : int = prim::Constant[value=10](), scope: M:: # /data/users/angelayi/pytorch/test/export/test_converter.py:746:34 %5 : int = prim::Constant[value=1](), scope: M:: %w.1 : int = prim::GetAttr[name="w"](%2), scope: M:: %7 : int = aten::mul(%w.1, %4), scope: M:: # /data/users/angelayi/pytorch/test/export/test_converter.py:746:25 = prim::SetAttr[name="w"](%2, %7), scope: M:: %h.1 : int = prim::GetAttr[name="h"](%2), scope: M:: %9 : int = aten::mul(%h.1, %3), scope: M:: # /data/users/angelayi/pytorch/test/export/test_converter.py:747:25 = prim::SetAttr[name="h"](%2, %9), scope: M:: %10 : bool = aten::gt(%y.1, %4), scope: M:: # /data/users/angelayi/pytorch/test/export/test_converter.py:749:19 %res.37 : Tensor = prim::If(%10), scope: M:: # /data/users/angelayi/pytorch/test/export/test_converter.py:749:16 block0(): %w.3 : int = prim::GetAttr[name="w"](%2), scope: M:: %res.1 : Tensor = aten::add(%x.1, %w.3, %5), scope: M:: # <string>:5:9 -> (%res.1) block1(): %h.3 : int = prim::GetAttr[name="h"](%2), scope: M:: %res.3 : Tensor = aten::add(%x.1, %h.3, %5), scope: M:: # <string>:5:9 -> (%res.3) %16 : bool = aten::lt(%y.1, %4), scope: M:: # /data/users/angelayi/pytorch/test/export/test_converter.py:754:19 %res : Tensor = prim::If(%16), scope: M:: # /data/users/angelayi/pytorch/test/export/test_converter.py:754:16 block0(): %w : int = prim::GetAttr[name="w"](%2), scope: M:: %res.15 : Tensor = aten::add(%res.37, %w, %5), scope: M:: # <string>:5:9 -> (%res.15) block1(): %h : int = prim::GetAttr[name="h"](%2), scope: M:: %res.21 : Tensor = aten::add(%res.37, %h, %5), scope: M:: # <string>:5:9 -> (%res.21) return (%res) ``` Test Plan: CI Differential Revision: D63399064 Pull Request resolved: #136648 Approved by: https://github.com/SherlockNoMad
diff --git a/test/export/test_converter.py b/test/export/test_converter.py
@@ -838,6 +838,32 @@ def forward(self, x: torch.Tensor):
                 orig_m(*inp),
             )
 
+    def test_convert_if_duplicate_attr_names(self):
+        class M(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.w = 1
+                self.h = 2
+
+            def forward(self, x: torch.Tensor, y: int):
+                self.w = self.w * 10
+                self.h = self.h * 20
+
+                if y > 10:
+                    res = self.w + x
+                else:
+                    res = self.h + x
+
+                if y < 10:
+                    res = self.w + res
+                else:
+                    res = self.h + res
+
+                return res
+
+        inp = (torch.ones(3), 5)
+        self._check_equal_ts_ep_converter(M(), inp, option=["script"])
+
     def test_ts2ep_converter_contains(self):
         class MIn(torch.nn.Module):
             def forward(self, x: torch.Tensor):
diff --git a/torch/_export/converter.py b/torch/_export/converter.py
@@ -260,7 +260,9 @@ def construct_fqn(ir, ref_map, name_map):
     return ".".join(reversed(name_list))
 
 
-def get_block_to_lifted_attrs(graph: torch._C.Graph) -> Dict[torch._C.Block, Set[str]]:
+def get_block_to_lifted_attrs(
+    graph: torch._C.Graph,
+) -> Tuple[Dict[torch._C.Block, Set[str]], Dict[str, str]]:
     """
     Perform two passes to get a mapping of blocks to a set of FQNs of its lifted attributes.
     When a graph has control flow, the graph will be divided into multiple blocks. We want to convert
@@ -272,7 +274,8 @@ def get_block_to_lifted_attrs(graph: torch._C.Graph) -> Dict[torch._C.Block, Set
         of the attributes used in the current block, and the lifted attributes of all its child blocks.
 
     Returns:
-        A mapping of blocks to a set of FQNs of its lifted attributes.
+        A mapping of blocks to a set of FQNs of its lifted attributes, and a
+        mapping of node names to the FQNs of its lifted attributes.
     """
 
     # A map from a block to its expected to be lifted arguments.
@@ -334,7 +337,7 @@ def _map_blocks_to_lifted_attrs(entry):
     _dfs_get_attr_dependency(graph)
     _map_blocks_to_lifted_attrs(graph)
 
-    return blocks_to_lifted_attrs
+    return blocks_to_lifted_attrs, node_to_attr_name
 
 
 def get_attribute_fqn_from_ts_node(
@@ -393,22 +396,28 @@ def __init__(
         blocks_to_lifted_attrs: Dict[torch._C.Block, Set[str]],
         name_to_non_tensor_attribute: Dict[str, Any],
         name_to_constant: Dict[str, Any],
+        name_to_attribute_fqn: Dict[str, str],
     ):
         self.ts_graph = ts_graph
+        # Mapping of parameter FQN to actual parameter value
         self.name_to_param = name_to_param
+        # Mapping of buffer FQN to actual buffer value
         self.name_to_buffer = name_to_buffer
 
         self.fx_graph: torch.fx.Graph = torch.fx.Graph()
         self.input_specs: List[InputSpec] = []
         self.output_specs: List[OutputSpec] = []
 
+        # Mapping of TS node name to converted FX node
         self.name_to_node: Dict[
             str, Union[torch.fx.Node, List[torch.fx.Node], Dict[Any, torch.fx.Node]]
         ] = {}
+        # Mapping of TS node name to constant value (int, str, TorchBind obj,
+        # tensor constants ...)
         self.name_to_constant: Dict[str, Any] = name_to_constant
 
         # Mapping from torchscript node output name to attribute fully qualified name
-        self.name_to_attribute_fqn: Dict[str, str] = {}
+        self.name_to_attribute_fqn: Dict[str, str] = name_to_attribute_fqn
 
         # Mapping from fully qualified name to real values or a fx graph node
         # During convert, this represents the current value of a non-tensor attribute
@@ -427,6 +436,8 @@ def __init__(
 
         self.subgraphs: Dict[str, torch.fx.GraphModule] = {}
 
+        # Mapping of block to list of attributes that need to be lifted for each
+        # block
         self.blocks_to_lifted_attrs = blocks_to_lifted_attrs
 
         # Populate methods for the standard operators.
@@ -467,8 +478,8 @@ def _convert_block_to_subgraph(self, node: torch._C.Node, arguments: List[str]):
                 self.blocks_to_lifted_attrs,
                 {},
                 self.name_to_constant,
+                self.name_to_attribute_fqn,
             )
-            subgraph_converter.name_to_attribute_fqn = self.name_to_attribute_fqn
 
             for block_arg in arguments:
                 normalized_block_arg_name = normalize_name(block_arg)
@@ -537,6 +548,8 @@ def get_fx_value_by_ir_value(self, value: torch._C.Value):
             if isinstance(self.name_to_constant[value_name], torch.ScriptObject):
                 return self.fx_graph.get_attr(value_name)
             return self.name_to_constant[value_name]
+        elif value_name in self.name_to_attribute_fqn:
+            return self.get_fx_value_by_fqn(self.name_to_attribute_fqn[value_name])
         else:
             raise ValueError(f"Input {value_name} not found")
 
@@ -1325,6 +1338,7 @@ def __init__(
         blocks_to_lifted_attrs: Dict[torch._C.Block, Set[str]],
         name_to_non_tensor_attribute: Dict[str, Any],
         name_to_constant: Dict[str, Any],
+        name_to_attribute_fqn: Dict[str, str],
     ):
         super().__init__(
             ts_graph,
@@ -1333,6 +1347,7 @@ def __init__(
             blocks_to_lifted_attrs,
             name_to_non_tensor_attribute,
             name_to_constant,
+            name_to_attribute_fqn,
         )
 
         # Data to keep track of unsupported nodes.
@@ -1427,7 +1442,9 @@ def convert(self) -> ExportedProgram:
         )
         log.info("TorchScript graph\n\n%s\n", self.ts_graph)
 
-        blocks_to_lifted_attrs = get_block_to_lifted_attrs(self.ts_graph)
+        blocks_to_lifted_attrs, name_to_attribute_fqn = get_block_to_lifted_attrs(
+            self.ts_graph
+        )
 
         graph_converter = TS2FXGraphConverter(
             self.ts_graph,
@@ -1436,6 +1453,7 @@ def convert(self) -> ExportedProgram:
             blocks_to_lifted_attrs,
             self.name_to_non_tensor_attributes,
             self.name_to_constant,
+            name_to_attribute_fqn,
         )
         gm = graph_converter.convert()
 
@@ -1464,7 +1482,9 @@ def convert(self) -> ExportedProgram:
 
     @disable_logging(log)
     def explain(self, print_output=True):
-        blocks_to_lifted_attrs = get_block_to_lifted_attrs(self.ts_graph)
+        blocks_to_lifted_attrs, name_to_attribute_fqn = get_block_to_lifted_attrs(
+            self.ts_graph
+        )
 
         graph_converter = ExplainTS2FXGraphConverter(
             self.ts_graph,
@@ -1473,6 +1493,7 @@ def explain(self, print_output=True):
             blocks_to_lifted_attrs,
             self.name_to_non_tensor_attributes,
             self.name_to_constant,
+            name_to_attribute_fqn,
         )
         graph_converter.explain()
         if len(graph_converter.unsupported_node_list) > 0:
