backend: (llvm) add ctypes conversion registry [STACKED PR]

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,27 +64,31 @@ 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
         engine=engine,  # pyright: ignore
         func=fn,
     )
 
-    return keepalive
-
 
 # JIT
 
 
 # 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

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)

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)