[Mlir-commits] [mlir] [mlir][xegpu] Fix operand layouts when rematerializing on layout conflict (PR #224730)

[Jianhui Li]

================
@@ -1652,6 +1652,36 @@ LogicalResult ResolveLayoutConflicts::assignResultLayout(OpResult &result) {
   return success();
 }
 
+// Recursively rematerialize a trivially-rematerializable producer subtree so
+// that the clone yields `layout`. Cloning only the producer would leave its
+// operands pointing at values that carry the original (conflicting) layout,
+// making the clone's operand and result layouts disagree. Stamp `layout` on the
+// clone and rematerialize each vector operand with the layout back-derived from
+// it.
+static Value rematerializeWithLayout(OpBuilder &builder, Value value,
+                                     xegpu::DistributeLayoutAttr layout) {
+  Operation *producerOp = value.getDefiningOp();
+  builder.setInsertionPointAfter(producerOp);
+  Operation *clone = builder.clone(*producerOp);
+  OpResult cloneResult = clone->getResult(0);
+  // Drop the inherited producer layout so the new layout takes effect.
----------------
Jianhui-Li wrote:

what do you mean "the new layout takes effect"? These layout on these rematerializable ops are temproary layout will be recovered in the beginning of distribution process anyway. 

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