| 134 | } |
| 135 | } |
| 136 | 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) { |
| 137 | if (command->op->type == OpType_MatMul) { |
| 138 | _gradForCommandMatMul(command, backwardMap, tensorNumber, extraInputs, extraOutputs, dstCommands); |
| 139 | return; |
| 140 | } |
| 141 | |
| 142 | if (command->op->type == OpType_UnaryOp) { |
| 143 | auto inputIndex = command->indexes[1]; |
| 144 | auto outputIndex = command->indexes[0]; |
| 145 | if (backwardMap.find(inputIndex) == backwardMap.end()) { |
| 146 | return; |
| 147 | } |
| 148 | MNN_ASSERT(backwardMap.find(outputIndex) != backwardMap.end()); |
| 149 | auto bpInput = backwardMap.find(outputIndex)->second; |
| 150 | auto bpOutput = backwardMap.find(inputIndex)->second; |
| 151 | |
| 152 | if (nullptr == command->op->main.value) { |
| 153 | std::unique_ptr<RegionCommandT> currentCommand(new RegionCommandT); |
| 154 | currentCommand->op.reset(new OpT); |
| 155 | currentCommand->op->type = OpType_UnaryOp; |
| 156 | currentCommand->fuse = BinaryOpOperation_ADD; |
| 157 | currentCommand->op->main.type = OpParameter_NONE; |
| 158 | currentCommand->indexes = {bpOutput, bpInput}; |
| 159 | currentCommand->view.resize(2); |
| 160 | currentCommand->view[0].reset(new ViewT); |
| 161 | currentCommand->view[1].reset(new ViewT); |
| 162 | *currentCommand->view[0] = *command->view[1]; |
| 163 | *currentCommand->view[1] = *command->view[0]; |
| 164 | currentCommand->size = command->size; |
| 165 | currentCommand->iterIndexes = {command->iterIndexes[1], command->iterIndexes[0]}; |
| 166 | currentCommand->steps = {command->steps[1], command->steps[0]}; |
| 167 | dstCommands.emplace_back(std::move(currentCommand)); |
| 168 | return; |
| 169 | } |
| 170 | FUNC_PRINT(1); |
| 171 | } |
| 172 | int inputSize = 0; |
| 173 | std::vector<VARP> inputs; |
| 174 | if (command->op->type == OpType_BinaryOp) { |
| 175 | inputSize = 2; |
| 176 | } |
| 177 | else if (command->op->type == OpType_UnaryOp) { |
| 178 | inputSize = 1; |
| 179 | } else { |
| 180 | MNN_ASSERT(false); |
| 181 | // TODO: Support MatMul |
| 182 | } |
| 183 | for (int i=0; i<inputSize; ++i) { |
| 184 | auto tempValue = _Const(0.0f, {2, 2}, NHWC); |
| 185 | inputs.emplace_back(tempValue); |
| 186 | } |
| 187 | VARP tempOutput = _Const(0.0f, {2, 2}, NHWC); |
| 188 | std::map<std::pair<EXPRP, int>, int> allTensors; |
| 189 | auto commandExpr = Expr::create(command->op.get(), inputs, 1); |
| 190 | allTensors.insert(std::make_pair(std::make_pair(commandExpr, 0), command->indexes[0])); |
| 191 | for (int i=0; i<inputSize; ++i) { |
| 192 | allTensors.insert(std::make_pair(inputs[i]->expr(), command->indexes[i+1])); |
| 193 | } |