[mlir][CPURunner] Avoid a crash in memrefCopy when called with empty shapes.

Differential Revision: https://reviews.llvm.org/D107346
This commit is contained in:
Adrian Kuegel 2021-08-03 15:18:01 +02:00
parent 11396641e4
commit 76fd3d4410
2 changed files with 11 additions and 1 deletions

View File

@ -47,6 +47,11 @@ memrefCopy(int64_t elemSize, UnrankedMemRefType<char> *srcArg,
DynamicMemRefType<char> dst(*dstArg);
int64_t rank = src.rank;
// Handle empty shapes -> nothing to copy.
for (int rankp = 0; rankp < rank; ++rankp)
if (src.sizes[rankp] == 0)
return;
char *srcPtr = src.data + src.offset * elemSize;
char *dstPtr = dst.data + dst.offset * elemSize;
@ -83,7 +88,7 @@ memrefCopy(int64_t elemSize, UnrankedMemRefType<char> *srcArg,
if (axis == 0)
return;
// Else, reset to 0 and undo the advancement of the linear index that
// this axis had. The continue with the axis one outer.
// this axis had. Then continue with the axis one outer.
indices[axis] = 0;
readIndex -= src.sizes[axis] * srcStrides[axis];
writeIndex -= dst.sizes[axis] * dstStrides[axis];

View File

@ -45,5 +45,10 @@ func @main() -> () {
// CHECK-NEXT: [1, 4]
// CHECK-NEXT: [2, 5]
%input_empty = memref.alloc() : memref<3x0x1xf32>
%copy_empty = memref.alloc() : memref<3x0x1xf32>
// Copying an empty shape should do nothing (and should not crash).
memref.copy %input_empty, %copy_empty : memref<3x0x1xf32> to memref<3x0x1xf32>
return
}