wehub-resource-sync
1858 lines
74 KiB
C++
1858 lines
74 KiB
C++
#include <set>
|
|
#include <map>
|
|
#include <queue>
|
|
#include <fstream>
|
|
#include <sstream>
|
|
#include "flatbuffers/flexbuffers.h"
|
|
#include <MNN/AutoTime.hpp>
|
|
#include <MNN/Interpreter.hpp>
|
|
#include <MNN/expr/Module.hpp>
|
|
#include <MNN/expr/ExprCreator.hpp>
|
|
#include <MNN/expr/ExecutorScope.hpp>
|
|
#include "core/MNNFileUtils.h"
|
|
#include "shape/SizeComputer.hpp"
|
|
#include "core/OpCommonUtils.hpp"
|
|
#include "core/Schedule.hpp"
|
|
#include "rapidjson/document.h"
|
|
#include <rapidjson/prettywriter.h>
|
|
#include "utils/InitNet.hpp"
|
|
#include "core/Command.hpp"
|
|
#include "geometry/GeometryComputer.hpp"
|
|
#include "geometry/GeometryComputerUtils.hpp"
|
|
|
|
using namespace MNN;
|
|
static bool gNeedOffline = false;
|
|
static int gMaxKVSize = 0;
|
|
static std::string gNPUName = "CoreML";
|
|
static std::string gOfflieSrc;
|
|
static std::string gOfflieDst;
|
|
static std::string gGraphName = "graph";
|
|
static std::string gCacheDir = "res";
|
|
static MNNForwardType gNPUType = MNN_FORWARD_NN;
|
|
static bool initConstTensorsNoAlloc(std::vector<std::shared_ptr<Tensor>>& tensors, const Net* net) {
|
|
bool valid = true;
|
|
tensors.resize(net->tensorName()->size());
|
|
// Set up const
|
|
for (int opIndex = 0; opIndex < net->oplists()->size(); ++opIndex) {
|
|
auto op = net->oplists()->GetAs<Op>(opIndex);
|
|
if (OpType_Const == op->type() || OpType_TrainableParam == op->type()) {
|
|
MNN_ASSERT(nullptr != op->outputIndexes());
|
|
auto index = op->outputIndexes()->data()[0];
|
|
tensors[index].reset(new Tensor);
|
|
TensorUtils::getDescribe(tensors[index].get())->index = index;
|
|
auto output = tensors[index].get();
|
|
if (op->type() == OpType_TrainableParam) {
|
|
TensorUtils::getDescribe(output)->usage = Tensor::InsideDescribe::TRAINABLE;
|
|
}
|
|
TensorUtils::getDescribe(output)->usage = Tensor::InsideDescribe::CONSTANT;
|
|
TensorUtils::getDescribe(output)->isMutable = false;
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
struct SubModuleInfo {
|
|
std::vector<int> opList;
|
|
std::vector<int> inputs;;
|
|
std::vector<int> outputs;
|
|
std::vector<uint8_t> tensorMask;
|
|
bool isBreak = false;
|
|
};
|
|
struct SubModuleIO {
|
|
std::vector<MNN::Express::VARP> inputs;
|
|
std::vector<MNN::Express::VARP> outputs;
|
|
std::vector<std::vector<int>> kvcache;
|
|
int seqLen = 0;
|
|
};
|
|
static void _computeTensorMask(SubModuleInfo& m, const Net* net) {
|
|
/**Compute All SubModule's inputs and outputs*/
|
|
// 0: not use, 1: input, 2: output, 3: mid, 4: valid output
|
|
m.tensorMask = std::vector<uint8_t>(net->tensorName()->size(), 0);
|
|
auto& tensorMask = m.tensorMask;
|
|
for (auto opIndex : m.opList) {
|
|
auto op = net->oplists()->GetAs<Op>(opIndex);
|
|
if (nullptr != op->inputIndexes()) {
|
|
for (int v=0; v<op->inputIndexes()->size(); ++v) {
|
|
auto index = op->inputIndexes()->data()[v];
|
|
tensorMask[index] = tensorMask[index] | 1;
|
|
}
|
|
}
|
|
if (nullptr != op->outputIndexes()) {
|
|
for (int v=0; v<op->outputIndexes()->size(); ++v) {
|
|
auto index = op->outputIndexes()->data()[v];
|
|
tensorMask[index] = tensorMask[index] | 2;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
static bool _npuSupportOp(const Op* op) {
|
|
if (gMaxKVSize > 0) {
|
|
return true;
|
|
}
|
|
if (op->type() == OpType_Attention) {
|
|
auto attn = op->main_as_AttentionParam();
|
|
if (nullptr != attn && attn->kv_cache()) {
|
|
return false;
|
|
}
|
|
}
|
|
if (op->type() == OpType_LinearAttention) {
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
static bool isBreakOp(const Op* op) {
|
|
bool isWhileControlflow = false;
|
|
if (op->type() == OpType_While && op->main_as_WhileParam() != nullptr) {
|
|
isWhileControlflow = true;
|
|
}
|
|
if (op->type() == OpType_If || isWhileControlflow || op->type() == OpType_Where || op->type() == OpType_Segment || op->type() == OpType_Unique || op->type() == OpType_NonMaxSuppressionV2) {
|
|
return true;
|
|
}
|
|
if (!_npuSupportOp(op)) {
|
|
return true;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
static std::vector<int> _collectNeededOps(const MNN::Net* net, const std::set<int>& inputIndexes, const std::set<int>& outputIndexes) {
|
|
// 0: not set, 1: output, 2:input
|
|
std::vector<int> tensorMask(net->tensorName()->size());
|
|
::memset(tensorMask.data(), 0, tensorMask.size() * sizeof(int));
|
|
|
|
// 0: use, 1: no use
|
|
std::vector<int> opMask(net->oplists()->size());
|
|
::memset(opMask.data(), 0, opMask.size() * sizeof(int));
|
|
|
|
// Set Initial Status
|
|
for (auto v : outputIndexes) {
|
|
tensorMask[v] = 1;
|
|
}
|
|
for (auto v : inputIndexes) {
|
|
// If both input/output, set as input
|
|
tensorMask[v] = 2;
|
|
}
|
|
bool change = false;
|
|
do {
|
|
change = false;
|
|
for (int i=0; i<opMask.size(); ++i) {
|
|
if (opMask[i] > 0) {
|
|
continue;
|
|
}
|
|
auto op = net->oplists()->GetAs<Op>(i);
|
|
if (nullptr != op->outputIndexes()) {
|
|
for (int j=0; j<op->outputIndexes()->size(); ++j) {
|
|
auto index = op->outputIndexes()->data()[j];
|
|
if (tensorMask[index] == 1) {
|
|
opMask[i] = 1;
|
|
change = true;
|
|
}
|
|
}
|
|
}
|
|
if (nullptr != op->inputIndexes() && opMask[i]) {
|
|
for (int j=0; j<op->inputIndexes()->size(); ++j) {
|
|
auto index = op->inputIndexes()->data()[j];
|
|
if (tensorMask[index] != 2) {
|
|
tensorMask[index] = 1;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
} while (change);
|
|
|
|
std::vector<int> ops;
|
|
for (int i=0; i<opMask.size(); ++i) {
|
|
if (opMask[i] > 0) {
|
|
auto op = net->oplists()->GetAs<Op>(i);
|
|
if (needComputeOp(op)) {
|
|
ops.emplace_back(i);
|
|
continue;
|
|
}
|
|
}
|
|
}
|
|
return ops;
|
|
}
|
|
|
|
// Global set to store extra break op indexes (ops between attention_mask and Attention)
|
|
static std::set<int> gExtraBreakOpIndexes;
|
|
|
|
// Find ops between attention_mask input and Attention/LinearAttention ops that should also be break ops
|
|
static std::set<int> _findMaskToAttentionOps(const Net* net, const std::set<int>& inputIndexes,
|
|
const std::set<int>& outputIndexes) {
|
|
std::set<int> extraBreakOps;
|
|
if (net->tensorName() == nullptr || net->oplists() == nullptr) {
|
|
return extraBreakOps;
|
|
}
|
|
|
|
// 1. Find tensor indexes corresponding to attention_mask inputs
|
|
std::set<int> maskTensorIndexes;
|
|
for (auto idx : inputIndexes) {
|
|
auto name = net->tensorName()->GetAsString(idx)->str();
|
|
if (name.find("attention_mask") != std::string::npos) {
|
|
maskTensorIndexes.insert(idx);
|
|
}
|
|
}
|
|
if (maskTensorIndexes.empty()) {
|
|
return extraBreakOps;
|
|
}
|
|
|
|
// 2. Collect all needed ops
|
|
auto selectOps = _collectNeededOps(net, inputIndexes, outputIndexes);
|
|
|
|
// 3. Build tensor -> producer op index mapping
|
|
std::map<int, int> tensorProducer;
|
|
for (auto opIdx : selectOps) {
|
|
auto op = net->oplists()->GetAs<Op>(opIdx);
|
|
if (op->outputIndexes() != nullptr) {
|
|
for (int j = 0; j < op->outputIndexes()->size(); ++j) {
|
|
tensorProducer[op->outputIndexes()->data()[j]] = opIdx;
|
|
}
|
|
}
|
|
}
|
|
|
|
// 4. Forward propagation: find all tensors that depend on attention_mask
|
|
// Stop propagation at Attention/LinearAttention ops
|
|
std::set<int> maskDependentTensors = maskTensorIndexes;
|
|
bool changed = true;
|
|
while (changed) {
|
|
changed = false;
|
|
for (auto opIdx : selectOps) {
|
|
auto op = net->oplists()->GetAs<Op>(opIdx);
|
|
if (op->inputIndexes() == nullptr || op->outputIndexes() == nullptr)
|
|
continue;
|
|
// Don't propagate through Attention/LinearAttention ops
|
|
if (op->type() == OpType_Attention || op->type() == OpType_LinearAttention)
|
|
continue;
|
|
|
|
bool dependsOnMask = false;
|
|
for (int j = 0; j < op->inputIndexes()->size(); ++j) {
|
|
if (maskDependentTensors.count(op->inputIndexes()->data()[j])) {
|
|
dependsOnMask = true;
|
|
break;
|
|
}
|
|
}
|
|
|
|
if (dependsOnMask) {
|
|
for (int j = 0; j < op->outputIndexes()->size(); ++j) {
|
|
if (maskDependentTensors.insert(op->outputIndexes()->data()[j]).second) {
|
|
changed = true;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// 5. For each Attention/LinearAttention break op, backward trace inputs that depend on mask
|
|
// Collect all intermediate ops on the path from attention_mask to Attention
|
|
for (auto opIdx : selectOps) {
|
|
auto op = net->oplists()->GetAs<Op>(opIdx);
|
|
if (op->type() != OpType_Attention && op->type() != OpType_LinearAttention)
|
|
continue;
|
|
if (!isBreakOp(op))
|
|
continue;
|
|
if (op->inputIndexes() == nullptr)
|
|
continue;
|
|
|
|
std::queue<int> bfsQueue;
|
|
std::set<int> visitedTensors;
|
|
|
|
for (int i = 0; i < op->inputIndexes()->size(); ++i) {
|
|
int tensorIdx = op->inputIndexes()->data()[i];
|
|
// This input depends on mask but is NOT the mask tensor itself
|
|
// => there are intermediate ops between mask and Attention
|
|
if (maskDependentTensors.count(tensorIdx) && !maskTensorIndexes.count(tensorIdx)) {
|
|
bfsQueue.push(tensorIdx);
|
|
}
|
|
}
|
|
|
|
while (!bfsQueue.empty()) {
|
|
int tensorIdx = bfsQueue.front();
|
|
bfsQueue.pop();
|
|
if (visitedTensors.count(tensorIdx))
|
|
continue;
|
|
visitedTensors.insert(tensorIdx);
|
|
|
|
if (maskTensorIndexes.count(tensorIdx))
|
|
continue; // Reached mask input, stop
|
|
|
|
auto it = tensorProducer.find(tensorIdx);
|
|
if (it == tensorProducer.end())
|
|
continue;
|
|
|
|
int producerOpIdx = it->second;
|
|
auto producerOp = net->oplists()->GetAs<Op>(producerOpIdx);
|
|
// Skip Const and TrainableParam ops
|
|
if (producerOp->type() == OpType_Const || producerOp->type() == OpType_TrainableParam)
|
|
continue;
|
|
|
|
extraBreakOps.insert(producerOpIdx);
|
|
|
|
// Continue tracing this op's inputs
|
|
if (producerOp->inputIndexes() != nullptr) {
|
|
for (int j = 0; j < producerOp->inputIndexes()->size(); ++j) {
|
|
int inputIdx = producerOp->inputIndexes()->data()[j];
|
|
if (maskDependentTensors.count(inputIdx)) {
|
|
bfsQueue.push(inputIdx);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if (!extraBreakOps.empty()) {
|
|
MNN_PRINT("Found %d extra break ops between attention_mask and Attention:\n", (int)extraBreakOps.size());
|
|
for (auto opIdx : extraBreakOps) {
|
|
auto op = net->oplists()->GetAs<Op>(opIdx);
|
|
if (op->name() != nullptr) {
|
|
MNN_PRINT(" Extra break op: %s (type: %s)\n", op->name()->c_str(), EnumNameOpType(op->type()));
|
|
}
|
|
}
|
|
}
|
|
|
|
return extraBreakOps;
|
|
}
|
|
|
|
static void _setInputOutputForOps(std::vector<std::shared_ptr<Tensor>>& allTensors, const std::vector<const Op*>& ops) {
|
|
std::set<int> inputIndexes;
|
|
std::set<int> outputIndexes;
|
|
// 0. deal virtual tensor for static model:
|
|
// when : A (Any_Op) -----> B (Raster_Op)
|
|
// the tensor will be like below:
|
|
// A_outputs : a_tensor
|
|
// B_inputs : b_tensor (virtual)
|
|
// b_tensor.describe.origin = a_tensor_ptr
|
|
// b_tensor is not a InputTensot, a_tensor is not a OutputTensor
|
|
// so add b_tensor to OutputIndexes, a_tensor to InputIndexes.
|
|
// 1. insert all output/input index in outputIndexes/inputIndexes
|
|
for (auto op : ops) {
|
|
if (nullptr != op->outputIndexes()) {
|
|
auto data = op->outputIndexes()->data();
|
|
for (int j = 0; j < op->outputIndexes()->size(); ++j) {
|
|
outputIndexes.insert(data[j]);
|
|
}
|
|
}
|
|
if (nullptr != op->inputIndexes()) {
|
|
auto data = op->inputIndexes()->data();
|
|
for (int j = 0; j < op->inputIndexes()->size(); ++j) {
|
|
inputIndexes.insert(data[j]);
|
|
}
|
|
}
|
|
MNN_ASSERT(OpType_Input != op->type());
|
|
}
|
|
// 2. the index in outputIndexes/inputIndexed but not in inputIndexes/outputIndexes is output/input
|
|
std::set<int> input;
|
|
std::set<int> output;
|
|
std::set_difference(outputIndexes.begin(), outputIndexes.end(), inputIndexes.begin(), inputIndexes.end(),
|
|
std::inserter(output, output.begin()));
|
|
std::set_difference(inputIndexes.begin(), inputIndexes.end(), outputIndexes.begin(), outputIndexes.end(),
|
|
std::inserter(input, input.begin()));
|
|
// 3. set usage for Tensor by index
|
|
for (auto index : input) {
|
|
auto des = TensorUtils::getDescribe(allTensors[index].get());
|
|
if (des->usage == Tensor::InsideDescribe::CONSTANT || des->usage == Tensor::InsideDescribe::TRAINABLE) {
|
|
continue;
|
|
}
|
|
des->usage = Tensor::InsideDescribe::INPUT;
|
|
}
|
|
for (auto index : output) {
|
|
auto des = TensorUtils::getDescribe(allTensors[index].get());
|
|
if (des->usage == Tensor::InsideDescribe::NORMAL) {
|
|
des->usage = TensorUsage::OUTPUT;
|
|
}
|
|
}
|
|
}
|
|
|
|
void _getConstData(const Net* net, std::vector<MNN::Express::VARP> inputs, const std::set<int>& inputIndexes,
|
|
const std::set<int>& outputIndexes,
|
|
std::map<int, std::tuple<int, int, std::vector<int>, std::vector<char>>>& constTensorData,
|
|
std::string srcpath) {
|
|
// set a backend and context to run resize
|
|
ScheduleConfig config;
|
|
config.type = MNN_FORWARD_CPU;
|
|
BackendConfig backendConfig;
|
|
backendConfig.precision = BackendConfig::Precision_High;
|
|
config.backendConfig = &backendConfig;
|
|
Backend::Info compute;
|
|
compute.type = config.type;
|
|
compute.numThread = config.numThread;
|
|
compute.user = config.backendConfig;
|
|
const RuntimeCreator* runtimeCreator(MNNGetExtraRuntimeCreator(compute.type));
|
|
std::unique_ptr<Runtime> runtime(runtimeCreator->onCreate(compute));
|
|
std::shared_ptr<Backend> backend(runtime->onCreate());
|
|
BackendConfig defaultConfig;
|
|
defaultConfig.flags = 4;
|
|
std::shared_ptr<Backend> defaultBackend(runtime->onCreate(&defaultConfig));
|
|
std::vector<std::shared_ptr<Tensor>> allTensors;
|
|
allTensors.resize(net->tensorName()->size());
|
|
ErrorCode code = NO_ERROR;
|
|
FileLoader loader((srcpath + ".weight").c_str());
|
|
initConstTensors(allTensors, net, defaultBackend.get(), code, &loader);
|
|
if (NO_ERROR != code) {
|
|
MNN_ERROR("Init tensor error code = %d\n", code);
|
|
return;
|
|
}
|
|
bool valid = initTensors(allTensors, net);
|
|
// set tensors' shape by inputConfig
|
|
std::map<std::string, MNN::Express::VARP> inputsMap;
|
|
for (int i = 0; i < inputs.size(); i++) {
|
|
auto name = inputs[i]->name();
|
|
inputsMap[name] = inputs[i];
|
|
}
|
|
for (int i = 0; i < allTensors.size(); i++) {
|
|
auto name = net->tensorName()->GetAsString(i)->str();
|
|
if (inputsMap.find(name) != inputsMap.end()) {
|
|
auto input = inputsMap[name];
|
|
auto info = input->getInfo();
|
|
auto& dims = info->dim;
|
|
allTensors[i]->buffer().dimensions = dims.size();
|
|
for (int j = 0; j < dims.size(); j++) {
|
|
allTensors[i]->setLength(j, dims[j]);
|
|
}
|
|
allTensors[i]->buffer().host =
|
|
(uint8_t*)MNNMemoryAllocAlign(info->size * sizeof(float), MNN_MEMORY_ALIGN_DEFAULT);
|
|
auto ptr = input->readMap<float>();
|
|
std::memcpy(allTensors[i]->buffer().host, ptr, info->size * sizeof(float));
|
|
}
|
|
}
|
|
std::vector<Schedule::OpCacheInfo> infos;
|
|
auto selectOps = _collectNeededOps(net, inputIndexes, outputIndexes);
|
|
{
|
|
std::vector<const Op*> ops;
|
|
for (int i = 0; i < selectOps.size(); i++) {
|
|
auto op = net->oplists()->GetAs<Op>(selectOps[i]);
|
|
if (needComputeOp(op)) {
|
|
ops.push_back(op);
|
|
}
|
|
}
|
|
initPipelineInfosFromOps(infos, ops, allTensors);
|
|
_setInputOutputForOps(allTensors, ops);
|
|
}
|
|
GeometryComputer::Context ctx(Interpreter::GeometryComputeMask::GEOMETRCOMPUTEMASK_ALL, defaultBackend);
|
|
// resize the session's info and store to buffer
|
|
std::vector<Tensor*> constTensors;
|
|
GeometryComputerUtils::buildConstantTensors(infos);
|
|
GeometryComputerUtils::shapeComputeAndGeometryTransform(runtime.get(), nullptr, infos, ctx, defaultBackend,
|
|
runtime->onGetCompilerType());
|
|
for (int i = 0; i < allTensors.size(); i++) {
|
|
auto index = TensorUtils::getDescribe(allTensors[i].get())->index;
|
|
auto iter = constTensorData.find(index);
|
|
if (iter != constTensorData.end()) {
|
|
auto ptr = allTensors[i]->host<char>();
|
|
if (ptr != nullptr) {
|
|
auto size = allTensors[i]->size();
|
|
auto shape = allTensors[i]->shape();
|
|
std::get<0>(iter->second) = TensorUtils::getDescribe(allTensors[i].get())->dimensionFormat;
|
|
std::get<1>(iter->second) = allTensors[i]->getType().code;
|
|
std::get<2>(iter->second).resize(0);
|
|
std::get<3>(iter->second).resize(size);
|
|
memcpy(std::get<3>(iter->second).data(), ptr, size);
|
|
for (int i = 0; i < shape.size(); ++i) {
|
|
std::get<2>(iter->second).push_back(shape[i]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
static void
|
|
_findAllConstTensorIndex(const Net* net, const std::set<int>& inputIndexes, const std::set<int>& outputIndexes,
|
|
std::shared_ptr<Schedule::ScheduleInfo> sharedConst, std::vector<int>& constOpId,
|
|
std::map<int, std::tuple<int, int, std::vector<int>, std::vector<char>>>* constTensorData) {
|
|
auto selectOps = _collectNeededOps(net, inputIndexes, outputIndexes);
|
|
std::set<int> constTensorIndex;
|
|
// 0: not used, 1: const, 2: output
|
|
std::vector<uint8_t> constMask(sharedConst->allTensors.size(), 0);
|
|
for (int i = 0; i < sharedConst->allTensors.size(); ++i) {
|
|
if (sharedConst->allTensors[i].get() != nullptr) {
|
|
constMask[i] = 1;
|
|
}
|
|
}
|
|
for (int v = 0; v < selectOps.size(); ++v) {
|
|
auto op = net->oplists()->GetAs<Op>(selectOps[v]);
|
|
if (nullptr == op->outputIndexes()) {
|
|
continue;
|
|
}
|
|
bool isConst = true;
|
|
if (nullptr != op->inputIndexes()) {
|
|
for (int i = 0; i < op->inputIndexes()->size(); ++i) {
|
|
auto index = op->inputIndexes()->data()[i];
|
|
if (constMask[index]) {
|
|
continue;
|
|
}
|
|
if (OpCommonUtils::opNeedContent(op, i)) {
|
|
isConst = false;
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
if (isConst) {
|
|
for (int i = 0; i < op->outputIndexes()->size(); ++i) {
|
|
auto index = op->outputIndexes()->data()[i];
|
|
constMask[index] = 1;
|
|
constTensorData->emplace(index, std::make_tuple(0, 0, std::vector<int>(0), std::vector<char>(0)));
|
|
constOpId.push_back(selectOps[v]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
static NetT* _replaceConstOp(const void* buffer, size_t bufferSize,
|
|
std::map<int, std::tuple<int, int, std::vector<int>, std::vector<char>>>& constTensorData,
|
|
std::vector<int>& constOpId) {
|
|
auto net = flatbuffers::GetRoot<Net>(buffer)->UnPack();
|
|
for (int i = 0; i < constOpId.size(); ++i) {
|
|
auto op = net->oplists[constOpId[i]].get();
|
|
auto index = op->outputIndexes[0];
|
|
auto iter = constTensorData.find(index);
|
|
if (iter == constTensorData.end()) {
|
|
continue;
|
|
}
|
|
auto name = op->name;
|
|
std::unique_ptr<MNN::OpT> newOp(new OpT);
|
|
newOp->type = OpType_Const;
|
|
newOp->name = name + "_Const";
|
|
auto blob = new BlobT;
|
|
blob->dataFormat = (MNN_DATA_FORMAT)std::get<0>(iter->second);
|
|
blob->dims = std::get<2>(iter->second);
|
|
if (std::get<1>(iter->second) == halide_type_float) {
|
|
blob->dataType = DataType_DT_FLOAT;
|
|
blob->float32s.resize(std::get<3>(iter->second).size() / 4);
|
|
memcpy(blob->float32s.data(), std::get<3>(iter->second).data(), std::get<3>(iter->second).size());
|
|
} else if (std::get<1>(iter->second) == halide_type_int) {
|
|
blob->dataType = DataType_DT_INT32;
|
|
blob->int32s.resize(std::get<3>(iter->second).size() / 4);
|
|
memcpy(blob->int32s.data(), std::get<3>(iter->second).data(), std::get<3>(iter->second).size());
|
|
} else {
|
|
continue;
|
|
}
|
|
newOp->main.value = blob;
|
|
newOp->main.type = OpParameter_Blob;
|
|
newOp->outputIndexes = {index};
|
|
net->oplists[constOpId[i]] = std::move(newOp);
|
|
}
|
|
auto oplist = std::move(net->oplists);
|
|
for (auto& op : oplist) {
|
|
if (nullptr != op.get()) {
|
|
net->oplists.emplace_back(std::move(op));
|
|
}
|
|
}
|
|
return net;
|
|
// flatbuffers::GetRoot<Net>(buffer)->Pack(builder, net);
|
|
// std::ofstream outputOs(dstMNN, std::ios::binary);
|
|
// outputOs.write((const char*)builder.GetBufferPointer(), builder.GetSize());
|
|
}
|
|
|
|
static std::vector<int> _findBreakIndex(const SubModuleInfo& info, const Net* net, std::shared_ptr<Schedule::ScheduleInfo> sharedConst) {
|
|
// 0: not used, 1: const, 2: output
|
|
std::vector<uint8_t> constMask(sharedConst->allTensors.size(), 0);
|
|
for (int i=0; i<sharedConst->allTensors.size(); ++i) {
|
|
if(sharedConst->allTensors[i].get() != nullptr) {
|
|
constMask[i] = 1;
|
|
}
|
|
}
|
|
for (int v = 0; v < info.opList.size(); ++v) {
|
|
auto op = net->oplists()->GetAs<Op>(info.opList[v]);
|
|
if (nullptr == op->outputIndexes()) {
|
|
continue;
|
|
}
|
|
bool isConst = true;
|
|
if (nullptr != op->inputIndexes()) {
|
|
for (int i=0; i<op->inputIndexes()->size(); ++i) {
|
|
auto index = op->inputIndexes()->data()[i];
|
|
if (constMask[index]) {
|
|
continue;
|
|
}
|
|
if (OpCommonUtils::opNeedContent(op, i)) {
|
|
isConst = false;
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
if (isConst) {
|
|
for (int i=0; i<op->outputIndexes()->size(); ++i) {
|
|
auto index = op->outputIndexes()->data()[i];
|
|
constMask[index] = 1;
|
|
}
|
|
}
|
|
}
|
|
std::vector<int> res;
|
|
// Check Break Index
|
|
for (int v = 0; v < info.opList.size(); ++v) {
|
|
auto op = net->oplists()->GetAs<Op>(info.opList[v]);
|
|
if (nullptr == op->outputIndexes() || nullptr == op->inputIndexes()) {
|
|
continue;
|
|
}
|
|
int inputNum = op->inputIndexes()->size();
|
|
auto dims = SizeComputer::needInputContent(op, inputNum);
|
|
for (auto index : dims) {
|
|
if (index < inputNum) {
|
|
if (constMask[op->inputIndexes()->data()[index]] != 1) {
|
|
res.emplace_back(v);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return res;
|
|
}
|
|
static std::vector<SubModuleInfo> _splitSubModuleForShapeConst(const std::vector<SubModuleInfo>& origin, const Net* net, std::shared_ptr<Schedule::ScheduleInfo> sharedConst) {
|
|
std::vector<SubModuleInfo> res;
|
|
for (auto& m : origin) {
|
|
if (m.isBreak) {
|
|
res.emplace_back(std::move(m));
|
|
continue;
|
|
}
|
|
auto breakIndexes = _findBreakIndex(m, net, sharedConst);
|
|
if (breakIndexes.size() > 0) {
|
|
int current = 0;
|
|
for (auto breakIndex : breakIndexes) {
|
|
// Split
|
|
if (breakIndex > current) {
|
|
SubModuleInfo m0;
|
|
m0.opList.insert(m0.opList.begin(), m.opList.begin() + current, m.opList.begin() + breakIndex);
|
|
res.emplace_back(std::move(m0));
|
|
}
|
|
SubModuleInfo m1;
|
|
m1.opList = {m.opList[breakIndex]};
|
|
res.emplace_back(std::move(m1));
|
|
current = breakIndex + 1;
|
|
}
|
|
if (current < m.opList.size()) {
|
|
SubModuleInfo m2;
|
|
m2.opList.insert(m2.opList.begin(), m.opList.begin() + current, m.opList.end());
|
|
res.emplace_back(std::move(m2));
|
|
}
|
|
} else {
|
|
res.emplace_back(std::move(m));
|
|
}
|
|
}
|
|
return res;
|
|
}
|
|
|
|
static bool _needSplitNet(const Net* net, const std::set<int>& inputIndexes, const std::set<int>& outputIndexes) {
|
|
auto selectOps = _collectNeededOps(net, inputIndexes, outputIndexes);
|
|
for (int si = 0; si < selectOps.size(); ++si) {
|
|
auto i = selectOps[si];
|
|
auto op = net->oplists()->GetAs<Op>(i);
|
|
if (isBreakOp(op) || gExtraBreakOpIndexes.count(i) > 0) {
|
|
return true;
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
static std::vector<SubModuleInfo> _createSubModuleInfo(const Net* net, const std::set<int>& inputIndexes, const std::set<int>& outputIndexes, const std::set<int>& noComputeIndexes, std::shared_ptr<Schedule::ScheduleInfo> sharedConst) {
|
|
std::vector<SubModuleInfo> submodule;
|
|
auto selectOps = _collectNeededOps(net, inputIndexes, outputIndexes);
|
|
|
|
// Separate the graph to serveral submodule
|
|
SubModuleInfo current;
|
|
for (int si=0; si<selectOps.size(); ++si) {
|
|
auto i = selectOps[si];
|
|
auto op = net->oplists()->GetAs<Op>(i);
|
|
if (isBreakOp(op) || gExtraBreakOpIndexes.count(i) > 0) {
|
|
// TODO: Don't need split segment
|
|
if (current.opList.size() > 0) {
|
|
// Not empty
|
|
submodule.emplace_back(std::move(current));
|
|
}
|
|
SubModuleInfo controlOp;
|
|
controlOp.opList = {i};
|
|
controlOp.isBreak = true;
|
|
if (nullptr != op->inputIndexes()) {
|
|
controlOp.inputs.resize(op->inputIndexes()->size());
|
|
::memcpy(controlOp.inputs.data(), op->inputIndexes()->data(), controlOp.inputs.size() * sizeof(int));
|
|
}
|
|
if (nullptr != op->outputIndexes()) {
|
|
controlOp.outputs.resize(op->outputIndexes()->size());
|
|
::memcpy(controlOp.outputs.data(), op->outputIndexes()->data(), controlOp.outputs.size() * sizeof(int));
|
|
}
|
|
submodule.emplace_back(std::move(controlOp));
|
|
continue;
|
|
}
|
|
current.opList.emplace_back(i);
|
|
}
|
|
if (!current.opList.empty()) {
|
|
submodule.emplace_back(std::move(current));
|
|
}
|
|
submodule = _splitSubModuleForShapeConst(submodule, net, sharedConst);
|
|
for (int moduleIndex=0; moduleIndex < submodule.size(); ++moduleIndex) {
|
|
auto& m = submodule[moduleIndex];
|
|
// Compute input / output
|
|
if (!m.isBreak) {
|
|
_computeTensorMask(m, net);
|
|
for (int i=0; i<m.tensorMask.size(); ++i) {
|
|
if (0 == m.tensorMask[i]) {
|
|
continue;
|
|
}
|
|
if (1 == m.tensorMask[i]) {
|
|
if (noComputeIndexes.find(i) != noComputeIndexes.end()) {
|
|
continue;
|
|
}
|
|
m.inputs.emplace_back(i);
|
|
continue;
|
|
}
|
|
if (2 == m.tensorMask[i]) {
|
|
m.outputs.emplace_back(i);
|
|
continue;
|
|
}
|
|
if (3 == m.tensorMask[i]) {
|
|
if (outputIndexes.find(i) != outputIndexes.end()) {
|
|
m.outputs.emplace_back(i);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
// Check if the module's input is valid
|
|
for (int i=0; i<m.inputs.size(); ++i) {
|
|
auto index = m.inputs[i];
|
|
if (inputIndexes.find(index) != inputIndexes.end()) {
|
|
continue;
|
|
}
|
|
if (noComputeIndexes.find(index) != noComputeIndexes.end()) {
|
|
continue;
|
|
}
|
|
bool find = false;
|
|
for (int sub=0; sub < moduleIndex; ++sub) {
|
|
for (auto out : submodule[sub].outputs) {
|
|
if (out == index) {
|
|
find = true;
|
|
break;
|
|
}
|
|
}
|
|
if (find) {
|
|
break;
|
|
}
|
|
}
|
|
if (find) {
|
|
continue;
|
|
}
|
|
// Find from module
|
|
for (int sub=0; sub < moduleIndex; ++sub) {
|
|
if (submodule[sub].tensorMask.empty()) {
|
|
continue;
|
|
}
|
|
if (submodule[sub].tensorMask[index] == 2) {
|
|
find = true;
|
|
break;
|
|
}
|
|
if (submodule[sub].tensorMask[index] == 3) {
|
|
submodule[sub].outputs.emplace_back(index);
|
|
submodule[sub].tensorMask[index] = 2;
|
|
find = true;
|
|
break;
|
|
}
|
|
}
|
|
if (!find) {
|
|
if (net->tensorName() != nullptr) {
|
|
MNN_PRINT("%d tensor [ %s ] is input but not found\n", index, net->tensorName()->GetAsString(index)->c_str());
|
|
}
|
|
}
|
|
MNN_ASSERT(find);
|
|
}
|
|
}
|
|
for (auto& m : submodule) {
|
|
m.tensorMask.clear();
|
|
}
|
|
// sort input and output
|
|
for (auto& m : submodule) {
|
|
std::sort(m.inputs.begin(), m.inputs.end());
|
|
std::sort(m.outputs.begin(), m.outputs.end());
|
|
}
|
|
return submodule;
|
|
}
|
|
static std::set<std::string> _getAttentionName(const void* buffer, size_t bufferSize) {
|
|
std::set<std::string> attentionNames;
|
|
auto net = flatbuffers::GetRoot<Net>(buffer);
|
|
if (nullptr == net->oplists()) {
|
|
return attentionNames;
|
|
}
|
|
for (int i=0; i<net->oplists()->size(); ++i) {
|
|
auto op = net->oplists()->GetAs<Op>(i);
|
|
if (op->type() == OpType_Attention) {
|
|
if (nullptr != op->main_as_AttentionParam()) {
|
|
if (op->main_as_AttentionParam()->kv_cache()) {
|
|
attentionNames.insert(op->name()->str());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return attentionNames;
|
|
}
|
|
static SubModuleIO _getSubModuleIO(std::vector<MNN::Express::VARP> inputs, const SubModuleInfo& info, const void* buffer, size_t bufferSize, std::string srcpath) {
|
|
// Deep clone output to let the module release
|
|
SubModuleIO io;
|
|
std::vector<std::string> inputNames(info.inputs.size());
|
|
std::vector<std::string> outputNames(info.outputs.size());
|
|
auto net = flatbuffers::GetRoot<Net>(buffer);
|
|
for (int i=0; i<info.inputs.size(); ++i) {
|
|
auto index = info.inputs[i];
|
|
inputNames[i] = net->tensorName()->GetAsString(index)->str();
|
|
}
|
|
for (int i=0; i<info.outputs.size(); ++i) {
|
|
auto index = info.outputs[i];
|
|
outputNames[i] = net->tensorName()->GetAsString(index)->str();
|
|
}
|
|
auto attentionNames = _getAttentionName(buffer, bufferSize);
|
|
MNN::ScheduleConfig config;
|
|
config.numThread = 1;
|
|
std::shared_ptr<MNN::Express::Executor::RuntimeManager> rtmgr(MNN::Express::Executor::RuntimeManager::createRuntimeManager(config));
|
|
rtmgr->setExternalFile((srcpath + ".weight").c_str());
|
|
rtmgr->setMode(MNN::Interpreter::Session_Debug);
|
|
std::shared_ptr<MNN::Express::Module> m(MNN::Express::Module::load(inputNames, outputNames, (const uint8_t*)buffer, bufferSize, rtmgr), MNN::Express::Module::destroy);
|
|
MNN::TensorCallBackWithInfo beforeCallBack = [&](const std::vector<MNN::Tensor*>& ntensors, const MNN::OperatorInfo* info) {
|
|
auto opName = info->name();
|
|
if (info->type() != "Attention") {
|
|
return true;
|
|
}
|
|
if (attentionNames.find(opName) != attentionNames.end()) {
|
|
auto query = ntensors[0];
|
|
auto key = ntensors[1];
|
|
auto value = ntensors[2];
|
|
int seq_len = query->length(1);
|
|
auto numHead = query->length(2);
|
|
auto headDim = query->length(3);
|
|
auto kvNumHead = key->length(2);
|
|
std::vector<int> kvDims = {kvNumHead, 1, 1, headDim};
|
|
io.kvcache.emplace_back(kvDims);
|
|
io.seqLen = seq_len;
|
|
}
|
|
return true;
|
|
};
|
|
MNN::TensorCallBackWithInfo callBack = [&](const std::vector<MNN::Tensor*>& ntensors, const MNN::OperatorInfo* info) {
|
|
return true;
|
|
};
|
|
MNN::Express::ExecutorScope::Current()->setCallBack(std::move(beforeCallBack), std::move(callBack));
|
|
auto outputs = m->onForward(inputs);
|
|
io.inputs = inputs;
|
|
io.outputs.resize(outputs.size());
|
|
for (int i=0; i<outputs.size(); ++i) {
|
|
io.outputs[i] = MNN::Express::_Clone(outputs[i], true);
|
|
}
|
|
return io;
|
|
}
|
|
|
|
static int _compileWholeModule(std::vector<std::string> inputNames, std::vector<std::string> outputNames,
|
|
std::vector<std::vector<MNN::Express::VARP>> inputs, const std::set<int>& inputIndexes,
|
|
const std::set<int>& outputIndexes, const void* buffer, size_t bufferSize,
|
|
std::string srcpath, std::string dstMNNPath) {
|
|
int npuIndex = 0;
|
|
std::vector<std::vector<MNN::Express::Variable::Info>> outputInfos;
|
|
std::vector<MNN::Express::Variable::Info> outputInfo;
|
|
std::map<std::string, std::vector<std::string>> merges;
|
|
std::vector<std::string> graphicNames;
|
|
auto path = gCacheDir + "/" + gGraphName + std::to_string(npuIndex);
|
|
if (!gOfflieDst.empty()) {
|
|
path += ("." + gOfflieDst);
|
|
}
|
|
std::vector<int> allInputShape;
|
|
for (int inputIndex = 0; inputIndex < inputs.size(); ++inputIndex) {
|
|
std::vector<MNN::Express::Variable::Info> inputInfos(inputs[inputIndex].size());
|
|
for (int i = 0; i < inputInfos.size(); ++i) {
|
|
inputInfos[i] = *inputs[inputIndex][i]->getInfo();
|
|
}
|
|
std::vector<int> currInputShape;
|
|
for (int i = 0; i < inputInfos.size(); i++) {
|
|
for (int j = 0; j < inputInfos[i].dim.size(); j++) {
|
|
currInputShape.emplace_back(inputInfos[i].dim[j]);
|
|
}
|
|
}
|
|
allInputShape.insert(allInputShape.end(), currInputShape.begin(), currInputShape.end());
|
|
|
|
std::string srcPath;
|
|
std::string graphicName;
|
|
if (inputIndex == 0) {
|
|
srcPath = gCacheDir + "/" + gGraphName + std::to_string(npuIndex);
|
|
graphicName = gGraphName + std::to_string(npuIndex);
|
|
} else {
|
|
srcPath = gCacheDir + "/" + gGraphName + std::to_string(inputIndex) + "_" + std::to_string(npuIndex);
|
|
graphicName = gGraphName + std::to_string(inputIndex) + "_" + std::to_string(npuIndex);
|
|
}
|
|
if (!gOfflieSrc.empty()) {
|
|
srcPath += ("." + gOfflieSrc);
|
|
}
|
|
if (merges.find(path) != merges.end()) {
|
|
merges[path].emplace_back(srcPath);
|
|
} else {
|
|
merges.insert(std::make_pair(path, std::vector<std::string>{srcPath}));
|
|
}
|
|
graphicNames.push_back(graphicName);
|
|
MNN::ScheduleConfig config;
|
|
config.type = gNPUType;
|
|
std::shared_ptr<MNN::Express::Executor::RuntimeManager> rtmgr(
|
|
MNN::Express::Executor::RuntimeManager::createRuntimeManager(config));
|
|
rtmgr->setExternalFile((srcpath + ".weight").c_str());
|
|
rtmgr->setCache(srcPath.c_str());
|
|
rtmgr->setHint(MNN::Interpreter::KVCACHE_SIZE_LIMIT, gMaxKVSize);
|
|
MNN::Express::Module::Config mdconfig;
|
|
mdconfig.shapeMutable = false;
|
|
std::shared_ptr<MNN::Express::Module> m(
|
|
MNN::Express::Module::load(inputNames, outputNames, (const uint8_t*)buffer, bufferSize, rtmgr, &mdconfig),
|
|
MNN::Express::Module::destroy);
|
|
auto outputs = m->onForward(inputs[inputIndex]);
|
|
std::vector<MNN::Express::Variable::Info> outputInfo(outputs.size());
|
|
for (int i = 0; i < outputInfo.size(); ++i) {
|
|
outputInfo[i] = *outputs[i]->getInfo();
|
|
}
|
|
outputInfos.emplace_back(outputInfo);
|
|
}
|
|
|
|
/** Fuse to Op*/
|
|
std::unique_ptr<MNN::OpT> op(new OpT);
|
|
for (int i = 0; i < inputs[0].size(); i++) {
|
|
op->inputIndexes.push_back(i);
|
|
}
|
|
for (int i = 0; i < outputInfos[0].size(); i++) {
|
|
op->outputIndexes.push_back(inputs[0].size() + i);
|
|
}
|
|
op->name = "qnn/plugin/op";
|
|
op->main.Reset();
|
|
op->type = MNN::OpType_Plugin;
|
|
op->main.type = MNN::OpParameter_Plugin;
|
|
op->main.value = new MNN::PluginT;
|
|
auto extra = op->main.AsPlugin();
|
|
extra->type = gNPUName;
|
|
std::unique_ptr<MNN::AttributeT> attr(new MNN::AttributeT);
|
|
attr->key = "path";
|
|
attr->s = path;
|
|
extra->attr.emplace_back(std::move(attr));
|
|
// Build name -> tensor index mapping to preserve inputNames order
|
|
auto net = flatbuffers::GetRoot<Net>(buffer);
|
|
std::map<std::string, int> nameToTensorIdx;
|
|
for (auto idx : inputIndexes) {
|
|
nameToTensorIdx[net->tensorName()->GetAsString(idx)->str()] = idx;
|
|
}
|
|
attr.reset(new MNN::AttributeT);
|
|
attr->key = "inputs";
|
|
attr->list.reset(new ListValueT);
|
|
attr->list->s.resize(inputNames.size());
|
|
for (int i = 0; i < inputNames.size(); i++) {
|
|
auto it = nameToTensorIdx.find(inputNames[i]);
|
|
MNN_ASSERT(it != nameToTensorIdx.end());
|
|
attr->list->s[i] = std::string("t") + std::to_string(it->second);
|
|
}
|
|
extra->attr.emplace_back(std::move(attr));
|
|
attr.reset(new AttributeT);
|
|
attr->key = "allGraphName";
|
|
attr->list.reset(new ListValueT);
|
|
attr->list->s = graphicNames;
|
|
extra->attr.emplace_back(std::move(attr));
|
|
|
|
// Build name -> tensor index mapping to preserve outputNames order
|
|
std::map<std::string, int> outNameToTensorIdx;
|
|
for (auto idx : outputIndexes) {
|
|
outNameToTensorIdx[net->tensorName()->GetAsString(idx)->str()] = idx;
|
|
}
|
|
attr.reset(new MNN::AttributeT);
|
|
attr->key = "outputs";
|
|
attr->list.reset(new ListValueT);
|
|
attr->list->s.resize(outputNames.size());
|
|
for (int i = 0; i < outputNames.size(); i++) {
|
|
auto it = outNameToTensorIdx.find(outputNames[i]);
|
|
MNN_ASSERT(it != outNameToTensorIdx.end());
|
|
attr->list->s[i] = std::string("t") + std::to_string(it->second);
|
|
}
|
|
extra->attr.emplace_back(std::move(attr));
|
|
attr.reset(new MNN::AttributeT);
|
|
attr->key = "allInputShape";
|
|
attr->list.reset(new ListValueT);
|
|
attr->list->i.insert(attr->list->i.end(), allInputShape.begin(), allInputShape.end());
|
|
extra->attr.emplace_back(std::move(attr));
|
|
|
|
for (int i = 0; i < outputInfos.size(); ++i) {
|
|
auto outputInfo = outputInfos[i];
|
|
for (int j = 0; j < outputInfo.size(); ++j) {
|
|
attr.reset(new MNN::AttributeT);
|
|
attr->key = "o_" + std::to_string(i) + "_" + std::to_string(j);
|
|
attr->tensor.reset(new BlobT);
|
|
attr->tensor->dataType = OpCommonUtils::convertDataType(outputInfo[j].type);
|
|
attr->tensor->dims = outputInfo[j].dim;
|
|
switch (outputInfo[j].order) {
|
|
case MNN::Express::NHWC:
|
|
attr->tensor->dataFormat = MNN_DATA_FORMAT_NHWC;
|
|
break;
|
|
case MNN::Express::NCHW:
|
|
attr->tensor->dataFormat = MNN_DATA_FORMAT_NCHW;
|
|
break;
|
|
case MNN::Express::NC4HW4:
|
|
attr->tensor->dataFormat = MNN_DATA_FORMAT_NC4HW4;
|
|
break;
|
|
default:
|
|
attr->tensor->dataFormat = MNN_DATA_FORMAT_NCHW;
|
|
break;
|
|
}
|
|
extra->attr.emplace_back(std::move(attr));
|
|
}
|
|
}
|
|
|
|
std::shared_ptr<MNN::NetT> dstNet(new NetT);
|
|
|
|
for (int i = 0; i < inputs[0].size(); ++i) {
|
|
auto inputInfos = *inputs[0][i]->getInfo();
|
|
std::unique_ptr<OpT> input(new OpT);
|
|
input->type = OpType_Input;
|
|
auto param(new InputT);
|
|
param->dims = inputInfos.dim;
|
|
|
|
input->main.type = OpParameter_Input;
|
|
input->main.value = param;
|
|
input->name = inputNames[i];
|
|
input->outputIndexes.push_back(i);
|
|
dstNet->oplists.emplace_back(std::move(input));
|
|
}
|
|
|
|
dstNet->tensorName = inputNames;
|
|
dstNet->tensorName.insert(dstNet->tensorName.end(), outputNames.begin(), outputNames.end());
|
|
dstNet->tensorName.push_back(op->name);
|
|
dstNet->outputName = outputNames;
|
|
|
|
std::unique_ptr<OpT> npuOp;
|
|
npuOp = std::move(op);
|
|
|
|
// Merge to dst
|
|
dstNet->oplists.emplace_back(std::move(npuOp));
|
|
|
|
// Store
|
|
flatbuffers::FlatBufferBuilder builder;
|
|
builder.Finish(Net::Pack(builder, dstNet.get()));
|
|
std::ofstream outputOs(dstMNNPath.c_str(), std::ios::binary);
|
|
outputOs.write((const char*)builder.GetBufferPointer(), builder.GetSize());
|
|
outputOs.close();
|
|
|
|
// Write Merge Info
|
|
rapidjson::Document resDocument;
|
|
resDocument.SetObject();
|
|
rapidjson::Value mergeMessages;
|
|
mergeMessages.SetObject();
|
|
for (auto& iter : merges) {
|
|
rapidjson::Value mergeSrc;
|
|
mergeSrc.SetArray();
|
|
for (auto& v : iter.second) {
|
|
rapidjson::Value vt;
|
|
vt.SetString(v.c_str(), resDocument.GetAllocator());
|
|
mergeSrc.GetArray().PushBack(vt, resDocument.GetAllocator());
|
|
}
|
|
rapidjson::Value key;
|
|
key.SetString(iter.first.c_str(), resDocument.GetAllocator());
|
|
mergeMessages.AddMember(key, mergeSrc, resDocument.GetAllocator());
|
|
}
|
|
{
|
|
rapidjson::Value type;
|
|
type.SetString(gNPUName.c_str(), resDocument.GetAllocator());
|
|
resDocument.AddMember("type", type, resDocument.GetAllocator());
|
|
}
|
|
resDocument.AddMember("merge", mergeMessages, resDocument.GetAllocator());
|
|
{
|
|
rapidjson::Value cachedir;
|
|
cachedir.SetString(gCacheDir.c_str(), resDocument.GetAllocator());
|
|
resDocument.AddMember("cache", cachedir, resDocument.GetAllocator());
|
|
}
|
|
rapidjson::StringBuffer buf;
|
|
rapidjson::PrettyWriter<rapidjson::StringBuffer> bufwriter(buf);
|
|
resDocument.Accept(bufwriter);
|
|
MNN_PRINT("Write config to npu_postreat.json\n");
|
|
std::ofstream os("npu_postreat.json");
|
|
os << buf.GetString();
|
|
return 0;
|
|
}
|
|
|
|
static std::unique_ptr<MNN::OpT> _compileSubModule(const SubModuleIO& io, SubModuleInfo& info, const void* buffer, size_t bufferSize, const std::string& path, std::string srcpath, const std::string& targetNpuPath, float& cpuTotal, float& npuTotal, int shapeIndex, std::string graphicName) {
|
|
std::vector<std::string> inputNames(info.inputs.size());
|
|
std::vector<std::string> outputNames(info.outputs.size());
|
|
auto net = flatbuffers::GetRoot<Net>(buffer);
|
|
for (int i=0; i<info.inputs.size(); ++i) {
|
|
auto index = info.inputs[i];
|
|
inputNames[i] = net->tensorName()->GetAsString(index)->str();
|
|
}
|
|
for (int i=0; i<info.outputs.size(); ++i) {
|
|
auto index = info.outputs[i];
|
|
outputNames[i] = net->tensorName()->GetAsString(index)->str();
|
|
}
|
|
/** Get Output shapes*/
|
|
std::vector<MNN::Express::Variable::Info> outputInfos(io.outputs.size());
|
|
for (int i=0; i<outputInfos.size(); ++i) {
|
|
outputInfos[i] = *io.outputs[i]->getInfo();
|
|
}
|
|
|
|
/** Make ML Model*/
|
|
do {
|
|
MNN::ScheduleConfig config;
|
|
config.type = gNPUType;
|
|
std::shared_ptr<MNN::Express::Executor::RuntimeManager> rtmgr(MNN::Express::Executor::RuntimeManager::createRuntimeManager(config));
|
|
rtmgr->setExternalFile((srcpath + ".weight").c_str());
|
|
rtmgr->setCache(path.c_str());
|
|
rtmgr->setHint(MNN::Interpreter::KVCACHE_SIZE_LIMIT, gMaxKVSize);
|
|
MNN::Express::Module::Config mdconfig;
|
|
mdconfig.shapeMutable = false;
|
|
std::shared_ptr<MNN::Express::Module> m(MNN::Express::Module::load(inputNames, outputNames, (const uint8_t*)buffer, bufferSize, rtmgr, &mdconfig), MNN::Express::Module::destroy);
|
|
auto predict = m->onForward(io.inputs);
|
|
if (gNeedOffline) {
|
|
break;
|
|
}
|
|
if (predict.size() != io.outputs.size()) {
|
|
MNN_ERROR("Failed to compile: %s\n", path.c_str());
|
|
info.isBreak = true;
|
|
return nullptr;
|
|
}
|
|
for (int i=0; i<predict.size(); ++i) {
|
|
auto error = io.outputs[i]-predict[i];
|
|
error = MNN::Express::_ReduceMax(MNN::Express::_Abs(MNN::Express::_Cast<float>(error)));
|
|
auto maxValue = MNN::Express::_ReduceMax(MNN::Express::_Abs(MNN::Express::_Cast<float>(io.outputs[i])))->readMap<float>()[0];
|
|
if (maxValue < 0.01f) {
|
|
maxValue = 0.01f;
|
|
}
|
|
auto errorf = error->readMap<float>()[0];
|
|
if (errorf / maxValue > 0.1f) {
|
|
MNN_ERROR("error = %f, max = %f for %s\n", errorf, maxValue, path.c_str());
|
|
info.isBreak = true;
|
|
return nullptr;
|
|
}
|
|
}
|
|
// Compare Speed
|
|
int testTime = 20;
|
|
MNN_PRINT("Start to Test speed for %d times\n", testTime);
|
|
MNN::Timer timer;
|
|
for (int i=0; i<testTime; ++i) {
|
|
predict = m->onForward(io.inputs);
|
|
((MNN::Tensor*)predict[0]->getTensor())->wait(MNN::Tensor::MAP_TENSOR_READ, true);
|
|
}
|
|
auto npuCost = timer.durationInUs();
|
|
MNN::ScheduleConfig configcpu;
|
|
configcpu.numThread = 4;
|
|
std::shared_ptr<MNN::Express::Executor::RuntimeManager> rtmgrCPU(MNN::Express::Executor::RuntimeManager::createRuntimeManager(configcpu));
|
|
rtmgrCPU->setExternalFile((srcpath + ".weight").c_str());
|
|
m.reset(MNN::Express::Module::load(inputNames, outputNames, (const uint8_t*)buffer, bufferSize, rtmgrCPU, &mdconfig), MNN::Express::Module::destroy);
|
|
predict = m->onForward(io.inputs);
|
|
timer.reset();
|
|
for (int i=0; i<testTime; ++i) {
|
|
predict = m->onForward(io.inputs);
|
|
((MNN::Tensor*)predict[0]->getTensor())->wait(MNN::Tensor::MAP_TENSOR_READ, true);
|
|
}
|
|
auto cpuCost = timer.durationInUs();
|
|
float npuF = (float)npuCost/ 1000.0f / testTime;
|
|
float cpuF = (float)cpuCost / 1000.0f / testTime;
|
|
|
|
MNN_PRINT("%s, Speed Compare: NPU: %f ms : CPU: %f ms\n", path.c_str(), npuF, cpuF);
|
|
cpuTotal += cpuF;
|
|
npuTotal += npuF;
|
|
} while (false);
|
|
|
|
/** Fuse to Op*/
|
|
std::unique_ptr<MNN::OpT> op(new OpT);
|
|
op->inputIndexes = info.inputs;
|
|
op->outputIndexes = info.outputs;
|
|
op->name = targetNpuPath;
|
|
op->main.Reset();
|
|
op->type = MNN::OpType_Plugin;
|
|
op->main.type = MNN::OpParameter_Plugin;
|
|
op->main.value = new MNN::PluginT;
|
|
auto extra = op->main.AsPlugin();
|
|
extra->type = gNPUName;
|
|
std::unique_ptr<MNN::AttributeT> attr(new MNN::AttributeT);
|
|
attr->key = "path";
|
|
attr->s = targetNpuPath;
|
|
extra->attr.emplace_back(std::move(attr));
|
|
attr.reset(new MNN::AttributeT);
|
|
attr->key = "inputs";
|
|
attr->list.reset(new ListValueT);
|
|
attr->list->s.resize(inputNames.size());
|
|
for (int i=0; i<inputNames.size(); ++i) {
|
|
// CoreML Backend will name tensor as t + index
|
|
attr->list->s[i] = std::string("t") + std::to_string(info.inputs[i]);
|
|
}
|
|
extra->attr.emplace_back(std::move(attr));
|
|
attr.reset(new AttributeT);
|
|
attr->key = "allGraphName";
|
|
attr->list.reset(new ListValueT);
|
|
attr->list->s = {graphicName};
|
|
extra->attr.emplace_back(std::move(attr));
|
|
if (io.kvcache.size() > 0) {
|
|
attr.reset(new MNN::AttributeT);
|
|
attr->key = "seq_len";
|
|
attr->list.reset(new ListValueT);
|
|
attr->list->i = {io.seqLen};
|
|
extra->attr.emplace_back(std::move(attr));
|
|
|
|
attr.reset(new MNN::AttributeT);
|
|
attr->key = "state";
|
|
attr->tensor.reset(new BlobT);
|
|
attr->tensor->dataType = DataType_DT_UINT8;
|
|
flexbuffers::Builder builder;
|
|
auto start = builder.StartMap();
|
|
builder.Int("number", io.kvcache.size() * 2);
|
|
builder.Int("max_length", gMaxKVSize);
|
|
builder.Int("axis", 2);
|
|
auto shapeStart = builder.StartVector("shape");
|
|
|
|
// Add State
|
|
for (int i=0; i<io.kvcache.size(); ++i) {
|
|
// Each KV has two state
|
|
for (int j=0; j<2; ++j) {
|
|
auto vecStart = builder.StartVector();
|
|
for (int v=0; v<io.kvcache[i].size(); ++v) {
|
|
builder.Add(io.kvcache[i][v]);
|
|
}
|
|
builder.EndVector(vecStart, false, false);
|
|
}
|
|
}
|
|
builder.EndVector(shapeStart, false, false);
|
|
|
|
builder.EndMap(start);
|
|
builder.Finish();
|
|
attr->tensor->uint8s = builder.GetBuffer();
|
|
if (false) {
|
|
// Try Read
|
|
auto ref = flexbuffers::GetRoot(attr->tensor->uint8s.data(), attr->tensor->uint8s.size());
|
|
auto refMap = ref.AsMap();
|
|
auto keys = refMap.Keys();
|
|
int readNumber = 0;
|
|
int maxLength = 0;
|
|
std::vector<std::vector<int>> stateShape;
|
|
for (int i=0; i<keys.size(); ++i) {
|
|
auto key = keys[i].AsKey();
|
|
if (std::string(key) == "number") {
|
|
readNumber = refMap.Values()[i].AsInt32();
|
|
continue;
|
|
}
|
|
if (std::string(key) == "max_length") {
|
|
maxLength = refMap.Values()[i].AsInt32();
|
|
continue;
|
|
}
|
|
if (std::string(key) == "shape") {
|
|
auto shapeVectors = refMap.Values()[i].AsVector();
|
|
for (int u=0; u<shapeVectors.size(); ++u) {
|
|
auto shapeV = shapeVectors[u].AsVector();
|
|
std::vector<int> shapes;
|
|
for (int v=0; v<shapeV.size(); ++v) {
|
|
shapes.emplace_back(shapeV[v].AsInt32());
|
|
}
|
|
stateShape.emplace_back(shapes);
|
|
}
|
|
continue;
|
|
}
|
|
}
|
|
FUNC_PRINT(1);
|
|
}
|
|
extra->attr.emplace_back(std::move(attr));
|
|
}
|
|
|
|
attr.reset(new MNN::AttributeT);
|
|
attr->key = "outputs";
|
|
attr->list.reset(new ListValueT);
|
|
attr->list->s.resize(outputNames.size());
|
|
for (int i=0; i<outputNames.size(); ++i) {
|
|
// CoreML Backend will name tensor as t + index
|
|
attr->list->s[i] = std::string("t") + std::to_string(info.outputs[i]);
|
|
}
|
|
extra->attr.emplace_back(std::move(attr));
|
|
attr.reset(new MNN::AttributeT);
|
|
attr->key = "allInputShape";
|
|
attr->list.reset(new ListValueT);
|
|
std::string inputsShapeStr = "";
|
|
for (int i = 0; i < io.inputs.size(); i++) {
|
|
for (int j = 0; j < io.inputs[i]->getInfo()->dim.size(); j++) {
|
|
attr->list->i.emplace_back(io.inputs[i]->getInfo()->dim[j]);
|
|
}
|
|
}
|
|
extra->attr.emplace_back(std::move(attr));
|
|
|
|
for (int i=0; i<outputInfos.size(); ++i) {
|
|
attr.reset(new MNN::AttributeT);
|
|
attr->key = "o_" + std::to_string(shapeIndex) + "_" + std::to_string(i);
|
|
attr->tensor.reset(new BlobT);
|
|
attr->tensor->dataType = OpCommonUtils::convertDataType( outputInfos[i].type);
|
|
attr->tensor->dims = outputInfos[i].dim;
|
|
switch(outputInfos[i].order) {
|
|
case MNN::Express::NHWC:
|
|
attr->tensor->dataFormat = MNN_DATA_FORMAT_NHWC;
|
|
break;
|
|
case MNN::Express::NCHW:
|
|
attr->tensor->dataFormat = MNN_DATA_FORMAT_NCHW;
|
|
break;
|
|
case MNN::Express::NC4HW4:
|
|
attr->tensor->dataFormat = MNN_DATA_FORMAT_NC4HW4;
|
|
break;
|
|
default:
|
|
attr->tensor->dataFormat = MNN_DATA_FORMAT_NCHW;
|
|
break;
|
|
}
|
|
extra->attr.emplace_back(std::move(attr));
|
|
}
|
|
return op;
|
|
}
|
|
static bool _fuse(MNN::NetT* net, MNN::NetT* srcNet) {
|
|
std::map<std::string, MNN::OpT*> dstPlugin;
|
|
for (auto& op : net->oplists) {
|
|
if (op->type == OpType_Plugin) {
|
|
dstPlugin.insert(std::make_pair(op->name, op.get()));
|
|
}
|
|
}
|
|
for (auto& op : srcNet->oplists) {
|
|
if (op->type != OpType_Plugin) {
|
|
continue;
|
|
}
|
|
auto iter = dstPlugin.find(op->name);
|
|
if (iter == dstPlugin.end()) {
|
|
MNN_ERROR("Can't find plugin: %s\n", op->name.c_str());
|
|
continue;
|
|
}
|
|
auto dst = iter->second->main.AsPlugin();
|
|
auto src = op->main.AsPlugin();
|
|
std::map<std::string, AttributeT*> dstKeys;
|
|
for (auto& dstAttr : dst->attr) {
|
|
dstKeys.insert(std::make_pair(dstAttr->key, dstAttr.get()));
|
|
}
|
|
for (auto&& srcAttr : src->attr) {
|
|
if (srcAttr->key == "inputs" || srcAttr->key == "outputs") {
|
|
// Don't fuse same one
|
|
continue;
|
|
}
|
|
auto dstIter = dstKeys.find(srcAttr->key);
|
|
if (dstIter == dstKeys.end()) {
|
|
dst->attr.emplace_back(std::move(srcAttr));
|
|
continue;
|
|
}
|
|
if (dstIter->second->list != nullptr && srcAttr->list != nullptr) {
|
|
dstIter->second->list->s.insert(dstIter->second->list->s.end(), srcAttr->list->s.begin(), srcAttr->list->s.end());
|
|
dstIter->second->list->i.insert(dstIter->second->list->i.end(), srcAttr->list->i.begin(), srcAttr->list->i.end());
|
|
}
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
static bool _reOrderOp(MNN::NetT* net) {
|
|
auto oplist = std::move(net->oplists);
|
|
std::set<int> validInputs;
|
|
do {
|
|
bool empty = true;
|
|
for (int i=0; i<oplist.size(); ++i) {
|
|
if (nullptr == oplist[i]) {
|
|
continue;
|
|
}
|
|
bool valid = true;
|
|
for (auto index : oplist[i]->inputIndexes) {
|
|
if (validInputs.find(index) == validInputs.end()) {
|
|
valid = false;
|
|
break;
|
|
}
|
|
}
|
|
if (valid) {
|
|
for (auto index : oplist[i]->outputIndexes) {
|
|
validInputs.insert(index);
|
|
}
|
|
net->oplists.emplace_back(std::move(oplist[i]));
|
|
oplist[i] = nullptr;
|
|
} else {
|
|
empty = false;
|
|
}
|
|
}
|
|
if (empty) {
|
|
break;
|
|
}
|
|
} while (true);
|
|
return true;
|
|
}
|
|
|
|
static bool _reIndexTensor(MNN::NetT* net) {
|
|
auto& mNet = net;
|
|
std::map<std::string, int> tensorNameIdx;
|
|
std::map<int, int> usefulTensorIndexMap;
|
|
std::vector<std::string> usefulTensorName;
|
|
// extraTensorDescribe reindex
|
|
for (int i = 0; i < mNet->tensorName.size(); i++) {
|
|
tensorNameIdx.insert(std::make_pair(mNet->tensorName[i], i));
|
|
}
|
|
for (int i = 0; i < mNet->extraTensorDescribe.size(); i++) {
|
|
auto name = mNet->extraTensorDescribe[i]->name;
|
|
auto iter = tensorNameIdx.find(name);
|
|
if (iter == tensorNameIdx.end()) {
|
|
mNet->extraTensorDescribe[i]->index = -1;
|
|
} else {
|
|
mNet->extraTensorDescribe[i]->index = iter->second;
|
|
}
|
|
}
|
|
|
|
std::vector<bool> tensorValid(mNet->tensorName.size(), false);
|
|
for (auto& op : mNet->oplists) {
|
|
for (auto index : op->inputIndexes) {
|
|
if (index < 0) {
|
|
continue; // optional input, ignore it
|
|
}
|
|
tensorValid[index] = true;
|
|
}
|
|
for (auto index : op->outputIndexes) {
|
|
tensorValid[index] = true;
|
|
}
|
|
}
|
|
|
|
for (int i = 0; i < tensorValid.size(); ++i) {
|
|
if (tensorValid[i]) {
|
|
usefulTensorIndexMap.insert(std::make_pair(i, usefulTensorName.size()));
|
|
usefulTensorName.push_back(mNet->tensorName[i]);
|
|
}
|
|
}
|
|
|
|
// Re index
|
|
for (auto& op : mNet->oplists) {
|
|
for (int i = 0; i < op->inputIndexes.size(); ++i) {
|
|
if (op->inputIndexes[i] < 0) {
|
|
continue;
|
|
}
|
|
auto iter = usefulTensorIndexMap.find(op->inputIndexes[i]);
|
|
op->inputIndexes[i] = iter->second;
|
|
}
|
|
for (int i = 0; i < op->outputIndexes.size(); ++i) {
|
|
auto iter = usefulTensorIndexMap.find(op->outputIndexes[i]);
|
|
op->outputIndexes[i] = iter->second;
|
|
}
|
|
}
|
|
|
|
mNet->tensorName = usefulTensorName;
|
|
for (auto iter = mNet->extraTensorDescribe.begin(); iter != mNet->extraTensorDescribe.end();) {
|
|
auto index = (*iter)->index;
|
|
if (usefulTensorIndexMap.find(index) == usefulTensorIndexMap.end()) {
|
|
iter = mNet->extraTensorDescribe.erase(iter);
|
|
continue;
|
|
}
|
|
(*iter)->index = usefulTensorIndexMap.find(index)->second;
|
|
iter++;
|
|
}
|
|
// Check dup name and modify
|
|
std::set<std::string> names;
|
|
std::set<std::string> tensorNames;
|
|
for (int i = 0; i < mNet->oplists.size(); ++i) {
|
|
auto& op = mNet->oplists[i];
|
|
auto opName = op->name;
|
|
if (opName.empty() || names.find(opName) != names.end()) {
|
|
std::ostringstream defaultName;
|
|
defaultName << EnumNameOpType(op->type);
|
|
defaultName << i;
|
|
op->name = defaultName.str();
|
|
#ifdef DEBUG
|
|
MNN_PRINT("%d op name is empty or dup, set to %s\n", i, op->name.c_str());
|
|
#endif
|
|
opName = op->name;
|
|
}
|
|
names.insert(opName);
|
|
for (auto output : op->outputIndexes) {
|
|
auto origin = net->tensorName[output];
|
|
if (origin.empty() || tensorNames.find(origin) != tensorNames.end()) {
|
|
std::ostringstream defaultName;
|
|
defaultName << output;
|
|
origin = defaultName.str();
|
|
net->tensorName[output] = origin;
|
|
}
|
|
tensorNames.insert(origin);
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
int main(int argc, const char* argv[]) {
|
|
if (argc < 3) {
|
|
MNN_PRINT("Usage: ./compilefornpu src.mnn dst.mnn npu.json\n");
|
|
return 0;
|
|
}
|
|
const char* srcMNN = argv[1];
|
|
const char* dstMNN = argv[2];
|
|
std::vector<std::string> inputNames;
|
|
std::vector<std::string> outputNames;
|
|
std::vector<std::vector<MNN::Express::VARP>> inputs;
|
|
std::set<std::string> skipOps;
|
|
|
|
if (argc >= 4) {
|
|
rapidjson::Document document;
|
|
std::ifstream fileNames(argv[3]);
|
|
std::ostringstream output;
|
|
output << fileNames.rdbuf();
|
|
auto outputStr = output.str();
|
|
document.Parse(outputStr.c_str());
|
|
if (document.HasParseError()) {
|
|
MNN_ERROR("Invalid json\n");
|
|
return 0;
|
|
}
|
|
gNPUName = document["type"].GetString();
|
|
if (gNPUName == "QNN") {
|
|
MNN_PRINT("Convert for QNN, QualComn's NPU\n");
|
|
gNPUType = MNN_CONVERT_QNN;
|
|
gNeedOffline = true;
|
|
gOfflieSrc = "";
|
|
gOfflieDst = "bin";
|
|
} else if (gNPUName == "MLDA") {
|
|
MNN_PRINT("Convert for MLDA, MTK's NPU\n");
|
|
gNPUType = MNN_CONVERT_NEUROPILOT;
|
|
gNeedOffline = true;
|
|
gOfflieSrc = "tflite";
|
|
gOfflieDst = "dla";
|
|
} else if (gNPUName == "CoreML") {
|
|
MNN_PRINT("Convert for CoreML, Apple's framework\n");
|
|
gNPUType = MNN_CONVERT_COREML;
|
|
gNeedOffline = true;
|
|
gOfflieSrc = "";
|
|
gOfflieDst = "";
|
|
} else {
|
|
MNN_PRINT("Use Native NPU compute\n");
|
|
}
|
|
if (document.HasMember("cache")) {
|
|
gCacheDir = document["cache"].GetString();
|
|
FUNC_PRINT_ALL(gCacheDir.c_str(), s);
|
|
MNNCreateDir(gCacheDir.c_str());
|
|
}
|
|
if (document.HasMember("graph_name")) {
|
|
gGraphName = document["graph_name"].GetString();
|
|
}
|
|
if (document.HasMember("skips")) {
|
|
auto skips = document["skips"].GetArray();
|
|
for (auto iter = skips.Begin(); iter != skips.End(); iter++) {
|
|
skipOps.insert(iter->GetString());
|
|
}
|
|
}
|
|
if (document.HasMember("KVCACHE_SIZE_LIMIT")) {
|
|
gMaxKVSize = document["KVCACHE_SIZE_LIMIT"].GetInt();
|
|
}
|
|
if (document.HasMember("testdir")) {
|
|
auto testdir = document["testdir"].GetArray();
|
|
for (auto iter = testdir.Begin(); iter != testdir.End(); iter++) {
|
|
std::string dirname = iter->GetString();
|
|
auto subinputs = MNN::Express::Variable::load((dirname + "/input.mnn").c_str());
|
|
if (subinputs.empty()) {
|
|
MNN_ERROR("Failed to load test inputs from %s/input.mnn\n", dirname.c_str());
|
|
return 1;
|
|
}
|
|
inputs.emplace_back(subinputs);
|
|
inputNames.clear();
|
|
for (int i=0; i<subinputs.size(); ++i) {
|
|
inputNames.emplace_back(subinputs[i]->name());
|
|
}
|
|
auto outputs = MNN::Express::Variable::load((dirname + "/output.mnn").c_str());
|
|
if (outputs.empty()) {
|
|
MNN_ERROR("Failed to load test outputs from %s/output.mnn\n", dirname.c_str());
|
|
return 1;
|
|
}
|
|
outputNames.clear();
|
|
for (int i=0; i<outputs.size(); ++i) {
|
|
outputNames.emplace_back(outputs[i]->name());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
if (outputNames.empty()) {
|
|
std::shared_ptr<MNN::Express::Module> m(MNN::Express::Module::load(inputNames, outputNames, srcMNN), MNN::Express::Module::destroy);
|
|
if (nullptr == m.get()) {
|
|
MNN_ERROR("Failed to load source module from %s\n", srcMNN);
|
|
return 1;
|
|
}
|
|
auto minfo = m->getInfo();
|
|
outputNames = minfo->outputNames;
|
|
inputNames = minfo->inputNames;
|
|
std::vector<MNN::Express::VARP> subinputs;
|
|
subinputs.resize(minfo->inputs.size());
|
|
for (int i=0; i<minfo->inputs.size(); ++i) {
|
|
auto& info = minfo->inputs[i];
|
|
auto varp = MNN::Express::_Input(info.dim, info.order, info.type);
|
|
varp->writeMap<void>();
|
|
subinputs[i] = varp;
|
|
}
|
|
inputs = {subinputs};
|
|
}
|
|
// Registor size computor
|
|
MNN::Express::Executor::getGlobalExecutor();
|
|
// Get Net struct
|
|
std::shared_ptr<MNN::Interpreter> netC(MNN::Interpreter::createFromFile(srcMNN), MNN::Interpreter::destroy);
|
|
auto bufferPair = netC->getModelBuffer();
|
|
std::shared_ptr<Schedule::ScheduleInfo> sharedConst;
|
|
auto buffer = bufferPair.first;
|
|
auto length = bufferPair.second;
|
|
auto net = GetNet(buffer);
|
|
std::map<std::string, int> tensorIndexMap;
|
|
for (int i=0; i<net->tensorName()->size(); ++i) {
|
|
auto tname = net->tensorName()->GetAsString(i)->str();
|
|
tensorIndexMap.insert(std::make_pair(tname, i));
|
|
}
|
|
// Extra Const Tensors
|
|
sharedConst.reset(new Schedule::ScheduleInfo);
|
|
std::vector<std::shared_ptr<Tensor>> allTensors;
|
|
sharedConst->allTensors.resize(net->tensorName()->size());
|
|
initConstTensorsNoAlloc(sharedConst->allTensors, net);
|
|
std::set<int> noneedComputeIndexes;
|
|
for (int i=0; i<sharedConst->allTensors.size(); ++i) {
|
|
if (sharedConst->allTensors[i].get() != nullptr) {
|
|
noneedComputeIndexes.insert(i);
|
|
}
|
|
}
|
|
|
|
std::set<int> inputIndexes;
|
|
std::set<int> outputIndexes;
|
|
std::map<std::string, int> outputIndexesMap;
|
|
for (int i=0; i<net->tensorName()->size(); ++i) {
|
|
auto tname = net->tensorName()->GetAsString(i)->str();
|
|
for (int j=0; j<inputNames.size(); ++j) {
|
|
if (tname == inputNames[j]) {
|
|
inputIndexes.emplace(i);
|
|
break;
|
|
}
|
|
}
|
|
for (int j=0; j<outputNames.size(); ++j) {
|
|
if (tname == outputNames[j]) {
|
|
outputIndexes.emplace(i);
|
|
outputIndexesMap.insert(std::make_pair(tname, i));
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
if (outputIndexesMap.size() != outputNames.size()) {
|
|
MNN_ERROR("PipelineModule:: Can't find enough output from the model, finded is:\n");
|
|
for (auto& iter : outputIndexesMap) {
|
|
MNN_ERROR("[ %s ] ", iter.first.c_str());
|
|
}
|
|
MNN_ERROR("\n");
|
|
}
|
|
auto firstInputIndex = inputIndexes;
|
|
std::set<int> firstOutputIndex;
|
|
for (int i=0; i<net->oplists()->size(); ++i) {
|
|
auto op = net->oplists()->GetAs<Op>(i);
|
|
if (skipOps.find(op->name()->str()) != skipOps.end()) {
|
|
MNN_PRINT("Skip %s op\n", op->name()->c_str());
|
|
auto outputSize = op->outputIndexes()->size();
|
|
for (int v=0; v<outputSize; ++v) {
|
|
firstOutputIndex.insert(op->outputIndexes()->data()[v]);
|
|
}
|
|
}
|
|
}
|
|
if (!firstOutputIndex.empty()) {
|
|
// Get New Inputs
|
|
std::vector<std::string> firstOutputNames;
|
|
|
|
// Compute New Input
|
|
for (auto output : firstOutputIndex) {
|
|
inputIndexes.insert(output);
|
|
firstOutputNames.emplace_back(net->tensorName()->GetAsString(output)->str());
|
|
}
|
|
std::vector<std::string> newInputNames;
|
|
for (auto index : inputIndexes) {
|
|
newInputNames.emplace_back(net->tensorName()->GetAsString(index)->str());
|
|
}
|
|
std::shared_ptr<MNN::Express::Module> m(MNN::Express::Module::load(inputNames, firstOutputNames, (const uint8_t*)bufferPair.first, bufferPair.second), MNN::Express::Module::destroy);
|
|
for (int i=0; i<inputs.size(); ++i) {
|
|
std::map<std::string, MNN::Express::VARP> vars;
|
|
for (int v=0; v<inputNames.size(); ++v) {
|
|
vars.insert(std::make_pair(inputNames[v], inputs[i][v]));
|
|
}
|
|
auto outputs = m->onForward(inputs[i]);
|
|
for (int v=0; v<firstOutputNames.size(); ++v) {
|
|
vars.insert(std::make_pair(firstOutputNames[v], outputs[v]));
|
|
}
|
|
inputs[i].clear();
|
|
for (int v=0; v<newInputNames.size(); ++v) {
|
|
inputs[i].emplace_back(vars[newInputNames[v]]);
|
|
}
|
|
}
|
|
inputNames = newInputNames;
|
|
}
|
|
|
|
// Find intermediate ops between attention_mask and Attention that should also be break ops
|
|
gExtraBreakOpIndexes = _findMaskToAttentionOps(net, inputIndexes, outputIndexes);
|
|
|
|
if (_needSplitNet(net, inputIndexes, outputIndexes) == false) {
|
|
return _compileWholeModule(inputNames, outputNames, inputs, inputIndexes, outputIndexes, bufferPair.first,
|
|
bufferPair.second, srcMNN, dstMNN);
|
|
}
|
|
|
|
std::vector<int> constOpId;
|
|
std::map<int, std::tuple<int, int, std::vector<int>, std::vector<char>>> constTensorData;
|
|
std::vector<NetT*> netReplace(inputs.size());
|
|
std::vector<flatbuffers::FlatBufferBuilder> BuilderTmp(inputs.size());
|
|
std::vector<NetT*> netVec(inputs.size());
|
|
|
|
_findAllConstTensorIndex(net, inputIndexes, outputIndexes, sharedConst, constOpId, &constTensorData);
|
|
for (int inputIndex = 0; inputIndex < inputs.size(); ++inputIndex) {
|
|
_getConstData(net, inputs[inputIndex], inputIndexes, outputIndexes, constTensorData, srcMNN);
|
|
netReplace[inputIndex] = _replaceConstOp(bufferPair.first, bufferPair.second, constTensorData, constOpId);
|
|
BuilderTmp[inputIndex].Finish(Net::Pack(BuilderTmp[inputIndex], netReplace[inputIndex]));
|
|
}
|
|
auto bufferPair0 = std::make_pair(BuilderTmp[0].GetBufferPointer(), BuilderTmp[0].GetSize());
|
|
auto net0 = GetNet(bufferPair0.first);
|
|
// Extra Const Tensors
|
|
sharedConst.reset(new Schedule::ScheduleInfo);
|
|
sharedConst->allTensors.resize(net->tensorName()->size());
|
|
initConstTensorsNoAlloc(sharedConst->allTensors, net0);
|
|
for (int i = 0; i < sharedConst->allTensors.size(); ++i) {
|
|
if (sharedConst->allTensors[i].get() != nullptr) {
|
|
noneedComputeIndexes.insert(i);
|
|
}
|
|
}
|
|
|
|
std::vector<bool> keepOp(net->oplists()->size(), false);
|
|
{
|
|
auto subModulesInfo =
|
|
_createSubModuleInfo(net0, inputIndexes, outputIndexes, noneedComputeIndexes, sharedConst);
|
|
for (int moduleIndex = 0; moduleIndex < subModulesInfo.size(); ++moduleIndex) {
|
|
auto moduleInfo = subModulesInfo[moduleIndex];
|
|
for (auto& index : moduleInfo.opList) {
|
|
keepOp[index] = true;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Split Module
|
|
auto subModulesInfo = _createSubModuleInfo(net0, inputIndexes, outputIndexes, noneedComputeIndexes, sharedConst);
|
|
// TODO: Insert pass to split submodule to npu and not npu
|
|
std::map<std::string, std::vector<std::string>> merges;
|
|
std::vector<std::shared_ptr<NetT>> allNets;
|
|
for (int inputIndex=0; inputIndex < inputs.size(); ++inputIndex) {
|
|
auto bufferPairTmp =
|
|
std::make_pair(BuilderTmp[inputIndex].GetBufferPointer(), BuilderTmp[inputIndex].GetSize());
|
|
auto net = GetNet(bufferPairTmp.first);
|
|
std::map<int, MNN::Express::VARP> stackes;
|
|
// Compute module's io
|
|
for (int i=0; i<net->tensorName()->size(); ++i) {
|
|
auto tname = net->tensorName()->GetAsString(i)->str();
|
|
for (int j=0; j<inputNames.size(); ++j) {
|
|
if (tname == inputNames[j]) {
|
|
stackes.insert(std::make_pair(i, inputs[inputIndex][j]));
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
std::vector<SubModuleIO> moduleIO(subModulesInfo.size());
|
|
for (int i=0; i<subModulesInfo.size(); ++i) {
|
|
auto& current = subModulesInfo[i];
|
|
std::vector<MNN::Express::VARP> subInputs;
|
|
for (auto index : current.inputs) {
|
|
subInputs.emplace_back(stackes[index]);
|
|
}
|
|
moduleIO[i] = _getSubModuleIO(subInputs, current, bufferPairTmp.first, bufferPairTmp.second, srcMNN);
|
|
for (int j=0; j<current.outputs.size(); ++j) {
|
|
stackes.insert(std::make_pair(current.outputs[j], moduleIO[i].outputs[j]));
|
|
}
|
|
}
|
|
for (int i=0; i<subModulesInfo.size(); ++i) {
|
|
if (subModulesInfo[i].isBreak) {
|
|
continue;
|
|
}
|
|
bool hasConvolution = false;
|
|
for (auto opIndex : subModulesInfo[i].opList) {
|
|
auto op = net->oplists()->GetAs<Op>(opIndex);
|
|
if (op->type() == OpType_Convolution) {
|
|
hasConvolution = true;
|
|
break;
|
|
}
|
|
}
|
|
if (!hasConvolution) {
|
|
subModulesInfo[i].isBreak = true;
|
|
}
|
|
}
|
|
// Compile NPU Module
|
|
std::vector<std::unique_ptr<OpT>> npuOps(subModulesInfo.size());
|
|
int npuIndex = 0;
|
|
float npuTotal = 0.0f;
|
|
float cpuTotal = 0.0f;
|
|
|
|
for (int i=0; i<subModulesInfo.size(); ++i) {
|
|
if (!subModulesInfo[i].isBreak) {
|
|
auto path = gCacheDir + "/" + gGraphName + std::to_string(npuIndex);
|
|
if (!gOfflieDst.empty()) {
|
|
path += ("." + gOfflieDst);
|
|
}
|
|
std::string srcPath;
|
|
std::string graphicName;
|
|
if (inputIndex == 0) {
|
|
srcPath = gCacheDir + "/" + gGraphName + std::to_string(npuIndex);
|
|
graphicName = gGraphName + std::to_string(npuIndex);
|
|
} else {
|
|
srcPath = gCacheDir + "/" + gGraphName + std::to_string(inputIndex) + "_" + std::to_string(npuIndex);
|
|
graphicName = gGraphName + std::to_string(inputIndex) + "_" + std::to_string(npuIndex);
|
|
}
|
|
if (!gOfflieSrc.empty()) {
|
|
srcPath += ("." + gOfflieSrc);
|
|
}
|
|
if (merges.find(path) != merges.end()) {
|
|
merges[path].emplace_back(srcPath);
|
|
} else {
|
|
merges.insert(std::make_pair(path, std::vector<std::string>{srcPath}));
|
|
}
|
|
npuOps[i] = std::move(_compileSubModule(moduleIO[i], subModulesInfo[i], bufferPairTmp.first,
|
|
bufferPairTmp.second, srcPath, srcMNN, path, cpuTotal, npuTotal,
|
|
inputIndex, graphicName));
|
|
npuIndex++;
|
|
}
|
|
}
|
|
MNN_PRINT("Total Speed Compare: NPU: %f ms : CPU: %f ms\n", npuTotal, cpuTotal);
|
|
// Merge to dst
|
|
std::shared_ptr<MNN::NetT> dstNet(flatbuffers::GetRoot<Net>(bufferPairTmp.first)->UnPack());
|
|
for (int i=0; i<keepOp.size(); ++i) {
|
|
if (dstNet->oplists[i]->inputIndexes.empty()) {
|
|
continue;
|
|
}
|
|
if (!keepOp[i]) {
|
|
dstNet->oplists[i].reset();
|
|
}
|
|
}
|
|
for (int moduleIndex=0; moduleIndex<subModulesInfo.size(); ++moduleIndex) {
|
|
auto moduleInfo = subModulesInfo[moduleIndex];
|
|
if (moduleInfo.isBreak) {
|
|
continue;
|
|
}
|
|
for (auto& index : moduleInfo.opList) {
|
|
dstNet->oplists[index].reset();
|
|
}
|
|
dstNet->oplists[moduleInfo.opList[0]] = std::move(npuOps[moduleIndex]);
|
|
}
|
|
auto oplist = std::move(dstNet->oplists);
|
|
for (auto& op : oplist) {
|
|
if (nullptr != op.get()) {
|
|
dstNet->oplists.emplace_back(std::move(op));
|
|
}
|
|
}
|
|
_reIndexTensor(dstNet.get());
|
|
_reOrderOp(dstNet.get());
|
|
allNets.emplace_back(std::move(dstNet));
|
|
}
|
|
// Fuse And Store
|
|
auto dstNet = allNets[0].get();
|
|
for (int i=1; i<allNets.size(); ++i) {
|
|
_fuse(dstNet, allNets[i].get());
|
|
allNets[i].reset();
|
|
}
|
|
flatbuffers::FlatBufferBuilder builder;
|
|
builder.Finish(Net::Pack(builder, dstNet));
|
|
std::ofstream outputOs(dstMNN, std::ios::binary);
|
|
outputOs.write((const char*)builder.GetBufferPointer(), builder.GetSize());
|
|
|
|
// Write Merge Info
|
|
rapidjson::Document resDocument;
|
|
resDocument.SetObject();
|
|
rapidjson::Value mergeMessages;
|
|
mergeMessages.SetObject();
|
|
for (auto& iter : merges) {
|
|
rapidjson::Value mergeSrc;
|
|
mergeSrc.SetArray();
|
|
for (auto& v : iter.second) {
|
|
rapidjson::Value vt;
|
|
vt.SetString(v.c_str(), resDocument.GetAllocator());
|
|
mergeSrc.GetArray().PushBack(vt, resDocument.GetAllocator());
|
|
}
|
|
rapidjson::Value key;
|
|
key.SetString(iter.first.c_str(), resDocument.GetAllocator());
|
|
mergeMessages.AddMember(key, mergeSrc, resDocument.GetAllocator());
|
|
}
|
|
{
|
|
rapidjson::Value type;
|
|
type.SetString(gNPUName.c_str(), resDocument.GetAllocator());
|
|
resDocument.AddMember("type", type, resDocument.GetAllocator());
|
|
}
|
|
resDocument.AddMember("merge", mergeMessages, resDocument.GetAllocator());
|
|
{
|
|
rapidjson::Value cachedir;
|
|
cachedir.SetString(gCacheDir.c_str(), resDocument.GetAllocator());
|
|
resDocument.AddMember("cache", cachedir, resDocument.GetAllocator());
|
|
}
|
|
rapidjson::StringBuffer buf;
|
|
rapidjson::PrettyWriter<rapidjson::StringBuffer> bufwriter(buf);
|
|
resDocument.Accept(bufwriter);
|
|
MNN_PRINT("Write config to npu_postreat.json\n");
|
|
std::ofstream os("npu_postreat.json");
|
|
os << buf.GetString();
|
|
|
|
return 0;
|
|
}
|