save

.vscode/launch.json

@@ -11,6 +11,8 @@
             "request": "launch",
             "program": "${workspaceFolder}/build/bin/thorin",
             "args": [
+                "-d",
+                "affine",
                 "-d",
                 "mem",
                 "-d",
@@ -21,8 +23,11 @@
                 "tool",
                 "-d",
                 "autodiff",
-                "${workspaceFolder}/lit/autodiff/simple_autodiff.thorin",
+                "${workspaceFolder}/lit/autodiff/arr.thorin",
+                //"${workspaceFolder}/lit/autodiff/simple_autodiff.thorin",
                 //"${file}",
+                "--output-ll",
+                "-",
                 "--output-thorin",
                 "-",
                 "-VVVV",

.vscode/settings.json

@@ -0,0 +1,62 @@
+{
+    "files.associations": {
+        "array": "cpp",
+        "atomic": "cpp",
+        "bit": "cpp",
+        "*.tcc": "cpp",
+        "bitset": "cpp",
+        "cctype": "cpp",
+        "cfenv": "cpp",
+        "chrono": "cpp",
+        "clocale": "cpp",
+        "cmath": "cpp",
+        "codecvt": "cpp",
+        "compare": "cpp",
+        "concepts": "cpp",
+        "cstddef": "cpp",
+        "cstdint": "cpp",
+        "cstdio": "cpp",
+        "cstdlib": "cpp",
+        "cstring": "cpp",
+        "ctime": "cpp",
+        "cwchar": "cpp",
+        "cwctype": "cpp",
+        "deque": "cpp",
+        "forward_list": "cpp",
+        "list": "cpp",
+        "map": "cpp",
+        "set": "cpp",
+        "unordered_map": "cpp",
+        "unordered_set": "cpp",
+        "vector": "cpp",
+        "exception": "cpp",
+        "functional": "cpp",
+        "initializer_list": "cpp",
+        "iomanip": "cpp",
+        "iosfwd": "cpp",
+        "iostream": "cpp",
+        "istream": "cpp",
+        "limits": "cpp",
+        "memory": "cpp",
+        "new": "cpp",
+        "numbers": "cpp",
+        "numeric": "cpp",
+        "optional": "cpp",
+        "ostream": "cpp",
+        "ranges": "cpp",
+        "ratio": "cpp",
+        "sstream": "cpp",
+        "stdexcept": "cpp",
+        "streambuf": "cpp",
+        "string": "cpp",
+        "string_view": "cpp",
+        "system_error": "cpp",
+        "type_traits": "cpp",
+        "tuple": "cpp",
+        "typeindex": "cpp",
+        "typeinfo": "cpp",
+        "utility": "cpp",
+        "variant": "cpp",
+        "any": "cpp"
+    }
+}

dialects/affine/passes/lower_for.cpp

@@ -35,18 +35,20 @@ const Def* LowerFor::rewrite(const Def* def) {
         { // construct for
             auto [iter, end, step, acc] = for_lam->vars<4>();
 
+
+
             // reduce the body to remove the cn parameter
             auto nom_body = body->as_nom<Lam>();
-            auto new_body = nom_body->stub(w, w.cn(w.sigma()), body->dbg());
-            new_body->set(nom_body->reduce(w.tuple({iter, acc, yield_lam})));
+            auto new_body = nom_body->stub(w, w.cn(acc->type()), body->dbg())->as<Lam>();
+            new_body->set(nom_body->reduce(w.tuple({iter, new_body->var(), yield_lam})));
 
             // break
-            auto if_else_cn = w.cn(w.sigma());
+            auto if_else_cn = w.cn(acc->type());
             auto if_else    = w.nom_lam(if_else_cn, nullptr);
-            if_else->app(false, brk, acc);
+            if_else->app(false, brk, if_else->var());
 
             auto cmp = core::op(core::icmp::ul, iter, end);
-            for_lam->branch(false, cmp, new_body, if_else, w.tuple());
+            for_lam->branch(false, cmp, new_body, if_else, acc);
         }
 
         DefArray for_args{for_ax->num_args() - 2, [&](size_t i) { return for_ax->arg(i); }};

dialects/autodiff/autodiff.cpp

@@ -18,6 +18,7 @@ extern "C" THORIN_EXPORT thorin::DialectInfo thorin_get_dialect_info() {
             [](thorin::PipelineBuilder& builder) {
                 builder.add_opt(110);
                 builder.add_opt(120);
+                builder.add_opt(300);
                 builder.extend_opt_phase(105, [](thorin::PassMan& man) { man.add<thorin::autodiff::AutoDiffEval>(); });
                 builder.extend_opt_phase(111, [](thorin::PassMan& man) {
                     // in theory only after partial eval (beta, ...)

dialects/autodiff/autodiff.h

@@ -127,6 +127,7 @@ inline const Def* op_autodiff(const Def* fun) {
 
 inline const Def* op_zero(const Def* A) {
     World& world = A->world();
+    A->dump();
     return world.app(world.ax<zero>(), A);
 }
 

dialects/autodiff/auxiliary/autodiff_aux.cpp

@@ -44,40 +44,6 @@ const Def* zero_pullback(const Def* E, const Def* A) {
 //  TODO: rename to op_tangent_type
 const Def* tangent_type_fun(const Def* ty) { return ty; }
 
-const Def* equip_mem(const Def* def){
-    auto& world = def->world();
-    auto memType = mem::type_mem(world);
-    if(match<mem::M>(def->proj(0))){
-        return def;
-    }
-
-    def->dump();
-    if(def->isa<Sigma>()){
-        size_t size = def->num_ops() + 1;
-        DefArray newOps(size, [&](size_t i){
-            return i == 0 ? memType : def->op(i - 1);
-        });
-
-        return world.sigma(newOps);
-    }else if(auto pack = def->isa<Pack>()){
-        auto count = as_lit(pack->shape());
-        DefArray newOps(count + 1, [&](size_t i){
-            return i == 0 ? memType : pack->body();
-        });
-
-        return world.sigma(newOps);
-    }else if(auto pack = def->isa<Arr>()){
-        auto count = as_lit(pack->shape());
-        DefArray newOps(count + 1, [&](size_t i){
-            return i == 0 ? memType : pack->body();
-        });
-
-        return world.sigma(newOps);
-    }else{
-        return world.sigma({memType, def});
-    }
-}
-
 /// computes pb type E* -> A*
 /// E - type of the expression (return type for a function)
 /// A - type of the argument (point of orientation resp. derivative - argument type for partial pullbacks)
@@ -186,7 +152,7 @@ const Def* zero_def(const Def* T) {
         auto shape = arr->shape();
         auto body  = arr->body();
         // auto inner_zero = zero_def(body);
-        auto inner_zero = world.app(world.ax<zero>(), body);
+        auto inner_zero = op_zero(body);
         auto zero_arr   = world.pack(shape, inner_zero);
         world.DLOG("zero_def for array of shape {} with type {}", shape, body);
         world.DLOG("zero_arr: {}", zero_arr);
@@ -206,14 +172,23 @@ const Def* zero_def(const Def* T) {
             return zero;
         }
     } else if (auto sig = T->isa<Sigma>()) {
-        DefArray ops(sig->ops(), [&](const Def* op) { return world.app(world.ax<zero>(), op); });
+        DefArray ops(sig->ops(), [&](const Def* op) { return op_zero(op); });
         return world.tuple(ops);
     }
+
+    if(match<mem::M>(T)){
+        return world.bot(mem::type_mem(world));
+    }
+
+    if(match<mem::Ptr>(T)){
+        auto lit_zero = world.lit_int_(64, 0);
+        return core::op_bitcast(T, lit_zero);
+    }
 
     // or return bot
     // assert(0);
     // or id => zero T
-    // return world.app(world.ax<zero>(), T);
+    T->dump();
     return nullptr;
 }
 
@@ -227,6 +202,121 @@ const Def* op_sum(const Def* T, DefArray defs) {
 
 namespace thorin {
 
+
+
+const Pi* cn_mem_wrap(const Pi* pi){
+    auto &world = pi->world();
+    auto dom = pi->dom();
+
+    const Pi* result;
+    if(pi->ret_pi()){    
+        auto arg = equip_mem(dom->proj(0));
+        auto ret_pi = cn_mem_wrap(dom->proj(1)->as<Pi>());
+        result = world.cn({arg, ret_pi});
+    }else{
+        auto arg = equip_mem(dom);
+        result = world.cn({arg});
+    }
+
+    return result;
+}
+
+const Def* lam_mem_wrap(const Def* lam){
+    auto &world = lam->world();
+    auto type = lam->type()->as<Pi>();
+    if(!match<mem::M>(type->dom(0)->proj(0))){
+
+        auto wrap = cn_mem_wrap(type);
+
+        auto mem_lam = world.nom_lam(wrap, world.dbg("memorized_" + lam->name()));
+        auto lam_return = world.nom_lam(type->ret_pi(), world.dbg("memorized_return_" + lam->name()));
+
+        auto mem_vars = mem_lam->var((nat_t)0)->projs();
+        auto mem = mem_vars[0];
+        auto vars = lam_return->vars();
+
+        auto compound = world.builder().add(mem).add(vars).tuple();
+        auto compound2 = world.builder().add(mem_vars.skip_front()).add(lam_return).tuple();
+
+        lam_return->set_body(
+            world.app(mem_lam->ret_var(), compound)
+        );
+
+        mem_lam->set_body(
+            world.app(lam, compound2)
+        );
+
+        mem_lam->set_filter(true);
+        lam_return->set_filter(true);
+        return mem_lam;
+    }
+
+    return lam;
+}
+
+const Def* get_mem(const Def* def){
+    auto& world = def->world();
+
+    if(match<mem::M>(def->type()->op(0))){
+        return def->proj(0);
+    } 
+
+    return nullptr;
+}
+
+const Def* remove_mem(const Def* def){
+    auto& world = def->world();
+
+    if(def->isa<Sigma>() && match<mem::M>(def->op(0))){
+        auto ops = def->ops();
+        return world.sigma(ops.skip_front(1));
+    }else if(match<mem::M>(def->type()->op(0))){
+        auto ops = def->projs();
+        return world.tuple(ops.skip_front(1));
+    } 
+
+    return def;
+}
+
+const Def* equip_mem(const Def* def){
+    auto& world = def->world();
+    auto memType = mem::type_mem(world);
+    if(match<mem::M>(def->proj(0))){
+        return def;
+    }
+
+    if(def->isa<Sigma>()){
+        size_t size = def->num_ops() + 1;
+        DefArray newOps(size, [&](size_t i){
+            return i == 0 ? memType : def->op(i - 1);
+        });
+
+        return world.sigma(newOps);
+    }else if(auto pack = def->isa<Pack>()){
+        auto count = as_lit(pack->shape());
+        DefArray newOps(count + 1, [&](size_t i){
+            return i == 0 ? memType : pack->body();
+        });
+
+        return world.sigma(newOps);
+    }else if(auto pack = def->isa<Arr>()){
+        auto count = as_lit(pack->shape());
+        DefArray newOps(count + 1, [&](size_t i){
+            return i == 0 ? memType : pack->body();
+        });
+
+        return world.sigma(newOps);
+    }else{
+        return world.sigma({memType, def});
+    }
+}
+
+const Def* continuation_codom(const Def* E) {
+    auto pi = E->as<Pi>();
+    assert(pi != NULL);
+    return pi->dom(1)->as<Pi>()->dom();
+}
+
 bool is_continuation_type(const Def* E) {
     if (auto pi = E->isa<Pi>()) { return pi->codom()->isa<Bot>(); }
     return false;
@@ -262,12 +352,6 @@ const Def* continuation_dom(const Def* E) {
     return pi->dom(0);
 }
 
-const Def* continuation_codom(const Def* E) {
-    auto pi = E->as<Pi>();
-    assert(pi != NULL);
-    return pi->dom(1)->as<Pi>()->dom();
-}
-
 /// high level
 /// f: B -> C
 /// g: A -> B
@@ -287,12 +371,22 @@ const Def* compose_continuation(const Def* f, const Def* g) {
     world.DLOG("compose g (A->B): {} : {}", g, g->type());
     assert(is_returning_continuation(f));
     assert(is_returning_continuation(g));
+
+    f = lam_mem_wrap(f);
+    g = lam_mem_wrap(g);
 
     auto F = f->type()->as<Pi>();
     auto G = g->type()->as<Pi>();
 
+    auto is_mem = match<mem::M>(F->dom(0)->proj(0));
+
+    F->dump();
+    G->dump();
+
     auto A = continuation_dom(G);
     auto B = continuation_codom(G);
+
+    auto B_hat = continuation_dom(F);
     auto C = continuation_codom(F);
     // better handled by application type checks
     // auto B2 = continuation_dom(F);
@@ -307,7 +401,7 @@ const Def* compose_continuation(const Def* f, const Def* g) {
     auto H     = world.cn({A, world.cn(C)});
     auto Hcont = world.cn(B);
 
-    auto h     = world.nom_lam(H, world.dbg("comp_" + f->name() + "_" + g->name()));
+    auto h     = world.nom_lam(H,  world.dbg("comp_" + f->name() + "_" + g->name()));
     auto hcont = world.nom_lam(Hcont, world.dbg("comp_" + f->name() + "_" + g->name() + "_cont"));
 
     h->app(true, g, {h->var((nat_t)0), hcont});

dialects/autodiff/auxiliary/autodiff_aux.h

@@ -22,6 +22,13 @@ const Def* op_sum(const Def* T, DefArray defs);
 
 namespace thorin {
 
+
+const Def* get_mem(const Def* def);
+const Def* equip_mem(const Def* def);
+const Def* remove_mem(const Def* def);
+const Def* lam_mem_wrap(const Def* lam);
+const Pi* cn_mem_wrap(const Pi* pi);
+
 bool is_continuation_type(const Def* E);
 bool is_continuation(const Def* e);
 // TODO: change name to returning_continuation

dialects/autodiff/auxiliary/autodiff_rewrite_inner.cpp

@@ -106,6 +106,12 @@ const Def* AutoDiffEval::augment_extract(const Extract* ext, Lam* f, Lam* f_diff
     auto aug_tuple = augment(tuple, f, f_diff);
     auto aug_index = augment(index, f, f_diff);
 
+    auto aug_ext              = world.extract(aug_tuple, aug_index);
+
+    if(match<mem::M>(ext->type())){
+        return aug_ext;
+    }
+
     // TODO: if not exists use:
     // e:T, b:B
     // b = e#i
@@ -125,16 +131,40 @@ const Def* AutoDiffEval::augment_extract(const Extract* ext, Lam* f, Lam* f_diff
         auto pb_fun   = world.nom_lam(pb_ty, world.dbg("extract_pb"));
         world.DLOG("Pullback: {} : {}", pb_fun, pb_fun->type());
         auto pb_tangent = pb_fun->var((nat_t)0, world.dbg("s"));
-        auto tuple_tan  = world.insert(op_zero(aug_tuple->type()), aug_index, pb_tangent, world.dbg("tup_s"));
-        pb_fun->app(true, tuple_pb,
-                    {
-                        tuple_tan,
-                        pb_fun->var(1) // ret_var but make sure to select correct one
-                    });
+
+        auto aug_tuple_type = aug_tuple->type();
+
+        auto rm_test = remove_mem(aug_tuple_type);
+
+        auto mem = get_mem(pb_tangent);
+        pb_tangent = remove_mem(pb_tangent);
+
+        //auto test = world.app(world.ax<mem::remem>(), mem);
+
+
+        //core::op(core::wrap::sub, core::WMode::none,  );
+
+        auto init = op_zero(rm_test);
+        const Def* tuple_tan;
+        if(rm_test != aug_tuple_type){
+            auto index_lit = as_lit(aug_index);
+            auto arity_lit = as_lit(aug_index->type()->as<Idx>()->size());
+
+            auto aug_index = world.lit_idx(index_lit - 1, arity_lit - 1);
+            if(arity_lit == 2){
+                tuple_tan = pb_tangent;
+            }else if(arity_lit >= 2){
+                tuple_tan  = world.insert(init, aug_index, pb_tangent, world.dbg("tup_s"));
+            }
+        }else{
+            tuple_tan  = world.insert(init, aug_index, pb_tangent, world.dbg("tup_s"));
+        }
+
+        auto arg = world.builder().add(mem).add(tuple_tan).tuple();
+
+        pb_fun->app(true, tuple_pb, {arg, pb_fun->var(1)});
         pb = pb_fun;
     }
-
-    auto aug_ext              = world.extract(aug_tuple, aug_index);
     partial_pullback[aug_ext] = pb;
 
     return aug_ext;
@@ -149,10 +179,16 @@ const Def* AutoDiffEval::augment_tuple(const Tuple* tup, Lam* f, Lam* f_diff) {
 
     auto projs = tup->projs();
     // TODO: should use ops instead?
-    DefArray aug_ops(projs.skip_front(isMem), [&](const Def* op) { return augment(op, f, f_diff); });
+    DefArray aug_ops(projs, [&](const Def* op) { 
+        return augment(op, f, f_diff); 
+    });
     auto aug_tup = world.tuple(aug_ops);
 
-    DefArray pbs(aug_ops, [&](const Def* op) { return partial_pullback[op]; });
+    DefArray pbs(aug_ops.skip_front(isMem), [&](const Def* op) { 
+        auto pb = partial_pullback[op];    
+        assert(pb);
+        return pb;
+    });
     world.DLOG("tuple pbs {,}", pbs);
     // create shadow pb
     auto shadow_pb           = world.tuple(pbs);
@@ -173,18 +209,21 @@ const Def* AutoDiffEval::augment_tuple(const Tuple* tup, Lam* f, Lam* f_diff) {
     // TODO: move op_cps2ds to direct dialect and merge then
     auto T = tangent_type_fun(f_arg_ty);
     auto mem = pb->var((nat_t)0)->proj(0);
-    auto sum = world.app(world.ax<zero>(), T);
+    //auto T_without_mem = remove_mem(T);
+    auto sum = autodiff_zero(mem, f);
 
     for(size_t i = 0 ; i < pbs.size() ; i++){
         auto re = direct::op_cps2ds_dep(pbs[i]);
+        auto codom = re->type()->as<Pi>()->codom();
+        auto isMem = match<mem::M>(codom->proj(0));
         auto extract = world.extract(pb_tangent, i + 1);
         auto app = world.app(re, {mem, extract});
         mem = world.extract(app, (nat_t)0);
-        mem->dump();
-        mem->dump();
         sum = world.app(world.app(world.ax<add>(), T), {sum, app});
     }
 
+    sum = world.insert(sum, (u64)0, mem);
+
     //DefArray tangents(pbs.size(), [&](nat_t i) { return world.app(direct::op_cps2ds_dep(pbs[i]), world.extract(pb_tangent, i)); });
     pb->app(true, pb->var(1), sum);
     partial_pullback[aug_tup] = pb;
@@ -241,6 +280,7 @@ const Def* AutoDiffEval::augment_app(const App* app, Lam* f, Lam* f_diff) {
     auto callee = app->callee();
     auto arg    = app->arg();
 
+    arg->type()->dump();
     auto aug_arg    = augment(arg, f, f_diff);
     auto aug_callee = augment(callee, f, f_diff);
     // auto arg_ppb    = partial_pullback[aug_arg];
@@ -406,53 +446,6 @@ const Def* AutoDiffEval::augment_app(const App* app, Lam* f, Lam* f_diff) {
         auto r_pb = c1->var(1);
         c1->app(true, aug_cont, {res, compose_continuation(e_pb, r_pb)});
 
-        // auto X = continuation_codom(g->type());
-        // // auto A = f_diff->var((nat_t)0);
-        // auto A = f_diff->type()->dom(0);
-        // // auto E = g_deriv->var((nat_t)0);
-        // auto E = g_deriv->type()->as<Pi>()->dom(0);
-        // world.DLOG("A (var f): {}", A);
-        // world.DLOG("E (var g): {}", E);
-        // world.DLOG("X (out g = out f): {}", X);
-        // // auto ret_g_deriv = g_deriv->var(1);
-        // auto ret_g_deriv_ty = g_deriv->type()->as<Pi>()->dom(1);
-        // world.DLOG("ret_g_deriv_ty: {} ", ret_g_deriv_ty);
-        // // auto ret_f_deriv=f_diff->var(1);
-        // // world.DLOG("ret_f_deriv: {} : {}", ret_f_deriv, ret_f_deriv->type());
-
-        // // TODO: better debug names
-        // auto c1_ty=ret_g_deriv_ty->as<Pi>();
-        // world.DLOG("c1_ty: (cn[X, cn[X+, cn E+]]) {}", c1_ty);
-        // auto c2_ty=aug_cont->type()->as<Pi>()->dom(2)->as<Pi>();
-        // world.DLOG("c2_ty: (cn[X+, cn A+]) {}", c2_ty);
-        // auto c3_ty=c1_ty->dom(2)->as<Pi>()->dom(2)->as<Pi>();
-        // world.DLOG("c3_ty: (cn E+) {}", c3_ty);
-        // auto c1 = world.nom_lam(c1_ty,world.dbg("c1"));
-        // auto c2 = world.nom_lam(c2_ty,world.dbg("c2"));
-        // auto c3 = world.nom_lam(c3_ty,world.dbg("c3"));
-
-        // c1->app(true,
-        //     aug_cont,
-        //     {
-        //         c1->var((nat_t)0),
-        //         c2
-        //     }
-        // );
-        // c2->app(true,
-        //     c1->var(1),
-        //     {
-        //         c2->var((nat_t)0),
-        //         c3
-        //     }
-        // );
-        // c3->app(true,
-        //     e_pb,
-        //     {
-        //         c3->var((nat_t)0),
-        //         c2->var(1)
-        //     }
-        // );
-
         auto aug_app = world.app(aug_callee, {real_aug_args, c1});
         world.DLOG("aug_app: {} : {}", aug_app, aug_app->type());
 
@@ -471,25 +464,102 @@ const Def* AutoDiffEval::augment_app(const App* app, Lam* f, Lam* f_diff) {
     assert(false && "should not be reached");
 }
 
+static const Def* tuple_of_types(const Def* t) {
+    auto& world = t->world();
+    if (auto sigma = t->isa<Sigma>()) return world.tuple(sigma->ops());
+    if (auto arr = t->isa<Arr>()) return world.pack(arr->shape(), arr->body());
+    return t;
+}
+
+const Def* op_lea(const Def* ptr, const Def* index, const Def* dbg = {}) {
+    auto &world = ptr->world();
+    auto [pointee, addr_space] = match<mem::Ptr>(ptr->type())->args<2>();
+    auto Ts                    = tuple_of_types(pointee);
+    return world.app(world.app(world.ax<mem::lea>(), {pointee->arity(), Ts, addr_space}), {ptr, index}, dbg);
+}
+
+const Def* op_load(const Def* mem, const Def* ptr, const Def* dbg = {}) {
+    mem->dump();
+    auto &world = mem->world();
+    auto [T, a] = match<mem::Ptr>(ptr->type())->args<2>();
+    return world.app(world.app(world.ax<mem::load>(), {T, a}), {mem, ptr}, dbg);
+}
+
+const Def* op_store(const Def* mem, const Def* ptr, const Def* val, const Def* dbg = {}) {
+    auto &world = mem->world();
+    auto [T, a] = match<mem::Ptr>(ptr->type())->args<2>();
+    return world.app(world.app(world.ax<mem::store>(), {T, a}), {mem, ptr, val}, dbg);
+}
 
 const Def* AutoDiffEval::augment_lea(const App* lea, Lam* f, Lam* f_diff) {
-    auto [ptr, as] = augment(lea->arg(), f, f_diff)->projs<2>();
+    auto &world = lea->world();
+
+    auto [arr_ptr, idx] = lea->arg()->projs<2>();
+
+    auto aug_ptr = augment(arr_ptr, f, f_diff);
+    auto aug_idx = augment(idx, f, f_diff);
+
+    auto aug_arg = augment(lea->arg(), f, f_diff);
+
+    //auto [ptr, idx] = aug_arg->projs<2>();
+    //auto ptr_def = ptr->as<App>();
+
+    auto ptr_ty = lea->type()->as<App>();
+    auto elem_ty = ptr_ty->arg(0);
+
+    auto [arg_ty, ret_pi] = f->type()->doms<2>();
+
+    auto pb_type = pullback_type(elem_ty, arg_ty);
+
+    auto lam = world.nom_lam(pb_type, world.dbg("test"));
 
-    auto pb = partial_pullback[ptr];
+    auto pb_arg = lam->var((nat_t)0);
+    auto pb_ret = lam->var((nat_t)1);
+    auto pb_mem = pb_arg->proj(0);
+    auto pb_s = pb_arg->proj(1);
 
-    lea->dump();
-    pb->dump();
-    lea->dump();
-    return lea;
+    auto gradient_array = shadow_gradient_array[aug_ptr];
+
+    auto gradient_lea = op_lea(gradient_array, aug_idx, world.dbg("gradient_array_lea"));
+    auto [gradient_mem, gradient] = op_load(pb_mem, gradient_lea,  world.dbg("gradient_array_load"))->projs<2>();
+    //auto [pb_s_mem, pb_s] = op_load(gradient_mem, pb_s_ptr)->projs<2>();
+    auto add = core::op(core::wrap::add, core::WMode::none, gradient, pb_s);
+    auto store_mem = op_store(gradient_mem, gradient_lea, add, world.dbg("add_to_gradient") );
+
+    //auto zero_ret = world.app(world.ax<zero>(), arg_ty);
+
+    auto default_zero = autodiff_zero(store_mem, f);
+    lam->set_body(world.app(pb_ret, {default_zero}));
+    lam->set_filter(true);
+
+    auto aug_lea = world.app(lea->callee(), aug_arg);
+    partial_pullback[aug_lea] = lam;
+    return aug_lea;
 }
 
 const Def* AutoDiffEval::augment_load(const App* load, Lam* f, Lam* f_diff) {
+    auto &world = load->world();
 
-    auto aug_arg = augment(load->arg(), f, f_diff);
+    auto [mem, lea] = load->arg()->projs<2>();
+
+    auto aug_mem = augment(mem, f, f_diff);
+    auto aug_lea = augment(lea, f, f_diff);
+
+    auto aug_load = op_load(aug_mem, aug_lea, world.dbg("aug_load"));
+
+    auto pb_lea = partial_pullback[aug_lea];
+    partial_pullback[aug_load] = pb_lea;
+
+    /*auto aug_arg = augment(load->arg(), f, f_diff);
+    auto pb = partial_pullback[aug_arg];
+    auto aug_load = world.app(load->callee(), aug_arg);
+
+
+    //auto [mem, ptr] = aug_arg->proj<2>();
 
     load->dump();
-    load->dump();
-    return load;
+    load->dump();*/
+    return aug_load;
 }
 
 const Def* AutoDiffEval::augment_store(const App* store, Lam* f, Lam* f_diff) {

dialects/autodiff/auxiliary/autodiff_rewrite_toplevel.cpp

@@ -2,6 +2,9 @@
 #include "dialects/autodiff/auxiliary/autodiff_aux.h"
 #include "dialects/autodiff/passes/autodiff_eval.h"
 
+#include "dialects/mem/autogen.h"
+#include "dialects/mem/mem.h"
+
 namespace thorin::autodiff {
 
 // void AutoDiffEval::create_shadow_id_pb(const Def* def) {
@@ -41,13 +44,74 @@ namespace thorin::autodiff {
 // //     // no structure => needs no structure pullback (base case also needs no str. pb because it is shallow)
 // }
 
+const Def* hoa(const Def* def, const Def* arg_ty){
+    if(auto arr = def->isa<Arr>()){
+        auto &world = def->world();
+        auto shape = arr->shape();
+        auto body = hoa(arr->body(), arg_ty);
+        return world.arr(shape, body);
+    }
+
+    auto pb_ty = pullback_type(def, arg_ty);
+    return pb_ty;
+}
+
+const Def* AutoDiffEval::autodiff_zero(const Def* mem, Lam* f) {
+    return autodiff_zero(mem, augmented[f->var()]->proj(0));
+}
+
+const Def* AutoDiffEval::autodiff_zero(const Def* mem, const Def* def) {
+    auto& world = def->world();
+
+    auto ty = def->type();
+
+    if (auto tup = def->isa<Tuple>()) {
+        DefArray ops(tup->ops(), [&](const Def* op) { return autodiff_zero(mem, op); });
+        return world.tuple(ops);
+    }
+/*
+    if (auto var = def->isa<Var>()) {
+        DefArray ops(var->projs(), [&](const Def* op) { return autodiff_zero(op); });
+        return world.tuple(ops);
+    }*/
+
+    if (auto app = ty->isa<App>()) {
+        auto callee = app->callee();
+        // auto args = app->args();
+        world.DLOG("app callee: {} : {} <{}>", callee, callee->type(), callee->node_name());
+        // TODO: can you directly match Tag::Int?
+        if (callee->isa<Idx>()) {
+            // auto size = app->arg(0);
+            auto zero = world.lit_idx(ty, 0, world.dbg("zero"));
+            // world.DLOG("zero_def for int of size {} is {}", size, zero);
+            world.DLOG("zero_def for int is {}", zero);
+            return zero;
+        }
+    } 
+
+    if(match<mem::M>(ty)){
+        return mem;
+    }
+
+    if(match<mem::Ptr>(ty)){
+        return shadow_gradient_array[def];
+    }
+
+    def->dump();
+    def->type()->dump();
+    return nullptr;
+}
+
 /// side effect: register pullback
 const Def* AutoDiffEval::derive_(const Def* def) {
     auto& world = def->world();
     if (auto lam = def->isa_nom<Lam>()) {
         world.DLOG("Derive lambda: {}", def);
         auto deriv_ty = autodiff_type_fun_pi(lam->type());
         auto deriv    = world.nom_lam(deriv_ty, world.dbg(lam->name() + "_deriv"));
+        auto memType = mem::type_mem(world);
+        auto deriv_inner    = world.nom_lam(world.cn({memType}), world.dbg(lam->name() + "_deriv_inner"));
+
 
         // pre register derivative
         // needed for recursion
@@ -78,8 +142,6 @@ const Def* AutoDiffEval::derive_(const Def* def) {
         //     id_pb_scalar
         // );
 
-        auto deriv_all_args  = deriv->var();
-        const Def* deriv_arg = deriv->var((nat_t)0, world.dbg("arg"));
         // R auto deriv_ret = deriv->var((nat_t)1, world.dbg("ret"));
         // R partial_pullback[deriv_arg] = id_pb;
 
@@ -102,9 +164,43 @@ const Def* AutoDiffEval::derive_(const Def* def) {
         //
         // but the DS/CPS special case has to be handled separately
 
+        const Def* var = deriv->var();
+
+        if(match<mem::M>(deriv->dom((nat_t) 0)->proj(0))){
+            auto vars = var->projs();
+            auto arg = vars[0];
+            current_mem = arg->proj(0);
+            auto args = arg->projs();
+            args[0] = deriv_inner->var();
+            arg = world.tuple(args);
+
+            for(auto var : deriv->var((nat_t)0)->projs()){
+                auto var_ty = var->type();
+                if(auto ptr = match<mem::Ptr>(var_ty)){
+
+                    auto [mem2, gradient_ptr] = mem::op_malloc(ptr->arg(0), current_mem, world.dbg("gradient_arr"))->projs<2>();
+                    auto pb_ty = hoa(ptr->arg(0), arg_ty);
+                    auto [mem3, pb_ptr] = mem::op_malloc(pb_ty, mem2, world.dbg("pullback_arr"))->projs<2>();
+                    current_mem = mem3;
+
+                    shadow_gradient_array[var] = gradient_ptr;
+                    shadow_pullback_array[var] = pb_ptr;
+                    var_ty->dump();
+                }
+            }
+
+            vars[0] = arg;
+            var = world.tuple(vars);
+        }
+
+        augmented[lam->var()] = var;
+
+        //auto deriv_all_args  = deriv->var();
+
         // TODO: check identity
         // could use identity tangent(arg_ty) = tangent(augment(arg_ty))
         // with deriv_arg->type() = augment(arg_ty)
+        const Def* deriv_arg = deriv->var((nat_t)0, world.dbg("arg"));
         auto arg_id_pb              = id_pullback(arg_ty);
         partial_pullback[deriv_arg] = arg_id_pb;
         // set no pullback to all_arg and return
@@ -115,11 +211,11 @@ const Def* AutoDiffEval::derive_(const Def* def) {
         auto ret_pb               = zero_pullback(lam->var(1)->type(), arg_ty);
         partial_pullback[ret_var] = ret_pb;
 
-        shadow_pullback[deriv_all_args] = world.tuple({arg_id_pb, ret_pb});
+        shadow_pullback[var] = world.tuple({arg_id_pb, ret_pb});
         world.DLOG("pullback for argument {} : {} is {} : {}", deriv_arg, deriv_arg->type(), arg_id_pb,
                    arg_id_pb->type());
-        world.DLOG("args shadow pb is {} : {}", shadow_pullback[deriv_all_args],
-                   shadow_pullback[deriv_all_args]->type());
+        world.DLOG("args shadow pb is {} : {}", shadow_pullback[var],
+                   shadow_pullback[var]->type());
 
         // TODO: remove as this is subsumed by lam->deriv
         // R const Def* lam_ret   = lam->var(1, world.dbg("ret"));
@@ -136,7 +232,7 @@ const Def* AutoDiffEval::derive_(const Def* def) {
         // TODO: transively remove
         // arguments (maybe not necessary)
         // auto src_arg = lam->var()
-        augmented[lam->var()] = deriv->var();
+
         world.DLOG("Associate vars {} with {}", lam->var(), deriv->var());
 
         // already contains the correct application of
@@ -146,15 +242,20 @@ const Def* AutoDiffEval::derive_(const Def* def) {
         //   this is needed for continuations (without closure conversion)
         //   but also essentially for the return continuation
 
+
         // reminder of types:
         // expression e: B
         //   implicit: e_fun: A -> B
         // partial pullback: e*: B* -> A*
         // partial derivative: e': B' × (B* -> A*)
         //   implicit: e'_fun: A' -> B' × (B* -> A*)
         auto new_body = augment(lam->body(), lam, deriv);
+        deriv_inner->set_filter(true);
+        deriv_inner->set_body(new_body);
+
         deriv->set_filter(true);
-        deriv->set_body(new_body);
+        deriv->set_body(world.app(deriv_inner, {current_mem}));
+        deriv->dump(10);
 
         return deriv;
     }

dialects/autodiff/normalizers.cpp

@@ -136,7 +136,7 @@ const Def* normalize_sum(const Def* type, const Def* callee, const Def* arg, con
         // R if(val == 1)
         // R     return arg;
         DefArray args = arg->projs(val);
-        auto sum      = world.app(world.ax<zero>(), T);
+        auto sum      = op_zero(T);
         // would also be handled by add zero
         if (val >= 1) { sum = args[0]; }
         for (auto i = 1; i < val; ++i) { sum = world.app(world.app(world.ax<add>(), T), {sum, args[i]}); }

dialects/autodiff/passes/autodiff_eval.cpp

@@ -35,14 +35,9 @@ const Def* AutoDiffEval::rewrite(const Def* def) {
         world().DLOG("found a autodiff::autodiff of {}", arg);
         // world.DLOG("found a autodiff::autodiff {} to {}",callee,arg);
 
-        if (arg->isa<Lam>()) {
-            // world.DLOG("found a autodiff::autodiff of a lambda");
-            return derive(arg);
-        }
-
+        assert(arg->isa<Lam>());
         // TODO: handle operators analogous 
-
-        assert(0);
+        def = derive(arg);
         return def;
     }
 

dialects/autodiff/passes/autodiff_eval.h

@@ -40,6 +40,9 @@ class AutoDiffEval : public RWPass<AutoDiffEval, Lam> {
     const Def* augment_load(const App*, Lam*, Lam*);
     const Def* augment_store(const App*, Lam*, Lam*);
 
+    const Def* autodiff_zero(const Def* mem, Lam* f);
+    const Def* autodiff_zero(const Def* mem, const Def* def);
+
     /// fills partial_pullback and shadow/structure pullback maps
     void create_shadow_id_pb(const Def*);
 
@@ -73,6 +76,10 @@ class AutoDiffEval : public RWPass<AutoDiffEval, Lam> {
 
     // TODO: remove?
     Def2Def app_pb;
+
+    Def2Def shadow_pullback_array;
+    Def2Def shadow_gradient_array;
+    const Def* current_mem;
 };
 
 } // namespace thorin::autodiff

dialects/core/be/ll/ll.cpp

@@ -221,7 +221,8 @@ std::string Emitter::convert_ret_pi(const Pi* pi) {
 void Emitter::start() {
     Super::start();
 
-    ostream() << "declare i8* @malloc(i64)" << '\n'; // HACK
+    ostream() << "declare i8* @calloc(i64)" << '\n'; // HACK
+    ostream() << "declare void @free(i8*)" << '\n'; // HACK
     // SJLJ intrinsics (GLIBC Versions)
     ostream() << "declare i32 @_setjmp(i8*) returns_twice" << '\n';
     ostream() << "declare void @longjmp(i8*, i32) noreturn" << '\n';
@@ -727,8 +728,16 @@ std::string Emitter::emit_bb(BB& bb, const Def* def) {
         emit_unsafe(malloc->arg(0));
         auto size  = emit(malloc->arg(1));
         auto ptr_t = convert(force<mem::Ptr>(def->proj(1)->type()));
-        bb.assign(name + ".i8", "call i8* @malloc(i64 {})", size);
+        bb.assign(name + ".i8", "call i8* @calloc(i64 {})", size);
         return bb.assign(name, "bitcast i8* {} to {}", name + ".i8", ptr_t);
+    }  else if (auto free = match<mem::free>(def)) {
+        emit_unsafe(free->arg(0));
+        auto ptr  = emit(free->arg(1));
+        auto ptr_t = convert(force<mem::Ptr>(free->arg(1)->type()));
+
+        bb.assign(name + ".i8", "bitcast {} {} to i8*", ptr_t, ptr);
+        bb.tail("call void @free(i8* {})", name + ".i8");
+        return {};
     } else if (auto mslot = match<mem::mslot>(def)) {
         emit_unsafe(mslot->arg(0));
         // TODO array with size
@@ -817,6 +826,9 @@ std::string Emitter::emit_bb(BB& bb, const Def* def) {
         print(vars_decls_, "{} = global {} {}\n", name, convert(pointee), init);
         return globals_[global] = name;
     }
+    def->dump();
+    def->dump();
+    def->dump();
 
     unreachable(); // not yet implemented
 }

dialects/mem/mem.h

@@ -109,6 +109,13 @@ inline const Def* op_malloc(const Def* type, const Def* mem, const Def* dbg) {
     return w.app(w.app(w.ax<malloc>(), {type, w.lit_nat_0()}), {mem, size}, dbg);
 }
 
+inline const Def* op_free(const Def* type, const Def* mem, const Def* ptr, const Def* dbg) {
+    World& w  = type->world();
+    auto ptr_ty = match<Ptr>(ptr->type())->as<App>();
+    auto pointee = ptr_ty->arg(0_s);
+    return w.app(w.app(w.ax<free>(), {pointee, w.lit_nat_0()}), {mem, ptr}, dbg);
+}
+
 inline const Def* op_mslot(const Def* type, const Def* mem, const Def* id, const Def* dbg) {
     World& w  = type->world();
     auto size = core::op(core::trait::size, type);

dialects/mem/mem.thorin

@@ -52,6 +52,11 @@
 /// The difference to %mem.alloc is that the `size` is automatically inferred.
 .ax %mem.malloc: Π [T: *, as: .Nat] -> [%mem.M, .Nat] -> [%mem.M, %mem.Ptr(T, as)];
 ///
+/// ### %mem.free
+///
+/// Frees memory of type `T` in address space `as`.
+.ax %mem.free: Π [T: *, as: .Nat] -> [%mem.M, %mem.Ptr(T, as)] -> [%mem.M];
+///
 /// ### %mem.mslot
 ///
 /// Reserves a memory slot for type `T` in address space `as`.

lit/autodiff/arr.thorin

@@ -7,35 +7,66 @@
 .import core;
 .import autodiff;
 .import mem;
+.import affine;
 
 .let i32 = .Idx 4294967296;
+.let arr_size = 100:.Nat;
+
+.cn printInteger [mem: %mem.M, val: i32, return : .Cn [%mem.M]];
 
 .cn f [[mem : %mem.M, a: %mem.Ptr («100:.Nat; i32», 0)], ret: .Cn [mem : %mem.M, i32]] = {
-    .let lea = %mem.lea (100:.Nat, <100:.Nat; i32>, 0) (a, 1:(.Idx 100));
-    .let (load_mem, load_val) = %mem.load (i32, 0) (mem, lea);
-    ret (load_mem, load_val)
+    .let lea1 = %mem.lea (100:.Nat, <100:.Nat; i32>, 0) (a, 0:(.Idx 100));
+    .let lea2 = %mem.lea (100:.Nat, <100:.Nat; i32>, 0) (a, 1:(.Idx 100));
+    .let (load_mem1, load_val1) = %mem.load (i32, 0) (mem, lea1);
+    .let (load_mem2, load_val2) = %mem.load (i32, 0) (load_mem1, lea2);
+    .let scaled = %core.wrap.mul (0:.Nat, 4294967296:.Nat) (load_val1, load_val2);
+    ret (load_mem2, scaled)
+};
+
+.cn init [mem: %mem.M, arr : %mem.Ptr (<<100:.Nat; i32>>, 0:.Nat), offset : i32, return : .Cn [%mem.M]] = {
+    .cn for_body [i : i32, mem : %mem.M, continue : .Cn [%mem.M]] = {
+        .let idx_100_i = %core.bitcast ((.Idx 100), i32) i;
+        .let lea = %mem.lea (arr_size, <arr_size; i32>, 0) (arr, idx_100_i);
+        .let add = %core.wrap.add (0, 4294967296) (offset, i);
+        .let store_mem = %mem.store (i32, 0) (mem, lea, add);
+        continue(store_mem)
+    };
+    %affine.For (4294967296, 1, (%mem.M)) (0:i32, 100:i32, 1:i32, (mem), for_body, return)
 };
 
 .cn .extern main [mem : %mem.M, argc : i32, argv : %mem.Ptr (%mem.Ptr (.Idx 256, 0:.Nat), 0:.Nat), return : .Cn [%mem.M, i32]] = {
 
     .cn ret_cont [[mem : %mem.M, a: i32], pb:.Cn[[%mem.M, i32],.Cn[[%mem.M, %mem.Ptr («100:.Nat; i32», 0)]]]] = {
         .cn pb_ret_cont [mem : %mem.M, a: %mem.Ptr («100:.Nat; i32», 0)] = {
-            return (mem, 99:i32)
+            .let lea1 = %mem.lea (100:.Nat, <100:.Nat; i32>, 0) (a, 0:(.Idx 100));
+
+            .let (load_mem1, load_val) = %mem.load (i32, 0) (mem, lea1);
+
+            .cn print_integer_callback2 [mem : %mem.M] = {
+                .let free_mem = %mem.free («100:.Nat; i32», 0) (mem, a);
+                return (free_mem, 0:i32)
+            };
+
+            .cn print_integer_callback [mem : %mem.M] = {
+                .let lea2 = %mem.lea (100:.Nat, <100:.Nat; i32>, 0) (a, 1:(.Idx 100));
+                .let (load_mem2, load_val) = %mem.load (i32, 0) (load_mem1, lea2);
+                printInteger ( load_mem2, load_val, print_integer_callback2 )
+            };
+
+            printInteger ( load_mem1, load_val, print_integer_callback )
         };
-        pb((mem, a), pb_ret_cont)
+        pb((mem, 1:i32), pb_ret_cont)
     };
 
     .let (alloc_mem, alloc_val) = %mem.alloc («100:.Nat; i32», 0) mem;
 
-    //.let (alloc_mem, alloc_val) = %mem.malloc (i32, 0) (mem, 4);
-    .let lea = %mem.lea (100:.Nat, <100:.Nat; i32>, 0) (alloc_val, 1:(.Idx 100));
-    .let store = %mem.store (i32, 0) (alloc_mem, lea, argc);
-
     .let f_diff = %autodiff.autodiff (.Cn [[%mem.M, %mem.Ptr («100:.Nat; i32», 0)],.Cn[%mem.M, i32]]) f;
 
-    f_diff ((store, alloc_val),ret_cont)
+    .cn init_callback [mem : %mem.M] = {
+        f_diff ((mem, alloc_val),ret_cont)
+    };
 
-    //f ((store, alloc_val),return)
+    init(alloc_mem, alloc_val, 0:i32, init_callback)
 };
 
 

lit/autodiff/arr2.thorin

@@ -0,0 +1,79 @@
+// RUN: rm -f %t.ll ; \
+// RUN: %thorin -d autodiff %s --output-ll %t.ll --output-thorin - | FileCheck %s
+
+// a call to a autodiff style function
+// ./build/bin/thorin -d debug -d direct -d autodiff ./lit/autodiff/multiply_autodiff.thorin --output-thorin - -VVVV
+
+.import core;
+.import autodiff;
+.import mem;
+.import affine;
+
+.let i32 = .Idx 4294967296;
+
+.cn printInteger [mem: %mem.M, val: i32, return : .Cn [%mem.M]];
+
+.cn f [[mem : %mem.M, a: %mem.Ptr («100:.Nat; i32», 0), b: %mem.Ptr («100:.Nat; i32», 0)], ret: .Cn [mem : %mem.M, i32]] = {
+    .cn for_exit acc :: [mem : %mem.M, i32, i32] = {
+        .let lea = %mem.lea (arr_size, <arr_size; i32>, 0) (ptr, %core.conv.u2u (arr_size, 4294967296) (%core.wrap.sub (0, 4294967296) (argc, 4:i32)));
+        .let (load_mem, val) = %mem.load (i32, 0) (mem, lea);
+        ret (load_mem2, scaled)
+    };
+
+    .cn for_body [i : i32, [mem : %mem.M, acc: i32], continue : .Cn [%mem.M]] = {
+        .let lea1 = %mem.lea (100:.Nat, <100:.Nat; i32>, 0) (a, i);
+        .let lea2 = %mem.lea (100:.Nat, <100:.Nat; i32>, 0) (a, i);
+        .let (load_mem1, load_val1) = %mem.load (i32, 0) (mem, lea1);
+        .let (load_mem2, load_val2) = %mem.load (i32, 0) (load_mem1, lea2);
+        .let scaled = %core.wrap.mul (0:.Nat, 4294967296:.Nat) (load_val1, load_val2);
+        continue (load_mem2, scaled)
+    };
+
+    %affine.For (4294967296, 3, (%mem.M)) (0:(.Idx 100), 100:(.Idx 100), 1:(.Idx 100), (alloc_mem, 0:i32), for_body, for_exit)
+};
+
+.cn .extern main [mem : %mem.M, argc : i32, argv : %mem.Ptr (%mem.Ptr (.Idx 256, 0:.Nat), 0:.Nat), return : .Cn [%mem.M, i32]] = {
+
+    .cn ret_cont [[mem : %mem.M, a: i32], pb:.Cn[[%mem.M, i32],.Cn[[%mem.M, %mem.Ptr («100:.Nat; i32», 0)]]]] = {
+        .cn pb_ret_cont [mem : %mem.M, a: %mem.Ptr («100:.Nat; i32», 0)] = {
+            .let lea1 = %mem.lea (100:.Nat, <100:.Nat; i32>, 0) (a, 0:(.Idx 100));
+
+            .let (load_mem1, load_val) = %mem.load (i32, 0) (mem, lea1);
+
+            .cn print_integer_callback2 [mem : %mem.M] = {
+                .let free_mem = %mem.free («100:.Nat; i32», 0) (mem, a);
+                return (free_mem, 0:i32)
+            };
+
+            .cn print_integer_callback [mem : %mem.M] = {
+                .let lea2 = %mem.lea (100:.Nat, <100:.Nat; i32>, 0) (a, 1:(.Idx 100));
+                .let (load_mem2, load_val) = %mem.load (i32, 0) (load_mem1, lea2);
+                printInteger ( load_mem2, load_val, print_integer_callback2 )
+            };
+
+            printInteger ( load_mem1, load_val, print_integer_callback )
+        };
+        pb((mem, 1:i32), pb_ret_cont)
+    };
+
+    .let (alloc_mem, alloc_val) = %mem.alloc («100:.Nat; i32», 0) mem;
+
+    .let lea1 = %mem.lea (100:.Nat, <100:.Nat; i32>, 0) (alloc_val, 0:(.Idx 100));
+    .let lea2 = %mem.lea (100:.Nat, <100:.Nat; i32>, 0) (alloc_val, 1:(.Idx 100));
+    .let store_mem = %mem.store (i32, 0) (alloc_mem, lea1, 11:i32);
+    .let store_mem2 = %mem.store (i32, 0) (store_mem, lea2, 42:i32);
+
+    .let f_diff = %autodiff.autodiff (.Cn [[%mem.M, %mem.Ptr («100:.Nat; i32», 0)],.Cn[%mem.M, i32]]) f;
+
+    f_diff ((store_mem2, alloc_val),ret_cont)
+};
+
+
+
+// CHECK-DAG: .cn .extern main _{{[0-9_]+}}::[mem_[[memId:[_0-9]*]]: %mem.M, (i32), %mem.Ptr (%mem.Ptr ((.Idx 256), 0:.Nat), 0:.Nat), return_[[returnId:[_0-9]*]]: .Cn [%mem.M, (i32)]] = {
+// CHECK-DAG: _[[appId:[_0-9]*]]: ⊥:★ = return_[[returnEtaId:[_0-9]*]] (mem_[[memId]], 42:(i32));
+// CHECK-DAG: _[[appId]]
+
+// CHECK-DAG: return_[[returnEtaId]] _[[returnEtaVarId:[0-9_]+]]: [%mem.M, (i32)] = {
+// CHECK-DAG: return_[[retAppId:[_0-9]*]]: ⊥:★ = return_[[returnId]] _[[returnEtaVarId]];
+// CHECK-DAG: return_[[retAppId]]

lit/autodiff/simple_autodiff.thorin

@@ -11,8 +11,6 @@
 .let I32 = .Idx 4294967296;
 
 .cn f [[mem : %mem.M, a:I32], ret: .Cn [%mem.M, I32]] = {
-    // is pack
-    // .let b = %core.wrap.mul (0:.Nat, 4294967296:.Nat) (a, a);
     .let b = %core.wrap.mul (0:.Nat, 4294967296:.Nat) (2:I32, a);
     ret (mem, b)
 };

lit/autodiff/test.thorin

@@ -0,0 +1,49 @@
+// RUN: rm -f %t.ll ; \
+// RUN: %thorin -d autodiff %s --output-ll %t.ll --output-thorin - | FileCheck %s
+
+// a call to a autodiff style function
+// ./build/bin/thorin -d debug -d direct -d autodiff ./lit/autodiff/multiply_autodiff.thorin --output-thorin - -VVVV
+
+.import core;
+.import autodiff;
+.import mem;
+
+.let i32 = .Idx 4294967296;
+
+.cn f [[mem : %mem.M, a: %mem.Ptr («100:.Nat; i32», 0)], ret: .Cn [mem : %mem.M, i32]] = {
+    .let lea = %mem.lea (100:.Nat, <100:.Nat; i32>, 0) (a, 1:(.Idx 100));
+    .let (load_mem, load_val) = %mem.load (i32, 0) (mem, lea);
+    .let scaled = %core.wrap.mul (0:.Nat, 4294967296:.Nat) (100:i32, load_val);
+    ret (load_mem, scaled)
+};
+
+.cn .extern main [mem : %mem.M, argc : i32, argv : %mem.Ptr (%mem.Ptr (.Idx 256, 0:.Nat), 0:.Nat), return : .Cn [%mem.M, i32]] = {
+
+    .cn ret_cont [[mem : %mem.M, a: i32], pb:.Cn[[%mem.M, i32],.Cn[[%mem.M, %mem.Ptr («100:.Nat; i32», 0)]]]] = {
+        .cn pb_ret_cont [mem : %mem.M, a: %mem.Ptr («100:.Nat; i32», 0)] = {
+            .let lea = %mem.lea (100:.Nat, <100:.Nat; i32>, 0) (a, 1:(.Idx 100));
+            .let (load_mem, load_val) = %mem.load (i32, 0) (mem, lea);
+            return (load_mem, load_val)
+        };
+        pb((mem, a), pb_ret_cont)
+    };
+
+    .let (alloc_mem, alloc_val) = %mem.alloc («100:.Nat; i32», 0) mem;
+
+    .let lea = %mem.lea (100:.Nat, <100:.Nat; i32>, 0) (alloc_val, 1:(.Idx 100));
+    .let store = %mem.store (i32, 0) (alloc_mem, lea, argc);
+
+    .let f_diff = %autodiff.autodiff (.Cn [[%mem.M, %mem.Ptr («100:.Nat; i32», 0)],.Cn[%mem.M, i32]]) f;
+
+    f_diff ((store, alloc_val),ret_cont)
+};
+
+
+
+// CHECK-DAG: .cn .extern main _{{[0-9_]+}}::[mem_[[memId:[_0-9]*]]: %mem.M, (i32), %mem.Ptr (%mem.Ptr ((.Idx 256), 0:.Nat), 0:.Nat), return_[[returnId:[_0-9]*]]: .Cn [%mem.M, (i32)]] = {
+// CHECK-DAG: _[[appId:[_0-9]*]]: ⊥:★ = return_[[returnEtaId:[_0-9]*]] (mem_[[memId]], 42:(i32));
+// CHECK-DAG: _[[appId]]
+
+// CHECK-DAG: return_[[returnEtaId]] _[[returnEtaVarId:[0-9_]+]]: [%mem.M, (i32)] = {
+// CHECK-DAG: return_[[retAppId:[_0-9]*]]: ⊥:★ = return_[[returnId]] _[[returnEtaVarId]];
+// CHECK-DAG: return_[[retAppId]]

thorin/CMakeLists.txt

@@ -23,6 +23,8 @@ add_library(libthorin
     tuple.h
     world.cpp
     world.h
+    builder.cpp
+    builder.h
     analyses/cfg.cpp
     analyses/cfg.h
     analyses/deptree.cpp

thorin/dump.cpp

@@ -108,13 +108,13 @@ std::ostream& operator<<(std::ostream& os, Inline u) {
 
     if (auto type = u->isa<Type>()) {
         auto level = as_lit(type->level()); // TODO other levels
-        return print(os, level == 0 ? "★" : "□");
+        return print(os, level == 0 ? "*" : "□");
     } else if (u->isa<Nat>()) {
         return print(os, ".Nat");
     } else if (auto bot = u->isa<Bot>()) {
-        return print(os, "⊥:{}", bot->type());
+        return print(os, ".bot:{}", bot->type());
     } else if (auto top = u->isa<Top>()) {
-        return print(os, "⊤:{}", top->type());
+        return print(os, ".top:{}", top->type());
     } else if (auto axiom = u->isa<Axiom>()) {
         const auto& name = axiom->name();
         return print(os, "{}{}", name[0] == '%' ? "" : "%", name);
@@ -129,8 +129,8 @@ std::ostream& operator<<(std::ostream& os, Inline u) {
     } else if (auto pi = u->isa<Pi>()) {
         if (pi->is_cn()) return print(os, ".Cn {}", pi->dom());
         if (auto nom = pi->isa_nom<Pi>(); nom && nom->var())
-            return print(os, "Π {}: {} → {}", nom->var(), pi->dom(), pi->codom());
-        return print(os, "Π {} → {}", pi->dom(), pi->codom());
+            return print(os, "Π {}: {} -> {}", nom->var(), pi->dom(), pi->codom());
+        return print(os, "Π {} -> {}", pi->dom(), pi->codom());
     } else if (auto lam = u->isa<Lam>()) {
         return print(os, "{}, {}", lam->filter(), lam->body());
     } else if (auto int_ = u->isa<Idx>()) {
@@ -153,12 +153,12 @@ std::ostream& operator<<(std::ostream& os, Inline u) {
         return tuple->type()->isa_nom() ? print(os, ":{}", tuple->type()) : os;
     } else if (auto arr = u->isa<Arr>()) {
         if (auto nom = arr->isa_nom<Arr>(); nom && nom->var())
-            return print(os, "«{}: {}; {}»", nom->var(), nom->shape(), nom->body());
-        return print(os, "«{}; {}»", arr->shape(), arr->body());
+            return print(os, "<<{}: {}; {}>>", nom->var(), nom->shape(), nom->body());
+        return print(os, "<<{}; {}>>", arr->shape(), arr->body());
     } else if (auto pack = u->isa<Pack>()) {
         if (auto nom = pack->isa_nom<Pack>(); nom && nom->var())
-            return print(os, "‹{}: {}; {}›", nom->var(), nom->shape(), nom->body());
-        return print(os, "‹{}; {}›", pack->shape(), pack->body());
+            return print(os, "<{}: {}; {}>", nom->var(), nom->shape(), nom->body());
+        return print(os, "<{}; {}>", pack->shape(), pack->body());
     } else if (auto proxy = u->isa<Proxy>()) {
         return print(os, ".proxy#{}#{} {, }", proxy->pass(), proxy->tag(), proxy->ops());
     } else if (auto bound = isa_bound(*u)) {

thorin/world.h

@@ -14,6 +14,7 @@
 #include "thorin/flags.h"
 #include "thorin/lattice.h"
 #include "thorin/tuple.h"
+#include "thorin/builder.h"
 
 #include "thorin/util/hash.h"
 #include "thorin/util/log.h"
@@ -443,6 +444,10 @@ class World {
     void write() const;                 ///< Same above but file name defaults to World::name.
     ///@}
 
+    Builder builder(){
+        return Builder(*this);
+    }
+
 private:
     /// @name put into sea of nodes
     ///@{