[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

374–377

382

383

When does this happen?

397

Avoid isa + cast : use dyn_cast instead.

405

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

mlir/test/Dialect/MemRef/canonicalize.mlir

957

Is this useful to check?

959

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.

971

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

383

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

365

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?

383

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

401

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

365

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.

383

Yeah, it was confusing for me too :)

401

Done. Looks much better now

[mamrami]

Updates after CR

[mehdi_amini]

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

365

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

365

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

413

mlir/test/Dialect/MemRef/canonicalize.mlir

1066

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

 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
+#include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallBitVector.h"
+#include "llvm/Support/Casting.h"
 
 using namespace mlir;
 using namespace mlir::memref;
 
 namespace {
 /// Idiomatic saturated operations on offsets, sizes and strides.
 namespace saturated_arith {
 struct Wrapper {
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. Note that the pattern is applied only on alloc and cast ops
+/// that are in the same block, to avoid folding a cast that appears inside a
+/// control flow.
+template <typename T>
+struct SimplifyCastedAlloc : public OpRewritePattern<T> {
+  using OpRewritePattern<T>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(T alloc,
+                                PatternRewriter &rewriter) const override {
+    // Valid users are ops of memref dialect. Folding the castOp is allowed only
+    // if result types of both alloc and cast are ranked, have static shape,
+    // strides and offsets. All memref.cast should have the same type, which
+    // will be the alloc's new type.
+    auto srcBaseMemRefType =
+        alloc.getMemref().getType().template cast<BaseMemRefType>();
+    MemRefType dstType;
+    Dialect *memrefDialect = alloc->getDialect();
+    for (auto user : alloc->getUsers()) {
+      if (user->getDialect() != memrefDialect)
+        return failure();
+      auto castUser = dyn_cast<CastOp>(user);
+      if (!castUser)
+        continue;
+      if (castUser->getBlock() != alloc->getBlock())
+        return failure();
+      auto dstBaseMemRefType =
+          castUser.getDest().getType().template cast<BaseMemRefType>();
+      if (!srcBaseMemRefType.hasRank() || !dstBaseMemRefType.hasRank())
+        return failure();
+      if (!srcBaseMemRefType.hasStaticShape() ||
+          !dstBaseMemRefType.hasStaticShape())
+        return failure();
+      auto currDstType = dstBaseMemRefType.template cast<MemRefType>();
+      auto [strides, offset] = getStridesAndOffset(currDstType);
+      if (offset == ShapedType::kDynamic ||
+          llvm::is_contained(strides, ShapedType::kDynamic))
+        return failure();
+      if (!dstType)
+        dstType = currDstType;
+      else if (dstType != currDstType)
+        return failure();
+    }
+
+    // If no cast operations were found, the pattern fails.
+    if (!dstType)
+      return failure();
+
+    // If the alloc type should not change, the pattern fails.
+    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

 // 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 @subview_rank_reduction(
 //  CHECK-SAME:     %[[arg0:.*]]: memref<1x384x384xf32>, %[[arg1:.*]]: index
 func.func @subview_rank_reduction(%arg0: memref<1x384x384xf32>, %idx: index)
     -> memref<?x?xf32, strided<[384, 1], offset: ?>> {
   %c1 = arith.constant 1 : index
   // CHECK: %[[subview:.*]] = memref.subview %[[arg0]][0, %[[arg1]], %[[arg1]]] [1, 1, %[[arg1]]] [1, 1, 1] : memref<1x384x384xf32> to memref<1x?xf32, strided<[384, 1], offset: ?>>
   // CHECK: %[[cast:.*]] = memref.cast %[[subview]] : memref<1x?xf32, strided<[384, 1], offset: ?>> to memref<?x?xf32, strided<[384, 1], offset: ?>>
   %0 = memref.subview %arg0[0, %idx, %idx] [1, %c1, %idx] [1, 1, 1]
       : memref<1x384x384xf32> to memref<?x?xf32, strided<[384, 1], offset: ?>>
   // CHECK: return %[[cast]]
   return %0 : memref<?x?xf32, strided<[384, 1], offset: ?>>
 }
+
+// -----
+
+// CHECK-LABEL:   func @fold_casted_alloc
+// CHECK:           %[[alloc:.*]] = memref.alloc() : memref<1x15x30x256xsi8, #map>
+// CHECK-NOT:       memref.cast
+
+#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]>>
+}
+
+// -----
+
+// CHECK-LABEL:   func @uncasted_alloc_is_unchanged
+// CHECK:           %[[alloc:.*]] = memref.alloc() : memref<1x15x30x256xsi8>
+
+#map = affine_map<(d0, d1, d2, d3) -> (d0 * 115200 + d1 * 7680 + d2 * 256 + d3)>
+func.func @uncasted_alloc_is_unchanged(%arg0: 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>
+  %transpose = memref.transpose %alloc (d0, d1, d2, d3) -> (d0, d3, d1, d2) : memref<1x15x30x256xsi8> to memref<1x256x15x30xsi8, strided<[115200, 1, 7680, 256]>>
+  return %transpose : memref<1x256x15x30xsi8, strided<[115200, 1, 7680, 256]>>
+}
+
+// -----
+
+// CHECK-LABEL:   func @non_memref_dialect_users
+// CHECK:           %[[alloc:.*]] = memref.alloc() : memref<1x15x30x256xsi8>
+// CHECK:           memref.cast
+
+#map = affine_map<(d0, d1, d2, d3) -> (d0 * 115200 + d1 * 7680 + d2 * 256 + d3)>
+func.func @non_memref_dialect_users() -> (memref<1x15x30x256xsi8, #map>, memref<1x15x30x256xsi8>) {
+  %alloc = memref.alloc() : memref<1x15x30x256xsi8>
+  %other_dialect_op = "my.op"(%alloc) : (memref<1x15x30x256xsi8>) -> (memref<1x15x30x256xsi8>)
+  %cast = memref.cast %alloc : memref<1x15x30x256xsi8> to memref<1x15x30x256xsi8, #map>
+  return %cast, %other_dialect_op : memref<1x15x30x256xsi8, #map>, memref<1x15x30x256xsi8>
+}
+
+// -----
+
+// CHECK-LABEL:   func @alloc_casted_to_dynamic_shape
+// CHECK:           %[[alloc:.*]] = memref.alloc() : memref<1x15x30x256xsi8>
+// CHECK:           memref.cast
+
+func.func @alloc_casted_to_dynamic_shape() -> (memref<?x15x30x256xsi8>) {
+  %alloc = memref.alloc() : memref<1x15x30x256xsi8>
+  %cast = memref.cast %alloc : memref<1x15x30x256xsi8> to memref<?x15x30x256xsi8>
+  return %cast : memref<?x15x30x256xsi8>
+}
+
+// -----
+
+// CHECK-LABEL:   func @alloc_casted_to_dynamic_stride
+// CHECK:           %[[alloc:.*]] = memref.alloc() : memref<1x15x30x256xsi8, strided<[115200, 1, 7680, 256]>>
+// CHECK:           memref.cast
+
+func.func @alloc_casted_to_dynamic_stride() -> (memref<1x15x30x256xsi8, strided<[?, 1, 7680, 256]>>) {
+  %alloc = memref.alloc() : memref<1x15x30x256xsi8, strided<[115200, 1, 7680, 256]>>
+  %cast = memref.cast %alloc : memref<1x15x30x256xsi8, strided<[115200, 1, 7680, 256]>> to memref<1x15x30x256xsi8, strided<[?, 1, 7680, 256]>>
+  return %cast : memref<1x15x30x256xsi8, strided<[?, 1, 7680, 256]>>
+}
+
+// -----
+
+// CHECK-LABEL:   func @casted_alloc_conflicts
+// CHECK:           %[[alloc:.*]] = memref.alloc() : memref<1x15x30x256xsi8>
+// CHECK:           memref.cast
+// CHECK:           memref.cast
+
+#map = affine_map<(d0, d1, d2, d3) -> (d0 * 115200 + d1 * 7680 + d2 * 256 + d3)>
+func.func @casted_alloc_conflicts() -> (memref<1x15x30x256xsi8, #map>, memref<?x15x30x256xsi8, #map>) {
+  %alloc = memref.alloc() : memref<1x15x30x256xsi8>
+  %cast = memref.cast %alloc : memref<1x15x30x256xsi8> to memref<1x15x30x256xsi8, #map>
+  %cast_0 = memref.cast %alloc : memref<1x15x30x256xsi8> to memref<?x15x30x256xsi8, #map>
+  return %cast, %cast_0 : memref<1x15x30x256xsi8, #map>, memref<?x15x30x256xsi8, #map>
+}
+// -----
+
+// CHECK-LABEL:   func @unranked_casted_alloc
+// CHECK:           %[[alloc:.*]] = memref.alloc() : memref<1x15x30x256xsi8>
+// CHECK:           memref.cast
+
+func.func @unranked_casted_alloc() -> (memref<*xsi8>) {
+  %alloc = memref.alloc() : memref<1x15x30x256xsi8>
+  %cast = memref.cast %alloc : memref<1x15x30x256xsi8> to memref<*xsi8>
+  return %cast : memref<*xsi8>
+}
+
+// -----
+
+
+// CHECK-LABEL:   func @cast_and_alloc_in_different_blocks
+// CHECK:           %[[VAL_0:.*]] = memref.alloc() : memref<1x15x30x256xsi8, #map>
+// CHECK:           %[[VAL_1:.*]] = memref.alloc() : memref<1x15x30x256xsi8, #map>
+
+#map = affine_map<(d0, d1, d2, d3) -> (d0 * 115200 + d1 * 7680 + d2 * 256 + d3)>
+func.func @cast_and_alloc_in_different_blocks(%cond: i1) -> () {
+  %alloc = memref.alloc() : memref<1x15x30x256xsi8, #map>
+  %alloc_1 = memref.alloc() : memref<1x15x30x256xsi8, #map>
+  memref.copy %alloc, %alloc_1 : memref<1x15x30x256xsi8, #map> to memref<1x15x30x256xsi8, #map>
+  scf.if %cond {
+    %cast = memref.cast %alloc : memref<1x15x30x256xsi8, #map> to memref<1x15x30x256xsi8>
+  }
+  return
+}
+
+// -----
+
+
+// CHECK-LABEL:   func @do_not_change_dynamic_alloc(
+// CHECK-SAME:                                 %{{.*}} %[[VAL_1:.*]]: index)
+// CHECK:           %[[alloc:.*]] = memref.alloc(%[[VAL_1]]) : memref<?x15x30x256xsi8>
+
+#map = affine_map<(d0, d1, d2, d3) -> (d0 * 115200 + d1 * 7680 + d2 * 256 + d3)>
+func.func @do_not_change_dynamic_alloc(%arg0: memref<?x15x30x256xsi8, #map>, %arg1: index) -> (memref<1x15x30x256xsi8, #map>, memref<?x256x15x30xsi8, strided<[115200, 1, 7680, 256]>>) {
+  %alloc = memref.alloc(%arg1) : memref<?x15x30x256xsi8>
+  memref.copy %arg0, %alloc : memref<?x15x30x256xsi8, #map> to memref<?x15x30x256xsi8>
+  %cast = memref.cast %alloc : memref<?x15x30x256xsi8> to memref<1x15x30x256xsi8, #map>
+  %transpose = memref.transpose %alloc (d0, d1, d2, d3) -> (d0, d3, d1, d2) : memref<?x15x30x256xsi8> to memref<?x256x15x30xsi8, strided<[115200, 1, 7680, 256]>>
+  return %cast, %transpose : memref<1x15x30x256xsi8, #map>, memref<?x256x15x30xsi8, strided<[115200, 1, 7680, 256]>>
+}