[Mlir-commits] [mlir] Extend MPI dialect (PR #123255)
Sergio Sánchez Ramírez
llvmlistbot at llvm.org
Thu Jan 16 15:31:28 PST 2025
https://github.com/mofeing created https://github.com/llvm/llvm-project/pull/123255
cc @tobiasgrosser @wsmoses
this PR adds some new ops and types to the MLIR MPI dialect. the goal is to get the minimum required ops here to get a project of us working, and if everything works well, continue adding ops to the mpi dialect on subsequent PRs until we achieve some level of compliance with the MPI standard.
>From 1fbae54306aff0c55ec544675bedb961574c2837 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, 16 Jan 2025 22:31:59 +0100
Subject: [PATCH 1/3] Add `MPI_Comm`, `MPI_Request`, `MPI_Status`, `MPI_Op`
type definitions
---
mlir/include/mlir/Dialect/MPI/IR/MPITypes.td | 30 ++++++++++++++++++++
1 file changed, 30 insertions(+)
diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPITypes.td b/mlir/include/mlir/Dialect/MPI/IR/MPITypes.td
index 87eefa719d45c0..1d96b49d16585b 100644
--- a/mlir/include/mlir/Dialect/MPI/IR/MPITypes.td
+++ b/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
>From dc84ca4b87412dcfa9c83c48ae9916663c7ca38e 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, 16 Jan 2025 23:02:23 +0100
Subject: [PATCH 2/3] Add `MPI_CommSize`, `MPI_ISend`, `MPI_IRecv` ops
---
mlir/include/mlir/Dialect/MPI/IR/MPIOps.td | 85 ++++++++++++++++++++++
1 file changed, 85 insertions(+)
diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td b/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
index 240fac5104c34f..8719b67cd7f5f0 100644
--- a/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
+++ b/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
>From 2ee10ab60bb793b9164b0795e6657ceccc4704ae 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, 16 Jan 2025 23:02:34 +0100
Subject: [PATCH 3/3] Fix typo
---
mlir/include/mlir/Dialect/MPI/IR/MPIOps.td | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td b/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
index 8719b67cd7f5f0..4be5a6dfea7777 100644
--- a/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
@@ -251,7 +251,7 @@ def MPI_RetvalCheckOp : MPI_Op<"retval_check", []> {
//===----------------------------------------------------------------------===//
-// RetvalCheckOp
+// ErrorClassOp
//===----------------------------------------------------------------------===//
def MPI_ErrorClassOp : MPI_Op<"error_class", []> {
More information about the Mlir-commits
mailing list