[Mlir-commits] [mlir] [MLIR][XeGPU] Rematerialize layout conversion sources after conflict resolution (PR #230669)
Nishant Patel
llvmlistbot at llvm.org
Fri Oct 9 14:01:54 PDT 2026
https://github.com/nbpatel created https://github.com/llvm/llvm-project/pull/230669
Layout conflict resolution inserts xegpu.convert_layout whenever there is a conflict. A separate step then replaces conversions of cheap producer chains (e.g. vector.step → arith.index_castui) by cloning the whole chain in the target layout, which fixes operand/result layout mismatches when an elementwise producer was cloned.
>From ba338efd0e14057818e30e3baa901590f4630d9b Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Fri, 18 Sep 2026 18:06:41 +0000
Subject: [PATCH 1/3] [mlir][xegpu] Fix operand layouts when rematerializing on
layout conflict
---
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 38 +++++++++++++++----
.../XeGPU/resolve-layout-conflicts.mlir | 25 ++++++++++++
2 files changed, 56 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index ca1c6501a4856a..00fec294b9ef53 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -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.
+ xegpu::removeLayoutAttr(cloneResult);
+ xegpu::setDistributeLayoutAttr(cloneResult, layout);
+ for (OpOperand &cloneOperand : clone->getOpOperands()) {
+ Value operandValue = cloneOperand.get();
+ if (!isa<VectorType>(operandValue.getType()) ||
+ !operandValue.getDefiningOp())
+ continue;
+ xegpu::DistributeLayoutAttr operandLayout =
+ xegpu::inferSourceLayoutFromResultForNonAnchorOp(cloneOperand, layout);
+ if (!operandLayout)
+ continue;
+ cloneOperand.set(
+ rematerializeWithLayout(builder, operandValue, operandLayout));
+ }
+ return cloneResult;
+}
+
LogicalResult
ResolveLayoutConflicts::resolveVectorConsumer(OpOperand &operand) {
Value vectorValue = operand.get();
@@ -1711,13 +1741,7 @@ ResolveLayoutConflicts::resolveVectorConsumer(OpOperand &operand) {
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);
+ operand.set(rematerializeWithLayout(builder, vectorValue, consumerLayout));
return success();
}
diff --git a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
index 99e3f3be31c188..4908875d57b426 100644
--- a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
+++ b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
@@ -408,6 +408,31 @@ gpu.func @step_muli_clone_via_anchor(%arg0: i64, %arg1: memref<32x32xf32>) kerne
gpu.return
}
+// Rematerializing a producer must also rematerialize its operands so the clone
+// stays internally consistent. Here `arith.index_castui` (elementwise) is cloned
+// for the second consumer layout; its `vector.step` operand must be cloned with
+// the matching layout too, otherwise the clone would have an operand and result
+// with different distributed shapes and the distribution pass would emit an
+// `arith.index_castui` with mismatched operand/result vector shapes.
+// CHECK-LABEL: gpu.func @step_index_castui_clone_via_anchor
+// CHECK: %[[STEP0:.*]] = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex>
+// CHECK: %[[STEP1:.*]] = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} : vector<32xindex>
+// CHECK: arith.index_castui %[[STEP0]] {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex> to vector<32xi32>
+// CHECK: arith.index_castui %[[STEP1]] {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} : vector<32xindex> to vector<32xi32>
+// CHECK-NOT: xegpu.convert_layout {{.*}} : vector<32xi32>
+gpu.func @step_index_castui_clone_via_anchor(%arg0: memref<32x32xi32>) kernel {
+ %step = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex>
+ %cast = arith.index_castui %step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex> to vector<32xi32>
+ %col2d = vector.shape_cast %cast {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>} : vector<32xi32> to vector<32x1xi32>
+ %colb = vector.broadcast %col2d {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} : vector<32x1xi32> to vector<32x32xi32>
+ %row2d = vector.shape_cast %cast {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} : vector<32xi32> to vector<1x32xi32>
+ %rowb = vector.broadcast %row2d {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} : vector<1x32xi32> to vector<32x32xi32>
+ %off = arith.addi %colb, %rowb {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} : vector<32x32xi32>
+ %tdesc = xegpu.create_nd_tdesc %arg0 : memref<32x32xi32> -> !xegpu.tensor_desc<32x32xi32>
+ xegpu.store_nd %off, %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>}> : vector<32x32xi32>, !xegpu.tensor_desc<32x32xi32>
+ 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]]
>From 70fd30444617e234dc1e827ad6ec481a7dcf90c3 Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Fri, 2 Oct 2026 04:06:24 +0000
Subject: [PATCH 2/3] Address feedback
---
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 37 ++++++-------------
.../XeGPU/resolve-layout-conflicts.mlir | 12 ++----
2 files changed, 15 insertions(+), 34 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 00fec294b9ef53..70d92fd8f4b9a4 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1652,34 +1652,16 @@ 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) {
+// Clones the trivially-rematerializable producer subtree of `value`.
+static OpResult cloneRematerializableSubtree(OpBuilder &builder, Value value) {
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.
- xegpu::removeLayoutAttr(cloneResult);
- xegpu::setDistributeLayoutAttr(cloneResult, layout);
- for (OpOperand &cloneOperand : clone->getOpOperands()) {
- Value operandValue = cloneOperand.get();
- if (!isa<VectorType>(operandValue.getType()) ||
- !operandValue.getDefiningOp())
- continue;
- xegpu::DistributeLayoutAttr operandLayout =
- xegpu::inferSourceLayoutFromResultForNonAnchorOp(cloneOperand, layout);
- if (!operandLayout)
- continue;
- cloneOperand.set(
- rematerializeWithLayout(builder, operandValue, operandLayout));
- }
- return cloneResult;
+ for (OpOperand &cloneOperand : clone->getOpOperands())
+ if (isa<VectorType>(cloneOperand.get().getType()))
+ cloneOperand.set(
+ cloneRematerializableSubtree(builder, cloneOperand.get()));
+ return clone->getResult(cast<OpResult>(value).getResultNumber());
}
LogicalResult
@@ -1741,7 +1723,10 @@ ResolveLayoutConflicts::resolveVectorConsumer(OpOperand &operand) {
producerOp && producerOp->getNumResults() == 1 &&
isa<OpResult>(vectorValue) &&
xegpu::isTriviallyRematerializable(producerOp)) {
- operand.set(rematerializeWithLayout(builder, vectorValue, consumerLayout));
+ OpResult cloneResult = cloneRematerializableSubtree(builder, vectorValue);
+ xegpu::removeLayoutAttr(cloneResult);
+ xegpu::setDistributeLayoutAttr(cloneResult, consumerLayout);
+ operand.set(cloneResult);
return success();
}
diff --git a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
index 4908875d57b426..cb0b254235aac3 100644
--- a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
+++ b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
@@ -408,15 +408,11 @@ gpu.func @step_muli_clone_via_anchor(%arg0: i64, %arg1: memref<32x32xf32>) kerne
gpu.return
}
-// Rematerializing a producer must also rematerialize its operands so the clone
-// stays internally consistent. Here `arith.index_castui` (elementwise) is cloned
-// for the second consumer layout; its `vector.step` operand must be cloned with
-// the matching layout too, otherwise the clone would have an operand and result
-// with different distributed shapes and the distribution pass would emit an
-// `arith.index_castui` with mismatched operand/result vector shapes.
+// Rematerializing `arith.index_castui` for the second consumer layout also
+// clones its `vector.step` operand, so the two casts do not share a producer.
// CHECK-LABEL: gpu.func @step_index_castui_clone_via_anchor
-// CHECK: %[[STEP0:.*]] = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex>
-// CHECK: %[[STEP1:.*]] = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} : vector<32xindex>
+// CHECK: %[[STEP0:.*]] = vector.step
+// CHECK: %[[STEP1:.*]] = vector.step
// CHECK: arith.index_castui %[[STEP0]] {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex> to vector<32xi32>
// CHECK: arith.index_castui %[[STEP1]] {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} : vector<32xindex> to vector<32xi32>
// CHECK-NOT: xegpu.convert_layout {{.*}} : vector<32xi32>
>From 33756957cd700dcf7f1b8d39cf743d099b5667c2 Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Wed, 7 Oct 2026 18:38:21 +0000
Subject: [PATCH 3/3] [MLIR][XeGPU] Rematerialize layout conversion sources as
a separate step after conflict resolution
---
.../XeGPU/Transforms/XeGPULayoutImpl.h | 20 +-
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 91 ++++++--
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 83 ++++---
.../rematerialize-conversion-sources.mlir | 205 ++++++++++++++++++
.../XeGPU/resolve-layout-conflicts.mlir | 77 ++-----
.../lib/Dialect/XeGPU/TestXeGPUTransforms.cpp | 32 +++
6 files changed, 383 insertions(+), 125 deletions(-)
create mode 100644 mlir/test/Dialect/XeGPU/rematerialize-conversion-sources.mlir
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index 8921380093e108..03cf2f17eda8de 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -18,7 +18,6 @@
#include "llvm/ADT/STLFunctionalExtras.h"
namespace mlir {
-
class VectorType;
class OpOperand;
class OpResult;
@@ -34,13 +33,15 @@ class TensorDescType;
} // namespace xegpu
namespace xegpu {
-
LogicalResult propagateLayouts(OpBuilder &builder, Operation *target,
LayoutKind layoutKind, unsigned indexBitWidth,
bool printOnly = false);
LogicalResult resolveLayoutConflicts(Operation *target);
+/// Replace layout conversions by recomputing eligible producer chains.
+void rematerializeConversionSources(OpBuilder &builder, Operation *target);
+
/// Sink `xegpu.convert_layout` ops past the elementwise operations they feed,
/// so that an elementwise op runs in the coarser layout.
void sinkElementwiseConversions(OpBuilder &builder, Operation *target);
@@ -465,14 +466,19 @@ completeDpasMxLaneLayoutFromInstData(DistributeLayoutAttr aLayout,
/// 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).
+/// Returns true if the vector produced by `op` can be recomputed in any layout
+/// by cloning a bounded chain of pure vector generators (`vector.step`, splat
+/// `arith.constant`, `vector.create_mask`, scalar `vector.broadcast`) and
+/// layout-preserving elementwise ops.
bool isTriviallyRematerializable(Operation *op);
+/// Clone a bounded chain of pure vector generators and layout-preserving
+/// elementwise ops in `layout`. Each producer is cloned once, immediately after
+/// its original; scalar operands remain shared. Return null if the chain is
+/// ineligible, without modifying the IR.
+Value rematerializeWithLayout(OpBuilder &builder, Value value,
+ DistributeLayoutAttr layout);
} // namespace xegpu
-
} // namespace mlir
#endif // MLIR_DIALECT_XEGPU_UTILS_XEGPUUTILS_H_
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 031d971cebe6cf..0a749efbba1f53 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -13,6 +13,7 @@
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/XeVMDialect.h"
@@ -21,12 +22,14 @@
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/IR/Builders.h"
+#include "mlir/IR/IRMapping.h"
#include "mlir/IR/Operation.h"
#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/DenseSet.h"
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/Support/FormatVariadic.h"
#include <cstdint>
@@ -132,25 +135,87 @@ 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).
-bool xegpu::isTriviallyRematerializable(Operation *op) {
- if (!op || op->getNumRegions() != 0)
- return false;
- if (!isMemoryEffectFree(op))
+static constexpr unsigned kMaxRematerializableChainOps = 16;
+
+/// Accept pure vector generators and elementwise ops that preserve vector
+/// shape. Generators have no vector operands; elementwise ops share their
+/// result layout with their vector operands, so the chain needs no layout
+/// inference.
+static bool isLayoutPreservingVectorProducer(Operation *op) {
+ if (op->getNumResults() != 1 || op->getNumRegions() != 0 ||
+ !isa<VectorType>(op->getResult(0).getType()) || !isMemoryEffectFree(op))
return false;
- for (Value v : op->getOperands()) {
- Operation *defOp = v.getDefiningOp();
- if (!defOp)
- return false;
- if (!isTriviallyRematerializable(defOp))
+ if (auto constant = dyn_cast<arith::ConstantOp>(op)) {
+ auto value = dyn_cast<DenseElementsAttr>(constant.getValue());
+ return value && value.isSplat();
+ }
+ if (isa<vector::StepOp, vector::CreateMaskOp>(op))
+ return true;
+ if (auto broadcast = dyn_cast<vector::BroadcastOp>(op))
+ return !isa<VectorType>(broadcast.getSource().getType());
+ return OpTrait::hasElementwiseMappableTraits(op);
+}
+
+/// Collect each eligible producer once, producers first, for cloning with an
+/// IRMapping. Reject chains exceeding the duplication limit or reaching an
+/// unsupported producer or vector block argument.
+static bool collectRematerializableChain(Operation *root,
+ SmallVectorImpl<Operation *> &chain) {
+ // Re-push each op below its producers to obtain an iterative post-order walk.
+ SmallVector<std::pair<Operation *, bool>> worklist{{root, false}};
+ DenseSet<Operation *> visited;
+ while (!worklist.empty()) {
+ auto [op, expanded] = worklist.pop_back_val();
+ if (expanded) {
+ chain.push_back(op);
+ continue;
+ }
+ if (!visited.insert(op).second)
+ continue;
+ if (visited.size() > kMaxRematerializableChainOps ||
+ !isLayoutPreservingVectorProducer(op))
return false;
+ worklist.emplace_back(op, true);
+ for (Value operand : op->getOperands()) {
+ if (!isa<VectorType>(operand.getType()))
+ continue;
+ Operation *producer = operand.getDefiningOp();
+ if (!producer)
+ return false;
+ if (!visited.contains(producer))
+ worklist.emplace_back(producer, false);
+ }
}
return true;
}
+bool xegpu::isTriviallyRematerializable(Operation *op) {
+ if (!op)
+ return false;
+ SmallVector<Operation *> chain;
+ return collectRematerializableChain(op, chain);
+}
+
+Value xegpu::rematerializeWithLayout(OpBuilder &builder, Value value,
+ DistributeLayoutAttr layout) {
+ Operation *producer = value.getDefiningOp();
+ SmallVector<Operation *> chain;
+ if (!producer || !collectRematerializableChain(producer, chain))
+ return nullptr;
+ // Clone producers first so vector operands map to clones; scalars stay
+ // shared.
+ OpBuilder::InsertionGuard guard(builder);
+ IRMapping mapping;
+ for (Operation *op : chain) {
+ builder.setInsertionPointAfter(op);
+ Operation *clone = builder.clone(*op, mapping);
+ OpResult cloneResult = clone->getResult(0);
+ removeLayoutAttr(cloneResult);
+ setDistributeLayoutAttr(cloneResult, layout);
+ }
+ return mapping.lookup(value);
+}
+
// 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) {
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 384e4a9af39df7..83adc7febef3b9 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -59,7 +59,6 @@ using namespace mlir;
using namespace mlir::dataflow;
namespace {
-
//===----------------------------------------------------------------------===//
// LayoutInfo
//===----------------------------------------------------------------------===//
@@ -1544,7 +1543,6 @@ void RunLayoutInfoPropagation::printAnalysisResult(llvm::raw_ostream &os) {
}
namespace {
-
//===----------------------------------------------------------------------===//
// ResolveLayoutConflicts
//===----------------------------------------------------------------------===//
@@ -1583,7 +1581,6 @@ struct ResolveLayoutConflicts {
LogicalResult resolveVectorConsumer(OpOperand &operand);
LogicalResult assignResultLayout(OpResult &result);
};
-
} // namespace
LogicalResult ResolveLayoutConflicts::run() {
@@ -1664,18 +1661,6 @@ LogicalResult ResolveLayoutConflicts::assignResultLayout(OpResult &result) {
return success();
}
-// Clones the trivially-rematerializable producer subtree of `value`.
-static OpResult cloneRematerializableSubtree(OpBuilder &builder, Value value) {
- Operation *producerOp = value.getDefiningOp();
- builder.setInsertionPointAfter(producerOp);
- Operation *clone = builder.clone(*producerOp);
- for (OpOperand &cloneOperand : clone->getOpOperands())
- if (isa<VectorType>(cloneOperand.get().getType()))
- cloneOperand.set(
- cloneRematerializableSubtree(builder, cloneOperand.get()));
- return clone->getResult(cast<OpResult>(value).getResultNumber());
-}
-
LogicalResult
ResolveLayoutConflicts::resolveVectorConsumer(OpOperand &operand) {
Value vectorValue = operand.get();
@@ -1725,23 +1710,6 @@ ResolveLayoutConflicts::resolveVectorConsumer(OpOperand &operand) {
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 more expensive than
- // recomputing a pure value generator.
- if (auto *producerOp = vectorValue.getDefiningOp();
- producerOp && producerOp->getNumResults() == 1 &&
- isa<OpResult>(vectorValue) &&
- xegpu::isTriviallyRematerializable(producerOp)) {
- OpResult cloneResult = cloneRematerializableSubtree(builder, vectorValue);
- 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(
@@ -1844,6 +1812,41 @@ static LogicalResult updateOpWithForwardFill(mlir::OpBuilder &builder,
return success();
}
+/// Recompute eligible producer chains in the conversion's target layout.
+///
+/// %step = vector.step {L1}
+/// %cast = arith.index_castui %step {L1}
+/// %cvt = xegpu.convert_layout %cast : L1 -> L2
+/// use(%cast) {L1}
+/// use(%cvt) {L2}
+///
+/// becomes
+///
+/// %step = vector.step {L1}
+/// %step' = vector.step {L2}
+/// %cast = arith.index_castui %step {L1}
+/// %cast' = arith.index_castui %step' {L2}
+/// use(%cast) {L1}
+/// use(%cast') {L2}
+///
+/// Clone the whole vector producer chain so the original and rematerialized
+/// values can carry independent layouts. For small, layout-preserving chains,
+/// recomputation can avoid cross-subgroup data movement through SLM.
+void xegpu::rematerializeConversionSources(OpBuilder &builder,
+ Operation *parentOp) {
+ SmallVector<xegpu::ConvertLayoutOp> conversions;
+ parentOp->walk(
+ [&](xegpu::ConvertLayoutOp convert) { conversions.push_back(convert); });
+ for (xegpu::ConvertLayoutOp convert : conversions) {
+ Value clone = xegpu::rematerializeWithLayout(builder, convert.getSource(),
+ convert.getTargetLayout());
+ if (!clone)
+ continue;
+ convert.getResult().replaceAllUsesWith(clone);
+ convert.erase();
+ }
+}
+
/// Optimize elementwise operations by sinking costly layout conversion.
///
/// %m = vector.create_mask ... {L1}
@@ -1928,17 +1931,11 @@ void xegpu::sinkElementwiseConversions(OpBuilder &builder,
continue;
if (rewiredOpIdxs.contains(operand.getOperandNumber()))
continue;
- Operation *definingOp = operandValue.getDefiningOp();
- assert(definingOp && xegpu::isTriviallyRematerializable(definingOp) &&
- "operand should have been rejected above");
// Rematerialize with uniform source layout.
- builder.setInsertionPointAfter(definingOp);
- Operation *clone = builder.clone(*definingOp);
- OpResult cloneResult =
- clone->getResult(cast<OpResult>(operandValue).getResultNumber());
- xegpu::removeLayoutAttr(cloneResult);
- xegpu::setDistributeLayoutAttr(cloneResult, uniformConvSrcLayout);
- operand.set(cloneResult);
+ Value clone = xegpu::rematerializeWithLayout(builder, operandValue,
+ uniformConvSrcLayout);
+ assert(clone && "operand should have been rejected above");
+ operand.set(clone);
}
// Run the op in the source layout and bridge its result back, so that
@@ -2008,7 +2005,6 @@ struct XeGPUPropagateLayoutPass final
: XeGPUPropagateLayoutBase(std::move(options)) {}
void runOnOperation() override;
};
-
} // namespace
LogicalResult xegpu::propagateLayouts(OpBuilder &builder, Operation *target,
@@ -2117,6 +2113,7 @@ void XeGPUPropagateLayoutPass::runOnOperation() {
signalPassFailure();
return;
}
+ xegpu::rematerializeConversionSources(builder, getOperation());
if (layoutKind == xegpu::LayoutKind::InstData)
xegpu::sinkElementwiseConversions(builder, getOperation());
}
diff --git a/mlir/test/Dialect/XeGPU/rematerialize-conversion-sources.mlir b/mlir/test/Dialect/XeGPU/rematerialize-conversion-sources.mlir
new file mode 100644
index 00000000000000..f90137b7fbb045
--- /dev/null
+++ b/mlir/test/Dialect/XeGPU/rematerialize-conversion-sources.mlir
@@ -0,0 +1,205 @@
+// RUN: mlir-opt --test-xegpu-rematerialize-conversion-sources -allow-unregistered-dialect -split-input-file %s | FileCheck %s
+
+// Row and column slices require different subgroup distributions.
+#row = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>
+#col = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>
+
+// Rematerialization preserves the original value for its other consumer.
+// CHECK-LABEL: gpu.func @step_conversion
+// CHECK: %[[STEP0:.*]] = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex>
+// CHECK-NEXT: %[[STEP1:.*]] = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} : vector<32xindex>
+// CHECK: arith.muli %[[STEP1]], %{{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} : vector<32xindex>
+// CHECK: vector.broadcast %[[STEP0]]
+// CHECK-NOT: xegpu.convert_layout
+gpu.module @test {
+gpu.func @step_conversion() kernel {
+ %cst32 = arith.constant {layout_result_0 = #col} dense<32> : vector<32xindex>
+ %step = vector.step {layout_result_0 = #row} : vector<32xindex>
+ %cvt = xegpu.convert_layout %step <{input_layout = #row, target_layout = #col}> : vector<32xindex>
+ %col = arith.muli %cvt, %cst32 {layout_result_0 = #col} : vector<32xindex>
+ %rowb = vector.broadcast %step {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} : vector<32xindex> to vector<32x32xindex>
+ "some_use"(%col, %rowb) : (vector<32xindex>, vector<32x32xindex>) -> ()
+ gpu.return
+}
+}
+
+// -----
+
+#row = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>
+#col = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>
+
+// The cloned elementwise chain shares no vector operands with the original.
+// CHECK-LABEL: gpu.func @elementwise_chain_conversion
+// CHECK-DAG: %[[CST1:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} dense<32> : vector<32xindex>
+// CHECK-DAG: %[[STEP1:.*]] = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} : vector<32xindex>
+// CHECK: %[[MUL1:.*]] = arith.muli %[[STEP1]], %[[CST1]] {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} : vector<32xindex>
+// CHECK: vector.shape_cast %[[MUL1]]
+// CHECK-NOT: xegpu.convert_layout
+gpu.module @test {
+gpu.func @elementwise_chain_conversion() kernel {
+ %cst32 = arith.constant {layout_result_0 = #row} dense<32> : vector<32xindex>
+ %step = vector.step {layout_result_0 = #row} : vector<32xindex>
+ %scaled = arith.muli %step, %cst32 {layout_result_0 = #row} : vector<32xindex>
+ %cvt = xegpu.convert_layout %scaled <{input_layout = #row, target_layout = #col}> : vector<32xindex>
+ %col2d = vector.shape_cast %cvt {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>} : vector<32xindex> to vector<32x1xindex>
+ %rowb = vector.broadcast %scaled {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} : vector<32xindex> to vector<32x32xindex>
+ "some_use"(%col2d, %rowb) : (vector<32x1xindex>, vector<32x32xindex>) -> ()
+ gpu.return
+}
+}
+
+// -----
+
+#row = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>
+#col = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>
+
+// The cloned cast and its step operand must have matching layouts.
+// CHECK-LABEL: gpu.func @cast_chain_conversion
+// CHECK: %[[STEP0:.*]] = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex>
+// CHECK-NEXT: %[[STEP1:.*]] = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} : vector<32xindex>
+// CHECK-NEXT: arith.index_castui %[[STEP0]] {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex> to vector<32xi32>
+// CHECK-NEXT: arith.index_castui %[[STEP1]] {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} : vector<32xindex> to vector<32xi32>
+// CHECK-NOT: xegpu.convert_layout
+gpu.module @test {
+gpu.func @cast_chain_conversion() kernel {
+ %step = vector.step {layout_result_0 = #row} : vector<32xindex>
+ %cast = arith.index_castui %step {layout_result_0 = #row} : vector<32xindex> to vector<32xi32>
+ %cvt = xegpu.convert_layout %cast <{input_layout = #row, target_layout = #col}> : vector<32xi32>
+ %col2d = vector.shape_cast %cvt {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>} : vector<32xi32> to vector<32x1xi32>
+ %row2d = vector.shape_cast %cast {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} : vector<32xi32> to vector<1x32xi32>
+ "some_use"(%col2d, %row2d) : (vector<32x1xi32>, vector<1x32xi32>) -> ()
+ gpu.return
+}
+}
+
+// -----
+
+#row = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>
+#col = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>
+
+// Shared producers are cloned only once.
+// CHECK-LABEL: gpu.func @shared_operand_cloned_once
+// CHECK-COUNT-2: vector.step
+// CHECK-NOT: vector.step
+// CHECK-NOT: xegpu.convert_layout
+// CHECK: gpu.return
+gpu.module @test {
+gpu.func @shared_operand_cloned_once() kernel {
+ %step = vector.step {layout_result_0 = #row} : vector<32xindex>
+ %sq = arith.muli %step, %step {layout_result_0 = #row} : vector<32xindex>
+ %cvt = xegpu.convert_layout %sq <{input_layout = #row, target_layout = #col}> : vector<32xindex>
+ "some_use"(%cvt) : (vector<32xindex>) -> ()
+ gpu.return
+}
+}
+
+// -----
+
+#row = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>
+#col = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>
+
+// Scalar block arguments remain shared because they carry no layout.
+// CHECK-LABEL: gpu.func @scalar_operand_shared
+// CHECK-SAME: (%[[N:.*]]: index)
+// CHECK: vector.create_mask %[[N]] {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xi1>
+// CHECK-NEXT: %[[MASK1:.*]] = vector.create_mask %[[N]] {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} : vector<32xi1>
+// CHECK: "some_use"(%[[MASK1]])
+// CHECK-NOT: xegpu.convert_layout
+gpu.module @test {
+gpu.func @scalar_operand_shared(%n: index) kernel {
+ %mask = vector.create_mask %n {layout_result_0 = #row} : vector<32xi1>
+ %cvt = xegpu.convert_layout %mask <{input_layout = #row, target_layout = #col}> : vector<32xi1>
+ "some_use"(%cvt) : (vector<32xi1>) -> ()
+ gpu.return
+}
+}
+
+// -----
+
+#row = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>
+#col = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>
+
+// All conversion users share one rematerialized value.
+// CHECK-LABEL: gpu.func @conversion_with_multiple_uses
+// CHECK: %[[STEP1:.*]] = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} : vector<32xindex>
+// CHECK: "use_a"(%[[STEP1]])
+// CHECK: "use_b"(%[[STEP1]])
+// CHECK-NOT: xegpu.convert_layout
+gpu.module @test {
+gpu.func @conversion_with_multiple_uses() kernel {
+ %step = vector.step {layout_result_0 = #row} : vector<32xindex>
+ %cvt = xegpu.convert_layout %step <{input_layout = #row, target_layout = #col}> : vector<32xindex>
+ "use_a"(%cvt) : (vector<32xindex>) -> ()
+ "use_b"(%cvt) : (vector<32xindex>) -> ()
+ gpu.return
+}
+}
+
+// -----
+
+#row = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>
+#col = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>
+
+// An ineligible producer prevents rematerialization of the chain.
+// CHECK-LABEL: gpu.func @opaque_source_kept
+// CHECK: xegpu.convert_layout
+gpu.module @test {
+gpu.func @opaque_source_kept() kernel {
+ %v = "some_op"() {layout_result_0 = #row} : () -> vector<32xindex>
+ %step = vector.step {layout_result_0 = #row} : vector<32xindex>
+ %sum = arith.addi %v, %step {layout_result_0 = #row} : vector<32xindex>
+ %cvt = xegpu.convert_layout %sum <{input_layout = #row, target_layout = #col}> : vector<32xindex>
+ "some_use"(%cvt) : (vector<32xindex>) -> ()
+ gpu.return
+}
+}
+
+// -----
+
+#row = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>
+
+// Shape-changing broadcasts cannot share one layout across the chain.
+// CHECK-LABEL: gpu.func @shape_changing_source_kept
+// CHECK: xegpu.convert_layout
+gpu.module @test {
+gpu.func @shape_changing_source_kept() kernel {
+ %step = vector.step {layout_result_0 = #row} : vector<32xindex>
+ %b = vector.broadcast %step {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} : vector<32xindex> to vector<32x32xindex>
+ %cvt = xegpu.convert_layout %b <{input_layout = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, target_layout = #xegpu.layout<sg_layout = [1, 32], sg_data = [32, 1]>}> : vector<32x32xindex>
+ "some_use"(%cvt) : (vector<32x32xindex>) -> ()
+ gpu.return
+}
+}
+
+// -----
+
+#row = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>
+#col = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>
+
+// A chain of 17 ops exceeds the duplication limit.
+// CHECK-LABEL: gpu.func @chain_too_long_kept
+// CHECK: xegpu.convert_layout
+gpu.module @test {
+gpu.func @chain_too_long_kept() kernel {
+ %s = vector.step {layout_result_0 = #row} : vector<32xindex>
+ %a1 = arith.addi %s, %s {layout_result_0 = #row} : vector<32xindex>
+ %a2 = arith.addi %a1, %a1 {layout_result_0 = #row} : vector<32xindex>
+ %a3 = arith.addi %a2, %a2 {layout_result_0 = #row} : vector<32xindex>
+ %a4 = arith.addi %a3, %a3 {layout_result_0 = #row} : vector<32xindex>
+ %a5 = arith.addi %a4, %a4 {layout_result_0 = #row} : vector<32xindex>
+ %a6 = arith.addi %a5, %a5 {layout_result_0 = #row} : vector<32xindex>
+ %a7 = arith.addi %a6, %a6 {layout_result_0 = #row} : vector<32xindex>
+ %a8 = arith.addi %a7, %a7 {layout_result_0 = #row} : vector<32xindex>
+ %a9 = arith.addi %a8, %a8 {layout_result_0 = #row} : vector<32xindex>
+ %a10 = arith.addi %a9, %a9 {layout_result_0 = #row} : vector<32xindex>
+ %a11 = arith.addi %a10, %a10 {layout_result_0 = #row} : vector<32xindex>
+ %a12 = arith.addi %a11, %a11 {layout_result_0 = #row} : vector<32xindex>
+ %a13 = arith.addi %a12, %a12 {layout_result_0 = #row} : vector<32xindex>
+ %a14 = arith.addi %a13, %a13 {layout_result_0 = #row} : vector<32xindex>
+ %a15 = arith.addi %a14, %a14 {layout_result_0 = #row} : vector<32xindex>
+ %a16 = arith.addi %a15, %a15 {layout_result_0 = #row} : vector<32xindex>
+ %cvt = xegpu.convert_layout %a16 <{input_layout = #row, target_layout = #col}> : vector<32xindex>
+ "some_use"(%cvt) : (vector<32xindex>) -> ()
+ gpu.return
+}
+}
diff --git a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
index 13341ac551ece1..e78b854bd0a749 100644
--- a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
+++ b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
@@ -154,18 +154,18 @@ func.func @bitcast_source_conflict() {
return
}
-// 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`
+// A conflicting constant accumulator requires a layout conversion.
// 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.slice<#xegpu.layout<inst_data = [16, 16]>, dims = [0]>}
+// CHECK-DAG: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [32]>}
// CHECK-SAME: dense<0.000000e+00> : vector<32xf16>
-// CHECK-NOT: xegpu.convert_layout %{{.*}} : vector<32xf16>
-// CHECK: %[[MR:.*]] = vector.multi_reduction <add>, %[[CVT0]], %[[CST]]
+// 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-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [16, 16]>, dims = [0]>}
// CHECK-SAME: [0] : vector<32x32xf16> to vector<32xf16>
// CHECK: return
@@ -356,16 +356,15 @@ 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, %arg1: memref<32x32xf32>) kernel {
+// Conflicting row and column uses of a step require a layout conversion.
+// CHECK-LABEL: gpu.func @step_conflict_via_anchor
+// CHECK: %[[STEP:.*]] = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex>
+// CHECK-NEXT: %[[CVT:.*]] = xegpu.convert_layout %[[STEP]]
+// CHECK-SAME: <{input_layout = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>, target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>}>
+// CHECK-SAME: : vector<32xindex>
+// CHECK-NEXT: arith.muli %[[CVT]], %{{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} : vector<32xindex>
+// CHECK: vector.broadcast %[[STEP]]
+gpu.func @step_conflict_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>
@@ -381,52 +380,6 @@ 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
-}
-
-// Rematerializing `arith.index_castui` for the second consumer layout also
-// clones its `vector.step` operand, so the two casts do not share a producer.
-// CHECK-LABEL: gpu.func @step_index_castui_clone_via_anchor
-// CHECK: %[[STEP0:.*]] = vector.step
-// CHECK: %[[STEP1:.*]] = vector.step
-// CHECK: arith.index_castui %[[STEP0]] {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex> to vector<32xi32>
-// CHECK: arith.index_castui %[[STEP1]] {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>, dims = [1]>} : vector<32xindex> to vector<32xi32>
-// CHECK-NOT: xegpu.convert_layout {{.*}} : vector<32xi32>
-gpu.func @step_index_castui_clone_via_anchor(%arg0: memref<32x32xi32>) kernel {
- %step = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex>
- %cast = arith.index_castui %step {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>, dims = [0]>} : vector<32xindex> to vector<32xi32>
- %col2d = vector.shape_cast %cast {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 1]>} : vector<32xi32> to vector<32x1xi32>
- %colb = vector.broadcast %col2d {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} : vector<32x1xi32> to vector<32x32xi32>
- %row2d = vector.shape_cast %cast {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} : vector<32xi32> to vector<1x32xi32>
- %rowb = vector.broadcast %row2d {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} : vector<1x32xi32> to vector<32x32xi32>
- %off = arith.addi %colb, %rowb {layout_result_0 = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>} : vector<32x32xi32>
- %tdesc = xegpu.create_nd_tdesc %arg0 : memref<32x32xi32> -> !xegpu.tensor_desc<32x32xi32>
- xegpu.store_nd %off, %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [32, 1], sg_data = [1, 32]>}> : vector<32x32xi32>, !xegpu.tensor_desc<32x32xi32>
- 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/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
index aacfd1a563d972..61b666a4b1a313 100644
--- a/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
+++ b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
@@ -361,6 +361,37 @@ struct TestXeGPUResolveLayoutConflicts
}
};
+struct TestXeGPURematerializeConversionSources
+ : public PassWrapper<TestXeGPURematerializeConversionSources,
+ OperationPass<gpu::GPUModuleOp>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+ TestXeGPURematerializeConversionSources)
+
+ StringRef getArgument() const final {
+ return "test-xegpu-rematerialize-conversion-sources";
+ }
+
+ StringRef getDescription() const final {
+ return "Test replacing XeGPU layout conversions of trivially "
+ "rematerializable values by recomputation.";
+ }
+
+ void getDependentDialects(::mlir::DialectRegistry ®istry) const override {
+ registry.insert<xegpu::XeGPUDialect>();
+ registry.insert<gpu::GPUDialect>();
+ registry.insert<vector::VectorDialect>();
+ }
+
+ TestXeGPURematerializeConversionSources() = default;
+ TestXeGPURematerializeConversionSources(
+ const TestXeGPURematerializeConversionSources &pass) = default;
+
+ void runOnOperation() override {
+ OpBuilder builder(&getContext());
+ xegpu::rematerializeConversionSources(builder, getOperation());
+ }
+};
+
struct TestXeGPUSinkElementwiseConversions
: public PassWrapper<TestXeGPUSinkElementwiseConversions,
OperationPass<gpu::GPUModuleOp>> {
@@ -588,6 +619,7 @@ void registerTestXeGPULowerings() {
PassRegistration<TestXeGPUSgToLaneDistribute>();
PassRegistration<TestXeGPUPropagateLayouts>();
PassRegistration<TestXeGPUResolveLayoutConflicts>();
+ PassRegistration<TestXeGPURematerializeConversionSources>();
PassRegistration<TestXeGPUSinkElementwiseConversions>();
PassRegistration<TestXeGPUArrayLengthOptimization>();
PassRegistration<TestXeGPUCoalesceGatherScatter>();
More information about the Mlir-commits
mailing list