[Mlir-commits] [mlir] [MLIR][XeGPU] Refactor layout propagation utilities (PR #179016)

[Jianhui Li]

================
@@ -646,88 +648,89 @@ void LayoutInfoPropagation::visitVectorMultiReductionOp(
     ArrayRef<LayoutInfoLattice *> operands,
     ArrayRef<const LayoutInfoLattice *> results) {
   // The layout of the result must be present.
-  LayoutInfo resultLayout = results[0]->getValue();
-  if (!resultLayout.isAssigned())
+  LayoutInfo resLayoutInfo = results[0]->getValue();
+  if (!resLayoutInfo.isAssigned())
     return;
-  // We only consider 2D -> 1D reductions at this point.
-  VectorType resultTy = llvm::dyn_cast<VectorType>(reduction.getDestType());
-  if (!resultTy || resultTy.getRank() != 1) {
-    reduction.emitWarning("Expecting output type to be 1D vector.");
-    return;
-  }
+
+  VectorType sourceTy =
+      llvm::dyn_cast<VectorType>(reduction.getSourceVectorType());
+  SmallVector<int64_t> reductionDims(reduction.getReductionDims().begin(),
+                                     reduction.getReductionDims().end());
+
   auto uArch = getUArch(xegpu::getChipStr(reduction).value_or(""));
-  // Given that the result is 1D, the layout of the operand should be 2D with
-  // default layout.
-  LayoutInfo operandLayout = getDefaultSIMTLayoutInfo(
-      reduction->getContext(), 2, uArch->getSubgroupSize());
-  propagateIfChanged(operands[0], operands[0]->meet(operandLayout));
+  auto consumerLayoutAttr =
+      dyn_cast<xegpu::DistributeLayoutAttr>(resLayoutInfo.get());
+
+  // The result layout represents the layout requirements of the operation.
+  // it is recorded to anchor layout or temporary layout.
+  // it must be honored for current op and may conflict with the layout
+  // propagated from consumer op, the conflict is resolved in later phase by
+  // converting the required result layout to the consumer layout
+  auto requiredResLayoutAttr = xegpu::setupMultiReductionResultLayout(
+      layoutKind, sourceTy, consumerLayoutAttr, reductionDims, uArch);
+  // resLayoutInfo.set(requiredResLayoutAttr);
+  xegpu::setTemporaryLayout(reduction->getResult(0), requiredResLayoutAttr);
----------------
Jianhui-Li wrote:

Setting temp layout attribute here doesn't impact the data flow analysis.  And the temp layout attribute only lives within the propagation pass and so exclusively used by the pass not interfered by other uses. Why you think it can be a source of bugs? 

I am against keeping a separate map but I don't think the current implementation is problematic. 

The second approach (let the analysis running naturally first) is interesting, how the second pass retrieve the required result layout?  

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