Fix bufferization of collapse_shape of subviews with size 1 dims.

[jreiffers]

Summary

Currently, there's an optimization that claims dimensions of size 1 are always
contiguous. This is not necessarily the case for subviews.

Input:

[
  [
    [0, 1],
    [2, 3]
  ],
  [
    [4, 5]
    [6, 7]
  ]
]

Subview:

  [
    [
      [0, 1],
    ],
    [
      [4, 5]
    ]
  ]

The old logic treats this subview as contiguous, when it is not.

[nicolasvasilache]

Thanks for fixing this @jreiffers

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

for (const ReassociationIndices &reassoc : llvm::reverse(reassociation)) {
     using saturated_arith::Wrapper;
     auto stride = Wrapper::stride(resultStrides[resultStrideIndex--]);
     for (int64_t idx : llvm::reverse(trailingReassocs)) {
       stride = stride * Wrapper::size(srcShape[idx]);
 
       // Both source and result stride must have the same static value. In that
       // case, we can be sure, that the dimensions are collapsible (because they
       // are contiguous).
-      //
-      // One special case is when the srcShape is `1`, in which case it can
-      // never produce non-contiguity.
-      if (srcShape[idx] == 1)
-        continue;
-
       // If `strict = false` (default during op verification), we accept cases
       // where one or both strides are dynamic. This is best effort: We reject
       // ops where obviously non-contiguous dims are collapsed, but accept ops
       // where we cannot be sure statically. Such ops may fail at runtime. See
       // the op documentation for details.
       auto srcStride = Wrapper::stride(srcStrides[idx - 1]);
       if (strict && (stride.saturated || srcStride.saturated))
         return failure();

mlir/test/Dialect/Tensor/bufferize.mlir

func.func @tensor.collapse_shape_of_slice4(%arg0: tensor<?x2x4xf32>, %offset: index, %size: index) -> tensor<8xf32> {
   // CHECK: memref.collapse_shape %{{.*}} [
   // CHECK-SAME: [0, 1, 2]] : memref<4x2x1xf32, strided<[8, 4, 1], offset: ?>> into memref<8xf32, strided<[4], offset: ?>>
   %ret = tensor.collapse_shape %0 [[0, 1, 2]] : tensor<4x2x1xf32> into tensor<8xf32>
   return %ret: tensor<8xf32>
 }
 
 // -----
 
+// CHECK-LABEL: func @tensor.collapse_shape_of_slice5(
+func.func @tensor.collapse_shape_of_slice5(%arg0: tensor<2x2x2xi64>) -> tensor<4xi64> {
+  // CHECK: %[[subview:.*]] = memref.subview %{{.*}} : memref<2x2x2xi64> to memref<2x1x2xi64, #{{.*}}>
+  %0 = tensor.extract_slice %arg0[0, 0, 0] [2, 1, 2] [1, 1, 1] : tensor<2x2x2xi64> to tensor<2x1x2xi64>
+
+  // This memref is not collapsible, so the buffer must be copied to get rid of
+  // the layout map.
+  // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<2x1x2xi64>
+  // CHECK: memref.copy %[[subview]], %[[alloc]]
+  // CHECK: memref.collapse_shape %[[alloc]] [
+  // CHECK-SAME: [0, 1, 2]] : memref<2x1x2xi64> into memref<4xi64>
+  %1 = tensor.collapse_shape %0 [[0, 1, 2]] : tensor<2x1x2xi64> into tensor<4xi64>
+  return %1 : tensor<4xi64>
+}
+
+// -----
+
 // CHECK-LABEL: func @tensor.reshape(
 //  CHECK-SAME:     %[[t1:.*]]: tensor<?x10xf32>
 func.func @tensor.reshape(%t1: tensor<?x10xf32>) -> tensor<2x2x5xf32> {
   // CHECK: %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<?x10xf32>
 
   // CHECK: %[[two:.*]] = arith.constant 2 : i64
   %two = arith.constant 2 : i64
   // CHECK: %[[five:.*]] = arith.constant 5 : i64