From 569bcdee704667d5ef5f9781897455a05cd28db4 Mon Sep 17 00:00:00 2001 From: hongwei03 Date: Fri, 18 Mar 2022 13:05:19 +0800 Subject: [PATCH 1/4] feat: Add converter files for torch::max Signed-off-by: hongwei03 --- core/conversion/converters/impl/max.cpp | 47 +++++++++++++++++++++++++ 1 file changed, 47 insertions(+) create mode 100644 core/conversion/converters/impl/max.cpp diff --git a/core/conversion/converters/impl/max.cpp b/core/conversion/converters/impl/max.cpp new file mode 100644 index 0000000000..83f5aefd80 --- /dev/null +++ b/core/conversion/converters/impl/max.cpp @@ -0,0 +1,47 @@ +#include "NvInfer.h" +#include "core/conversion/converters/converters.h" +#include "core/conversion/tensorcontainer/TensorContainer.h" +#include "core/util/prelude.h" +#include "torch/torch.h" + +#include +#include + +namespace torch_tensorrt { +namespace core { +namespace conversion { +namespace converters { +namespace impl { +namespace { +auto max_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern( + {"aten::max.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto k = 1; + auto dim = args[1].unwrapToInt(); + auto largest = true; + auto selfDim = util::toVec(self->getDimensions()); + if (dim < 0) { + dim = selfDim.size() + dim; + } + uint32_t shiftDim = 1 << dim; + + auto TopKOperation = largest ? (nvinfer1::TopKOperation::kMAX) : (nvinfer1::TopKOperation::kMIN); + + auto new_layer = ctx->net->addTopK(*self, TopKOperation, 1, shiftDim); + TORCHTRT_CHECK(new_layer, "Unable to create max layer from node: " << *n); + + auto out0 = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0)); + auto out1 = ctx->AssociateValueAndTensor(n->outputs()[1], new_layer->getOutput(1)); + + LOG_DEBUG("Output tensor(0) shape: " << out0->getDimensions()); + LOG_DEBUG("Output tensor(1) shape: " << out1->getDimensions()); + + return true; + }}); +} // namespace +} // namespace impl +} // namespace converters +} // namespace conversion +} // namespace core +} // namespace torch_tensorrt From f628acadc5d2d0a9b093b8d9827e61b0fcc7315f Mon Sep 17 00:00:00 2001 From: Yiran Liu Date: Sat, 19 Mar 2022 12:26:13 +0800 Subject: [PATCH 2/4] feat: Add converter files for torch::max Signed-off-by: hongwei03 --- core/conversion/converters/impl/max.cpp | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/core/conversion/converters/impl/max.cpp b/core/conversion/converters/impl/max.cpp index 83f5aefd80..f9cd9f1954 100644 --- a/core/conversion/converters/impl/max.cpp +++ b/core/conversion/converters/impl/max.cpp @@ -17,17 +17,13 @@ auto max_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().patter {"aten::max.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)", [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { auto self = args[0].ITensorOrFreeze(ctx); - auto k = 1; auto dim = args[1].unwrapToInt(); - auto largest = true; auto selfDim = util::toVec(self->getDimensions()); if (dim < 0) { dim = selfDim.size() + dim; } uint32_t shiftDim = 1 << dim; - - auto TopKOperation = largest ? (nvinfer1::TopKOperation::kMAX) : (nvinfer1::TopKOperation::kMIN); - + auto TopKOperation = nvinfer1::TopKOperation::kMAX; auto new_layer = ctx->net->addTopK(*self, TopKOperation, 1, shiftDim); TORCHTRT_CHECK(new_layer, "Unable to create max layer from node: " << *n); From 0b5673be1cad3b1aad9d638e29f6641b18ab569d Mon Sep 17 00:00:00 2001 From: hongwei03 Date: Wed, 23 Mar 2022 11:48:54 +0800 Subject: [PATCH 3/4] add path of max.cpp to BUILD Signed-off-by: hongwei03 --- core/conversion/converters/BUILD | 1 + 1 file changed, 1 insertion(+) diff --git a/core/conversion/converters/BUILD b/core/conversion/converters/BUILD index 8b5338c08b..f9948eea7d 100755 --- a/core/conversion/converters/BUILD +++ b/core/conversion/converters/BUILD @@ -79,6 +79,7 @@ cc_library( "impl/squeeze.cpp", "impl/stack.cpp", "impl/topk.cpp", + "impl/max.cpp", "impl/unary.cpp", "impl/unsqueeze.cpp", ], From dd7a44e1f94f5644d23d9f3591f6b79de08f3b32 Mon Sep 17 00:00:00 2001 From: hongwei03 Date: Wed, 23 Mar 2022 14:25:33 +0800 Subject: [PATCH 4/4] feat: Add converter files for torch::max Signed-off-by: hongwei03 --- core/conversion/converters/impl/max.cpp | 36 ++++++++++++------------- 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/core/conversion/converters/impl/max.cpp b/core/conversion/converters/impl/max.cpp index f9cd9f1954..adc8d06ed0 100644 --- a/core/conversion/converters/impl/max.cpp +++ b/core/conversion/converters/impl/max.cpp @@ -14,27 +14,27 @@ namespace converters { namespace impl { namespace { auto max_registrations TORCHTRT_UNUSED = RegisterNodeConversionPatterns().pattern( - {"aten::max.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)", - [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { - auto self = args[0].ITensorOrFreeze(ctx); - auto dim = args[1].unwrapToInt(); - auto selfDim = util::toVec(self->getDimensions()); - if (dim < 0) { - dim = selfDim.size() + dim; - } - uint32_t shiftDim = 1 << dim; - auto TopKOperation = nvinfer1::TopKOperation::kMAX; - auto new_layer = ctx->net->addTopK(*self, TopKOperation, 1, shiftDim); - TORCHTRT_CHECK(new_layer, "Unable to create max layer from node: " << *n); + {"aten::max.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)", + [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool { + auto self = args[0].ITensorOrFreeze(ctx); + auto dim = args[1].unwrapToInt(); + auto selfDim = util::toVec(self->getDimensions()); + if (dim < 0) { + dim = selfDim.size() + dim; + } + uint32_t shiftDim = 1 << dim; + auto TopKOperation = nvinfer1::TopKOperation::kMAX; + auto new_layer = ctx->net->addTopK(*self, TopKOperation, 1, shiftDim); + TORCHTRT_CHECK(new_layer, "Unable to create max layer from node: " << *n); - auto out0 = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0)); - auto out1 = ctx->AssociateValueAndTensor(n->outputs()[1], new_layer->getOutput(1)); + auto out0 = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0)); + auto out1 = ctx->AssociateValueAndTensor(n->outputs()[1], new_layer->getOutput(1)); - LOG_DEBUG("Output tensor(0) shape: " << out0->getDimensions()); - LOG_DEBUG("Output tensor(1) shape: " << out1->getDimensions()); + LOG_DEBUG("Output tensor(0) shape: " << out0->getDimensions()); + LOG_DEBUG("Output tensor(1) shape: " << out1->getDimensions()); - return true; - }}); + return true; + }}); } // namespace } // namespace impl } // namespace converters