MCPcopy Create free account
hub / github.com/alibaba/MNN / buildKernel

Method buildKernel

codegen/SourceModule.cpp:91–162  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

89}
90
91InOutTensors SourceModule::buildKernel(std::vector<Node*> nodes, int idx) {
92 VarScope scope;
93 std::sort(nodes.begin(), nodes.end(), [](Node* x, Node* y) { return x->topoIndex < y->topoIndex; });
94 for (auto& node : nodes) {
95 mOpName.append(opStr(node->cmd->op));
96 }
97
98 mKernelName = "kernel_" + std::to_string(idx);
99 // 0. gen kernel macro
100 std::string kernelMacro = mTarget->macro();
101 // 1. gen kernel body
102 std::stringstream kernelBody;
103 kernelBody << "{\n";
104 down();
105 kernelBody << getIndent() << "OFFSET_CHECK;\n";
106 std::string offset = "offset";
107 std::unordered_map<MNN::Tensor*, std::string> cacheMap;
108 bool singleConvertRaster = false;
109 for (auto& node : nodes) {
110 auto cmd = node->cmd;
111 std::vector<std::string> inputs(cmd->inputs.size());
112 for (int i = 0; i < cmd->inputs.size(); i++) {
113 if(cmd->op->type() == MNN::OpType_Raster) {
114 singleConvertRaster = true;
115 }
116 auto t = cmd->inputs[i];
117 if (scope.hasVar(t)) {
118 inputs[i] = scope.getVar(t);
119 } else {
120 inputs[i] = scope.addVar(t);
121 std::string code;
122 if ((cmd->inputs[i]->shape().empty() || cmd->inputs[i]->elementSize() == 1) &&
123 TensorUtils::getDescribe(cmd->inputs[i])->usage == Tensor::InsideDescribe::CONSTANT) {
124 float val = cmd->inputs[i]->host<float>()[0];
125 code = mTarget->type() + inputs[i] + "=" + mTarget->number(val);
126 } else {
127 if (cmd->inputs[i]->elementSize() == 1) {
128 code = mTarget->loadscalar(scope.addInput(t), inputs[i]);
129 } else {
130 code = mTarget->load(scope.addInput(t), offset, cmd, inputs[i]);
131 }
132 }
133 kernelBody << getIndent() << code << ";\n";
134 }
135 scope.setUse(t);
136 }
137 auto tmpVar = scope.addVar(cmd->outputs[0]);
138 kernelBody << getIndent() << mTarget->type() << tmpVar << ";\n";
139 std::string computeCode = mTarget->codegen(inputs, cmd, tmpVar);
140 kernelBody << getIndent() << computeCode << ";\n";
141 }
142 scope.computeOutput();
143 auto res = scope.getIOTensors();
144 for (auto t : res.second) {
145 kernelBody << getIndent() << mTarget->store(scope.addOutput(t), offset, scope.getVar(t));
146 }
147 up();
148 kernelBody << "}\n";

Callers 1

codegenFunction · 0.45

Calls 15

opStrFunction · 0.85
appendMethod · 0.80
hasVarMethod · 0.80
getVarMethod · 0.80
addVarMethod · 0.80
shapeMethod · 0.80
addInputMethod · 0.80
setUseMethod · 0.80
computeOutputMethod · 0.80
getIOTensorsMethod · 0.80
addOutputMethod · 0.80
getIOMethod · 0.80

Tested by

no test coverage detected