Skip to content

Commit

Permalink
[mlir][crunner] fix bug in memref copy for rank 0
Browse files Browse the repository at this point in the history
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
  • Loading branch information
aartbik committed Jul 15, 2021
1 parent 9805afd commit 04bddb6
Showing 1 changed file with 8 additions and 3 deletions.
11 changes: 8 additions & 3 deletions mlir/lib/ExecutionEngine/CRunnerUtils.cpp
Expand Up @@ -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;
Expand Down

0 comments on commit 04bddb6

Please sign in to comment.