mlir/Affine: slightly improve the affine scalrep pass

[rohany]

Summary

This commit improves the affine scalrep pass in two ways:

• Allows for some more dead memrefs to be eliminated from functions
• Uses a bit of aliasing information to forward more loads and stores between memrefs allocated within functions.

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

for (unsigned d = nsLoops + 1; d > minSurroundingLoops; d--) {
       // A dependence failure or the presence of a dependence implies a
       // side effect.
       if (!noDependence(result))
         return true;
     }
     // No side effect was seen.
     return false;
   }
+
+  // We don't have complete aliasing information yet (the TODO below),
+  // but we have some information we can work with around aliasing based
+  // on where the memrefs are coming from. If the two memrefs we are 
+  // considering aren't the same, and at least one of them was allocated
+  // within this scope, then the two memrefs can't alias, since allocation
+  // produces a fresh memory range.
+  if (srcAccess.memref != destAccess.memref) {
+    auto isLocallyAllocated = [](Value memref) {
+      auto *defOp = memref.getDefiningOp();
+      return defOp && hasSingleEffect<MemoryEffects::Allocate>(defOp, memref);
+    };
+    if (isLocallyAllocated(srcAccess.memref) || isLocallyAllocated(destAccess.memref)) {
+      return false;
+    }
+  }
+
   // TODO: Check here if the memrefs alias: there is no side effect if
   // `srcAccess.memref` and `destAccess.memref` don't alias.
   return true;
 }
 
 template <typename EffectType, typename T>
 bool mlir::hasNoInterveningEffect(Operation *start, T memOp) {
   auto isLocallyAllocated = [](Value memref) {
void mlir::affineScalarReplace(func::FuncOp f, DominanceInfo &domInfo,
   // 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();
 
+  // The above operation only adds stores that are at least loaded
+  // from once to memrefsToErase. Memrefs that only have store operations
+  // performed on them wouldn't be added to memrefsToErase to be considered
+  // for deletion below, so we add them here.
+  f.walk([&](AffineWriteOpInterface storeOp) {
+    memrefsToErase.insert(storeOp.getMemRef());
+  });
+
   // 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

func.func @cross_block() {
   %true = arith.constant true
   affine.store %true, %alloc_99[%c10] : memref<13xi1>
   cf.br ^bb1(%alloc_83 : memref<1x13xf32>)
 ^bb1(%35: memref<1x13xf32>):
   // CHECK: affine.load
   %69 = affine.load %alloc_99[%c10] : memref<13xi1>
   return
 }
+
+// Ensure that memrefs with just store operations get eliminated.
+// CHECK-LABEL : func @dead_stores() {
+func.func @dead_stores() {
+  %c1 = arith.constant 1.0 : f64
+  %alloc = memref.alloc() : memref<100xf64>
+  affine.for %i = 0 to 100 {
+    affine.store %c1, %alloc[%i] : memref<100xf64>
+  }
+  return
+  // CHECK: %{{.*}} = arith.constant
+  // CHECK-NEXT: affine.for %{{.*}} = 0 to 100 {
+  // CHECK-NEXT: }
+  // CHECK-NEXT: return
+}
+
+// Check that we can utilize aliasing information available from
+// local allocations against function parameters.
+
+// CHECK-LABEL : func @aliasing_local_params
+// CHECK-SAME ([[ARG_0_:%.+]]: memref<100xf64>) -> memref<100xf64> {
+func.func @aliasing_local_params(%arg0: memref<100xf64>) -> memref<100xf64> {
+  %c1 = arith.constant 1.0 : f64
+  %alloc = memref.alloc() : memref<100xf64>
+  affine.for %i = 0 to 100 {
+    %0 = affine.load %alloc[%i] : memref<100xf64>
+    affine.store %c1, %arg0[%i] : memref<100xf64>
+    %1 = affine.load %alloc[%i] : memref<100xf64>
+    %2 = arith.addf %1, %c1 : f64
+    affine.store %2, %alloc[%i] : memref<100xf64>
+  }
+  // TODO (rohany): I want to check in the affine.store below that we're storing to
+  //  the argument 0, but for some reason the filecheck test keeps telling me that
+  //  ARG_0_ is undefined...
+  return %alloc : memref<100xf64>
+  // CHECK: [[C1:%.+]] = arith.constant
+  // CHECK-NEXT: [[AL:%.+]] = memref.alloc
+  // CHECK-NEXT: affine.for [[I:%.+]] = 0 to 100 {
+  // CHECK-NEXT:   [[T0:%.+]] = affine.load [[AL]][[[I]]] : memref<100xf64>
+  // CHECK-NEXT:   affine.store
+  // CHECK-NEXT:   [[T2:%.+]] = arith.addf [[T0]], [[C1]] : f64
+  // CHECK-NEXT:   affine.store [[T2]], [[AL]][[[I]]] : memref<100xf64>
+  // CHECK-NEXT: }
+  // CHECK-NEXT: return [[AL]] : memref<100xf64>
+}