Activate OpenMP translation in MLIR execution engine CAPI.

[rafaelubalmw]

Summary

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.

[mehdi_amini]

Can this be tested?

[rafaelubalmw]

Added unit test

The test for this feature creates an MLIR execution engine using the C API with
an MLIR module containing 'omp' ops. The creation of the execution engine instance
would fail if OpenMP translation had not been activated as expected.

[rafaelubalmw]

Hi @mehdi_amini - I currently don't have commit access to the LLVM repo. If you could merge this approved change, that'd be great. Thanks!

mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp

 //===- ExecutionEngine.cpp - C API for MLIR JIT ---------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir-c/ExecutionEngine.h"
 #include "mlir/CAPI/ExecutionEngine.h"
 #include "mlir/CAPI/IR.h"
 #include "mlir/CAPI/Support.h"
 #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"
 
 using namespace mlir;
 
 extern "C" MlirExecutionEngine
 mlirExecutionEngineCreate(MlirModule op, int optLevel, int numPaths,
                           const MlirStringRef *sharedLibPaths,
                           bool enableObjectDump) {
   static bool initOnce = [] {
     llvm::InitializeNativeTarget();
     llvm::InitializeNativeTargetAsmParser(); // needed for inline_asm
     llvm::InitializeNativeTargetAsmPrinter();
     return true;
   }();
   (void)initOnce;
 
   auto &ctx = *unwrap(op)->getContext();
   mlir::registerBuiltinDialectTranslation(ctx);
   mlir::registerLLVMDialectTranslation(ctx);
+  mlir::registerOpenMPDialectTranslation(ctx);
 
   auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();
   if (!tmBuilderOrError) {
     llvm::errs() << "Failed to create a JITTargetMachineBuilder for the host\n";
     return MlirExecutionEngine{nullptr};
   }
   auto tmOrError = tmBuilderOrError->createTargetMachine();
   if (!tmOrError) {

mlir/test/CAPI/execution_engine.c

"}"));
   }
   // CHECK: Input: 42 Result: 84
   printf("Input: %d Result: %d\n", input, result);
   mlirExecutionEngineDestroy(jit);
   mlirModuleDestroy(module);
   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
 #define STRINGIFY(x) _STRINGIFY(x)
 #define TEST(test)                                                             \
   printf("Running test '" STRINGIFY(test) "'\n");                              \
   test();
 
   TEST(testSimpleExecution);
+  TEST(testOmpCreation);
   return 0;
 }