[Mlir-commits] [mlir] [MLIR][XeGPU] Prefer the nearer consumer's layout in LayoutInfo::meet (PR #208365)
Jianhui Li
llvmlistbot at llvm.org
Wed Jul 8 19:06:10 PDT 2026
https://github.com/Jianhui-Li created https://github.com/llvm/llvm-project/pull/208365
**Summary**
During backward layout propagation a value can be demanded by multiple consumers with conflicting
layouts. Previously LayoutInfo::meet kept whichever demand was assigned first, so the winner depended on
the dataflow worklist's visitation order — nondeterministic and often not what we want.
This changes meet to prefer the layout demanded by the consumer nearer to the producer in program order.
This tends to preserve a consumer's layout as far up the def chain as possible, reducing layout
conversions. It is a heuristic hint, not a proven optimum.
- Each op gets a program-order index from a pre-order walk (matching printed-IR order), so a use inside
an scf.for body is nearer than a use after the loop.
- LayoutInfo carries the demanding op's index in a new programOrder field. visitOperation sets a
file-scoped currentProgramOrder, and the single-arg LayoutInfo constructor stamps it automatically.
- meet keeps the smaller programOrder; ties keep lhs.
- Removes a dead LayoutInfo::transpose declaration.
- Adds truncf_prefers_nearer_user in propagate-layout-inst-data.mlir.
- Updates the multiple-use tests in propagate-layout.mlir (scatter_ops_preserve_load_perm_layout,
binary_op_multiple_uses, if_multiple_uses) to reflect that the nearer consumer's layout now wins.
- a minor layout fix for simple_mxfp_gemm.mlir
>From c589918d9f4111e049d6165a56fb7fc2bfa6d344 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 8 Jul 2026 23:30:49 +0000
Subject: [PATCH 1/4] [MLIR][XeGPU] Prefer slice layout in LayoutInfo::meet
During backward layout propagation a value can receive different
candidate layouts from its consumers. Previously LayoutInfo::meet kept
whichever layout was assigned first, so a value feeding both a
slice-producing consumer (e.g. a broadcast whose inferred source layout
is a slice) and a plain-layout consumer (e.g. a dpas_mx scale operand)
could end up with the plain layout depending on visitation order.
meet() now prefers the slice layout when both sides are assigned
(falling back to lhs when both/neither are slices). For that preference
to take effect, operator== compares the stored layout (via isEqualTo)
rather than only the assigned bit; otherwise the dataflow framework
treats the refined result as a no-op and discards it. The refinement
only moves a plain layout toward a slice layout, so it stays monotonic.
Adds a propagate-layout-inst-data test and updates the a-operand
lane_data in the simple_mxfp_gemm integration test to match the HW
constraint.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 25 +++++++++++-----
.../XeGPU/propagate-layout-inst-data.mlir | 29 +++++++++++++++++++
.../Dialect/XeGPU/WG/simple_mxfp_gemm.mlir | 2 +-
3 files changed, 47 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 64d0d8063b7ff..ff9239b637a10 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -71,12 +71,13 @@ 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 side carrying a slice layout.
+/// If both (or neither) are slice layouts, prefer the lhs (current
+/// state) so an already assigned unique layout is not changed.
struct LayoutInfo {
private:
@@ -86,10 +87,14 @@ struct LayoutInfo {
LayoutInfo() = default;
LayoutInfo(const xegpu::DistributeLayoutAttr &layout) : storage(layout) {}
- // 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.
bool operator==(const LayoutInfo &other) const {
- return this->isAssigned() == other.isAssigned();
+ if (isAssigned() != other.isAssigned())
+ return false;
+ if (!isAssigned())
+ return true;
+ return storage.isEqualTo(other.storage);
}
static LayoutInfo meet(const LayoutInfo &lhs, const LayoutInfo &rhs);
@@ -141,6 +146,10 @@ void LayoutInfo::print(raw_ostream &os) const {
LayoutInfo LayoutInfo::meet(const LayoutInfo &lhs, const LayoutInfo &rhs) {
if (!lhs.isAssigned())
return rhs;
+ if (!rhs.isAssigned())
+ return lhs;
+ if (!lhs.isSliceLayout() && rhs.isSliceLayout())
+ return rhs;
return lhs;
}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 515c59db72819..6bf7fff810006 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -653,3 +653,32 @@ func.func @complete_dpas_mx_inst_data(%arg0: vector<16x1024xf8E5M2>, %arg1: vect
return
}
}
+
+// -----
+// A value with two consumers: one back-propagates a plain layout (the store_nd
+// of %trunc), the other back-propagates a slice layout (the broadcast/transpose
+// chain feeding the reduction result). `meet` must prefer the slice layout, and
+// the inst_data / lane_layout / lane_data fields must be preserved on it.
+gpu.module @test {
+ // CHECK-LABEL: truncf_prefers_slice
+ // CHECK: %[[TRUNC:.*]] = arith.truncf
+ // CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [4, 8, 4], lane_layout = [4, 1, 4], lane_data = [1, 1, 1], order = [0, 2, 1]>, dims = [0]>}
+ // CHECK-SAME: : vector<32x4xbf16> to vector<32x4xf8E8M0FNU>
+ gpu.func @truncf_prefers_slice(%src: memref<32x128xbf16>, %dst_red: memref<32x128xf8E8M0FNU>,
+ %dst_plain: memref<32x4xf8E8M0FNU>) kernel {
+ %cst = arith.constant dense<0xFF80> : vector<32x4xbf16>
+ %tdesc = xegpu.create_nd_tdesc %src : memref<32x128xbf16> -> !xegpu.tensor_desc<32x128xbf16>
+ %load = xegpu.load_nd %tdesc[0, 0] : !xegpu.tensor_desc<32x128xbf16> -> vector<32x128xbf16>
+ %sc1 = vector.shape_cast %load : vector<32x128xbf16> to vector<32x4x32xbf16>
+ %red = vector.multi_reduction <maximumf>, %sc1, %cst [2] : vector<32x4x32xbf16> to vector<32x4xbf16>
+ %trunc = arith.truncf %red : vector<32x4xbf16> to vector<32x4xf8E8M0FNU>
+ %bcast = vector.broadcast %trunc : vector<32x4xf8E8M0FNU> to vector<32x32x4xf8E8M0FNU>
+ %bcast2 = vector.transpose %bcast, [1, 2, 0] : vector<32x32x4xf8E8M0FNU> to vector<32x4x32xf8E8M0FNU>
+ %sc2 = vector.shape_cast %bcast2 : vector<32x4x32xf8E8M0FNU> to vector<32x128xf8E8M0FNU>
+ %tdesc_red = xegpu.create_nd_tdesc %dst_red : memref<32x128xf8E8M0FNU> -> !xegpu.tensor_desc<32x128xf8E8M0FNU>
+ xegpu.store_nd %sc2, %tdesc_red[0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<32x128xf8E8M0FNU>, !xegpu.tensor_desc<32x128xf8E8M0FNU>
+ %tdesc_plain = xegpu.create_nd_tdesc %dst_plain : memref<32x4xf8E8M0FNU> -> !xegpu.tensor_desc<32x4xf8E8M0FNU>
+ xegpu.store_nd %trunc, %tdesc_plain[0, 0] <{layout = #xegpu.layout<inst_data = [8, 2], lane_layout = [8, 1], lane_data = [1, 1]>}> : vector<32x4xf8E8M0FNU>, !xegpu.tensor_desc<32x4xf8E8M0FNU>
+ gpu.return
+ }
+}
diff --git a/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm.mlir b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm.mlir
index f584ede003bd4..9bfb783299055 100644
--- a/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm.mlir
+++ b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm.mlir
@@ -8,7 +8,7 @@
// XFAIL: *
// Note: layouts used by dpas_mx need to match HW constaint. Otherwise dpas_mx is not unrolled.
-#a = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 1024], inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 1]>
+#a = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 1024], inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>
#b_packed = #xegpu.layout<sg_layout = [2, 2], sg_data = [512, 16], inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>
#b = #xegpu.layout<sg_layout = [2, 2], sg_data = [1024, 16], inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>
#c = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 16], inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
>From c31ad5f6b4f3d97c50c91e62737b39cd07a23ecf Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 9 Jul 2026 01:55:14 +0000
Subject: [PATCH 2/4] [MLIR][XeGPU] Prefer nearer consumer's layout in
LayoutInfo::meet
When a value is demanded by multiple consumers during backward layout
propagation, prefer the layout of the consumer that is nearer to the
producer in program order. This tends to preserve a consumer's layout as
far up the def chain as possible, reducing layout conversions. It is a
hint, not an optimum.
Each op is assigned a program-order index via a pre-order walk (matching
printed-IR order), so a use inside an scf.for body is nearer than a use
after the loop. LayoutInfo carries the demanding op's index in a new
`programOrder` field, stamped by the single-argument constructor from a
file-scoped `currentProgramOrder` that visitOperation sets, so the ~30
meet() call sites are unchanged. `programOrder` is not propagated (each
visited op stamps its own index) and is excluded from operator==.
This replaces the earlier slice-layout preference in meet(), which is
now redundant: distinct users always have distinct indices, so program
order decides every real conflict. Also removes the dead
LayoutInfo::transpose declaration.
Updates the multiple-uses propagate-layout tests to reflect that the
nearer consumer's layout now wins.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 75 +++++++++++++++++--
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 20 ++---
2 files changed, 77 insertions(+), 18 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index ff9239b637a10..72b10e1d3265f 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -30,6 +30,7 @@
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/SmallVector.h"
@@ -38,6 +39,7 @@
#include "llvm/Support/Debug.h"
#include "llvm/Support/LogicalResult.h"
#include "llvm/Support/raw_ostream.h"
+#include <limits>
namespace mlir {
namespace xegpu {
@@ -58,6 +60,14 @@ namespace {
// LayoutInfo
//===----------------------------------------------------------------------===//
+/// Program-order index of the op currently being visited by the backward
+/// analysis. `visitOperation` sets this before dispatching, and the
+/// single-argument LayoutInfo constructor stamps it onto every demand pushed to
+/// an operand, so the ~30 `operand->meet(LayoutInfo(...))` call sites need no
+/// change. A larger index means farther from the producer; the sentinel max
+/// marks demands with no associated op (e.g. exit state).
+static int64_t currentProgramOrder = std::numeric_limits<int64_t>::max();
+
/// Helper class for tracking the analysis state of an mlir value. For layout
/// propagation, the analysis state is simply the distribution layout of
/// each value. The distribution layout information is encapsulated using
@@ -75,20 +85,34 @@ namespace {
/// assigned with the same layout.
/// 3) The meet operator works as follows:
/// - If only one side is assigned, return that side.
-/// - If both sides are assigned, prefer the side carrying a slice layout.
-/// If both (or neither) are slice layouts, prefer the lhs (current
-/// state) so an already assigned unique layout is not changed.
+/// - 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 are both unassigned, or both assigned
- // with the same layout.
+ // 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 {
if (isAssigned() != other.isAssigned())
return false;
@@ -105,8 +129,6 @@ struct LayoutInfo {
bool isAssigned() const { return storage != nullptr; }
- LayoutInfo transpose(ArrayRef<int64_t> permutation) const;
-
SmallVector<int> getLaneLayout() const;
SmallVector<int> getLaneData() const;
@@ -135,6 +157,11 @@ struct LayoutInfo {
void set(const xegpu::DistributeLayoutAttr &layout) { storage = layout; }
};
+// Stamp every demand pushed by the current op with that op's program-order
+// index so `meet` can prefer the nearest consumer.
+LayoutInfo::LayoutInfo(const xegpu::DistributeLayoutAttr &layout)
+ : storage(layout), programOrder(currentProgramOrder) {}
+
void LayoutInfo::print(raw_ostream &os) const {
if (isAssigned()) {
os << storage;
@@ -148,7 +175,10 @@ LayoutInfo LayoutInfo::meet(const LayoutInfo &lhs, const LayoutInfo &rhs) {
return rhs;
if (!rhs.isAssigned())
return lhs;
- if (!lhs.isSliceLayout() && rhs.isSliceLayout())
+ // Prefer the demand from the user nearer to the producer in program order.
+ // Distinct users always have distinct indices, so this decides every
+ // real conflict; on a tie (same op, or both unknown) keep lhs.
+ if (rhs.programOrder < lhs.programOrder)
return rhs;
return lhs;
}
@@ -185,6 +215,15 @@ class LayoutInfoPropagation
private:
xegpu::LayoutKind layoutKind;
unsigned indexBitWidth;
+
+ // Program-order index of every op, built lazily on first use via a pre-order
+ // walk of the top-level module/function (matching printed-IR order). Used to
+ // tell which consumer of a value is nearer to its producer.
+ DenseMap<Operation *, int64_t> programOrder;
+ // Returns the program-order index of `op`, populating `programOrder` from
+ // `op`'s top-level ancestor on first call.
+ int64_t getProgramOrder(Operation *op);
+
void visitDpasOp(xegpu::DpasOp dpas, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results);
@@ -300,9 +339,29 @@ class LayoutInfoPropagation
};
} // namespace
+int64_t LayoutInfoPropagation::getProgramOrder(Operation *op) {
+ auto it = programOrder.find(op);
+ if (it != programOrder.end())
+ return it->second;
+ // First time we see this op's tree: number every op under its top-level
+ // ancestor in pre-order (i.e. printed-IR order). Nested ops (e.g. inside an
+ // scf.for body) get an index between their parent and the parent's next
+ // sibling, so a use inside a loop is "nearer" than a use after it.
+ Operation *root = op;
+ while (root->getParentOp())
+ root = root->getParentOp();
+ int64_t counter = 0;
+ root->walk<WalkOrder::PreOrder>(
+ [&](Operation *o) { programOrder[o] = counter++; });
+ return programOrder.lookup(op);
+}
+
LogicalResult LayoutInfoPropagation::visitOperation(
Operation *op, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) {
+ // Stamp demands pushed by this op with its program-order index so `meet` can
+ // prefer the nearest consumer.
+ currentProgramOrder = getProgramOrder(op);
TypeSwitch<Operation *>(op)
.Case(
[&](xegpu::DpasOp dpasOp) { visitDpasOp(dpasOp, operands, results); })
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index 1a741ba21e128..49006e1ef9507 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -219,8 +219,8 @@ func.func @scatter_ops_custom_perm_layout(%src: memref<256xf16>) {
gpu.module @test {
// CHECK-LABEL: func.func @scatter_ops_preserve_load_perm_layout(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<256xf16>) {
-// CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [8], lane_data = [1]>} dense<true> : vector<16xi1>
-// CHECK: %[[OFFSETS:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [8], lane_data = [1]>} dense<12> : vector<16xindex>
+// CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<16xi1>
+// CHECK: %[[OFFSETS:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<12> : vector<16xindex>
// CHECK: %[[LOAD_VEC:.*]] = xegpu.load %[[ARG0]][%[[OFFSETS]]], %[[MASK]]
// CHECK-SAME: memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16xf16>
// CHECK: %[[ADD_RES:.*]] = arith.addf %[[LOAD_VEC]], %[[LOAD_VEC]] {layout_result_0 = #xegpu.layout<lane_layout = [8], lane_data = [1]>} : vector<16xf16>
@@ -341,10 +341,10 @@ func.func @binary_op_one_use(%arg0: !xegpu.tensor_desc<8x16xf16>, %arg1: !xegpu.
gpu.module @test {
// CHECK-LABEL: func.func @binary_op_multiple_uses(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>,
-// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>,
+// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>>,
// CHECK-SAME: %[[ARG2:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>,
// CHECK-SAME: %[[ARG3:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>) {
-// CHECK: %[[T2:.*]] = arith.addf %{{.*}}, %{{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : vector<16x16xf16>
+// CHECK: %[[T2:.*]] = arith.addf %{{.*}}, %{{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>} : vector<16x16xf16>
// CHECK: %[[T3:.*]] = xegpu.dpas %{{.*}}, %[[T2]] {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : vector<8x16xf16>, vector<16x16xf16> -> vector<8x16xf32>
// CHECK-NEXT: xegpu.store_nd %[[T3]], %[[ARG2]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK-NEXT: xegpu.store_nd %[[T2]], %[[ARG3]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<16x16xf16>, !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
@@ -429,18 +429,18 @@ func.func @if_single_use(%arg0: !xegpu.tensor_desc<8x16xf16>, %arg1: !xegpu.tens
gpu.module @test {
// CHECK-LABEL: func.func @if_multiple_uses(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>,
-// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>,
+// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>>,
// CHECK-SAME: %[[ARG2:[0-9a-zA-Z]+]]: i1, %[[ARG3:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>,
// CHECK-SAME: %[[ARG4:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>) {
// CHECK: %[[T1:.*]] = scf.if %[[ARG2]] -> (vector<16x16xf16>) {
-// CHECK-NEXT: %[[T3:.*]] = xegpu.load_nd %[[ARG1]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<16x16xf16>
+// CHECK-NEXT: %[[T3:.*]] = xegpu.load_nd %[[ARG1]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<16x16xf16>
// CHECK-NEXT: scf.yield %[[T3]] : vector<16x16xf16>
// CHECK-NEXT: } else {
-// CHECK-NEXT: %[[T4:.*]] = xegpu.load_nd %[[ARG1]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> :
-// CHECK-SAME: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<16x16xf16>
+// CHECK-NEXT: %[[T4:.*]] = xegpu.load_nd %[[ARG1]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}> :
+// CHECK-SAME: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<16x16xf16>
// CHECK-NEXT: scf.yield %[[T4]] : vector<16x16xf16>
-// CHECK-NEXT: } {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+// CHECK-NEXT: } {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
func.func @if_multiple_uses(%arg0: !xegpu.tensor_desc<8x16xf16>, %arg1: !xegpu.tensor_desc<16x16xf16>, %arg2: i1, %arg3: !xegpu.tensor_desc<8x16xf32>, %arg4: !xegpu.tensor_desc<16x16xf16>) {
%0 = xegpu.load_nd %arg0[0, 0] : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
%1 = scf.if %arg2 -> (vector<16x16xf16>) {
>From 3c4ada8b1f50d0d9c48f3c8543aa0662f1f2164c Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 9 Jul 2026 01:58:24 +0000
Subject: [PATCH 3/4] [MLIR][XeGPU] Reframe propagate-layout test as
nearer-user preference
Rename truncf_prefers_slice to truncf_prefers_nearer_user and update its
comment: on this branch %trunc keeps the broadcast chain's slice layout
because that consumer is nearer in program order, not because meet()
prefers slice layouts.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../Dialect/XeGPU/propagate-layout-inst-data.mlir | 12 ++++++------
1 file changed, 6 insertions(+), 6 deletions(-)
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 6bf7fff810006..47471b890d552 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -655,16 +655,16 @@ func.func @complete_dpas_mx_inst_data(%arg0: vector<16x1024xf8E5M2>, %arg1: vect
}
// -----
-// A value with two consumers: one back-propagates a plain layout (the store_nd
-// of %trunc), the other back-propagates a slice layout (the broadcast/transpose
-// chain feeding the reduction result). `meet` must prefer the slice layout, and
-// the inst_data / lane_layout / lane_data fields must be preserved on it.
+// %trunc has two consumers: the broadcast/transpose chain (nearer in program
+// order, back-propagates a slice layout) and the store_nd of %trunc (farther,
+// back-propagates a plain layout). `meet` must keep the nearer consumer's
+// layout, preserving its inst_data / lane_layout / lane_data fields.
gpu.module @test {
- // CHECK-LABEL: truncf_prefers_slice
+ // CHECK-LABEL: truncf_prefers_nearer_user
// CHECK: %[[TRUNC:.*]] = arith.truncf
// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [4, 8, 4], lane_layout = [4, 1, 4], lane_data = [1, 1, 1], order = [0, 2, 1]>, dims = [0]>}
// CHECK-SAME: : vector<32x4xbf16> to vector<32x4xf8E8M0FNU>
- gpu.func @truncf_prefers_slice(%src: memref<32x128xbf16>, %dst_red: memref<32x128xf8E8M0FNU>,
+ gpu.func @truncf_prefers_nearer_user(%src: memref<32x128xbf16>, %dst_red: memref<32x128xf8E8M0FNU>,
%dst_plain: memref<32x4xf8E8M0FNU>) kernel {
%cst = arith.constant dense<0xFF80> : vector<32x4xbf16>
%tdesc = xegpu.create_nd_tdesc %src : memref<32x128xbf16> -> !xegpu.tensor_desc<32x128xbf16>
>From 9c4de99b6094f502c82dcd411b8eca2c1c18141d Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 9 Jul 2026 02:00:09 +0000
Subject: [PATCH 4/4] git format
---
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 14 +++++++-------
1 file changed, 7 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 72b10e1d3265f..75e8db468f361 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -320,18 +320,18 @@ class LayoutInfoPropagation
visitOperation(Operation *op, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results) override;
- void visitBranchOperand(OpOperand &operand) override {};
+ void visitBranchOperand(OpOperand &operand) override{};
- void visitCallOperand(OpOperand &operand) override {};
+ void visitCallOperand(OpOperand &operand) override{};
void
visitNonControlFlowArguments(RegionSuccessor &successor,
- ArrayRef<BlockArgument> arguments) override {};
+ ArrayRef<BlockArgument> arguments) override{};
- void visitExternalCall(CallOpInterface call,
- ArrayRef<LayoutInfoLattice *> operands,
- ArrayRef<const LayoutInfoLattice *> results) override {
- };
+ void
+ visitExternalCall(CallOpInterface call,
+ ArrayRef<LayoutInfoLattice *> operands,
+ ArrayRef<const LayoutInfoLattice *> results) override{};
void setToExitState(LayoutInfoLattice *lattice) override {
(void)lattice->meet(LayoutInfo());
More information about the Mlir-commits
mailing list