[Mlir-commits] [mlir] [mlir] Initial patch to add an MPI dialect (PR #68892)
Anton Lydike
llvmlistbot at llvm.org
Mon Nov 6 05:57:44 PST 2023
https://github.com/AntonLydike updated https://github.com/llvm/llvm-project/pull/68892
>From 5e0ff463f0adf03939a5f0b92821cc07767b15b6 Mon Sep 17 00:00:00 2001
From: Anton Lydike <me at antonlydike.de>
Date: Thu, 12 Oct 2023 14:55:52 +0100
Subject: [PATCH] [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 | 39 ++++++
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, 305 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 1c4569ecfa58485..9788e24e4a1d91a 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 000000000000000..f33061b2d87cffc
--- /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 000000000000000..3b7b76240d375f1
--- /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 000000000000000..0cffe8edda44312
--- /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 000000000000000..e3d0af069636a0f
--- /dev/null
+++ b/mlir/include/mlir/Dialect/MPI/IR/MPIBase.td
@@ -0,0 +1,39 @@
+//===- 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 a targetable dialect, that 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.
+
+ The current version of the dialect models a very limited part of the MPI Library,
+ often simplifying operations. This was a conscious decision that enables us to
+ incrementally add features as they are needed.
+
+ 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 000000000000000..70e6ec1b623cfdf
--- /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 00f400aab5d50a0..42087994d0f0c89 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 ®istry) {
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 68776a695cac4d4..c72107939cf42bb 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 000000000000000..f33061b2d87cffc
--- /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 000000000000000..870f034e5ccca73
--- /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 000000000000000..19ed5d8207af459
--- /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 000000000000000..84b27c2dae419b2
--- /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 000000000000000..9ce405a445f9d6c
--- /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 000000000000000..a03d44b9b270133
--- /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
+}
More information about the Mlir-commits
mailing list