-
Notifications
You must be signed in to change notification settings - Fork 372
feat!: Updating versions of CUDA, cuDNN, TensorRT and PyTorch #389
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.
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
Conversation
BREAKING CHANGE: PyTorch version has been bumped to 1.8.0 Default CUDA version is CUDA 11.1 TensorRT version is TensorRT 7.2.3.4 cuDNN version is now cuDNN 8.1 Signed-off-by: Naren Dasan <naren@narendasan.com> Signed-off-by: Naren Dasan <narens@nvidia.com>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Code conforms to Python style guidelines
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There are some changes that do not conform to C++ style guidelines:
diff --git a/workspace/tests/core/conversion/converters/test_activation.cpp b/tmp/changes.txt
index f3f3897..053b5c5 100644
--- a/workspace/tests/core/conversion/converters/test_activation.cpp
+++ b/tmp/changes.txt
@@ -85,8 +85,7 @@ TEST(Converters, ATenHardTanhConvertsCorrectly) {
params = trtorch::core::conversion::get_named_params(g->inputs(), {});
auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
- ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0],
- trt_results[0], 2e-6));
+ ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6));
}
TEST(Converters, ATenHardTanhCustomRangeConvertsCorrectly) {
ERROR: Some files do not conform to style guidelines
models to fail while trace models work. Seems to be down to the fact that the two create different graphs and script was having issues with aten::linear Signed-off-by: Naren Dasan <naren@narendasan.com> Signed-off-by: Naren Dasan <narens@nvidia.com>
preparation for 1.9 removing the preprocessing function from the backend class Signed-off-by: Naren Dasan <naren@narendasan.com> Signed-off-by: Naren Dasan <narens@nvidia.com>
Signed-off-by: Naren Dasan <naren@narendasan.com> Signed-off-by: Naren Dasan <narens@nvidia.com>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There are some changes that do not conform to C++ style guidelines:
diff --git a/workspace/tests/core/conversion/converters/test_activation.cpp b/tmp/changes.txt
index f3f3897..053b5c5 100644
--- a/workspace/tests/core/conversion/converters/test_activation.cpp
+++ b/tmp/changes.txt
@@ -85,8 +85,7 @@ TEST(Converters, ATenHardTanhConvertsCorrectly) {
params = trtorch::core::conversion::get_named_params(g->inputs(), {});
auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
- ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0],
- trt_results[0], 2e-6));
+ ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6));
}
TEST(Converters, ATenHardTanhCustomRangeConvertsCorrectly) {
ERROR: Some files do not conform to style guidelines
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Code conforms to Python style guidelines
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There are some changes that do not conform to C++ style guidelines:
diff --git a/workspace/tests/core/conversion/converters/test_activation.cpp b/tmp/changes.txt
index f3f3897..053b5c5 100644
--- a/workspace/tests/core/conversion/converters/test_activation.cpp
+++ b/tmp/changes.txt
@@ -85,8 +85,7 @@ TEST(Converters, ATenHardTanhConvertsCorrectly) {
params = trtorch::core::conversion::get_named_params(g->inputs(), {});
auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
- ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0],
- trt_results[0], 2e-6));
+ ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6));
}
TEST(Converters, ATenHardTanhCustomRangeConvertsCorrectly) {
ERROR: Some files do not conform to style guidelines
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Code conforms to Python style guidelines
@peri044 can you take a look. make sure to delete and redownload all the modules from |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Code conforms to Python style guidelines
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Code conforms to C++ style guidelines
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM. Had some minor questions. Tested it on my local and it runs fine.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There are some changes that do not conform to Python style guidelines:
Reformatting /workspace/docsrc/conf.py
Reformatting /workspace/tests/modules/hub.py
Reformatting /workspace/tests/py/test_api_dla.py
Reformatting /workspace/tests/py/test_multi_gpu.py
Reformatting /workspace/tests/py/test_to_backend_api.py
Reformatting /workspace/tests/py/model_test_case.py
Reformatting /workspace/tests/py/test_api.py
Reformatting /workspace/tests/py/test_ptq_trt_calibrator.py
Reformatting /workspace/tests/py/test_ptq_to_backend.py
Reformatting /workspace/tests/py/test_ptq_dataloader_calibrator.py
Reformatting /workspace/examples/custom_converters/elu_converter/setup.py
Reformatting /workspace/examples/custom_converters/elu_model.py
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There are some changes that do not conform to C++ style guidelines:
diff --git a/workspace/tests/core/conversion/converters/test_interpolate.cpp b/tmp/changes.txt
index e8ff41e..6f6d2fd 100644
--- a/workspace/tests/core/conversion/converters/test_interpolate.cpp
+++ b/tmp/changes.txt
@@ -9,7 +9,7 @@
const auto graph = graph_src; \
\
auto g = std::make_shared<torch::jit::Graph>(); \
- torch::jit::parseIR(graph, g.get()); \
+ torch::jit::parseIR(graph, g.get()); \
\
auto in = at::randint(1, 10, input_shape, {at::kCUDA}); \
auto jit_in = at::clone(in); \
@@ -28,7 +28,7 @@
const auto graph = graph_src; \
\
auto g = std::make_shared<torch::jit::Graph>(); \
- torch::jit::parseIR(graph, g.get()); \
+ torch::jit::parseIR(graph, g.get()); \
\
auto in = at::randint(1, 10, input_shape, {at::kCUDA}); \
auto jit_in = at::clone(in); \
@@ -48,7 +48,7 @@
const auto graph = graph_src; \
\
auto g = std::make_shared<torch::jit::Graph>(); \
- torch::jit::parseIR(graph, g.get()); \
+ torch::jit::parseIR(graph, g.get()); \
\
auto in = at::randint(1, 10, input_shape, {at::kCUDA}); \
auto jit_in = at::clone(in); \
diff --git a/workspace/tests/core/conversion/converters/test_unary.cpp b/tmp/changes.txt
index 4306888..2910f3a 100644
--- a/workspace/tests/core/conversion/converters/test_unary.cpp
+++ b/tmp/changes.txt
@@ -19,7 +19,7 @@ std::string gen_test_graph(const std::string& unary) {
const auto graph = gen_test_graph(#unary); \
\
auto g = std::make_shared<torch::jit::Graph>(); \
- torch::jit::parseIR(graph, g.get()); \
+ torch::jit::parseIR(graph, g.get()); \
\
float offset = 0; \
if (strcmp(#name, "Acosh") == 0) \
ERROR: Some files do not conform to style guidelines
Signed-off-by: Naren Dasan <naren@narendasan.com> Signed-off-by: Naren Dasan <narens@nvidia.com>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Code conforms to Python style guidelines
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Code conforms to C++ style guidelines
Signed-off-by: Naren Dasan <naren@narendasan.com> Signed-off-by: Naren Dasan <narens@nvidia.com>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There are some changes that do not conform to Python style guidelines:
Reformatting /workspace/tests/modules/hub.py
Reformatting /workspace/tests/py/test_api_dla.py
Reformatting /workspace/tests/py/test_multi_gpu.py
Reformatting /workspace/tests/py/test_to_backend_api.py
Reformatting /workspace/tests/py/model_test_case.py
Reformatting /workspace/tests/py/test_api.py
Reformatting /workspace/tests/py/test_ptq_trt_calibrator.py
Reformatting /workspace/tests/py/test_ptq_to_backend.py
Reformatting /workspace/tests/py/test_ptq_dataloader_calibrator.py
Reformatting /workspace/cpp/ptq/training/vgg16/vgg16.py
Reformatting /workspace/cpp/ptq/training/vgg16/export_ckpt.py
Reformatting /workspace/cpp/ptq/training/vgg16/main.py
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There are some changes that do not conform to C++ style guidelines:
diff --git a/workspace/core/conversion/converters/impl/expand.cpp b/tmp/changes.txt
index 16b5426..c620178 100644
--- a/workspace/core/conversion/converters/impl/expand.cpp
+++ b/tmp/changes.txt
@@ -196,91 +196,94 @@ bool add_expand_dynamic(
auto expand_registrations TRTORCH_UNUSED =
RegisterNodeConversionPatterns()
- .pattern({"aten::expand(Tensor(a) self, int[] size, *, bool implicit=False) -> (Tensor(a))",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto in = args[0].ITensor();
- auto input_dims = in->getDimensions();
- auto expanded_size = args[1].unwrapToIntList();
- auto expandedDims = util::toDims(expanded_size);
- LOG_DEBUG("(expand layer) Expand input from " << input_dims << " to " << expandedDims);
- if (ctx->input_is_dynamic) {
- at::Tensor thExpanded_size = torch::tensor(expanded_size.vec(), torch::kInt32);
- auto expandedDimsTensor = tensor_to_const(ctx, thExpanded_size);
- return add_expand_dynamic(ctx, n, in, expandedDimsTensor, expandedDims, true);
- } else {
- return add_expand(ctx, n, in, expandedDims);
- }
- }})
- .pattern({"aten::expand_as(Tensor(a) self, Tensor other) -> (Tensor(a))",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto in = args[0].ITensor();
- auto input_dims = in->getDimensions();
- auto targetTensor = args[1].ITensor();
- auto targetDims = targetTensor->getDimensions();
- LOG_DEBUG("(expand_as layer) Expand input from " << input_dims << " to " << targetDims);
- if (ctx->input_is_dynamic) {
- return add_expand_dynamic(
- ctx, n, in, ctx->net->addShape(*targetTensor)->getOutput(0), targetDims, false);
- } else {
- return add_expand(ctx, n, in, targetDims);
- }
- }})
- .pattern({"aten::repeat(Tensor self, int[] repeats) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto in = args[0].ITensor();
- auto input_dims = in->getDimensions();
- auto repeats = args[1].unwrapToIntList().vec();
- int repeats_rank = repeats.size();
- TRTORCH_CHECK(
- repeats_rank >= input_dims.nbDims,
- "Number of repeat dimensions cannot be smaller than number of input dimensions");
- auto num_expand_dims = repeats_rank - input_dims.nbDims;
+ .pattern(
+ {"aten::expand(Tensor(a) self, int[] size, *, bool implicit=False) -> (Tensor(a))",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto in = args[0].ITensor();
+ auto input_dims = in->getDimensions();
+ auto expanded_size = args[1].unwrapToIntList();
+ auto expandedDims = util::toDims(expanded_size);
+ LOG_DEBUG("(expand layer) Expand input from " << input_dims << " to " << expandedDims);
+ if (ctx->input_is_dynamic) {
+ at::Tensor thExpanded_size = torch::tensor(expanded_size.vec(), torch::kInt32);
+ auto expandedDimsTensor = tensor_to_const(ctx, thExpanded_size);
+ return add_expand_dynamic(ctx, n, in, expandedDimsTensor, expandedDims, true);
+ } else {
+ return add_expand(ctx, n, in, expandedDims);
+ }
+ }})
+ .pattern(
+ {"aten::expand_as(Tensor(a) self, Tensor other) -> (Tensor(a))",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto in = args[0].ITensor();
+ auto input_dims = in->getDimensions();
+ auto targetTensor = args[1].ITensor();
+ auto targetDims = targetTensor->getDimensions();
+ LOG_DEBUG("(expand_as layer) Expand input from " << input_dims << " to " << targetDims);
+ if (ctx->input_is_dynamic) {
+ return add_expand_dynamic(
+ ctx, n, in, ctx->net->addShape(*targetTensor)->getOutput(0), targetDims, false);
+ } else {
+ return add_expand(ctx, n, in, targetDims);
+ }
+ }})
+ .pattern(
+ {"aten::repeat(Tensor self, int[] repeats) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto in = args[0].ITensor();
+ auto input_dims = in->getDimensions();
+ auto repeats = args[1].unwrapToIntList().vec();
+ int repeats_rank = repeats.size();
+ TRTORCH_CHECK(
+ repeats_rank >= input_dims.nbDims,
+ "Number of repeat dimensions cannot be smaller than number of input dimensions");
+ auto num_expand_dims = repeats_rank - input_dims.nbDims;
- if (ctx->input_is_dynamic) {
- int input_rank = input_dims.nbDims;
- int output_rank = repeats_rank;
- auto new_input_shape_tensor = concat(output_rank, input_rank, ctx, in);
+ if (ctx->input_is_dynamic) {
+ int input_rank = input_dims.nbDims;
+ int output_rank = repeats_rank;
+ auto new_input_shape_tensor = concat(output_rank, input_rank, ctx, in);
- // Add a reshape layer to expand dims
- auto shuffle = ctx->net->addShuffle(*in);
- shuffle->setInput(1, *new_input_shape_tensor);
- in = shuffle->getOutput(0);
- } else {
- if (num_expand_dims > 0) {
- nvinfer1::Dims reshape_dims;
- reshape_dims.nbDims = repeats.size();
- for (int i = 0; i < num_expand_dims; i++) {
- reshape_dims.d[i] = 1;
- }
- for (int i = 0; i < input_dims.nbDims; i++) {
- reshape_dims.d[num_expand_dims + i] = input_dims.d[i];
- }
- // Add a reshape layer to expand dims
- auto reshape_layer = ctx->net->addShuffle(*in);
- reshape_layer->setReshapeDimensions(reshape_dims);
- in = reshape_layer->getOutput(0);
- LOG_DEBUG("Input reshaped to : " << in->getDimensions() << " from " << input_dims);
- }
- LOG_DEBUG("Repeats: " << repeats);
- }
+ // Add a reshape layer to expand dims
+ auto shuffle = ctx->net->addShuffle(*in);
+ shuffle->setInput(1, *new_input_shape_tensor);
+ in = shuffle->getOutput(0);
+ } else {
+ if (num_expand_dims > 0) {
+ nvinfer1::Dims reshape_dims;
+ reshape_dims.nbDims = repeats.size();
+ for (int i = 0; i < num_expand_dims; i++) {
+ reshape_dims.d[i] = 1;
+ }
+ for (int i = 0; i < input_dims.nbDims; i++) {
+ reshape_dims.d[num_expand_dims + i] = input_dims.d[i];
+ }
+ // Add a reshape layer to expand dims
+ auto reshape_layer = ctx->net->addShuffle(*in);
+ reshape_layer->setReshapeDimensions(reshape_dims);
+ in = reshape_layer->getOutput(0);
+ LOG_DEBUG("Input reshaped to : " << in->getDimensions() << " from " << input_dims);
+ }
+ LOG_DEBUG("Repeats: " << repeats);
+ }
- // Concat across all repeat axes.
- // TODO: Implementation might not be performant. Explore other strategies to improve performance.
- for (int i = repeats.size() - 1; i >= 0; --i) {
- std::vector<nvinfer1::ITensor*> tensors_vec;
- for (int j = 0; j < repeats[i]; j++) {
- tensors_vec.push_back(in);
- }
- auto concat_layer = ctx->net->addConcatenation(tensors_vec.data(), tensors_vec.size());
- concat_layer->setAxis(i);
- in = concat_layer->getOutput(0);
- }
+ // Concat across all repeat axes.
+ // TODO: Implementation might not be performant. Explore other strategies to improve performance.
+ for (int i = repeats.size() - 1; i >= 0; --i) {
+ std::vector<nvinfer1::ITensor*> tensors_vec;
+ for (int j = 0; j < repeats[i]; j++) {
+ tensors_vec.push_back(in);
+ }
+ auto concat_layer = ctx->net->addConcatenation(tensors_vec.data(), tensors_vec.size());
+ concat_layer->setAxis(i);
+ in = concat_layer->getOutput(0);
+ }
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], in);
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], in);
- LOG_DEBUG("Repeat layer output tensor shape: " << out->getDimensions());
- return true;
- }});
+ LOG_DEBUG("Repeat layer output tensor shape: " << out->getDimensions());
+ return true;
+ }});
} // namespace
} // namespace impl
diff --git a/workspace/core/conversion/converters/impl/activation.cpp b/tmp/changes.txt
index 64edeed..908dcb9 100644
--- a/workspace/core/conversion/converters/impl/activation.cpp
+++ b/tmp/changes.txt
@@ -42,24 +42,25 @@ convert(tanh, kTANH);
auto acthardtanh TRTORCH_UNUSED =
RegisterNodeConversionPatterns()
- .pattern({"aten::hardtanh(Tensor self, Scalar min_val=-1, Scalar max_val=1) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto in = args[0].ITensorOrFreeze(ctx);
- auto min = args[1].unwrapToDouble();
- auto max = args[2].unwrapToDouble();
+ .pattern(
+ {"aten::hardtanh(Tensor self, Scalar min_val=-1, Scalar max_val=1) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto in = args[0].ITensorOrFreeze(ctx);
+ auto min = args[1].unwrapToDouble();
+ auto max = args[2].unwrapToDouble();
- auto new_layer = ctx->net->addActivation(*in, nvinfer1::ActivationType::kCLIP);
- TRTORCH_CHECK(new_layer, "Unable to create layer for aten::hardtanh");
+ auto new_layer = ctx->net->addActivation(*in, nvinfer1::ActivationType::kCLIP);
+ TRTORCH_CHECK(new_layer, "Unable to create layer for aten::hardtanh");
- new_layer->setAlpha(min);
- new_layer->setBeta(max);
+ new_layer->setAlpha(min);
+ new_layer->setBeta(max);
- new_layer->setName(util::node_info(n).c_str());
- auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
+ new_layer->setName(util::node_info(n).c_str());
+ auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
- LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
- return true;
- }})
+ LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
+ return true;
+ }})
.pattern({// TODO: Remove after functionalization
"aten::hardtanh_(Tensor(a!) self, Scalar min_val=-1, Scalar max_val=1) -> (Tensor(a!))",
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
@@ -125,49 +126,52 @@ auto acthardtanh TRTORCH_UNUSED =
LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
return true;
}})
- .pattern({"aten::leaky_relu(Tensor self, Scalar negative_slope=0.01) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto self = args[0].ITensorOrFreeze(ctx);
- auto negative_slopeScalar = args[1].unwrapToScalar().to<float>();
+ .pattern(
+ {"aten::leaky_relu(Tensor self, Scalar negative_slope=0.01) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto negative_slopeScalar = args[1].unwrapToScalar().to<float>();
- auto new_layer = ctx->net->addActivation(*self, nvinfer1::ActivationType::kLEAKY_RELU);
- new_layer->setAlpha(negative_slopeScalar);
+ auto new_layer = ctx->net->addActivation(*self, nvinfer1::ActivationType::kLEAKY_RELU);
+ new_layer->setAlpha(negative_slopeScalar);
- new_layer->setName(util::node_info(n).c_str());
- auto out_tensor = new_layer->getOutput(0);
- out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], out_tensor);
- LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
- return true;
- }})
- .pattern({"aten::leaky_relu_(Tensor(a!) self, Scalar negative_slope=0.01) -> Tensor(a!)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto self = args[0].ITensorOrFreeze(ctx);
- auto negative_slopeScalar = args[1].unwrapToScalar().to<float>();
+ new_layer->setName(util::node_info(n).c_str());
+ auto out_tensor = new_layer->getOutput(0);
+ out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], out_tensor);
+ LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::leaky_relu_(Tensor(a!) self, Scalar negative_slope=0.01) -> Tensor(a!)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto negative_slopeScalar = args[1].unwrapToScalar().to<float>();
- auto new_layer = ctx->net->addActivation(*self, nvinfer1::ActivationType::kLEAKY_RELU);
- new_layer->setAlpha(negative_slopeScalar);
+ auto new_layer = ctx->net->addActivation(*self, nvinfer1::ActivationType::kLEAKY_RELU);
+ new_layer->setAlpha(negative_slopeScalar);
- new_layer->setName(util::node_info(n).c_str());
- auto out_tensor = new_layer->getOutput(0);
- out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], out_tensor);
- LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
- return true;
- }})
- .pattern({"aten::elu(Tensor self, Scalar alpha=1, Scalar scale=1, Scalar input_scale=1) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto in = args[0].ITensorOrFreeze(ctx);
- auto alpha = args[1].unwrapToDouble();
+ new_layer->setName(util::node_info(n).c_str());
+ auto out_tensor = new_layer->getOutput(0);
+ out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], out_tensor);
+ LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::elu(Tensor self, Scalar alpha=1, Scalar scale=1, Scalar input_scale=1) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto in = args[0].ITensorOrFreeze(ctx);
+ auto alpha = args[1].unwrapToDouble();
- auto new_layer = ctx->net->addActivation(*in, nvinfer1::ActivationType::kELU);
- TRTORCH_CHECK(new_layer, "Unable to create layer for aten::elu");
- new_layer->setAlpha(alpha);
+ auto new_layer = ctx->net->addActivation(*in, nvinfer1::ActivationType::kELU);
+ TRTORCH_CHECK(new_layer, "Unable to create layer for aten::elu");
+ new_layer->setAlpha(alpha);
- new_layer->setName(trtorch::core::util::node_info(n).c_str());
+ new_layer->setName(trtorch::core::util::node_info(n).c_str());
- auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
- LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
- return true;
- }});
+ auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
+ LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
+ return true;
+ }});
} // namespace
} // namespace impl
diff --git a/workspace/core/conversion/converters/impl/select.cpp b/tmp/changes.txt
index c883b69..1ef5157 100644
--- a/workspace/core/conversion/converters/impl/select.cpp
+++ b/tmp/changes.txt
@@ -65,111 +65,114 @@ bool add_split(ConversionCtx* ctx, const torch::jit::Node* n, args& args, bool s
auto select_registrations TRTORCH_UNUSED =
RegisterNodeConversionPatterns()
- .pattern({"aten::select.int(Tensor(a) self, int dim, int index) -> (Tensor(a))",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto in = args[0].ITensor();
- auto axis = args[1].unwrapToInt();
- auto ind = (int32_t)args[2].unwrapToInt();
-
- // index to access needs to be an at::Tensor
- at::Tensor indices = torch::tensor({ind}).to(torch::kI32);
- auto weights = Weights(ctx, indices);
-
- // IConstantLayer to convert indices from Weights to ITensor
- auto const_layer = ctx->net->addConstant(weights.shape, weights.data);
- TRTORCH_CHECK(const_layer, "Unable to create constant layer from node: " << *n);
- auto const_out = const_layer->getOutput(0);
-
- // IGatherLayer takes in input tensor, the indices, and the axis
- // of input tensor to take indices from
- auto gather_layer = ctx->net->addGather(*in, *const_out, axis);
- TRTORCH_CHECK(gather_layer, "Unable to create gather layer from node: " << *n);
- auto gather_out = gather_layer->getOutput(0);
-
- // IShuffleLayer removes redundant dimensions
- auto shuffle_layer = ctx->net->addShuffle(*gather_out);
- TRTORCH_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
- shuffle_layer->setReshapeDimensions(util::unpadDims(gather_out->getDimensions()));
- shuffle_layer->setName(util::node_info(n).c_str());
- auto shuffle_out = shuffle_layer->getOutput(0);
-
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle_out);
-
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
-
- return true;
- }})
- .pattern({"aten::narrow(Tensor(a) self, int dim, int start, int length) -> Tensor(a)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto in = args[0].ITensor();
- auto axis = args[1].unwrapToInt();
- auto start = (int32_t)args[2].unwrapToInt();
- auto length = (int32_t)args[3].unwrapToInt();
-
- // index to access needs to be an at::Tensor
- at::Tensor indices = torch::arange(start, start + length, 1).to(torch::kI32);
- auto weights = Weights(ctx, indices);
-
- // IConstantLayer to convert indices from Weights to ITensor
- auto const_layer = ctx->net->addConstant(weights.shape, weights.data);
- TRTORCH_CHECK(const_layer, "Unable to create constant layer from node: " << *n);
- auto const_out = const_layer->getOutput(0);
-
- // IGatherLayer takes in input tensor, the indices, and the axis
- // of input tensor to take indices from
- auto gather_layer = ctx->net->addGather(*in, *const_out, axis);
- TRTORCH_CHECK(gather_layer, "Unable to create gather layer from node: " << *n);
- auto gather_out = gather_layer->getOutput(0);
-
- // IShuffleLayer removes redundant dimensions
- auto shuffle_layer = ctx->net->addShuffle(*gather_out);
- TRTORCH_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
- shuffle_layer->setReshapeDimensions(util::unpadDims(gather_out->getDimensions()));
- shuffle_layer->setName(util::node_info(n).c_str());
- auto shuffle_out = shuffle_layer->getOutput(0);
-
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle_out);
-
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
-
- return true;
- }})
- .pattern({"aten::narrow.Tensor(Tensor(a) self, int dim, Tensor start, int length) -> Tensor(a)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto in = args[0].ITensor();
- auto axis = args[1].unwrapToInt();
- torch::Tensor start = args[2].IValue()->toTensor().to(torch::kI32);
- int32_t startIdx = start.item().to<int32_t>();
- auto length = (int32_t)args[3].unwrapToInt();
-
- // index to access needs to be an at::Tensor
- at::Tensor indices = torch::arange(startIdx, startIdx + length, 1).to(torch::kI32);
- auto weights = Weights(ctx, indices);
-
- // IConstantLayer to convert indices from Weights to ITensor
- auto const_layer = ctx->net->addConstant(weights.shape, weights.data);
- TRTORCH_CHECK(const_layer, "Unable to create constant layer from node: " << *n);
- auto const_out = const_layer->getOutput(0);
-
- // IGatherLayer takes in input tensor, the indices, and the axis
- // of input tensor to take indices from
- auto gather_layer = ctx->net->addGather(*in, *const_out, axis);
- TRTORCH_CHECK(gather_layer, "Unable to create gather layer from node: " << *n);
- auto gather_out = gather_layer->getOutput(0);
-
- // IShuffleLayer removes redundant dimensions
- auto shuffle_layer = ctx->net->addShuffle(*gather_out);
- TRTORCH_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
- shuffle_layer->setReshapeDimensions(util::unpadDims(gather_out->getDimensions()));
- shuffle_layer->setName(util::node_info(n).c_str());
- auto shuffle_out = shuffle_layer->getOutput(0);
-
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle_out);
-
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
-
- return true;
- }})
+ .pattern(
+ {"aten::select.int(Tensor(a) self, int dim, int index) -> (Tensor(a))",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto in = args[0].ITensor();
+ auto axis = args[1].unwrapToInt();
+ auto ind = (int32_t)args[2].unwrapToInt();
+
+ // index to access needs to be an at::Tensor
+ at::Tensor indices = torch::tensor({ind}).to(torch::kI32);
+ auto weights = Weights(ctx, indices);
+
+ // IConstantLayer to convert indices from Weights to ITensor
+ auto const_layer = ctx->net->addConstant(weights.shape, weights.data);
+ TRTORCH_CHECK(const_layer, "Unable to create constant layer from node: " << *n);
+ auto const_out = const_layer->getOutput(0);
+
+ // IGatherLayer takes in input tensor, the indices, and the axis
+ // of input tensor to take indices from
+ auto gather_layer = ctx->net->addGather(*in, *const_out, axis);
+ TRTORCH_CHECK(gather_layer, "Unable to create gather layer from node: " << *n);
+ auto gather_out = gather_layer->getOutput(0);
+
+ // IShuffleLayer removes redundant dimensions
+ auto shuffle_layer = ctx->net->addShuffle(*gather_out);
+ TRTORCH_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
+ shuffle_layer->setReshapeDimensions(util::unpadDims(gather_out->getDimensions()));
+ shuffle_layer->setName(util::node_info(n).c_str());
+ auto shuffle_out = shuffle_layer->getOutput(0);
+
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle_out);
+
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+
+ return true;
+ }})
+ .pattern(
+ {"aten::narrow(Tensor(a) self, int dim, int start, int length) -> Tensor(a)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto in = args[0].ITensor();
+ auto axis = args[1].unwrapToInt();
+ auto start = (int32_t)args[2].unwrapToInt();
+ auto length = (int32_t)args[3].unwrapToInt();
+
+ // index to access needs to be an at::Tensor
+ at::Tensor indices = torch::arange(start, start + length, 1).to(torch::kI32);
+ auto weights = Weights(ctx, indices);
+
+ // IConstantLayer to convert indices from Weights to ITensor
+ auto const_layer = ctx->net->addConstant(weights.shape, weights.data);
+ TRTORCH_CHECK(const_layer, "Unable to create constant layer from node: " << *n);
+ auto const_out = const_layer->getOutput(0);
+
+ // IGatherLayer takes in input tensor, the indices, and the axis
+ // of input tensor to take indices from
+ auto gather_layer = ctx->net->addGather(*in, *const_out, axis);
+ TRTORCH_CHECK(gather_layer, "Unable to create gather layer from node: " << *n);
+ auto gather_out = gather_layer->getOutput(0);
+
+ // IShuffleLayer removes redundant dimensions
+ auto shuffle_layer = ctx->net->addShuffle(*gather_out);
+ TRTORCH_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
+ shuffle_layer->setReshapeDimensions(util::unpadDims(gather_out->getDimensions()));
+ shuffle_layer->setName(util::node_info(n).c_str());
+ auto shuffle_out = shuffle_layer->getOutput(0);
+
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle_out);
+
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+
+ return true;
+ }})
+ .pattern(
+ {"aten::narrow.Tensor(Tensor(a) self, int dim, Tensor start, int length) -> Tensor(a)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto in = args[0].ITensor();
+ auto axis = args[1].unwrapToInt();
+ torch::Tensor start = args[2].IValue()->toTensor().to(torch::kI32);
+ int32_t startIdx = start.item().to<int32_t>();
+ auto length = (int32_t)args[3].unwrapToInt();
+
+ // index to access needs to be an at::Tensor
+ at::Tensor indices = torch::arange(startIdx, startIdx + length, 1).to(torch::kI32);
+ auto weights = Weights(ctx, indices);
+
+ // IConstantLayer to convert indices from Weights to ITensor
+ auto const_layer = ctx->net->addConstant(weights.shape, weights.data);
+ TRTORCH_CHECK(const_layer, "Unable to create constant layer from node: " << *n);
+ auto const_out = const_layer->getOutput(0);
+
+ // IGatherLayer takes in input tensor, the indices, and the axis
+ // of input tensor to take indices from
+ auto gather_layer = ctx->net->addGather(*in, *const_out, axis);
+ TRTORCH_CHECK(gather_layer, "Unable to create gather layer from node: " << *n);
+ auto gather_out = gather_layer->getOutput(0);
+
+ // IShuffleLayer removes redundant dimensions
+ auto shuffle_layer = ctx->net->addShuffle(*gather_out);
+ TRTORCH_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
+ shuffle_layer->setReshapeDimensions(util::unpadDims(gather_out->getDimensions()));
+ shuffle_layer->setName(util::node_info(n).c_str());
+ auto shuffle_out = shuffle_layer->getOutput(0);
+
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle_out);
+
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+
+ return true;
+ }})
.pattern(
{"aten::embedding(Tensor weight, Tensor indices, int padding_idx=-1, bool scale_grad_by_freq=False, bool sparse=False) -> (Tensor)",
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
@@ -223,24 +226,27 @@ auto select_registrations TRTORCH_UNUSED =
return true;
}})
- .pattern({"aten::split(Tensor self, int[] split_sizes, int dim=0) -> (Tensor[])",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- add_split(ctx, n, args, true);
- LOG_DEBUG("Converted split op into a list of IValues");
- return true;
- }})
- .pattern({"aten::split.Tensor(Tensor(a) self, int split_size, int dim=0) -> (Tensor[])",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- add_split(ctx, n, args, false);
- LOG_DEBUG("Converted split op into a list of IValues");
- return true;
- }})
- .pattern({"aten::split_with_sizes(Tensor(a) self, int[] split_sizes, int dim=0) -> (Tensor[])",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- add_split(ctx, n, args, true);
- LOG_DEBUG("Converted split op into a list of IValues");
- return true;
- }});
+ .pattern(
+ {"aten::split(Tensor self, int[] split_sizes, int dim=0) -> (Tensor[])",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ add_split(ctx, n, args, true);
+ LOG_DEBUG("Converted split op into a list of IValues");
+ return true;
+ }})
+ .pattern(
+ {"aten::split.Tensor(Tensor(a) self, int split_size, int dim=0) -> (Tensor[])",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ add_split(ctx, n, args, false);
+ LOG_DEBUG("Converted split op into a list of IValues");
+ return true;
+ }})
+ .pattern(
+ {"aten::split_with_sizes(Tensor(a) self, int[] split_sizes, int dim=0) -> (Tensor[])",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ add_split(ctx, n, args, true);
+ LOG_DEBUG("Converted split op into a list of IValues");
+ return true;
+ }});
} // namespace
} // namespace impl
diff --git a/workspace/core/conversion/converters/impl/shuffle.cpp b/tmp/changes.txt
index a9669a7..000d4a1 100644
--- a/workspace/core/conversion/converters/impl/shuffle.cpp
+++ b/tmp/changes.txt
@@ -38,95 +38,99 @@ static auto shuffle_registrations TRTORCH_UNUSED =
LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
return true;
}})
- .pattern({"aten::reshape(Tensor self, int[] shape) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto in = args[0].ITensorOrFreeze(ctx);
- auto in_shape = util::toVec(in->getDimensions());
- std::vector<int64_t> new_shape;
- if (ctx->input_is_dynamic) {
- TRTORCH_THROW_ERROR("Resize is currently not support in dynamic input shape compilation");
- } else {
- new_shape = torch::reshape(torch::rand(in_shape), args[1].unwrapToIntList().vec()).sizes().vec();
- }
-
- auto shuffle = ctx->net->addShuffle(*in);
- TRTORCH_CHECK(shuffle, "Unable to create shuffle layer from node: " << *n);
- shuffle->setReshapeDimensions(util::toDims(new_shape));
- shuffle->setName(util::node_info(n).c_str());
-
- auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
-
- return true;
- }})
- .pattern({"aten::view(Tensor(a) self, int[] size) -> (Tensor(a))",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto in = args[0].ITensorOrFreeze(ctx);
- auto in_shape = util::toVec(in->getDimensions());
-
- auto shuffle = ctx->net->addShuffle(*in);
- TRTORCH_CHECK(shuffle, "Unable to create shuffle layer from node: " << *n);
- shuffle->setReshapeDimensions(util::toDims(args[1].unwrapToIntList().vec()));
- shuffle->setName(util::node_info(n).c_str());
-
- auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
-
- return true;
- }})
- .pattern({"aten::permute(Tensor(a) self, int[] dims) -> (Tensor(a))",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto in = args[0].ITensorOrFreeze(ctx);
- auto in_shape = util::toVec(in->getDimensions());
- auto new_order = args[1].unwrapToIntList().vec();
-
- LOG_DEBUG("Shuffle to: " << util::toDims(new_order));
-
- auto shuffle = ctx->net->addShuffle(*in);
- TRTORCH_CHECK(shuffle, "Unable to create shuffle layer from node: " << *n);
- nvinfer1::Permutation permute;
- std::copy(new_order.begin(), new_order.end(), permute.order);
- shuffle->setSecondTranspose(permute);
- shuffle->setName(util::node_info(n).c_str());
-
- auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
-
- return true;
- }})
- .pattern({"aten::transpose.int(Tensor(a) self, int dim0, int dim1) -> (Tensor(a))",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto in = args[0].ITensorOrFreeze(ctx);
- auto in_shape = util::toVec(in->getDimensions());
- auto ndims = in_shape.size();
- auto dim0 = args[1].unwrapToInt();
- auto dim1 = args[2].unwrapToInt();
-
- std::vector<int64_t> new_order;
- for (size_t i = 0; i < ndims; i++) {
- new_order.push_back(i);
- }
- dim0 = dim0 < 0 ? (dim0 + ndims) : dim0;
- dim1 = dim1 < 0 ? (dim1 + ndims) : dim1;
- auto tmp = dim0;
- new_order[dim0] = new_order[dim1];
- new_order[dim1] = tmp;
-
- LOG_DEBUG("Shuffle to: " << util::toDims(new_order));
-
- auto shuffle = ctx->net->addShuffle(*in);
- TRTORCH_CHECK(shuffle, "Unable to create shuffle layer from node: " << *n);
- nvinfer1::Permutation permute;
- std::copy(new_order.begin(), new_order.end(), permute.order);
-
- shuffle->setSecondTranspose(permute);
- shuffle->setName(util::node_info(n).c_str());
-
- auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
-
- return true;
- }});
+ .pattern(
+ {"aten::reshape(Tensor self, int[] shape) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto in = args[0].ITensorOrFreeze(ctx);
+ auto in_shape = util::toVec(in->getDimensions());
+ std::vector<int64_t> new_shape;
+ if (ctx->input_is_dynamic) {
+ TRTORCH_THROW_ERROR("Resize is currently not support in dynamic input shape compilation");
+ } else {
+ new_shape = torch::reshape(torch::rand(in_shape), args[1].unwrapToIntList().vec()).sizes().vec();
+ }
+
+ auto shuffle = ctx->net->addShuffle(*in);
+ TRTORCH_CHECK(shuffle, "Unable to create shuffle layer from node: " << *n);
+ shuffle->setReshapeDimensions(util::toDims(new_shape));
+ shuffle->setName(util::node_info(n).c_str());
+
+ auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
+
+ return true;
+ }})
+ .pattern(
+ {"aten::view(Tensor(a) self, int[] size) -> (Tensor(a))",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto in = args[0].ITensorOrFreeze(ctx);
+ auto in_shape = util::toVec(in->getDimensions());
+
+ auto shuffle = ctx->net->addShuffle(*in);
+ TRTORCH_CHECK(shuffle, "Unable to create shuffle layer from node: " << *n);
+ shuffle->setReshapeDimensions(util::toDims(args[1].unwrapToIntList().vec()));
+ shuffle->setName(util::node_info(n).c_str());
+
+ auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
+
+ return true;
+ }})
+ .pattern(
+ {"aten::permute(Tensor(a) self, int[] dims) -> (Tensor(a))",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto in = args[0].ITensorOrFreeze(ctx);
+ auto in_shape = util::toVec(in->getDimensions());
+ auto new_order = args[1].unwrapToIntList().vec();
+
+ LOG_DEBUG("Shuffle to: " << util::toDims(new_order));
+
+ auto shuffle = ctx->net->addShuffle(*in);
+ TRTORCH_CHECK(shuffle, "Unable to create shuffle layer from node: " << *n);
+ nvinfer1::Permutation permute;
+ std::copy(new_order.begin(), new_order.end(), permute.order);
+ shuffle->setSecondTranspose(permute);
+ shuffle->setName(util::node_info(n).c_str());
+
+ auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
+
+ return true;
+ }})
+ .pattern(
+ {"aten::transpose.int(Tensor(a) self, int dim0, int dim1) -> (Tensor(a))",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto in = args[0].ITensorOrFreeze(ctx);
+ auto in_shape = util::toVec(in->getDimensions());
+ auto ndims = in_shape.size();
+ auto dim0 = args[1].unwrapToInt();
+ auto dim1 = args[2].unwrapToInt();
+
+ std::vector<int64_t> new_order;
+ for (size_t i = 0; i < ndims; i++) {
+ new_order.push_back(i);
+ }
+ dim0 = dim0 < 0 ? (dim0 + ndims) : dim0;
+ dim1 = dim1 < 0 ? (dim1 + ndims) : dim1;
+ auto tmp = dim0;
+ new_order[dim0] = new_order[dim1];
+ new_order[dim1] = tmp;
+
+ LOG_DEBUG("Shuffle to: " << util::toDims(new_order));
+
+ auto shuffle = ctx->net->addShuffle(*in);
+ TRTORCH_CHECK(shuffle, "Unable to create shuffle layer from node: " << *n);
+ nvinfer1::Permutation permute;
+ std::copy(new_order.begin(), new_order.end(), permute.order);
+
+ shuffle->setSecondTranspose(permute);
+ shuffle->setName(util::node_info(n).c_str());
+
+ auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
+
+ return true;
+ }});
} // namespace
} // namespace impl
diff --git a/workspace/core/conversion/converters/impl/element_wise.cpp b/tmp/changes.txt
index b95d712..ec5f0d8 100644
--- a/workspace/core/conversion/converters/impl/element_wise.cpp
+++ b/tmp/changes.txt
@@ -85,677 +85,692 @@ nvinfer1::ITensor* clamp_util(
auto element_wise_registrations TRTORCH_UNUSED =
RegisterNodeConversionPatterns()
- .pattern({"aten::add.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> "
- "Tensor",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- // Should implement self + alpha * other
- auto self = args[0].ITensorOrFreeze(ctx);
- auto other = args[1].ITensorOrFreeze(ctx);
- auto scalar = args[2].unwrapToScalar().to<float>();
-
- if (1 != scalar) {
- auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
- auto scaleLayer = add_elementwise(
- ctx,
- nvinfer1::ElementWiseOperation::kPROD,
- other,
- alphaTensor,
- util::node_info(n) + std::string("_AlphaMultiplier"));
- TRTORCH_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n);
- other = scaleLayer->getOutput(0);
- }
-
- auto add =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUM, self, other, util::node_info(n));
- TRTORCH_CHECK(add, "Unable to create add layer from node: " << *n);
-
- add->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], add->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::add_.Tensor(Tensor(a!) self, Tensor other, *, Scalar "
- "alpha=1) -> (Tensor(a!))",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- // Should implement self + alpha * other
- auto self = args[0].ITensorOrFreeze(ctx);
- auto other = args[1].ITensorOrFreeze(ctx);
- auto scalar = args[2].unwrapToScalar().to<float>();
-
- if (1 != scalar) {
- auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
- auto scaleLayer = add_elementwise(
- ctx,
- nvinfer1::ElementWiseOperation::kPROD,
- other,
- alphaTensor,
- util::node_info(n) + std::string("_AlphaMultiplier"));
- TRTORCH_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n);
- other = scaleLayer->getOutput(0);
- }
-
- auto add =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUM, self, other, util::node_info(n));
- TRTORCH_CHECK(add, "Unable to create add layer from node: " << *n);
-
- add->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], add->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::add.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- // Should implement self + alpha * other
- auto self = args[0].ITensorOrFreeze(ctx);
- auto otherScalar = args[2].unwrapToScalar().to<float>() * args[1].unwrapToScalar().to<float>();
- auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
-
- auto add =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUM, self, other, util::node_info(n));
- TRTORCH_CHECK(add, "Unable to create add layer from node: " << *n);
-
- add->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], add->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::clamp(Tensor self, Scalar? min=None, Scalar? max=None) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- // Compute min(max(min_threshold, input), max_threshold)
- auto self = args[0].ITensorOrFreeze(ctx);
- auto clamp_layer_out = self;
-
- if (args[1].isIValue() && args[1].IValue()->isScalar() && args[2].isIValue() &&
- args[2].IValue()->isScalar()) {
- auto alpha = args[1].unwrapToScalar().to<float>();
- auto beta = args[2].unwrapToScalar().to<float>();
- auto clip_layer = ctx->net->addActivation(*self, nvinfer1::ActivationType::kCLIP);
- TRTORCH_CHECK(clip_layer, "Unable to create clip layer for node: " << *n);
- clip_layer->setAlpha(alpha);
- clip_layer->setBeta(beta);
- clamp_layer_out = clip_layer->getOutput(0);
- } else if (args[1].isIValue() && args[1].IValue()->isScalar()) {
- auto limit = args[1].unwrapToScalar().to<float>();
- clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMAX, "_max");
- } else if (args[2].isIValue() && args[2].IValue()->isScalar()) {
- auto limit = args[2].unwrapToScalar().to<float>();
- clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMIN, "_min");
- }
-
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out);
- LOG_DEBUG("Clamp layer output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::clamp_min(Tensor self, Scalar min) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- // Compute min(max(min_threshold, input), max_threshold)
- auto self = args[0].ITensorOrFreeze(ctx);
- auto clamp_layer_out = self;
- if (args[1].isIValue() && args[1].IValue()->isScalar()) {
- auto limit = args[1].unwrapToScalar().to<float>();
- clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMAX, "_max");
- }
-
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out);
- LOG_DEBUG("clamp_min layer output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::clamp_max(Tensor self, Scalar max) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- // Compute min(max(min_threshold, input), max_threshold)
- auto self = args[0].ITensorOrFreeze(ctx);
- auto clamp_layer_out = self;
- if (args[1].isIValue() && args[1].IValue()->isScalar()) {
- auto limit = args[1].unwrapToScalar().to<float>();
- clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMIN, "_min");
- }
-
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out);
- LOG_DEBUG("clamp_max layer output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::sub.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> "
- "Tensor",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- // Should implement self - alpha * other
- auto self = args[0].ITensorOrFreeze(ctx);
- auto scalar = args[2].unwrapToScalar().to<float>();
- auto other = args[1].ITensorOrFreeze(ctx);
-
- if (1 != scalar) {
- auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
- auto scaleLayer = add_elementwise(
- ctx,
- nvinfer1::ElementWiseOperation::kPROD,
- other,
- alphaTensor,
- util::node_info(n) + std::string("_AlphaMultiplier"));
- TRTORCH_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n);
- other = scaleLayer->getOutput(0);
- }
-
- auto sub =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, self, other, util::node_info(n));
- TRTORCH_CHECK(sub, "Unable to create sub layer from node: " << *n);
-
- sub->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::sub_.Tensor(Tensor(a!) self, Tensor other, *, Scalar "
- "alpha=1) -> (Tensor(a!))",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- // Should implement self - alpha * other
- auto self = args[0].ITensorOrFreeze(ctx);
- auto scalar = args[2].unwrapToScalar().to<float>();
- auto other = args[1].ITensorOrFreeze(ctx);
-
- if (1 != scalar) {
- auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
- auto scaleLayer = add_elementwise(
- ctx,
- nvinfer1::ElementWiseOperation::kPROD,
- other,
- alphaTensor,
- util::node_info(n) + std::string("_AlphaMultiplier"));
- TRTORCH_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n);
- other = scaleLayer->getOutput(0);
- }
-
- auto sub =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, self, other, util::node_info(n));
- TRTORCH_CHECK(sub, "Unable to create sub layer from node: " << *n);
-
- sub->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::rsub.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- // Should implement other - alpha * self
- auto self = args[0].ITensorOrFreeze(ctx);
- auto otherScalar = args[1].unwrapToScalar().to<float>();
- auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
- auto scalar = args[2].unwrapToScalar().to<float>();
-
- if (1 != scalar) {
- auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
- auto scaleLayer = add_elementwise(
- ctx,
- nvinfer1::ElementWiseOperation::kPROD,
- self,
- alphaTensor,
- util::node_info(n) + std::string("_AlphaMultiplier"));
- TRTORCH_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n);
- self = scaleLayer->getOutput(0);
- }
-
- auto rsub =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, other, self, util::node_info(n));
- TRTORCH_CHECK(rsub, "Unable to create rsub layer from node: " << *n);
-
- rsub->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], rsub->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::rsub.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- // Should implement other - alpha * self
- auto self = args[0].ITensorOrFreeze(ctx);
- auto other = args[1].ITensorOrFreeze(ctx);
- auto scalar = args[2].unwrapToScalar().to<float>();
-
- if (1 != scalar) {
- auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
- auto scaleLayer = add_elementwise(
- ctx,
- nvinfer1::ElementWiseOperation::kPROD,
- self,
- alphaTensor,
- util::node_info(n) + std::string("_AlphaMultiplier"));
- TRTORCH_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n);
- self = scaleLayer->getOutput(0);
- }
-
- auto rsub =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, other, self, util::node_info(n));
- TRTORCH_CHECK(rsub, "Unable to create rsub layer from node: " << *n);
-
- rsub->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], rsub->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::div.Tensor(Tensor self, Tensor other) -> Tensor",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- // Should implement self / other
- auto self = args[0].ITensorOrFreeze(ctx);
- auto other = args[1].ITensorOrFreeze(ctx);
- auto div =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n));
-
- TRTORCH_CHECK(div, "Unable to create div layer from node: " << *n);
-
- div->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0));
-
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .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>();
- auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
- auto div =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n));
- TRTORCH_CHECK(div, "Unable to create div layer from node: " << *n);
-
- div->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::div_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- // TODO: Remove with functionalization
- auto self = args[0].ITensorOrFreeze(ctx);
- auto other = args[1].ITensorOrFreeze(ctx);
- auto div =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n));
-
- TRTORCH_CHECK(div, "Unable to create div layer from node: " << *n);
-
- div->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0));
-
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::div_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto self = args[0].ITensorOrFreeze(ctx);
- auto otherScalar = args[1].unwrapToScalar().to<float>();
- auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
- auto div =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n));
- TRTORCH_CHECK(div, "Unable to create div layer from node: " << *n);
-
- div->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::mul.Tensor(Tensor self, Tensor other) -> Tensor",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- // Should implement self * other
- auto self = args[0].ITensorOrFreeze(ctx);
- auto other = args[1].ITensorOrFreeze(ctx);
- auto mul =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, other, util::node_info(n));
- TRTORCH_CHECK(mul, "Unable to create mul layer from node: " << *n);
-
- mul->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], mul->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::mul.Scalar(Tensor self, Scalar other) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- // TODO: Remove with functionalization
- auto self = args[0].ITensorOrFreeze(ctx);
- auto otherScalar = args[1].unwrapToScalar().to<float>();
- auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
- auto mul =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, other, util::node_info(n));
- TRTORCH_CHECK(mul, "Unable to create mul layer from node: " << *n);
-
- mul->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], mul->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::mul_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- // TODO: Remove with functionalization
- auto self = args[0].ITensorOrFreeze(ctx);
- auto other = args[1].ITensorOrFreeze(ctx);
- auto mul =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, other, util::node_info(n));
- TRTORCH_CHECK(mul, "Unable to create mul layer from node: " << *n);
-
- mul->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], mul->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::ne.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 equal = add_elementwise(
- ctx,
- nvinfer1::ElementWiseOperation::kEQUAL,
- self,
- other,
- util::node_info(n) + std::string("is_equal"));
- TRTORCH_CHECK(equal, "Unable to create elementwise equal layer from node: " << *n);
- // XOR with ones negates and produces not_equal result
- auto options = torch::TensorOptions().dtype(torch::kFloat32);
- auto ones = at::full({1}, 1, {options});
- auto ones_tensor = tensor_to_const(ctx, ones);
- nvinfer1::IIdentityLayer* cast_layer = ctx->net->addIdentity(*ones_tensor);
- cast_layer->setOutputType(0, nvinfer1::DataType::kBOOL);
-
- auto sub = add_elementwise(
- ctx,
- nvinfer1::ElementWiseOperation::kXOR,
- cast_layer->getOutput(0),
- equal->getOutput(0),
- util::node_info(n));
- TRTORCH_CHECK(sub, "Unable to create ne (not equal) layer from node: " << *n);
-
- sub->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0));
- LOG_DEBUG("Not equal layer output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::ne.Scalar(Tensor self, Scalar other) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto self = args[0].ITensorOrFreeze(ctx);
- auto scalar = args[1].unwrapToScalar().to<float>();
- auto scalar_tensor = tensor_to_const(ctx, torch::tensor({scalar}));
- auto equal = add_elementwise(
- ctx,
- nvinfer1::ElementWiseOperation::kEQUAL,
- self,
- scalar_tensor,
- util::node_info(n) + std::string("is_equal"));
- TRTORCH_CHECK(equal, "Unable to create elementwise equal layer from node: " << *n);
- // XOR with ones negates and produces not_equal result
- auto options = torch::TensorOptions().dtype(torch::kFloat32);
- auto ones = at::full({1}, 1, {options});
- auto ones_tensor = tensor_to_const(ctx, ones);
- nvinfer1::IIdentityLayer* cast_layer = ctx->net->addIdentity(*ones_tensor);
- cast_layer->setOutputType(0, nvinfer1::DataType::kBOOL);
-
- auto sub = add_elementwise(
- ctx,
- nvinfer1::ElementWiseOperation::kXOR,
- cast_layer->getOutput(0),
- equal->getOutput(0),
- util::node_info(n));
- TRTORCH_CHECK(sub, "Unable to create ne (not equal) layer from node: " << *n);
-
- sub->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0));
- LOG_DEBUG("Not equal layer output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::pow.Tensor_Tensor(Tensor self, Tensor exponent) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto self = args[0].ITensorOrFreeze(ctx);
- auto exponent = args[1].ITensorOrFreeze(ctx);
- auto pow =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPOW, self, exponent, util::node_info(n));
- TRTORCH_CHECK(pow, "Unable to create Power layer from node: " << *n);
-
- pow->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], pow->getOutput(0));
-
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::pow.Tensor_Scalar(Tensor self, Scalar exponent) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto self = args[0].ITensorOrFreeze(ctx);
- auto exponentScalar = args[1].unwrapToScalar().to<float>();
- auto exponent = tensor_to_const(ctx, torch::tensor({exponentScalar}));
- auto pow =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPOW, self, exponent, util::node_info(n));
- TRTORCH_CHECK(pow, "Unable to create Power layer from node: " << *n);
-
- pow->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], pow->getOutput(0));
-
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::floor_divide(Tensor self, Tensor other) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- // TODO: Remove with functionalization
- auto self = args[0].ITensorOrFreeze(ctx);
- auto other = args[1].ITensorOrFreeze(ctx);
- auto floor_divide = add_elementwise(
- ctx, nvinfer1::ElementWiseOperation::kFLOOR_DIV, self, other, util::node_info(n));
- TRTORCH_CHECK(floor_divide, "Unable to create floor_divide layer from node: " << *n);
-
- floor_divide->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], floor_divide->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::floor_divide.Scalar(Tensor self, Scalar other) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- // TODO: Remove with functionalization
- auto self = args[0].ITensorOrFreeze(ctx);
- auto otherScalar = args[1].unwrapToScalar().to<float>();
- auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
- auto floor_divide = add_elementwise(
- ctx, nvinfer1::ElementWiseOperation::kFLOOR_DIV, self, other, util::node_info(n));
- TRTORCH_CHECK(floor_divide, "Unable to create floor_divide layer from node: " << *n);
-
- floor_divide->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], floor_divide->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::max.other(Tensor self, Tensor other) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- // TODO: Remove with functionalization
- auto self = args[0].ITensorOrFreeze(ctx);
- auto other = args[1].ITensorOrFreeze(ctx);
- auto max =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kMAX, self, other, util::node_info(n));
- TRTORCH_CHECK(max, "Unable to create max layer from node: " << *n);
-
- max->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], max->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::min.other(Tensor self, Tensor other) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- // TODO: Remove with functionalization
- auto self = args[0].ITensorOrFreeze(ctx);
- auto other = args[1].ITensorOrFreeze(ctx);
- auto min =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kMIN, self, other, util::node_info(n));
- TRTORCH_CHECK(min, "Unable to create min layer from node: " << *n);
-
- min->setName(util::node_info(n).c_str());
- auto out = ctx->AssociateValueAndTensor(n->outputs()[0], min->getOutput(0));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::gt.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 gt =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n));
- 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));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::gt.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>();
- auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
- auto gt =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n));
- 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));
- 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));
- 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));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::lt.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>();
- auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
- auto lt =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n));
- 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));
- 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));
- 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));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::eq.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>();
- auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
- auto eq =
- add_elementwise(ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n));
- 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));
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::ge.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 greater = add_elementwise(
- ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n) + "_greater");
- TRTORCH_CHECK(greater, "Unable to create Greater layer from node: " << *n);
-
- auto equal = add_elementwise(
- ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal");
- TRTORCH_CHECK(equal, "Unable to create Equal layer from node: " << *n);
-
- auto or_op = ctx->net->addElementWise(
- *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));
-
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::ge.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>();
- auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
-
- auto greater = add_elementwise(
- ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n) + "_greater");
- TRTORCH_CHECK(greater, "Unable to create Greater layer from node: " << *n);
-
- auto equal = add_elementwise(
- ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal");
- TRTORCH_CHECK(equal, "Unable to create Equal layer from node: " << *n);
-
- auto or_op = ctx->net->addElementWise(
- *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));
-
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::le.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 less = add_elementwise(
- ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n) + "_less");
- TRTORCH_CHECK(less, "Unable to create Less layer from node: " << *n);
-
- auto equal = add_elementwise(
- ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal");
- TRTORCH_CHECK(equal, "Unable to create Equal layer from node: " << *n);
-
- auto or_op = ctx->net->addElementWise(
- *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));
-
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }})
- .pattern({"aten::le.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>();
- auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
-
- auto less = add_elementwise(
- ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n) + "_less");
- TRTORCH_CHECK(less, "Unable to create Less layer from node: " << *n);
-
- auto equal = add_elementwise(
- ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal");
- TRTORCH_CHECK(equal, "Unable to create Equal layer from node: " << *n);
-
- auto or_op = ctx->net->addElementWise(
- *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));
-
- LOG_DEBUG("Output tensor shape: " << out->getDimensions());
- return true;
- }});
+ .pattern(
+ {"aten::add.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> "
+ "Tensor",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ // Should implement self + alpha * other
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto other = args[1].ITensorOrFreeze(ctx);
+ auto scalar = args[2].unwrapToScalar().to<float>();
+
+ if (1 != scalar) {
+ auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
+ auto scaleLayer = add_elementwise(
+ ctx,
+ nvinfer1::ElementWiseOperation::kPROD,
+ other,
+ alphaTensor,
+ util::node_info(n) + std::string("_AlphaMultiplier"));
+ TRTORCH_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n);
+ other = scaleLayer->getOutput(0);
+ }
+
+ auto add = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUM, self, other, util::node_info(n));
+ TRTORCH_CHECK(add, "Unable to create add layer from node: " << *n);
+
+ add->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], add->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::add_.Tensor(Tensor(a!) self, Tensor other, *, Scalar "
+ "alpha=1) -> (Tensor(a!))",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ // Should implement self + alpha * other
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto other = args[1].ITensorOrFreeze(ctx);
+ auto scalar = args[2].unwrapToScalar().to<float>();
+
+ if (1 != scalar) {
+ auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
+ auto scaleLayer = add_elementwise(
+ ctx,
+ nvinfer1::ElementWiseOperation::kPROD,
+ other,
+ alphaTensor,
+ util::node_info(n) + std::string("_AlphaMultiplier"));
+ TRTORCH_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n);
+ other = scaleLayer->getOutput(0);
+ }
+
+ auto add = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUM, self, other, util::node_info(n));
+ TRTORCH_CHECK(add, "Unable to create add layer from node: " << *n);
+
+ add->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], add->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::add.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ // Should implement self + alpha * other
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto otherScalar = args[2].unwrapToScalar().to<float>() * args[1].unwrapToScalar().to<float>();
+ auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+
+ auto add = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUM, self, other, util::node_info(n));
+ TRTORCH_CHECK(add, "Unable to create add layer from node: " << *n);
+
+ add->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], add->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::clamp(Tensor self, Scalar? min=None, Scalar? max=None) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ // Compute min(max(min_threshold, input), max_threshold)
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto clamp_layer_out = self;
+
+ if (args[1].isIValue() && args[1].IValue()->isScalar() && args[2].isIValue() &&
+ args[2].IValue()->isScalar()) {
+ auto alpha = args[1].unwrapToScalar().to<float>();
+ auto beta = args[2].unwrapToScalar().to<float>();
+ auto clip_layer = ctx->net->addActivation(*self, nvinfer1::ActivationType::kCLIP);
+ TRTORCH_CHECK(clip_layer, "Unable to create clip layer for node: " << *n);
+ clip_layer->setAlpha(alpha);
+ clip_layer->setBeta(beta);
+ clamp_layer_out = clip_layer->getOutput(0);
+ } else if (args[1].isIValue() && args[1].IValue()->isScalar()) {
+ auto limit = args[1].unwrapToScalar().to<float>();
+ clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMAX, "_max");
+ } else if (args[2].isIValue() && args[2].IValue()->isScalar()) {
+ auto limit = args[2].unwrapToScalar().to<float>();
+ clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMIN, "_min");
+ }
+
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out);
+ LOG_DEBUG("Clamp layer output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::clamp_min(Tensor self, Scalar min) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ // Compute min(max(min_threshold, input), max_threshold)
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto clamp_layer_out = self;
+ if (args[1].isIValue() && args[1].IValue()->isScalar()) {
+ auto limit = args[1].unwrapToScalar().to<float>();
+ clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMAX, "_max");
+ }
+
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out);
+ LOG_DEBUG("clamp_min layer output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::clamp_max(Tensor self, Scalar max) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ // Compute min(max(min_threshold, input), max_threshold)
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto clamp_layer_out = self;
+ if (args[1].isIValue() && args[1].IValue()->isScalar()) {
+ auto limit = args[1].unwrapToScalar().to<float>();
+ clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMIN, "_min");
+ }
+
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out);
+ LOG_DEBUG("clamp_max layer output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::sub.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> "
+ "Tensor",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ // Should implement self - alpha * other
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto scalar = args[2].unwrapToScalar().to<float>();
+ auto other = args[1].ITensorOrFreeze(ctx);
+
+ if (1 != scalar) {
+ auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
+ auto scaleLayer = add_elementwise(
+ ctx,
+ nvinfer1::ElementWiseOperation::kPROD,
+ other,
+ alphaTensor,
+ util::node_info(n) + std::string("_AlphaMultiplier"));
+ TRTORCH_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n);
+ other = scaleLayer->getOutput(0);
+ }
+
+ auto sub = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, self, other, util::node_info(n));
+ TRTORCH_CHECK(sub, "Unable to create sub layer from node: " << *n);
+
+ sub->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::sub_.Tensor(Tensor(a!) self, Tensor other, *, Scalar "
+ "alpha=1) -> (Tensor(a!))",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ // Should implement self - alpha * other
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto scalar = args[2].unwrapToScalar().to<float>();
+ auto other = args[1].ITensorOrFreeze(ctx);
+
+ if (1 != scalar) {
+ auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
+ auto scaleLayer = add_elementwise(
+ ctx,
+ nvinfer1::ElementWiseOperation::kPROD,
+ other,
+ alphaTensor,
+ util::node_info(n) + std::string("_AlphaMultiplier"));
+ TRTORCH_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n);
+ other = scaleLayer->getOutput(0);
+ }
+
+ auto sub = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, self, other, util::node_info(n));
+ TRTORCH_CHECK(sub, "Unable to create sub layer from node: " << *n);
+
+ sub->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::rsub.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ // Should implement other - alpha * self
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto otherScalar = args[1].unwrapToScalar().to<float>();
+ auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+ auto scalar = args[2].unwrapToScalar().to<float>();
+
+ if (1 != scalar) {
+ auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
+ auto scaleLayer = add_elementwise(
+ ctx,
+ nvinfer1::ElementWiseOperation::kPROD,
+ self,
+ alphaTensor,
+ util::node_info(n) + std::string("_AlphaMultiplier"));
+ TRTORCH_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n);
+ self = scaleLayer->getOutput(0);
+ }
+
+ auto rsub = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, other, self, util::node_info(n));
+ TRTORCH_CHECK(rsub, "Unable to create rsub layer from node: " << *n);
+
+ rsub->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], rsub->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::rsub.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ // Should implement other - alpha * self
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto other = args[1].ITensorOrFreeze(ctx);
+ auto scalar = args[2].unwrapToScalar().to<float>();
+
+ if (1 != scalar) {
+ auto alphaTensor = tensor_to_const(ctx, torch::tensor({scalar}));
+ auto scaleLayer = add_elementwise(
+ ctx,
+ nvinfer1::ElementWiseOperation::kPROD,
+ self,
+ alphaTensor,
+ util::node_info(n) + std::string("_AlphaMultiplier"));
+ TRTORCH_CHECK(scaleLayer, "Unable to create alpha*input layer from node: " << *n);
+ self = scaleLayer->getOutput(0);
+ }
+
+ auto rsub = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kSUB, other, self, util::node_info(n));
+ TRTORCH_CHECK(rsub, "Unable to create rsub layer from node: " << *n);
+
+ rsub->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], rsub->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::div.Tensor(Tensor self, Tensor other) -> Tensor",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ // Should implement self / other
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto other = args[1].ITensorOrFreeze(ctx);
+ auto div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n));
+
+ TRTORCH_CHECK(div, "Unable to create div layer from node: " << *n);
+
+ div->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0));
+
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .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>();
+ auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+ auto div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n));
+ TRTORCH_CHECK(div, "Unable to create div layer from node: " << *n);
+
+ div->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::div_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ // TODO: Remove with functionalization
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto other = args[1].ITensorOrFreeze(ctx);
+ auto div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n));
+
+ TRTORCH_CHECK(div, "Unable to create div layer from node: " << *n);
+
+ div->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0));
+
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::div_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto otherScalar = args[1].unwrapToScalar().to<float>();
+ auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+ auto div = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kDIV, self, other, util::node_info(n));
+ TRTORCH_CHECK(div, "Unable to create div layer from node: " << *n);
+
+ div->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], div->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::mul.Tensor(Tensor self, Tensor other) -> Tensor",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ // Should implement self * other
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto other = args[1].ITensorOrFreeze(ctx);
+ auto mul = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, other, util::node_info(n));
+ TRTORCH_CHECK(mul, "Unable to create mul layer from node: " << *n);
+
+ mul->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], mul->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::mul.Scalar(Tensor self, Scalar other) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ // TODO: Remove with functionalization
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto otherScalar = args[1].unwrapToScalar().to<float>();
+ auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+ auto mul = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, other, util::node_info(n));
+ TRTORCH_CHECK(mul, "Unable to create mul layer from node: " << *n);
+
+ mul->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], mul->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::mul_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ // TODO: Remove with functionalization
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto other = args[1].ITensorOrFreeze(ctx);
+ auto mul = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPROD, self, other, util::node_info(n));
+ TRTORCH_CHECK(mul, "Unable to create mul layer from node: " << *n);
+
+ mul->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], mul->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::ne.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 equal = add_elementwise(
+ ctx,
+ nvinfer1::ElementWiseOperation::kEQUAL,
+ self,
+ other,
+ util::node_info(n) + std::string("is_equal"));
+ TRTORCH_CHECK(equal, "Unable to create elementwise equal layer from node: " << *n);
+ // XOR with ones negates and produces not_equal result
+ auto options = torch::TensorOptions().dtype(torch::kFloat32);
+ auto ones = at::full({1}, 1, {options});
+ auto ones_tensor = tensor_to_const(ctx, ones);
+ nvinfer1::IIdentityLayer* cast_layer = ctx->net->addIdentity(*ones_tensor);
+ cast_layer->setOutputType(0, nvinfer1::DataType::kBOOL);
+
+ auto sub = add_elementwise(
+ ctx,
+ nvinfer1::ElementWiseOperation::kXOR,
+ cast_layer->getOutput(0),
+ equal->getOutput(0),
+ util::node_info(n));
+ TRTORCH_CHECK(sub, "Unable to create ne (not equal) layer from node: " << *n);
+
+ sub->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0));
+ LOG_DEBUG("Not equal layer output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::ne.Scalar(Tensor self, Scalar other) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto scalar = args[1].unwrapToScalar().to<float>();
+ auto scalar_tensor = tensor_to_const(ctx, torch::tensor({scalar}));
+ auto equal = add_elementwise(
+ ctx,
+ nvinfer1::ElementWiseOperation::kEQUAL,
+ self,
+ scalar_tensor,
+ util::node_info(n) + std::string("is_equal"));
+ TRTORCH_CHECK(equal, "Unable to create elementwise equal layer from node: " << *n);
+ // XOR with ones negates and produces not_equal result
+ auto options = torch::TensorOptions().dtype(torch::kFloat32);
+ auto ones = at::full({1}, 1, {options});
+ auto ones_tensor = tensor_to_const(ctx, ones);
+ nvinfer1::IIdentityLayer* cast_layer = ctx->net->addIdentity(*ones_tensor);
+ cast_layer->setOutputType(0, nvinfer1::DataType::kBOOL);
+
+ auto sub = add_elementwise(
+ ctx,
+ nvinfer1::ElementWiseOperation::kXOR,
+ cast_layer->getOutput(0),
+ equal->getOutput(0),
+ util::node_info(n));
+ TRTORCH_CHECK(sub, "Unable to create ne (not equal) layer from node: " << *n);
+
+ sub->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], sub->getOutput(0));
+ LOG_DEBUG("Not equal layer output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::pow.Tensor_Tensor(Tensor self, Tensor exponent) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto exponent = args[1].ITensorOrFreeze(ctx);
+ auto pow =
+ add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPOW, self, exponent, util::node_info(n));
+ TRTORCH_CHECK(pow, "Unable to create Power layer from node: " << *n);
+
+ pow->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], pow->getOutput(0));
+
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::pow.Tensor_Scalar(Tensor self, Scalar exponent) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto exponentScalar = args[1].unwrapToScalar().to<float>();
+ auto exponent = tensor_to_const(ctx, torch::tensor({exponentScalar}));
+ auto pow =
+ add_elementwise(ctx, nvinfer1::ElementWiseOperation::kPOW, self, exponent, util::node_info(n));
+ TRTORCH_CHECK(pow, "Unable to create Power layer from node: " << *n);
+
+ pow->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], pow->getOutput(0));
+
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::floor_divide(Tensor self, Tensor other) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ // TODO: Remove with functionalization
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto other = args[1].ITensorOrFreeze(ctx);
+ auto floor_divide =
+ add_elementwise(ctx, nvinfer1::ElementWiseOperation::kFLOOR_DIV, self, other, util::node_info(n));
+ TRTORCH_CHECK(floor_divide, "Unable to create floor_divide layer from node: " << *n);
+
+ floor_divide->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], floor_divide->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::floor_divide.Scalar(Tensor self, Scalar other) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ // TODO: Remove with functionalization
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto otherScalar = args[1].unwrapToScalar().to<float>();
+ auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+ auto floor_divide =
+ add_elementwise(ctx, nvinfer1::ElementWiseOperation::kFLOOR_DIV, self, other, util::node_info(n));
+ TRTORCH_CHECK(floor_divide, "Unable to create floor_divide layer from node: " << *n);
+
+ floor_divide->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], floor_divide->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::max.other(Tensor self, Tensor other) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ // TODO: Remove with functionalization
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto other = args[1].ITensorOrFreeze(ctx);
+ auto max = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kMAX, self, other, util::node_info(n));
+ TRTORCH_CHECK(max, "Unable to create max layer from node: " << *n);
+
+ max->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], max->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::min.other(Tensor self, Tensor other) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ // TODO: Remove with functionalization
+ auto self = args[0].ITensorOrFreeze(ctx);
+ auto other = args[1].ITensorOrFreeze(ctx);
+ auto min = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kMIN, self, other, util::node_info(n));
+ TRTORCH_CHECK(min, "Unable to create min layer from node: " << *n);
+
+ min->setName(util::node_info(n).c_str());
+ auto out = ctx->AssociateValueAndTensor(n->outputs()[0], min->getOutput(0));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::gt.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 gt =
+ add_elementwise(ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n));
+ 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));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::gt.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>();
+ auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+ auto gt =
+ add_elementwise(ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n));
+ 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));
+ 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));
+ 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));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::lt.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>();
+ auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+ auto lt = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n));
+ 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));
+ 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));
+ 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));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::eq.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>();
+ auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+ auto eq = add_elementwise(ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n));
+ 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));
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::ge.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 greater = add_elementwise(
+ ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n) + "_greater");
+ TRTORCH_CHECK(greater, "Unable to create Greater layer from node: " << *n);
+
+ auto equal = add_elementwise(
+ ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal");
+ TRTORCH_CHECK(equal, "Unable to create Equal layer from node: " << *n);
+
+ auto or_op = ctx->net->addElementWise(
+ *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));
+
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::ge.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>();
+ auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+
+ auto greater = add_elementwise(
+ ctx, nvinfer1::ElementWiseOperation::kGREATER, self, other, util::node_info(n) + "_greater");
+ TRTORCH_CHECK(greater, "Unable to create Greater layer from node: " << *n);
+
+ auto equal = add_elementwise(
+ ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal");
+ TRTORCH_CHECK(equal, "Unable to create Equal layer from node: " << *n);
+
+ auto or_op = ctx->net->addElementWise(
+ *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));
+
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::le.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 less = add_elementwise(
+ ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n) + "_less");
+ TRTORCH_CHECK(less, "Unable to create Less layer from node: " << *n);
+
+ auto equal = add_elementwise(
+ ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal");
+ TRTORCH_CHECK(equal, "Unable to create Equal layer from node: " << *n);
+
+ auto or_op = ctx->net->addElementWise(
+ *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));
+
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::le.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>();
+ auto other = tensor_to_const(ctx, torch::tensor({otherScalar}));
+
+ auto less = add_elementwise(
+ ctx, nvinfer1::ElementWiseOperation::kLESS, self, other, util::node_info(n) + "_less");
+ TRTORCH_CHECK(less, "Unable to create Less layer from node: " << *n);
+
+ auto equal = add_elementwise(
+ ctx, nvinfer1::ElementWiseOperation::kEQUAL, self, other, util::node_info(n) + "_equal");
+ TRTORCH_CHECK(equal, "Unable to create Equal layer from node: " << *n);
+
+ auto or_op = ctx->net->addElementWise(
+ *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));
+
+ LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+ return true;
+ }});
} // namespace
} // namespace impl
diff --git a/workspace/core/conversion/converters/impl/reduce.cpp b/tmp/changes.txt
index 101e9ac..bce7be9 100644
--- a/workspace/core/conversion/converters/impl/reduce.cpp
+++ b/tmp/changes.txt
@@ -11,25 +11,25 @@ namespace {
auto reduce_registrations TRTORCH_UNUSED =
RegisterNodeConversionPatterns()
- .pattern({"aten::mean(Tensor self, *, ScalarType? dtype=None) -> (Tensor)",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto in_tensor = args[0].ITensorOrFreeze(ctx);
- auto in_dims = util::toVec(in_tensor->getDimensions());
- LOG_WARNING("Mean Converter disregards dtype");
+ .pattern(
+ {"aten::mean(Tensor self, *, ScalarType? dtype=None) -> (Tensor)",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto in_tensor = args[0].ITensorOrFreeze(ctx);
+ auto in_dims = util::toVec(in_tensor->getDimensions());
+ LOG_WARNING("Mean Converter disregards dtype");
- uint32_t axis_mask = (uint32_t)(((uint64_t)1 << in_dims.size()) - 1);
+ uint32_t axis_mask = (uint32_t)(((uint64_t)1 << in_dims.size()) - 1);
- auto mean_layer =
- ctx->net->addReduce(*in_tensor, nvinfer1::ReduceOperation::kAVG, axis_mask, false);
+ auto mean_layer = ctx->net->addReduce(*in_tensor, nvinfer1::ReduceOperation::kAVG, axis_mask, false);
- TRTORCH_CHECK(mean_layer, "Unable to create mean layer from node: " << *n);
+ TRTORCH_CHECK(mean_layer, "Unable to create mean layer from node: " << *n);
- mean_layer->setName(util::node_info(n).c_str());
- auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], mean_layer->getOutput(0));
+ mean_layer->setName(util::node_info(n).c_str());
+ auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], mean_layer->getOutput(0));
- LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
- return true;
- }})
+ LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
+ return true;
+ }})
.pattern(
{"aten::mean.dim(Tensor self, int[] dim, bool keepdim=False, *, int? dtype=None) -> (Tensor)",
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
@@ -65,24 +65,25 @@ auto reduce_registrations TRTORCH_UNUSED =
LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
return true;
}})
- .pattern({"aten::sum(Tensor self, *, ScalarType? dtype=None) -> Tensor",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto in_tensor = args[0].ITensorOrFreeze(ctx);
- auto in_dims = util::toVec(in_tensor->getDimensions());
- LOG_WARNING("Sum Converter disregards dtype");
+ .pattern(
+ {"aten::sum(Tensor self, *, ScalarType? dtype=None) -> Tensor",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto in_tensor = args[0].ITensorOrFreeze(ctx);
+ auto in_dims = util::toVec(in_tensor->getDimensions());
+ LOG_WARNING("Sum Converter disregards dtype");
- uint32_t axis_mask = (uint32_t)(((uint64_t)1 << in_dims.size()) - 1);
+ uint32_t axis_mask = (uint32_t)(((uint64_t)1 << in_dims.size()) - 1);
- auto sum_layer = ctx->net->addReduce(*in_tensor, nvinfer1::ReduceOperation::kSUM, axis_mask, false);
+ auto sum_layer = ctx->net->addReduce(*in_tensor, nvinfer1::ReduceOperation::kSUM, axis_mask, false);
- TRTORCH_CHECK(sum_layer, "Unable to create sum layer from node: " << *n);
+ TRTORCH_CHECK(sum_layer, "Unable to create sum layer from node: " << *n);
- sum_layer->setName(util::node_info(n).c_str());
- auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], sum_layer->getOutput(0));
+ sum_layer->setName(util::node_info(n).c_str());
+ auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], sum_layer->getOutput(0));
- LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
- return true;
- }})
+ LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
+ return true;
+ }})
.pattern(
{"aten::sum.dim_IntList(Tensor self, int[1] dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor",
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
@@ -122,66 +123,67 @@ auto reduce_registrations TRTORCH_UNUSED =
LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
return true;
}})
- .pattern({"aten::prod(Tensor self, *, ScalarType? dtype=None) -> Tensor",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto in_tensor = args[0].ITensorOrFreeze(ctx);
- auto in_dims = util::toVec(in_tensor->getDimensions());
- LOG_WARNING("Prod Converter disregards dtype");
+ .pattern(
+ {"aten::prod(Tensor self, *, ScalarType? dtype=None) -> Tensor",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto in_tensor = args[0].ITensorOrFreeze(ctx);
+ auto in_dims = util::toVec(in_tensor->getDimensions());
+ LOG_WARNING("Prod Converter disregards dtype");
- uint32_t axis_mask = (uint32_t)(((uint64_t)1 << in_dims.size()) - 1);
+ uint32_t axis_mask = (uint32_t)(((uint64_t)1 << in_dims.size()) - 1);
- auto prod_layer =
- ctx->net->addReduce(*in_tensor, nvinfer1::ReduceOperation::kPROD, axis_mask, false);
+ auto prod_layer = ctx->net->addReduce(*in_tensor, nvinfer1::ReduceOperation::kPROD, axis_mask, false);
- TRTORCH_CHECK(prod_layer, "Unable to create sum layer from node: " << *n);
+ TRTORCH_CHECK(prod_layer, "Unable to create sum layer from node: " << *n);
- prod_layer->setName(util::node_info(n).c_str());
- auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], prod_layer->getOutput(0));
+ prod_layer->setName(util::node_info(n).c_str());
+ auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], prod_layer->getOutput(0));
- LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
- return true;
- }})
- .pattern({"aten::prod.dim_int(Tensor self, int dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto in_tensor = args[0].ITensorOrFreeze(ctx);
- auto dim = args[1].unwrapToInt();
- LOG_DEBUG("Dim to reduce:" << dim); // Some abuse of toDim but just for debug info
+ LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::prod.dim_int(Tensor self, int dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto in_tensor = args[0].ITensorOrFreeze(ctx);
+ auto dim = args[1].unwrapToInt();
+ LOG_DEBUG("Dim to reduce:" << dim); // Some abuse of toDim but just for debug info
- uint32_t axis_mask = 1 << dim;
- LOG_DEBUG("Axis Mask" << std::bitset<32>(axis_mask));
+ uint32_t axis_mask = 1 << dim;
+ LOG_DEBUG("Axis Mask" << std::bitset<32>(axis_mask));
- auto keepdim = args[2].unwrapToBool();
- LOG_DEBUG("Keep dims :" << keepdim);
+ auto keepdim = args[2].unwrapToBool();
+ LOG_DEBUG("Keep dims :" << keepdim);
- LOG_WARNING("Prod converter disregards dtype");
- auto prod_layer =
- ctx->net->addReduce(*in_tensor, nvinfer1::ReduceOperation::kPROD, axis_mask, keepdim);
+ LOG_WARNING("Prod converter disregards dtype");
+ auto prod_layer = ctx->net->addReduce(*in_tensor, nvinfer1::ReduceOperation::kPROD, axis_mask, keepdim);
- TRTORCH_CHECK(prod_layer, "Unable to create mean layer from node: " << *n);
+ TRTORCH_CHECK(prod_layer, "Unable to create mean layer from node: " << *n);
- prod_layer->setName(util::node_info(n).c_str());
- auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], prod_layer->getOutput(0));
+ prod_layer->setName(util::node_info(n).c_str());
+ auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], prod_layer->getOutput(0));
- LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
- return true;
- }})
- .pattern({"aten::max(Tensor self) -> Tensor",
- [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
- auto in_tensor = args[0].ITensorOrFreeze(ctx);
- auto in_dims = util::toVec(in_tensor->getDimensions());
+ LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
+ return true;
+ }})
+ .pattern(
+ {"aten::max(Tensor self) -> Tensor",
+ [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+ auto in_tensor = args[0].ITensorOrFreeze(ctx);
+ auto in_dims = util::toVec(in_tensor->getDimensions());
- uint32_t axis_mask = (uint32_t)(((uint64_t)1 << in_dims.size()) - 1);
+ uint32_t axis_mask = (uint32_t)(((uint64_t)1 << in_dims.size()) - 1);
- auto max_layer = ctx->net->addReduce(*in_tensor, nvinfer1::ReduceOperation::kMAX, axis_mask, false);
+ auto max_layer = ctx->net->addReduce(*in_tensor, nvinfer1::ReduceOperation::kMAX, axis_mask, false);
- TRTORCH_CHECK(max_layer, "Unable to create max layer from node: " << *n);
+ TRTORCH_CHECK(max_layer, "Unable to create max layer from node: " << *n);
- max_layer->setName(util::node_info(n).c_str());
- auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], max_layer->getOutput(0));
+ max_layer->setName(util::node_info(n).c_str());
+ auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], max_layer->getOutput(0));
- LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
- return true;
- }})
+ LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
+ return true;
+ }})
.pattern(
{"aten::min(Tensor self) -> Tensor", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
auto in_tensor = args[0].ITensorOrFreeze(ctx);
diff --git a/workspace/core/conversion/converters/impl/plugins/interpolate_plugin.cpp b/tmp/changes.txt
index 5689d88..ff9ba86 100644
--- a/workspace/core/conversion/converters/impl/plugins/interpolate_plugin.cpp
+++ b/tmp/changes.txt
@@ -254,7 +254,8 @@ int InterpolatePlugin::enqueue(
cudaStream_t stream) {
#if NV_TENSORRT_MAJOR < 7 || (NV_TENSORRT_MAJOR == 7 && NV_TENSORRT_MINOR < 1)
at::Tensor input = at::from_blob((void*)inputs[0], util::toVec(inputDesc->dims), [](void*) {}, tensor_options_);
- at::Tensor output = at::from_blob(outputs[0], util::volume(outputDesc->dims), [](void*) {}, tensor_options_);
+ at::Tensor output = at::from_blob(
+ outputs[0], util::volume(outputDesc->dims), [](void*) {}, tensor_options_);
at::cuda::CUDAStream torch_stream = at::cuda::getStreamFromPool();
at::cuda::CUDAStreamGuard torch_guard(torch_stream);
diff --git a/workspace/core/conversion/evaluators/prim.cpp b/tmp/changes.txt
index 38d6530..ee2275c 100644
--- a/workspace/core/conversion/evaluators/prim.cpp
+++ b/tmp/changes.txt
@@ -21,235 +21,245 @@ namespace {
auto prim_registrations =
RegisterNodeEvaluators()
- .evaluator({torch::jit::prim::Constant,
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- if (n->output()->type()->kind() == at::FunctionType::Kind) {
- return {};
- }
- return evaluators::toIValue(n->output());
- }})
- .evaluator({torch::jit::prim::NumToTensor,
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- return at::scalar_to_tensor(args.at(n->input(0)).IValue()->toScalar());
- }})
- .evaluator({torch::jit::prim::ListUnpack,
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- // Outputs is an IValue which has list of tensors which can be found in ctx->evaluated_value_map
- const torch::jit::IValue* outputs = args.at(n->input()).IValue();
- auto outputVec = outputs->toList().vec();
- return std::move(c10::ivalue::Tuple::create(outputVec));
- }})
- .evaluator({torch::jit::prim::ListConstruct,
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- const auto num_inputs = n->inputs().size();
- if (constTypesOnly(args)) {
- c10::ListTypePtr lt = n->output()->type()->expect<c10::ListType>();
- if (torch::jit::IntType::get() == lt->getElementType()) {
- c10::List<int64_t> list;
- list.reserve(num_inputs);
- for (auto in : n->inputs()) {
- list.emplace_back(std::move(args.at(in).unwrapToInt()));
- }
- return c10::optional<torch::jit::IValue>(std::move(torch::jit::IValue(list)));
- } else if (torch::jit::FloatType::get() == lt->getElementType()) {
- c10::List<double> list;
- list.reserve(num_inputs);
- for (auto in : n->inputs()) {
- list.emplace_back(std::move(args.at(in).unwrapToDouble()));
- }
- return c10::optional<torch::jit::IValue>(std::move(torch::jit::IValue(list)));
- } else if (lt->getElementType() == torch::jit::BoolType::get()) {
- c10::List<bool> list;
- list.reserve(num_inputs);
- for (auto in : n->inputs()) {
- list.emplace_back(std::move(args.at(in).unwrapToBool()));
- }
- return c10::optional<torch::jit::IValue>(std::move(torch::jit::IValue(list)));
- } else if (lt->getElementType()->isSubtypeOf(torch::jit::TensorType::get())) {
- c10::List<at::Tensor> list;
- list.reserve(num_inputs);
- for (auto in : n->inputs()) {
- if (args.at(in).isIValue()) {
- list.emplace_back(std::move(args.at(in).unwrapToTensor()));
- }
- }
- return c10::optional<torch::jit::IValue>(std::move(torch::jit::IValue(list)));
- } else {
- c10::TypePtr elementType = lt->getElementType();
- auto list = c10::impl::GenericList(elementType);
- list.reserve(num_inputs);
- for (auto in : n->inputs()) {
- list.emplace_back(std::move(*(args.at(in).IValue())));
- }
- return c10::optional<torch::jit::IValue>(std::move(torch::jit::IValue(list)));
- }
- } else {
- c10::ListTypePtr lt = n->output()->type()->expect<c10::ListType>();
- c10::TypePtr elementType = lt->getElementType();
- auto list = c10::impl::GenericList(elementType);
- list.reserve(num_inputs);
- for (auto in : n->inputs()) {
- if (args.at(in).isITensor()) {
- auto tensor_holder = TensorContainer();
- tensor_holder.hold_tensor(args.at(in).ITensor());
- auto ival = c10::IValue(std::move(c10::make_intrusive<TensorContainer>(tensor_holder)));
- list.emplace_back(std::move(ival));
- } else {
- list.emplace_back(std::move(args.at(in).unwrapToTensor()));
- }
- }
- return c10::optional<torch::jit::IValue>(std::move(torch::jit::IValue(list)));
- }
- }})
- .evaluator({c10::Symbol::fromQualString("prim::min"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- if (n->inputs().size() == 1) {
- auto a = args.at(n->input(0)).unwrapToIntList();
- int64_t min = std::numeric_limits<int64_t>::max();
+ .evaluator(
+ {torch::jit::prim::Constant,
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ if (n->output()->type()->kind() == at::FunctionType::Kind) {
+ return {};
+ }
+ return evaluators::toIValue(n->output());
+ }})
+ .evaluator(
+ {torch::jit::prim::NumToTensor,
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ return at::scalar_to_tensor(args.at(n->input(0)).IValue()->toScalar());
+ }})
+ .evaluator(
+ {torch::jit::prim::ListUnpack,
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ // Outputs is an IValue which has list of tensors which can be found in ctx->evaluated_value_map
+ const torch::jit::IValue* outputs = args.at(n->input()).IValue();
+ auto outputVec = outputs->toList().vec();
+ return std::move(c10::ivalue::Tuple::create(outputVec));
+ }})
+ .evaluator(
+ {torch::jit::prim::ListConstruct,
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ const auto num_inputs = n->inputs().size();
+ if (constTypesOnly(args)) {
+ c10::ListTypePtr lt = n->output()->type()->expect<c10::ListType>();
+ if (torch::jit::IntType::get() == lt->getElementType()) {
+ c10::List<int64_t> list;
+ list.reserve(num_inputs);
+ for (auto in : n->inputs()) {
+ list.emplace_back(std::move(args.at(in).unwrapToInt()));
+ }
+ return c10::optional<torch::jit::IValue>(std::move(torch::jit::IValue(list)));
+ } else if (torch::jit::FloatType::get() == lt->getElementType()) {
+ c10::List<double> list;
+ list.reserve(num_inputs);
+ for (auto in : n->inputs()) {
+ list.emplace_back(std::move(args.at(in).unwrapToDouble()));
+ }
+ return c10::optional<torch::jit::IValue>(std::move(torch::jit::IValue(list)));
+ } else if (lt->getElementType() == torch::jit::BoolType::get()) {
+ c10::List<bool> list;
+ list.reserve(num_inputs);
+ for (auto in : n->inputs()) {
+ list.emplace_back(std::move(args.at(in).unwrapToBool()));
+ }
+ return c10::optional<torch::jit::IValue>(std::move(torch::jit::IValue(list)));
+ } else if (lt->getElementType()->isSubtypeOf(torch::jit::TensorType::get())) {
+ c10::List<at::Tensor> list;
+ list.reserve(num_inputs);
+ for (auto in : n->inputs()) {
+ if (args.at(in).isIValue()) {
+ list.emplace_back(std::move(args.at(in).unwrapToTensor()));
+ }
+ }
+ return c10::optional<torch::jit::IValue>(std::move(torch::jit::IValue(list)));
+ } else {
+ c10::TypePtr elementType = lt->getElementType();
+ auto list = c10::impl::GenericList(elementType);
+ list.reserve(num_inputs);
+ for (auto in : n->inputs()) {
+ list.emplace_back(std::move(*(args.at(in).IValue())));
+ }
+ return c10::optional<torch::jit::IValue>(std::move(torch::jit::IValue(list)));
+ }
+ } else {
+ c10::ListTypePtr lt = n->output()->type()->expect<c10::ListType>();
+ c10::TypePtr elementType = lt->getElementType();
+ auto list = c10::impl::GenericList(elementType);
+ list.reserve(num_inputs);
+ for (auto in : n->inputs()) {
+ if (args.at(in).isITensor()) {
+ auto tensor_holder = TensorContainer();
+ tensor_holder.hold_tensor(args.at(in).ITensor());
+ auto ival = c10::IValue(std::move(c10::make_intrusive<TensorContainer>(tensor_holder)));
+ list.emplace_back(std::move(ival));
+ } else {
+ list.emplace_back(std::move(args.at(in).unwrapToTensor()));
+ }
+ }
+ return c10::optional<torch::jit::IValue>(std::move(torch::jit::IValue(list)));
+ }
+ }})
+ .evaluator(
+ {c10::Symbol::fromQualString("prim::min"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ if (n->inputs().size() == 1) {
+ auto a = args.at(n->input(0)).unwrapToIntList();
+ int64_t min = std::numeric_limits<int64_t>::max();
- for (size_t i = 0; i < a.size(); i++) {
- if (a[i] < min) {
- min = a[i];
- }
- }
+ for (size_t i = 0; i < a.size(); i++) {
+ if (a[i] < min) {
+ min = a[i];
+ }
+ }
- return min;
- } else if (n->inputs().size() == 2) {
- if (args.at(n->input(0)).IValue()->isInt()) {
- auto a = args.at(n->input(0)).unwrapToInt();
- if (args.at(n->input(1)).IValue()->isInt()) {
- auto b = args.at(n->input(1)).unwrapToInt();
- return a < b ? a : b;
- } else if (args.at(n->input(1)).IValue()->isDouble()) {
- auto b = args.at(n->input(1)).unwrapToDouble();
- return a < b ? a : b;
- } else {
- TRTORCH_THROW_ERROR(
- "Unimplemented data type for " << n->kind().toQualString() << " evaluator b arg: "
- << args.at(n->input(1)).IValue()->type()->str());
- return {};
- }
- } else if (args.at(n->input(0)).IValue()->isDouble()) {
- auto a = args.at(n->input(0)).unwrapToDouble();
- if (args.at(n->input(1)).IValue()->isInt()) {
- auto b = args.at(n->input(1)).unwrapToInt();
- return a < b ? a : b;
- } else if (args.at(n->input(1)).IValue()->isDouble()) {
- auto b = args.at(n->input(1)).unwrapToDouble();
- return a < b ? a : b;
- } else {
- TRTORCH_THROW_ERROR(
- "Unimplemented data type for " << n->kind().toQualString() << " evaluator b arg: "
- << args.at(n->input(1)).IValue()->type()->str());
- return {};
- }
- } else {
- TRTORCH_THROW_ERROR(
- "Unimplemented data type for " << n->kind().toQualString() << " evaluator a arg: "
- << args.at(n->input(0)).IValue()->type()->str());
- return {};
- }
- } else {
- TRTORCH_THROW_ERROR("Unimplemented " << n->kind().toQualString() << " evaluator case");
- return {};
- }
- },
- EvalOptions().validSchemas({
- "prim::min.self_int(int[] self) -> (int)",
- "prim::min.bool(bool a, bool b) -> (bool)",
- "prim::min.int(int a, int b) -> (bool)",
- "prim::min.float(float a, float b) -> (bool)",
- "prim::min.int_float(int a, float b) -> (bool)",
- "prim::min.float_int(float a, int b) -> (bool)",
- })})
- .evaluator({c10::Symbol::fromQualString("prim::max"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- if (n->inputs().size() == 1) {
- auto a = args.at(n->input(0)).unwrapToIntList();
- int64_t max = std::numeric_limits<int64_t>::min();
+ return min;
+ } else if (n->inputs().size() == 2) {
+ if (args.at(n->input(0)).IValue()->isInt()) {
+ auto a = args.at(n->input(0)).unwrapToInt();
+ if (args.at(n->input(1)).IValue()->isInt()) {
+ auto b = args.at(n->input(1)).unwrapToInt();
+ return a < b ? a : b;
+ } else if (args.at(n->input(1)).IValue()->isDouble()) {
+ auto b = args.at(n->input(1)).unwrapToDouble();
+ return a < b ? a : b;
+ } else {
+ TRTORCH_THROW_ERROR(
+ "Unimplemented data type for " << n->kind().toQualString() << " evaluator b arg: "
+ << args.at(n->input(1)).IValue()->type()->str());
+ return {};
+ }
+ } else if (args.at(n->input(0)).IValue()->isDouble()) {
+ auto a = args.at(n->input(0)).unwrapToDouble();
+ if (args.at(n->input(1)).IValue()->isInt()) {
+ auto b = args.at(n->input(1)).unwrapToInt();
+ return a < b ? a : b;
+ } else if (args.at(n->input(1)).IValue()->isDouble()) {
+ auto b = args.at(n->input(1)).unwrapToDouble();
+ return a < b ? a : b;
+ } else {
+ TRTORCH_THROW_ERROR(
+ "Unimplemented data type for " << n->kind().toQualString() << " evaluator b arg: "
+ << args.at(n->input(1)).IValue()->type()->str());
+ return {};
+ }
+ } else {
+ TRTORCH_THROW_ERROR(
+ "Unimplemented data type for " << n->kind().toQualString() << " evaluator a arg: "
+ << args.at(n->input(0)).IValue()->type()->str());
+ return {};
+ }
+ } else {
+ TRTORCH_THROW_ERROR("Unimplemented " << n->kind().toQualString() << " evaluator case");
+ return {};
+ }
+ },
+ EvalOptions().validSchemas({
+ "prim::min.self_int(int[] self) -> (int)",
+ "prim::min.bool(bool a, bool b) -> (bool)",
+ "prim::min.int(int a, int b) -> (bool)",
+ "prim::min.float(float a, float b) -> (bool)",
+ "prim::min.int_float(int a, float b) -> (bool)",
+ "prim::min.float_int(float a, int b) -> (bool)",
+ })})
+ .evaluator(
+ {c10::Symbol::fromQualString("prim::max"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ if (n->inputs().size() == 1) {
+ auto a = args.at(n->input(0)).unwrapToIntList();
+ int64_t max = std::numeric_limits<int64_t>::min();
- for (size_t i = 0; i < a.size(); i++) {
- if (a[i] > max) {
- max = a[i];
- }
- }
+ for (size_t i = 0; i < a.size(); i++) {
+ if (a[i] > max) {
+ max = a[i];
+ }
+ }
- return max;
- } else if (n->inputs().size() == 2) {
- if (args.at(n->input(0)).IValue()->isInt()) {
- auto a = args.at(n->input(0)).unwrapToInt();
- if (args.at(n->input(1)).IValue()->isInt()) {
- auto b = args.at(n->input(1)).unwrapToInt();
- return a > b ? a : b;
- } else if (args.at(n->input(1)).IValue()->isDouble()) {
- auto b = args.at(n->input(1)).unwrapToDouble();
- return a > b ? a : b;
- } else {
- TRTORCH_THROW_ERROR(
- "Unimplemented data type for " << n->kind().toQualString() << " evaluator b arg: "
- << args.at(n->input(1)).IValue()->type()->str());
- return {};
- }
- } else if (args.at(n->input(0)).IValue()->isDouble()) {
- auto a = args.at(n->input(0)).unwrapToDouble();
- if (args.at(n->input(1)).IValue()->isInt()) {
- auto b = args.at(n->input(1)).unwrapToInt();
- return a > b ? a : b;
- } else if (args.at(n->input(1)).IValue()->isDouble()) {
- auto b = args.at(n->input(1)).unwrapToDouble();
- return a > b ? a : b;
- } else {
- TRTORCH_THROW_ERROR(
- "Unimplemented data type for " << n->kind().toQualString() << " evaluator b arg: "
- << args.at(n->input(1)).IValue()->type()->str());
- return {};
- }
- } else {
- TRTORCH_THROW_ERROR(
- "Unimplemented data type for " << n->kind().toQualString() << " evaluator a arg: "
- << args.at(n->input(0)).IValue()->type()->str());
- return {};
- }
- } else {
- TRTORCH_THROW_ERROR("Unimplemented " << n->kind().toQualString() << " evaluator case");
- return {};
- }
- },
- EvalOptions().validSchemas({
- "prim::max.self_int(int[] self) -> (int)",
- "prim::max.bool(bool a, bool b) -> (bool)",
- "prim::max.int(int a, int b) -> (bool)",
- "prim::max.float(float a, float b) -> (bool)",
- "prim::max.int_float(int a, float b) -> (bool)",
- "prim::max.float_int(float a, int b) -> (bool)",
- })})
- .evaluator({c10::Symbol::fromQualString("prim::shape"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- LOG_WARNING("There may be undefined behavior using dynamic shape and prim::shape");
- auto tensor_var = args.at(n->input(0));
- if (tensor_var.isITensor()) {
- auto tensor = tensor_var.ITensor();
- return util::toVec(tensor->getDimensions());
- } else {
- auto tensor = tensor_var.unwrapToTensor();
- return tensor.sizes();
- }
- },
- EvalOptions().validSchemas({"prim::shape(Tensor a) -> (int[])"})})
- .evaluator({c10::Symbol::fromQualString("prim::unchecked_cast"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- return *(args.at(n->input(0)).IValue());
- }})
- .evaluator({c10::Symbol::fromQualString("prim::Uninitialized"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- return c10::IValue::uninitialized();
- }})
- .evaluator({c10::Symbol::fromQualString("prim::RaiseException"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- auto exception = args.at(n->input(0)).IValue();
- TRTORCH_THROW_ERROR("Error from TorchScript: " << *exception);
- return {};
- }});
+ return max;
+ } else if (n->inputs().size() == 2) {
+ if (args.at(n->input(0)).IValue()->isInt()) {
+ auto a = args.at(n->input(0)).unwrapToInt();
+ if (args.at(n->input(1)).IValue()->isInt()) {
+ auto b = args.at(n->input(1)).unwrapToInt();
+ return a > b ? a : b;
+ } else if (args.at(n->input(1)).IValue()->isDouble()) {
+ auto b = args.at(n->input(1)).unwrapToDouble();
+ return a > b ? a : b;
+ } else {
+ TRTORCH_THROW_ERROR(
+ "Unimplemented data type for " << n->kind().toQualString() << " evaluator b arg: "
+ << args.at(n->input(1)).IValue()->type()->str());
+ return {};
+ }
+ } else if (args.at(n->input(0)).IValue()->isDouble()) {
+ auto a = args.at(n->input(0)).unwrapToDouble();
+ if (args.at(n->input(1)).IValue()->isInt()) {
+ auto b = args.at(n->input(1)).unwrapToInt();
+ return a > b ? a : b;
+ } else if (args.at(n->input(1)).IValue()->isDouble()) {
+ auto b = args.at(n->input(1)).unwrapToDouble();
+ return a > b ? a : b;
+ } else {
+ TRTORCH_THROW_ERROR(
+ "Unimplemented data type for " << n->kind().toQualString() << " evaluator b arg: "
+ << args.at(n->input(1)).IValue()->type()->str());
+ return {};
+ }
+ } else {
+ TRTORCH_THROW_ERROR(
+ "Unimplemented data type for " << n->kind().toQualString() << " evaluator a arg: "
+ << args.at(n->input(0)).IValue()->type()->str());
+ return {};
+ }
+ } else {
+ TRTORCH_THROW_ERROR("Unimplemented " << n->kind().toQualString() << " evaluator case");
+ return {};
+ }
+ },
+ EvalOptions().validSchemas({
+ "prim::max.self_int(int[] self) -> (int)",
+ "prim::max.bool(bool a, bool b) -> (bool)",
+ "prim::max.int(int a, int b) -> (bool)",
+ "prim::max.float(float a, float b) -> (bool)",
+ "prim::max.int_float(int a, float b) -> (bool)",
+ "prim::max.float_int(float a, int b) -> (bool)",
+ })})
+ .evaluator(
+ {c10::Symbol::fromQualString("prim::shape"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ LOG_WARNING("There may be undefined behavior using dynamic shape and prim::shape");
+ auto tensor_var = args.at(n->input(0));
+ if (tensor_var.isITensor()) {
+ auto tensor = tensor_var.ITensor();
+ return util::toVec(tensor->getDimensions());
+ } else {
+ auto tensor = tensor_var.unwrapToTensor();
+ return tensor.sizes();
+ }
+ },
+ EvalOptions().validSchemas({"prim::shape(Tensor a) -> (int[])"})})
+ .evaluator(
+ {c10::Symbol::fromQualString("prim::unchecked_cast"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ return *(args.at(n->input(0)).IValue());
+ }})
+ .evaluator(
+ {c10::Symbol::fromQualString("prim::Uninitialized"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ return c10::IValue::uninitialized();
+ }})
+ .evaluator(
+ {c10::Symbol::fromQualString("prim::RaiseException"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ auto exception = args.at(n->input(0)).IValue();
+ TRTORCH_THROW_ERROR("Error from TorchScript: " << *exception);
+ return {};
+ }});
}
} // namespace evaluators
} // namespace conversion
diff --git a/workspace/core/conversion/evaluators/aten.cpp b/tmp/changes.txt
index 0e5c810..cf0b67f 100644
--- a/workspace/core/conversion/evaluators/aten.cpp
+++ b/tmp/changes.txt
@@ -114,204 +114,214 @@ DEFINE_TWO_INPUT_SIMPLE_EVALUATOR(
auto aten_registrations TRTORCH_UNUSED =
RegisterNodeEvaluators()
- .evaluator({c10::Symbol::fromQualString("aten::zeros"),
- // aten::zeros(int[] size, *, int? dtype=None, int? layout=None,
- // Device? device=None, bool? pin_memory=None) -> (Tensor)
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- auto options = torch::TensorOptions().layout(torch::kStrided).device(torch::kCUDA);
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::zeros"),
+ // aten::zeros(int[] size, *, int? dtype=None, int? layout=None,
+ // Device? device=None, bool? pin_memory=None) -> (Tensor)
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ auto options = torch::TensorOptions().layout(torch::kStrided).device(torch::kCUDA);
- // Input 1 here is the dtype
- if (!args.at(n->input(1)).isNone() && !args.at(n->input(1)).IValue()->isNone()) {
- options = options.dtype(c10::ScalarType(args.at(n->input(1)).unwrapToInt()));
- }
+ // Input 1 here is the dtype
+ if (!args.at(n->input(1)).isNone() && !args.at(n->input(1)).IValue()->isNone()) {
+ options = options.dtype(c10::ScalarType(args.at(n->input(1)).unwrapToInt()));
+ }
- auto out_tensor = torch::zeros(args.at(n->input(0)).unwrapToIntList().vec(), options);
- return out_tensor;
- }})
- .evaluator({c10::Symbol::fromQualString("aten::slice"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- c10::List<c10::IValue> list = args.at(n->input(0)).IValue()->to<c10::List<c10::IValue>>();
- int64_t start = args.at(n->input(1)).unwrapToInt();
- int64_t end = args.at(n->input(2)).unwrapToInt();
- int64_t step = args.at(n->input(3)).unwrapToInt();
+ auto out_tensor = torch::zeros(args.at(n->input(0)).unwrapToIntList().vec(), options);
+ return out_tensor;
+ }})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::slice"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ c10::List<c10::IValue> list = args.at(n->input(0)).IValue()->to<c10::List<c10::IValue>>();
+ int64_t start = args.at(n->input(1)).unwrapToInt();
+ int64_t end = args.at(n->input(2)).unwrapToInt();
+ int64_t step = args.at(n->input(3)).unwrapToInt();
- const int64_t list_size = list.size();
+ const int64_t list_size = list.size();
- // clamp start and end to the bounds of the list
- const auto normalized_start = std::max((int64_t)0, normalizeIndex(start, list_size));
- const auto normalized_end = std::min(list_size, normalizeIndex(end, list_size));
+ // clamp start and end to the bounds of the list
+ const auto normalized_start = std::max((int64_t)0, normalizeIndex(start, list_size));
+ const auto normalized_end = std::min(list_size, normalizeIndex(end, list_size));
- auto sliced_list = c10::impl::GenericList(list.elementType());
- if (normalized_end <= normalized_start) {
- // early exit if the slice is trivially empty
- return sliced_list;
- }
+ auto sliced_list = c10::impl::GenericList(list.elementType());
+ if (normalized_end <= normalized_start) {
+ // early exit if the slice is trivially empty
+ return sliced_list;
+ }
- sliced_list.reserve(normalized_end - normalized_start);
+ sliced_list.reserve(normalized_end - normalized_start);
- for (auto i = normalized_start; i < normalized_end;) {
- sliced_list.push_back(list.get(i));
- i += step;
- }
+ for (auto i = normalized_start; i < normalized_end;) {
+ sliced_list.push_back(list.get(i));
+ i += step;
+ }
- return sliced_list;
- },
- EvalOptions().validSchemas(
- {"aten::slice.t(t[] l, int start, int end=9223372036854775807, int step=1) -> (t[])"})})
- .evaluator({c10::Symbol::fromQualString("aten::len"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- c10::List<c10::IValue> list = args.at(n->input(0)).IValue()->to<c10::List<c10::IValue>>();
- return static_cast<int64_t>(list.size());
- },
- EvalOptions().validSchemas({"aten::len.t(t[] a) -> (int)"})})
- .evaluator({c10::Symbol::fromQualString("aten::size"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- LOG_WARNING("There may be undefined behavior using dynamic shape and aten::size");
- auto tensor_var = args.at(n->input(0));
- if (n->inputs().size() == 1) {
- if (tensor_var.isITensor()) {
- auto tensor = tensor_var.ITensor();
- return util::toVec(tensor->getDimensions());
- } else {
- auto tensor = tensor_var.unwrapToTensor();
- return tensor.sizes();
- }
- } else {
- auto dim = args.at(n->input(1)).unwrapToInt();
- if (tensor_var.isITensor()) {
- auto tensor = tensor_var.ITensor();
- return util::toVec(tensor->getDimensions())[dim];
- } else {
- auto tensor = tensor_var.unwrapToTensor();
- return tensor.sizes()[dim];
- }
- }
- },
- EvalOptions().validSchemas(
- {"aten::size(Tensor self) -> (int[])", "aten::size.int(Tensor self, int dim) -> (int)"})})
- .evaluator({c10::Symbol::fromQualString("aten::__getitem__"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- auto list = args.at(n->input(0)).IValue()->to<c10::List<c10::IValue>>();
- auto idx = args.at(n->input(1)).unwrapToInt();
+ return sliced_list;
+ },
+ EvalOptions().validSchemas(
+ {"aten::slice.t(t[] l, int start, int end=9223372036854775807, int step=1) -> (t[])"})})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::len"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ c10::List<c10::IValue> list = args.at(n->input(0)).IValue()->to<c10::List<c10::IValue>>();
+ return static_cast<int64_t>(list.size());
+ },
+ EvalOptions().validSchemas({"aten::len.t(t[] a) -> (int)"})})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::size"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ LOG_WARNING("There may be undefined behavior using dynamic shape and aten::size");
+ auto tensor_var = args.at(n->input(0));
+ if (n->inputs().size() == 1) {
+ if (tensor_var.isITensor()) {
+ auto tensor = tensor_var.ITensor();
+ return util::toVec(tensor->getDimensions());
+ } else {
+ auto tensor = tensor_var.unwrapToTensor();
+ return tensor.sizes();
+ }
+ } else {
+ auto dim = args.at(n->input(1)).unwrapToInt();
+ if (tensor_var.isITensor()) {
+ auto tensor = tensor_var.ITensor();
+ return util::toVec(tensor->getDimensions())[dim];
+ } else {
+ auto tensor = tensor_var.unwrapToTensor();
+ return tensor.sizes()[dim];
+ }
+ }
+ },
+ EvalOptions().validSchemas(
+ {"aten::size(Tensor self) -> (int[])", "aten::size.int(Tensor self, int dim) -> (int)"})})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::__getitem__"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ auto list = args.at(n->input(0)).IValue()->to<c10::List<c10::IValue>>();
+ auto idx = args.at(n->input(1)).unwrapToInt();
- const int64_t list_size = list.size();
- const int64_t normalized_idx = normalizeIndex(idx, list_size);
- TRTORCH_CHECK(
- normalized_idx >= 0 || normalized_idx < list_size,
- "List index out of range (aten::__getitem__)");
- return list.get(normalized_idx);
- },
- EvalOptions().validSchemas({
- "aten::__getitem__.t(t[](a) list, int idx) -> (t(*))",
- })})
- .evaluator({c10::Symbol::fromQualString("aten::append"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- auto list = args.at(n->input(0)).IValue()->to<c10::List<c10::IValue>>();
- auto el = args.at(n->input(1)).IValue();
+ const int64_t list_size = list.size();
+ const int64_t normalized_idx = normalizeIndex(idx, list_size);
+ TRTORCH_CHECK(
+ normalized_idx >= 0 || normalized_idx < list_size, "List index out of range (aten::__getitem__)");
+ return list.get(normalized_idx);
+ },
+ EvalOptions().validSchemas({
+ "aten::__getitem__.t(t[](a) list, int idx) -> (t(*))",
+ })})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::append"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ auto list = args.at(n->input(0)).IValue()->to<c10::List<c10::IValue>>();
+ auto el = args.at(n->input(1)).IValue();
- list.push_back(std::move(*el));
- return list;
- },
- EvalOptions().validSchemas({
- "aten::append.t(t[](a!) self, t(c -> *) el) -> (t[](a!))",
- })})
- .evaluator({c10::Symbol::fromQualString("aten::neg"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- auto el = args.at(n->input(0)).unwrapToInt();
+ list.push_back(std::move(*el));
+ return list;
+ },
+ EvalOptions().validSchemas({
+ "aten::append.t(t[](a!) self, t(c -> *) el) -> (t[](a!))",
+ })})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::neg"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ auto el = args.at(n->input(0)).unwrapToInt();
- return el * -1;
- },
- EvalOptions().validSchemas({
- "aten::neg.int(int a) -> (int)",
- })})
- .evaluator({c10::Symbol::fromQualString("aten::add"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- if (args.at(n->input(0)).IValue()->isInt()) {
- auto a = args.at(n->input(0)).unwrapToInt();
- auto b = args.at(n->input(1)).unwrapToInt();
- return a + b;
- } else if (args.at(n->input(0)).IValue()->isDouble()) {
- auto a = args.at(n->input(0)).unwrapToDouble();
- auto b = args.at(n->input(1)).unwrapToDouble();
- return a + b;
- } else {
- TRTORCH_THROW_ERROR(
- "Unimplemented data type for aten::add evaluator: "
- << args.at(n->input(0)).IValue()->type()->str());
- return {};
- }
- },
- EvalOptions().validSchemas(
- {"aten::add.int(int a, int b) -> (int)", "aten::add.float(float a, float b) -> (float)"})})
- .evaluator({c10::Symbol::fromQualString("aten::add_"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- if (args.at(n->input(0)).IValue()->isList()) {
- auto a = args.at(n->input(0)).IValue()->toListRef();
- auto b = args.at(n->input(1)).IValue()->toListRef();
+ return el * -1;
+ },
+ EvalOptions().validSchemas({
+ "aten::neg.int(int a) -> (int)",
+ })})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::add"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ if (args.at(n->input(0)).IValue()->isInt()) {
+ auto a = args.at(n->input(0)).unwrapToInt();
+ auto b = args.at(n->input(1)).unwrapToInt();
+ return a + b;
+ } else if (args.at(n->input(0)).IValue()->isDouble()) {
+ auto a = args.at(n->input(0)).unwrapToDouble();
+ auto b = args.at(n->input(1)).unwrapToDouble();
+ return a + b;
+ } else {
+ TRTORCH_THROW_ERROR(
+ "Unimplemented data type for aten::add evaluator: "
+ << args.at(n->input(0)).IValue()->type()->str());
+ return {};
+ }
+ },
+ EvalOptions().validSchemas(
+ {"aten::add.int(int a, int b) -> (int)", "aten::add.float(float a, float b) -> (float)"})})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::add_"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ if (args.at(n->input(0)).IValue()->isList()) {
+ auto a = args.at(n->input(0)).IValue()->toListRef();
+ auto b = args.at(n->input(1)).IValue()->toListRef();
- c10::ListTypePtr lt = n->output()->type()->expect<c10::ListType>();
- c10::TypePtr elementType = lt->getElementType();
+ c10::ListTypePtr lt = n->output()->type()->expect<c10::ListType>();
+ c10::TypePtr elementType = lt->getElementType();
- auto merged = c10::impl::GenericList(elementType);
- merged.reserve(a.size() + b.size());
+ auto merged = c10::impl::GenericList(elementType);
+ merged.reserve(a.size() + b.size());
- for (auto each : a) {
- merged.emplace_back(each);
- }
+ for (auto each : a) {
+ merged.emplace_back(each);
+ }
- for (auto each : b) {
- merged.emplace_back(each);
- }
+ for (auto each : b) {
+ merged.emplace_back(each);
+ }
- return merged;
- } else {
- TRTORCH_THROW_ERROR(
- "Unimplemented data type for aten::add_ evaluator: "
- << args.at(n->input(0)).IValue()->type()->str());
- return {};
- }
- },
- EvalOptions().validSchemas({"aten::add_.t(t[](a!) self, t[] b) -> (t[])"})})
- .evaluator({c10::Symbol::fromQualString("aten::mul"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- if (args.at(n->input(0)).IValue()->isInt()) {
- auto a = args.at(n->input(0)).unwrapToInt();
- auto b = args.at(n->input(1)).unwrapToInt();
- return a * b;
- } else if (args.at(n->input(0)).IValue()->isDouble()) {
- auto a = args.at(n->input(0)).unwrapToDouble();
- auto b = args.at(n->input(1)).unwrapToDouble();
- return a * b;
- } else {
- TRTORCH_THROW_ERROR(
- "Unimplemented data type for aten::mul evaluator: "
- << args.at(n->input(0)).IValue()->type()->str());
- return {};
- }
- },
- EvalOptions().validSchemas(
- {"aten::mul.int(int a, int b) -> (int)", "aten::mul.float(float a, float b) -> (float)"})})
- .evaluator({c10::Symbol::fromQualString("aten::sub"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- if (args.at(n->input(0)).IValue()->isInt()) {
- auto a = args.at(n->input(0)).unwrapToInt();
- auto b = args.at(n->input(1)).unwrapToInt();
- return a - b;
- } else if (args.at(n->input(0)).IValue()->isDouble()) {
- auto a = args.at(n->input(0)).unwrapToDouble();
- auto b = args.at(n->input(1)).unwrapToDouble();
- return a - b;
- } else {
- TRTORCH_THROW_ERROR(
- "Unimplemented data type for aten::sub evaluator: "
- << args.at(n->input(0)).IValue()->type()->str());
- return {};
- }
- },
- EvalOptions().validSchemas({
- "aten::sub.float(float a, float b) -> (float)",
- "aten::sub.int(int a, int b) -> (int)",
- })})
+ return merged;
+ } else {
+ TRTORCH_THROW_ERROR(
+ "Unimplemented data type for aten::add_ evaluator: "
+ << args.at(n->input(0)).IValue()->type()->str());
+ return {};
+ }
+ },
+ EvalOptions().validSchemas({"aten::add_.t(t[](a!) self, t[] b) -> (t[])"})})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::mul"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ if (args.at(n->input(0)).IValue()->isInt()) {
+ auto a = args.at(n->input(0)).unwrapToInt();
+ auto b = args.at(n->input(1)).unwrapToInt();
+ return a * b;
+ } else if (args.at(n->input(0)).IValue()->isDouble()) {
+ auto a = args.at(n->input(0)).unwrapToDouble();
+ auto b = args.at(n->input(1)).unwrapToDouble();
+ return a * b;
+ } else {
+ TRTORCH_THROW_ERROR(
+ "Unimplemented data type for aten::mul evaluator: "
+ << args.at(n->input(0)).IValue()->type()->str());
+ return {};
+ }
+ },
+ EvalOptions().validSchemas(
+ {"aten::mul.int(int a, int b) -> (int)", "aten::mul.float(float a, float b) -> (float)"})})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::sub"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ if (args.at(n->input(0)).IValue()->isInt()) {
+ auto a = args.at(n->input(0)).unwrapToInt();
+ auto b = args.at(n->input(1)).unwrapToInt();
+ return a - b;
+ } else if (args.at(n->input(0)).IValue()->isDouble()) {
+ auto a = args.at(n->input(0)).unwrapToDouble();
+ auto b = args.at(n->input(1)).unwrapToDouble();
+ return a - b;
+ } else {
+ TRTORCH_THROW_ERROR(
+ "Unimplemented data type for aten::sub evaluator: "
+ << args.at(n->input(0)).IValue()->type()->str());
+ return {};
+ }
+ },
+ EvalOptions().validSchemas({
+ "aten::sub.float(float a, float b) -> (float)",
+ "aten::sub.int(int a, int b) -> (int)",
+ })})
.evaluator(
{c10::Symbol::fromQualString("aten::Bool"),
[](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
@@ -329,213 +339,225 @@ auto aten_registrations TRTORCH_UNUSED =
}
},
EvalOptions().validSchemas({"aten::Bool.int(int a) -> (bool)", "aten::Bool.float(float b) -> (bool)"})})
- .evaluator({c10::Symbol::fromQualString("aten::Float"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- if (args.at(n->input(0)).IValue()->isInt()) {
- auto a = args.at(n->input(0)).unwrapToInt();
- return (float)a;
- } else if (args.at(n->input(0)).IValue()->isDouble()) {
- auto a = args.at(n->input(0)).unwrapToDouble();
- return (float)a;
- } else if (args.at(n->input(0)).IValue()->isBool()) {
- auto a = args.at(n->input(0)).unwrapToBool();
- return (float)a;
- } else {
- TRTORCH_THROW_ERROR(
- "Unimplemented data type for aten::Float evaluator: "
- << args.at(n->input(0)).IValue()->type()->str());
- return {};
- }
- },
- EvalOptions().validSchemas({
- "aten::Float.Scalar(Scalar a) -> float",
- "aten::Float.int(int a) -> float",
- "aten::Float.bool(bool a) -> float",
- })})
- .evaluator({c10::Symbol::fromQualString("aten::__not__"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- auto el = args.at(n->input(0)).unwrapToBool();
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::Float"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ if (args.at(n->input(0)).IValue()->isInt()) {
+ auto a = args.at(n->input(0)).unwrapToInt();
+ return (float)a;
+ } else if (args.at(n->input(0)).IValue()->isDouble()) {
+ auto a = args.at(n->input(0)).unwrapToDouble();
+ return (float)a;
+ } else if (args.at(n->input(0)).IValue()->isBool()) {
+ auto a = args.at(n->input(0)).unwrapToBool();
+ return (float)a;
+ } else {
+ TRTORCH_THROW_ERROR(
+ "Unimplemented data type for aten::Float evaluator: "
+ << args.at(n->input(0)).IValue()->type()->str());
+ return {};
+ }
+ },
+ EvalOptions().validSchemas({
+ "aten::Float.Scalar(Scalar a) -> float",
+ "aten::Float.int(int a) -> float",
+ "aten::Float.bool(bool a) -> float",
+ })})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::__not__"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ auto el = args.at(n->input(0)).unwrapToBool();
- return !el;
- },
- EvalOptions().validSchemas({
- "aten::__not__(bool self) -> bool",
- })})
- .evaluator({c10::Symbol::fromQualString("aten::__is__"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- auto self = args.at(n->input(0)).IValue();
- auto obj = args.at(n->input(1)).IValue();
+ return !el;
+ },
+ EvalOptions().validSchemas({
+ "aten::__not__(bool self) -> bool",
+ })})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::__is__"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ auto self = args.at(n->input(0)).IValue();
+ auto obj = args.at(n->input(1)).IValue();
- return self->isSameIdentity(*obj);
- },
- EvalOptions().validSchemas({
- "aten::__is__(t1 self, t2 obj) -> bool",
- })})
- .evaluator({c10::Symbol::fromQualString("aten::__isnot__"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- auto self = args.at(n->input(0)).IValue();
- auto obj = args.at(n->input(1)).IValue();
+ return self->isSameIdentity(*obj);
+ },
+ EvalOptions().validSchemas({
+ "aten::__is__(t1 self, t2 obj) -> bool",
+ })})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::__isnot__"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ auto self = args.at(n->input(0)).IValue();
+ auto obj = args.at(n->input(1)).IValue();
- return !self->isSameIdentity(*obj);
- },
- EvalOptions().validSchemas({
- "aten::__isnot__(t1 self, t2 obj) -> bool",
- })})
- .evaluator({c10::Symbol::fromQualString("aten::numel"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- LOG_WARNING("There may be undefined behavior using dynamic shape and aten::numel");
- auto tensor_var = args.at(n->input(0));
- if (tensor_var.isITensor()) {
- auto tensor = tensor_var.ITensor();
- return util::volume(tensor->getDimensions());
- } else {
- auto tensor = tensor_var.unwrapToTensor();
- return tensor.numel();
- }
- },
- EvalOptions().validSchemas({
- "aten::numel(Tensor self) -> int",
- })})
- .evaluator({c10::Symbol::fromQualString("aten::t"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- auto tensor_var = args.at(n->input(0));
- if (tensor_var.IValue()->isTensor()) {
- auto tensor = tensor_var.unwrapToTensor();
- return tensor.t();
- } else {
- TRTORCH_THROW_ERROR("Unimplemented data type for aten::t evaluator: ITensor");
- return {};
- }
- },
- EvalOptions().validSchemas({
- "aten::t(Tensor self) -> Tensor",
- })})
- .evaluator({c10::Symbol::fromQualString("aten::dim"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- auto tensor_var = args.at(n->input(0));
- if (tensor_var.isITensor()) {
- auto tensor = tensor_var.ITensor();
- return tensor->getDimensions().nbDims;
- } else {
- auto tensor = tensor_var.unwrapToTensor();
- return tensor.dim();
- }
- },
- EvalOptions().validSchemas({
- "aten::dim(Tensor self) -> int",
- })})
- .evaluator({c10::Symbol::fromQualString("aten::div"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- if (args.at(n->input(0)).IValue()->isInt()) {
- auto a = args.at(n->input(0)).unwrapToInt();
- auto b = args.at(n->input(1)).unwrapToInt();
- return static_cast<double>(a) / static_cast<double>(b);
- } else if (args.at(n->input(0)).IValue()->isDouble()) {
- auto a = args.at(n->input(0)).unwrapToDouble();
- auto b = args.at(n->input(1)).unwrapToDouble();
- return a / b;
- } else {
- TRTORCH_THROW_ERROR(
- "Unimplemented data type for aten::div evaluator: "
- << args.at(n->input(0)).IValue()->type()->str());
- return {};
- }
- },
- EvalOptions().validSchemas({
- "aten::div.float(float a, float b) -> (float)",
- "aten::div.int(int a, int b) -> (float)",
- })})
- .evaluator({c10::Symbol::fromQualString("aten::floordiv"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- if (args.at(n->input(0)).IValue()->isInt()) {
- auto a = args.at(n->input(0)).unwrapToInt();
- auto b = args.at(n->input(1)).unwrapToInt();
- return std::floor(a / b);
- } else if (args.at(n->input(0)).IValue()->isDouble()) {
- auto a = args.at(n->input(0)).unwrapToDouble();
- auto b = args.at(n->input(1)).unwrapToDouble();
- return std::floor(a / b);
- } else {
- TRTORCH_THROW_ERROR(
- "Unimplemented data type for aten::floordiv evaluator: "
- << args.at(n->input(0)).IValue()->type()->str());
- return {};
- }
- },
- EvalOptions().validSchemas({
- "aten::floordiv.float(float a, float b) -> (int)",
- "aten::floordiv.int(int a, int b) -> (int)",
- })})
- .evaluator({c10::Symbol::fromQualString("aten::floor"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- auto el = args.at(n->input(0)).unwrapToDouble();
+ return !self->isSameIdentity(*obj);
+ },
+ EvalOptions().validSchemas({
+ "aten::__isnot__(t1 self, t2 obj) -> bool",
+ })})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::numel"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ LOG_WARNING("There may be undefined behavior using dynamic shape and aten::numel");
+ auto tensor_var = args.at(n->input(0));
+ if (tensor_var.isITensor()) {
+ auto tensor = tensor_var.ITensor();
+ return util::volume(tensor->getDimensions());
+ } else {
+ auto tensor = tensor_var.unwrapToTensor();
+ return tensor.numel();
+ }
+ },
+ EvalOptions().validSchemas({
+ "aten::numel(Tensor self) -> int",
+ })})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::t"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ auto tensor_var = args.at(n->input(0));
+ if (tensor_var.IValue()->isTensor()) {
+ auto tensor = tensor_var.unwrapToTensor();
+ return tensor.t();
+ } else {
+ TRTORCH_THROW_ERROR("Unimplemented data type for aten::t evaluator: ITensor");
+ return {};
+ }
+ },
+ EvalOptions().validSchemas({
+ "aten::t(Tensor self) -> Tensor",
+ })})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::dim"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ auto tensor_var = args.at(n->input(0));
+ if (tensor_var.isITensor()) {
+ auto tensor = tensor_var.ITensor();
+ return tensor->getDimensions().nbDims;
+ } else {
+ auto tensor = tensor_var.unwrapToTensor();
+ return tensor.dim();
+ }
+ },
+ EvalOptions().validSchemas({
+ "aten::dim(Tensor self) -> int",
+ })})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::div"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ if (args.at(n->input(0)).IValue()->isInt()) {
+ auto a = args.at(n->input(0)).unwrapToInt();
+ auto b = args.at(n->input(1)).unwrapToInt();
+ return static_cast<double>(a) / static_cast<double>(b);
+ } else if (args.at(n->input(0)).IValue()->isDouble()) {
+ auto a = args.at(n->input(0)).unwrapToDouble();
+ auto b = args.at(n->input(1)).unwrapToDouble();
+ return a / b;
+ } else {
+ TRTORCH_THROW_ERROR(
+ "Unimplemented data type for aten::div evaluator: "
+ << args.at(n->input(0)).IValue()->type()->str());
+ return {};
+ }
+ },
+ EvalOptions().validSchemas({
+ "aten::div.float(float a, float b) -> (float)",
+ "aten::div.int(int a, int b) -> (float)",
+ })})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::floordiv"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ if (args.at(n->input(0)).IValue()->isInt()) {
+ auto a = args.at(n->input(0)).unwrapToInt();
+ auto b = args.at(n->input(1)).unwrapToInt();
+ return std::floor(a / b);
+ } else if (args.at(n->input(0)).IValue()->isDouble()) {
+ auto a = args.at(n->input(0)).unwrapToDouble();
+ auto b = args.at(n->input(1)).unwrapToDouble();
+ return std::floor(a / b);
+ } else {
+ TRTORCH_THROW_ERROR(
+ "Unimplemented data type for aten::floordiv evaluator: "
+ << args.at(n->input(0)).IValue()->type()->str());
+ return {};
+ }
+ },
+ EvalOptions().validSchemas({
+ "aten::floordiv.float(float a, float b) -> (int)",
+ "aten::floordiv.int(int a, int b) -> (int)",
+ })})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::floor"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ auto el = args.at(n->input(0)).unwrapToDouble();
- return static_cast<int64_t>(std::floor(el));
- },
- EvalOptions().validSchemas({
- "aten::floor.float(float a) -> (int)",
- })})
- .evaluator({c10::Symbol::fromQualString("aten::warn"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- auto warning = args.at(n->input(0)).IValue();
- LOG_WARNING("Warning from TorchScript: " << *warning);
- return {};
- },
- EvalOptions()})
- .evaluator({c10::Symbol::fromQualString("aten::arange"),
- [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
- int input_size = n->inputs().size();
- int scalar_count = 0;
- for (int i = 0; i < input_size; i++) {
- if (args.at(n->input(i)).IValue()->isScalar()) {
- scalar_count += 1;
- }
- }
- if (scalar_count == 1) {
- if (args.at(n->input(0)).IValue()->isInt()) {
- int end_scalar = args.at(n->input(0)).unwrapToInt();
- return torch::arange(end_scalar);
- } else if (args.at(n->input(0)).IValue()->isDouble()) {
- float end_scalar = args.at(n->input(0)).unwrapToScalar().to<float>();
- return torch::arange(end_scalar);
- }
- } else if (scalar_count == 2) {
- if (args.at(n->input(0)).IValue()->isDouble() || args.at(n->input(1)).IValue()->isDouble()) {
- float start_scalar = args.at(n->input(0)).unwrapToScalar().to<float>();
- float end_scalar = args.at(n->input(1)).unwrapToScalar().to<float>();
- return torch::arange(start_scalar, end_scalar);
- } else {
- int start_scalar = args.at(n->input(0)).unwrapToInt();
- int end_scalar = args.at(n->input(1)).unwrapToInt();
- return torch::arange(start_scalar, end_scalar);
- }
- } else if (scalar_count == 3) {
- if (args.at(n->input(0)).IValue()->isDouble() || args.at(n->input(1)).IValue()->isDouble() ||
- args.at(n->input(2)).IValue()->isDouble()) {
- float start_scalar = args.at(n->input(0)).unwrapToScalar().to<float>();
- float end_scalar = args.at(n->input(1)).unwrapToScalar().to<float>();
- float step_scalar = args.at(n->input(2)).unwrapToScalar().to<float>();
- return torch::arange(start_scalar, end_scalar, step_scalar);
- } else {
- int start_scalar = args.at(n->input(0)).unwrapToInt();
- int end_scalar = args.at(n->input(1)).unwrapToInt();
- int step_scalar = args.at(n->input(2)).unwrapToInt();
- return torch::arange(start_scalar, end_scalar, step_scalar);
- }
- } else {
- TRTORCH_THROW_ERROR(
- "Invalid input argument size for aten::arange, input argument size: " << input_size);
- }
- return {};
- },
- EvalOptions().validSchemas({
- R"SIG(aten::arange(Scalar end, *, int? dtype=None, int? layout=None,
+ return static_cast<int64_t>(std::floor(el));
+ },
+ EvalOptions().validSchemas({
+ "aten::floor.float(float a) -> (int)",
+ })})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::warn"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ auto warning = args.at(n->input(0)).IValue();
+ LOG_WARNING("Warning from TorchScript: " << *warning);
+ return {};
+ },
+ EvalOptions()})
+ .evaluator(
+ {c10::Symbol::fromQualString("aten::arange"),
+ [](const torch::jit::Node* n, kwargs& args) -> c10::optional<torch::jit::IValue> {
+ int input_size = n->inputs().size();
+ int scalar_count = 0;
+ for (int i = 0; i < input_size; i++) {
+ if (args.at(n->input(i)).IValue()->isScalar()) {
+ scalar_count += 1;
+ }
+ }
+ if (scalar_count == 1) {
+ if (args.at(n->input(0)).IValue()->isInt()) {
+ int end_scalar = args.at(n->input(0)).unwrapToInt();
+ return torch::arange(end_scalar);
+ } else if (args.at(n->input(0)).IValue()->isDouble()) {
+ float end_scalar = args.at(n->input(0)).unwrapToScalar().to<float>();
+ return torch::arange(end_scalar);
+ }
+ } else if (scalar_count == 2) {
+ if (args.at(n->input(0)).IValue()->isDouble() || args.at(n->input(1)).IValue()->isDouble()) {
+ float start_scalar = args.at(n->input(0)).unwrapToScalar().to<float>();
+ float end_scalar = args.at(n->input(1)).unwrapToScalar().to<float>();
+ return torch::arange(start_scalar, end_scalar);
+ } else {
+ int start_scalar = args.at(n->input(0)).unwrapToInt();
+ int end_scalar = args.at(n->input(1)).unwrapToInt();
+ return torch::arange(start_scalar, end_scalar);
+ }
+ } else if (scalar_count == 3) {
+ if (args.at(n->input(0)).IValue()->isDouble() || args.at(n->input(1)).IValue()->isDouble() ||
+ args.at(n->input(2)).IValue()->isDouble()) {
+ float start_scalar = args.at(n->input(0)).unwrapToScalar().to<float>();
+ float end_scalar = args.at(n->input(1)).unwrapToScalar().to<float>();
+ float step_scalar = args.at(n->input(2)).unwrapToScalar().to<float>();
+ return torch::arange(start_scalar, end_scalar, step_scalar);
+ } else {
+ int start_scalar = args.at(n->input(0)).unwrapToInt();
+ int end_scalar = args.at(n->input(1)).unwrapToInt();
+ int step_scalar = args.at(n->input(2)).unwrapToInt();
+ return torch::arange(start_scalar, end_scalar, step_scalar);
+ }
+ } else {
+ TRTORCH_THROW_ERROR(
+ "Invalid input argument size for aten::arange, input argument size: " << input_size);
+ }
+ return {};
+ },
+ EvalOptions().validSchemas({
+ R"SIG(aten::arange(Scalar end, *, int? dtype=None, int? layout=None,
Device? device=None, bool? pin_memory=None) -> (Tensor))SIG",
- R"SIG(aten::arange.start(Scalar start, Scalar end, *, ScalarType? dtype=None,
+ R"SIG(aten::arange.start(Scalar start, Scalar end, *, ScalarType? dtype=None,
Layout? layout=None, Device? device=None, bool? pin_memory=None) -> (Tensor))SIG",
- R"SIG(aten::arange.start_step(Scalar start, Scalar end, Scalar step, *, ScalarType? dtype=None,
+ R"SIG(aten::arange.start_step(Scalar start, Scalar end, Scalar step, *, ScalarType? dtype=None,
Layout? layout=None, Device? device=None, bool? pin_memory=None) -> (Tensor))SIG",
- })});
+ })});
} // namespace
} // namespace evaluators
} // namespace conversion
diff --git a/workspace/core/lowering/register_trt_placeholder_ops.cpp b/tmp/changes.txt
index b15ce5c..bafb2c3 100644
--- a/workspace/core/lowering/register_trt_placeholder_ops.cpp
+++ b/tmp/changes.txt
@@ -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) {}, aliasAnalysisFromSchema()),
+ Operator(
+ "trt::const(Tensor val) -> Tensor",
+ [](Stack* stack) {},
+ aliasAnalysisFromSchema()),
});
} // namespace jit
diff --git a/workspace/cpp/api/include/trtorch/ptq.h b/tmp/changes.txt
index 8efc802..bf9fe91 100644
--- a/workspace/cpp/api/include/trtorch/ptq.h
+++ b/tmp/changes.txt
@@ -159,7 +159,7 @@ class Int8Calibrator : Algorithm {
*
* @return nvinfer1::IInt8Calibrator*
*/
- operator nvinfer1::IInt8Calibrator*() {
+ operator nvinfer1::IInt8Calibrator *() {
return reinterpret_cast<nvinfer1::IInt8Calibrator*>(this);
}
@@ -274,7 +274,7 @@ class Int8CacheCalibrator : Algorithm {
*
* @return nvinfer1::IInt8Calibrator*
*/
- operator nvinfer1::IInt8Calibrator*() {
+ operator nvinfer1::IInt8Calibrator *() {
return reinterpret_cast<nvinfer1::IInt8Calibrator*>(this);
}
diff --git a/workspace/py/trtorch/csrc/util.h b/tmp/changes.txt
index 22515f5..db5f131 100644
--- a/workspace/py/trtorch/csrc/util.h
+++ b/tmp/changes.txt
@@ -19,8 +19,8 @@ template <typename T>
py::function getOverload(const T* self, const std::string& overloadName) {
py::function overload = py::get_override(self, overloadName.c_str());
if (!overload) {
- std::string msg{"Method: " + overloadName +
- " was not overriden. Please provide an implementation for this method."};
+ std::string msg{
+ "Method: " + overloadName + " was not overriden. Please provide an implementation for this method."};
LOG_ERROR(msg);
}
return overload;
ERROR: Some files do not conform to style guidelines
everything Signed-off-by: Naren Dasan <naren@narendasan.com> Signed-off-by: Naren Dasan <narens@nvidia.com>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Code conforms to Python style guidelines
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There are some changes that do not conform to C++ style guidelines:
diff --git a/workspace/core/conversion/converters/impl/plugins/interpolate_plugin.cpp b/tmp/changes.txt
index 5689d88..ff9ba86 100644
--- a/workspace/core/conversion/converters/impl/plugins/interpolate_plugin.cpp
+++ b/tmp/changes.txt
@@ -254,7 +254,8 @@ int InterpolatePlugin::enqueue(
cudaStream_t stream) {
#if NV_TENSORRT_MAJOR < 7 || (NV_TENSORRT_MAJOR == 7 && NV_TENSORRT_MINOR < 1)
at::Tensor input = at::from_blob((void*)inputs[0], util::toVec(inputDesc->dims), [](void*) {}, tensor_options_);
- at::Tensor output = at::from_blob(outputs[0], util::volume(outputDesc->dims), [](void*) {}, tensor_options_);
+ at::Tensor output = at::from_blob(
+ outputs[0], util::volume(outputDesc->dims), [](void*) {}, tensor_options_);
at::cuda::CUDAStream torch_stream = at::cuda::getStreamFromPool();
at::cuda::CUDAStreamGuard torch_guard(torch_stream);
diff --git a/workspace/core/lowering/register_trt_placeholder_ops.cpp b/tmp/changes.txt
index b15ce5c..bafb2c3 100644
--- a/workspace/core/lowering/register_trt_placeholder_ops.cpp
+++ b/tmp/changes.txt
@@ -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) {}, aliasAnalysisFromSchema()),
+ Operator(
+ "trt::const(Tensor val) -> Tensor",
+ [](Stack* stack) {},
+ aliasAnalysisFromSchema()),
});
} // namespace jit
ERROR: Some files do not conform to style guidelines
Signed-off-by: Naren Dasan <naren@narendasan.com> Signed-off-by: Naren Dasan <narens@nvidia.com>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Code conforms to Python style guidelines
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Code conforms to C++ style guidelines
Description
Updates dependency library versions:
Type of change
Please delete options that are not relevant and/or add your own.
Checklist: