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

[Antonio Frighetto via llvm-commits]

================
@@ -2830,37 +2831,48 @@ class AllocaSliceRewriter : public InstVisitor<AllocaSliceRewriter, bool> {
   ///   store <2 x float> %val2, ptr %alloca+16          ; offset 4-5
   ///   store <2 x float> %val1, ptr %alloca+8           ; offset 2-3
   ///   store <2 x float> %val3, ptr %alloca+24          ; offset 6-7
+  ///   %r = load <8 x float>, ptr %alloca
   ///
-  /// After:
-  ///   %alloca = alloca <8 x float>
-  ///   %shuffle0 = shufflevector %val0, %val1, <4 x i32> <i32 0, i32 1, i32 2,
-  ///                 i32 3>
-  ///   %shuffle1 = shufflevector %val2, %val3, <4 x i32> <i32 0, i32 1, i32 2,
-  ///                 i32 3>
-  ///   %shuffle2 = shufflevector %shuffle0, %shuffle1, <8 x i32> <i32 0, i32 1,
-  ///                 i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
-  ///   store %shuffle2, ptr %alloca
+  ///   After: tree of shufflevectors producing <8 x float> directly.
+  ///
+  /// Pattern 2 (init + RMW, possibly multi-round):
+  ///   A single full-width init store, followed by partial loads and
+  ///   partial stores that read-modify-write the alloca one or more
+  ///   times, optionally followed by a full-width load. The only
+  ///   structural requirement is that the distinct [begin, end) ranges
+  ///   touched by the partial loads and stores, taken together, tile
+  ///   the alloca disjointly.
+  ///
+  ///   We keep a map from each slice range to the SSA value that
+  ///   currently lives there, `SliceValues[r] -> Value*`:
+  ///     - initialize each entry to the corresponding piece of the
+  ///       init store's value (via a shufflevector picking the
+  ///       range's elements out of the init value),
+  ///     - walk partial loads and stores in block order,
+  ///     - for a partial load at range r: RAUW with `SliceValues[r]`,
+  ///     - for a partial store at range r: update `SliceValues[r]` to
+  ///       the stored value and drop the store.
+  ///   At the end, the final `SliceValues[r]` entries are tree-merged
+  ///   (in range order) into a single store to the alloca, and the
+  ///   optional full-width load is replaced by a load of the alloca.
   ///
-  /// The optimization looks for partitions that:
-  /// 1. Have no overlapping split slice tails
-  /// 2. Contain non-overlapping stores that cover the entire alloca
-  /// 3. Have exactly one load that reads the complete alloca structure and not
-  ///    in the middle of the stores (TODO: maybe we can relax the constraint
-  ///    about reading the entire alloca structure)
+  ///   Because the ranges are disjoint by construction, a store at one
+  ///   range cannot affect another range's tracked value, so a single
+  ///   block-order walk correctly tracks the memory state at each
+  ///   range. The algorithm handles multi-round RMW, partial loads
+  ///   and stores interleaved in any order, read-only slices (the
+  ///   tracked value stays at the init extract), and write-only
+  ///   slices (the tracked value never flows into a load).
   ///
   /// \param P The partition to analyze and potentially rewrite
-  /// \return An optional vector of values that were deleted during the rewrite
-  ///         process, or std::nullopt if the partition cannot be optimized
-  ///         using tree-structured merge
+  /// \return An optional vector of values that were deleted during the
+  ///         rewrite, or std::nullopt if the partition cannot be optimized.
   std::optional<SmallVector<Value *, 4>>
   rewriteTreeStructuredMerge(Partition &P) {
----------------
antoniofrighetto wrote:

I found a bit harder to go through this function. I would like to ask, as a follow up, a bit of a more principled decomposition, possibly reducing the amount of lambda this routine has. I believe that, having dedicate helper for collecting merge candidates, pattern validations, and building the shufflevectors tree / rewriting would improve readability.

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