[MLIR][Affine] Improve load elimination

[rikhuijzer]

Summary

Common Subexpression Eliminiation (CSE) was applied via loadCSE to remove redundant loads from code. However, this process failed in cases where there were redundant stores in between the redundant loads. In those cases, the redundant stores would be seen as having an intervening effect during the loadCSE step.

This patch fixes that by moving the load CSE behind the store CSE.

Fixes #62639.

[bondhugula]

Thanks for improving this. Please add a commit summary -- along the lines of your code comment at L1112.

mlir/lib/Dialect/Affine/Utils/Utils.cpp

1003

loadB should not dominate loadA

1114

as intervening -> as having an intervening

[rikhuijzer]

Update comment

[rikhuijzer]

Make comments inside loadCSE more clear

[rikhuijzer]

In D154769#4483279, @bondhugula wrote:

Thanks for improving this. Please add a commit summary -- along the lines of your code comment at L1112.

I've updated the patch description in order to let that automatically update the commit summary. However, I just saw after committing that the commit message wasn't automatically updated. Apologies for that. I think it's not worth fixing this anymore now. I'll pay more attention next time!

[bondhugula]

In D154769#4483362, @rikhuijzer wrote:

In D154769#4483279, @bondhugula wrote:

Thanks for improving this. Please add a commit summary -- along the lines of your code comment at L1112.

I've updated the patch description in order to let that automatically update the commit summary. However, I just saw after committing that the commit message wasn't automatically updated. Apologies for that. I think it's not worth fixing this anymore now. I'll pay more attention next time!

No problem. In the future, you can do an arc diff --edit --verbatim -- that would update the commit summary on the commit as well as Phab here in one shot.

[rikhuijzer]

In D154769#4489440, @bondhugula wrote:

In D154769#4483362, @rikhuijzer wrote:

In D154769#4483279, @bondhugula wrote:

Thanks for improving this. Please add a commit summary -- along the lines of your code comment at L1112.

I've updated the patch description in order to let that automatically update the commit summary. However, I just saw after committing that the commit message wasn't automatically updated. Apologies for that. I think it's not worth fixing this anymore now. I'll pay more attention next time!

No problem. In the future, you can do an arc diff --edit --verbatim -- that would update the commit summary on the commit as well as Phab here in one shot.

Yes, I was already wondering about that. Thank you

mlir/lib/Dialect/Affine/Utils/Utils.cpp

 
 /// Attempt to eliminate loadOp by replacing it with a value stored into memory
 /// which the load is guaranteed to retrieve. This check involves three
 /// components: 1) The store and load must be on the same location 2) The store
 /// must dominate (and therefore must always occur prior to) the load 3) No
 /// other operations will overwrite the memory loaded between the given load
 /// and store.  If such a value exists, the replaced `loadOp` will be added to
 /// `loadOpsToErase` and its memref will be added to `memrefsToErase`.
-static LogicalResult forwardStoreToLoad(
-    AffineReadOpInterface loadOp, SmallVectorImpl<Operation *> &loadOpsToErase,
-    SmallPtrSetImpl<Value> &memrefsToErase, DominanceInfo &domInfo) {
+static void forwardStoreToLoad(AffineReadOpInterface loadOp,
+                               SmallVectorImpl<Operation *> &loadOpsToErase,
+                               SmallPtrSetImpl<Value> &memrefsToErase,
+                               DominanceInfo &domInfo) {
 
   // The store op candidate for forwarding that satisfies all conditions
   // to replace the load, if any.
   Operation *lastWriteStoreOp = nullptr;
 
   for (auto *user : loadOp.getMemRef().getUsers()) {
     auto storeOp = dyn_cast<AffineWriteOpInterface>(user);
     if (!storeOp)
for (auto *user : loadOp.getMemRef().getUsers()) {
 
     // We now have a candidate for forwarding.
     assert(lastWriteStoreOp == nullptr &&
            "multiple simultaneous replacement stores");
     lastWriteStoreOp = storeOp;
   }
 
   if (!lastWriteStoreOp)
-    return failure();
+    return;
 
   // Perform the actual store to load forwarding.
   Value storeVal =
       cast<AffineWriteOpInterface>(lastWriteStoreOp).getValueToStore();
   // Check if 2 values have the same shape. This is needed for affine vector
   // loads and stores.
   if (storeVal.getType() != loadOp.getValue().getType())
-    return failure();
+    return;
   loadOp.getValue().replaceAllUsesWith(storeVal);
   // Record the memref for a later sweep to optimize away.
   memrefsToErase.insert(loadOp.getMemRef());
   // Record this to erase later.
   loadOpsToErase.push_back(loadOp);
-  return success();
 }
 
 template bool
 mlir::affine::hasNoInterveningEffect<mlir::MemoryEffects::Read,
                                      affine::AffineReadOpInterface>(
     mlir::Operation *, affine::AffineReadOpInterface);
 
 // This attempts to find stores which have no impact on the final result.
static void loadCSE(AffineReadOpInterface loadA,
   for (auto *user : loadA.getMemRef().getUsers()) {
     auto loadB = dyn_cast<AffineReadOpInterface>(user);
     if (!loadB || loadB == loadA)
       continue;
 
     MemRefAccess srcAccess(loadB);
     MemRefAccess destAccess(loadA);
 
-    // 1. The accesses have to be to the same location.
+    // 1. The accesses should be to be to the same location.
     if (srcAccess != destAccess) {
       continue;
     }
 
-    // 2. The store has to dominate the load op to be candidate.
+    // 2. loadB should dominate loadA.
     if (!domInfo.dominates(loadB, loadA))
       continue;
 
-    // 3. There is no write between loadA and loadB.
+    // 3. There should not be a write between loadA and loadB.
     if (!affine::hasNoInterveningEffect<MemoryEffects::Write>(
             loadB.getOperation(), loadA))
       continue;
 
     // Check if two values have the same shape. This is needed for affine vector
     // loads.
     if (loadB.getValue().getType() != loadA.getValue().getType())
       continue;
void mlir::affine::affineScalarReplace(func::FuncOp f, DominanceInfo &domInfo,
   // Load op's whose results were replaced by those forwarded from stores.
   SmallVector<Operation *, 8> opsToErase;
 
   // A list of memref's that are potentially dead / could be eliminated.
   SmallPtrSet<Value, 4> memrefsToErase;
 
   // Walk all load's and perform store to load forwarding.
   f.walk([&](AffineReadOpInterface loadOp) {
-    if (failed(
-            forwardStoreToLoad(loadOp, opsToErase, memrefsToErase, domInfo))) {
-      loadCSE(loadOp, opsToErase, domInfo);
-    }
+    forwardStoreToLoad(loadOp, opsToErase, memrefsToErase, domInfo);
   });
-
-  // Erase all load op's whose results were replaced with store fwd'ed ones.
   for (auto *op : opsToErase)
     op->erase();
   opsToErase.clear();
 
   // Walk all store's and perform unused store elimination
   f.walk([&](AffineWriteOpInterface storeOp) {
     findUnusedStore(storeOp, opsToErase, postDomInfo);
   });
-  // Erase all store op's which don't impact the program
   for (auto *op : opsToErase)
     op->erase();
+  opsToErase.clear();
 
   // Check if the store fwd'ed memrefs are now left with only stores and
   // deallocs and can thus be completely deleted. Note: the canonicalize pass
   // should be able to do this as well, but we'll do it here since we collected
   // these anyway.
   for (auto memref : memrefsToErase) {
     // If the memref hasn't been locally alloc'ed, skip.
     Operation *defOp = memref.getDefiningOp();
     if (!defOp || !hasSingleEffect<MemoryEffects::Allocate>(defOp, memref))
       // TODO: if the memref was returned by a 'call' operation, we
       // could still erase it if the call had no side-effects.
       continue;
     if (llvm::any_of(memref.getUsers(), [&](Operation *ownerOp) {
           return !isa<AffineWriteOpInterface>(ownerOp) &&
                  !hasSingleEffect<MemoryEffects::Free>(ownerOp, memref);
         }))
       continue;
 
     // Erase all stores, the dealloc, and the alloc on the memref.
     for (auto *user : llvm::make_early_inc_range(memref.getUsers()))
       user->erase();
     defOp->erase();
   }
+
+  // To eliminate as many loads as possible, run load CSE after eliminating
+  // stores. Otherwise, some stores are wrongly seen as having an intervening
+  // effect.
+  f.walk([&](AffineReadOpInterface loadOp) {
+    loadCSE(loadOp, opsToErase, domInfo);
+  });
+  for (auto *op : opsToErase)
+    op->erase();
 }
 
 // Perform the replacement in `op`.
 LogicalResult mlir::affine::replaceAllMemRefUsesWith(
     Value oldMemRef, Value newMemRef, Operation *op,
     ArrayRef<Value> extraIndices, AffineMap indexRemap,
     ArrayRef<Value> extraOperands, ArrayRef<Value> symbolOperands,
     bool allowNonDereferencingOps) {

mlir/test/Dialect/Affine/scalrep.mlir

// CHECK: affine.for %[[J:.*]] =
       %w = affine.load %A[%j, %i] : memref<100x100xf32>
       // CHECK-NEXT: "foo"
       "foo" (%u, %v, %w) : (f32, f32, f32) -> ()
     }
   }
   return
 }
 
+// CHECK-LABEL: func @elim_load_after_store
+func.func @elim_load_after_store(%arg0: memref<100xf32>, %arg1: memref<100xf32>) {
+  %alloc = memref.alloc() : memref<1xf32>
+  %alloc_0 = memref.alloc() : memref<1xf32>
+  // CHECK: affine.for
+  affine.for %arg2 = 0 to 100 {
+    // CHECK: affine.load
+    %0 = affine.load %arg0[%arg2] : memref<100xf32>
+    %1 = affine.load %arg0[%arg2] : memref<100xf32>
+    // CHECK: arith.addf
+    %2 = arith.addf %0, %1 : f32
+    affine.store %2, %alloc_0[0] : memref<1xf32>
+    %3 = affine.load %arg0[%arg2] : memref<100xf32>
+    %4 = affine.load %alloc_0[0] : memref<1xf32>
+    // CHECK-NEXT: arith.addf
+    %5 = arith.addf %3, %4 : f32
+    affine.store %5, %alloc[0] : memref<1xf32>
+    %6 = affine.load %arg0[%arg2] : memref<100xf32>
+    %7 = affine.load %alloc[0] : memref<1xf32>
+    %8 = arith.addf %6, %7 : f32
+    affine.store %8, %arg1[%arg2] : memref<100xf32>
+  }
+  return
+}
+
 // The test checks for value forwarding from vector stores to vector loads.
 // The value loaded from %in can directly be stored to %out by eliminating
 // store and load from %tmp.
 func.func @vector_forwarding(%in : memref<512xf32>, %out : memref<512xf32>) {
   %tmp = memref.alloc() : memref<512xf32>
   affine.for %i = 0 to 16 {
     %ld0 = affine.vector_load %in[32*%i] : memref<512xf32>, vector<32xf32>
     affine.vector_store %ld0, %tmp[32*%i] : memref<512xf32>, vector<32xf32>