fix: Error on aten::div with truncation

- `aten::div` with truncation on integer tensor inputs currently throws
an error if both inputs are integer type, as the TRT unary operations
for absolute value and floor do not apply to Int32 or Bool types
- For absolute value, this is a legitimate bug as `aten::abs` is
functional for integer types
- For the floor operation, `aten::floor` does not explicitly support
integer inputs, and `torch.floor()` does not work with Int32 inputs by
default. However, `torch.div(..., rounding_mode="trunc")` with integer
tensors does return an integer value, and so the corollary Torch-TRT
converter should behave similarly
- Modified `aten:abs` converter logic to be a utility, as it is used in
multiple locations
- Added regression test to ensure truncation divide with two integer
tensors is functional

Author: gs-olive

core/conversion/converters/converter_util.cpp

@@ -156,6 +156,37 @@ nvinfer1::ILayer* add_elementwise(
   return ele;
 }
 
+nvinfer1::ILayer* add_absolute_value(
+    ConversionCtx* ctx,
+    const torch::jit::Node* n,
+    nvinfer1::ITensor* self,
+    const std::string& name) {
+  nvinfer1::ILayer* absolute_value;
+
+  // Check if TRT Unary ops support the input type
+  bool unary_supported_input = (self->getType() == nvinfer1::DataType::kFLOAT) ||
+      (self->getType() == nvinfer1::DataType::kHALF) || (self->getType() == nvinfer1::DataType::kINT8);
+  if (unary_supported_input) {
+    absolute_value = ctx->net->addUnary(*self, nvinfer1::UnaryOperation::kABS);
+    TORCHTRT_CHECK(absolute_value, "Unable to create abs layer from node: " << *n);
+    absolute_value->setName(name.c_str());
+  } else {
+    LOG_GRAPH(
+        "Tensor is of unsupported type "
+        << self->getType() << " for IUnaryLayer::kABS. Using backup implementation via IElementWise (max(x, -x)");
+    // 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(self->getType()));
+    auto neg_one_const = tensor_to_const(ctx, neg_one);
+    auto neg_layer = add_elementwise(
+        ctx, nvinfer1::ElementWiseOperation::kPROD, self, neg_one_const, util::node_info(n) + std::string("_Negation"));
+    TORCHTRT_CHECK(neg_layer, "Unable to create prod layer from node: " << *n);
+    absolute_value = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kMAX, self, neg_layer->getOutput(0), name);
+    TORCHTRT_CHECK(absolute_value, "Unable to create max layer from node: " << *n);
+  }
+
+  return absolute_value;
+}
+
 nvinfer1::ITensor* applyIdentityOp(ConversionCtx* ctx, nvinfer1::ITensor* tensor, const std::string& tensor_name) {
   auto id_layer = ctx->net->addIdentity(*tensor);
   auto id_out_tensor = id_layer->getOutput(0);

core/conversion/converters/converter_util.h

@@ -42,6 +42,12 @@ nvinfer1::ILayer* add_elementwise(
     nvinfer1::ITensor* other,
     const std::string& name);
 
+nvinfer1::ILayer* add_absolute_value(
+    ConversionCtx* ctx,
+    const torch::jit::Node* n,
+    nvinfer1::ITensor* self,
+    const std::string& name);
+
 // Apply an identity operation on a tensor. Used in the case where an input is an output to a network.
 nvinfer1::ITensor* applyIdentityOp(ConversionCtx* ctx, nvinfer1::ITensor* tensor, const std::string& name);
 

core/conversion/converters/impl/element_wise.cpp

@@ -326,8 +326,24 @@ auto element_wise_registrations TORCHTRT_UNUSED =
                } else if (rounding_mode == "trunc") {
                  // trunc = floor(abs(div)) * sign(div)
                  auto tmp_div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, "tmp_div");
-                 auto abs = ctx->net->addUnary(*tmp_div->getOutput(0), nvinfer1::UnaryOperation::kABS);
-                 auto floor = ctx->net->addUnary(*abs->getOutput(0), nvinfer1::UnaryOperation::kFLOOR);
+                 auto abs = add_absolute_value(ctx, n, tmp_div->getOutput(0), util::node_info(n) + "_absolute_val");
+
+                 // In this case, we allow the floor unary on non-TRT Unary types, as it is needed for this
+                 // specific function. Floor applied to non-float types equates to identity
+                 nvinfer1::ILayer* floor;
+                 if ((abs->getOutput(0)->getType() == nvinfer1::DataType::kINT32) ||
+                     (abs->getOutput(0)->getType() == nvinfer1::DataType::kBOOL)) {
+                   LOG_GRAPH(
+                       "Tensor is of unsupported type " << abs->getOutput(0)->getType()
+                                                        << " for IUnaryLayer::kFLOOR. Using identity instead.");
+                   floor = ctx->net->addIdentity(*abs->getOutput(0));
+                   TORCHTRT_CHECK(floor, "Unable to create identity layer from node: " << *n);
+                 } else {
+                   floor = ctx->net->addUnary(*abs->getOutput(0), nvinfer1::UnaryOperation::kFLOOR);
+                   TORCHTRT_CHECK(floor, "Unable to create floor layer from node: " << *n);
+                 }
+                 floor->setName((util::node_info(n) + "_floor").c_str());
+
                  auto sign = ctx->net->addUnary(*tmp_div->getOutput(0), nvinfer1::UnaryOperation::kSIGN);
                  div = add_elementwise(
                      ctx,

core/conversion/converters/impl/unary.cpp

@@ -13,40 +13,10 @@ 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].ITensorOrFreeze(ctx);
-       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 {
-         LOG_GRAPH(
-             "Tensor is of unsupported type "
-             << in->getType() << " for IUnaryLayer::kABS. Using backup implementation via IElementWise (max(x, -x)");
-         // 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;
-       }
+       auto abs_layer = add_absolute_value(ctx, n, in, util::node_info(n));
+       auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], abs_layer->getOutput(0));
+       LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
+       return true;
      }});
 
 auto reciprocal_registration TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern(

tests/core/conversion/converters/test_element_wise.cpp

@@ -4,6 +4,7 @@
 #include "gtest/gtest.h"
 #include "tests/util/util.h"
 #include "torch/csrc/jit/ir/irparser.h"
+#include "torch/torch.h"
 
 void pointwise_test_helper(
     std::string graph_ir,
@@ -235,6 +236,29 @@ TEST(Converters, ATenDivRoundingNoneConvertsCorrectly) {
   pointwise_test_helper(graph, false, true, {4, 3}, {3, 4, 3}, true);
 }
 
+TEST(Converters, ATenDivRoundingTruncWithIntsConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor, %1 : Tensor):
+        %trunc : str = prim::Constant[value="trunc"]()
+        %out : Tensor = aten::div(%0, %1, %trunc)
+        return (%out))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  // Avoid divide-by-zero issues by making denominator >= 1
+  auto in_0 = at::randint(-5, 5, {4, 1, 7, 8}, {at::kCUDA}).to(torch::kInt32);
+  auto in_1 = at::randint(1, 10, {4, 1, 7, 8}, {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_0, in_1});
+
+  params = torch_tensorrt::core::ir::get_static_params(g->inputs(), {});
+  auto trt_results = torch_tensorrt::tests::util::RunGraphEngine(g, params, {in_0, in_1});
+
+  ASSERT_TRUE(torch_tensorrt::tests::util::exactlyEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0])));
+}
+
 TEST(Converters, ATenPowTensorConvertsCorrectly) {
   const auto graph = R"IR(
        graph(%x.1 : Tensor, %x2.1 : Tensor):