From 1dc533cea4ae6b9f72f7df53516c9f5c152ce35a Mon Sep 17 00:00:00 2001
From: Nicolas Vasilache <nicolas.vasilache@gmail.com>
Date: Mon, 19 Apr 2021 19:30:29 +0000
Subject: [PATCH] [mlir][python] ExecutionEngine can dump to object file

Differential Revision: https://reviews.llvm.org/D100786
---
 mlir/include/mlir-c/ExecutionEngine.h             |  5 +++++
 mlir/lib/Bindings/Python/ExecutionEngine.cpp      | 10 +++++++++-
 mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp |  5 +++++
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir-c/ExecutionEngine.h b/mlir/include/mlir-c/ExecutionEngine.h
index 5210f108ee6b..4a5d6ad9f648 100644
--- a/mlir/include/mlir-c/ExecutionEngine.h
+++ b/mlir/include/mlir-c/ExecutionEngine.h
@@ -67,6 +67,11 @@ MLIR_CAPI_EXPORTED void
 mlirExecutionEngineRegisterSymbol(MlirExecutionEngine jit, MlirStringRef name,
                                   void *sym);
 
+/// Dump as an object in `fileName`.
+MLIR_CAPI_EXPORTED void
+mlirExecutionEngineDumpToObjectFile(MlirExecutionEngine jit,
+                                    MlirStringRef fileName);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/mlir/lib/Bindings/Python/ExecutionEngine.cpp b/mlir/lib/Bindings/Python/ExecutionEngine.cpp
index 0e8ae8b38b3c..b5c8dde7563f 100644
--- a/mlir/lib/Bindings/Python/ExecutionEngine.cpp
+++ b/mlir/lib/Bindings/Python/ExecutionEngine.cpp
@@ -93,5 +93,13 @@ void mlir::python::populateExecutionEngineSubmodule(py::module &m) {
                 mlirStringRefCreate(name.c_str(), name.size()),
                 reinterpret_cast<void *>(sym));
           },
-          "Lookup function `func` in the ExecutionEngine.");
+          "Lookup function `func` in the ExecutionEngine.")
+      .def(
+          "dump_to_object_file",
+          [](PyExecutionEngine &executionEngine, const std::string &fileName) {
+            mlirExecutionEngineDumpToObjectFile(
+                executionEngine.get(),
+                mlirStringRefCreate(fileName.c_str(), fileName.size()));
+          },
+          "Dump ExecutionEngine to an object file.");
 }
diff --git a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp
index 345eac2193d8..36f24ed88965 100644
--- a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp
@@ -66,3 +66,8 @@ extern "C" void mlirExecutionEngineRegisterSymbol(MlirExecutionEngine jit,
     return symbolMap;
   });
 }
+
+extern "C" void mlirExecutionEngineDumpToObjectFile(MlirExecutionEngine jit,
+                                                    MlirStringRef name) {
+  unwrap(jit)->dumpToObjectFile(unwrap(name));
+}