From 300d7c27841dc199b8b2ce708ae535f561afc2de Mon Sep 17 00:00:00 2001
From: Prem Chintalapudi <prem.chintalapudi@gmail.com>
Date: Mon, 18 Apr 2022 21:50:48 -0400
Subject: [PATCH] Fix aarch64 build and remove broken dump_method_asm_impl

---
 src/jitlayers.cpp | 38 +++++++++++++++++++-------------------
 src/jitlayers.h   | 19 +++++++++++++------
 2 files changed, 32 insertions(+), 25 deletions(-)

diff --git a/src/jitlayers.cpp b/src/jitlayers.cpp
index 6256590ce6624..5ac6f93a3b53c 100644
--- a/src/jitlayers.cpp
+++ b/src/jitlayers.cpp
@@ -413,6 +413,7 @@ extern "C" JL_DLLEXPORT
 jl_value_t *jl_dump_method_asm_impl(jl_method_instance_t *mi, size_t world,
         char raw_mc, char getwrapper, const char* asm_variant, const char *debuginfo, char binary)
 {
+#ifndef JL_COMPILE_ON_DEMAND
     // printing via disassembly
     jl_code_instance_t *codeinst = jl_generate_fptr(mi, world);
     if (codeinst) {
@@ -461,6 +462,7 @@ jl_value_t *jl_dump_method_asm_impl(jl_method_instance_t *mi, size_t world,
         if (specfptr != 0)
             return jl_dump_fptr_asm(specfptr, raw_mc, asm_variant, debuginfo, binary);
     }
+#endif
 
     // whatever, that didn't work - use the assembler output instead
     void *F = jl_get_llvmf_defn(mi, world, getwrapper, true, jl_default_cgparams);
@@ -678,21 +680,18 @@ RTDyldMemoryManager* createRTDyldMemoryManager(void);
 // A simple forwarding class, since OrcJIT v2 needs a unique_ptr, while we have a shared_ptr
 class ForwardingMemoryManager : public RuntimeDyld::MemoryManager {
 private:
-#ifndef JL_COMPILE_ON_DEMAND
-    std::shared_ptr<RuntimeDyld::MemoryManager> MemMgr_;
+#ifdef JL_USE_POOLED_MEMMGRS
+    typedef JuliaOJIT::MemMgrPoolT::OwningResource MemMgrT_;
 #else
-    JuliaOJIT::MemMgrPoolT::OwningResource MemMgr_;
+    typedef std::shared_ptr<RuntimeDyld::MemoryManager> MemMgrT_;
 #endif
 
+    MemMgrT_ MemMgr_;
+
 public:
-    ForwardingMemoryManager(
-#ifndef JL_COMPILE_ON_DEMAND
-        std::shared_ptr<RuntimeDyld::MemoryManager> MemMgr
-#else
-        decltype(MemMgr_) MemMgr
-#endif
-    ) : MemMgr_(std::move(MemMgr)) {}
-#ifdef JL_COMPILE_ON_DEMAND
+    ForwardingMemoryManager(MemMgrT_ MemMgr) : MemMgr_(std::move(MemMgr)) {}
+
+#ifdef JL_USE_POOLED_MEMMGRS
 #define MemMgr (*MemMgr_)
 #else
 #define MemMgr MemMgr_
@@ -728,7 +727,7 @@ class ForwardingMemoryManager : public RuntimeDyld::MemoryManager {
     }
     virtual bool finalizeMemory(std::string *ErrMsg = nullptr) override {
         auto error = MemMgr->finalizeMemory(ErrMsg);
-#ifdef JL_COMPILE_ON_DEMAND
+#ifdef JL_USE_POOLED_MEMMGRS
         MemMgr_.reset();
 #endif
         return error;
@@ -1019,6 +1018,11 @@ JuliaOJIT::JuliaOJIT()
     GlobalJD(ES.createBareJITDylib("JuliaGlobals")),
     JD(ES.createBareJITDylib("JuliaOJIT")),
     ContextPool([](){ return orc::ThreadSafeContext(std::make_unique<LLVMContext>()); }),
+#ifdef JL_COMPILE_ON_DEMAND
+    //TODO set an actual COD error handler address
+    LCTM(cantFail(orc::createLocalLazyCallThroughManager(TM->getTargetTriple(), ES, 0))),
+#endif
+
 #ifdef JL_USE_JITLINK
     // TODO: Port our memory management optimisations to JITLink instead of using the
     // default InProcessMemoryManager.
@@ -1027,8 +1031,7 @@ JuliaOJIT::JuliaOJIT()
 # else
     ObjectLayer(ES, cantFail(jitlink::InProcessMemoryManager::Create())),
 # endif
-#else
-#ifndef JL_COMPILE_ON_DEMAND
+#elif !defined(JL_USE_POOLED_MEMMGRS)
     MemMgr(createRTDyldMemoryManager()),
     ObjectLayer(
             ES,
@@ -1038,8 +1041,6 @@ JuliaOJIT::JuliaOJIT()
             }
         ),
 #else
-    //TODO set an actual COD error handler address
-    LCTM(cantFail(orc::createLocalLazyCallThroughManager(TM->getTargetTriple(), ES, 0))),
     MemMgrs([](){ return std::unique_ptr<RTDyldMemoryManager>(createRTDyldMemoryManager()); }),
     ObjectLayer(
             ES,
@@ -1048,7 +1049,6 @@ JuliaOJIT::JuliaOJIT()
                 return result;
             }
         ),
-#endif
 #endif
     Pipelines{
         std::make_unique<PipelineT>(ObjectLayer, *TM, 0),
@@ -1078,7 +1078,7 @@ JuliaOJIT::JuliaOJIT()
         [this](orc::MaterializationResponsibility &MR,
                const object::ObjectFile &Object,
                const RuntimeDyld::LoadedObjectInfo &LO) {
-#ifdef JL_COMPILE_ON_DEMAND
+#ifdef JL_USE_POOLED_MEMMGRS
             auto MemMgr = nullptr;
 #endif
             registerRTDyldJITObject(Object, LO, MemMgr);
@@ -1304,7 +1304,7 @@ size_t getRTDyldMemoryManagerTotalBytes(RTDyldMemoryManager *mm);
 
 size_t JuliaOJIT::getTotalBytes() const
 {
-#ifndef JL_COMPILE_ON_DEMAND
+#ifndef JL_USE_POOLED_MEMMGRS
     return getRTDyldMemoryManagerTotalBytes(MemMgr.get());
 #else
     return 0;
diff --git a/src/jitlayers.h b/src/jitlayers.h
index cf43a447d43a6..df1aa7db731e9 100644
--- a/src/jitlayers.h
+++ b/src/jitlayers.h
@@ -54,6 +54,11 @@
 #define JL_COMPILE_ON_DEMAND
 #endif
 
+//Pool memory managers because non-JITLink can't handle recursive codegen
+#if defined(JL_COMPILE_ON_DEMAND) && !defined(JL_USE_JITLINK)
+#define JL_USE_POOLED_MEMMGRS
+#endif
+
 using namespace llvm;
 
 extern "C" jl_cgparams_t jl_default_cgparams;
@@ -345,7 +350,7 @@ class JuliaOJIT {
     template<typename ResourceT, size_t max = 0>
     using QueuedResourcePool = ResourcePool<ResourceT, max, std::queue<ResourceT>>;
 
-#ifdef JL_COMPILE_ON_DEMAND
+#ifdef JL_USE_POOLED_MEMMGRS
     typedef QueuedResourcePool<std::unique_ptr<RTDyldMemoryManager>> MemMgrPoolT;
 #endif
     struct PipelineT {
@@ -456,14 +461,16 @@ class JuliaOJIT {
 
     QueuedResourcePool<orc::ThreadSafeContext> ContextPool;
 
-#ifndef JL_USE_JITLINK
-#ifndef JL_COMPILE_ON_DEMAND
-    const std::shared_ptr<RTDyldMemoryManager> MemMgr;
-#else
+#ifdef JL_COMPILE_ON_DEMAND
     std::unique_ptr<orc::LazyCallThroughManager> LCTM;
-    MemMgrPoolT MemMgrs;
 #endif
+
+#ifdef JL_USE_POOLED_MEMMGRS
+    MemMgrPoolT MemMgrs;
+#elif !defined(JL_USE_JITLINK)
+    const std::shared_ptr<RTDyldMemoryManager> MemMgr;
 #endif
+
     ObjLayerT ObjectLayer;
     const std::array<std::unique_ptr<PipelineT>, 4> Pipelines;
     OptSelLayerT OptSelLayer;