[mlir][MemRef] Return 0 for the canonical strided layout expr of a 0d memref

There can't be any strides, and the offset for the canonical expr is
always 0. Fixes #55229.

Differential Revision: https://reviews.llvm.org/D124795
This commit is contained in:
Benjamin Kramer 2022-05-02 21:18:37 +02:00
parent 1d712c3818
commit dccb73318a
2 changed files with 16 additions and 4 deletions

View file

@ -945,13 +945,11 @@ MemRefType mlir::canonicalizeStridedLayout(MemRefType t) {
AffineExpr mlir::makeCanonicalStridedLayoutExpr(ArrayRef<int64_t> sizes,
ArrayRef<AffineExpr> exprs,
MLIRContext *context) {
assert(!sizes.empty() && !exprs.empty() &&
"expected non-empty sizes and exprs");
// Size 0 corner case is useful for canonicalizations.
if (llvm::is_contained(sizes, 0))
if (sizes.empty() || llvm::is_contained(sizes, 0))
return getAffineConstantExpr(0, context);
assert(!exprs.empty() && "expected exprs");
auto maps = AffineMap::inferFromExprList(exprs);
assert(!maps.empty() && "Expected one non-empty map");
unsigned numDims = maps[0].getNumDims(), nSymbols = maps[0].getNumSymbols();

View file

@ -1052,6 +1052,20 @@ func.func @memref_copy_contiguous(%in: memref<16x2xi32>, %offset: index) {
// -----
// CHECK-LABEL: func @memref_copy_0d_offset
#map0 = affine_map<(d0) -> (d0 + 1)>
#map1 = affine_map<() -> (1)>
func.func @memref_copy_0d_offset(%in: memref<2xi32>) {
%buf = memref.alloc() : memref<i32>
%sub = memref.subview %in[1] [1] [1] : memref<2xi32> to memref<1xi32, #map0>
%scalar = memref.collapse_shape %sub [] : memref<1xi32, #map0> into memref<i32, #map1>
memref.copy %scalar, %buf : memref<i32, #map1> to memref<i32>
// CHECK: llvm.intr.memcpy
return
}
// -----
// CHECK-LABEL: func @memref_copy_noncontiguous
#map = affine_map<(d0, d1)[s0] -> (d0 * 2 + s0 + d1)>
func.func @memref_copy_noncontiguous(%in: memref<16x2xi32>, %offset: index) {