mlir/test/CAPI/execution_engine.c

   1 //===- execution_engine.c - Test for the C bindings for the MLIR JIT-------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM
   4 // Exceptions.
   5 // See https://llvm.org/LICENSE.txt for license information.
   6 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   7 //
   8 //===----------------------------------------------------------------------===//
   9
  10 /* RUN: mlir-capi-execution-engine-test 2>&1 | FileCheck %s
  11  */
  12 /* REQUIRES: host-supports-jit
  13  */
  14
  15 #include "mlir-c/Conversion.h"
  16 #include "mlir-c/ExecutionEngine.h"
  17 #include "mlir-c/IR.h"
  18 #include "mlir-c/RegisterEverything.h"
  19
  20 #include <assert.h>
  21 #include <math.h>
  22 #include <stdio.h>
  23 #include <stdlib.h>
  24 #include <string.h>
  25
  26 static void registerAllUpstreamDialects(MlirContext ctx) {
  27   MlirDialectRegistry registry = mlirDialectRegistryCreate();
  28   mlirRegisterAllDialects(registry);
  29   mlirContextAppendDialectRegistry(ctx, registry);
  30   mlirDialectRegistryDestroy(registry);
  31 }
  32
  33 void lowerModuleToLLVM(MlirContext ctx, MlirModule module) {
  34   MlirPassManager pm = mlirPassManagerCreate(ctx);
  35   MlirOpPassManager opm = mlirPassManagerGetNestedUnder(
  36       pm, mlirStringRefCreateFromCString("func.func"));
  37   mlirPassManagerAddOwnedPass(pm, mlirCreateConversionConvertFuncToLLVMPass());
  38   mlirOpPassManagerAddOwnedPass(
  39       opm, mlirCreateConversionArithToLLVMConversionPass());
  40   MlirLogicalResult status =
  41       mlirPassManagerRunOnOp(pm, mlirModuleGetOperation(module));
  42   if (mlirLogicalResultIsFailure(status)) {
  43     fprintf(stderr, "Unexpected failure running pass pipeline\n");
  44     exit(2);
  45   }
  46   mlirPassManagerDestroy(pm);
  47 }
  48
  49 // CHECK-LABEL: Running test 'testSimpleExecution'
  50 void testSimpleExecution(void) {
  51   MlirContext ctx = mlirContextCreate();
  52   registerAllUpstreamDialects(ctx);
  53
  54   MlirModule module = mlirModuleCreateParse(
  55       ctx, mlirStringRefCreateFromCString(
  56                // clang-format off
  57 "module {                                                                    \n"
  58 #ifdef __s390__
  59 "  func.func @add(%arg0 : i32) -> (i32 {llvm.signext}) attributes { llvm.emit_c_interface } {     \n"
  60 #else
  61 "  func.func @add(%arg0 : i32) -> i32 attributes { llvm.emit_c_interface } {     \n"
  62 #endif
  63 "    %res = arith.addi %arg0, %arg0 : i32                                        \n"
  64 "    return %res : i32                                                           \n"
  65 "  }                                                                             \n"
  66 "}"));
  67   // clang-format on
  68   lowerModuleToLLVM(ctx, module);
  69   mlirRegisterAllLLVMTranslations(ctx);
  70   MlirExecutionEngine jit = mlirExecutionEngineCreate(
  71       module, /*optLevel=*/2, /*numPaths=*/0, /*sharedLibPaths=*/NULL,
  72       /*enableObjectDump=*/false);
  73   if (mlirExecutionEngineIsNull(jit)) {
  74     fprintf(stderr, "Execution engine creation failed");
  75     exit(2);
  76   }
  77   int input = 42;
  78   int result = -1;
  79   void *args[2] = {&input, &result};
  80   if (mlirLogicalResultIsFailure(mlirExecutionEngineInvokePacked(
  81           jit, mlirStringRefCreateFromCString("add"), args))) {
  82     fprintf(stderr, "Execution engine creation failed");
  83     abort();
  84   }
  85   // CHECK: Input: 42 Result: 84
  86   printf("Input: %d Result: %d\n", input, result);
  87   mlirExecutionEngineDestroy(jit);
  88   mlirModuleDestroy(module);
  89   mlirContextDestroy(ctx);
  90 }
  91
  92 // CHECK-LABEL: Running test 'testOmpCreation'
  93 void testOmpCreation(void) {
  94   MlirContext ctx = mlirContextCreate();
  95   registerAllUpstreamDialects(ctx);
  96
  97   MlirModule module = mlirModuleCreateParse(
  98       ctx, mlirStringRefCreateFromCString(
  99                // clang-format off
 100 "module {                                                                       \n"
 101 "  func.func @main() attributes { llvm.emit_c_interface } {                     \n"
 102 "    %0 = arith.constant 0 : i32                                                \n"
 103 "    %1 = arith.constant 1 : i32                                                \n"
 104 "    %2 = arith.constant 2 : i32                                                \n"
 105 "    omp.parallel {                                                             \n"
 106 "      omp.wsloop {                                                             \n"
 107 "        omp.loop_nest (%3) : i32 = (%0) to (%2) step (%1) {                    \n"
 108 "          omp.yield                                                            \n"
 109 "        }                                                                      \n"
 110 "      }                                                                        \n"
 111 "      omp.terminator                                                           \n"
 112 "    }                                                                          \n"
 113 "    llvm.return                                                                \n"
 114 "  }                                                                            \n"
 115 "}                                                                              \n"
 116       ));
 117   // clang-format on
 118   lowerModuleToLLVM(ctx, module);
 119
 120   // At this point all operations in the MLIR module have been lowered to the
 121   // 'llvm' dialect except 'omp' operations. The goal of this test is
 122   // guaranteeing that the execution engine C binding has registered OpenMP
 123   // translations and therefore does not fail when it encounters 'omp' ops.
 124   // We don't attempt to run the engine, since that would force us to link
 125   // against the OpenMP library.
 126   MlirExecutionEngine jit = mlirExecutionEngineCreate(
 127       module, /*optLevel=*/2, /*numPaths=*/0, /*sharedLibPaths=*/NULL,
 128       /*enableObjectDump=*/false);
 129   if (mlirExecutionEngineIsNull(jit)) {
 130     fprintf(stderr, "Engine creation failed with OpenMP");
 131     exit(2);
 132   }
 133   // CHECK: Engine creation succeeded with OpenMP
 134   printf("Engine creation succeeded with OpenMP\n");
 135   mlirExecutionEngineDestroy(jit);
 136   mlirModuleDestroy(module);
 137   mlirContextDestroy(ctx);
 138 }
 139
 140 int main(void) {
 141
 142 #define _STRINGIFY(x) #x
 143 #define STRINGIFY(x) _STRINGIFY(x)
 144 #define TEST(test)                                                             \
 145   printf("Running test '" STRINGIFY(test) "'\n");                              \
 146   test();
 147
 148   TEST(testSimpleExecution);
 149   TEST(testOmpCreation);
 150   return 0;
 151 }