[Mlir-commits] [mlir] [MLIR] MPI: add `count` optional argument to communication ops (PR #129095)

[Sergio Sánchez Ramírez]

https://github.com/mofeing created https://github.com/llvm/llvm-project/pull/129095

cc @hhkit @wsmoses @JBlaschke

having a `count` input argument that replicates the same name argument of the MPI routines like `MPI_Send` is useful. specially if the size of the `memref` is unknown (e.g. `memref<?xf32>`) or we just want to send a part of the full data.

this is still WIP

@fschlimb @tobiasgrosser @AntonLydike @Dinistro i fear that the custom assembly format is gonna get very difficult to parse. in such case, is it okay if we ditch it temporarily and use the explicit format instead?

>From 1def0dbc676783f998ba375e33fe1dcd7c0a66fc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20S=C3=A1nchez=20Ram=C3=ADrez?=
 <sergio.sanchez.ramirez+git at bsc.es>
Date: Thu, 27 Feb 2025 19:14:13 +0100
Subject: [PATCH] [MLIR] MPI: add `count` optional argument to communication
 ops

---
 mlir/include/mlir/Dialect/MPI/IR/MPIOps.td | 23 +++++++++++++---------
 1 file changed, 14 insertions(+), 9 deletions(-)

diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td b/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
index db28bd09678f8..e29a66a54a6da 100644
--- a/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
@@ -101,14 +101,15 @@ def MPI_SendOp : MPI_Op<"send", []> {
 
   let arguments = (
     ins AnyMemRef : $ref,
+    Optional<I32> : $count,
     I32 : $tag,
     I32 : $dest
   );
 
   let results = (outs Optional<MPI_Retval>:$retval);
 
-  let assemblyFormat = "`(` $ref `,` $tag `,` $dest `)` attr-dict `:` "
-                       "type($ref) `,` type($tag) `,` type($dest)"
+  let assemblyFormat = "`(` $ref (`,` $count^?) `,` $tag `,` $dest `)` attr-dict `:` "
+                       "type($ref) (`,` type($count)^?) `,` type($tag) `,` type($dest)"
                        "(`->` type($retval)^)?";
   let hasCanonicalizer = 1;
 }
@@ -134,6 +135,7 @@ def MPI_ISendOp : MPI_Op<"isend", []> {
 
   let arguments = (
     ins AnyMemRef : $ref,
+    Optional<I32> : $count,
     I32 : $tag,
     I32 : $rank
   );
@@ -143,8 +145,8 @@ def MPI_ISendOp : MPI_Op<"isend", []> {
     MPI_Request : $req
   );
 
-  let assemblyFormat = "`(` $ref `,` $tag `,` $rank `)` attr-dict "
-                       "`:` type($ref) `,` type($tag) `,` type($rank) "
+  let assemblyFormat = "`(` $ref (`,` $count^?) `,` $tag `,` $rank `)` attr-dict "
+                       "`:` type($ref) (`,` type($count)^?) `,` type($tag) `,` type($rank) "
                        "`->` type(results)";
   let hasCanonicalizer = 1;
 }
@@ -172,13 +174,14 @@ def MPI_RecvOp : MPI_Op<"recv", []> {
 
   let arguments = (
     ins AnyMemRef : $ref,
+    Optional<I32> : $count,
     I32 : $tag, I32 : $source
   );
 
   let results = (outs Optional<MPI_Retval>:$retval);
 
-  let assemblyFormat = "`(` $ref `,` $tag `,` $source `)` attr-dict `:` "
-                       "type($ref) `,` type($tag) `,` type($source)"
+  let assemblyFormat = "`(` $ref (`,` $count^?) `,` $tag `,` $source `)` attr-dict `:` "
+                       "type($ref) (`,` type($count)^?) `,` type($tag) `,` type($source)"
                        "(`->` type($retval)^)?";
   let hasCanonicalizer = 1;
 }
@@ -204,6 +207,7 @@ def MPI_IRecvOp : MPI_Op<"irecv", []> {
 
   let arguments = (
     ins AnyMemRef : $ref,
+    Optional<I32> : $count,
     I32 : $tag,
     I32 : $rank
   );
@@ -213,8 +217,8 @@ def MPI_IRecvOp : MPI_Op<"irecv", []> {
     MPI_Request : $req
   );
 
-  let assemblyFormat = "`(` $ref `,` $tag `,` $rank `)` attr-dict `:`"
-                       "type($ref) `,` type($tag) `,` type($rank) `->`"
+  let assemblyFormat = "`(` $ref (`,` $count^?) `,` $tag `,` $rank `)` attr-dict `:`"
+                       "type($ref) (`,` type($count)^?) `,` type($tag) `,` type($rank) `->`"
                        "type(results)";
   let hasCanonicalizer = 1;
 }
@@ -244,12 +248,13 @@ def MPI_AllReduceOp : MPI_Op<"allreduce", []> {
   let arguments = (
     ins AnyMemRef : $sendbuf,
     AnyMemRef : $recvbuf,
+    Optional<I32> : $count,
     MPI_OpClassAttr : $op
   );
 
   let results = (outs Optional<MPI_Retval>:$retval);
 
-  let assemblyFormat = "`(` $sendbuf `,` $recvbuf `,` $op `)` attr-dict `:`"
+  let assemblyFormat = "`(` $sendbuf `,` $recvbuf (`,` $count^?) `,` $op `)` attr-dict `:`"
                        "type($sendbuf) `,` type($recvbuf)"
                        "(`->` type($retval)^)?";
 }