[llvm] [OpenMPIRBuilder] Added `createTeams` (PR #65785)
via llvm-commits
llvm-commits at lists.llvm.org
Fri Sep 8 11:00:05 PDT 2023
https://github.com/shraiysh created https://github.com/llvm/llvm-project/pull/65785:
This patch adds a generator for the teams construct. The generated IR looks like the following:
```
current_fn() {
...
call @__kmpc_fork_teams(ptr @ident, i32 num_args, ptr @outlined_omp_teams, ...args)
...
}
outlined_omp_teams(ptr %global_tid, ptr %bound_tid, ...args) {
; teams body
}
```
It does this by first generating the body in the current function. Then we outline the body in a temporary function. We then create the @outlined_omp_teams function and embed the temporary outlined function in this function. We then emit the call to runtime function.
>From 67aa8afc08f5d26f3c55f43802586513fc520fc0 Mon Sep 17 00:00:00 2001
From: Shraiysh Vaishay <shraiysh.vaishay at amd.com>
Date: Fri, 8 Sep 2023 12:56:56 -0500
Subject: [PATCH] [OpenMPIRBuilder] Added `createTeams`
This patch adds a generator for the teams construct. The generated IR looks like the following:
```
current_fn() {
...
call @__kmpc_fork_teams(ptr @ident, i32 num_args, ptr @outlined_omp_teams, ...args)
...
}
outlined_omp_teams(ptr %global_tid, ptr %bound_tid, ...args) {
; teams body
}
```
It does this by first generating the body in the current function. Then we outline the
body in a temporary function. We then create the @outlined_omp_teams function and embed
the temporary outlined function in this function. We then emit the call to runtime
function.
---
.../llvm/Frontend/OpenMP/OMPIRBuilder.h | 7 +
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 141 ++++++++++++++++++
.../Frontend/OpenMPIRBuilderTest.cpp | 73 +++++++++
3 files changed, 221 insertions(+)
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
index b6460267c08aaa1..037e0a5662be5bb 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
@@ -2005,6 +2005,13 @@ class OpenMPIRBuilder {
/// \param Loc The insert and source location description.
void createTargetDeinit(const LocationDescription &Loc);
+ /// Generator for `#omp teams`
+ ///
+ /// \param Loc The location where the teams construct was encountered.
+ /// \param BodyGenCB Callback that will generate the region code.
+ InsertPointTy createTeams(const LocationDescription &Loc,
+ BodyGenCallbackTy BodyGenCB);
+
///}
private:
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 6cab70df269c170..86bf1133e2b1401 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -6106,6 +6106,147 @@ void OpenMPIRBuilder::loadOffloadInfoMetadata(Module &M) {
}
}
+OpenMPIRBuilder::InsertPointTy
+OpenMPIRBuilder::createTeams(const LocationDescription &Loc,
+ BodyGenCallbackTy BodyGenCB) {
+ if (!updateToLocation(Loc))
+ return Loc.IP;
+
+ uint32_t SrcLocStrSize;
+ Constant *SrcLocStr = getOrCreateSrcLocStr(Loc, SrcLocStrSize);
+ Value *Ident = getOrCreateIdent(SrcLocStr, SrcLocStrSize);
+
+ // Splitting a basic block expects a terminator. Hence, creating an
+ // unreachable instruction, which will be deleted later.
+ UnreachableInst *UI = Builder.CreateUnreachable();
+ BasicBlock *CurrentBasicBlock = Builder.GetInsertBlock();
+
+ // The current basic block is split into four basic blocks. After outlining,
+ // they will be mapped as follows:
+ // ```
+ // def current_fn() {
+ // current_basic_block:
+ // br label %teams.exit
+ // teams.exit:
+ // ; instructions after teams
+ // }
+ // def outlined_fn() {
+ // teams.alloca:
+ // br label %teams.body
+ // teams.body:
+ // ; instructions within teams body
+ // }
+ // ```
+ BasicBlock *AllocaBB = CurrentBasicBlock->splitBasicBlock(UI, "teams.alloca");
+ BasicBlock *BodyBB = AllocaBB->splitBasicBlock(UI, "teams.body");
+ BasicBlock *ExitBB = BodyBB->splitBasicBlock(UI, "teams.exit");
+
+ UI->eraseFromParent();
+
+ // Generate the body of teams.
+ InsertPointTy CodeGenIP(BodyBB, BodyBB->begin());
+ InsertPointTy AllocaIP(AllocaBB, AllocaBB->begin());
+ BodyGenCB(AllocaIP, CodeGenIP);
+
+ OutlineInfo OI;
+ OI.EntryBB = AllocaBB;
+ OI.ExitBB = ExitBB;
+ OI.PostOutlineCB = [this, Ident](Function &OutlinedFn) {
+ // The input IR here looks like the following-
+ // ```
+ // func @current_fn() {
+ // outlined_fn(%args)
+ // }
+ // func @outlined_fn(%args) {
+ // ; teams body
+ // }
+ // ```
+ //
+ // This is changed to the following-
+ //
+ // ```
+ // func @current_fn() {
+ // runtime_call(..., wrapper_fn, ...)
+ // }
+ // func @wrapper_fn(..., %args) {
+ // ; teams body
+ // }
+ // ```
+
+ // The outlined function has different inputs than what is expected from it.
+ // So, a wrapper function with expected signature is created and the
+ // required arguments are passed to the outlined function. The stale call
+ // instruction in current function will be replaced with a new call
+ // instruction for runtime call with the wrapper function. The outlined
+ // function is then inlined in the wrapper function and the call from the
+ // current function is removed.
+
+ assert(OutlinedFn.getNumUses() == 1 &&
+ "there must be a single user for the outlined function");
+ CallInst *StaleCI = cast<CallInst>(OutlinedFn.user_back());
+ assert(StaleCI && "Error while outlining - no CallInst user found for the "
+ "outlined function.");
+ OutlinedFn.addFnAttr(Attribute::AttrKind::AlwaysInline);
+
+ // Create the wrapper function.
+ Builder.SetInsertPoint(StaleCI);
+ SmallVector<Type *> WrapperArgTys{Builder.getPtrTy(), Builder.getPtrTy()};
+ for (auto &Arg : OutlinedFn.args()) {
+ WrapperArgTys.push_back(Arg.getType());
+ }
+ FunctionCallee WrapperFuncVal = M.getOrInsertFunction(
+ (Twine(OutlinedFn.getName()) + ".teams").str(),
+ FunctionType::get(Builder.getVoidTy(), WrapperArgTys, false));
+ Function *WrapperFunc = dyn_cast<Function>(WrapperFuncVal.getCallee());
+ WrapperFunc->getArg(0)->setName("global_tid");
+ WrapperFunc->getArg(1)->setName("bound_tid");
+ WrapperFunc->getArg(2)->setName("data");
+
+ // Emit the body of the wrapper function - just a call to outlined function
+ // and return statement.
+ BasicBlock *WrapperEntryBB =
+ BasicBlock::Create(M.getContext(), "entrybb", WrapperFunc);
+ Builder.SetInsertPoint(WrapperEntryBB);
+ SmallVector<Value *> Args;
+ for (size_t ArgIndex = 2; ArgIndex < WrapperFunc->arg_size(); ArgIndex++) {
+ Args.push_back(WrapperFunc->getArg(ArgIndex));
+ }
+ CallInst *OutlinedFnCall = Builder.CreateCall(&OutlinedFn, Args);
+ Builder.CreateRetVoid();
+
+ // Call to the runtime function for teams in the current function.
+ Builder.SetInsertPoint(StaleCI);
+ Args = {Ident, Builder.getInt32(StaleCI->arg_size()), WrapperFunc};
+ for (Use &Arg : StaleCI->args()) {
+ Args.push_back(Arg);
+ }
+ Builder.CreateCall(getOrCreateRuntimeFunctionPtr(
+ omp::RuntimeFunction::OMPRTL___kmpc_fork_teams),
+ Args);
+ StaleCI->eraseFromParent();
+
+ // Inlining the outlined teams function in the wrapper. This wrapper is the
+ // argument for the runtime call.
+ assert(OutlinedFn.getNumUses() == 1 &&
+ "More than one use for the outlined function found. Expected only "
+ "one use.");
+ InlineFunctionInfo IFI;
+ InlineResult IR = InlineFunction(*OutlinedFnCall, IFI);
+ LLVM_DEBUG(if (!IR.isSuccess()) {
+ dbgs() << "Attempt to merge the outlined function in the wrapper failed: "
+ << IR.getFailureReason() << "\n";
+ });
+ assert(IR.isSuccess() && "Inlining outlined omp teams failed");
+ OutlinedFn.eraseFromParent();
+ };
+
+ addOutlineInfo(std::move(OI));
+
+ Builder.SetInsertPoint(ExitBB);
+
+ return Builder.saveIP();
+}
+
bool OffloadEntriesInfoManager::empty() const {
return OffloadEntriesTargetRegion.empty() &&
OffloadEntriesDeviceGlobalVar.empty();
diff --git a/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp b/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
index 5870457956b5433..b189559b3430461 100644
--- a/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
+++ b/llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp
@@ -6033,4 +6033,77 @@ TEST_F(OpenMPIRBuilderTest, createGPUOffloadEntry) {
EXPECT_TRUE(Fn->hasFnAttribute(Attribute::MustProgress));
}
+TEST_F(OpenMPIRBuilderTest, createTeams) {
+ using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
+ OpenMPIRBuilder OMPBuilder(*M);
+ OMPBuilder.initialize();
+ F->setName("func");
+ IRBuilder<> Builder(BB);
+
+ AllocaInst *ValPtr32 = Builder.CreateAlloca(Builder.getInt32Ty());
+ AllocaInst *ValPtr128 = Builder.CreateAlloca(Builder.getInt128Ty());
+ Value *Val128 =
+ Builder.CreateLoad(Builder.getInt128Ty(), ValPtr128, "bodygen.load");
+
+ auto BodyGenCB = [&](InsertPointTy AllocaIP, InsertPointTy CodeGenIP) {
+ Builder.restoreIP(AllocaIP);
+ AllocaInst *Local128 = Builder.CreateAlloca(Builder.getInt128Ty(), nullptr,
+ "bodygen.alloca128");
+
+ Builder.restoreIP(CodeGenIP);
+ // Loading and storing captured pointer and values
+ Builder.CreateStore(Val128, Local128);
+ Value *Val32 = Builder.CreateLoad(ValPtr32->getAllocatedType(), ValPtr32,
+ "bodygen.load32");
+
+ LoadInst *PrivLoad128 = Builder.CreateLoad(
+ Local128->getAllocatedType(), Local128, "bodygen.local.load128");
+ Value *Cmp = Builder.CreateICmpNE(
+ Val32, Builder.CreateTrunc(PrivLoad128, Val32->getType()));
+ Instruction *ThenTerm, *ElseTerm;
+ SplitBlockAndInsertIfThenElse(Cmp, CodeGenIP.getBlock()->getTerminator(),
+ &ThenTerm, &ElseTerm);
+ };
+
+ OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP(), DL});
+ Builder.restoreIP(OMPBuilder.createTeams(Builder, BodyGenCB));
+ OMPBuilder.finalize();
+ Builder.CreateRetVoid();
+
+ EXPECT_FALSE(verifyModule(*M, &errs()));
+
+ CallInst *TeamsForkCall = dyn_cast<CallInst>(
+ OMPBuilder.getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_fork_teams)
+ ->user_back());
+
+ // Verify the Ident argument
+ GlobalVariable *Ident = cast<GlobalVariable>(TeamsForkCall->getArgOperand(0));
+ ASSERT_NE(Ident, nullptr);
+ EXPECT_TRUE(Ident->hasInitializer());
+ Constant *Initializer = Ident->getInitializer();
+ GlobalVariable *SrcStrGlob =
+ cast<GlobalVariable>(Initializer->getOperand(4)->stripPointerCasts());
+ ASSERT_NE(SrcStrGlob, nullptr);
+ ConstantDataArray *SrcSrc =
+ dyn_cast<ConstantDataArray>(SrcStrGlob->getInitializer());
+ ASSERT_NE(SrcSrc, nullptr);
+
+ // Verify the outlined function signature.
+ Function *OutlinedFn =
+ dyn_cast<Function>(TeamsForkCall->getArgOperand(2)->stripPointerCasts());
+ ASSERT_NE(OutlinedFn, nullptr);
+ EXPECT_FALSE(OutlinedFn->isDeclaration());
+ EXPECT_TRUE(OutlinedFn->arg_size() >= 3);
+ EXPECT_EQ(OutlinedFn->getArg(0)->getType(), Builder.getPtrTy()); // global_tid
+ EXPECT_EQ(OutlinedFn->getArg(1)->getType(), Builder.getPtrTy()); // bound_tid
+ EXPECT_EQ(OutlinedFn->getArg(2)->getType(),
+ Builder.getPtrTy()); // captured args
+
+ // Check for TruncInst and ICmpInst in the outlined function.
+ EXPECT_TRUE(any_of(instructions(OutlinedFn),
+ [](Instruction &inst) { return isa<TruncInst>(&inst); }));
+ EXPECT_TRUE(any_of(instructions(OutlinedFn),
+ [](Instruction &inst) { return isa<ICmpInst>(&inst); }));
+}
+
} // namespace
More information about the llvm-commits
mailing list