Optimization selection layer

Author: pchintalapudi

src/jitlayers.cpp

@@ -516,6 +516,33 @@ OptimizerResultT JuliaOJIT::OptimizerT::operator()(orc::ThreadSafeModule TSM, or
     return Expected<orc::ThreadSafeModule>{std::move(TSM)};
 }
 
+void JuliaOJIT::OptSelLayerT::emit(std::unique_ptr<orc::MaterializationResponsibility> R, orc::ThreadSafeModule TSM) {
+    size_t optlevel = ~0ull;
+    TSM.withModuleDo([&](Module &M) {
+        if (jl_generating_output()) {
+            optlevel = 0;
+        }
+        else {
+            optlevel = std::max(static_cast<int>(jl_options.opt_level), 0);
+            size_t optlevel_min = std::max(static_cast<int>(jl_options.opt_level_min), 0);
+            for (auto &F : M.functions()) {
+                if (!F.getBasicBlockList().empty()) {
+                    Attribute attr = F.getFnAttribute("julia-optimization-level");
+                    StringRef val = attr.getValueAsString();
+                    if (val != "") {
+                        size_t ol = (size_t)val[0] - '0';
+                        if (ol >= 0 && ol < optlevel)
+                            optlevel = ol;
+                    }
+                }
+            }
+            optlevel = std::min(std::max(optlevel, optlevel_min), this->count);
+        }
+    });
+    assert(optlevel != ~0ull && "Failed to select a valid optimization level!");
+    this->optimizers[optlevel].emit(std::move(R), std::move(TSM));
+}
+
 void jl_register_jit_object(const object::ObjectFile &debugObj,
                             std::function<uint64_t(const StringRef &)> getLoadAddress,
                             std::function<void *(void *)> lookupWriteAddress);
@@ -819,7 +846,8 @@ JuliaOJIT::JuliaOJIT(TargetMachine &TM, LLVMContext *LLVMCtx)
         {ES, CompileLayer1, OptimizerT(this, PM1, 1)},
         {ES, CompileLayer2, OptimizerT(this, PM2, 2)},
         {ES, CompileLayer3, OptimizerT(this, PM3, 3)},
-    }
+    },
+    OptSelLayer(OptimizeLayers)
 {
 #ifdef JL_USE_JITLINK
 # if defined(_OS_DARWIN_) && defined(LLVM_SHLIB)
@@ -919,30 +947,8 @@ void JuliaOJIT::addModule(std::unique_ptr<Module> M)
         }
     }
 #endif
-
-    int optlevel;
-    int optlevel_min;
-    if (jl_generating_output()) {
-        optlevel = 0;
-    }
-    else {
-        optlevel = jl_options.opt_level;
-        optlevel_min = jl_options.opt_level_min;
-        for (auto &F : M->functions()) {
-            if (!F.getBasicBlockList().empty()) {
-                Attribute attr = F.getFnAttribute("julia-optimization-level");
-                StringRef val = attr.getValueAsString();
-                if (val != "") {
-                    int ol = (int)val[0] - '0';
-                    if (ol >= 0 && ol < optlevel)
-                        optlevel = ol;
-                }
-            }
-        }
-        optlevel = std::max(optlevel, optlevel_min);
-    }
     // TODO: what is the performance characteristics of this?
-    cantFail(OptimizeLayers[std::max(std::min(optlevel, 3), 0)].add(JD, orc::ThreadSafeModule(std::move(M), TSCtx)));
+    cantFail(OptSelLayer.add(JD, orc::ThreadSafeModule(std::move(M), TSCtx)));
     // force eager compilation (for now), due to memory management specifics
     // (can't handle compilation recursion)
     for (auto Name : NewExports)

src/jitlayers.h

@@ -180,6 +180,16 @@ using CompilerResultT = Expected<std::unique_ptr<llvm::MemoryBuffer>>;
 using OptimizerResultT = Expected<orc::ThreadSafeModule>;
 
 class JuliaOJIT {
+public:
+#ifdef JL_USE_JITLINK
+    typedef orc::ObjectLinkingLayer ObjLayerT;
+#else
+    typedef orc::RTDyldObjectLinkingLayer ObjLayerT;
+#endif
+    typedef orc::IRCompileLayer CompileLayerT;
+    typedef orc::IRTransformLayer OptimizeLayerT;
+    typedef object::OwningBinary<object::ObjectFile> OwningObj;
+private:
     struct OptimizerT {
         OptimizerT(JuliaOJIT *pjit, legacy::PassManager &PM, int optlevel) : optlevel(optlevel), PM(PM), jit(*pjit) {}
 
@@ -193,15 +203,21 @@ class JuliaOJIT {
     template <typename ObjT, typename LoadResult>
     void registerObject(const ObjT &Obj, const LoadResult &LO);
 
+    struct OptSelLayerT : orc::IRLayer {
+
+        template<size_t N>
+        OptSelLayerT(OptimizeLayerT (&optimizers)[N]) : orc::IRLayer(optimizers[0].getExecutionSession(), optimizers[0].getManglingOptions()), optimizers(optimizers), count(N) {
+            static_assert(N > 0, "Expected array with at least one optimizer!");
+        }
+
+        void emit(std::unique_ptr<orc::MaterializationResponsibility> R, orc::ThreadSafeModule TSM) override;
+
+        private:
+        OptimizeLayerT *optimizers;
+        size_t count;
+    };
+
 public:
-#ifdef JL_USE_JITLINK
-    typedef orc::ObjectLinkingLayer ObjLayerT;
-#else
-    typedef orc::RTDyldObjectLinkingLayer ObjLayerT;
-#endif
-    typedef orc::IRCompileLayer CompileLayerT;
-    typedef orc::IRTransformLayer OptimizeLayerT;
-    typedef object::OwningBinary<object::ObjectFile> OwningObj;
 
     JuliaOJIT(TargetMachine &TM, LLVMContext *Ctx);
 
@@ -251,6 +267,7 @@ class JuliaOJIT {
     CompileLayerT CompileLayer2;
     CompileLayerT CompileLayer3;
     OptimizeLayerT OptimizeLayers[4];
+    OptSelLayerT OptSelLayer;
 
     DenseMap<void*, std::string> ReverseLocalSymbolTable;
 };