[MLIR][XeGPU] Fix Layout collapse dims out of bounds (#193661)

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

---------

Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>

Author: Jianhui-Li

mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp

@@ -615,7 +615,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,

mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir

@@ -373,3 +373,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
+  }
+}