Support MPI (#752)

* Register MPI symbols on load

* ops

* Update ext/ReactantMPIExt/Overrides.jl

Co-authored-by: Paul Berg <[email protected]>

* Update ext/ReactantMPIExt/Overrides.jl

* register MPI constants

* Fix MPI specializations

* fix some symbol registration

* refactor MPI Ops

* Add functionality for parsing single operations (Julia code)

* Update `Ops.comm_rank` to use handwritten MLIR injection

* comment

* Update `Ops.comm_size`

* fixes

* Refactor `Ops.barrier`

* Refactor to `try_inject_to_top_block!`

* Refactor MLIR injection

* Refactor MPI constante registration

* Fix type inference in `Ops.hlo_call` on empty args

* Fix MLIR of `Ops.comm_rank`

* Fix MLIR injection C-functions

* Go back to `Cint` for registering symbols

* Add `tryinjectop!`

* Add `tryinject!`, `inject!` methods

* Update `comm_rank`

* Update `mlirOperationInject`, `mlirOperationParse`

* Add `verify` flag to `tryinject!`, `parse(::Operation)`

* Update `Ops.comm_rank`

* Update `comm_rank`, `comm_size`, `barrier`, `wait`

* Implement `Ops.allreduce`

* Implement `Ops.send`

* Remove comment

* Update `Ops.wait`

* Remove `comm` argument from `Comm_size`, `Barrier` overrides

* Fix `Ops.comm_size`

* Fix `Ops.barrier`

* Fixes and renames

* Override `MPI.Allreduce!`

* Fix conversion of MPI constants to word-size type

* Comment unused MPI datatypes

* small fixes

* Implement `MPI.Recv!`

* Test MPI

* Update src/mlir/IR/Operation.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Fix `mpiexec` symbol import

* Fix typo

* Init and Finalize on MPI tests

* Fix changes introduces in "feat: IR inject functions (#1217)"

* return errcode for `send` and `recv!`

* Add tests for Send and Recv!

* Chipping away at isend

* Finish off Ops.isend. Could be working now

* Add default constructor for TracedRequest

* Clean up mpi.jl tests

* TracedRequest must be mutable struct? Also make TracedRequest optional
arg in MPI.Isend

* Scratching together an implementation of make_tracer for TracedRequest

* Additions to TracedRequest and associated functions/uses

* compile_xla/mlir might be working properly with TracedRequest now

* Add create_result for MPI.Request, fix make_tracer

* Comment these out, for now we'll assume no Request needs to cross the
compile boundary

* Add Irecv! and friends. Seems to potentially be working

* Fix overrides Irecv!

* Commit debug testing stuff
I will of course remove these before any PR is submitted, just useful to
have tracked for dev'ing across multiple machines

* Update test-isend-irecv.jl

* Clean up formatting

* Send/Recv altogether works as long as message size is large enough
If the message size is small (<64KB) Send buffers the message and
doesn't block. This seems to screw things up somehow

* Add TODO

* Add test where Send and Recv! are compiled together in one function

* Update tests, send/recv-compiled-together test now working

* Cleanup testing mess

* Fix typo in Ops.irecv!(), return error code from Wait
Isend/Irecv test still not working

* Update isend/irecv tests

* irecv!: Need to return recvbuf and modify mlir data I think

* Fix typo in MPI wait

* Add launcher.jl

* Don't need this until we need to let a Request cross compile barrier

* Set size of status to max necessary hopefully

* Add Isend/Irecv/Wait unit test

* Cleanup return vals in tests

* No need to retun error codes from Send/Recv! now that we have this
is_pure stuff

* No need to return errorcode from wait either

* Remove set_default_backend("cpu") from mpi.jl - Not sure that's the
right way/place to do it

* Use tryinject instead of inject in wait for MPI_COMM_WORLD

* UPdate tests

* Delete debug files

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Use the already dlopened libmpi

* Remove function convert_julia_type_to_mpi_datatype(), use MPI.Dataype()
instead

* Cleanup/remove some comments

* Cleanup some comments

* set default backend cpu for mpi tests, then set it back

* Add test for MPI Barrier

* Remove comment

* Comment out MPI.Init and MPI.Finalize in Overrides, not implemented yet

* Don't register MPI_ERR_RMA_RANGE so CI green since not defined for
windows

* Update ext/ReactantMPIExt/ReactantMPIExt.jl

Co-authored-by: Avik Pal <[email protected]>

* Revert "Don't register MPI_ERR_RMA_RANGE so CI green since not defined for"

This reverts commit 6feb05d.

* Update ext/ReactantMPIExt/ReactantMPIExt.jl

---------

Co-authored-by: Paul Berg <[email protected]>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: romanlee <[email protected]>
Co-authored-by: Roman Lee <[email protected]>
Co-authored-by: Roman Lee <[email protected]>
Co-authored-by: Mosè Giordano <[email protected]>
Co-authored-by: Roman Lee <[email protected]>
Co-authored-by: Avik Pal <[email protected]>

Author: 7 people