[Mlir-commits] [mlir] 33bb824 - [MLIR][XeGPU] Rematerialize layout conversion sources after conflict resolution (#230669)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Oct 23 20:46:06 PDT 2026
Author: Nishant Patel
Date: 2026-10-23T20:46:01-07:00
New Revision: 33bb8249a9edddb8fcbf5002d9a012ffd772dd03
URL: https://github.com/llvm/llvm-project/commit/33bb8249a9edddb8fcbf5002d9a012ffd772dd03
DIFF: https://github.com/llvm/llvm-project/commit/33bb8249a9edddb8fcbf5002d9a012ffd772dd03.diff
LOG: [MLIR][XeGPU] Rematerialize layout conversion sources after conflict resolution (#230669)
This PR separates layout conflict resolution from rematerialization. The
resolver now inserts xegpu.convert_layout for conflicts ....a following
step clones eligible, cheap vector producer chains in the target layout
and replaces those conversions. This avoids SLM round trips.
Added:
mlir/test/Dialect/XeGPU/rematerialize-conversion-sources.mlir
Modified:
mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index 8921380093e108..d587dd2bde7303 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 of cheap producer chains by recomputation.
+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,18 @@ 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` is cheap to recompute in any
+/// layout, i.e. `rematerialize` would succeed: 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 without layout attributes. 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 rematerialize(OpBuilder &builder, Value value);
} // 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 649231342e51bd..8e260fc03fdec8 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,84 @@ 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::rematerialize(OpBuilder &builder, Value value) {
+ 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);
+ removeLayoutAttrs(clone);
+ }
+ 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 3629e122b1bf92..69f927fc0d4943 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() {
@@ -1713,26 +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)) {
- 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(
@@ -1835,6 +1812,42 @@ static LogicalResult updateOpWithForwardFill(mlir::OpBuilder &builder,
return success();
}
+/// Replace layout conversions of cheap producer chains by recomputation.
+///
+/// %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
+/// %cast = arith.index_castui %step {L1}
+/// %cast' = arith.index_castui %step'
+/// use(%cast) {L1}
+/// use(%cast') {L2}
+///
+/// Clone the whole vector producer chain so the original and rematerialized
+/// values can serve consumers with
diff erent layouts. Clones carry no layout
+/// attributes; layout-dependent passes recover them from the use points. 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::rematerialize(builder, convert.getSource());
+ if (!clone)
+ continue;
+ convert.getResult().replaceAllUsesWith(clone);
+ convert.erase();
+ }
+}
+
/// Optimize elementwise operations by sinking costly layout conversion.
///
/// %m = vector.create_mask ... {L1}
@@ -1999,7 +2012,6 @@ struct XeGPUPropagateLayoutPass final
: XeGPUPropagateLayoutBase(std::move(options)) {}
void runOnOperation() override;
};
-
} // namespace
LogicalResult xegpu::propagateLayouts(OpBuilder &builder, Operation *target,
@@ -2108,6 +2120,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..8d66079652a460
--- /dev/null
+++ b/mlir/test/Dialect/XeGPU/rematerialize-conversion-sources.mlir
@@ -0,0 +1,204 @@
+// RUN: mlir-opt --test-xegpu-rematerialize-conversion-sources -allow-unregistered-dialect -split-input-file %s | FileCheck %s
+
+// Row and column slices require
diff erent 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 : 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 dense<32> : vector<32xindex>
+// CHECK-DAG: %[[STEP1:.*]] = vector.step : vector<32xindex>
+// CHECK: %[[MUL1:.*]] = arith.muli %[[STEP1]], %[[CST1]] : 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]>
+
+// 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 : 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]] : 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]] : 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 : 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 d822b487e53880..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
diff erent 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,31 +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
-}
-
// 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