[Mlir-commits] [mlir] [MLIR][XeGPU] Prefer the nearer consumer's layout in LayoutInfo::meet (PR #208365)

[Artem Kroviakov]

================
@@ -71,25 +81,44 @@ namespace {
 /// Given this, LayoutInfo  satisifies the following properties:
 ///  1) A LayoutInfo value can be in one of two states - `assigned` or `not
 ///  assigned`.
-///  2) Two LayoutInfo values are equal if they are both assigned or
-///  both not assigned. The concrete value of assigned state does not matter.
+///  2) Two LayoutInfo values are equal if they are both not assigned, or both
+///  assigned with the same layout.
 ///  3) The meet operator works as follows:
-///     - If current state is assigned, return the current state. (already
-///     a unique layout is assigned. don't change it)
-///     - Otherwise, return the other state.
+///     - If only one side is assigned, return that side.
+///     - If both sides are assigned, prefer the layout demanded by the user
+///       that is nearer to the producer in program order (smaller
+///       `programOrder`); on a tie keep lhs.
+///
+/// The `programOrder` field records the program-order index of the consumer op
+/// that demanded the layout (see `currentProgramOrder`). During this backward
+/// analysis a value can be demanded by several users; keeping the nearest one
+/// tends to preserve a consumer's layout as far up the def chain as possible,
+/// minimizing layout conversions. This is a hint, not an optimum.
+/// `programOrder` is never propagated up the chain - each visited op stamps its
+/// own index - so it is deliberately excluded from `operator==`.
 
 struct LayoutInfo {
 private:
   xegpu::DistributeLayoutAttr storage = nullptr;
+  // Program-order index of the consumer op that demanded this layout. Smaller
+  // means nearer to the producer. Unassigned/unknown demands sort last.
+  int64_t programOrder = std::numeric_limits<int64_t>::max();
 
 public:
   LayoutInfo() = default;
-  LayoutInfo(const xegpu::DistributeLayoutAttr &layout) : storage(layout) {}
+  LayoutInfo(const xegpu::DistributeLayoutAttr &layout);
+  LayoutInfo(const xegpu::DistributeLayoutAttr &layout, int64_t programOrder)
+      : storage(layout), programOrder(programOrder) {}
 
-  // Two lattice values are equal if they have `some` layout. The actual
-  // content of the layout does not matter.
+  // Two lattice values are equal if they are both unassigned, or both assigned
+  // with the same layout. `programOrder` is intentionally excluded: it is not
+  // propagated, so a pure order refinement must not be reported as a change.
   bool operator==(const LayoutInfo &other) const {
-    return this->isAssigned() == other.isAssigned();
+    if (isAssigned() != other.isAssigned())
+      return false;
+    if (!isAssigned())
----------------
akroviakov wrote:

The comment above does not describe this condition

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