From 04bddb6cc7c405f1a82ee1d94f96596c2cb387d9 Mon Sep 17 00:00:00 2001 From: Aart Bik Date: Wed, 14 Jul 2021 21:01:00 -0700 Subject: [PATCH] [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 to memref 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 --- mlir/lib/ExecutionEngine/CRunnerUtils.cpp | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/mlir/lib/ExecutionEngine/CRunnerUtils.cpp b/mlir/lib/ExecutionEngine/CRunnerUtils.cpp index d4ebc46aa47d..0207232871f8 100644 --- a/mlir/lib/ExecutionEngine/CRunnerUtils.cpp +++ b/mlir/lib/ExecutionEngine/CRunnerUtils.cpp @@ -47,13 +47,18 @@ memrefCopy(int64_t elemSize, UnrankedMemRefType *srcArg, DynamicMemRefType 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(alloca(sizeof(int64_t) * rank)); int64_t *srcStrides = static_cast(alloca(sizeof(int64_t) * rank)); int64_t *dstStrides = static_cast(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;