| 47 | } |
| 48 | |
| 49 | BackwardGraphWithClosure::BackwardGraphWithClosure( |
| 50 | std::shared_ptr<OptimizedBackwardGraphResult> backward_graph, |
| 51 | std::shared_ptr<OpDef> op, Span<ValueRef> inputs, Span<ValueRef> outputs) |
| 52 | : backward_graph(backward_graph), |
| 53 | output_mask_offset(inputs.size()), |
| 54 | grad_mask_offset(inputs.size() + outputs.size()), |
| 55 | op(op) { |
| 56 | auto& save_for_backward = backward_graph->save_for_backward; |
| 57 | mgb_assert(save_for_backward.size() == inputs.size() + 2 * outputs.size()); |
| 58 | size_t count = std::count_if( |
| 59 | save_for_backward.begin(), save_for_backward.end(), ranges::identity{}); |
| 60 | if (!backward_graph->precomp.empty()) { |
| 61 | SmallVector<ValueRef> inputs_and_outputs(inputs.size() + outputs.size()); |
| 62 | auto it = inputs_and_outputs.begin(); |
| 63 | for (auto&& input : inputs) { |
| 64 | *it++ = input; |
| 65 | } |
| 66 | for (auto&& output : outputs) { |
| 67 | *it++ = output; |
| 68 | } |
| 69 | auto precomp = imperative::apply(backward_graph->precomp, inputs_and_outputs); |
| 70 | closure.reserve(precomp.size() + count); |
| 71 | std::copy(precomp.begin(), precomp.end(), std::back_inserter(closure)); |
| 72 | } else { |
| 73 | closure.reserve(count); |
| 74 | } |
| 75 | for (size_t i = 0; i < inputs.size(); ++i) { |
| 76 | if (save_for_backward[i]) { |
| 77 | closure.push_back(inputs[i]); |
| 78 | } |
| 79 | } |
| 80 | for (size_t i = 0; i < outputs.size(); ++i) { |
| 81 | if (save_for_backward[inputs.size() + i]) { |
| 82 | closure.push_back(outputs[i]); |
| 83 | } |
| 84 | } |
| 85 | if (outputs.size() > 1) { |
| 86 | output_descs.reserve(outputs.size()); |
| 87 | for (auto&& output : outputs) { |
| 88 | auto symbolic_shape = imperative::apply(*GetVarShape::make(), output)[0]; |
| 89 | output_descs.push_back({symbolic_shape, output.dtype(), output.device()}); |
| 90 | } |
| 91 | } |
| 92 | } |
| 93 | void BackwardGraphWithClosure::operator()( |
| 94 | Span<ValueRef> grads, std::function<void(size_t, ValueRef)> receiver) { |
| 95 | ValueRef args[closure.size() + grads.size()]; |