[MLIR][Python] expose translate_module_to_llvmir

[makslevental]

This PR exposes translate_module_to_llvmir in the Python bindings.

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[llvmbot]

@llvm/pr-subscribers-mlir-llvm

@llvm/pr-subscribers-mlir

Author: Maksim Levental (makslevental)

Changes

This PR exposes translate_module_to_llvmir in the Python bindings. Note, I added mlirTranslateModuleToLLVMIRToString instead of doing something like

 m.def( "translate_module_to_llvmir", [](MlirOperation module) { LLVMContextRef llvmCtx = LLVMContextCreate(); LLVMModuleRef llvmModule = mlirTranslateModuleToLLVMIR(module, llvmCtx); char *llvmir = LLVMPrintModuleToString(llvmModule); LLVMDisposeModule(llvmModule); LLVMContextDispose(llvmCtx); return llvmir; },

because there seems to be absolutely no way to prevent symbol DCE in the transitive dep MLIRCAPITarget which is the target that actually links LLVMCore.

---

Full diff: https://github.com/llvm/llvm-project/pull/163881.diff

5 Files Affected:

• (modified) mlir/include/mlir-c/Target/LLVMIR.h (+3)
• (modified) mlir/lib/Bindings/Python/DialectLLVM.cpp (+12-1)
• (modified) mlir/lib/CAPI/Target/LLVMIR.cpp (+9)
• (modified) mlir/python/CMakeLists.txt (+2)
• (modified) mlir/test/python/dialects/llvm.py (+19)

diff --git a/mlir/include/mlir-c/Target/LLVMIR.h b/mlir/include/mlir-c/Target/LLVMIR.h index b5f948961e898..ec5ae88d485dc 100644 --- a/mlir/include/mlir-c/Target/LLVMIR.h +++ b/mlir/include/mlir-c/Target/LLVMIR.h @@ -33,6 +33,9 @@ extern "C" { MLIR_CAPI_EXPORTED LLVMModuleRef mlirTranslateModuleToLLVMIR(MlirOperation module, LLVMContextRef context); +MLIR_CAPI_EXPORTED char * +mlirTranslateModuleToLLVMIRToString(MlirOperation module); + struct MlirTypeFromLLVMIRTranslator { void *ptr; }; diff --git a/mlir/lib/Bindings/Python/DialectLLVM.cpp b/mlir/lib/Bindings/Python/DialectLLVM.cpp index 38de4a0e329a0..870a713b8edcb 100644 --- a/mlir/lib/Bindings/Python/DialectLLVM.cpp +++ b/mlir/lib/Bindings/Python/DialectLLVM.cpp @@ -11,6 +11,7 @@ #include "mlir-c/Dialect/LLVM.h" #include "mlir-c/IR.h" #include "mlir-c/Support.h" +#include "mlir-c/Target/LLVMIR.h" #include "mlir/Bindings/Python/Diagnostics.h" #include "mlir/Bindings/Python/Nanobind.h" #include "mlir/Bindings/Python/NanobindAdaptors.h" @@ -24,7 +25,7 @@ using namespace mlir; using namespace mlir::python; using namespace mlir::python::nanobind_adaptors; -static void populateDialectLLVMSubmodule(const nanobind::module_ &m) { +static void populateDialectLLVMSubmodule(nanobind::module_ &m) { //===--------------------------------------------------------------------===// // StructType @@ -154,6 +155,16 @@ static void populateDialectLLVMSubmodule(const nanobind::module_ &m) { .def_property_readonly("address_space", [](MlirType type) { return mlirLLVMPointerTypeGetAddressSpace(type); }); + + m.def( + "translate_module_to_llvmir", + [](MlirOperation module) { + return mlirTranslateModuleToLLVMIRToString(module); + }, + // clang-format off + nb::sig("def translate_module_to_llvmir(module: " MAKE_MLIR_PYTHON_QUALNAME("ir.Operation") ") -> str"), + // clang-format on + "module"_a, nb::rv_policy::take_ownership); } NB_MODULE(_mlirDialectsLLVM, m) { diff --git a/mlir/lib/CAPI/Target/LLVMIR.cpp b/mlir/lib/CAPI/Target/LLVMIR.cpp index 1c1912aec0f2f..00229dffafb61 100644 --- a/mlir/lib/CAPI/Target/LLVMIR.cpp +++ b/mlir/lib/CAPI/Target/LLVMIR.cpp @@ -34,6 +34,15 @@ LLVMModuleRef mlirTranslateModuleToLLVMIR(MlirOperation module, return moduleRef; } +char *mlirTranslateModuleToLLVMIRToString(MlirOperation module) { + LLVMContextRef llvmCtx = LLVMContextCreate(); + LLVMModuleRef llvmModule = mlirTranslateModuleToLLVMIR(module, llvmCtx); + char *llvmir = LLVMPrintModuleToString(llvmModule); + LLVMDisposeModule(llvmModule); + LLVMContextDispose(llvmCtx); + return llvmir; +} + DEFINE_C_API_PTR_METHODS(MlirTypeFromLLVMIRTranslator, mlir::LLVM::TypeFromLLVMIRTranslator) diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt index 9f5246de6bda0..a64d2e49c974f 100644 --- a/mlir/python/CMakeLists.txt +++ b/mlir/python/CMakeLists.txt @@ -591,6 +591,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPILLVM + # Misnomer - this is only the LLVMIR translation target. + MLIRCAPITarget ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind diff --git a/mlir/test/python/dialects/llvm.py b/mlir/test/python/dialects/llvm.py index d9ffdeb65bfd4..8ea0fddee3f7c 100644 --- a/mlir/test/python/dialects/llvm.py +++ b/mlir/test/python/dialects/llvm.py @@ -150,3 +150,22 @@ def testIntrinsics(): result = llvm.intr_memset(alloca, c_0, c_128, False) # CHECK: "llvm.intr.memset"(%[[ALLOCA]], %[[CST0]], %[[CST128]]) <{isVolatile = false}> : (!llvm.ptr, i8, i32) -> () print(result) + + +# CHECK-LABEL: testTranslateToLLVMIR +@constructAndPrintInModule +def testTranslateToLLVMIR(): + with Context(), Location.unknown(): + module = Module.parse( + """\ + llvm.func @add(%arg0: i64, %arg1: i64) -> i64 {  + %0 = llvm.add %arg0, %arg1 : i64  + llvm.return %0 : i64  + } + """ + ) + # CHECK: define i64 @add(i64 %0, i64 %1) { + # CHECK: %3 = add i64 %0, %1 + # CHECK: ret i64 %3 + # CHECK: } + print(llvm.translate_module_to_llvmir(module.operation)) 

[makslevental]

Note, I added mlirTranslateModuleToLLVMIRToString instead of doing something like

 m.def( "translate_module_to_llvmir", [](MlirOperation module) { LLVMContextRef llvmCtx = LLVMContextCreate(); LLVMModuleRef llvmModule = mlirTranslateModuleToLLVMIR(module, llvmCtx); char *llvmir = LLVMPrintModuleToString(llvmModule); LLVMDisposeModule(llvmModule); LLVMContextDispose(llvmCtx); return llvmir; },

because there seems to be absolutely no way to prevent symbol DCE (of e.g. LLVMContextCreate, LLVMPrintModuleToString) in the transitive dep MLIRCAPITarget which is the target that actually links LLVMCore.

[PragmaTwice]

It might be easier if we had an official LLVM Python bindings for this.

[makslevental]

It might be easier if we had an official LLVM Python bindings for this.

Lol I agree but I'm not gonna even start to try to propose anything along the lines of that - and anyway I have bindings in eudsl: https://github.com/llvm/eudsl/tree/main/projects/eudsl-llvmpy

[superbobry]

LGTM, but it does feel a bit ad-hoc.