[Mlir-commits] [mlir] [MLIR][XeGPU] Add layout propagation for `xegpu.store_matrix` (PR #174952)

[Jianhui Li]

================
@@ -1100,6 +1107,28 @@ void LayoutInfoPropagation::visitStoreScatterOp(
     propagateIfChanged(operands[3], operands[3]->meet(maskLayout));
 }
 
+void LayoutInfoPropagation::visitStoreMatrixOp(
+    xegpu::StoreMatrixOp storeMatrix, ArrayRef<LayoutInfoLattice *> operands,
+    ArrayRef<const LayoutInfoLattice *> results) {
+  Value operand = storeMatrix.getData();
+  unsigned index =
+      std::distance(storeMatrix.operand_begin(),
+                    llvm::find(storeMatrix->getOperands(), operand));
+
+  auto uArch = getUArch(getChipStr(storeMatrix).value_or(""));
+  const int subgroupSize = uArch->getSubgroupSize();
+  SmallVector<int> instData = {1, 8};
----------------
Jianhui-Li wrote:

>There is a [test that uses [sg_size, 1]](https://github.com/llvm/llvm-project/blob/a927d37f854468d76f4e3f664d5720128b263932/mlir/test/Dialect/XeGPU/subgroup-distribute.mlir#L335) layout for store_matrix.
There is a [test contradicting](https://github.com/llvm/llvm-project/blob/a927d37f854468d76f4e3f664d5720128b263932/mlir/test/Dialect/XeGPU/subgroup-distribute.mlir#L318) the notion of a single 32-bit dist unit for store_matrix (lane data is [2,1] for f32), this is not a valid lane_data for scattered stores.
nit: Both test actually needs some fixes although not the scope of this PR. The first one uses subgroup io but the 2d tile is not contiguous as required by the lowering target store1d. The second one use lane_data [2, 1], but 2 elements in the first dim are not contiguous which is required by chunked store (the lowering target).  

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