[Mlir-commits] [mlir] #208902: Added xegpu load/store flattening (PR #210592)

[llvmlistbot at llvm.org]

https://github.com/Naaz30 created https://github.com/llvm/llvm-project/pull/210592

None

>From 54feb26eb7c8a206e3440d3eddbfddcb612ff92f Mon Sep 17 00:00:00 2001
From: Naazni Yahya <naazniyahya at Naaznis-MacBook-Air.local>
Date: Sat, 18 Jul 2026 02:05:32 +0530
Subject: [PATCH] [MLIR][Conversion][XeGPU][XeVM] Added xegpu load/store
 flattening(#208902)

---
 .../lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index 6144d7c0c1a15..645fb6494e19c 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -719,9 +719,11 @@ class LoadStoreMatrixToXeVMPattern : public OpConversionPattern<OpType> {
       // Some transforms may leave unit dimension in the 2D vector, adaptors do
       // not catch it for results.
       if (auto vecType = dyn_cast<VectorType>(resType)) {
-        assert(llvm::count_if(vecType.getShape(),
-                              [](int64_t d) { return d != 1; }) <= 1 &&
-               "Expected either 1D vector or nD with unit dimensions");
+        // Flatten to 1D
+        // Accepts genuine multi-dim tiles 
+        // (e.g. a non-distributed `vector<4x8xf32>` result) by reinterpreting the whole
+        // tile as one flat, contiguous run. This is only valid when the
+        // underlying `mem_desc` region is actually contiguous in memory
         resType = VectorType::get({vecType.getNumElements()},
                                   vecType.getElementType());
       }
@@ -787,9 +789,11 @@ class LoadStoreMatrixToXeVMPattern : public OpConversionPattern<OpType> {
 
     if (valOrResVecTy.getNumElements() >= 1) {
       auto chipOpt = xegpu::getChipStr(op);
-      if (!chipOpt ||
-          (*chipOpt != "pvc" && *chipOpt != "bmg" && *chipOpt != "cri")) {
-        // the lowering for chunk load only works for pvc, bmg or cri
+      // Only reject an explicitly unsupported chip; a missing gpu.module/target
+      // (chipOpt == nullopt) is no longer treated as a hard failure, so this
+      // path also works for standalone/non-GPU-kernel-context IR.
+      if (chipOpt && *chipOpt != "pvc" && *chipOpt != "bmg" &&
+          *chipOpt != "cri") {
         return rewriter.notifyMatchFailure(
             op, "The lowering is specific to pvc, bmg or cri.");
       }
@@ -1626,4 +1630,4 @@ void mlir::populateXeGPUToXeVMConversionPatterns(
   patterns.add<DpasMxToXeVMPattern>(typeConverter, patterns.getContext());
   patterns.add<ExtfToXeVMPattern, TruncfToXeVMPattern>(typeConverter,
                                                        patterns.getContext());
-}
+}
\ No newline at end of file