[Mlir-commits] [mlir] 4712ca8 - [MLIR][XeGPU] Recover temporary layout from Anchor Layout (#191947)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Apr 20 20:54:29 PDT 2026
Author: Jianhui Li
Date: 2026-04-20T20:54:24-07:00
New Revision: 4712ca81420ee95a955b3cc2e01ad5b9ae49b409
URL: https://github.com/llvm/llvm-project/commit/4712ca81420ee95a955b3cc2e01ad5b9ae49b409
DIFF: https://github.com/llvm/llvm-project/commit/4712ca81420ee95a955b3cc2e01ad5b9ae49b409.diff
LOG: [MLIR][XeGPU] Recover temporary layout from Anchor Layout (#191947)
This PR refactor the recoverTemporaryLayout() method so that the
temporary layout is recovered from anchor layout, not from any user
specified temporary layout.
Added:
mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir
Modified:
mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
mlir/test/Dialect/XeGPU/sg-to-wi-experimental-unit.mlir
mlir/test/Dialect/XeGPU/sg-to-wi-experimental.mlir
mlir/test/Dialect/XeGPU/subgroup-distribute.mlir
mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops.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 2172a24bb7a59..83eb939cf1bec 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -61,6 +61,11 @@ void removeLayoutAttr(const T &operandOrResult);
/// applied recursively to the contained operations
void removeLayoutAttrs(Operation *op);
+/// Removes the temporary layout attributes for each OpOperand and OpResult of
+/// the given operation. Recursive for contained operations if the given
+/// operation contains regions.
+void removeTemporaryLayoutAttrs(Operation *op);
+
/// Updates the NamedAttribute sequence by dropping sg-layout and
/// sg-data information from any DistributeLayoutAttr found.
SmallVector<NamedAttribute>
@@ -118,6 +123,11 @@ DistributeLayoutAttr
inferMaskOffsetLayoutForScatterIO(DistributeLayoutAttr payloadLayout,
int chunkSize);
+/// Infers the source layout attribute for an operand using result layout
+/// attribute
+DistributeLayoutAttr
+inferSourceLayoutFromResult(OpOperand &operand, DistributeLayoutAttr resLayout);
+
/// Sets up layout for Multi-Reduction operations by creating a SliceAttr for
/// the result.
///
@@ -192,7 +202,7 @@ setupDpasLayout(LayoutKind layoutKind, VectorType aTy, VectorType bTy,
/// Gets the expected layout for a given consumer operand. This will check if
/// the owning operation of the consumer operand is one of the special layout
/// users and determine the expected layout accordingly.
-xegpu::DistributeLayoutAttr getConsumerLayoutAt(OpOperand &operand);
+DistributeLayoutAttr getConsumerLayoutAt(OpOperand &operand);
} // namespace xegpu
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index ffbd3b497aae8..7d48315eec6ff 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -18,12 +18,15 @@
#include "mlir/Dialect/LLVMIR/XeVMDialect.h"
#include "mlir/Dialect/SCF/Transforms/Patterns.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/ValueRange.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Transforms/DialectConversion.h"
+#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/Support/FormatVariadic.h"
#include <cstdint>
#include <numeric>
@@ -80,30 +83,225 @@ xegpu::dropInstDataOnAttrs(ArrayRef<NamedAttribute> attrs) {
return out;
}
-// Attach layout attributes to all vector-type operands of operations within
-// the given operation's region. Reports an error if any vector operand lacks
-// a layout attribute.
-bool xegpu::recoverTemporaryLayouts(Operation *rootOp) {
- auto result = rootOp->walk([&](Operation *op) {
- for (OpOperand &operand : op->getOpOperands()) {
- // Layouts are needed for vector type only.
- if (!isa<VectorType>(operand.get().getType()))
- continue;
- // Skip block arguments since they don't have defining ops to attach
- // layout attributes to.
- if (isa<BlockArgument>(operand.get()))
+// the walkRegionBackward() is a recursive function
+// the input rootOp is the function operation, which is also a region op.
+// it recursively processes the region op in reverse topological order.
+static void walkRegionBackward(Region ®ion,
+ llvm::function_ref<void(Operation *)> visit) {
+
+ // Use post-order traversal to process blocks in reverse topological order.
+ // This ensures that use blocks are visited before def blocks, which is
+ // required for backward layout propagation.
+ if (region.empty())
+ return;
+ llvm::ReversePostOrderTraversal<Region *> rpot(®ion);
+ SmallVector<Block *> blocks(rpot.begin(), rpot.end());
+ for (Block *block : llvm::reverse(blocks)) {
+ // ops: back -> front
+ for (Operation &op : llvm::reverse(*block)) {
+ // make sure we first visit inside the region op (so yield op first)
+ // and then move to region op itself
+ // Regions are iterated in forward order so that for multi-region ops
+ // like scf.while, earlier regions (e.g., "before/cond") are processed
+ // first. This ensures that when a later region's terminator (e.g., "do"
+ // yield) needs the layout of an earlier region's block args, those
+ // layouts are already available from use points.
+ for (Region &nested : op.getRegions())
+ walkRegionBackward(nested, visit);
+
+ visit(&op);
+ }
+ }
+}
+
+static xegpu::DistributeLayoutAttr getLayoutFromUsePoints(Value result) {
+ xegpu::DistributeLayoutAttr layout = nullptr;
+ for (OpOperand &use : result.getUses()) {
+ if (auto tmpLayout = xegpu::getDistributeLayoutAttr(use)) {
+ if (!layout)
+ layout = tmpLayout;
+ break;
+ }
+ }
+ return layout;
+}
+
+// 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;
+
+ OpResult result = op->getResult(0);
+ xegpu::DistributeLayoutAttr resLayout = getLayoutFromUsePoints(result);
+ Type resultType = result.getType();
+
+ // recover layout for tensor Descriptor type, which is a special case since
+ // its layout is not stored as an attribute but encoded in the type itself.
+ // For vector type, we attach the layout as an attribute to op.
+ if (auto tensorDescTy = dyn_cast<xegpu::TensorDescType>(resultType)) {
+ auto layout = tensorDescTy.getLayoutAttr();
+ if (!layout) {
+ auto typeWithLayout = xegpu::TensorDescType::get(
+ tensorDescTy.getContext(), tensorDescTy.getShape(),
+ tensorDescTy.getElementType(), tensorDescTy.getEncoding(), resLayout);
+ result.setType(typeWithLayout);
+ }
+ }
+ if (isa<VectorType>(resultType) && resLayout)
+ xegpu::setTemporaryLayout(result, resLayout);
+
+ for (OpOperand &opr : op->getOpOperands()) {
+ xegpu::DistributeLayoutAttr operandLayout =
+ xegpu::inferSourceLayoutFromResult(opr, resLayout);
+ if (isa<VectorType>(opr.get().getType()) && operandLayout)
+ xegpu::setTemporaryLayout(opr, operandLayout);
+ }
+}
+
+// Propagate layout from region op results and sibling region block args
+// to yield/condition operands. For each successor of this terminator:
+// - Parent successor: propagate from parent op's result layouts (use points).
+// - Region successor: propagate from target region's block arg layouts (use
+// points), e.g., scf.yield in "after/do" region propagates to "before/cond"
+// block args.
+static void propagateRegionResultsToYieldOperands(
+ mlir::RegionBranchTerminatorOpInterface yieldOp) {
+ auto regionBranchOp =
+ dyn_cast<RegionBranchOpInterface>(yieldOp->getParentOp());
+ if (!regionBranchOp)
+ return;
+
+ SmallVector<RegionSuccessor> successors;
+ SmallVector<Attribute> operandAttrs(yieldOp->getNumOperands(), nullptr);
+ yieldOp.getSuccessorRegions(operandAttrs, successors);
+
+ for (const RegionSuccessor &successor : successors) {
+ OperandRange succOps = yieldOp.getSuccessorOperands(successor);
+ if (succOps.empty())
+ continue;
+ unsigned beginIdx = succOps.getBeginOperandIndex();
+ ValueRange successorInputs = regionBranchOp.getSuccessorInputs(successor);
+ unsigned count = std::min<unsigned>(succOps.size(), successorInputs.size());
+
+ for (unsigned i = 0; i < count; ++i) {
+ xegpu::DistributeLayoutAttr layout;
+ if (successor.isParent()) {
+ // For parent successor, get layout from external use points of the
+ // parent op's results.
+ layout = getLayoutFromUsePoints(regionBranchOp->getResult(i));
+ if (layout)
+ xegpu::setTemporaryLayout(regionBranchOp->getResult(i), layout);
+ } else {
+ // For region successor, get layout from the target region's block
+ // arg use points (e.g., "before/cond" region args for scf.while
+ // "after/do" yield).
+ layout = getLayoutFromUsePoints(successorInputs[i]);
+ }
+ if (!layout)
continue;
- auto layout = xegpu::getDistributeLayoutAttr(operand.get());
- if (!layout) {
- op->emitWarning("Could not find layout attribute for operand ")
- << operand.getOperandNumber() << " of operation " << op->getName();
+ if (isa<VectorType>(succOps[i].getType()))
+ xegpu::setTemporaryLayout(yieldOp->getOpOperand(beginIdx + i), layout);
+ }
+ }
+}
+
+// Propagate layout from region arguments to region op's init operands. This
+// sets the temporary layout for region arguments and init operands.
+static void propagateRegionArgsToInits(mlir::RegionBranchOpInterface regionOp) {
+ // Iterate all regions of the region op. For each block argument that has a
+ // layout (determined from its use points), trace back to find the
+ // corresponding init operand of the regionOp and set the layout on it.
+ // This works generically for scf.for, scf.while, and other
+ // RegionBranchOpInterface ops.
+ for (Region ®ion : regionOp->getRegions()) {
+ RegionSuccessor regionSuccessor(®ion);
+ // Use getSuccessorInputs to get the block arguments that correspond to
+ // predecessor operands. This correctly handles ops like scf.for where
+ // the induction variable is a block arg but not a successor input.
+ ValueRange successorInputs = regionOp.getSuccessorInputs(regionSuccessor);
+ for (auto [inputIdx, regionArg] : llvm::enumerate(successorInputs)) {
+ auto layout = getLayoutFromUsePoints(regionArg);
+ if (!layout)
continue;
+
+ // Recover layout for tensor_desc block args by updating the type.
+ if (auto tensorDescTy =
+ dyn_cast<xegpu::TensorDescType>(regionArg.getType())) {
+ if (!tensorDescTy.getLayoutAttr()) {
+ auto typeWithLayout = xegpu::TensorDescType::get(
+ tensorDescTy.getContext(), tensorDescTy.getShape(),
+ tensorDescTy.getElementType(), tensorDescTy.getEncoding(),
+ layout);
+ regionArg.setType(typeWithLayout);
+ }
+ }
+
+ // Find all predecessor values that flow into this block argument.
+ SmallVector<Value> predValues;
+ regionOp.getPredecessorValues(regionSuccessor, inputIdx, predValues);
+ for (Value predVal : predValues) {
+ // Match predecessor value to an operand of the regionOp.
+ for (OpOperand &operand : regionOp->getOpOperands()) {
+ if (operand.get() == predVal)
+ xegpu::setTemporaryLayout(operand, layout);
+ }
}
- xegpu::setTemporaryLayout(operand, layout);
}
- return WalkResult::advance();
+ }
+}
+
+// Prerequisite for Layout Recovery
+// It relies on the following invariant:
+// 1. there is no layout conflict between
diff erent uses of the same definition.
+// 2. each definition has a well-defined layout requirement at its use point.
+// - Every definition must have at least one use that appears after it in
+// topological order.
+// - TODO: If a definition has no such use (e.g., a loop result or region
+// output), an explicit convert_layout operation is inserted to create a
+// use.
+// - Only the result of convert_layout is permitted to have no subsequent
+// use.
+//
+// The recovery proceeds by scanning the operation in reverse topological order
+// as follows:
+// For regular operations: First the result layouts are propagated from uses.
+// Then the result layouts are propagated to operands.
+//
+// For region operations (e.g., loops):
+// - When backward propagation reaches a region op, it sets the layout of
+// the region op’s results according to use points like regular ops.
+// - Then, the result layouts (such as a loop output) are propagated to
+// their corresponding operands in the yield.
+// - When backward propagation reaches the first operation inside the
+// region, the pass examines the region op’s initialization list,
+// propagating from region arguments to the corresponding initialization
+// operands.
+// - This ensures that layouts are consistently propagated
+// across region boundaries while preserving a single well-defined use for
+// each definition at the region-op level.
+bool xegpu::recoverTemporaryLayouts(Operation *rootOp) {
+ auto processFunc = [&](Region &body, StringRef funcName) {
+ walkRegionBackward(body, [&](Operation *op) {
+ if (auto regionOp = dyn_cast<mlir::RegionBranchOpInterface>(op)) {
+ propagateRegionArgsToInits(regionOp);
+ } else if (auto yieldOp =
+ dyn_cast<mlir::RegionBranchTerminatorOpInterface>(op)) {
+ propagateRegionResultsToYieldOperands(yieldOp);
+ } else if (!dyn_cast<xegpu::AnchorLayoutInterface>(op)) {
+ propagateResultsToRegularOperands(op);
+ }
+ });
+ };
+
+ rootOp->walk([&](func::FuncOp func) {
+ processFunc(func.getBody(), func.getSymName());
});
- return !result.wasInterrupted();
+ rootOp->walk([&](gpu::GPUFuncOp func) {
+ processFunc(func.getBody(), func.getName());
+ });
+
+ return true;
}
template <typename T, typename>
@@ -135,6 +333,18 @@ void xegpu::removeLayoutAttrs(Operation *op) {
});
}
+void xegpu::removeTemporaryLayoutAttrs(Operation *op) {
+ op->walk([&](Operation *nestOp) {
+ SmallVector<StringAttr> attrsToRemove;
+ for (auto namedAttr : nestOp->getDiscardableAttrs()) {
+ if (isa<xegpu::DistributeLayoutAttr>(namedAttr.getValue()))
+ attrsToRemove.push_back(namedAttr.getName());
+ }
+ for (auto attrName : attrsToRemove)
+ nestOp->removeDiscardableAttr(attrName);
+ });
+}
+
/// Infers the source layout attribute for a broadcast operation given the
/// result layout attribute, result shape, source shape.
xegpu::DistributeLayoutAttr
@@ -1122,20 +1332,19 @@ xegpu::setupDpasLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
return std::nullopt;
}
-xegpu::DistributeLayoutAttr xegpu::getConsumerLayoutAt(OpOperand &operand) {
+xegpu::DistributeLayoutAttr
+xegpu::inferSourceLayoutFromResult(OpOperand &operand,
+ xegpu::DistributeLayoutAttr resLayout) {
+ if (!resLayout)
+ return nullptr;
Operation *op = operand.getOwner();
unsigned idx = operand.getOperandNumber();
- xegpu::DistributeLayoutAttr resLayout;
- if (op->getNumResults() == 1)
- resLayout = xegpu::getDistributeLayoutAttr(op->getResult(0));
// For vector::BroadcastOp, infer the source layout from the result layout.
if (auto broadcast = dyn_cast<vector::BroadcastOp>(op)) {
- if (!resLayout)
- return xegpu::DistributeLayoutAttr();
auto srcTy = dyn_cast<VectorType>(broadcast.getSourceType());
if (!srcTy)
- return xegpu::DistributeLayoutAttr();
+ return nullptr;
return xegpu::inferBroadcastSourceLayout(
resLayout, broadcast.getResultVectorType().getShape(),
srcTy.getShape());
@@ -1145,8 +1354,6 @@ xegpu::DistributeLayoutAttr xegpu::getConsumerLayoutAt(OpOperand &operand) {
// using reduction dims. Acc operand is expected to have the same layout as
// the result.
if (auto reduction = dyn_cast<vector::MultiDimReductionOp>(op)) {
- if (!resLayout)
- return xegpu::DistributeLayoutAttr();
if (idx == 0) {
SmallVector<int64_t> reductionDims(reduction.getReductionDims());
return xegpu::inferMultiReductionSourceLayout(resLayout, reductionDims);
@@ -1155,17 +1362,12 @@ xegpu::DistributeLayoutAttr xegpu::getConsumerLayoutAt(OpOperand &operand) {
return resLayout;
}
- if (auto reduction = dyn_cast<vector::ReductionOp>(op)) {
- if (!resLayout)
- return xegpu::DistributeLayoutAttr();
+ if (auto reduction = dyn_cast<vector::ReductionOp>(op))
return xegpu::inferReductionSourceLayout(resLayout);
- }
// For vector::BitCastOp, infer source layout from result layout using
// element type bitwidths.
if (auto bitcast = dyn_cast<vector::BitCastOp>(op)) {
- if (!resLayout)
- return xegpu::DistributeLayoutAttr();
int resElemBitWidth =
bitcast.getResultVectorType().getElementType().getIntOrFloatBitWidth();
int srcElemBitWidth =
@@ -1177,8 +1379,6 @@ xegpu::DistributeLayoutAttr xegpu::getConsumerLayoutAt(OpOperand &operand) {
// For vector::ShapeCastOp, infer source layout from result layout using
// shapes.
if (auto shapeCast = dyn_cast<vector::ShapeCastOp>(op)) {
- if (!resLayout)
- return xegpu::DistributeLayoutAttr();
return xegpu::inferShapeCastSourceLayout(
resLayout, shapeCast.getResultVectorType().getShape(),
shapeCast.getSourceVectorType().getShape());
@@ -1187,12 +1387,11 @@ xegpu::DistributeLayoutAttr xegpu::getConsumerLayoutAt(OpOperand &operand) {
// For vector::InsertStridedSliceOp, infer source layout from result layout.
// Dest vector must have the same layout as the result.
if (auto insertSlice = dyn_cast<vector::InsertStridedSliceOp>(op)) {
- if (!resLayout)
- return xegpu::DistributeLayoutAttr();
- if (idx == 0)
+ if (idx == 0) {
return xegpu::inferInsertStridedSliceSourceLayout(
resLayout, insertSlice.getDestVectorType().getShape(),
insertSlice.getSourceVectorType().getShape());
+ }
if (idx == 1)
return resLayout;
}
@@ -1200,20 +1399,29 @@ xegpu::DistributeLayoutAttr xegpu::getConsumerLayoutAt(OpOperand &operand) {
// For vector::TransposeOp, infer source layout from result layout using
// permutation.
if (auto transpose = dyn_cast<vector::TransposeOp>(op)) {
- if (!resLayout)
- return xegpu::DistributeLayoutAttr();
return xegpu::inferTransposeSourceLayout(resLayout,
transpose.getPermutation());
}
+ // For vector::ExtractStridedSliceOp, simply return result layout
+ if (dyn_cast<vector::ExtractStridedSliceOp>(op))
+ return resLayout;
// For elementwise operations, all operands must have the same layout as the
// result.
- if (OpTrait::hasElementwiseMappableTraits(op) && op->getNumResults() == 1) {
- if (!resLayout)
- return xegpu::DistributeLayoutAttr();
+ if (OpTrait::hasElementwiseMappableTraits(op) && op->getNumResults() == 1)
return resLayout;
- }
- // TODO: Handle more cases as needed here.
+
+ return nullptr;
+}
+
+xegpu::DistributeLayoutAttr xegpu::getConsumerLayoutAt(OpOperand &operand) {
+ Operation *op = operand.getOwner();
+ xegpu::DistributeLayoutAttr resLayout;
+ if (op->getNumResults() == 1)
+ resLayout = xegpu::getDistributeLayoutAttr(op->getResult(0));
+ auto inferredOperandLayout = inferSourceLayoutFromResult(operand, resLayout);
+ if (inferredOperandLayout)
+ return inferredOperandLayout;
// By default, assume no layout conflict and return the current layout of
// the operand.
return xegpu::getDistributeLayoutAttr(operand.get());
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
index 3496756e8a6d3..8ade936724480 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
@@ -598,16 +598,7 @@ struct XeGPUPeepHoleOptimizerPass final
RewritePatternSet emptyPatterns(ctx);
(void)applyPatternsGreedily(getOperation(), std::move(emptyPatterns));
- // Remove the temporary layout after all patterns are applied.
- getOperation()->walk([](Operation *op) {
- SmallVector<StringAttr> attrsToRemove;
- for (auto namedAttr : op->getDiscardableAttrs()) {
- if (isa<xegpu::DistributeLayoutAttr>(namedAttr.getValue()))
- attrsToRemove.push_back(namedAttr.getName());
- }
- for (auto attrName : attrsToRemove)
- op->removeDiscardableAttr(attrName);
- });
+ xegpu::removeTemporaryLayoutAttrs(getOperation());
}
};
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 686cb20e1976e..43998ed41f7aa 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -1652,16 +1652,9 @@ LogicalResult xegpu::resolveLayoutConflicts(Operation *target) {
}
void XeGPUPropagateLayoutPass::runOnOperation() {
- // Clean up temporary layout attributes
- getOperation()->walk([](Operation *op) {
- SmallVector<StringAttr> attrsToRemove;
- for (auto namedAttr : op->getDiscardableAttrs()) {
- if (isa<xegpu::DistributeLayoutAttr>(namedAttr.getValue()))
- attrsToRemove.push_back(namedAttr.getName());
- }
- for (auto attrName : attrsToRemove)
- op->removeDiscardableAttr(attrName);
- });
+
+ xegpu::removeTemporaryLayoutAttrs(getOperation());
+
xegpu::LayoutKind layoutKind;
if (this->layoutKind == "lane") {
layoutKind = xegpu::LayoutKind::Lane;
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
index b086a6571ddb4..c153db431c035 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
@@ -1609,6 +1609,8 @@ void XeGPUSgToWiDistributeExperimentalPass::runOnOperation() {
}
});
}
+
+ xegpu::removeTemporaryLayoutAttrs(getOperation());
}
void xegpu::populateXeGPUSgToWiDistributeTypeConversions(
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
index 9459164e4d48b..71fff428cceb9 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
@@ -2280,4 +2280,6 @@ void XeGPUSubgroupDistributePass::runOnOperation() {
op->erase();
return WalkResult::advance();
});
+
+ xegpu::removeTemporaryLayoutAttrs(getOperation());
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index 1e867b9c34069..e083507173d31 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -498,9 +498,6 @@ struct WgToSgVectorBroadcastOp
VectorType newResultType =
VectorType::get(sgShape, resultType.getElementType());
- if (!layout.isDistributable(SmallVector<int64_t>(wgShape)))
- return failure();
-
SmallVector<Value> newBroadcastOps;
auto distSource = adaptor.getOperands().front();
int numDistributions = count / distSource.size();
@@ -1778,17 +1775,5 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
applyPartialConversion(getOperation(), target, std::move(patterns))))
return signalPassFailure();
- // Remove layout attributes from SCF ops
- getOperation()->walk([](Operation *op) {
- if (!isa<RegionBranchOpInterface, RegionBranchTerminatorOpInterface>(op))
- return;
-
- SmallVector<StringAttr> attrsToRemove;
- for (auto namedAttr : op->getDiscardableAttrs()) {
- if (isa<xegpu::DistributeLayoutAttr>(namedAttr.getValue()))
- attrsToRemove.push_back(namedAttr.getName());
- }
- for (auto attrName : attrsToRemove)
- op->removeDiscardableAttr(attrName);
- });
+ xegpu::removeTemporaryLayoutAttrs(getOperation());
}
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 9098eb7e4815b..2d1ce6eea17aa 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -191,13 +191,27 @@ xegpu::getDistributeLayoutAttr(const OpOperand &opr) {
if (idx == 0)
return layout;
- // For store operations (StoreScatterOp, StoreNdOp, StoreMatrixOp),
+ // For StoreNdOp and StoreMatrixOp,
// the layout is valid for the first two operands: value and memref/tdesc.
- // For other operations, the layout applies to the first operand only.
- if (isa<xegpu::StoreScatterOp, xegpu::StoreNdOp, xegpu::StoreMatrixOp>(
- op) &&
- (idx < 2))
+ if (isa<xegpu::StoreNdOp, xegpu::StoreMatrixOp>(op) && (idx < 2))
+ return layout;
+
+ if (isa<xegpu::StoreScatterOp>(op)) {
+ xegpu::StoreScatterOp store(op);
+ int chunkSize = store.getChunkSize().value_or(1);
+ if (layout && idx >= 2 && chunkSize > 1)
+ return layout.dropDims(llvm::to_vector(
+ llvm::seq<int64_t>(layout.getRank() - 1, layout.getRank())));
+ return layout;
+ }
+ if (isa<xegpu::LoadGatherOp>(op)) {
+ xegpu::LoadGatherOp load(op);
+ int chunkSize = load.getChunkSize().value_or(1);
+ if (layout && idx >= 1 && chunkSize > 1)
+ return layout.dropDims(llvm::to_vector(
+ llvm::seq<int64_t>(layout.getRank() - 1, layout.getRank())));
return layout;
+ }
}
std::string layoutName = xegpu::getTemporaryLayoutName(opr);
diff --git a/mlir/test/Dialect/XeGPU/sg-to-wi-experimental-unit.mlir b/mlir/test/Dialect/XeGPU/sg-to-wi-experimental-unit.mlir
index 4c3727388831b..08b334ddec3fc 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-wi-experimental-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-wi-experimental-unit.mlir
@@ -136,6 +136,11 @@ gpu.func @elementwise() {
%3 = arith.addf %0, %2
{layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
: vector<16x16xf32>
+ %cl3 = xegpu.convert_layout %3
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+ }> : vector<16x16xf32>
gpu.return
}
@@ -145,6 +150,11 @@ gpu.func @elementwise() {
gpu.func @arith_constant() {
%0 = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
dense<1.0> : vector<16x16xf32>
+ %cl0 = xegpu.convert_layout %0
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+ }> : vector<16x16xf32>
gpu.return
}
@@ -363,6 +373,11 @@ gpu.func @vector_multi_reduction_dim1_distributed_dim1_reduction(%laneid: index)
layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [1]>
}
[1] : vector<2x16xf32> to vector<2xf32>
+ %cl1 = xegpu.convert_layout %1
+ <{
+ input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [1]>,
+ target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [1]>
+ }> : vector<2xf32>
gpu.return
}
@@ -428,6 +443,11 @@ gpu.func @vector_multi_reduction_dim0_distributed_dim0_reduction(%laneid: index)
layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>, dims = [0]>
}
[0] : vector<16x2xf32> to vector<2xf32>
+ %cl1 = xegpu.convert_layout %1
+ <{
+ input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>, dims = [0]>,
+ target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>, dims = [0]>
+ }> : vector<2xf32>
gpu.return
}
@@ -453,6 +473,11 @@ gpu.func @vector_multi_reduction_dim1_distributed_dim0_reduction(%laneid: index)
layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>
}
[0] : vector<4x16xf32> to vector<16xf32>
+ %cl1 = xegpu.convert_layout %1
+ <{
+ input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>,
+ target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>
+ }> : vector<16xf32>
gpu.return
}
@@ -478,6 +503,11 @@ gpu.func @vector_multi_reduction_dim0_distributed_dim1_reduction(%laneid: index)
layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>, dims = [1]>
}
[1] : vector<16x12xf32> to vector<16xf32>
+ %cl1 = xegpu.convert_layout %1
+ <{
+ input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>, dims = [1]>,
+ target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>, dims = [1]>
+ }> : vector<16xf32>
gpu.return
}
@@ -495,6 +525,11 @@ gpu.func @vector_transpose() {
layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
}
: vector<16x2xf32> to vector<2x16xf32>
+ %transpose2 = xegpu.convert_layout %transpose
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+ }> : vector<2x16xf32>
gpu.return
}
@@ -527,6 +562,11 @@ gpu.func @create_mask_1d(%m0: index) {
%mask = vector.create_mask %m0
{layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>}
: vector<16xi1>
+ %mask_cl = xegpu.convert_layout %mask
+ <{
+ input_layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>,
+ target_layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>
+ }> : vector<16xi1>
gpu.return
}
@@ -542,6 +582,11 @@ gpu.func @constant_mask_1d() {
%mask = vector.constant_mask [4]
{layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>}
: vector<16xi1>
+ %mask_cl = xegpu.convert_layout %mask
+ <{
+ input_layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>,
+ target_layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>
+ }> : vector<16xi1>
gpu.return
}
@@ -563,6 +608,11 @@ gpu.func @create_mask_2d(%m0: index, %m1: index) {
%mask = vector.create_mask %m0, %m1
{layout_result_0 = #xegpu.layout<lane_layout = [8, 2], lane_data = [1, 1]>}
: vector<8x4xi1>
+ %mask_cl = xegpu.convert_layout %mask
+ <{
+ input_layout = #xegpu.layout<lane_layout = [8, 2], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [8, 2], lane_data = [1, 1]>
+ }> : vector<8x4xi1>
gpu.return
}
@@ -585,6 +635,11 @@ gpu.func @constant_mask_2d() {
%mask = vector.constant_mask [2, 3]
{layout_result_0 = #xegpu.layout<lane_layout = [8, 2], lane_data = [1, 1]>}
: vector<8x4xi1>
+ %mask_cl = xegpu.convert_layout %mask
+ <{
+ input_layout = #xegpu.layout<lane_layout = [8, 2], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [8, 2], lane_data = [1, 1]>
+ }> : vector<8x4xi1>
gpu.return
}
@@ -615,6 +670,11 @@ gpu.func @vector_multi_reduction_3d_leading_unit_dim_lane_local() {
layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [1]>
}
[1] : vector<1x16x32xf32> to vector<1x32xf32>
+ %cl1 = xegpu.convert_layout %1
+ <{
+ input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [1]>,
+ target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [1]>
+ }> : vector<1x32xf32>
gpu.return
}
@@ -648,6 +708,11 @@ gpu.func @vector_multi_reduction_3d_leading_unit_dim_cross_lane() {
layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16, 1], lane_data = [1, 1, 1]>, dims = [1]>
}
[1] : vector<1x16x2xf32> to vector<1x2xf32>
+ %cl1 = xegpu.convert_layout %1
+ <{
+ input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16, 1], lane_data = [1, 1, 1]>, dims = [1]>,
+ target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16, 1], lane_data = [1, 1, 1]>, dims = [1]>
+ }> : vector<1x2xf32>
gpu.return
}
@@ -660,6 +725,11 @@ gpu.func @vector_extract_from_2d() {
%0 = vector.extract %src[0]
{layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>}
: vector<16xf32> from vector<4x16xf32>
+ %cl0 = xegpu.convert_layout %0
+ <{
+ input_layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>,
+ target_layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>
+ }> : vector<16xf32>
gpu.return
}
@@ -672,6 +742,11 @@ gpu.func @vector_extract_from_2d_offset2() {
%0 = vector.extract %src[2]
{layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>}
: vector<16xf32> from vector<8x16xf32>
+ %cl0 = xegpu.convert_layout %0
+ <{
+ input_layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>,
+ target_layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>
+ }> : vector<16xf32>
gpu.return
}
@@ -687,6 +762,11 @@ gpu.func @vector_insert_into_2d() {
%0 = vector.insert %val, %dst[0]
{layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
: vector<16xf32> into vector<4x16xf32>
+ %cl0 = xegpu.convert_layout %0
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+ }> : vector<4x16xf32>
gpu.return
}
@@ -702,6 +782,11 @@ gpu.func @vector_insert_into_2d_offset2() {
%0 = vector.insert %val, %dst[2]
{layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
: vector<16xf32> into vector<8x16xf32>
+ %cl0 = xegpu.convert_layout %0
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+ }> : vector<8x16xf32>
gpu.return
}
@@ -715,6 +800,11 @@ gpu.func @vector_extract_strided_slice_distributed_dim_fully_extracted() {
layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
}
: vector<24x16xf32> to vector<8x16xf32>
+ %cl1 = xegpu.convert_layout %1
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+ }> : vector<8x16xf32>
gpu.return
}
@@ -728,6 +818,11 @@ gpu.func @vector_extract_strided_slice_inner_distributed() {
layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
}
: vector<24x64xf32> to vector<8x16xf32>
+ %cl1 = xegpu.convert_layout %1
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+ }> : vector<8x16xf32>
gpu.return
}
@@ -741,6 +836,11 @@ gpu.func @vector_extract_strided_slice_outer_distributed() {
layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>
}
: vector<32x16xf32> to vector<16x16xf32>
+ %cl1 = xegpu.convert_layout %1
+ <{
+ input_layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>
+ }> : vector<16x16xf32>
gpu.return
}
@@ -754,6 +854,11 @@ gpu.func @vector_extract_strided_slice_1d() {
layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>
}
: vector<64xf32> to vector<32xf32>
+ %cl1 = xegpu.convert_layout %1
+ <{
+ input_layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>,
+ target_layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>
+ }> : vector<32xf32>
gpu.return
}
@@ -767,6 +872,11 @@ gpu.func @vector_extract_strided_slice_partial_offsets() {
layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
}
: vector<24x16xf32> to vector<8x16xf32>
+ %cl1 = xegpu.convert_layout %1
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+ }> : vector<8x16xf32>
gpu.return
}
@@ -783,6 +893,11 @@ gpu.func @vector_insert_strided_slice_distributed_dim_fully_inserted() {
layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
}
: vector<16x16xf32> into vector<64x16xf32>
+ %cl2 = xegpu.convert_layout %2
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+ }> : vector<64x16xf32>
gpu.return
}
@@ -799,6 +914,11 @@ gpu.func @vector_insert_strided_slice_inner_distributed() {
layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
}
: vector<16x16xf32> into vector<64x32xf32>
+ %cl2 = xegpu.convert_layout %2
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+ }> : vector<64x32xf32>
gpu.return
}
@@ -815,6 +935,11 @@ gpu.func @vector_insert_strided_slice_outer_distributed() {
layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>
}
: vector<16x16xf32> into vector<48x32xf32>
+ %cl2 = xegpu.convert_layout %2
+ <{
+ input_layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>
+ }> : vector<48x32xf32>
gpu.return
}
@@ -831,6 +956,11 @@ gpu.func @vector_insert_strided_slice_1d() {
layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>
}
: vector<16xf32> into vector<48xf32>
+ %cl2 = xegpu.convert_layout %2
+ <{
+ input_layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>,
+ target_layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>
+ }> : vector<48xf32>
gpu.return
}
@@ -847,6 +977,11 @@ gpu.func @vector_insert_strided_slice_
diff erent_ranks() {
layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
}
: vector<16xf32> into vector<64x16xf32>
+ %cl2 = xegpu.convert_layout %2
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+ }> : vector<64x16xf32>
gpu.return
}
@@ -965,6 +1100,11 @@ gpu.func @elementwise_wrap_around_dim() {
: () -> vector<16x1xf16>
%1 = arith.negf %0 {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
: vector<16x1xf16>
+ %cl1 = xegpu.convert_layout %1
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+ }> : vector<16x1xf16>
gpu.return
}
}
@@ -979,6 +1119,11 @@ gpu.module @xevm_module {
// CHECK: %[[VEC:.*]] = vector.from_elements %[[REM2]] : vector<1xindex>
gpu.func @vector_step_slice() {
%0 = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>, dims = [0, 1, 2]>} : vector<16xindex>
+ %cl0 = xegpu.convert_layout %0
+ <{
+ input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>, dims = [0, 1, 2]>,
+ target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>, dims = [0, 1, 2]>
+ }> : vector<16xindex>
gpu.return
}
}
@@ -989,6 +1134,11 @@ gpu.module @xevm_module {
// CHECK: %[[VEC:.*]] = vector.from_elements %{{.*}} : vector<1xindex>
gpu.func @vector_step_slice_unit() {
%0 = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>, dims = [0, 1, 3]>} : vector<1xindex>
+ %cl0 = xegpu.convert_layout %0
+ <{
+ input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>, dims = [0, 1, 3]>,
+ target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 1, 16], lane_data = [1, 1, 1, 1]>, dims = [0, 1, 3]>
+ }> : vector<1xindex>
gpu.return
}
}
@@ -1006,6 +1156,11 @@ gpu.module @xevm_module {
// CHECK: %[[VEC:.*]] = vector.from_elements %[[V0]], %[[V1]], %[[V2]], %[[V3]] : vector<4xindex>
gpu.func @vector_step_slice_multi_dist() {
%0 = vector.step {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [2, 4, 2], lane_data = [1, 2, 1]>, dims = [0, 2]>} : vector<16xindex>
+ %cl0 = xegpu.convert_layout %0
+ <{
+ input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [2, 4, 2], lane_data = [1, 2, 1]>, dims = [0, 2]>,
+ target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [2, 4, 2], lane_data = [1, 2, 1]>, dims = [0, 2]>
+ }> : vector<16xindex>
gpu.return
}
}
@@ -1023,6 +1178,11 @@ gpu.func @vector_shapecast_rank_increasing() {
layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
}
: vector<16xf32> to vector<1x16xf32>
+ %cast_cl = xegpu.convert_layout %cast
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+ }> : vector<1x16xf32>
gpu.return
}
}
@@ -1040,6 +1200,11 @@ gpu.func @vector_shapecast_rank_reducing() {
layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>
}
: vector<1x16xf32> to vector<16xf32>
+ %cast_cl = xegpu.convert_layout %cast
+ <{
+ input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>,
+ target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>
+ }> : vector<16xf32>
gpu.return
}
}
@@ -1057,6 +1222,11 @@ gpu.func @vector_shapecast_rank_increasing_without_slicing_layout() {
layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
}
: vector<16xf32> to vector<1x16xf32>
+ %cast_cl = xegpu.convert_layout %cast
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+ }> : vector<1x16xf32>
gpu.return
}
}
@@ -1083,6 +1253,11 @@ gpu.module @xevm_module {
gpu.func @constant_wrap_around_dim() {
%0 = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
dense<1.0> : vector<16x1xf16>
+ %cl0 = xegpu.convert_layout %0
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+ }> : vector<16x1xf16>
gpu.return
}
}
@@ -1109,7 +1284,8 @@ gpu.module @xevm_module {
gpu.func @vector_broadcast_2d_to_2d_noop(%laneid: index) {
%0 = "some_op"() {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : () -> vector<16x1xf16>
%1 = vector.broadcast %0 {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : vector<16x1xf16> to vector<16x16xf16>
- "some_use"(%1) : (vector<16x16xf16>) -> ()
+ %2 = xegpu.convert_layout %1 <{input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<16x16xf16>
+ "some_use"(%2) : (vector<16x16xf16>) -> ()
gpu.return
}
}
@@ -1171,6 +1347,11 @@ gpu.func @vector_multi_reduction_1d_to_scalar() {
layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [16], lane_data = [1]>, dims = [0]>
}
[0] : vector<32xf32> to f32
+ %cl1 = xegpu.convert_layout %1
+ <{
+ input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [16], lane_data = [1]>, dims = [0]>,
+ target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [16], lane_data = [1]>, dims = [0]>
+ }> : f32
gpu.return
}
}
diff --git a/mlir/test/Dialect/XeGPU/sg-to-wi-experimental.mlir b/mlir/test/Dialect/XeGPU/sg-to-wi-experimental.mlir
index babb01c131792..c4401515414b6 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-wi-experimental.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-wi-experimental.mlir
@@ -235,28 +235,24 @@ gpu.func @gemm_with_postop(%arg0: memref<1024x1024xbf16>, %arg1: memref<1024x102
gpu.module @xevm_module{
gpu.func @load_dpas_postop_store(%arg0: memref<8x16xf16>, %arg1: memref<16x16xf16>, %arg2: memref<8x16xf32>) {
%c0 = arith.constant 0 : index
- %cst = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} dense<0.0> : vector<8x16xf32>
+ %cst = arith.constant dense<0.0> : vector<8x16xf32>
%0 = xegpu.create_nd_tdesc %arg0 : memref<8x16xf16>
-> !xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
%1 = xegpu.load_nd %0[%c0, %c0]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
- layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
+ {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
!xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x16xf16>
%2 = xegpu.create_nd_tdesc %arg1: memref<16x16xf16>
-> !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>>
%3 = xegpu.load_nd %2[%c0, %c0]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>,
- layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
+ {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>>
-> vector<16x16xf16>
%4 = xegpu.dpas %1, %3, %cst
{layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>,
- layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
- layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
: vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
%5 = math.exp %4
- {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
: vector<8x16xf32>
%6 = xegpu.create_nd_tdesc %arg2 : memref<8x16xf32> ->
!xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
@@ -279,7 +275,7 @@ gpu.module @xevm_module{
// CHECK-NEXT: scf.yield %[[LD_CAST]] : vector<1x8xf16>
// CHECK-NEXT: } else {
// CHECK-NEXT: scf.yield %[[CST]] : vector<1x8xf16>
-// CHECK-NEXT: } {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>}
+// CHECK-NEXT: }
// CHECK-NEXT: %[[IF_CAST:.*]] = vector.shape_cast %[[IF]] : vector<1x8xf16> to vector<8xf16>
// CHECK-NEXT: xegpu.store %[[IF_CAST]], %{{.*}}[%[[OFFSET]]], %[[MASK]] <{chunk_size = 8 : i64}>
// CHECK-SAME: vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
@@ -289,16 +285,13 @@ gpu.module @xevm_module{
%offset = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<12> : vector<16xindex>
%loaded = scf.if %pred -> (vector<16x8xf16>) {
%3 = xegpu.load %src[%offset], %1 <{chunk_size=8}> {
- layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>,
- layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>
+ layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>
} : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
scf.yield %3 : vector<16x8xf16>
} else {
- %3 = arith.constant {
- layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>
- } dense<12.> : vector<16x8xf16>
+ %3 = arith.constant dense<12.> : vector<16x8xf16>
scf.yield %3 : vector<16x8xf16>
- } { layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]> }
+ }
xegpu.store %loaded, %src[%offset], %1 <{chunk_size=8}> {layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>} : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
gpu.return
}
@@ -318,12 +311,11 @@ gpu.module @xevm_module{
gpu.module @xevm_module{
gpu.func @scatter_ops_scf_non_yield(%src: memref<256xf16>) {
%pred = llvm.mlir.poison : i1
- %1 = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<1>: vector<16xi1>
- %offset = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<12> : vector<16xindex>
+ %1 = arith.constant dense<1>: vector<16xi1>
+ %offset = arith.constant dense<12> : vector<16xindex>
scf.if %pred {
%3 = xegpu.load %src[%offset], %1 <{chunk_size=8}> {
- layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>,
- layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>
+ layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>
} : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
xegpu.store %3, %src[%offset], %1 <{chunk_size=8}> {layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>} : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
}
diff --git a/mlir/test/Dialect/XeGPU/subgroup-distribute.mlir b/mlir/test/Dialect/XeGPU/subgroup-distribute.mlir
index c3cdc79d9f70e..285669cae7174 100644
--- a/mlir/test/Dialect/XeGPU/subgroup-distribute.mlir
+++ b/mlir/test/Dialect/XeGPU/subgroup-distribute.mlir
@@ -10,7 +10,7 @@
// CHECK: %[[T1:.*]] = xegpu.load_nd %[[T0]][%{{.*}}] <{packed}> : !xegpu.tensor_desc<16x16xf16> -> vector<16xf16>
// CHECK-DAG: %[[T4:.*]] = xegpu.dpas %[[T3]], %[[T1]] : vector<8xf16>, vector<16xf16> -> vector<8xf32>
// CHECK: %[[T5:.*]] = vector.shape_cast %[[T4]] : vector<8xf32> to vector<8x1xf32>
-// CHECK: %[[T6:.*]] = math.exp %[[T5]] {{{.*}}} : vector<8x1xf32>
+// CHECK: %[[T6:.*]] = math.exp %[[T5]] : vector<8x1xf32>
// CHECK-DAG: %[[T8:.*]] = vector.shape_cast %[[T6]] : vector<8x1xf32> to vector<8xf32>
// CHECK-DAG: %[[T7:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32>
// CHECK: xegpu.store_nd %[[T8]], %[[T7]][{{.*}}] : vector<8xf32>, !xegpu.tensor_desc<8x16xf32>
@@ -423,16 +423,17 @@ gpu.module @xevm_test {
// CHECK: %[[VEC_RED:.*]] = vector.broadcast %[[VEC_RED_3]] : f32 to vector<1xf32>
// CHECK: xegpu.store %[[VEC_RED]], %arg1[%[[CST]]], %[[CST_0]] : vector<1xf32>, memref<256xf32>, vector<1xindex>, vector<1xi1>
gpu.func @vector_reduce_2d(%arg0: memref<4x16xf32>, %arg1: memref<256xf32>) {
- %cst = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0, 1]>} 1.000000e+00 : f32
+ %cst = arith.constant 1.000000e+00 : f32
%0 = xegpu.create_nd_tdesc %arg0 : memref<4x16xf32> -> !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
%1 = xegpu.load_nd %0[0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : !xegpu.tensor_desc<4x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<4x16xf32>
- %2 = vector.broadcast %cst {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} : f32 to vector<16xf32>
- %3 = vector.multi_reduction <add>, %1, %2 {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} [0] : vector<4x16xf32> to vector<16xf32>
- %4 = vector.reduction <add>, %3 {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0, 1]>} : vector<16xf32> into f32
- %5 = vector.broadcast %4 {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} : f32 to vector<16xf32>
- %cst_0 = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<0> : vector<16xindex>
- %cst_1 = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<16xi1>
- xegpu.store %5, %arg1[%cst_0], %cst_1 <{layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>}> : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
+ %2 = vector.broadcast %cst : f32 to vector<16xf32>
+ %3 = vector.multi_reduction <add>, %1, %2 [0] : vector<4x16xf32> to vector<16xf32>
+ %4 = vector.reduction <add>, %3 : vector<16xf32> into f32
+ %40 = xegpu.convert_layout %4 <{input_layout = #xegpu.slice<#xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>, dims = [0]>, target_layout = #xegpu.slice<#xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>, dims = [0]>}>: f32
+ %5 = vector.broadcast %40 : f32 to vector<16xf32>
+ %cst_0 = arith.constant dense<0> : vector<16xindex>
+ %cst_1 = arith.constant dense<true> : vector<16xi1>
+ xegpu.store %5, %arg1[%cst_0], %cst_1 <{layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
gpu.return
}
}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir b/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir
new file mode 100644
index 0000000000000..a00d6d7bb3b14
--- /dev/null
+++ b/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir
@@ -0,0 +1,150 @@
+// RUN: mlir-opt -test-xegpu-recover-temporary-layouts -split-input-file %s | FileCheck %s
+
+// -----
+// Test scf.for: Recovery should propagate layout from the store_nd consumer
+// of the loop result back to the scf.for result, scf.yield operands, and
+// the arith.constant init value. Tensor desc types start without layouts
+// and only anchor ops (load_nd, store_nd, dpas) carry layout attributes.
+
+gpu.module @test_for {
+// CHECK-LABEL: gpu.func @for_basic
+gpu.func @for_basic(%arg0: memref<8x128xf16>, %arg1: memref<128x16xf16>, %arg2: memref<8x16xf32>) {
+ %c0 = arith.constant 0 : index
+ %c128 = arith.constant 128 : index
+ %c16 = arith.constant 16 : index
+ // CHECK: xegpu.create_nd_tdesc
+ // CHECK-SAME: -> !xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
+ %0 = xegpu.create_nd_tdesc %arg0[%c0, %c0] : memref<8x128xf16>
+ -> !xegpu.tensor_desc<8x16xf16>
+ // CHECK: xegpu.create_nd_tdesc
+ // CHECK-SAME: -> !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>>
+ %1 = xegpu.create_nd_tdesc %arg1[%c0, %c0] : memref<128x16xf16>
+ -> !xegpu.tensor_desc<16x16xf16>
+ // Recovery propagates layout from dpas (via store_nd) back to arith.constant.
+ // CHECK: arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ // CHECK-SAME: dense<0.000000e+00> : vector<8x16xf32>
+ %cst = arith.constant dense<0.000000e+00> : vector<8x16xf32>
+ // CHECK: scf.for
+ %2 = scf.for %arg3 = %c0 to %c128 step %c16
+ iter_args(%arg6 = %cst) -> (vector<8x16xf32>) {
+ %4 = xegpu.load_nd %0 {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
+ %5 = xegpu.load_nd %1 {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
+ : !xegpu.tensor_desc<16x16xf16> -> vector<16x16xf16>
+ %6 = xegpu.dpas %4, %5, %arg6
+ {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>,
+ layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ : vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
+ // Recovery propagates layout to scf.yield vector operand.
+ // CHECK: scf.yield {layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ scf.yield %6
+ : vector<8x16xf32>
+ // Recovery sets layout_result_0 on the scf.for for the vector result.
+ // CHECK: layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+ }
+ // CHECK: xegpu.create_nd_tdesc
+ // CHECK-SAME: -> !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
+ %3 = xegpu.create_nd_tdesc %arg2[%c0, %c0] : memref<8x16xf32>
+ -> !xegpu.tensor_desc<8x16xf32>
+ xegpu.store_nd %2, %3 {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32>
+ gpu.return
+}
+}
+
+// -----
+// Test scf.while: Recovery should propagate layout from the store_nd consumer
+// of the while result back through scf.condition (which branches to parent).
+// The scf.yield in the "do" region branches to the "before" region (not
+// parent), so propagateRegionResultsToYieldOperands skips it.
+
+gpu.module @test_while {
+// CHECK-LABEL: gpu.func @while_basic
+gpu.func @while_basic(%arg0: memref<1024xf32>, %arg1: memref<1024xf32>) {
+ %c1_i32 = arith.constant 1 : i32
+ %c10_i32 = arith.constant 10 : i32
+ %c0_i32 = arith.constant 0 : i32
+ // CHECK: xegpu.create_nd_tdesc
+ // CHECK-SAME: -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>>
+ %0 = xegpu.create_nd_tdesc %arg0[0] : memref<1024xf32>
+ -> !xegpu.tensor_desc<256xf32>
+ %1 = xegpu.load_nd %0 {layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}
+ : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
+ // CHECK: xegpu.create_nd_tdesc
+ // CHECK-SAME: -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>>
+ %2 = xegpu.create_nd_tdesc %arg1[0] : memref<1024xf32>
+ -> !xegpu.tensor_desc<256xf32>
+
+ // CHECK: scf.while
+ %3:2 = scf.while (%arg2 = %1, %arg3 = %c0_i32)
+ : (vector<256xf32>, i32) -> (vector<256xf32>, i32) {
+ %4 = arith.cmpi slt, %arg3, %c10_i32 : i32
+ // Recovery propagates layout to scf.condition vector operand.
+ // CHECK: scf.condition
+ // CHECK-SAME: {layout_operand_1 = #xegpu.layout<sg_layout = [16], sg_data = [16]>}
+ scf.condition(%4) %arg2, %arg3 : vector<256xf32>, i32
+ } do {
+ ^bb0(%arg2: vector<256xf32>, %arg3: i32):
+ xegpu.store_nd %arg2, %2 {layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}
+ : vector<256xf32>, !xegpu.tensor_desc<256xf32>
+ %4 = arith.addi %arg3, %c1_i32 : i32
+ %5 = xegpu.update_nd_offset %0, [256]
+ : !xegpu.tensor_desc<256xf32>
+ %6 = xegpu.load_nd %5 {layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}
+ : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
+ // Recovery propagates layout to scf.yield in the "do" region via
+ // sibling region propagation (from "before" region arg back to "do" yield).
+ // CHECK: scf.yield {layout_operand_0 = #xegpu.layout<sg_layout = [16], sg_data = [16]>}
+ scf.yield %6, %4 : vector<256xf32>, i32
+ // Recovery sets layout_result_0 on the scf.while for the vector result.
+ // CHECK: } attributes {layout_operand_0 = #xegpu.layout<sg_layout = [16], sg_data = [16]>,
+ // CHECK-SAME: layout_result_0 = #xegpu.layout<sg_layout = [16], sg_data = [16]>}
+ }
+ xegpu.store_nd %3#0, %2 {layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}
+ : vector<256xf32>, !xegpu.tensor_desc<256xf32>
+ gpu.return
+}
+}
+
+// -----
+// Test scf.if: Recovery should propagate layout from the dpas consumer of the
+// if result back to the scf.if result and both yield operands.
+
+gpu.module @test_if {
+// CHECK-LABEL: gpu.func @if_basic
+gpu.func @if_basic(
+ %arg0: !xegpu.tensor_desc<8x16xf16>,
+ %arg1: !xegpu.tensor_desc<16x16xf16>,
+ %arg2: i1,
+ %arg3: !xegpu.tensor_desc<8x16xf32>) {
+ %0 = xegpu.load_nd %arg0 {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
+ // CHECK: scf.if
+ %1 = scf.if %arg2 -> (vector<16x16xf16>) {
+ %3 = xegpu.load_nd %arg1 {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
+ : !xegpu.tensor_desc<16x16xf16>
+ -> vector<16x16xf16>
+ // Recovery propagates layout to scf.yield operand in "then" region.
+ // CHECK: scf.yield {layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
+ scf.yield %3 : vector<16x16xf16>
+ } else {
+ %3 = xegpu.load_nd %arg1 {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
+ : !xegpu.tensor_desc<16x16xf16>
+ -> vector<16x16xf16>
+ // Recovery propagates layout to scf.yield operand in "else" region.
+ // CHECK: scf.yield {layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
+ scf.yield %3 : vector<16x16xf16>
+ // Recovery sets layout_result_0 on the scf.if for the vector result.
+ // CHECK: } {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
+ }
+ %2 = xegpu.dpas %0, %1
+ {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>,
+ layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ : vector<8x16xf16>, vector<16x16xf16> -> vector<8x16xf32>
+ xegpu.store_nd %2, %arg3 {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32>
+ gpu.return
+}
+}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops.mlir
index bbdffa0986962..3bc43b780ade2 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops.mlir
@@ -405,7 +405,7 @@ gpu.module @test_distribution {
// CHECK-LABEL: gpu.func @vector_reduce_scalar_cross_sg
// CHECK-SAME: (%[[ARG0:.*]]: memref<32x32xf32>)
- // CHECK-DAG: %[[CST:.*]] = arith.constant {{.*}} 0.000000e+00 : f32
+ // CHECK-DAG: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[LOAD:.*]] = xegpu.load_nd %{{.*}} : !xegpu.tensor_desc<8x8xf32> -> vector<8x8xf32>
// CHECK-DAG: %[[CST_ACC:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[LOCAL:.*]] = vector.multi_reduction <add>, %[[LOAD]], %[[CST_ACC]] [0, 1] : vector<8x8xf32> to f32
diff --git a/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
index 3394d63dcbbdc..a3d2560cedf63 100644
--- a/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
+++ b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
@@ -203,6 +203,38 @@ class TestStepOpPattern : public OpConversionPattern<vector::StepOp> {
}
};
+struct TestXeGPURecoverTemporaryLayouts
+ : public PassWrapper<TestXeGPURecoverTemporaryLayouts,
+ OperationPass<gpu::GPUModuleOp>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestXeGPURecoverTemporaryLayouts)
+
+ StringRef getArgument() const final {
+ return "test-xegpu-recover-temporary-layouts";
+ }
+
+ StringRef getDescription() const final {
+ return "Test the implementation of XeGPU temporary layout recovery";
+ }
+
+ void getDependentDialects(::mlir::DialectRegistry ®istry) const override {
+ registry.insert<arith::ArithDialect>();
+ registry.insert<memref::MemRefDialect>();
+ registry.insert<xegpu::XeGPUDialect>();
+ registry.insert<vector::VectorDialect>();
+ registry.insert<gpu::GPUDialect>();
+ }
+
+ TestXeGPURecoverTemporaryLayouts() = default;
+ TestXeGPURecoverTemporaryLayouts(const TestXeGPURecoverTemporaryLayouts &pass)
+ : PassWrapper(pass) {}
+
+ void runOnOperation() override {
+ Operation *op = getOperation();
+ if (!xegpu::recoverTemporaryLayouts(op))
+ signalPassFailure();
+ }
+};
+
struct TestXeGPUSGDistribute
: public PassWrapper<TestXeGPUSGDistribute,
OperationPass<gpu::GPUModuleOp>> {
@@ -457,6 +489,7 @@ namespace test {
void registerTestXeGPULowerings() {
PassRegistration<TestXeGPUUnrollingPatterns>();
PassRegistration<TestXeGPULayoutInterface>();
+ PassRegistration<TestXeGPURecoverTemporaryLayouts>();
PassRegistration<TestXeGPUSGDistribute>();
PassRegistration<TestXeGPUSgToWiDistributeExperimental>();
PassRegistration<TestXeGPUMoveFuncBodyToWarpOp>();
More information about the Mlir-commits
mailing list