mirror of
https://github.com/LLVMParty/llvm-nanobind
synced 2026-06-21 13:43:38 +00:00
Duncan Ogilvie
84 lines
2.3 KiB
Python
84 lines
2.3 KiB
Python
"""Compile LLVM IR in-process and call it from Python.
|
|
|
|
This example uses the LLVM-C ORC LLJIT wrapper. ``jit.add_module(mod)``
|
|
transfers the module into the JIT and invalidates the Python ``Module`` wrapper
|
|
on success.
|
|
|
|
Run from the repository root with:
|
|
|
|
uv run python examples/jit_add.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ctypes
|
|
from typing import cast
|
|
|
|
import llvm
|
|
|
|
|
|
CALLBACK_TYPE = ctypes.CFUNCTYPE(ctypes.c_int32, ctypes.c_int32)
|
|
|
|
|
|
@CALLBACK_TYPE
|
|
def py_add_ten(value: int) -> int:
|
|
return value + 10
|
|
|
|
|
|
def add_functions_to_module(ctx: llvm.Context, mod: llvm.Module) -> None:
|
|
i32 = ctx.types.i32
|
|
|
|
add_fn = mod.add_function("add_i32", ctx.types.function(i32, [i32, i32]))
|
|
lhs = add_fn.get_param(0)
|
|
rhs = add_fn.get_param(1)
|
|
add_entry = add_fn.append_basic_block("entry")
|
|
with add_entry.create_builder() as builder:
|
|
builder.ret(builder.add(lhs, rhs, "sum"))
|
|
|
|
callback_decl = mod.add_function("py_add_ten", ctx.types.function(i32, [i32]))
|
|
call_python_fn = mod.add_function("call_python", ctx.types.function(i32, [i32]))
|
|
arg = call_python_fn.get_param(0)
|
|
callback_entry = call_python_fn.append_basic_block("entry")
|
|
with callback_entry.create_builder() as builder:
|
|
result = builder.call(callback_decl, [arg], "result")
|
|
builder.ret(result)
|
|
|
|
|
|
def run_jit() -> tuple[int, int]:
|
|
with llvm.JIT.host() as jit:
|
|
jit.add_symbol("py_add_ten", py_add_ten)
|
|
|
|
with llvm.create_context() as ctx:
|
|
with ctx.create_module("jit_example") as mod:
|
|
add_functions_to_module(ctx, mod)
|
|
assert mod.verify(), mod.verification_error
|
|
jit.add_module(mod)
|
|
|
|
add_i32 = jit.ctypes_function(
|
|
"add_i32",
|
|
restype=ctypes.c_int32,
|
|
argtypes=[ctypes.c_int32, ctypes.c_int32],
|
|
)
|
|
call_python = jit.ctypes_function(
|
|
"call_python",
|
|
restype=ctypes.c_int32,
|
|
argtypes=[ctypes.c_int32],
|
|
)
|
|
|
|
return cast(int, add_i32(40, 2)), cast(int, call_python(5))
|
|
|
|
|
|
def main() -> None:
|
|
try:
|
|
add_result, callback_result = run_jit()
|
|
except llvm.LLVMError as exc:
|
|
print(f"skipped: host JIT is unavailable: {exc}")
|
|
return
|
|
|
|
print(f"add_i32(40, 2) = {add_result}")
|
|
print(f"call_python(5) = {callback_result}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|