Thorin Inline Optimization

cli/main.cpp

@@ -98,16 +98,18 @@ int main(int argc, char** argv) {
         }
 
         // we always need core and mem, as long as we are not in bootstrap mode..
-        if (!os[H]) dialect_plugins.insert(dialect_plugins.end(), {"core", "mem"});
+        if (!os[H]) dialect_plugins.insert(dialect_plugins.end(), {"core", "mem", "compile"});
 
         std::vector<Dialect> dialects;
         thorin::Backends backends;
         thorin::Normalizers normalizers;
+        thorin::Passes passes;
         if (!dialect_plugins.empty()) {
             for (const auto& dialect : dialect_plugins) {
                 dialects.push_back(Dialect::load(dialect, dialect_paths));
                 dialects.back().register_backends(backends);
                 dialects.back().register_normalizers(normalizers);
+                dialects.back().register_passes(passes);
             }
         }
 
@@ -130,13 +132,13 @@ int main(int argc, char** argv) {
         if (os[H]) parser.bootstrap(*os[H]);
 
         PipelineBuilder builder;
-        for (const auto& dialect : dialects) { dialect.register_passes(builder); }
+        for (const auto& dialect : dialects) { dialect.add_passes(builder); }
 
         // clang-format off
         switch (opt) {
             case 0:                             break;
             case 1: Phase::run<Cleanup>(world); break;
-            case 2: optimize(world, builder);   break;
+            case 2: optimize(world, passes, builder);   break;
             default: errln("error: illegal optimization level '{}'", opt);
         }
         // clang-format on

dialects/CMakeLists.txt

@@ -37,20 +37,6 @@ add_thorin_dialect(autodiff
     INSTALL
 )
 
-add_thorin_dialect(core
-    SOURCES
-        core/core.cpp
-        core/core.h
-        core/normalizers.cpp
-        core/be/ll/ll.cpp
-        core/be/ll/ll.h
-    DEPENDS
-        math
-        mem
-        clos
-    INSTALL
-)
-
 add_thorin_dialect(clos
     SOURCES
         clos/clos.cpp
@@ -77,6 +63,30 @@ add_thorin_dialect(clos
     INSTALL
 )
 
+add_thorin_dialect(compile
+    SOURCES
+        compile/compile.cpp
+        compile/compile.h
+        compile/normalizers.cpp
+        compile/passes/debug_print.cpp
+        compile/passes/debug_print.h
+    INSTALL
+)
+
+add_thorin_dialect(core
+    SOURCES
+        core/core.cpp
+        core/core.h
+        core/normalizers.cpp
+        core/be/ll/ll.cpp
+        core/be/ll/ll.h
+    DEPENDS
+        math
+        mem
+        clos
+    INSTALL
+)
+
 add_thorin_dialect(demo
     SOURCES
         demo/demo.cpp
@@ -118,6 +128,8 @@ add_thorin_dialect(mem
         mem/passes/rw/alloc2malloc.h
         mem/passes/rw/remem_elim.cpp
         mem/passes/rw/remem_elim.h
+    DEPENDS
+        compile
     HEADER_DEPENDS
         core
     INSTALL

dialects/affine/affine.cpp

@@ -5,12 +5,16 @@
 
 #include "thorin/dialects.h"
 
+// #include "thorin/pass/pipelinebuilder.h"
+
 #include "dialects/affine/passes/lower_for.h"
 
+using namespace thorin;
+
 extern "C" THORIN_EXPORT thorin::DialectInfo thorin_get_dialect_info() {
     return {"affine",
             [](thorin::PipelineBuilder& builder) {
                 builder.extend_opt_phase([](thorin::PassMan& man) { man.add<thorin::affine::LowerFor>(); });
             },
-            nullptr, nullptr};
+            [](Passes& passes) { register_pass<affine::lower_for_pass, affine::LowerFor>(passes); }, nullptr, nullptr};
 }

dialects/affine/affine.thorin

@@ -22,3 +22,11 @@
     .Cn [start: .Idx m, stop: .Idx m, step: .Idx m, init: «i: n; Ts#i»,
         body: .Cn [iter: .Idx m, acc: «i: n; Ts#i», yield: .Cn [«i: n; Ts#i»]], 
         exit: .Cn [«i: n; Ts#i»]];
+///
+/// ## Compilation Passes and Phases
+///
+/// ### %affine.lower_for_pass
+/// 
+/// Loweres the %affine.For operation to recursive function calls.
+/// 
+.ax %affine.lower_for_pass: %compile.Pass;

dialects/autodiff/autodiff.cpp

@@ -35,5 +35,11 @@ extern "C" THORIN_EXPORT thorin::DialectInfo thorin_get_dialect_info() {
                     man.add<thorin::autodiff::AutoDiffExternalCleanup>();
                 });
             },
+            [](Passes& passes) {
+                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);
+            },
             nullptr, [](Normalizers& normalizers) { autodiff::register_normalizers(normalizers); }};
 }

dialects/autodiff/autodiff.thorin

@@ -5,6 +5,7 @@
 /// ## Dependencies
 ///
 .import mem;
+.import compile;
 /// For derivatives:
 .import core;
 ///
@@ -42,6 +43,34 @@
 .ax %autodiff.add: Π [T:*] -> [T,T] -> T, normalize_add;
 .ax %autodiff.sum: Π [n:.Nat,T:*] -> «n; T» -> T, normalize_sum;
 ///
+/// ## Compilation Passes and Phases
+/// 
+/// ### Passes
+/// 
+.ax %autodiff.ad_eval_pass: %compile.Pass;
+.ax %autodiff.ad_zero_pass: %compile.Pass;
+.ax %autodiff.ad_zero_cleanup_pass: %compile.Pass;
+.ax %autodiff.ad_ext_cleanup_pass: %compile.Pass;
+/// 
+/// ### Phases
+/// 
+.let ad_cleanup_phase = {
+    %compile.phases_to_phase (⊤:.Nat)
+        (
+            (%compile.passes_to_phase 1 %autodiff.ad_zero_cleanup_pass),
+            (%compile.passes_to_phase 1 %autodiff.ad_ext_cleanup_pass)
+        )
+};
+.let ad_phase = {
+    %compile.phases_to_phase (⊤:.Nat)
+        (
+            optimization_phase,
+            (%compile.passes_to_phase 1 %autodiff.ad_eval_pass),
+            (%compile.passes_to_phase 1 %autodiff.ad_zero_pass),
+            ad_cleanup_phase
+        )
+};
+/// 
 /// ## Registered translations
 /// 
 /// In this section, we define translations for axioms of other dialects.

dialects/clos/clos.cpp

@@ -179,5 +179,5 @@ extern "C" THORIN_EXPORT DialectInfo thorin_get_dialect_info() {
 
                 builder.extend_opt_phase(base++, [](PassMan& man) { man.add<clos::LowerTypedClosWrapper>(); });
             },
-            nullptr, [](Normalizers& normalizers) { clos::register_normalizers(normalizers); }};
+            nullptr, nullptr, [](Normalizers& normalizers) { clos::register_normalizers(normalizers); }};
 }

dialects/compile/compile.cpp

@@ -0,0 +1,128 @@
+#include "dialects/compile/compile.h"
+
+#include <thorin/config.h>
+#include <thorin/dialects.h>
+#include <thorin/pass/pass.h>
+
+#include "thorin/pass/fp/beta_red.h"
+#include "thorin/pass/fp/eta_exp.h"
+#include "thorin/pass/fp/eta_red.h"
+#include "thorin/pass/fp/tail_rec_elim.h"
+#include "thorin/pass/pipelinebuilder.h"
+#include "thorin/pass/rw/lam_spec.h"
+#include "thorin/pass/rw/partial_eval.h"
+#include "thorin/pass/rw/ret_wrap.h"
+#include "thorin/pass/rw/scalarize.h"
+
+#include "dialects/compile/autogen.h"
+#include "dialects/compile/passes/debug_print.h"
+
+using namespace thorin;
+
+void addPhases(DefVec& phases, World& world, Passes& passes, PipelineBuilder& builder) {
+    for (auto phase : phases) {
+        auto [phase_def, phase_args] = collect_args(phase);
+        world.DLOG("phase: {}", phase_def);
+        if (auto phase_ax = phase_def->isa<Axiom>()) {
+            auto flag = phase_ax->flags();
+            if (passes.contains(flag)) {
+                auto phase_fun = passes[flag];
+                phase_fun(world, builder, phase);
+            } else {
+                world.WLOG("phase '{}' not found", phase_ax->name());
+            }
+        } else {
+            world.WLOG("phase '{}' is not an axiom", phase_def);
+        }
+    }
+}
+
+void addPasses(World& world, PipelineBuilder& builder, Passes& passes, DefVec& pass_list) {
+    // Concept: We create a convention that passes register in the pipeline using append_**in**_last.
+    // This pass then calls the registered passes in the order they were registered in the last phase.
+
+    // We create a new dummy phase in which the passes should be inserted.
+    builder.append_phase_end([](Pipeline&) {});
+
+    for (auto pass : pass_list) {
+        auto [pass_def, pass_args] = collect_args(pass);
+        world.DLOG("pass: {}", pass_def);
+        if (auto pass_ax = pass_def->isa<Axiom>()) {
+            auto flag = pass_ax->flags();
+            if (passes.contains(flag)) {
+                auto pass_fun = passes[flag];
+                pass_fun(world, builder, pass);
+            } else {
+                world.WLOG("pass '{}' not found", pass_ax->name());
+            }
+        } else {
+            world.WLOG("pass '{}' is not an axiom", pass_def);
+        }
+    }
+}
+
+extern "C" THORIN_EXPORT thorin::DialectInfo thorin_get_dialect_info() {
+    return {"compile", nullptr,
+            [](Passes& passes) {
+                auto debug_phase_flag    = flags_t(Axiom::Base<thorin::compile::debug_phase>);
+                passes[debug_phase_flag] = [](World& world, PipelineBuilder& builder, const Def* app) {
+                    world.DLOG("Generate debug_phase: {}", app);
+                    int level = (int)(app->as<App>()->arg(0)->as<Lit>()->get<u64>());
+                    world.DLOG("  Level: {}", level);
+                    // TODO: add a debug pass to the pipeline
+                    builder.append_pass_after_end([=](PassMan& man) {
+                        man.add<thorin::compile::DebugPrint>(level);
+                        // man.add<thorin::compile::DebugPrint>(level);
+                        // man.add<thorin::compile::DebugPrint>(42);
+                    });
+                };
+
+                auto pass_phase_flag    = flags_t(Axiom::Base<thorin::compile::pass_phase>);
+                passes[pass_phase_flag] = [&](World& world, PipelineBuilder& builder, const Def* app) {
+                    auto [ax, pass_list] = collect_args(app->as<App>()->arg());
+                    addPasses(world, builder, passes, pass_list);
+                };
+
+                auto combined_phase_flag    = flags_t(Axiom::Base<thorin::compile::combined_phase>);
+                passes[combined_phase_flag] = [&](World& world, PipelineBuilder& builder, const Def* app) {
+                    auto [ax, phase_list] = collect_args(app->as<App>()->arg());
+                    addPhases(phase_list, world, passes, builder);
+                };
+
+                passes[flags_t(Axiom::Base<thorin::compile::passes_to_phase>)] =
+                    [&](World& world, PipelineBuilder& builder, const Def* app) {
+                        auto pass_array = app->as<App>()->arg()->projs();
+                        DefVec pass_list;
+                        for (auto pass : pass_array) { pass_list.push_back(pass); }
+                        addPasses(world, builder, passes, pass_list);
+                    };
+
+                passes[flags_t(Axiom::Base<thorin::compile::phases_to_phase>)] =
+                    [&](World& world, PipelineBuilder& builder, const Def* app) {
+                        auto phase_array = app->as<App>()->arg()->projs();
+                        DefVec phase_list;
+                        for (auto phase : phase_array) { phase_list.push_back(phase); }
+                        addPhases(phase_list, world, passes, builder);
+                    };
+
+                passes[flags_t(Axiom::Base<thorin::compile::pipe>)] = [&](World& world, PipelineBuilder& builder,
+                                                                          const Def* app) {
+                    auto [ax, phases] = collect_args(app);
+                    addPhases(phases, world, passes, builder);
+                };
+
+                register_pass<compile::partial_eval_pass, PartialEval>(passes);
+                register_pass<compile::beta_red_pass, BetaRed>(passes);
+                register_pass<compile::eta_red_pass, EtaRed>(passes);
+
+                register_pass<compile::scalerize_no_arg_pass, Scalerize>(passes);
+                register_pass<compile::tail_rec_elim_no_arg_pass, TailRecElim>(passes);
+                register_pass<compile::lam_spec_pass, LamSpec>(passes);
+                register_pass<compile::ret_wrap_pass, RetWrap>(passes);
+
+                register_pass_with_arg<compile::eta_exp_pass, EtaExp, EtaRed>(passes);
+                register_pass_with_arg<compile::scalerize_pass, Scalerize, EtaExp>(passes);
+                register_pass_with_arg<compile::tail_rec_elim_pass, TailRecElim, EtaRed>(passes);
+            },
+            nullptr, [](Normalizers& normalizers) { compile::register_normalizers(normalizers); }};
+}

dialects/compile/compile.h

@@ -0,0 +1,7 @@
+#pragma once
+
+#include <thorin/world.h>
+
+#include "dialects/compile/autogen.h"
+
+namespace thorin::compile {} // namespace thorin::compile