[mlir][memref] Fix memref.copy of scalar memref

[herhut]

Summary

Also fix a memory leak in the test while at it.

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

struct MemRefCopyOpLowering : public ConvertOpToLLVMPattern<memref::CopyOp> {
   lowerToMemCopyIntrinsic(memref::CopyOp op, OpAdaptor adaptor,
                           ConversionPatternRewriter &rewriter) const {
     auto loc = op.getLoc();
     auto srcType = op.source().getType().dyn_cast<MemRefType>();
 
     MemRefDescriptor srcDesc(adaptor.source());
 
     // Compute number of elements.
-    Value numElements;
+    Value numElements = rewriter.create<LLVM::ConstantOp>(
+        loc, getIndexType(), rewriter.getIndexAttr(1));
     for (int pos = 0; pos < srcType.getRank(); ++pos) {
       auto size = srcDesc.size(rewriter, loc, pos);
-      numElements = numElements
-                        ? rewriter.create<LLVM::MulOp>(loc, numElements, size)
-                        : size;
+      numElements = rewriter.create<LLVM::MulOp>(loc, numElements, size);
     }
+
     // Get element size.
     auto sizeInBytes = getSizeInBytes(loc, srcType.getElementType(), rewriter);
     // Compute total.
     Value totalSize =
         rewriter.create<LLVM::MulOp>(loc, numElements, sizeInBytes);
 
     Value srcBasePtr = srcDesc.alignedPtr(rewriter, loc);
     MemRefDescriptor targetDesc(adaptor.target());

mlir/test/mlir-cpu-runner/copy.mlir

 // RUN: mlir-opt %s -convert-scf-to-std -convert-arith-to-llvm -convert-memref-to-llvm -convert-std-to-llvm -reconcile-unrealized-casts \
 // RUN: | mlir-cpu-runner -e main -entry-point-result=void \
 // RUN: -shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext,%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext \
 // RUN: | FileCheck %s
 
 func private @print_memref_f32(memref<*xf32>) attributes { llvm.emit_c_interface }
 
 func @main() -> () {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
+  %c42 = arith.constant 42.0 : f32
 
   // Initialize input.
   %input = memref.alloc() : memref<2x3xf32>
   %dim_x = memref.dim %input, %c0 : memref<2x3xf32>
   %dim_y = memref.dim %input, %c1 : memref<2x3xf32>
   scf.parallel (%i, %j) = (%c0, %c0) to (%dim_x, %dim_y) step (%c1, %c1) {
     %prod = arith.muli %i,  %dim_y : index
     %val = arith.addi %prod, %j : index
func @main() -> () {
   memref.copy %input_empty, %copy_empty : memref<3x0x1xf32> to memref<3x0x1xf32>
 
   %input_empty_casted = memref.reinterpret_cast %input_empty to offset: [0], sizes: [0, 3, 1], strides: [3, 1, 1]
     : memref<3x0x1xf32> to memref<0x3x1xf32, offset: 0, strides: [3, 1, 1]>
   %copy_empty_casted = memref.alloc() : memref<0x3x1xf32>
   // Copying a casted empty shape should do nothing (and should not crash).
   memref.copy %input_empty_casted, %copy_empty_casted : memref<0x3x1xf32, offset: 0, strides: [3, 1, 1]> to memref<0x3x1xf32>
 
+  %scalar = memref.alloc() : memref<f32>
+  memref.store %c42, %scalar[] : memref<f32>
+  %scalar_copy = memref.alloc() : memref<f32>
+  memref.copy %scalar, %scalar_copy : memref<f32> to memref<f32>
+  %unranked_scalar_copy = memref.cast %scalar_copy : memref<f32> to memref<*xf32>
+  call @print_memref_f32(%unranked_scalar_copy) : (memref<*xf32>) -> ()
+  // CHECK: rank = 0 offset = 0 sizes = [] strides = []
+  // CHECK-NEXT [42]
+
   memref.dealloc %copy_empty : memref<3x0x1xf32>
+  memref.dealloc %copy_empty_casted : memref<0x3x1xf32>
   memref.dealloc %input_empty : memref<3x0x1xf32>
   memref.dealloc %copy_two : memref<3x2xf32>
   memref.dealloc %copy : memref<2x3xf32>
   memref.dealloc %input : memref<2x3xf32>
+  memref.dealloc %scalar : memref<f32>
+  memref.dealloc %scalar_copy : memref<f32>
   return
 }