maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[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
This commit is contained in:
Stephan Herhut 2021-11-29 14:39:51 +01:00
parent fc53eb69c2
commit 95f34e318c
2 changed files with 20 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
View File

@ -1209,6 +1209,12 @@ static void print(OpAsmPrinter &p, CollapseShapeOp op) {
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 @@ static bool isReshapableDimBand(unsigned dim, unsigned extent,
// 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.

14
mlir/test/Dialect/MemRef/ops.mlir
View File

@ -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_static(%arg0: memref<3x4x5xf32>,
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 @@ func @expand_collapse_shape_dynamic(%arg0: memref<?x?x?xf32>,
%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 @@ func @expand_collapse_shape_dynamic(%arg0: memref<?x?x?xf32>,
// 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>) {