wehub-resource-sync
123 lines
5.0 KiB
C++
123 lines
5.0 KiB
C++
//
|
|
// NPUStridedSlice.cpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2019/09/07.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "NPUStridedSlice.hpp"
|
|
#include "NPUBackend.hpp"
|
|
|
|
using namespace std;
|
|
|
|
namespace MNN {
|
|
|
|
NPUStridedSlice::NPUStridedSlice(Backend *b, const Op *op, const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) : MNN::NPUCommonExecution(b,op)
|
|
{
|
|
isConst1 = TensorUtils::getDescribe(inputs[1])->usage==Tensor::InsideDescribe::Usage::CONSTANT;
|
|
isConst2 = TensorUtils::getDescribe(inputs[2])->usage==Tensor::InsideDescribe::Usage::CONSTANT;
|
|
isConst3 = false;
|
|
isConst4 = false;
|
|
auto opName = mOp->name()->str();
|
|
Tensor *begin = inputs[1];
|
|
Tensor *end = inputs[2];
|
|
if(isConst1) {
|
|
mConst_b = hiai::op::Const(opName + "_b_const");
|
|
ge::TensorDesc fdesc(ge::Shape({begin->elementSize()}), ge::DT_INT32);
|
|
ge::TensorPtr filter = std::make_shared<ge::Tensor>();
|
|
filter->SetTensorDesc(fdesc);
|
|
filter->SetData((uint8_t*)begin->host<int32_t>(), begin->elementSize()*sizeof(int32_t));
|
|
mConst_b.set_attr_value(filter);
|
|
}
|
|
if(isConst2) {
|
|
mConst_e = hiai::op::Const(opName + "_e_const");
|
|
ge::TensorDesc fdesc(ge::Shape({end->elementSize()}), ge::DT_INT32);
|
|
ge::TensorPtr filter = std::make_shared<ge::Tensor>();
|
|
filter->SetTensorDesc(fdesc);
|
|
filter->SetData((uint8_t *)end->host<int32_t>(), end->elementSize()*sizeof(int32_t));
|
|
mConst_e.set_attr_value(filter);
|
|
}
|
|
auto parameter = mOp->main_as_StridedSliceParam();
|
|
beginMask = convertMask(begin, parameter->beginMask(),1);
|
|
endMask = convertMask(begin, parameter->endMask(),1);
|
|
ellipsisMask = parameter->ellipsisMask(); //框架未使用
|
|
newAxisMask = parameter->newAxisMask();
|
|
shrinkAxisMask = convertMask(begin, parameter->shrinkAxisMask());
|
|
}
|
|
|
|
ErrorCode NPUStridedSlice::onResize(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
|
|
mNpuBackend->setNetworkInput(inputs, mOp);
|
|
|
|
auto opName = mOp->name()->str();
|
|
Tensor *axis = nullptr;
|
|
Tensor *strides = nullptr;
|
|
if (inputs.size() > 3) {
|
|
axis = inputs[3];
|
|
isConst3 = TensorUtils::getDescribe(inputs[3])->usage==Tensor::InsideDescribe::Usage::CONSTANT;
|
|
}
|
|
if (inputs.size() > 4) {
|
|
strides = inputs[4];
|
|
isConst4 = TensorUtils::getDescribe(inputs[4])->usage==Tensor::InsideDescribe::Usage::CONSTANT;
|
|
}
|
|
if (isConst3) {
|
|
vector<int32_t> axisdims;
|
|
vector<int32_t> axisdims1;
|
|
for (int32_t i = 0; i < axis->elementSize(); i++) {
|
|
axisdims.push_back(i);
|
|
if (count(axisdims1.begin(), axisdims1.end(), axis->host<int32_t>()[i]) == 0) {
|
|
axisdims1.push_back(axis->host<int32_t>()[i]);
|
|
}
|
|
}
|
|
mConst_a = hiai::op::Const(opName + "_a_const");
|
|
ge::TensorDesc fdesc(ge::Shape({axis->elementSize()}), ge::DT_INT32);
|
|
ge::TensorPtr filter = std::make_shared<ge::Tensor>();
|
|
filter->SetTensorDesc(fdesc);
|
|
if (axisdims1.size() != axisdims.size() || (axisdims.size() == 1 && axisdims1[0] == 1)) {
|
|
filter->SetData((uint8_t*)axisdims.data(), axis->elementSize()*sizeof(int32_t));
|
|
} else {
|
|
filter->SetData((uint8_t*)axisdims1.data(), axis->elementSize()*sizeof(int32_t));
|
|
}
|
|
mConst_a.set_attr_value(filter);
|
|
}
|
|
if (isConst4) {
|
|
mConst_s = hiai::op::Const(opName + "_s_const");
|
|
ge::TensorDesc fdesc(ge::Shape({strides->elementSize()}), ge::DT_INT32);
|
|
ge::TensorPtr filter = std::make_shared<ge::Tensor>();
|
|
filter->SetTensorDesc(fdesc);
|
|
filter->SetData((uint8_t*)strides->host<int32_t>(), strides->elementSize()*sizeof(int32_t));
|
|
mConst_s.set_attr_value(filter);
|
|
} else {
|
|
vector<int32_t> axisdims;
|
|
for (int32_t i = 0; i < axis->elementSize(); i++) {
|
|
axisdims.push_back(1);
|
|
}
|
|
mConst_s = hiai::op::Const(opName + "_s_const");
|
|
ge::TensorDesc fdesc(ge::Shape({axis->elementSize()}), ge::DT_INT32);
|
|
ge::TensorPtr filter = std::make_shared<ge::Tensor>();
|
|
filter->SetTensorDesc(fdesc);
|
|
filter->SetData((uint8_t*)axisdims.data(), axis->elementSize()*sizeof(int32_t));
|
|
mConst_s.set_attr_value(filter);
|
|
}
|
|
shared_ptr<hiai::op::StridedSliceV2> stride_slice(new hiai::op::StridedSliceV2(opName));
|
|
|
|
auto inputIndex = mOp->inputIndexes()->data()[0];
|
|
auto iops = mNpuBackend->mGrapMap[inputIndex]; // x
|
|
auto xOp = iops.back().first;
|
|
(*stride_slice)
|
|
.set_input_x(*xOp.get())
|
|
.set_input_begin(mConst_b)
|
|
.set_input_end(mConst_e);
|
|
if (isConst3) {
|
|
(*stride_slice).set_input_axes(mConst_a);
|
|
}
|
|
(*stride_slice).set_input_strides(mConst_s);
|
|
mNpuBackend->setOutputOps(mOp, {stride_slice}, outputs);
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
NPUCreatorRegister<TypedCreator<NPUStridedSlice>> __stride_slice_op(OpType_StridedSlice);
|
|
|
|
} // namespace MNN
|