[MLIR][Python] expose translate_module_to_llvmir

Author: makslevental

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

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,19 @@ 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;
+        return mlirTranslateModuleToLLVMIRToString(module);
+      },
+      "module"_a, nb::rv_policy::take_ownership);
 }
 
 NB_MODULE(_mlirDialectsLLVM, m) {

mlir/lib/CAPI/Target/LLVMIR.cpp

@@ -20,6 +20,14 @@
 
 using namespace mlir;
 
+extern "C" {
+// Prevent linker DCE.
+[[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);
@@ -34,6 +42,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)
 

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