wehub-resource-sync
1447 lines
54 KiB
C++
1447 lines
54 KiB
C++
//
|
|
// FuseTransformerC4.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2026/06/23.
|
|
//
|
|
|
|
#include "../PostTreatUtils.hpp"
|
|
#include <map>
|
|
#include <set>
|
|
#include <string>
|
|
#include <unordered_map>
|
|
#include <vector>
|
|
|
|
using namespace MNN;
|
|
|
|
namespace {
|
|
|
|
static bool sameConsumers(const std::unordered_map<int, std::vector<int>>& consumers, int tensor, int opIdx) {
|
|
auto iter = consumers.find(tensor);
|
|
return iter != consumers.end() && iter->second.size() == 1 && iter->second[0] == opIdx;
|
|
}
|
|
|
|
static std::unique_ptr<OpT> cloneOp(const OpT* op) {
|
|
if (op == nullptr) {
|
|
return nullptr;
|
|
}
|
|
flatbuffers::FlatBufferBuilder builder(1024);
|
|
builder.Finish(Op::Pack(builder, op));
|
|
return std::unique_ptr<OpT>(flatbuffers::GetRoot<Op>(builder.GetBufferPointer())->UnPack());
|
|
}
|
|
|
|
static std::vector<std::unique_ptr<OpT>> cloneOps(const std::vector<std::unique_ptr<OpT>>& ops) {
|
|
std::vector<std::unique_ptr<OpT>> copy;
|
|
copy.reserve(ops.size());
|
|
for (auto& op : ops) {
|
|
copy.emplace_back(cloneOp(op.get()));
|
|
}
|
|
return copy;
|
|
}
|
|
|
|
struct PreConvertMatch {
|
|
bool valid;
|
|
int reshapeIdx;
|
|
int convertIdx;
|
|
int convertOut;
|
|
int inputTensor;
|
|
int hiddenSize;
|
|
std::vector<int> convUsers;
|
|
PreConvertMatch() : valid(false), reshapeIdx(-1), convertIdx(-1), convertOut(-1), inputTensor(-1), hiddenSize(0) {}
|
|
};
|
|
|
|
struct PostConvertMatch {
|
|
bool valid;
|
|
int convIdx;
|
|
int convertIdx;
|
|
int reshapeIdx;
|
|
int convOut;
|
|
int reshapeOut;
|
|
PostConvertMatch() : valid(false), convIdx(-1), convertIdx(-1), reshapeIdx(-1), convOut(-1), reshapeOut(-1) {}
|
|
};
|
|
|
|
struct RopeInputMatch {
|
|
bool valid;
|
|
int input;
|
|
int channel;
|
|
std::vector<int> removeIndexes;
|
|
RopeInputMatch() : valid(false), input(-1), channel(0) {}
|
|
};
|
|
|
|
class TransformerC4Graph {
|
|
public:
|
|
TransformerC4Graph(std::vector<std::unique_ptr<OpT>>& ops, std::vector<std::string>& tensors)
|
|
: mOps(ops), mTensors(tensors) {}
|
|
|
|
bool run() {
|
|
if (!canFuseAsKnownTransformerGraph()) {
|
|
return false;
|
|
}
|
|
return runFusePipeline();
|
|
}
|
|
|
|
private:
|
|
bool runFusePipeline() {
|
|
bool changed = false;
|
|
changed |= fuseAttentionOutputC4();
|
|
changed |= fuseMulSilu();
|
|
changed |= fuseMlpOutputC4();
|
|
changed |= fuseRoPEInputC4();
|
|
changed |= fuseAttentionValueC4();
|
|
changed |= fuseBinaryLayerNormC4();
|
|
changed |= fuseHiddenStateC4();
|
|
changed |= fuseBinaryLayerNormC4();
|
|
return changed;
|
|
}
|
|
|
|
std::vector<std::unique_ptr<OpT>>& mOps;
|
|
std::vector<std::string>& mTensors;
|
|
std::unordered_map<int, int> mProducer;
|
|
std::unordered_map<int, std::vector<int>> mConsumers;
|
|
|
|
void rebuildMaps() {
|
|
mProducer.clear();
|
|
mConsumers.clear();
|
|
for (int i = 0; i < (int)mOps.size(); ++i) {
|
|
auto op = mOps[i].get();
|
|
if (op == nullptr) {
|
|
continue;
|
|
}
|
|
for (auto output : op->outputIndexes) {
|
|
mProducer[output] = i;
|
|
}
|
|
for (auto input : op->inputIndexes) {
|
|
mConsumers[input].push_back(i);
|
|
}
|
|
}
|
|
}
|
|
|
|
int producerOf(int tensor) const {
|
|
auto iter = mProducer.find(tensor);
|
|
if (iter == mProducer.end()) {
|
|
return -1;
|
|
}
|
|
return iter->second;
|
|
}
|
|
|
|
bool replaceInput(int opIdx, int oldTensor, int newTensor) {
|
|
if (opIdx < 0 || opIdx >= (int)mOps.size()) {
|
|
return false;
|
|
}
|
|
bool changed = false;
|
|
auto& inputs = mOps[opIdx]->inputIndexes;
|
|
for (int i = 0; i < (int)inputs.size(); ++i) {
|
|
if (inputs[i] == oldTensor) {
|
|
inputs[i] = newTensor;
|
|
changed = true;
|
|
}
|
|
}
|
|
return changed;
|
|
}
|
|
|
|
int buildTensor(const std::string& name) {
|
|
for (int i = 0; i < (int)mTensors.size(); ++i) {
|
|
if (mTensors[i] == name) {
|
|
return i;
|
|
}
|
|
}
|
|
int index = (int)mTensors.size();
|
|
mTensors.push_back(name);
|
|
return index;
|
|
}
|
|
|
|
int buildUniqueTensor(const std::string& prefix) {
|
|
int suffix = (int)mTensors.size();
|
|
while (true) {
|
|
auto name = prefix + "_" + std::to_string(suffix++);
|
|
bool exists = false;
|
|
for (int i = 0; i < (int)mTensors.size(); ++i) {
|
|
if (mTensors[i] == name) {
|
|
exists = true;
|
|
break;
|
|
}
|
|
}
|
|
if (exists) {
|
|
continue;
|
|
}
|
|
int index = (int)mTensors.size();
|
|
mTensors.push_back(name);
|
|
return index;
|
|
}
|
|
}
|
|
|
|
void removeOps(const std::set<int>& removeIndexes) {
|
|
if (removeIndexes.empty()) {
|
|
return;
|
|
}
|
|
std::vector<std::unique_ptr<OpT>> newOps;
|
|
newOps.reserve(mOps.size() - removeIndexes.size());
|
|
for (int i = 0; i < (int)mOps.size(); ++i) {
|
|
if (removeIndexes.find(i) == removeIndexes.end()) {
|
|
newOps.emplace_back(std::move(mOps[i]));
|
|
}
|
|
}
|
|
mOps.swap(newOps);
|
|
}
|
|
|
|
static bool isReshape(OpT* op) {
|
|
return op != nullptr && op->type == OpType_Reshape && op->main.type == OpParameter_Reshape &&
|
|
op->main.AsReshape() != nullptr;
|
|
}
|
|
|
|
static bool isConvert(OpT* op, MNN_DATA_FORMAT source, MNN_DATA_FORMAT dest) {
|
|
if (op == nullptr || op->type != OpType_ConvertTensor || op->main.type != OpParameter_TensorConvertInfo ||
|
|
op->main.AsTensorConvertInfo() == nullptr) {
|
|
return false;
|
|
}
|
|
auto convert = op->main.AsTensorConvertInfo();
|
|
return convert->source == source && convert->dest == dest;
|
|
}
|
|
|
|
static bool isConvolution(OpT* op) {
|
|
return op != nullptr && op->type == OpType_Convolution && op->main.type == OpParameter_Convolution2D &&
|
|
op->main.AsConvolution2D() != nullptr;
|
|
}
|
|
|
|
static bool isLayerNorm(OpT* op) {
|
|
return op != nullptr && op->type == OpType_LayerNorm && op->main.type == OpParameter_LayerNorm &&
|
|
op->main.AsLayerNorm() != nullptr;
|
|
}
|
|
|
|
static bool isBinaryOp(OpT* op, BinaryOpOperation opType) {
|
|
return op != nullptr && op->type == OpType_BinaryOp && op->main.type == OpParameter_BinaryOp &&
|
|
op->main.AsBinaryOp() != nullptr && op->main.AsBinaryOp()->opType == opType;
|
|
}
|
|
|
|
static int convolutionOutputCount(OpT* op) {
|
|
if (!isConvolution(op) || op->main.AsConvolution2D()->common == nullptr) {
|
|
return 0;
|
|
}
|
|
return op->main.AsConvolution2D()->common->outputCount;
|
|
}
|
|
|
|
static bool dimsEqual(const std::vector<int>& dims, const std::vector<int>& expected) { return dims == expected; }
|
|
|
|
int singleConsumer(int tensor) const {
|
|
auto iter = mConsumers.find(tensor);
|
|
if (iter == mConsumers.end() || iter->second.size() != 1) {
|
|
return -1;
|
|
}
|
|
return iter->second[0];
|
|
}
|
|
|
|
std::vector<int> constInt32s(int tensorIdx) const {
|
|
std::vector<int> values;
|
|
int opIdx = producerOf(tensorIdx);
|
|
if (opIdx < 0) {
|
|
return values;
|
|
}
|
|
auto op = mOps[opIdx].get();
|
|
if (op == nullptr || op->type != OpType_Const || op->main.type != OpParameter_Blob ||
|
|
op->main.AsBlob() == nullptr) {
|
|
return values;
|
|
}
|
|
auto blob = op->main.AsBlob();
|
|
if (blob->int32s.empty()) {
|
|
return values;
|
|
}
|
|
values.reserve(blob->int32s.size());
|
|
for (auto value : blob->int32s) {
|
|
values.push_back(value);
|
|
}
|
|
return values;
|
|
}
|
|
|
|
int findConstIntTensor(int value) const {
|
|
for (auto& opPtr : mOps) {
|
|
auto op = opPtr.get();
|
|
if (op == nullptr || op->outputIndexes.size() != 1) {
|
|
continue;
|
|
}
|
|
auto values = constInt32s(op->outputIndexes[0]);
|
|
if (values.size() == 1 && values[0] == value) {
|
|
return op->outputIndexes[0];
|
|
}
|
|
}
|
|
return -1;
|
|
}
|
|
|
|
RoPEParamT* ropeParam(OpT* op) const {
|
|
if (op == nullptr || op->type != OpType_RoPE || op->main.type != OpParameter_RoPEParam) {
|
|
return nullptr;
|
|
}
|
|
return op->main.AsRoPEParam();
|
|
}
|
|
|
|
std::unique_ptr<OpT> makeReshape(const std::string& name, int input, int output, const std::vector<int>& dims) {
|
|
std::unique_ptr<OpT> op(new OpT);
|
|
op->type = OpType_Reshape;
|
|
op->name = name;
|
|
op->inputIndexes = {input};
|
|
op->outputIndexes = {output};
|
|
op->main.type = OpParameter_Reshape;
|
|
op->main.value = new ReshapeT;
|
|
op->main.AsReshape()->dims = dims;
|
|
op->main.AsReshape()->dimType = MNN_DATA_FORMAT_NCHW;
|
|
op->defaultDimentionFormat = MNN_DATA_FORMAT_NCHW;
|
|
return op;
|
|
}
|
|
|
|
std::unique_ptr<OpT> makeConvert(const std::string& name, int input, int output, MNN_DATA_FORMAT source,
|
|
MNN_DATA_FORMAT dest) {
|
|
std::unique_ptr<OpT> op(new OpT);
|
|
op->type = OpType_ConvertTensor;
|
|
op->name = name;
|
|
op->inputIndexes = {input};
|
|
op->outputIndexes = {output};
|
|
op->main.type = OpParameter_TensorConvertInfo;
|
|
op->main.value = new TensorConvertInfoT;
|
|
op->main.AsTensorConvertInfo()->source = source;
|
|
op->main.AsTensorConvertInfo()->dest = dest;
|
|
op->defaultDimentionFormat = MNN_DATA_FORMAT_NCHW;
|
|
return op;
|
|
}
|
|
|
|
bool fuseAttentionOutputC4() {
|
|
rebuildMaps();
|
|
std::set<int> removeIndexes;
|
|
for (int idx = 0; idx < (int)mOps.size(); ++idx) {
|
|
auto op = mOps[idx].get();
|
|
if (op->type != OpType_Attention || op->main.type != OpParameter_AttentionParam ||
|
|
op->main.AsAttentionParam() == nullptr || op->main.AsAttentionParam()->output_c4 ||
|
|
op->outputIndexes.size() != 1) {
|
|
continue;
|
|
}
|
|
int attentionOut = op->outputIndexes[0];
|
|
auto reshapeUsers = mConsumers[attentionOut];
|
|
if (reshapeUsers.size() != 1) {
|
|
continue;
|
|
}
|
|
int reshapeIdx = reshapeUsers[0];
|
|
auto reshape = mOps[reshapeIdx].get();
|
|
if (!isReshape(reshape) || reshape->outputIndexes.size() != 1) {
|
|
continue;
|
|
}
|
|
auto dims = reshape->main.AsReshape()->dims;
|
|
if (dims.size() != 4 || dims[2] != 1 || dims[3] != 1) {
|
|
continue;
|
|
}
|
|
auto convertUsers = mConsumers[reshape->outputIndexes[0]];
|
|
if (convertUsers.size() != 1) {
|
|
continue;
|
|
}
|
|
int convertIdx = convertUsers[0];
|
|
auto convert = mOps[convertIdx].get();
|
|
if (!isConvert(convert, MNN_DATA_FORMAT_NCHW, MNN_DATA_FORMAT_NC4HW4) ||
|
|
convert->outputIndexes.size() != 1) {
|
|
continue;
|
|
}
|
|
auto convUsers = mConsumers[convert->outputIndexes[0]];
|
|
if (convUsers.size() != 1 || mOps[convUsers[0]]->type != OpType_Convolution) {
|
|
continue;
|
|
}
|
|
op->main.AsAttentionParam()->output_c4 = true;
|
|
op->outputIndexes = convert->outputIndexes;
|
|
removeIndexes.insert(reshapeIdx);
|
|
removeIndexes.insert(convertIdx);
|
|
}
|
|
bool changed = !removeIndexes.empty();
|
|
removeOps(removeIndexes);
|
|
return changed;
|
|
}
|
|
|
|
bool fuseMulSilu() {
|
|
rebuildMaps();
|
|
std::unordered_map<int, int> consumerCount;
|
|
for (auto& iter : mConsumers) {
|
|
consumerCount[iter.first] = (int)iter.second.size();
|
|
}
|
|
std::set<int> removeIndexes;
|
|
for (int idx = 0; idx < (int)mOps.size(); ++idx) {
|
|
auto op = mOps[idx].get();
|
|
if (op->type != OpType_BinaryOp || op->main.type != OpParameter_BinaryOp ||
|
|
op->main.AsBinaryOp()->opType != BinaryOpOperation_MUL || op->inputIndexes.size() != 2) {
|
|
continue;
|
|
}
|
|
int siluPos = -1;
|
|
OpT* siluOp = nullptr;
|
|
for (int i = 0; i < 2; ++i) {
|
|
int producerIdx = producerOf(op->inputIndexes[i]);
|
|
if (producerIdx < 0 || consumerCount[op->inputIndexes[i]] != 1) {
|
|
continue;
|
|
}
|
|
auto candidate = mOps[producerIdx].get();
|
|
if (candidate->type == OpType_UnaryOp && candidate->main.type == OpParameter_UnaryOp &&
|
|
candidate->main.AsUnaryOp()->opType == UnaryOpOperation_SILU) {
|
|
siluPos = i;
|
|
siluOp = candidate;
|
|
break;
|
|
}
|
|
}
|
|
if (siluOp == nullptr || siluOp->inputIndexes.empty() || siluOp->outputIndexes.empty()) {
|
|
continue;
|
|
}
|
|
int gateInput = siluOp->inputIndexes[0];
|
|
int upInput = op->inputIndexes[1 - siluPos];
|
|
int gateProducerIdx = producerOf(gateInput);
|
|
int upProducerIdx = producerOf(upInput);
|
|
if (gateProducerIdx < 0 || upProducerIdx < 0) {
|
|
continue;
|
|
}
|
|
auto gateProducer = mOps[gateProducerIdx].get();
|
|
auto upProducer = mOps[upProducerIdx].get();
|
|
if (!isReshape(gateProducer) || !isReshape(upProducer) ||
|
|
gateProducer->main.AsReshape()->dims != upProducer->main.AsReshape()->dims) {
|
|
continue;
|
|
}
|
|
op->inputIndexes = {upInput, gateInput};
|
|
op->main.AsBinaryOp()->opType = BinaryOpOperation_MUL_SILU;
|
|
int siluIdx = producerOf(siluOp->outputIndexes[0]);
|
|
if (siluIdx >= 0) {
|
|
removeIndexes.insert(siluIdx);
|
|
}
|
|
}
|
|
bool changed = !removeIndexes.empty();
|
|
removeOps(removeIndexes);
|
|
return changed;
|
|
}
|
|
|
|
PostConvertMatch matchLinearPost(int tensorIdx) const {
|
|
PostConvertMatch match;
|
|
int reshapeIdx = producerOf(tensorIdx);
|
|
if (reshapeIdx < 0) {
|
|
return match;
|
|
}
|
|
auto reshape = mOps[reshapeIdx].get();
|
|
if (!isReshape(reshape) || reshape->inputIndexes.size() != 1 || reshape->main.AsReshape()->dims.size() != 3) {
|
|
return match;
|
|
}
|
|
int convertIdx = producerOf(reshape->inputIndexes[0]);
|
|
if (convertIdx < 0) {
|
|
return match;
|
|
}
|
|
auto convert = mOps[convertIdx].get();
|
|
if (!isConvert(convert, MNN_DATA_FORMAT_NC4HW4, MNN_DATA_FORMAT_NCHW) || convert->inputIndexes.size() != 1) {
|
|
return match;
|
|
}
|
|
int convIdx = producerOf(convert->inputIndexes[0]);
|
|
if (convIdx < 0 || mOps[convIdx]->type != OpType_Convolution) {
|
|
return match;
|
|
}
|
|
match.valid = true;
|
|
match.convIdx = convIdx;
|
|
match.convertIdx = convertIdx;
|
|
match.reshapeIdx = reshapeIdx;
|
|
match.convOut = convert->inputIndexes[0];
|
|
match.reshapeOut = tensorIdx;
|
|
return match;
|
|
}
|
|
|
|
bool fuseMlpOutputC4() {
|
|
rebuildMaps();
|
|
std::set<int> removeIndexes;
|
|
for (int idx = 0; idx < (int)mOps.size(); ++idx) {
|
|
auto op = mOps[idx].get();
|
|
if (op->type != OpType_BinaryOp || op->main.type != OpParameter_BinaryOp ||
|
|
op->main.AsBinaryOp()->opType != BinaryOpOperation_MUL_SILU || op->inputIndexes.size() != 2 ||
|
|
op->outputIndexes.size() != 1) {
|
|
continue;
|
|
}
|
|
auto upPost = matchLinearPost(op->inputIndexes[0]);
|
|
auto gatePost = matchLinearPost(op->inputIndexes[1]);
|
|
if (!upPost.valid || !gatePost.valid) {
|
|
continue;
|
|
}
|
|
auto downUsers = mConsumers[op->outputIndexes[0]];
|
|
if (downUsers.size() != 1) {
|
|
continue;
|
|
}
|
|
int downReshapeIdx = downUsers[0];
|
|
auto downReshape = mOps[downReshapeIdx].get();
|
|
if (!isReshape(downReshape) || downReshape->outputIndexes.size() != 1) {
|
|
continue;
|
|
}
|
|
auto downConvertUsers = mConsumers[downReshape->outputIndexes[0]];
|
|
if (downConvertUsers.size() != 1) {
|
|
continue;
|
|
}
|
|
int downConvertIdx = downConvertUsers[0];
|
|
auto downConvert = mOps[downConvertIdx].get();
|
|
if (!isConvert(downConvert, MNN_DATA_FORMAT_NCHW, MNN_DATA_FORMAT_NC4HW4) ||
|
|
downConvert->outputIndexes.size() != 1) {
|
|
continue;
|
|
}
|
|
auto downConvUsers = mConsumers[downConvert->outputIndexes[0]];
|
|
if (downConvUsers.size() != 1 || mOps[downConvUsers[0]]->type != OpType_Convolution) {
|
|
continue;
|
|
}
|
|
op->inputIndexes = {mOps[upPost.convIdx]->outputIndexes[0], mOps[gatePost.convIdx]->outputIndexes[0]};
|
|
op->outputIndexes = downConvert->outputIndexes;
|
|
removeIndexes.insert(upPost.convertIdx);
|
|
removeIndexes.insert(upPost.reshapeIdx);
|
|
removeIndexes.insert(gatePost.convertIdx);
|
|
removeIndexes.insert(gatePost.reshapeIdx);
|
|
removeIndexes.insert(downReshapeIdx);
|
|
removeIndexes.insert(downConvertIdx);
|
|
}
|
|
bool changed = !removeIndexes.empty();
|
|
removeOps(removeIndexes);
|
|
return changed;
|
|
}
|
|
|
|
bool fuseBinaryLayerNormC4() {
|
|
rebuildMaps();
|
|
std::set<int> removeIndexes;
|
|
std::set<int> fusedLayerNorms;
|
|
for (int idx = 0; idx < (int)mOps.size(); ++idx) {
|
|
auto op = mOps[idx].get();
|
|
if (!isBinaryOp(op, BinaryOpOperation_ADD) || op->inputIndexes.size() != 2 ||
|
|
op->outputIndexes.size() != 1) {
|
|
continue;
|
|
}
|
|
int layerNormIdx = -1;
|
|
auto users = mConsumers[op->outputIndexes[0]];
|
|
for (int i = 0; i < (int)users.size(); ++i) {
|
|
int userIdx = users[i];
|
|
if (fusedLayerNorms.find(userIdx) != fusedLayerNorms.end()) {
|
|
continue;
|
|
}
|
|
auto user = mOps[userIdx].get();
|
|
if (!isLayerNorm(user) || user->inputIndexes != op->outputIndexes || user->outputIndexes.size() != 1) {
|
|
continue;
|
|
}
|
|
layerNormIdx = userIdx;
|
|
break;
|
|
}
|
|
if (layerNormIdx < 0 || layerNormIdx <= idx) {
|
|
continue;
|
|
}
|
|
bool hasEarlierUser = false;
|
|
for (int userIdx : users) {
|
|
if (userIdx != layerNormIdx && userIdx < layerNormIdx) {
|
|
hasEarlierUser = true;
|
|
break;
|
|
}
|
|
}
|
|
if (hasEarlierUser) {
|
|
continue;
|
|
}
|
|
auto layerNorm = mOps[layerNormIdx].get();
|
|
int normOutput = layerNorm->outputIndexes[0];
|
|
layerNorm->inputIndexes = op->inputIndexes;
|
|
layerNorm->outputIndexes = {op->outputIndexes[0], normOutput};
|
|
layerNorm->defaultDimentionFormat = MNN_DATA_FORMAT_NC4HW4;
|
|
removeIndexes.insert(idx);
|
|
fusedLayerNorms.insert(layerNormIdx);
|
|
}
|
|
bool changed = !removeIndexes.empty();
|
|
removeOps(removeIndexes);
|
|
return changed;
|
|
}
|
|
|
|
RopeInputMatch matchRoPEInput(int tensorIdx, int ropeIdx) const {
|
|
RopeInputMatch match;
|
|
std::vector<int> removeChain;
|
|
int current = tensorIdx;
|
|
while (true) {
|
|
int opIdx = producerOf(current);
|
|
if (opIdx < 0) {
|
|
return match;
|
|
}
|
|
auto op = mOps[opIdx].get();
|
|
if (!isReshape(op)) {
|
|
break;
|
|
}
|
|
if (removeChain.empty()) {
|
|
if (!sameConsumers(mConsumers, current, ropeIdx)) {
|
|
return match;
|
|
}
|
|
} else if (mConsumers.find(current) == mConsumers.end() || mConsumers.at(current).size() != 1) {
|
|
return match;
|
|
}
|
|
if (op->inputIndexes.empty()) {
|
|
return match;
|
|
}
|
|
removeChain.push_back(opIdx);
|
|
current = op->inputIndexes[0];
|
|
}
|
|
if (removeChain.empty()) {
|
|
return match;
|
|
}
|
|
int convertIdx = producerOf(current);
|
|
if (convertIdx < 0) {
|
|
return match;
|
|
}
|
|
auto convert = mOps[convertIdx].get();
|
|
if (!isConvert(convert, MNN_DATA_FORMAT_NC4HW4, MNN_DATA_FORMAT_NCHW) ||
|
|
mConsumers.find(current) == mConsumers.end() || mConsumers.at(current).size() != 1 ||
|
|
convert->inputIndexes.size() != 1 || convert->outputIndexes.size() != 1) {
|
|
return match;
|
|
}
|
|
int convIdx = producerOf(convert->inputIndexes[0]);
|
|
if (convIdx < 0) {
|
|
return match;
|
|
}
|
|
auto conv = mOps[convIdx].get();
|
|
if (conv->type != OpType_Convolution || conv->outputIndexes.size() != 1 ||
|
|
!sameConsumers(mConsumers, conv->outputIndexes[0], convertIdx)) {
|
|
return match;
|
|
}
|
|
int outputCount = convolutionOutputCount(conv);
|
|
auto postReshape = mOps[removeChain.back()].get();
|
|
auto postDims = postReshape->main.AsReshape()->dims;
|
|
if (outputCount <= 0 || postDims.size() != 3 || postDims[0] != 1 || postDims[2] != outputCount ||
|
|
outputCount % 4 != 0) {
|
|
return match;
|
|
}
|
|
match.valid = true;
|
|
match.input = conv->outputIndexes[0];
|
|
match.channel = outputCount;
|
|
match.removeIndexes = removeChain;
|
|
match.removeIndexes.push_back(convertIdx);
|
|
return match;
|
|
}
|
|
|
|
bool fuseRoPEInputC4() {
|
|
rebuildMaps();
|
|
std::set<int> removeIndexes;
|
|
for (int idx = 0; idx < (int)mOps.size(); ++idx) {
|
|
auto op = mOps[idx].get();
|
|
auto param = ropeParam(op);
|
|
if (param == nullptr || op->inputIndexes.size() < 2 || param->num_head > 0 || param->kv_num_head > 0 ||
|
|
param->head_dim > 0) {
|
|
continue;
|
|
}
|
|
auto qMatch = matchRoPEInput(op->inputIndexes[0], idx);
|
|
auto kMatch = matchRoPEInput(op->inputIndexes[1], idx);
|
|
if (!qMatch.valid || !kMatch.valid) {
|
|
continue;
|
|
}
|
|
int headDim = param->rope_cut_head_dim;
|
|
if (headDim <= 0 || qMatch.channel % headDim != 0 || kMatch.channel % headDim != 0 || headDim % 4 != 0) {
|
|
continue;
|
|
}
|
|
op->inputIndexes[0] = qMatch.input;
|
|
op->inputIndexes[1] = kMatch.input;
|
|
param->num_head = qMatch.channel / headDim;
|
|
param->kv_num_head = kMatch.channel / headDim;
|
|
param->head_dim = headDim;
|
|
removeIndexes.insert(qMatch.removeIndexes.begin(), qMatch.removeIndexes.end());
|
|
removeIndexes.insert(kMatch.removeIndexes.begin(), kMatch.removeIndexes.end());
|
|
}
|
|
bool changed = !removeIndexes.empty();
|
|
removeOps(removeIndexes);
|
|
return changed;
|
|
}
|
|
|
|
RopeInputMatch matchAttentionValue(int tensorIdx, int attentionIdx) const {
|
|
RopeInputMatch match;
|
|
std::vector<int> removeChain;
|
|
int current = tensorIdx;
|
|
while (true) {
|
|
int opIdx = producerOf(current);
|
|
if (opIdx < 0) {
|
|
return match;
|
|
}
|
|
auto op = mOps[opIdx].get();
|
|
if (!isReshape(op)) {
|
|
break;
|
|
}
|
|
if (removeChain.empty()) {
|
|
if (!sameConsumers(mConsumers, current, attentionIdx)) {
|
|
return match;
|
|
}
|
|
} else if (mConsumers.find(current) == mConsumers.end() || mConsumers.at(current).size() != 1) {
|
|
return match;
|
|
}
|
|
if (op->inputIndexes.empty()) {
|
|
return match;
|
|
}
|
|
removeChain.push_back(opIdx);
|
|
current = op->inputIndexes[0];
|
|
}
|
|
if (removeChain.empty()) {
|
|
return match;
|
|
}
|
|
int convertIdx = producerOf(current);
|
|
if (convertIdx < 0) {
|
|
return match;
|
|
}
|
|
auto convert = mOps[convertIdx].get();
|
|
if (!isConvert(convert, MNN_DATA_FORMAT_NC4HW4, MNN_DATA_FORMAT_NCHW) ||
|
|
mConsumers.find(current) == mConsumers.end() || mConsumers.at(current).size() != 1 ||
|
|
convert->inputIndexes.size() != 1 || convert->outputIndexes.size() != 1) {
|
|
return match;
|
|
}
|
|
int convIdx = producerOf(convert->inputIndexes[0]);
|
|
if (convIdx < 0) {
|
|
return match;
|
|
}
|
|
auto conv = mOps[convIdx].get();
|
|
if (conv->type != OpType_Convolution || conv->outputIndexes.size() != 1 ||
|
|
!sameConsumers(mConsumers, conv->outputIndexes[0], convertIdx)) {
|
|
return match;
|
|
}
|
|
int outputCount = convolutionOutputCount(conv);
|
|
auto postReshape = mOps[removeChain.back()].get();
|
|
auto postDims = postReshape->main.AsReshape()->dims;
|
|
if (outputCount <= 0 || postDims.size() != 3 || postDims[0] != 1 || postDims[2] != outputCount ||
|
|
outputCount % 4 != 0) {
|
|
return match;
|
|
}
|
|
match.valid = true;
|
|
match.input = conv->outputIndexes[0];
|
|
match.removeIndexes = removeChain;
|
|
match.removeIndexes.push_back(convertIdx);
|
|
return match;
|
|
}
|
|
|
|
bool fuseAttentionValueC4() {
|
|
rebuildMaps();
|
|
std::set<int> removeIndexes;
|
|
for (int idx = 0; idx < (int)mOps.size(); ++idx) {
|
|
auto op = mOps[idx].get();
|
|
if (op->type != OpType_Attention || op->inputIndexes.size() < 3) {
|
|
continue;
|
|
}
|
|
auto match = matchAttentionValue(op->inputIndexes[2], idx);
|
|
if (!match.valid) {
|
|
continue;
|
|
}
|
|
op->inputIndexes[2] = match.input;
|
|
removeIndexes.insert(match.removeIndexes.begin(), match.removeIndexes.end());
|
|
}
|
|
bool changed = !removeIndexes.empty();
|
|
removeOps(removeIndexes);
|
|
return changed;
|
|
}
|
|
|
|
PreConvertMatch matchPreConvertFromConv(int convIdx) const {
|
|
PreConvertMatch match;
|
|
if (convIdx < 0 || convIdx >= (int)mOps.size()) {
|
|
return match;
|
|
}
|
|
auto conv = mOps[convIdx].get();
|
|
if (!isConvolution(conv) || conv->inputIndexes.empty()) {
|
|
return match;
|
|
}
|
|
int convertIdx = producerOf(conv->inputIndexes[0]);
|
|
if (convertIdx < 0) {
|
|
return match;
|
|
}
|
|
auto convert = mOps[convertIdx].get();
|
|
if (!isConvert(convert, MNN_DATA_FORMAT_NCHW, MNN_DATA_FORMAT_NC4HW4) || convert->inputIndexes.size() != 1 ||
|
|
convert->outputIndexes.size() != 1 || convert->outputIndexes[0] != conv->inputIndexes[0]) {
|
|
return match;
|
|
}
|
|
int reshapeIdx = producerOf(convert->inputIndexes[0]);
|
|
if (reshapeIdx < 0) {
|
|
return match;
|
|
}
|
|
auto reshape = mOps[reshapeIdx].get();
|
|
if (!isReshape(reshape) || reshape->inputIndexes.size() != 1 || reshape->outputIndexes.size() != 1 ||
|
|
reshape->outputIndexes[0] != convert->inputIndexes[0]) {
|
|
return match;
|
|
}
|
|
const auto& dims = reshape->main.AsReshape()->dims;
|
|
if (dims.size() != 4 || dims[0] != -1 || dims[1] <= 0 || dims[2] != 1 || dims[3] != 1 || dims[1] % 4 != 0) {
|
|
return match;
|
|
}
|
|
int convertOut = convert->outputIndexes[0];
|
|
if (!sameConsumers(mConsumers, reshape->outputIndexes[0], convertIdx)) {
|
|
return match;
|
|
}
|
|
auto consumers = mConsumers.find(convertOut);
|
|
if (consumers != mConsumers.end()) {
|
|
for (int userIdx : consumers->second) {
|
|
if (isConvolution(mOps[userIdx].get())) {
|
|
match.convUsers.push_back(userIdx);
|
|
}
|
|
}
|
|
}
|
|
if (match.convUsers.empty()) {
|
|
return match;
|
|
}
|
|
match.valid = true;
|
|
match.reshapeIdx = reshapeIdx;
|
|
match.convertIdx = convertIdx;
|
|
match.convertOut = convertOut;
|
|
match.inputTensor = reshape->inputIndexes[0];
|
|
match.hiddenSize = dims[1];
|
|
return match;
|
|
}
|
|
|
|
std::vector<PreConvertMatch> findPreConvertMatchesFromInput(int inputTensor, int hiddenSize) const {
|
|
std::vector<PreConvertMatch> matches;
|
|
std::set<int> seenConvert;
|
|
auto consumers = mConsumers.find(inputTensor);
|
|
if (consumers == mConsumers.end()) {
|
|
return matches;
|
|
}
|
|
for (int reshapeIdx : consumers->second) {
|
|
auto reshape = mOps[reshapeIdx].get();
|
|
if (!isReshape(reshape) || reshape->inputIndexes.size() != 1 || reshape->outputIndexes.size() != 1) {
|
|
continue;
|
|
}
|
|
const auto& dims = reshape->main.AsReshape()->dims;
|
|
if (!dimsEqual(dims, std::vector<int>{-1, hiddenSize, 1, 1})) {
|
|
continue;
|
|
}
|
|
int convertIdx = singleConsumer(reshape->outputIndexes[0]);
|
|
if (convertIdx < 0 || seenConvert.find(convertIdx) != seenConvert.end()) {
|
|
continue;
|
|
}
|
|
auto convert = mOps[convertIdx].get();
|
|
if (!isConvert(convert, MNN_DATA_FORMAT_NCHW, MNN_DATA_FORMAT_NC4HW4) ||
|
|
convert->inputIndexes.size() != 1 || convert->outputIndexes.size() != 1) {
|
|
continue;
|
|
}
|
|
PreConvertMatch match;
|
|
match.reshapeIdx = reshapeIdx;
|
|
match.convertIdx = convertIdx;
|
|
match.convertOut = convert->outputIndexes[0];
|
|
match.inputTensor = inputTensor;
|
|
match.hiddenSize = hiddenSize;
|
|
auto convUsers = mConsumers.find(match.convertOut);
|
|
if (convUsers == mConsumers.end()) {
|
|
continue;
|
|
}
|
|
for (int userIdx : convUsers->second) {
|
|
if (isConvolution(mOps[userIdx].get())) {
|
|
match.convUsers.push_back(userIdx);
|
|
}
|
|
}
|
|
if (match.convUsers.empty()) {
|
|
continue;
|
|
}
|
|
match.valid = true;
|
|
seenConvert.insert(convertIdx);
|
|
matches.push_back(match);
|
|
}
|
|
return matches;
|
|
}
|
|
|
|
bool findPreConvertForConv(const std::vector<PreConvertMatch>& matches, int convIdx, PreConvertMatch* match) const {
|
|
for (auto& candidate : matches) {
|
|
for (int userIdx : candidate.convUsers) {
|
|
if (userIdx == convIdx) {
|
|
if (match != nullptr) {
|
|
*match = candidate;
|
|
}
|
|
return true;
|
|
}
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
PostConvertMatch matchPostConvertFromConv(int convIdx, int hiddenSize) const {
|
|
PostConvertMatch match;
|
|
if (convIdx < 0 || convIdx >= (int)mOps.size()) {
|
|
return match;
|
|
}
|
|
auto conv = mOps[convIdx].get();
|
|
if (!isConvolution(conv) || conv->outputIndexes.size() != 1) {
|
|
return match;
|
|
}
|
|
int convOut = conv->outputIndexes[0];
|
|
int convertIdx = singleConsumer(convOut);
|
|
if (convertIdx < 0) {
|
|
return match;
|
|
}
|
|
auto convert = mOps[convertIdx].get();
|
|
if (!isConvert(convert, MNN_DATA_FORMAT_NC4HW4, MNN_DATA_FORMAT_NCHW) || convert->inputIndexes.size() != 1 ||
|
|
convert->inputIndexes[0] != convOut || convert->outputIndexes.size() != 1) {
|
|
return match;
|
|
}
|
|
int reshapeIdx = singleConsumer(convert->outputIndexes[0]);
|
|
if (reshapeIdx < 0) {
|
|
return match;
|
|
}
|
|
auto reshape = mOps[reshapeIdx].get();
|
|
if (!isReshape(reshape) || reshape->inputIndexes.size() != 1 ||
|
|
reshape->inputIndexes[0] != convert->outputIndexes[0] || reshape->outputIndexes.size() != 1 ||
|
|
!dimsEqual(reshape->main.AsReshape()->dims, std::vector<int>{1, -1, hiddenSize})) {
|
|
return match;
|
|
}
|
|
match.valid = true;
|
|
match.convIdx = convIdx;
|
|
match.convertIdx = convertIdx;
|
|
match.reshapeIdx = reshapeIdx;
|
|
match.convOut = convOut;
|
|
match.reshapeOut = reshape->outputIndexes[0];
|
|
return match;
|
|
}
|
|
|
|
int findAddConsumerWithInputs(int tensor, int otherTensor) const {
|
|
auto consumers = mConsumers.find(tensor);
|
|
if (consumers == mConsumers.end()) {
|
|
return -1;
|
|
}
|
|
int matched = -1;
|
|
for (int userIdx : consumers->second) {
|
|
auto user = mOps[userIdx].get();
|
|
if (!isBinaryOp(user, BinaryOpOperation_ADD) || user->inputIndexes.size() != 2 ||
|
|
user->outputIndexes.size() != 1) {
|
|
continue;
|
|
}
|
|
bool hasTensor = user->inputIndexes[0] == tensor || user->inputIndexes[1] == tensor;
|
|
bool hasOther = user->inputIndexes[0] == otherTensor || user->inputIndexes[1] == otherTensor;
|
|
if (!hasTensor || !hasOther) {
|
|
continue;
|
|
}
|
|
if (matched >= 0) {
|
|
return -1;
|
|
}
|
|
matched = userIdx;
|
|
}
|
|
return matched;
|
|
}
|
|
|
|
int findLayerNormConsumer(int tensor) const {
|
|
auto consumers = mConsumers.find(tensor);
|
|
if (consumers == mConsumers.end()) {
|
|
return -1;
|
|
}
|
|
int matched = -1;
|
|
for (int userIdx : consumers->second) {
|
|
auto user = mOps[userIdx].get();
|
|
if (!isLayerNorm(user) || user->inputIndexes.size() != 1 || user->inputIndexes[0] != tensor ||
|
|
user->outputIndexes.size() != 1) {
|
|
continue;
|
|
}
|
|
if (matched >= 0) {
|
|
return -1;
|
|
}
|
|
matched = userIdx;
|
|
}
|
|
return matched;
|
|
}
|
|
|
|
int findLayerNormConsumerWithInputs(int tensor, int otherTensor) const {
|
|
auto consumers = mConsumers.find(tensor);
|
|
if (consumers == mConsumers.end()) {
|
|
return -1;
|
|
}
|
|
int matched = -1;
|
|
for (int userIdx : consumers->second) {
|
|
auto user = mOps[userIdx].get();
|
|
if (!isLayerNorm(user) || user->inputIndexes.size() != 2 || user->outputIndexes.size() != 2) {
|
|
continue;
|
|
}
|
|
bool hasTensor = user->inputIndexes[0] == tensor || user->inputIndexes[1] == tensor;
|
|
bool hasOther = user->inputIndexes[0] == otherTensor || user->inputIndexes[1] == otherTensor;
|
|
if (!hasTensor || !hasOther) {
|
|
continue;
|
|
}
|
|
if (matched >= 0) {
|
|
return -1;
|
|
}
|
|
matched = userIdx;
|
|
}
|
|
return matched;
|
|
}
|
|
|
|
int findSingleConvolutionConsumer(int tensor) const {
|
|
auto consumers = mConsumers.find(tensor);
|
|
if (consumers == mConsumers.end()) {
|
|
return -1;
|
|
}
|
|
int matched = -1;
|
|
for (int userIdx : consumers->second) {
|
|
if (!isConvolution(mOps[userIdx].get())) {
|
|
continue;
|
|
}
|
|
if (matched >= 0) {
|
|
return -1;
|
|
}
|
|
matched = userIdx;
|
|
}
|
|
return matched;
|
|
}
|
|
|
|
int findSliceConsumer(int tensor) const {
|
|
auto consumers = mConsumers.find(tensor);
|
|
if (consumers == mConsumers.end()) {
|
|
return -1;
|
|
}
|
|
int matched = -1;
|
|
for (int userIdx : consumers->second) {
|
|
auto user = mOps[userIdx].get();
|
|
if (user == nullptr || (user->type != OpType_Slice && user->type != OpType_StridedSlice) ||
|
|
user->inputIndexes.empty() || user->inputIndexes[0] != tensor || user->outputIndexes.size() != 1) {
|
|
continue;
|
|
}
|
|
if (matched >= 0) {
|
|
return -1;
|
|
}
|
|
matched = userIdx;
|
|
}
|
|
return matched;
|
|
}
|
|
struct HiddenBlockPlan {
|
|
int blockIdx;
|
|
int hiddenSize;
|
|
int blockHidden;
|
|
int blockInputSource;
|
|
int inputLnIdx;
|
|
int inputLnOut;
|
|
int blockReshapeIdx;
|
|
std::vector<PreConvertMatch> attentionPre;
|
|
PostConvertMatch attentionPost;
|
|
int attentionAddIdx;
|
|
int attentionAddOut;
|
|
bool attentionAddFused;
|
|
int postLnIdx;
|
|
int postLnOut;
|
|
PreConvertMatch gatePre;
|
|
PreConvertMatch upPre;
|
|
PostConvertMatch downPost;
|
|
int mlpAddIdx;
|
|
int mlpAddOut;
|
|
HiddenBlockPlan()
|
|
: blockIdx(0),
|
|
hiddenSize(0),
|
|
blockHidden(-1),
|
|
blockInputSource(-1),
|
|
inputLnIdx(-1),
|
|
inputLnOut(-1),
|
|
blockReshapeIdx(-1),
|
|
attentionAddIdx(-1),
|
|
attentionAddOut(-1),
|
|
attentionAddFused(false),
|
|
postLnIdx(-1),
|
|
postLnOut(-1),
|
|
mlpAddIdx(-1),
|
|
mlpAddOut(-1) {}
|
|
};
|
|
|
|
bool matchMlpFromPostLayerNorm(HiddenBlockPlan* plan) const {
|
|
auto preMatches = findPreConvertMatchesFromInput(plan->postLnOut, plan->hiddenSize);
|
|
if (preMatches.empty()) {
|
|
return false;
|
|
}
|
|
int mulSiluIdx = -1;
|
|
PreConvertMatch firstPre;
|
|
PreConvertMatch secondPre;
|
|
for (int idx = 0; idx < (int)mOps.size(); ++idx) {
|
|
auto op = mOps[idx].get();
|
|
if (!isBinaryOp(op, BinaryOpOperation_MUL_SILU) || op->inputIndexes.size() != 2 ||
|
|
op->outputIndexes.size() != 1) {
|
|
continue;
|
|
}
|
|
int firstConv = producerOf(op->inputIndexes[0]);
|
|
int secondConv = producerOf(op->inputIndexes[1]);
|
|
PreConvertMatch candidateFirst;
|
|
PreConvertMatch candidateSecond;
|
|
if (firstConv < 0 || secondConv < 0 || firstConv == secondConv ||
|
|
!findPreConvertForConv(preMatches, firstConv, &candidateFirst) ||
|
|
!findPreConvertForConv(preMatches, secondConv, &candidateSecond)) {
|
|
continue;
|
|
}
|
|
if (mulSiluIdx >= 0) {
|
|
return false;
|
|
}
|
|
mulSiluIdx = idx;
|
|
firstPre = candidateFirst;
|
|
secondPre = candidateSecond;
|
|
}
|
|
if (mulSiluIdx < 0) {
|
|
return false;
|
|
}
|
|
auto mulSilu = mOps[mulSiluIdx].get();
|
|
int downConvIdx = findSingleConvolutionConsumer(mulSilu->outputIndexes[0]);
|
|
if (downConvIdx < 0) {
|
|
return false;
|
|
}
|
|
auto downPost = matchPostConvertFromConv(downConvIdx, plan->hiddenSize);
|
|
if (!downPost.valid) {
|
|
return false;
|
|
}
|
|
int mlpAddIdx = findAddConsumerWithInputs(downPost.reshapeOut, plan->attentionAddOut);
|
|
if (mlpAddIdx < 0) {
|
|
return false;
|
|
}
|
|
plan->gatePre = firstPre;
|
|
plan->upPre = secondPre;
|
|
plan->downPost = downPost;
|
|
plan->mlpAddIdx = mlpAddIdx;
|
|
plan->mlpAddOut = mOps[mlpAddIdx]->outputIndexes[0];
|
|
return true;
|
|
}
|
|
|
|
bool matchHiddenBlockFromAttention(int attentionIdx, HiddenBlockPlan* plan) const {
|
|
if (attentionIdx < 0 || attentionIdx >= (int)mOps.size() || plan == nullptr) {
|
|
return false;
|
|
}
|
|
auto attention = mOps[attentionIdx].get();
|
|
if (attention == nullptr || attention->type != OpType_Attention || attention->inputIndexes.size() < 3 ||
|
|
attention->outputIndexes.size() != 1) {
|
|
return false;
|
|
}
|
|
int ropeIdx = producerOf(attention->inputIndexes[0]);
|
|
if (ropeIdx < 0 || ropeIdx != producerOf(attention->inputIndexes[1])) {
|
|
return false;
|
|
}
|
|
auto rope = mOps[ropeIdx].get();
|
|
if (rope == nullptr || rope->type != OpType_RoPE || rope->inputIndexes.size() < 2) {
|
|
return false;
|
|
}
|
|
int qConvIdx = producerOf(rope->inputIndexes[0]);
|
|
int kConvIdx = producerOf(rope->inputIndexes[1]);
|
|
int vConvIdx = producerOf(attention->inputIndexes[2]);
|
|
if (qConvIdx < 0 || kConvIdx < 0 || vConvIdx < 0) {
|
|
return false;
|
|
}
|
|
auto qPre = matchPreConvertFromConv(qConvIdx);
|
|
auto kPre = matchPreConvertFromConv(kConvIdx);
|
|
auto vPre = matchPreConvertFromConv(vConvIdx);
|
|
if (!qPre.valid || !kPre.valid || !vPre.valid || qPre.inputTensor != kPre.inputTensor ||
|
|
qPre.inputTensor != vPre.inputTensor || qPre.hiddenSize != kPre.hiddenSize ||
|
|
qPre.hiddenSize != vPre.hiddenSize) {
|
|
return false;
|
|
}
|
|
int inputLnIdx = producerOf(qPre.inputTensor);
|
|
if (inputLnIdx < 0) {
|
|
return false;
|
|
}
|
|
auto inputLn = mOps[inputLnIdx].get();
|
|
if (!isLayerNorm(inputLn) || inputLn->inputIndexes.size() != 1 || inputLn->outputIndexes.size() != 1 ||
|
|
inputLn->outputIndexes[0] != qPre.inputTensor) {
|
|
return false;
|
|
}
|
|
int oConvIdx = findSingleConvolutionConsumer(attention->outputIndexes[0]);
|
|
if (oConvIdx < 0) {
|
|
return false;
|
|
}
|
|
auto attentionPost = matchPostConvertFromConv(oConvIdx, qPre.hiddenSize);
|
|
if (!attentionPost.valid) {
|
|
return false;
|
|
}
|
|
int blockHidden = inputLn->inputIndexes[0];
|
|
int attentionAddIdx = findAddConsumerWithInputs(attentionPost.reshapeOut, blockHidden);
|
|
int attentionAddOut = -1;
|
|
int postLnIdx = -1;
|
|
bool attentionAddFused = false;
|
|
if (attentionAddIdx >= 0) {
|
|
attentionAddOut = mOps[attentionAddIdx]->outputIndexes[0];
|
|
postLnIdx = findLayerNormConsumer(attentionAddOut);
|
|
if (postLnIdx < 0) {
|
|
return false;
|
|
}
|
|
} else {
|
|
postLnIdx = findLayerNormConsumerWithInputs(attentionPost.reshapeOut, blockHidden);
|
|
if (postLnIdx < 0) {
|
|
return false;
|
|
}
|
|
attentionAddFused = true;
|
|
attentionAddOut = mOps[postLnIdx]->outputIndexes[0];
|
|
}
|
|
|
|
plan->hiddenSize = qPre.hiddenSize;
|
|
plan->blockHidden = blockHidden;
|
|
plan->inputLnIdx = inputLnIdx;
|
|
plan->inputLnOut = qPre.inputTensor;
|
|
plan->attentionPre.push_back(qPre);
|
|
plan->attentionPre.push_back(kPre);
|
|
plan->attentionPre.push_back(vPre);
|
|
plan->attentionPost = attentionPost;
|
|
plan->attentionAddIdx = attentionAddIdx;
|
|
plan->attentionAddOut = attentionAddOut;
|
|
plan->attentionAddFused = attentionAddFused;
|
|
plan->postLnIdx = postLnIdx;
|
|
plan->postLnOut = attentionAddFused ? mOps[postLnIdx]->outputIndexes[1] : mOps[postLnIdx]->outputIndexes[0];
|
|
int blockHiddenProducer = producerOf(blockHidden);
|
|
if (blockHiddenProducer >= 0 && isReshape(mOps[blockHiddenProducer].get()) &&
|
|
!mOps[blockHiddenProducer]->inputIndexes.empty()) {
|
|
plan->blockReshapeIdx = blockHiddenProducer;
|
|
plan->blockInputSource = mOps[blockHiddenProducer]->inputIndexes[0];
|
|
} else {
|
|
plan->blockInputSource = blockHidden;
|
|
}
|
|
return matchMlpFromPostLayerNorm(plan);
|
|
}
|
|
|
|
bool buildHiddenBlockChain(std::vector<HiddenBlockPlan>* plans) const {
|
|
if (plans == nullptr) {
|
|
return false;
|
|
}
|
|
std::vector<HiddenBlockPlan> candidates;
|
|
for (int idx = 0; idx < (int)mOps.size(); ++idx) {
|
|
HiddenBlockPlan plan;
|
|
if (!matchHiddenBlockFromAttention(idx, &plan)) {
|
|
continue;
|
|
}
|
|
candidates.push_back(plan);
|
|
}
|
|
if (candidates.empty()) {
|
|
return false;
|
|
}
|
|
std::set<int> mlpOutputs;
|
|
for (auto& plan : candidates) {
|
|
mlpOutputs.insert(plan.mlpAddOut);
|
|
}
|
|
int start = -1;
|
|
for (int i = 0; i < (int)candidates.size(); ++i) {
|
|
if (mlpOutputs.find(candidates[i].blockInputSource) != mlpOutputs.end()) {
|
|
continue;
|
|
}
|
|
if (start >= 0) {
|
|
return false;
|
|
}
|
|
start = i;
|
|
}
|
|
if (start < 0) {
|
|
return false;
|
|
}
|
|
std::set<int> visited;
|
|
int current = start;
|
|
while (current >= 0) {
|
|
if (visited.find(current) != visited.end()) {
|
|
return false;
|
|
}
|
|
visited.insert(current);
|
|
candidates[current].blockIdx = (int)plans->size();
|
|
plans->push_back(candidates[current]);
|
|
|
|
int next = -1;
|
|
int currentOut = candidates[current].mlpAddOut;
|
|
for (int i = 0; i < (int)candidates.size(); ++i) {
|
|
if (visited.find(i) != visited.end() || candidates[i].blockInputSource != currentOut) {
|
|
continue;
|
|
}
|
|
if (next >= 0) {
|
|
return false;
|
|
}
|
|
next = i;
|
|
}
|
|
current = next;
|
|
}
|
|
return plans->size() == candidates.size();
|
|
}
|
|
|
|
bool validateHiddenPlans(std::vector<HiddenBlockPlan>* plans) {
|
|
rebuildMaps();
|
|
if (!buildHiddenBlockChain(plans) || plans->empty()) {
|
|
return false;
|
|
}
|
|
int hiddenSize = (*plans)[0].hiddenSize;
|
|
if (hiddenSize <= 0 || hiddenSize % 4 != 0) {
|
|
return false;
|
|
}
|
|
for (auto& plan : *plans) {
|
|
if (plan.hiddenSize != hiddenSize) {
|
|
return false;
|
|
}
|
|
}
|
|
return findSliceConsumer(plans->back().mlpAddOut) >= 0;
|
|
}
|
|
|
|
bool isCurrentGraphKnownTransformerC4Candidate() {
|
|
std::vector<HiddenBlockPlan> plans;
|
|
if (!validateHiddenPlans(&plans)) {
|
|
return false;
|
|
}
|
|
int attentionCount = 0;
|
|
int sequenceMixerCount = 0;
|
|
for (auto& opPtr : mOps) {
|
|
auto op = opPtr.get();
|
|
if (op == nullptr) {
|
|
continue;
|
|
}
|
|
if (op->type == OpType_Attention) {
|
|
++attentionCount;
|
|
++sequenceMixerCount;
|
|
} else if (op->type == OpType_LinearAttention) {
|
|
++sequenceMixerCount;
|
|
}
|
|
}
|
|
return attentionCount > 0 && attentionCount == (int)plans.size() && sequenceMixerCount == attentionCount;
|
|
}
|
|
|
|
bool canFuseAsKnownTransformerGraph() const {
|
|
auto dryRunOps = cloneOps(mOps);
|
|
auto dryRunTensors = mTensors;
|
|
TransformerC4Graph dryRunGraph(dryRunOps, dryRunTensors);
|
|
dryRunGraph.fuseAttentionOutputC4();
|
|
dryRunGraph.fuseMulSilu();
|
|
dryRunGraph.fuseMlpOutputC4();
|
|
dryRunGraph.fuseRoPEInputC4();
|
|
dryRunGraph.fuseAttentionValueC4();
|
|
dryRunGraph.fuseBinaryLayerNormC4();
|
|
return dryRunGraph.isCurrentGraphKnownTransformerC4Candidate();
|
|
}
|
|
|
|
bool fuseHiddenStateC4() {
|
|
std::vector<HiddenBlockPlan> plans;
|
|
if (!validateHiddenPlans(&plans)) {
|
|
return false;
|
|
}
|
|
int hiddenSize = plans[0].hiddenSize;
|
|
int firstHidden = plans[0].blockHidden;
|
|
int firstLayerNormIdx = plans[0].inputLnIdx;
|
|
int currentC4 = plans.back().mlpAddOut;
|
|
int finalSliceIdx = findSliceConsumer(currentC4);
|
|
if (finalSliceIdx < 0) {
|
|
return false;
|
|
}
|
|
int finalLayerNormIdx = -1;
|
|
int lmHeadPreReshapeIdx = -1;
|
|
int lmHeadPreConvertIdx = -1;
|
|
int lmHeadConvIdx = -1;
|
|
int constAxis0 = findConstIntTensor(0);
|
|
bool useC4Tail = false;
|
|
auto finalSlice = mOps[finalSliceIdx].get();
|
|
if (constAxis0 >= 0 && finalSlice->outputIndexes.size() == 1 && finalSlice->inputIndexes.size() >= 4) {
|
|
auto finalLayerNormUsers = mConsumers[finalSlice->outputIndexes[0]];
|
|
if (finalLayerNormUsers.size() == 1) {
|
|
finalLayerNormIdx = finalLayerNormUsers[0];
|
|
auto finalLayerNorm = mOps[finalLayerNormIdx].get();
|
|
if (isLayerNorm(finalLayerNorm) && finalLayerNorm->outputIndexes.size() == 1) {
|
|
auto preReshapeUsers = mConsumers[finalLayerNorm->outputIndexes[0]];
|
|
if (preReshapeUsers.size() == 1) {
|
|
lmHeadPreReshapeIdx = preReshapeUsers[0];
|
|
auto preReshape = mOps[lmHeadPreReshapeIdx].get();
|
|
if (isReshape(preReshape) && preReshape->outputIndexes.size() == 1) {
|
|
auto preConvertUsers = mConsumers[preReshape->outputIndexes[0]];
|
|
if (preConvertUsers.size() == 1) {
|
|
lmHeadPreConvertIdx = preConvertUsers[0];
|
|
auto preConvert = mOps[lmHeadPreConvertIdx].get();
|
|
if (isConvert(preConvert, MNN_DATA_FORMAT_NCHW, MNN_DATA_FORMAT_NC4HW4) &&
|
|
preConvert->outputIndexes.size() == 1) {
|
|
auto lmHeadUsers = mConsumers[preConvert->outputIndexes[0]];
|
|
if (lmHeadUsers.size() == 1 && isConvolution(mOps[lmHeadUsers[0]].get())) {
|
|
lmHeadConvIdx = lmHeadUsers[0];
|
|
useC4Tail = true;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
int initialReshape = buildUniqueTensor("FuseTransformerC4_pre_reshape");
|
|
int initialC4 = buildUniqueTensor("FuseTransformerC4_pre_convert");
|
|
int finalNchw4 = -1;
|
|
int finalNchw = -1;
|
|
|
|
std::map<int, std::vector<std::unique_ptr<OpT>>> insertBefore;
|
|
insertBefore[firstLayerNormIdx].push_back(
|
|
makeReshape(mTensors[initialReshape], firstHidden, initialReshape, {-1, hiddenSize, 1, 1}));
|
|
insertBefore[firstLayerNormIdx].push_back(
|
|
makeConvert(mTensors[initialC4], initialReshape, initialC4, MNN_DATA_FORMAT_NCHW, MNN_DATA_FORMAT_NC4HW4));
|
|
if (!useC4Tail) {
|
|
finalNchw4 = buildUniqueTensor("FuseTransformerC4_post_convert");
|
|
finalNchw = buildUniqueTensor("FuseTransformerC4_post_reshape");
|
|
insertBefore[finalSliceIdx].push_back(
|
|
makeConvert(mTensors[finalNchw4], currentC4, finalNchw4, MNN_DATA_FORMAT_NC4HW4, MNN_DATA_FORMAT_NCHW));
|
|
insertBefore[finalSliceIdx].push_back(
|
|
makeReshape(mTensors[finalNchw], finalNchw4, finalNchw, {1, -1, hiddenSize}));
|
|
}
|
|
|
|
std::set<int> removeIndexes;
|
|
currentC4 = initialC4;
|
|
for (auto& plan : plans) {
|
|
replaceInput(plan.inputLnIdx, plan.blockHidden, currentC4);
|
|
if (plan.blockIdx > 0 && plan.blockReshapeIdx >= 0) {
|
|
removeIndexes.insert(plan.blockReshapeIdx);
|
|
}
|
|
for (auto& pre : plan.attentionPre) {
|
|
for (int convIdx : pre.convUsers) {
|
|
replaceInput(convIdx, pre.convertOut, plan.inputLnOut);
|
|
}
|
|
removeIndexes.insert(pre.reshapeIdx);
|
|
removeIndexes.insert(pre.convertIdx);
|
|
}
|
|
|
|
if (plan.attentionAddFused) {
|
|
auto& layerNormInputs = mOps[plan.postLnIdx]->inputIndexes;
|
|
layerNormInputs = {currentC4, plan.attentionPost.convOut};
|
|
mOps[plan.postLnIdx]->defaultDimentionFormat = MNN_DATA_FORMAT_NC4HW4;
|
|
} else {
|
|
replaceInput(plan.attentionAddIdx, plan.blockHidden, currentC4);
|
|
replaceInput(plan.attentionAddIdx, plan.attentionPost.reshapeOut, plan.attentionPost.convOut);
|
|
auto& addInputs = mOps[plan.attentionAddIdx]->inputIndexes;
|
|
if (addInputs.size() != 2 || addInputs[0] != currentC4) {
|
|
addInputs = {currentC4, plan.attentionPost.convOut};
|
|
}
|
|
}
|
|
removeIndexes.insert(plan.attentionPost.convertIdx);
|
|
removeIndexes.insert(plan.attentionPost.reshapeIdx);
|
|
|
|
for (int convIdx : plan.gatePre.convUsers) {
|
|
replaceInput(convIdx, plan.gatePre.convertOut, plan.postLnOut);
|
|
}
|
|
removeIndexes.insert(plan.gatePre.reshapeIdx);
|
|
removeIndexes.insert(plan.gatePre.convertIdx);
|
|
for (int convIdx : plan.upPre.convUsers) {
|
|
replaceInput(convIdx, plan.upPre.convertOut, plan.postLnOut);
|
|
}
|
|
removeIndexes.insert(plan.upPre.reshapeIdx);
|
|
removeIndexes.insert(plan.upPre.convertIdx);
|
|
|
|
replaceInput(plan.mlpAddIdx, plan.downPost.reshapeOut, plan.downPost.convOut);
|
|
auto& add1Inputs = mOps[plan.mlpAddIdx]->inputIndexes;
|
|
if (add1Inputs.size() != 2 || add1Inputs[0] != plan.attentionAddOut) {
|
|
add1Inputs = {plan.attentionAddOut, plan.downPost.convOut};
|
|
}
|
|
removeIndexes.insert(plan.downPost.convertIdx);
|
|
removeIndexes.insert(plan.downPost.reshapeIdx);
|
|
currentC4 = plan.mlpAddOut;
|
|
}
|
|
if (useC4Tail) {
|
|
auto finalLayerNorm = mOps[finalLayerNormIdx].get();
|
|
auto lmHeadPreConvert = mOps[lmHeadPreConvertIdx].get();
|
|
mOps[finalSliceIdx]->inputIndexes[0] = currentC4;
|
|
mOps[finalSliceIdx]->inputIndexes[3] = constAxis0;
|
|
mOps[finalSliceIdx]->defaultDimentionFormat = MNN_DATA_FORMAT_NC4HW4;
|
|
finalLayerNorm->defaultDimentionFormat = MNN_DATA_FORMAT_NC4HW4;
|
|
replaceInput(lmHeadConvIdx, lmHeadPreConvert->outputIndexes[0], finalLayerNorm->outputIndexes[0]);
|
|
removeIndexes.insert(lmHeadPreReshapeIdx);
|
|
removeIndexes.insert(lmHeadPreConvertIdx);
|
|
} else {
|
|
mOps[finalSliceIdx]->inputIndexes[0] = finalNchw;
|
|
}
|
|
|
|
std::vector<std::unique_ptr<OpT>> newOps;
|
|
for (int idx = 0; idx < (int)mOps.size(); ++idx) {
|
|
auto insertIter = insertBefore.find(idx);
|
|
if (insertIter != insertBefore.end()) {
|
|
for (auto& insertOp : insertIter->second) {
|
|
newOps.emplace_back(std::move(insertOp));
|
|
}
|
|
}
|
|
if (removeIndexes.find(idx) == removeIndexes.end()) {
|
|
newOps.emplace_back(std::move(mOps[idx]));
|
|
}
|
|
}
|
|
mOps.swap(newOps);
|
|
return true;
|
|
}
|
|
};
|
|
|
|
} // namespace
|
|
|
|
class FuseTransformerC4 : public PostConverter {
|
|
public:
|
|
virtual bool onExecute(std::unique_ptr<MNN::NetT>& net) const override {
|
|
if (net == nullptr) {
|
|
return true;
|
|
}
|
|
TransformerC4Graph mainGraph(net->oplists, net->tensorName);
|
|
mainGraph.run();
|
|
for (auto& subgraph : net->subgraphs) {
|
|
if (subgraph == nullptr) {
|
|
continue;
|
|
}
|
|
TransformerC4Graph subGraph(subgraph->nodes, subgraph->tensors);
|
|
subGraph.run();
|
|
}
|
|
return true;
|
|
}
|
|
};
|
|
|
|
static PostConverterRegister<FuseTransformerC4> __l("FuseTransformerC4");
|