[mlir][memref] Add better support for identity layouts in memref.collapse_shape canonicalizer

[herhut]

Summary

When computing the new type of a collapse_shape operation, we need to at least take into account whether the type has an identity layout, in which case we can easily support dynamic strides. Otherwise, the canonicalizer creates invalid IR.

Longer term, both the verifier and the canoncializer need to be extended to support the general case.

[nicolasvasilache]

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

1354–1364

You should short-circuit here, it does not make much sense to pass isIdentityLayout to the isReshapableDimBand.

[herhut]

Pull size checking out of helper.

[herhut]

PTAL

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

1354–1364

I can do that, too.

[bkramer]

looks good, thanks

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

 /// MemRefType.
 static MemRefType
 computeReshapeCollapsedType(MemRefType type,
                             ArrayRef<AffineMap> reassociation) {
   auto sizes = type.getShape();
   AffineExpr offset;
   SmallVector<AffineExpr, 4> strides;
   auto status = getStridesAndOffset(type, strides, offset);
+  auto isIdentityLayout = type.getLayout().isIdentity();
   (void)status;
   assert(succeeded(status) && "expected strided memref");
 
   SmallVector<int64_t, 4> newSizes;
   newSizes.reserve(reassociation.size());
   SmallVector<AffineExpr, 4> newStrides;
   newStrides.reserve(reassociation.size());
 
   // Use the fact that reassociation is valid to simplify the logic: only use
   // each map's rank.
   assert(isReassociationValid(reassociation) && "invalid reassociation");
   unsigned currentDim = 0;
   for (AffineMap m : reassociation) {
     unsigned dim = m.getNumResults();
     int64_t size = 1;
     AffineExpr stride = strides[currentDim + dim - 1];
-    if (!isReshapableDimBand(currentDim, dim, sizes, strides)) {
+    if (isIdentityLayout ||
+        isReshapableDimBand(currentDim, dim, sizes, strides)) {
+      for (unsigned d = 0; d < dim; ++d) {
+        int64_t currentSize = sizes[currentDim + d];
+        if (ShapedType::isDynamic(currentSize)) {
           size = ShapedType::kDynamicSize;
-      stride = AffineExpr();
+          break;
+        }
+        size *= currentSize;
+      }
     } else {
-      for (unsigned d = 0; d < dim; ++d)
-        size *= sizes[currentDim + d];
+      size = ShapedType::kDynamicSize;
+      stride = AffineExpr();
     }
     newSizes.push_back(size);
     newStrides.push_back(stride);
     currentDim += dim;
   }
 
   // Early-exit: if `type` is contiguous, the result must be contiguous.
   if (canonicalizeStridedLayout(type).getLayout().isIdentity())

mlir/test/Dialect/MemRef/canonicalize.mlir

 // CHECK:           return %[[DYNAMIC]] : memref<?x?xf32>
 // CHECK:         }
 func @collapse_after_memref_cast_type_change(%arg0 : memref<?x512x1x1xf32>) -> memref<?x?xf32> {
   %dynamic = memref.cast %arg0: memref<?x512x1x1xf32> to memref<?x?x?x?xf32>
   %collapsed = memref.collapse_shape %dynamic [[0], [1, 2, 3]] : memref<?x?x?x?xf32> into memref<?x?xf32>
   return %collapsed : memref<?x?xf32>
 }
 
+// -----
+
 // CHECK-LABEL:   func @collapse_after_memref_cast(
 // CHECK-SAME:      %[[INPUT:.*]]: memref<?x512x1x?xf32>) -> memref<?x?xf32> {
 // CHECK:           %[[COLLAPSED:.*]] = memref.collapse_shape %[[INPUT]]
 // CHECK_SAME:        {{\[\[}}0], [1, 2, 3]] : memref<?x512x1x?xf32> into memref<?x?xf32>
 // CHECK:           return %[[COLLAPSED]] : memref<?x?xf32>
 func @collapse_after_memref_cast(%arg0 : memref<?x512x1x?xf32>) -> memref<?x?xf32> {
   %dynamic = memref.cast %arg0: memref<?x512x1x?xf32> to memref<?x?x?x?xf32>
   %collapsed = memref.collapse_shape %dynamic [[0], [1, 2, 3]] : memref<?x?x?x?xf32> into memref<?x?xf32>
   return %collapsed : memref<?x?xf32>
 }
 
 // -----
 
+// CHECK-LABEL:   func @collapse_after_memref_cast_type_change_dynamic(
+// CHECK-SAME:      %[[INPUT:.*]]: memref<1x1x1x?xi64>) -> memref<?x?xi64> {
+// CHECK:           %[[COLLAPSED:.*]] = memref.collapse_shape %[[INPUT]]
+// CHECK_SAME:        {{\[\[}}0, 1, 2], [3]] : memref<1x1x1x?xi64> into memref<1x?xi64>
+// CHECK:           %[[DYNAMIC:.*]] = memref.cast %[[COLLAPSED]] :
+// CHECK-SAME:         memref<1x?xi64> to memref<?x?xi64>
+// CHECK:           return %[[DYNAMIC]] : memref<?x?xi64>
+func @collapse_after_memref_cast_type_change_dynamic(%arg0: memref<1x1x1x?xi64>) -> memref<?x?xi64> {
+  %casted = memref.cast %arg0 : memref<1x1x1x?xi64> to memref<1x1x?x?xi64>
+  %collapsed = memref.collapse_shape %casted [[0, 1, 2], [3]] : memref<1x1x?x?xi64> into memref<?x?xi64>
+  return %collapsed : memref<?x?xi64>
+}
+
+// -----
+
 func @reduced_memref(%arg0: memref<2x5x7x1xf32>, %arg1 :index)
     -> memref<1x4x1xf32, affine_map<(d0, d1, d2)[s0] -> (d0 * 35 + s0 + d1 * 7 + d2)>> {
   %c0 = arith.constant 0 : index
   %c5 = arith.constant 5 : index
   %c4 = arith.constant 4 : index
   %c2 = arith.constant 2 : index
   %c1 = arith.constant 1 : index
   %0 = memref.subview %arg0[%arg1, %arg1, %arg1, 0] [%c1, %c4, %c1, 1] [1, 1, 1, 1]