[Mlir-commits] [mlir] [mlir][XeGPU][VectorToXeGPU] Propagate vector layouts to xegpu ops (PR #163071)

[Dmitry Chigarev]

================
@@ -374,22 +388,30 @@ static Value computeOffsets(PatternRewriter &rewriter, OpType gatScatOp,
         arith::AddIOp::create(rewriter, loc, baseOffset, offsetContrib);
   }
   Value indices = gatScatOp.getIndices();
+  // Extract indices layout and propagate it to all 'vector' ops created here
+  auto indicesLayout = mlir::xegpu::getDistributeLayoutAttr(indices);
   VectorType vecType = cast<VectorType>(indices.getType());
----------------
dchigarev wrote:

I'm wondering whether I understand the idea of the `-xegpu-propagate-layout` pass correctly.

Based on your comment and on my initial impression (and on [this docstring)](https://github.com/llvm/llvm-project/blob/83f751a001fa57dd8eaa60f08ed420e90170e11d/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp#L849-L850), the pass should propagate user's layouts or set a default one, for example:
```mlir
// if I set a custom layout for `xegpu.store`, I would expect the same layout
// to propagate to the result of `arith.select` (producer of the operand_0)
%res = arith.select %mask, %other, %3 : vector<16x16xi1>, vector<16x16xf16>
xegpu.store %res, %src[%offset], %1 {
    layout_operand_0 = #xegpu.layout<lane_layout = [4, 4], lane_data = [1, 2]>}
    : vector<16x16xf16>, memref<256xf16>, vector<16x16xindex>, vector<16x16xi1>

// however in reality it applies "defaultSIMTLayout" ignoring my custom layout:
%res = arith.select %mask, %other, %3 {
    layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>}
    : vector<16x16xi1>, vector<16x16xf16>
xegpu.store %res, %src[%offset], %1  {
    layout_operand_0 = #xegpu.layout<lane_layout = [4, 4], lane_data = [1, 2]>} 
    : vector<16x16xf16>, memref<256xf16>, vector<16x16xindex>, vector<16x16xi1>
```
Based on the logic from the pass, it always applies default-simt layout in case of results/producers for [store_scatter](https://github.com/llvm/llvm-project/blob/83f751a001fa57dd8eaa60f08ed420e90170e11d/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp#L745-L746) (and probably for load_gather as well). Is this intended behavior or under "we should use the propagation pass" you meant, that we should improve the pass to also consider custom user layouts?

cc @charithaintc 

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