[Mlir-commits] [mlir] db7cc03 - Activate OpenMP translation in MLIR execution engine CAPI.

[Mehdi Amini]

Author: Rafael Ubal Tena
Date: 2023-06-05T11:56:50-07:00
New Revision: db7cc0348c9a0aa44b23948949afef488d5083dd

URL: https://github.com/llvm/llvm-project/commit/db7cc0348c9a0aa44b23948949afef488d5083dd
DIFF: https://github.com/llvm/llvm-project/commit/db7cc0348c9a0aa44b23948949afef488d5083dd.diff

LOG: Activate OpenMP translation in MLIR execution engine CAPI.

We've observed that the MLIR Jit Engine fails when the `omp` dialect is used due to a failure to register OpenMP-related translations. This small patch addresses this issue.

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D151577

Added: 
    

Modified: 
    mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp
    mlir/test/CAPI/execution_engine.c

Removed: 
    


################################################################################
diff  --git a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp
index 1075ec46030c0..067cf677ee348 100644
--- a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp
@@ -13,6 +13,7 @@
 #include "mlir/ExecutionEngine/OptUtils.h"
 #include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.h"
 #include "llvm/ExecutionEngine/Orc/Mangling.h"
 #include "llvm/Support/TargetSelect.h"
 
@@ -33,6 +34,7 @@ mlirExecutionEngineCreate(MlirModule op, int optLevel, int numPaths,
   auto &ctx = *unwrap(op)->getContext();
   mlir::registerBuiltinDialectTranslation(ctx);
   mlir::registerLLVMDialectTranslation(ctx);
+  mlir::registerOpenMPDialectTranslation(ctx);
 
   auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();
   if (!tmBuilderOrError) {

diff  --git a/mlir/test/CAPI/execution_engine.c b/mlir/test/CAPI/execution_engine.c
index 2b4d448a65ef6..38a8fb8c3e213 100644
--- a/mlir/test/CAPI/execution_engine.c
+++ b/mlir/test/CAPI/execution_engine.c
@@ -85,6 +85,52 @@ void testSimpleExecution(void) {
   mlirContextDestroy(ctx);
 }
 
+// CHECK-LABEL: Running test 'testOmpCreation'
+void testOmpCreation(void) {
+  MlirContext ctx = mlirContextCreate();
+  registerAllUpstreamDialects(ctx);
+
+  MlirModule module = mlirModuleCreateParse(
+      ctx, mlirStringRefCreateFromCString(
+               // clang-format off
+"module {                                                                       \n"
+"  func.func @main() attributes { llvm.emit_c_interface } {                     \n"
+"    %0 = arith.constant 0 : i32                                                \n"
+"    %1 = arith.constant 1 : i32                                                \n"
+"    %2 = arith.constant 2 : i32                                                \n"
+"    omp.parallel {                                                             \n"
+"      omp.wsloop for (%3) : i32 = (%0) to (%2) step (%1) {                     \n"
+"        omp.yield                                                              \n"
+"      }                                                                        \n"
+"      omp.terminator                                                           \n"
+"    }                                                                          \n"
+"    llvm.return                                                                \n"
+"  }                                                                            \n"
+"}                                                                              \n"
+      ));
+  // clang-format on
+  lowerModuleToLLVM(ctx, module);
+
+  // At this point all operations in the MLIR module have been lowered to the
+  // 'llvm' dialect except 'omp' operations. The goal of this test is
+  // guaranteeing that the execution engine C binding has registered OpenMP
+  // translations and therefore does not fail when it encounters 'omp' ops.
+  // We don't attempt to run the engine, since that would force us to link
+  // against the OpenMP library.
+  MlirExecutionEngine jit = mlirExecutionEngineCreate(
+      module, /*optLevel=*/2, /*numPaths=*/0, /*sharedLibPaths=*/NULL,
+      /*enableObjectDump=*/false);
+  if (mlirExecutionEngineIsNull(jit)) {
+    fprintf(stderr, "Engine creation failed with OpenMP");
+    exit(2);
+  }
+  // CHECK: Engine creation succeeded with OpenMP
+  printf("Engine creation succeeded with OpenMP\n");
+  mlirExecutionEngineDestroy(jit);
+  mlirModuleDestroy(module);
+  mlirContextDestroy(ctx);
+}
+
 int main(void) {
 
 #define _STRINGIFY(x) #x
@@ -94,5 +140,6 @@ int main(void) {
   test();
 
   TEST(testSimpleExecution);
+  TEST(testOmpCreation);
   return 0;
 }