[Mlir-commits] [mlir] [MLIR][XeGPU] Clean up stale convert_layout on single-element vector in peephole (PR #194043)

[Nishant Patel]

================
@@ -449,17 +449,49 @@ class MultiRed2dOpPattern
     auto loc = reductionOp.getLoc();
     auto acc = reductionOp.getAcc();
 
-    // If the result is scalar after reduction, look for consumer
-    // convert_layout op and remove it. The layout propagation pass will
-    // re-install it properly after the decomposition.
-    Type resultType = reductionOp.getResult().getType();
-    if (resultType.isIntOrFloat()) {
-      for (auto &use : reductionOp.getResult().getUses()) {
-        if (auto convertLayoutOp =
-                llvm::dyn_cast<xegpu::ConvertLayoutOp>(use.getOwner())) {
-          rewriter.replaceOp(convertLayoutOp, reductionOp.getResult());
-          break;
-        }
+    // The decomposition below splits the 2D reduction into an intra-lane
+    // then a cross-lane 1D reduction. If a consumer xegpu.convert_layout
+    // exists on the reduction result, its input_layout was stamped by
+    // layout propagation against the original 2D reduction's slice and
+    // is therefore stale once we replace the producer with two 1D
+    // reductions.
+    //
+    // Hence insert a NEW xegpu.convert_layout between the decomposed
+    // reduction result and the existing convert_layout. The new op
+    // bridges from the natural post-decomposition producer layout
+    // to the layout that the existing convert_layout currently expects on its
+    // input. The existing convert_layout is left untouched.
+    xegpu::ConvertLayoutOp consumerConvertOp;
+    for (auto &use : reductionOp.getResult().getUses()) {
+      if (auto convertLayoutOp =
+              llvm::dyn_cast<xegpu::ConvertLayoutOp>(use.getOwner())) {
+        consumerConvertOp = convertLayoutOp;
+        break;
+      }
+    }
+    xegpu::DistributeLayoutAttr postDecompLayout;
+    if (consumerConvertOp) {
+      // Derive the source vector's layout.
+      xegpu::DistributeLayoutAttr srcLayoutForCvt;
+      if (auto resSlice = dyn_cast_if_present<xegpu::SliceAttr>(resLayout))
+        srcLayoutForCvt = resSlice.getParent();
+      if (!srcLayoutForCvt)
+        srcLayoutForCvt =
----------------
nbpatel wrote:

if the consumer is not convert layout why will it even be an issue? because the layout (temporary) gets dropped anyway at the end of the pass and if the consumer has anchor layout...isn't it the duty of the propagation pass after peephole to take care of it?

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