[llvm] [NVPTX][AtomicExpandPass] Complete support for AtomicRMW in NVPTX (PR #176015)

[Alex MacLean via llvm-commits]

================
@@ -7066,40 +7066,70 @@ NVPTXTargetLowering::shouldExpandAtomicRMWInIR(const AtomicRMWInst *AI) const {
 bool NVPTXTargetLowering::shouldInsertFencesForAtomic(
     const Instruction *I) const {
   auto *CI = dyn_cast<AtomicCmpXchgInst>(I);
+  auto *RI = dyn_cast<AtomicRMWInst>(I);
   // When CAS bitwidth is not supported on the hardware, the CAS is emulated
-  // using a retry loop that uses a higher-bitwidth monotonic CAS. We enforce
-  // the memory order using explicit fences around the retry loop.
-  // The memory order of natively supported CAS operations can be enforced
-  // by lowering to an atom.cas with the right memory synchronizing effect.
-  // However, atom.cas only supports relaxed, acquire, release and acq_rel.
-  // So we also use explicit fences for enforcing memory order for
-  // seq_cast CAS with natively-supported bitwidths.
-  return CI &&
-         (cast<IntegerType>(CI->getCompareOperand()->getType())->getBitWidth() <
-              STI.getMinCmpXchgSizeInBits() ||
-          CI->getMergedOrdering() == AtomicOrdering::SequentiallyConsistent);
+  // using a retry loop that uses a higher-bitwidth monotonic CAS. Similarly, if
+  // the atomicrmw operation is not supported on hardware, we emulate it with a
+  // cmpxchg loop. In such cases, we enforce the memory order using explicit
+  // fences around the retry loop.
+  // The memory order of natively supported CAS or RMW operations can be
+  // enforced by lowering to an `atom.<op>` instr with the right memory
+  // synchronization effect.  However, atom only supports relaxed, acquire,
+  // release and acq_rel.  So we also use explicit fences to enforce memory
+  // order in seq_cast CAS or RMW instructions that can be lowered as acq_rel.
+  if (CI)
+    return (cast<IntegerType>(CI->getCompareOperand()->getType())
+                ->getBitWidth() < STI.getMinCmpXchgSizeInBits()) ||
+           CI->getMergedOrdering() == AtomicOrdering::SequentiallyConsistent;
+  if (RI)
+    return shouldExpandAtomicRMWInIR(RI) == AtomicExpansionKind::CmpXChg ||
+           RI->getOrdering() == AtomicOrdering::SequentiallyConsistent;
+  return false;
 }
 
 AtomicOrdering NVPTXTargetLowering::atomicOperationOrderAfterFenceSplit(
     const Instruction *I) const {
-  auto *CI = dyn_cast<AtomicCmpXchgInst>(I);
-  bool BitwidthSupportedAndIsSeqCst =
-      CI && CI->getMergedOrdering() == AtomicOrdering::SequentiallyConsistent &&
+  // Only lower to atom.<op>.acquire if the operation is not emulated, and its
+  // ordering is seq_cst. This produces a sequence of the form:
+  // fence.sc
+  // atom.<op>.acquire
+  // Instead of
+  // fence.sc
+  // atom.<op>
+  // fence.acquire
+  // The two-instruction sequence is weaker than the alternative, but guarantees
+  // seq_cst ordering.
+  //
+  // In all other cases, lower to atom.<op>.relaxed
+  const auto *CI = dyn_cast<AtomicCmpXchgInst>(I);
+  const auto *RI = dyn_cast<AtomicRMWInst>(I);
+  AtomicOrdering Ordering;
----------------
AlexMaclean wrote:

Nit: as above, I think this would be a little clearer if the dyn_casts were in if checks, and the result were directly returned rather than being stored in `Ordering` first. 

https://github.com/llvm/llvm-project/pull/176015