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

[Artem Kroviakov]

================
@@ -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>
----------------
akroviakov wrote:

Corrected the test. The new version fails on main with
```
LLVM ERROR: dimensions being collapsed must be adjacent in order
```

This happens because the propagation shifts the reduced dims to the last walk order, so an expanding shape cast (e.g., `...x64` -> `...x8x8`) that feeds an inner-most reduction gets non-adjacent walk orders for the expansion result, failing the collapseDims.

This PR changes the order inference logic for reductions to set the order of the reduced dim as an inner dim to the outer neighbor.

We cannot generally know the "true" order without looking at the reduction input producer, and the reduction result inference is not designed to query reduction input producers, so this PR sets a more reasonable rule that the reduced dim was adjacent in walk order to its non-reduced neighbor dim.

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