MNN/express/module/StaticModule.hpp

67 lines
2.3 KiB
C++

//
// StaticModule.hpp
// MNN
//
// Created by MNN on 2020/09/10.
// Copyright © 2018, Alibaba Group Holding Limited
//
#ifndef StaticModule_hpp
#define StaticModule_hpp
#include <MNN/expr/Module.hpp>
#include <MNN/expr/Executor.hpp>
#include "core/Schedule.hpp"
#include "core/Session.hpp"
namespace MNN {
class Session;
class Backend;
struct BufferStorage;
namespace Express {
class StaticModule : public Module {
public:
StaticModule(std::vector<int> inputs, std::vector<int> outputs, std::vector<std::shared_ptr<BufferStorage>>&& buffer, Schedule::ScheduleInfo&& scheduleInfo, std::shared_ptr<Schedule::ScheduleInfo> sharedConst, Session::ModeGroup&& mode, std::shared_ptr<Executor::RuntimeManager> rt, const Module::Config& config);
virtual ~ StaticModule();
virtual std::vector<Express::VARP> onForward(const std::vector<Express::VARP>& inputs) override;
virtual void onClearCache() override;
virtual int onOptimize(Interpreter::SessionMode stage) override;
const Session* getSession() const { return mSession.get(); }
private:
ErrorCode _resize(const std::vector<Express::VARP>& inputs);
ErrorCode _execute();
StaticModule() = default;
void resetInputOutputs();
Module* clone(CloneContext* ctx) const override;
struct Resource {
std::vector<int> mInputs;
std::vector<int> mOutputs;
int mOutputNumbers;
Backend::Info mBnInfo;
BackendConfig mBnConfig;
// First: outputIndex, Second: outputTensor Index
std::vector<int> mOutputFromTensor;
// First: outputIndex, Second: input var index
std::vector<std::pair<int, int>> mOutputFromInput;
bool mUseContentInputs = false;
std::shared_ptr<Schedule::ScheduleInfo> mSharedConst;
Session::ModeGroup mModes;
std::vector<std::shared_ptr<BufferStorage>> mBuffer;
std::vector<bool> mInputNeedCPU;
};
std::shared_ptr<Session> mSession;
std::vector<Tensor*> mInputTensors;
std::vector<std::pair<Tensor*, MNNForwardType>> mPrevInputTensor;
std::vector<Tensor*> mOutputTensors;
std::shared_ptr<Resource> mResource;
bool mShapeInferSeperate = false;
std::vector<MNN::Express::VARP> mOutputVars;
std::shared_ptr<MNN::Express::Executor::RuntimeManager> mRuntimeManager;
};
}
}
#endif