Integrate with Reactant?

[ChrisRackauckas]

From your slides someone sent me, it seems you are wanting some kernel fusion. The Reactant.jl project is doing this, so it seems like cuNumeric.jl should be trying to integrate with Reactant, or the other way around?

[ejmeitz]

Yes, this is definitely something we've been thinking about how to incorporate. In fact, the NVIDIA team implemented a similar MLIR workflow here specifically for kernel fusion. Reactant is more mature, more powerful and easier to integrate with than what is done in that paper, but (probably) lacks the logic to handle distributed systems.

Our first pass to achieve kernel fusion takes advantage of our ability to launch distributed CUDA kernels given the PTX code. We can use CUDA.jl to fuse broadcast expressions and run the generated kernels. This method will still not be as good as Reactant, but is better than nothing.

Long term I think cuNumeric.jl should integrate Reactant. I do not fully understand Reactant or its capabilities, but the obvious thing to me is to write a macro the user can inject to have Reactant optimize that line/function. Reactant will then emit PTX that we can run. It would be good to discuss more with @wsmoses.

[ejmeitz]

If its possible for Reactant to optimize things like mul!(C, A, B) without actually caring about the implementation of mul! that could be really powerful for us as well. For us this just calls into a C++ library, but if Reactant is able to understand that this is just a MatMul and do optimizations then there is definitely room for deeper integration.

I assume this is possible since mul! on a Julia array also calls into a foreign library (BLAS). The key for us is that Reactant has to emit a calls to the functions we support.

[wsmoses]

Yeah Reactant understands this is directly a matmul (and in fact can optimize the surrounding context, e.g. recogtnizing you're performing multiple matmuls and change to a batched gemm, and other algebraic/tensor simplifications).

[ejmeitz]

Ok sounds promising. The issue is we do not support things like batched gemm. Reactant would need to only emit calls to functions supported by cuNumeric. For us calling mul! actually calls a function that runs on all your GPUs, so Reactant would have to respect our limited API.

If the problem is embarrassingly parallel and then we can probably relax this constraint.

[wsmoses]

yeah presumably in this case after running algebraic optimizations (like perhaps you multiply by the identity, only use a slice of the result, etc), we would run a compiler transformation that maps from our tensor IR to primitives you support.

[ejmeitz]

Cool, to start I think we'll try the macro version where we only apply reactant to places the user deems "safe".

Is there a simple way to get PTX (or whatever executable Reactant makes) from Reactant? In a similar manner to CUDA.code_ptx. I think this is all we should need to start.

[krasow]

Combing Reactant's novel MLIR optimiziations and Legate's novel task scheduling for distributed task execution could be an interesting outcome. We would likely benchmark against the experimental cuPyNumeric work done by Rohan and the legate team @ ASPLOS 2025. As we mentioned, we would like to integrate more with the existing Julia ecosystem to achieve strong performance.