[MLIR][Affine] Fix fusion crash from memory space int assumption (#127032)

Fix fusion crash from memory space int assumption from assumption on int
attr-based memory spaces.

Fixes: #118759

Author: bondhugula

mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp

@@ -339,14 +339,14 @@ static Value createPrivateMemRef(AffineForOp forOp,
   auto eltSize = getMemRefIntOrFloatEltSizeInBytes(oldMemRefType);
   assert(eltSize && "memrefs with size elt types expected");
   uint64_t bufSize = *eltSize * *numElements;
-  unsigned newMemSpace;
+  Attribute newMemSpace;
   if (bufSize <= localBufSizeThreshold && fastMemorySpace.has_value()) {
-    newMemSpace = *fastMemorySpace;
+    newMemSpace = b.getI64IntegerAttr(*fastMemorySpace);
   } else {
-    newMemSpace = oldMemRefType.getMemorySpaceAsInt();
+    newMemSpace = oldMemRefType.getMemorySpace();
   }
   auto newMemRefType = MemRefType::get(newShape, oldMemRefType.getElementType(),
-                                       {}, newMemSpace);
+                                       /*map=*/AffineMap(), newMemSpace);
 
   // Create new private memref for fused loop 'forOp'. 'newShape' is always
   // a constant shape.

mlir/test/Dialect/Affine/loop-fusion-4.mlir

@@ -358,6 +358,8 @@ func.func @same_memref_load_multiple_stores(%producer : memref<32xf32>, %produce
   return
 }
 
+// -----
+
 #map = affine_map<()[s0] -> (s0 + 5)>
 #map1 = affine_map<()[s0] -> (s0 + 17)>
 
@@ -391,3 +393,35 @@ func.func @memref_index_type() {
   // PRODUCER-CONSUMER-MAXIMAL: return
   return
 }
+
+// -----
+
+#map = affine_map<(d0) -> (d0)>
+#map1 =affine_map<(d0) -> (d0 + 1)>
+
+// Test non-integer memory spaces.
+
+// PRODUCER-CONSUMER-LABEL: func @non_int_memory_space
+func.func @non_int_memory_space() {
+  %alloc = memref.alloc() : memref<256x8xf32, #spirv.storage_class<StorageBuffer>>
+  affine.for %arg0 = 0 to 64 {
+    affine.for %arg1 = 0 to 8 {
+      %0 = affine.apply #map(%arg1)
+      %1 = affine.load %alloc[%arg0, %0] : memref<256x8xf32, #spirv.storage_class<StorageBuffer>>
+      affine.store %1, %alloc[%arg0, %arg1] : memref<256x8xf32, #spirv.storage_class<StorageBuffer>>
+    }
+  }
+  affine.for %arg0 = 16 to 32 {
+    affine.for %arg1 = 0 to 8 {
+      %0 = affine.apply #map(%arg1)
+      %1 = affine.load %alloc[%arg0, %0] : memref<256x8xf32, #spirv.storage_class<StorageBuffer>>
+      affine.store %1, %alloc[%arg0, %arg1] : memref<256x8xf32, #spirv.storage_class<StorageBuffer>>
+    }
+  }
+  // Fused nest.
+  // PRODUCER-CONSUMER-NEXT: memref.alloc()
+  // PRODUCER-CONSUMER-NEXT: memref.alloc()
+  // PRODUCER-CONSUMER-NEXT: affine.for %{{.*}} = 16 to 32
+  // PRODUCER-CONSUMER-NEXT:   affine.for %{{.*}} = 0 to 8
+  return
+}