[JIT] UseVariadicOp handles multiple lists (pytorch#66288)

Summary:
Pull Request resolved: pytorch#66288

This change makes it so `UseVariadicOp` can transform ops with many Tensor list inputs.

Input pattern:
```
%output : Type = op(%list_1, %arg_1, %list_2, %list_3)
```
Output pattern:
```
%output : Type = variadic_op(%list_11, ..., %list_1N, %arg_1, %list_21, ..., %list_2M, %list_31, ..., %list_3K, N, M, K)
```
The length of each list is passed at the end of the variadic op so that the op implementation can process the inputs appropriately. This also frees us from needing to update `hasVarArgs` in static runtime each time we add a variadic op.

This diff also makes `UseVariadicOp` more robust. Before, `list_idx` was passed as an argument. Now, `VariadicUpdater` determines `list_idx` from the node's schema.

Test Plan:
Existing variadic ops do not break:
`buck test caffe2/benchmarks/static_runtime:static_runtime_cpptest`

Reviewed By: d1jang

Differential Revision: D31450811

fbshipit-source-id: 808fcc3ae8940b9e602586f38f8cf9154c9a6462

Author: Mike Iovine

torch/csrc/jit/passes/variadic_ops.cpp

@@ -2,24 +2,46 @@
 
 #include <torch/csrc/jit/ir/alias_analysis.h>
 #include <torch/csrc/jit/jit_log.h>
+#include <torch/csrc/jit/passes/constant_pooling.h>
 #include <torch/csrc/jit/passes/remove_mutation.h>
 
 namespace torch {
 namespace jit {
 
 namespace {
 
+std::vector<size_t> identifyListArgIndices(const c10::FunctionSchema& schema) {
+  std::vector<size_t> list_indices;
+  const auto& args = schema.arguments();
+  for (const auto i : c10::irange(args.size())) {
+    auto list_type = args[i].type()->castRaw<ListType>();
+    if (list_type && list_type->getElementType()->castRaw<TensorType>()) {
+      list_indices.push_back(i);
+    }
+  }
+  return list_indices;
+}
+
+bool isTensorListConstruct(Node* node) {
+  if (node->kind() != prim::ListConstruct) {
+    return false;
+  }
+  const auto type = node->output()->type()->castRaw<ListType>();
+  TORCH_CHECK(type != nullptr);
+  const auto& elem_type = type->getElementType();
+  return elem_type->castRaw<TensorType>();
+}
+
 class VariadicUpdater {
  public:
-  explicit VariadicUpdater(
+  VariadicUpdater(
       std::shared_ptr<Graph> graph,
       NodeKind op,
-      NodeKind variadic_op,
-      size_t list_idx = 0)
+      NodeKind variadic_op)
       : graph_(std::move(graph)),
+        alias_db_(graph_),
         op_(op),
-        variadic_op_(variadic_op),
-        list_idx_(list_idx) {}
+        variadic_op_(variadic_op) {}
 
   bool run() {
     collectOpNodes(graph_->block());
@@ -31,90 +53,151 @@ class VariadicUpdater {
   }
 
  private:
+  void recordSchema(Node* op_node) {
+    const auto& schema = op_node->schema();
+    auto it = schema_to_list_indices_.find(schema.name());
+    if (it == schema_to_list_indices_.end()) {
+      schema_to_list_indices_.emplace(
+          schema.overload_name(), identifyListArgIndices(schema));
+    }
+  }
+
+  const std::vector<size_t>& getListIndices(Node* op_node) const {
+    const auto& schema = op_node->schema();
+    auto it = schema_to_list_indices_.find(schema.overload_name());
+    TORCH_CHECK(it != schema_to_list_indices_.end());
+    return it->second;
+  }
+
   void collectOpNodes(Block* block) {
     for (auto node : block->nodes()) {
       if (node->kind() == op_) {
         op_nodes_.push_back(node);
+        recordSchema(node);
       }
       for (Block* b : node->blocks()) {
         collectOpNodes(b);
       }
     }
   }
 
-  bool replaceWithVariadicOp(Node* op_node) {
+  bool allListInputsAreValid(Node* op_node) {
     const size_t num_inputs = op_node->inputs().size();
-    TORCH_CHECK(list_idx_ < num_inputs);
-    if (op_node->input(list_idx_)->node()->kind() != prim::ListConstruct) {
-      return false;
+    for (const auto list_idx : getListIndices(op_node)) {
+      TORCH_CHECK(list_idx < num_inputs);
+      const auto list = op_node->input(list_idx)->node();
+      // We do not transform ops whose list input can not be moved to the
+      // position before op. This in turn implies that there is some mutation
+      // of the input list before op.
+      if (!isTensorListConstruct(list) ||
+          !alias_db_.couldMoveBeforeTopologically(list, op_node)) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  void insertAllInputsBetween(
+      std::vector<Value*>& inputs,
+      Node* node,
+      size_t start_idx,
+      size_t end_idx) const {
+    const size_t num_inputs = node->inputs().size();
+    TORCH_CHECK(start_idx <= end_idx && end_idx <= num_inputs);
+    inputs.insert(
+        inputs.end(),
+        node->inputs().begin() + start_idx,
+        node->inputs().begin() + end_idx);
+  }
+
+  void insertIntegerInput(std::vector<Value*>& inputs, size_t input) {
+    auto constant = graph_->create(prim::Constant);
+    constant->output()->setType(c10::IntType::get());
+    constant->i_(attr::value, input);
+    graph_->prependNode(constant);
+    inputs.push_back(constant->output());
+  }
+
+  void deleteOpNodeAndLists(Node* op_node) {
+    // Collect the lists before we destroy op_node
+    std::vector<Node*> lists;
+    const auto& list_indices = getListIndices(op_node);
+    lists.reserve(list_indices.size());
+    for (const size_t list_idx : list_indices) {
+      auto* list = op_node->input(list_idx)->node();
+      lists.push_back(list);
     }
-    auto list = op_node->input(list_idx_)->node();
-    const size_t list_len = list->inputs().size();
 
-    // We do not transform ops whose list input can not be moved to the
-    // position before op. This in turn implies that there is some mutation
-    // of the input list before op.
-    if (!getOrCreateAliasDb()->couldMoveBeforeTopologically(list, op_node)) {
+    GRAPH_UPDATE("Deleting\n", *op_node);
+    op_node->destroy();
+    for (auto* list : lists) {
+      if (!list->hasUses()) {
+        GRAPH_UPDATE("Deleting\n", *list);
+        list->destroy();
+      }
+    }
+  }
+
+  bool replaceWithVariadicOp(Node* op_node) {
+    if (!allListInputsAreValid(op_node)) {
       return false;
     }
 
-    // Construct new inputs
     std::vector<Value*> inputs;
-    inputs.reserve(num_inputs + list_len - 1);
-    inputs.insert(
-        inputs.end(),
-        op_node->inputs().begin(),
-        op_node->inputs().begin() + list_idx_);
-    inputs.insert(inputs.end(), list->inputs().begin(), list->inputs().end());
-    inputs.insert(
-        inputs.end(),
-        op_node->inputs().begin() + list_idx_ + 1,
-        op_node->inputs().end());
+    size_t cur_idx = 0;
+    std::vector<size_t> list_lens;
+    for (const size_t list_idx : getListIndices(op_node)) {
+      insertAllInputsBetween(inputs, op_node, cur_idx, list_idx);
+      const auto list = op_node->input(list_idx)->node();
+      const auto list_len = list->inputs().size();
+      list_lens.push_back(list_len);
+      insertAllInputsBetween(inputs, list, 0, list_len);
+      cur_idx = list_idx + 1;
+    }
+    insertAllInputsBetween(inputs, op_node, cur_idx, op_node->inputs().size());
+
+    // We insert these extra integers at the end of the argument list only if we
+    // have more than one variadic list (the information is redundant when there
+    // is only one list because the interpreter knows how many arguments there
+    // are).
+    if (list_lens.size() > 1) {
+      for (const size_t list_len : list_lens) {
+        insertIntegerInput(inputs, list_len);
+      }
+    }
 
     auto var_op_node = op_node->owningGraph()->create(variadic_op_, inputs);
     var_op_node->output()->setType(op_node->output()->type());
     GRAPH_UPDATE("Adding\n", *var_op_node);
     var_op_node->insertBefore(op_node);
     GRAPH_UPDATE("Replacing\n", *op_node, "with\n", *var_op_node);
     op_node->output()->replaceAllUsesWith(var_op_node->output());
-    GRAPH_UPDATE("Deleting\n", *op_node);
-    op_node->destroy();
-    if (!list->hasUses()) {
-      GRAPH_UPDATE("Deleting\n", *list);
-      list->destroy();
-    }
+    deleteOpNodeAndLists(op_node);
     return true;
   }
 
-  AliasDb* getOrCreateAliasDb() {
-    if (!aliasDb_) {
-      aliasDb_ = std::make_unique<AliasDb>(graph_);
-    }
-    return aliasDb_.get();
-  }
-
   std::shared_ptr<Graph> graph_;
-  std::unique_ptr<AliasDb> aliasDb_ = nullptr;
-
   std::vector<Node*> op_nodes_;
 
+  AliasDb alias_db_;
+
   NodeKind op_;
   NodeKind variadic_op_;
 
-  size_t list_idx_;
+  std::unordered_map<std::string, std::vector<size_t>> schema_to_list_indices_;
 };
 
 } // namespace
 
 bool UseVariadicOp(
     const std::shared_ptr<Graph>& graph,
     NodeKind op,
-    NodeKind variadic_op,
-    size_t list_idx) {
+    NodeKind variadic_op) {
   const std::string pass_name = std::string("variadic ") + op.toQualString();
   GRAPH_DUMP("Before " + pass_name, graph);
-  bool changed = VariadicUpdater(graph, op, variadic_op, list_idx).run();
+  bool changed = VariadicUpdater(graph, op, variadic_op).run();
   if (changed) {
+    ConstantPooling(graph);
     GRAPH_DUMP("After " + pass_name, graph);
   }
   return changed;
@@ -123,14 +206,13 @@ bool UseVariadicOp(
 bool RemoveListMutationAndUseVariadicOp(
     const std::shared_ptr<Graph>& graph,
     NodeKind op,
-    NodeKind variadic_op,
-    size_t list_idx) {
+    NodeKind variadic_op) {
   bool changed_in_last_iter = true;
   bool changed = false;
   while (changed_in_last_iter) {
     changed_in_last_iter = RemoveListMutation(graph);
     changed_in_last_iter =
-        UseVariadicOp(graph, op, variadic_op, list_idx) || changed_in_last_iter;
+        UseVariadicOp(graph, op, variadic_op) || changed_in_last_iter;
     changed = changed || changed_in_last_iter;
   }
   return changed;
@@ -139,15 +221,13 @@ bool RemoveListMutationAndUseVariadicOp(
 bool UseVariadicCat(const std::shared_ptr<Graph>& graph) {
   return UseVariadicOp(graph, aten::cat, prim::VarConcat);
 }
-
 bool RemoveListMutationAndUseVariadicCat(const std::shared_ptr<Graph>& graph) {
   return RemoveListMutationAndUseVariadicOp(graph, aten::cat, prim::VarConcat);
 }
 
 bool UseVariadicStack(const std::shared_ptr<Graph>& graph) {
   return UseVariadicOp(graph, aten::stack, prim::VarStack);
 }
-
 bool RemoveListMutationAndUseVariadicStack(
     const std::shared_ptr<Graph>& graph) {
   return RemoveListMutationAndUseVariadicOp(graph, aten::stack, prim::VarStack);

torch/csrc/jit/passes/variadic_ops.h

@@ -5,31 +5,27 @@
 namespace torch {
 namespace jit {
 
-// Replaces the `aten::cat` ops in the given graph with variadic cat ops.
-// Returns true if the graph is modified.
-TORCH_API bool UseVariadicCat(const std::shared_ptr<Graph>& graph);
-
-TORCH_API bool RemoveListMutationAndUseVariadicCat(
-    const std::shared_ptr<Graph>& graph);
-
-// Replaces the `aten::stack` ops in the given graph with variadic cat ops.
-// Returns true if the graph is modified.
-TORCH_API bool UseVariadicStack(const std::shared_ptr<Graph>& graph);
-
-TORCH_API bool RemoveListMutationAndUseVariadicStack(
-    const std::shared_ptr<Graph>& graph);
-
+// Try to replace an op that takes a list input with another op that takes a
+// variadic number of arguments.
 TORCH_API bool UseVariadicOp(
     const std::shared_ptr<Graph>& graph,
     NodeKind op,
-    NodeKind variadic_op,
-    size_t list_idx = 0);
+    NodeKind variadic_op);
 
 TORCH_API bool RemoveListMutationAndUseVariadicOp(
     const std::shared_ptr<Graph>& graph,
     NodeKind op,
-    NodeKind variadic_op,
-    size_t list_idx = 0);
+    NodeKind variadic_op);
+
+// Convenient functions for replacing aten::stack/aten::cat with their
+// variadic versions.
+TORCH_API bool UseVariadicCat(const std::shared_ptr<Graph>& graph);
+TORCH_API bool RemoveListMutationAndUseVariadicCat(
+    const std::shared_ptr<Graph>& graph);
+
+TORCH_API bool UseVariadicStack(const std::shared_ptr<Graph>& graph);
+TORCH_API bool RemoveListMutationAndUseVariadicStack(
+    const std::shared_ptr<Graph>& graph);
 
 } // namespace jit
 } // namespace torch

torch/csrc/jit/runtime/static/impl.cpp

@@ -95,8 +95,7 @@ void OptimizeGraph(
         graph,
         c10::Symbol::fromQualString("fb::sigrid_transforms_torch_bind"),
         c10::Symbol::fromQualString(
-            "fb::variadic_sigrid_transforms_torch_bind"),
-        1 /* list_idx */);
+            "fb::variadic_sigrid_transforms_torch_bind"));
     FuseSignLog1P(graph);
 
     // TODO: we can avoid this guard by moving operations