Reactant integration

[AntonOresten]

cutile-python has JAX support so I was thinking that cuTile.jl could potentially support Reactant.jl.

I'm not that familiar with current CUDA.jl integration in Reactant. It might be that cuTile already runs and compiles inside Reactant, but there may still be room for further integration, especially considering Reactant's existing CUDA Tile MLIR dialect (mentioned below).

I let Astra think a little about this

cuTile.jl + Reactant integration

The same approach as cuTile Python’s JAX integration fits cuTile.jl + Reactant. I’d expose it through cuTile’s existing @cuda backend=cuTile syntax.

The intended user experience would look like this—proposed integration, not currently working code:

using CUDA, cuTile, Reactant
import cuTile as ct

function double_kernel!(out, x, block::Int)
    i = ct.bid(1)
    tile = ct.load(x; index=i, shape=(block,))
    ct.store(out; index=i, tile=tile .* 2)
    return
end

function double_and_sin(x)
    out = similar(x)
    block = 128

    @cuda backend=cuTile blocks=cld(length(x), block) double_kernel!(
        out, x, ct.Constant(block)
    )

    return sin.(out)
end

x = Reactant.to_rarray(rand(Float32, 1024))
f = Reactant.@compile double_and_sin(x)
y = f(x)

Here, similar(x) supplies the output shape and element type. The extension would record the kernel’s write to out as a new Reactant value, letting subsequent operations consume it. You wouldn’t necessarily need Python-style OutputPlaceholder objects.

The implementation would have four parts:

1. Intercept cuTile launches during Reactant tracing. cuTile already implements CUDA.jl’s kernel_convert and kernel_compile backend hooks. A Reactant extension would adapt traced arrays using their shape/type metadata and record a launch instead of executing it.

2. Compile the kernel with cuTile.jl. Its compile_or_lookup already produces a cached cubin without requiring a CUDA context. The bridge would use that compilation path, targeting Reactant’s execution device, and embed the binary, entry-point name, and launch metadata in the graph.

3. Emit an XLA custom call. This is the direct equivalent of Python’s cutile_launch: input/output buffers, declared aliases, and launch attributes. Layouts should match Julia’s column-major arrays; Reactant’s existing julia_callback lowering demonstrates those layout declarations.

4. Launch through a native FFI handler on XLA’s stream. The handler would load/cache the cubin and pack cuTile.jl’s arguments. That packing needs its own implementation: Julia’s launch ABI flattens arrays into pointers, sizes, and strides, recursively flattens structs, and appends an implicit kernel-state seed. Reusing the Python handler unchanged wouldn’t suffice.

Reactant’s existing CUDA integration doesn’t automatically cover this. It intercepts CUDA.cufunction, imports LLVM code, and emits enzymexla.kernel_call. cuTile uses its own compiler and produces Tile IR → cubin, so it needs a separate path.

There are two possible depths of integration:

Approach	What it provides
Cubin + XLA FFI	Closest to the JAX integration; preserves cuTile’s compilation and kernel behavior.
Import Tile IR into Reactant	Exposes kernel operations to compiler passes, but requires additional lowering and transformation support.

Reactant already has CUDA Tile dialect bindings, which makes the second direction interesting, but those bindings alone don’t provide cuTile.jl integration.

I’d start with the cubin/FFI bridge. Differentiation would need explicit derivative rules or further compiler work: Reactant’s documented kernel differentiation currently relies on raising kernels into tensor operations, which an opaque binary call doesn’t enable. See the Reactant kernel documentation.

[maleadt]

cc @mofeing