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

[aartbik]

Summary

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.

[herhut]

Thanks for fixing this!

Do you need the support for arbitrary strides in your use case? I have also planned to add specialized code generation for the case that the maps are identity maps and one can just use a memcpy. I just did not have a use for that, yet.

[aartbik]

In D106036#2879217, @herhut wrote:

Do you need the support for arbitrary strides in your use case? I have also planned to add specialized code generation for the case that the maps are identity maps and one can just use a memcpy. I just did not have a use for that, yet.

Not right away, but in the long run, probably yes!. Thanks!

mlir/lib/ExecutionEngine/CRunnerUtils.cpp

 
 extern "C" MLIR_CRUNNERUTILS_EXPORT void
 memrefCopy(int64_t elemSize, UnrankedMemRefType<char> *srcArg,
            UnrankedMemRefType<char> *dstArg) {
   DynamicMemRefType<char> src(*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;
     srcStrides[rankp] = src.strides[rankp] * elemSize;
     dstStrides[rankp] = dst.strides[rankp] * elemSize;
   }
 
   int64_t readIndex = 0, writeIndex = 0;