[mlir] Add a canonicalization pattern of casted AllocOp & AllocaOp

[mamrami]

Summary

In some cases the casted buffer could be allocated in the destination's type.
By rebuilding the allocating op with the matching type, the memref.cast is not needed.

[mamrami]

clang-format

[mehdi_amini]

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

378–381

386

387

When does this happen?

401

Avoid isa + cast : use dyn_cast instead.

409

Can we do all of this in a single traversal of alloc->getUsers() instead of 3?

mlir/test/Dialect/MemRef/canonicalize.mlir

947

Is this useful to check?

949

Same here, I'm not sure this is necessary to check here. Checking the alloc and the NOT: cast seems like capturing the intent of the test I think.

961

Can you add negative tests? Basically we should have a case for each of the return failure() you have in the pattern I think.

[mamrami]

Address CR

[mamrami]

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

387

When castOp result is an UnrankedMemRefType, it cannot be casted to MemRefType.
I changed the code to make it more readable and clear.

[mamrami]

@mehdi_amini - Thanks for the detailed CR!!
And sorry for the delay..

[mehdi_amini]

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

369

Is this valid for all operations in the memref dialect?

Also what about something like:

%alloc = memref.alloc(...)  : memref<?x?xf32>
%memref.copy %alloc, ...
if (%cond) {
  %casted = memref.cast %alloc : memref<?x?xf32> to memref<4x4xf32>
  ...
}

The control flow should prevent the folding, because the cast may not be reachable I think?

387

Ah! Forgot about unranked... It does not help that for Tensor the hierarchy is named with a different convention...

405

This logic seems overly complex because the interaction between isValidUser and this (they both check for CastOp for example`, and getting the result MemrefType of the cast use.

What does it look like if we inline isValidUser here instead?

[mamrami]

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

369

I believe you are right!
I think the cast and the alloc should be in the same block.
I'll add this if you agree it's a suffice requirement.

387

Yeah, it was confusing for me too :)

405

Done. Looks much better now

[mamrami]

Updates after CR

[mehdi_amini]

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

369

I think post-dominate is the condition strictly speaking, but "same block" is a good start for now.

[• generadora1999]

This comment was removed by mehdi_amini.

[mamrami]

CR + rebase

[mamrami]

rebase

[mamrami]

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

369

Added a test with scf.if and blocked it.
Also blocked dynamic/unranked alloc that is casted to static, and added a test for dynamic alloc.

[mehdi_amini]

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

417

mlir/test/Dialect/MemRef/canonicalize.mlir

1041

Can you add // CHECK: memref.cast?
I'm afraid right now it is just deleted, so you may have to add a user of the cast.

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

LogicalResult matchAndRewrite(T alloc,
 
     for (Operation *user : llvm::make_early_inc_range(alloc->getUsers()))
       rewriter.eraseOp(user);
 
     rewriter.eraseOp(alloc);
     return success();
   }
 };
+
+/// Fold casted alloc/alloca operations so they will have the destination's
+/// type, and the memref.cast will not be needed. This is allowed if the cast
+/// does not lose information, of rank and size. In some cases the allocated
+/// buffer type cannot change, for example if it is used by a func.call. Thus
+/// the pattern is limited to alloc/alloca ops that are used by ops from the
+/// memref dialect.
+template <typename T>
+struct SimplifyCastedAlloc : public OpRewritePattern<T> {
+  using OpRewritePattern<T>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(T alloc,
+                                PatternRewriter &rewriter) const override {
+    if (llvm::none_of(alloc->getUsers(),
+                      [&](Operation *op) { return isa<CastOp>(op); }))
+      return failure();
+    auto isValidUser = [&](Operation *op) -> bool {
+      if (!isa<memref::MemRefDialect>(op->getDialect()))
+        return false;
+      auto castUser = dyn_cast<CastOp>(op);
+      if (!castUser)
+        return true;
+      MemRefType currDstMemRefType =
+          cast<CastOp>(op).getDest().getType().dyn_cast<MemRefType>();
+      if (!currDstMemRefType)
+        return false;
+      if (!currDstMemRefType.hasStaticShape())
+        return false;
+      auto [strides, offset] = getStridesAndOffset(currDstMemRefType);
+      return offset != ShapedType::kDynamic &&
+             !llvm::is_contained(strides, ShapedType::kDynamic);
+    };
+    if (!llvm::all_of(alloc->getUsers(), isValidUser))
+      return failure();
+    // All memref.cast should have the same type, which will be the alloc's new
+    // type.
+    MemRefType dstType;
+    for (auto user : alloc->getUsers()) {
+      if (!isa<CastOp>(user))
+        continue;
+      MemRefType currDstType =
+          cast<CastOp>(user).getDest().getType().template cast<MemRefType>();
+      if (!dstType)
+        dstType = currDstType;
+      else if (dstType != currDstType)
+        return failure();
+    }
+    if (alloc.getMemref().getType().template cast<MemRefType>() == dstType)
+      return failure();
+    rewriter.replaceOpWithNewOp<T>(alloc, dstType, alloc.getDynamicSizes(),
+                                   alloc.getSymbolOperands(),
+                                   alloc.getAlignmentAttr());
+    return success();
+  }
+};
 } // namespace
 
 void AllocOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                           MLIRContext *context) {
-  results.add<SimplifyAllocConst<AllocOp>, SimplifyDeadAlloc<AllocOp>>(context);
+  results.add<SimplifyAllocConst<AllocOp>, SimplifyDeadAlloc<AllocOp>,
+              SimplifyCastedAlloc<AllocOp>>(context);
 }
 
 void AllocaOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                            MLIRContext *context) {
-  results.add<SimplifyAllocConst<AllocaOp>, SimplifyDeadAlloc<AllocaOp>>(
-      context);
+  results.add<SimplifyAllocConst<AllocaOp>, SimplifyDeadAlloc<AllocaOp>,
+              SimplifyCastedAlloc<AllocaOp>>(context);
+
 }
 
 //===----------------------------------------------------------------------===//
 // ReallocOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult ReallocOp::verify() {
   auto sourceType = llvm::cast<MemRefType>(getOperand(0).getType());

mlir/test/Dialect/MemRef/canonicalize.mlir

 func.func private @ub_negative_alloc_size() -> memref<?x?x?xi1> {
   %idx1 = index.constant 1
   %c-2 = arith.constant -2 : index
   %c15 = arith.constant 15 : index
 // CHECK:   %[[ALLOC:.*]] = memref.alloc(%c-2) : memref<15x?x1xi1>
   %alloc = memref.alloc(%c15, %c-2, %idx1) : memref<?x?x?xi1>
   return %alloc : memref<?x?x?xi1>
 }
+
+// -----
+
+// CHECK-LABEL:   func @fold_casted_alloc(
+// CHECK-SAME:                            %[[arg:.*]]: memref<1x15x30x256xsi8, #map>) -> (memref<1x15x30x256xsi8, #map>, memref<1x256x15x30xsi8, strided<[115200, 1, 7680, 256]>>) {
+// CHECK:           %[[alloc:.*]] = memref.alloc() : memref<1x15x30x256xsi8, #map>
+// CHECK:           memref.copy %[[arg]], %[[alloc]] : memref<1x15x30x256xsi8, #map> to memref<1x15x30x256xsi8, #map>
+// CHECK-NOT:       memref.cast
+// CHECK:           %[[transpose:.*]] = memref.transpose %[[alloc]] (d0, d1, d2, d3) -> (d0, d3, d1, d2) : memref<1x15x30x256xsi8, #map> to memref<1x256x15x30xsi8, strided<[115200, 1, 7680, 256]>>
+// CHECK:           return %[[alloc]], %[[transpose]] : memref<1x15x30x256xsi8, #map>, memref<1x256x15x30xsi8, strided<[115200, 1, 7680, 256]>>
+
+#map = affine_map<(d0, d1, d2, d3) -> (d0 * 115200 + d1 * 7680 + d2 * 256 + d3)>
+func.func @fold_casted_alloc(%arg0: memref<1x15x30x256xsi8, #map>) -> (memref<1x15x30x256xsi8, #map>, memref<1x256x15x30xsi8, strided<[115200, 1, 7680, 256]>>) {
+  %alloc = memref.alloc() : memref<1x15x30x256xsi8>
+  memref.copy %arg0, %alloc : memref<1x15x30x256xsi8, #map> to memref<1x15x30x256xsi8>
+  %cast = memref.cast %alloc : memref<1x15x30x256xsi8> to memref<1x15x30x256xsi8, #map>
+  %transpose = memref.transpose %alloc (d0, d1, d2, d3) -> (d0, d3, d1, d2) : memref<1x15x30x256xsi8> to memref<1x256x15x30xsi8, strided<[115200, 1, 7680, 256]>>
+  return %cast, %transpose : memref<1x15x30x256xsi8, #map>, memref<1x256x15x30xsi8, strided<[115200, 1, 7680, 256]>>
+}