[MLIR] Eliminate unnecessary affine stores

[wsmoses]

Summary

Deduce circumstances where an affine store could not possibly be read by an operation (such as an affine load), and if so, eliminate the store

[wsmoses]

Fix rebase errors

[mehdi_amini]

Typo in the description: Deduce circumstances where an affine load ... ; this last word was meant to be store I think. Looks good to me in general!

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

72

The name is slightly misleading, as this isn't "removing" anything, what about findUnusedStore instead?

274

[wsmoses]

Address comments

[ftynse]

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

267

351–353

Rebase gone wrong?

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

struct AffineScalarReplacement
   LogicalResult forwardStoreToLoad(AffineReadOpInterface loadOp,
                                    SmallVectorImpl<Operation *> &loadOpsToErase,
                                    SmallPtrSetImpl<Value> &memrefsToErase,
                                    DominanceInfo &domInfo);
 
   void loadCSE(AffineReadOpInterface loadOp,
                SmallVectorImpl<Operation *> &loadOpsToErase,
                DominanceInfo &domInfo);
+
+  void findUnusedStore(AffineWriteOpInterface storeOp,
+                       SmallVectorImpl<Operation *> &storeOpsToErase,
+                       SmallPtrSetImpl<Value> &memrefsToErase,
+                       PostDominanceInfo &postDominanceInfo);
 };
 
 } // 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() {
std::function<void(Operation *, Operation *)> recur =
                 todoBlocks.push_back(succ);
           }
         }
       };
   recur(start, memOp);
   return !hasSideEffect;
 }
 
+// This attempts to find stores which have no impact on the final result.
+// A writing op writeA will be eliminated if there exists an op writeB if
+// 1) writeA and writeB have mathematically equivalent affine access functions.
+// 2) writeB postdominates writeA.
+// 3) There is no potential read between writeA and writeB.
+void AffineScalarReplacement::findUnusedStore(
+    AffineWriteOpInterface writeA, SmallVectorImpl<Operation *> &opsToErase,
+    SmallPtrSetImpl<Value> &memrefsToErase,
+    PostDominanceInfo &postDominanceInfo) {
+
+  for (Operation *user : writeA.getMemRef().getUsers()) {
+    // Only consider writing operations.
+    auto writeB = dyn_cast<AffineWriteOpInterface>(user);
+    if (!writeB)
+      continue;
+
+    // The operations must be distinct.
+    if (writeB == writeA)
+      continue;
+
+    // Both operations must lie in the same region.
+    if (writeB->getParentRegion() != writeA->getParentRegion())
+      continue;
+
+    // Both operations must write to the same memory.
+    MemRefAccess srcAccess(writeB);
+    MemRefAccess destAccess(writeA);
+
+    if (srcAccess != destAccess)
+      continue;
+
+    // writeB must postdominate writeA.
+    if (!postDominanceInfo.postDominates(writeB, writeA))
+      continue;
+
+    // There cannot be an operation which reads from memory between
+    // the two writes.
+    if (!hasNoInterveningEffect<MemoryEffects::Read>(writeA, writeB))
+      continue;
+
+    opsToErase.push_back(writeA);
+    break;
+  }
+}
+
 /// 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`.
 LogicalResult AffineScalarReplacement::forwardStoreToLoad(
for (auto *user : loadOp.getMemRef().getUsers()) {
     if (!domInfo.dominates(storeOp, loadOp))
       continue;
 
     // 3. Ensure there is no intermediate operation which could replace the
     // value in memory.
     if (!hasNoInterveningEffect<MemoryEffects::Write>(storeOp, loadOp))
       continue;
 
     // We now have a candidate for forwarding.
     assert(lastWriteStoreOp == nullptr &&
            "multiple simulataneous replacement stores");
     lastWriteStoreOp = storeOp;
   }
 
   if (!lastWriteStoreOp)
     return failure();
 
   // Perform the actual store to load forwarding.
   Value storeVal =
void AffineScalarReplacement::runOnFunction() {
 
   // 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;
 
   auto &domInfo = getAnalysis<DominanceInfo>();
+  auto &postDomInfo = getAnalysis<PostDominanceInfo>();
 
   // 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);
     }
   });
 
   // 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, memrefsToErase, postDomInfo);
+  });
+  // Erase all store op's which don't impact the program
+  for (auto *op : opsToErase)
+    op->erase();
 
   // Check if the store fwd'ed memrefs are now left with only stores 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 alloc'ed in this function, skip.
     Operation *defOp = memref.getDefiningOp();
     if (!defOp || !isa<memref::AllocOp>(defOp))

mlir/test/Dialect/Affine/scalrep.mlir

 // CHECK-NEXT:      %{{.*}} = addf %{{.*}}, %{{.*}} : f32
 // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // CHECK-NEXT:  %{{.*}} = affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
 // CHECK-NEXT:  return %{{.*}} : f32
 }
 
+// CHECK-LABEL: func @redundant_store_elim
+
+func @redundant_store_elim(%out : memref<512xf32>) {
+  %cf1 = constant 1.0 : f32
+  %cf2 = constant 2.0 : f32
+  affine.for %i = 0 to 16 {
+    affine.store %cf1, %out[32*%i] : memref<512xf32>
+    affine.store %cf2, %out[32*%i] : memref<512xf32>
+  }
+  return
+}
+
+// CHECK: affine.for
+// CHECK-NEXT:   affine.store
+// CHECK-NEXT: }
+
+// CHECK-LABEL: func @redundant_store_elim_fail
+
+func @redundant_store_elim_fail(%out : memref<512xf32>) {
+  %cf1 = constant 1.0 : f32
+  %cf2 = constant 2.0 : f32
+  affine.for %i = 0 to 16 {
+    affine.store %cf1, %out[32*%i] : memref<512xf32>
+    "test.use"(%out) : (memref<512xf32>) -> ()
+    affine.store %cf2, %out[32*%i] : memref<512xf32>
+  }
+  return
+}
+
+// CHECK: affine.for
+// CHECK-NEXT:   affine.store
+// CHECK-NEXT:   "test.use"
+// CHECK-NEXT:   affine.store
+// CHECK-NEXT: }
+