[Mlir-commits] [mlir] [mlir][XeGPU] Fix 2D reduction dim ordering for nested slice layouts (PR #230850)

[Jianhui Li]

================
@@ -512,6 +512,82 @@ gpu.module @xevm_test {
   }
 }
 
+// -----
+// Nested slice layout: source effective lane layout is [16, 1], so dim 1
+// (intra-lane) reduces first, then the cross-lane dim 0.
+// CHECK-LABEL: gpu.func @nested_slice_reduce_order(
+// CHECK-SAME:    %[[ARG0:[0-9a-zA-Z]+]]: vector<4x16x8xi32>) -> i32 {
+// CHECK:         %[[ACC_VEC:.*]] = arith.constant dense<0> : vector<16xi32>
+// CHECK:         %[[ACC_SCALAR:.*]] = arith.constant 0 : i32
+// CHECK:         %[[EXTRACTED:.*]] = vector.extract %[[ARG0]][0] : vector<16x8xi32> from vector<4x16x8xi32>
+// CHECK:         %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[EXTRACTED]], %[[ACC_VEC]] [1] : vector<16x8xi32> to vector<16xi32>
+// CHECK:         %[[REDUCE_2:.*]] = vector.multi_reduction <add>, %[[REDUCE_1]], %[[ACC_SCALAR]] [0] : vector<16xi32> to i32
+// CHECK:         xegpu.convert_layout %[[REDUCE_2]]
+#root = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [1, 1, 1]>
+#src = #xegpu.slice<#root, dims = [0]>
+#res = #xegpu.slice<#src, dims = [0, 1]>
+gpu.module @xevm_module {
+  gpu.func @nested_slice_reduce_order(%arg: vector<4x16x8xi32>) -> i32 {
+    %v = vector.extract %arg[0]
----------------
Jianhui-Li wrote:

Does the test still fail without this special layout? 
Note that we don't expect extract to carry a special layout, so the test doesn't appear very typical use case at current form. 

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