[MLIR] Strengthen affine scalar replacement's findUnusedStore

[bondhugula]

Summary

Strengthen affine scalar replacement's unused store elimination by
checking for non-aliasing memrefs for scalar replacement's "intervening
ops" check.

[bondhugula]

Improve comments.

[bondhugula]

Rebase.

[wsmoses]

mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp

18

This feels like it shares a common goal with https://mlir.llvm.org/doxygen/LocalAliasAnalysis_8cpp_source.html#l00319 and similar, is there a reason why that infra cannot be used here?

30

For my own edification, is this an assumption we can make? I agree if the individual load or store is to a different address space, or the allocation itself, then this applies. But this could check between intermediate values.

I don't think this exists in MLIR (at least upstream to my knowledge) but suppose there were a addrspacecast(addrspacecast(a : memref<f32,0>, 3), 0).

65

It would be good to also check if they both are getglobal and come from different symbol's

[bondhugula]

mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp

18

I should definitely explore using that far more general framework. This one is light-weight and covers more common/lower hanging cases (for eg. LocalAliasAnalysis doesn't know about global memrefs -- but we can just extend LocalAliasAnalysis). I couldn't find a single user of LocalAliasAnalysis in the codebase. Are you aware of something that I might have missed?

30

There isn't an address space cast abstraction on memref's and I don't think one should be added -- it doesn't make sense, destroys higher-level abstractions, and only gets in the way of analysis and optimization. (For low-level LLVM things including pointers, I noticed addrcast ops exist.)

I actually already raised the following PR putting forward such a constraint on the memref type's doc: https://reviews.llvm.org/D121270
So, we should be relying on the design here if there is a consensus on D121270.

65

Sure, good point.

mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h

-//===- MemRefUtils.h - MemRef transformation utilities ----------*- C++ -*-===//
+//===- MemRefUtils.h - MemRef analysis/transformation utilities -*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
-// This header file defines prototypes for various transformation utilities for
-// the MemRefOps dialect. These are not passes by themselves but are used
-// either by passes, optimization sequences, or in turn by other transformation
-// utilities.
+// This header file defines prototypes for various analysis and transformation
+// utilities for the MemRefOps dialect and memref types. These are not passes by
+// themselves but are used either by passes, optimization sequences, or in turn
+// by other transformation utilities.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_MEMREF_UTILS_MEMREFUTILS_H
 #define MLIR_DIALECT_MEMREF_UTILS_MEMREFUTILS_H
 
+namespace mlir {
+
+class AliasResult;
+class Value;
+
+namespace memref {
+
+/// Returns the result on whether `memrefA` and `memrefB` alias.
+//  TODO: This should be replaced with a proper alias analysis framework when it
+//  is available.
+AliasResult memrefsAlias(Value memrefA, Value memrefB);
+
+} // namespace memref
+
+} // namespace mlir
+
 #endif // MLIR_DIALECT_MEMREF_UTILS_MEMREFUTILS_H

mlir/lib/Dialect/Affine/Utils/CMakeLists.txt

 add_mlir_dialect_library(MLIRAffineUtils
   LoopFusionUtils.cpp
   LoopUtils.cpp
   Utils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Affine
 
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRAffineAnalysis
   MLIRAnalysis
   MLIRMemRef
+  MLIRMemRefUtils
   MLIRTransformUtils
   )

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

 //===- Utils.cpp ---- Utilities for affine dialect transformation ---------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements miscellaneous transformation utilities for the Affine
 // dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/Utils.h"
 
+#include "mlir/Analysis/AliasAnalysis.h"
 #include "mlir/Dialect/Affine/Analysis/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Dialect/Affine/LoopUtils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 #define DEBUG_TYPE "affine-utils"
 
if (auto memEffect = dyn_cast<MemoryEffectOpInterface>(op)) {
               hasSideEffect = true;
               return;
             }
           }
 
           // No side effect was seen, simply return.
           return;
         }
+        // There is no side effect if the memrefs do not alias.
+        if (memref::memrefsAlias(srcAccess.memref, destAccess.memref) ==
+            AliasResult::NoAlias)
+          return;
       }
+
+      // We have an op with a memory effect and we cannot prove if it
+      // intervenes.
       hasSideEffect = true;
       return;
     }
 
     if (op->hasTrait<OpTrait::HasRecursiveSideEffects>()) {
       // Recurse into the regions for this op and check whether the internal
       // operations may have the side effect `EffectType` on memOp.
       for (Region &region : op->getRegions())

mlir/lib/Dialect/MemRef/Utils/CMakeLists.txt

 add_mlir_dialect_library(MLIRMemRefUtils
   MemRefUtils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/inlude/mlir/Dialect/MemRefDialect
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRMemRef
+  MLIRViewLikeInterface
 )
 

mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp

 //===- MemRefUtils.cpp - Utilities to support the MemRef dialect ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
-// This file implements utilities for the MemRef dialect.
+// This file implements utilities to support the MemRef op dialect.
 //
 //===----------------------------------------------------------------------===//
-
 #include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
+
+#include "mlir/Analysis/AliasAnalysis.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+
+using namespace mlir;
+
+AliasResult mlir::memref::memrefsAlias(Value memrefA, Value memrefB) {
+  auto memrefAType = memrefA.getType().dyn_cast<BaseMemRefType>();
+  auto memrefBType = memrefB.getType().dyn_cast<BaseMemRefType>();
+  assert(
+      memrefAType && memrefBType &&
+      "memrefsAlias method is to be used only for ranked or unranked memrefs");
+
+  if (memrefA == memrefB)
+    return AliasResult::MustAlias;
+
+  // Memrefs in two different memory spaces cannot alias.
+  if (memrefAType.getMemorySpace() != memrefBType.getMemorySpace())
+    return AliasResult::NoAlias;
+
+  Operation *aDefOp = memrefA.getDefiningOp();
+  Operation *bDefOp = memrefB.getDefiningOp();
+
+  auto isAllocLikeOp = [](Operation *op) {
+    auto memOp = dyn_cast<MemoryEffectOpInterface>(op);
+    return memOp && memOp.hasEffect<MemoryEffects::Allocate>();
+  };
+
+  // If one is a locally alloc'ed and the other is a block argument, they
+  // don't alias. If both are locally alloc'ed, they can't alias.
+  if (aDefOp && isAllocLikeOp(aDefOp) && (!bDefOp || isAllocLikeOp(bDefOp)))
+    return AliasResult::NoAlias;
+  if (bDefOp && isAllocLikeOp(bDefOp) && !aDefOp)
+    return AliasResult::NoAlias;
+
+  // Both are block arguments; we do not know if they alias.
+  if (!aDefOp && !bDefOp)
+    return AliasResult::MayAlias;
+
+  // If they are aliases of another, use the original memref. Recursion here is
+  // bounded by SSA use-def list length.
+  if (auto viewA = dyn_cast_or_null<ViewLikeOpInterface>(aDefOp))
+    return memrefsAlias(viewA.getViewSource(), memrefB);
+  if (auto viewB = dyn_cast_or_null<ViewLikeOpInterface>(bDefOp))
+    return memrefsAlias(memrefA, viewB.getViewSource());
+
+  // If one is a global memref and the other is alloc-like, they can't
+  // alias.
+  if ((aDefOp && isAllocLikeOp(aDefOp) &&
+       isa_and_nonnull<memref::GetGlobalOp>(bDefOp)) ||
+      (bDefOp && isAllocLikeOp(bDefOp) &&
+       isa_and_nonnull<memref::GetGlobalOp>(aDefOp)))
+    return AliasResult::NoAlias;
+
+  // We do not know if they alias; return false conservatively.
+  return AliasResult::MayAlias;
+}

mlir/test/Dialect/Affine/scalrep.mlir

func @redundant_store_elim(%out : 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 = arith.constant 1.0 : f32
   %cf2 = arith.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>
   }

 }
 
 // CHECK:  %[[pi:.+]] = arith.constant 3.140000e+00 : f64
 // CHECK:  %{{.*}} = scf.if %arg2 -> (f64) {
 // CHECK:        scf.yield %{{.*}} : f64
 // CHECK:      } else {
 // CHECK:        scf.yield %[[pi]] : f64
 // CHECK:      }
+
+// CHECK-LABEL: func @dead_stores_intervening_non_aliasing
+func @dead_stores_intervening_non_aliasing(%arg0: memref<3x4096x4096xf32>, %buf: memref<4096x4092xf32, 3>) {
+  %cst = arith.constant 3.750000e-01 : f32
+  %cst_0 = arith.constant 2.500000e-01 : f32
+  %cst_1 = arith.constant 6.250000e-02 : f32
+  %in = memref.expand_shape %arg0 [[0], [1], [2, 3]] : memref<3x4096x4096xf32> into memref<3x4096x4096x1xf32>
+  %out = memref.alloc() : memref<3x4096x4092x1xf32>
+  // All stores to %out except the last one will be dead post forwarding.
+  affine.for %arg1 = 0 to 4096 {
+    affine.for %arg2 = 0 to 4092 {
+      %2 = affine.load %in[0, %arg1, %arg2, 0] : memref<3x4096x4096x1xf32>
+      %3 = affine.load %out[0, %arg1, %arg2, 0] : memref<3x4096x4092x1xf32>
+      %4 = arith.mulf %2, %cst_1 : f32
+      %5 = arith.addf %3, %4 : f32
+      affine.store %5, %out[0, %arg1, %arg2, 0] : memref<3x4096x4092x1xf32>
+      %6 = affine.load %in[0, %arg1, %arg2 + 1, 0] : memref<3x4096x4096x1xf32>
+      affine.load %buf[%arg1, %arg2] : memref<4096x4092xf32, 3>
+      %7 = affine.load %out[0, %arg1, %arg2, 0] : memref<3x4096x4092x1xf32>
+      %8 = arith.mulf %6, %cst_0 : f32
+      %9 = arith.addf %7, %8 : f32
+      affine.store %9, %out[0, %arg1, %arg2, 0] : memref<3x4096x4092x1xf32>
+      %10 = affine.load %in[0, %arg1, %arg2 + 2, 0] : memref<3x4096x4096x1xf32>
+      %11 = affine.load %out[0, %arg1, %arg2, 0] : memref<3x4096x4092x1xf32>
+      %12 = arith.mulf %10, %cst : f32
+      %13 = arith.addf %11, %12 : f32
+      affine.store %13, %out[0, %arg1, %arg2, 0] : memref<3x4096x4092x1xf32>
+      %14 = affine.load %in[0, %arg1, %arg2 + 3, 0] : memref<3x4096x4096x1xf32>
+      %15 = affine.load %out[0, %arg1, %arg2, 0] : memref<3x4096x4092x1xf32>
+      %16 = arith.mulf %14, %cst_0 : f32
+      %17 = arith.addf %15, %16 : f32
+      affine.store %17, %out[0, %arg1, %arg2, 0] : memref<3x4096x4092x1xf32>
+      %18 = affine.load %in[0, %arg1, %arg2 + 4, 0] : memref<3x4096x4096x1xf32>
+      %19 = affine.load %out[0, %arg1, %arg2, 0] : memref<3x4096x4092x1xf32>
+      %20 = arith.mulf %18, %cst_1 : f32
+      %21 = arith.addf %19, %20 : f32
+      affine.store %21, %out[0, %arg1, %arg2, 0] : memref<3x4096x4092x1xf32>
+      // All but the last affine.store is eliminated. This relies on being able
+      // to determine that %out does not alias with %in and %buf.
+      // CHECK:     affine.store
+      // CHECK-NOT: affine.store
+    }
+  }
+  // CHECK: return
+  return
+}