Skip to content

Backend

backend

JIT execution of LLVM IR modules.

The xdsl.jit.llvm.backend.LLVMJITBackend class compiles xDSL modules written in the LLVM dialect to native code, using llvmlite’s MCJIT execution engine and makes their signatures callable from Python.

To achieve this, the backend first derives the C function signature from the IR and then applies LLVM optimizations at level 2 by default.

For example, take the following add.mlir file, followed by a jit.py file which compiles and executes it from the Python API:

builtin.module {
  llvm.func @main() -> i32 {
    %left = llvm.mlir.constant(21 : i32) : i32
    %right = llvm.mlir.constant(21 : i32) : i32
    %result = llvm.add %left, %right : i32
    llvm.return %result : i32
  }
}

# /// script
# requires-python = ">=3.11"
# dependencies = ["xdsl[llvm]"]
# ///
import sys
from xdsl.context import Context
from xdsl.parser import Parser
from xdsl.dialects import builtin, llvm
from xdsl.jit.llvm.backend import LLVMJITBackend

ctx = Context()
ctx.load_dialect(builtin.Builtin)
ctx.load_dialect(llvm.LLVM)

with open(sys.argv[1]) as f:
    module = Parser(ctx, f.read()).parse_module()

backend = LLVMJITBackend(lowering=())
main = backend.jit(module, "main", ctx)

print(main.c_func())

uv run jit.py add.mlir  # 42

We pass lowering=() to LLVMJITBackend, because the module is already in the LLVM dialect and doesn’t need additional lowering tooling.

LLVMRawJITFunc dataclass

Bases: RawJITFunc

:class:RawJITFunc that retains LLVM MCJIT runtime objects.

The engine and related objects must remain referenced for as long as c_func may be called.

Source code in xdsl/jit/llvm/backend.py 
80
81
82
83
84
85
86
87
88
89
90
91
92
@dataclass(slots=True)
class LLVMRawJITFunc(RawJITFunc):
    """
    :class:`RawJITFunc` that retains LLVM MCJIT runtime objects.

    The engine and related objects must remain referenced for as long as
    ``c_func`` may be called.
    """

    target: Target
    target_machine: TargetMachine
    backing_mod: ModuleRef
    engine: ExecutionEngine

target: Target instance-attribute

target_machine: TargetMachine instance-attribute

backing_mod: ModuleRef instance-attribute

engine: ExecutionEngine instance-attribute

__init__(c_func_type: CFuncSignature, c_func: CFunc, target: Target, target_machine: TargetMachine, backing_mod: ModuleRef, engine: ExecutionEngine) -> None

LLVMJITBackend

Bases: JITBackend

:class:JITBackend using xDSL's LLVM converter and llvmlite MCJIT.

Runs :attr:lowering, requires symbol to name an llvm.FuncOp, then converts the module and exposes the entry point through CFFI ABI mode.

Source code in xdsl/jit/llvm/backend.py 
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
class LLVMJITBackend(JITBackend):
    """
    :class:`JITBackend` using xDSL's LLVM converter and llvmlite MCJIT.

    Runs :attr:`lowering`, requires ``symbol`` to name an ``llvm.FuncOp``, then
    converts the module and exposes the entry point through CFFI ABI mode.
    """

    lowering: tuple[ModulePass, ...]
    """Pass pipeline applied before resolving ``symbol``."""

    opt_level: Literal[0, 1, 2, 3]
    """LLVM optimization level, applied to code generation and the IR pipeline."""

    def __init__(
        self,
        lowering: tuple[ModulePass, ...] = (
            MLIROptPass(
                arguments=("--convert-arith-to-llvm", "--convert-func-to-llvm"),
                generic=True,
            ),
        ),
        *,
        opt_level: Literal[0, 1, 2, 3] = 2,
    ):
        """Construct the backend with the given ``lowering`` and ``opt_level``."""
        super().__init__()
        register_builtin_types(self.c_type_context)
        register_llvm_types(self.c_type_context)
        self.lowering = lowering
        self.opt_level = opt_level

    def jit(
        self,
        mlir_module: builtin.ModuleOp,
        symbol: str,
        ir_context: Context,
    ) -> LLVMRawJITFunc:
        """Lower ``mlir_module``, bind ``symbol``, and return an :class:`LLVMRawJITFunc`."""
        # `jit` may be called more than once against the same context
        if llvm.LLVM.name not in ir_context.registered_dialect_names:
            ir_context.load_dialect(llvm.LLVM)
        PassPipeline(self.lowering).apply(ir_context, mlir_module)
        func_op = SymbolTable.lookup_symbol(mlir_module, symbol)
        if func_op is None:
            raise JITException(f"No symbol to JIT compile: {symbol}")
        if not isinstance(func_op, llvm.FuncOp):
            raise JITException(
                f"Symbol {symbol} is a {func_op.name}, not an llvm.func: "
                "the lowering must leave it in the LLVM dialect"
            )
        c_func_type = to_c_func_type(self.c_type_context, func_op.function_type)
        target, target_machine = _create_target_machine(opt_level=self.opt_level)
        llvm_module = convert_module(
            mlir_module,
            fallback_target_triple=target_machine.triple,
            data_layout=str(target_machine.target_data),
        )
        return _compile_module(
            llvm_module,
            symbol,
            c_func_type,
            target=target,
            target_machine=target_machine,
            opt_level=self.opt_level,
        )

lowering: tuple[ModulePass, ...] = lowering instance-attribute

Pass pipeline applied before resolving symbol.

opt_level: Literal[0, 1, 2, 3] = opt_level instance-attribute

LLVM optimization level, applied to code generation and the IR pipeline.

__init__(lowering: tuple[ModulePass, ...] = (MLIROptPass(arguments=('--convert-arith-to-llvm', '--convert-func-to-llvm'), generic=True),), *, opt_level: Literal[0, 1, 2, 3] = 2)

Construct the backend with the given lowering and opt_level.

Source code in xdsl/jit/llvm/backend.py 
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
def __init__(
    self,
    lowering: tuple[ModulePass, ...] = (
        MLIROptPass(
            arguments=("--convert-arith-to-llvm", "--convert-func-to-llvm"),
            generic=True,
        ),
    ),
    *,
    opt_level: Literal[0, 1, 2, 3] = 2,
):
    """Construct the backend with the given ``lowering`` and ``opt_level``."""
    super().__init__()
    register_builtin_types(self.c_type_context)
    register_llvm_types(self.c_type_context)
    self.lowering = lowering
    self.opt_level = opt_level

jit(mlir_module: builtin.ModuleOp, symbol: str, ir_context: Context) -> LLVMRawJITFunc

Lower mlir_module, bind symbol, and return an :class:LLVMRawJITFunc.

Source code in xdsl/jit/llvm/backend.py 
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
def jit(
    self,
    mlir_module: builtin.ModuleOp,
    symbol: str,
    ir_context: Context,
) -> LLVMRawJITFunc:
    """Lower ``mlir_module``, bind ``symbol``, and return an :class:`LLVMRawJITFunc`."""
    # `jit` may be called more than once against the same context
    if llvm.LLVM.name not in ir_context.registered_dialect_names:
        ir_context.load_dialect(llvm.LLVM)
    PassPipeline(self.lowering).apply(ir_context, mlir_module)
    func_op = SymbolTable.lookup_symbol(mlir_module, symbol)
    if func_op is None:
        raise JITException(f"No symbol to JIT compile: {symbol}")
    if not isinstance(func_op, llvm.FuncOp):
        raise JITException(
            f"Symbol {symbol} is a {func_op.name}, not an llvm.func: "
            "the lowering must leave it in the LLVM dialect"
        )
    c_func_type = to_c_func_type(self.c_type_context, func_op.function_type)
    target, target_machine = _create_target_machine(opt_level=self.opt_level)
    llvm_module = convert_module(
        mlir_module,
        fallback_target_triple=target_machine.triple,
        data_layout=str(target_machine.target_data),
    )
    return _compile_module(
        llvm_module,
        symbol,
        c_func_type,
        target=target,
        target_machine=target_machine,
        opt_level=self.opt_level,
    )

is_native_triple(module_triple: str, native_triple: str) -> bool

Source code in xdsl/jit/llvm/backend.py 
112
113
114
115
116
117
118
119
120
121
122
123
def is_native_triple(module_triple: str, native_triple: str) -> bool:
    # LLVM spells one target several ways: arm64 for aarch64, macosx for darwin
    if module_triple in ("", "unknown-unknown-unknown"):
        return True
    module_parts = llvmlite.binding.get_triple_parts(module_triple)
    native_parts = llvmlite.binding.get_triple_parts(native_triple)
    return (
        module_parts.Arch == native_parts.Arch
        and module_parts.SubArch in ("", native_parts.SubArch)
        and module_parts.ObjectFormat == native_parts.ObjectFormat
        and module_parts.Env in ("unknown", native_parts.Env)
    )