Function MGB_IMPL_OPR_GRAD

src/jit/impl/executor_opr.cpp:483–626

anonymous namespace

};
}  // anonymous namespace
MGB_IMPL_OPR_GRAD(JITExecutor) {
    VarNodeArray grad_inputs;
    for (auto input : opr.input())
        grad_inputs.push_back(input);
    mgb_assert(out_grad[0]);
    grad_inputs.push_back(opr.output(0));
    grad_inputs.push_back(out_grad[0]);
    auto fwd_igraph_ptr = opr.internal_graph_ptr();
    auto output_ph = JITPlaceholder::make(
            fwd_igraph_ptr->output(), fwd_igraph_ptr->placeholders().size());
    auto og_ph = JITPlaceholder::make(
            out_grad[0], fwd_igraph_ptr->placeholders().size() + 1);
    auto loss = opr::VirtualLoss::make({fwd_igraph_ptr->output()}, {og_ph});
    auto gx = cg::grad(
            loss, fwd_igraph_ptr->placeholders()[wrt_idx]->output(0), false, false);
    if (!gx.node()) {
        return nullptr;
    }
    if (gx.node()->owner_opr()->same_type<opr::InvalidGrad>()) {
        return opr::InvalidGrad::make(opr, wrt_idx);
    }
    // early return if grad expression is single node
    for (size_t i = 0; i < fwd_igraph_ptr->placeholders().size(); ++i) {
        if (gx.node() == fwd_igraph_ptr->placeholders()[i]->output(0)) {
            return grad_inputs[i];
        }
    }
    if (gx.node() == og_ph.node()) {
        return out_grad[0];
    }
    if (gx.node() == fwd_igraph_ptr->output()) {
        return opr.output(0);
    }
    if (auto imm = gopt::try_cast_as_op<opr::ImmutableTensor>(gx.node()->owner_opr())) {
        HostTensorND hval{grad_inputs[0]->comp_node()};
        hval.copy_from(imm->value()).sync();
        return opr::ImmutableTensor::make(*imm->owner_graph(), hval).node();
    }

    // replace output var in internal graph with output placeholder, so
    // we could forward opr.output(computeed by forward JITExecutor) into
    // placeholder to avoid redundant computation
    InternalGraphRewriter rewriter{gx.node()};
    rewriter.iter([&rewriter, &fwd_igraph_ptr, &output_ph](cg::OperatorNodeBase* opr) {
        if (opr == fwd_igraph_ptr->output()->owner_opr()) {
            rewriter.replace_var(opr->output(0), output_ph.node());
            return;
        }
        rewriter.auto_replace_outputs(opr);
    });

    auto expand_into_origin_graph = [&rewriter](
                                            cg::OperatorNodeBase* opr,
                                            const VarNodeArray& grad_inputs) {
        if (auto ph = gopt::try_cast_as_op<JITPlaceholder>(opr)) {
            rewriter.replace_var(opr->output(0), grad_inputs.at(ph->input_id()));
            return;
        }