From 1b3718bd12a48775724446dd1475d0d1d1d33d65 Mon Sep 17 00:00:00 2001
From: Roland Conybeare <rconybeare@gmail.com>
Date: Sat, 15 Jun 2024 17:13:59 -0400
Subject: [PATCH] xo-jit: + machgen_current_module(): generate machine code!

---
 include/xo/jit/Jit.hpp | 22 +++++++++++++++++++++-
 src/jit/Jit.cpp        | 32 +++++++++++++++++++++++++++-----
 2 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/include/xo/jit/Jit.hpp b/include/xo/jit/Jit.hpp
index 3a8e4c14..14b368d3 100644
--- a/include/xo/jit/Jit.hpp
+++ b/include/xo/jit/Jit.hpp
@@ -98,6 +98,13 @@ namespace xo {
 
             llvm::Value * codegen(ref::brw<Expression> expr);
 
+            // ----- jit online execution -----
+
+            /** add IR code in current module to JIT,
+             *  so that its available for execution
+             **/
+            void machgen_current_module();
+
             llvm::orc::ExecutorAddr lookup_symbol(const std::string & x);
 
             virtual void display(std::ostream & os) const;
@@ -116,6 +123,9 @@ namespace xo {
             /* iniitialize native builder (i.e. for platform we're running on) */
             static void init_once();
 
+            /** (re)create pipeline to turn expressions into llvm IR code **/
+            void recreate_llvm_ir_pipeline();
+
         private:
             // ----- this part adapted from LLVM 19.0 KaleidoscopeJIT.hpp [wip] -----
 
@@ -133,6 +143,14 @@ namespace xo {
 
             // ----- this part adapted from kaleidoscope.cpp -----
 
+            /** everything bleow represents a pipeline
+             *  that takes expressions, and turns them into llvm IR.
+             *
+             *  llvm IR can be added to running JIT by calling
+             *    kal_jit_.addModule()
+             *  Note that this makes the module itself unavailable to us
+             **/
+
             /** owns + manages core "global" llvm data,
              *  including type- and constant- unique-ing tables.
              *
@@ -142,7 +160,9 @@ namespace xo {
             std::unique_ptr<llvm::LLVMContext> llvm_cx_;
             /** builder for intermediate-representation objects **/
             std::unique_ptr<llvm::IRBuilder<>> llvm_ir_builder_;
-            /** a module (aka library) being prepared by llvm.
+            /** a module (1:1 with library) being prepared by llvm.
+             *  IR-level -- does not contain machine code
+             *
              *  - function names are unique within a module.
              **/
             std::unique_ptr<llvm::Module> llvm_module_;
diff --git a/src/jit/Jit.cpp b/src/jit/Jit.cpp
index 7d53a313..df517969 100644
--- a/src/jit/Jit.cpp
+++ b/src/jit/Jit.cpp
@@ -98,7 +98,7 @@ namespace xo {
                  llvm::DataLayout dl
 #endif
             )
-            : kal_jit_{std::move(kal_jit)},
+            : kal_jit_{std::move(kal_jit)}
 #ifdef NOT_USING
             jit_es_(std::move(jit_es)),
               jit_data_layout_(std::move(dl)),
@@ -114,9 +114,6 @@ namespace xo {
               jit_our_dynamic_lib_(this->jit_es_->createBareJITDylib("<xojitlib>")), /*was MainJD*/
 #endif
 
-              llvm_cx_{std::make_unique<llvm::LLVMContext>()},
-              llvm_ir_builder_{std::make_unique<llvm::IRBuilder<>>(*llvm_cx_)},
-              llvm_module_{std::make_unique<llvm::Module>("xojit", *llvm_cx_)}
         {
 #ifdef NOT_USING
             jit_our_dynamic_lib_.addGenerator
@@ -129,6 +126,16 @@ namespace xo {
             }
 #endif
 
+            this->recreate_llvm_ir_pipeline();
+        }
+
+        void
+        Jit::recreate_llvm_ir_pipeline()
+        {
+            llvm_cx_ = std::make_unique<llvm::LLVMContext>();
+            llvm_ir_builder_ = std::make_unique<llvm::IRBuilder<>>(*llvm_cx_);
+            llvm_module_ = std::make_unique<llvm::Module>("xojit", *llvm_cx_);
+
             llvm_module_->setDataLayout(kal_jit_->getDataLayout());
 
             if (!llvm_cx_.get()) {
@@ -162,7 +169,7 @@ namespace xo {
             llvm_pass_builder.registerModuleAnalyses(*llvm_mamgr_);
             llvm_pass_builder.registerFunctionAnalyses(*llvm_famgr_);
             llvm_pass_builder.crossRegisterProxies(*llvm_lamgr_, *llvm_famgr_, *llvm_cgamgr_, *llvm_mamgr_);
-        }
+        } /*recreate_llvm_ir_pipeline*/
 
         const std::string &
         Jit::target_triple() const {
@@ -465,6 +472,21 @@ namespace xo {
             return nullptr;
         } /*codegen*/
 
+        void
+        Jit::machgen_current_module()
+        {
+            static llvm::ExitOnError llvm_exit_on_err;
+
+            auto tracker = kal_jit_->getMainJITDylib().createResourceTracker();
+
+            auto ts_module = llvm::orc::ThreadSafeModule(std::move(llvm_module_),
+                                                         std::move(llvm_cx_));
+
+            llvm_exit_on_err(kal_jit_->addModule(std::move(ts_module), tracker));
+
+            this->recreate_llvm_ir_pipeline();
+        }
+
         llvm::orc::ExecutorAddr
         Jit::lookup_symbol(const std::string & sym)
         {