wehub-resource-sync
408 lines
20 KiB
C++
408 lines
20 KiB
C++
//
|
|
// LoopGrad.cpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on b'2022/10/10'.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "OpGrad.hpp"
|
|
using namespace std;
|
|
namespace MNN {
|
|
using namespace MNN::Express;
|
|
|
|
class LoopGrad : public OpGrad {
|
|
public:
|
|
static std::vector<VARP> _getGradExpr(EXPRP expr, std::vector<VARP> inputs, VARP tempOutput) {
|
|
auto gradValue = OpGrad::get(expr->get()->type())->onGrad(expr, {tempOutput});
|
|
MNN_ASSERT(gradValue.size() == inputs.size());
|
|
return gradValue;
|
|
}
|
|
static void _gradForCommandMatMul(const RegionCommandT* command, const std::map<int, int>& backwardMap, int tensorNumber, std::map<int, VARP>& extraInputs, std::map<int, VARP>& extraOutputs, std::vector<std::unique_ptr<RegionCommandT>>& dstCommands) {
|
|
auto AIndex = command->indexes[1];
|
|
auto BIndex = command->indexes[2];
|
|
auto CIndex = command->indexes[0];
|
|
auto transA = command->op->main.AsMatMul()->transposeA;
|
|
auto transB = command->op->main.AsMatMul()->transposeB;
|
|
int e = command->size[0];
|
|
int l = command->size[1];
|
|
int h = command->size[2];
|
|
|
|
if (backwardMap.find(CIndex) == backwardMap.end()) {
|
|
return;
|
|
}
|
|
std::map<int, int> originIndexMap;
|
|
for (int i=0; i<command->indexes.size(); ++i) {
|
|
auto index = command->indexes[i];
|
|
originIndexMap.insert(std::make_pair(index, i));
|
|
auto bkIter = backwardMap.find(index);
|
|
if (bkIter != backwardMap.end()) {
|
|
originIndexMap.insert(std::make_pair(bkIter->second, i));
|
|
}
|
|
}
|
|
auto CDiffIndex = backwardMap.find(CIndex)->second;
|
|
if (backwardMap.find(AIndex) != backwardMap.end()) {
|
|
auto ADiffIndex = backwardMap.find(AIndex)->second;
|
|
// Compute A Diff
|
|
std::unique_ptr<RegionCommandT> currentCommand(new RegionCommandT);
|
|
currentCommand->op.reset(new OpT);
|
|
currentCommand->op->type = OpType_MatMul;
|
|
currentCommand->op->main.value = new MatMulT;
|
|
currentCommand->op->main.type = OpParameter_MatMul;
|
|
currentCommand->indexes = {ADiffIndex, CDiffIndex, BIndex};
|
|
currentCommand->view.resize(3);
|
|
for (int j=0; j<currentCommand->view.size(); ++j) {
|
|
currentCommand->view[j].reset(new ViewT);
|
|
}
|
|
currentCommand->iterIndexes.resize(3);
|
|
currentCommand->steps.resize(3);
|
|
for (int j=0; j<currentCommand->indexes.size(); ++j) {
|
|
// Compute output info
|
|
auto index = currentCommand->indexes[j];
|
|
currentCommand->iterIndexes[j] = command->iterIndexes[originIndexMap[index]];
|
|
currentCommand->steps[j] = command->steps[originIndexMap[index]];
|
|
*currentCommand->view[j] = *command->view[originIndexMap[index]];
|
|
}
|
|
// TODO: Optimize fuse option
|
|
currentCommand->fuse = BinaryOpOperation_ADD;
|
|
|
|
// Reorder the view's stride by size order change to e, h, l
|
|
std::vector<int> order = {0, 2, 1};
|
|
currentCommand->size = {e, h, l};
|
|
for (int j=0; j<currentCommand->indexes.size(); ++j) {
|
|
auto view = currentCommand->view[j].get();
|
|
auto originStride = view->stride;
|
|
for (int k=0; k<3; ++k) {
|
|
view->stride[k] = originStride[order[k]];
|
|
}
|
|
}
|
|
// Set Transpose Info by stride
|
|
auto dstMatMulParam = currentCommand->op->main.AsMatMul();
|
|
{
|
|
auto transAView = currentCommand->view[1].get();
|
|
auto transBView = currentCommand->view[1].get();
|
|
dstMatMulParam->transposeA = transAView->stride[1] != 1;
|
|
dstMatMulParam->transposeB = transBView->stride[1] == 1;
|
|
}
|
|
dstCommands.emplace_back(std::move(currentCommand));
|
|
}
|
|
|
|
if (backwardMap.find(BIndex) != backwardMap.end()) {
|
|
auto BDiffIndex = backwardMap.find(BIndex)->second;
|
|
// Compute A Diff
|
|
std::unique_ptr<RegionCommandT> currentCommand(new RegionCommandT);
|
|
currentCommand->op.reset(new OpT);
|
|
currentCommand->op->type = OpType_MatMul;
|
|
currentCommand->op->main.value = new MatMulT;
|
|
currentCommand->op->main.type = OpParameter_MatMul;
|
|
currentCommand->indexes = {BDiffIndex, CDiffIndex, AIndex};
|
|
currentCommand->view.resize(3);
|
|
for (int j=0; j<currentCommand->view.size(); ++j) {
|
|
currentCommand->view[j].reset(new ViewT);
|
|
}
|
|
currentCommand->iterIndexes.resize(3);
|
|
currentCommand->steps.resize(3);
|
|
for (int j=0; j<currentCommand->indexes.size(); ++j) {
|
|
// Compute output info
|
|
auto index = currentCommand->indexes[j];
|
|
currentCommand->iterIndexes[j] = command->iterIndexes[originIndexMap[index]];
|
|
currentCommand->steps[j] = command->steps[originIndexMap[index]];
|
|
*currentCommand->view[j] = *command->view[originIndexMap[index]];
|
|
}
|
|
// TODO: Optimize fuse option
|
|
currentCommand->fuse = BinaryOpOperation_ADD;
|
|
|
|
// Reorder the view's stride by size order change to e, h, l
|
|
std::vector<int> order = {2, 0, 1};
|
|
currentCommand->size = {h, e, l};
|
|
for (int j=0; j<currentCommand->indexes.size(); ++j) {
|
|
auto view = currentCommand->view[j].get();
|
|
auto originStride = view->stride;
|
|
for (int k=0; k<3; ++k) {
|
|
view->stride[k] = originStride[order[k]];
|
|
}
|
|
}
|
|
// Set Transpose Info by stride
|
|
auto dstMatMulParam = currentCommand->op->main.AsMatMul();
|
|
{
|
|
auto transAView = currentCommand->view[1].get();
|
|
auto transBView = currentCommand->view[1].get();
|
|
dstMatMulParam->transposeA = transAView->stride[1] != 1;
|
|
dstMatMulParam->transposeB = transBView->stride[1] == 1;
|
|
}
|
|
dstCommands.emplace_back(std::move(currentCommand));
|
|
}
|
|
}
|
|
static void _gradForCommand(const RegionCommandT* command, const std::map<int, int>& backwardMap, int tensorNumber, std::map<int, VARP>& extraInputs, std::map<int, VARP>& extraOutputs, std::vector<std::unique_ptr<RegionCommandT>>& dstCommands) {
|
|
if (command->op->type == OpType_MatMul) {
|
|
_gradForCommandMatMul(command, backwardMap, tensorNumber, extraInputs, extraOutputs, dstCommands);
|
|
return;
|
|
}
|
|
|
|
if (command->op->type == OpType_UnaryOp) {
|
|
auto inputIndex = command->indexes[1];
|
|
auto outputIndex = command->indexes[0];
|
|
if (backwardMap.find(inputIndex) == backwardMap.end()) {
|
|
return;
|
|
}
|
|
MNN_ASSERT(backwardMap.find(outputIndex) != backwardMap.end());
|
|
auto bpInput = backwardMap.find(outputIndex)->second;
|
|
auto bpOutput = backwardMap.find(inputIndex)->second;
|
|
|
|
if (nullptr == command->op->main.value) {
|
|
std::unique_ptr<RegionCommandT> currentCommand(new RegionCommandT);
|
|
currentCommand->op.reset(new OpT);
|
|
currentCommand->op->type = OpType_UnaryOp;
|
|
currentCommand->fuse = BinaryOpOperation_ADD;
|
|
currentCommand->op->main.type = OpParameter_NONE;
|
|
currentCommand->indexes = {bpOutput, bpInput};
|
|
currentCommand->view.resize(2);
|
|
currentCommand->view[0].reset(new ViewT);
|
|
currentCommand->view[1].reset(new ViewT);
|
|
*currentCommand->view[0] = *command->view[1];
|
|
*currentCommand->view[1] = *command->view[0];
|
|
currentCommand->size = command->size;
|
|
currentCommand->iterIndexes = {command->iterIndexes[1], command->iterIndexes[0]};
|
|
currentCommand->steps = {command->steps[1], command->steps[0]};
|
|
dstCommands.emplace_back(std::move(currentCommand));
|
|
return;
|
|
}
|
|
FUNC_PRINT(1);
|
|
}
|
|
int inputSize = 0;
|
|
std::vector<VARP> inputs;
|
|
if (command->op->type == OpType_BinaryOp) {
|
|
inputSize = 2;
|
|
}
|
|
else if (command->op->type == OpType_UnaryOp) {
|
|
inputSize = 1;
|
|
} else {
|
|
MNN_ASSERT(false);
|
|
// TODO: Support MatMul
|
|
}
|
|
for (int i=0; i<inputSize; ++i) {
|
|
auto tempValue = _Const(0.0f, {2, 2}, NHWC);
|
|
inputs.emplace_back(tempValue);
|
|
}
|
|
VARP tempOutput = _Const(0.0f, {2, 2}, NHWC);
|
|
std::map<std::pair<EXPRP, int>, int> allTensors;
|
|
auto commandExpr = Expr::create(command->op.get(), inputs, 1);
|
|
allTensors.insert(std::make_pair(std::make_pair(commandExpr, 0), command->indexes[0]));
|
|
for (int i=0; i<inputSize; ++i) {
|
|
allTensors.insert(std::make_pair(inputs[i]->expr(), command->indexes[i+1]));
|
|
}
|
|
std::map<int, int> originIndexMap;
|
|
for (int i=0; i<command->indexes.size(); ++i) {
|
|
auto index = command->indexes[i];
|
|
originIndexMap.insert(std::make_pair(index, i));
|
|
auto bkIter = backwardMap.find(index);
|
|
if (bkIter != backwardMap.end()) {
|
|
originIndexMap.insert(std::make_pair(bkIter->second, i));
|
|
}
|
|
}
|
|
|
|
auto gradOutputIter = backwardMap.find(command->indexes[0]);
|
|
MNN_ASSERT(gradOutputIter != backwardMap.end());
|
|
allTensors.insert(std::make_pair(tempOutput->expr(), gradOutputIter->second));
|
|
|
|
auto backwardVars = _getGradExpr(commandExpr, inputs, tempOutput);
|
|
for (int i=0; i<inputSize; ++i) {
|
|
if (nullptr != backwardVars[i]) {
|
|
auto gradInputIter = backwardMap.find(command->indexes[1 + i]);
|
|
if (gradInputIter != backwardMap.end()) {
|
|
if (backwardVars[i]->expr().first->get() == nullptr) {
|
|
// Make Copy Command
|
|
auto inputIndexIter = allTensors.find(backwardVars[i]->expr());
|
|
MNN_ASSERT(inputIndexIter != allTensors.end());
|
|
std::unique_ptr<RegionCommandT> currentCommand(new RegionCommandT);
|
|
currentCommand->op.reset(new OpT);
|
|
currentCommand->op->type = OpType_UnaryOp;
|
|
currentCommand->indexes.resize(2);
|
|
currentCommand->indexes[0] = gradInputIter->second;
|
|
currentCommand->indexes[1] = inputIndexIter->second;
|
|
currentCommand->view.resize(2);
|
|
for (int j=0; j<currentCommand->view.size(); ++j) {
|
|
currentCommand->view[j].reset(new ViewT);
|
|
}
|
|
currentCommand->iterIndexes.resize(2);
|
|
currentCommand->steps.resize(2);
|
|
// Compute output info
|
|
currentCommand->iterIndexes[0] = command->iterIndexes[originIndexMap[currentCommand->indexes[0]]];
|
|
currentCommand->steps[0] = command->steps[originIndexMap[currentCommand->indexes[0]]];
|
|
*currentCommand->view[0] = *command->view[originIndexMap[currentCommand->indexes[0]]];
|
|
currentCommand->size = command->size;
|
|
// TODO: Optimize fuse option
|
|
currentCommand->fuse = BinaryOpOperation_ADD;
|
|
*currentCommand->view[1] = *command->view[originIndexMap[currentCommand->indexes[1]]];
|
|
currentCommand->iterIndexes[1] = command->iterIndexes[originIndexMap[currentCommand->indexes[1]]];
|
|
currentCommand->steps[1] = command->steps[originIndexMap[currentCommand->indexes[1]]];
|
|
dstCommands.emplace_back(std::move(currentCommand));
|
|
} else {
|
|
allTensors.insert(std::make_pair(backwardVars[i]->expr(), gradInputIter->second));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
auto exprLists = Variable::getExecuteOrder(backwardVars);
|
|
for (int i=0; i<exprLists.size(); ++i) {
|
|
auto curExpr = exprLists[i];
|
|
// Find tensor or make cache
|
|
MNN_ASSERT(1 == curExpr->outputSize());
|
|
auto iter = allTensors.find(std::make_pair(curExpr, 0));
|
|
MNN_ASSERT(iter != allTensors.end());
|
|
if (nullptr == curExpr->get()) {
|
|
continue;
|
|
}
|
|
// Make Command
|
|
std::unique_ptr<RegionCommandT> currentCommand(new RegionCommandT);
|
|
currentCommand->op.reset(curExpr->get()->UnPack());
|
|
currentCommand->indexes.resize(curExpr->inputs().size() + curExpr->outputSize());
|
|
currentCommand->indexes[0] = iter->second;
|
|
currentCommand->view.resize(curExpr->inputs().size() + curExpr->outputSize());
|
|
for (int j=0; j<currentCommand->view.size(); ++j) {
|
|
currentCommand->view[j].reset(new ViewT);
|
|
}
|
|
currentCommand->iterIndexes.resize(curExpr->inputs().size() + curExpr->outputSize());
|
|
currentCommand->steps.resize(curExpr->inputs().size() + curExpr->outputSize());
|
|
// Compute output info
|
|
currentCommand->iterIndexes[0] = command->iterIndexes[originIndexMap[iter->second]];
|
|
currentCommand->steps[0] = command->steps[originIndexMap[iter->second]];
|
|
*currentCommand->view[0] = *command->view[originIndexMap[iter->second]];
|
|
currentCommand->size = command->size;
|
|
// TODO: Optimize fuse option
|
|
currentCommand->fuse = BinaryOpOperation_ADD;
|
|
for (int j=0; j<curExpr->inputs().size(); ++j) {
|
|
iter = allTensors.find(curExpr->inputs()[j]->expr());
|
|
if (iter == allTensors.end()) {
|
|
MNN_ASSERT(false);
|
|
}
|
|
currentCommand->indexes[1 + j] = iter->second;
|
|
MNN_ASSERT(originIndexMap.find(iter->second) != originIndexMap.end());
|
|
*currentCommand->view[1 + j] = *command->view[originIndexMap[iter->second]];
|
|
currentCommand->iterIndexes[1 + j] = command->iterIndexes[originIndexMap[iter->second]];
|
|
currentCommand->steps[1 + j] = command->steps[originIndexMap[iter->second]];
|
|
}
|
|
|
|
dstCommands.emplace_back(std::move(currentCommand));
|
|
}
|
|
MNN_ASSERT(dstCommands.size() > 0);
|
|
}
|
|
|
|
virtual std::vector<Express::VARP> onGrad(Express::EXPRP expr,
|
|
const std::vector<Express::VARP>& backwardOutput) override {
|
|
auto inputs = expr->inputs();
|
|
std::vector<VARP> result(inputs.size(), nullptr);
|
|
auto op = expr->get();
|
|
if (op->main_type() != OpParameter_LoopParam) {
|
|
return result;
|
|
}
|
|
std::unique_ptr<LoopParamT> srcParam(op->main_as_LoopParam()->UnPack());
|
|
MNN_ASSERT(srcParam->inputIndexes.size() == inputs.size());
|
|
MNN_ASSERT(srcParam->outputIndexes.size() == backwardOutput.size());
|
|
std::unique_ptr<LoopParamT> dstParam(new LoopParamT);
|
|
dstParam->loopNumber = srcParam->loopNumber;
|
|
dstParam->inputIndexes.resize(srcParam->inputIndexes.size() + srcParam->outputIndexes.size() + backwardOutput.size());
|
|
int dstParamTensorSrcInputOffset = 0;
|
|
int dstParamTensorSrcOutputOffset = srcParam->inputIndexes.size();
|
|
int dstParamTensorDstOutputOffset = srcParam->tensorNumber;
|
|
int dstParamTensorDstInputOffset = srcParam->tensorNumber + backwardOutput.size();
|
|
::memcpy(dstParam->inputIndexes.data() + 0, srcParam->inputIndexes.data(), srcParam->inputIndexes.size() * sizeof(int));
|
|
::memcpy(dstParam->inputIndexes.data() + srcParam->inputIndexes.size(), srcParam->outputIndexes.data(), srcParam->outputIndexes.size() * sizeof(int));
|
|
std::map<int, int> backwardMap;
|
|
for (int i=0; i<backwardOutput.size(); ++i) {
|
|
dstParam->inputIndexes[srcParam->inputIndexes.size() + srcParam->outputIndexes.size() + i] = dstParamTensorDstOutputOffset + i;
|
|
backwardMap.insert(std::make_pair(srcParam->outputIndexes[i], dstParamTensorDstOutputOffset + i));
|
|
}
|
|
dstParam->tensorNumber = srcParam->tensorNumber + backwardOutput.size();
|
|
dstParam->outputIndexes.clear();
|
|
std::vector<int> gradInputs;
|
|
for (int i=0; i<inputs.size(); ++i) {
|
|
// Only need compute grad for float tensor
|
|
if (inputs[i]->getInfo()->type.code != halide_type_float) {
|
|
continue;
|
|
}
|
|
gradInputs.emplace_back(i);
|
|
auto curNumber = dstParam->tensorNumber;
|
|
dstParam->outputIndexes.emplace_back(curNumber);
|
|
dstParam->tensorNumber++;
|
|
backwardMap.insert(std::make_pair(srcParam->inputIndexes[i], curNumber));
|
|
}
|
|
// Clear zero firstly for backward inputs
|
|
for (auto index : srcParam->inputIndexes) {
|
|
auto iter = backwardMap.find(index);
|
|
if (iter == backwardMap.end()) {
|
|
continue;
|
|
}
|
|
std::unique_ptr<RegionCommandT> zeroCmd(new RegionCommandT);
|
|
zeroCmd->indexes = {iter->second};
|
|
dstParam->initCommand.emplace_back(std::move(zeroCmd));
|
|
}
|
|
std::map<int, VARP> extraVarps;
|
|
std::map<int, VARP> extraOutputVarps;
|
|
for (int i=0; i<srcParam->commands.size(); ++i) {
|
|
auto reverseI = (int)srcParam->commands.size() - 1 - i;
|
|
auto cmd = srcParam->commands[reverseI].get();
|
|
_gradForCommand(cmd, backwardMap, dstParam->tensorNumber, extraVarps, extraOutputVarps, dstParam->commands);
|
|
}
|
|
// Make Op
|
|
std::vector<VARP> loopInputs(inputs.size() + 2 * expr->outputSize() + extraVarps.size());
|
|
for (int i=0; i<inputs.size(); ++i) {
|
|
loopInputs[i] = inputs[i];
|
|
}
|
|
for (int i=0; i<expr->outputSize(); ++i) {
|
|
loopInputs[i + inputs.size()] = Variable::create(expr, i);
|
|
}
|
|
for (int i=0; i<expr->outputSize(); ++i) {
|
|
loopInputs[i + inputs.size() + expr->outputSize()] = backwardOutput[i];
|
|
}
|
|
for (auto& iter : extraVarps) {
|
|
loopInputs[inputs.size() + 2 * expr->outputSize() + iter.first] = iter.second;
|
|
}
|
|
MNN_ASSERT(dstParam->commands.size() > 0);
|
|
MNN_ASSERT(dstParam->outputIndexes.size() == gradInputs.size());
|
|
for (int i=0; i<gradInputs.size(); ++i) {
|
|
auto info = inputs[gradInputs[i]]->getInfo();
|
|
std::unique_ptr<MNN::TensorDescribeT> describe(new TensorDescribeT);
|
|
describe->index = dstParam->outputIndexes[i];
|
|
describe->blob.reset(new BlobT);
|
|
describe->blob->dims = info->dim;
|
|
describe->blob->dataType = DataType_DT_FLOAT;
|
|
switch (info->order) {
|
|
case MNN::Express::NCHW:
|
|
describe->blob->dataFormat = MNN_DATA_FORMAT_NCHW;
|
|
break;
|
|
case MNN::Express::NHWC:
|
|
describe->blob->dataFormat = MNN_DATA_FORMAT_NHWC;
|
|
break;
|
|
case MNN::Express::NC4HW4:
|
|
describe->blob->dataFormat = MNN_DATA_FORMAT_NC4HW4;
|
|
break;
|
|
default:
|
|
break;
|
|
}
|
|
dstParam->extraTensorInfos.emplace_back(std::move(describe));
|
|
}
|
|
|
|
std::unique_ptr<OpT> loopOp(new OpT);
|
|
loopOp->type = OpType_While;
|
|
loopOp->main.type = OpParameter_LoopParam;
|
|
loopOp->main.value = dstParam.release();
|
|
|
|
auto gradExpr = Expr::create(std::move(loopOp), loopInputs, gradInputs.size());
|
|
for (int i=0; i<gradInputs.size(); ++i) {
|
|
result[gradInputs[i]] = Variable::create(gradExpr, i);
|
|
}
|
|
return result;
|
|
}
|
|
};
|
|
|
|
static void _create() {
|
|
static LoopGrad _c;
|
|
OpGrad::insert(OpType_While, &_c);
|
|
|
|
}
|
|
|
|
REGISTER_GRAD(LoopGrad_cpp, _create);
|
|
};
|
|
|