wehub-resource-sync
47 lines
1.6 KiB
C++
47 lines
1.6 KiB
C++
#include "QNNTopKV2.hpp"
|
|
|
|
namespace MNN {
|
|
namespace QNN {
|
|
#ifdef ENABLE_QNN_ONLINE_FINALIZE
|
|
|
|
ErrorCode QNNTopKV2::onEncode(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
|
|
MNN_ASSERT(inputs.size() == 1 && outputs.size() == 1);
|
|
auto topk = mOp->main_as_TopKV2();
|
|
int k = inputs[0]->shape()[inputs[0]->dimensions() - 1];
|
|
bool largest = topk->largest();
|
|
|
|
this->createParamScalar("k", (uint32_t)k);
|
|
this->createParamScalar("largest", (bool)largest);
|
|
|
|
mNodeType = "TopK";
|
|
{
|
|
CLEAR_BEFORE_ADDING_NODE;
|
|
|
|
mNodeType = "TopK";
|
|
|
|
mInputs.push_back(*(mBackend->getNativeTensor(inputs[0])));
|
|
mParams.push_back(*(mParamScalarWrappers[0]->getNativeParam())); // k
|
|
mParams.push_back(*(mParamScalarWrappers[1]->getNativeParam())); // largest
|
|
mOutputs.push_back(*(mBackend->getNativeTensor(outputs[0])));
|
|
mOutputs.push_back(*(mBackend->getNativeTensor(outputs[1])));
|
|
|
|
mBackend->addNodeToGraph(mOpConfigVersion, mNodeName.c_str(), mPackageName.c_str(), mNodeType.c_str(), mParams,
|
|
mInputs, mOutputs);
|
|
}
|
|
|
|
return NO_ERROR;
|
|
}
|
|
|
|
class QNNTopKV2Creator : 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 {
|
|
return new QNNTopKV2(backend, op);
|
|
}
|
|
};
|
|
|
|
REGISTER_QNN_OP_CREATOR(QNNTopKV2Creator, OpType_TopKV2)
|
|
#endif
|
|
} // end namespace QNN
|
|
} // end namespace MNN
|