[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

379–382

387

388

When does this happen?

402

Avoid isa + cast : use dyn_cast instead.

410

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

388

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

370

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?

388

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

406

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

370

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.

388

Yeah, it was confusing for me too :)

406

Done. Looks much better now

[mamrami]

Updates after CR

[mehdi_amini]

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

370

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

370

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

418

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

 #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 "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.
+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. CastOps are allowed only if
+    // their result type is ranked, has static shape, strides and offsets.
+    // All memref.cast should have the same type, which will be the alloc's new
+    // type.
+    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;
+      auto dstBaseMemRefType =
+          castUser.getDest().getType().template cast<BaseMemRefType>();
+      if (!dstBaseMemRefType.hasRank())
+        return failure();
+      if (!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

 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:           %[[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>
+}