Support MPI

[mofeing]

This PR...

• Registers MPI routine symbol address when MPI.jl gets loaded
• Specializes MPI.jl methods to be traced by Reactant

unresolved questions

• how can we represent MPI_Request with tensor and stablehlo types?
• mmm stablehlo.custom_call has a backend attribute that could be useful during lowering; e.g. if we want to lower to NCCL instead of MPI, since both have a similar API, we could potentially add our own custom c-functions that use NCCL but adapt them to MPI-like API
• @wsmoses can we create @cfunctions in Julia and pass them to the symbol table? some MPI routines might need a lil bit of adaption and writing them in Julia would be easier, faster (and also, would use the correct symbols from MPI.jl-loaded libmpi)

tested

• MPI.Comm_rank
•  MPI.Comm_size
• MPI.Barrier
  • Waiting for feat: add side_effect and backend_config to jit_call Enzyme-JAX#425 because we need side_effect=true on enzymexla.jit_call
• MPI.Wait
• MPI.Send
• MPI.Isend
• MPI.Recv
• MPI.Irecv
• MPI.Allreduce

to do

• MPI communicators
• sharding
• more MPI routines
• custom reduction operators

cc @JBlaschke @hhkit

[wsmoses]

you won't, instead you'll emit something like

function send_wrap(%arg : memref<axb>) {
    mpi.send %arg
}

function main() {
    ...
    enzymexla.jit_call @set_wrap(%x : tensor<...>)
}

And then lower-jit will convert into a custom call. however you will need to define a lowering of mpi.send into a corresponding MPI_Send call [which will use the symbol you just registered here]

Re CUDA though we also need to ensure we are sync'd wrt the current custream which you can get via enzymexla.get_stream

[mofeing]

mmm from our last discussion on this a couple of weeks ago, i understood that we would emit this

function main() {
    ...
    mpi.send(%arg0, ...)
    ...
}

and it would get lowered to

function send_wrap(%arg : memref<axb>) {
    llvm.call <0xffff> (%arg)
}

function main() {
    ...
    enzymexla.jit_call @send_wrap(%x : tensor<...>)
    ...
}

which will finally lower to the following with the enzymexla.jit pass

function main() {
    ...
    stablehlo.custom_call @mpi_send_wrap(%x : tensor<...>)
    ...
}

is this correct or do we need to emit the enzymexla.jit_call directly from Reactant?

ahh or do you mean that any wrapping we need to do around MPI should be done in this way?

Re CUDA though we also need to ensure we are sync'd wrt the current custream which you can get via enzymexla.get_stream

okay, this will probably be required for NCCL

[mofeing]

The PR is ready for review. The missing MPI routines are waiting to other PRs or need some fix, but they can be added later.

@giordano MPI testset result gets printed multiple times due to them being run on all ranks but I guess is not a problem?

[giordano]

MPI testset result gets printed multiple times due to them being run on all ranks but I guess is not a problem?

That's unfortunate. Just don't use @testset? We don't in https://github.com/JuliaParallel/MPI.jl/tree/5ef7fef6d6c3e2ab2ad380f346c77235f47213bf/test

[giordano]

Suggested change

	run(`$(mpiexec()) -n $nranks $(Base.julia_cmd()) integration/mpi.jl`)
	run(`$(mpiexec()) -n $nranks $(Base.julia_cmd()) --startup-file=no $(joinpath(@__DIR__, "integration", "mpi.jl")`)