Fix CollapsedLayoutMap for dim size 1 case

[cathyzhyi]

Summary

This change fixes CollapsedLayoutMap for cases where the collapsed
dims are size 1. The cases where inner most dims are size 1 and
noncontiguous can be represented by the strided form and therefore can
be allowed. For such cases, the new stride should be of the next entry
in an association whose dimension is not size 1. If the next entry is
dynamic, it's not possible to decide which stride to use at compilation
time and the stride is set to dynamic.

[springerm]

nice!

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

1825–1827

Nit: I would rewrite it as follows:

The result stride of a reassociation group is the stride of the last entry of the reassociation. (TODO: Should be the minimum stride in the reassociation because strides are not necessarily sorted. E.g., when using memref.transpose.) Dimensions of size 1 should be skipped, because their strides are meaningless and could have any arbitrary value.

1838

Nit: add a comment:

Dynamically-sized dims may turn out to be dims of size 1 at runtime, so the corresponding stride may have to be skipped. (See above comment.) Therefore, the result stride cannot be statically determined and must be dynamic.

mlir/test/Dialect/Tensor/bufferize.mlir

384

these must be func.func now

[cathyzhyi]

Address comments

[cathyzhyi]

Fix test cases

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

 computeCollapsedLayoutMap(MemRefType srcType,
                           ArrayRef<ReassociationIndices> reassociation) {
   int64_t srcOffset;
   SmallVector<int64_t> srcStrides;
   auto srcShape = srcType.getShape();
   if (failed(getStridesAndOffset(srcType, srcStrides, srcOffset)))
     return failure();
 
   // The result strides are exactly the strides of the last entry of each
-  // reassociation.
+  // reassociation. The only exception is when the dim of last entry is of size
+  // 1, the result stride should be of the next entry which isn't size 1.
+  // TODO: Consider memref.transpose case where the strides are not sorted.
   SmallVector<int64_t> resultStrides;
   resultStrides.reserve(reassociation.size());
-  for (ReassociationIndices reassoc : reassociation)
-    resultStrides.push_back(srcStrides[reassoc.back()]);
+  for (const ReassociationIndices &reassoc : reassociation) {
+    ArrayRef<int64_t> ref = llvm::makeArrayRef(reassoc);
+    while (srcShape[ref.back()] == 1 && ref.size() > 1)
+      ref = ref.drop_back();
+    if (!ShapedType::isDynamic(srcShape[ref.back()]) || ref.size() == 1)
+      resultStrides.push_back(srcStrides[ref.back()]);
+    else
+      resultStrides.push_back(ShapedType::kDynamicStrideOrOffset);
+  }
 
   // Validate that each reassociation group is contiguous.
   unsigned resultStrideIndex = resultStrides.size() - 1;
   for (ReassociationIndices reassoc : llvm::reverse(reassociation)) {
     auto trailingReassocs = ArrayRef<int64_t>(reassoc).drop_front();
     using saturated_arith::Wrapper;
     auto stride = Wrapper::stride(resultStrides[resultStrideIndex--]);
     for (int64_t idx : llvm::reverse(trailingReassocs)) {

mlir/test/Dialect/MemRef/canonicalize.mlir

 // CHECK-LABEL: func @compose_collapse_of_collapse
 //       CHECK:   memref.collapse_shape %{{.*}} {{\[}}[0, 1, 2], [3, 4]]
 //   CHECK-NOT:   memref.collapse_shape
 
 // -----
 
 func @do_not_compose_collapse_of_expand_non_identity_layout(
     %arg0: memref<?x?xf32, offset : 0, strides : [?, 1]>)
-    -> memref<?xf32> {
+    -> memref<?xf32, offset : 0, strides : [?]> {
   %1 = memref.expand_shape %arg0 [[0, 1], [2]] :
     memref<?x?xf32, offset : 0, strides : [?, 1]> into
     memref<?x4x?xf32, offset : 0, strides : [?, ?, 1]>
   %2 = memref.collapse_shape %1 [[0, 1, 2]] :
     memref<?x4x?xf32, offset : 0, strides : [?, ?, 1]> into
-    memref<?xf32>
-  return %2 : memref<?xf32>
+    memref<?xf32, offset : 0, strides : [?]>
+  return %2 : memref<?xf32, offset : 0, strides : [?]>
 }
 // CHECK-LABEL: func @do_not_compose_collapse_of_expand_non_identity_layout
 // CHECK: expand
 // CHECK: collapse
 
 // -----
 
 func @compose_expand_of_expand(%arg0 : memref<?x?xf32>)

mlir/test/Dialect/Tensor/bufferize.mlir

 // RUN: mlir-opt %s -tensor-bufferize | FileCheck %s
 
  // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
  // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 20 + s0 + d1)>
  // CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3)[s0] -> (d0 * 140 + d1 * 20 + d2 * 5 + d3 + s0)>
-// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0) -> (d0 + 1)>
-// CHECK-DAG: #[[$MAP4:.*]] = affine_map<() -> (1)>
-// CHECK-DAG: #[[$MAP5:.*]] = affine_map<(d0, d1) -> (d0 * 2 + d1)>
+ // CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1) -> (d0 * 2 + d1)>
+ // CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0) -> (d0 * 2)>
+ // CHECK-DAG: #[[$MAP5:.*]] = affine_map<(d0) -> (d0 + 1)>
+ // CHECK-DAG: #[[$MAP6:.*]] = affine_map<() -> (1)>
+ // CHECK-DAG: #[[$MAP7:.*]] = affine_map<(d0, d1, d2)[s0] -> (d0 * 8 + s0 + d1 * 4 + d2)>
+ // CHECK-DAG: #[[$MAP8:.*]] = affine_map<(d0)[s0] -> (d0 * 4 + s0)>
+ // CHECK-DAG: #[[$MAP9:.*]] = affine_map<(d0, d1, d2)[s0, s1] -> (d0 * s1 + s0 + d1 * 4 + d2)>
+ // CHECK-DAG: #[[$MAP10:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
 
 // CHECK-LABEL:   func @dim(
 // CHECK-SAME:              %[[TENSOR:.*]]: tensor<f32>,
 // CHECK-SAME:              %[[INDEX:.*]]: index) -> index {
 // CHECK:           %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]] : memref<f32>
 // CHECK:           %[[EXTENT:.*]] = memref.dim %[[MEMREF]], %[[INDEX]] : memref<f32>
 // CHECK:           return %[[EXTENT]] : index
 func @dim(%arg0: tensor<f32>, %arg1: index) -> index {
func @tensor.expand_shape_of_slice(
   // CHECK-SAME: [0, 1], [2, 3]] : memref<?x10xf32, #[[$MAP1]]> into memref<?x7x2x5xf32, #[[$MAP2]]>
   %1 = tensor.expand_shape %0 [[0, 1], [2, 3]] :
       tensor<?x10xf32> into tensor<?x7x2x5xf32>
   // CHECK: %[[r:.*]] = bufferization.to_tensor %[[expanded]]
   // CHECK: return %[[r]]
   return %1 : tensor<?x7x2x5xf32>
 }
 
-// CHECK-LABEL: func @tensor.expand_shape_of_slice2(
-//  CHECK-SAME:     %[[t1:.*]]: tensor<1x2xf32>
-func @tensor.expand_shape_of_slice2(%t1: tensor<1x2xf32>) -> tensor<1xf32> {
-  // CHECK: memref.subview {{.*}} : memref<1x2xf32> to memref<1x1xf32, #[[$MAP5]]>
-  %0 = tensor.extract_slice %t1[0, 0][1, 1][1, 1] : tensor<1x2xf32> to tensor<1x1xf32>
-  // CHECK: memref.collapse_shape %{{.*}} [
-  // CHECK-SAME: [0, 1]] : memref<1x1xf32, #[[$MAP5]]> into memref<1xf32>
-  %1 = tensor.collapse_shape %0 [[0, 1]] : tensor<1x1xf32> into tensor<1xf32>
-  return %1 : tensor<1xf32>
-}
-
 // CHECK-LABEL: func @tensor.collapse_shape(
 //  CHECK-SAME:     %[[t1:.*]]: tensor<2x?x?xf32>
 func @tensor.collapse_shape(%t1: tensor<2x?x?xf32>) -> tensor<?x?xf32> {
   // CHECK: %[[m1:.*]] = bufferization.to_memref %[[t1]] : memref<2x?x?xf32>
   // CHECK: %[[collapsed:.*]] = memref.collapse_shape %[[m1]] [
   // CHECK-SAME: [0, 1], [2]] : memref<2x?x?xf32> into memref<?x?xf32>
   %0 = tensor.collapse_shape %t1 [[0, 1], [2]]
       : tensor<2x?x?xf32> into tensor<?x?xf32>
func @tensor.collapse_shape_to_scalar(%t1: tensor<1x1x1xf32>) -> tensor<f32> {
 
   // CHECK: %[[r:.*]] = bufferization.to_tensor %[[collapsed]]
   // CHECK: return %[[r]]
   return %0 : tensor<f32>
 }
 
 // CHECK-LABEL: func @tensor.collapse_shape_of_slice(
 func @tensor.collapse_shape_of_slice(%arg0: tensor<2xi32>) -> tensor<i32> {
-  // CHECK: memref.subview %{{.*}}[1] [1] [1] : memref<2xi32> to memref<1xi32, #[[$MAP3]]>
+  // CHECK: memref.subview %{{.*}}[1] [1] [1] : memref<2xi32> to memref<1xi32, #[[$MAP5]]>
   %0 = tensor.extract_slice %arg0[1] [1] [1] : tensor<2xi32> to tensor<1xi32>
-  // CHECK: memref.collapse_shape %{{.*}} [] : memref<1xi32, #[[$MAP3]]> into memref<i32, #[[$MAP4]]>
+  // CHECK: memref.collapse_shape %{{.*}} [] : memref<1xi32, #[[$MAP5]]> into memref<i32, #[[$MAP6]]>
   %1 = tensor.collapse_shape %0 [] : tensor<1xi32> into tensor<i32>
   return %1 : tensor<i32>
 }
+
+// CHECK-LABEL: func @tensor.collapse_shape_of_slice2(
+//  CHECK-SAME:     %[[t1:.*]]: tensor<1x2xf32>
+func @tensor.collapse_shape_of_slice2(%t1: tensor<1x2xf32>) -> tensor<1xf32> {
+  // CHECK: memref.subview {{.*}} : memref<1x2xf32> to memref<1x1xf32, #[[$MAP3]]>
+  %0 = tensor.extract_slice %t1[0, 0][1, 1][1, 1] : tensor<1x2xf32> to tensor<1x1xf32>
+  // CHECK: memref.collapse_shape %{{.*}} [
+  // CHECK-SAME: [0, 1]] : memref<1x1xf32, #[[$MAP3]]> into memref<1xf32, #[[$MAP4]]>
+  %1 = tensor.collapse_shape %0 [[0, 1]] : tensor<1x1xf32> into tensor<1xf32>
+  return %1 : tensor<1xf32>
+}
+
+// CHECK-LABEL:   func @tensor.collapse_shape_of_slice3(
+//  CHECK-SAME:     %[[t1:.*]]: tensor<?x2x4xf32>,
+// CHECK-SAME:      %[[OFFSET:.*]]: index) -> tensor<8xf32> {
+func @tensor.collapse_shape_of_slice3(%arg0: tensor<?x2x4xf32>, %offset: index, %size: index) -> tensor<8xf32> {
+  // CHECK: memref.subview %{{.*}} : memref<?x2x4xf32> to memref<4x2x1xf32, #[[$MAP7]]>
+  %0 = tensor.extract_slice %arg0[0, 0, %offset] [4, 2, 1] [1, 1, 1] : tensor<?x2x4xf32> to tensor<4x2x1xf32>
+  // CHECK: memref.collapse_shape %{{.*}} [
+  // CHECK-SAME: [0, 1, 2]] : memref<4x2x1xf32, #[[$MAP7]]> into memref<8xf32, #[[$MAP8]]>
+  %ret = tensor.collapse_shape %0 [[0, 1, 2]] : tensor<4x2x1xf32> into tensor<8xf32>
+  return %ret: tensor<8xf32>
+}
+
+// CHECK-LABEL:   func @tensor.collapse_shape_of_slice4(
+//  CHECK-SAME:     %[[t1:.*]]: tensor<?x?x4xf32>,
+// CHECK-SAME:      %[[OFFSET:.*]]: index,
+// CHECK-SAME:      %[[SIZE:.*]]: index) -> tensor<?xf32> {
+func @tensor.collapse_shape_of_slice4(%arg0: tensor<?x?x4xf32>, %offset: index, %size: index) -> tensor<?xf32> {
+  // CHECK: memref.subview %{{.*}} : memref<?x?x4xf32> to memref<4x?x1xf32, #[[$MAP9]]>
+  %0 = tensor.extract_slice %arg0[0, 0, %offset] [4, %size, 1] [1, 1, 1] : tensor<?x?x4xf32> to tensor<4x?x1xf32>
+  // CHECK: memref.collapse_shape %{{.*}} [
+  // CHECK-SAME: [0, 1, 2]] : memref<4x?x1xf32, #[[$MAP9]]> into memref<?xf32, #[[$MAP10]]>
+  %ret = tensor.collapse_shape %0 [[0, 1, 2]] : tensor<4x?x1xf32> into tensor<?xf32>
+  return %ret: tensor<?xf32>
+}