| 475 | /*********************** Logical Tensor Impl ***********************/ |
| 476 | |
| 477 | EncodedSubgraph ProxyGraph::make_backward_graph( |
| 478 | const OpDef& opdef, const SmallVector<LogicalTensorDesc>& input_descs, |
| 479 | const SmallVector<bool>& input_requires_grad, |
| 480 | const SmallVector<bool>& output_has_grad) { |
| 481 | using op_t = OperatorNodeBase*; |
| 482 | using var_t = VarNode*; |
| 483 | using vars_t = VarNodeArray; |
| 484 | auto inputs = make_input_place_holders(input_descs); |
| 485 | auto outputs = OpDef::apply_on_var_node(opdef, inputs); |
| 486 | SmallVector<LogicalTensorDesc> output_descs; |
| 487 | for (auto&& i : outputs) { |
| 488 | output_descs.push_back({TensorLayout{i->dtype()}, i->comp_node()}); |
| 489 | } |
| 490 | GradContext<op_t, var_t> grad_context{[&](VarNode* lhs, VarNode* rhs) -> VarNode* { |
| 491 | auto add = opr::Elemwise::Mode::ADD; |
| 492 | return opr::Elemwise::make(VarNodeArray{lhs, rhs}, add).node(); |
| 493 | }}; |
| 494 | cg::DepOprIter iter{[&](OperatorNodeBase* op) { |
| 495 | grad_context.record_expr(op, op->input(), op->output()); |
| 496 | }}; |
| 497 | for (size_t i = 0; i < inputs.size(); ++i) { |
| 498 | auto& input = inputs[i]; |
| 499 | iter.set_visited(input->owner_opr()); |
| 500 | if (input_requires_grad[i]) { |
| 501 | grad_context.mark_require_grad(input); |
| 502 | } |
| 503 | } |
| 504 | for (auto&& output : outputs) { |
| 505 | iter.add(output); |
| 506 | } |
| 507 | auto output_grads = make_input_place_holders(output_descs); |
| 508 | for (size_t i = 0; i < outputs.size(); ++i) { |
| 509 | if (!output_has_grad[i]) { |
| 510 | output_grads[i] = nullptr; |
| 511 | } |
| 512 | } |
| 513 | auto compute_input_grads = [&](op_t op, vars_t inputs, vars_t outputs, |
| 514 | vars_t output_grads) { |
| 515 | auto* gfunc = cg::lookup_grad_func(op->dyn_typeinfo()); |
| 516 | vars_t input_grads(inputs.size(), nullptr); |
| 517 | bool any_grad = false; |
| 518 | for (auto&& output_grad : output_grads) { |
| 519 | if (output_grad) { |
| 520 | any_grad = true; |
| 521 | } |
| 522 | } |
| 523 | if (!gfunc || !any_grad) { |
| 524 | return input_grads; |
| 525 | } |
| 526 | Maybe<VarNodeArray> grad_results; |
| 527 | auto&& input_requires_grad = grad_context.get_require_grads(inputs); |
| 528 | for (size_t i = 0; i < inputs.size(); ++i) { |
| 529 | VarNode* grad; |
| 530 | if (grad_results.valid()) { |
| 531 | grad = grad_results.val()[i]; |
| 532 | } else { |
| 533 | mgb_assert(gfunc, "could not find grad function"); |
| 534 | auto res = (*gfunc)(op, i, output_grads); |
nothing calls this directly
no test coverage detected