[Mlir-commits] [mlir] 76fd3d4 - [mlir][CPURunner] Avoid a crash in memrefCopy when called with empty shapes.

[Adrian Kuegel]

Author: Adrian Kuegel
Date: 2021-08-03T16:02:01+02:00
New Revision: 76fd3d4410c1dc637944a930a3ce760836b4d765

URL: https://github.com/llvm/llvm-project/commit/76fd3d4410c1dc637944a930a3ce760836b4d765
DIFF: https://github.com/llvm/llvm-project/commit/76fd3d4410c1dc637944a930a3ce760836b4d765.diff

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

Differential Revision: https://reviews.llvm.org/D107346

Added: 
    

Modified: 
    mlir/lib/ExecutionEngine/CRunnerUtils.cpp
    mlir/test/mlir-cpu-runner/copy.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/ExecutionEngine/CRunnerUtils.cpp b/mlir/lib/ExecutionEngine/CRunnerUtils.cpp
index 0207232871f87..f8ca6854c2d6a 100644
--- a/mlir/lib/ExecutionEngine/CRunnerUtils.cpp
+++ b/mlir/lib/ExecutionEngine/CRunnerUtils.cpp
@@ -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];

diff  --git a/mlir/test/mlir-cpu-runner/copy.mlir b/mlir/test/mlir-cpu-runner/copy.mlir
index f1fe1f07264cb..2ef1d3138b206 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
 }