[Mlir-commits] [mlir] [MLIR][XeGPU] Promote memref.alloca to SLM in convert-vector-to-xegpu (PR #197978)

[Igor Zamyatin]

================
@@ -946,9 +947,86 @@ struct ContractionLowering : public OpRewritePattern<vector::ContractionOp> {
   }
 };
 
+// Returns `memrefTy` with its memory space replaced by `newMemSpace`.
+static MemRefType withMemorySpace(MemRefType memrefTy, Attribute newMemSpace) {
+  return MemRefType::get(memrefTy.getShape(), memrefTy.getElementType(),
+                         memrefTy.getLayout(), newMemSpace);
+}
+
+// Rewrite every `memref.alloca` not already in shared local memory (SLM) to
+// be in SLM (address space 3), and propagate the new memory space through
+// memref-producing aliasing users (e.g. memref.cast, memref.subview,
+// memref.expand_shape, ...). Consumers that take a memref operand but
+// produce a non-memref result (e.g. vector.transfer_read, vector.load) are
+// left untouched: their operand type simply reflects the new memory space.
+//
+// This makes `xegpu.load_matrix`/`xegpu.store_matrix` lowering work end-to-end
+// for IR coming from bufferization, which by default assigns memory space 0/1
+// to allocations.
+static void promoteAllocasToSLM(Operation *root) {
+  MLIRContext *ctx = root->getContext();
+  Attribute slmAttr = IntegerAttr::get(IntegerType::get(ctx, 64), 3);
----------------
Garra1980 wrote:

Can xegpu::MemorySpaceAttr be used here?

https://github.com/llvm/llvm-project/pull/197978