Matrix dialect

dialects/CMakeLists.txt

@@ -23,13 +23,12 @@ add_thorin_plugin(autodiff
         autodiff/passes/autodiff_zero.h
         autodiff/passes/autodiff_zero_cleanup.cpp
         autodiff/passes/autodiff_zero_cleanup.h
-        autodiff/passes/autodiff_ext_cleanup.cpp
-        autodiff/passes/autodiff_ext_cleanup.h
         autodiff/auxiliary/autodiff_aux.cpp
         autodiff/auxiliary/autodiff_aux.h
         autodiff/auxiliary/autodiff_rewrite_inner.cpp
         autodiff/auxiliary/autodiff_rewrite_toplevel.cpp
         autodiff/normalizers.cpp
+        compile/passes/internal_cleanup.cpp
     DEPENDS
         mem
         core
@@ -127,6 +126,27 @@ add_thorin_plugin(math
     INSTALL
 )
 
+add_thorin_plugin(matrix
+    SOURCES
+        matrix/matrix.cpp
+        matrix/matrix.h
+        matrix/normalizers.cpp
+        matrix/passes/lower_matrix_highlevel.cpp
+        matrix/passes/lower_matrix_highlevel.h
+        matrix/passes/lower_matrix_mediumlevel.cpp
+        matrix/passes/lower_matrix_mediumlevel.h
+        matrix/passes/lower_matrix_lowlevel.cpp
+        matrix/passes/lower_matrix_lowlevel.h
+        compile/passes/internal_cleanup.cpp
+    DEPENDS
+        refly
+        direct
+        affine
+        core
+        mem
+        compile
+    INSTALL
+)
 add_thorin_plugin(mem
     SOURCES
         mem/mem.cpp

dialects/affine/affine.h

@@ -14,9 +14,16 @@ inline const Def* fn_for(World& w, Defs params) {
 
 /// Returns a fully applied affine_for axiom.
 /// See documentation for %affine.For axiom in @ref affine.
-inline const Def*
-op_for(World& w, const Def* begin, const Def* end, const Def* step, Defs inits, const Def* body, const Def* brk) {
+// clang-format off
+inline const Def* op_for(World& w,
+                         Ref begin,
+                         Ref end,
+                         Ref step,
+                         Defs inits,
+                         Ref body,
+                         Ref brk) {
     DefArray types(inits.size(), [&](size_t i) { return inits[i]->type(); });
     return w.app(fn_for(w, types), {begin, end, step, w.tuple(inits), body, brk});
 }
+// clang-format on
 } // namespace thorin::affine

dialects/autodiff/autodiff.cpp

@@ -5,9 +5,9 @@
 #include <thorin/pass/pipelinebuilder.h>
 
 #include "dialects/autodiff/passes/autodiff_eval.h"
-#include "dialects/autodiff/passes/autodiff_ext_cleanup.h"
 #include "dialects/autodiff/passes/autodiff_zero.h"
 #include "dialects/autodiff/passes/autodiff_zero_cleanup.h"
+#include "dialects/compile/passes/internal_cleanup.h"
 #include "dialects/direct/passes/ds2cps.h"
 
 using namespace thorin;
@@ -18,7 +18,7 @@ extern "C" THORIN_EXPORT Plugin thorin_get_plugin() {
                 register_pass<autodiff::ad_eval_pass, autodiff::AutoDiffEval>(passes);
                 register_pass<autodiff::ad_zero_pass, autodiff::AutoDiffZero>(passes);
                 register_pass<autodiff::ad_zero_cleanup_pass, autodiff::AutoDiffZeroCleanup>(passes);
-                register_pass<autodiff::ad_ext_cleanup_pass, autodiff::AutoDiffExternalCleanup>(passes);
+                register_pass<autodiff::ad_ext_cleanup_pass, compile::InternalCleanup>(passes, "internal_diff_");
             },
             nullptr, [](Normalizers& normalizers) { autodiff::register_normalizers(normalizers); }};
 }

dialects/autodiff/rules.rs

@@ -493,7 +493,7 @@ zip*:
 
 zip*:
     λ S. split(
-        mapReduce
+        map_reduce
             (λ (m,n,s).
                 let (_, f*) = f' (m,n);
                 f* s

dialects/clos/clos.cpp

@@ -139,11 +139,11 @@ extern "C" THORIN_EXPORT Plugin thorin_get_plugin() {
     return {"clos",
             [](Passes& passes) {
                 register_pass<clos::clos_conv_prep_pass, clos::ClosConvPrep>(passes, nullptr);
-                register_pass<clos::clos_conv_pass, ClosConvWrapper>(passes);
+                register_pass<clos::clos_conv_pass, clos::ClosConvWrapper>(passes);
                 register_pass<clos::branch_clos_pass, clos::BranchClosElim>(passes);
                 register_pass<clos::lower_typed_clos_prep_pass, clos::LowerTypedClosPrep>(passes);
                 register_pass<clos::clos2sjlj_pass, clos::Clos2SJLJ>(passes);
-                register_pass<clos::lower_typed_clos_pass, LowerTypedClosWrapper>(passes);
+                register_pass<clos::lower_typed_clos_pass, clos::LowerTypedClosWrapper>(passes);
                 // TODO:; remove after ho_codegen merge
                 passes[flags_t(Axiom::Base<clos::eta_red_bool_pass>)] = [&](World&, PipelineBuilder& builder, Ref app) {
                     auto bb      = app->as<App>()->arg();

dialects/clos/pass/rw/phase_wrapper.h

@@ -6,7 +6,7 @@
 #include "dialects/clos/phase/clos_conv.h"
 #include "dialects/clos/phase/lower_typed_clos.h"
 
-using namespace thorin;
+namespace thorin::clos {
 
 class ClosConvWrapper : public RWPass<ClosConvWrapper, Lam> {
 public:
@@ -23,3 +23,5 @@ class LowerTypedClosWrapper : public RWPass<LowerTypedClosWrapper, Lam> {
 
     void prepare() override { clos::LowerTypedClos(world()).run(); }
 };
+
+}

dialects/compile/compile.thorin

@@ -44,6 +44,7 @@
 .ax %compile.clos_plugin    : %compile.Plugin;
 .ax %compile.direct_plugin  : %compile.Plugin;
 .ax %compile.refly_plugin   : %compile.Plugin;
+.ax %compile.matrix_plugin  : %compile.Plugin;
 ///
 /// ### %opt.is_loaded
 /// 

dialects/core/be/ll/ll.cpp

@@ -192,7 +192,7 @@ std::string Emitter::convert(const Def* type) {
 
 std::string Emitter::convert_ret_pi(const Pi* pi) {
     auto dom = mem::strip_mem_ty(pi->dom());
-    if (dom == world().sigma()) { return "void"; }
+    if (dom == world().sigma()) return "void";
     return convert(dom);
 }
 
@@ -261,9 +261,8 @@ void Emitter::finalize(const Scope& scope) {
             print(func_impls_, "{}:\n", lam->unique_name());
 
             ++tab;
-            for (const auto& part : bb.parts) {
+            for (const auto& part : bb.parts)
                 for (const auto& line : part) tab.print(func_impls_, "{}\n", line.str());
-            }
             --tab;
             func_impls_ << std::endl;
         }
@@ -305,7 +304,26 @@ void Emitter::emit_epilogue(Lam* lam) {
             }
         }
     } else if (auto ex = app->callee()->isa<Extract>(); ex && app->callee_type()->is_basicblock()) {
-        emit_unsafe(app->arg());
+        // emit_unsafe(app->arg());
+        // A call to an extract like constructed for conditionals (else,then)#cond (args)
+        // TODO: we can not rely on the structure of the extract (it might be a nested extract)
+        for (auto callee_def : ex->tuple()->projs()) {
+            // dissect the tuple of lambdas
+            auto callee = callee_def->as_nom<Lam>();
+            // each callees type should agree with the argument type (should be checked by type checking).
+            // Especially, the number of vars should be the number of arguments.
+            // TODO: does not hold for complex arguments that are not tuples.
+            assert(callee->num_vars() == app->num_args());
+            for (size_t i = 0, e = callee->num_vars(); i != e; ++i) {
+                // emits the arguments one by one (TODO: handle together like before)
+                if (auto arg = emit_unsafe(app->arg(i)); !arg.empty()) {
+                    auto phi = callee->var(i);
+                    assert(!match<mem::M>(phi->type()));
+                    lam2bb_[callee].phis[phi].emplace_back(arg, id(lam, true));
+                    locals_[phi] = id(phi);
+                }
+            }
+        }
 
         auto c = emit(ex->index());
         if (ex->tuple()->num_projs() == 2) {
@@ -344,9 +362,8 @@ void Emitter::emit_epilogue(Lam* lam) {
 
         std::vector<std::string> args;
         auto app_args = app->args();
-        for (auto arg : app_args.skip_back()) {
+        for (auto arg : app_args.skip_back())
             if (auto v_arg = emit_unsafe(arg); !v_arg.empty()) args.emplace_back(convert(arg->type()) + " " + v_arg);
-        }
 
         if (app->args().back()->isa<Bot>()) {
             // TODO: Perhaps it'd be better to simply η-wrap this prior to the BE...
@@ -725,8 +742,6 @@ std::string Emitter::emit_bb(BB& bb, const Def* def) {
         auto [v_i, t_i] = emit_gep_index(i);
 
         return bb.assign(name, "getelementptr inbounds {}, {} {}, i64 0, {} {}", t_pointee, t_ptr, v_ptr, t_i, v_i);
-    } else if (match<core::trait>(def)) {
-        unreachable();
     } else if (auto malloc = match<mem::malloc>(def)) {
         declare("i8* @malloc(i64)");
 

dialects/direct/passes/cps2ds.cpp

@@ -46,6 +46,7 @@ void CPS2DS::rewrite_lam(Lam* lam) {
 }
 
 const Def* CPS2DS::rewrite_body(const Def* def) {
+    if (!def) return nullptr;
     if (auto i = rewritten_.find(def); i != rewritten_.end()) return i->second;
     auto new_def    = rewrite_body_(def);
     rewritten_[def] = new_def;

dialects/matrix/matrix.cpp

@@ -0,0 +1,26 @@
+
+#include "dialects/matrix/matrix.h"
+
+#include <thorin/plugin.h>
+#include <thorin/pass/pass.h>
+
+#include "dialects/compile/passes/internal_cleanup.h"
+#include "dialects/matrix/passes/lower_matrix_highlevel.h"
+#include "dialects/matrix/passes/lower_matrix_lowlevel.h"
+#include "dialects/matrix/passes/lower_matrix_mediumlevel.h"
+#include "dialects/refly/refly.h"
+
+using namespace thorin;
+
+extern "C" THORIN_EXPORT Plugin thorin_get_plugin() {
+    return {"matrix",
+            [](Passes& passes) {
+                register_pass<matrix::lower_matrix_high_level_map_reduce, thorin::matrix::LowerMatrixHighLevelMapRed>(
+                    passes);
+                register_pass<matrix::lower_matrix_medium_level, thorin::matrix::LowerMatrixMediumLevel>(passes);
+                register_phase<matrix::lower_matrix_low_level, thorin::matrix::LowerMatrixLowLevel>(passes);
+                register_pass<matrix::internal_map_reduce_cleanup, thorin::compile::InternalCleanup>(passes,
+                                                                                                     INTERNAL_PREFIX);
+            },
+            nullptr, [](Normalizers& normalizers) { matrix::register_normalizers(normalizers); }};
+}

dialects/matrix/matrix.h

@@ -0,0 +1,31 @@
+#pragma once
+
+#include <thorin/pass/pipelinebuilder.h>
+#include <thorin/world.h>
+
+#include "dialects/matrix/autogen.h"
+#include "dialects/mem/mem.h"
+
+namespace thorin::matrix {
+
+#define INTERNAL_PREFIX "internal_mapRed_"
+
+/// %mat.zero: Π [n: .Nat, S: «n; .Nat», m: .Nat] -> %mat.Mat (n,S,(.Idx m));
+inline const Def* zero_int(World& w, Ref n, Ref S, Ref mem, nat_t m) {
+    // TODO: use thorin definition by name
+    return w.app(w.ax<matrix::constMat>(), {n, S, w.type_idx(m), mem, w.lit_idx(m, 0)});
+}
+
+inline const Def* op_read(Ref mem, Ref matrix, Ref idx) {
+    auto& world = matrix->world();
+    auto mat_ty = match<Mat>(matrix->type());
+    if (!mat_ty) return matrix;
+    assert(mat_ty);
+    world.DLOG("matrix read: {}[{}]", matrix, idx);
+    world.DLOG(" matrix type: {}", matrix->type());
+    auto [n, S, T] = mat_ty->args<3>();
+    world.DLOG(" (n,S,T): {}, {}, {}", n, S, T);
+    return world.app(world.app(world.ax<read>(), {n, S, T}), {mem, matrix, idx});
+}
+
+} // namespace thorin::matrix