fix: Repair EliminateExceptions lowering pass

- Update EliminateExceptions to use `replaceAllUsesDominatedByNodeWith`
instead of `replaceAllUsesWith` to avoid issue with invalid IR causing
program halting

Author: gs-olive

core/lowering/lowering.cpp

@@ -104,6 +104,7 @@ void LowerGraph(std::shared_ptr<torch::jit::Graph>& g, std::vector<torch::jit::I
   torch::jit::InlineFunctionalGraphs(g);
   torch::jit::PeepholeOptimize(g, false);
   torch::jit::FuseLinear(g);
+  passes::EliminateExceptionsNew(g);
   if (!lower_info.disable_cse) {
     torch::jit::EliminateCommonSubexpression(g);
   }

core/lowering/passes/exception_elimination.cpp

@@ -108,6 +108,57 @@ void EliminateExceptionOrPassPattern(std::shared_ptr<Graph> graph) {
   }
 }
 
+/*
+  Below is a fork of the torch::jit::EliminateExceptions pass, with node replacement
+  using replaceAllUsesDominatedByNodeWith instead of replaceAllUsesWith,
+  so as to not invalidate the IR in challenging cases, such as nested Ifs
+
+  Original Source from which it was adapted:
+  https://github.com/pytorch/pytorch/blob/c29ab84115f40614d04e4557ea2e1ac40b7aa75c/torch/csrc/jit/passes/remove_exceptions.cpp
+*/
+
+bool certainlyThrows(Block* block) {
+  // A block certainly throws an exception if it contains
+  // the prim::RaiseException operation
+  for (Node* n : block->nodes()) {
+    if (n->kind() == prim::RaiseException) {
+      return true;
+    }
+  }
+  return false;
+}
+
+void EliminateExceptionsNew(Block* block) {
+  auto graph = block->owningGraph();
+  // Generate false and true constant placeholders
+  Value* false_const = graph->insertConstant(IValue(false));
+  Value* true_const = graph->insertConstant(IValue(true));
+
+  // For each prim::If node, if either block certainly throws an exception
+  // Replace all uses of the node input with the logical opposite
+  for (Node* n : block->nodes()) {
+    if (n->kind() == prim::If) {
+      Block* true_block = n->blocks()[0];
+      Block* false_block = n->blocks()[1];
+
+      if (certainlyThrows(true_block)) {
+        n->input(0)->replaceAllUsesDominatedByNodeWith(n, false_const);
+      } else if (certainlyThrows(false_block)) {
+        n->input(0)->replaceAllUsesDominatedByNodeWith(n, true_const);
+      }
+    }
+
+    // Inspect and replace all instances within subblocks of the current node
+    for (Block* subblock : n->blocks()) {
+      EliminateExceptionsNew(subblock);
+    }
+  }
+}
+
+void EliminateExceptionsNew(std::shared_ptr<Graph>& graph) {
+  EliminateExceptionsNew(graph->block());
+}
+
 } // namespace passes
 } // namespace lowering
 } // namespace core

core/lowering/passes/passes.h

@@ -20,6 +20,7 @@ void Conv3DToConvolution(std::shared_ptr<torch::jit::Graph>& graph);
 void ConvTransposed3DToConvolution(std::shared_ptr<torch::jit::Graph>& graph);
 void FuseAddMMBranches(std::shared_ptr<torch::jit::Graph> graph);
 void LinearToAddMM(std::shared_ptr<torch::jit::Graph>& graph);
+void EliminateExceptionsNew(std::shared_ptr<torch::jit::Graph>& graph);
 void EliminateExceptionOrPassPattern(std::shared_ptr<torch::jit::Graph> graph);
 void ReduceToOperation(std::shared_ptr<torch::jit::Graph>& graph);
 void ReduceGelu(std::shared_ptr<torch::jit::Graph>& graph);