[MLIR][Affine] Fix fusion crash from memory space int assumption

[bondhugula]

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

Fixes: #118759, #108369

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-affine

Author: Uday Bondhugula (bondhugula)

Changes

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

Fixes: #118759

---

Full diff: https://github.com/llvm/llvm-project/pull/127032.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp (+4-4)
• (modified) mlir/test/Dialect/Affine/loop-fusion-4.mlir (+30)

diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
index 30019447d94e8..090ee7829b593 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
+++ b/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.
diff --git a/mlir/test/Dialect/Affine/loop-fusion-4.mlir b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
index 788d7f9470530..6773a0294072b 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-4.mlir
@@ -391,3 +391,33 @@ 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:      affine.for %{{.*}} = 16 to 32
+  // PRODUCER-CONSUMER-NEXT:   affine.for %{{.*}} = 0 to 8
+  return
+}

[patel-vimal]

LGTM. Thanks.

[arnab-polymage]

LGTM