[Mlir-commits] [mlir] aa3cabe - [mlir][memref] Fix memref.copy of scalar memref

[Stephan Herhut]

Author: Stephan Herhut
Date: 2022-01-14T16:13:12+01:00
New Revision: aa3cabe3cbe8aec83e8966a5c95085f30f3ece64

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

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

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

Differential Revision: https://reviews.llvm.org/D117314

Added: 
    

Modified: 
    mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
    mlir/test/mlir-cpu-runner/copy.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
index 0bc6eb923d164..72eeca1662221 100644
--- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
+++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp
@@ -723,13 +723,13 @@ struct MemRefCopyOpLowering : public ConvertOpToLLVMPattern<memref::CopyOp> {
     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.

diff  --git a/mlir/test/mlir-cpu-runner/copy.mlir b/mlir/test/mlir-cpu-runner/copy.mlir
index 8581f135bc851..35cedc46e2b15 100644
--- a/mlir/test/mlir-cpu-runner/copy.mlir
+++ b/mlir/test/mlir-cpu-runner/copy.mlir
@@ -8,6 +8,7 @@ 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>
@@ -56,10 +57,22 @@ func @main() -> () {
   // 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
 }