Skip to content

Convert the output type of bool operation to input type #350

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

Closed
wants to merge 1 commit into from
Closed

Convert the output type of bool operation to input type #350

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
78 changes: 41 additions & 37 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 @@ -9,6 +9,14 @@ namespace converters {
namespace impl {
namespace {

nvinfer1::IIdentityLayer* convert_to_op_input_type(nvinfer1::ILayer* op, ConversionCtx* ctx, const torch::jit::Node* n) {
auto identity = ctx->net->addIdentity(*op->getOutput(0));
identity->setOutputType(0, op->getInput(0)->getType());
TRTORCH_CHECK(identity, "Unable to create Identity layer from node: " << *n);
identity->setName(util::node_info(n).c_str());
return identity;
}

nvinfer1::ILayer* add_elementwise(
ConversionCtx* ctx,
nvinfer1::ElementWiseOperation op,
Expand Down Expand Up @@ -340,7 +348,7 @@ auto element_wise_registrations TRTORCH_UNUSED =
LOG_DEBUG("Output tensor shape: " << out->getDimensions());
return true;
}})
.pattern({"aten::div_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)",
.pattern({"aten::div_.Scalar(Tensor self, Scalar other) -> (Tensor)",
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
auto self = args[0].ITensorOrFreeze(ctx);
auto otherScalar = args[1].unwrapToScalar().to<float>();
Expand Down Expand Up @@ -535,12 +543,12 @@ auto element_wise_registrations TRTORCH_UNUSED =
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
auto self = args[0].ITensorOrFreeze(ctx);
auto other = args[1].ITensorOrFreeze(ctx);
auto gt =
add_elementwise(ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n));
auto gt = add_elementwise(
ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n) + "_gt");
TRTORCH_CHECK(gt, "Unable to create greater layer from node: " << *n);

gt->setName(util::node_info(n).c_str());
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], gt->getOutput(0));
auto identity = convert_to_op_input_type(gt, ctx, n);
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], identity->getOutput(0));
LOG_DEBUG("Output tensor shape: " << out->getDimensions());
return true;
}})
Expand All @@ -549,25 +557,25 @@ auto element_wise_registrations TRTORCH_UNUSED =
auto self = args[0].ITensorOrFreeze(ctx);
auto otherScalar = args[1].unwrapToScalar().to<float>();
auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
auto gt =
add_elementwise(ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n));
auto gt = add_elementwise(
ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n) + "_gt");
TRTORCH_CHECK(gt, "Unable to create greater layer from node: " << *n);

gt->setName(util::node_info(n).c_str());
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], gt->getOutput(0));
auto identity = convert_to_op_input_type(gt, ctx, n);
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], identity->getOutput(0));
LOG_DEBUG("Output tensor shape: " << out->getDimensions());
return true;
}})
.pattern({"aten::lt.Tensor(Tensor self, Tensor other) -> (Tensor)",
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
auto self = args[0].ITensorOrFreeze(ctx);
auto other = args[1].ITensorOrFreeze(ctx);
auto lt =
add_elementwise(ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n));
auto lt = add_elementwise(
ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n) + "_lt");
TRTORCH_CHECK(lt, "Unable to create less layer from node: " << *n);

lt->setName(util::node_info(n).c_str());
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], lt->getOutput(0));
auto identity = convert_to_op_input_type(lt, ctx, n);
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], identity->getOutput(0));
LOG_DEBUG("Output tensor shape: " << out->getDimensions());
return true;
}})
Expand All @@ -576,25 +584,25 @@ auto element_wise_registrations TRTORCH_UNUSED =
auto self = args[0].ITensorOrFreeze(ctx);
auto otherScalar = args[1].unwrapToScalar().to<float>();
auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
auto lt =
add_elementwise(ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n));
auto lt = add_elementwise(
ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n) + "_lt");
TRTORCH_CHECK(lt, "Unable to create less layer from node: " << *n);

lt->setName(util::node_info(n).c_str());
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], lt->getOutput(0));
auto identity = convert_to_op_input_type(lt, ctx, n);
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], identity->getOutput(0));
LOG_DEBUG("Output tensor shape: " << out->getDimensions());
return true;
}})
.pattern({"aten::eq.Tensor(Tensor self, Tensor other) -> (Tensor)",
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
auto self = args[0].ITensorOrFreeze(ctx);
auto other = args[1].ITensorOrFreeze(ctx);
auto eq =
add_elementwise(ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n));
auto eq = add_elementwise(
ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_eq");
TRTORCH_CHECK(eq, "Unable to create equal layer from node: " << *n);

eq->setName(util::node_info(n).c_str());
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], eq->getOutput(0));
auto identity = convert_to_op_input_type(eq, ctx, n);
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], identity->getOutput(0));
LOG_DEBUG("Output tensor shape: " << out->getDimensions());
return true;
}})
Expand All @@ -603,12 +611,12 @@ auto element_wise_registrations TRTORCH_UNUSED =
auto self = args[0].ITensorOrFreeze(ctx);
auto otherScalar = args[1].unwrapToScalar().to<float>();
auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
auto eq =
add_elementwise(ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n));
auto eq = add_elementwise(
ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_eq");
TRTORCH_CHECK(eq, "Unable to create equal layer from node: " << *n);

eq->setName(util::node_info(n).c_str());
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], eq->getOutput(0));
auto identity = convert_to_op_input_type(eq, ctx, n);
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], identity->getOutput(0));
LOG_DEBUG("Output tensor shape: " << out->getDimensions());
return true;
}})
Expand All @@ -629,9 +637,8 @@ auto element_wise_registrations TRTORCH_UNUSED =
*greater->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR);

TRTORCH_CHECK(or_op, "Unable to create Or layer from node: " << *n);
or_op->setName(util::node_info(n).c_str());
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0));

auto identity = convert_to_op_input_type(or_op, ctx, n);
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], identity->getOutput(0));
LOG_DEBUG("Output tensor shape: " << out->getDimensions());
return true;
}})
Expand All @@ -653,9 +660,8 @@ auto element_wise_registrations TRTORCH_UNUSED =
*greater->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR);

TRTORCH_CHECK(or_op, "Unable to create Or layer from node: " << *n);
or_op->setName(util::node_info(n).c_str());
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0));

auto identity = convert_to_op_input_type(or_op, ctx, n);
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], identity->getOutput(0));
LOG_DEBUG("Output tensor shape: " << out->getDimensions());
return true;
}})
Expand All @@ -676,9 +682,8 @@ auto element_wise_registrations TRTORCH_UNUSED =
*less->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR);

TRTORCH_CHECK(or_op, "Unable to create Or layer from node: " << *n);
or_op->setName(util::node_info(n).c_str());
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0));

auto identity = convert_to_op_input_type(or_op, ctx, n);
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], identity->getOutput(0));
LOG_DEBUG("Output tensor shape: " << out->getDimensions());
return true;
}})
Expand All @@ -700,9 +705,8 @@ auto element_wise_registrations TRTORCH_UNUSED =
*less->getOutput(0), *equal->getOutput(0), nvinfer1::ElementWiseOperation::kOR);

TRTORCH_CHECK(or_op, "Unable to create Or layer from node: " << *n);
or_op->setName(util::node_info(n).c_str());
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], or_op->getOutput(0));

auto identity = convert_to_op_input_type(or_op, ctx, n);
auto out = ctx->AssociateValueAndTensor(n->outputs()[0], identity->getOutput(0));
LOG_DEBUG("Output tensor shape: " << out->getDimensions());
return true;
}});
Expand Down