[jit] [shape analysis] Move constant tensors out of fused subgraphs during generalization (pytorch#70320)

Summary:
Pull Request resolved: pytorch#70320

ghstack-source-id: 146514368

Test Plan: `buck test mode/dev-nosan //caffe2/test/cpp/jit:jit`

Reviewed By: eellison

Differential Revision: D33280508

fbshipit-source-id: fe4291d7c49f0a498b330de96b698e99f6f6a505

Author: navahgar

test/cpp/jit/test_shape_analysis.cpp

@@ -7,6 +7,7 @@
 #include <torch/csrc/jit/ir/ir.h>
 #include <torch/csrc/jit/ir/ir_views.h>
 #include <torch/csrc/jit/ir/irparser.h>
+#include <torch/csrc/jit/passes/constant_propagation.h>
 #include <torch/csrc/jit/passes/symbolic_shape_runtime_fusion.h>
 #include <torch/csrc/jit/passes/utils/subgraph_utils.h>
 #include <torch/csrc/jit/runtime/graph_iterator.h>
@@ -246,5 +247,46 @@ TEST(ShapeAnalysisTest, DynamicShapesFusion) {
   }
 }
 
+TEST(ShapeAnalysisTest, MovingConstantOutOfFusionGroups) {
+  std::shared_ptr<Graph> subgraph = std::make_shared<Graph>();
+  const auto graph_string = R"IR(
+      graph(%x.1 : Tensor):
+        %none : NoneType = prim::Constant()
+        %size1 : int = prim::Constant[value=1]()
+        %size10 : int = prim::Constant[value=10]()
+        %sizes : int[] = prim::ListConstruct(%size10, %size1)
+        %device : Device = prim::Constant[value="cpu"]()
+        %10 : Tensor = aten::ones(%sizes, %none, %none, %device, %none)
+        %3 : Tensor = aten::tanh(%x.1)
+        %29 : Tensor = aten::mul(%3, %10)
+        return (%29))IR";
+  torch::jit::parseIR(graph_string, subgraph.get());
+  ConstantPropagation(subgraph);
+
+  std::shared_ptr<Graph> g = std::make_shared<Graph>();
+  auto x_inp = g->addInput("x_inp");
+  auto x_type = TensorType::create(at::rand({10, 5}));
+  x_inp->setType(x_type);
+  subgraph->inputs().at(0)->setType(x_type);
+  auto output = g->insertNode(g->create(prim::TensorExprGroup))->output();
+  output->node()->addInput(x_inp);
+  output->node()->g_(attr::Subgraph, subgraph);
+
+  auto success = GenerateGuard(output->node());
+  TORCH_INTERNAL_ASSERT(success);
+
+  // Check that the constants have been moved out of the fused graph.
+  // This should result in not have any conditionals other than the one
+  // checking the result of TensorExprDynamicGuard.
+  testing::FileCheck()
+      .check("TensorExprDynamicGuard")
+      ->check_next("prim::If")
+      ->check_not("prim::If") // no other IFs due to constants.
+      ->check("TensorExprGroup")
+      ->check("block1")
+      ->check("FallbackGraph")
+      ->run(*g);
+}
+
 } // namespace jit
 } // namespace torch

torch/csrc/jit/passes/symbolic_shape_runtime_fusion.cpp

@@ -106,9 +106,58 @@ bool TryGeneralizeInputDimensionsToSymbolicShapes(
   return true;
 }
 
+void moveConstantTensorsOutOfSubgraph(
+    Node* tensorexpr_graph_node,
+    std::shared_ptr<Graph> tensorexpr_graph) {
+  auto parent = tensorexpr_graph_node->owningGraph();
+
+  auto env = [&](Value* v) {
+    TORCH_INTERNAL_ASSERT(
+        false,
+        "this should never happen since constant nodes do not have any inputs",
+        v->debugName());
+    return v;
+  };
+
+  WithInsertPoint wip(tensorexpr_graph_node);
+  std::vector<Node*> to_destroy;
+  for (auto node : tensorexpr_graph->nodes()) {
+    if (node->kind() == prim::Constant) {
+      if (!node->output()->type()->cast<TensorType>()) {
+        continue;
+      }
+
+      // copy the constant and insert that copy into the parent graph.
+      auto copy = parent->createClone(node, env);
+      parent->insertNode(copy);
+
+      // add a new input to the te subgraph and replace the uses of the
+      // constant with this input.
+      auto new_const = tensorexpr_graph->addInput();
+      new_const->setType(node->output()->type());
+      node->output()->replaceAllUsesWith(new_const);
+
+      // add the copy as input to the te node
+      tensorexpr_graph_node->addInput(copy->output());
+
+      to_destroy.push_back(node);
+    }
+  }
+
+  for (auto n : to_destroy) {
+    n->destroy();
+  }
+}
+
 bool GenerateGuard(Node* tensorexpr_graph_node, bool add_composed_op) {
   auto tensorexpr_graph = SubgraphUtils::getSubgraph(tensorexpr_graph_node);
 
+  // Move constant tensors from the subgraph to the outer scope.
+  // This is necessary because symbolic shape analysis does not handle the
+  // case of broadcast(constant, symbolic_shape) well and that results in poor
+  // performance.
+  moveConstantTensorsOutOfSubgraph(tensorexpr_graph_node, tensorexpr_graph);
+
   // Generalize Inputs
   if (!TryGeneralizeInputDimensionsToSymbolicShapes(tensorexpr_graph)) {
     return false;