[mlir][affine] Support index element in data copy generation pass

[Lewuathe]

Summary

Support index element in memref or vector type in the data copy generation pass. The original issue was reported in https://github.com/llvm/llvm-project/issues/61310.

[Lewuathe]

Resolved by https://reviews.llvm.org/D146495

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

 }
 
 static unsigned getMemRefEltSizeInBytes(MemRefType memRefType) {
   auto elementType = memRefType.getElementType();
 
   unsigned sizeInBits;
   if (elementType.isIntOrFloat()) {
     sizeInBits = elementType.getIntOrFloatBitWidth();
+  } else if (elementType.isIndex()) {
+    sizeInBits = IndexType::kInternalStorageBitWidth;
   } else {
     auto vectorType = elementType.cast<VectorType>();
     sizeInBits =
         vectorType.getElementTypeBitWidth() * vectorType.getNumElements();
   }
   return llvm::divideCeil(sizeInBits, 8);
 }
 

 /// Returns the size of memref data in bytes if it's statically shaped,
 /// std::nullopt otherwise.  If the element of the memref has vector type, takes
 /// into account size of the vector as well.
 //  TODO: improve/complete this when we have target data.
 std::optional<uint64_t> mlir::getMemRefSizeInBytes(MemRefType memRefType) {
   if (!memRefType.hasStaticShape())
     return std::nullopt;
   auto elementType = memRefType.getElementType();
-  if (!elementType.isIntOrFloat() && !elementType.isa<VectorType>())
+  if (!elementType.isIntOrFloat() && !elementType.isa<VectorType>() &&
+      !elementType.isIndex())
     return std::nullopt;
 
   uint64_t sizeInBytes = getMemRefEltSizeInBytes(memRefType);
   for (unsigned i = 0, e = memRefType.getRank(); i < e; i++) {
     sizeInBytes = sizeInBytes * memRefType.getDimSize(i);
   }
   return sizeInBytes;
 }

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

 // RUN: mlir-opt %s -split-input-file -affine-data-copy-generate="generate-dma=false fast-mem-space=0 skip-non-unit-stride-loops" | FileCheck %s
 // Small buffer size to trigger fine copies.
 // RUN: mlir-opt %s -split-input-file -affine-data-copy-generate="generate-dma=false fast-mem-space=0 fast-mem-capacity=1" | FileCheck --check-prefix=CHECK-SMALL %s
+// RUN: mlir-opt %s -split-input-file -affine-data-copy-generate | FileCheck --check-prefix=CHECK-INDEX %s
 
 // Test affine data copy with a memref filter. We use a test pass that invokes
 // affine data copy utility on the input loop nest.
 // '-test-affine-data-copy-memref-filter' passes the first memref found in an
 // affine.load op in the innermost loop as a filter.
 // RUN: mlir-opt %s -split-input-file -test-affine-data-copy='memref-filter' | FileCheck %s --check-prefix=FILTER
 // RUN: mlir-opt %s -split-input-file -test-affine-data-copy='for-memref-region' | FileCheck %s --check-prefix=MEMREF_REGION
 
func.func @empty_loops(%arg0: memref<1024x1024xf64>) {
   }
   affine.for %i = 0 to -16 {
     affine.load %arg0[0, %i] : memref<1024x1024xf64>
   }
   return
   // CHECK-NOT:    memref.alloc
   // CHECK:        return
 }
+
+// -----
+func.func @index_element_memref(%arg0: memref<1x2x4x8xindex>) {
+  // Not to fail with an index element memref
+  affine.for %arg1 = 0 to 1 {
+    affine.for %arg2 = 0 to 2 {
+      affine.for %arg3 = 0 to 4 {
+        affine.for %arg4 = 0 to 8 {
+          affine.store %arg4, %arg0[%arg1, %arg2, %arg3, %arg4] : memref<1x2x4x8xindex>
+        }
+      }
+    }
+  }
+  return
+}
+
+// CHECK-INDEX-LABEL: func @index_element_memref
+// CHECK-INDEX:    %[[ALLOC:.*]] = memref.alloc() : memref<1x2x4x8xindex, 1>
+// CHECK-INDEX-NEXT:    %[[ALLOC1:.*]] = memref.alloc() : memref<1xi32>
+// CHECK-INDEX-NEXT:    affine.for %{{.*}} = 0 to 1 {
+// CHECK-INDEX-NEXT:      affine.for %{{.*}} = 0 to 2 {
+// CHECK-INDEX-NEXT:        affine.for %{{.*}} = 0 to 4 {
+// CHECK-INDEX-NEXT:          affine.for %{{.*}} = 0 to 8 {
+// CHECK-INDEX-NEXT:            affine.store %{{.*}}, %[[ALLOC]][%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref<1x2x4x8xindex, 1>
+// CHECK-INDEX-NEXT:          }
+// CHECK-INDEX-NEXT:        }
+// CHECK-INDEX-NEXT:      }
+// CHECK-INDEX-NEXT:    }
+
+// CHECK-INDEX:    memref.dealloc %[[ALLOC1]] : memref<1xi32>
+// CHECK-INDEX-NEXT:    memref.dealloc %[[ALLOC]] : memref<1x2x4x8xindex, 1>
+// CHECK-INDEX-NEXT:    return