wehub-resource-sync
106 lines
3.6 KiB
C++
106 lines
3.6 KiB
C++
//
|
|
// QNNSoftmax.cpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on b'2025/04/10'.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "QNNSoftmax.hpp"
|
|
|
|
namespace MNN {
|
|
namespace QNN {
|
|
#ifdef ENABLE_QNN_ONLINE_FINALIZE
|
|
|
|
ErrorCode QNNSoftmax::onEncode(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
|
|
this->createParamScalar("axis", (uint32_t) inputs[0]->dimensions() - 1);
|
|
|
|
if (mAxis != inputs[0]->dimensions() - 1) {
|
|
return this->onEncodePermute(inputs, outputs);
|
|
}
|
|
|
|
mNodeType = "Softmax";
|
|
this->addNodeCommon(inputs, outputs);
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
ErrorCode QNNSoftmax::onEncodePermute(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
|
|
// Create resources.
|
|
int dim = inputs[0]->dimensions();
|
|
std::vector<uint32_t> permData(dim, 0);
|
|
for (int i = 0; i < dim; i++) {
|
|
if (i == mAxis) {
|
|
permData[i] = dim - 1;
|
|
continue;
|
|
}
|
|
if (i == dim - 1) {
|
|
permData[i] = mAxis;
|
|
continue;
|
|
}
|
|
permData[i] = i;
|
|
}
|
|
|
|
std::vector<uint32_t> shapeStageTensor = getNHWCShape(inputs[0]);
|
|
{
|
|
uint32_t temp = shapeStageTensor[dim - 1];
|
|
shapeStageTensor[dim - 1] = shapeStageTensor[mAxis];
|
|
shapeStageTensor[mAxis] = temp;
|
|
}
|
|
|
|
this->createParamTensor("perm", QNN_DATATYPE_UINT_32, {(uint32_t) dim}, (void *) permData.data(), "before"); // mParamTensorWrappers[0], permBefore
|
|
this->createParamTensor("perm", QNN_DATATYPE_UINT_32, {(uint32_t) dim}, (void *) permData.data(), "after"); // mParamTensorWrappers[1], permAfter
|
|
|
|
Qnn_DataType_t qnnDataType = mBackend->getNativeTensor(inputs[0])->v1.dataType;
|
|
this->createStageTensor("stageInput", qnnDataType, shapeStageTensor, inputs[0]); // mTempTensorWrappers[0], stage input
|
|
this->createStageTensor("stageOutput", qnnDataType, shapeStageTensor, outputs[0]); // mTempTensorWrappers[1], stage output
|
|
|
|
// Add nodes.
|
|
{
|
|
this->addNodeCommonPermute("PermBefore",
|
|
*(mBackend->getNativeTensor(inputs[0])),
|
|
*(mParamTensorWrappers[0]->getNativeParam()),
|
|
*(mTempTensorWrappers[0]->getNativeTensor()));
|
|
}
|
|
|
|
{
|
|
CLEAR_BEFORE_ADDING_NODE;
|
|
|
|
std::string name = mNodeName + "_Softmax";
|
|
mNodeType = "Softmax";
|
|
mInputs.push_back(*(mTempTensorWrappers[0]->getNativeTensor()));
|
|
mParams.push_back(*(mParamScalarWrappers[0]->getNativeParam()));
|
|
mOutputs.push_back(*(mTempTensorWrappers[1]->getNativeTensor()));
|
|
|
|
mBackend->addNodeToGraph(mOpConfigVersion, name.c_str(), mPackageName.c_str(), mNodeType.c_str(), mParams, mInputs, mOutputs);
|
|
}
|
|
|
|
{
|
|
this->addNodeCommonPermute("PermAfter",
|
|
*(mTempTensorWrappers[1]->getNativeTensor()),
|
|
*(mParamTensorWrappers[1]->getNativeParam()),
|
|
*(mBackend->getNativeTensor(outputs[0])));
|
|
}
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
class QNNSoftmaxCreator : public QnnBackend::Creator {
|
|
public:
|
|
virtual QNNCommonExecution * onCreate(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs, const MNN::Op* op,
|
|
Backend* backend) const override {
|
|
const auto softmaxParam = op->main_as_Axis();
|
|
int axis = softmaxParam->axis();
|
|
if (axis < 0) {
|
|
axis = inputs[0]->dimensions() + axis;
|
|
}
|
|
|
|
return new QNNSoftmax(backend, op, axis);
|
|
}
|
|
};
|
|
|
|
REGISTER_QNN_OP_CREATOR(QNNSoftmaxCreator, OpType_Softmax)
|
|
#endif
|
|
} // end namespace QNN
|
|
} // end namespace MNN
|