[mlir][crunner] fix bug in memref copy for rank 0

While replacing linalg.copy with the more desired memref.copy
I found a bug in the support library for rank 0 memref copying.
The code would loop for something like the following, since there
is code for no-rank and rank > 0, but rank == 0 was unexpected.

  memref.copy %0, %1: memref<f32> to memref<f32>

Note that a "regression test" for this will follow using the
sparse compiler migration to memref.copy which exercises this
case many times.

Reviewed By: herhut

Differential Revision: https://reviews.llvm.org/D106036
This commit is contained in:
Aart Bik 2021-07-14 21:01:00 -07:00
parent 9805afdfea
commit 04bddb6cc7
1 changed files with 8 additions and 3 deletions

View File

@ -47,13 +47,18 @@ memrefCopy(int64_t elemSize, UnrankedMemRefType<char> *srcArg,
DynamicMemRefType<char> dst(*dstArg);
int64_t rank = src.rank;
char *srcPtr = src.data + src.offset * elemSize;
char *dstPtr = dst.data + dst.offset * elemSize;
if (rank == 0) {
memcpy(dstPtr, srcPtr, elemSize);
return;
}
int64_t *indices = static_cast<int64_t *>(alloca(sizeof(int64_t) * rank));
int64_t *srcStrides = static_cast<int64_t *>(alloca(sizeof(int64_t) * rank));
int64_t *dstStrides = static_cast<int64_t *>(alloca(sizeof(int64_t) * rank));
char *srcPtr = src.data + src.offset * elemSize;
char *dstPtr = dst.data + dst.offset * elemSize;
// Initialize index and scale strides.
for (int rankp = 0; rankp < rank; ++rankp) {
indices[rankp] = 0;