[OMPIRBuilder] Add support for atomic compare

[tianshilei1992]

Summary

This patch adds the support for atomic compare in OMPIRBuilder.

[tianshilei1992]

fix typo

[tianshilei1992]

rebase

[tianshilei1992]

add test

[jdoerfert]

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

1347

Brief documentation, prefer an enum class, potentially in OpenMPConstants.h

1372

Op is missing above.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

3531

Once you use a conditional and once a ternary op. Maybe use a helper lambda instead.
Signedness has to come from the user.

[nikic]

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

3488

Should be X.ElemTy.

[jdoerfert]

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

3488

I have to start paying attention to the getPointerElementType uses.

[tianshilei1992]

fix comments

[tianshilei1992]

rebase and remove unused arguments

[tianshilei1992]

rebase and remove unused comments

[jdoerfert]

LG, minor things to address

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

3490

assertions don't go into LLVM_DEBUG.

3504

assertions should have a message

3523

SMin, SMax

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

3078

Add SMax test too.

[tianshilei1992]

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

3523

There is no value called SMin or SMax.

[tianshilei1992]

fix comments

llvm/include/llvm/Frontend/OpenMP/OMPConstants.h

enum class AddressSpace : unsigned {
   Shared = 3,
   Constant = 4,
   Local = 5,
 };
 
 /// \note This needs to be kept in sync with interop.h enum kmp_interop_type_t.:
 enum class OMPInteropType { Unknown, Target, TargetSync };
 
+/// Atomic compare operations. Currently OpenMP only supports ==, >, and <.
+enum class OMPAtomicCompareOp : unsigned { EQ, MIN, MAX };
+
 } // end namespace omp
 
 } // end namespace llvm
 
 #endif // LLVM_FRONTEND_OPENMP_OMPCONSTANTS_H

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

private:
   /// \param XOld the value of the atomic memory address to use for update
   /// \param IRB reference to the IRBuilder to use
   ///
   /// \returns Value to update X to.
   using AtomicUpdateCallbackTy =
       const function_ref<Value *(Value *XOld, IRBuilder<> &IRB)>;
 
 private:
-  enum AtomicKind { Read, Write, Update, Capture };
+  enum AtomicKind { Read, Write, Update, Capture, Compare };
 
   /// Determine whether to emit flush or not
   ///
   /// \param Loc    The insert and source location description.
   /// \param AO     The required atomic ordering
   /// \param AK     The OpenMP atomic operation kind used.
   ///
   /// \returns    wether a flush was emitted or not
public:
   /// \return Insertion point after generated atomic capture IR.
   InsertPointTy
   createAtomicCapture(const LocationDescription &Loc, Instruction *AllocIP,
                       AtomicOpValue &X, AtomicOpValue &V, Value *Expr,
                       AtomicOrdering AO, AtomicRMWInst::BinOp RMWOp,
                       AtomicUpdateCallbackTy &UpdateOp, bool UpdateExpr,
                       bool IsPostfixUpdate, bool IsXBinopExpr);
 
+  /// Emit atomic compare for constructs: --- Only scalar data types
+  /// cond-update-atomic:
+  /// x = x ordop expr ? expr : x;
+  /// x = expr ordop x ? expr : x;
+  /// x = x == e ? d : x;
+  /// x = e == x ? d : x; (this one is not in the spec)
+  /// cond-update-stmt:
+  /// if (x ordop expr) { x = expr; }
+  /// if (expr ordop x) { x = expr; }
+  /// if (x == e) { x = d; }
+  /// if (e == x) { x = d; } (this one is not in the spec)
+  ///
+  /// \param Loc          The insert and source location description.
+  /// \param AllocIP      Instruction to create AllocaInst before.
+  /// \param X            The target atomic pointer to be updated.
+  /// \param E            The expected value ('e') for forms that use an
+  ///                     equality comparison or an expression ('expr') for
+  ///                     forms that use 'ordop' (logically an atomic maximum or
+  ///                     minimum).
+  /// \param D            The desired value for forms that use an equality
+  ///                     comparison. If forms that use 'ordop', it should be
+  ///                     \p nullptr.
+  /// \param AO           Atomic ordering of the generated atomic instructions.
+  /// \param OP           Atomic compare operation. It can only be ==, <, or >.
+  /// \param IsXBinopExpr True if the conditional statement is in the form where
+  ///                     x is on LHS. It only matters for < or >.
+  /// \param IsSignedOp   If the operation is signed or unsigned. It only
+  ///                     matters for < or >.
+  ///
+  /// \return Insertion point after generated atomic capture IR.
+  InsertPointTy createAtomicCompare(const LocationDescription &Loc,
+                                    Instruction *AllocIP, AtomicOpValue &X,
+                                    Value *E, Value *D, AtomicOrdering AO,
+                                    omp::OMPAtomicCompareOp Op,
+                                    bool IsXBinopExpr, bool IsSignedOp);
+
   /// Create the control flow structure of a canonical OpenMP loop.
   ///
   /// The emitted loop will be disconnected, i.e. no edge to the loop's
   /// preheader and no terminator in the AfterBB. The OpenMPIRBuilder's
   /// IRBuilder location is not preserved.
   ///
   /// \param DL        DebugLoc used for the instructions in the skeleton.
   /// \param TripCount Value to be used for the trip count.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

bool OpenMPIRBuilder::checkAndEmitFlushAfterAtomic(
   case Read:
     if (AO == AtomicOrdering::Acquire || AO == AtomicOrdering::AcquireRelease ||
         AO == AtomicOrdering::SequentiallyConsistent) {
       FlushAO = AtomicOrdering::Acquire;
       Flush = true;
     }
     break;
   case Write:
+  case Compare:
   case Update:
     if (AO == AtomicOrdering::Release || AO == AtomicOrdering::AcquireRelease ||
         AO == AtomicOrdering::SequentiallyConsistent) {
       FlushAO = AtomicOrdering::Release;
       Flush = true;
     }
     break;
   case Capture:
OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createAtomicCapture(
 
   Value *CapturedVal = (IsPostfixUpdate ? Result.first : Result.second);
   Builder.CreateStore(CapturedVal, V.Var, V.IsVolatile);
 
   checkAndEmitFlushAfterAtomic(Loc, AO, AtomicKind::Capture);
   return Builder.saveIP();
 }
 
+OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createAtomicCompare(
+    const LocationDescription &Loc, Instruction *AllocIP, AtomicOpValue &X,
+    Value *E, Value *D, AtomicOrdering AO, OMPAtomicCompareOp Op,
+    bool IsXBinopExpr, bool IsSignedOp) {
+  if (!updateToLocation(Loc))
+    return Loc.IP;
+
+  LLVM_DEBUG({
+    assert(X.Var->getType()->isPointerTy() &&
+           "OMP atomic expects a pointer to target memory");
+    assert((X.ElemTy->isFloatingPointTy() || X.ElemTy->isIntegerTy() ||
+            X.ElemTy->isPointerTy()) &&
+           "OMP atomic compare expected a scalar type");
+  });
+
+  if (Op == OMPAtomicCompareOp::EQ) {
+    unsigned Addrspace = cast<PointerType>(X.Var->getType())->getAddressSpace();
+    IntegerType *IntCastTy =
+        IntegerType::get(M.getContext(), X.ElemTy->getScalarSizeInBits());
+    Value *XAddr =
+        X.ElemTy->isIntegerTy()
+            ? X.Var
+            : Builder.CreateBitCast(X.Var, IntCastTy->getPointerTo(Addrspace));
+    AtomicOrdering Failure = AtomicCmpXchgInst::getStrongestFailureOrdering(AO);
+    // We don't need the result for now.
+    (void)Builder.CreateAtomicCmpXchg(XAddr, E, D, MaybeAlign(), AO, Failure);
+  } else {
+    assert(Op == OMPAtomicCompareOp::MAX || Op == OMPAtomicCompareOp::MIN);
+    assert(X.ElemTy->isIntegerTy() &&
+           "max and min operators only support integer type");
+
+    // Reverse the ordop as the OpenMP forms are different from LLVM forms.
+    // Let's take max as example.
+    // OpenMP form:
+    // x = x > expr ? expr : x;
+    // LLVM form:
+    // *ptr = *ptr > val ? *ptr : val;
+    // We need to transform to LLVM form.
+    // x = x <= expr ? x : expr;
+    AtomicRMWInst::BinOp NewOp;
+    if (IsXBinopExpr) {
+      if (IsSignedOp)
+        NewOp = Op == OMPAtomicCompareOp::MAX ? AtomicRMWInst::Min
+                                              : AtomicRMWInst::Max;
+      else
+        NewOp = Op == OMPAtomicCompareOp::MAX ? AtomicRMWInst::UMin
+                                              : AtomicRMWInst::UMax;
+    } else {
+      if (IsSignedOp)
+        NewOp = Op == OMPAtomicCompareOp::MAX ? AtomicRMWInst::Max
+                                              : AtomicRMWInst::Min;
+      else
+        NewOp = Op == OMPAtomicCompareOp::MAX ? AtomicRMWInst::UMax
+                                              : AtomicRMWInst::UMin;
+    }
+    // We dont' need the result for now.
+    (void)Builder.CreateAtomicRMW(NewOp, X.Var, E, MaybeAlign(), AO);
+  }
+
+  checkAndEmitFlushAfterAtomic(Loc, AO, AtomicKind::Compare);
+
+  return Builder.saveIP();
+}
+
 GlobalVariable *
 OpenMPIRBuilder::createOffloadMapnames(SmallVectorImpl<llvm::Constant *> &Names,
                                        std::string VarName) {
   llvm::Constant *MapNamesArrayInit = llvm::ConstantArray::get(
       llvm::ArrayType::get(
           llvm::Type::getInt8Ty(M.getContext())->getPointerTo(), Names.size()),
       Names);
   auto *MapNamesArrayGlobal = new llvm::GlobalVariable(

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

TEST_F(OpenMPIRBuilderTest, OMPAtomicCapture) {
   EXPECT_NE(St, nullptr);
   EXPECT_EQ(St->getPointerOperand(), VVal);
 
   Builder.CreateRetVoid();
   OMPBuilder.finalize();
   EXPECT_FALSE(verifyModule(*M, &errs()));
 }
 
+TEST_F(OpenMPIRBuilderTest, OMPAtomicCompare) {
+  OpenMPIRBuilder OMPBuilder(*M);
+  OMPBuilder.initialize();
+  F->setName("func");
+  IRBuilder<> Builder(BB);
+
+  OpenMPIRBuilder::LocationDescription Loc({Builder.saveIP(), DL});
+
+  LLVMContext &Ctx = M->getContext();
+  IntegerType *Int32 = Type::getInt32Ty(Ctx);
+  AllocaInst *XVal = Builder.CreateAlloca(Int32);
+  XVal->setName("x");
+  StoreInst *Init =
+      Builder.CreateStore(ConstantInt::get(Type::getInt32Ty(Ctx), 0U), XVal);
+
+  OpenMPIRBuilder::AtomicOpValue X = {XVal, Int32, false, false};
+  AtomicOrdering AO = AtomicOrdering::Monotonic;
+  ConstantInt *Expr = ConstantInt::get(Type::getInt32Ty(Ctx), 1U);
+  ConstantInt *D = ConstantInt::get(Type::getInt32Ty(Ctx), 1U);
+  OMPAtomicCompareOp OpMax = OMPAtomicCompareOp::MAX;
+  OMPAtomicCompareOp OpEQ = OMPAtomicCompareOp::EQ;
+
+  BasicBlock *EntryBB = BB;
+  Instruction *AllocIP = EntryBB->getFirstNonPHI();
+
+  Builder.restoreIP(OMPBuilder.createAtomicCompare(
+      Builder, AllocIP, X, Expr, nullptr, AO, OpMax, true, true));
+  Builder.restoreIP(OMPBuilder.createAtomicCompare(
+      Builder, AllocIP, X, Expr, nullptr, AO, OpMax, false, false));
+  Builder.restoreIP(OMPBuilder.createAtomicCompare(Builder, AllocIP, X, Expr, D,
+                                                   AO, OpEQ, true, true));
+
+  EXPECT_EQ(EntryBB->getParent()->size(), 1U);
+  EXPECT_EQ(EntryBB->size(), 5U);
+
+  AtomicRMWInst *ARWM1 = dyn_cast<AtomicRMWInst>(Init->getNextNode());
+  EXPECT_NE(ARWM1, nullptr);
+  EXPECT_EQ(ARWM1->getPointerOperand(), XVal);
+  EXPECT_EQ(ARWM1->getValOperand(), Expr);
+  EXPECT_EQ(ARWM1->getOperation(), AtomicRMWInst::Min);
+
+  AtomicRMWInst *ARWM2 = dyn_cast<AtomicRMWInst>(ARWM1->getNextNode());
+  EXPECT_NE(ARWM2, nullptr);
+  EXPECT_EQ(ARWM2->getPointerOperand(), XVal);
+  EXPECT_EQ(ARWM2->getValOperand(), Expr);
+  EXPECT_EQ(ARWM2->getOperation(), AtomicRMWInst::UMax);
+
+  AtomicCmpXchgInst *AXCHG = dyn_cast<AtomicCmpXchgInst>(ARWM2->getNextNode());
+  EXPECT_NE(AXCHG, nullptr);
+  EXPECT_EQ(AXCHG->getPointerOperand(), XVal);
+  EXPECT_EQ(AXCHG->getCompareOperand(), Expr);
+  EXPECT_EQ(AXCHG->getNewValOperand(), D);
+
+  Builder.CreateRetVoid();
+  OMPBuilder.finalize();
+  EXPECT_FALSE(verifyModule(*M, &errs()));
+}
+
 /// Returns the single instruction of InstTy type in BB that uses the value V.
 /// If there is more than one such instruction, returns null.
 template <typename InstTy>
 static InstTy *findSingleUserInBlock(Value *V, BasicBlock *BB) {
   InstTy *Result = nullptr;
   for (User *U : V->users()) {
     auto *Inst = dyn_cast<InstTy>(U);
     if (!Inst || Inst->getParent() != BB)