wehub-resource-sync
36 lines
976 B
C++
36 lines
976 B
C++
//
|
|
// NPUSoftmax.cpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2019/09/07.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "NPUSoftmax.hpp"
|
|
#include "NPUBackend.hpp"
|
|
|
|
using namespace std;
|
|
|
|
namespace MNN {
|
|
|
|
NPUSoftmax::NPUSoftmax(Backend *b, const Op *op, const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) : MNN::NPUCommonExecution(b,op) {}
|
|
|
|
ErrorCode NPUSoftmax::onResize(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
|
|
mNpuBackend->setNetworkInput(inputs, mOp);
|
|
|
|
auto opName = mOp->name()->str();
|
|
auto param = mOp->main_as_Axis();
|
|
|
|
auto xOp = mNpuBackend->getInputOps(mOp);
|
|
shared_ptr<hiai::op::Softmax> softmax(new hiai::op::Softmax(opName));
|
|
|
|
(*softmax).set_input_x(*xOp.get()).set_attr_axis(param->axis());
|
|
|
|
mNpuBackend->setOutputOps(mOp, {softmax}, outputs);
|
|
return NO_ERROR;
|
|
}
|
|
|
|
NPUCreatorRegister<TypedCreator<NPUSoftmax>> __softmax_op(OpType_Softmax);
|
|
|
|
} // namespace MNN
|