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

[Jianhui Li]

================
@@ -616,16 +636,34 @@ struct GatherLowering : public OpRewritePattern<vector::GatherOp> {
         computeOffsets(rewriter, gatherOp, meta.first, meta.second);
     Value flatMemref = memrefToIndexPtr(gatherOp, rewriter);
 
+    auto layoutRes = xegpu::getDistributeLayoutAttr(gatherOp.getResult());
----------------
Jianhui-Li wrote:

At the entry XeGPU IR, we need to set the layout to anchor op, and rely on the propagation pass to propagate the layout to all values. The anchor op include: load/store, load_matrix/store_matrix, load_nd/store_nd, and dpas, convert_layout. The layout describes the layout for the tensor tile: for memory operation, it describes the memory operand (like memref, mem_desc, nd_tdesc), for dpas, it describes the vector operands. 

The layout propagated are temporarily attributes and can be lost during IR transformation or lowering. But the anchor op defined layout attributes are permanents and we rely on these attributes being set up properly so able to recover in case user compose xegpu passes with their own passes. 

The example you show above actually looks fine to me, as long as the propogation rules decides to drop one dimension for the mask and index and keep the tensor tile layout as 2d.  It does look a bit complex, but I don't expect this kind of layout is exposed to user since chunkload should be lower-level representation for memory coalesce or load_matrix lowering. 

So I still suggest changing this PR and just set the layout attribute according to the tensor tile's layout for anchor ops only and remove the attributes setting for other ops and operands. 

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