Address comments on PR

- Update utility name to add_abs for conciseness
- Refactor absolute value utility to return ITensor*
- Update logging level for certain debug messages

Author: gs-olive

core/conversion/converters/converter_util.cpp

@@ -156,7 +156,7 @@ nvinfer1::ILayer* add_elementwise(
   return ele;
 }
 
-nvinfer1::ILayer* add_absolute_value(
+nvinfer1::ITensor* add_abs(
     ConversionCtx* ctx,
     const torch::jit::Node* n,
     nvinfer1::ITensor* self,
@@ -185,7 +185,7 @@ nvinfer1::ILayer* add_absolute_value(
     TORCHTRT_CHECK(absolute_value_layer, "Unable to create max layer from node: " << *n);
   }
 
-  return absolute_value_layer;
+  return absolute_value_layer->getOutput(0);
 }
 
 nvinfer1::ITensor* applyIdentityOp(ConversionCtx* ctx, nvinfer1::ITensor* tensor, const std::string& tensor_name) {

core/conversion/converters/converter_util.h

@@ -35,14 +35,16 @@ nvinfer1::ITensor* addUnpadding(
     bool trailing = true,
     bool use_zeros = true);
 
+// TODO: Change add_elementwise schema to output nvinfer1::ITensor* instead, for consistency with other utils
+// Need to change schema and usage in all calling functions
 nvinfer1::ILayer* add_elementwise(
     ConversionCtx* ctx,
     nvinfer1::ElementWiseOperation op,
     nvinfer1::ITensor* self,
     nvinfer1::ITensor* other,
     const std::string& name);
 
-nvinfer1::ILayer* add_absolute_value(
+nvinfer1::ITensor* add_abs(
     ConversionCtx* ctx,
     const torch::jit::Node* n,
     nvinfer1::ITensor* self,

core/conversion/converters/impl/element_wise.cpp

@@ -326,31 +326,27 @@ 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 = add_absolute_value(ctx, n, tmp_div->getOutput(0), util::node_info(n) + "_absolute_val");
+                 auto abs = add_abs(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()
+                 nvinfer1::ITensor* floor;
+
+                 if ((abs->getType() == nvinfer1::DataType::kINT32) || (abs->getType() == nvinfer1::DataType::kBOOL)) {
+                   LOG_DEBUG(
+                       "Tensor is of unsupported type " << abs->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);
+                   floor = abs;
                  } else {
-                   floor = ctx->net->addUnary(*abs->getOutput(0), nvinfer1::UnaryOperation::kFLOOR);
-                   TORCHTRT_CHECK(floor, "Unable to create floor layer from node: " << *n);
+                   auto floor_layer = ctx->net->addUnary(*abs, nvinfer1::UnaryOperation::kFLOOR);
+                   TORCHTRT_CHECK(floor_layer, "Unable to create floor layer from node: " << *n);
+                   floor_layer->setName((util::node_info(n) + "_floor").c_str());
+                   floor = floor_layer->getOutput(0);
                  }
-                 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,
-                     nvinfer1::ElementWiseOperation::kPROD,
-                     floor->getOutput(0),
-                     sign->getOutput(0),
-                     util::node_info(n));
+                     ctx, nvinfer1::ElementWiseOperation::kPROD, floor, sign->getOutput(0), util::node_info(n));
                } else {
                  div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n));
                }

core/conversion/converters/impl/unary.cpp

@@ -13,8 +13,8 @@ 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);
-       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));
+       auto abs_tensor = add_abs(ctx, n, in, util::node_info(n));
+       auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], abs_tensor);
        LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
        return true;
      }});