[mlir][memref] Lower copy of memrefs with outer size-1 dims to intrinsic memcpy.

[olegshyshkov]

Summary

With this change, more memref.copy will be lowered to the efficient memcpy. For example,

memref.copy %subview, %alloc : memref<1x576xf32, strided<[704, 1]>> to memref<1x576xf32>

[bkramer]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1067

dimensions

[nicolasvasilache]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1067

While you are at it, please make this a reusable util on the memref type.

1067

s/make this a/extract this as a

[olegshyshkov]

typoes

[olegshyshkov]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1067

Yes, I already planned to do that in a follow-up. I didn't want to move code and do functional changes in the same commit.

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

auto isStaticShapeAndContiguousRowMajor = [](MemRefType type) {
       if (!type.hasStaticShape())
         return false;
 
       SmallVector<int64_t> strides;
       int64_t offset;
       if (failed(getStridesAndOffset(type, strides, offset)))
         return false;
 
+      // MemRef is contiguous if outer dimensions are size-1 and inner
+      // dimensions have unit strides.
       int64_t runningStride = 1;
-      for (unsigned i = strides.size(); i > 0; --i) {
-        if (strides[i - 1] != runningStride)
-          return false;
-        runningStride *= type.getDimSize(i - 1);
+      int64_t curDim = strides.size() - 1;
+      // Finds all inner dimensions with unit strides.
+      while (curDim >= 0 && strides[curDim] == runningStride) {
+        runningStride *= type.getDimSize(curDim);
+        --curDim;
+      }
+
+      // Check if other dimensions are size-1.
+      while (curDim >= 0 && type.getDimSize(curDim) == 1) {
+        --curDim;
       }
-      return true;
+
+      // All dims are unit-strided or size-1.
+      return curDim < 0;
     };
 
     auto isContiguousMemrefType = [&](BaseMemRefType type) {
       auto memrefType = dyn_cast<mlir::MemRefType>(type);
       // We can use memcpy for memrefs if they have an identity layout or are
       // contiguous with an arbitrary offset. Ignore empty memrefs, which is a
       // special case handled by memrefCopy.
       return memrefType &&

mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir

func.func @memref_copy_ranked() {
   // CHECK: "llvm.intr.memcpy"([[GEP2]], [[GEP1]], [[SIZE]], [[VOLATILE]]) : (!llvm.ptr, !llvm.ptr, i64, i1) -> ()
   return
 }
 
 
 // -----
 
 // CHECK-LABEL: func @memref_copy_contiguous
-func.func @memref_copy_contiguous(%in: memref<16x2xi32>, %offset: index) {
+func.func @memref_copy_contiguous(%in: memref<16x4xi32>, %offset: index) {
   %buf = memref.alloc() : memref<1x2xi32>
-  %sub = memref.subview %in[%offset, 0] [1, 2] [1, 1] : memref<16x2xi32> to memref<1x2xi32, strided<[2, 1], offset: ?>>
-  memref.copy %sub, %buf : memref<1x2xi32, strided<[2, 1], offset: ?>> to memref<1x2xi32>
+  %sub = memref.subview %in[%offset, 0] [1, 2] [1, 1] : memref<16x4xi32> to memref<1x2xi32, strided<[4, 1], offset: ?>>
+  memref.copy %sub, %buf : memref<1x2xi32, strided<[4, 1], offset: ?>> to memref<1x2xi32>
   // Skip the memref descriptor of the alloc.
   // CHECK: llvm.insertvalue {{%.*}}, {{%.*}}[4, 1]
   // Get the memref for the subview.
-  // CHECK: %[[SUBVIEW:.*]] = memref.subview %{{.*}}[%{{.*}}, 0] [1, 2] [1, 1] : memref<16x2xi32> to memref<1x2xi32, strided<[2, 1], offset: ?>>
-  // CHECK: %[[DESC:.*]] = builtin.unrealized_conversion_cast %[[SUBVIEW]] : memref<1x2xi32, strided<[2, 1], offset: ?>> to !llvm.struct<(ptr
+  // CHECK: %[[SUBVIEW:.*]] = memref.subview %{{.*}}[%{{.*}}, 0] [1, 2] [1, 1] : memref<16x4xi32> to memref<1x2xi32, strided<[4, 1], offset: ?>>
+  // CHECK: %[[DESC:.*]] = builtin.unrealized_conversion_cast %[[SUBVIEW]] : memref<1x2xi32, strided<[4, 1], offset: ?>> to !llvm.struct<(ptr
   // CHECK: [[EXTRACT0:%.*]] = llvm.extractvalue %[[DESC]][3, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
   // CHECK: [[MUL1:%.*]] = llvm.mul {{.*}}, [[EXTRACT0]] : i64
   // CHECK: [[EXTRACT1:%.*]] = llvm.extractvalue %[[DESC]][3, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
   // CHECK: [[MUL2:%.*]] = llvm.mul [[MUL1]], [[EXTRACT1]] : i64
   // CHECK: [[NULL:%.*]] = llvm.mlir.null : !llvm.ptr
   // CHECK: [[GEP:%.*]] = llvm.getelementptr [[NULL]][1] : (!llvm.ptr) -> !llvm.ptr, i32
   // CHECK: [[PTRTOINT:%.*]] = llvm.ptrtoint [[GEP]] : !llvm.ptr to i64
   // CHECK: [[SIZE:%.*]] = llvm.mul [[MUL2]], [[PTRTOINT]] : i64