Add guard elimination cases for operators encountered on an RL workload.

resistor · resistor · commit a31cb905724a · 2020-03-19T16:38:21.000-07:00
diff --git a/torch/csrc/jit/passes/guard_elimination.cpp b/torch/csrc/jit/passes/guard_elimination.cpp
@@ -222,6 +222,7 @@ struct GuardElimination {
     const static auto no_exceptions = std::unordered_set<size_t>{};
     switch (n->kind()) {
       case aten::add:
+      case aten::add_:
       case aten::sub:
       case aten::mul:
       case aten::div:
@@ -279,16 +280,32 @@ struct GuardElimination {
       case aten::addcmul:
       case aten::where:
       case aten::_cast_Float:
+      case aten::_cast_Long:
       case aten::_sigmoid_backward:
       case aten::_tanh_backward:
       case aten::__and__:
       case aten::__or__:
       case aten::__xor__:
       case aten::__lshift__:
       case aten::__rshift__:
+      case aten::bitwise_not:
+      case aten::bitwise_and:
+      case aten::bitwise_or:
+      case aten::bitwise_xor:
         return checkInputs(n, no_exceptions, true);
+      case aten::softmax:
+        return checkInputs(n, std::unordered_set<size_t>{1}, true);
+      case aten::multinomial:
+        return checkInputs(n, std::unordered_set<size_t>{2, 3}, false);
+      case aten::flatten:
+      case aten::argmax:
+      case aten::squeeze:
       case aten::avg_pool2d:
         return checkInputs(n, no_exceptions, false);
+      case aten::conv1d:
+      case aten::conv2d:
+      case aten::conv3d:
+        return checkInputs(n, std::unordered_set<size_t>{2, 6}, false);
       case aten::slice:
         return !n->input(0)->type()->expect<TensorType>()->isSummarized() &&
             // check that the dimension argument is constant
