diff --git a/tests/filecheck/projects/pyjit/two_plus_two.py b/tests/filecheck/projects/pyjit/two_plus_two.py index 43b024d505..f92e7b205e 100644 --- a/tests/filecheck/projects/pyjit/two_plus_two.py +++ b/tests/filecheck/projects/pyjit/two_plus_two.py @@ -1,9 +1,9 @@ # RUN: python %s | filecheck %s from collections.abc import Callable -from ctypes import CFUNCTYPE, c_double +from ctypes import CFUNCTYPE from dataclasses import dataclass -from typing import Generic, ParamSpec +from typing import Any, Generic, ParamSpec import llvmlite import llvmlite.binding @@ -12,7 +12,10 @@ from xdsl.backend.llvm.convert import convert_module from xdsl.dialects import arith, builtin, func, llvm +from xdsl.dialects.builtin import ModuleOp from xdsl.frontend.pyast.context import PyASTContext +from xdsl.jit.llvm.c_type_context import CTypeContext, register_builtin_ctypes +from xdsl.traits import SymbolTable from xdsl.transforms.desymref import FrontendDesymrefyPass from xdsl.transforms.mlir_opt import MLIROptPass @@ -44,10 +47,12 @@ def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R: return self.func(*args, **kwargs) -# TODO: support automatic conversion of types -def mcjit_f64_f64_f64_binary( - llvm_module: llvm_ir.Module, symbol: str -) -> McJitKeepalive[[float, float], float]: +def mcjit_compile( + xdsl_module: ModuleOp, + symbol: str, + ctype_ctx: CTypeContext, +) -> McJitKeepalive[..., Any]: + llvm_module: llvm_ir.Module = convert_module(xdsl_module) llvm_ir_text = str(llvm_module) llvmlite.binding.initialize_native_target() # pyright: ignore llvmlite.binding.initialize_native_asmprinter() # pyright: ignore @@ -59,11 +64,16 @@ def mcjit_f64_f64_f64_binary( engine.finalize_object() # pyright: ignore engine.run_static_constructors() # pyright: ignore + func_op = SymbolTable.lookup_symbol(xdsl_module, symbol) + assert isinstance(func_op, llvm.FuncOp) + ret_ctype = ctype_ctx.to_ctype(func_op.function_type.output) + arg_ctypes = [ctype_ctx.to_ctype(t) for t in func_op.function_type.inputs] + fn_type = CFUNCTYPE(ret_ctype, *arg_ctypes) + func_ptr = engine.get_function_address(symbol) # pyright: ignore - fn_type = CFUNCTYPE(c_double, c_double, c_double) fn = fn_type(func_ptr) # pyright: ignore - keepalive = McJitKeepalive( + return McJitKeepalive( target=target, # pyright: ignore target_machine=target_machine, # pyright: ignore backing_mod=backing_mod, # pyright: ignore @@ -71,8 +81,6 @@ def mcjit_f64_f64_f64_binary( func=fn, ) - return keepalive - # JIT @@ -80,6 +88,7 @@ def mcjit_f64_f64_f64_binary( # TODO: support extending the JIT with more functionality class JITContext: pyast_ctx: PyASTContext + ctype_ctx: CTypeContext def __init__(self): ctx = PyASTContext(post_transforms=[FrontendDesymrefyPass(), convert_to_llvm]) @@ -91,12 +100,13 @@ def __init__(self): ctx.register_dialect(func.Func) self.pyast_ctx = ctx - def jit( - self, func: Callable[[float, float], float] - ) -> McJitKeepalive[[float, float], float]: + self.ctype_ctx = CTypeContext() + register_builtin_ctypes(self.ctype_ctx) + + # TODO: pair with a Python-type registry so unsupported P/R fail here, not deep in MCJIT. + def jit(self, func: Callable[P, R]) -> McJitKeepalive[P, R]: parsed_program = self.pyast_ctx.parse_program(func) - module = convert_module(parsed_program.module) - return mcjit_f64_f64_f64_binary(module, parsed_program.name) + return mcjit_compile(parsed_program.module, parsed_program.name, self.ctype_ctx) # Test diff --git a/tests/jit/llvm/test_c_type_context.py b/tests/jit/llvm/test_c_type_context.py new file mode 100644 index 0000000000..8fb34bdbed --- /dev/null +++ b/tests/jit/llvm/test_c_type_context.py @@ -0,0 +1,103 @@ +import ctypes + +import pytest + +from xdsl.dialects.builtin import ( + Float16Type, + Float32Type, + Float64Type, + IndexType, + IntAttr, + IntegerType, + NoneType, + Signedness, +) +from xdsl.dialects.llvm import LLVMPointerType, LLVMVoidType +from xdsl.ir import Attribute +from xdsl.jit.llvm.c_type_context import CTypeContext, register_builtin_ctypes +from xdsl.utils.exceptions import LLVMTranslationException + + +def test_register_and_resolve(): + ctx = CTypeContext() + ctx.register_ctype(Float32Type, lambda _: ctypes.c_float) + assert ctx.to_ctype(Float32Type()) is ctypes.c_float + + +def test_converter_receives_the_attribute(): + ctx = CTypeContext() + seen: list[IntegerType] = [] + + def converter(t: IntegerType): + seen.append(t) + return ctypes.c_int32 + + ctx.register_ctype(IntegerType, converter) + int_t = IntegerType(32) + assert ctx.to_ctype(int_t) is ctypes.c_int32 + assert seen == [int_t] + + +def test_resolve_unregistered_raises(): + ctx = CTypeContext() + with pytest.raises(LLVMTranslationException, match="No ctypes mapping"): + ctx.to_ctype(Float32Type()) + + +def test_re_register_overwrites(): + ctx = CTypeContext() + ctx.register_ctype(Float32Type, lambda _: ctypes.c_float) + ctx.register_ctype(Float32Type, lambda _: ctypes.c_double) + assert ctx.to_ctype(Float32Type()) is ctypes.c_double + + +def test_two_contexts_are_independent(): + a = CTypeContext() + b = CTypeContext() + a.register_ctype(Float32Type, lambda _: ctypes.c_float) + with pytest.raises(LLVMTranslationException): + b.to_ctype(Float32Type()) + + +@pytest.fixture +def ctx() -> CTypeContext: + c = CTypeContext() + register_builtin_ctypes(c) + return c + + +@pytest.mark.parametrize( + "type_attr, expected", + [ + (IntegerType(1), ctypes.c_bool), + (IntegerType(8), ctypes.c_int8), + (IntegerType(16), ctypes.c_int16), + (IntegerType(32), ctypes.c_int32), + (IntegerType(64), ctypes.c_int64), + (IntegerType(32, Signedness.SIGNED), ctypes.c_int32), + (IntegerType(32, Signedness.UNSIGNED), ctypes.c_int32), + (Float32Type(), ctypes.c_float), + (Float64Type(), ctypes.c_double), + (LLVMPointerType(), ctypes.c_void_p), + (LLVMPointerType(IntAttr(1)), ctypes.c_void_p), + (LLVMVoidType(), None), + (NoneType(), None), + ], +) +def test_builtin_resolve(ctx: CTypeContext, type_attr: Attribute, expected: object): + assert ctx.to_ctype(type_attr) is expected + + +@pytest.mark.parametrize( + "type_attr, match", + [ + (IntegerType(0), "integer of width 0"), + (IntegerType(17), "integer of width 17"), + (IntegerType(128), "integer of width 128"), + (Float16Type(), "No ctypes mapping for type"), + (IndexType(), "No ctypes mapping for type"), + ], +) +def test_builtin_unsupported(ctx: CTypeContext, type_attr: Attribute, match: str): + with pytest.raises(LLVMTranslationException, match=match): + ctx.to_ctype(type_attr) diff --git a/xdsl/jit/__init__.py b/xdsl/jit/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/xdsl/jit/llvm/__init__.py b/xdsl/jit/llvm/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/xdsl/jit/llvm/c_type_context.py b/xdsl/jit/llvm/c_type_context.py new file mode 100644 index 0000000000..ab6ef66761 --- /dev/null +++ b/xdsl/jit/llvm/c_type_context.py @@ -0,0 +1,60 @@ +import ctypes +from collections.abc import Callable +from typing import Any + +from xdsl.dialects.builtin import Float32Type, Float64Type, IntegerType, NoneType +from xdsl.dialects.llvm import LLVMPointerType, LLVMVoidType +from xdsl.ir import Attribute +from xdsl.utils.exceptions import LLVMTranslationException + + +class CTypeContext: + """Registry of xDSL attribute classes to ctypes converters.""" + + registry: dict[type[Attribute], Callable[[Any], Any]] + """Map from an xDSL attribute class to a converter producing its ctypes type.""" + + def __init__(self) -> None: + self.registry = {} + + def register_ctype( + self, + attr_type: type[Attribute], + converter: Callable[[Any], Any], + ) -> None: + self.registry[attr_type] = converter + + def to_ctype(self, type_attr: Attribute) -> Any: + try: + converter = self.registry[type(type_attr)] + except KeyError: + raise LLVMTranslationException(f"No ctypes mapping for type: {type_attr}") + return converter(type_attr) + + +_INT_CTYPE_BY_WIDTH: dict[int, Any] = { + 1: ctypes.c_bool, + 8: ctypes.c_int8, + 16: ctypes.c_int16, + 32: ctypes.c_int32, + 64: ctypes.c_int64, +} + + +def _int_to_ctype(type_attr: IntegerType) -> Any: + width = type_attr.width.data + try: + return _INT_CTYPE_BY_WIDTH[width] + except KeyError: + raise LLVMTranslationException( + f"No ctypes mapping for integer of width {width}" + ) + + +def register_builtin_ctypes(ctx: CTypeContext) -> None: + ctx.register_ctype(Float32Type, lambda _: ctypes.c_float) + ctx.register_ctype(Float64Type, lambda _: ctypes.c_double) + ctx.register_ctype(IntegerType, _int_to_ctype) + ctx.register_ctype(LLVMPointerType, lambda _: ctypes.c_void_p) + ctx.register_ctype(LLVMVoidType, lambda _: None) + ctx.register_ctype(NoneType, lambda _: None)