Enzyme icon indicating copy to clipboard operation
Enzyme copied to clipboard

initial rust mpi support

Open ZuseZ4 opened this issue 1 year ago • 7 comments

ZuseZ4 avatar Jul 30 '24 20:07 ZuseZ4

We need to teach Enzyme to recognize

fn main( ) {
%5 = load ptr %MPI_SUM
%6 = call diffe_test_mpi_call(..., %5)
}

test_mpi_call(... ptr %5) {
%7 = call MPI_Allreduce(..., %5)
}

Because right now Enzyme yells about unknown mpi_allreduce op %5, I assume we already do that to e.g. recognize enyzme_const.

ZuseZ4 avatar Jul 30 '24 22:07 ZuseZ4

See my comment on the other thread, but no that should not be the way it is handled. Speciically if it is a generic reduction op, the derivative code can similarly be arbitrary

wsmoses avatar Jul 30 '24 22:07 wsmoses

Given that this has a test that passes at least in release mode, can we merge this (and maybe create an issue with further TODOs, to make sure the progress here doesn't get lost?

ZuseZ4 avatar Sep 09 '24 18:09 ZuseZ4

I don’t think there’s a test in this PR, no? (From my earlier comment which doesn’t look resolved the llvm test doesn’t test this?)

wsmoses avatar Sep 09 '24 19:09 wsmoses

Is it appropriate to add a test that pulls in rsmpi? That would require building a Rust toolchain using this Enzyme and then cargo test (which would build rsmpi, assuming an MPI implementation is available somewhere in the CI environment). I'm supposing that "soon", people will be getting their Enzyme-enable Rust from rustup, and that will increase interest in support for differentiable rsmpi, etc.

Is there a stand-alone test we could use here instead?

jedbrown avatar Mar 05 '25 19:03 jedbrown

As a small note, I wouldn't recommend building rustc anymore, rustc is providing dist(ributable) builds. Enzyme could download those, rm libenzyme.so and replace it by one that was just build on the latest commit. The compiler explorer is already using the dist artifacts (for now I still create them manually), but I will deprecate those soon in favour of the rustc ones: https://github.com/EnzymeAD/rust/releases

ZuseZ4 avatar Mar 05 '25 20:03 ZuseZ4

yeah we can pull in something. probably I'd recommend we make a separate repo for rust specific CI jobs (like how julia CI is inside of Enzyme.jl).

Before then, we should get these tests fixed up and merged ofc (which it looks like is just addressing some simplification comments above). Something like https://github.com/EnzymeAD/Enzyme/blob/8371758747eec72371626e719499d7bddc22dcad/enzyme/test/Enzyme/ReverseMode/mpi.ll#L26 but ofc with rsmpi (so we're not checking all of the other misc instructions)

wsmoses avatar Mar 08 '25 20:03 wsmoses