Add target attributes to LLVM IR Module getting optimized

ffi/newpassmanagers.cpp

@@ -52,13 +52,51 @@ LLVMPY_CreateNewModulePassManager() {
     return llvm::wrap(new PassManager<Module>());
 }
 
+void LLVMPY_AddTargetAttributes(Module *M, std::string &triple,
+                                std::string &cpu, std::string &features) {
+    if (!triple.empty())
+        M->setTargetTriple(triple);
+
+    if (!cpu.empty() || !features.empty()) {
+        for (Function &F : *M) {
+            auto &Ctx = F.getContext();
+            AttributeList Attrs = F.getAttributes();
+            AttrBuilder NewAttrs(Ctx);
+
+            if (!cpu.empty() && !F.hasFnAttribute("target-cpu"))
+                NewAttrs.addAttribute("target-cpu", cpu);
+
+            if (!features.empty()) {
+                StringRef OldFeatures =
+                    F.getFnAttribute("target-features").getValueAsString();
+
+                if (OldFeatures.empty())
+                    NewAttrs.addAttribute("target-features", features);
+                else {
+                    SmallString<256> Appended(OldFeatures);
+                    Appended.push_back(',');
+                    Appended.append(features);
+                    NewAttrs.addAttribute("target-features", Appended);
+                }
+            }
+            // Override with new attributes
+            F.setAttributes(Attrs.addFnAttributes(Ctx, NewAttrs));
+        }
+    }
+}
+
 API_EXPORT(void)
 LLVMPY_RunNewModulePassManager(LLVMModulePassManagerRef MPMRef,
-                               LLVMModuleRef mod, LLVMPassBuilderRef PBRef) {
+                               LLVMModuleRef mod, LLVMPassBuilderRef PBRef,
+                               const char *triple, const char *cpu,
+                               const char *features) {
 
     ModulePassManager *MPM = llvm::unwrap(MPMRef);
     Module *M = llvm::unwrap(mod);
     PassBuilder *PB = llvm::unwrap(PBRef);
+    std::string triple_str = std::string(triple);
+    std::string cpu_str = std::string(cpu);
+    std::string features_str = std::string(features);
 
     // TODO: Make these set(able) by user
     bool DebugLogging = false;
@@ -84,6 +122,9 @@ LLVMPY_RunNewModulePassManager(LLVMModulePassManagerRef MPMRef,
     PB->registerCGSCCAnalyses(CGAM);
     PB->registerModuleAnalyses(MAM);
     PB->crossRegisterProxies(LAM, FAM, CGAM, MAM);
+
+    LLVMPY_AddTargetAttributes(M, triple_str, cpu_str, features_str);
+
     MPM->run(*M, MAM);
 }
 

ffi/targets.cpp

@@ -265,6 +265,21 @@ LLVMPY_CreateTargetMachineData(LLVMTargetMachineRef TM) {
         new llvm::DataLayout(llvm::unwrap(TM)->createDataLayout()));
 }
 
+API_EXPORT(const char *)
+LLVMPY_getTargetTriple(LLVMTargetMachineRef TM) {
+    return llvm::unwrap(TM)->getTargetTriple().getTriple().c_str();
+}
+
+API_EXPORT(const char *)
+LLVMPY_getTargetCPU(LLVMTargetMachineRef TM) {
+    return llvm::unwrap(TM)->getTargetCPU().data();
+}
+
+API_EXPORT(const char *)
+LLVMPY_getTargetFeatureString(LLVMTargetMachineRef TM) {
+    return llvm::unwrap(TM)->getTargetFeatureString().data();
+}
+
 API_EXPORT(void)
 LLVMPY_AddAnalysisPasses(LLVMTargetMachineRef TM, LLVMPassManagerRef PM) {
     LLVMAddAnalysisPasses(TM, PM);

llvmlite/binding/newpassmanagers.py

@@ -1,4 +1,4 @@
-from ctypes import c_bool, c_int, c_size_t
+from ctypes import c_bool, c_int, c_size_t, c_char_p
 from enum import IntFlag
 from llvmlite.binding import ffi
 
@@ -35,7 +35,11 @@ def __init__(self, ptr=None):
         super().__init__(ptr)
 
     def run(self, module, pb):
-        ffi.lib.LLVMPY_RunNewModulePassManager(self, module, pb)
+        ffi.lib.LLVMPY_RunNewModulePassManager(
+            self, module, pb,
+            pb._target_triple.encode('utf8'),
+            pb._target_cpu.encode('utf8'),
+            pb._target_features.encode('utf8'))
 
     def add_verifier(self):
         ffi.lib.LLVMPY_AddVerifierPass(self)
@@ -214,6 +218,9 @@ def __init__(self, tm, pto):
         super().__init__(ffi.lib.LLVMPY_CreatePassBuilder(tm, pto))
         self._pto = pto
         self._tm = tm
+        self._target_triple = tm.get_target_triple()
+        self._target_cpu = tm.get_target_cpu()
+        self._target_features = tm.get_target_feature_string()
 
     def getModulePassManager(self):
         return ModulePassManager(
@@ -240,7 +247,8 @@ def _dispose(self):
 
 ffi.lib.LLVMPY_RunNewModulePassManager.argtypes = [
     ffi.LLVMModulePassManagerRef, ffi.LLVMModuleRef,
-    ffi.LLVMPassBuilderRef,]
+    ffi.LLVMPassBuilderRef,
+    c_char_p, c_char_p, c_char_p,]
 
 ffi.lib.LLVMPY_AddVerifierPass.argtypes = [ffi.LLVMModulePassManagerRef,]
 ffi.lib.LLVMPY_AddAAEvalPass_module.argtypes = [ffi.LLVMModulePassManagerRef,]

llvmlite/binding/targets.py

@@ -376,6 +376,18 @@ def _emit_to_memory(self, module, use_object=False):
         finally:
             ffi.lib.LLVMPY_DisposeMemoryBuffer(mb)
 
+    def get_target_triple(self):
+        s = ffi.lib.LLVMPY_getTargetTriple(self)
+        return _decode_string(s)
+
+    def get_target_cpu(self):
+        s = ffi.lib.LLVMPY_getTargetCPU(self)
+        return _decode_string(s)
+
+    def get_target_feature_string(self):
+        s = ffi.lib.LLVMPY_getTargetFeatureString(self)
+        return _decode_string(s)
+
     @property
     def target_data(self):
         return TargetData(ffi.lib.LLVMPY_CreateTargetMachineData(self))
@@ -511,6 +523,15 @@ def has_svml():
 
 ffi.lib.LLVMPY_DisposeMemoryBuffer.argtypes = [ffi.LLVMMemoryBufferRef]
 
+ffi.lib.LLVMPY_getTargetTriple.argtypes = [ffi.LLVMTargetMachineRef]
+ffi.lib.LLVMPY_getTargetTriple.restype = c_char_p
+
+ffi.lib.LLVMPY_getTargetCPU.argtypes = [ffi.LLVMTargetMachineRef]
+ffi.lib.LLVMPY_getTargetCPU.restype = c_char_p
+
+ffi.lib.LLVMPY_getTargetFeatureString.argtypes = [ffi.LLVMTargetMachineRef]
+ffi.lib.LLVMPY_getTargetFeatureString.restype = c_char_p
+
 ffi.lib.LLVMPY_CreateTargetMachineData.argtypes = [
     ffi.LLVMTargetMachineRef,
 ]

llvmlite/tests/test_binding.py

@@ -3021,6 +3021,15 @@ def test_size_level_constraints(self):
 
 
 class NewPassManagerMixin(object):
+    def icelake_pass_builder(self, speed_level=0, size_level=0):
+        llvm.initialize_all_targets()
+        target = llvm.Target.from_triple("x86_64-unknown-unknown")
+        tm = target.create_target_machine(
+            cpu='skylake-avx512',
+            features="+avx,+aes,+sahf,+pclmul,-xop,+crc32", opt=3)
+
+        pto = llvm.create_pipeline_tuning_options(speed_level, size_level)
+        return llvm.create_pass_builder(tm, pto)
 
     def pb(self, speed_level=0, size_level=0):
         tm = self.target_machine(jit=False)
@@ -3127,6 +3136,21 @@ def test_add_passes(self):
         mpm.add_jump_threading_pass()
         mpm.add_refprune_pass()
 
+    def test_target(self):
+        pb = self.icelake_pass_builder(3, 0)
+        mpm = pb.getModulePassManager()
+        mod = llvm.parse_assembly("define void @foo() { entry: ret void }")
+        mpm.run(mod, pb)
+
+        # Check for target triple
+        self.assertIn('x86_64-unknown-unknown', str(mod))
+
+        # Check for target-cpu
+        self.assertIn('skylake-avx512', str(mod))
+
+        # Check for target-features
+        self.assertIn("+avx,+aes,+sahf,+pclmul,-xop,+crc32", str(mod))
+
 
 class TestNewFunctionPassManager(BaseTest, NewPassManagerMixin):
     def pm(self):