[Mlir-commits] [mlir] [MLIR][Affine] Fix fusion crash from memory space int assumption (PR #127032)

[Uday Bondhugula]

https://github.com/bondhugula updated https://github.com/llvm/llvm-project/pull/127032

>From 340d7e4a23f43591e662017b3531526ee571247a Mon Sep 17 00:00:00 2001
From: Uday Bondhugula <uday at polymagelabs.com>
Date: Thu, 13 Feb 2025 14:38:30 +0530
Subject: [PATCH] [MLIR][Affine] Fix fusion crash from memory space int
 assumption

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

Fixes: https://github.com/llvm/llvm-project/issues/118759
---
 .../Dialect/Affine/Transforms/LoopFusion.cpp  |  8 ++---
 mlir/test/Dialect/Affine/loop-fusion-4.mlir   | 34 +++++++++++++++++++
 2 files changed, 38 insertions(+), 4 deletions(-)

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..c8686c8cdc193 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-4.mlir
+++ b/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:      affine.for %{{.*}} = 16 to 32
+  // PRODUCER-CONSUMER-NEXT:   affine.for %{{.*}} = 0 to 8
+  return
+}