[MLIR][Python] expose translate_module_to_llvmir (#163881)

This PR exposes `translate_module_to_llvmir` in the Python bindings.
This commit is contained in:
Maksim Levental
2025-10-20 12:14:52 -04:00
committed by GitHub
parent 8c826066e9
commit 5a112dedff
5 changed files with 45 additions and 1 deletions

View File

@@ -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;
};

View File

@@ -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) {

View File

@@ -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)

View File

@@ -601,6 +601,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

View File

@@ -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))