[OpenMP][IRBuilder] Add final clause to task

[shraiysh]

Summary

This patch adds final clause to OpenMP IR Builder.

[Meinersbur]

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

627

Please document the new parameter in the doxygen comment.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

1259

[style] LLVM naming conventions have parameters start with an upper case letter.

Also I don't see the point of the Bool suffix. Consider IsFinal or Final.

1340

[style]

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

4590

[style]

4611

[👍] Great use of any_of

[shraiysh]

Addressed comments. Thanks for the review @Meinersbur.

[shraiysh]

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

4611

Thank you! :)

[Meinersbur]

LGTM with a request to describe the llvm::Value argument more precisely.

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

627

Since IsFinal is not a bool, you cannot pass true/false like for Tied. What is passed is the llvm::Value that decides whether to create a (non-)final task at runtime.

[shraiysh]

Thanks for pointing it out @Meinersbur. Somehow I mistook it for
a bool value while addressing comments earlier. Corrected the
description now.

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

public:
   void createTaskyield(const LocationDescription &Loc);
 
   /// Generator for `#omp task`
   ///
   /// \param Loc The location where the task construct was encountered.
   /// \param AllocaIP The insertion point to be used for alloca instructions.
   /// \param BodyGenCB Callback that will generate the region code.
   /// \param Tied True if the task is tied, false if the task is untied.
   InsertPointTy createTask(const LocationDescription &Loc,
                            InsertPointTy AllocaIP, BodyGenCallbackTy BodyGenCB,
-                           bool Tied = true);
+                           bool Tied = true, Value *finalBool = nullptr);
 
   /// Functions used to generate reductions. Such functions take two Values
   /// representing LHS and RHS of the reduction, respectively, and a reference
   /// to the value that is updated to refer to the reduction result.
   using ReductionGenTy =
       function_ref<InsertPointTy(InsertPointTy, Value *, Value *, Value *&)>;
 
   /// Functions used to generate atomic reductions. Such functions take two

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

void OpenMPIRBuilder::createTaskyield(const LocationDescription &Loc) {
   if (!updateToLocation(Loc))
     return;
   emitTaskyieldImpl(Loc);
 }
 
 OpenMPIRBuilder::InsertPointTy
 OpenMPIRBuilder::createTask(const LocationDescription &Loc,
                             InsertPointTy AllocaIP, BodyGenCallbackTy BodyGenCB,
-                            bool Tied) {
+                            bool Tied, Value *finalBool) {
   if (!updateToLocation(Loc))
     return InsertPointTy();
 
   // The current basic block is split into four basic blocks. After outlining,
   // they will be mapped as follows:
   // ```
   // def current_fn() {
   //   current_basic_block:
OpenMPIRBuilder::createTask(const LocationDescription &Loc,
   BasicBlock *TaskBodyBB = splitBB(Builder, /*CreateBranch=*/true, "task.body");
   BasicBlock *TaskAllocaBB =
       splitBB(Builder, /*CreateBranch=*/true, "task.alloca");
 
   OutlineInfo OI;
   OI.EntryBB = TaskAllocaBB;
   OI.OuterAllocaBB = AllocaIP.getBlock();
   OI.ExitBB = TaskExitBB;
-  OI.PostOutlineCB = [this, &Loc, Tied](Function &OutlinedFn) {
+  OI.PostOutlineCB = [this, &Loc, Tied, finalBool](Function &OutlinedFn) {
     // The input IR here looks like the following-
     // ```
     // func @current_fn() {
     //   outlined_fn(%args)
     // }
     // func @outlined_fn(%args) { ... }
     // ```
     //
OI.PostOutlineCB = [this, &Loc, Tied, finalBool](Function &OutlinedFn) {
     // Arguments - `loc_ref` (Ident) and `gtid` (ThreadID)
     // call.
     uint32_t SrcLocStrSize;
     Constant *SrcLocStr = getOrCreateSrcLocStr(Loc, SrcLocStrSize);
     Value *Ident = getOrCreateIdent(SrcLocStr, SrcLocStrSize);
     Value *ThreadID = getOrCreateThreadID(Ident);
 
     // Argument - `flags`
-    // If task is tied, then (Flags & 1) == 1.
-    // If task is untied, then (Flags & 1) == 0.
+    // Task is tied iff (Flags & 1) == 1.
+    // Task is untied iff (Flags & 1) == 0.
+    // Task is final iff (Flags & 2) == 2.
+    // Task is not final iff (Flags & 2) == 0.
     // TODO: Handle the other flags.
     Value *Flags = Builder.getInt32(Tied);
+    if (finalBool) {
+      Value *finalFlag = Builder.CreateSelect(finalBool, Builder.getInt32(2),
+                                              Builder.getInt32(0));
+      Flags = Builder.CreateOr(finalFlag, Flags);
+    }
 
     // Argument - `sizeof_kmp_task_t` (TaskSize)
     // Tasksize refers to the size in bytes of kmp_task_t data structure
     // including private vars accessed in task.
     Value *TaskSize = Builder.getInt64(0);
     if (HasTaskData) {
       AllocaInst *ArgStructAlloca =
           dyn_cast<AllocaInst>(StaleCI->getArgOperand(0));

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

TEST_F(OpenMPIRBuilderTest, CreateTaskUntied) {
   ASSERT_NE(TaskAllocCall, nullptr);
   ConstantInt *Flags = dyn_cast<ConstantInt>(TaskAllocCall->getArgOperand(2));
   ASSERT_NE(Flags, nullptr);
   EXPECT_EQ(Flags->getZExtValue() & 1U, 0U);
 
   EXPECT_FALSE(verifyModule(*M, &errs()));
 }
 
+TEST_F(OpenMPIRBuilderTest, CreateTaskFinal) {
+  using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
+  OpenMPIRBuilder OMPBuilder(*M);
+  OMPBuilder.initialize();
+  F->setName("func");
+  IRBuilder<> Builder(BB);
+  auto BodyGenCB = [&](InsertPointTy AllocaIP, InsertPointTy CodeGenIP) {};
+  IRBuilderBase::InsertPoint AllocaIP = Builder.saveIP();
+  BasicBlock *BodyBB = splitBB(Builder, /*CreateBranch=*/true, "alloca.split");
+  Builder.SetInsertPoint(BodyBB);
+  Value *finalBool = Builder.CreateICmp(
+      CmpInst::Predicate::ICMP_EQ, F->getArg(0),
+      ConstantInt::get(Type::getInt32Ty(M->getContext()), 0U));
+  OpenMPIRBuilder::LocationDescription Loc(Builder.saveIP(), DL);
+  Builder.restoreIP(OMPBuilder.createTask(Loc, AllocaIP, BodyGenCB,
+                                          /*Tied=*/false, finalBool));
+  OMPBuilder.finalize();
+  Builder.CreateRetVoid();
+
+  // Check for the `Tied` argument
+  CallInst *TaskAllocCall = dyn_cast<CallInst>(
+      OMPBuilder.getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_omp_task_alloc)
+          ->user_back());
+  ASSERT_NE(TaskAllocCall, nullptr);
+  BinaryOperator *OrInst =
+      dyn_cast<BinaryOperator>(TaskAllocCall->getArgOperand(2));
+  ASSERT_NE(OrInst, nullptr);
+  EXPECT_EQ(OrInst->getOpcode(), BinaryOperator::BinaryOps::Or);
+
+  // One of the arguments to `or` instruction is the tied flag, which is equal
+  // to zero.
+  EXPECT_TRUE(any_of(OrInst->operands(), [](Value *op) {
+    if (ConstantInt *TiedValue = dyn_cast<ConstantInt>(op))
+      return TiedValue->getSExtValue() == 0;
+    return false;
+  }));
+
+  // One of the arguments to `or` instruction is the final condition.
+  EXPECT_TRUE(any_of(OrInst->operands(), [finalBool](Value *op) {
+    if (SelectInst *Select = dyn_cast<SelectInst>(op)) {
+      ConstantInt *TrueValue = dyn_cast<ConstantInt>(Select->getTrueValue());
+      ConstantInt *FalseValue = dyn_cast<ConstantInt>(Select->getFalseValue());
+      if (!TrueValue || !FalseValue)
+        return false;
+      return Select->getCondition() == finalBool &&
+             TrueValue->getSExtValue() == 2 && FalseValue->getSExtValue() == 0;
+    }
+    return false;
+  }));
+
+  EXPECT_FALSE(verifyModule(*M, &errs()));
+}
+
 } // namespace