refactor(//core/lowering): Make logic a bit clearer in EE pass

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

core/lowering/passes/exception_elimination.cpp

@@ -44,10 +44,10 @@ struct ExceptionOrPassPatternElimination {
     bool arm1_starts_with_exception = (*arm1_start)->kind() == prim::RaiseException;
     bool arm2_starts_with_exception = (*arm2_start)->kind() == prim::RaiseException;
 
-    if (!arm1_starts_with_exception && !arm2_starts_with_exception) {
+    //if (!arm1_starts_with_exception && !arm2_starts_with_exception) {
       // Neither arm matches the pattern
-      return false;
-    }
+   //   return false;
+    //}
 
     /// Check if this Node hosts a pattern like so:
     ///  = prim::If(%5958)
@@ -57,14 +57,12 @@ struct ExceptionOrPassPatternElimination {
     ///   block1():
     ///    -> ()
     if (arm1_starts_with_exception) {
-      if ((*(++arm1_start))->kind() != prim::Return) {
+      if ((*(++arm1_start))->kind() == prim::Return) {
         // Make sure that block0 is solely just the exception and the return
-        return false;
-      }
-
-      if ((*(arm2_start))->kind() != prim::Return) {
-        // Make sure that block1 is solely the return
-        return false;
+        if ((*(arm2_start))->kind() == prim::Return) {
+          // Make sure that block1 is solely the return
+          return true;
+        }
       }
     }
 
@@ -76,25 +74,23 @@ struct ExceptionOrPassPatternElimination {
     ///     = prim::RaiseException(%45)
     ///    -> ()
     if (arm2_starts_with_exception) {
-      if ((*(++arm2_start))->kind() != prim::Return) {
+      if ((*(++arm2_start))->kind() == prim::Return) {
         // Make sure that block1 is solely just the exception and the return
-        return false;
-      }
-
-      if ((*(arm1_start))->kind() != prim::Return) {
-        // Make sure that block0 is solely the return
-        return false;
+        if ((*(arm1_start))->kind() == prim::Return) {
+          // Make sure that block0 is solely the return
+          return true;
+        }
       }
     }
 
-    return true;
+    return false;
   }
 
   void findExceptionOrPassNodes(Block* b) {
     for (auto it = b->nodes().begin(); it != b->nodes().end(); it++) {
       auto n = *it;
       if (n->kind() == prim::If && isExceptionOrPassNode(n)) {
-        LOG_GRAPH("Found that node " << *n << "  is an exception or pass node (EliminateChecks)" << std::endl);
+        LOG_ERROR("Found that node " << *n << "  is an exception or pass node (EliminateChecks)" << std::endl);
         it.destroyCurrent();
       }
     }
@@ -107,6 +103,9 @@ struct ExceptionOrPassPatternElimination {
 void EliminateExceptionOrPassPattern(std::shared_ptr<Graph> graph) {
   ExceptionOrPassPatternElimination eppe(std::move(graph));
   eppe.run();
+  if (graph) {
+    LOG_ERROR("Post Eliminate Exception or Pass Patterns: " << *graph);
+  }
 }
 
 } // namespace passes

tests/core/lowering/test_exception_elimination_pass.cpp

@@ -44,7 +44,7 @@ TEST(LoweringPasses, EliminateExceptionOrPassPattern_Block0) {
   auto if_block0 = if_node->addBlock();
   auto exception_node = g->create(torch::jit::prim::RaiseException, {except_val, none_const_val}, 0);
   if_block0->appendNode(exception_node);
-  auto if_block1 = if_node->addBlock();
+  /*auto if_block1 =*/ if_node->addBlock();
   g->insertNode(if_node);
   auto cat_node = g->create(torch::jit::aten::cat, {list_node->output(), zero_const_val});
   g->insertNode(cat_node);
@@ -97,7 +97,7 @@ TEST(LoweringPasses, EliminateExceptionOrPassPattern_Block1) {
   bool_node->output()->setType(torch::jit::BoolType::get());
   g->insertNode(bool_node);
   auto if_node = g->create(torch::jit::prim::If, {bool_node->output()}, 0);
-  auto if_block0 = if_node->addBlock();
+  /*auto if_block0 = */if_node->addBlock();
   auto if_block1 = if_node->addBlock();
   auto exception_node = g->create(torch::jit::prim::RaiseException, {except_val, none_const_val}, 0);
   if_block1->appendNode(exception_node);
@@ -154,7 +154,7 @@ TEST(LoweringPasses, EliminateExceptionOrPassPattern_Negative) {
   auto if_block0 = if_node->addBlock();
   auto append_node = g->create(torch::jit::aten::append, {list_node->output(), y});
   if_block0->appendNode(append_node);
-  auto if_block1 = if_node->addBlock();
+  /*auto if_block1 = */if_node->addBlock();
   g->insertNode(if_node);
   auto cat_node = g->create(torch::jit::aten::cat, {list_node->output(), zero_const_val});
   g->insertNode(cat_node);