[MLIR] Fix unchecked use of memref memory space attr in affine data copy generate (#116763)

Fix unchecked use of memref memory space attr in affine data copy
generate. In the case of memory accesses without a memory space
attribute or those other than integer attributes, the pass treats them
as slow memory spaces.

Fixes #116536

Author: bondhugula

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

@@ -2304,21 +2304,26 @@ mlir::affine::affineDataCopyGenerate(Block::iterator begin, Block::iterator end,
 
   // Walk this range of operations  to gather all memory regions.
   block->walk(begin, end, [&](Operation *opInst) {
+    Value memref;
+    MemRefType memrefType;
     // Gather regions to allocate to buffers in faster memory space.
     if (auto loadOp = dyn_cast<AffineLoadOp>(opInst)) {
-      if ((filterMemRef.has_value() && filterMemRef != loadOp.getMemRef()) ||
-          (loadOp.getMemRefType().getMemorySpaceAsInt() !=
-           copyOptions.slowMemorySpace))
-        return;
+      memref = loadOp.getMemRef();
+      memrefType = loadOp.getMemRefType();
     } else if (auto storeOp = dyn_cast<AffineStoreOp>(opInst)) {
-      if ((filterMemRef.has_value() && filterMemRef != storeOp.getMemRef()) ||
-          storeOp.getMemRefType().getMemorySpaceAsInt() !=
-              copyOptions.slowMemorySpace)
-        return;
-    } else {
-      // Neither load nor a store op.
-      return;
+      memref = storeOp.getMemRef();
+      memrefType = storeOp.getMemRefType();
     }
+    // Neither load nor a store op.
+    if (!memref)
+      return;
+
+    auto memorySpaceAttr =
+        dyn_cast_or_null<IntegerAttr>(memrefType.getMemorySpace());
+    if ((filterMemRef.has_value() && filterMemRef != memref) ||
+        (memorySpaceAttr &&
+         memrefType.getMemorySpaceAsInt() != copyOptions.slowMemorySpace))
+      return;
 
     // Compute the MemRefRegion accessed.
     auto region = std::make_unique<MemRefRegion>(opInst->getLoc());

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

@@ -333,3 +333,23 @@ func.func @index_elt_type(%arg0: memref<1x2x4x8xindex>) {
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 8
   return
 }
+
+#map = affine_map<(d0) -> (d0 + 1)>
+
+// CHECK-LABEL: func @arbitrary_memory_space
+func.func @arbitrary_memory_space() {
+  %alloc = memref.alloc() : memref<256x8xi8, #spirv.storage_class<StorageBuffer>>
+  affine.for %arg0 = 0 to 32 step 4 {
+    %0 = affine.apply #map(%arg0)
+    affine.for %arg1 = 0 to 8 step 2 {
+      %1 = affine.apply #map(%arg1)
+      affine.for %arg2 = 0 to 8 step 2 {
+        // CHECK: memref.alloc() : memref<1x7xi8>
+        %2 = affine.apply #map(%arg2)
+        %3 = affine.load %alloc[%0, %1] : memref<256x8xi8, #spirv.storage_class<StorageBuffer>>
+        affine.store %3, %alloc[%0, %2] : memref<256x8xi8, #spirv.storage_class<StorageBuffer>>
+      }
+    }
+  }
+  return
+}