[mlir][xegpu] propagate-layout infers a source layout for a splitting vector.shape_cast that changes per-subgroup element ownership
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
## Summary
`xegpu-propagate-layout=layout-kind=subgroup` infers the source layout of a one-to-many vector.shape_cast by collapsing the result layout's sg_layout and sg_data independently (collapseDims). When the result layout distributes an inner dimension while each subgroup's data tile spans an outer dimension, the per-subgroup element set is non-contiguous in the flattened source, but the inferred 1-D layout is contiguous. No convert_layout is inserted and everything verifies, so the two sides of a data-movement-free reshape disagree about which elements a subgroup owns.
```mlir
gpu.module @test {
gpu.func @shape_cast_expand(%dst: memref<2x4xf32>) kernel {
%src = arith.constant dense<0.0> : vector<8xf32>
%result = vector.shape_cast %src : vector<8xf32> to vector<2x4xf32>
%td = xegpu.create_nd_tdesc %dst : memref<2x4xf32> -> !xegpu.tensor_desc<2x4xf32>
xegpu.store_nd %result, %td[0, 0]
<{layout = #xegpu.layout}>
: vector<2x4xf32>, !xegpu.tensor_desc<2x4xf32>
gpu.return
}
}
```
```
mlir-opt -xevm-attach-target='chip=cri' -xegpu-propagate-layout='layout-kind=subgroup' repro.mlir
produces on the source:
#xegpu.layout
```
## Expected
The result layout gives subgroup 1 the 2×2 tile based at (0, 2), i.e. elements (0,2),(0,3),(1,2),(1,3) = flattened {2, 3, 6, 7}. A row-major shape_cast moves no data, so the source layout must assign subgroup 1 exactly those elements. sg_layout=[2], sg_data=[2] does (two distribution units, round-robin). The inferred sg_data=[4] assigns it {4, 5, 6, 7} which is two of four elements wrong.
## Configuration
LLVM/MLIR revision: `31a1c03fe67cccd917f11e5ab50cc425770c1400`
Build: `cmake -G Ninja -DLLVM_ENABLE_PROJECTS=mlir -DLLVM_TARGETS_TO_BUILD=Native -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=ON `
Contributor guide
Research direction
Start by running the provided mlir-opt command on repro.mlir and inspect the xegpu-propagate-layout path for one-to-many vector.shape_cast operations, especially collapseDims. Trace how sg_layout and sg_data are inferred from the result layout. Done means the source layout preserves subgroup 1's elements {2, 3, 6, 7}, rather than inferring the contiguous {4, 5, 6, 7} ownership.
Written by the indexing model from the issue text.
Assessment
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 42/100