Merge branch 'main' into dynamo_refactor

Author: peri044

core/lowering/lowering.cpp

@@ -9,7 +9,6 @@
 #include "torch/csrc/jit/passes/lower_graph.h"
 #include "torch/csrc/jit/passes/lower_tuples.h"
 #include "torch/csrc/jit/passes/peephole.h"
-#include "torch/csrc/jit/passes/remove_exceptions.h"
 #include "torch/csrc/jit/passes/remove_mutation.h"
 
 #include "core/lowering/lowering.h"
@@ -105,7 +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);
-  torch::jit::EliminateExceptions(g);
+  passes::EliminateExceptionsSafe(g);
   if (!lower_info.disable_cse) {
     torch::jit::EliminateCommonSubexpression(g);
   }

core/lowering/passes/exception_elimination.cpp

@@ -1,5 +1,6 @@
 #include "torch/csrc/jit/ir/alias_analysis.h"
 #include "torch/csrc/jit/jit_log.h"
+#include "torch/csrc/jit/passes/constant_pooling.h"
 #include "torch/csrc/jit/passes/constant_propagation.h"
 #include "torch/csrc/jit/passes/dead_code_elimination.h"
 #include "torch/csrc/jit/passes/guard_elimination.h"
@@ -108,6 +109,71 @@ 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 EliminateExceptionsSafe(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 input conditional 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];
+      bool removed_exception = false;
+      Value* input_value_replacement;
+
+      // If the block throws an exception, replace input with logical opposite
+      if (certainlyThrows(true_block)) {
+        removed_exception = true;
+        input_value_replacement = false_const;
+      } else if (certainlyThrows(false_block)) {
+        removed_exception = true;
+        input_value_replacement = true_const;
+      }
+
+      // Log node and perform input replacement
+      if (removed_exception) {
+        LOG_WARNING("Detected and removing exception in TorchScript IR for node: " << util::node_info(n));
+        n->insertInput(0, input_value_replacement);
+        n->removeInput(1);
+      }
+    }
+
+    // Inspect and replace all instances within subblocks of the current node
+    for (Block* subblock : n->blocks()) {
+      EliminateExceptionsSafe(subblock);
+    }
+  }
+}
+
+void EliminateExceptionsSafe(std::shared_ptr<Graph>& graph) {
+  EliminateExceptionsSafe(graph->block());
+  ConstantPropagation(graph);
+  ConstantPooling(graph);
+}
+
 } // 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 EliminateExceptionsSafe(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);

py/torch_tensorrt/__init__.py

@@ -94,7 +94,7 @@ def _find_lib(name, paths):
 
 from torch_tensorrt import fx
 
-if version.parse(torch.__version__) >= version.parse("2.1.dev"):
+if version.parse(sanitized_torch_version()) >= version.parse("2.1.dev"):
     from torch_tensorrt import dynamo
     from torch_tensorrt.dynamo import backend
 

py/torch_tensorrt/_util.py

@@ -30,3 +30,11 @@ def get_build_info() -> str:
 
 def set_device(gpu_id):
     _C.set_device(gpu_id)
+
+
+def sanitized_torch_version() -> str:
+    return (
+        torch.__version__
+        if ".nv" not in torch.__version__
+        else torch.__version__.split(".nv")[0]
+    )

py/torch_tensorrt/dynamo/backend/backends.py

@@ -88,13 +88,14 @@ def _pretraced_backend(
             )
             return gm.forward
         else:
-            raise AssertionError(
+            logger.critical(
                 "Halting compilation on build failure since "
                 + "pass_through_build_failures was specified as True. "
                 + "To return the default Torch implementation and avoid "
                 + "halting compilation on engine build failures, "
                 + "specify pass_through_build_failures=False."
             )
+            raise
 
 
 def _compile_module(

py/torch_tensorrt/fx/test/passes/test_fix_reshape_batch_dim.py

@@ -11,6 +11,7 @@
 from torch.testing._internal.common_utils import run_tests, TestCase
 from torch_tensorrt.fx.passes.lower_basic_pass import fix_reshape_batch_dim
 from torch_tensorrt.fx.tracer.acc_tracer import acc_tracer
+from torch_tensorrt._util import sanitized_torch_version
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -43,7 +44,9 @@ def forward(self, x, y):
     %reshape : [num_users=1] = call_function[target=torch_tensorrt.fx.tracer.acc_tracer.acc_ops.reshape](args = (), kwargs = {input: %y, acc_out_ty: ((%getitem_1, -1, 3), None, None, None, None, None, None)})
     return reshape
 """
-        if version.parse(torch.__version__) < version.parse("2.1.0.dev20230620"):
+        if version.parse(sanitized_torch_version()) < version.parse(
+            "2.1.0.dev20230620"
+        ):
             expected_graph = expected_graph.replace("num_users", "#users")
 
         assert (

py/torch_tensorrt/fx/test/passes/test_remove_duplicate_output_args.py

@@ -8,6 +8,8 @@
 import torch.nn as nn
 
 import torch_tensorrt.fx.passes.remove_duplicate_output_args as dedup
+from torch_tensorrt._util import sanitized_torch_version
+
 from torch.testing._internal.common_utils import run_tests, TestCase
 
 _LOGGER = logging.getLogger(__name__)
@@ -57,7 +59,9 @@ def is_leaf_module(self, m, qn):
     return add
 """.strip()
 
-        if version.parse(torch.__version__) < version.parse("2.1.0.dev20230620"):
+        if version.parse(sanitized_torch_version()) < version.parse(
+            "2.1.0.dev20230620"
+        ):
             ttop_graph_expected = ttop_graph_expected.replace("num_users", "#users")
 
         assert (
@@ -71,7 +75,9 @@ def is_leaf_module(self, m, qn):
     return (x,)
 """.strip()
 
-        if version.parse(torch.__version__) < version.parse("2.1.0.dev20230620"):
+        if version.parse(sanitized_torch_version()) < version.parse(
+            "2.1.0.dev20230620"
+        ):
             ttop_a_graph_expected = ttop_a_graph_expected.replace("num_users", "#users")
 
         assert (

py/torch_tensorrt/fx/tracer/dispatch_tracer/aten_tracer.py

@@ -3,10 +3,11 @@
 from contextlib import contextmanager
 from typing import Any, Callable, Dict, Generator, List, Optional, Set, Tuple, Union
 from packaging import version
+from torch_tensorrt._util import sanitized_torch_version
 
 import torch
 
-if version.parse(torch.__version__) >= version.parse("2.dev"):
+if version.parse(sanitized_torch_version()) >= version.parse("2.dev"):
     import torch._dynamo as torchdynamo
 
 from torch.fx.passes.infra.pass_base import PassResult

py/torch_tensorrt/fx/utils.py

@@ -12,7 +12,7 @@
     replace_op_with_indices,
     run_const_fold,
 )
-
+from torch_tensorrt._util import sanitized_torch_version
 from .types import Shape, TRTDataType
 
 
@@ -160,7 +160,7 @@ def nested_decorator(f: Callable):
         def function_wrapper(*args, **kwargs):
             # Parse minimum and current Torch versions
             min_version = version.parse(min_torch_version)
-            current_version = version.parse(torch.__version__)
+            current_version = version.parse(sanitized_torch_version())
 
             if current_version < min_version:
                 raise AssertionError(