From fede9643cfad7d39b26fc64fbbe8cd899508d021 Mon Sep 17 00:00:00 2001 From: Yahya Jabary Date: Tue, 16 Jun 2026 14:38:31 +0200 Subject: [PATCH 1/6] backend: (llvm) add ctypes conversion registry Add CTypeContext, a registry mapping xDSL attribute classes to either ctypes types or per-instance converter callables. register_builtin_ctypes installs the standard scalar mappings for the builtin and llvm dialects. Wire the registry into the pyjit prototype so the CFUNCTYPE signature is derived from llvm.func types instead of hardcoded. --- tests/backend/llvm/test_convert_ctypes.py | 103 ++++++++++++++++++ .../filecheck/projects/pyjit/two_plus_two.py | 41 ++++--- xdsl/backend/llvm/convert_ctypes.py | 60 ++++++++++ 3 files changed, 189 insertions(+), 15 deletions(-) create mode 100644 tests/backend/llvm/test_convert_ctypes.py create mode 100644 xdsl/backend/llvm/convert_ctypes.py diff --git a/tests/backend/llvm/test_convert_ctypes.py b/tests/backend/llvm/test_convert_ctypes.py new file mode 100644 index 0000000000..f5fb7d31f4 --- /dev/null +++ b/tests/backend/llvm/test_convert_ctypes.py @@ -0,0 +1,103 @@ +import ctypes + +import pytest + +from xdsl.backend.llvm.convert_ctypes import CTypeContext, register_builtin_ctypes +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.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/tests/filecheck/projects/pyjit/two_plus_two.py b/tests/filecheck/projects/pyjit/two_plus_two.py index 43b024d505..359ce814c7 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 @@ -11,7 +11,9 @@ from typing_extensions import TypeVar from xdsl.backend.llvm.convert import convert_module +from xdsl.backend.llvm.convert_ctypes import CTypeContext, register_builtin_ctypes from xdsl.dialects import arith, builtin, func, llvm +from xdsl.dialects.builtin import ModuleOp from xdsl.frontend.pyast.context import PyASTContext from xdsl.transforms.desymref import FrontendDesymrefyPass from xdsl.transforms.mlir_opt import MLIROptPass @@ -44,10 +46,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 +63,19 @@ def mcjit_f64_f64_f64_binary( engine.finalize_object() # pyright: ignore engine.run_static_constructors() # pyright: ignore + func_op = next( + op + for op in xdsl_module.ops + if isinstance(op, llvm.FuncOp) and op.sym_name.data == symbol + ) + 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 +83,6 @@ def mcjit_f64_f64_f64_binary( func=fn, ) - return keepalive - # JIT @@ -80,6 +90,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 +102,12 @@ 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) + + 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/xdsl/backend/llvm/convert_ctypes.py b/xdsl/backend/llvm/convert_ctypes.py new file mode 100644 index 0000000000..ab6ef66761 --- /dev/null +++ b/xdsl/backend/llvm/convert_ctypes.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) From e835614a04624d90ba53c1929c9afbfd1fac2a9b Mon Sep 17 00:00:00 2001 From: Yahya Jabary Date: Wed, 17 Jun 2026 17:22:10 +0200 Subject: [PATCH 2/6] backend: (llvm) move convert_ctypes to xdsl.jit.llvm Sasha's suggestion: hang JIT infrastructure off a top-level xdsl.jit package so future targets (WASM, in-xDSL JIT) don't have to live under xdsl.backend. The LLVM-specific layer becomes xdsl.jit.llvm. --- tests/filecheck/projects/pyjit/two_plus_two.py | 2 +- tests/{backend => jit}/llvm/test_convert_ctypes.py | 2 +- xdsl/jit/__init__.py | 0 xdsl/jit/llvm/__init__.py | 0 xdsl/{backend => jit}/llvm/convert_ctypes.py | 0 5 files changed, 2 insertions(+), 2 deletions(-) rename tests/{backend => jit}/llvm/test_convert_ctypes.py (97%) create mode 100644 xdsl/jit/__init__.py create mode 100644 xdsl/jit/llvm/__init__.py rename xdsl/{backend => jit}/llvm/convert_ctypes.py (100%) diff --git a/tests/filecheck/projects/pyjit/two_plus_two.py b/tests/filecheck/projects/pyjit/two_plus_two.py index 359ce814c7..ab4e173ffb 100644 --- a/tests/filecheck/projects/pyjit/two_plus_two.py +++ b/tests/filecheck/projects/pyjit/two_plus_two.py @@ -11,10 +11,10 @@ from typing_extensions import TypeVar from xdsl.backend.llvm.convert import convert_module -from xdsl.backend.llvm.convert_ctypes import CTypeContext, register_builtin_ctypes 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.convert_ctypes import CTypeContext, register_builtin_ctypes from xdsl.transforms.desymref import FrontendDesymrefyPass from xdsl.transforms.mlir_opt import MLIROptPass diff --git a/tests/backend/llvm/test_convert_ctypes.py b/tests/jit/llvm/test_convert_ctypes.py similarity index 97% rename from tests/backend/llvm/test_convert_ctypes.py rename to tests/jit/llvm/test_convert_ctypes.py index f5fb7d31f4..dc81888c72 100644 --- a/tests/backend/llvm/test_convert_ctypes.py +++ b/tests/jit/llvm/test_convert_ctypes.py @@ -2,7 +2,6 @@ import pytest -from xdsl.backend.llvm.convert_ctypes import CTypeContext, register_builtin_ctypes from xdsl.dialects.builtin import ( Float16Type, Float32Type, @@ -15,6 +14,7 @@ ) from xdsl.dialects.llvm import LLVMPointerType, LLVMVoidType from xdsl.ir import Attribute +from xdsl.jit.llvm.convert_ctypes import CTypeContext, register_builtin_ctypes from xdsl.utils.exceptions import LLVMTranslationException 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/backend/llvm/convert_ctypes.py b/xdsl/jit/llvm/convert_ctypes.py similarity index 100% rename from xdsl/backend/llvm/convert_ctypes.py rename to xdsl/jit/llvm/convert_ctypes.py From 22c493f2a1817e85c8fbbea13d2b9805916ec4de Mon Sep 17 00:00:00 2001 From: Yahya Jabary Date: Wed, 17 Jun 2026 18:11:32 +0200 Subject: [PATCH 3/6] testing: (pyjit) look up entry FuncOp via SymbolTable --- tests/filecheck/projects/pyjit/two_plus_two.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/tests/filecheck/projects/pyjit/two_plus_two.py b/tests/filecheck/projects/pyjit/two_plus_two.py index ab4e173ffb..5c813b8c19 100644 --- a/tests/filecheck/projects/pyjit/two_plus_two.py +++ b/tests/filecheck/projects/pyjit/two_plus_two.py @@ -15,6 +15,7 @@ from xdsl.dialects.builtin import ModuleOp from xdsl.frontend.pyast.context import PyASTContext from xdsl.jit.llvm.convert_ctypes import CTypeContext, register_builtin_ctypes +from xdsl.traits import SymbolTable from xdsl.transforms.desymref import FrontendDesymrefyPass from xdsl.transforms.mlir_opt import MLIROptPass @@ -63,11 +64,8 @@ def mcjit_compile( engine.finalize_object() # pyright: ignore engine.run_static_constructors() # pyright: ignore - func_op = next( - op - for op in xdsl_module.ops - if isinstance(op, llvm.FuncOp) and op.sym_name.data == symbol - ) + 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) From ec4753d4a12337eb10d5e85d7eccb7784e993e3c Mon Sep 17 00:00:00 2001 From: Yahya Jabary Date: Wed, 17 Jun 2026 18:15:14 +0200 Subject: [PATCH 4/6] testing: (pyjit) TODO for future Python-type registry --- tests/filecheck/projects/pyjit/two_plus_two.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/filecheck/projects/pyjit/two_plus_two.py b/tests/filecheck/projects/pyjit/two_plus_two.py index 5c813b8c19..b1dd075bdb 100644 --- a/tests/filecheck/projects/pyjit/two_plus_two.py +++ b/tests/filecheck/projects/pyjit/two_plus_two.py @@ -103,6 +103,7 @@ def __init__(self): 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) return mcjit_compile(parsed_program.module, parsed_program.name, self.ctype_ctx) From 257a62cf77e2ee745c4d1d5506e606606b5d4764 Mon Sep 17 00:00:00 2001 From: Yahya Jabary Date: Wed, 17 Jun 2026 18:22:10 +0200 Subject: [PATCH 5/6] jit: (llvm) rename convert_ctypes to ctype_context --- tests/filecheck/projects/pyjit/two_plus_two.py | 2 +- .../jit/llvm/{test_convert_ctypes.py => test_ctype_context.py} | 2 +- xdsl/jit/llvm/{convert_ctypes.py => ctype_context.py} | 0 3 files changed, 2 insertions(+), 2 deletions(-) rename tests/jit/llvm/{test_convert_ctypes.py => test_ctype_context.py} (97%) rename xdsl/jit/llvm/{convert_ctypes.py => ctype_context.py} (100%) diff --git a/tests/filecheck/projects/pyjit/two_plus_two.py b/tests/filecheck/projects/pyjit/two_plus_two.py index b1dd075bdb..45c992ec07 100644 --- a/tests/filecheck/projects/pyjit/two_plus_two.py +++ b/tests/filecheck/projects/pyjit/two_plus_two.py @@ -14,7 +14,7 @@ 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.convert_ctypes import CTypeContext, register_builtin_ctypes +from xdsl.jit.llvm.ctype_context import CTypeContext, register_builtin_ctypes from xdsl.traits import SymbolTable from xdsl.transforms.desymref import FrontendDesymrefyPass from xdsl.transforms.mlir_opt import MLIROptPass diff --git a/tests/jit/llvm/test_convert_ctypes.py b/tests/jit/llvm/test_ctype_context.py similarity index 97% rename from tests/jit/llvm/test_convert_ctypes.py rename to tests/jit/llvm/test_ctype_context.py index dc81888c72..f90686d829 100644 --- a/tests/jit/llvm/test_convert_ctypes.py +++ b/tests/jit/llvm/test_ctype_context.py @@ -14,7 +14,7 @@ ) from xdsl.dialects.llvm import LLVMPointerType, LLVMVoidType from xdsl.ir import Attribute -from xdsl.jit.llvm.convert_ctypes import CTypeContext, register_builtin_ctypes +from xdsl.jit.llvm.ctype_context import CTypeContext, register_builtin_ctypes from xdsl.utils.exceptions import LLVMTranslationException diff --git a/xdsl/jit/llvm/convert_ctypes.py b/xdsl/jit/llvm/ctype_context.py similarity index 100% rename from xdsl/jit/llvm/convert_ctypes.py rename to xdsl/jit/llvm/ctype_context.py From 4034d1b088da31c510f7a3015f23e1486c1c968a Mon Sep 17 00:00:00 2001 From: Yahya Jabary Date: Wed, 17 Jun 2026 18:45:59 +0200 Subject: [PATCH 6/6] jit: (llvm) rename ctype_context to c_type_context --- tests/filecheck/projects/pyjit/two_plus_two.py | 2 +- .../jit/llvm/{test_ctype_context.py => test_c_type_context.py} | 2 +- xdsl/jit/llvm/{ctype_context.py => c_type_context.py} | 0 3 files changed, 2 insertions(+), 2 deletions(-) rename tests/jit/llvm/{test_ctype_context.py => test_c_type_context.py} (97%) rename xdsl/jit/llvm/{ctype_context.py => c_type_context.py} (100%) diff --git a/tests/filecheck/projects/pyjit/two_plus_two.py b/tests/filecheck/projects/pyjit/two_plus_two.py index 45c992ec07..f92e7b205e 100644 --- a/tests/filecheck/projects/pyjit/two_plus_two.py +++ b/tests/filecheck/projects/pyjit/two_plus_two.py @@ -14,7 +14,7 @@ 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.ctype_context import CTypeContext, register_builtin_ctypes +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 diff --git a/tests/jit/llvm/test_ctype_context.py b/tests/jit/llvm/test_c_type_context.py similarity index 97% rename from tests/jit/llvm/test_ctype_context.py rename to tests/jit/llvm/test_c_type_context.py index f90686d829..8fb34bdbed 100644 --- a/tests/jit/llvm/test_ctype_context.py +++ b/tests/jit/llvm/test_c_type_context.py @@ -14,7 +14,7 @@ ) from xdsl.dialects.llvm import LLVMPointerType, LLVMVoidType from xdsl.ir import Attribute -from xdsl.jit.llvm.ctype_context import CTypeContext, register_builtin_ctypes +from xdsl.jit.llvm.c_type_context import CTypeContext, register_builtin_ctypes from xdsl.utils.exceptions import LLVMTranslationException diff --git a/xdsl/jit/llvm/ctype_context.py b/xdsl/jit/llvm/c_type_context.py similarity index 100% rename from xdsl/jit/llvm/ctype_context.py rename to xdsl/jit/llvm/c_type_context.py