Reactant integration
- Dominant language
- Julia
- Stars
- 86
- Forks
- 12
- Avg merge
- 8h 24m
- Merged PRs (30d)
- 7
Description
cutile-python has [JAX support](https://docs.nvidia.com/cuda/cutile-python/interoperability.html#jax-ffi) so I was thinking that cuTile.jl could potentially support [Reactant.jl](https://github.com/EnzymeAD/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](https://github.com/EnzymeAD/Reactant.jl/blob/main/src/mlir/Dialects/CUDATile.jl) (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**:
```julia
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](https://github.com/JuliaGPU/cuTile.jl/blob/main/src/launch.jl). 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`](https://github.com/JuliaGPU/cuTile.jl/blob/main/src/launch.jl) 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](https://github.com/EnzymeAD/Reactant.jl/blob/main/src/Ops.jl) 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](https://github.com/JuliaGPU/cuTile.jl/blob/main/src/launch.jl) 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`](https://github.com/EnzymeAD/Reactant.jl/blob/main/ext/ReactantCUDAExt.jl), 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](https://github.com/EnzymeAD/Reactant.jl/blob/main/src/mlir/Dialects/CUDATile.jl), 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](https://github.com/EnzymeAD/Reactant.jl/blob/main/docs/src/tutorials/kernels.md).
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with cuTile.jl's launch.jl, especially kernel_convert, kernel_compile, and compile_or_lookup, then compare ReactantCUDAExt.jl, CUDATile.jl, and Ops.jl. Read Reactant's kernel documentation and the linked cuTile Python JAX integration before deciding between a cubin/FFI bridge and Tile IR integration. Done means a documented, tested Reactant path that traces, compiles, and executes a cuTile kernel.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- julia
- Domain
- backend, compilers
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100