Skip to content

Commit

Permalink
[mlir][OpenMP][IRBuilder] Add support for nowait on single construct
Browse files Browse the repository at this point in the history 
This patch adds the nowait parameter to `createSingle` in
OpenMPIRBuilder and handling for IR generation from OpenMP Dialect.

Also added tests for the same.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D122371
  • Loading branch information
@shraiysh
shraiysh committed Mar 24, 2022
1 parent dc58c3b commit 8722c12
Show file tree
Hide file tree
Showing 5 changed files with 155 additions and 11 deletions.
4 changes: 3 additions & 1 deletion llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -867,12 +867,14 @@ class OpenMPIRBuilder {
/// \param Loc The source location description.
/// \param BodyGenCB Callback that will generate the region code.
/// \param FiniCB Callback to finalize variable copies.
/// \param IsNowait If false, a barrier is emitted.
/// \param DidIt Local variable used as a flag to indicate 'single' thread
///
/// \returns The insertion position *after* the single call.
InsertPointTy createSingle(const LocationDescription &Loc,
BodyGenCallbackTy BodyGenCB,
FinalizeCallbackTy FiniCB, llvm::Value *DidIt);
FinalizeCallbackTy FiniCB, bool IsNowait,
llvm::Value *DidIt);

/// Generator for '#omp master'
///
Expand Down
20 changes: 13 additions & 7 deletions llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2739,10 +2739,9 @@ OpenMPIRBuilder::createCopyPrivate(const LocationDescription &Loc,
return Builder.saveIP();
}

OpenMPIRBuilder::InsertPointTy
OpenMPIRBuilder::createSingle(const LocationDescription &Loc,
BodyGenCallbackTy BodyGenCB,
FinalizeCallbackTy FiniCB, llvm::Value *DidIt) {
OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createSingle(
const LocationDescription &Loc, BodyGenCallbackTy BodyGenCB,
FinalizeCallbackTy FiniCB, bool IsNowait, llvm::Value *DidIt) {

if (!updateToLocation(Loc))
return Loc.IP;
Expand Down Expand Up @@ -2770,9 +2769,16 @@ OpenMPIRBuilder::createSingle(const LocationDescription &Loc,
// .... single region ...
// __kmpc_end_single
// }

return EmitOMPInlinedRegion(OMPD, EntryCall, ExitCall, BodyGenCB, FiniCB,
/*Conditional*/ true, /*hasFinalize*/ true);
// __kmpc_barrier

EmitOMPInlinedRegion(OMPD, EntryCall, ExitCall, BodyGenCB, FiniCB,
/*Conditional*/ true,
/*hasFinalize*/ true);
if (!IsNowait)
createBarrier(LocationDescription(Builder.saveIP(), Loc.DL),
omp::Directive::OMPD_unknown, /* ForceSimpleCall */ false,
/* CheckCancelFlag */ false);
return Builder.saveIP();
}

OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createCritical(
Expand Down
104 changes: 102 additions & 2 deletions llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2821,8 +2821,8 @@ TEST_F(OpenMPIRBuilderTest, SingleDirective) {
EXPECT_NE(IPBB->end(), IP.getPoint());
};

Builder.restoreIP(
OMPBuilder.createSingle(Builder, BodyGenCB, FiniCB, /*DidIt*/ nullptr));
Builder.restoreIP(OMPBuilder.createSingle(
Builder, BodyGenCB, FiniCB, /*IsNowait*/ false, /*DidIt*/ nullptr));
Value *EntryBBTI = EntryBB->getTerminator();
EXPECT_NE(EntryBBTI, nullptr);
EXPECT_TRUE(isa<BranchInst>(EntryBBTI));
Expand Down Expand Up @@ -2854,6 +2854,106 @@ TEST_F(OpenMPIRBuilderTest, SingleDirective) {
EXPECT_EQ(SingleEndCI->arg_size(), 2U);
EXPECT_TRUE(isa<GlobalVariable>(SingleEndCI->getArgOperand(0)));
EXPECT_EQ(SingleEndCI->getArgOperand(1), SingleEntryCI->getArgOperand(1));

bool FoundBarrier = false;
for (auto &FI : *ExitBB) {
Instruction *cur = &FI;
if (auto CI = dyn_cast<CallInst>(cur)) {
if (CI->getCalledFunction()->getName() == "__kmpc_barrier") {
FoundBarrier = true;
break;
}
}
}
EXPECT_TRUE(FoundBarrier);
}

TEST_F(OpenMPIRBuilderTest, SingleDirectiveNowait) {
using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
OpenMPIRBuilder OMPBuilder(*M);
OMPBuilder.initialize();
F->setName("func");
IRBuilder<> Builder(BB);

OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP(), DL});

AllocaInst *PrivAI = nullptr;

BasicBlock *EntryBB = nullptr;
BasicBlock *ExitBB = nullptr;
BasicBlock *ThenBB = nullptr;

auto BodyGenCB = [&](InsertPointTy AllocaIP, InsertPointTy CodeGenIP,
BasicBlock &FiniBB) {
if (AllocaIP.isSet())
Builder.restoreIP(AllocaIP);
else
Builder.SetInsertPoint(&*(F->getEntryBlock().getFirstInsertionPt()));
PrivAI = Builder.CreateAlloca(F->arg_begin()->getType());
Builder.CreateStore(F->arg_begin(), PrivAI);

llvm::BasicBlock *CodeGenIPBB = CodeGenIP.getBlock();
llvm::Instruction *CodeGenIPInst = &*CodeGenIP.getPoint();
EXPECT_EQ(CodeGenIPBB->getTerminator(), CodeGenIPInst);

Builder.restoreIP(CodeGenIP);

// collect some info for checks later
ExitBB = FiniBB.getUniqueSuccessor();
ThenBB = Builder.GetInsertBlock();
EntryBB = ThenBB->getUniquePredecessor();

// simple instructions for body
Value *PrivLoad =
Builder.CreateLoad(PrivAI->getAllocatedType(), PrivAI, "local.use");
Builder.CreateICmpNE(F->arg_begin(), PrivLoad);
};

auto FiniCB = [&](InsertPointTy IP) {
BasicBlock *IPBB = IP.getBlock();
EXPECT_NE(IPBB->end(), IP.getPoint());
};

Builder.restoreIP(OMPBuilder.createSingle(
Builder, BodyGenCB, FiniCB, /*IsNowait*/ true, /*DidIt*/ nullptr));
Value *EntryBBTI = EntryBB->getTerminator();
EXPECT_NE(EntryBBTI, nullptr);
EXPECT_TRUE(isa<BranchInst>(EntryBBTI));
BranchInst *EntryBr = cast<BranchInst>(EntryBB->getTerminator());
EXPECT_TRUE(EntryBr->isConditional());
EXPECT_EQ(EntryBr->getSuccessor(0), ThenBB);
EXPECT_EQ(ThenBB->getUniqueSuccessor(), ExitBB);
EXPECT_EQ(EntryBr->getSuccessor(1), ExitBB);

CmpInst *CondInst = cast<CmpInst>(EntryBr->getCondition());
EXPECT_TRUE(isa<CallInst>(CondInst->getOperand(0)));

CallInst *SingleEntryCI = cast<CallInst>(CondInst->getOperand(0));
EXPECT_EQ(SingleEntryCI->arg_size(), 2U);
EXPECT_EQ(SingleEntryCI->getCalledFunction()->getName(), "__kmpc_single");
EXPECT_TRUE(isa<GlobalVariable>(SingleEntryCI->getArgOperand(0)));

CallInst *SingleEndCI = nullptr;
for (auto &FI : *ThenBB) {
Instruction *cur = &FI;
if (isa<CallInst>(cur)) {
SingleEndCI = cast<CallInst>(cur);
if (SingleEndCI->getCalledFunction()->getName() == "__kmpc_end_single")
break;
SingleEndCI = nullptr;
}
}
EXPECT_NE(SingleEndCI, nullptr);
EXPECT_EQ(SingleEndCI->arg_size(), 2U);
EXPECT_TRUE(isa<GlobalVariable>(SingleEndCI->getArgOperand(0)));
EXPECT_EQ(SingleEndCI->getArgOperand(1), SingleEntryCI->getArgOperand(1));

for (auto &FI : *ExitBB) {
Instruction *cur = &FI;
if (auto CI = dyn_cast<CallInst>(cur)) {
EXPECT_FALSE(CI->getCalledFunction()->getName() == "__kmpc_barrier");
}
}
}

TEST_F(OpenMPIRBuilderTest, OMPAtomicReadFlt) {
Expand Down
2 changes: 1 addition & 1 deletion mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -652,7 +652,7 @@ convertOmpSingle(omp::SingleOp &singleOp, llvm::IRBuilderBase &builder,
};
auto finiCB = [&](InsertPointTy codeGenIP) {};
builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createSingle(
ompLoc, bodyCB, finiCB, /*DidIt=*/nullptr));
ompLoc, bodyCB, finiCB, singleOp.nowait(), /*DidIt=*/nullptr));
return bodyGenStatus;
}

Expand Down
36 changes: 36 additions & 0 deletions mlir/test/Target/LLVMIR/openmp-llvm.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1828,15 +1828,51 @@ llvm.func @repeated_successor(%arg0: i64, %arg1: i64, %arg2: i64, %arg3: i1) {
// CHECK-LABEL: @single
// CHECK-SAME: (i32 %[[x:.*]], i32 %[[y:.*]], i32* %[[zaddr:.*]])
llvm.func @single(%x: i32, %y: i32, %zaddr: !llvm.ptr<i32>) {
// CHECK: %[[a:.*]] = sub i32 %[[x]], %[[y]]
%a = llvm.sub %x, %y : i32
// CHECK: store i32 %[[a]], i32* %[[zaddr]]
llvm.store %a, %zaddr : !llvm.ptr<i32>
// CHECK: call i32 @__kmpc_single
omp.single {
// CHECK: %[[z:.*]] = add i32 %[[x]], %[[y]]
%z = llvm.add %x, %y : i32
// CHECK: store i32 %[[z]], i32* %[[zaddr]]
llvm.store %z, %zaddr : !llvm.ptr<i32>
// CHECK: call void @__kmpc_end_single
// CHECK: call void @__kmpc_barrier
omp.terminator
}
// CHECK: %[[b:.*]] = mul i32 %[[x]], %[[y]]
%b = llvm.mul %x, %y : i32
// CHECK: store i32 %[[b]], i32* %[[zaddr]]
llvm.store %b, %zaddr : !llvm.ptr<i32>
// CHECK: ret void
llvm.return
}

// -----

// CHECK-LABEL: @single_nowait
// CHECK-SAME: (i32 %[[x:.*]], i32 %[[y:.*]], i32* %[[zaddr:.*]])
llvm.func @single_nowait(%x: i32, %y: i32, %zaddr: !llvm.ptr<i32>) {
// CHECK: %[[a:.*]] = sub i32 %[[x]], %[[y]]
%a = llvm.sub %x, %y : i32
// CHECK: store i32 %[[a]], i32* %[[zaddr]]
llvm.store %a, %zaddr : !llvm.ptr<i32>
// CHECK: call i32 @__kmpc_single
omp.single nowait {
// CHECK: %[[z:.*]] = add i32 %[[x]], %[[y]]
%z = llvm.add %x, %y : i32
// CHECK: store i32 %[[z]], i32* %[[zaddr]]
llvm.store %z, %zaddr : !llvm.ptr<i32>
// CHECK: call void @__kmpc_end_single
// CHECK-NOT: call void @__kmpc_barrier
omp.terminator
}
// CHECK: %[[t:.*]] = mul i32 %[[x]], %[[y]]
%t = llvm.mul %x, %y : i32
// CHECK: store i32 %[[t]], i32* %[[zaddr]]
llvm.store %t, %zaddr : !llvm.ptr<i32>
// CHECK: ret void
llvm.return
}

0 comments on commit 8722c12

Please sign in to comment.