Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: Error with aten::div when using truncation with Int32 tensor inputs #1442

Merged
merged 1 commit into from
Nov 18, 2022
Merged

fix: Error with aten::div when using truncation with Int32 tensor inputs #1442

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 32 additions & 0 deletions core/conversion/converters/converter_util.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,38 @@ nvinfer1::ILayer* add_elementwise(
return ele;
}

nvinfer1::ITensor* add_abs(
ConversionCtx* ctx,
const torch::jit::Node* n,
nvinfer1::ITensor* self,
const std::string& name) {
nvinfer1::ILayer* absolute_value_layer;

// 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_layer = ctx->net->addUnary(*self, nvinfer1::UnaryOperation::kABS);
TORCHTRT_CHECK(absolute_value_layer, "Unable to create abs layer from node: " << *n);
absolute_value_layer->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_layer =
add_elementwise(ctx, nvinfer1::ElementWiseOperation::kMAX, self, neg_layer->getOutput(0), name);
TORCHTRT_CHECK(absolute_value_layer, "Unable to create max layer from node: " << *n);
}

return absolute_value_layer->getOutput(0);
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Switched function schema to return the output of the absolute value layer and not the layer itself

}

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);
Expand Down
8 changes: 8 additions & 0 deletions core/conversion/converters/converter_util.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,13 +35,21 @@ nvinfer1::ITensor* addUnpadding(
bool trailing = true,
bool use_zeros = true);

// TODO: Change add_elementwise schema to output nvinfer1::ITensor* instead of nvinfer1::ILayer*,
// for consistency with other utils. Need to change schema and usage in all calling contexts
nvinfer1::ILayer* add_elementwise(
ConversionCtx* ctx,
nvinfer1::ElementWiseOperation op,
nvinfer1::ITensor* self,
nvinfer1::ITensor* other,
const std::string& name);

nvinfer1::ITensor* add_abs(
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);

Expand Down
26 changes: 19 additions & 7 deletions core/conversion/converters/impl/element_wise.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -326,15 +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 = ctx->net->addUnary(*tmp_div->getOutput(0), nvinfer1::UnaryOperation::kABS);
auto floor = ctx->net->addUnary(*abs->getOutput(0), nvinfer1::UnaryOperation::kFLOOR);
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::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 = abs;
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instead of using the identity function, the floor output tensor is just set to the absolute value result, to avoid unnecessary computation.

} else {
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);
}

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));
}
Expand Down
38 changes: 4 additions & 34 deletions core/conversion/converters/impl/unary.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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_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;
}});

auto reciprocal_registration TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern(
Expand Down
24 changes: 24 additions & 0 deletions tests/core/conversion/converters/test_element_wise.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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):
Expand Down