Merge remote-tracking branch 'origin/master' into squashed_collections

Author: narendasan

.circleci/config.yml

@@ -747,7 +747,7 @@ parameters:
   # Nightly platform config
   torch-nightly-build:
     type: string
-    default: "1.13.0.dev20220715+cu113"
+    default: "1.13.0.dev20220731+cu113"
   torch-nightly-build-index:
     type: string
     default: "https://download.pytorch.org/whl/nightly/cu113"

core/conversion/converters/impl/unary.cpp

@@ -1,13 +1,55 @@
 #include "core/conversion/converters/converters.h"
 #include "core/util/prelude.h"
 
+#include "torch/torch.h"
+
 namespace torch_tensorrt {
 namespace core {
 namespace conversion {
 namespace converters {
 namespace impl {
 namespace {
 
+
+auto abs_registration TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern(
+  {"aten::abs (Tensor self) -> Tensor",
+  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+    auto in = args[0].ITensor();
+    bool unary_supported_input = in->getType() == nvinfer1::DataType::kFLOAT
+     || in->getType() == nvinfer1::DataType::kHALF
+     || in->getType() == nvinfer1::DataType::kINT8;
+    if(unary_supported_input){
+      auto unary_layer = ctx->net->addUnary(*in, nvinfer1::UnaryOperation::kABS);
+      TORCHTRT_CHECK(unary_layer, "Unable to create abs layer from node: " << *n);
+      unary_layer->setName(util::node_info(n).c_str());
+      auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], unary_layer->getOutput(0));
+      LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
+      return true;
+    }
+    else{
+      //For types not supported by kABS, use an elementwise implementation abs(x) = max(x, -1 * x)
+      at::Tensor neg_one = torch::full({1}, -1).to(util::TRTDataTypeToScalarType(in->getType()));
+      auto neg_one_const = tensor_to_const(ctx, neg_one);
+      auto neg_layer = add_elementwise(
+             ctx,
+             nvinfer1::ElementWiseOperation::kPROD,
+             in,
+             neg_one_const,
+             util::node_info(n) + std::string("_Negation"));
+      TORCHTRT_CHECK(neg_layer, "Unable to create prod layer from node: " << *n);
+      auto max_layer = add_elementwise(
+             ctx,
+             nvinfer1::ElementWiseOperation::kMAX,
+             in,
+             neg_layer->getOutput(0),
+             util::node_info(n) + std::string("_Max"));
+      TORCHTRT_CHECK(max_layer, "Unable to create max layer from node: " << *n);
+      auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], max_layer->getOutput(0));
+      LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
+      return true;
+    }
+  }});
+
 #define convert(unary, trt_type)                                                               \
   auto unary##_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern(       \
       {"aten::" #unary "(Tensor self) -> Tensor",                                              \
@@ -32,7 +74,6 @@ convert(asin, kASIN);
 convert(sinh, kSINH);
 convert(tan, kTAN);
 convert(atan, kATAN);
-convert(abs, kABS);
 convert(floor, kFLOOR);
 convert(reciprocal, kRECIP);
 convert(log, kLOG);

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

py/torch_tensorrt/fx/test/converters/acc_op/test_convolution.py

@@ -144,7 +144,8 @@ def forward(self, x):
             ("tuple_parameters", 1, (1, 1, 1), (1, 1, 1)),
             param("non_zero_padding", 1, padding=1),
             param("dilation", 1, dilation=2),
-            param("groups", 1, groups=3),
+            # TODO TRT 8.4.1 will trigger issue with this test. T127981773
+            # param("groups", 1, groups=3),
         ]
     )
     def test_conv3d(

py/torch_tensorrt/fx/test/converters/acc_op/test_type_as.py

@@ -1,4 +1,5 @@
 import torch
+import unittest
 import torch_tensorrt.fx.tracer.acc_tracer.acc_ops as acc_ops
 from torch.testing._internal.common_utils import run_tests
 from torch_tensorrt.fx.tools.common_fx2trt import AccTestCase, InputTensorSpec
@@ -102,7 +103,8 @@ def forward(self, input):
             expected_ops={acc_ops.to_dtype},
             precision=LowerPrecision.FP16,
         )
-
+        
+    @unittest.skip("Does not pass in TRT 8.4.1 T127981773")
     def test_type_tensor_with_dynamic_shape_four_dimensions(self):
         class Type_as(torch.nn.Module):
             def forward(self, input):

tests/core/conversion/converters/test_unary.cpp

@@ -1,4 +1,5 @@
 #include <string>
+#include "torch/torch.h"
 #include "core/compiler.h"
 #include "gtest/gtest.h"
 #include "tests/util/util.h"
@@ -14,6 +15,22 @@ std::string gen_test_graph(const std::string& unary) {
 }
 } // namespace
 
+TEST(Converters, ATenAbsIntConvertsCorrectly) {
+  const auto graph = gen_test_graph("abs");
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  auto in = at::tensor({-1, 1, -2, 2, -3, 3}, {at::kCUDA}).to(torch::kInt32);
+  auto params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto jit_results = torch_tensorrt::tests::util::RunGraph(g, params, {in});  
+
+  in = at::clone(in);
+  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in});
+
+  ASSERT_TRUE(torch_tensorrt::tests::util::exactlyEqual(jit_results[0], trt_results[0]));
+}
+
 #define test_unary(unary, name)                                                                  \
   TEST(Converters, ATen##name##ConvertsCorrectly) {                                              \
     const auto graph = gen_test_graph(#unary);                                                   \