Skip to content

Commit fd30503

Browse files
fschlimbfschlimb
committed
lowering mpi.comm_split
1 parent 2003e17 commit fd30503

File tree

4 files changed

+86
-9
lines changed

4 files changed

+86
-9
lines changed

‎mlir/include/mlir/Dialect/MPI/IR/MPIOps.td

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ def MPI_CommSizeOp : MPI_Op<"comm_size", []> {
100100
// CommSplitOp
101101
//===----------------------------------------------------------------------===//
102102

103-
def MPI_CommSplit : MPI_Op<"comm_split", []> {
103+
def MPI_CommSplitOp : MPI_Op<"comm_split", []> {
104104
let summary = "Partition the group associated to the given communicator into "
105105
"disjoint subgroups";
106106
let description = [{
@@ -121,7 +121,6 @@ def MPI_CommSplit : MPI_Op<"comm_split", []> {
121121
);
122122

123123
let assemblyFormat = "`(` $comm `,` $color `,` $key `)` attr-dict `:` "
124-
"type($color) `,` type($key) `->` "
125124
"type(results)";
126125
}
127126

‎mlir/lib/Conversion/MPIToLLVM/MPIToLLVM.cpp

Lines changed: 57 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -335,6 +335,58 @@ struct CommWorldOpLowering : public ConvertOpToLLVMPattern<mpi::CommWorldOp> {
335335
}
336336
};
337337

338+
//===----------------------------------------------------------------------===//
339+
// CommSplitOpLowering
340+
//===----------------------------------------------------------------------===//
341+
342+
struct CommSplitOpLowering : public ConvertOpToLLVMPattern<mpi::CommSplitOp> {
343+
using ConvertOpToLLVMPattern::ConvertOpToLLVMPattern;
344+
345+
LogicalResult
346+
matchAndRewrite(mpi::CommSplitOp op, OpAdaptor adaptor,
347+
ConversionPatternRewriter &rewriter) const override {
348+
// grab a reference to the global module op:
349+
auto moduleOp = op->getParentOfType<ModuleOp>();
350+
auto mpiTraits = MPIImplTraits::get(moduleOp);
351+
Type i32 = rewriter.getI32Type();
352+
Type ptrType = LLVM::LLVMPointerType::get(op->getContext());
353+
Location loc = op.getLoc();
354+
355+
// get communicator
356+
Value comm = mpiTraits->castComm(loc, rewriter, adaptor.getComm());
357+
auto one = rewriter.create<LLVM::ConstantOp>(loc, i32, 1);
358+
auto outPtr =
359+
rewriter.create<LLVM::AllocaOp>(loc, ptrType, comm.getType(), one);
360+
361+
// int MPI_Comm_split(MPI_Comm comm, int color, int key, MPI_Comm * newcomm)
362+
auto funcType =
363+
LLVM::LLVMFunctionType::get(i32, {comm.getType(), i32, i32, ptrType});
364+
// get or create function declaration:
365+
LLVM::LLVMFuncOp funcDecl = getOrDefineFunction(moduleOp, loc, rewriter,
366+
"MPI_Comm_split", funcType);
367+
368+
auto callOp = rewriter.create<LLVM::CallOp>(
369+
loc, funcDecl,
370+
ValueRange{comm, adaptor.getColor(), adaptor.getKey(),
371+
outPtr.getRes()});
372+
373+
// load the communicator into a register
374+
auto res = rewriter.create<LLVM::LoadOp>(loc, i32, outPtr.getResult());
375+
376+
// if retval is checked, replace uses of retval with the results from the
377+
// call op
378+
SmallVector<Value> replacements;
379+
if (op.getRetval())
380+
replacements.push_back(callOp.getResult());
381+
382+
// replace op
383+
replacements.push_back(res.getRes());
384+
rewriter.replaceOp(op, replacements);
385+
386+
return success();
387+
}
388+
};
389+
338390
//===----------------------------------------------------------------------===//
339391
// CommRankOpLowering
340392
//===----------------------------------------------------------------------===//
@@ -367,7 +419,7 @@ struct CommRankOpLowering : public ConvertOpToLLVMPattern<mpi::CommRankOp> {
367419
LLVM::LLVMFuncOp initDecl = getOrDefineFunction(
368420
moduleOp, loc, rewriter, "MPI_Comm_rank", rankFuncType);
369421

370-
// replace init with function call
422+
// replace with function call
371423
auto one = rewriter.create<LLVM::ConstantOp>(loc, i32, 1);
372424
auto rankptr = rewriter.create<LLVM::AllocaOp>(loc, ptrType, i32, one);
373425
auto callOp = rewriter.create<LLVM::CallOp>(
@@ -543,8 +595,10 @@ void mpi::populateMPIToLLVMConversionPatterns(LLVMTypeConverter &converter,
543595
converter.addConversion([](mpi::CommType type) {
544596
return IntegerType::get(type.getContext(), 64);
545597
});
546-
patterns.add<CommRankOpLowering, CommWorldOpLowering, FinalizeOpLowering,
547-
InitOpLowering, SendOpLowering, RecvOpLowering>(converter);
598+
patterns
599+
.add<CommRankOpLowering, CommSplitOpLowering, CommWorldOpLowering,
600+
FinalizeOpLowering, InitOpLowering, SendOpLowering, RecvOpLowering>(
601+
converter);
548602
}
549603

550604
void mpi::registerConvertMPIToLLVMInterface(DialectRegistry &registry) {

‎mlir/test/Conversion/MPIToLLVM/ops.mlir

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
// COM: Test MPICH ABI
44
// CHECK: module attributes {mpi.dlti = #dlti.map<"MPI:Implementation" = "MPICH">} {
55
// CHECK: llvm.func @MPI_Finalize() -> i32
6+
// CHECK: llvm.func @MPI_Comm_split(i32, i32, i32, !llvm.ptr) -> i32
67
// CHECK: llvm.func @MPI_Recv(!llvm.ptr, i32, i32, i32, i32, i32, !llvm.ptr) -> i32
78
// CHECK: llvm.func @MPI_Send(!llvm.ptr, i32, i32, i32, i32, i32) -> i32
89
// CHECK: llvm.func @MPI_Comm_rank(i32, !llvm.ptr) -> i32
@@ -75,6 +76,17 @@ module attributes { mpi.dlti = #dlti.map<"MPI:Implementation" = "MPICH"> } {
7576
// CHECK: [[v47:%.*]] = llvm.inttoptr [[v46]] : i64 to !llvm.ptr
7677
// CHECK: [[v48:%.*]] = llvm.call @MPI_Recv([[v41]], [[v43]], [[v44]], [[v12]], [[v12]], [[comm_4]], [[v47]]) : (!llvm.ptr, i32, i32, i32, i32, i32, !llvm.ptr) -> i32
7778
%2 = mpi.recv(%arg0, %rank, %rank, %comm) : memref<100xf32>, i32, i32 -> !mpi.retval
79+
80+
// CHECK: [[v51:%.*]] = llvm.mlir.constant(10 : i32) : i32
81+
%color = arith.constant 10 : i32
82+
// CHECK: [[v52:%.*]] = llvm.mlir.constant(22 : i32) : i32
83+
%key = arith.constant 22 : i32
84+
// CHECK: [[v53:%.*]] = llvm.trunc [[comm]] : i64 to i32
85+
// CHECK: [[v54:%.*]] = llvm.mlir.constant(1 : i32) : i32
86+
// CHECK: [[v55:%.*]] = llvm.alloca [[v54]] x i32 : (i32) -> !llvm.ptr
87+
// CHECK: [[v56:%.*]] = llvm.call @MPI_Comm_split([[v53]], [[v51]], [[v52]], [[v55]]) : (i32, i32, i32, !llvm.ptr) -> i32
88+
// CHECK: [[v57:%.*]] = llvm.load [[v55]] : !llvm.ptr -> i32
89+
%split = mpi.comm_split(%comm, %color, %key) : !mpi.comm
7890

7991
// CHECK: [[v49:%.*]] = llvm.call @MPI_Finalize() : () -> i32
8092
%3 = mpi.finalize : !mpi.retval
@@ -88,6 +100,7 @@ module attributes { mpi.dlti = #dlti.map<"MPI:Implementation" = "MPICH"> } {
88100
// COM: Test OpenMPI ABI
89101
// CHECK: module attributes {mpi.dlti = #dlti.map<"MPI:Implementation" = "OpenMPI">} {
90102
// CHECK: llvm.func @MPI_Finalize() -> i32
103+
// CHECK: llvm.func @MPI_Comm_split(!llvm.ptr, i32, i32, !llvm.ptr) -> i32
91104
// CHECK: llvm.func @MPI_Recv(!llvm.ptr, i32, !llvm.ptr, i32, i32, !llvm.ptr, !llvm.ptr) -> i32
92105
// CHECK: llvm.func @MPI_Send(!llvm.ptr, i32, !llvm.ptr, i32, i32, !llvm.ptr) -> i32
93106
// CHECK: llvm.mlir.global external @ompi_mpi_float() {addr_space = 0 : i32} : !llvm.struct<"ompi_predefined_datatype_t", opaque>
@@ -163,6 +176,17 @@ module attributes { mpi.dlti = #dlti.map<"MPI:Implementation" = "OpenMPI"> } {
163176
// CHECK: [[v48:%.*]] = llvm.call @MPI_Recv([[v41]], [[v43]], [[v44]], [[v12]], [[v12]], [[v45]], [[v47]]) : (!llvm.ptr, i32, !llvm.ptr, i32, i32, !llvm.ptr, !llvm.ptr) -> i32
164177
%2 = mpi.recv(%arg0, %rank, %rank, %comm) : memref<100xf32>, i32, i32 -> !mpi.retval
165178

179+
// CHECK: [[v51:%.*]] = llvm.mlir.constant(10 : i32) : i32
180+
%color = arith.constant 10 : i32
181+
// CHECK: [[v52:%.*]] = llvm.mlir.constant(22 : i32) : i32
182+
%key = arith.constant 22 : i32
183+
// CHECK: [[v53:%.*]] = llvm.inttoptr [[comm]] : i64 to !llvm.ptr
184+
// CHECK: [[v54:%.*]] = llvm.mlir.constant(1 : i32) : i32
185+
// CHECK: [[v55:%.*]] = llvm.alloca [[v54]] x !llvm.ptr : (i32) -> !llvm.ptr
186+
// CHECK: [[v56:%.*]] = llvm.call @MPI_Comm_split([[v53]], [[v51]], [[v52]], [[v55]]) : (!llvm.ptr, i32, i32, !llvm.ptr) -> i32
187+
// CHECK: [[v57:%.*]] = llvm.load [[v55]] : !llvm.ptr -> i32
188+
%split = mpi.comm_split(%comm, %color, %key) : !mpi.comm
189+
166190
// CHECK: [[v49:%.*]] = llvm.call @MPI_Finalize() : () -> i32
167191
%3 = mpi.finalize : !mpi.retval
168192

‎mlir/test/Dialect/MPI/ops.mlir

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,11 @@ func.func @mpi_test(%ref : memref<100xf32>) -> () {
2323
// CHECK-NEXT: [[vretval_1:%.*]], [[vsize_2:%.*]] = mpi.comm_size([[v1]]) : !mpi.retval, i32
2424
%retval_0, %size_1 = mpi.comm_size(%comm) : !mpi.retval, i32
2525

26-
// CHECK-NEXT: [[vnewcomm:%.*]] = mpi.comm_split([[v1]], [[vrank]], [[vrank]]) : i32, i32 -> !mpi.comm
27-
%new_comm = mpi.comm_split(%comm, %rank, %rank) : i32, i32 -> !mpi.comm
26+
// CHECK-NEXT: [[vnewcomm:%.*]] = mpi.comm_split([[v1]], [[vrank]], [[vrank]]) : !mpi.comm
27+
%new_comm = mpi.comm_split(%comm, %rank, %rank) : !mpi.comm
2828

29-
// CHECK-NEXT: [[vretval_3:%.*]], [[vnewcomm_4:%.*]] = mpi.comm_split([[v1]], [[vrank]], [[vrank]]) : i32, i32 -> !mpi.retval, !mpi.comm
30-
%retval_1, %new_comm_1 = mpi.comm_split(%comm, %rank, %rank) : i32, i32 -> !mpi.retval, !mpi.comm
29+
// CHECK-NEXT: [[vretval_3:%.*]], [[vnewcomm_4:%.*]] = mpi.comm_split([[v1]], [[vrank]], [[vrank]]) : !mpi.retval, !mpi.comm
30+
%retval_1, %new_comm_1 = mpi.comm_split(%comm, %rank, %rank) : !mpi.retval, !mpi.comm
3131

3232
// CHECK-NEXT: mpi.send([[varg0]], [[vrank]], [[vrank]], [[v1]]) : memref<100xf32>, i32, i32
3333
mpi.send(%ref, %rank, %rank, %comm) : memref<100xf32>, i32, i32

0 commit comments

Comments
 (0)