Enzyme
Enzyme copied to clipboard
initial rust mpi support
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.
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
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?
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?)
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?
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
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)