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

llvmlistbot at llvm.org llvmlistbot at llvm.org
Wed Jul 8 19:06:49 PDT 2026


llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Jianhui Li (Jianhui-Li)

<details>
<summary>Changes</summary>

 **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

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


4 Files Affected:

- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp (+86-18) 
- (modified) mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir (+29) 
- (modified) mlir/test/Dialect/XeGPU/propagate-layout.mlir (+10-10) 
- (modified) mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm.mlir (+1-1) 


``````````diff
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 64d0d8063b7ff..75e8db468f361 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
@@ -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())
+      return true;
+    return storage.isEqualTo(other.storage);
   }
 
   static LayoutInfo meet(const LayoutInfo &lhs, const LayoutInfo &rhs);
@@ -100,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;
@@ -130,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;
@@ -141,6 +173,13 @@ 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;
+  // 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;
 }
 
@@ -176,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);
 
@@ -272,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());
@@ -291,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-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 515c59db72819..47471b890d552 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
 }
 }
+
+// -----
+// %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_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_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>
+    %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/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>) {
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]>

``````````

</details>


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


More information about the Mlir-commits mailing list