[OpenMPIRBuilder] Detect and fix ambiguous InsertPoints for createSections.

[Meinersbur]

Summary

Follow-up on D117226 for createSections.

[Meinersbur]

ping

[peixin]

Added @shraiysh and @NimishMishra . They may be more familiar to this code.

[shraiysh]

LGTM. Can we please rebase this and fix the same in convertOmpAtomicUpdate and convertOmpAtomicCapture too?

[shraiysh]

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

74–75

[Suggestion] We probably don't need to do this, because if a conversion does not require an alloca (for example, barrier), we will be creating a basicblock unnecessarily. So, creating this on-demand seems okay to me.

[Meinersbur]

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

74–75

Making decision based on the current position of an IRBuilder creates heaps of problems and unpredictability. Special cases like this one here need to be tested and maintained, and still easily introduce bugs. An unconditional dedicated alloca block avoids all these problems.

An extra basic block on the other side is insignificant. It will be gone as soon as simplify-cfg runs.

However, this seems to be controversial, so I removed the TODO for now.

[Meinersbur]

• Rebase
• Fix convertOmpAtomicUpdate, convertOmpAtomicCapture, and convertOmpOrdered

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

if (!updateToLocation(Loc))
     return;
   emitTaskyieldImpl(Loc);
 }
 
 OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createSections(
     const LocationDescription &Loc, InsertPointTy AllocaIP,
     ArrayRef<StorableBodyGenCallbackTy> SectionCBs, PrivatizeCallbackTy PrivCB,
     FinalizeCallbackTy FiniCB, bool IsCancellable, bool IsNowait) {
+  assert(!isConflictIP(AllocaIP, Loc.IP) && "Dedicated IP allocas required");
+
   if (!updateToLocation(Loc))
     return Loc.IP;
 
   auto FiniCBWrapper = [&](InsertPointTy IP) {
     if (IP.getBlock()->end() != IP.getPoint())
       return FiniCB(IP);
     // This must be done otherwise any nested constructs using FinalizeOMPRegion
     // will fail because that function requires the Finalization Basic Block to

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

 
 TEST_F(OpenMPIRBuilderTest, CreateSectionsSimple) {
   using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
   using BodyGenCallbackTy = llvm::OpenMPIRBuilder::StorableBodyGenCallbackTy;
   OpenMPIRBuilder OMPBuilder(*M);
   OMPBuilder.initialize();
   F->setName("func");
   IRBuilder<> Builder(BB);
+
+  BasicBlock *EnterBB = BasicBlock::Create(Ctx, "sections.enter", F);
+  Builder.CreateBr(EnterBB);
+  Builder.SetInsertPoint(EnterBB);
   OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP(), DL});
+
   llvm::SmallVector<BodyGenCallbackTy, 4> SectionCBVector;
   llvm::SmallVector<BasicBlock *, 4> CaseBBs;
 
   auto FiniCB = [&](InsertPointTy IP) {};
   auto SectionCB = [&](InsertPointTy AllocaIP, InsertPointTy CodeGenIP,
                        BasicBlock &FiniBB) {
     Builder.restoreIP(CodeGenIP);
     Builder.CreateBr(&FiniBB);

 TEST_F(OpenMPIRBuilderTest, CreateSectionsNoWait) {
   using InsertPointTy = OpenMPIRBuilder::InsertPointTy;
   using BodyGenCallbackTy = llvm::OpenMPIRBuilder::StorableBodyGenCallbackTy;
   OpenMPIRBuilder OMPBuilder(*M);
   OMPBuilder.initialize();
   F->setName("func");
   IRBuilder<> Builder(BB);
 
+  BasicBlock *EnterBB = BasicBlock::Create(Ctx, "sections.enter", F);
+  Builder.CreateBr(EnterBB);
+  Builder.SetInsertPoint(EnterBB);
   OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP(), DL});
+
   IRBuilder<>::InsertPoint AllocaIP(&F->getEntryBlock(),
                                     F->getEntryBlock().getFirstInsertionPt());
   llvm::SmallVector<BodyGenCallbackTy, 4> SectionCBVector;
   auto PrivCB = [](InsertPointTy AllocaIP, InsertPointTy CodeGenIP,
                    llvm::Value &, llvm::Value &Val,
                    llvm::Value *&ReplVal) { return CodeGenIP; };
   auto FiniCB = [&](InsertPointTy IP) {};
 

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

findAllocaInsertPoint(llvm::IRBuilderBase &builder,
   WalkResult walkResult = moduleTranslation.stackWalk<OpenMPAllocaStackFrame>(
       [&](const OpenMPAllocaStackFrame &frame) {
         allocaInsertPoint = frame.allocaInsertPoint;
         return WalkResult::interrupt();
       });
   if (walkResult.wasInterrupted())
     return allocaInsertPoint;
 
-  // Otherwise, insert to the entry block of the surrounding function.
+  // Otherwise, insert to the entry block of the surrounding function.s
+  // If the current IRBuilder InsertPoint is the function's entry, it cannot
+  // also be used for alloca insertion which would result in insertion order
+  // confusion. Create a new BasicBlock for the Builder and use the entry block
+  // for the allocs.
+  // TODO: Create a dedicated alloca BasicBlock at function creation such that
+  // we do not need to move the current InertPoint here.
+  if (builder.GetInsertBlock() ==
+      &builder.GetInsertBlock()->getParent()->getEntryBlock()) {
+    assert(builder.GetInsertPoint() == builder.GetInsertBlock()->end() &&
+           "Assuming end of basic block");
+    llvm::BasicBlock *entryBB = llvm::BasicBlock::Create(
+        builder.getContext(), "entry", builder.GetInsertBlock()->getParent(),
+        builder.GetInsertBlock()->getNextNode());
+    builder.CreateBr(entryBB);
+    builder.SetInsertPoint(entryBB);
+  }
+
   llvm::BasicBlock &funcEntryBlock =
       builder.GetInsertBlock()->getParent()->getEntryBlock();
   return llvm::OpenMPIRBuilder::InsertPointTy(
       &funcEntryBlock, funcEntryBlock.getFirstInsertionPt());
 }
 
 /// Converts the given region that appears within an OpenMP dialect operation to
 /// LLVM IR, creating a branch from the `sourceBlock` to the entry block of the
convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
   if (auto numThreadsVar = opInst.num_threads_var())
     numThreads = moduleTranslation.lookupValue(numThreadsVar);
   auto pbKind = llvm::omp::OMP_PROC_BIND_default;
   if (auto bind = opInst.proc_bind_val())
     pbKind = getProcBindKind(*bind);
   // TODO: Is the Parallel construct cancellable?
   bool isCancellable = false;
 
-  // Ensure that the BasicBlock for the the parallel region is sparate from the
-  // function entry which we may need to insert allocas.
-  if (builder.GetInsertBlock() ==
-      &builder.GetInsertBlock()->getParent()->getEntryBlock()) {
-    assert(builder.GetInsertPoint() == builder.GetInsertBlock()->end() &&
-           "Assuming end of basic block");
-    llvm::BasicBlock *entryBB =
-        llvm::BasicBlock::Create(builder.getContext(), "parallel.entry",
-                                 builder.GetInsertBlock()->getParent(),
-                                 builder.GetInsertBlock()->getNextNode());
-    builder.CreateBr(entryBB);
-    builder.SetInsertPoint(entryBB);
-  }
+  llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
+      findAllocaInsertPoint(builder, moduleTranslation);
   llvm::OpenMPIRBuilder::LocationDescription ompLoc(
       builder.saveIP(), builder.getCurrentDebugLocation());
   builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createParallel(
-      ompLoc, findAllocaInsertPoint(builder, moduleTranslation), bodyGenCB,
-      privCB, finiCB, ifCond, numThreads, pbKind, isCancellable));
+      ompLoc, allocaIP, bodyGenCB, privCB, finiCB, ifCond, numThreads, pbKind,
+      isCancellable));
 
   return bodyGenStatus;
 }
 
 /// Converts an OpenMP 'master' operation into LLVM IR using OpenMPIRBuilder.
 static LogicalResult
 convertOmpMaster(Operation &opInst, llvm::IRBuilderBase &builder,
                  LLVM::ModuleTranslation &moduleTranslation) {
convertOmpOrdered(Operation &opInst, llvm::IRBuilderBase &builder,
   size_t indexVecValues = 0;
   while (indexVecValues < vecValues.size()) {
     SmallVector<llvm::Value *> storeValues;
     storeValues.reserve(numLoops);
     for (unsigned i = 0; i < numLoops; i++) {
       storeValues.push_back(vecValues[indexVecValues]);
       indexVecValues++;
     }
+    llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
+        findAllocaInsertPoint(builder, moduleTranslation);
     builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createOrderedDepend(
-        ompLoc, findAllocaInsertPoint(builder, moduleTranslation), numLoops,
-        storeValues, ".cnt.addr", isDependSource));
+        ompLoc, allocaIP, numLoops, storeValues, ".cnt.addr", isDependSource));
   }
   return success();
 }
 
 /// Converts an OpenMP 'ordered_region' operation into LLVM IR using
 /// OpenMPIRBuilder.
 static LogicalResult
 convertOmpOrderedRegion(Operation &opInst, llvm::IRBuilderBase &builder,
auto privCB = [&](InsertPointTy, InsertPointTy codeGenIP, llvm::Value &,
     replacementValue = &vPtr;
     return codeGenIP;
   };
 
   // TODO: Perform finalization actions for variables. This has to be
   // called for variables which have destructors/finalizers.
   auto finiCB = [&](InsertPointTy codeGenIP) {};
 
+  llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
+      findAllocaInsertPoint(builder, moduleTranslation);
   llvm::OpenMPIRBuilder::LocationDescription ompLoc(
       builder.saveIP(), builder.getCurrentDebugLocation());
   builder.restoreIP(moduleTranslation.getOpenMPBuilder()->createSections(
-      ompLoc, findAllocaInsertPoint(builder, moduleTranslation), sectionCBs,
-      privCB, finiCB, false, sectionsOp.nowait()));
+      ompLoc, allocaIP, sectionCBs, privCB, finiCB, false,
+      sectionsOp.nowait()));
   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

 
 // -----
 
 // Check IR generation for simple empty sections. This only checks the overall
 // shape of the IR, detailed checking is done by the OpenMPIRBuilder.
 
 // CHECK-LABEL: @omp_sections_trivial
 llvm.func @omp_sections_trivial() -> () {
+  // CHECK:   br label %[[ENTRY:[a-zA-Z_.]+]]
+
+  // CHECK: [[ENTRY]]:
   // CHECK:   br label %[[PREHEADER:.*]]
 
   // CHECK: [[PREHEADER]]:
   // CHECK:   %{{.*}} = call i32 @__kmpc_global_thread_num({{.*}})
   // CHECK:   call void @__kmpc_for_static_init_4u({{.*}})
   // CHECK:   br label %[[HEADER:.*]]
 
   // CHECK: [[HEADER]]: