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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-gpu

Author: Naaz30

<details>
<summary>Changes</summary>



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


1 Files Affected:

- (modified) mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp (+11-7) 


``````````diff
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

``````````

</details>


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