Merge pull request #1033 from NVIDIA/fixing_master

Fixing some issues with the master branch

Author: narendasan

core/conversion/evaluators/aten.cpp

@@ -342,15 +342,20 @@ auto aten_registrations TORCHTRT_UNUSED =
                         auto a = args.at(n->input(0)).unwrapToDouble();
                         auto b = args.at(n->input(1)).unwrapToDouble();
                         return a + b;
+                      } else if (args.at(n->input(0)).IValue()->isString()) {
+                        auto a = args.at(n->input(0)).unwrapToString();
+                        auto b = args.at(n->input(1)).unwrapToString();
+                        return a + b;
                       } else {
                         TORCHTRT_THROW_ERROR(
                             "Unimplemented data type for aten::add evaluator: "
                             << args.at(n->input(0)).IValue()->type()->str());
                         return {};
                       }
                     },
-                    EvalOptions().validSchemas(
-                        {"aten::add.int(int a, int b) -> (int)", "aten::add.float(float a, float b) -> (float)"})})
+                    EvalOptions().validSchemas({"aten::add.int(int a, int b) -> (int)",
+                                                "aten::add.float(float a, float b) -> (float)",
+                                                "aten::add.str(str a, str b) -> (str)"})})
         .evaluator({c10::Symbol::fromQualString("aten::add_"),
                     [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
                       if (args.at(n->input(0)).IValue()->isList()) {

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) {
-      // Neither arm matches the pattern
-      return false;
-    }
+    // if (!arm1_starts_with_exception && !arm2_starts_with_exception) {
+    // Neither arm matches the pattern
+    //   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

core/lowering/register_trt_placeholder_ops.cpp

@@ -10,7 +10,10 @@ c10::AliasAnalysisKind aliasAnalysisFromSchema() {
 RegisterOperators trt_placeholder_ops_reg({
     /// Op marks a Tensor to be conveted from an Torch Tensor
     /// to a TRT constant Tensor
-    Operator("trt::const(Tensor val) -> Tensor", [](Stack& stack) { /*noop*/ }, aliasAnalysisFromSchema()),
+    Operator(
+        "trt::const(Tensor val) -> Tensor",
+        [](Stack& stack) { /*noop*/ },
+        aliasAnalysisFromSchema()),
 });
 
 } // namespace jit

core/partitioning/partitioning.cpp

@@ -25,8 +25,10 @@ TEST_P(AccuracyTests, FP16AccuracyIsClose) {
   }
   torch::Tensor jit_accuracy = (jit_correct / jit_total) * 100;
 
-  std::vector<std::vector<int64_t>> input_shape = {{32, 3, 32, 32}};
-  auto compile_spec = torch_tensorrt::ts::CompileSpec({input_shape});
+  std::vector<int64_t> input_shape = {32, 3, 32, 32};
+  auto input = torch_tensorrt::Input(input_shape);
+  input.dtype = torch::kF16;
+  auto compile_spec = torch_tensorrt::ts::CompileSpec({input});
   compile_spec.enabled_precisions.insert(torch::kF16);
 
   auto trt_mod = torch_tensorrt::ts::compile(mod, compile_spec);

tests/accuracy/test_fp16_accuracy.cpp

@@ -42,15 +42,17 @@ TEST(LoweringPasses, EliminateExceptionOrPassPattern_Block0) {
   g->insertNode(bool_node);
   auto if_node = g->create(torch::jit::prim::If, {bool_node->output()}, 0);
   auto if_block0 = if_node->addBlock();
-  auto exception_node = g->create(torch::jit::prim::RaiseException, {except_val}, 0);
+  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);
   g->registerOutput(cat_node->output());
 
+  std::cout << "Source Graph: " << *g << std::endl;
   torch_tensorrt::core::lowering::passes::EliminateExceptionOrPassPattern(g);
+  std::cout << "Modified Graph: " << *g << std::endl;
   for (auto node : g->nodes()) {
     EXPECT_NE(node, if_node);
   }
@@ -95,16 +97,18 @@ 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}, 0);
+  auto exception_node = g->create(torch::jit::prim::RaiseException, {except_val, none_const_val}, 0);
   if_block1->appendNode(exception_node);
   g->insertNode(if_node);
   auto cat_node = g->create(torch::jit::aten::cat, {list_node->output(), zero_const_val});
   g->insertNode(cat_node);
   g->registerOutput(cat_node->output());
 
+  std::cout << "Source Graph: " << *g << std::endl;
   torch_tensorrt::core::lowering::passes::EliminateExceptionOrPassPattern(g);
+  std::cout << "Modified Graph: " << *g << std::endl;
   for (auto node : g->nodes()) {
     EXPECT_NE(node, if_node);
   }
@@ -150,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);