1#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
2#  See https://llvm.org/LICENSE.txt for license information.
3#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
4
5# Simply a wrapper around the extension module of the same name.
6from ._mlir_libs import _mlirExecutionEngine as _execution_engine
7import ctypes
8
9__all__ = [
10  "ExecutionEngine",
11]
12
13class ExecutionEngine(_execution_engine.ExecutionEngine):
14
15  def lookup(self, name):
16    """Lookup a function emitted with the `llvm.emit_c_interface`
17    attribute and returns a ctype callable.
18    Raise a RuntimeError if the function isn't found.
19    """
20    func = self.raw_lookup("_mlir_ciface_" + name)
21    if not func:
22      raise RuntimeError("Unknown function " + name)
23    prototype = ctypes.CFUNCTYPE(None, ctypes.c_void_p)
24    return prototype(func)
25
26  def invoke(self, name, *ctypes_args):
27    """Invoke a function with the list of ctypes arguments.
28    All arguments must be pointers.
29    Raise a RuntimeError if the function isn't found.
30    """
31    func = self.lookup(name)
32    packed_args = (ctypes.c_void_p * len(ctypes_args))()
33    for argNum in range(len(ctypes_args)):
34      packed_args[argNum] = ctypes.cast(ctypes_args[argNum], ctypes.c_void_p)
35    func(packed_args)
36
37  def register_runtime(self, name, ctypes_callback):
38    """Register a runtime function available to the jitted code
39    under the provided `name`. The `ctypes_callback` must be a
40    `CFuncType` that outlives the execution engine.
41    """
42    callback = ctypes.cast(ctypes_callback, ctypes.c_void_p)
43    self.raw_register_runtime("_mlir_ciface_" + name, callback)
44