[MLIR] Extend MPI dialect

mlir/include/mlir/Dialect/MPI/IR/MPIOps.td

@@ -59,6 +59,28 @@ def MPI_CommRankOp : MPI_Op<"comm_rank", []> {
   let assemblyFormat = "attr-dict `:` type(results)";
 }
 
+//===----------------------------------------------------------------------===//
+// CommSizeOp
+//===----------------------------------------------------------------------===//
+
+def MPI_CommSizeOp : MPI_Op<"comm_size", []> {
+  let summary = "Get the size of the group associated to the communicator, equivalent to "
+                "`MPI_Comm_size(MPI_COMM_WORLD, &size)`";
+  let description = [{
+    Communicators other than `MPI_COMM_WORLD` are not supported for now.
+
+    This operation can optionally return an `!mpi.retval` value that can be used
+    to check for errors.
+  }];
+
+  let results = (
+    outs Optional<MPI_Retval> : $retval,
+    I32 : $size
+  );
+
+  let assemblyFormat = "attr-dict `:` type(results)";
+}
+
 //===----------------------------------------------------------------------===//
 // SendOp
 //===----------------------------------------------------------------------===//
@@ -87,6 +109,37 @@ def MPI_SendOp : MPI_Op<"send", []> {
   let hasCanonicalizer = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// ISendOp
+//===----------------------------------------------------------------------===//
+
+// TODO what about request handler?
+// NOTE datatype & count args are implicit by the type of the first argument (i.e. memref of eltype)
+// NOTE other communicators not yet supported by the `mpi` dialect
+def MPI_ISendOp : MPI_Op<"isend", []> {
+  let summary =
+      "Equivalent to `MPI_Isend(ptr, size, dtype, dest, tag, MPI_COMM_WORLD)`";
+  let description = [{
+    MPI_Isend begins a non-blocking send of `size` elements of type `dtype` to rank
+    `dest`. The `tag` value and communicator enables the library to determine 
+    the matching of multiple sends and receives between the same ranks.
+
+    Communicators other than `MPI_COMM_WORLD` are not supprted for now.
+
+    This operation can optionally return an `!mpi.retval` value that can be used
+    to check for errors.
+  }];
+
+  let arguments = (ins AnyMemRef : $ref, I32 : $tag, I32 : $rank);
+
+  let results = (outs Optional<MPI_Retval>:$retval);
+
+  let assemblyFormat = "`(` $ref `,` $tag `,` $rank `)` attr-dict `:` "
+                       "type($ref) `,` type($tag) `,` type($rank)"
+                       "(`->` type($retval)^)?";
+  let hasCanonicalizer = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // RecvOp
 //===----------------------------------------------------------------------===//
@@ -118,6 +171,38 @@ def MPI_RecvOp : MPI_Op<"recv", []> {
   let hasCanonicalizer = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// IRecvOp
+//===----------------------------------------------------------------------===//
+
+// TODO same as MPI_ISendOp
+def MPI_IRecvOp : MPI_Op<"irecv", []> {
+  let summary = "Equivalent to `MPI_Irecv(ptr, size, dtype, dest, tag, "
+                "MPI_COMM_WORLD, MPI_STATUS_IGNORE)`";
+  let description = [{
+    MPI_Irecv begins a non-blocking receive of `size` elements of type `dtype` 
+    from rank `dest`. The `tag` value and communicator enables the library to 
+    determine the matching of multiple sends and receives between the same 
+    ranks.
+
+    Communicators other than `MPI_COMM_WORLD` are not supprted for now.
+    The MPI_Status is set to `MPI_STATUS_IGNORE`, as the status object 
+    is not yet ported to MLIR.
+
+    This operation can optionally return an `!mpi.retval` value that can be used
+    to check for errors.
+  }];
+
+  let arguments = (ins AnyMemRef : $ref, I32 : $tag, I32 : $rank);
+
+  let results = (outs Optional<MPI_Retval>:$retval);
+
+  let assemblyFormat = "`(` $ref `,` $tag `,` $rank `)` attr-dict `:` "
+                       "type($ref) `,` type($tag) `,` type($rank)"
+                       "(`->` type($retval)^)?";
+  let hasCanonicalizer = 1;
+}
+
 
 //===----------------------------------------------------------------------===//
 // FinalizeOp
@@ -166,7 +251,7 @@ def MPI_RetvalCheckOp : MPI_Op<"retval_check", []> {
 
 
 //===----------------------------------------------------------------------===//
-// RetvalCheckOp
+// ErrorClassOp
 //===----------------------------------------------------------------------===//
 
 def MPI_ErrorClassOp : MPI_Op<"error_class", []> {

mlir/include/mlir/Dialect/MPI/IR/MPITypes.td

@@ -40,4 +40,34 @@ def MPI_Retval : MPI_Type<"Retval", "retval"> {
   }];
 }
 
+// TODO
+def MPI_Comm : MPI_Type<"Comm", "comm"> {
+  let summary = "..."
+  let description = [{
+    This type represents a handler to the MPI communicator.
+  }]
+}
+
+// TODO
+def MPI_Request : MPI_Type<"Request", "request"> {
+  let summary = "..."
+  let description = [{
+    This type represents a handler to an asynchronous requests.
+  }]
+}
+
+// TODO
+def MPI_Status : MPI_Type<"Status", "status"> {
+  let summary = "";
+  let description = [{
+  }];
+}
+
+// TODO
+def MPI_Op : MPI_Type<"Op", "op"> {
+  let summary = "";
+  let description = [{
+  }];
+}
+
 #endif // MLIR_DIALECT_MPI_IR_MPITYPES_TD