diff --git a/mlir/lib/ExecutionEngine/CRunnerUtils.cpp b/mlir/lib/ExecutionEngine/CRunnerUtils.cpp index 0207232871f874..f8ca6854c2d6a0 100644 --- a/mlir/lib/ExecutionEngine/CRunnerUtils.cpp +++ b/mlir/lib/ExecutionEngine/CRunnerUtils.cpp @@ -47,6 +47,11 @@ memrefCopy(int64_t elemSize, UnrankedMemRefType *srcArg, DynamicMemRefType 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 *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]; diff --git a/mlir/test/mlir-cpu-runner/copy.mlir b/mlir/test/mlir-cpu-runner/copy.mlir index f1fe1f07264cbe..2ef1d3138b2068 100644 --- a/mlir/test/mlir-cpu-runner/copy.mlir +++ b/mlir/test/mlir-cpu-runner/copy.mlir @@ -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 }