[Mlir-commits] [mlir] [MLIR][XeGPU] Fix Layout collapse dims out of bounds (PR #193661)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-gpu

Author: Jianhui Li (Jianhui-Li)

<details>
<summary>Changes</summary>

Fix a bug in LayoutAttr::collapseDims() implementation. 

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


2 Files Affected:

- (modified) mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp (+1-1) 
- (modified) mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir (+26) 


``````````diff
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index 811b09b011e47..9d83d858d351a 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -599,7 +599,7 @@ DistributeLayoutAttr LayoutAttr::collapseDims(SmallVector<int64_t> dimGroup) {
     // say we have orderVec = {5, 3, 2, 1, 0}
     // Create indices [0, 1, 2, 3, 4]
     SmallVector<size_t> indices =
-        llvm::to_vector(llvm::seq<size_t>(0, orderAttr.size()));
+        llvm::to_vector(llvm::seq<size_t>(0, origOrder.size()));
 
     // Sort indices based on corresponding values
     llvm::sort(indices,
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 5a95185c8de48..65bbfe2f0b100 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -391,3 +391,29 @@ gpu.module @test{
     gpu.return
   }
 }
+
+// -----
+// CHECK-LABEL: gpu.func @shape_cast_collapse_dims_with_order
+gpu.module @test_collapse_dims [#xevm.target<O = 3, chip = "pvc">] {
+  gpu.func @shape_cast_collapse_dims_with_order(%arg0: memref<32x32xf32>) {
+    // CHECK: %[[STEP:.*]] = vector.step {{.*}} : vector<1024xindex>
+    %0 = vector.step : vector<1024xindex>
+
+    // Shape cast from 1D to 2D triggers collapseDims in layout propagation
+    // CHECK: %[[CAST:.*]] = vector.shape_cast %[[STEP]] {{.*}} : vector<1024xindex> to vector<32x32xindex>
+    %1 = vector.shape_cast %0 : vector<1024xindex> to vector<32x32xindex>
+
+    // Anchor the layout with a store operation
+    %ptr = memref.extract_aligned_pointer_as_index %arg0 : memref<32x32xf32> -> index
+    %ptr_i64 = arith.index_cast %ptr : index to i64
+    %mask = arith.constant dense<true> : vector<32x32xi1>
+    %data = arith.constant dense<0.0> : vector<32x32xf32>
+
+    // CHECK: xegpu.store {{.*}} <{{{.*}}layout = #xegpu.layout<inst_data = [32, 32]>{{.*}}}> :
+    xegpu.store %data, %ptr_i64[%1], %mask {
+      layout = #xegpu.layout<inst_data = [32, 32]>
+    } : vector<32x32xf32>, i64, vector<32x32xindex>, vector<32x32xi1>
+
+    gpu.return
+  }
+}

``````````

</details>


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