eval.cpp

Go to the documentation of this file.

#include "mim/plug/autodiff/phase/eval.h"
 
#include <mim/lam.h>
 
#include <mim/plug/cps/cps.h>
 
#include "mim/plug/autodiff/autodiff.h"
 
using namespace std::literals;
 
namespace mim::plug::autodiff::phase {
 
/// The high level view is:
/// ```
/// f: B -> C
/// g: A -> B
/// f o g := λ x. f(g(x)) : A -> C
/// ```
/// In CPS the types look like:
/// ```
/// f:  Cn[B, Cn C]
/// g:  Cn[A, Cn B]
/// h = f o g
/// h:  Cn[A, cn C]
/// h = λ (a ret_h) = g (a, h')
/// h': Cn B
/// h'= λ b = f (b, ret_h)
/// ```
static const Def* compose_cn(const Def* f, const Def* g) {
    auto& world = f->world();
    auto F      = f->type()->as<Pi>();
    auto G      = g->type()->as<Pi>();
 
    assert(Pi::isa_returning(F));
    assert(Pi::isa_returning(G));
 
    auto A = G->dom(2, 0);
    auto B = G->ret_dom();
    auto C = F->ret_dom();
    // The type check of codom G = dom F is better handled by the application type checking
 
    world.log().d("compose f: {}: {}, g: {}: {}", f, F, g, G);
    world.log().d("    A = {}, B = {}, C = {}", A, B, C);
 
    auto name  = "comp_"s + f->sym().str() + "_" + g->sym().str();
    auto h     = world.mut_fun(A, C)->set(name);
    auto hcont = world.mut_con(B)->set(name + "_cont");
 
    h->app(true, g, {h->var(2, 0), hcont});
 
    auto hcont_var = hcont->var(); // Warning: not var(0) => only one var => normalization flattens tuples down here.
    hcont->app(true, f, {hcont_var, h->var(2, 1) /* ret_var */});
 
    return h;
}
 
// TODO: maybe use template (https://codereview.stackexchange.com/questions/141961/memoization-via-template) to memoize
const Def* Eval::augment(const Def* def, Lam* f, Lam* f_diff) {
    if (auto res = fe::lookup(augmented, def)) return res;
    return augmented[def] = augment_(def, f, f_diff);
}
 
const Def* Eval::derive(const Def* def) {
    if (auto res = fe::lookup(derived, def)) return res;
    return derived[def] = derive_(def);
}
 
const Def* Eval::rewrite_imm_App(const App* app) {
    if (auto ad_app = Axm::isa<ad>(app); ad_app) {
        // callee = autodiff T
        // arg = function of type T
        //   (or operator)
        // Rewrite the argument into the new world first; the whole derivation then works on that copy.
        auto arg = rewrite(ad_app->arg());
 
        if (arg->isa<Lam>()) return derive(arg);
 
        // TODO: handle operators analogous
 
        assert(0 && "not implemented");
        return arg;
    }
 
    return RWPhase::rewrite_imm_App(app);
}
 
/*
 * toplevel derivation
 */
 
/// Additionally to the derivation, the pullback is registered and the maps are initialized.
const Def* Eval::derive_(const Def* def) {
    auto lam      = def->as_mut<Lam>(); // TODO check if mutable
    auto deriv_ty = autodiff_type_fun_pi(lam->type());
    auto deriv    = world().mut_lam(deriv_ty)->set(lam->sym().str() + "_deriv");
 
    // We first pre-register the derivatives.
    // This knowledge is needed for recursion.
    // (Alternatively, we could also use projections out the variables instead of pre-partial-pullback
    // initialization.)
    derived[lam] = deriv;
 
    auto [arg_ty, ret_pi] = lam->type()->doms<2>();
    auto deriv_all_args   = deriv->var();
    const Def* deriv_arg  = deriv->var(2, 0)->set("arg");
 
    // We generate the shadow pullbacks dynamically to save work and avoid code duplication.
    // Only the toplevel pullback for arguments and return continuation is special cased.
 
    // TODO: check identity: could use identity tangent(arg_ty) = tangent(augment(arg_ty)) with deriv_arg->type() =
    // augment(arg_ty) We give the argument the identity pullback.
    auto arg_id_pb              = id_pullback(arg_ty);
    partial_pullback[deriv_arg] = arg_id_pb;
    // The return continuation has to formally exist but should never be directly accessed.
    auto ret_var              = deriv->var(2, 1);
    auto ret_pb               = zero_pullback(lam->var(2, 1)->type(), arg_ty);
    partial_pullback[ret_var] = ret_pb;
 
    shadow_pullback[deriv_all_args] = world().tuple({arg_id_pb, ret_pb});
 
    // We pre-register the augment replacements.
    // The function and its variables are replaced by their new derived versions.
    // TODO: maybe leave out function call (duplication with derived)
    augmented[def]        = deriv;
    augmented[lam->var()] = deriv->var();
 
    // already contains the correct application of
    // deriv->ret_var() by specification
    // f : cn[R] has a partial derivative (exception to closed rule)
    // f': cn[R, cn[R, cn[A]]]
    //   this is needed for continuations (without closure conversion)
    //   but also essentially for the return continuation
 
    // Here a reminder of types:
    // The expression `e: B` has the implicit function `e_fun: A -> B`
    // The partial pullback is then `e*: B* -> A*`
    // The derivatived version is `e': B' × (B* -> A*)` which is an application of `e'_fun: A' -> B' × (B* -> A*)`
    auto new_body = augment(lam->body(), lam, deriv);
    deriv->set(true, new_body);
 
    return deriv;
}
 
/*
 * augment
 */
 
const Def* Eval::augment_lit(const Lit* lit, Lam* f, Lam*) {
    auto pb               = zero_pullback(lit->type(), f->dom(2, 0));
    partial_pullback[lit] = pb;
    return lit;
}
 
const Def* Eval::pullback(const Def* aug, Lam* f) {
    // A singleton has no tangent, and its sole inhabitant is shared by all functions.
    if (world().inhabitant(aug->type()) == aug) return zero_pullback(aug->type(), f->dom(2, 0));
    return partial_pullback[aug];
}
 
const Def* Eval::augment_var(const Var* var, Lam*, Lam*) {
    assert(augmented.count(var));
    assert(partial_pullback.count(augmented[var]));
    return var;
}
 
const Def* Eval::augment_lam(Lam* lam, Lam* f, Lam* f_diff) {
    // TODO: we need partial pullbacks for tuples (higher-order / ret-cont application)
    // also for higher-order args, ret_cont (at another point)
    // the pullback is not important but formally required by tuple rule
    if (augmented.count(lam)) {
        // We already know the function:
        // * recursion
        // * higher order arguments
        // * new encounter of previous function
        return augmented[lam];
    }
    // TODO: better fix (another pass as analysis?)
    // TODO: handle open functions
    if (Lam::isa_basicblock(lam) || lam->sym().view().contains("ret") || lam->sym().view().contains("_cont")) {
        // A open continuation behaves the same as return:
        // ```
        // cont: Cn[X]
        // cont': Cn[X,Cn[X,A]]
        // ```
        // There is dependency on the closed function context.
        // (All derivatives are with respect to the arguments of a closed function.)
 
        auto cont_dom             = lam->type()->dom(); // not only 0 but all
        auto pb_ty                = pullback_type(cont_dom, f->dom(2, 0));
        auto aug_dom              = autodiff_type_fun(cont_dom);
        auto aug_lam              = world().mut_con({aug_dom, pb_ty})->set("aug_"s + lam->sym().str());
        auto aug_var              = aug_lam->var(2, 0);
        augmented[lam->var()]     = aug_var;
        augmented[lam]            = aug_lam; // TODO: only one of these two
        derived[lam]              = aug_lam;
        auto pb                   = aug_lam->var(2, 1);
        partial_pullback[aug_var] = pb;
        // We are still in same closed function.
        auto new_body = augment(lam->body(), f, f_diff);
        // TODO we also need to rewrite the filter
        aug_lam->set(lam->filter(), new_body);
 
        auto lam_pb               = zero_pullback(lam->type(), f->dom(2, 0));
        partial_pullback[aug_lam] = lam_pb;
 
        return aug_lam;
    }
    // Some general function in the program needs to be differentiated.
    // The old pass emitted a new `autodiff.ad` application here and relied on the PassMan to revisit it;
    // as a Phase we derive eagerly instead (derive() pre-registers itself, so recursion terminates).
    auto aug_lam = derive(lam);
    // TODO: directly more association here? => partly inline op_autodiff
    return aug_lam;
}
 
const Def* Eval::augment_extract(const Extract* ext, Lam* f, Lam* f_diff) {
    auto tuple = ext->tuple();
    auto index = ext->index();
 
    auto aug_tuple = augment(tuple, f, f_diff);
    auto aug_index = augment(index, f, f_diff);
 
    const Def* pb;
    if (shadow_pullback.count(aug_tuple)) {
        auto shadow_tuple_pb = shadow_pullback[aug_tuple];
        pb                   = world().extract(shadow_tuple_pb, aug_index);
    } else {
        // ```
        // e:T, b:B
        // b = e#i
        // b* = \lambda (s:B). e* (insert s at i in (zero T))
        // ```
        auto tuple_pb = pullback(aug_tuple, f);
        assert(tuple_pb);
        auto pb_ty      = pullback_type(ext->type(), f->dom(2, 0));
        auto pb_fun     = world().mut_lam(pb_ty)->set("extract_pb");
        auto pb_tangent = pb_fun->var(2, 0)->set("s");
        auto tuple_tan  = world().insert(world().call<zero>(aug_tuple->type()), aug_index, pb_tangent)->set("tup_s");
        pb_fun->app(true, tuple_pb, {tuple_tan, pb_fun->var(2, 1) /* ret_var but make sure to select correct one */});
        pb = pb_fun;
    }
 
    auto aug_ext              = world().extract(aug_tuple, aug_index);
    partial_pullback[aug_ext] = pb;
 
    return aug_ext;
}
 
const Def* Eval::augment_tuple(const Tuple* tup, Lam* f, Lam* f_diff) {
    // TODO: should use ops instead?
    auto aug_ops = tup->projs([&](const Def* op) -> const Def* { return augment(op, f, f_diff); });
    auto aug_tup = world().tuple(aug_ops);
 
    auto pbs = DefVec(Defs(aug_ops), [&](const Def* op) { return pullback(op, f); });
    // shadow pb = tuple of pbs
    auto shadow_pb           = world().tuple(pbs);
    shadow_pullback[aug_tup] = shadow_pb;
 
    // ```
    // \lambda (s:[E0,...,Em]).
    //    sum (m,A)
    //      ((cps2ds e0*) (s#0), ..., (cps2ds em*) (s#m))
    // ```
    auto pb_ty = pullback_type(tup->type(), f->dom(2, 0));
    auto pb    = world().mut_lam(pb_ty)->set("tup_pb");
 
    auto pb_tangent = pb->var(2, 0)->set("tup_s");
 
    auto tangents = DefVec(
        pbs.size(), [&](nat_t i) { return world().app(cps::op_cps2ds_dep(pbs[i]), world().extract(pb_tangent, i)); });
    pb->app(true, pb->var(2, 1),
            // summed up tangents
            op_sum(tangent_type_fun(f->dom(2, 0)), tangents));
    partial_pullback[aug_tup] = pb;
 
    return aug_tup;
}
 
const Def* Eval::augment_pack(const Pack* pack, Lam* f, Lam* f_diff) {
    auto arity = pack->arity(); // TODO: arity vs shape
    auto body  = pack->body();
 
    auto aug_arity = augment_(arity, f, f_diff);
    auto aug_body  = augment(body, f, f_diff);
 
    auto aug_pack = world().pack(aug_arity, aug_body);
 
    // TODO: or use scale axm
    auto body_pb = pullback(aug_body, f);
    assert(body_pb && "pack pullback should exists");
    auto pb_pack              = world().pack(aug_arity, body_pb);
    shadow_pullback[aug_pack] = pb_pack;
 
    auto pb_type = pullback_type(pack->type(), f->dom(2, 0));
    auto pb      = world().mut_lam(pb_type)->set("pack_pb");
 
    auto f_arg_ty_diff = tangent_type_fun(f->dom(2, 0));
    auto app_pb        = world().mut_pack(world().arr(aug_arity, f_arg_ty_diff))->set_shape(aug_arity);
 
    // TODO: special case for const width (special tuple)
 
    // <i:n, cps2ds body_pb (s#i)>
    app_pb->set_body(world().app(cps::op_cps2ds_dep(body_pb), world().extract(pb->var(2, 0), app_pb->var())));
 
    auto sumup = world().app(world().annex<sum>(), {aug_arity, f_arg_ty_diff});
 
    pb->app(true, pb->var(2, 1), world().app(sumup, app_pb));
 
    partial_pullback[aug_pack] = pb;
 
    return aug_pack;
}
 
const Def* Eval::augment_app(const App* app, Lam* f, Lam* f_diff) {
    auto callee = app->callee();
    auto arg    = app->arg();
 
    auto aug_arg    = augment(arg, f, f_diff);
    auto aug_callee = augment(callee, f, f_diff);
 
    // TODO: move down to if(!is_cont(callee))
    if (!Pi::isa_cn(callee->type()) && Pi::isa_cn(aug_callee->type())) aug_callee = cps::op_cps2ds_dep(aug_callee);
 
    // nested (inner application)
    if (app->type()->isa<Pi>()) {
        auto aug_app = world().app(aug_callee, aug_arg);
        // We do not add a pullback as the pullback is bundled in the cps call or returned by the ds call
        return aug_app;
    }
 
    // continuation (ret, if, ...)
    if (Pi::isa_basicblock(callee->type())) {
        // TODO: check if function (not operator)
        // The original function is an open function (return cont / continuation) of type `Cn[E]`
        // The augmented function `aug_callee` looks like a function but is not really a function has the type `Cn[E,
        // Cn[E, Cn[A]]]`
 
        // ret(e) => ret'(e, e*)
 
        auto arg_pb  = pullback(aug_arg, f);
        auto aug_app = world().app(aug_callee, {aug_arg, arg_pb});
        return aug_app;
    }
 
    // ds function
    if (!Pi::isa_cn(callee->type())) {
        auto aug_app = world().app(aug_callee, aug_arg);
 
        // The calle is ds function (e.g. operator (or its partial application))
        auto [aug_res, fun_pb] = aug_app->projs<2>();
        // We compose `fun_pb` with `argument_pb` to get the result pb
        // TODO: combine case with cps function case
        auto arg_pb = pullback(aug_arg, f);
        assert(arg_pb);
        // `fun_pb: out_tan -> arg_tan`
        // `arg_pb: arg_tan -> fun_tan`
        auto res_pb               = compose_cn(arg_pb, fun_pb);
        partial_pullback[aug_res] = res_pb;
        return aug_res;
    }
 
    // TODO: dest with a function such that f args != g args
    // normal function app
    // ```
    // g: cn[E, cn X]
    // g(args,cont)
    // g': cn[E, cn[X, cn[X, cn E]]]
    // g'(aug_args, ____)
    // ```
    // At this point g_deriv might still be "autodiff ... g".
    auto g_deriv = aug_callee;
 
    auto [real_aug_args, aug_cont] = aug_arg->projs<2>();
    auto e_pb                      = pullback(real_aug_args, f);
 
    // TODO: better debug names
    auto ret_g_deriv_ty = g_deriv->type()->as<Pi>()->dom(2, 1);
    auto c1_ty          = ret_g_deriv_ty->as<Pi>();
    auto c1             = world().mut_lam(c1_ty)->set("c1");
    auto res            = c1->var(2, 0);
    auto r_pb           = c1->var(2, 1);
    c1->app(true, aug_cont, {res, compose_cn(e_pb, r_pb)});
 
    auto aug_app = world().app(aug_callee, {real_aug_args, c1});
 
    // The result is * => no pb needed, no composition needed.
    return aug_app;
}
 
/// Rewrites the given definition in a lambda environment.
const Def* Eval::augment_(const Def* def, Lam* f, Lam* f_diff) {
    // Applications are continuations, operators, or full functions.
    if (auto app = def->isa<App>()) {
        return augment_app(app, f, f_diff);
    } else if (auto ext = def->isa<Extract>()) {
        return augment_extract(ext, f, f_diff);
    } else if (auto var = def->isa<Var>()) {
        return augment_var(var, f, f_diff);
    } else if (auto lam = def->isa_mut<Lam>()) {
        return augment_lam(lam, f, f_diff);
    } else if (auto lam = def->isa<Lam>()) {
        log().e("cannot augment an immutable Lam: {}", lam);
        assert(false && "can not handle non-mutable lambdas");
    } else if (auto lit = def->isa<Lit>()) {
        return augment_lit(lit, f, f_diff);
    } else if (auto tup = def->isa<Tuple>()) {
        return augment_tuple(tup, f, f_diff);
    } else if (auto pack = def->isa<Pack>()) {
        // TODO: handle mut packs (dependencies in the pack) (=> see paper about vectors)
        return augment_pack(pack, f, f_diff);
    } else if (auto ax = def->isa<Axm>()) {
        auto diff_name = ax->sym().str();
        fe::find_and_replace(diff_name, ".", "_");
        diff_name = "autodiff.diff." + diff_name;
 
        // Look the derivative up in the old world; rewrite() below maps it into the new one.
        auto old_diff_fun = old_world().annex(old_world().sym(diff_name));
        if (!old_diff_fun) {
            auto expected_type = autodiff_type_fun(ax->type());
            log().e("no derivative `{}`; expected type: {}", diff_name, expected_type);
            assert(false && "unhandled axm");
        }
        // TODO: why does this cause a depth error?
        return rewrite(old_diff_fun);
    }
 
    // TODO: handle Pi for axm app
    // TODO: remaining (lambda, axm)
 
    log().e("cannot augment {} node {}: {}", def->node_name(), def, def->type());
    assert(false && "augment not implemented on this def");
    std::unreachable();
}
 
} // namespace mim::plug::autodiff::phase