isReshapableDimBand: ignore strides of unit dims.

[Benoit]

Summary

See the included lit test: without this isReshapableDimBand change, it
fails to compile as it expects a different, dynamic-shape result memref
type:

mlir/test/Dialect/MemRef/ops.mlir:200:8: error: 'memref.collapse_shape' op expected collapsed type to be 'memref<?xf32, affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>>', but got 'memref<64xf32, affine_map<(d0)[s0] -> (d0 + s0)>>'
  %1 = memref.collapse_shape %arg [[0, 1, 2, 3]] : memref<1x1x8x8xf32, affine_map<(d0, d1, d2, d3)[s0, s1, s2] -> (d0 * s0 + d1 * s1 + s2 + d2 * 8 + d3)>> into memref<64xf32, affine_map<(d0)[s0] -> (d0 + s0)>>
       ^

Another thing is changed by this diff: dim -> idx in the
isDynamicSize check. I wasn't able to observe a practical impact of this
change on any test case, but it feels more correct? sizes[idx+1] is what
is being multiplied below, and if really dim had been intended, this
would have been moved out of the for loop, where dim is a constant.

[Benoit]

newline at end of file; minor wording update

[mravishankar]

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

1232

Do we need the special check for sizes[idx] != 1. If a size is 1 then strides[idx+1] == strides[idx] which should be covered by the other part of th &&` check,

[Benoit]

https://reviews.llvm.org/D114702 landed a few days ago and fixes the dim vs idx issue that I was reporting here! Thanks @herhut !

The part that remains here is the sizes[idx+1]!=1 check -- I will update the diff.

[Benoit]

Rebase

[nicolasvasilache]

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

1209

While you're touching this, could you please s/extent/numDims ?
Seems confusing now that I read it months later..

1232

Hmm, can't tell whether this is true or not without digging deeper.
This is a sign this needs to be reimplemented more simply.

Please use BuiltinTypes.h

AffineExpr makeCanonicalStridedLayoutExpr(ArrayRef<int64_t> sizes,
                                          MLIRContext *context);

Then extract the strides and offset and do a simple comparison of strides.drop_front(dim).take_front(numDims).

Please add missing helpers to allow these things to compose properly at the API level.

[Benoit]

Thank you all for weighing in. After discussion with @nicolasvasilache and @ThomasRaoux today I switch to dropping the unit dims, which caused my need for the present diff, by inserting rank-reducing memref.subview in the pattern that is what generating the problematic memref.collapse_shape, see https://reviews.llvm.org/D114993 . So I can now close this diff, but thanks all for the conversation.

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

 }
 
 static void print(OpAsmPrinter &p, CollapseShapeOp op) {
   ::mlir::printReshapeOp<CollapseShapeOp>(p, op);
 }
 
 /// Detect whether memref dims [dim, dim + extent) can be reshaped without
 /// copies.
 static bool isReshapableDimBand(unsigned dim, unsigned extent,
                                 ArrayRef<int64_t> sizes,
                                 ArrayRef<AffineExpr> strides) {
   // Bands of extent one can be reshaped, as they are not reshaped at all.
   if (extent == 1)
     return true;
   // Otherwise, the size of the first dimension needs to be known.
   if (ShapedType::isDynamic(sizes[dim]))
     return false;
   assert(sizes.size() == strides.size() && "mismatched ranks");
   // off by 1 indexing to avoid out of bounds
   //                       V
   for (auto idx = dim, e = dim + extent; idx + 1 < e; ++idx) {
     // Only bands of static shapes are reshapable. This is due to the fact that
     // there is no relation between dynamic sizes and dynamic strides: we do not
     // have enough information to know whether a "-1" size corresponds to the
     // proper symbol in the AffineExpr of a stride.
     if (ShapedType::isDynamic(sizes[idx + 1]))
       return false;
     // TODO: Refine this by passing the proper nDims and nSymbols so we can
     // simplify on the fly and catch more reshapable cases.
-    if (strides[idx] != strides[idx + 1] * sizes[idx + 1])
+    // The size!=1 condition here means that we allow any stride for unit dims,
+    // as strides of unit dims should not make a practical difference.
+    if (sizes[idx] != 1 && strides[idx] != strides[idx + 1] * sizes[idx + 1])
       return false;
   }
   return true;
 }
 
 /// Compute the MemRefType obtained by applying the `reassociation` (which is
 /// expected to be valid) to `type`.
 /// If `type` is Contiguous MemRefType, this always produce a contiguous

mlir/test/Dialect/MemRef/ops.mlir

func @collapse_shape_to_dynamic
   (%arg0: memref<?x?x?x4x?xf32>) -> memref<?x?x?xf32> {
   %0 = memref.collapse_shape %arg0 [[0], [1], [2, 3, 4]] :
     memref<?x?x?x4x?xf32> into memref<?x?x?xf32>
   return %0 : memref<?x?x?xf32>
 }
 //      CHECK: func @collapse_shape_to_dynamic
 //      CHECK:   memref.collapse_shape
 // CHECK-SAME:    [0], [1], [2, 3, 4]
+
+func @collapse_static_shape_with_unit_dims_with_dynamic_strides(%arg: memref<1x1x8x8xf32, affine_map<(d0, d1, d2, d3)[s0, s1, s2] -> (d0 * s0 + d1 * s1 + s2 + d2 * 8 + d3)>>) -> memref<64xf32, affine_map<(d0)[s0] -> (d0 + s0)>> {
+  %1 = memref.collapse_shape %arg [[0, 1, 2, 3]] : memref<1x1x8x8xf32, affine_map<(d0, d1, d2, d3)[s0, s1, s2] -> (d0 * s0 + d1 * s1 + s2 + d2 * 8 + d3)>> into memref<64xf32, affine_map<(d0)[s0] -> (d0 + s0)>>
+  return %1 : memref<64xf32, affine_map<(d0)[s0] -> (d0 + s0)>>
+}
+
+//      CHECK: func @collapse_static_shape_with_unit_dims_with_dynamic_strides
+//      CHECK:   memref.collapse_shape
+// CHECK-SAME:     into memref<64xf32,