Map speed level, size level to appropriate OptimizationLevel

numba · May 9, 2024 · 445ebb4 · 445ebb4
1 parent 4ea3d86
commit 445ebb4
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 50 deletions.
diff --git a/ffi/newpassmanagers.cpp b/ffi/newpassmanagers.cpp
@@ -24,15 +24,17 @@ DEFINE_SIMPLE_CONVERSION_FUNCTIONS(ModulePassManager, LLVMModulePassManagerRef)
 
 struct OpaqueFunctionPassManager;
 typedef OpaqueFunctionPassManager *LLVMFunctionPassManagerRef;
-DEFINE_SIMPLE_CONVERSION_FUNCTIONS(FunctionPassManager, LLVMFunctionPassManagerRef)
+DEFINE_SIMPLE_CONVERSION_FUNCTIONS(FunctionPassManager,
+                                   LLVMFunctionPassManagerRef)
 
 struct OpaquePassBuilder;
 typedef OpaquePassBuilder *LLVMPassBuilderRef;
 DEFINE_SIMPLE_CONVERSION_FUNCTIONS(PassBuilder, LLVMPassBuilderRef)
 
 struct OpaquePipelineTuningOptions;
 typedef OpaquePipelineTuningOptions *LLVMPipelineTuningOptionsRef;
-DEFINE_SIMPLE_CONVERSION_FUNCTIONS(PipelineTuningOptions, LLVMPipelineTuningOptionsRef)
+DEFINE_SIMPLE_CONVERSION_FUNCTIONS(PipelineTuningOptions,
+                                   LLVMPipelineTuningOptionsRef)
 
 static TargetMachine *unwrap(LLVMTargetMachineRef P) {
     return reinterpret_cast<TargetMachine *>(P);
@@ -251,50 +253,61 @@ LLVMPY_CreatePassBuilder(LLVMTargetMachineRef TM,
     return llvm::wrap(new PassBuilder(target, *pt));
 }
 
+static OptimizationLevel mapLevel(int speed_level, int size_level) {
+    switch (size_level) {
+    case 0:
+        switch (speed_level) {
+        case 0:
+            return OptimizationLevel::O0;
+        case 1:
+            return OptimizationLevel::O1;
+        case 2:
+            return OptimizationLevel::O2;
+        case 3:
+            return OptimizationLevel::O3;
+        default:
+            llvm_unreachable("Invalid optimization level");
+        }
+    case 1:
+        if (speed_level == 1)
+            return OptimizationLevel::Os;
+        llvm_unreachable("Invalid optimization level for size level 1");
+    case 2:
+        if (speed_level == 2)
+            return OptimizationLevel::Oz;
+        llvm_unreachable("Invalid optimization level for size level 2");
+    default:
+        llvm_unreachable("Invalid size level");
+        break;
+    }
+}
+
 API_EXPORT(LLVMModulePassManagerRef)
-LLVMPY_buildPerModuleDefaultPipeline(LLVMPassBuilderRef PBref, int opt_level) {
+LLVMPY_buildPerModuleDefaultPipeline(LLVMPassBuilderRef PBref, int speed_level,
+                                     int size_level) {
 
     PassBuilder *PB = llvm::unwrap(PBref);
-    ModulePassManager *MPM;
-    if (opt_level == 0) {
-        MPM = new ModulePassManager(
-            PB->buildO0DefaultPipeline(OptimizationLevel::O0));
-        return llvm::wrap(MPM);
+    OptimizationLevel OL = mapLevel(speed_level, size_level);
+    if (OL == OptimizationLevel::O0) {
+        return llvm::wrap(
+            new ModulePassManager(PB->buildO0DefaultPipeline(OL)));
     }
 
-    OptimizationLevel lvl;
-    if (opt_level == 1)
-        lvl = OptimizationLevel::O1;
-    else if (opt_level == 2)
-        lvl = OptimizationLevel::O2;
-    else
-        lvl = OptimizationLevel::O3;
-
-    MPM = new ModulePassManager(PB->buildPerModuleDefaultPipeline(lvl));
-    return llvm::wrap(MPM);
+    return llvm::wrap(
+        new ModulePassManager(PB->buildPerModuleDefaultPipeline(OL)));
 }
 
 API_EXPORT(LLVMFunctionPassManagerRef)
 LLVMPY_buildFunctionSimplificationPipeline(LLVMPassBuilderRef PBref,
-                                           int opt_level) {
+                                           int speed_level, int size_level) {
 
     PassBuilder *PB = llvm::unwrap(PBref);
-    FunctionPassManager *FPM;
-    if (opt_level == 0) {
-        FPM = new FunctionPassManager();
-        return llvm::wrap(FPM);
-    }
-
-    OptimizationLevel lvl;
-    if (opt_level == 1)
-        lvl = OptimizationLevel::O1;
-    else if (opt_level == 2)
-        lvl = OptimizationLevel::O2;
-    else
-        lvl = OptimizationLevel::O3;
+    OptimizationLevel OL = mapLevel(speed_level, size_level);
+    if (OL == OptimizationLevel::O0)
+        return llvm::wrap(new FunctionPassManager());
 
-    FPM = new FunctionPassManager(
-        PB->buildFunctionSimplificationPipeline(lvl, ThinOrFullLTOPhase::None));
+    FunctionPassManager *FPM = new FunctionPassManager(
+        PB->buildFunctionSimplificationPipeline(OL, ThinOrFullLTOPhase::None));
     return llvm::wrap(FPM);
 }
 

diff --git a/llvmlite/binding/newpassmangers.py b/llvmlite/binding/newpassmangers.py
@@ -95,8 +95,8 @@ class PipelineTuningOptions(ffi.ObjectRef):
     def __init__(self, ptr=None):
         if ptr is None:
             ptr = ffi.lib.LLVMPY_CreatePipelineTuningOptions()
-            self._opt_level = 3
-            # self._size_level = 0
+            self._opt_level = 2
+            self._size_level = 0
         ffi.ObjectRef.__init__(self, ptr)
 
     @property
@@ -152,19 +152,17 @@ def opt_level(self):
     def opt_level(self, level):
         self._opt_level = level
 
-    # FIXME: Model optimization level combining opt level and
-    # size level similar to llvm
-    # @property
-    # def size_level(self):
-    #     """
-    #     Whether and how much to optimize for size.
-    #     An integer between 0 and 2.
-    #     """
-    #     return self._size_level
+    @property
+    def size_level(self):
+        """
+        Whether and how much to optimize for size.
+        An integer between 0 and 2.
+        """
+        return self._size_level
 
-    # @size_level.setter
-    # def size_level(self, size):
-    #     self._size_level = size
+    @size_level.setter
+    def size_level(self, size):
+        self._size_level = size
 
     def _dispose(self):
         ffi.lib.LLVMPY_DisposePipelineTuningOptions(self)
@@ -182,13 +180,13 @@ def __init__(self, tm, pto, ptr=None):
     def getNewModulePassManager(self):
         return NewModulePassManager(
             ffi.lib.LLVMPY_buildPerModuleDefaultPipeline(
-                self, self._pto.opt_level)
+                self, self._pto.opt_level, self._pto.size_level)
         )
 
     def getNewFunctionPassManager(self):
         return NewFunctionPassManger(
             ffi.lib.LLVMPY_buildFunctionSimplificationPipeline(
-                self, self._pto.opt_level)
+                self, self._pto.opt_level, self._pto.size_level)
         )
 
     def _dispose(self):