[Mlir-commits] [mlir] [MLIR][XeGPU] Improve multi-reduction layout inference (PR #224618)

[Jianhui Li]

================
@@ -487,6 +487,27 @@ gpu.module @test {
   }
 }
 
+// -----
+gpu.module @test {
+  // CHECK-LABEL: reduce_split_transpose
+  // CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [8, 1, 1], sg_data = [8, 8, 8], order = [0, 2, 1]>, dims = [2]>}
+  // CHECK: xegpu.create_nd_tdesc {{.*}} -> !xegpu.tensor_desc<64x64xf16, #xegpu.layout<sg_layout = [8, 1], sg_data = [8, 64], order = [0, 1]>>
+  // CHECK: %[[SPLIT:.*]] = vector.shape_cast %{{.*}} {layout_result_0 = #xegpu.layout<sg_layout = [8, 1, 1], sg_data = [8, 8, 8], order = [0, 2, 1]>} : vector<64x64xf16> to vector<64x8x8xf16>
----------------
Jianhui-Li wrote:

Can you enhance this test case to show that the order inferred from multi-reduction can't be collapsed by shape_cast? If I understand the code correctly, that is the issue this PR try to fix? 

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