#include <string>
#include <vector>
#include <fstream>
#include <unordered_map>
#include "core/TensorUtils.hpp"
#include "MNN_generated.h"
#include "AST.hpp"
using namespace AST;
using namespace MNN;
std::vector<std::vector<Node*>> spliteNodes(std::vector<Node*>& nodes) {
std::vector<std::vector<Node*>> res;
res.push_back(std::vector<Node*>());
for (auto node : nodes) {
if (isElemWise(node->cmd->op->type())) {
res.back().push_back(node);
} else {
if (res.back().empty()) {
res.back().push_back(node);
} else {
res.push_back({node});
}
res.push_back(std::vector<Node*>());
}
}
if (res.back().empty()) {
res.pop_back();
}
return res;
}
class SourceModule::CPUPluginFunction {
public:
CPUPluginFunction(std::vector<Node*>& nodes, int idx) : nodes(nodes) {
sort(nodes.begin(), nodes.end(), [](Node* x, Node* y) { return x->topoIndex < y->topoIndex; });
std::unique_ptr<ListExpr> list(new ListExpr);
auto subNodes = spliteNodes(nodes);
for (auto subNode : subNodes) {
if (subNode.size() > 1 || subNode.back()->cmd->op->type() != OpType_Raster) {
auto loop = addElemwiseLoop(subNode);
list->push_back(std::move(loop));
} else {
auto raster = addRaster(subNode);
list->push_back(std::move(raster));
}
}
auto proto = std::make_unique<PrototypeAST>("kernel_" + std::to_string(idx), inputs.size(), outputs.size());
function = std::make_unique<FunctionAST>(std::move(proto), std::move(list));
}
std::string codegen(SourceTarget* target) {
return function->codegen(target);
}
std::vector<Tensor*> getInputs() { return inputs; }
std::vector<Tensor*> getOutputs() { return outputs; }
private:
std::unique_ptr<Expr> addElemwiseLoop(std::vector<Node*>& nodes) {
std::map<const Tensor*, int> varShape;
std::unordered_map<Tensor*, std::unique_ptr<Expr>> outMap;
for (auto& node : nodes) {
auto cmd = node->cmd;
std::vector<std::unique_ptr<Expr>> in(cmd->inputs.size());
for (int i = 0; i < cmd->inputs.size(); i++) {
if (outMap.find(cmd->inputs[i]) == outMap.end()) {
auto inputExpr = getExprByTensor(cmd->inputs[i], true);
int size = cmd->inputs[i]->elementSize();
varShape[cmd->inputs[i]] = size;
if (size > 1) {
in[i] = std::make_unique<SubscriptExpr>(std::move(inputExpr), "i");
} else {
in[i] = std::make_unique<SubscriptExpr>(std::move(inputExpr), 0);
}
} else {
in[i] = std::move(outMap[cmd->inputs[i]]);
outMap.erase(cmd->inputs[i]);
}
}
switch (cmd->op->type()) {
case MNN::OpType_BinaryOp:
{
auto type = static_cast<MNN::BinaryOpOperation>(cmd->op->main_as_BinaryOp()->opType());
outMap[cmd->outputs[0]] = std::make_unique<BinaryExpr>(type, std::move(in[0]), std::move(in[1]));
break;
}
case MNN::OpType_Eltwise:
{
std::map<EltwiseType, MNN::BinaryOpOperation> elemToBinary = {
{EltwiseType_PROD, BinaryOpOperation_MUL},
{EltwiseType_SUM, BinaryOpOperation_ADD},
{EltwiseType_MAXIMUM, BinaryOpOperation_MAXIMUM},
{EltwiseType_SUB, BinaryOpOperation_SUB}
};
auto type = elemToBinary[cmd->op->main_as_Eltwise()->type()];
auto tmp = std::make_unique<BinaryExpr>(type, std::move(in[0]), std::move(in[1]));
for (int i = 2; i < cmd->inputs.size(); i++) {
tmp = std::make_unique<BinaryExpr>(type, std::move(tmp), std::move(in[i]));
}
outMap[cmd->outputs[0]] = std::move(tmp);
break;
}
case MNN::OpType_UnaryOp:
{
auto unary = cmd->op->main_as_UnaryOp();
auto type = unary->opType();
outMap[cmd->outputs[0]] = std::make_unique<UnaryExpr>(type, std::move(in[0]));
break;
}
case MNN::OpType_ReLU6:
{
auto relu6 = cmd->op->main_as_Relu6();
float minv = relu6->minValue();
float maxv = relu6->maxValue();
outMap[cmd->outputs[0]] = std::make_unique<ReluExpr>(minv, maxv, std::move(in[0]));
break;
}
case MNN::OpType_ReLU:
{
auto relu = cmd->op->main_as_Relu();
float slope = relu->slope();
outMap[cmd->outputs[0]] = std::make_unique<ReluExpr>(slope, 0, std::move(in[0]));
break;
}
default:
break;
}
}
std::unique_ptr<Expr> content;
for (auto& iter : outMap) {
auto outputExpr = getExprByTensor(iter.first, false);
auto output = std::make_unique<SubscriptExpr>(std::move(outputExpr), "i");
varShape[iter.first] = iter.first->elementSize();
content = std::make_unique<AssignExpr>(std::move(output), std::move(iter.second));
}
int size = -1;
for (auto& iter : varShape) {
if (iter.second > 1) {
if (size > 1 && iter.second != size) {
MNN_ERROR("size not equal!\n");
exit(0);
} else {
size = iter.second;
}
}
}
auto start = std::make_unique<NumberExpr>(0);
auto end = std::make_unique<NumberExpr>(size);
auto step = std::make_unique<NumberExpr>(1);
auto loop = std::make_unique<ForExpr>("i", std::move(start), std::move(end), std::move(step), std::move(content));
return loop;
}
std::unique_ptr<Expr> addRaster(std::vector<Node*>& nodes) {
auto node = nodes.back();
auto input = node->cmd->inputs[0];
auto output = node->cmd->outputs[0];
auto des = TensorUtils::getDescribe(input);
std::string foots[3] = { "i", "j", "k" };
auto getOffset = [&foots](int strides[], int offset) {
std::unique_ptr<Expr> steps[3];
for (int i = 0; i < 3; i++) {
auto stride = std::make_unique<NumberExpr>(strides[i]);
auto foot = std::make_unique<VariableExpr>(foots[i]);
steps[i] = std::make_unique<BinaryExpr>(MNN::BinaryOpOperation_MUL, std::move(foot), std::move(stride));
}
auto res = std::make_unique<BinaryExpr>(MNN::BinaryOpOperation_ADD, std::move(steps[1]), std::move(steps[2]));
res = std::make_unique<BinaryExpr>(MNN::BinaryOpOperation_ADD, std::move(steps[0]), std::move(res));
return std::make_unique<BinaryExpr>(MNN::BinaryOpOperation_ADD, std::move(res), std::make_unique<NumberExpr>(offset));
};
std::unique_ptr<ListExpr> list(new ListExpr);
for (auto& reg : des->regions) {
auto inputExpr = getExprByTensor(reg.origin, true);
auto outputExpr = getExprByTensor(output, false);
auto srcPtr = std::make_unique<SubscriptExpr>(std::move(inputExpr), getOffset(reg.src.stride, reg.src.offset));
auto dstPtr = std::make_unique<SubscriptExpr>(std::move(outputExpr), getOffset(reg.dst.stride, reg.dst.offset));
std::unique_ptr<Expr> content = std::make_unique<AssignExpr>(std::move(dstPtr), std::move(srcPtr));
for (int i = 2; i >= 0; i--) {
auto start = std::make_unique<NumberExpr>(0);
auto end = std::make_unique<NumberExpr>(reg.size[i]);
auto step = std::make_unique<NumberExpr>(1);
content = std::make_unique<ForExpr>(foots[i], std::move(start), std::move(end), std::move(step), std::move(content));
}
list->push_back(std::move(content));
}
return list;
}
std::unique_ptr<Expr> getExprByTensor(Tensor* t, bool read) {
if (inputMap.find(t) != inputMap.end()) {
return std::make_unique<SubscriptExpr>("inputs", inputMap[t]);
}
if (outputMap.find(t) != outputMap.end()) {
return std::make_unique<SubscriptExpr>("outputs", outputMap[t]);
}
if (read) {
int idx = inputs.size();
inputs.push_back(t);
inputMap[t] = idx;
return std::make_unique<SubscriptExpr>("inputs", idx);
} else {
int idx = outputs.size();
outputs.push_back(t);
outputMap[t] = idx;
return std::make_unique<SubscriptExpr>("outputs", idx);
}
}
private:
std::vector<Node*> nodes;
std::vector<Tensor*> inputs;
std::vector<Tensor*> outputs;
std::unordered_map<const Tensor*, int> inputMap;
std::unordered_map<const Tensor*, int> outputMap;
std::unique_ptr<FunctionAST> function;
};
void CPUPluginModule::codegen(LLVMTarget* target) {
for (int i = 0; i < getFunctionNum(); i++) {
functions[i]->codegen(target);
}
}
void CPUPluginModule::codegen() {
std::ofstream headerFile("./kernel.h");
std::ofstream sourceFile("./kernel.c");
if (!headerFile.is_open()) {
return;
}
headerFile << "extern \"C\" {\n";
#ifdef MNN_CODEGEN_LLVM
std::unique_ptr<LLVMTarget> llvm(new LLVMTarget(name));
#endif
#ifdef MNN_CODEGEN_C
sourceFile << "#include \"math.h\"\n";
std::unique_ptr<SourceTarget> source(new CTarget(name));
#endif
for (int i = 0; i < getFunctionNum(); i++) {
headerFile << "void kernel_" + std::to_string(i) + "(float**, float**);\n";
#ifdef MNN_CODEGEN_C
sourceFile << functions[i]->codegen(source.get());
#endif
#ifdef MNN_CODEGEN_LLVM
functions[i]->codegen(llvm.get());
#endif
}
headerFile << "}\n";
headerFile << "void (*kernels[])(float**, float**) = {\n";
for (int i = 0; i < getFunctionNum(); i++) {
headerFile << "&kernel_" + std::to_string(i) + ",\n";
}
headerFile << "};\n";
#ifdef MNN_CODEGEN_LLVM
std::error_code EC;
llvm::raw_fd_ostream OS("kernel.bc", EC, sys::fs::F_None);
WriteBitcodeToFile(*llvm->getModule(), OS);
OS.flush();
#endif
}
InOutTensors CPUPluginModule::addFunction(std::vector<Node*> nodes) {
std::unique_ptr<CPUPluginFunction> func(new CPUPluginFunction(nodes, getFunctionNum()));
auto res = std::make_pair<std::vector<MNN::Tensor*>, std::vector<MNN::Tensor*>>(func->getInputs(), func->getOutputs());
functions.emplace_back(std::move(func));
return res;
}
CPUPluginModule::CPUPluginModule() {}
CPUPluginModule::CPUPluginModule(std::string name) : name(name) {}
CPUPluginModule::~CPUPluginModule() = default;
CPUPluginModule::CPUPluginModule(CPUPluginModule&& m) = default;
CPUPluginModule& CPUPluginModule::operator=(CPUPluginModule&& m) = default;