[mlir][OpenMP][IRBuilder] Add support for nowait on single construct

[shraiysh]

Summary

This patch adds the nowait parameter to createSingle in
OpenMPIRBuilder and handling for IR generation from OpenMP Dialect.

Also added tests for the same.

[probinson]

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

2954

It seems this EXPECT_FALSE is never executed, causing it to be flagged in my "Rotten Green Tests" work. Perhaps the loop could be rewritten to be more similar to the ThenBB check?

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

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

InsertPointTy createCopyPrivate(const LocationDescription &Loc,
                                   llvm::Value *BufSize, llvm::Value *CpyBuf,
                                   llvm::Value *CpyFn, llvm::Value *DidIt);
 
   /// Generator for '#omp single'
   ///
   /// \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'
   ///
   /// \param Loc The insert and source location description.
   /// \param BodyGenCB Callback that will generate the region code.
   /// \param FiniCB Callback to finalize variable copies.
   ///
   /// \returns The insertion position *after* the master.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

OpenMPIRBuilder::createCopyPrivate(const LocationDescription &Loc,
   Value *Args[] = {Ident, ThreadId, BufSize, CpyBuf, CpyFn, DidItLD};
 
   Function *Fn = getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_copyprivate);
   Builder.CreateCall(Fn, Args);
 
   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;
 
   // If needed (i.e. not null), initialize `DidIt` with 0
   if (DidIt) {
     Builder.CreateStore(Builder.getInt32(0), DidIt);
   }
OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createSingle(
   Function *ExitRTLFn = getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_end_single);
   Instruction *ExitCall = Builder.CreateCall(ExitRTLFn, Args);
 
   // generates the following:
   // if (__kmpc_single()) {
   //    .... single region ...
   //    __kmpc_end_single
   // }
+  // __kmpc_barrier
 
-  return EmitOMPInlinedRegion(OMPD, EntryCall, ExitCall, BodyGenCB, FiniCB,
-                              /*Conditional*/ true, /*hasFinalize*/ true);
+  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(
     const LocationDescription &Loc, BodyGenCallbackTy BodyGenCB,
     FinalizeCallbackTy FiniCB, StringRef CriticalName, Value *HintInst) {
 
   if (!updateToLocation(Loc))
     return Loc.IP;

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

auto BodyGenCB = [&](InsertPointTy AllocaIP, InsertPointTy CodeGenIP,
     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, /*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));
   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);
if (isa<CallInst>(cur)) {
         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));
+
+  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) {
   OpenMPIRBuilder OMPBuilder(*M);
   OMPBuilder.initialize();
   F->setName("func");
   IRBuilder<> Builder(BB);
 

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

convertOmpSingle(omp::SingleOp &singleOp, llvm::IRBuilderBase &builder,
   auto bodyCB = [&](InsertPointTy allocaIP, InsertPointTy codegenIP,
                     llvm::BasicBlock &continuationBB) {
     convertOmpOpRegions(singleOp.region(), "omp.single.region",
                         *codegenIP.getBlock(), continuationBB, builder,
                         moduleTranslation, bodyGenStatus);
   };
   auto finiCB = [&](InsertPointTy codeGenIP) {};
   builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createSingle(
-      ompLoc, bodyCB, finiCB, /*DidIt=*/nullptr));
+      ompLoc, bodyCB, finiCB, singleOp.nowait(), /*DidIt=*/nullptr));
   return bodyGenStatus;
 }
 
 /// Converts an OpenMP workshare loop into LLVM IR using OpenMPIRBuilder.
 static LogicalResult
 convertOmpWsLoop(Operation &opInst, llvm::IRBuilderBase &builder,
                  LLVM::ModuleTranslation &moduleTranslation) {
   auto loop = cast<omp::WsLoopOp>(opInst);

mlir/test/Target/LLVMIR/openmp-llvm.mlir

llvm.func @repeated_successor(%arg0: i64, %arg1: i64, %arg2: i64, %arg3: i1) {
   llvm.return
 }
 
 // -----
 
 // 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
 }