From 4ed4b253f6bd988365eecb955590bcd7f2f72081 Mon Sep 17 00:00:00 2001
From: Anton Lydike <me@antonlydike.de>
Date: Thu, 12 Oct 2023 14:55:52 +0100
Subject: [PATCH 1/6] [mlir] Initial patch to add an MPI dialect

---
 mlir/include/mlir/Dialect/CMakeLists.txt      |   1 +
 mlir/include/mlir/Dialect/MPI/CMakeLists.txt  |   1 +
 .../mlir/Dialect/MPI/IR/CMakeLists.txt        |   3 +
 mlir/include/mlir/Dialect/MPI/IR/MPI.h        |  28 ++++
 mlir/include/mlir/Dialect/MPI/IR/MPIBase.td   |  36 +++++
 mlir/include/mlir/Dialect/MPI/IR/MPIOps.td    | 124 ++++++++++++++++++
 mlir/include/mlir/InitAllDialects.h           |   2 +
 mlir/lib/Dialect/CMakeLists.txt               |   1 +
 mlir/lib/Dialect/MPI/CMakeLists.txt           |   1 +
 mlir/lib/Dialect/MPI/IR/CMakeLists.txt        |  16 +++
 mlir/lib/Dialect/MPI/IR/MPIDialect.cpp        |  27 ++++
 mlir/lib/Dialect/MPI/IR/MPIOps.cpp            |  19 +++
 mlir/test/Dialect/MPI/invalid.mlir            |  22 ++++
 mlir/test/Dialect/MPI/ops.mlir                |  21 +++
 14 files changed, 302 insertions(+)
 create mode 100644 mlir/include/mlir/Dialect/MPI/CMakeLists.txt
 create mode 100644 mlir/include/mlir/Dialect/MPI/IR/CMakeLists.txt
 create mode 100644 mlir/include/mlir/Dialect/MPI/IR/MPI.h
 create mode 100644 mlir/include/mlir/Dialect/MPI/IR/MPIBase.td
 create mode 100644 mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
 create mode 100644 mlir/lib/Dialect/MPI/CMakeLists.txt
 create mode 100644 mlir/lib/Dialect/MPI/IR/CMakeLists.txt
 create mode 100644 mlir/lib/Dialect/MPI/IR/MPIDialect.cpp
 create mode 100644 mlir/lib/Dialect/MPI/IR/MPIOps.cpp
 create mode 100644 mlir/test/Dialect/MPI/invalid.mlir
 create mode 100644 mlir/test/Dialect/MPI/ops.mlir

diff --git a/mlir/include/mlir/Dialect/CMakeLists.txt b/mlir/include/mlir/Dialect/CMakeLists.txt
index 1c4569ecfa5848..9788e24e4a1d91 100644
--- a/mlir/include/mlir/Dialect/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/CMakeLists.txt
@@ -21,6 +21,7 @@ add_subdirectory(Math)
 add_subdirectory(MemRef)
 add_subdirectory(Mesh)
 add_subdirectory(MLProgram)
+add_subdirectory(MPI)
 add_subdirectory(NVGPU)
 add_subdirectory(OpenACC)
 add_subdirectory(OpenACCMPCommon)
diff --git a/mlir/include/mlir/Dialect/MPI/CMakeLists.txt b/mlir/include/mlir/Dialect/MPI/CMakeLists.txt
new file mode 100644
index 00000000000000..f33061b2d87cff
--- /dev/null
+++ b/mlir/include/mlir/Dialect/MPI/CMakeLists.txt
@@ -0,0 +1 @@
+add_subdirectory(IR)
diff --git a/mlir/include/mlir/Dialect/MPI/IR/CMakeLists.txt b/mlir/include/mlir/Dialect/MPI/IR/CMakeLists.txt
new file mode 100644
index 00000000000000..3b7b76240d375f
--- /dev/null
+++ b/mlir/include/mlir/Dialect/MPI/IR/CMakeLists.txt
@@ -0,0 +1,3 @@
+set(LLVM_TARGET_DEFINITIONS MPIOps.td)
+add_mlir_dialect(MPIOps mpi)
+add_mlir_doc(MPIOps MPIOps Dialects/ -gen-dialect-doc)
diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPI.h b/mlir/include/mlir/Dialect/MPI/IR/MPI.h
new file mode 100644
index 00000000000000..0cffe8edda4431
--- /dev/null
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPI.h
@@ -0,0 +1,28 @@
+//===- MPI.h - MPI dialect ----------------------------*- C++-*-==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_DIALECT_MPI_IR_MPI_H_
+#define MLIR_DIALECT_MPI_IR_MPI_H_
+
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpImplementation.h"
+
+//===----------------------------------------------------------------------===//
+// MPIDialect
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/MPI/IR/MPIOpsDialect.h.inc"
+
+//===----------------------------------------------------------------------===//
+// MPI Dialect Operations
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/MPI/IR/MPIOps.h.inc"
+
+#endif // MLIR_DIALECT_MPI_IR_MPI_H_
diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPIBase.td b/mlir/include/mlir/Dialect/MPI/IR/MPIBase.td
new file mode 100644
index 00000000000000..d42a5d0030fcbe
--- /dev/null
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPIBase.td
@@ -0,0 +1,36 @@
+//===- MPIBase.td - Base defs for mpi dialect --*- tablegen -*-==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MPI_BASE
+#define MPI_BASE
+
+include "mlir/IR/OpBase.td"
+
+def MPI_Dialect : Dialect {
+  let name = "mpi";
+  let cppNamespace = "::mlir::mpi";
+  let description = [{
+    This dialect models the Message Passing Interface (MPI), version 4.0. It is
+    meant to serve as an interfacing dialect that is targeted by higher-level dialects.
+    The MPI dialect itself can be lowered to multiple MPI implementations and hide
+    differences in ABI. The dialect models the functions of the MPI specification as
+    close to 1:1 as possible while preserving SSA value semantics where it makes sense,
+    and uses `memref` types instead of bare pointers.
+
+    This dialect is under active development, and while stability is an
+    eventual goal, it is not guaranteed at this juncture. Given the early state,
+    it is recommended to inquire further prior to using this dialect.
+
+    For an in-depth documentation of the MPI library interface, please refer to official documentation
+    such as the [OpenMPI online documentation](https://www.open-mpi.org/doc/current/).
+  }];
+
+  let usePropertiesForAttributes = 1;
+}
+
+#endif // MPI_BASE
diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td b/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
new file mode 100644
index 00000000000000..70e6ec1b623cfd
--- /dev/null
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
@@ -0,0 +1,124 @@
+//===- MPI.td - Message Passing Interface Ops ---------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MPI_OPS
+#define MPI_OPS
+
+include "mlir/Dialect/MPI/IR/MPIBase.td"
+
+class MPI_Op<string mnemonic, list<Trait> traits = []>
+    : Op<MPI_Dialect, mnemonic, traits>;
+
+//===----------------------------------------------------------------------===//
+// InitOp
+//===----------------------------------------------------------------------===//
+
+def MPI_InitOp : MPI_Op<"init", [
+
+]> {
+  let summary =
+      "Initialize the MPI library, equivalent to `MPI_Init(NULL, NULL)`";
+  let description = [{
+    This operation must preceed most MPI calls (except for very few exceptions,
+    please consult with the MPI specification on these).
+
+    Passing &argc, &argv is not supported currently.
+    Inspecting the functions return value (error code) is also not supported.
+  }];
+
+  let assemblyFormat = "attr-dict";
+}
+
+//===----------------------------------------------------------------------===//
+// CommRankOp
+//===----------------------------------------------------------------------===//
+
+def MPI_CommRankOp : MPI_Op<"comm_rank", [
+
+]> {
+  let summary = "Get the current rank, equivalent to "
+                "`MPI_Comm_rank(MPI_COMM_WORLD, &rank)`";
+  let description = [{
+    Communicators other than `MPI_COMM_WORLD` are not supprted for now.
+    Inspecting the functions return value (error code) is also not supported.
+  }];
+
+  let results = (outs I32 : $result);
+
+  let assemblyFormat = "attr-dict `:` type($result)";
+}
+
+//===----------------------------------------------------------------------===//
+// SendOp
+//===----------------------------------------------------------------------===//
+
+def MPI_SendOp : MPI_Op<"send", [
+
+]> {
+  let summary =
+      "Equivalent to `MPI_Send(ptr, size, dtype, dest, tag, MPI_COMM_WORLD)`";
+  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.
+    Inspecting the functions return value (error code) is also not supported.
+  }];
+
+  let arguments = (ins AnyMemRef : $ref, I32 : $tag, I32 : $rank);
+
+  let assemblyFormat = "`(` $ref `,` $tag `,` $rank `)` attr-dict `:` "
+                       "type($ref) `,` type($tag) `,` type($rank)";
+}
+
+//===----------------------------------------------------------------------===//
+// RecvOp
+//===----------------------------------------------------------------------===//
+
+def MPI_RecvOp : MPI_Op<"recv", [
+
+]> {
+  let summary = "Equivalent to `MPI_Recv(ptr, size, dtype, dest, tag, "
+                "MPI_COMM_WORLD, 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.
+    The MPI_Status is set to `MPI_STATUS_IGNORE`, as the status object is not yet ported to MLIR.
+    Inspecting the functions return value (error code) is also not supported.
+  }];
+
+  let arguments = (ins AnyMemRef : $ref, I32 : $tag, I32 : $rank);
+
+  let assemblyFormat = "`(` $ref `,` $tag `,` $rank `)` attr-dict `:` "
+                       "type($ref) `,` type($tag) `,` type($rank)";
+}
+
+//===----------------------------------------------------------------------===//
+// FinalizeOp
+//===----------------------------------------------------------------------===//
+
+def MPI_FinalizeOp : MPI_Op<"finalize", [
+
+]> {
+  let summary = "Finalize the MPI library, equivalent to `MPI_Finalize()`";
+  let description = [{
+    This function cleans up the MPI state. Afterwards, no MPI methods may be invoked
+    (excpet for MPI_Get_version, MPI_Initialized, and MPI_Finalized).
+    Notably, MPI_Init cannot be called again in the same program.
+
+    Inspecting the functions return value (error code) is not supported.
+  }];
+
+  let assemblyFormat = "attr-dict";
+}
+
+#endif // MPI_OPS
diff --git a/mlir/include/mlir/InitAllDialects.h b/mlir/include/mlir/InitAllDialects.h
index 00f400aab5d50a..42087994d0f0c8 100644
--- a/mlir/include/mlir/InitAllDialects.h
+++ b/mlir/include/mlir/InitAllDialects.h
@@ -48,6 +48,7 @@
 #include "mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/Transforms/TilingInterfaceImpl.h"
 #include "mlir/Dialect/MLProgram/IR/MLProgram.h"
+#include "mlir/Dialect/MPI/IR/MPI.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/IR/MemRefMemorySlot.h"
@@ -120,6 +121,7 @@ inline void registerAllDialects(DialectRegistry &registry) {
                   memref::MemRefDialect,
                   mesh::MeshDialect,
                   ml_program::MLProgramDialect,
+                  mpi::MPIDialect,
                   nvgpu::NVGPUDialect,
                   NVVM::NVVMDialect,
                   omp::OpenMPDialect,
diff --git a/mlir/lib/Dialect/CMakeLists.txt b/mlir/lib/Dialect/CMakeLists.txt
index 68776a695cac4d..c72107939cf42b 100644
--- a/mlir/lib/Dialect/CMakeLists.txt
+++ b/mlir/lib/Dialect/CMakeLists.txt
@@ -21,6 +21,7 @@ add_subdirectory(Math)
 add_subdirectory(MemRef)
 add_subdirectory(Mesh)
 add_subdirectory(MLProgram)
+add_subdirectory(MPI)
 add_subdirectory(NVGPU)
 add_subdirectory(OpenACC)
 add_subdirectory(OpenACCMPCommon)
diff --git a/mlir/lib/Dialect/MPI/CMakeLists.txt b/mlir/lib/Dialect/MPI/CMakeLists.txt
new file mode 100644
index 00000000000000..f33061b2d87cff
--- /dev/null
+++ b/mlir/lib/Dialect/MPI/CMakeLists.txt
@@ -0,0 +1 @@
+add_subdirectory(IR)
diff --git a/mlir/lib/Dialect/MPI/IR/CMakeLists.txt b/mlir/lib/Dialect/MPI/IR/CMakeLists.txt
new file mode 100644
index 00000000000000..870f034e5ccca7
--- /dev/null
+++ b/mlir/lib/Dialect/MPI/IR/CMakeLists.txt
@@ -0,0 +1,16 @@
+add_mlir_dialect_library(MLIRMPIDialect
+  MPIOps.cpp
+  MPIDialect.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/MPI
+
+  DEPENDS
+  MLIRMPIOpsIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRDialect
+  MLIRIR
+  MLIRInferTypeOpInterface
+  MLIRSideEffectInterfaces
+  )
diff --git a/mlir/lib/Dialect/MPI/IR/MPIDialect.cpp b/mlir/lib/Dialect/MPI/IR/MPIDialect.cpp
new file mode 100644
index 00000000000000..19ed5d8207af45
--- /dev/null
+++ b/mlir/lib/Dialect/MPI/IR/MPIDialect.cpp
@@ -0,0 +1,27 @@
+//===- MPIDialect.cpp - MPI dialect implementation ------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/MPI/IR/MPI.h"
+#include "mlir/IR/DialectImplementation.h"
+
+using namespace mlir;
+using namespace mlir::mpi;
+
+//===----------------------------------------------------------------------===//
+/// Tablegen Definitions
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/MPI/IR/MPIOpsDialect.cpp.inc"
+
+void mpi::MPIDialect::initialize() {
+
+  addOperations<
+#define GET_OP_LIST
+#include "mlir/Dialect/MPI/IR/MPIOps.cpp.inc"
+      >();
+}
diff --git a/mlir/lib/Dialect/MPI/IR/MPIOps.cpp b/mlir/lib/Dialect/MPI/IR/MPIOps.cpp
new file mode 100644
index 00000000000000..84b27c2dae419b
--- /dev/null
+++ b/mlir/lib/Dialect/MPI/IR/MPIOps.cpp
@@ -0,0 +1,19 @@
+//===- MPIOps.cpp - MPI dialect ops implementation ------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/MPI/IR/MPI.h"
+
+using namespace mlir;
+using namespace mlir::mpi;
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/MPI/IR/MPIOps.cpp.inc"
diff --git a/mlir/test/Dialect/MPI/invalid.mlir b/mlir/test/Dialect/MPI/invalid.mlir
new file mode 100644
index 00000000000000..9ce405a445f9d6
--- /dev/null
+++ b/mlir/test/Dialect/MPI/invalid.mlir
@@ -0,0 +1,22 @@
+// RUN: mlir-opt -split-input-file -verify-diagnostics %s
+
+// expected-error @+1 {{op result #0 must be 32-bit signless integer, but got 'i64'}}
+%rank = mpi.comm_rank : i64
+
+// -----
+
+func.func @mpi_test(%ref : !llvm.ptr, %rank: i32) -> () {
+    // expected-error @+1 {{invalid kind of type specified}}
+    mpi.send(%ref, %rank, %rank) : !llvm.ptr, i32, i32
+
+    return
+}
+
+// -----
+
+func.func @mpi_test(%ref : !llvm.ptr, %rank: i32) -> () {
+    // expected-error @+1 {{invalid kind of type specified}}
+    mpi.recv(%ref, %rank, %rank) : !llvm.ptr, i32, i32
+
+    return
+}
diff --git a/mlir/test/Dialect/MPI/ops.mlir b/mlir/test/Dialect/MPI/ops.mlir
new file mode 100644
index 00000000000000..a03d44b9b27013
--- /dev/null
+++ b/mlir/test/Dialect/MPI/ops.mlir
@@ -0,0 +1,21 @@
+// RUN: mlir-opt %s | mlir-opt | FileCheck %s
+// RUN: mlir-opt %s --mlir-print-op-generic | mlir-opt | FileCheck %s
+
+func.func @mpi_test(%ref : memref<100xf32>) -> () {
+    // CHECK: mpi.init
+    mpi.init
+
+    // CHECK-NEXT: mpi.comm_rank : i32
+    %rank = mpi.comm_rank : i32
+
+    // CHECK-NEXT: mpi.send(%arg0, %0, %0) : memref<100xf32>, i32, i32
+    mpi.send(%ref, %rank, %rank) : memref<100xf32>, i32, i32
+
+    // CHECK-NEXT: mpi.recv(%arg0, %0, %0) : memref<100xf32>, i32, i32
+    mpi.recv(%ref, %rank, %rank) : memref<100xf32>, i32, i32
+
+    // CHECK-NEXT: mpi.finalize
+    mpi.finalize
+
+    func.return
+}

From 4d574df3243bb873c64a4ea9e0d179b69d54b39c Mon Sep 17 00:00:00 2001
From: Anton Lydike <me@antonlydike.de>
Date: Wed, 13 Dec 2023 21:05:18 +0100
Subject: [PATCH 2/6] minor breaking version

---
 .../mlir/Dialect/MPI/IR/CMakeLists.txt        |  26 ++-
 mlir/include/mlir/Dialect/MPI/IR/MPI.h        |  10 +-
 mlir/include/mlir/Dialect/MPI/IR/MPI.td       | 178 ++++++++++++++++++
 mlir/include/mlir/Dialect/MPI/IR/MPIBase.td   |  36 ----
 mlir/include/mlir/Dialect/MPI/IR/MPIOps.td    |  20 +-
 mlir/include/mlir/Dialect/MPI/IR/MPITypes.td  |  41 ++++
 mlir/lib/Dialect/MPI/IR/CMakeLists.txt        |   3 +
 mlir/lib/Dialect/MPI/IR/MPIDialect.cpp        |  22 ++-
 8 files changed, 284 insertions(+), 52 deletions(-)
 create mode 100644 mlir/include/mlir/Dialect/MPI/IR/MPI.td
 delete mode 100644 mlir/include/mlir/Dialect/MPI/IR/MPIBase.td
 create mode 100644 mlir/include/mlir/Dialect/MPI/IR/MPITypes.td

diff --git a/mlir/include/mlir/Dialect/MPI/IR/CMakeLists.txt b/mlir/include/mlir/Dialect/MPI/IR/CMakeLists.txt
index 3b7b76240d375f..08222dd5f20af4 100644
--- a/mlir/include/mlir/Dialect/MPI/IR/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/MPI/IR/CMakeLists.txt
@@ -1,3 +1,25 @@
+add_mlir_dialect(MPI mpi)
+add_mlir_doc(MPI MPI Dialects/ -gen-dialect-doc)
+
+# Add MPI operations
 set(LLVM_TARGET_DEFINITIONS MPIOps.td)
-add_mlir_dialect(MPIOps mpi)
-add_mlir_doc(MPIOps MPIOps Dialects/ -gen-dialect-doc)
+mlir_tablegen(MPIOps.h.inc -gen-op-decls)
+mlir_tablegen(MPIOps.cpp.inc -gen-op-defs)
+add_public_tablegen_target(MLIRMPIOpsIncGen)
+add_dependencies(mlir-generic-headers MLIRMPIOpsIncGen)
+
+# Add MPI types
+set(LLVM_TARGET_DEFINITIONS MPITypes.td)
+mlir_tablegen(MPITypesGen.h.inc -gen-typedef-decls)
+mlir_tablegen(MPITypesGen.cpp.inc -gen-typedef-defs)
+add_public_tablegen_target(MLIRMPITypesIncGen)
+add_dependencies(mlir-generic-headers MLIRMPITypesIncGen)
+
+# Add MPI attributes
+set(LLVM_TARGET_DEFINITIONS MPI.td)
+mlir_tablegen(MPIEnums.h.inc -gen-enum-decls)
+mlir_tablegen(MPIEnums.cpp.inc -gen-enum-defs)
+mlir_tablegen(MPIAttributes.h.inc -gen-attrdef-decls)
+mlir_tablegen(MPIAttributes.cpp.inc -gen-attrdef-defs)
+add_public_tablegen_target(MLIRMPIAttrsIncGen)
+add_dependencies(mlir-generic-headers MLIRMPIAttrsIncGen)
diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPI.h b/mlir/include/mlir/Dialect/MPI/IR/MPI.h
index 0cffe8edda4431..503f29fe1cb769 100644
--- a/mlir/include/mlir/Dialect/MPI/IR/MPI.h
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPI.h
@@ -18,9 +18,13 @@
 
 #include "mlir/Dialect/MPI/IR/MPIOpsDialect.h.inc"
 
-//===----------------------------------------------------------------------===//
-// MPI Dialect Operations
-//===----------------------------------------------------------------------===//
+#define GET_TYPEDEF_CLASSES
+#include "mlir/Dialect/MPI/IR/MPITypesGen.h.inc"
+
+#include "mlir/Dialect/MPI/IR/MPIEnums.h.inc"
+
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/MPI/IR/MPIAttributes.h.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/MPI/IR/MPIOps.h.inc"
diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPI.td b/mlir/include/mlir/Dialect/MPI/IR/MPI.td
new file mode 100644
index 00000000000000..c73e621619a459
--- /dev/null
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPI.td
@@ -0,0 +1,178 @@
+//===- MPIBase.td - Base defs for mpi dialect --*- tablegen -*-==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MPI_MLIR_IR_BASE
+#define MPI_MLIR_IR_BASE
+
+include "mlir/IR/AttrTypeBase.td"
+include "mlir/IR/OpBase.td"
+include "mlir/IR/EnumAttr.td"
+
+def MPI_Dialect : Dialect {
+  let name = "mpi";
+  let cppNamespace = "::mlir::mpi";
+  let description = [{
+    This dialect models the Message Passing Interface (MPI), version 4.0. It is
+    meant to serve as an interfacing dialect that is targeted by higher-level dialects.
+    The MPI dialect itself can be lowered to multiple MPI implementations and hide
+    differences in ABI. The dialect models the functions of the MPI specification as
+    close to 1:1 as possible while preserving SSA value semantics where it makes sense,
+    and uses `memref` types instead of bare pointers.
+
+    This dialect is under active development, and while stability is an
+    eventual goal, it is not guaranteed at this juncture. Given the early state,
+    it is recommended to inquire further prior to using this dialect.
+
+    For an in-depth documentation of the MPI library interface, please refer to official documentation
+    such as the [OpenMPI online documentation](https://www.open-mpi.org/doc/current/).
+  }];
+
+  let usePropertiesForAttributes = 1;
+  let useDefaultAttributePrinterParser = 1;
+  let useDefaultTypePrinterParser = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// Error classes enum:
+//===----------------------------------------------------------------------===//
+
+def MpiCodeSuccess : I32EnumAttrCase<"MPI_SUCCESS", 0, "MPI_SUCCESS">;
+def MpiCodeErrAccess : I32EnumAttrCase<"MPI_ERR_ACCESS", 1, "MPI_ERR_ACCESS">;
+def MpiCodeErrAmode : I32EnumAttrCase<"MPI_ERR_AMODE", 2, "MPI_ERR_AMODE">;
+def MpiCodeErrArg : I32EnumAttrCase<"MPI_ERR_ARG", 3, "MPI_ERR_ARG">;
+def MpiCodeErrAssert : I32EnumAttrCase<"MPI_ERR_ASSERT", 4, "MPI_ERR_ASSERT">;
+def MpiCodeErrBadFile : I32EnumAttrCase<"MPI_ERR_BAD_FILE", 5, "MPI_ERR_BAD_FILE">;
+def MpiCodeErrBase : I32EnumAttrCase<"MPI_ERR_BASE", 6, "MPI_ERR_BASE">;
+def MpiCodeErrBuffer : I32EnumAttrCase<"MPI_ERR_BUFFER", 7, "MPI_ERR_BUFFER">;
+def MpiCodeErrComm : I32EnumAttrCase<"MPI_ERR_COMM", 8, "MPI_ERR_COMM">;
+def MpiCodeErrConversion : I32EnumAttrCase<"MPI_ERR_CONVERSION", 9, "MPI_ERR_CONVERSION">;
+def MpiCodeErrCount : I32EnumAttrCase<"MPI_ERR_COUNT", 10, "MPI_ERR_COUNT">;
+def MpiCodeErrDims : I32EnumAttrCase<"MPI_ERR_DIMS", 11, "MPI_ERR_DIMS">;
+def MpiCodeErrDisp : I32EnumAttrCase<"MPI_ERR_DISP", 12, "MPI_ERR_DISP">;
+def MpiCodeErrDupDatarep : I32EnumAttrCase<"MPI_ERR_DUP_DATAREP", 13, "MPI_ERR_DUP_DATAREP">;
+def MpiCodeErrErrhandler : I32EnumAttrCase<"MPI_ERR_ERRHANDLER", 14, "MPI_ERR_ERRHANDLER">;
+def MpiCodeErrFile : I32EnumAttrCase<"MPI_ERR_FILE", 15, "MPI_ERR_FILE">;
+def MpiCodeErrFileExists : I32EnumAttrCase<"MPI_ERR_FILE_EXISTS", 16, "MPI_ERR_FILE_EXISTS">;
+def MpiCodeErrFileInUse : I32EnumAttrCase<"MPI_ERR_FILE_IN_USE", 17, "MPI_ERR_FILE_IN_USE">;
+def MpiCodeErrGroup : I32EnumAttrCase<"MPI_ERR_GROUP", 18, "MPI_ERR_GROUP">;
+def MpiCodeErrInfo : I32EnumAttrCase<"MPI_ERR_INFO", 19, "MPI_ERR_INFO">;
+def MpiCodeErrInfoKey : I32EnumAttrCase<"MPI_ERR_INFO_KEY", 20, "MPI_ERR_INFO_KEY">;
+def MpiCodeErrInfoNokey : I32EnumAttrCase<"MPI_ERR_INFO_NOKEY", 21, "MPI_ERR_INFO_NOKEY">;
+def MpiCodeErrInfoValue : I32EnumAttrCase<"MPI_ERR_INFO_VALUE", 22, "MPI_ERR_INFO_VALUE">;
+def MpiCodeErrInStatus : I32EnumAttrCase<"MPI_ERR_IN_STATUS", 23, "MPI_ERR_IN_STATUS">;
+def MpiCodeErrIntern : I32EnumAttrCase<"MPI_ERR_INTERN", 24, "MPI_ERR_INTERN">;
+def MpiCodeErrIo : I32EnumAttrCase<"MPI_ERR_IO", 25, "MPI_ERR_IO">;
+def MpiCodeErrKeyval : I32EnumAttrCase<"MPI_ERR_KEYVAL", 26, "MPI_ERR_KEYVAL">;
+def MpiCodeErrLocktype : I32EnumAttrCase<"MPI_ERR_LOCKTYPE", 27, "MPI_ERR_LOCKTYPE">;
+def MpiCodeErrName : I32EnumAttrCase<"MPI_ERR_NAME", 28, "MPI_ERR_NAME">;
+def MpiCodeErrNoMem : I32EnumAttrCase<"MPI_ERR_NO_MEM", 29, "MPI_ERR_NO_MEM">;
+def MpiCodeErrNoSpace : I32EnumAttrCase<"MPI_ERR_NO_SPACE", 30, "MPI_ERR_NO_SPACE">;
+def MpiCodeErrNoSuchFile : I32EnumAttrCase<"MPI_ERR_NO_SUCH_FILE", 31, "MPI_ERR_NO_SUCH_FILE">;
+def MpiCodeErrNotSame : I32EnumAttrCase<"MPI_ERR_NOT_SAME", 32, "MPI_ERR_NOT_SAME">;
+def MpiCodeErrOp : I32EnumAttrCase<"MPI_ERR_OP", 33, "MPI_ERR_OP">;
+def MpiCodeErrOther : I32EnumAttrCase<"MPI_ERR_OTHER", 34, "MPI_ERR_OTHER">;
+def MpiCodeErrPending : I32EnumAttrCase<"MPI_ERR_PENDING", 35, "MPI_ERR_PENDING">;
+def MpiCodeErrPort : I32EnumAttrCase<"MPI_ERR_PORT", 36, "MPI_ERR_PORT">;
+def MpiCodeErrProcAborted : I32EnumAttrCase<"MPI_ERR_PROC_ABORTED", 37, "MPI_ERR_PROC_ABORTED">;
+def MpiCodeErrQuota : I32EnumAttrCase<"MPI_ERR_QUOTA", 38, "MPI_ERR_QUOTA">;
+def MpiCodeErrRank : I32EnumAttrCase<"MPI_ERR_RANK", 39, "MPI_ERR_RANK">;
+def MpiCodeErrReadOnly : I32EnumAttrCase<"MPI_ERR_READ_ONLY", 40, "MPI_ERR_READ_ONLY">;
+def MpiCodeErrRequest : I32EnumAttrCase<"MPI_ERR_REQUEST", 41, "MPI_ERR_REQUEST">;
+def MpiCodeErrRmaAttach : I32EnumAttrCase<"MPI_ERR_RMA_ATTACH", 42, "MPI_ERR_RMA_ATTACH">;
+def MpiCodeErrRmaConflict : I32EnumAttrCase<"MPI_ERR_RMA_CONFLICT", 43, "MPI_ERR_RMA_CONFLICT">;
+def MpiCodeErrRmaFlavor : I32EnumAttrCase<"MPI_ERR_RMA_FLAVOR", 44, "MPI_ERR_RMA_FLAVOR">;
+def MpiCodeErrRmaRange : I32EnumAttrCase<"MPI_ERR_RMA_RANGE", 45, "MPI_ERR_RMA_RANGE">;
+def MpiCodeErrRmaShared : I32EnumAttrCase<"MPI_ERR_RMA_SHARED", 46, "MPI_ERR_RMA_SHARED">;
+def MpiCodeErrRmaSync : I32EnumAttrCase<"MPI_ERR_RMA_SYNC", 47, "MPI_ERR_RMA_SYNC">;
+def MpiCodeErrRoot : I32EnumAttrCase<"MPI_ERR_ROOT", 48, "MPI_ERR_ROOT">;
+def MpiCodeErrService : I32EnumAttrCase<"MPI_ERR_SERVICE", 49, "MPI_ERR_SERVICE">;
+def MpiCodeErrSession : I32EnumAttrCase<"MPI_ERR_SESSION", 50, "MPI_ERR_SESSION">;
+def MpiCodeErrSize : I32EnumAttrCase<"MPI_ERR_SIZE", 51, "MPI_ERR_SIZE">;
+def MpiCodeErrSpawn : I32EnumAttrCase<"MPI_ERR_SPAWN", 52, "MPI_ERR_SPAWN">;
+def MpiCodeErrTag : I32EnumAttrCase<"MPI_ERR_TAG", 53, "MPI_ERR_TAG">;
+def MpiCodeErrTopology : I32EnumAttrCase<"MPI_ERR_TOPOLOGY", 54, "MPI_ERR_TOPOLOGY">;
+def MpiCodeErrTruncate : I32EnumAttrCase<"MPI_ERR_TRUNCATE", 55, "MPI_ERR_TRUNCATE">;
+def MpiCodeErrType : I32EnumAttrCase<"MPI_ERR_TYPE", 56, "MPI_ERR_TYPE">;
+def MpiCodeErrUnknown : I32EnumAttrCase<"MPI_ERR_UNKNOWN", 57, "MPI_ERR_UNKNOWN">;
+def MpiCodeErrUnsupportedDatarep : I32EnumAttrCase<"MPI_ERR_UNSUPPORTED_DATAREP", 58, "MPI_ERR_UNSUPPORTED_DATAREP">;
+def MpiCodeErrUnsupportedOperation : I32EnumAttrCase<"MPI_ERR_UNSUPPORTED_OPERATION", 59, "MPI_ERR_UNSUPPORTED_OPERATION">;
+def MpiCodeErrValueTooLarge : I32EnumAttrCase<"MPI_ERR_VALUE_TOO_LARGE", 60, "MPI_ERR_VALUE_TOO_LARGE">;
+def MpiCodeErrWin : I32EnumAttrCase<"MPI_ERR_WIN", 61, "MPI_ERR_WIN">;
+def MpiCodeErrLastcode : I32EnumAttrCase<"MPI_ERR_LASTCODE", 62, "MPI_ERR_LASTCODE">;
+
+def MpiErrorClasses : I32EnumAttr<"MpiErrorClass", 
+                                "MPI error class name",
+  [ MpiCodeSuccess
+    ,MpiCodeErrAccess
+    ,MpiCodeErrAmode
+    ,MpiCodeErrArg
+    ,MpiCodeErrAssert
+    ,MpiCodeErrBadFile
+    ,MpiCodeErrBase
+    ,MpiCodeErrBuffer
+    ,MpiCodeErrComm
+    ,MpiCodeErrConversion
+    ,MpiCodeErrCount
+    ,MpiCodeErrDims
+    ,MpiCodeErrDisp
+    ,MpiCodeErrDupDatarep
+    ,MpiCodeErrErrhandler
+    ,MpiCodeErrFile
+    ,MpiCodeErrFileExists
+    ,MpiCodeErrFileInUse
+    ,MpiCodeErrGroup
+    ,MpiCodeErrInfo
+    ,MpiCodeErrInfoKey
+    ,MpiCodeErrInfoNokey
+    ,MpiCodeErrInfoValue
+    ,MpiCodeErrInStatus
+    ,MpiCodeErrIntern
+    ,MpiCodeErrIo
+    ,MpiCodeErrKeyval
+    ,MpiCodeErrLocktype
+    ,MpiCodeErrName
+    ,MpiCodeErrNoMem
+    ,MpiCodeErrNoSpace
+    ,MpiCodeErrNoSuchFile
+    ,MpiCodeErrNotSame
+    ,MpiCodeErrOp
+    ,MpiCodeErrOther
+    ,MpiCodeErrPending
+    ,MpiCodeErrPort
+    ,MpiCodeErrProcAborted
+    ,MpiCodeErrQuota
+    ,MpiCodeErrRank
+    ,MpiCodeErrReadOnly
+    ,MpiCodeErrRequest
+    ,MpiCodeErrRmaAttach
+    ,MpiCodeErrRmaConflict
+    ,MpiCodeErrRmaFlavor
+    ,MpiCodeErrRmaRange
+    ,MpiCodeErrRmaShared
+    ,MpiCodeErrRmaSync
+    ,MpiCodeErrRoot
+    ,MpiCodeErrService
+    ,MpiCodeErrSession
+    ,MpiCodeErrSize
+    ,MpiCodeErrSpawn
+    ,MpiCodeErrTag
+    ,MpiCodeErrTopology
+    ,MpiCodeErrTruncate
+    ,MpiCodeErrType
+    ,MpiCodeErrUnknown
+    ,MpiCodeErrUnsupportedDatarep
+    ,MpiCodeErrUnsupportedOperation
+    ,MpiCodeErrValueTooLarge
+    ,MpiCodeErrWin
+    ,MpiCodeErrLastcode]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::mpi";
+}
+
+
+#endif // MPI_MLIR_IR_BASE
diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPIBase.td b/mlir/include/mlir/Dialect/MPI/IR/MPIBase.td
deleted file mode 100644
index d42a5d0030fcbe..00000000000000
--- a/mlir/include/mlir/Dialect/MPI/IR/MPIBase.td
+++ /dev/null
@@ -1,36 +0,0 @@
-//===- MPIBase.td - Base defs for mpi dialect --*- tablegen -*-==//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef MPI_BASE
-#define MPI_BASE
-
-include "mlir/IR/OpBase.td"
-
-def MPI_Dialect : Dialect {
-  let name = "mpi";
-  let cppNamespace = "::mlir::mpi";
-  let description = [{
-    This dialect models the Message Passing Interface (MPI), version 4.0. It is
-    meant to serve as an interfacing dialect that is targeted by higher-level dialects.
-    The MPI dialect itself can be lowered to multiple MPI implementations and hide
-    differences in ABI. The dialect models the functions of the MPI specification as
-    close to 1:1 as possible while preserving SSA value semantics where it makes sense,
-    and uses `memref` types instead of bare pointers.
-
-    This dialect is under active development, and while stability is an
-    eventual goal, it is not guaranteed at this juncture. Given the early state,
-    it is recommended to inquire further prior to using this dialect.
-
-    For an in-depth documentation of the MPI library interface, please refer to official documentation
-    such as the [OpenMPI online documentation](https://www.open-mpi.org/doc/current/).
-  }];
-
-  let usePropertiesForAttributes = 1;
-}
-
-#endif // MPI_BASE
diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td b/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
index 70e6ec1b623cfd..a6b5c71c779a63 100644
--- a/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
@@ -1,4 +1,4 @@
-//===- MPI.td - Message Passing Interface Ops ---------*- tablegen -*-===//
+//===- MPI.td - Message Passing Interface Ops --------------*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -6,10 +6,11 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MPI_OPS
-#define MPI_OPS
+#ifndef MPI_MLIR_IR_MPIOPS
+#define MPI_MLIR_IR_MPIOPS
 
-include "mlir/Dialect/MPI/IR/MPIBase.td"
+include "mlir/Dialect/MPI/IR/MPI.td"
+include "mlir/Dialect/MPI/IR/MPITypes.td"
 
 class MPI_Op<string mnemonic, list<Trait> traits = []>
     : Op<MPI_Dialect, mnemonic, traits>;
@@ -18,9 +19,7 @@ class MPI_Op<string mnemonic, list<Trait> traits = []>
 // InitOp
 //===----------------------------------------------------------------------===//
 
-def MPI_InitOp : MPI_Op<"init", [
-
-]> {
+def MPI_InitOp : MPI_Op<"init", []> {
   let summary =
       "Initialize the MPI library, equivalent to `MPI_Init(NULL, NULL)`";
   let description = [{
@@ -31,7 +30,9 @@ def MPI_InitOp : MPI_Op<"init", [
     Inspecting the functions return value (error code) is also not supported.
   }];
 
-  let assemblyFormat = "attr-dict";
+  let results = (outs Optional<MPI_Retval>:$retval);
+
+  let assemblyFormat = "attr-dict (`:` type($retval)^)?";
 }
 
 //===----------------------------------------------------------------------===//
@@ -102,6 +103,7 @@ def MPI_RecvOp : MPI_Op<"recv", [
                        "type($ref) `,` type($tag) `,` type($rank)";
 }
 
+
 //===----------------------------------------------------------------------===//
 // FinalizeOp
 //===----------------------------------------------------------------------===//
@@ -121,4 +123,4 @@ def MPI_FinalizeOp : MPI_Op<"finalize", [
   let assemblyFormat = "attr-dict";
 }
 
-#endif // MPI_OPS
+#endif // MPI_MLIR_IR_MPIOPS
diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPITypes.td b/mlir/include/mlir/Dialect/MPI/IR/MPITypes.td
new file mode 100644
index 00000000000000..362b3bac5d9af3
--- /dev/null
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPITypes.td
@@ -0,0 +1,41 @@
+//===- MPITypes.td - Message Passing Interface types -------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares the Message Passing Interface dialect types.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MPI_MLIR_IR_MPITYPES
+#define MPI_MLIR_IR_MPITYPES
+
+include "mlir/IR/AttrTypeBase.td"
+include "mlir/Dialect/MPI/IR/MPI.td"
+
+//===----------------------------------------------------------------------===//
+// NPI Types
+//===----------------------------------------------------------------------===//
+
+class MPI_Type<string name, string typeMnemonic, list<Trait> traits = []>
+    : TypeDef<MPI_Dialect, name, traits> {
+  let mnemonic = typeMnemonic;
+}
+
+//===----------------------------------------------------------------------===//
+// pdl::AttributeType
+//===----------------------------------------------------------------------===//
+
+def MPI_Retval : MPI_Type<"Retval", "retval"> {
+  let summary = "The return value of an MPI call.";
+  let description = [{
+    This type represents a value returned from an MPI call. This value can be
+    MPI_SUCCESS, MPI_ERR_IN_STATUS, or any error code.
+  }];
+}
+
+
+#endif // MPI_MLIR_IR_MPITYPES
\ No newline at end of file
diff --git a/mlir/lib/Dialect/MPI/IR/CMakeLists.txt b/mlir/lib/Dialect/MPI/IR/CMakeLists.txt
index 870f034e5ccca7..ed5b21216ec4e8 100644
--- a/mlir/lib/Dialect/MPI/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/MPI/IR/CMakeLists.txt
@@ -6,7 +6,10 @@ add_mlir_dialect_library(MLIRMPIDialect
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/MPI
 
   DEPENDS
+  MLIRMPIIncGen
   MLIRMPIOpsIncGen
+  MLIRMPITypesIncGen
+  MLIRMPIAttrsIncGen
 
   LINK_LIBS PUBLIC
   MLIRDialect
diff --git a/mlir/lib/Dialect/MPI/IR/MPIDialect.cpp b/mlir/lib/Dialect/MPI/IR/MPIDialect.cpp
index 19ed5d8207af45..e062d62437a5b8 100644
--- a/mlir/lib/Dialect/MPI/IR/MPIDialect.cpp
+++ b/mlir/lib/Dialect/MPI/IR/MPIDialect.cpp
@@ -18,10 +18,28 @@ using namespace mlir::mpi;
 
 #include "mlir/Dialect/MPI/IR/MPIOpsDialect.cpp.inc"
 
-void mpi::MPIDialect::initialize() {
-
+void MPIDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/MPI/IR/MPIOps.cpp.inc"
       >();
+
+  addTypes<
+#define GET_TYPEDEF_LIST
+#include "mlir/Dialect/MPI/IR/MPITypesGen.cpp.inc"
+      >();
+
+  addAttributes<
+#define GET_ATTRDEF_LIST
+#include "mlir/Dialect/MPI/IR/MPIAttributes.cpp.inc"
+      >();
 }
+
+
+#define GET_TYPEDEF_CLASSES
+#include "mlir/Dialect/MPI/IR/MPITypesGen.cpp.inc"
+
+#include "mlir/Dialect/MPI/IR/MPIEnums.cpp.inc"
+
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/MPI/IR/MPIAttributes.cpp.inc"

From 1f43fae02472ecaf067327a53df3cfd9e1982a74 Mon Sep 17 00:00:00 2001
From: Anton Lydike <me@antonlydike.de>
Date: Thu, 14 Dec 2023 09:23:31 +0100
Subject: [PATCH 3/6] clean up a lot of stuff

---
 mlir/include/mlir/Dialect/MPI/IR/CMakeLists.txt |  4 ++--
 mlir/include/mlir/Dialect/MPI/IR/MPI.h          |  4 ++--
 mlir/include/mlir/Dialect/MPI/IR/MPI.td         |  9 ++++++---
 mlir/include/mlir/Dialect/MPI/IR/MPIOps.td      |  1 -
 mlir/include/mlir/Dialect/MPI/IR/MPITypes.td    |  8 ++++----
 mlir/lib/Dialect/MPI/IR/MPIDialect.cpp          | 15 ++++++++++++---
 6 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/mlir/include/mlir/Dialect/MPI/IR/CMakeLists.txt b/mlir/include/mlir/Dialect/MPI/IR/CMakeLists.txt
index 08222dd5f20af4..dc4b7a9087e609 100644
--- a/mlir/include/mlir/Dialect/MPI/IR/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/MPI/IR/CMakeLists.txt
@@ -19,7 +19,7 @@ add_dependencies(mlir-generic-headers MLIRMPITypesIncGen)
 set(LLVM_TARGET_DEFINITIONS MPI.td)
 mlir_tablegen(MPIEnums.h.inc -gen-enum-decls)
 mlir_tablegen(MPIEnums.cpp.inc -gen-enum-defs)
-mlir_tablegen(MPIAttributes.h.inc -gen-attrdef-decls)
-mlir_tablegen(MPIAttributes.cpp.inc -gen-attrdef-defs)
+mlir_tablegen(MPIAttrDefs.h.inc -gen-attrdef-decls)
+mlir_tablegen(MPIAttrDefs.cpp.inc -gen-attrdef-defs)
 add_public_tablegen_target(MLIRMPIAttrsIncGen)
 add_dependencies(mlir-generic-headers MLIRMPIAttrsIncGen)
diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPI.h b/mlir/include/mlir/Dialect/MPI/IR/MPI.h
index 503f29fe1cb769..1f71f3d7d70809 100644
--- a/mlir/include/mlir/Dialect/MPI/IR/MPI.h
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPI.h
@@ -16,7 +16,7 @@
 // MPIDialect
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Dialect/MPI/IR/MPIOpsDialect.h.inc"
+#include "mlir/Dialect/MPI/IR/MPIDialect.h.inc"
 
 #define GET_TYPEDEF_CLASSES
 #include "mlir/Dialect/MPI/IR/MPITypesGen.h.inc"
@@ -24,7 +24,7 @@
 #include "mlir/Dialect/MPI/IR/MPIEnums.h.inc"
 
 #define GET_ATTRDEF_CLASSES
-#include "mlir/Dialect/MPI/IR/MPIAttributes.h.inc"
+#include "mlir/Dialect/MPI/IR/MPIAttrDefs.h.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/MPI/IR/MPIOps.h.inc"
diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPI.td b/mlir/include/mlir/Dialect/MPI/IR/MPI.td
index c73e621619a459..8587fb42566c2c 100644
--- a/mlir/include/mlir/Dialect/MPI/IR/MPI.td
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPI.td
@@ -6,8 +6,8 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MPI_MLIR_IR_BASE
-#define MPI_MLIR_IR_BASE
+#ifndef MLIR_DIALECT_MPI_IR_MPI
+#define MLIR_DIALECT_MPI_IR_MPI
 
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/OpBase.td"
@@ -41,6 +41,8 @@ def MPI_Dialect : Dialect {
 // Error classes enum:
 //===----------------------------------------------------------------------===//
 
+/*
+
 def MpiCodeSuccess : I32EnumAttrCase<"MPI_SUCCESS", 0, "MPI_SUCCESS">;
 def MpiCodeErrAccess : I32EnumAttrCase<"MPI_ERR_ACCESS", 1, "MPI_ERR_ACCESS">;
 def MpiCodeErrAmode : I32EnumAttrCase<"MPI_ERR_AMODE", 2, "MPI_ERR_AMODE">;
@@ -174,5 +176,6 @@ def MpiErrorClasses : I32EnumAttr<"MpiErrorClass",
   let cppNamespace = "::mlir::mpi";
 }
 
+*/
 
-#endif // MPI_MLIR_IR_BASE
+#endif // MLIR_DIALECT_MPI_IR_MPI
diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td b/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
index a6b5c71c779a63..9423628a0f696c 100644
--- a/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
@@ -27,7 +27,6 @@ def MPI_InitOp : MPI_Op<"init", []> {
     please consult with the MPI specification on these).
 
     Passing &argc, &argv is not supported currently.
-    Inspecting the functions return value (error code) is also not supported.
   }];
 
   let results = (outs Optional<MPI_Retval>:$retval);
diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPITypes.td b/mlir/include/mlir/Dialect/MPI/IR/MPITypes.td
index 362b3bac5d9af3..4cf07ca37267f2 100644
--- a/mlir/include/mlir/Dialect/MPI/IR/MPITypes.td
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPITypes.td
@@ -10,11 +10,11 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MPI_MLIR_IR_MPITYPES
-#define MPI_MLIR_IR_MPITYPES
+#ifndef MLIR_DIALECT_MPI_IR_MPITYPES
+#define MLIR_DIALECT_MPI_IR_MPITYPES
 
 include "mlir/IR/AttrTypeBase.td"
-include "mlir/Dialect/MPI/IR/MPI.td"
+include "MPI.td"
 
 //===----------------------------------------------------------------------===//
 // NPI Types
@@ -38,4 +38,4 @@ def MPI_Retval : MPI_Type<"Retval", "retval"> {
 }
 
 
-#endif // MPI_MLIR_IR_MPITYPES
\ No newline at end of file
+#endif // MLIR_DIALECT_MPI_IR_MPITYPES
\ No newline at end of file
diff --git a/mlir/lib/Dialect/MPI/IR/MPIDialect.cpp b/mlir/lib/Dialect/MPI/IR/MPIDialect.cpp
index e062d62437a5b8..5ebf0434de20d8 100644
--- a/mlir/lib/Dialect/MPI/IR/MPIDialect.cpp
+++ b/mlir/lib/Dialect/MPI/IR/MPIDialect.cpp
@@ -8,6 +8,7 @@
 
 #include "mlir/Dialect/MPI/IR/MPI.h"
 #include "mlir/IR/DialectImplementation.h"
+#include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::mpi;
@@ -16,7 +17,9 @@ using namespace mlir::mpi;
 /// Tablegen Definitions
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Dialect/MPI/IR/MPIOpsDialect.cpp.inc"
+#include "mlir/Dialect/MPI/IR/MPI.cpp.inc"
+
+#include "mlir/Dialect/MPI/IR/MPIDialect.cpp.inc"
 
 void MPIDialect::initialize() {
   addOperations<
@@ -31,10 +34,13 @@ void MPIDialect::initialize() {
 
   addAttributes<
 #define GET_ATTRDEF_LIST
-#include "mlir/Dialect/MPI/IR/MPIAttributes.cpp.inc"
+#include "mlir/Dialect/MPI/IR/MPIAttrDefs.cpp.inc"
       >();
 }
 
+//===----------------------------------------------------------------------===//
+// TableGen'd dialect, type, and op definitions
+//===----------------------------------------------------------------------===//
 
 #define GET_TYPEDEF_CLASSES
 #include "mlir/Dialect/MPI/IR/MPITypesGen.cpp.inc"
@@ -42,4 +48,7 @@ void MPIDialect::initialize() {
 #include "mlir/Dialect/MPI/IR/MPIEnums.cpp.inc"
 
 #define GET_ATTRDEF_CLASSES
-#include "mlir/Dialect/MPI/IR/MPIAttributes.cpp.inc"
+#include "mlir/Dialect/MPI/IR/MPIAttrDefs.cpp.inc"
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/MPI/IR/MPIOps.cpp.inc"

From aa8638f65cc73cbb46af333344231856559f0f3f Mon Sep 17 00:00:00 2001
From: Anton Lydike <me@antonlydike.de>
Date: Thu, 14 Dec 2023 10:44:48 +0100
Subject: [PATCH 4/6] now further along the journey

---
 mlir/include/mlir/Dialect/MPI/IR/MPI.h        |  3 ++-
 mlir/include/mlir/Dialect/MPI/IR/MPI.td       | 10 ++++-----
 mlir/include/mlir/Dialect/MPI/IR/MPIOps.td    | 21 +++++++++++++++++++
 mlir/lib/Dialect/MPI/IR/CMakeLists.txt        |  2 +-
 .../MPI/IR/{MPIDialect.cpp => MPI.cpp}        |  4 +++-
 mlir/lib/Dialect/MPI/IR/MPIOps.cpp            |  2 ++
 mlir/test/Dialect/MPI/ops.mlir                |  8 +++----
 7 files changed, 37 insertions(+), 13 deletions(-)
 rename mlir/lib/Dialect/MPI/IR/{MPIDialect.cpp => MPI.cpp} (91%)

diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPI.h b/mlir/include/mlir/Dialect/MPI/IR/MPI.h
index 1f71f3d7d70809..f06b911ce3fe31 100644
--- a/mlir/include/mlir/Dialect/MPI/IR/MPI.h
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPI.h
@@ -1,4 +1,4 @@
-//===- MPI.h - MPI dialect ----------------------------*- C++-*-==//
+//===- MPI.h - MPI dialect ----------------------------------------*- C++-*-==//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -8,6 +8,7 @@
 #ifndef MLIR_DIALECT_MPI_IR_MPI_H_
 #define MLIR_DIALECT_MPI_IR_MPI_H_
 
+#include "mlir/Bytecode/BytecodeOpInterface.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPI.td b/mlir/include/mlir/Dialect/MPI/IR/MPI.td
index 8587fb42566c2c..afff8b90d6b57d 100644
--- a/mlir/include/mlir/Dialect/MPI/IR/MPI.td
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPI.td
@@ -1,4 +1,4 @@
-//===- MPIBase.td - Base defs for mpi dialect --*- tablegen -*-==//
+//===- MPIBase.td - Base defs for mpi dialect ---------------*- tablegen -*-==//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -33,7 +33,7 @@ def MPI_Dialect : Dialect {
   }];
 
   let usePropertiesForAttributes = 1;
-  let useDefaultAttributePrinterParser = 1;
+  // let useDefaultAttributePrinterParser = 1;
   let useDefaultTypePrinterParser = 1;
 }
 
@@ -41,7 +41,7 @@ def MPI_Dialect : Dialect {
 // Error classes enum:
 //===----------------------------------------------------------------------===//
 
-/*
+
 
 def MpiCodeSuccess : I32EnumAttrCase<"MPI_SUCCESS", 0, "MPI_SUCCESS">;
 def MpiCodeErrAccess : I32EnumAttrCase<"MPI_ERR_ACCESS", 1, "MPI_ERR_ACCESS">;
@@ -107,7 +107,7 @@ def MpiCodeErrValueTooLarge : I32EnumAttrCase<"MPI_ERR_VALUE_TOO_LARGE", 60, "MP
 def MpiCodeErrWin : I32EnumAttrCase<"MPI_ERR_WIN", 61, "MPI_ERR_WIN">;
 def MpiCodeErrLastcode : I32EnumAttrCase<"MPI_ERR_LASTCODE", 62, "MPI_ERR_LASTCODE">;
 
-def MpiErrorClasses : I32EnumAttr<"MpiErrorClass", 
+def MpiErrorClass : I32EnumAttr<"MpiErrorClassEnum", 
                                 "MPI error class name",
   [ MpiCodeSuccess
     ,MpiCodeErrAccess
@@ -176,6 +176,4 @@ def MpiErrorClasses : I32EnumAttr<"MpiErrorClass",
   let cppNamespace = "::mlir::mpi";
 }
 
-*/
-
 #endif // MLIR_DIALECT_MPI_IR_MPI
diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td b/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
index 9423628a0f696c..ca280fee4865ce 100644
--- a/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
@@ -122,4 +122,25 @@ def MPI_FinalizeOp : MPI_Op<"finalize", [
   let assemblyFormat = "attr-dict";
 }
 
+
+//===----------------------------------------------------------------------===//
+// FinalizeOp
+//===----------------------------------------------------------------------===//
+
+def MPI_RetvalCheck : MPI_Op<"retval_check", [
+
+]> {
+  let summary = "Check an MPI return value against an error class";
+  let description = [{
+    
+  }];
+
+  let arguments = (
+    ins MPI_Retval:$val,
+    MpiErrorClass:$errclass
+  );
+
+  let assemblyFormat = "$val `=` $errclass attr-dict";
+}
+
 #endif // MPI_MLIR_IR_MPIOPS
diff --git a/mlir/lib/Dialect/MPI/IR/CMakeLists.txt b/mlir/lib/Dialect/MPI/IR/CMakeLists.txt
index ed5b21216ec4e8..29d0b2379da747 100644
--- a/mlir/lib/Dialect/MPI/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/MPI/IR/CMakeLists.txt
@@ -1,6 +1,6 @@
 add_mlir_dialect_library(MLIRMPIDialect
   MPIOps.cpp
-  MPIDialect.cpp
+  MPI.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/MPI
diff --git a/mlir/lib/Dialect/MPI/IR/MPIDialect.cpp b/mlir/lib/Dialect/MPI/IR/MPI.cpp
similarity index 91%
rename from mlir/lib/Dialect/MPI/IR/MPIDialect.cpp
rename to mlir/lib/Dialect/MPI/IR/MPI.cpp
index 5ebf0434de20d8..6c5f69febcd63d 100644
--- a/mlir/lib/Dialect/MPI/IR/MPIDialect.cpp
+++ b/mlir/lib/Dialect/MPI/IR/MPI.cpp
@@ -1,4 +1,4 @@
-//===- MPIDialect.cpp - MPI dialect implementation ------------===//
+//===- MPI.cpp - MPI dialect implementation -------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -7,6 +7,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/MPI/IR/MPI.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "llvm/ADT/TypeSwitch.h"
 
diff --git a/mlir/lib/Dialect/MPI/IR/MPIOps.cpp b/mlir/lib/Dialect/MPI/IR/MPIOps.cpp
index 84b27c2dae419b..5827fabe664d4c 100644
--- a/mlir/lib/Dialect/MPI/IR/MPIOps.cpp
+++ b/mlir/lib/Dialect/MPI/IR/MPIOps.cpp
@@ -7,6 +7,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/MPI/IR/MPI.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
 
 using namespace mlir;
 using namespace mlir::mpi;
diff --git a/mlir/test/Dialect/MPI/ops.mlir b/mlir/test/Dialect/MPI/ops.mlir
index a03d44b9b27013..ed0d2485712e23 100644
--- a/mlir/test/Dialect/MPI/ops.mlir
+++ b/mlir/test/Dialect/MPI/ops.mlir
@@ -2,16 +2,16 @@
 // RUN: mlir-opt %s --mlir-print-op-generic | mlir-opt | FileCheck %s
 
 func.func @mpi_test(%ref : memref<100xf32>) -> () {
-    // CHECK: mpi.init
-    mpi.init
+    // CHECK: %0 = mpi.init : !mpi.retval
+    %err = mpi.init : !mpi.retval
 
     // CHECK-NEXT: mpi.comm_rank : i32
     %rank = mpi.comm_rank : i32
 
-    // CHECK-NEXT: mpi.send(%arg0, %0, %0) : memref<100xf32>, i32, i32
+    // CHECK-NEXT: mpi.send(%arg0, %1, %1) : memref<100xf32>, i32, i32
     mpi.send(%ref, %rank, %rank) : memref<100xf32>, i32, i32
 
-    // CHECK-NEXT: mpi.recv(%arg0, %0, %0) : memref<100xf32>, i32, i32
+    // CHECK-NEXT: mpi.recv(%arg0, %1, %1) : memref<100xf32>, i32, i32
     mpi.recv(%ref, %rank, %rank) : memref<100xf32>, i32, i32
 
     // CHECK-NEXT: mpi.finalize

From 77b3cadff9e21d575fde4c9ffe7c65ddf7f27998 Mon Sep 17 00:00:00 2001
From: Anton Lydike <me@antonlydike.de>
Date: Thu, 14 Dec 2023 17:04:22 +0100
Subject: [PATCH 5/6] fixed first version

---
 mlir/include/mlir/Dialect/MPI/IR/MPI.td    |  8 ++-
 mlir/include/mlir/Dialect/MPI/IR/MPIOps.td | 76 ++++++++++++++++++----
 mlir/test/Dialect/MPI/ops.mlir             | 19 ++++--
 3 files changed, 82 insertions(+), 21 deletions(-)

diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPI.td b/mlir/include/mlir/Dialect/MPI/IR/MPI.td
index afff8b90d6b57d..f109260cdac59f 100644
--- a/mlir/include/mlir/Dialect/MPI/IR/MPI.td
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPI.td
@@ -33,7 +33,7 @@ def MPI_Dialect : Dialect {
   }];
 
   let usePropertiesForAttributes = 1;
-  // let useDefaultAttributePrinterParser = 1;
+  let useDefaultAttributePrinterParser = 1;
   let useDefaultTypePrinterParser = 1;
 }
 
@@ -107,7 +107,7 @@ def MpiCodeErrValueTooLarge : I32EnumAttrCase<"MPI_ERR_VALUE_TOO_LARGE", 60, "MP
 def MpiCodeErrWin : I32EnumAttrCase<"MPI_ERR_WIN", 61, "MPI_ERR_WIN">;
 def MpiCodeErrLastcode : I32EnumAttrCase<"MPI_ERR_LASTCODE", 62, "MPI_ERR_LASTCODE">;
 
-def MpiErrorClass : I32EnumAttr<"MpiErrorClassEnum", 
+def MpiErrorClassEnum : I32EnumAttr<"MpiErrorClassEnum", 
                                 "MPI error class name",
   [ MpiCodeSuccess
     ,MpiCodeErrAccess
@@ -176,4 +176,8 @@ def MpiErrorClass : I32EnumAttr<"MpiErrorClassEnum",
   let cppNamespace = "::mlir::mpi";
 }
 
+def MpiErrorClassAttr : EnumAttr<MPI_Dialect, MpiErrorClassEnum, "errclass"> {
+  let assemblyFormat = "`<` $value `>`";
+}
+
 #endif // MLIR_DIALECT_MPI_IR_MPI
diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td b/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
index ca280fee4865ce..1b96eb65a0d4bf 100644
--- a/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPIOps.td
@@ -27,6 +27,9 @@ def MPI_InitOp : MPI_Op<"init", []> {
     please consult with the MPI specification on these).
 
     Passing &argc, &argv is not supported currently.
+
+    This operation can optionally return an `!mpi.retval` value that can be used
+    to check for errors.
   }];
 
   let results = (outs Optional<MPI_Retval>:$retval);
@@ -45,12 +48,17 @@ def MPI_CommRankOp : MPI_Op<"comm_rank", [
                 "`MPI_Comm_rank(MPI_COMM_WORLD, &rank)`";
   let description = [{
     Communicators other than `MPI_COMM_WORLD` are not supprted for now.
-    Inspecting the functions return value (error code) is also not supported.
+
+    This operation can optionally return an `!mpi.retval` value that can be used
+    to check for errors.
   }];
 
-  let results = (outs I32 : $result);
+  let results = (
+    outs Optional<MPI_Retval> : $retval,
+    I32 : $rank
+  );
 
-  let assemblyFormat = "attr-dict `:` type($result)";
+  let assemblyFormat = "attr-dict `:` type(results)";
 }
 
 //===----------------------------------------------------------------------===//
@@ -68,13 +76,18 @@ def MPI_SendOp : MPI_Op<"send", [
     multiple sends and receives between the same ranks.
 
     Communicators other than `MPI_COMM_WORLD` are not supprted for now.
-    Inspecting the functions return value (error code) is also not supported.
+
+    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($ref) `,` type($tag) `,` type($rank)"
+                       "(`->` type($retval)^)?";
 }
 
 //===----------------------------------------------------------------------===//
@@ -93,13 +106,18 @@ def MPI_RecvOp : MPI_Op<"recv", [
 
     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.
-    Inspecting the functions return value (error code) is also not supported.
+
+    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($ref) `,` type($tag) `,` type($rank)"
+                       "(`->` type($retval)^)?";
 }
 
 
@@ -116,18 +134,21 @@ def MPI_FinalizeOp : MPI_Op<"finalize", [
     (excpet for MPI_Get_version, MPI_Initialized, and MPI_Finalized).
     Notably, MPI_Init cannot be called again in the same program.
 
-    Inspecting the functions return value (error code) is not supported.
+    This operation can optionally return an `!mpi.retval` value that can be used
+    to check for errors.
   }];
 
-  let assemblyFormat = "attr-dict";
+  let results = (outs Optional<MPI_Retval>:$retval);
+
+  let assemblyFormat = "attr-dict (`:` type($retval)^)?";
 }
 
 
 //===----------------------------------------------------------------------===//
-// FinalizeOp
+// RetvalCheckOp
 //===----------------------------------------------------------------------===//
 
-def MPI_RetvalCheck : MPI_Op<"retval_check", [
+def MPI_RetvalCheckOp : MPI_Op<"retval_check", [
 
 ]> {
   let summary = "Check an MPI return value against an error class";
@@ -137,10 +158,39 @@ def MPI_RetvalCheck : MPI_Op<"retval_check", [
 
   let arguments = (
     ins MPI_Retval:$val,
-    MpiErrorClass:$errclass
+    MpiErrorClassAttr:$errclass
+  );
+
+  let results = (
+    outs I1:$res
+  );
+
+  let assemblyFormat = "$val `=` $errclass attr-dict `:` type($res)";
+}
+
+
+
+//===----------------------------------------------------------------------===//
+// RetvalCheckOp
+//===----------------------------------------------------------------------===//
+
+def MPI_ErrorClassOp : MPI_Op<"error_class", [
+
+]> {
+  let summary = "Get the error class from an error code, equivalent to the `MPI_Error_class` functoin.";
+  let description = [{
+    
+  }];
+
+  let arguments = (
+    ins MPI_Retval:$val
+  );
+
+  let results = (
+    outs MPI_Retval:$errclass
   );
 
-  let assemblyFormat = "$val `=` $errclass attr-dict";
+  let assemblyFormat = "$val attr-dict `:` type($val)";
 }
 
 #endif // MPI_MLIR_IR_MPIOPS
diff --git a/mlir/test/Dialect/MPI/ops.mlir b/mlir/test/Dialect/MPI/ops.mlir
index ed0d2485712e23..11595a9850d963 100644
--- a/mlir/test/Dialect/MPI/ops.mlir
+++ b/mlir/test/Dialect/MPI/ops.mlir
@@ -5,17 +5,24 @@ func.func @mpi_test(%ref : memref<100xf32>) -> () {
     // CHECK: %0 = mpi.init : !mpi.retval
     %err = mpi.init : !mpi.retval
 
-    // CHECK-NEXT: mpi.comm_rank : i32
-    %rank = mpi.comm_rank : i32
+    // CHECK-NEXT: %retval, %rank = mpi.comm_rank : !mpi.retval, i32
+    %retval, %rank = mpi.comm_rank : !mpi.retval, i32
 
-    // CHECK-NEXT: mpi.send(%arg0, %1, %1) : memref<100xf32>, i32, i32
+    // CHECK-NEXT: mpi.send(%arg0, %rank, %rank) : memref<100xf32>, i32, i32
     mpi.send(%ref, %rank, %rank) : memref<100xf32>, i32, i32
 
-    // CHECK-NEXT: mpi.recv(%arg0, %1, %1) : memref<100xf32>, i32, i32
+    // CHECK-NEXT: mpi.recv(%arg0, %rank, %rank) : memref<100xf32>, i32, i32
     mpi.recv(%ref, %rank, %rank) : memref<100xf32>, i32, i32
 
-    // CHECK-NEXT: mpi.finalize
-    mpi.finalize
+    // CHECK-NEXT: %1 = mpi.finalize : !mpi.retval
+    %rval = mpi.finalize : !mpi.retval
 
+    // CHECK-NEXT: %2 = mpi.retval_check %retval = <MPI_SUCCESS> : i1
+    %res = mpi.retval_check %retval = <MPI_SUCCESS> : i1
+
+    // CHECK-NEXT: %3 = mpi.error_class %0 : !mpi.retval
+    %errclass = mpi.error_class %err : !mpi.retval
+
+    // CHECK-NEXT: return
     func.return
 }

From 6e361de8ddf322646d07eae063d53cae307a9101 Mon Sep 17 00:00:00 2001
From: Anton Lydike <me@antonlydike.de>
Date: Thu, 14 Dec 2023 17:56:57 +0100
Subject: [PATCH 6/6] add more tests, it's ready now

---
 mlir/include/mlir/Dialect/MPI/IR/MPITypes.td |  2 +-
 mlir/test/Dialect/MPI/invalid.mlir           | 28 ++++++++++++++++++++
 mlir/test/Dialect/MPI/ops.mlir               | 14 +++++++---
 3 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Dialect/MPI/IR/MPITypes.td b/mlir/include/mlir/Dialect/MPI/IR/MPITypes.td
index 4cf07ca37267f2..1f1ae3144e6cea 100644
--- a/mlir/include/mlir/Dialect/MPI/IR/MPITypes.td
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPITypes.td
@@ -30,7 +30,7 @@ class MPI_Type<string name, string typeMnemonic, list<Trait> traits = []>
 //===----------------------------------------------------------------------===//
 
 def MPI_Retval : MPI_Type<"Retval", "retval"> {
-  let summary = "The return value of an MPI call.";
+  let summary = "MPI function call return value";
   let description = [{
     This type represents a value returned from an MPI call. This value can be
     MPI_SUCCESS, MPI_ERR_IN_STATUS, or any error code.
diff --git a/mlir/test/Dialect/MPI/invalid.mlir b/mlir/test/Dialect/MPI/invalid.mlir
index 9ce405a445f9d6..1da154c7a58126 100644
--- a/mlir/test/Dialect/MPI/invalid.mlir
+++ b/mlir/test/Dialect/MPI/invalid.mlir
@@ -20,3 +20,31 @@ func.func @mpi_test(%ref : !llvm.ptr, %rank: i32) -> () {
 
     return
 }
+
+// -----
+
+func.func @mpi_test(%ref : memref<100xf32>, %rank: i32) -> () {
+    // expected-error @+1 {{'mpi.recv' op result #0 must be MPI function call return value, but got 'i32'}}
+    %res = mpi.recv(%ref, %rank, %rank) : memref<100xf32>, i32, i32 -> i32
+
+    return
+}
+
+// -----
+
+func.func @mpi_test(%ref : memref<100xf32>, %rank: i32) -> () {
+    // expected-error @+1 {{'mpi.send' op result #0 must be MPI function call return value, but got 'i32'}}
+    %res = mpi.send(%ref, %rank, %rank) : memref<100xf32>, i32, i32 -> i32
+
+    return
+}
+
+// -----
+
+func.func @mpi_test(%retval: !mpi.retval) -> () {
+    // expected-error @+2 {{custom op 'mpi.retval_check' expected ::mlir::mpi::MpiErrorClassEnum}}
+    // expected-error @+1 {{custom op 'mpi.retval_check' failed to parse MpiErrorClassAttr parameter 'value'}}
+    %res = mpi.retval_check %retval = <MPI_ERR_DOES_NOT_EXIST>
+
+    return
+}
diff --git a/mlir/test/Dialect/MPI/ops.mlir b/mlir/test/Dialect/MPI/ops.mlir
index 11595a9850d963..f7ffa23f7e82ea 100644
--- a/mlir/test/Dialect/MPI/ops.mlir
+++ b/mlir/test/Dialect/MPI/ops.mlir
@@ -2,6 +2,8 @@
 // RUN: mlir-opt %s --mlir-print-op-generic | mlir-opt | FileCheck %s
 
 func.func @mpi_test(%ref : memref<100xf32>) -> () {
+    // Note: the !mpi.retval result is optional on all operations except mpi.error_class
+
     // CHECK: %0 = mpi.init : !mpi.retval
     %err = mpi.init : !mpi.retval
 
@@ -11,16 +13,22 @@ func.func @mpi_test(%ref : memref<100xf32>) -> () {
     // CHECK-NEXT: mpi.send(%arg0, %rank, %rank) : memref<100xf32>, i32, i32
     mpi.send(%ref, %rank, %rank) : memref<100xf32>, i32, i32
 
+    // CHECK-NEXT: %1 = mpi.send(%arg0, %rank, %rank) : memref<100xf32>, i32, i32 -> !mpi.retval
+    %err2 = mpi.send(%ref, %rank, %rank) : memref<100xf32>, i32, i32 -> !mpi.retval
+
     // CHECK-NEXT: mpi.recv(%arg0, %rank, %rank) : memref<100xf32>, i32, i32
     mpi.recv(%ref, %rank, %rank) : memref<100xf32>, i32, i32
 
-    // CHECK-NEXT: %1 = mpi.finalize : !mpi.retval
+    // CHECK-NEXT: %2 = mpi.recv(%arg0, %rank, %rank) : memref<100xf32>, i32, i32 -> !mpi.retval
+    %err3 = mpi.recv(%ref, %rank, %rank) : memref<100xf32>, i32, i32 -> !mpi.retval
+
+    // CHECK-NEXT: %3 = mpi.finalize : !mpi.retval
     %rval = mpi.finalize : !mpi.retval
 
-    // CHECK-NEXT: %2 = mpi.retval_check %retval = <MPI_SUCCESS> : i1
+    // CHECK-NEXT: %4 = mpi.retval_check %retval = <MPI_SUCCESS> : i1
     %res = mpi.retval_check %retval = <MPI_SUCCESS> : i1
 
-    // CHECK-NEXT: %3 = mpi.error_class %0 : !mpi.retval
+    // CHECK-NEXT: %5 = mpi.error_class %0 : !mpi.retval
     %errclass = mpi.error_class %err : !mpi.retval
 
     // CHECK-NEXT: return