[Mlir-commits] [mlir] [MLIR] Extend MPI dialect (PR #123255)

[Tobias Grosser]

Sergio =?utf-8?q?Sánchez_Ramírez?=,
Sergio =?utf-8?q?Sánchez_Ramírez?=,
Sergio =?utf-8?q?Sánchez_Ramírez?=,
Sergio =?utf-8?q?Sánchez_Ramírez?=,
Sergio =?utf-8?q?Sánchez_Ramírez?=,
Sergio =?utf-8?q?Sánchez_Ramírez?=,
Sergio =?utf-8?q?Sánchez_Ramírez?=,
Sergio =?utf-8?q?Sánchez_Ramírez?=,
Sergio =?utf-8?q?Sánchez_Ramírez?=,
Sergio =?utf-8?q?Sánchez_Ramírez?=,
Sergio =?utf-8?q?Sánchez_Ramírez?=,
Sergio =?utf-8?q?Sánchez_Ramírez?=,
Sergio =?utf-8?q?Sánchez_Ramírez?=,
Sergio =?utf-8?q?Sánchez_Ramírez?=,
Sergio =?utf-8?q?Sánchez_Ramírez?=,
Sergio =?utf-8?q?Sánchez_Ramírez?=,
Sergio =?utf-8?q?Sánchez_Ramírez?=,
Sergio =?utf-8?q?Sánchez_Ramírez?Message-ID:
In-Reply-To: <llvm.org/llvm/llvm-project/pull/123255 at github.com>


================
@@ -65,59 +118,214 @@ def MPI_CommRankOp : MPI_Op<"comm_rank", []> {
 
 def MPI_SendOp : MPI_Op<"send", []> {
   let summary =
-      "Equivalent to `MPI_Send(ptr, size, dtype, dest, tag, MPI_COMM_WORLD)`";
+      "Equivalent to `MPI_Send(ptr, size, dtype, dest, tag, comm)`";
   let description = [{
     MPI_Send performs a 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.
+    If communicator is not specified, `MPI_COMM_WORLD` is used by default.
 
     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 arguments = (
+    ins AnyMemRef : $ref,
+    I32 : $tag,
+    I32 : $rank,
+    Optional<MPI_Comm> : $comm
+  );
 
   let results = (outs Optional<MPI_Retval>:$retval);
 
-  let assemblyFormat = "`(` $ref `,` $tag `,` $rank `)` attr-dict `:` "
+  let assemblyFormat = "`(` $ref `,` $tag `,` $rank (`,` $comm)? `)` attr-dict `:` "
                        "type($ref) `,` type($tag) `,` type($rank)"
                        "(`->` type($retval)^)?";
   let hasCanonicalizer = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// ISendOp
+//===----------------------------------------------------------------------===//
+
+def MPI_ISendOp : MPI_Op<"isend", []> {
+  let summary =
+      "Equivalent to `MPI_Isend(ptr, size, dtype, dest, tag, comm)`";
+  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.
+
+    If communicator is not specified, `MPI_COMM_WORLD` is used by default.
+
+    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,
+    Optional<MPI_Comm> : $comm
+  );
+
+  let results = (outs Optional<MPI_Retval>:$retval, MPI_Request : $req);
+
+  let assemblyFormat = "`(` $ref `,` $tag `,` $rank (`,` $comm)?`)` attr-dict "
+                       "`:` type($ref) `,` type($tag) `,` type($rank) "
+                       "(`,` type($comm))? `->` (type($retval) `,` ^)? type($req)";
+  let hasCanonicalizer = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // RecvOp
 //===----------------------------------------------------------------------===//
 
 def MPI_RecvOp : MPI_Op<"recv", []> {
   let summary = "Equivalent to `MPI_Recv(ptr, size, dtype, dest, tag, "
-                "MPI_COMM_WORLD, MPI_STATUS_IGNORE)`";
+                "comm, MPI_STATUS_IGNORE)`";
   let description = [{
     MPI_Recv performs a 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.
+    If communicator is not specified, `MPI_COMM_WORLD` is used by default.
     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 arguments = (
+    ins AnyMemRef : $ref,
+    I32 : $tag, I32 : $rank,
+    Optional<MPI_Comm> : $comm
+  );
 
   let results = (outs Optional<MPI_Retval>:$retval);
 
-  let assemblyFormat = "`(` $ref `,` $tag `,` $rank `)` attr-dict `:` "
-                       "type($ref) `,` type($tag) `,` type($rank)"
-                       "(`->` type($retval)^)?";
+  let assemblyFormat = "`(` $ref `,` $tag `,` $rank (`,` $comm)?`)` attr-dict "
+                       "`:` type($ref) `,` type($tag) `,` type($rank) "
+                       "(`,` type($comm))? (`->` type($retval)^)?";
+  let hasCanonicalizer = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// IRecvOp
+//===----------------------------------------------------------------------===//
+
+def MPI_IRecvOp : MPI_Op<"irecv", []> {
+  let summary = "Equivalent to `MPI_Irecv(ptr, size, dtype, dest, tag, "
+                "comm, &req)`";
+  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.
+
+    If communicator is not specified, `MPI_COMM_WORLD` is used by default.
+
+    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,
+    Optional<MPI_Comm> : $comm
+  );
+
+  let results = (outs Optional<MPI_Retval>:$retval, MPI_Request : $req);
+
+  let assemblyFormat = "`(` $ref `,` $tag `,` $rank (`,` $comm)?`)` attr-dict "
+                       "`:` type($ref) `,` type($tag) `,` type($rank)"
+                       "(`,` type($comm))? `->` (type($retval) `,` ^)? type($req)";
   let hasCanonicalizer = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// AllReduceOp
+//===----------------------------------------------------------------------===//
+
+def MPI_AllReduceOp : MPI_Op<"allreduce", []> {
+  let summary = "Equivalent to `MPI_Allreduce(sendbuf, recvbuf, op, comm)`";
+  let description = [{
+    MPI_Allreduce performs a reduction operation on the values in the sendbuf
+    array and stores the result in the recvbuf array. The operation is 
+    performed across all processes in the communicator.
+
+    If communicator is not specified, `MPI_COMM_WORLD` is used by default.
+
+    This operation can optionally return an `!mpi.retval` value that can be used
+    to check for errors.
+  }];
+
+  let arguments = (
+    ins AnyMemRef : $sendbuf,
+    AnyMemRef : $recvbuf,
+    Optional<MPI_Comm> : $comm
+  );
+
+  let regions = (region SizedRegion<1>:$op);
----------------
tobiasgrosser wrote:

https://mlir.llvm.org/docs/Traits/#isolatedfromabove

https://github.com/llvm/llvm-project/pull/123255