diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp --- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp +++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp @@ -695,6 +695,23 @@ // 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(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; @@ -1084,6 +1101,14 @@ 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 diff --git a/mlir/test/Dialect/Affine/scalrep.mlir b/mlir/test/Dialect/Affine/scalrep.mlir --- a/mlir/test/Dialect/Affine/scalrep.mlir +++ b/mlir/test/Dialect/Affine/scalrep.mlir @@ -888,3 +888,48 @@ %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> +}