Add LowerSimpleTuples for freeze tuples (#57915)

FloydHsiu · facebook-github-bot · commit f4a5730a6b89 · 2021-05-12T19:07:20.000-07:00
Summary: Fixes #57698 Follow the suggestion mentioned in #57698 add a call to LowerSimpleTuples after the call: https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/passes/freeze_module.cpp#L89. Pull Request resolved: #57915 Reviewed By: agolynski Differential Revision: D28387310 Pulled By: nikithamalgifb fbshipit-source-id: 5becb608c5352240b30dfdf03a821d0297e9609c
diff --git a/test/jit/test_freezing.py b/test/jit/test_freezing.py
@@ -1416,6 +1416,35 @@ def forward(self, input):
         self.assertEqual(model(inp), script_model(inp))
         FileCheck().check_not("GetAttr").run(script_model.graph)
 
+    def test_freeze_module_with_tupleoutput_submodule(self):
+        class SubModule(nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                return (x + 1, x + 2)
+
+        class TestModule(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.sub = SubModule()
+
+            def forward(self, x):
+                y1, y2 = self.sub(x)
+                return y1 + y2
+
+        m = torch.jit.script(TestModule())
+        m = m.eval()
+        mf = torch.jit.freeze(m)
+        inp = torch.randn(2, 2)
+        expected = m.forward(inp)
+        output = mf.forward(inp)
+        # Check if prim::TupleConstruct and prim::TupleUnpack
+        # Don't exist in frozen graph
+        FileCheck().check_not("prim::TupleConstruct").run(mf.graph)
+        FileCheck().check_not("prim::TupleUnpack").run(mf.graph)
+        self.assertEqual(output, expected)
+
 class TestFrozenOptimizations(JitTestCase):
     def setUp(self):
         self.default_dtype = torch.get_default_dtype()
diff --git a/torch/csrc/jit/passes/freeze_module.cpp b/torch/csrc/jit/passes/freeze_module.cpp
@@ -5,6 +5,7 @@
 #include <torch/csrc/jit/ir/alias_analysis.h>
 #include <torch/csrc/jit/passes/clear_profiling.h>
 #include <torch/csrc/jit/passes/inliner.h>
+#include <torch/csrc/jit/passes/lower_tuples.h>
 #include <torch/csrc/jit/runtime/graph_executor_impl.h>
 
 #include <stack>
@@ -88,6 +89,7 @@ class AttributePropagator {
     auto applyOptimizations = [](std::shared_ptr<Graph>& subgraph) {
       runOptimization(
           subgraph, /* unroll? */ false, /* const_prop_user_classes? */ false);
+      LowerSimpleTuples(subgraph);
     };
 
     for (auto function : preservedMethods_) {
