[Mlir-commits] [mlir] [mlir][XeGPU] Add an SLM round-trip fallback for sg-to-lane convert_layout (PR #227884)

[Sang Ik Lee]

silee2 wrote:

Following up on the open thread about `lane_data` spanning multiple rows - I think the concern is real, but it is not about the number of load/stores. It is that a *single* `store_matrix`/`load_matrix` is asked to move a block that is **not contiguous** in the row-major scratch tile.

The pattern creates the `mem_desc` with a null `mem_layout` (row-major), and `LoadStoreMatrixToXeVMPattern` lowers these ops to one flat LLVM access at the linearized *start* offset, with no per-dimension stride walk. Minimal demonstration:

```mlir
gpu.module @xevm_module {
gpu.func @tiny(%a: memref<64xi8, 3>, %v: vector<2x1xf32>) kernel {
  %c1 = arith.constant 1 : index
  %c2 = arith.constant 2 : index
  %md = xegpu.create_mem_desc %a : memref<64xi8, 3> -> !xegpu.mem_desc<4x4xf32>
  xegpu.store_matrix %v, %md[%c1, %c2] : vector<2x1xf32>, !xegpu.mem_desc<4x4xf32>, index, index
  gpu.return
}
}
```

```sh
mlir-opt --xevm-attach-target='module=xevm_* chip=cri' tiny.mlir \
| mlir-opt --pass-pipeline="builtin.module(gpu.module(convert-xegpu-to-xevm,canonicalize))"
```

```mlir
%2 = arith.addi %1, %c24_i32 : i32              // byte 24 = element 6 = (1,2)
llvm.store %0, %3 : vector<2xf32>, !llvm.ptr<3> // writes (1,2) and (1,3)
```

A `2x1` block at `[1,2]` is elements `(1,2)` and `(2,2)`, i.e. linear 6 and 10. The lowering writes 6 and 7.

Both directions are reachable from the new pattern:

```mlir
// store side mis-addressed
input_layout  = #xegpu.layout<lane_layout = [16, 1], lane_data = [2, 1]>
target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>   // vector<32x16xf32>
// -> xegpu.store_matrix ... : vector<2x1xf32>, !xegpu.mem_desc<128x16xf32>
// -> llvm.store ... : vector<2xf32>

// load side mis-addressed
input_layout  = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
target_layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [2, 1]>   // vector<32x16xf32>
// -> xegpu.load_matrix ... -> vector<2x1xf32>
// -> llvm.load ... -> vector<2xf32>
```

I checked that no specialized pattern claims these: lane layouts differ and it is not a pure outer-dim shrink, so the `gpu.shuffle` path declines, and `lane_data` differs, so the `lane_shuffle` repack path declines. The SLM fallback is the one that fires.

A payload with two non-unit dims (e.g. `target_layout` `lane_data = [2, 32]`) hits an assert instead of wrong code:

```
XeGPUToXeVM.cpp:933: Assertion `llvm::count_if(vecType.getShape(), [](int64_t d) { return d != 1; }) <= 1
  && "Expected either 1D vector or nD with unit dimensions"' failed.
```

## Attribution

To be clear, **this is not caused by the PR**. The root cause is in `LoadStoreMatrixToXeVMPattern` (`mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp:912`), which discards the payload shape (`:932-938`) and emits a single flat `llvm.load`/`llvm.store` (`:1009-1022`), so it implicitly assumes the block occupies `getNumElements()` consecutive elements. `MemDescType::getLinearOffsets` is fine; the bug is in the consumer, which never reuses the strides for the remaining elements. The pre-existing `SgToLaneLoadMatrix` already emits the same op shape and is already mislowered, so this is independently reachable today and deserves its own fix.

Worth noting the dialect already spells out this precondition for `subgroup_block_io` in `IsValidMatrixOpParams` (`XeGPUOps.cpp:150`): `isLaneDataContiguous`, "accessed data must be contiguous and coalesced". It was just never applied to the plain path, which is the path this pattern uses.

## Why it still matters for this PR

I verified the before/after on the `lane_data = [2,1]` case above. At baseline (PR reverted) it is `error: failed to legalize operation 'xegpu.convert_layout'` - compilation stops. With the PR the fallback fires and the result is silently wrong data. So a clean compile error becomes a silent miscompile, and `lane_data = [2,1]` is the ordinary packed/VNNI B-operand shape rather than an exotic input.

Since the pattern cannot fix the addressing itself, two options:

1. **Guard and decline**, reusing the dialect's own `isLaneDataContiguous` test on both layouts (i.e. `lane_data` must be `[1, ..., 1, k]`). That restores the pre-PR behavior for this class, and the guard can be deleted when `LoadStoreMatrixToXeVMPattern` is fixed. I tried this locally: the three reachable bad cases decline and the two new tests plus the full `mlir/test/Dialect/XeGPU/` suite still pass.

2. **Emit only contiguous transfers** - decompose each `lane_data` block into its innermost contiguous runs, one `store_matrix`/`load_matrix` per run instead of one per distribution unit. For `lane_data = [2,1]` that is two `vector<1x1xf32>` transfers at `[r,c]` and `[r+1,c]`. This is correct under the current lowering, needs no `XeGPUToXeVM` change, and also avoids the assert above. More ops, but probably acceptable for a last-resort fallback that already pays for an SLM round-trip.

Either way, could the PR description state the limit explicitly? As written, the "any pair of layouts which can both distribute the value" claim holds only for `lane_data` of the form `[1, ..., 1, k]`, which is what both new tests use.


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