[MLIR][Python] expose translate_module_to_llvmir

Author: makslevental

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,18 @@ 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) {
+ LLVMContextRef llvmCtx = LLVMContextCreate();
+ LLVMModuleRef llvmModule = mlirTranslateModuleToLLVMIR(module, llvmCtx);
+ char *llvmir = LLVMPrintModuleToString(llvmModule);
+ LLVMDisposeModule(llvmModule);
+ LLVMContextDispose(llvmCtx);
+ return llvmir;
+ },
+ "module"_a, nb::rv_policy::take_ownership);
 }
 
 NB_MODULE(_mlirDialectsLLVM, m) {

mlir/lib/CAPI/Target/LLVMIR.cpp

@@ -20,6 +20,12 @@
 
 using namespace mlir;
 
+[[maybe_unused]] auto *_LLVMContextCreate = &LLVMContextCreate;
+[[maybe_unused]] auto *_LLVMPrintModuleToString =
+ &LLVMPrintModuleToString;
+[[maybe_unused]] auto *_LLVMDisposeModule = &LLVMDisposeModule;
+[[maybe_unused]] auto *_LLVMContextDispose = &LLVMContextDispose;
+
 LLVMModuleRef mlirTranslateModuleToLLVMIR(MlirOperation module,
  LLVMContextRef context) {
  Operation *moduleOp = unwrap(module);

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

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