[mlir][memref] Fix bug in verification of memref.collapse_shape
The verifier computed an illegal type with negative dimension size when collapsing partially static memrefs.
Differential Revision: https://reviews.llvm.org/D114702
GitOrigin-RevId: 95f34e318c469806879a0cd1a6c5290901ed12df
diff --git a/lib/Dialect/MemRef/IR/MemRefOps.cpp b/lib/Dialect/MemRef/IR/MemRefOps.cpp
index f806cf5..ac36d1d 100644
--- a/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -1209,6 +1209,12 @@
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
@@ -1217,7 +1223,7 @@
// 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[dim + 1]))
+ 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.
diff --git a/test/Dialect/MemRef/ops.mlir b/test/Dialect/MemRef/ops.mlir
index 9a5f3fe..620d93d 100644
--- a/test/Dialect/MemRef/ops.mlir
+++ b/test/Dialect/MemRef/ops.mlir
@@ -5,6 +5,7 @@
// CHECK-DAG: #[[$strided3D:.*]] = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2 + d2)>
// CHECK-DAG: #[[$strided2DOFF0:.*]] = affine_map<(d0, d1)[s0] -> (d0 * s0 + d1)>
// CHECK-DAG: #[[$strided3DOFF0:.*]] = affine_map<(d0, d1, d2)[s0, s1] -> (d0 * s0 + d1 * s1 + d2)>
+// CHECK-DAG: #[[$strided2D42:.*]] = affine_map<(d0, d1) -> (d0 * 42 + d1)>
// CHECK-LABEL: func @memref_reinterpret_cast
func @memref_reinterpret_cast(%in: memref<?xf32>)
@@ -143,7 +144,8 @@
func @expand_collapse_shape_dynamic(%arg0: memref<?x?x?xf32>,
%arg1: memref<?x?x?xf32, offset : 0, strides : [?, ?, 1]>,
- %arg2: memref<?x?x?xf32, offset : ?, strides : [?, ?, 1]>) {
+ %arg2: memref<?x?x?xf32, offset : ?, strides : [?, ?, 1]>,
+ %arg3: memref<?x42xf32, offset : 0, strides : [42, 1]>) {
%0 = memref.collapse_shape %arg0 [[0, 1], [2]] :
memref<?x?x?xf32> into memref<?x?xf32>
%r0 = memref.expand_shape %0 [[0, 1], [2]] :
@@ -160,6 +162,12 @@
%r2 = memref.expand_shape %2 [[0, 1], [2]] :
memref<?x?xf32, offset : ?, strides : [?, 1]> into
memref<?x4x?xf32, offset : ?, strides : [?, ?, 1]>
+ %3 = memref.collapse_shape %arg3 [[0, 1]] :
+ memref<?x42xf32, offset : 0, strides : [42, 1]> into
+ memref<?xf32, offset : 0, strides : [1]>
+ %r3 = memref.expand_shape %3 [[0, 1]] :
+ memref<?xf32, offset : 0, strides : [1]> into
+ memref<?x42xf32, offset : 0, strides : [42, 1]>
return
}
// CHECK-LABEL: func @expand_collapse_shape_dynamic
@@ -175,6 +183,10 @@
// CHECK-SAME: memref<?x?x?xf32, #[[$strided3D]]> into memref<?x?xf32, #[[$strided2D]]>
// CHECK: memref.expand_shape {{.*}} {{\[}}[0, 1], [2]]
// CHECK-SAME: memref<?x?xf32, #[[$strided2D]]> into memref<?x4x?xf32, #[[$strided3D]]>
+// CHECK: memref.collapse_shape {{.*}} {{\[}}[0, 1]]
+// CHECK-SAME: memref<?x42xf32, #[[$strided2D42]]> into memref<?xf32>
+// CHECK: memref.expand_shape {{.*}} {{\[}}[0, 1]]
+// CHECK-SAME: memref<?xf32> into memref<?x42xf32, #[[$strided2D42]]>
func @expand_collapse_shape_zero_dim(%arg0 : memref<1x1xf32>, %arg1 : memref<f32>)
-> (memref<f32>, memref<1x1xf32>) {
