From 6a61eceb224029d9f61add4668ff40b2d377f38d Mon Sep 17 00:00:00 2001 From: Liu_Xuu Date: Thu, 22 Jul 2021 10:29:09 +0800 Subject: [PATCH] [MSLITE] add mul, slice, pool, pad ops for tensorrt --- .../tensorrt/op/elementwise_tensorrt.cc | 10 ++ .../src/delegate/tensorrt/op/pad_tensorrt.cc | 105 ++++++++++++++++ .../src/delegate/tensorrt/op/pad_tensorrt.h | 41 +++++++ .../src/delegate/tensorrt/op/pool_tensorrt.cc | 116 ++++++++++++++++++ .../src/delegate/tensorrt/op/pool_tensorrt.h | 40 ++++++ .../delegate/tensorrt/op/shuffle_tensorrt.cc | 4 +- .../delegate/tensorrt/op/slice_tensorrt.cc | 68 ++++++++++ .../src/delegate/tensorrt/op/slice_tensorrt.h | 37 ++++++ .../delegate/tensorrt/tensorrt_delegate.cc | 7 ++ .../src/delegate/tensorrt/tensorrt_subgraph.h | 3 +- .../src/delegate/tensorrt/tensorrt_utils.cc | 6 +- .../src/delegate/tensorrt/tensorrt_utils.h | 2 +- 12 files changed, 432 insertions(+), 7 deletions(-) create mode 100644 mindspore/lite/src/delegate/tensorrt/op/pad_tensorrt.cc create mode 100644 mindspore/lite/src/delegate/tensorrt/op/pad_tensorrt.h create mode 100644 mindspore/lite/src/delegate/tensorrt/op/pool_tensorrt.cc create mode 100644 mindspore/lite/src/delegate/tensorrt/op/pool_tensorrt.h create mode 100644 mindspore/lite/src/delegate/tensorrt/op/slice_tensorrt.cc create mode 100644 mindspore/lite/src/delegate/tensorrt/op/slice_tensorrt.h diff --git a/mindspore/lite/src/delegate/tensorrt/op/elementwise_tensorrt.cc b/mindspore/lite/src/delegate/tensorrt/op/elementwise_tensorrt.cc index de059c789ab..8f0f2fa2894 100644 --- a/mindspore/lite/src/delegate/tensorrt/op/elementwise_tensorrt.cc +++ b/mindspore/lite/src/delegate/tensorrt/op/elementwise_tensorrt.cc @@ -26,6 +26,7 @@ int ElementWiseTensorRT::IsSupport(const schema::Primitive *primitive, {schema::PrimitiveType_PowFusion, nvinfer1::ElementWiseOperation::kPOW}, {schema::PrimitiveType_DivFusion, nvinfer1::ElementWiseOperation::kDIV}, {schema::PrimitiveType_SubFusion, nvinfer1::ElementWiseOperation::kSUB}, + {schema::PrimitiveType_MulFusion, nvinfer1::ElementWiseOperation::kPROD}, }; auto iter_op = element_wise_ops.find(this->type_); if (iter_op != element_wise_ops.end()) { @@ -151,6 +152,15 @@ nvinfer1::ITensor *ElementWiseTensorRT::AddActivation(nvinfer1::INetworkDefiniti activation = sub_op->activation_type(); break; } + case schema::PrimitiveType_MulFusion: { + auto mul_op = op_primitive_->value_as_MulFusion(); + if (mul_op == nullptr) { + MS_LOG(ERROR) << "MulFusion convert failed."; + return nullptr; + } + activation = mul_op->activation_type(); + break; + } default: MS_LOG(DEBUG) << "no activation need for: " << op_name_; } diff --git a/mindspore/lite/src/delegate/tensorrt/op/pad_tensorrt.cc b/mindspore/lite/src/delegate/tensorrt/op/pad_tensorrt.cc new file mode 100644 index 00000000000..d5565765c98 --- /dev/null +++ b/mindspore/lite/src/delegate/tensorrt/op/pad_tensorrt.cc @@ -0,0 +1,105 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include "src/delegate/tensorrt/op/pad_tensorrt.h" +#include "src/delegate/tensorrt/tensorrt_utils.h" + +namespace mindspore::lite { +int PadTensorRT::IsSupport(const mindspore::schema::Primitive *primitive, + const std::vector &in_tensors, + const std::vector &out_tensors) { + if (in_tensors.size() != 2 && in_tensors.size() != 3) { + MS_LOG(ERROR) << "Unsupported input tensor size, size is " << in_tensors.size(); + return RET_ERROR; + } + if (out_tensors.size() != 1) { + MS_LOG(ERROR) << "Unsupported output tensor size, size is " << out_tensors.size(); + return RET_ERROR; + } + if (in_tensors_[1].Data() == nullptr) { + MS_LOG(ERROR) << "invalid pad tensor for: " << op_name_; + return RET_ERROR; + } + auto pad_primitive = this->GetPrimitive()->value_as_PadFusion(); + if (pad_primitive == nullptr) { + MS_LOG(ERROR) << "convert PadFusion failed: " << op_name_; + return RET_ERROR; + } + schema::PaddingMode padding_mode = pad_primitive->padding_mode(); + if (padding_mode != schema::PaddingMode::PaddingMode_CONSTANT) { + MS_LOG(ERROR) << "Unsupported padding mode: " << pad_primitive << ", for op: " << op_name_; + return RET_ERROR; + } + constant_value_ = pad_primitive->constant_value(); + return RET_OK; +} + +int PadTensorRT::AddInnerOp(nvinfer1::INetworkDefinition *network) { + mindspore::MSTensor &pad_tensor = in_tensors_[1]; + int element_cnt = std::accumulate(pad_tensor.Shape().begin(), pad_tensor.Shape().end(), 1, std::multiplies()); + if (element_cnt != tensorrt_in_tensors_[0]->getDimensions().nbDims * 2) { + MS_LOG(ERROR) << "pad tensor cnt is invalid. cnt: " << element_cnt + << ", input tensor dims cnt: " << tensorrt_in_tensors_[0]->getDimensions().nbDims; + return RET_ERROR; + } + // transpose: NHWC->NCHW + nvinfer1::IShuffleLayer *transpose_layer_in = NHWC2NCHW(network, *tensorrt_in_tensors_[0]); + if (transpose_layer_in == nullptr) { + MS_LOG(ERROR) << "transpose: NHWC->NCHW failed"; + return RET_ERROR; + } + transpose_layer_in->setName((op_name_ + "_transpose2NCHW").c_str()); + + // trt 6 only support 2D padding + const int *padding_data = reinterpret_cast(in_tensors_[1].Data().get()); + nvinfer1::IPaddingLayer *padding_layer = nullptr; + if (element_cnt == index_NHWC_ * 2) { + // NHWC only support pad at HW index + // 0: N_pre, 1: N_post, 2: H_pre, 3: H_post, 4: W_pre, 5: W_post, 6: C_pre, 7: C_post + if (*padding_data != 0 || *(padding_data + 1) != 0 || *(padding_data + 6) != 0 || *(padding_data + 7) != 0) { + MS_LOG(WARNING) << "tensorrt padding only support pad at HW index, unsupported padding value of: " << op_name_; + } + nvinfer1::DimsHW prePadding{*(padding_data + 2), *(padding_data + 4)}; + nvinfer1::DimsHW postPadding{*(padding_data + 3), *(padding_data + 5)}; + MS_LOG(INFO) << "prePadding: " << *(padding_data + 2) << ", " << *(padding_data + 4); + MS_LOG(INFO) << "postPadding: " << *(padding_data + 3) << ", " << *(padding_data + 5); + + padding_layer = network->addPadding(*transpose_layer_in->getOutput(0), prePadding, postPadding); + } else { + MS_LOG(ERROR) << "need check for pad_tensor dims: " << op_name_ + << ", pad_tensor ElementNum: " << pad_tensor.ElementNum(); + return RET_ERROR; + } + if (padding_layer == nullptr) { + MS_LOG(ERROR) << "add padding layer failed for " << op_name_; + return RET_ERROR; + } + padding_layer->setName(op_name_.c_str()); + + // transpose: NCHW->NHWC + nvinfer1::IShuffleLayer *transpose_layer_out = NCHW2NHWC(network, *padding_layer->getOutput(0)); + if (transpose_layer_out == nullptr) { + MS_LOG(ERROR) << "op action convert failed"; + return RET_ERROR; + } + transpose_layer_out->setName((op_name_ + "_transpose2NHWC").c_str()); + + this->AddInnerOutTensors(transpose_layer_out->getOutput(0)); + return RET_OK; +} +} // namespace mindspore::lite diff --git a/mindspore/lite/src/delegate/tensorrt/op/pad_tensorrt.h b/mindspore/lite/src/delegate/tensorrt/op/pad_tensorrt.h new file mode 100644 index 00000000000..6a7017cfe5b --- /dev/null +++ b/mindspore/lite/src/delegate/tensorrt/op/pad_tensorrt.h @@ -0,0 +1,41 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef MINDSPORE_LITE_SRC_DELEGATE_TENSORRT_OP_PAD_TENSORRT_H_ +#define MINDSPORE_LITE_SRC_DELEGATE_TENSORRT_OP_PAD_TENSORRT_H_ +#include +#include +#include "src/delegate/tensorrt/op/tensorrt_op.h" + +namespace mindspore::lite { +class PadTensorRT : public TensorRTOp { + public: + PadTensorRT(const schema::Primitive *primitive, const std::vector &in_tensors, + const std::vector &out_tensors, const std::string &name) + : TensorRTOp(primitive, in_tensors, out_tensors, name) {} + + ~PadTensorRT() override = default; + + int AddInnerOp(nvinfer1::INetworkDefinition *network) override; + + int IsSupport(const schema::Primitive *primitive, const std::vector &in_tensors, + const std::vector &out_tensors) override; + + private: + const int index_NHWC_ = 4; + float constant_value_ = 0.0f; +}; +} // namespace mindspore::lite +#endif // MINDSPORE_LITE_SRC_DELEGATE_TENSORRT_OP_PAD_TENSORRT_H_ diff --git a/mindspore/lite/src/delegate/tensorrt/op/pool_tensorrt.cc b/mindspore/lite/src/delegate/tensorrt/op/pool_tensorrt.cc new file mode 100644 index 00000000000..4263755c2fc --- /dev/null +++ b/mindspore/lite/src/delegate/tensorrt/op/pool_tensorrt.cc @@ -0,0 +1,116 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "src/delegate/tensorrt/op/pool_tensorrt.h" +#include "src/delegate/tensorrt/op/activation_tensorrt.h" +#include "src/delegate/tensorrt/tensorrt_utils.h" + +namespace mindspore::lite { +int PoolTensorRT::IsSupport(const mindspore::schema::Primitive *primitive, + const std::vector &in_tensors, + const std::vector &out_tensors) { + if (in_tensors.size() != 1) { + MS_LOG(ERROR) << "Unsupported input tensor size, size is " << in_tensors.size(); + return RET_ERROR; + } + if (out_tensors.size() != 1) { + MS_LOG(ERROR) << "Unsupported output tensor size, size is " << out_tensors.size(); + return RET_ERROR; + } + return RET_OK; +} + +int PoolTensorRT::AddInnerOp(nvinfer1::INetworkDefinition *network) { + const schema::AvgPoolFusion *pool_primitive = this->GetPrimitive()->value_as_AvgPoolFusion(); + if (pool_primitive == nullptr) { + MS_LOG(ERROR) << "convert PoolFusion failed: " << op_name_; + return RET_ERROR; + } + if (tensorrt_in_tensors_.size() != 1) { + MS_LOG(ERROR) << "invalid input tensor size: " << tensorrt_in_tensors_.size(); + return RET_ERROR; + } + // transpose: NHWC->NCHW + nvinfer1::IShuffleLayer *transpose_layer_in = NHWC2NCHW(network, *tensorrt_in_tensors_[0]); + if (transpose_layer_in == nullptr) { + MS_LOG(ERROR) << "transpose: NHWC->NCHW failed"; + return RET_ERROR; + } + transpose_layer_in->setName((op_name_ + "_transpose2NCHW").c_str()); + + // pooling layer + nvinfer1::PoolingType pooling_type = nvinfer1::PoolingType::kAVERAGE; + auto kernel_size = pool_primitive->kernel_size(); + if (kernel_size == nullptr) { + MS_LOG(ERROR) << "get kernel size failed: " << op_name_; + return RET_ERROR; + } + std::vector kernel_size_val = std::vector(kernel_size->begin(), kernel_size->end()); + nvinfer1::Dims windowSize = lite::ConvertCudaDims(kernel_size_val); + nvinfer1::IPoolingLayer *pooling_layer = + network->addPoolingNd(*transpose_layer_in->getOutput(0), pooling_type, windowSize); + if (pooling_layer == nullptr) { + MS_LOG(ERROR) << "addPoolingNd failed for TensorRT."; + return RET_ERROR; + } + AddParams(pool_primitive, pooling_layer); + pooling_layer->setName(op_name_.c_str()); + + // add activation + nvinfer1::ILayer *activation_layer = nullptr; + if (pool_primitive->activation_type() == schema::ActivationType::ActivationType_NO_ACTIVATION) { + activation_layer = pooling_layer; + } else { + activation_layer = + ActivationTensorRT::AddActivation(network, pool_primitive->activation_type(), 0, pooling_layer->getOutput(0)); + if (activation_layer == nullptr) { + MS_LOG(ERROR) << "addActivation for pool failed"; + return RET_ERROR; + } + activation_layer->setName((op_name_ + "_activation").c_str()); + } + // transpose: NCHW->NHWC + nvinfer1::IShuffleLayer *transpose_layer_out = NCHW2NHWC(network, *activation_layer->getOutput(0)); + if (transpose_layer_out == nullptr) { + MS_LOG(ERROR) << "op action convert failed"; + return RET_ERROR; + } + transpose_layer_out->setName((op_name_ + "_transpose2NHWC").c_str()); + this->AddInnerOutTensors(transpose_layer_out->getOutput(0)); + return RET_OK; +} + +void PoolTensorRT::AddParams(const schema::AvgPoolFusion *primitive, nvinfer1::IPoolingLayer *pooling_layer) { + auto stride = primitive->strides(); + std::vector stride_val = std::vector(stride->begin(), stride->end()); + nvinfer1::Dims stride_dims = ConvertCudaDims(stride_val); + pooling_layer->setStrideNd(stride_dims); + + schema::PadMode pad_mode = primitive->pad_mode(); + if (pad_mode == schema::PadMode::PadMode_SAME) { + pooling_layer->setPaddingMode(nvinfer1::PaddingMode::kSAME_UPPER); + } + auto padding = primitive->pad(); + if (padding != nullptr) { + auto padding_val = std::vector(padding->begin(), padding->end()); + nvinfer1::Dims dims{}; + dims.nbDims = 2; + dims.d[0] = padding_val[1]; + dims.d[1] = padding_val[2]; + pooling_layer->setPaddingNd(dims); + } +} +} // namespace mindspore::lite diff --git a/mindspore/lite/src/delegate/tensorrt/op/pool_tensorrt.h b/mindspore/lite/src/delegate/tensorrt/op/pool_tensorrt.h new file mode 100644 index 00000000000..9a823a6075e --- /dev/null +++ b/mindspore/lite/src/delegate/tensorrt/op/pool_tensorrt.h @@ -0,0 +1,40 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef MINDSPORE_LITE_SRC_DELEGATE_TENSORRT_OP_POOL_TENSORRT_H_ +#define MINDSPORE_LITE_SRC_DELEGATE_TENSORRT_OP_POOL_TENSORRT_H_ +#include +#include +#include "src/delegate/tensorrt/op/tensorrt_op.h" + +namespace mindspore::lite { +class PoolTensorRT : public TensorRTOp { + public: + PoolTensorRT(const schema::Primitive *primitive, const std::vector &in_tensors, + const std::vector &out_tensors, const std::string &name) + : TensorRTOp(primitive, in_tensors, out_tensors, name) {} + + ~PoolTensorRT() override = default; + + int AddInnerOp(nvinfer1::INetworkDefinition *network) override; + + int IsSupport(const schema::Primitive *primitive, const std::vector &in_tensors, + const std::vector &out_tensors) override; + + private: + void AddParams(const schema::AvgPoolFusion *primitive, nvinfer1::IPoolingLayer *pooling_layer); +}; +} // namespace mindspore::lite +#endif // MINDSPORE_LITE_SRC_DELEGATE_TENSORRT_OP_POOL_TENSORRT_H_ diff --git a/mindspore/lite/src/delegate/tensorrt/op/shuffle_tensorrt.cc b/mindspore/lite/src/delegate/tensorrt/op/shuffle_tensorrt.cc index 90b3cdd773f..d5d21cf9270 100644 --- a/mindspore/lite/src/delegate/tensorrt/op/shuffle_tensorrt.cc +++ b/mindspore/lite/src/delegate/tensorrt/op/shuffle_tensorrt.cc @@ -190,8 +190,8 @@ int ShuffleTensorRT::AddReshapeOp(nvinfer1::IShuffleLayer *shuffle_layer) { MS_LOG(ERROR) << "AddReshapeOp size of in tensort needs check: " << in_tensors_.size(); return RET_ERROR; } - mindspore::MSTensor shape_tensor = in_tensors_[1]; - nvinfer1::Dims reshape_dims = ConvertCudaDims(shape_tensor.MutableData(), shape_tensor.ElementNum()); + mindspore::MSTensor &shape_tensor = in_tensors_[1]; + nvinfer1::Dims reshape_dims = ConvertCudaDims(shape_tensor.Data().get(), shape_tensor.ElementNum()); int ret = InferReshapeDims(tensorrt_in_tensors_[0]->getDimensions(), &reshape_dims); if (ret != RET_OK) { MS_LOG(ERROR) << "invalid dims for reshape " << op_name_; diff --git a/mindspore/lite/src/delegate/tensorrt/op/slice_tensorrt.cc b/mindspore/lite/src/delegate/tensorrt/op/slice_tensorrt.cc new file mode 100644 index 00000000000..4946fa0b501 --- /dev/null +++ b/mindspore/lite/src/delegate/tensorrt/op/slice_tensorrt.cc @@ -0,0 +1,68 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "src/delegate/tensorrt/op/slice_tensorrt.h" +#include "src/delegate/tensorrt/tensorrt_utils.h" + +namespace mindspore::lite { +int SliceTensorRT::IsSupport(const mindspore::schema::Primitive *primitive, + const std::vector &in_tensors, + const std::vector &out_tensors) { + if (in_tensors.size() != 4 && in_tensors.size() != 5) { + MS_LOG(ERROR) << "Unsupported input tensor size, size is " << in_tensors.size(); + return RET_ERROR; + } + if (out_tensors.size() != 1) { + MS_LOG(ERROR) << "Unsupported output tensor size, size is " << out_tensors.size(); + return RET_ERROR; + } + if (in_tensors_[1].Data() == nullptr) { + MS_LOG(ERROR) << "invalid pad tensor for: " << op_name_; + return RET_ERROR; + } + return RET_OK; +} + +int SliceTensorRT::AddInnerOp(nvinfer1::INetworkDefinition *network) { + auto slice_primitive = this->GetPrimitive()->value_as_StridedSlice(); + if (slice_primitive == nullptr) { + MS_LOG(ERROR) << "convert StridedSlice failed: " << op_name_; + return RET_ERROR; + } + const mindspore::MSTensor &begin = in_tensors_[1]; + // mindspore::MSTensor &end = in_tensors_[2]; + const mindspore::MSTensor &stride = in_tensors_[3]; + + nvinfer1::Dims start_dims = lite::ConvertCudaDims(begin.Data().get(), begin.ElementNum()); + nvinfer1::Dims size_dims = lite::ConvertCudaDims(out_tensors_[0].Shape()); + nvinfer1::Dims stride_dims = lite::ConvertCudaDims(stride.Data().get(), stride.ElementNum()); + + nvinfer1::ISliceLayer *slice_layer = network->addSlice(*tensorrt_in_tensors_[0], start_dims, size_dims, stride_dims); + if (slice_layer == nullptr) { + MS_LOG(ERROR) << "add Slice op failed for TensorRT: " << op_name_; + return RET_ERROR; + } + slice_layer->setName(op_name_.c_str()); + nvinfer1::ITensor *out_tensor = slice_layer->getOutput(0); + if (out_tensor == nullptr) { + MS_LOG(ERROR) << "output tensor create failed"; + return RET_ERROR; + } + out_tensor->setName(out_tensors_[0].Name().c_str()); + this->AddInnerOutTensors(out_tensor); + return RET_OK; +} +} // namespace mindspore::lite diff --git a/mindspore/lite/src/delegate/tensorrt/op/slice_tensorrt.h b/mindspore/lite/src/delegate/tensorrt/op/slice_tensorrt.h new file mode 100644 index 00000000000..7bedfaf2adf --- /dev/null +++ b/mindspore/lite/src/delegate/tensorrt/op/slice_tensorrt.h @@ -0,0 +1,37 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef MINDSPORE_LITE_SRC_DELEGATE_TENSORRT_OP_SLICE_TENSORRT_H_ +#define MINDSPORE_LITE_SRC_DELEGATE_TENSORRT_OP_SLICE_TENSORRT_H_ +#include +#include +#include "src/delegate/tensorrt/op/tensorrt_op.h" + +namespace mindspore::lite { +class SliceTensorRT : public TensorRTOp { + public: + SliceTensorRT(const schema::Primitive *primitive, const std::vector &in_tensors, + const std::vector &out_tensors, const std::string &name) + : TensorRTOp(primitive, in_tensors, out_tensors, name) {} + + ~SliceTensorRT() override = default; + + int AddInnerOp(nvinfer1::INetworkDefinition *network) override; + + int IsSupport(const schema::Primitive *primitive, const std::vector &in_tensors, + const std::vector &out_tensors) override; +}; +} // namespace mindspore::lite +#endif // MINDSPORE_LITE_SRC_DELEGATE_TENSORRT_OP_SLICE_TENSORRT_H_ diff --git a/mindspore/lite/src/delegate/tensorrt/tensorrt_delegate.cc b/mindspore/lite/src/delegate/tensorrt/tensorrt_delegate.cc index b055d2b1b3a..4965a6c1059 100644 --- a/mindspore/lite/src/delegate/tensorrt/tensorrt_delegate.cc +++ b/mindspore/lite/src/delegate/tensorrt/tensorrt_delegate.cc @@ -33,6 +33,9 @@ #include "src/delegate/tensorrt/op/unary_tensorrt.h" #include "src/delegate/tensorrt/op/matmul_tensorrt.h" #include "src/delegate/tensorrt/op/scale_tensorrt.h" +#include "src/delegate/tensorrt/op/slice_tensorrt.h" +#include "src/delegate/tensorrt/op/pool_tensorrt.h" +#include "src/delegate/tensorrt/op/pad_tensorrt.h" namespace mindspore::lite { bool IsHardwareSupport() { @@ -76,12 +79,16 @@ int TensorRTDelegate::Init() { {schema::PrimitiveType_DivFusion, GetTensorRTOp}, {schema::PrimitiveType_PowFusion, GetTensorRTOp}, {schema::PrimitiveType_AddFusion, GetTensorRTOp}, + {schema::PrimitiveType_MulFusion, GetTensorRTOp}, {schema::PrimitiveType_Eltwise, GetTensorRTOp}, {schema::PrimitiveType_Transpose, GetTensorRTOp}, {schema::PrimitiveType_ReduceFusion, GetTensorRTOp}, {schema::PrimitiveType_Sqrt, GetTensorRTOp}, {schema::PrimitiveType_MatMul, GetTensorRTOp}, {schema::PrimitiveType_ScaleFusion, GetTensorRTOp}, + {schema::PrimitiveType_StridedSlice, GetTensorRTOp}, + {schema::PrimitiveType_AvgPoolFusion, GetTensorRTOp}, + {schema::PrimitiveType_PadFusion, GetTensorRTOp}, }; return RET_OK; } diff --git a/mindspore/lite/src/delegate/tensorrt/tensorrt_subgraph.h b/mindspore/lite/src/delegate/tensorrt/tensorrt_subgraph.h index 4cdaea99d3a..80ed386df7d 100644 --- a/mindspore/lite/src/delegate/tensorrt/tensorrt_subgraph.h +++ b/mindspore/lite/src/delegate/tensorrt/tensorrt_subgraph.h @@ -37,7 +37,8 @@ class TensorRTSubGraph : public kernel::Kernel { trt_specific_weight_nodes_ = { schema::PrimitiveType_Conv2DFusion, schema::PrimitiveType_ReduceFusion, schema::PrimitiveType_Transpose, schema::PrimitiveType_Gather, schema::PrimitiveType_Reshape, schema::PrimitiveType_PowFusion, - schema::PrimitiveType_DivFusion, schema::PrimitiveType_MatMul, schema::PrimitiveType_ScaleFusion}; + schema::PrimitiveType_DivFusion, schema::PrimitiveType_MatMul, schema::PrimitiveType_ScaleFusion, + schema::PrimitiveType_MulFusion, schema::PrimitiveType_StridedSlice, schema::PrimitiveType_PadFusion}; } ~TensorRTSubGraph() override; diff --git a/mindspore/lite/src/delegate/tensorrt/tensorrt_utils.cc b/mindspore/lite/src/delegate/tensorrt/tensorrt_utils.cc index c44f2125173..230c35c829d 100644 --- a/mindspore/lite/src/delegate/tensorrt/tensorrt_utils.cc +++ b/mindspore/lite/src/delegate/tensorrt/tensorrt_utils.cc @@ -37,11 +37,11 @@ nvinfer1::Dims ConvertCudaDims(int data, size_t size) { return dims; } -nvinfer1::Dims ConvertCudaDims(void *data, size_t size) { +nvinfer1::Dims ConvertCudaDims(const void *data, int64_t size) { nvinfer1::Dims dims{}; dims.nbDims = size; - int *dims_data = reinterpret_cast(data); - for (size_t i = 0; i < size; i++) { + const int *dims_data = reinterpret_cast(data); + for (int i = 0; i < size; i++) { dims.d[i] = *(dims_data + i); } return dims; diff --git a/mindspore/lite/src/delegate/tensorrt/tensorrt_utils.h b/mindspore/lite/src/delegate/tensorrt/tensorrt_utils.h index 0355835148c..aacaed8534a 100644 --- a/mindspore/lite/src/delegate/tensorrt/tensorrt_utils.h +++ b/mindspore/lite/src/delegate/tensorrt/tensorrt_utils.h @@ -34,7 +34,7 @@ struct ActivationParams { nvinfer1::Dims ConvertCudaDims(const std::vector &shape); // Convert Tensor data to Cuda dims. -nvinfer1::Dims ConvertCudaDims(void *data, size_t size); +nvinfer1::Dims ConvertCudaDims(const void *data, int64_t size); nvinfer1::Dims ConvertCudaDims(int data, size_t size);