[llvm] [SROA] Extend tree-structured merge to handle init + RMW pattern (PR #194441)

[Antonio Frighetto via llvm-commits]

================
@@ -2895,143 +2915,339 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
 
     for (Slice &S : P) {
       auto *User = cast<Instruction>(S.getUse()->getUser());
+      // A "full-width" slice spans the entire alloca; it's either the single
+      // init store (Pattern 2) or the single final load (both patterns).
+      bool IsFullWidth = (S.beginOffset() == NewAllocaBeginOffset &&
+                          S.endOffset() == NewAllocaEndOffset);
       if (auto *LI = dyn_cast<LoadInst>(User)) {
-        // Do not handle the case if
-        //   1. There is more than one load
-        //   2. The load is volatile
-        //   3. The load does not read the entire alloca structure
-        //   4. The load does not meet the conditions in the helper function
-        if (TheLoad || !IsTypeValidForTreeStructuredMerge(LI->getType()) ||
-            S.beginOffset() != NewAllocaBeginOffset ||
-            S.endOffset() != NewAllocaEndOffset || LI->isVolatile())
+        // Only handle simple (non-volatile, non-atomic) loads.
+        if (!LI->isSimple() ||
+            !IsTypeValidForTreeStructuredMerge(LI->getType()))
           return std::nullopt;
-        TheLoad = LI;
+        if (IsFullWidth) {
+          // We accept at most one full-width load (the "final" load, after
+          // all the partial stores).
+          if (FullLoad)
+            return std::nullopt;
+          FullLoad = LI;
+        } else {
+          // Partial load (RMW pattern only).
+          LoadInfos.push_back({LI, S.beginOffset(), S.endOffset()});
+        }
       } else if (auto *SI = dyn_cast<StoreInst>(User)) {
         // Do not handle the case if
         //   1. The store does not meet the conditions in the helper function
-        //   2. The store is volatile
-        //   3. The total store size is not a multiple of the allocated element
-        //   type size
-        if (!IsTypeValidForTreeStructuredMerge(
-                SI->getValueOperand()->getType()) ||
-            SI->isVolatile())
+        //   2. The store is not simple — we drop stores as part of the
+        //      rewrite, so volatile stores (which must be kept) and atomic
+        //      stores (which carry memory-ordering semantics) are unsound
+        //      to replace with SSA bookkeeping.
+        //   3. The total store size is not a multiple of the allocated
+        //      element type size (required so the tree merge can produce a
+        //      vector whose element type matches the alloca).
+        if (!SI->isSimple() || !IsTypeValidForTreeStructuredMerge(
+                                   SI->getValueOperand()->getType()))
           return std::nullopt;
-        auto *VecTy = cast<FixedVectorType>(SI->getValueOperand()->getType());
-        unsigned NumElts = VecTy->getNumElements();
-        unsigned EltSize = DL.getTypeSizeInBits(VecTy->getElementType());
+        auto *StVecTy = cast<FixedVectorType>(SI->getValueOperand()->getType());
+        unsigned NumElts = StVecTy->getNumElements();
+        unsigned EltSize = DL.getTypeSizeInBits(StVecTy->getElementType());
         if (NumElts * EltSize % AllocatedEltTySize != 0)
           return std::nullopt;
-        StoreInfos.emplace_back(SI, S.beginOffset(), S.endOffset(),
-                                SI->getValueOperand());
+        if (IsFullWidth) {
+          // At most one full-width store is allowed — it's the init store
+          // for the RMW pattern.
+          if (InitStore)
+            return std::nullopt;
+          InitStore = SI;
+        } else {
+          StoreInfos.emplace_back(SI, S.beginOffset(), S.endOffset(),
+                                  SI->getValueOperand());
+        }
       } else {
-        // If we have instructions other than load and store, we cannot do the
-        // tree structured merge
+        // If we have instructions other than load and store, we cannot do
+        // the tree structured merge.
         return std::nullopt;
       }
     }
-    // If we do not have any load, we cannot do the tree structured merge
-    if (!TheLoad)
+
+    // Classify the pattern by looking at what we collected:
+    //   Pattern 1 (stores-only): only partial stores + exactly one full load.
+    //   Pattern 2 (RMW): one full init store + partial loads + partial stores
+    //     (+ optional full final load). RMW also needs VecTy to be set
+    //     because we use getIndex() to convert byte offsets to element
+    //     indices, which requires a promoted vector alloca.
+    bool IsRMWPattern = InitStore && VecTy && !LoadInfos.empty();
+    bool IsStoresOnlyPattern = !InitStore && FullLoad && LoadInfos.empty();
+    if (!IsRMWPattern && !IsStoresOnlyPattern)
----------------
antoniofrighetto wrote:

I'm not sure if it may be common for real GPU kernels, could we have a pattern with an init store + partial stores, and no partial loads? This would be rejected?

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