Add static_runtime::fused_equally_split (#2)

Summary:
Pull Request resolved: pytorch/pytorch-canary#2

Pull Request resolved: pytorch#66881

Adds `static_runtime::fused_equally_split` operator and removes `is_fused` logic from original operator. Modifies `FuseUnpackListV2` to map `fb::equally_split` to this new operator.

Test Plan:
```
adityapillai@5960 /data/sandcastle/boxes/fbsource/fbcode 1m 13s
❯ buck test //caffe2/benchmarks/static_runtime/fb:test_fb_operators
```
and sandcastle
strange_what_could_go_wrong

Reviewed By: mikeiovine

Differential Revision: D31742293

fbshipit-source-id: 60b35589c8817719b005d49811f575b6590d1c39

Author: pilleye

benchmarks/static_runtime/test_utils.cc

@@ -277,6 +277,18 @@ void testStaticRuntime(
   compareTensorLists(args_tensors, args_copy, use_allclose, use_equalnan);
 }
 
+bool hasProcessedNodeWithName(
+    torch::jit::StaticModule& smodule,
+    const char* name) {
+  for (torch::jit::ProcessedNode& pnode : smodule.runtime().nodes()) {
+    auto op_name = pnode.node()->kind().toQualString();
+    if (strcmp(op_name, name) == 0) {
+      return true;
+    }
+  }
+  return false;
+}
+
 } // namespace test
 } // namespace jit
 } // namespace torch

benchmarks/static_runtime/test_utils.h

@@ -6,6 +6,7 @@
 #include <vector>
 
 #include <torch/csrc/jit/ir/ir.h>
+#include <torch/csrc/jit/runtime/static/impl.h>
 
 namespace c10 {
 struct IValue;
@@ -30,6 +31,8 @@ void testStaticRuntime(
 
 std::shared_ptr<Graph> getGraphFromIR(const std::string& ir);
 
+bool hasProcessedNodeWithName(torch::jit::StaticModule& smodule, const char *name);
+
 } // namespace test
 } // namespace jit
 } // namespace torch

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

@@ -347,6 +347,9 @@ TORCH_LIBRARY_FRAGMENT(static_runtime, m) {
   m.def(torch::schema(
       "static_runtime::VarTupleUnpack(...) -> ...",
       c10::AliasAnalysisKind::CONSERVATIVE));
+  m.def(torch::schema(
+      "static_runtime::fused_equally_split(Tensor input, int num_split, int dim) -> ...",
+      c10::AliasAnalysisKind::PURE_FUNCTION));
 }
 
 void FuseSignLog1P(std::shared_ptr<torch::jit::Graph>& graph) {
@@ -529,13 +532,11 @@ void FuseListUnpack(std::shared_ptr<torch::jit::Graph>& graph) {
   const std::vector<Value*> graph_outputs(
       graph->outputs().begin(), graph->outputs().end());
   auto nodes = graph->nodes();
-  std::vector<Node*> equally_splits_to_remove;
   for (auto it = nodes.begin(); it != nodes.end(); ++it) {
     Node* node = *it;
     const std::string node_qual_string = node->kind().toQualString();
     if (node_qual_string == "fb::sigrid_transforms" ||
         node_qual_string == "fb::sigrid_transforms_torch_bind" ||
-        node_qual_string == "fb::equally_split" ||
         node_qual_string == "fb::gather_ranges_to_dense" ||
         node_qual_string == "fb::gather_ranges_to_dense_v2" ||
         node_qual_string == "fb::variadic_sigrid_transforms_torch_bind") {
@@ -577,22 +578,9 @@ void FuseListUnpack(std::shared_ptr<torch::jit::Graph>& graph) {
       it_next.destroyCurrent(); // remove list_unpack
 
       node->eraseOutput(0);
-
-      if (node_qual_string == "fb::equally_split" &&
-          node->outputs().size() == 1) {
-        // This captures a case of `y = fb::equally_split(x, 1, _)` where y
-        // becomes just an alias of x.
-        // If this case is found, replace y with x to avoid executing this op.
-        equally_splits_to_remove.push_back(node);
-      }
     }
   }
 
-  for (Node* node : equally_splits_to_remove) {
-    node->output(0)->replaceAllUsesWith(node->input(0));
-    node->destroy();
-  }
-
 #ifndef NDEBUG
   graph->lint();
   AliasDb db2(graph);
@@ -601,7 +589,9 @@ void FuseListUnpack(std::shared_ptr<torch::jit::Graph>& graph) {
 }
 
 void FuseListUnpackV2(std::shared_ptr<torch::jit::Graph>& graph) {
-  const FastMap<c10::Symbol, c10::Symbol> unfused_to_fused = {};
+  const FastMap<c10::Symbol, c10::Symbol> unfused_to_fused = {
+      {c10::Symbol::fromQualString("fb::equally_split"),
+       c10::Symbol::fromQualString("static_runtime::fused_equally_split")}};
 
   AliasDb alias_db(
       graph, /*isFrozen=*/false, /*enablePreciseTupleContainerAnalysis=*/true);