Add API for unexecuted op (#43629)

Elias Ellison · facebook-github-bot · commit d1c4d75c14d8 · 2020-08-27T14:34:05.000-07:00
Summary: Pull Request resolved: #43629 We have a few places where we count the size a block / subgraph - it's nice to have a shared API to ignore operators that are not executed in the optimized graph (will be used when i add a new profiling node in PR ^^) Test Plan: Imported from OSS Reviewed By: bertmaher Differential Revision: D23358807 Pulled By: eellison fbshipit-source-id: 62c745d9025de94bdafd9f748f7c5a8574cace3f
diff --git a/torch/csrc/jit/ir/ir.h b/torch/csrc/jit/ir/ir.h
@@ -435,6 +435,11 @@ struct TORCH_API Node {
   bool isNondeterministic() const;
   bool hasSideEffects() const;
 
+  // instructions lowered by the interpreter and not run in the optimized graph
+  bool notExecutedOp() const {
+    return kind_ == prim::Constant || kind_ == prim::profile;
+  }
+
   // Graphs
 
   // Note [Topological invariant]
diff --git a/torch/csrc/jit/passes/create_autodiff_subgraphs.cpp b/torch/csrc/jit/passes/create_autodiff_subgraphs.cpp
@@ -228,8 +228,7 @@ class SubgraphSlicer {
     size_t i = 0;
     for (auto it = subgraph->nodes().begin(); it != subgraph->nodes().end();
          ++it) {
-      // constants are not interpreted as instructions, ignore them
-      i += it->kind() != prim::Constant;
+      i += !it->notExecutedOp();
       if (i >= minSubgraphSize_) {
         return false;
       }
diff --git a/torch/csrc/jit/passes/loop_unrolling.cpp b/torch/csrc/jit/passes/loop_unrolling.cpp
@@ -35,10 +35,13 @@ bool isForLoop(Node* node) {
 int64_t limitedBlockSize(Block* body, int64_t limit) {
   auto it = body->nodes().begin();
   auto end = body->nodes().end();
-  for (int64_t i = 0; i < limit; ++i, ++it) {
+  for (int64_t i = 0; i < limit; ++it) {
     for (Block* subblock : it->blocks()) {
       i += limitedBlockSize(subblock, limit - i);
     }
+    if (!it->notExecutedOp()) {
+      ++i;
+    }
     if (it == end) {
       return i;
     }
