Skip to content

Commit 76fd3d4

Browse files
committed
[mlir][CPURunner] Avoid a crash in memrefCopy when called with empty shapes.
Differential Revision: https://reviews.llvm.org/D107346
1 parent 1139664 commit 76fd3d4

File tree

2 files changed

+11
-1
lines changed

2 files changed

+11
-1
lines changed

mlir/lib/ExecutionEngine/CRunnerUtils.cpp

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,11 @@ memrefCopy(int64_t elemSize, UnrankedMemRefType<char> *srcArg,
4747
DynamicMemRefType<char> dst(*dstArg);
4848

4949
int64_t rank = src.rank;
50+
// Handle empty shapes -> nothing to copy.
51+
for (int rankp = 0; rankp < rank; ++rankp)
52+
if (src.sizes[rankp] == 0)
53+
return;
54+
5055
char *srcPtr = src.data + src.offset * elemSize;
5156
char *dstPtr = dst.data + dst.offset * elemSize;
5257

@@ -83,7 +88,7 @@ memrefCopy(int64_t elemSize, UnrankedMemRefType<char> *srcArg,
8388
if (axis == 0)
8489
return;
8590
// Else, reset to 0 and undo the advancement of the linear index that
86-
// this axis had. The continue with the axis one outer.
91+
// this axis had. Then continue with the axis one outer.
8792
indices[axis] = 0;
8893
readIndex -= src.sizes[axis] * srcStrides[axis];
8994
writeIndex -= dst.sizes[axis] * dstStrides[axis];

mlir/test/mlir-cpu-runner/copy.mlir

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,5 +45,10 @@ func @main() -> () {
4545
// CHECK-NEXT: [1, 4]
4646
// CHECK-NEXT: [2, 5]
4747

48+
%input_empty = memref.alloc() : memref<3x0x1xf32>
49+
%copy_empty = memref.alloc() : memref<3x0x1xf32>
50+
// Copying an empty shape should do nothing (and should not crash).
51+
memref.copy %input_empty, %copy_empty : memref<3x0x1xf32> to memref<3x0x1xf32>
52+
4853
return
4954
}

0 commit comments

Comments
 (0)