[OpenMP][IRBuilder] Handle floats for atomic update and fix AllocaIP for update/capture

[shraiysh]

Summary

This patch fixes createAtomicUpdate for lowering with float types.
Test added for the same.

This patch also changes the alloca argument for createAtomicUpdate and
createAtomicCapture from Instruction* to InsertPointTy. This is in
line with the other functions of the OpenMPIRBuilder class which take
AllocaIP as an InsertPointTy.

[shraiysh]

This is actually a WIP. (Apologies for adding all the reviewers so early). As I am working on the lowering for omp.atomic.update in the OpenMP Dialect (MLIR), I might find other issues with the createAtomicUpdate function along the way. I will ping this thread once it is ready for review. Thanks!

[Meinersbur]

I added a [WIP] into the title. Please remove it once you think it is ready.

[shraiysh]

Fix createAtomicUpdate for lowering float types. This fix is now ready for review.

[Meinersbur]

The title is "Fix AllocIP for atomic update and capture", but reading the summary the main motivation seems to be to fix createAtomicUpdate for floats.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

3301–3303

Is this move relevant?

If this is due to AllocaIP and Loc.IP being the same, you could take the occasion and check that they are different since is would be ambiguous which one is "first". See isConflictIP in D117226.

[shraiysh]

Addressed comments. Thanks for the review @Meinersbur.

[shraiysh]

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

3301–3303

Good point. I have added the check for ambiguous IPs in the createAtomicUpdate function now. This move is however unrelated. This is just to do generate all the ContBB instructions later without having to switch IPs for creating alloca. The move is simply to avoid RestoreIP and SaveIP.

[Meinersbur]

LGTM

[shraiysh]

Rebase with main.

[shraiysh]

The failures are unrelated and are timeout failures. Should I go ahead with the merge @Meinersbur?

[shraiysh]

Rebase with main. Hope the builds pass.

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

private:
   /// \returns    wether a flush was emitted or not
   bool checkAndEmitFlushAfterAtomic(const LocationDescription &Loc,
                                     AtomicOrdering AO, AtomicKind AK);
 
   /// Emit atomic update for constructs: X = X BinOp Expr ,or X = Expr BinOp X
   /// For complex Operations: X = UpdateOp(X) => CmpExch X, old_X, UpdateOp(X)
   /// Only Scalar data types.
   ///
-  /// \param AllocIP    Instruction to create AllocaInst before.
+  /// \param AllocaIP   The insertion point to be used for alloca
+  ///                   instructions.
   /// \param X          The target atomic pointer to be updated
   /// \param Expr       The value to update X with.
   /// \param AO         Atomic ordering of the generated atomic
   ///                   instructions.
   /// \param RMWOp      The binary operation used for update. If
   ///                   operation is not supported by atomicRMW,
   ///                   or belong to {FADD, FSUB, BAD_BINOP}.
   ///                   Then a `cmpExch` based  atomic will be generated.
   /// \param UpdateOp   Code generator for complex expressions that cannot be
   ///                   expressed through atomicrmw instruction.
   /// \param VolatileX       true if \a X volatile?
   /// \param IsXBinopExpr true if \a X is Left H.S. in Right H.S. part of the
   ///                     update expression, false otherwise.
   ///                     (e.g. true for X = X BinOp Expr)
   ///
   /// \returns A pair of the old value of X before the update, and the value
   ///          used for the update.
-  std::pair<Value *, Value *> emitAtomicUpdate(Instruction *AllocIP, Value *X,
+  std::pair<Value *, Value *> emitAtomicUpdate(InsertPointTy AllocaIP, Value *X,
                                                Value *Expr, AtomicOrdering AO,
                                                AtomicRMWInst::BinOp RMWOp,
                                                AtomicUpdateCallbackTy &UpdateOp,
                                                bool VolatileX,
                                                bool IsXBinopExpr);
 
   /// Emit the binary op. described by \p RMWOp, using \p Src1 and \p Src2 .
   ///
InsertPointTy createAtomicWrite(const LocationDescription &Loc,
                                   AtomicOpValue &X, Value *Expr,
                                   AtomicOrdering AO);
 
   /// Emit atomic update for constructs: X = X BinOp Expr ,or X = Expr BinOp X
   /// For complex Operations: X = UpdateOp(X) => CmpExch X, old_X, UpdateOp(X)
   /// Only Scalar data types.
   ///
   /// \param Loc      The insert and source location description.
-  /// \param AllocIP  Instruction to create AllocaInst before.
+  /// \param AllocaIP The insertion point to be used for alloca instructions.
   /// \param X        The target atomic pointer to be updated
   /// \param Expr     The value to update X with.
   /// \param AO       Atomic ordering of the generated atomic instructions.
   /// \param RMWOp    The binary operation used for update. If operation
   ///                 is  not supported by atomicRMW, or belong to
   ///                 {FADD, FSUB, BAD_BINOP}. Then a `cmpExch` based
   ///                 atomic will be generated.
   /// \param UpdateOp   Code generator for complex expressions that cannot be
   ///                   expressed through atomicrmw instruction.
   /// \param IsXBinopExpr true if \a X is Left H.S. in Right H.S. part of the
   ///                     update expression, false otherwise.
   ///                     (e.g. true for X = X BinOp Expr)
   ///
   /// \return Insertion point after generated atomic update IR.
   InsertPointTy createAtomicUpdate(const LocationDescription &Loc,
-                                   Instruction *AllocIP, AtomicOpValue &X,
+                                   InsertPointTy AllocaIP, AtomicOpValue &X,
                                    Value *Expr, AtomicOrdering AO,
                                    AtomicRMWInst::BinOp RMWOp,
                                    AtomicUpdateCallbackTy &UpdateOp,
                                    bool IsXBinopExpr);
 
   /// Emit atomic update for constructs: --- Only Scalar data types
   /// V = X; X = X BinOp Expr ,
   /// X = X BinOp Expr; V = X,
   /// V = X; X = Expr BinOp X,
   /// X = Expr BinOp X; V = X,
   /// V = X; X = UpdateOp(X),
   /// X = UpdateOp(X); V = X,
   ///
   /// \param Loc        The insert and source location description.
-  /// \param AllocIP    Instruction to create AllocaInst before.
+  /// \param AllocaIP   The insertion point to be used for alloca instructions.
   /// \param X          The target atomic pointer to be updated
   /// \param V          Memory address where to store captured value
   /// \param Expr       The value to update X with.
   /// \param AO         Atomic ordering of the generated atomic instructions
   /// \param RMWOp      The binary operation used for update. If
   ///                   operation is not supported by atomicRMW, or belong to
   ///                   {FADD, FSUB, BAD_BINOP}. Then a cmpExch based
   ///                   atomic will be generated.
   /// \param UpdateOp   Code generator for complex expressions that cannot be
   ///                   expressed through atomicrmw instruction.
   /// \param UpdateExpr true if X is an in place update of the form
   ///                   X = X BinOp Expr or X = Expr BinOp X
   /// \param IsXBinopExpr true if X is Left H.S. in Right H.S. part of the
   ///                     update expression, false otherwise.
   ///                     (e.g. true for X = X BinOp Expr)
   /// \param IsPostfixUpdate true if original value of 'x' must be stored in
   ///                        'v', not an updated one.
   ///
   /// \return Insertion point after generated atomic capture IR.
   InsertPointTy
-  createAtomicCapture(const LocationDescription &Loc, Instruction *AllocIP,
+  createAtomicCapture(const LocationDescription &Loc, InsertPointTy AllocaIP,
                       AtomicOpValue &X, AtomicOpValue &V, Value *Expr,
                       AtomicOrdering AO, AtomicRMWInst::BinOp RMWOp,
                       AtomicUpdateCallbackTy &UpdateOp, bool UpdateExpr,
                       bool IsPostfixUpdate, bool IsXBinopExpr);
 
   /// Create the control flow structure of a canonical OpenMP loop.
   ///
   /// The emitted loop will be disconnected, i.e. no edge to the loop's

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

if (XElemTy->isIntegerTy()) {
     XSt->setAtomic(AO);
   }
 
   checkAndEmitFlushAfterAtomic(Loc, AO, AtomicKind::Write);
   return Builder.saveIP();
 }
 
 OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createAtomicUpdate(
-    const LocationDescription &Loc, Instruction *AllocIP, AtomicOpValue &X,
+    const LocationDescription &Loc, InsertPointTy AllocaIP, AtomicOpValue &X,
     Value *Expr, AtomicOrdering AO, AtomicRMWInst::BinOp RMWOp,
     AtomicUpdateCallbackTy &UpdateOp, bool IsXBinopExpr) {
   if (!updateToLocation(Loc))
     return Loc.IP;
 
   LLVM_DEBUG({
     Type *XTy = X.Var->getType();
     assert(XTy->isPointerTy() &&
            "OMP Atomic expects a pointer to target memory");
     Type *XElemTy = XTy->getPointerElementType();
     assert((XElemTy->isFloatingPointTy() || XElemTy->isIntegerTy() ||
             XElemTy->isPointerTy()) &&
            "OMP atomic update expected a scalar type");
     assert((RMWOp != AtomicRMWInst::Max) && (RMWOp != AtomicRMWInst::Min) &&
            (RMWOp != AtomicRMWInst::UMax) && (RMWOp != AtomicRMWInst::UMin) &&
            "OpenMP atomic does not support LT or GT operations");
   });
 
-  emitAtomicUpdate(AllocIP, X.Var, Expr, AO, RMWOp, UpdateOp, X.IsVolatile,
+  emitAtomicUpdate(AllocaIP, X.Var, Expr, AO, RMWOp, UpdateOp, X.IsVolatile,
                    IsXBinopExpr);
   checkAndEmitFlushAfterAtomic(Loc, AO, AtomicKind::Update);
   return Builder.saveIP();
 }
 
 Value *OpenMPIRBuilder::emitRMWOpAsInstruction(Value *Src1, Value *Src2,
                                                AtomicRMWInst::BinOp RMWOp) {
   switch (RMWOp) {
Value *OpenMPIRBuilder::emitRMWOpAsInstruction(Value *Src1, Value *Src2,
   case AtomicRMWInst::UMax:
   case AtomicRMWInst::UMin:
     llvm_unreachable("Unsupported atomic update operation");
   }
   llvm_unreachable("Unsupported atomic update operation");
 }
 
 std::pair<Value *, Value *>
-OpenMPIRBuilder::emitAtomicUpdate(Instruction *AllocIP, Value *X, Value *Expr,
+OpenMPIRBuilder::emitAtomicUpdate(InsertPointTy AllocaIP, Value *X, Value *Expr,
                                   AtomicOrdering AO, AtomicRMWInst::BinOp RMWOp,
                                   AtomicUpdateCallbackTy &UpdateOp,
                                   bool VolatileX, bool IsXBinopExpr) {
   Type *XElemTy = X->getType()->getPointerElementType();
 
   bool DoCmpExch =
       ((RMWOp == AtomicRMWInst::BAD_BINOP) || (RMWOp == AtomicRMWInst::FAdd)) ||
       (RMWOp == AtomicRMWInst::FSub) ||
if (XElemTy->isIntegerTy() && !DoCmpExch) {
     BasicBlock *CurBB = Builder.GetInsertBlock();
     Instruction *CurBBTI = CurBB->getTerminator();
     CurBBTI = CurBBTI ? CurBBTI : Builder.CreateUnreachable();
     BasicBlock *ExitBB =
         CurBB->splitBasicBlock(CurBBTI, X->getName() + ".atomic.exit");
     BasicBlock *ContBB = CurBB->splitBasicBlock(CurBB->getTerminator(),
                                                 X->getName() + ".atomic.cont");
     ContBB->getTerminator()->eraseFromParent();
     Builder.SetInsertPoint(ContBB);
     llvm::PHINode *PHI = Builder.CreatePHI(OldVal->getType(), 2);
     PHI->addIncoming(OldVal, CurBB);
+    Builder.restoreIP(AllocaIP);
     AllocaInst *NewAtomicAddr = Builder.CreateAlloca(XElemTy);
     NewAtomicAddr->setName(X->getName() + "x.new.val");
-    NewAtomicAddr->moveBefore(AllocIP);
+    Builder.SetInsertPoint(ContBB);
     IntegerType *NewAtomicCastTy =
         IntegerType::get(M.getContext(), XElemTy->getScalarSizeInBits());
     bool IsIntTy = XElemTy->isIntegerTy();
     Value *NewAtomicIntAddr =
         (IsIntTy)
             ? NewAtomicAddr
             : Builder.CreateBitCast(NewAtomicAddr,
                                     NewAtomicCastTy->getPointerTo(Addrspace));
if (UnreachableInst *ExitTI =
       Builder.SetInsertPoint(ExitTI);
     }
   }
 
   return Res;
 }
 
 OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createAtomicCapture(
-    const LocationDescription &Loc, Instruction *AllocIP, AtomicOpValue &X,
+    const LocationDescription &Loc, InsertPointTy AllocaIP, AtomicOpValue &X,
     AtomicOpValue &V, Value *Expr, AtomicOrdering AO,
     AtomicRMWInst::BinOp RMWOp, AtomicUpdateCallbackTy &UpdateOp,
     bool UpdateExpr, bool IsPostfixUpdate, bool IsXBinopExpr) {
   if (!updateToLocation(Loc))
     return Loc.IP;
 
   LLVM_DEBUG({
     Type *XTy = X.Var->getType();
     assert(XTy->isPointerTy() &&
            "OMP Atomic expects a pointer to target memory");
     Type *XElemTy = XTy->getPointerElementType();
     assert((XElemTy->isFloatingPointTy() || XElemTy->isIntegerTy() ||
             XElemTy->isPointerTy()) &&
            "OMP atomic capture expected a scalar type");
     assert((RMWOp != AtomicRMWInst::Max) && (RMWOp != AtomicRMWInst::Min) &&
            "OpenMP atomic does not support LT or GT operations");
   });
 
   // If UpdateExpr is 'x' updated with some `expr` not based on 'x',
   // 'x' is simply atomically rewritten with 'expr'.
   AtomicRMWInst::BinOp AtomicOp = (UpdateExpr ? RMWOp : AtomicRMWInst::Xchg);
-  std::pair<Value *, Value *> Result = emitAtomicUpdate(
-      AllocIP, X.Var, Expr, AO, AtomicOp, UpdateOp, X.IsVolatile, IsXBinopExpr);
+  std::pair<Value *, Value *> Result =
+      emitAtomicUpdate(AllocaIP, X.Var, Expr, AO, AtomicOp, UpdateOp,
+                       X.IsVolatile, IsXBinopExpr);
 
   Value *CapturedVal = (IsPostfixUpdate ? Result.first : Result.second);
   Builder.CreateStore(CapturedVal, V.Var, V.IsVolatile);
 
   checkAndEmitFlushAfterAtomic(Loc, AO, AtomicKind::Capture);
   return Builder.saveIP();
 }
 

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

TEST_F(OpenMPIRBuilderTest, OMPAtomicUpdate) {
   OpenMPIRBuilder::AtomicOpValue X = {XVal, false, false};
   AtomicOrdering AO = AtomicOrdering::Monotonic;
   ConstantInt *ConstVal = ConstantInt::get(Type::getInt32Ty(Ctx), 1U);
   Value *Expr = nullptr;
   AtomicRMWInst::BinOp RMWOp = AtomicRMWInst::Sub;
   bool IsXLHSInRHSPart = false;
 
   BasicBlock *EntryBB = BB;
-  Instruction *AllocIP = EntryBB->getFirstNonPHI();
+  OpenMPIRBuilder::InsertPointTy AllocaIP(EntryBB,
+                                          EntryBB->getFirstInsertionPt());
   Value *Sub = nullptr;
 
   auto UpdateOp = [&](Value *Atomic, IRBuilder<> &IRB) {
     Sub = IRB.CreateSub(ConstVal, Atomic);
     return Sub;
   };
   Builder.restoreIP(OMPBuilder.createAtomicUpdate(
-      Builder, AllocIP, X, Expr, AO, RMWOp, UpdateOp, IsXLHSInRHSPart));
+      Builder, AllocaIP, X, Expr, AO, RMWOp, UpdateOp, IsXLHSInRHSPart));
   BasicBlock *ContBB = EntryBB->getSingleSuccessor();
   BranchInst *ContTI = dyn_cast<BranchInst>(ContBB->getTerminator());
   EXPECT_NE(ContTI, nullptr);
   BasicBlock *EndBB = ContTI->getSuccessor(0);
   EXPECT_TRUE(ContTI->isConditional());
   EXPECT_EQ(ContTI->getSuccessor(1), ContBB);
   EXPECT_NE(EndBB, nullptr);
 
TEST_F(OpenMPIRBuilderTest, OMPAtomicCapture) {
   AtomicOrdering AO = AtomicOrdering::Monotonic;
   ConstantInt *Expr = ConstantInt::get(Type::getInt32Ty(Ctx), 1U);
   AtomicRMWInst::BinOp RMWOp = AtomicRMWInst::Add;
   bool IsXLHSInRHSPart = true;
   bool IsPostfixUpdate = true;
   bool UpdateExpr = true;
 
   BasicBlock *EntryBB = BB;
-  Instruction *AllocIP = EntryBB->getFirstNonPHI();
+  OpenMPIRBuilder::InsertPointTy AllocaIP(EntryBB,
+                                          EntryBB->getFirstInsertionPt());
 
   // integer update - not used
   auto UpdateOp = [&](Value *Atomic, IRBuilder<> &IRB) { return nullptr; };
 
   Builder.restoreIP(OMPBuilder.createAtomicCapture(
-      Builder, AllocIP, X, V, Expr, AO, RMWOp, UpdateOp, UpdateExpr,
+      Builder, AllocaIP, X, V, Expr, AO, RMWOp, UpdateOp, UpdateExpr,
       IsPostfixUpdate, IsXLHSInRHSPart));
   EXPECT_EQ(EntryBB->getParent()->size(), 1U);
   AtomicRMWInst *ARWM = dyn_cast<AtomicRMWInst>(Init->getNextNode());
   EXPECT_NE(ARWM, nullptr);
   EXPECT_EQ(ARWM->getPointerOperand(), XVal);
   EXPECT_EQ(ARWM->getOperation(), RMWOp);
   StoreInst *St = dyn_cast<StoreInst>(ARWM->user_back());
   EXPECT_NE(St, nullptr);