Lower mpi

@dataclass(frozen=True)
class MpiLibraryInfo:
    """
    This object is meant to capture characteristics of a specific MPI implementations.

    It holds magic values, sizes of structs, field offsets and much more.

    We need these as we currently cannot load these library headers into the programs we want to lower,
    therefore we need to generate our own external stubs and load magic values directly.

    This way of doing it is inherently fragile, but we don't know of any better way.
    We plan to include a C file that automagically extracts all this information from MPI headers.
    You can see the current C file used in this PR: https://github.com/xdslproject/xdsl/pull/526
    You can see the status of OpenMPI support here: https://github.com/xdslproject/xdsl/issues/523

    These defaults have been extracted from MPICH 3.3a2. We would highly suggest
    running the mpi-info.c file yourself with your version of the library!
    """

    # MPI_Datatype
    MPI_Datatype_size: int = 4
    MPI_CHAR: int = 0x4C000101
    MPI_SIGNED_CHAR: int = 0x4C000118
    MPI_UNSIGNED_CHAR: int = 0x4C000102
    MPI_BYTE: int = 0x4C00010D
    MPI_WCHAR: int = 0x4C00040E
    MPI_SHORT: int = 0x4C000203
    MPI_UNSIGNED_SHORT: int = 0x4C000204
    MPI_INT: int = 0x4C000405
    MPI_UNSIGNED: int = 0x4C000406
    MPI_LONG: int = 0x4C000807
    MPI_UNSIGNED_LONG: int = 0x4C000808
    MPI_FLOAT: int = 0x4C00040A
    MPI_DOUBLE: int = 0x4C00080B
    MPI_LONG_DOUBLE: int = 0x4C00100C
    MPI_LONG_LONG_INT: int = 0x4C000809
    MPI_UNSIGNED_LONG_LONG: int = 0x4C000819
    MPI_LONG_LONG: int = 0x4C000809

    # MPI_Op
    MPI_Op_size: int = 4
    MPI_MAX: int = 0x58000001
    MPI_MIN: int = 0x58000002
    MPI_SUM: int = 0x58000003
    MPI_PROD: int = 0x58000004
    MPI_LAND: int = 0x58000005
    MPI_BAND: int = 0x58000006
    MPI_LOR: int = 0x58000007
    MPI_BOR: int = 0x58000008
    MPI_LXOR: int = 0x58000009
    MPI_BXOR: int = 0x5800000A
    MPI_MINLOC: int = 0x5800000B
    MPI_MAXLOC: int = 0x5800000C
    MPI_REPLACE: int = 0x5800000D
    MPI_NO_OP: int = 0x5800000E

    # MPI_Comm
    MPI_Comm_size: int = 4
    MPI_COMM_WORLD: int = 0x44000000
    MPI_COMM_SELF: int = 0x44000001

    # MPI_Request
    MPI_Request_size: int = 4
    MPI_REQUEST_NULL = 0x2C000000

    # MPI_Status
    MPI_Status_size: int = 20
    MPI_STATUS_IGNORE: int = 0x00000001
    MPI_STATUSES_IGNORE: int = 0x00000001
    MPI_Status_field_MPI_SOURCE: int = (
        8  # offset of field MPI_SOURCE in struct MPI_Status
    )
    MPI_Status_field_MPI_TAG: int = 12  # offset of field MPI_TAG in struct MPI_Status
    MPI_Status_field_MPI_ERROR: int = (
        16  # offset of field MPI_ERROR in struct MPI_Status
    )

    # In place MPI All reduce
    MPI_IN_PLACE: int = -1

class LowerMpiInit(_MPIToLLVMRewriteBase):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: mpi.InitOp, rewriter: PatternRewriter, /):
        rewriter.replace(op, *self.lower(op))

    def lower(self, op: mpi.InitOp) -> tuple[list[Operation], list[SSAValue | None]]:
        """
        We currently don't model any argument passing to `MPI_Init()` and pass two nullptrs.
        """
        return [
            nullptr := llvm.ZeroOp(result_types=[llvm.LLVMPointerType()]),
            func.CallOp(self._mpi_name(op), [nullptr, nullptr], [i32]),
        ], []

class LowerMpiFinalize(_MPIToLLVMRewriteBase):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: mpi.FinalizeOp, rewriter: PatternRewriter, /):
        rewriter.replace(op, *self.lower(op))

    def lower(
        self, op: mpi.FinalizeOp
    ) -> tuple[list[Operation], list[SSAValue | None]]:
        """
        Relatively straight forward lowering of mpi.finalize operation.
        """
        return [
            func.CallOp(self._mpi_name(op), [], [i32]),
        ], []

class LowerMpiWait(_MPIToLLVMRewriteBase):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: mpi.WaitOp, rewriter: PatternRewriter, /):
        rewriter.replace(op, *self.lower(op))

    def lower(self, op: mpi.WaitOp) -> tuple[list[Operation], list[SSAValue | None]]:
        """
        Relatively straight forward lowering of mpi.wait operation.
        """
        ops, new_results, res = self._emit_mpi_status_objs(len(op.results))
        return [
            *ops,
            func.CallOp(self._mpi_name(op), [op.request, res], [i32]),
        ], new_results

class LowerMpiWaitall(_MPIToLLVMRewriteBase):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: mpi.WaitallOp, rewriter: PatternRewriter, /):
        rewriter.replace(op, *self.lower(op))

    def lower(self, op: mpi.WaitallOp) -> tuple[list[Operation], list[SSAValue | None]]:
        """
        Relatively straight forward lowering of mpi.waitall operation.
        """

        ops, new_results, res = self._emit_mpi_status_objs(len(op.results))
        return [
            *ops,
            func.CallOp(self._mpi_name(op), [op.count, op.requests, res], [i32]),
        ], new_results

class LowerMpiReduce(_MPIToLLVMRewriteBase):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: mpi.ReduceOp, rewriter: PatternRewriter, /):
        rewriter.replace(op, *self.lower(op))

    def lower(self, op: mpi.ReduceOp) -> tuple[list[Operation], list[SSAValue | None]]:
        """
        Lowers the MPI Reduce operation
        """

        return [
            comm_global := arith.ConstantOp.from_int_and_width(
                self.info.MPI_COMM_WORLD, i32
            ),
            mpi_op := self._emit_mpi_operation_load(op.operationtype),
            func.CallOp(
                self._mpi_name(op),
                [
                    op.send_buffer,
                    op.recv_buffer,
                    op.count,
                    op.datatype,
                    mpi_op,
                    op.root,
                    comm_global,
                ],
                [],
            ),
        ], []

class LowerMpiAllreduce(_MPIToLLVMRewriteBase):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: mpi.AllreduceOp, rewriter: PatternRewriter, /):
        rewriter.replace(op, *self.lower(op))

    def lower(
        self, op: mpi.AllreduceOp
    ) -> tuple[list[Operation], list[SSAValue | None]]:
        """
        Lowers the MPI Allreduce operation
        """

        # Send buffer is optional (if not provided then call using MPI_IN_PLACE)
        has_send_buffer = op.send_buffer is not None

        comm_global = arith.ConstantOp.from_int_and_width(self.info.MPI_COMM_WORLD, i32)
        mpi_op = self._emit_mpi_operation_load(op.operationtype)

        operations = [comm_global, mpi_op]

        send_buffer_op: SSAValue | Operation
        if has_send_buffer:
            assert op.send_buffer is not None
            send_buffer_op = op.send_buffer
        else:
            send_buffer_op = arith.ConstantOp.from_int_and_width(
                self.info.MPI_IN_PLACE, i64
            )
            operations.append(send_buffer_op)

        return [
            *operations,
            func.CallOp(
                self._mpi_name(op),
                [
                    send_buffer_op,
                    op.recv_buffer,
                    op.count,
                    op.datatype,
                    mpi_op,
                    comm_global,
                ],
                [],
            ),
        ], []

class LowerMpiBcast(_MPIToLLVMRewriteBase):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: mpi.BcastOp, rewriter: PatternRewriter, /):
        rewriter.replace(op, *self.lower(op))

    def lower(self, op: mpi.BcastOp) -> tuple[list[Operation], list[SSAValue | None]]:
        """
        Lowers the MPI Bcast operation
        """

        return [
            comm_global := arith.ConstantOp.from_int_and_width(
                self.info.MPI_COMM_WORLD, i32
            ),
            func.CallOp(
                self._mpi_name(op),
                [op.buffer, op.count, op.datatype, op.root, comm_global],
                [],
            ),
        ], []

class LowerMpiIsend(_MPIToLLVMRewriteBase):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: mpi.IsendOp, rewriter: PatternRewriter, /):
        rewriter.replace(op, *self.lower(op))

    def lower(self, op: mpi.IsendOp) -> tuple[list[Operation], list[SSAValue | None]]:
        """
        This method lowers mpi.isend

        int MPI_Isend(const void *buf, int count, MPI_Datatype datatype, int dest,
              int tag, MPI_Comm comm, MPI_Request *request)
        """

        return [
            comm_global := arith.ConstantOp.from_int_and_width(
                self.info.MPI_COMM_WORLD, i32
            ),
            func.CallOp(
                self._mpi_name(op),
                [
                    op.buffer,
                    op.count,
                    op.datatype,
                    op.dest,
                    op.tag,
                    comm_global,
                    op.request,
                ],
                [i32],
            ),
        ], []

class LowerMpiIrecv(_MPIToLLVMRewriteBase):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: mpi.IrecvOp, rewriter: PatternRewriter, /):
        rewriter.replace(op, *self.lower(op))

    def lower(self, op: mpi.IrecvOp) -> tuple[list[Operation], list[SSAValue | None]]:
        """
        This method lowers mpi.irecv operations

        int MPI_Irecv(void *buf, int count, MPI_Datatype datatype, int source, int tag,
              MPI_Comm comm, MPI_Request *request)
        """

        return [
            comm_global := arith.ConstantOp.from_int_and_width(
                self.info.MPI_COMM_WORLD, i32
            ),
            func.CallOp(
                self._mpi_name(op),
                [
                    op.buffer,
                    op.count,
                    op.datatype,
                    op.source,
                    op.tag,
                    comm_global,
                    op.request,
                ],
                [i32],
            ),
        ], []

class LowerMpiSend(_MPIToLLVMRewriteBase):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: mpi.SendOp, rewriter: PatternRewriter, /):
        rewriter.replace(op, *self.lower(op))

    def lower(self, op: mpi.SendOp) -> tuple[list[Operation], list[SSAValue | None]]:
        """
        This method lowers mpi.send operations

        MPI_Send signature:

        int MPI_Send(const void *buf, int count, MPI_Datatype datatype, int dest,
                 int tag, MPI_Comm comm)
        """

        return [
            comm_global := arith.ConstantOp.from_int_and_width(
                self.info.MPI_COMM_WORLD, i32
            ),
            func.CallOp(
                self._mpi_name(op),
                [op.buffer, op.count, op.datatype, op.dest, op.tag, comm_global],
                [i32],
            ),
        ], []

class LowerMpiRecv(_MPIToLLVMRewriteBase):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: mpi.RecvOp, rewriter: PatternRewriter, /):
        rewriter.replace(op, *self.lower(op))

    def lower(self, op: mpi.RecvOp) -> tuple[list[Operation], list[SSAValue | None]]:
        """
        This method lowers mpi.recv operations

        MPI_Recv signature:

        int MPI_Recv(void *buf, int count, MPI_Datatype datatype, int source, int tag,
             MPI_Comm comm, MPI_Status *status)
        """

        mpi_status_ops, new_results, status = self._emit_mpi_status_objs(
            len(op.results)
        )

        return [
            *mpi_status_ops,
            comm_global := arith.ConstantOp.from_int_and_width(
                self.info.MPI_COMM_WORLD, i32
            ),
            func.CallOp(
                self._mpi_name(op),
                [
                    op.buffer,
                    op.count,
                    op.datatype,
                    op.source,
                    op.tag,
                    comm_global,
                    status,
                ],
                [i32],
            ),
        ], new_results

class LowerMpiUnwrapMemRefOp(_MPIToLLVMRewriteBase):
    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: mpi.UnwrapMemRefOp, rewriter: PatternRewriter, /):
        rewriter.replace(op, *self.lower(op))

    def lower(
        self, op: mpi.UnwrapMemRefOp
    ) -> tuple[list[Operation], list[SSAValue | None]]:
        count_ops, count_ssa_val = self._emit_memref_counts(op.ref)
        extract_ptr_ops, ptr = self._memref_get_llvm_ptr(op.ref)

        elem_type = cast(MemRefType[mpi.AnyNumericType], op.ref.type).element_type

        return [
            *extract_ptr_ops,
            *count_ops,
            dtype := mpi.GetDtypeOp(elem_type),
        ], [ptr.results[0], count_ssa_val, dtype.result]