[Mlir-commits] [mlir] 04bddb6 - [mlir][crunner] fix bug in memref copy for rank 0

[Aart Bik]

Author: Aart Bik
Date: 2021-07-15T01:15:07-07:00
New Revision: 04bddb6cc7c405f1a82ee1d94f96596c2cb387d9

URL: https://github.com/llvm/llvm-project/commit/04bddb6cc7c405f1a82ee1d94f96596c2cb387d9
DIFF: https://github.com/llvm/llvm-project/commit/04bddb6cc7c405f1a82ee1d94f96596c2cb387d9.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/ExecutionEngine/CRunnerUtils.cpp

Removed: 
    


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