[TensorExpr] Hook Fuser Pass to JIT opt-limit utility.

Mikhail Zolotukhin · Mikhail Zolotukhin · commit d7d5365ce737 · 2021-01-13T19:15:36.000-08:00
That new feature allows to bisect the pass easily by hard-stopping it
after a given number of hits.

[ghstack-poisoned]
diff --git a/torch/csrc/jit/passes/tensorexpr_fuser.cpp b/torch/csrc/jit/passes/tensorexpr_fuser.cpp
@@ -5,6 +5,7 @@
 #include <torch/csrc/jit/codegen/fuser/interface.h>
 #include <torch/csrc/jit/ir/alias_analysis.h>
 #include <torch/csrc/jit/jit_log.h>
+#include <torch/csrc/jit/jit_opt_limit.h>
 #include <torch/csrc/jit/passes/common_subexpression_elimination.h>
 #include <torch/csrc/jit/passes/constant_pooling.h>
 #include <torch/csrc/jit/passes/dead_code_elimination.h>
@@ -975,6 +976,13 @@ class TensorExprFuser {
 
     REQ(tensorexpr::isSupported(node));
     REQ(typesAreSupported(node));
+
+    // A hook to optimizations limitter to allow bisecting the pass
+    auto allowed = JIT_OPT_LIMIT();
+    if (!allowed) {
+      return false;
+    }
+
     return true;
   }
 
