DSL Protocols

FlyDSL uses three structural protocols at compilation and storage boundaries:

• JitArgument describes how a Python value enters a compiled JIT function.

• DslType converts DSL values to and from MLIR SSA values.

• Storable defines a DSL type’s byte layout and typed memory access.

They are capabilities, not base classes. A type implements a protocol by providing its hooks, and may implement any combination of the three.

This page connects Python arguments, DSL values, and typed memory. The detailed rules for composites, and storage and allocation live in their respective pages.

JitArgument

Hook	Meaning
value.__get_ir_types__()	MLIR argument types contributed to the function signature
value.__cache_signature__()	stable, hashable argument contribution to the compilation cache key
value.__c_abi_spec__()	ordered (ctype, fill) entries for the packed C interface

The IR types and final ABI slots must describe the same arguments in corresponding order, and that order must also agree with DslType reconstruction. Cache signatures should include compilation-relevant metadata such as type, dtype, rank, address space, alignment, and static layout, but exclude run-time values and addresses.

Each C-ABI fill(argument, storage) callback writes the current Python argument into preallocated ctypes storage. It must read each call’s value from argument, rather than capture that value when the specification is built. __c_abi_spec__() must match the final ABI produced after the types from __get_ir_types__() pass through MLIR lowering and llvm.emit_c_interface. This is not necessarily a one-to-one mapping: one MLIR argument may lower to multiple packed ABI slots. Implementing this hook therefore requires an understanding of the relevant MLIR-to-LLVM lowering and the generated C interface; a mismatch in slot type or order passes incorrect data to the compiled function.

JitArgumentRegistry.register(py_type, dsl_type=...) registers an adapter for a raw Python type. The adapter defines the host contract; dsl_type defines the value visible while tracing. A value implementing both JitArgument and DslType can be used directly without an adapter.

DslType

Hook	Meaning
value.__extract_to_ir_values__()	convert a DSL value into ordered ir.Values
Type.__construct_from_ir_values__(values, exemplar=None)	rebuild the DSL value from those SSA values

The two hooks form an ordered round trip. Reconstruction accepts replacement SSA values with the same count, order, and types as extraction produces, such as a new function’s block arguments. exemplar carries metadata not encoded by the MLIR type; for example, Vector uses it to retain its logical shape and element type.

A zero-length result is valid and is used by Constexpr values with no run-time representation.

Storable

Hook	Meaning
Type.__dsl_size_of__()	static byte size
Type.__dsl_align_of__()	required byte alignment
Type.__peek_from_ptr__(ptr)	materialize a typed value from an address
Type.__poke_into_ptr__(ptr, value)	write a typed value to an address

The free functions dsl_size_of, dsl_align_of, peek_from_ptr, and poke_into_ptr dispatch to these hooks. Storage[T] and the allocators use the same protocol; see Storage and Allocator.

Python-to-MLIR boundary

A Python argument and its corresponding DSL value may have different types. For example, a host tensor is represented by a JitArgument adapter at the Python boundary and by an fx.Tensor argument inside the JIT function.

The protocols connect these representations:

• JitArgumentRegistry associates a Python type with its host adapter and destination DSL type. A value implementing both JitArgument and DslType can be used directly.

• __get_ir_types__() supplies the ordered MLIR parameter types. DslType.__construct_from_ir_values__() rebuilds the corresponding DSL value from the function’s block arguments, using the adapter as an exemplar when metadata is needed.

• The type sequence returned by __get_ir_types__() must match the SSA value types expected by the corresponding DslType, in count, order, and MLIR type.

• At execution, __c_abi_spec__() supplies the ABI slots and callbacks that fill them from the current Python argument. These slots must match the lowered C interface of those MLIR parameters.

Python calls @jit
    │
    └─ adapt ─────> JitArgument
                      │
          compilation │ __cache_signature__ -> cache key
                      │ __get_ir_types__     -> JIT MLIR signature
                      │
                      v
              JIT function arguments
                      │ DslType.__construct_from_ir_values__
                      v
              DSL values in JIT body
                      │ call @kernel: trace its body
                      ├─ get_ir_types --------------------> gpu.func signature
                      ├─ __construct_from_ir_values__ ----> kernel DSL values
                      └─ .launch(): __extract_to_ir_values__ -> gpu.launch_func operands

             execution │ __c_abi_spec__ / fill
                       v
               packed C ABI -> lowered JIT function

MLIR lowering compiles the generated host func.func and device gpu.func into executable code. A cache hit reuses the compiled artifact without tracing again; fill still reads the current arguments on every call. Parameters annotated Constexpr[T] are supplied directly to tracing and affect specialization, but contribute no run-time MLIR parameters or C-ABI slots.

JitArgument is used only at the Python-to-MLIR boundary. The jit-to-kernel boundary is already inside MLIR construction and therefore uses DslType, not the host C ABI.

Protocol support

Representative type	DslType	JitArgument	Storable
supported scalar Numeric, such as fx.Int32	yes	yes	yes
fx.Vector, fx.Pointer, traced fx.Tensor	yes	no	no
storable specialized fx.Vector[...] / fx.Pointer[...]	yes	no	yes
TorchTensorJitArg, DLTensorJitArg	no	yes	no
fx.Stream	yes	yes	no
fx.Array[E, N]	yes	no	yes
@fx.struct	recursive	recursive	recursive
@fx.union	no	no	recursive
fx.Constexpr[T]	zero SSA values	special JIT handling	no

For a struct, each protocol is supported when all non-Constexpr fields support it. The fields are processed recursively in declaration order.

The table describes usable capabilities. All three protocols are @runtime_checkable, but structural isinstance and issubclass checks confirm only that the hooks exist, including hooks that raise for unsupported operations. Implementations remain responsible for correct ordering, MLIR types, cache keys, and byte layout.