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

[Artem Kroviakov]

================
@@ -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};
----------------
akroviakov wrote:

This was a placeholder until inst_data for scattered ops https://github.com/llvm/llvm-project/pull/172845 is clarified.
 
The question is the following. Since the store_matrix lowers to scattered ops, it seems to be meaningful to apply [the same restrictions](https://github.com/llvm/llvm-project/blob/be50b78afea0ffd3cfba4a2106fb312fd39bc694/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp#L313) for layout/inst_data. But for matrix ops, we assume 2D payload, so it would not be [1, sg_size], but [sg_size, packed_elements]. 

I see 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. 

I suppose the flexibility of store_matrix allows it to act similar to a scattered store with the default layout, so the propagation should be fine with these conservative assumptions for default layout (i.e., 2D scattered store = store matrix)? 

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