[MLIR][MemRef] Only allow fold of cast for the pointer operand, not the value

[wsmoses]

Summary

Currently canonicalizations of a store and a cast try to fold all casts into the store.

In the case where the operand being stored is itself a cast, this is illegal as the type of the value being stored
will change. This PR fixes this by not checking the value for folding with a cast.

Depends on https://reviews.llvm.org/D103828

[mehdi_amini]

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

948–950

Could also use make_filter_range instead:

for (OpOperand &operand : make_filter_range(op->getOpOperands(),
                    [] (OpOperand &operand) { return operand.get() != ignore; })) {

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

 
 //===----------------------------------------------------------------------===//
 // Common canonicalization pattern support logic
 //===----------------------------------------------------------------------===//
 
 /// This is a common class used for patterns of the form
 /// "someop(memrefcast) -> someop".  It folds the source of any memref.cast
 /// into the root operation directly.
-static LogicalResult foldMemRefCast(Operation *op) {
+static LogicalResult foldMemRefCast(Operation *op, Value ignore = nullptr) {
   bool folded = false;
   for (OpOperand &operand : op->getOpOperands()) {
     auto cast = operand.get().getDefiningOp<memref::CastOp>();
-    if (cast && !cast.getOperand().getType().isa<UnrankedMemRefType>()) {
+    if (cast && operand.get() != ignore &&
+        !cast.getOperand().getType().isa<UnrankedMemRefType>()) {
       operand.set(cast.getOperand());
       folded = true;
     }
   }
   return success(folded);
 }
 
 //===----------------------------------------------------------------------===//

 void AffineStoreOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                 MLIRContext *context) {
   results.add<SimplifyAffineOp<AffineStoreOp>>(context);
 }
 
 LogicalResult AffineStoreOp::fold(ArrayRef<Attribute> cstOperands,
                                   SmallVectorImpl<OpFoldResult> &results) {
   /// store(memrefcast) -> store
-  return foldMemRefCast(*this);
+  return foldMemRefCast(*this, getValueToStore());
 }
 
 //===----------------------------------------------------------------------===//
 // AffineMinMaxOpBase
 //===----------------------------------------------------------------------===//
 
 template <typename T>
 static LogicalResult verifyAffineMinMaxOp(T op) {

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

 
 //===----------------------------------------------------------------------===//
 // Common canonicalization pattern support logic
 //===----------------------------------------------------------------------===//
 
 /// This is a common class used for patterns of the form
 /// "someop(memrefcast) -> someop".  It folds the source of any memref.cast
 /// into the root operation directly.
-static LogicalResult foldMemRefCast(Operation *op) {
+static LogicalResult foldMemRefCast(Operation *op, Value inner = nullptr) {
   bool folded = false;
   for (OpOperand &operand : op->getOpOperands()) {
     auto cast = operand.get().getDefiningOp<CastOp>();
-    if (cast && !cast.getOperand().getType().isa<UnrankedMemRefType>()) {
+    if (cast && operand.get() != inner &&
+        !cast.getOperand().getType().isa<UnrankedMemRefType>()) {
       operand.set(cast.getOperand());
       folded = true;
     }
   }
   return success(folded);
 }
 
 //===----------------------------------------------------------------------===//
if (op.getNumOperands() != 2 + op.getMemRefType().getRank())
     return op.emitOpError("store index operand count not equal to memref rank");
 
   return success();
 }
 
 LogicalResult StoreOp::fold(ArrayRef<Attribute> cstOperands,
                             SmallVectorImpl<OpFoldResult> &results) {
   /// store(memrefcast) -> store
-  return foldMemRefCast(*this);
+  return foldMemRefCast(*this, getValueToStore());
 }
 
 //===----------------------------------------------------------------------===//
 // SubViewOp
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// Helpers to write more idiomatic operations.

mlir/test/Dialect/Affine/canonicalize.mlir

affine.for %i = 0 to 1024 {
 
     // Map remains the same, but operand changes on composition.
     %copy = affine.apply affine_map<(d0) -> (d0)> (%i)
     %1 = affine.vector_load %A[%copy] : memref<1024xf32>, vector<8xf32>
     "prevent.dce"(%1) : (vector<8xf32>) -> ()
   }
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @no_fold_of_store
+//  CHECK:   %[[cst]] = memref.cast %arg to memref<?xi8>
+//  CHECK:   affine.store %[[cst]]
+func @no_fold_of_store(%arg : memref<32xi8>, %holder: memref<memref<?xi8>>) {
+  %0 = memref.cast %arg to memref<?xi8>
+  affine.store %0, %holder[] : memref<?xi8>
+  return
+}
+

mlir/test/Dialect/MemRef/canonicalize.mlir

 //       CHECK:   return %[[SIZE]] : index
 func @dim_of_sized_view(%arg : memref<?xi8>, %size: index) -> index {
   %c0 = constant 0 : index
   %0 = memref.reinterpret_cast %arg to offset: [0], sizes: [%size], strides: [0] : memref<?xi8> to memref<?xi8>
   %1 = memref.dim %0, %c0 : memref<?xi8>
   return %1 : index
 }
 
+// -----
+
+// CHECK-LABEL: func @no_fold_of_store
+//  CHECK:   %[[cst]] = memref.cast %arg to memref<?xi8>
+//  CHECK:   memref.store %[[cst]]
+func @no_fold_of_store(%arg : memref<32xi8>, %holder: memref<memref<?xi8>>) {
+  %0 = memref.cast %arg to memref<?xi8>
+  memref.store %0, %holder[] : memref<?xi8>
+  return
+}