Commit 76fd3d44 authored by Adrian Kuegel's avatar Adrian Kuegel
Browse files

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

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

mlir/lib/ExecutionEngine/CRunnerUtils.cpp

+6 −1
Original line number Diff line number Diff line
@@ -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];

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

+5 −0
Original line number Diff line number Diff line
@@ -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

}