[Mlir-commits] [mlir] [MLIR][XeGPU] Clone trivial operations with multiple consumers in layout propagation to avoid layout conflict (PR #197514)
Nishant Patel
llvmlistbot at llvm.org
Tue May 26 09:13:02 PDT 2026
https://github.com/nbpatel updated https://github.com/llvm/llvm-project/pull/197514
>From 0abe6110ae3bd7f935fe5bb87273b47205d389cf Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Fri, 1 May 2026 21:19:55 +0000
Subject: [PATCH 1/6] Fix recoverTemporaryLayouts
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 78 +++++++++++++++++++
1 file changed, 78 insertions(+)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index f91e80823c2e9..62caffb136998 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -25,7 +25,9 @@
#include "mlir/IR/ValueRange.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Transforms/DialectConversion.h"
+#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/Support/FormatVariadic.h"
#include <cstdint>
@@ -126,12 +128,88 @@ static xegpu::DistributeLayoutAttr getLayoutFromUsePoints(Value result) {
return layout;
}
+// Returns true if `op` is safe and cheap to clone (no side effects, no
+// regions, and all operands are themselves trivially rematerializable, e.g.
+// block-arg-free pure value generators such as `vector.step`, splat
+// `arith.constant`, or `vector.create_mask` whose operands are constants).
+static bool isTriviallyRematerializable(Operation *op) {
+ if (!op || op->getNumRegions() != 0)
+ return false;
+ if (!isMemoryEffectFree(op))
+ return false;
+ for (Value v : op->getOperands()) {
+ Operation *defOp = v.getDefiningOp();
+ if (!defOp)
+ return false;
+ if (!isTriviallyRematerializable(defOp))
+ return false;
+ }
+ return true;
+}
+
+// Backward layout propagation assumes a single well-defined layout per def at
+// all its use points. Upstream value-numbering passes (e.g. CSE) can merge
+// pure value generators with no operands and no layout-bearing attributes
+// (such as two identical `vector.step` ops), producing a single SSA value
+// whose distinct consumers later require *different* layouts. Bridging two
+// such layouts at distribution time can force a cross-subgroup data movement
+// through SLM. To preserve the single-layout-per-def invariant without paying
+// that cost, clone the producer once per distinct required layout and rewrite
+// the offending uses, but only when the producer is trivially
+// rematerializable.
+static void splitOnConflictingUseLayouts(Operation *op) {
+ if (op->getNumResults() != 1)
+ return;
+ OpResult result = op->getResult(0);
+ if (!isa<VectorType>(result.getType()) || result.use_empty())
+ return;
+
+ // Bucket uses by required layout. Uses without a recorded layout are
+ // attached to the first bucket that gets created so they stay on the
+ // original op.
+ llvm::MapVector<mlir::Attribute, SmallVector<OpOperand *>> buckets;
+ SmallVector<OpOperand *> unlabeled;
+ for (OpOperand &use : result.getUses()) {
+ if (auto l = xegpu::getDistributeLayoutAttr(use))
+ buckets[l].push_back(&use);
+ else
+ unlabeled.push_back(&use);
+ }
+ if (buckets.size() <= 1)
+ return;
+ if (!isTriviallyRematerializable(op))
+ return;
+
+ // Keep the first bucket (and any unlabeled uses) on the original op.
+ // Clone the op for each remaining bucket and rewire its uses.
+ OpBuilder builder(op);
+ bool first = true;
+ for (auto &kv : buckets) {
+ if (first) {
+ first = false;
+ continue;
+ }
+ Operation *clone = builder.clone(*op);
+ Value newRes = clone->getResult(0);
+ for (OpOperand *use : kv.second)
+ use->set(newRes);
+ }
+}
+
// For regular operations: First the result layouts are propagated from uses.
// Then the result layouts are propagated to uses (operands).
static void propagateResultsToRegularOperands(Operation *op) {
if (op->getNumResults() == 0 || op->getNumResults() > 1)
return;
+ // If multiple uses demand distinct layouts and op is cheap to
+ // rematerialize, clone per layout so the recovered IR has a single
+ // well-defined layout per def. This avoids inserting `convert_layout`
+ // (and the SLM round-trip its WG-level lowering would entail) for cases
+ // like CSE-merged `vector.step` feeding broadcasts with conflicting
+ // distributions.
+ splitOnConflictingUseLayouts(op);
+
OpResult result = op->getResult(0);
xegpu::DistributeLayoutAttr resLayout = getLayoutFromUsePoints(result);
Type resultType = result.getType();
>From 32d1b647720eca8b895efda8028a77bed498299c Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Tue, 12 May 2026 20:07:55 +0000
Subject: [PATCH 2/6] Add test case
---
.../Dialect/XeGPU/xegpu-recover-layout.mlir | 27 +++++++++++++++++++
1 file changed, 27 insertions(+)
diff --git a/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir b/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir
index e2a4897fac519..abf1863941a88 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir
@@ -146,3 +146,30 @@ gpu.func @if_basic(
gpu.return
}
}
+
+// -----
+// Test splitOnConflictingUseLayouts: a single `vector.step` (a trivially
+// rematerializable pure value generator, like one produced by CSE) whose two
+// uses require different *workgroup* layouts must be cloned so each clone has
+// a single well-defined layout. Without this, layout recovery would have to
+// insert a bridging `xegpu.convert_layout`, whose WG-to-SG lowering would
+// entail an SLM round-trip to redistribute data across subgroups.
+
+gpu.module @test_step_split {
+// CHECK-LABEL: gpu.func @step_split_on_conflicting_sg_layouts
+gpu.func @step_split_on_conflicting_sg_layouts(
+ %arg0: memref<1024xf32>, %arg1: memref<1024xf32>) {
+ %mask = arith.constant dense<true> : vector<128xi1>
+ // CHECK-DAG: vector.step {layout_result_0 = #xegpu.layout<sg_layout = [16], sg_data = [8]>} : vector<128xindex>
+ // CHECK-DAG: vector.step {layout_result_0 = #xegpu.layout<sg_layout = [8], sg_data = [16]>} : vector<128xindex>
+ // CHECK-NOT: xegpu.convert_layout
+ %off = vector.step : vector<128xindex>
+ %ld_a = xegpu.load %arg0[%off], %mask
+ <{layout = #xegpu.layout<sg_layout = [16], sg_data = [8]>}>
+ : memref<1024xf32>, vector<128xindex>, vector<128xi1> -> vector<128xf32>
+ %ld_b = xegpu.load %arg1[%off], %mask
+ <{layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>}>
+ : memref<1024xf32>, vector<128xindex>, vector<128xi1> -> vector<128xf32>
+ gpu.return
+}
+}
>From 578c3a4ce6f65d6999687de7e9b674d10da8fefe Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Mon, 18 May 2026 20:52:09 +0000
Subject: [PATCH 3/6] Add the clone in layout propagation
---
.../XeGPU/Transforms/XeGPULayoutImpl.h | 6 ++
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 60 +------------------
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 20 +++++++
.../XeGPU/resolve-layout-conflicts.mlir | 35 +++++++++--
.../Dialect/XeGPU/xegpu-recover-layout.mlir | 26 --------
5 files changed, 57 insertions(+), 90 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index 299f9f18e3be6..3e307437a3147 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -239,6 +239,12 @@ setupDpasMxLayout(LayoutKind layoutKind, VectorType aTy, VectorType bTy,
/// users and determine the expected layout accordingly.
DistributeLayoutAttr getConsumerLayoutAt(OpOperand &operand);
+/// Returns true if `op` is safe and cheap to clone: it has no side effects,
+/// no regions, and all of its operands are themselves trivially
+/// rematerializable (e.g. `vector.step`, splat `arith.constant`, or
+/// `vector.create_mask` whose operands are constants).
+bool isTriviallyRematerializable(Operation *op);
+
} // namespace xegpu
} // namespace mlir
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 7381dd77ad63b..a519f1008c891 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -27,7 +27,6 @@
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Transforms/DialectConversion.h"
-#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/Support/FormatVariadic.h"
#include <cstdint>
@@ -130,7 +129,7 @@ static xegpu::DistributeLayoutAttr getLayoutFromUsePoints(Value result) {
// regions, and all operands are themselves trivially rematerializable, e.g.
// block-arg-free pure value generators such as `vector.step`, splat
// `arith.constant`, or `vector.create_mask` whose operands are constants).
-static bool isTriviallyRematerializable(Operation *op) {
+bool xegpu::isTriviallyRematerializable(Operation *op) {
if (!op || op->getNumRegions() != 0)
return false;
if (!isMemoryEffectFree(op))
@@ -145,55 +144,6 @@ static bool isTriviallyRematerializable(Operation *op) {
return true;
}
-// Backward layout propagation assumes a single well-defined layout per def at
-// all its use points. Upstream value-numbering passes (e.g. CSE) can merge
-// pure value generators with no operands and no layout-bearing attributes
-// (such as two identical `vector.step` ops), producing a single SSA value
-// whose distinct consumers later require *different* layouts. Bridging two
-// such layouts at distribution time can force a cross-subgroup data movement
-// through SLM. To preserve the single-layout-per-def invariant without paying
-// that cost, clone the producer once per distinct required layout and rewrite
-// the offending uses, but only when the producer is trivially
-// rematerializable.
-static void splitOnConflictingUseLayouts(Operation *op) {
- if (op->getNumResults() != 1)
- return;
- OpResult result = op->getResult(0);
- if (!isa<VectorType>(result.getType()) || result.use_empty())
- return;
-
- // Bucket uses by required layout. Uses without a recorded layout are
- // attached to the first bucket that gets created so they stay on the
- // original op.
- llvm::MapVector<mlir::Attribute, SmallVector<OpOperand *>> buckets;
- SmallVector<OpOperand *> unlabeled;
- for (OpOperand &use : result.getUses()) {
- if (auto l = xegpu::getDistributeLayoutAttr(use))
- buckets[l].push_back(&use);
- else
- unlabeled.push_back(&use);
- }
- if (buckets.size() <= 1)
- return;
- if (!isTriviallyRematerializable(op))
- return;
-
- // Keep the first bucket (and any unlabeled uses) on the original op.
- // Clone the op for each remaining bucket and rewire its uses.
- OpBuilder builder(op);
- bool first = true;
- for (auto &kv : buckets) {
- if (first) {
- first = false;
- continue;
- }
- Operation *clone = builder.clone(*op);
- Value newRes = clone->getResult(0);
- for (OpOperand *use : kv.second)
- use->set(newRes);
- }
-}
-
// For regular operations: First the result layouts are propagated from uses.
// Then the result layouts are propagated to uses (operands).
static void propagateResultsToRegularOperands(Operation *op) {
@@ -202,14 +152,6 @@ static void propagateResultsToRegularOperands(Operation *op) {
if (op->getNumResults() > 1 && !isa<vector::DeinterleaveOp>(op))
return;
- // If multiple uses demand distinct layouts and op is cheap to
- // rematerialize, clone per layout so the recovered IR has a single
- // well-defined layout per def. This avoids inserting `convert_layout`
- // (and the SLM round-trip its WG-level lowering would entail) for cases
- // like CSE-merged `vector.step` feeding broadcasts with conflicting
- // distributions.
- splitOnConflictingUseLayouts(op);
-
OpResult result = op->getResult(0);
xegpu::DistributeLayoutAttr resLayout = getLayoutFromUsePoints(result);
Type resultType = result.getType();
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 9c63beda281ad..2fceebafcce17 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1560,6 +1560,26 @@ ResolveLayoutConflicts::resolveVectorConsumer(OpOperand &operand) {
if (consumerLayout.isEqualTo(producerLayout))
return success();
+ // If the producer is trivially rematerializable (e.g. `vector.step`, splat
+ // `arith.constant`), clone it and stamp the consumer's expected layout on
+ // the clone instead of inserting a `xegpu.convert_layout`. The convert
+ // would otherwise lower to a cross-subgroup data movement through SLM at
+ // WG-to-SG distribution time, which is strictly more expensive than
+ // recomputing a pure value generator.
+ if (auto *producerOp = vectorValue.getDefiningOp();
+ producerOp && producerOp->getNumResults() == 1 &&
+ isa<OpResult>(vectorValue) &&
+ xegpu::isTriviallyRematerializable(producerOp)) {
+ builder.setInsertionPointAfter(producerOp);
+ Operation *clone = builder.clone(*producerOp);
+ OpResult cloneResult = clone->getResult(0);
+ // Drop the inherited producer layout so the new layout takes effect
+ xegpu::removeLayoutAttr(cloneResult);
+ xegpu::setDistributeLayoutAttr(cloneResult, consumerLayout);
+ operand.set(cloneResult);
+ return success();
+ }
+
// Insert a convert_layout op to resolve the conflict.
builder.setInsertionPointAfterValue(vectorValue);
auto convertOp = xegpu::ConvertLayoutOp::create(
diff --git a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
index 860cbc2112d20..6a410ec612880 100644
--- a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
+++ b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
@@ -154,17 +154,18 @@ func.func @bitcast_source_conflict() -> vector<32x32xf16> {
return %1 : vector<32x32xf16>
}
+// The accumulator `arith.constant` is trivially rematerializable, so the
+// conflict resolver clones it with the consumer's expected layout instead
+// of inserting an `xegpu.convert_layout`
// CHECK-LABEL: func.func @multireduction_source_conflict
// CHECK-DAG: %[[V0:.*]] = "some_op"() {layout_result_0 = #xegpu.layout<inst_data = [32, 16]>} : () -> vector<32x32xf16>
// CHECK-DAG: %[[CVT0:.*]] = xegpu.convert_layout %[[V0]]
// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [32, 16]>, target_layout = #xegpu.layout<inst_data = [16, 16]>}>
// CHECK-SAME: : vector<32x32xf16>
-// CHECK-DAG: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [32]>}
+// CHECK-DAG: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [16, 16]>, dims = [0]>}
// CHECK-SAME: dense<0.000000e+00> : vector<32xf16>
-// CHECK-DAG: %[[CVT1:.*]] = xegpu.convert_layout %[[CST]]
-// CHECK-SAME: <{input_layout = #xegpu.layout<inst_data = [32]>, target_layout = #xegpu.slice<#xegpu.layout<inst_data = [16, 16]>, dims = [0]>}>
-// CHECK-SAME: : vector<32xf16>
-// CHECK: %[[MR:.*]] = vector.multi_reduction <add>, %[[CVT0]], %[[CVT1]]
+// CHECK-NOT: xegpu.convert_layout %{{.*}} : vector<32xf16>
+// CHECK: %[[MR:.*]] = vector.multi_reduction <add>, %[[CVT0]], %[[CST]]
// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [16, 16]>, dims = [0]>}
// CHECK-SAME: [0] : vector<32x32xf16> to vector<32xf16>
// CHECK: return %[[MR]] : vector<32xf16>
@@ -278,6 +279,30 @@ func.func @convert_layout() {
return
}
+// A `vector.step` reaches two consumers that need different slice layouts of
+// the anchor's sg_layout (one along dim 0, one along dim 1). The conflict is
+// trivially rematerializable, so the resolver clones the step with the second
+// layout instead of inserting an `xegpu.convert_layout` (which would lower to
+// a cross-subgroup SLM round-trip at WG-to-SG distribution time).
+// CHECK-LABEL: gpu.func @step_clone_via_anchor
+// CHECK-DAG: vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex>
+// CHECK-DAG: vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} : vector<32xindex>
+// CHECK-NOT: xegpu.convert_layout {{.*}} : vector<32xindex>
+gpu.func @step_clone_via_anchor(%arg0: i64) kernel {
+ %mask = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} dense<true> : vector<32x32xi1>
+ %cst32 = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} dense<32> : vector<32xindex>
+ %step = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex>
+ %col = arith.muli %step, %cst32 {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} : vector<32xindex>
+ %col2d = vector.shape_cast %col {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>} : vector<32xindex> to vector<32x1xindex>
+ %colb = vector.broadcast %col2d {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} : vector<32x1xindex> to vector<32x32xindex>
+ %rowb = vector.broadcast %step {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} : vector<32xindex> to vector<32x32xindex>
+ %off = arith.addi %colb, %rowb {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} : vector<32x32xindex>
+ %v = xegpu.load %arg0[%off], %mask <{layout = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>}>
+ : i64, vector<32x32xindex>, vector<32x32xi1> -> vector<32x32xf32>
+ "consume"(%v) : (vector<32x32xf32>) -> ()
+ gpu.return
+}
+
// CHECK-LABEL: func.func @extract_source_conflict_with_order
// CHECK-DAG: %[[V0:.*]] = "some_op"() {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1], order = [2, 3, 0, 1]>} : () -> vector<2x4x16x32xf16>
// CHECK-DAG: %[[CVT:.*]] = xegpu.convert_layout %[[V0]]
diff --git a/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir b/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir
index abf1863941a88..a943e68bdcf41 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir
@@ -147,29 +147,3 @@ gpu.func @if_basic(
}
}
-// -----
-// Test splitOnConflictingUseLayouts: a single `vector.step` (a trivially
-// rematerializable pure value generator, like one produced by CSE) whose two
-// uses require different *workgroup* layouts must be cloned so each clone has
-// a single well-defined layout. Without this, layout recovery would have to
-// insert a bridging `xegpu.convert_layout`, whose WG-to-SG lowering would
-// entail an SLM round-trip to redistribute data across subgroups.
-
-gpu.module @test_step_split {
-// CHECK-LABEL: gpu.func @step_split_on_conflicting_sg_layouts
-gpu.func @step_split_on_conflicting_sg_layouts(
- %arg0: memref<1024xf32>, %arg1: memref<1024xf32>) {
- %mask = arith.constant dense<true> : vector<128xi1>
- // CHECK-DAG: vector.step {layout_result_0 = #xegpu.layout<sg_layout = [16], sg_data = [8]>} : vector<128xindex>
- // CHECK-DAG: vector.step {layout_result_0 = #xegpu.layout<sg_layout = [8], sg_data = [16]>} : vector<128xindex>
- // CHECK-NOT: xegpu.convert_layout
- %off = vector.step : vector<128xindex>
- %ld_a = xegpu.load %arg0[%off], %mask
- <{layout = #xegpu.layout<sg_layout = [16], sg_data = [8]>}>
- : memref<1024xf32>, vector<128xindex>, vector<128xi1> -> vector<128xf32>
- %ld_b = xegpu.load %arg1[%off], %mask
- <{layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>}>
- : memref<1024xf32>, vector<128xindex>, vector<128xi1> -> vector<128xf32>
- gpu.return
-}
-}
>From 99b3650f39b681bfc0b67172a1f6632af0af9dcb Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Mon, 18 May 2026 21:17:51 +0000
Subject: [PATCH 4/6] Remove new line
---
mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp | 1 -
mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp | 2 +-
mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir | 1 -
3 files changed, 1 insertion(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index a519f1008c891..bf5d260db9e18 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -151,7 +151,6 @@ static void propagateResultsToRegularOperands(Operation *op) {
return;
if (op->getNumResults() > 1 && !isa<vector::DeinterleaveOp>(op))
return;
-
OpResult result = op->getResult(0);
xegpu::DistributeLayoutAttr resLayout = getLayoutFromUsePoints(result);
Type resultType = result.getType();
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 2fceebafcce17..75f27d6922ea2 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1564,7 +1564,7 @@ ResolveLayoutConflicts::resolveVectorConsumer(OpOperand &operand) {
// `arith.constant`), clone it and stamp the consumer's expected layout on
// the clone instead of inserting a `xegpu.convert_layout`. The convert
// would otherwise lower to a cross-subgroup data movement through SLM at
- // WG-to-SG distribution time, which is strictly more expensive than
+ // WG-to-SG distribution time, which is more expensive than
// recomputing a pure value generator.
if (auto *producerOp = vectorValue.getDefiningOp();
producerOp && producerOp->getNumResults() == 1 &&
diff --git a/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir b/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir
index a943e68bdcf41..e2a4897fac519 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir
@@ -146,4 +146,3 @@ gpu.func @if_basic(
gpu.return
}
}
-
>From 2ba4f4c7c536d0a756fbd8781cb822d7e7ec5d0c Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Mon, 18 May 2026 21:45:59 +0000
Subject: [PATCH 5/6] update test
---
mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir | 5 +++--
1 file changed, 3 insertions(+), 2 deletions(-)
diff --git a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
index 6a410ec612880..88401373f38e1 100644
--- a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
+++ b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
@@ -288,7 +288,7 @@ func.func @convert_layout() {
// CHECK-DAG: vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex>
// CHECK-DAG: vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} : vector<32xindex>
// CHECK-NOT: xegpu.convert_layout {{.*}} : vector<32xindex>
-gpu.func @step_clone_via_anchor(%arg0: i64) kernel {
+gpu.func @step_clone_via_anchor(%arg0: i64, %arg1: memref<32x32xf32>) kernel {
%mask = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} dense<true> : vector<32x32xi1>
%cst32 = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} dense<32> : vector<32xindex>
%step = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex>
@@ -299,7 +299,8 @@ gpu.func @step_clone_via_anchor(%arg0: i64) kernel {
%off = arith.addi %colb, %rowb {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} : vector<32x32xindex>
%v = xegpu.load %arg0[%off], %mask <{layout = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>}>
: i64, vector<32x32xindex>, vector<32x32xi1> -> vector<32x32xf32>
- "consume"(%v) : (vector<32x32xf32>) -> ()
+ %tdesc = xegpu.create_nd_tdesc %arg1 : memref<32x32xf32> -> !xegpu.tensor_desc<32x32xf32, #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>>
+ xegpu.store_nd %v, %tdesc[0, 0] : vector<32x32xf32>, !xegpu.tensor_desc<32x32xf32, #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>>
gpu.return
}
>From 538610aac6843122a8fc1cd2ff4fc24833fe1357 Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Tue, 26 May 2026 15:50:20 +0000
Subject: [PATCH 6/6] Add test
---
.../XeGPU/resolve-layout-conflicts.mlir | 25 +++++++++++++++++++
1 file changed, 25 insertions(+)
diff --git a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
index 8e816311ece44..40e907be6e4a1 100644
--- a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
+++ b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
@@ -304,6 +304,31 @@ gpu.func @step_clone_via_anchor(%arg0: i64, %arg1: memref<32x32xf32>) kernel {
gpu.return
}
+// Test that when a derived value (arith.muli) has a layout conflict on its
+// uses, and the derived value's producer chain is trivially rematerializable
+// (because step and constant stride are themselves trivially rematerializable),
+// the resolver clones the arith.muli instead of inserting an xegpu.convert_layout.
+// CHECK-LABEL: gpu.func @step_muli_clone_via_anchor
+// CHECK-DAG: vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex>
+// CHECK-DAG: arith.muli {{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex>
+// CHECK-DAG: arith.muli {{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} : vector<32xindex>
+// CHECK-NOT: xegpu.convert_layout {{.*}} : vector<32xindex>
+gpu.func @step_muli_clone_via_anchor(%arg0: i64, %arg1: memref<32x32xf32>) kernel {
+ %mask = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} dense<true> : vector<32x32xi1>
+ %cst32 = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} dense<32> : vector<32xindex>
+ %step = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex>
+ %scaled = arith.muli %step, %cst32 {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex>
+ %col2d = vector.shape_cast %scaled {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>} : vector<32xindex> to vector<32x1xindex>
+ %colb = vector.broadcast %col2d {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} : vector<32x1xindex> to vector<32x32xindex>
+ %rowb = vector.broadcast %scaled {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} : vector<32xindex> to vector<32x32xindex>
+ %off = arith.addi %colb, %rowb {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} : vector<32x32xindex>
+ %v = xegpu.load %arg0[%off], %mask <{layout = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>}>
+ : i64, vector<32x32xindex>, vector<32x32xi1> -> vector<32x32xf32>
+ %tdesc = xegpu.create_nd_tdesc %arg1 : memref<32x32xf32> -> !xegpu.tensor_desc<32x32xf32>
+ xegpu.store_nd %v, %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>}> : vector<32x32xf32>, !xegpu.tensor_desc<32x32xf32>
+ gpu.return
+}
+
// CHECK-LABEL: func.func @extract_source_conflict_with_order
// CHECK-DAG: %[[V0:.*]] = "some_op"() {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1], order = [2, 3, 0, 1]>} : () -> vector<2x4x16x32xf16>
// CHECK-DAG: %[[CVT:.*]] = xegpu.convert_layout %[[V0]]
More information about the Mlir-commits
mailing list