Mpi

@irdl_attr_definition
class OperationType(ParametrizedAttribute, TypeAttribute):
    """
    This type represents the MPI_Op type.

    They are used by the reduction MPI functions
    """

    name = "mpi.operation"

    op_str: StringAttr

class MpiOp:
    """
    A collection of MPI_Op types used for
    """

    MPI_MAX = OperationType(StringAttr("MPI_MAX"))
    MPI_MIN = OperationType(StringAttr("MPI_MIN"))
    MPI_SUM = OperationType(StringAttr("MPI_SUM"))
    MPI_PROD = OperationType(StringAttr("MPI_PROD"))
    MPI_LAND = OperationType(StringAttr("MPI_LAND"))
    MPI_BAND = OperationType(StringAttr("MPI_BAND"))
    MPI_LOR = OperationType(StringAttr("MPI_LOR"))
    MPI_BOR = OperationType(StringAttr("MPI_BOR"))
    MPI_LXOR = OperationType(StringAttr("MPI_LXOR"))
    MPI_BXOR = OperationType(StringAttr("MPI_BXOR"))
    MPI_MINLOC = OperationType(StringAttr("MPI_MINLOC"))
    MPI_MAXLOC = OperationType(StringAttr("MPI_MAXLOC"))
    MPI_REPLACE = OperationType(StringAttr("MPI_REPLACE"))
    MPI_NO_OP = OperationType(StringAttr("MPI_NO_OP"))

@irdl_attr_definition
class RequestType(ParametrizedAttribute, TypeAttribute):
    """
    This type represents the MPI_Request type.

    They are used by the asynchronous MPI functions
    """

    name = "mpi.request"

@irdl_attr_definition
class StatusType(ParametrizedAttribute, TypeAttribute):
    """
    This type represents the MPI_Status type.

    It's a struct containing status information for requests.
    """

    name = "mpi.status"

@irdl_attr_definition
class DataType(ParametrizedAttribute, TypeAttribute):
    """
    This type represents MPI_Datatype
    """

    name = "mpi.datatype"

@irdl_attr_definition
class VectorType(ParametrizedAttribute, TypeAttribute, Generic[_VectorT]):
    """
    This type holds multiple MPI types
    """

    name = "mpi.vector"
    wrapped_type: _VectorT

class StatusTypeField(Enum):
    """
    This enum lists all fields in the MPI_Status struct
    """

    MPI_SOURCE = "MPI_SOURCE"
    MPI_TAG = "MPI_TAG"
    MPI_ERROR = "MPI_ERROR"

class MPIBaseOp(IRDLOperation, ABC):
    """
    Base class for MPI Operations
    """

    pass

@irdl_op_definition
class ReduceOp(MPIBaseOp):
    """
    This wraps the MPI_Reduce function (blocking reduction)

    See external [documentation](https://www.mpich.org/static/docs/v4.1/www3/MPI_Reduce.html).

    ## The MPI_Reduce Function Docs:

    int MPI_Reduce(const void *sendbuf, void *recvbuf, int count,
               MPI_Datatype datatype, MPI_Op op, int root, MPI_Comm comm)

        sendbuf: address of send buffer (choice)
        recvbuf: address of receive buffer (choice)
        count: number of elements in send buffer (non-negative integer)
        datatype: data type of elements of send buffer (handle)
        op: reduce operation (handle)
        root: rank of root process (integer)
        comm: communicator (handle)

    ## Our Abstraction:

        - We omit the possibility of using multiple communicators, defaulting
          to MPI_COMM_WORLD
    """

    name = "mpi.reduce"

    send_buffer = operand_def(Attribute)
    recv_buffer = operand_def(Attribute)
    count = operand_def(i32)
    datatype = operand_def(DataType)
    operationtype = attr_def(OperationType)
    root = operand_def(i32)

    def __init__(
        self,
        send_buffer: SSAValue | Operation,
        recv_buffer: SSAValue | Operation,
        count: SSAValue | Operation,
        datatype: SSAValue | Operation,
        operationtype: OperationType,
        root: SSAValue | Operation,
    ):
        return super().__init__(
            operands=[send_buffer, recv_buffer, count, datatype, root],
            attributes={"operationtype": operationtype},
            result_types=[],
        )

@irdl_op_definition
class AllreduceOp(MPIBaseOp):
    """
    This wraps the MPI_Allreduce function (blocking all reduction)

    See external [documentation](https://www.mpich.org/static/docs/v4.1/www3/MPI_Allreduce.html).

    ## The MPI_Allreduce Function Docs:

    int MPI_Allreduce(const void *sendbuf, void *recvbuf, int count,
                  MPI_Datatype datatype, MPI_Op op, MPI_Comm comm)

        sendbuf: address of send buffer (choice)
        recvbuf: address of receive buffer (choice)
        count: number of elements in send buffer (non-negative integer)
        datatype: data type of elements of send buffer (handle)
        op: reduce operation (handle)
        comm: communicator (handle)

    ## Our Abstraction:

        - We omit the possibility of using multiple communicators, defaulting
          to MPI_COMM_WORLD
    """

    name = "mpi.allreduce"

    send_buffer = opt_operand_def(Attribute)
    recv_buffer = operand_def(Attribute)
    count = operand_def(i32)
    datatype = operand_def(DataType)
    operationtype = attr_def(OperationType)

    def __init__(
        self,
        send_buffer: SSAValue | Operation | None,
        recv_buffer: SSAValue | Operation,
        count: SSAValue | Operation,
        datatype: SSAValue | Operation,
        operationtype: OperationType,
    ):
        operands_to_add: Sequence[
            SSAValue | Operation | Sequence[SSAValue | Operation]
        ] = []

        if send_buffer is None:
            operands_to_add = [[], recv_buffer, count, datatype]
        else:
            operands_to_add = [[send_buffer], recv_buffer, count, datatype]

        return super().__init__(
            operands=operands_to_add,
            attributes={"operationtype": operationtype},
            result_types=[],
        )

@irdl_op_definition
class BcastOp(MPIBaseOp):
    """
    This wraps the MPI_Bcast function (blocking broadcast)

    See external [documentation](https://www.mpich.org/static/docs/v4.1/www3/MPI_Bcast.html).

    ## The MPI_Bcast Function Docs:

    int MPI_Bcast(void *buffer, int count, MPI_Datatype datatype, int root,
              MPI_Comm comm)

        buffer: starting address of buffer (choice)
        count: number of elements in send buffer (non-negative integer)
        datatype: data type of elements of send buffer (handle)
        root: rank of broadcast root (integer)
        comm: communicator (handle)

    ## Our Abstraction:

        - We omit the possibility of using multiple communicators, defaulting
          to MPI_COMM_WORLD
    """

    name = "mpi.bcast"

    buffer = operand_def(Attribute)
    count = operand_def(i32)
    datatype = operand_def(DataType)
    root = operand_def(i32)

    def __init__(
        self,
        buffer: SSAValue | Operation,
        count: SSAValue | Operation,
        datatype: SSAValue | Operation,
        root: SSAValue | Operation,
    ):
        return super().__init__(
            operands=[buffer, count, datatype, root],
            result_types=[],
        )

@irdl_op_definition
class IsendOp(MPIBaseOp):
    """
    This wraps the MPI_Isend function (nonblocking send)

    See external [documentation](https://www.mpich.org/static/docs/v4.1/www3/MPI_Isend.html).

    ## The MPI_Isend Function Docs:

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

        - buf: Initial address of send buffer (choice).
        - count: Number of elements in send buffer (integer).
        - datatype: Datatype of each send buffer element (handle).
        - dest: Rank of destination (integer).
        - tag: Message tag (integer).
        - comm: Communicator (handle).

    ## Our Abstraction:

        - We omit the possibility of using multiple communicators, defaulting
          to MPI_COMM_WORLD
    """

    name = "mpi.isend"

    buffer = operand_def(Attribute)
    count = operand_def(i32)
    datatype = operand_def(DataType)
    dest = operand_def(i32)
    tag = operand_def(i32)
    request = operand_def(RequestType)

    def __init__(
        self,
        buffer: SSAValue | Operation,
        count: SSAValue | Operation,
        datatype: SSAValue | Operation,
        dest: SSAValue | Operation,
        tag: SSAValue | Operation,
        request: SSAValue | Operation,
    ):
        return super().__init__(
            operands=[buffer, count, datatype, dest, tag, request],
            result_types=[],
        )

@irdl_op_definition
class SendOp(MPIBaseOp):
    """
    This wraps the MPI_Send function (blocking send)

    See external [documentation](https://www.mpich.org/static/docs/v4.1/www3/MPI_Send.html).

    ## The MPI_Send Function Docs:

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

        - buf: Initial address of send buffer (choice).
        - count: Number of elements in send buffer (non-negative integer).
        - datatype: Datatype of each send buffer element (handle).
        - dest: Rank of destination (integer).
        - tag: Message tag (integer).
        - comm: Communicator (handle).

    ## Our Abstraction:

        - We omit the possibility of using multiple communicators, defaulting
          to MPI_COMM_WORLD
    """

    name = "mpi.send"

    buffer = operand_def(Attribute)
    count = operand_def(i32)
    datatype = operand_def(DataType)
    dest = operand_def(i32)
    tag = operand_def(i32)

    def __init__(
        self,
        buffer: SSAValue | Operation,
        count: SSAValue | Operation,
        datatype: SSAValue | Operation,
        dest: SSAValue | Operation,
        tag: SSAValue | Operation,
    ):
        return super().__init__(
            operands=[buffer, count, datatype, dest, tag], result_types=[]
        )

@irdl_op_definition
class IrecvOp(MPIBaseOp):
    """
    This wraps the MPI_Irecv function (nonblocking receive).

    See external [documentation](https://www.mpich.org/static/docs/v4.1/www3/MPI_Irecv.html).

    ## The MPI_Irecv Function Docs:

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

        - buf: Initial address of receive buffer (choice).
        - count: Number of elements in receive buffer (integer).
        - datatype: Datatype of each receive buffer element (handle).
        - source: Rank of source (integer).
        - tag: Message tag (integer).
        - comm: Communicator (handle).
        - request: Communication request (handle).

    ## Our Abstractions:

        - We omit the possibility of using multiple communicators, defaulting
          to MPI_COMM_WORLD
    """

    name = "mpi.irecv"

    buffer = operand_def(Attribute)
    count = operand_def(i32)
    datatype = operand_def(DataType)
    source = operand_def(i32)
    tag = operand_def(i32)
    request = operand_def(RequestType)

    def __init__(
        self,
        buffer: SSAValue | Operation,
        count: SSAValue | Operation,
        datatype: SSAValue | Operation,
        source: SSAValue | Operation,
        tag: SSAValue | Operation,
        request: SSAValue | Operation,
    ):
        return super().__init__(
            operands=[buffer, count, datatype, source, tag, request],
            result_types=[],
        )

@irdl_op_definition
class RecvOp(MPIBaseOp):
    """
    This wraps the MPI_Recv function (blocking receive).

    See external [documentation](https://www.mpich.org/static/docs/v4.1/www3/MPI_Recv.html).

    ## The MPI_Recv Function Docs:

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

        - buf: Initial address of receive buffer (choice).
        - count: Number of elements in receive buffer (integer).
        - datatype: Datatype of each receive buffer element (handle).
        - source: Rank of source (integer).
        - tag: Message tag (integer).
        - comm: Communicator (handle).
        - status: status object (Status).

    ## Our Abstractions:

        - We omit the possibility of using multiple communicators, defaulting
          to MPI_COMM_WORLD
    """

    name = "mpi.recv"

    buffer = operand_def(Attribute)
    count = operand_def(i32)
    datatype = operand_def(DataType)
    source = operand_def(i32)
    tag = operand_def(i32)

    status = opt_result_def(StatusType)

    def __init__(
        self,
        buffer: SSAValue | Operation,
        count: SSAValue | Operation,
        datatype: SSAValue | Operation,
        source: SSAValue | Operation,
        tag: SSAValue | Operation,
        ignore_status: bool = True,
    ):
        return super().__init__(
            operands=[buffer, count, datatype, source, tag],
            result_types=[[]] if ignore_status else [[StatusType()]],
        )

@irdl_op_definition
class TestOp(MPIBaseOp):
    """
    Class for wrapping the MPI_Test function (test for completion of request)

    See external [documentation](https://www.mpich.org/static/docs/v4.1/www3/MPI_Test.html).

    ## The MPI_Test Function Docs:

    int MPI_Test(MPI_Request *request, int *flag, MPI_Status *status)

        - request: Communication request (handle)
        - flag: true if operation completed (logical)
        - status: Status object (Status)
    """

    name = "mpi.test"

    request = operand_def(RequestType)

    flag = result_def(t_bool)
    status = result_def(StatusType)

    def __init__(self, request: Operand):
        return super().__init__(operands=[request], result_types=[t_bool, StatusType()])

@irdl_op_definition
class WaitOp(MPIBaseOp):
    """
    Class for wrapping the MPI_Wait function (blocking wait for request)

    See external [documentation](https://www.mpich.org/static/docs/v4.1/www3/MPI_Wait.html).

    ## The MPI_Test Function Docs:

    int MPI_Wait(MPI_Request *request, MPI_Status *status)

        - request: Request (handle)
        - status: Status object (Status)
    """

    name = "mpi.wait"

    request = operand_def(RequestType)
    status = opt_result_def(StatusType)

    def __init__(self, request: Operand, ignore_status: bool = True):
        result_types: list[list[Attribute]] = [[StatusType()]]
        if ignore_status:
            result_types = [[]]

        return super().__init__(operands=[request], result_types=result_types)

@irdl_op_definition
class WaitallOp(MPIBaseOp):
    """
    Class for wrapping the MPI_Waitall function (blocking wait for requests)

    See external [documentation](https://www.mpich.org/static/docs/v4.1/www3/MPI_Waitall.html).

    ## The MPI_Test Function Docs:

    int MPI_Waitall(int count, MPI_Request array_of_requests[],
                MPI_Status *array_of_statuses)

        - count: Number of handles
        - array_of_requests: Request handles
        - array_of_statuses: Status objects
    """

    name = "mpi.waitall"

    requests = operand_def(VectorType[RequestType])
    count = operand_def(i32)
    statuses = opt_result_def(VectorType[StatusType])

    def __init__(self, requests: Operand, count: Operand, ignore_status: bool = True):
        result_types: list[list[Attribute]] = [[VectorType(StatusType())]]
        if ignore_status:
            result_types = [[]]

        return super().__init__(operands=[requests, count], result_types=result_types)

@irdl_op_definition
class GetStatusFieldOp(MPIBaseOp):
    """
    Accessors for the MPI_Status struct

    This allows access to the three integer properties in the struct called
        - MPI_SOURCE
        - MPI_TAG
        - MPI_ERROR

    All fields are of type int.
    """

    name = "mpi.status.get"

    status = operand_def(StatusType)

    field = attr_def(StringAttr)

    result = result_def(i32)

    def __init__(self, status_obj: Operand, field: StatusTypeField):
        return super().__init__(
            operands=[status_obj],
            attributes={"field": StringAttr(field.value)},
            result_types=[i32],
        )

@irdl_op_definition
class CommRankOp(MPIBaseOp):
    """
    Represents the MPI_Comm_size(MPI_Comm comm, int *rank) function call which returns
    the rank of the communicator

    Currently limited to COMM_WORLD
    """

    name = "mpi.comm.rank"

    rank = result_def(i32)

    def __init__(self):
        return super().__init__(result_types=[i32])

@irdl_op_definition
class CommSizeOp(MPIBaseOp):
    """
    Represents the MPI_Comm_size(MPI_Comm comm, int *size) function call which returns
    the size of the communicator

    Currently limited to COMM_WORLD
    """

    name = "mpi.comm.size"

    size = result_def(i32)

    def __init__(self):
        return super().__init__(result_types=[i32])

@irdl_op_definition
class InitOp(MPIBaseOp):
    """
    This represents a bare MPI_Init call with both args being nullptr
    """

    name = "mpi.init"

@irdl_op_definition
class FinalizeOp(MPIBaseOp):
    """
    This represents an MPI_Finalize call with both args being nullptr
    """

    name = "mpi.finalize"

@irdl_op_definition
class UnwrapMemRefOp(MPIBaseOp):
    """
    This Op can be used as a helper to get memrefs into MPI calls.

    It takes any MemRef as input, and returns an llvm.ptr, element count and MPI_Datatype.
    """

    name = "mpi.unwrap_memref"

    ref = operand_def(MemRefType[AnyNumericType])

    ptr = result_def(llvm.LLVMPointerType)
    len = result_def(i32)
    type = result_def(DataType)

    def __init__(self, ref: SSAValue | Operation):
        return super().__init__(
            operands=[ref],
            result_types=[llvm.LLVMPointerType(), i32, DataType()],
        )