[mlir][affine] Fix violation for uniqueness of value in scalrep pass.

[Lewuathe]

Summary

Scalrep pass expects the values in the given affine map are unique. But this expectation is not satisfied in case of the failure to eliminate redundant store operations before forwarding store to load. We should be able to eliminate all unused store operations before forwarding store to load in the scalrep pass to make sure the uniqueness of the value. This is the case reported in https://github.com/llvm/llvm-project/issues/59461.

[bondhugula]

The crash at https://github.com/llvm/llvm-project/issues/59461 is a bug that's just hidden with this change. Unused store elimination is currently done after store-to-load forward because that way you have fewer stores to look at (this should significantly speed up things on large bodies). In any case, this doesn't fix the bug that caused the assertion in https://github.com/llvm/llvm-project/issues/59461 - the input to store-to-load forwarding there is still a valid one with those two stores.

[bondhugula]

As described here https://github.com/llvm/llvm-project/issues/59461#issuecomment-1372924955, this isn't the right fix; it only hides the issue. I should be able to post revisions to fix this later today.

[bondhugula]

This fixes the assertion from addAffineIfOpDomain: https://reviews.llvm.org/D141250
And https://reviews.llvm.org/D141255 fixes the scalrep crash exposed and produces the desired simplification.

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

 void mlir::affineScalarReplace(func::FuncOp f, DominanceInfo &domInfo,
                                PostDominanceInfo &postDomInfo) {
   // 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 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();
+
   // 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, 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
   // 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))

mlir/test/Dialect/Affine/scalrep.mlir

 // CHECK:       %[[C7:.*]] = arith.constant 7.000000e+00 : f32
 // CHECK-NEXT:  affine.parallel (%{{.*}}) = (0) to (10) {
 // CHECK-NEXT:    affine.for %{{.*}} = 0 to 10 {
 // CHECK-NEXT:      arith.addf %[[C7]], %[[C7]] : f32
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // CHECK-NEXT:  return
 }
+
+// CHECK-LABEL: func @redundant_store_op() -> memref<15xi1> {
+func.func @redundant_store_op() -> memref<15xi1> {
+  %c1 = arith.constant 1 : index
+  %alloc = memref.alloc() : memref<15xi1>
+  %true = arith.constant true
+  affine.store %true, %alloc[%c1] : memref<15xi1>
+  affine.store %true, %alloc[%c1] : memref<15xi1>
+  %c2 = arith.constant 2 : index
+  affine.if affine_set<(d0, d1) : ((d0 + 1) mod 8 >= 0, d0 * -8 >= 0)>(%c1, %c1){
+    %load1 = affine.load %alloc[%c1] : memref<15xi1>
+  }
+  %m1 = memref.alloc() : memref<15xi1>
+  return %m1 : memref<15xi1>
+// CHECK:      %[[C1:.*]] = arith.constant 1 : index
+// CHECK-NEXT: %true = arith.constant true
+// CHECK-NEXT: %[[C2:.*]] = arith.constant 2 : index
+// CHECK-NEXT: affine.if #set(%[[C1]], %[[C1]]) {
+// CHECK-NEXT: }
+// CHECK-NEXT: %[[MEM:.*]] = memref.alloc() : memref<15xi1>
+// CHECK-NEXT: return %[[MEM]] : memref<15xi1>
+}