[MLIR] Make store to load fwd condition less conservative

[bondhugula]

Summary

Make store to load fwd condition for -memref-dataflow-opt less
conservative. Post dominance info is not really needed. Add additional
check for common cases.

[bondhugula]

Fix stale comments.

[vinayaka-polymage]

Looks great. Minor comments -

mlir/lib/Transforms/MemRefDataFlowOpt.cpp

44–45 ↗

(On Diff #351634)

Nit: op's => ops (unrelated to this change)

Should the comment change to
does not dominate all store => does not dominate any store ?

[bondhugula]

Fix comments.

[bondhugula]

Make the test case simpler/direct.

[ayzhuang]

Looks great! Minor comments:

mlir/lib/Transforms/MemRefDataFlowOpt.cpp

47 ↗

(On Diff #351651)

dominate any store op

267 ↗

(On Diff #351651)

dominate any op in depSrcStores

[vinayaka-polymage]

Looks good, saves PostDominanceInfo calculation.
Minor: comments at some places need to change to does not dominate any ... at a couple of places as pointed out by @ayzhuang .

[bondhugula]

Address reviewer comments. Rebase on tip.

[bondhugula]

Thank you for the reviews!

mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp

 // 1) store/load and load need to have mathematically equivalent affine access
 // functions (checked after full composition of load/store operands); this
 // implies that they access the same single memref element for all iterations of
 // the common surrounding loop,
 //
 // 2) the store/load op should dominate the load op,
 //
 // 3) among all op's that satisfy both (1) and (2), for store to load
-// forwarding, the one that postdominates all store op's that have a dependence
-// into the load, is provably the last writer to the particular memref location
-// being loaded at the load op, and its store value can be forwarded to the
-// load; for load CSE, any op that postdominates all store op's that have a
-// dependence into the load can be forwarded and the first one found is chosen.
-// Note that the only dependences that are to be considered are those that are
-// satisfied at the block* of the innermost common surrounding loop of the
-// <store/load, load> being considered.
+// forwarding, the one that does not dominate any store op that has a
+// dependence into the load, is provably the last writer to the particular
+// memref location being loaded at the load op, and its store value can be
+// forwarded to the load; for load CSE, any op that does not dominate any store
+// op that have a dependence into the load can be forwarded and the first one
+// found is chosen. Note that the only dependences that are to be considered are
+// those that are satisfied at the block* of the innermost common surrounding
+// loop of the <store/load, load> being considered.
 //
 // (* A dependence being satisfied at a block: a dependence that is satisfied by
 // virtue of the destination operation appearing textually / lexically after
 // the source operation within the body of a 'affine.for' operation; thus, a
 // dependence is always either satisfied by a loop or by a block).
 //
 // The above conditions are simple to check, sufficient, and powerful for most
 // cases in practice - they are sufficient, but not necessary --- since they
struct AffineScalarReplacement
     : public AffineScalarReplacementBase<AffineScalarReplacement> {
   void runOnFunction() override;
 
   LogicalResult forwardStoreToLoad(AffineReadOpInterface loadOp);
   void loadCSE(AffineReadOpInterface loadOp);
 
   // A list of memref's that are potentially dead / could be eliminated.
   SmallPtrSet<Value, 4> memrefsToErase;
-  // Load op's whose results were replaced by those forwarded from stores
+  // Load ops whose results were replaced by those forwarded from stores
   // dominating stores or loads..
   SmallVector<Operation *, 8> loadOpsToErase;
 
   DominanceInfo *domInfo = nullptr;
-  PostDominanceInfo *postDomInfo = nullptr;
 };
 
 } // end anonymous namespace
 
 /// Creates a pass to perform optimizations relying on memref dataflow such as
 /// store to load forwarding, elimination of dead stores, and dead allocs.
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createAffineScalarReplacementPass() {
AffineScalarReplacement::forwardStoreToLoad(AffineReadOpInterface loadOp) {
   }
 
   // The list of store op candidates for forwarding that satisfy conditions
   // (1) and (2) above - they will be filtered later when checking (3).
   SmallVector<Operation *, 8> fwdingCandidates;
 
   // Store ops that have a dependence into the load (even if they aren't
   // forwarding candidates). Each forwarding candidate will be checked for a
-  // post-dominance on these. 'fwdingCandidates' are a subset of depSrcStores.
+  // dominance on these. 'fwdingCandidates' are a subset of depSrcStores.
   SmallVector<Operation *, 8> depSrcStores;
 
   for (auto *storeOp : storeOps) {
     if (!storeMayReachLoad(storeOp, loadOp, minSurroundingLoops))
       continue;
 
     // Stores that *may* be reaching the load.
     depSrcStores.push_back(storeOp);
for (auto *storeOp : storeOps) {
     // 2. The store has to dominate the load op to be candidate.
     if (!domInfo->dominates(storeOp, loadOp))
       continue;
 
     // We now have a candidate for forwarding.
     fwdingCandidates.push_back(storeOp);
   }
 
-  // 3. Of all the store op's that meet the above criteria, the store that
-  // postdominates all 'depSrcStores' (if one exists) is the unique store
-  // providing the value to the load, i.e., provably the last writer to that
-  // memref loc.
-  // Note: this can be implemented in a cleaner way with postdominator tree
-  // traversals. Consider this for the future if needed.
+  // 3. Of all the store ops that meet the above criteria, the store op
+  // that does not dominate any of the ops in 'depSrcStores' (if such exists)
+  // will not have any of those latter ops on its paths to `loadOp`. It would
+  // thus be the unique store providing the value to the load. This condition is
+  // however conservative for eg:
+  //
+  // for ... {
+  //    store
+  //    load
+  //    store
+  //    load
+  // }
+  //
   Operation *lastWriteStoreOp = nullptr;
   for (auto *storeOp : fwdingCandidates) {
     if (llvm::all_of(depSrcStores, [&](Operation *depStore) {
-          return postDomInfo->postDominates(storeOp, depStore);
+          return !domInfo->properlyDominates(storeOp, depStore);
         })) {
       lastWriteStoreOp = storeOp;
       break;
     }
   }
   if (!lastWriteStoreOp)
     return failure();
 
AffineScalarReplacement::forwardStoreToLoad(AffineReadOpInterface loadOp) {
   return success();
 }
 
 // The load to load forwarding / redundant load elimination is similar to the
 // store to load forwarding.
 // loadA will be be replaced with loadB if:
 // 1) loadA and loadB have mathematically equivalent affine access functions.
 // 2) loadB dominates loadA.
-// 3) loadB postdominates all the store op's that have a dependence into loadA.
+// 3) loadB does not dominate any of the store ops that have a dependence into
+//    loadA.
 void AffineScalarReplacement::loadCSE(AffineReadOpInterface loadOp) {
   // The list of load op candidates for forwarding that satisfy conditions
   // (1) and (2) above - they will be filtered later when checking (3).
   SmallVector<Operation *, 8> fwdingCandidates;
   SmallVector<Operation *, 8> storeOps;
   unsigned minSurroundingLoops = getNestingDepth(loadOp);
   MemRefAccess memRefAccess(loadOp);
   // First pass over the use list to get 1) the minimum number of surrounding
for (auto *storeOp : storeOps) {
     if (!storeMayReachLoad(storeOp, loadOp, minSurroundingLoops))
       continue;
 
     // Stores that *may* be reaching the load.
     depSrcStores.push_back(storeOp);
   }
 
   // 3. Of all the load op's that meet the above criteria, return the first load
-  // found that postdominates all 'depSrcStores' and has the same shape as the
-  // load to be replaced (if one exists). The shape check is needed for affine
-  // vector loads.
+  // found that does not dominate any op in 'depSrcStores' and has the same
+  // shape as the load to be replaced (if one exists). The shape check is needed
+  // for affine vector loads.
   Operation *firstLoadOp = nullptr;
   Value oldVal = loadOp.getValue();
   for (auto *loadOp : fwdingCandidates) {
     if (llvm::all_of(depSrcStores,
                      [&](Operation *depStore) {
-                       return postDomInfo->postDominates(loadOp, depStore);
+                       return !domInfo->properlyDominates(loadOp, depStore);
                      }) &&
         cast<AffineReadOpInterface>(loadOp).getValue().getType() ==
             oldVal.getType()) {
       firstLoadOp = loadOp;
       break;
     }
   }
   if (!firstLoadOp)
void AffineScalarReplacement::runOnFunction() {
   // Only supports single block functions at the moment.
   FuncOp f = getFunction();
   if (!llvm::hasSingleElement(f)) {
     markAllAnalysesPreserved();
     return;
   }
 
   domInfo = &getAnalysis<DominanceInfo>();
-  postDomInfo = &getAnalysis<PostDominanceInfo>();
 
   loadOpsToErase.clear();
   memrefsToErase.clear();
 
   // Walk all load's and perform store to load forwarding and loadCSE.
   f.walk([&](AffineReadOpInterface loadOp) {
     // Do store to load forwarding first, if no success, try loadCSE.
     if (failed(forwardStoreToLoad(loadOp)))

mlir/test/Dialect/Affine/scalrep.mlir

affine.for %i = 0 to 16 {
     // CHECK:       affine.vector_load
     %ld1 = affine.vector_load %in[32*%i] : memref<512xf32>, vector<32xf32>
     %add = addf %ld0, %ld1 : vector<32xf32>
     affine.vector_store %ld1, %out[32*%i] : memref<512xf32>, vector<32xf32>
   }
   return
 }
 
+// CHECK-LABEL: func @reduction_multi_store
+func @reduction_multi_store() -> memref<1xf32> {
+  %A = memref.alloc() : memref<1xf32>
+  %cf0 = constant 0.0 : f32
+  %cf5 = constant 5.0 : f32
+
+ affine.store %cf0, %A[0] : memref<1xf32>
+  affine.for %i = 0 to 100 step 2 {
+    %l = affine.load %A[0] : memref<1xf32>
+    %s = addf %l, %cf5 : f32
+    // Store to load forwarding from this store should happen.
+    affine.store %s, %A[0] : memref<1xf32>
+    %m = affine.load %A[0] : memref<1xf32>
+   "test.foo"(%m) : (f32) -> ()
+  }
+
+// CHECK:       affine.for
+// CHECK:         affine.load
+// CHECK:         affine.store %[[S:.*]],
+// CHECK-NEXT:    "test.foo"(%[[S]])
+
+  return %A : memref<1xf32>
+}
+
 // CHECK-LABEL: func @vector_load_affine_apply_store_load
 func @vector_load_affine_apply_store_load(%in : memref<512xf32>, %out : memref<512xf32>) {
   %cf1 = constant 1: index
   affine.for %i = 0 to 15 {
     // CHECK:       affine.vector_load
     %ld0 = affine.vector_load %in[32*%i] : memref<512xf32>, vector<32xf32>
     %idx = affine.apply affine_map<(d0) -> (d0 + 1)> (%i)
     affine.vector_store %ld0, %in[32*%idx] : memref<512xf32>, vector<32xf32>
     // CHECK-NOT:   affine.vector_load
     %ld1 = affine.vector_load %in[32*%i] : memref<512xf32>, vector<32xf32>
     %add = addf %ld0, %ld1 : vector<32xf32>
     affine.vector_store %ld1, %out[32*%i] : memref<512xf32>, vector<32xf32>
   }
   return
 }