Skip to content

Commit 3f6ab73

Browse files
committed
very initial MPI dialect
1 parent 985a72b commit 3f6ab73

File tree

15 files changed

+424
-0
lines changed

15 files changed

+424
-0
lines changed

flake.lock

Lines changed: 104 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

flake.nix

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
{
2+
inputs = {
3+
nix-vscode-extensions.url = "github:nix-community/nix-vscode-extensions";
4+
flake-utils.follows = "nix-vscode-extensions/flake-utils";
5+
nixpkgs.follows = "nix-vscode-extensions/nixpkgs";
6+
};
7+
8+
outputs = inputs:
9+
inputs.flake-utils.lib.eachDefaultSystem
10+
(system:
11+
let
12+
pkgs = inputs.nixpkgs.legacyPackages.${system};
13+
extensions = inputs.nix-vscode-extensions.extensions.${system};
14+
inherit (pkgs) vscode-with-extensions vscodium;
15+
16+
packages.default =
17+
vscode-with-extensions.override {
18+
vscode = vscodium;
19+
vscodeExtensions = with extensions.vscode-marketplace; [
20+
# ms-vscode.cpptools
21+
llvm-vs-code-extensions.vscode-clangd
22+
llvm-vs-code-extensions.vscode-mlir
23+
eamodio.gitlens
24+
]
25+
++ [ pkgs.vscode-extensions.ms-vscode.cpptools ];
26+
};
27+
28+
devShells.default = pkgs.mkShell {
29+
buildInputs = [
30+
packages.default
31+
pkgs.git pkgs.ripgrep pkgs.gdb
32+
pkgs.cmake pkgs.ninja pkgs.clang pkgs.lld
33+
pkgs.fish pkgs.neovim pkgs.git pkgs.mold pkgs.libclang
34+
];
35+
shellHook = ''
36+
printf "VSCodium with extensions:\n"
37+
'';
38+
};
39+
in
40+
{
41+
inherit packages devShells;
42+
});
43+
}
44+

mlir/include/mlir/Dialect/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ add_subdirectory(LLVMIR)
2020
add_subdirectory(Math)
2121
add_subdirectory(MemRef)
2222
add_subdirectory(MLProgram)
23+
add_subdirectory(MPI)
2324
add_subdirectory(NVGPU)
2425
add_subdirectory(OpenACC)
2526
add_subdirectory(OpenMP)
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
add_subdirectory(IR)
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
set(LLVM_TARGET_DEFINITIONS MPIOps.td)
2+
add_mlir_dialect(MPIOps mpi)
3+
add_mlir_doc(MPIOps MPIOps Dialects/ -gen-dialect-doc)
Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
//===- MPI.h - MPI dialect ----------------------------*- C++-*-==//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
#ifndef MLIR_DIALECT_MPI_IR_MPI_H_
9+
#define MLIR_DIALECT_MPI_IR_MPI_H_
10+
11+
#include "mlir/IR/Dialect.h"
12+
#include "mlir/IR/FunctionInterfaces.h"
13+
#include "mlir/IR/OpDefinition.h"
14+
#include "mlir/IR/OpImplementation.h"
15+
#include "mlir/IR/RegionKindInterface.h"
16+
#include "mlir/IR/SymbolTable.h"
17+
#include "mlir/Interfaces/CallInterfaces.h"
18+
#include "mlir/Interfaces/ControlFlowInterfaces.h"
19+
#include "mlir/Interfaces/SideEffectInterfaces.h"
20+
21+
//===----------------------------------------------------------------------===//
22+
// MPIDialect
23+
//===----------------------------------------------------------------------===//
24+
25+
#include "mlir/Dialect/MPI/IR/MPIOpsDialect.h.inc"
26+
27+
//===----------------------------------------------------------------------===//
28+
// MPI Dialect Operations
29+
//===----------------------------------------------------------------------===//
30+
31+
#define GET_OP_CLASSES
32+
#include "mlir/Dialect/MPI/IR/MPIOps.h.inc"
33+
34+
#endif // MLIR_DIALECT_MPI_IR_MPI_H_
Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
//===- MPIBase.td - Base defs for mpi dialect --*- tablegen -*-==//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#ifndef MPI_BASE
10+
#define MPI_BASE
11+
12+
include "mlir/IR/OpBase.td"
13+
14+
def MPI_Dialect : Dialect {
15+
let name = "mpi";
16+
let cppNamespace = "::mlir::mpi";
17+
let description = [{
18+
TODO
19+
20+
This dialect is under active development, and while stability is an
21+
eventual goal, it is not guaranteed at this juncture. Given the early state,
22+
it is recommended to inquire further prior to using this dialect.
23+
}];
24+
let dependentDialects = ["arith::ArithDialect"];
25+
26+
let usePropertiesForAttributes = 1;
27+
}
28+
29+
#endif // MPI_BASE
Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
//===- MPI.td - Message Passing Interface Ops ---------*- tablegen -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
9+
#ifndef MPI_OPS
10+
#define MPI_OPS
11+
12+
include "mlir/Dialect/MPI/IR/MPIBase.td"
13+
include "mlir/Interfaces/CallInterfaces.td"
14+
include "mlir/Interfaces/ControlFlowInterfaces.td"
15+
include "mlir/Interfaces/SideEffectInterfaces.td"
16+
include "mlir/IR/FunctionInterfaces.td"
17+
include "mlir/IR/RegionKindInterface.td"
18+
include "mlir/IR/SymbolInterfaces.td"
19+
20+
class MPI_Op<string mnemonic, list<Trait> traits = []> :
21+
Op<MPI_Dialect, mnemonic, traits>;
22+
23+
//===----------------------------------------------------------------------===//
24+
// InitOp
25+
//===----------------------------------------------------------------------===//
26+
27+
def MPI_InitOp : MPI_Op<"init", [
28+
29+
]> {
30+
let summary = "Initialize the MPI library, equivalent to `MPI_Init(NULL, NULL)`";
31+
let description = [{
32+
Passing &argc, &argv is not supported currently, also inspecting the return value is not supported.
33+
}];
34+
35+
let assemblyFormat = "attr-dict";
36+
}
37+
38+
//===----------------------------------------------------------------------===//
39+
// CommRankOp
40+
//===----------------------------------------------------------------------===//
41+
42+
def MPI_CommRankOp : MPI_Op<"comm_rank", [
43+
44+
]> {
45+
let summary = "Get the current rank, equivalent to `MPI_Comm_rank(MPI_COMM_WORLD, &rank)`";
46+
let description = [{
47+
Communicators other than `MPI_COMM_WORLD` are not supprted for now.
48+
}];
49+
50+
let results = (outs
51+
I32:$result
52+
);
53+
54+
let assemblyFormat = "attr-dict `:` type($result)";
55+
}
56+
57+
58+
//===----------------------------------------------------------------------===//
59+
// SendOp
60+
//===----------------------------------------------------------------------===//
61+
62+
def MPI_SendOp : MPI_Op<"send", [
63+
64+
]> {
65+
let summary = "Equivalent to `MPI_Send(ptr, size, dtype, dest, tag, MPI_COMM_WORLD)`";
66+
let description = [{
67+
Communicators other than `MPI_COMM_WORLD` are not supprted for now.
68+
}];
69+
70+
let arguments = (ins
71+
AnyMemRef:$ref,
72+
I32:$tag,
73+
I32:$rank
74+
);
75+
76+
let assemblyFormat = "attr-dict $ref `:` type($ref) `,` $tag `:` type($tag) `,` $rank `:` type($rank)";
77+
}
78+
79+
//===----------------------------------------------------------------------===//
80+
// RecvOp
81+
//===----------------------------------------------------------------------===//
82+
83+
def MPI_RecvOp : MPI_Op<"recv", [
84+
85+
]> {
86+
let summary = "Equivalent to `MPI_Recv(ptr, size, dtype, dest, tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE)`";
87+
let description = [{
88+
Communicators other than `MPI_COMM_WORLD` are not supprted for now.
89+
}];
90+
91+
let arguments = (ins
92+
AnyMemRef:$ref,
93+
I32:$tag,
94+
I32:$rank
95+
);
96+
97+
let assemblyFormat = "attr-dict $ref `:` type($ref) `,` $tag `:` type($tag) `,` $rank `:` type($rank)";
98+
}
99+
100+
//===----------------------------------------------------------------------===//
101+
// FinalizeOp
102+
//===----------------------------------------------------------------------===//
103+
104+
def MPI_FinalizeOp : MPI_Op<"finalize", [
105+
106+
]> {
107+
let summary = "Finalize the MPI library, equivalent to `MPI_Finalize()`";
108+
let description = [{
109+
110+
}];
111+
112+
let assemblyFormat = "attr-dict";
113+
}
114+
115+
#endif // MPI_OPS

mlir/include/mlir/InitAllDialects.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
#include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h"
4949
#include "mlir/Dialect/Linalg/Transforms/TilingInterfaceImpl.h"
5050
#include "mlir/Dialect/MLProgram/IR/MLProgram.h"
51+
#include "mlir/Dialect/MPI/IR/MPI.h"
5152
#include "mlir/Dialect/Math/IR/Math.h"
5253
#include "mlir/Dialect/MemRef/IR/MemRef.h"
5354
#include "mlir/Dialect/MemRef/IR/MemRefMemorySlot.h"
@@ -117,6 +118,7 @@ inline void registerAllDialects(DialectRegistry &registry) {
117118
math::MathDialect,
118119
memref::MemRefDialect,
119120
ml_program::MLProgramDialect,
121+
mpi::MPIDialect,
120122
nvgpu::NVGPUDialect,
121123
NVVM::NVVMDialect,
122124
omp::OpenMPDialect,

mlir/lib/Dialect/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ add_subdirectory(LLVMIR)
2020
add_subdirectory(Math)
2121
add_subdirectory(MemRef)
2222
add_subdirectory(MLProgram)
23+
add_subdirectory(MPI)
2324
add_subdirectory(NVGPU)
2425
add_subdirectory(OpenACC)
2526
add_subdirectory(OpenMP)

0 commit comments

Comments
 (0)