[llvm] [AtomicExpand] Add elementwise expansion (PR #196464)

[Harrison Hao via llvm-commits]

================
@@ -382,30 +385,41 @@ bool AtomicExpandImpl::processAtomicInstr(Instruction *I) {
   }
 
   if (auto *RMWI = dyn_cast<AtomicRMWInst>(I)) {
+    bool IsElementwiseExpand =
+        RMWI->isElementwise() &&
+        TLI->shouldExpandAtomicRMWInIR(RMWI) ==
+            TargetLoweringBase::AtomicExpansionKind::Expand;
+
     if (!atomicSizeSupported(TLI, RMWI)) {
-      expandAtomicRMWToLibcall(RMWI);
-      return true;
+      // Elementwise expansion may split an atomicrmw into smaller, supported
+      // atomic sizes.
+      if (!IsElementwiseExpand) {
+        expandAtomicRMWToLibcall(RMWI);
+        return true;
+      }
     }
 
     bool MadeChange = false;
-    if (TLI->shouldCastAtomicRMWIInIR(RMWI) ==
-        TargetLoweringBase::AtomicExpansionKind::CastToInteger) {
+    if (!IsElementwiseExpand &&
+        TLI->shouldCastAtomicRMWIInIR(RMWI) ==
+            TargetLoweringBase::AtomicExpansionKind::CastToInteger) {
       RMWI = convertAtomicXchgToIntegerType(RMWI);
       MadeChange = true;
     }
 
     MadeChange |= tryInsertFencesForAtomic(
-        RMWI,
-        isReleaseOrStronger(RMWI->getOrdering()) ||
-            isAcquireOrStronger(RMWI->getOrdering()),
+        RMWI, isStrongerThanMonotonic(RMWI->getOrdering()),
         TLI->atomicOperationOrderAfterFenceSplit(RMWI));
 
     // There are two different ways of expanding RMW instructions:
     // - into a load if it is idempotent
     // - into a Cmpxchg/LL-SC loop otherwise
     // we try them in that order.
-    MadeChange |= (isIdempotentRMW(RMWI) && simplifyIdempotentRMW(RMWI)) ||
-                  tryExpandAtomicRMW(RMWI);
+    if (IsElementwiseExpand)
----------------
harrisonGPU wrote:

Should we return false from `isIdempotentRMW()` and `simplifyIdempotentRMW()` when RMWI is elementwise, instead of handling it here?

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