[mlir][affine] Skip point-wise copy for scalar memref

[Lewuathe]

Summary

Point wise copy does not consider scalar typed memref and it does not generate any affine.for on top. It causes segmentation fault. We can skip the point-wise copy in case of scalar memref. Issue: https://github.com/llvm/llvm-project/issues/61167

[Lewuathe]

@bondhugula @nicolasvasilache @dcaballe Sorry for bothering you but can I ask you to review this change at your convenience?

mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp

static LogicalResult generateCopy(
 
   // Create fully composed affine maps for each memref.
   auto memAffineMap = b.getMultiDimIdentityMap(memIndices.size());
   fullyComposeAffineMapAndOperands(&memAffineMap, &memIndices);
   auto bufAffineMap = b.getMultiDimIdentityMap(bufIndices.size());
   fullyComposeAffineMapAndOperands(&bufAffineMap, &bufIndices);
 
   if (!copyOptions.generateDma) {
+    // Unable to copy generation for scala memref due to lack of access point.
+    if (rank == 0)
+      return failure();
     // Point-wise copy generation.
     auto copyNest =
         generatePointWiseCopy(loc, memref, fastMemRef, lbMaps,
                               /*lbOperands=*/regionSymbols, ubMaps,
                               /*ubOperands=*/regionSymbols, fastBufOffsets,
                               /*isCopyOut=*/region.isWrite(), b);
 
     // Record this so that we can skip it from yet another copy.

mlir/test/Dialect/Affine/affine-data-copy.mlir

func.func @index_elt_type(%arg0: memref<1x2x4x8xindex>) {
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 4
   // CHECK-NEXT:      affine.for %{{.*}} = 0 to 8
 
   // CHECK:     affine.for %{{.*}} = 0 to 2
   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 4
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 8
   return
 }
+
+// CHECK-LABEL: func @scalar_memref_copy_without_dma
+func.func @scalar_memref_copy_without_dma() {
+    %false = arith.constant false
+    %4 = memref.alloc() {alignment = 128 : i64} : memref<i1>
+    affine.store %false, %4[] : memref<i1>
+
+    // CHECK: %[[FALSE:.*]] = arith.constant false
+    // CHECK: %[[MEMREF:.*]] = memref.alloc() {alignment = 128 : i64} : memref<i1>
+    // CHECK: affine.store %[[FALSE]], %[[MEMREF]][] : memref<i1>
+    return
+}
+
+// CHECK-LABEL: func @scalar_memref_copy_in_loop
+func.func @scalar_memref_copy_in_loop(%3:memref<480xi1>) {
+  %false = arith.constant false
+  %4 = memref.alloc() {alignment = 128 : i64} : memref<i1>
+  affine.store %false, %4[] : memref<i1>
+  %5 = memref.alloc() {alignment = 128 : i64} : memref<i1>
+  memref.copy %4, %5 : memref<i1> to memref<i1>
+  affine.for %arg0 = 0 to 480 {
+    %11 = affine.load %3[%arg0] : memref<480xi1>
+    %12 = affine.load %5[] : memref<i1>
+    %13 = arith.cmpi slt, %11, %12 : i1
+    %14 = arith.select %13, %11, %12 : i1
+    affine.store %14, %5[] : memref<i1>
+  }
+
+  // CHECK: %[[FALSE:.*]] = arith.constant false
+  // CHECK: %[[MEMREF:.*]] = memref.alloc() {alignment = 128 : i64} : memref<i1>
+  // CHECK: affine.store %[[FALSE]], %[[MEMREF]][] : memref<i1>
+  // CHECK: %[[TARGET:.*]] = memref.alloc() {alignment = 128 : i64} : memref<i1>
+  // CHECK: memref.copy %alloc, %[[TARGET]] : memref<i1> to memref<i1>
+  // CHECK: %[[FAST_MEMREF:.*]] = memref.alloc() : memref<480xi1>
+  // CHECK: affine.for %{{.*}} = 0 to 480 {
+  // CHECK:   %{{.*}} = affine.load %arg0[%{{.*}}] : memref<480xi1>
+  // CHECK:   affine.store %{{.*}}, %[[FAST_MEMREF]][%{{.*}}] : memref<480xi1>
+  // CHECK: }
+  // CHECK: affine.for %arg1 = 0 to 480 {
+  // CHECK:   %[[L0:.*]] = affine.load %[[FAST_MEMREF]][%arg1] : memref<480xi1>
+  // CHECK:   %[[L1:.*]] = affine.load %[[TARGET]][] : memref<i1>
+  // CHECK:   %[[CMPI:.*]] = arith.cmpi slt, %[[L0]], %[[L1]] : i1
+  // CHECK:   %[[SELECT:.*]] = arith.select %[[CMPI]], %[[L0]], %[[L1]] : i1
+  // CHECK:   affine.store %[[SELECT]], %[[TARGET]][] : memref<i1>
+  // CHECK: }
+  // CHECK: memref.dealloc %[[FAST_MEMREF]] : memref<480xi1>
+  return
+}
 No newline at end of file