initial rust mpi support

[ZuseZ4]

No description provided.

[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.

[wsmoses]

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]

This is unused atm so this won't test any thing. Can you make this a minimal runnable case. Maybe a different mpi fn?

Also get rid of the other stuff like enzyme_type etc

[ZuseZ4]

@jedbrown I can't run things, can you test if some other function works on this branch?

[jedbrown]

I don't know how to resolve this.

source_id: DefId(0:12 ~ dot_enzyme[f53a]::dot_parallel)
num_fnc_args: 5
input_activity.len(): 5
error: <unknown>:0:0: in function preprocess__ZN10dot_enzyme12dot_parallel17hd37f1f8a2c8de07dE double (ptr, ptr, i64, ptr, i64): Enzyme: cannot compute with global variable that doesn't have marked shadow global
@RSMPI_SUM = external local_unnamed_addr global ptr

This is the relevant code: rsmpi/rsmpi@840e01c#diff-9a676b0d0c142cd1e89e8174ddb007db982d8602bd374a04e40e9f6a421acaebR216-R228

Run with

$ RUSTFLAGS='-Z unstable-options' cargo +enzyme r --example=dot_enzyme --release

[ZuseZ4]

Me neither, but can you share the module you got from ENZYME_OPT=1? Then I can experiment around to see if I find the right changes.

[wsmoses]

@jedbrown you need to add RSMPI_SUM to the ActivityAnalysis.cpp code [the message (poorly) warns that the global variable is differentiable, but Enzyme is unable to determine a differentiable version of the global. Of course it makes no sense to differentiate wrt MPI_SUM so we can mark that in ActivityAnalysis.cpp]

[jedbrown]

Thanks. Why didn't ompi_mpi_op_sum need to be there?

[wsmoses]

I think one has an extra pointer indirection causing a load which needs to be analyzed

[wsmoses]

yeah in this case, now the issue is that the op of MPI_allreduce cannot be detected (the earlier check if the argument was a literal global no longer applies, since this is a load of RSMPI_SUM). Changing the MPI_Allreduce check to consider something along the lines of:

if (LI = dyn_cast<LoadInst>(...)))
  if (auto GV = dyn_cast<GlobalVariable>(LI->getPointerOperand()))
    if (GV->getName() == "RSMPI_SUM")
      legal = true;

[jedbrown]

I pushed that (cb2d739; assuming it was the correct place), but now I have

source_id: DefId(0:12 ~ dot_enzyme[42dc]::dot_parallel)
num_fnc_args: 5
input_activity.len(): 5
error: <unknown>:0:0: in function preprocess__ZN10dot_enzyme12dot_parallel17h2f3ed146b457ca09E double (ptr, ptr, i64, ptr, i64): Enzyme: cannot compute with global variable that doesn't have marked shadow global
@RSMPI_COMM_WORLD = external local_unnamed_addr global ptr

[wsmoses]

okay that's the same as the first issue [a global which cannot be proven non-differentiable]. ActivityAnalysis.cpp is again the right place to add that.

[jedbrown]

Thanks for the hand-holding. I think this output is correct now.

$ RUSTFLAGS='-Z unstable-options' cargo +enzyme mpirun -n 2 --example=dot_enzyme --release | sort
[0] bx: [0.0, 2.0, 4.0, 6.0, 8.0], by: [0.0, 2.0, 4.0, 6.0, 8.0]
[0] local: 30
[1] bx: [200.0, 202.0, 204.0, 206.0, 208.0], by: [20.0, 22.0, 24.0, 26.0, 28.0]
[1] local: 6130
global: 6160
global: 6160

[jedbrown]

Note that --release is required, otherwise I see

invertedPointers:
   invertedPointers[ptr %0] =   <badref> = load ptr, ptr %4, align 8
end invertedPointers
  <badref> = load ptr, ptr %4, align 8
rustc: /home/jed/src/rust-enzyme/src/tools/enzyme/enzyme/Enzyme/GradientUtils.cpp:8489: virtual void InvertedPointerVH::deleted(): Assertion `0 && "erasing something in invertedPointers map"' failed.
error: could not compile `mpi` (example "dot_enzyme")

[ZuseZ4]

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?

[wsmoses]

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?)

[jedbrown]

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?

[ZuseZ4]

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

[wsmoses]

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

Enzyme/enzyme/test/Enzyme/ReverseMode/mpi.ll

Line 26 in 8371758

%call = call i32 @MPI_Send(i8* %1, i32 1, %struct.ompi_datatype_t* bitcast (%struct.ompi_predefined_datatype_t* @ompi_mpi_real to %struct.ompi_datatype_t*), i32 %numprocsuiv, i32 %etiquette, %struct.ompi_communicator_t* bitcast (%struct.ompi_predefined_communicator_t* @ompi_mpi_comm_world to %struct.ompi_communicator_t*)) #4

but ofc with rsmpi (so we're not checking all of the other misc instructions)