31 changes: 31 additions & 0 deletions mlir/lib/Bindings/Python/mlir/execution_engine.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
# 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

# Simply a wrapper around the extension module of the same name.
from . import _cext
import ctypes

class ExecutionEngine(_cext.execution_engine.ExecutionEngine):

def lookup(self, name):
"""Lookup a function emitted with the `llvm.emit_c_interface`
attribute and returns a ctype callable.
Raise a RuntimeError if the function isn't found.
"""
func = self.raw_lookup("_mlir_ciface_" + name)
if not func:
raise RuntimeError("Unknown function " + name)
prototype = ctypes.CFUNCTYPE(None, ctypes.c_void_p)
return prototype(func)

def invoke(self, name, *ctypes_args):
"""Invoke a function with the list of ctypes arguments.
All arguments must be pointers.
Raise a RuntimeError if the function isn't found.
"""
func = self.lookup(name)
packed_args = (ctypes.c_void_p * len(ctypes_args))()
for argNum in range(len(ctypes_args)):
packed_args[argNum] = ctypes.cast(ctypes_args[argNum], ctypes.c_void_p)
func(packed_args)
4 changes: 3 additions & 1 deletion mlir/lib/CAPI/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
add_subdirectory(Dialect)
add_subdirectory(Conversion)
add_subdirectory(ExecutionEngine)
add_subdirectory(IR)
add_subdirectory(Registration)
add_subdirectory(Dialect)
add_subdirectory(Transforms)


Expand Down
7 changes: 7 additions & 0 deletions mlir/lib/CAPI/Conversion/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
add_mlir_public_c_api_library(MLIRCAPIConversion
Passes.cpp

LINK_LIBS PUBLIC
${conversion_libs}
)
26 changes: 26 additions & 0 deletions mlir/lib/CAPI/Conversion/Passes.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
//===- Conversion.cpp - C API for Conversion Passes -----------------------===//
//
// 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/Conversion/Passes.h"
#include "mlir/CAPI/Pass.h"
#include "mlir/Pass/Pass.h"

// Must include the declarations as they carry important visibility attributes.
#include "mlir/Conversion/Passes.capi.h.inc"

using namespace mlir;

#ifdef __cplusplus
extern "C" {
#endif

#include "mlir/Conversion/Passes.capi.cpp.inc"

#ifdef __cplusplus
}
#endif
7 changes: 7 additions & 0 deletions mlir/lib/CAPI/ExecutionEngine/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
# Main API shared library.
add_mlir_public_c_api_library(MLIRCEXECUTIONENGINE
ExecutionEngine.cpp

LINK_LIBS PUBLIC
MLIRExecutionEngine
)
56 changes: 56 additions & 0 deletions mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
//===- 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/Target/LLVMIR.h"
#include "llvm/Support/TargetSelect.h"

using namespace mlir;

extern "C" MlirExecutionEngine mlirExecutionEngineCreate(MlirModule op) {
static bool init_once = [] {
llvm::InitializeNativeTarget();
llvm::InitializeNativeTargetAsmPrinter();
return true;
}();
(void)init_once;

mlir::registerLLVMDialectTranslation(*unwrap(op)->getContext());
auto jitOrError = ExecutionEngine::create(unwrap(op));
if (!jitOrError) {
consumeError(jitOrError.takeError());
return MlirExecutionEngine{nullptr};
}
return wrap(jitOrError->release());
}

extern "C" void mlirExecutionEngineDestroy(MlirExecutionEngine jit) {
delete (unwrap(jit));
}

extern "C" MlirLogicalResult
mlirExecutionEngineInvokePacked(MlirExecutionEngine jit, MlirStringRef name,
void **arguments) {
const std::string ifaceName = ("_mlir_ciface_" + unwrap(name)).str();
llvm::Error error = unwrap(jit)->invokePacked(
ifaceName, MutableArrayRef<void *>{arguments, (size_t)0});
if (error)
return wrap(failure());
return wrap(success());
}

extern "C" void *mlirExecutionEngineLookup(MlirExecutionEngine jit,
MlirStringRef name) {
auto expectedFPtr = unwrap(jit)->lookup(unwrap(name));
if (!expectedFPtr)
return nullptr;
return reinterpret_cast<void *>(*expectedFPtr);
}
5 changes: 5 additions & 0 deletions mlir/lib/CAPI/Registration/Registration.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -10,9 +10,14 @@

#include "mlir/CAPI/IR.h"
#include "mlir/InitAllDialects.h"
#include "mlir/Target/LLVMIR.h"

void mlirRegisterAllDialects(MlirContext context) {
mlir::registerAllDialects(*unwrap(context));
// TODO: we may not want to eagerly load here.
unwrap(context)->loadAllAvailableDialects();
}

void mlirRegisterAllLLVMTranslations(MlirContext context) {
mlir::registerLLVMDialectTranslation(*unwrap(context));
}
99 changes: 99 additions & 0 deletions mlir/test/Bindings/Python/execution_engine.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,99 @@
# RUN: %PYTHON %s 2>&1 | FileCheck %s

import gc, sys
from mlir.ir import *
from mlir.passmanager import *
from mlir.execution_engine import *

# Log everything to stderr and flush so that we have a unified stream to match
# errors/info emitted by MLIR to stderr.
def log(*args):
print(*args, file=sys.stderr)
sys.stderr.flush()

def run(f):
log("\nTEST:", f.__name__)
f()
gc.collect()
assert Context._get_live_count() == 0

# Verify capsule interop.
# CHECK-LABEL: TEST: testCapsule
def testCapsule():
with Context():
module = Module.parse(r"""
llvm.func @none() {
llvm.return
}
""")
execution_engine = ExecutionEngine(module)
execution_engine_capsule = execution_engine._CAPIPtr
# CHECK: mlir.execution_engine.ExecutionEngine._CAPIPtr
log(repr(execution_engine_capsule))
execution_engine._testing_release()
execution_engine1 = ExecutionEngine._CAPICreate(execution_engine_capsule)
# CHECK: _mlir.execution_engine.ExecutionEngine
log(repr(execution_engine1))

run(testCapsule)

# Test invalid ExecutionEngine creation
# CHECK-LABEL: TEST: testInvalidModule
def testInvalidModule():
with Context():
# Builtin function
module = Module.parse(r"""
func @foo() { return }
""")
# CHECK: Got RuntimeError: Failure while creating the ExecutionEngine.
try:
execution_engine = ExecutionEngine(module)
except RuntimeError as e:
log("Got RuntimeError: ", e)

run(testInvalidModule)

def lowerToLLVM(module):
import mlir.conversions
pm = PassManager.parse("convert-std-to-llvm")
pm.run(module)
return module

# Test simple ExecutionEngine execution
# CHECK-LABEL: TEST: testInvokeVoid
def testInvokeVoid():
with Context():
module = Module.parse(r"""
func @void() attributes { llvm.emit_c_interface } {
return
}
""")
execution_engine = ExecutionEngine(lowerToLLVM(module))
# Nothing to check other than no exception thrown here.
execution_engine.invoke("void")

run(testInvokeVoid)


# Test argument passing and result with a simple float addition.
# CHECK-LABEL: TEST: testInvokeFloatAdd
def testInvokeFloatAdd():
with Context():
module = Module.parse(r"""
func @add(%arg0: f32, %arg1: f32) -> f32 attributes { llvm.emit_c_interface } {
%add = std.addf %arg0, %arg1 : f32
return %add : f32
}
""")
execution_engine = ExecutionEngine(lowerToLLVM(module))
# Prepare arguments: two input floats and one result.
# Arguments must be passed as pointers.
c_float_p = ctypes.c_float * 1
arg0 = c_float_p(42.)
arg1 = c_float_p(2.)
res = c_float_p(-1.)
execution_engine.invoke("add", arg0, arg1, res)
# CHECK: 42.0 + 2.0 = 44.0
log("{0} + {1} = {2}".format(arg0[0], arg1[0], res[0]))

run(testInvokeFloatAdd)
14 changes: 14 additions & 0 deletions mlir/test/CAPI/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
set(LLVM_OPTIONAL_SOURCES
execution_engine.c
ir.c
pass.c
)
Expand Down Expand Up @@ -28,3 +29,16 @@ target_link_libraries(mlir-capi-pass-test
PRIVATE
MLIRPublicAPI
)

add_llvm_executable(mlir-capi-execution-engine-test
execution_engine.c
DEPENDS
MLIRConversionPassIncGen
)
llvm_update_compile_flags(mlir-capi-execution-engine-test)

get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
target_link_libraries(mlir-capi-execution-engine-test
PRIVATE
MLIRPublicAPI
)
81 changes: 81 additions & 0 deletions mlir/test/CAPI/execution_engine.c
Original file line number Diff line number Diff line change
@@ -0,0 +1,81 @@
//===- execution_engine.c - Test for the C bindings for the 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
//
//===----------------------------------------------------------------------===//

/* RUN: mlir-capi-execution-engine-test 2>&1 | FileCheck %s
*/

#include "mlir-c/Conversion.h"
#include "mlir-c/ExecutionEngine.h"
#include "mlir-c/IR.h"
#include "mlir-c/Registration.h"

#include <assert.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

void lowerModuleToLLVM(MlirContext ctx, MlirModule module) {
MlirPassManager pm = mlirPassManagerCreate(ctx);
mlirPassManagerAddOwnedPass(pm, mlirCreateConversionConvertStandardToLLVM());
MlirLogicalResult status = mlirPassManagerRun(pm, module);
if (mlirLogicalResultIsFailure(status)) {
fprintf(stderr, "Unexpected failure running pass pipeline\n");
exit(2);
}
mlirPassManagerDestroy(pm);
}

// CHECK-LABEL: Running test 'testSimpleExecution'
void testSimpleExecution() {
MlirContext ctx = mlirContextCreate();
mlirRegisterAllDialects(ctx);
MlirModule module = mlirModuleCreateParse(
ctx, mlirStringRefCreateFromCString(
// clang-format off
"module { \n"
" func @add(%arg0 : i32) -> i32 attributes { llvm.emit_c_interface } { \n"
" %res = std.addi %arg0, %arg0 : i32 \n"
" return %res : i32 \n"
" } \n"
"}"));
// clang-format on
lowerModuleToLLVM(ctx, module);
mlirRegisterAllLLVMTranslations(ctx);
MlirExecutionEngine jit = mlirExecutionEngineCreate(module);
if (mlirExecutionEngineIsNull(jit)) {
fprintf(stderr, "Execution engine creation failed");
exit(2);
}
int input = 42;
int result = -1;
void *args[2] = {&input, &result};
if (mlirLogicalResultIsFailure(mlirExecutionEngineInvokePacked(
jit, mlirStringRefCreateFromCString("add"), args))) {
fprintf(stderr, "Execution engine creation failed");
abort();
}
// CHECK: Input: 42 Result: 84
printf("Input: %d Result: %d\n", input, result);
mlirExecutionEngineDestroy(jit);
mlirModuleDestroy(module);
mlirContextDestroy(ctx);
}

int main() {

#define _STRINGIFY(x) #x
#define STRINGIFY(x) _STRINGIFY(x)
#define TEST(test) \
printf("Running test '" STRINGIFY(test) "'\n"); \
test();

TEST(testSimpleExecution);
return 0;
}
2 changes: 2 additions & 0 deletions mlir/test/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ configure_lit_site_cfg(

set(MLIR_TEST_DEPENDS
FileCheck count not
mlir-capi-execution-engine-test
mlir-capi-ir-test
mlir-capi-pass-test
mlir-cpu-runner
Expand Down Expand Up @@ -118,6 +119,7 @@ if(MLIR_BINDINGS_PYTHON_ENABLED)
MLIRBindingsPythonExtension
MLIRBindingsPythonTestOps
MLIRTransformsBindingsPythonExtension
MLIRConversionsBindingsPythonExtension
)
endif()

Expand Down