start of real tests

Author: NeuralCoder3

dialects/autodiff/auxiliary/autodiff_aux.cpp

@@ -140,12 +140,13 @@ const Def* autodiff_type_fun(const Def* ty) {
         auto body_ad = autodiff_type_fun(body);
         if (!body_ad) return nullptr;
         return world.arr(shape, body_ad);
-    }
-    if (auto sig = ty->isa<Sigma>()) {
+    } else if (auto sig = ty->isa<Sigma>()) {
         // TODO: nom sigma
         DefArray ops(sig->ops(), [&](const Def* op) { return autodiff_type_fun(op); });
         world.DLOG("ops: {,}", ops);
         return world.sigma(ops);
+    } else if (auto real = match<core::Real>(ty)) {
+        return ty;
     }
     // mem
     if (match<mem::M>(ty)) return ty;
@@ -171,6 +172,11 @@ const Def* zero_def(const Def* T) {
         auto zero = world.lit_idx(T, 0, world.dbg("zero"));
         world.DLOG("zero_def for int is {}", zero);
         return zero;
+    } else if (auto real = match<core::Real>(T)) {
+        auto width = as_lit<nat_t>(real->arg());
+        auto zero  = core::lit_real(T->world(), width, 0.0);
+        world.DLOG("zero_def for real is {}", zero);
+        return zero;
     } else if (auto sig = T->isa<Sigma>()) {
         DefArray ops(sig->ops(), [&](const Def* op) { return world.app(world.ax<zero>(), op); });
         return world.tuple(ops);

dialects/autodiff/normalizers.cpp

@@ -84,6 +84,14 @@ const Def* normalize_add(const Def* type, const Def* callee, const Def* arg, con
             world.app(world.app(world.ax(core::wrap::add), {world.lit_nat_0(), world.lit_nat(width)}), {a, b});
         world.DLOG("int add {} : {}", int_add, world.iinfer(int_add));
         return int_add;
+    } else if (auto real = match<core::Real>(T)) {
+        auto width = as_lit<nat_t>(real->arg());
+        world.DLOG("width {}", width);
+        auto real_add =
+            world.app(world.app(world.ax(core::rop::add), {world.lit_nat_0(), world.lit_nat(width)}), {a, b});
+        world.DLOG("real add {} : {}", real_add, real_add->type());
+        return real_add;
+
     } else if (auto app = T->isa<App>()) {
         auto callee = app->callee();
         assert(0 && "not handled");

lit/autodiff/square_real.thorin

@@ -0,0 +1,67 @@
+// RUN: rm -f %t.ll ; \
+// RUN: %thorin -d tool -d direct -d autodiff %s --output-ll %t.ll --output-thorin - | FileCheck %s
+
+.import core;
+.import autodiff;
+.import mem;
+
+.let _32 = 32;
+.let R32 = %core.Real _32;
+
+
+.lam .extern internal_diff_core_rop_mul
+    ![m:.Nat, w:.Nat] -> 
+    (.Cn[[%core.Real w, %core.Real w], .Cn[%core.Real w, .Cn[%core.Real w, .Cn[%core.Real w, %core.Real w]]]])
+    = { 
+    .cn inner_mul_deriv_cps
+        ![[a:%core.Real w, b:%core.Real w], ret:.Cn[%core.Real w, .Cn[%core.Real w, .Cn[%core.Real w, %core.Real w]]]]@(.tt)
+        = { 
+        .let result = %core.rop.mul (m,w) (a,b);
+        .cn mul_pb ![s:(%core.Real w), pb_ret:(.Cn [%core.Real w, %core.Real w])]@(.tt) = { 
+            .let lhs = %core.rop.mul (m,w) (s,b);
+            .let rhs = %core.rop.mul (m,w) (s,a);
+            pb_ret (lhs, rhs)
+        };
+        ret (result,mul_pb)
+    };
+    inner_mul_deriv_cps
+};
+
+
+
+
+.cn f [a:R32, ret: .Cn [R32]] = {
+    .let b = %core.rop.mul (0, _32) (a, a);
+    ret b
+};
+
+// .cn g [a:R32, ret: .Cn [R32, .Cn [R32, .Cn [R32]]]] = {
+//     .let f_diff = %autodiff.autodiff (.Cn [R32,.Cn[R32]]) f;
+//     f_diff (a,ret)
+// };
+
+.cn .extern test [[], ret: .Cn [R32]] = {
+    // .cn inner [r:R32, pb:.Cn [R32, .Cn [R32]]] = {
+    //     .cn pb_rec p::[R32] = {
+    //         // ret (5.0:R32)
+    //         ret r
+    //     };
+    //     pb (1.0:R32, pb_rec)
+    // };
+    // g (5.0:R32, inner)
+
+
+    //     .cn pb_rec p::[R32] = {
+    //         ret p
+    //     };
+
+    // f (5.0:R32, pb_rec)
+
+
+    // ret (6.0:R32)
+
+    .let b = %core.rop.mul (0, _32) (2.0:R32, 3.0:R32);
+    ret b
+};
+
+// CHECK-DAG: return{{.*}}176484