[Mlir-commits] [mlir] [MLIR][XeGPU] Use context-aware type converter in WgToSgDistribute and Blocking pass (PR #194685)
Nishant Patel
llvmlistbot at llvm.org
Thu Jun 4 11:55:41 PDT 2026
https://github.com/nbpatel updated https://github.com/llvm/llvm-project/pull/194685
>From 890475a634653cb52ff388632d1b0642a5db9b2d Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Tue, 28 Apr 2026 17:03:55 +0000
Subject: [PATCH 1/9] Add context aware type converter
---
.../Dialect/XeGPU/Transforms/Transforms.h | 6 +
.../mlir/Dialect/XeGPU/Utils/XeGPUUtils.h | 44 ++-
.../XeGPU/Transforms/XeGPUBlocking.cpp | 122 ++++---
.../XeGPUSgToWiDistributeExperimental.cpp | 73 +----
.../Transforms/XeGPUWgToSgDistribute.cpp | 219 ++++---------
mlir/lib/Dialect/XeGPU/Utils/CMakeLists.txt | 2 +
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 301 +++++++++---------
.../test/Dialect/XeGPU/xegpu-wg-to-sg-rr.mlir | 2 +-
mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir | 25 +-
9 files changed, 355 insertions(+), 439 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
index fe989ebb17059..51052efe33130 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
@@ -67,6 +67,12 @@ void populateXeGPUPeepHoleOptimizerPatterns(RewritePatternSet &patterns);
void populateXeGPUSubgroupDistributePatterns(RewritePatternSet &patterns);
/// Appends patterns for moving function body into gpu.warp_execute_on_lane0 op.
void populateXeGPUMoveFuncBodyToWarpOpPatterns(RewritePatternSet &patterns);
+/// Define the type conversions needed for XeGPU workgroup to subgroup
+/// distribution. This includes a context-aware 1:N conversion for VectorType
+/// (using the distribute layout attribute on the Value) and a 1:N conversion
+/// for TensorDescType.
+void populateXeGPUWgToSgDistributeTypeConversions(TypeConverter &converter,
+ Operation *topLevelOp);
/// Appends patterns for XeGPU workgroup to subgroup distribution into
/// `patterns`.
void populateXeGPUWgToSgDistributePatterns(RewritePatternSet &patterns);
diff --git a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
index 1b594f17e15ec..9bf2b751320f4 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
@@ -12,8 +12,12 @@
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpDefinition.h"
+#include "llvm/ADT/SetVector.h"
+#include <functional>
+
namespace mlir {
+class UnrealizedConversionCastOp;
class VectorType;
class OpOperand;
class OpResult;
@@ -94,17 +98,6 @@ Value createVectorWithShapeFromValues(OpBuilder &builder, Location loc,
ValueRange values,
ArrayRef<int64_t> shape);
-/// Do type conversion for SCF structural ops, e.g., scf.for using SCF structure
-/// type convertion patterns. Since VectorType cannot carry the layout
-/// attribute, which is needed to guide the type conversion for XeGPU, they are
-/// first converted into RankedTensorType, where the layout attribute can be
-/// attached. And then upstream SCF structural type conversion patterns are
-/// applied with the provided converter.
-/// TODO: This is a temporary solution. We should refactor it when context-aware
-/// type conversion is available.
-void doSCFStructuralTypeConversionWithTensorType(Operation *op,
- TypeConverter converter);
-
/// Retrieves the chip string from the XeVM target attribute of the parent
/// GPU module operation. Returns the chip identifier if found, or nullopt
/// if no GPU module parent or XeVM target attribute exists.
@@ -234,6 +227,35 @@ bool matchUnitDimExpansion(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
bool matchSplitDimExpansion(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
SmallVector<SmallVector<int64_t>> &splitDimGroups);
+/// Callback type for computing sub-shape and count for 1:N VectorType
+/// conversion. Given a VectorType and its DistributeLayoutAttr, returns
+/// (subShape, count). A count <= 0 means no conversion is needed.
+using SubShapeAndCountFn = std::function<std::pair<SmallVector<int64_t>, int>(
+ VectorType, DistributeLayoutAttr)>;
+
+/// Adds source (N:1) and target (1:1) materializations using
+/// UnrealizedConversionCastOp to the given TypeConverter.
+void addSCFStructuralMaterializations(TypeConverter &converter);
+
+/// Pre-computes block argument type mappings for SCF loop ops and adds a
+/// context-aware 1:N VectorType conversion to the TypeConverter.
+/// Pre-computation is needed because during structural type conversion
+/// (especially scf.while), blocks may be detached from their parent region,
+/// making Block::getParent() crash (LLVM ilist assertion). The
+/// `getSubShapeAndCount` callback computes (subShape, count) for a VectorType
+/// and its layout; count <= 0 means no conversion needed.
+void addContextAwareVectorTypeConversion(
+ TypeConverter &converter, Operation *topLevelOp,
+ SubShapeAndCountFn getSubShapeAndCount);
+
+/// Cleans up UnrealizedConversionCastOps inserted during SCF structural type
+/// conversion. Folds cancelling N:1->1:N and 1:N->N:1 cast chains (inserting
+/// vector.shape_cast when shapes differ but element counts match), and
+/// erases dead casts. Casts in `existingCasts` are preserved.
+void cleanupUnrealizedConversionCasts(
+ Operation *root,
+ const llvm::SmallSetVector<UnrealizedConversionCastOp, 8> &existingCasts);
+
} // namespace xegpu
} // namespace mlir
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
index 98c9dc3f5e53a..f2c1b066460a1 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
@@ -9,6 +9,8 @@
#include "mlir/Dialect/XeGPU/Transforms/Passes.h"
#include "mlir/Dialect/Index/IR/IndexDialect.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/SCF/Transforms/Patterns.h"
#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Transforms/Transforms.h"
@@ -19,6 +21,7 @@
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SetVector.h"
#include "llvm/Support/DebugLog.h"
namespace mlir {
@@ -272,51 +275,80 @@ void XeGPUBlockingPass::runOnOperation() {
return std::make_pair(tileShape, count);
};
- // Perform type conversion for SCF control folow ops
- TypeConverter converter;
- converter.addConversion([](Type type) -> Type { return type; });
- converter.addConversion(
- [&](RankedTensorType type,
- SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {
- Type elemTy = type.getElementType();
- ArrayRef<int64_t> shape = type.getShape();
-
- auto layout =
- llvm::dyn_cast_if_present<xegpu::LayoutAttr>(type.getEncoding());
- if (layout && layout.isForWorkgroup())
- return failure();
-
- int count;
- SmallVector<int64_t> subShape;
- std::tie(subShape, count) = getTileShapeAndCount(shape, layout);
- auto newTy = VectorType::get(subShape, elemTy);
- result.append(count, newTy);
- return success();
- });
- converter.addConversion(
- [&](xegpu::TensorDescType type,
- SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {
- Type elemTy = type.getElementType();
- ArrayRef<int64_t> shape = type.getShape();
-
- xegpu::DistributeLayoutAttr layout = type.getLayoutAttr();
- if (layout && layout.isForWorkgroup())
- return failure();
-
- int count;
- SmallVector<int64_t> subShape;
- std::tie(subShape, count) = getTileShapeAndCount(shape, layout);
-
- if (layout)
- layout = layout.dropInstData();
-
- auto newTy = xegpu::TensorDescType::get(
- type.getContext(), subShape, elemTy, type.getEncoding(), layout);
- result.append(count, newTy);
- return success();
- });
-
- xegpu::doSCFStructuralTypeConversionWithTensorType(op, converter);
+ // Perform context-aware type conversion for SCF structural ops.
+ // Inspects Values to find inst_data layout information for 1:N conversion.
+ llvm::SmallSetVector<UnrealizedConversionCastOp, 8> existingCasts;
+ op->walk(
+ [&](UnrealizedConversionCastOp castOp) { existingCasts.insert(castOp); });
+
+ {
+ TypeConverter converter;
+ converter.addConversion([](Type type) -> Type { return type; });
+
+ // TensorDescType 1:N converter (type-based, layout is in the type).
+ converter.addConversion(
+ [&](xegpu::TensorDescType type,
+ SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {
+ Type elemTy = type.getElementType();
+ ArrayRef<int64_t> shape = type.getShape();
+
+ xegpu::DistributeLayoutAttr layout = type.getLayoutAttr();
+ if (layout && layout.isForWorkgroup())
+ return failure();
+
+ int count;
+ SmallVector<int64_t> subShape;
+ std::tie(subShape, count) = getTileShapeAndCount(shape, layout);
+
+ if (layout)
+ layout = layout.dropInstData();
+
+ auto newTy = xegpu::TensorDescType::get(
+ type.getContext(), subShape, elemTy, type.getEncoding(), layout);
+ result.append(count, newTy);
+ return success();
+ });
+
+ // Context-aware 1:N conversion for VectorType based on inst_data.
+ xegpu::addContextAwareVectorTypeConversion(
+ converter, op,
+ [&](VectorType vecTy, xegpu::DistributeLayoutAttr layout)
+ -> std::pair<SmallVector<int64_t>, int> {
+ if (layout.isForWorkgroup())
+ return {{}, 0};
+ auto instData = layout.getEffectiveInstDataAsInt();
+ if (instData.empty())
+ return {{}, 0};
+ int count =
+ computeProduct(vecTy.getShape()) / computeProduct(instData);
+ if (count <= 1)
+ return {{}, 0};
+ return {SmallVector<int64_t>(instData), count};
+ });
+ xegpu::addSCFStructuralMaterializations(converter);
+ // Blocking runs SCF conversion separately (not combined with XeGPU
+ // patterns), so it also needs a 1:N target materialization.
+ converter.addTargetMaterialization(
+ [](mlir::OpBuilder &builder, mlir::TypeRange types,
+ mlir::ValueRange inputs, mlir::Location loc) -> SmallVector<Value> {
+ auto castOp =
+ UnrealizedConversionCastOp::create(builder, loc, types, inputs);
+ return SmallVector<Value>(castOp.getResults());
+ });
+
+ ConversionTarget target(*ctx);
+ target.addLegalOp<UnrealizedConversionCastOp>();
+ target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
+
+ RewritePatternSet scfPatterns(ctx);
+ scf::populateSCFStructuralTypeConversionsAndLegality(converter, scfPatterns,
+ target);
+ if (failed(applyPartialConversion(op, target, std::move(scfPatterns))))
+ return signalPassFailure();
+
+ // Fold cancelling cast chains and erase dead casts.
+ xegpu::cleanupUnrealizedConversionCasts(op, existingCasts);
+ }
xegpu::UnrollOptions options;
options.setFilterConstraint(
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
index c153db431c035..e1cf70094b1c3 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
@@ -1529,19 +1529,11 @@ void XeGPUSgToWiDistributeExperimentalPass::runOnOperation() {
// Perform a structural type conversion to convert structural ops to have WI
// types. This will insert UnrealizedConversionCastOps to make the IR
// valid.
- auto materializeCast = [&](mlir::OpBuilder &builder, mlir::Type type,
- mlir::ValueRange inputs,
- mlir::Location loc) -> mlir::Value {
- UnrealizedConversionCastOp castOp =
- UnrealizedConversionCastOp::create(builder, loc, type, inputs);
- return castOp.getResult(0);
- };
{
ConversionTarget target(getContext());
TypeConverter typeConverter;
RewritePatternSet patterns(&getContext());
- typeConverter.addSourceMaterialization(materializeCast);
- typeConverter.addTargetMaterialization(materializeCast);
+ xegpu::addSCFStructuralMaterializations(typeConverter);
xegpu::populateXeGPUSgToWiDistributeTypeConversions(typeConverter);
scf::populateSCFStructuralTypeConversionsAndLegality(typeConverter,
patterns, target);
@@ -1550,67 +1542,8 @@ void XeGPUSgToWiDistributeExperimentalPass::runOnOperation() {
target.addLegalOp<UnrealizedConversionCastOp>();
(void)applyPartialConversion(root, target, std::move(patterns));
}
- // Structural type conversion can generate some redundant
- // UnrealizedConversionCastOps to materialize the SG type from type converted
- // WI type. These are redundant at this point and can be eliminated by
- // inserting shape casts instead.
- // Example:
- // %1 = UnrealizedConversionCastOp %0 : vector<16x1xf32> to vector<16x16xf32>
- // %2 = UnrealizedConversionCastOp %1 : vector<16x16xf32> to vector<16xf32>
- // This can be replaced with:
- // %2 = vector.shape_cast %0 : vector<16x1xf32> to vector<16xf32>
- OpBuilder builder(root);
- root->walk([&](UnrealizedConversionCastOp op) {
- // If this op existed before, nothing to do.
- if (existingCasts.contains(op))
- return;
- // number of inputs and outputs must be 1.
- if (op.getNumOperands() != 1 || op.getNumResults() != 1)
- return;
- // Both input and output types must be vector types.
- auto singleInput = op.getInputs()[0];
- auto inputTy = dyn_cast<VectorType>(singleInput.getType());
- auto outputTy = dyn_cast<VectorType>(op.getResult(0).getType());
- if (!inputTy || !outputTy)
- return;
-
- // Check if the defining op of the input is also an
- // UnrealizedConversionCastOp and it has a single user (which is this
- // op).
- auto definingOp = singleInput.getDefiningOp<UnrealizedConversionCastOp>();
- if (!definingOp || !definingOp->hasOneUse())
- return;
- auto inputOfDefiningOp = definingOp.getInputs()[0];
- // If the input of the defining op and output type are both vector types
- // have same number of elements, insert a shape cast.
- auto inputOfDefiningOpTy =
- dyn_cast<VectorType>(inputOfDefiningOp.getType());
- if (inputOfDefiningOpTy &&
- inputOfDefiningOpTy.getNumElements() == outputTy.getNumElements()) {
- builder.setInsertionPoint(op);
- auto shapeCast = vector::ShapeCastOp::create(builder, op.getLoc(),
- outputTy, inputOfDefiningOp);
- op.replaceAllUsesWith(ValueRange{shapeCast.getResult()});
- return;
- }
- });
- // At this point, we will have some dead UnrealizedConversionCastOps. Just
- // erase them.
- bool changed = true;
- while (changed) {
- changed = false;
- root->walk([&](UnrealizedConversionCastOp op) {
- // Skip existing casts.
- if (existingCasts.contains(op))
- return;
- if (op.use_empty()) {
- op.erase();
- changed = true;
- }
- });
- }
-
- xegpu::removeTemporaryLayoutAttrs(getOperation());
+ // Fold cancelling cast chains and erase dead casts.
+ xegpu::cleanupUnrealizedConversionCasts(root, existingCasts);
}
void xegpu::populateXeGPUSgToWiDistributeTypeConversions(
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index 8aa0758943cd1..dc725bcb26e5d 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -22,6 +22,7 @@
#include "mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
#include "mlir/Transforms/DialectConversion.h"
+#include "llvm/ADT/SetVector.h"
#include <optional>
namespace mlir {
@@ -581,82 +582,6 @@ struct WgToSgConvertLayoutOp
}
};
-// Handles UnrealizedConversionCastOp generated during
-// SCFStructuralTypeConversions (step 1). This op may appear as either a
-// target or source materialization for Vector values, e.g.:
-// 1. unrealized_cast %1 : vector<256xf32> to vector<16xf32>, ...
-// 2. unrealized_cast %1 : vector<16xf32>, ... to vector<256xf32>
-// it could be either 1:N or N:1 cast. In both cases, the pattern
-// simply forwards the inputs to the outputs using 1:1 or 1:N interface.
-// for example, the following scf::forOp
-// ```
-// %for = scf.for ... iter_args(%arg1 = %0)->(vector<128x128xf16>) {
-// %n = use(%arg1): vector<128x128xf16>
-// scf.yield %n : vector<128x128xf16>
-// }
-// ```
-// Could be converted to:
-// ```
-// %1 = unrealized_conversion_cast %0
-// : vector<128x128xf16> to vector<16x16xf16>, vector<16x16xf16>
-// %for:2 = scf.for ... iter_args(%arg1 = %1#1, %arg2 = %1#2)
-// -> (vector<16x16xf16>, vector<16x16xf16) {
-// %m = unrealized_conversion_cast %arg1, %arg2
-// : vector<16x16xf16>, vector<16x16xf16> to vector<128x128xf16>
-// %n = use(%m): vector<128x128xf16>
-// %b = unrealized_conversion_cast %n
-// : vector<128x128xf16> to vector<16x16xf16>, vector<16x16xf16>
-// scf.yield %b#1, %b#2 : vector<16x16xf16>, vector<16x16xf16>
-// }
-// %cast = unrealized_conversion_cast %for:2
-// : vector<16x16xf16>, vector<16x16xf16> to vector<128x128xf16>
-// ```
-// TODO: remove it when context-aware type converter is ready.
-struct UnrealizedConversionCastOpPattern
- : public OpConversionPattern<mlir::UnrealizedConversionCastOp> {
- using OpConversionPattern<
- mlir::UnrealizedConversionCastOp>::OpConversionPattern;
-
- mlir::LogicalResult
- matchAndRewrite(mlir::UnrealizedConversionCastOp op, OneToNOpAdaptor adaptor,
- ConversionPatternRewriter &rewriter) const override {
- SmallVector<Value> inputs = xegpu::flattenValues(adaptor.getInputs());
-
- auto inputTy = dyn_cast<VectorType>(inputs[0].getType());
- auto outputTy = dyn_cast<VectorType>(op->getOpResult(0).getType());
-
- if (!inputTy || !outputTy || !llvm::all_equal(op->getResultTypes()) ||
- !llvm::all_equal(ValueRange(inputs).getTypes()))
- return failure();
-
- // Handles the case "cast %1 : vector<256xf32> to vector<16xf32>, ...".
- // It is generated by source materialization (e.g., inits to scf forOp).
- // The input values provided by the adaptor should already be distributed,
- // and their types should correspond exactly to the result types of the
- // operation.
- if (op.getNumOperands() == 1 &&
- llvm::equal(ValueRange(inputs).getTypes(), op->getResultTypes())) {
- rewriter.replaceOp(op, inputs);
- return success();
- }
-
- // Handles the case "cast %1 : vector<16xf32>, ... to vector<256xf32>".
- // It is generated by target materialization (e.g., arguments/results
- // of scf forOp). All input values must have the same vector type, and
- // their shape must be evenly divisible by the output vector's shape
- // (determined by the nature of the workgroup to subgroup distribution).
- // TODO: it is not safe to do such forward, since such N:1 cast could be
- // from others.
- if (op.getNumResults() == 1 &&
- computeShapeRatio(outputTy.getShape(), inputTy.getShape())) {
- rewriter.replaceOpWithMultiple(op, {inputs});
- return success();
- }
-
- return mlir::failure();
- }
-};
-
// This pattern distributes arith.constant op into subgroup-level constants
struct WgToSgArithConstantOp : public OpConversionPattern<arith::ConstantOp> {
using OpConversionPattern<arith::ConstantOp>::OpConversionPattern;
@@ -1407,10 +1332,48 @@ using WgToSgVectorCreateMaskOp = WgToSgVectorMaskOp<vector::CreateMaskOp>;
namespace mlir {
namespace xegpu {
+void populateXeGPUWgToSgDistributeTypeConversions(TypeConverter &converter,
+ Operation *topLevelOp) {
+ // Pass through all types by default.
+ converter.addConversion([](Type type) -> Type { return type; });
+
+ // For TensorDescType, convert WG-level tensor descs to N SG-level descs.
+ converter.addConversion(
+ [](xegpu::TensorDescType type,
+ SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {
+ xegpu::DistributeLayoutAttr layout = type.getLayoutAttr();
+ if (!layout || !layout.isForWorkgroup())
+ return std::nullopt;
+
+ Type elemTy = type.getElementType();
+ ArrayRef<int64_t> shape = type.getShape();
+
+ int count;
+ SmallVector<int64_t> subShape;
+ std::tie(subShape, count) = getSgShapeAndCount(shape, layout);
+
+ layout = layout.dropSgLayoutAndData();
+
+ auto newTy = xegpu::TensorDescType::get(
+ type.getContext(), subShape, elemTy, type.getEncoding(), layout);
+ result.append(count, newTy);
+ return success();
+ });
+
+ // Context-aware 1:N conversion for VectorType based on sg_layout/sg_data.
+ xegpu::addContextAwareVectorTypeConversion(
+ converter, topLevelOp,
+ [](VectorType vecTy, xegpu::DistributeLayoutAttr layout)
+ -> std::pair<SmallVector<int64_t>, int> {
+ if (!layout.isForWorkgroup())
+ return {{}, 0};
+ return getSgShapeAndCount(vecTy.getShape(), layout);
+ });
+}
+
void populateXeGPUWgToSgDistributePatterns(RewritePatternSet &patterns) {
patterns.add<WgToSgCreateNdOp, WgToSgLoadNdOp, WgToSgStoreNdOp, WgToSgDpasOp,
- WgToSgPrefetchNdOp, UnrealizedConversionCastOpPattern,
- WgToSgElementwiseOp, WgToSgVectorBroadcastOp,
+ WgToSgPrefetchNdOp, WgToSgElementwiseOp, WgToSgVectorBroadcastOp,
WgToSgConvertLayoutOp, WgToSgArithConstantOp, WgToSgLoadGatherOp,
WgToSgStoreScatterOp, WgToSgLoadMatrixOp, WgToSgStoreMatrixOp,
WgToSgVectorStepOp, WgToSgVectorShapeCastOp,
@@ -1436,78 +1399,22 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
return;
}
- // Track existing UnrealizedConversionCastOps
- SmallVector<Operation *> existingCastOps;
- getOperation()->walk([&](UnrealizedConversionCastOp castOp) {
- existingCastOps.push_back(castOp.getOperation());
- });
-
- {
- // Step 1: Apply SCFStructuralTypeConversions to SCF operations with
- // VectorType operands. This first converts such operands to
- // RankedTensorType, propagates the layout attribute into the encoding
- // attribute, and finally converts the RankedTensorType to VectorType based
- // on the encoding.
-
- TypeConverter converter;
- converter.addConversion([&](Type type) -> Type { return type; });
- converter.addConversion(
- [&](RankedTensorType type,
- SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {
- // Only convert RankedTensorTypes that carry an XeGPU layout encoding.
- // Plain tensors (e.g. tensor<?xi32>) have no XeGPU encoding and must
- // not be converted: VectorType does not support dynamic dimensions.
- auto encoding = dyn_cast_if_present<xegpu::DistributeLayoutAttr>(
- type.getEncoding());
- if (!encoding)
- return std::nullopt;
-
- Type elemTy = type.getElementType();
- ArrayRef<int64_t> shape = type.getShape();
-
- int count;
- SmallVector<int64_t> subShape;
- std::tie(subShape, count) = getSgShapeAndCount(shape, encoding);
-
- auto newTy = VectorType::get(subShape, elemTy);
- result.append(count, newTy);
- return success();
- });
-
- xegpu::doSCFStructuralTypeConversionWithTensorType(getOperation(),
- converter);
- }
+ // Collect existing UnrealizedConversionCastOps. These must be preserved.
+ llvm::SmallSetVector<UnrealizedConversionCastOp, 8> existingCasts;
+ getOperation()->walk(
+ [&](UnrealizedConversionCastOp castOp) { existingCasts.insert(castOp); });
- // Step 2: Perform workgroup to subgroup distribution for TensorDesc values,
- // as well as XeGPU, Arith, and Vector operations.
+ // Perform workgroup to subgroup distribution for TensorDesc and Vector
+ // values, as well as XeGPU, Arith, and Vector operations. Uses a
+ // context-aware type converter that inspects Values to retrieve the
+ // distribute layout attribute for 1:N type conversion.
MLIRContext *ctx = &getContext();
RewritePatternSet patterns(ctx);
ConversionTarget target(*ctx);
TypeConverter converter;
- converter.addConversion([&](Type type) -> Type { return type; });
- converter.addConversion(
- [&](xegpu::TensorDescType type,
- SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {
- xegpu::DistributeLayoutAttr layout = type.getLayoutAttr();
- // Only convert WG-level tensor descs. SG-level or layout-less types
- // are already legal and should pass through unchanged.
- if (!layout || !layout.isForWorkgroup())
- return std::nullopt;
-
- Type elemTy = type.getElementType();
- ArrayRef<int64_t> shape = type.getShape();
-
- int count;
- SmallVector<int64_t> subShape;
- std::tie(subShape, count) = getSgShapeAndCount(shape, layout);
-
- layout = layout.dropSgLayoutAndData();
-
- auto newTy = xegpu::TensorDescType::get(
- type.getContext(), subShape, elemTy, type.getEncoding(), layout);
- result.append(count, newTy);
- return success();
- });
+ xegpu::addSCFStructuralMaterializations(converter);
+ xegpu::populateXeGPUWgToSgDistributeTypeConversions(converter,
+ getOperation());
auto getTensorDescType = [](Operation *op) -> xegpu::TensorDescType {
if (auto createOp = dyn_cast<xegpu::CreateNdDescOp>(op))
@@ -1613,10 +1520,7 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
return isLegal(layout);
});
- target.addDynamicallyLegalOp<UnrealizedConversionCastOp>(
- [=](UnrealizedConversionCastOp op) {
- return llvm::is_contained(existingCastOps, op.getOperation());
- });
+ target.addLegalOp<UnrealizedConversionCastOp>();
target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
@@ -1627,5 +1531,20 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
applyPartialConversion(getOperation(), target, std::move(patterns))))
return signalPassFailure();
- xegpu::removeTemporaryLayoutAttrs(getOperation());
+ // Fold cancelling cast chains and erase dead casts.
+ xegpu::cleanupUnrealizedConversionCasts(getOperation(), existingCasts);
+
+ // 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);
+ });
}
diff --git a/mlir/lib/Dialect/XeGPU/Utils/CMakeLists.txt b/mlir/lib/Dialect/XeGPU/Utils/CMakeLists.txt
index d9bf4a1461c27..6cd94aa5e65d0 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/CMakeLists.txt
+++ b/mlir/lib/Dialect/XeGPU/Utils/CMakeLists.txt
@@ -8,6 +8,8 @@ add_mlir_dialect_library(MLIRXeGPUUtils
MLIRIR
MLIRSCFTransforms
MLIRGPUDialect
+ MLIRTransformUtils
+ MLIRVectorDialect
MLIRXeVMDialect
MLIRXeGPUDialect
)
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 2d1ce6eea17aa..c46f492ec07fe 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -11,13 +11,16 @@
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/XeVMDialect.h"
-#include "mlir/Dialect/SCF/Transforms/Patterns.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h"
#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/ValueRange.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
@@ -432,153 +435,6 @@ Value xegpu::createVectorWithShapeFromValues(OpBuilder &builder, Location loc,
return result;
}
-void xegpu::doSCFStructuralTypeConversionWithTensorType(
- Operation *op, TypeConverter converter) {
- MLIRContext *context = op->getContext();
-
- auto materializeCast = [](OpBuilder &builder, Type type, ValueRange inputs,
- Location loc) -> Value {
- return UnrealizedConversionCastOp::create(builder, loc, type, inputs)
- .getResult(0);
- };
-
- { // convert VectorType to RankedTensorType for SCF Structural ops
- TypeConverter converter;
- converter.addConversion([](Type type) -> Type { return type; });
- converter.addConversion([](VectorType type) -> Type {
- return RankedTensorType::get(type.getShape(), type.getElementType());
- });
- converter.addSourceMaterialization(materializeCast);
- converter.addTargetMaterialization(materializeCast);
-
- mlir::ConversionTarget target(*context);
- target.addLegalOp<UnrealizedConversionCastOp>();
-
- mlir::RewritePatternSet patterns(context);
- scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,
- target);
- (void)mlir::applyPartialConversion(op, target, std::move(patterns));
- }
-
- { // propagate the layout attribute to RankedTensorType by checking
- // BuiltInUnrealizedCastOps
- // for VectorType to RankedTensorType cast.
- op->walk([](UnrealizedConversionCastOp castOp) {
- if (castOp.getNumOperands() != 1 || castOp.getNumResults() != 1)
- return WalkResult::skip();
-
- Value input = castOp.getInputs()[0];
- Value result = castOp.getResults()[0];
- auto inputTy = dyn_cast<VectorType>(input.getType());
- auto resultTy = dyn_cast<RankedTensorType>(result.getType());
-
- // Only look at ops casting from VectorType to RankedTensorType
- if (!inputTy || !resultTy)
- return WalkResult::skip();
-
- xegpu::DistributeLayoutAttr layout =
- xegpu::getDistributeLayoutAttr(input);
- if (!layout)
- return WalkResult::skip();
-
- RankedTensorType newTy = resultTy.cloneWithEncoding(layout);
- result.setType(newTy);
-
- // update the arguments if user is a LoopLike op.
- for (OpOperand &use : result.getUses()) {
- if (auto loop = dyn_cast<LoopLikeOpInterface>(use.getOwner())) {
- BlockArgument arg = loop.getTiedLoopRegionIterArg(&use);
- arg.setType(newTy);
- }
- // whileOp has two regions, the BlockArgument of the after region
- // is not exposed by LoopLikeOpInterface
- if (auto whileOp = dyn_cast<scf::WhileOp>(use.getOwner())) {
- unsigned idx = use.getOperandNumber();
- BlockArgument arg = whileOp.getAfterArguments()[idx];
- arg.setType(newTy);
- }
- }
- return WalkResult::advance();
- });
-
- // using yieldOp as anchor to update the result type of its ParentOp
- op->walk([](scf::YieldOp yieldOp) {
- Operation *parentOp = yieldOp->getParentOp();
- for (OpResult r : parentOp->getOpResults()) {
- unsigned idx = r.getResultNumber();
- Type resultTy = r.getType();
- Type yieldTy = yieldOp.getResults()[idx].getType();
- if (isa<RankedTensorType>(resultTy) && yieldTy != resultTy)
- r.setType(yieldTy);
- }
- });
- }
-
- { // perform the conversion from RankedTensorType to VectorType based on the
- // DistributeLayoutAttr
-
- // Handle the UnrealizedConversionCastOp introduced by the first step.
- // For vector->RankedTensorType, it will simply forward the inputs.
- // For RankedTensorType->vector, it will update the inputs with the
- // one from the adaptor.
- class UnrealizedConversionCastOpPattern
- : public OpConversionPattern<mlir::UnrealizedConversionCastOp> {
- using OpConversionPattern<
- mlir::UnrealizedConversionCastOp>::OpConversionPattern;
-
- mlir::LogicalResult
- matchAndRewrite(mlir::UnrealizedConversionCastOp op,
- OneToNOpAdaptor adaptor,
- ConversionPatternRewriter &rewriter) const override {
- auto inputs = op.getOperands();
- auto outputs = op.getOutputs();
-
- if (inputs.size() != 1 || outputs.size() != 1)
- return failure();
-
- auto inputTy = inputs[0].getType();
- auto outputTy = outputs[0].getType();
-
- if (isa<VectorType>(inputTy) && isa<RankedTensorType>(outputTy)) {
- rewriter.replaceOpWithMultiple(op, adaptor.getInputs());
- return success();
- }
-
- if (isa<RankedTensorType>(inputTy) && isa<VectorType>(outputTy)) {
- SmallVector<Value> values = xegpu::flattenValues(adaptor.getInputs());
- auto newOp = UnrealizedConversionCastOp::create(rewriter, op.getLoc(),
- outputTy, values);
- rewriter.replaceOp(op, newOp);
- return success();
- }
- return failure();
- }
- };
-
- converter.addSourceMaterialization(materializeCast);
- converter.addTargetMaterialization([&](OpBuilder &builder, TypeRange type,
- ValueRange inputs, Location loc) {
- return UnrealizedConversionCastOp::create(builder, loc, type, inputs)
- .getResults();
- });
-
- mlir::ConversionTarget target(*context);
- target.addDynamicallyLegalOp<UnrealizedConversionCastOp>(
- [](UnrealizedConversionCastOp op) {
- auto isTensorTy = [](Type type) {
- return isa<RankedTensorType>(type);
- };
- return llvm::none_of(op->getOperandTypes(), isTensorTy) &&
- llvm::none_of(op->getResultTypes(), isTensorTy);
- });
- mlir::RewritePatternSet patterns(context);
- patterns.insert<UnrealizedConversionCastOpPattern>(context);
- scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,
- target);
- (void)mlir::applyPartialConversion(op, target, std::move(patterns));
- }
-}
-
std::optional<std::string> xegpu::getChipStr(Operation *op) {
auto gpuModuleOp = op->getParentOfType<gpu::GPUModuleOp>();
@@ -961,3 +817,152 @@ bool xegpu::matchSplitDimExpansion(
}
return srcIdx == src.size();
}
+
+//===----------------------------------------------------------------------===//
+// Context-aware type conversion utilities
+//===----------------------------------------------------------------------===//
+
+void xegpu::addSCFStructuralMaterializations(TypeConverter &converter) {
+ auto materializeCast = [](OpBuilder &builder, Type type, ValueRange inputs,
+ Location loc) -> Value {
+ return UnrealizedConversionCastOp::create(builder, loc, type, inputs)
+ .getResult(0);
+ };
+ // Source materialization: N:1 (N converted values -> 1 original value).
+ converter.addSourceMaterialization(materializeCast);
+ // Target materialization: 1:1 (single value type conversion).
+ converter.addTargetMaterialization(materializeCast);
+}
+
+void xegpu::addContextAwareVectorTypeConversion(
+ TypeConverter &converter, Operation *topLevelOp,
+ SubShapeAndCountFn getSubShapeAndCount) {
+ // Pre-compute 1:N type mappings for scf.while block arguments only.
+ // During scf.while structural conversion, blocks are detached from their
+ // parent region before convertBlockSignature is called. Block::getParent()
+ // crashes on detached blocks (LLVM ilist assertion), so we cannot look up
+ // layout attributes at that point. Other SCF ops (scf.for, scf.if) keep
+ // blocks attached during conversion.
+ auto whileArgTypeMap = std::make_shared<DenseMap<Value, SmallVector<Type>>>();
+ auto recordBlockArgTypes = [&](Value init, BlockArgument arg) {
+ auto vecTy = dyn_cast<VectorType>(init.getType());
+ if (!vecTy)
+ return;
+ auto layout = xegpu::getDistributeLayoutAttr(init);
+ if (!layout)
+ return;
+ auto [subShape, count] = getSubShapeAndCount(vecTy, layout);
+ if (count <= 0)
+ return;
+ auto newTy = VectorType::get(subShape, vecTy.getElementType());
+ SmallVector<Type> types(count, newTy);
+ (*whileArgTypeMap)[arg] = std::move(types);
+ };
+ topLevelOp->walk([&](scf::WhileOp whileOp) {
+ // "before" region block arguments.
+ for (auto [init, arg] :
+ llvm::zip(whileOp.getInits(), whileOp.getBeforeArguments()))
+ recordBlockArgTypes(init, arg);
+ // "after" region block arguments.
+ for (auto [init, arg] :
+ llvm::zip(whileOp.getInits(), whileOp.getAfterArguments()))
+ recordBlockArgTypes(init, arg);
+ });
+
+ // Context-aware 1:N conversion for VectorType. For scf.while block
+ // arguments, uses the pre-computed map. For all other Values, retrieves
+ // the layout directly via getDistributeLayoutAttr.
+ converter.addConversion(
+ [whileArgTypeMap, getSubShapeAndCount](
+ Value v,
+ SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {
+ if (!isa<VectorType>(v.getType()))
+ return std::nullopt;
+
+ // Check pre-computed map first (for scf.while block args).
+ if (isa<BlockArgument>(v)) {
+ auto it = whileArgTypeMap->find(v);
+ if (it != whileArgTypeMap->end()) {
+ result.append(it->second.begin(), it->second.end());
+ return success();
+ }
+ }
+
+ // For OpResults and other block arguments (scf.for, scf.if, etc.),
+ // retrieve the layout directly.
+ auto layout = xegpu::getDistributeLayoutAttr(v);
+ if (!layout)
+ return std::nullopt;
+
+ auto vecType = cast<VectorType>(v.getType());
+ auto [subShape, count] = getSubShapeAndCount(vecType, layout);
+ if (count <= 0)
+ return std::nullopt;
+
+ auto newTy = VectorType::get(subShape, vecType.getElementType());
+ result.append(count, newTy);
+ return success();
+ });
+}
+
+void xegpu::cleanupUnrealizedConversionCasts(
+ Operation *root,
+ const llvm::SmallSetVector<UnrealizedConversionCastOp, 8> &existingCasts) {
+ OpBuilder builder(root);
+ root->walk([&](UnrealizedConversionCastOp op) {
+ if (existingCasts.contains(op))
+ return;
+ // Handle N:1 cast (N >= 1) where all inputs come from a single 1:N cast.
+ if (op.getNumResults() == 1 && op.getNumOperands() >= 1) {
+ auto defOp =
+ op.getInputs()[0].getDefiningOp<UnrealizedConversionCastOp>();
+ if (defOp && !existingCasts.contains(defOp) &&
+ defOp.getNumOperands() == 1 &&
+ defOp.getNumResults() == op.getNumOperands() &&
+ llvm::all_of(op.getInputs(),
+ [&](Value v) { return v.getDefiningOp() == defOp; })) {
+ Value orig = defOp.getInputs()[0];
+ auto origTy = dyn_cast<VectorType>(orig.getType());
+ auto resTy = dyn_cast<VectorType>(op.getResult(0).getType());
+ if (origTy && resTy &&
+ origTy.getNumElements() == resTy.getNumElements() &&
+ origTy != resTy) {
+ builder.setInsertionPoint(op);
+ auto shapeCast =
+ vector::ShapeCastOp::create(builder, op.getLoc(), resTy, orig);
+ op.replaceAllUsesWith(ValueRange{shapeCast.getResult()});
+ } else {
+ op.replaceAllUsesWith(ValueRange{orig});
+ }
+ }
+ return;
+ }
+ // Handle 1:N cast where the single input comes from an N:1 cast.
+ if (op.getNumOperands() == 1 && op.getNumResults() > 1) {
+ auto defOp =
+ op.getInputs()[0].getDefiningOp<UnrealizedConversionCastOp>();
+ if (defOp && !existingCasts.contains(defOp) &&
+ defOp.getNumResults() == 1 &&
+ defOp.getNumOperands() == op.getNumResults() &&
+ llvm::equal(ValueRange(defOp.getInputs()).getTypes(),
+ op->getResultTypes())) {
+ op.replaceAllUsesWith(defOp.getInputs());
+ }
+ return;
+ }
+ });
+
+ // Erase dead casts iteratively.
+ bool changed = true;
+ while (changed) {
+ changed = false;
+ root->walk([&](UnrealizedConversionCastOp op) {
+ if (existingCasts.contains(op))
+ return;
+ if (op.use_empty()) {
+ op.erase();
+ changed = true;
+ }
+ });
+ }
+}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-rr.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-rr.mlir
index 17a5db6b8401d..415c7455e67ae 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-rr.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-rr.mlir
@@ -304,7 +304,7 @@ gpu.module @test_distribution {
%4 = arith.addi %arg3, %c1_i32 : i32
%6 = xegpu.load_nd %0[%c256] {layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>} : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>> -> vector<256xf32>
scf.yield %6, %4 : vector<256xf32>, i32
- }
+ } attributes {layout_result_0 = #xegpu.layout<sg_layout = [8], sg_data = [16]>}
gpu.return
}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
index f2cc05808ed12..a368d782615dc 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
@@ -405,19 +405,16 @@ 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: %[[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
- // CHECK-DAG: %[[BCAST:.*]] = vector.broadcast %[[LOCAL]] : f32 to vector<1x1xf32>
- // CHECK-DAG: %[[ALLOCA:.*]] = memref.alloca() : memref<64xi8, 3>
- // CHECK-DAG: %[[MEM_DESC:.*]] = xegpu.create_mem_desc %[[ALLOCA]] : memref<64xi8, 3> -> !xegpu.mem_desc<4x4xf32>
- // CHECK-DAG: xegpu.store_matrix %[[BCAST]], %[[MEM_DESC]]{{.*}} : vector<1x1xf32>, !xegpu.mem_desc<4x4xf32>
- // CHECK-DAG: gpu.barrier
- // CHECK-DAG: %[[LOAD_SLM:.*]] = xegpu.load_matrix %[[MEM_DESC]]{{.*}} -> vector<4x4xf32>
- // CHECK-DAG: %[[CST_FINAL:.*]] = arith.constant 0.000000e+00 : f32
- // CHECK-DAG: %[[FINAL:.*]] = vector.multi_reduction <add>, %[[LOAD_SLM]], %[[CST_FINAL]] [0, 1] : vector<4x4xf32> to f32
- // CHECK-DAG: arith.addf %[[FINAL]], %[[CST]] : f32
+ // CHECK: %[[LOAD:.*]] = xegpu.load_nd %{{.*}}[{{%.*}}, {{%.*}}] : !xegpu.tensor_desc<8x8xf32> -> vector<8x8xf32>
+ // CHECK: %[[LOCAL:.*]] = vector.multi_reduction <add>, %[[LOAD]], %{{.*}} [0, 1] : vector<8x8xf32> to f32
+ // CHECK: %[[BCAST:.*]] = vector.broadcast %[[LOCAL]] : f32 to vector<1x1xf32>
+ // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<64xi8, 3>
+ // CHECK: %[[MEM_DESC:.*]] = xegpu.create_mem_desc %[[ALLOCA]] : memref<64xi8, 3> -> !xegpu.mem_desc<4x4xf32>
+ // CHECK: xegpu.store_matrix %[[BCAST]], %[[MEM_DESC]]{{.*}} : vector<1x1xf32>, !xegpu.mem_desc<4x4xf32>
+ // CHECK: gpu.barrier
+ // CHECK: %[[LOAD_SLM:.*]] = xegpu.load_matrix %[[MEM_DESC]]{{.*}} -> vector<4x4xf32>
+ // CHECK: %[[FINAL:.*]] = vector.multi_reduction <add>, %[[LOAD_SLM]], %{{.*}} [0, 1] : vector<4x4xf32> to f32
+ // CHECK: arith.addf %[[FINAL]], %{{.*}} : f32
gpu.func @vector_reduce_scalar_cross_sg(%src: memref<32x32xf32>) {
%cst = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [4, 4], sg_data = [8, 8]>, dims = [0, 1]>} 0.0 : f32
%tdesc = xegpu.create_nd_tdesc %src : memref<32x32xf32>
@@ -1030,7 +1027,7 @@ gpu.module @test_distribution {
%4 = arith.addi %arg3, %c1_i32 : i32
%6 = xegpu.load_nd %0[%c256] {layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>} : !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>> -> vector<256xf32>
scf.yield %6, %4 : vector<256xf32>, i32
- }
+ } attributes {layout_result_0 = #xegpu.layout<sg_layout = [16], sg_data = [16]>}
gpu.return
}
>From 94fa2ecece879dc55d31a9619caf2cbd39bbad6a Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Fri, 1 May 2026 21:49:23 +0000
Subject: [PATCH 2/9] Clean up
---
.../XeGPUSgToWiDistributeExperimental.cpp | 1 +
.../XeGPU/Transforms/XeGPUWgToSgDistribute.cpp | 15 +--------------
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 13 ++++++++-----
3 files changed, 10 insertions(+), 19 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
index e1cf70094b1c3..3e9e4d8c09305 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
@@ -1544,6 +1544,7 @@ void XeGPUSgToWiDistributeExperimentalPass::runOnOperation() {
}
// Fold cancelling cast chains and erase dead casts.
xegpu::cleanupUnrealizedConversionCasts(root, existingCasts);
+ xegpu::removeTemporaryLayoutAttrs(getOperation());
}
void xegpu::populateXeGPUSgToWiDistributeTypeConversions(
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index dc725bcb26e5d..72479e3436ceb 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -1533,18 +1533,5 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
// Fold cancelling cast chains and erase dead casts.
xegpu::cleanupUnrealizedConversionCasts(getOperation(), existingCasts);
-
- // 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 c46f492ec07fe..52978f6102c35 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -859,14 +859,17 @@ void xegpu::addContextAwareVectorTypeConversion(
(*whileArgTypeMap)[arg] = std::move(types);
};
topLevelOp->walk([&](scf::WhileOp whileOp) {
- // "before" region block arguments.
+ // "before" region block arguments correspond to the `inits` operands.
for (auto [init, arg] :
llvm::zip(whileOp.getInits(), whileOp.getBeforeArguments()))
recordBlockArgTypes(init, arg);
- // "after" region block arguments.
- for (auto [init, arg] :
- llvm::zip(whileOp.getInits(), whileOp.getAfterArguments()))
- recordBlockArgTypes(init, arg);
+ // "after" region block arguments correspond to the operands of the
+ // embedded `scf.condition` op (not the `inits`). In general the two
+ // type lists may differ.
+ scf::ConditionOp condOp = whileOp.getConditionOp();
+ for (auto [condArg, arg] :
+ llvm::zip(condOp.getArgs(), whileOp.getAfterArguments()))
+ recordBlockArgTypes(condArg, arg);
});
// Context-aware 1:N conversion for VectorType. For scf.while block
>From 88f4bfdd2932131d1b1b246e29936981f6f8e784 Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Thu, 21 May 2026 02:46:48 +0000
Subject: [PATCH 3/9] Address feedback
---
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 14 +++++++
.../XeGPU/sg-to-wi-experimental-unit.mlir | 1 +
.../Dialect/XeGPU/sg-to-wi-experimental.mlir | 41 +++++++++++++++++++
mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir | 28 +++++++++++++
4 files changed, 84 insertions(+)
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index f81746404b7fc..44d05eed6e654 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -936,6 +936,20 @@ void xegpu::addContextAwareVectorTypeConversion(
void xegpu::cleanupUnrealizedConversionCasts(
Operation *root,
const llvm::SmallSetVector<UnrealizedConversionCastOp, 8> &existingCasts) {
+ // Structural type conversion can generate some redundant
+ // UnrealizedConversionCastOps to materialize the original type from the
+ // type converted (sub-tile) type. These are redundant at this point and
+ // can be eliminated by either folding the cancelling cast chain or, when
+ // the original and final shapes differ but their element counts match,
+ // inserting a vector.shape_cast instead.
+ //
+ // Example (shape differs but element count matches -> shape_cast):
+ // %1 = UnrealizedConversionCastOp %0 : vector<16x1xf32>
+ // to vector<16x16xf32>
+ // %2 = UnrealizedConversionCastOp %1 : vector<16x16xf32>
+ // to vector<16xf32>
+ // becomes:
+ // %2 = vector.shape_cast %0 : vector<16x1xf32> to vector<16xf32>
OpBuilder builder(root);
root->walk([&](UnrealizedConversionCastOp op) {
if (existingCasts.contains(op))
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 0e890ab625b0a..0192a91e8adbd 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-wi-experimental-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-wi-experimental-unit.mlir
@@ -1243,4 +1243,5 @@ gpu.func @vector_multi_reduction_1d_to_scalar() {
}> : 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 c8a9530641951..5e21e3e66078d 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-wi-experimental.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-wi-experimental.mlir
@@ -495,3 +495,44 @@ gpu.module @xevm_module {
gpu.return
}
}
+
+// -----
+// Exercises the shape_cast-emission branch of cleanupUnrealizedConversionCasts:
+// the result of scf.for is a 2D anchor (vector<16x16xf32>) but the distributed
+// type is 1D (vector<16xf32>). The cleanup replaces the N:1 / 1:N
+// UnrealizedConversionCast pair around the scf.for result with a single
+// vector.shape_cast (shapes differ but element counts match), and folds away
+// the in-loop cast pair where shapes already coincide.
+// CHECK-LABEL: gpu.func @scf_for_shape_cast_cleanup
+// CHECK: %[[CST:.*]] = arith.constant dense<0.000000e+00> : vector<16x1xf32>
+// CHECK: %[[LOAD:.*]] = xegpu.load_nd {{.*}} -> vector<16xf32>
+// CHECK: %[[CAST_IN:.*]] = vector.shape_cast %[[LOAD]] : vector<16xf32> to vector<16x1xf32>
+// CHECK: %[[FOR:.*]] = scf.for {{.*}} iter_args(%[[ACC:.*]] = %[[CST]]) -> (vector<16x1xf32>)
+// CHECK: %[[ADD:.*]] = arith.addf %[[ACC]], %[[CAST_IN]] : vector<16x1xf32>
+// CHECK: scf.yield %[[ADD]] : vector<16x1xf32>
+// CHECK: %[[CAST_OUT:.*]] = vector.shape_cast %[[FOR]] : vector<16x1xf32> to vector<16xf32>
+// CHECK: xegpu.store_nd %[[CAST_OUT]]
+// CHECK-NOT: builtin.unrealized_conversion_cast
+gpu.module @xevm_module {
+ gpu.func @scf_for_shape_cast_cleanup(%arg0: memref<16x16xf32>, %arg1: memref<16x16xf32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c16 = arith.constant 16 : index
+ %cst = arith.constant dense<0.0> : vector<16x16xf32>
+ %td = xegpu.create_nd_tdesc %arg0 : memref<16x16xf32> -> !xegpu.tensor_desc<16x16xf32>
+ %ld = xegpu.load_nd %td[%c0, %c0]
+ {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ : !xegpu.tensor_desc<16x16xf32> -> vector<16x16xf32>
+ %r = scf.for %i = %c0 to %c16 step %c1 iter_args(%acc = %cst) -> (vector<16x16xf32>) {
+ %add = arith.addf %acc, %ld
+ {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ : vector<16x16xf32>
+ scf.yield %add : vector<16x16xf32>
+ } {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ %td2 = xegpu.create_nd_tdesc %arg1 : memref<16x16xf32> -> !xegpu.tensor_desc<16x16xf32>
+ xegpu.store_nd %r, %td2[%c0, %c0]
+ {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ : vector<16x16xf32>, !xegpu.tensor_desc<16x16xf32>
+ gpu.return
+ }
+}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
index 50bccb9c300ff..8dc91f1246d84 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
@@ -1204,6 +1204,34 @@ gpu.module @test_distribution {
gpu.return
}
+ // CHECK-LABEL: gpu.func @scf_while_multi_vector_iter_args
+ gpu.func @scf_while_multi_vector_iter_args(%arg0: memref<1024x1024xf32>, %arg1: memref<1024x1024xf32>) {
+ %c1_i32 = arith.constant 1 : i32
+ %c10_i32 = arith.constant 10 : i32
+ %c0_i32 = arith.constant 0 : i32
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<1024x1024xf32> -> !xegpu.tensor_desc<128x128xf32>
+ %1 = xegpu.load_nd %0[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} : !xegpu.tensor_desc<128x128xf32> -> vector<128x128xf32>
+ %2 = xegpu.create_nd_tdesc %arg1 : memref<1024x1024xf32> -> !xegpu.tensor_desc<128x128xf32>
+ %3 = xegpu.load_nd %2[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} : !xegpu.tensor_desc<128x128xf32> -> vector<128x128xf32>
+
+ // CHECK: scf.while {{.*}} : (vector<16x16xf32>, vector<16x16xf32>, i32) -> (vector<16x16xf32>, vector<16x16xf32>, i32)
+ %4:3 = scf.while (%arg2 = %1, %arg3 = %3, %arg4 = %c0_i32) : (vector<128x128xf32>, vector<128x128xf32>, i32) -> (vector<128x128xf32>, vector<128x128xf32>, i32) {
+ %cond = arith.cmpi slt, %arg4, %c10_i32 : i32
+ // CHECK: scf.condition{{.*}} : vector<16x16xf32>, vector<16x16xf32>, i32
+ scf.condition(%cond) %arg2, %arg3, %arg4 : vector<128x128xf32>, vector<128x128xf32>, i32
+ } do {
+ // CHECK: (%{{.*}}: vector<16x16xf32>, %{{.*}}: vector<16x16xf32>, %{{.*}}: i32)
+ ^bb0(%arg2: vector<128x128xf32>, %arg3: vector<128x128xf32>, %arg4: i32):
+ %nx = arith.addi %arg4, %c1_i32 : i32
+ %ld0 = xegpu.load_nd %0[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} : !xegpu.tensor_desc<128x128xf32> -> vector<128x128xf32>
+ %ld1 = xegpu.load_nd %2[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} : !xegpu.tensor_desc<128x128xf32> -> vector<128x128xf32>
+ scf.yield %ld0, %ld1, %nx : vector<128x128xf32>, vector<128x128xf32>, i32
+ }
+ xegpu.store_nd %4#0, %2[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32>
+ xegpu.store_nd %4#1, %2[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32>
+ gpu.return
+ }
+
gpu.func @scf_if(%arg0: memref<1024xf32>, %arg1: memref<1024xf32>) {
%c10 = arith.constant 10 : index
%id = gpu.subgroup_id : index
>From 527561b88f0a3672ae3a30bcdda39d42186fbb32 Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Wed, 27 May 2026 15:13:06 +0000
Subject: [PATCH 4/9] Address feedback
---
.../mlir/Dialect/XeGPU/Utils/XeGPUUtils.h | 7 ++--
.../XeGPU/Transforms/XeGPUBlocking.cpp | 22 ++++++------
.../Transforms/XeGPUSgToLaneDistribute.cpp | 10 +++++-
.../Transforms/XeGPUWgToSgDistribute.cpp | 13 +++++--
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 36 +++++++++----------
5 files changed, 50 insertions(+), 38 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
index 9bf2b751320f4..9ce938b816e02 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
@@ -229,14 +229,11 @@ bool matchSplitDimExpansion(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
/// Callback type for computing sub-shape and count for 1:N VectorType
/// conversion. Given a VectorType and its DistributeLayoutAttr, returns
-/// (subShape, count). A count <= 0 means no conversion is needed.
+/// (subShape, count). The returned count is always >= 1; callers treat
+/// count == 1 as "no split needed".
using SubShapeAndCountFn = std::function<std::pair<SmallVector<int64_t>, int>(
VectorType, DistributeLayoutAttr)>;
-/// Adds source (N:1) and target (1:1) materializations using
-/// UnrealizedConversionCastOp to the given TypeConverter.
-void addSCFStructuralMaterializations(TypeConverter &converter);
-
/// Pre-computes block argument type mappings for SCF loop ops and adds a
/// context-aware 1:N VectorType conversion to the TypeConverter.
/// Pre-computation is needed because during structural type conversion
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
index 62a207283d1cb..97737f10ce749 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
@@ -380,6 +380,7 @@ void XeGPUBlockingPass::runOnOperation() {
tileShape = layout.getEffectiveInstDataAsInt();
count = computeProduct(shape) / computeProduct(tileShape);
}
+ assert(count >= 1 && "count must be at least 1");
return std::make_pair(tileShape, count);
};
@@ -423,17 +424,18 @@ void XeGPUBlockingPass::runOnOperation() {
[&](VectorType vecTy, xegpu::DistributeLayoutAttr layout)
-> std::pair<SmallVector<int64_t>, int> {
if (layout.isForWorkgroup())
- return {{}, 0};
- auto instData = layout.getEffectiveInstDataAsInt();
- if (instData.empty())
- return {{}, 0};
- int count =
- computeProduct(vecTy.getShape()) / computeProduct(instData);
- if (count <= 1)
- return {{}, 0};
- return {SmallVector<int64_t>(instData), count};
+ return {SmallVector<int64_t>(vecTy.getShape()), 1};
+ return getTileShapeAndCount(vecTy.getShape(), layout);
});
- xegpu::addSCFStructuralMaterializations(converter);
+ // Source (N:1) and target (1:1) materializations using
+ // UnrealizedConversionCastOp.
+ auto materializeCast = [](OpBuilder &builder, Type type, ValueRange inputs,
+ Location loc) -> Value {
+ return UnrealizedConversionCastOp::create(builder, loc, type, inputs)
+ .getResult(0);
+ };
+ converter.addSourceMaterialization(materializeCast);
+ converter.addTargetMaterialization(materializeCast);
// Blocking runs SCF conversion separately (not combined with XeGPU
// patterns), so it also needs a 1:N target materialization.
converter.addTargetMaterialization(
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index b7b7b28ba210c..9789b7b56ba81 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1712,7 +1712,15 @@ void XeGPUSgToLaneDistributePass::runOnOperation() {
ConversionTarget target(getContext());
TypeConverter typeConverter;
RewritePatternSet patterns(&getContext());
- xegpu::addSCFStructuralMaterializations(typeConverter);
+ // Source (N:1) and target (1:1) materializations using
+ // UnrealizedConversionCastOp.
+ auto materializeCast = [](OpBuilder &builder, Type type, ValueRange inputs,
+ Location loc) -> Value {
+ return UnrealizedConversionCastOp::create(builder, loc, type, inputs)
+ .getResult(0);
+ };
+ typeConverter.addSourceMaterialization(materializeCast);
+ typeConverter.addTargetMaterialization(materializeCast);
xegpu::populateXeGPUSgToLaneDistributeTypeConversions(typeConverter);
scf::populateSCFStructuralTypeConversionsAndLegality(typeConverter,
patterns, target);
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index 91b4fd7ca8baa..d72c2f2ae6e1e 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -59,6 +59,7 @@ getSgShapeAndCount(ArrayRef<int64_t> shape,
return std::make_pair(sgShape, count);
auto sgData = layout.getEffectiveSgDataAsInt();
count = computeProduct(distributedShape.value()) / computeProduct(sgData);
+ assert(count >= 1 && "count must be at least 1");
return std::make_pair(sgData, count);
}
@@ -1507,7 +1508,7 @@ void populateXeGPUWgToSgDistributeTypeConversions(TypeConverter &converter,
[](VectorType vecTy, xegpu::DistributeLayoutAttr layout)
-> std::pair<SmallVector<int64_t>, int> {
if (!layout.isForWorkgroup())
- return {{}, 0};
+ return {SmallVector<int64_t>(vecTy.getShape()), 1};
return getSgShapeAndCount(vecTy.getShape(), layout);
});
}
@@ -1555,7 +1556,15 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
RewritePatternSet patterns(ctx);
ConversionTarget target(*ctx);
TypeConverter converter;
- xegpu::addSCFStructuralMaterializations(converter);
+ // Source (N:1) and target (1:1) materializations using
+ // UnrealizedConversionCastOp.
+ auto materializeCast = [](OpBuilder &builder, Type type, ValueRange inputs,
+ Location loc) -> Value {
+ return UnrealizedConversionCastOp::create(builder, loc, type, inputs)
+ .getResult(0);
+ };
+ converter.addSourceMaterialization(materializeCast);
+ converter.addTargetMaterialization(materializeCast);
xegpu::populateXeGPUWgToSgDistributeTypeConversions(converter,
getOperation());
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 0ef376349bfe3..e5a50b0591dcf 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -852,27 +852,21 @@ bool xegpu::matchSplitDimExpansion(
// Context-aware type conversion utilities
//===----------------------------------------------------------------------===//
-void xegpu::addSCFStructuralMaterializations(TypeConverter &converter) {
- auto materializeCast = [](OpBuilder &builder, Type type, ValueRange inputs,
- Location loc) -> Value {
- return UnrealizedConversionCastOp::create(builder, loc, type, inputs)
- .getResult(0);
- };
- // Source materialization: N:1 (N converted values -> 1 original value).
- converter.addSourceMaterialization(materializeCast);
- // Target materialization: 1:1 (single value type conversion).
- converter.addTargetMaterialization(materializeCast);
-}
-
void xegpu::addContextAwareVectorTypeConversion(
TypeConverter &converter, Operation *topLevelOp,
SubShapeAndCountFn getSubShapeAndCount) {
- // Pre-compute 1:N type mappings for scf.while block arguments only.
- // During scf.while structural conversion, blocks are detached from their
- // parent region before convertBlockSignature is called. Block::getParent()
- // crashes on detached blocks (LLVM ilist assertion), so we cannot look up
- // layout attributes at that point. Other SCF ops (scf.for, scf.if) keep
- // blocks attached during conversion.
+ // Pre-compute 1:N type mappings for scf.while block arguments.
+ //
+ // Block-arg layouts ARE available in the IR (layout recovery propagates
+ // them onto region block args). The reason we cannot rely on the regular
+ // `getDistributeLayoutAttr(v)` lookup during scf.while conversion is
+ // structural, not informational: `scf::WhileOpConversion` detaches the
+ // before/after blocks from their parent region before invoking
+ // `convertSignatureBlock`. At that point, looking up a BlockArgument's
+ // layout walks `v.getParentBlock()->getParent()`, which trips an LLVM
+ // ilist assertion on detached blocks. The pre-compute is purely a
+ // detached-block workaround; for scf.for / scf.if (whose blocks stay
+ // attached during conversion) the direct lookup works fine.
auto whileArgTypeMap = std::make_shared<DenseMap<Value, SmallVector<Type>>>();
auto recordBlockArgTypes = [&](Value init, BlockArgument arg) {
auto vecTy = dyn_cast<VectorType>(init.getType());
@@ -882,7 +876,8 @@ void xegpu::addContextAwareVectorTypeConversion(
if (!layout)
return;
auto [subShape, count] = getSubShapeAndCount(vecTy, layout);
- if (count <= 0)
+ assert(count >= 1 && "getSubShapeAndCount must return count >= 1");
+ if (count <= 1)
return;
auto newTy = VectorType::get(subShape, vecTy.getElementType());
SmallVector<Type> types(count, newTy);
@@ -929,7 +924,8 @@ void xegpu::addContextAwareVectorTypeConversion(
auto vecType = cast<VectorType>(v.getType());
auto [subShape, count] = getSubShapeAndCount(vecType, layout);
- if (count <= 0)
+ assert(count >= 1 && "getSubShapeAndCount must return count >= 1");
+ if (count <= 1)
return std::nullopt;
auto newTy = VectorType::get(subShape, vecType.getElementType());
>From 3dccb5634dd9085b41124c4db654a4a1c1673277 Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Wed, 27 May 2026 16:53:10 +0000
Subject: [PATCH 5/9] Address feedback
---
mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h | 5 +++--
mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp | 3 +--
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 6 ++----
3 files changed, 6 insertions(+), 8 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
index 9ce938b816e02..3f5e5fbcffe9d 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
@@ -229,8 +229,9 @@ bool matchSplitDimExpansion(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
/// Callback type for computing sub-shape and count for 1:N VectorType
/// conversion. Given a VectorType and its DistributeLayoutAttr, returns
-/// (subShape, count). The returned count is always >= 1; callers treat
-/// count == 1 as "no split needed".
+/// (subShape, count). A count <= 0 signals "no conversion needed"; a
+/// count >= 1 produces `count` copies of `subShape` (count == 1 is a
+/// 1:1 shape-changing conversion).
using SubShapeAndCountFn = std::function<std::pair<SmallVector<int64_t>, int>(
VectorType, DistributeLayoutAttr)>;
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index d72c2f2ae6e1e..b379300b13bac 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -59,7 +59,6 @@ getSgShapeAndCount(ArrayRef<int64_t> shape,
return std::make_pair(sgShape, count);
auto sgData = layout.getEffectiveSgDataAsInt();
count = computeProduct(distributedShape.value()) / computeProduct(sgData);
- assert(count >= 1 && "count must be at least 1");
return std::make_pair(sgData, count);
}
@@ -1508,7 +1507,7 @@ void populateXeGPUWgToSgDistributeTypeConversions(TypeConverter &converter,
[](VectorType vecTy, xegpu::DistributeLayoutAttr layout)
-> std::pair<SmallVector<int64_t>, int> {
if (!layout.isForWorkgroup())
- return {SmallVector<int64_t>(vecTy.getShape()), 1};
+ return {{}, 0};
return getSgShapeAndCount(vecTy.getShape(), layout);
});
}
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index e5a50b0591dcf..0ccc843b9a96d 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -876,8 +876,7 @@ void xegpu::addContextAwareVectorTypeConversion(
if (!layout)
return;
auto [subShape, count] = getSubShapeAndCount(vecTy, layout);
- assert(count >= 1 && "getSubShapeAndCount must return count >= 1");
- if (count <= 1)
+ if (count <= 0)
return;
auto newTy = VectorType::get(subShape, vecTy.getElementType());
SmallVector<Type> types(count, newTy);
@@ -924,8 +923,7 @@ void xegpu::addContextAwareVectorTypeConversion(
auto vecType = cast<VectorType>(v.getType());
auto [subShape, count] = getSubShapeAndCount(vecType, layout);
- assert(count >= 1 && "getSubShapeAndCount must return count >= 1");
- if (count <= 1)
+ if (count <= 0)
return std::nullopt;
auto newTy = VectorType::get(subShape, vecType.getElementType());
>From 76e08d6e78659c0bc65748b787060aa5f6a16ab4 Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Wed, 27 May 2026 17:55:13 +0000
Subject: [PATCH 6/9] Rename to xegpu::addVectorTypeConversion
---
mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h | 2 +-
mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp | 2 +-
mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp | 2 +-
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 2 +-
4 files changed, 4 insertions(+), 4 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
index 3f5e5fbcffe9d..c577fb265522b 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
@@ -242,7 +242,7 @@ using SubShapeAndCountFn = std::function<std::pair<SmallVector<int64_t>, int>(
/// making Block::getParent() crash (LLVM ilist assertion). The
/// `getSubShapeAndCount` callback computes (subShape, count) for a VectorType
/// and its layout; count <= 0 means no conversion needed.
-void addContextAwareVectorTypeConversion(
+void addVectorTypeConversion(
TypeConverter &converter, Operation *topLevelOp,
SubShapeAndCountFn getSubShapeAndCount);
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
index 97737f10ce749..5b9e6cdf9b188 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
@@ -419,7 +419,7 @@ void XeGPUBlockingPass::runOnOperation() {
});
// Context-aware 1:N conversion for VectorType based on inst_data.
- xegpu::addContextAwareVectorTypeConversion(
+ xegpu::addVectorTypeConversion(
converter, op,
[&](VectorType vecTy, xegpu::DistributeLayoutAttr layout)
-> std::pair<SmallVector<int64_t>, int> {
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index b379300b13bac..f0c474a8d7261 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -1502,7 +1502,7 @@ void populateXeGPUWgToSgDistributeTypeConversions(TypeConverter &converter,
});
// Context-aware 1:N conversion for VectorType based on sg_layout/sg_data.
- xegpu::addContextAwareVectorTypeConversion(
+ xegpu::addVectorTypeConversion(
converter, topLevelOp,
[](VectorType vecTy, xegpu::DistributeLayoutAttr layout)
-> std::pair<SmallVector<int64_t>, int> {
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 0ccc843b9a96d..b75a6d1c2a793 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -852,7 +852,7 @@ bool xegpu::matchSplitDimExpansion(
// Context-aware type conversion utilities
//===----------------------------------------------------------------------===//
-void xegpu::addContextAwareVectorTypeConversion(
+void xegpu::addVectorTypeConversion(
TypeConverter &converter, Operation *topLevelOp,
SubShapeAndCountFn getSubShapeAndCount) {
// Pre-compute 1:N type mappings for scf.while block arguments.
>From 834c8722d2a1cd5e55490832e9009cc09cbd3cfe Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Wed, 27 May 2026 21:22:05 +0000
Subject: [PATCH 7/9] more clean up
---
.../Dialect/XeGPU/Transforms/Transforms.h | 4 +-
.../mlir/Dialect/XeGPU/Utils/XeGPUUtils.h | 43 +++++++++------
.../XeGPU/Transforms/XeGPUBlocking.cpp | 20 +++----
.../Transforms/XeGPUSgToLaneDistribute.cpp | 52 +++++++++----------
.../Transforms/XeGPUWgToSgDistribute.cpp | 22 ++++----
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 47 ++++++++++-------
.../Dialect/XeGPU/sg-to-lane-distribute.mlir | 38 ++++++++++++++
.../lib/Dialect/XeGPU/TestXeGPUTransforms.cpp | 2 +-
8 files changed, 144 insertions(+), 84 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
index 301faff41329d..7fb024c434570 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h
@@ -81,14 +81,14 @@ void populateXeGPUWgToSgDistributePatterns(RewritePatternSet &patterns);
/// Define only the type conversions needed for XeGPU subgroup to lane
/// distribution.
void populateXeGPUSgToLaneDistributeTypeConversions(
- TypeConverter &typeConverter);
+ TypeConverter &typeConverter, Operation *topLevelOp);
/// Defines type conversions and legality for XeGPU subgroup to lane
/// distribution and appends the required conversion patterns into `patterns`.
/// Appends patterns for XeGPU subgroup to lane distribution into
/// `patterns`.
void populateXeGPUSgToLaneDistributeTypeConversionAndLegality(
TypeConverter &typeConverter, RewritePatternSet &patterns,
- ConversionTarget &target);
+ ConversionTarget &target, Operation *topLevelOp);
/// Collect a set of patterns to unroll xegpu operations to a smaller shapes.
/// Users can control whether an operation to be unrolled or not, as well as
diff --git a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
index c577fb265522b..600c6e6a55f80 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
@@ -227,24 +227,37 @@ bool matchUnitDimExpansion(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
bool matchSplitDimExpansion(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
SmallVector<SmallVector<int64_t>> &splitDimGroups);
-/// Callback type for computing sub-shape and count for 1:N VectorType
-/// conversion. Given a VectorType and its DistributeLayoutAttr, returns
-/// (subShape, count). A count <= 0 signals "no conversion needed"; a
-/// count >= 1 produces `count` copies of `subShape` (count == 1 is a
-/// 1:1 shape-changing conversion).
+/// Callback type for computing sub-shape and count for 1:N (or 1:1
+/// shape-changing) VectorType conversion. Given a VectorType and its
+/// DistributeLayoutAttr, returns (subShape, count). A count <= 0 signals
+/// "no conversion needed"; count == 1 is a 1:1 shape-changing conversion;
+/// count > 1 produces `count` copies of `subShape`.
using SubShapeAndCountFn = std::function<std::pair<SmallVector<int64_t>, int>(
VectorType, DistributeLayoutAttr)>;
-/// Pre-computes block argument type mappings for SCF loop ops and adds a
-/// context-aware 1:N VectorType conversion to the TypeConverter.
-/// Pre-computation is needed because during structural type conversion
-/// (especially scf.while), blocks may be detached from their parent region,
-/// making Block::getParent() crash (LLVM ilist assertion). The
-/// `getSubShapeAndCount` callback computes (subShape, count) for a VectorType
-/// and its layout; count <= 0 means no conversion needed.
-void addVectorTypeConversion(
- TypeConverter &converter, Operation *topLevelOp,
- SubShapeAndCountFn getSubShapeAndCount);
+/// Pre-computes VectorType mappings for `scf.while` block arguments under
+/// `topLevelOp` (1:1 shape-changing or 1:N). Pre-computation is needed because
+/// during structural type conversion `scf::WhileOpConversion` detaches the
+/// before/after blocks from their parent region before invoking
+/// `convertSignatureBlock`. At that point, looking up a BlockArgument's layout
+/// via `getDistributeLayoutAttr` walks `v.getParentBlock()->getParent()`,
+/// which trips an LLVM ilist assertion on detached blocks. For scf.for /
+/// scf.if (whose blocks stay attached) the direct lookup works, so they don't
+/// need this pre-computation.
+DenseMap<Value, SmallVector<Type>>
+precomputeWhileBlockArgTypes(Operation *topLevelOp,
+ SubShapeAndCountFn getSubShapeAndCount);
+
+/// Adds a context-aware VectorType conversion to `converter` (1:1
+/// shape-changing or 1:N, depending on `getSubShapeAndCount`'s returned
+/// count). `getSubShapeAndCount` computes (subShape, count) for a VectorType
+/// and its layout; count <= 0 means no conversion needed. `whileArgTypes`
+/// (typically obtained from `precomputeWhileBlockArgTypes`) provides the
+/// pre-computed types for `scf.while` block arguments; pass an empty map if
+/// the IR has no `scf.while` ops.
+void addVectorTypeConversion(TypeConverter &converter,
+ SubShapeAndCountFn getSubShapeAndCount,
+ DenseMap<Value, SmallVector<Type>> whileArgTypes);
/// Cleans up UnrealizedConversionCastOps inserted during SCF structural type
/// conversion. Folds cancelling N:1->1:N and 1:N->N:1 cast chains (inserting
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
index 5b9e6cdf9b188..444c0f32222a1 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
@@ -418,15 +418,17 @@ void XeGPUBlockingPass::runOnOperation() {
return success();
});
- // Context-aware 1:N conversion for VectorType based on inst_data.
- xegpu::addVectorTypeConversion(
- converter, op,
- [&](VectorType vecTy, xegpu::DistributeLayoutAttr layout)
- -> std::pair<SmallVector<int64_t>, int> {
- if (layout.isForWorkgroup())
- return {SmallVector<int64_t>(vecTy.getShape()), 1};
- return getTileShapeAndCount(vecTy.getShape(), layout);
- });
+ // Context-aware VectorType conversion based on inst_data (1:1
+ // shape-changing or 1:N).
+ auto getSubShapeAndCount = [&](VectorType vecTy,
+ xegpu::DistributeLayoutAttr layout)
+ -> std::pair<SmallVector<int64_t>, int> {
+ return getTileShapeAndCount(vecTy.getShape(), layout);
+ };
+ auto whileArgTypes =
+ xegpu::precomputeWhileBlockArgTypes(op, getSubShapeAndCount);
+ xegpu::addVectorTypeConversion(converter, getSubShapeAndCount,
+ std::move(whileArgTypes));
// Source (N:1) and target (1:1) materializations using
// UnrealizedConversionCastOp.
auto materializeCast = [](OpBuilder &builder, Type type, ValueRange inputs,
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 9789b7b56ba81..4dd56a5774f7a 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1721,11 +1721,11 @@ void XeGPUSgToLaneDistributePass::runOnOperation() {
};
typeConverter.addSourceMaterialization(materializeCast);
typeConverter.addTargetMaterialization(materializeCast);
- xegpu::populateXeGPUSgToLaneDistributeTypeConversions(typeConverter);
+ xegpu::populateXeGPUSgToLaneDistributeTypeConversions(typeConverter, root);
scf::populateSCFStructuralTypeConversionsAndLegality(typeConverter,
patterns, target);
xegpu::populateXeGPUSgToLaneDistributeTypeConversionAndLegality(
- typeConverter, patterns, target);
+ typeConverter, patterns, target, root);
target.addLegalOp<UnrealizedConversionCastOp>();
(void)applyPartialConversion(root, target, std::move(patterns));
}
@@ -1735,13 +1735,10 @@ void XeGPUSgToLaneDistributePass::runOnOperation() {
}
void xegpu::populateXeGPUSgToLaneDistributeTypeConversions(
- TypeConverter &typeConverter) {
- // Any type other than TensorDescType and VectorType are legal as is.
- typeConverter.addConversion([](Type type) -> std::optional<Type> {
- if (!isa<TensorDescType, VectorType>(type))
- return type;
- return std::nullopt;
- });
+ TypeConverter &typeConverter, Operation *topLevelOp) {
+ // Pass through any type by default; more specific conversions registered
+ // below override this for TensorDescType and (distributing) VectorType.
+ typeConverter.addConversion([](Type type) -> Type { return type; });
// For TensorDescType, drop the layout attribute if any.
typeConverter.addConversion([](TensorDescType type) -> Type {
if (type.getLayoutAttr()) {
@@ -1749,29 +1746,28 @@ void xegpu::populateXeGPUSgToLaneDistributeTypeConversions(
}
return type;
});
- // For VectorType, check if there is a distribute layout attribute on the
- // value. If so, convert to the distributed vector type based on the layout.
- typeConverter.addConversion([](Value v) -> std::optional<Type> {
- auto type = v.getType();
- // If value is not vector type, nothing to do.
- if (!isa<VectorType>(type))
- return std::nullopt;
- auto layout = xegpu::getDistributeLayoutAttr(v);
- if (!layout || !layout.isForSubgroup())
- return type;
- // Vector type is distributed based on lane layout.
- auto newTyOrFailure =
- getDistVecTypeBasedOnLaneLayout(layout, cast<VectorType>(type));
- if (failed(newTyOrFailure))
- return type;
- return *newTyOrFailure;
- });
+ // For VectorType, distribute based on the lane layout (1:1 shape-changing
+ // conversion). Uses xegpu::addVectorTypeConversion with a pre-computed
+ // map for scf.while block args (see precomputeWhileBlockArgTypes for the
+ // detached-block rationale).
+ auto getSubShapeAndCount = [](VectorType vecTy,
+ xegpu::DistributeLayoutAttr layout)
+ -> std::pair<SmallVector<int64_t>, int> {
+ auto distTyOrFailure = getDistVecTypeBasedOnLaneLayout(layout, vecTy);
+ if (failed(distTyOrFailure))
+ return {{}, 0};
+ return {SmallVector<int64_t>(distTyOrFailure->getShape()), 1};
+ };
+ auto whileArgTypes =
+ xegpu::precomputeWhileBlockArgTypes(topLevelOp, getSubShapeAndCount);
+ xegpu::addVectorTypeConversion(typeConverter, getSubShapeAndCount,
+ std::move(whileArgTypes));
}
void xegpu::populateXeGPUSgToLaneDistributeTypeConversionAndLegality(
TypeConverter &typeConverter, RewritePatternSet &patterns,
- ConversionTarget &target) {
- populateXeGPUSgToLaneDistributeTypeConversions(typeConverter);
+ ConversionTarget &target, Operation *topLevelOp) {
+ populateXeGPUSgToLaneDistributeTypeConversions(typeConverter, topLevelOp);
// CreateNdDescOp is legal only if its result type has no layout attribute.
target.addDynamicallyLegalOp<xegpu::CreateNdDescOp>(
[&](xegpu::CreateNdDescOp op) { return !op.getType().getLayoutAttr(); });
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index f0c474a8d7261..2c8e8e77010f5 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -1501,15 +1501,19 @@ void populateXeGPUWgToSgDistributeTypeConversions(TypeConverter &converter,
return success();
});
- // Context-aware 1:N conversion for VectorType based on sg_layout/sg_data.
- xegpu::addVectorTypeConversion(
- converter, topLevelOp,
- [](VectorType vecTy, xegpu::DistributeLayoutAttr layout)
- -> std::pair<SmallVector<int64_t>, int> {
- if (!layout.isForWorkgroup())
- return {{}, 0};
- return getSgShapeAndCount(vecTy.getShape(), layout);
- });
+ // Context-aware VectorType conversion based on sg_layout/sg_data
+ // (1:1 shape-changing or 1:N).
+ auto getSubShapeAndCount = [](VectorType vecTy,
+ xegpu::DistributeLayoutAttr layout)
+ -> std::pair<SmallVector<int64_t>, int> {
+ if (!layout.isForWorkgroup())
+ return {{}, 0};
+ return getSgShapeAndCount(vecTy.getShape(), layout);
+ };
+ auto whileArgTypes =
+ xegpu::precomputeWhileBlockArgTypes(topLevelOp, getSubShapeAndCount);
+ xegpu::addVectorTypeConversion(converter, getSubShapeAndCount,
+ std::move(whileArgTypes));
}
void populateXeGPUWgToSgDistributePatterns(RewritePatternSet &patterns) {
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index b75a6d1c2a793..2eb6a154ae0be 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -852,22 +852,22 @@ bool xegpu::matchSplitDimExpansion(
// Context-aware type conversion utilities
//===----------------------------------------------------------------------===//
-void xegpu::addVectorTypeConversion(
- TypeConverter &converter, Operation *topLevelOp,
- SubShapeAndCountFn getSubShapeAndCount) {
- // Pre-compute 1:N type mappings for scf.while block arguments.
- //
- // Block-arg layouts ARE available in the IR (layout recovery propagates
- // them onto region block args). The reason we cannot rely on the regular
- // `getDistributeLayoutAttr(v)` lookup during scf.while conversion is
- // structural, not informational: `scf::WhileOpConversion` detaches the
- // before/after blocks from their parent region before invoking
- // `convertSignatureBlock`. At that point, looking up a BlockArgument's
- // layout walks `v.getParentBlock()->getParent()`, which trips an LLVM
- // ilist assertion on detached blocks. The pre-compute is purely a
- // detached-block workaround; for scf.for / scf.if (whose blocks stay
- // attached during conversion) the direct lookup works fine.
- auto whileArgTypeMap = std::make_shared<DenseMap<Value, SmallVector<Type>>>();
+// Pre-computes block argument type mappings for scf.while ops.
+//
+// Block-arg layouts ARE available in the IR (layout recovery propagates
+// them onto region block args). The reason we cannot rely on the regular
+// `getDistributeLayoutAttr(v)` lookup during scf.while conversion is
+// structural, not informational: `scf::WhileOpConversion` detaches the
+// before/after blocks from their parent region before invoking
+// `convertSignatureBlock`. At that point, looking up a BlockArgument's
+// layout walks `v.getParentBlock()->getParent()`, which trips an LLVM
+// ilist assertion on detached blocks. The pre-compute is purely a
+// detached-block workaround; for scf.for / scf.if (whose blocks stay
+// attached during conversion) the direct lookup works fine.
+DenseMap<Value, SmallVector<Type>>
+xegpu::precomputeWhileBlockArgTypes(Operation *topLevelOp,
+ SubShapeAndCountFn getSubShapeAndCount) {
+ DenseMap<Value, SmallVector<Type>> whileArgTypes;
auto recordBlockArgTypes = [&](Value init, BlockArgument arg) {
auto vecTy = dyn_cast<VectorType>(init.getType());
if (!vecTy)
@@ -880,7 +880,7 @@ void xegpu::addVectorTypeConversion(
return;
auto newTy = VectorType::get(subShape, vecTy.getElementType());
SmallVector<Type> types(count, newTy);
- (*whileArgTypeMap)[arg] = std::move(types);
+ whileArgTypes[arg] = std::move(types);
};
topLevelOp->walk([&](scf::WhileOp whileOp) {
// "before" region block arguments correspond to the `inits` operands.
@@ -895,10 +895,17 @@ void xegpu::addVectorTypeConversion(
llvm::zip(condOp.getArgs(), whileOp.getAfterArguments()))
recordBlockArgTypes(condArg, arg);
});
+ return whileArgTypes;
+}
- // Context-aware 1:N conversion for VectorType. For scf.while block
- // arguments, uses the pre-computed map. For all other Values, retrieves
- // the layout directly via getDistributeLayoutAttr.
+void xegpu::addVectorTypeConversion(
+ TypeConverter &converter, SubShapeAndCountFn getSubShapeAndCount,
+ DenseMap<Value, SmallVector<Type>> whileArgTypes) {
+ // Context-aware VectorType conversion (1:1 shape-changing or 1:N). For
+ // scf.while block arguments, uses the pre-computed map. For all other
+ // Values, retrieves the layout directly via getDistributeLayoutAttr.
+ auto whileArgTypeMap = std::make_shared<DenseMap<Value, SmallVector<Type>>>(
+ std::move(whileArgTypes));
converter.addConversion(
[whileArgTypeMap, getSubShapeAndCount](
Value v,
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
index 9aa72d9179cf0..fa9897770a08e 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
@@ -536,3 +536,41 @@ gpu.module @xevm_module {
gpu.return
}
}
+
+// -----
+// CHECK-LABEL: gpu.func @scf_while_lane_distribute
+// CHECK: scf.while {{.*}} : (vector<16x1xf32>, i32) -> (vector<16x1xf32>, i32)
+// CHECK: scf.condition{{.*}} : vector<16x1xf32>, i32
+// CHECK: (%{{.*}}: vector<16x1xf32>, %{{.*}}: i32)
+// CHECK: scf.yield %{{.*}}, %{{.*}} : vector<16x1xf32>, i32
+gpu.module @xevm_module {
+ gpu.func @scf_while_lane_distribute(%arg0: memref<16x16xf32>, %arg1: memref<16x16xf32>) {
+ %c0 = arith.constant 0 : index
+ %c0_i32 = arith.constant 0 : i32
+ %c10_i32 = arith.constant 10 : i32
+ %c1_i32 = arith.constant 1 : i32
+ %td0 = xegpu.create_nd_tdesc %arg0 : memref<16x16xf32> -> !xegpu.tensor_desc<16x16xf32>
+ %ld0 = xegpu.load_nd %td0[%c0, %c0]
+ {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ : !xegpu.tensor_desc<16x16xf32> -> vector<16x16xf32>
+ %td1 = xegpu.create_nd_tdesc %arg1 : memref<16x16xf32> -> !xegpu.tensor_desc<16x16xf32>
+ %r:2 = scf.while (%arg2 = %ld0, %arg3 = %c0_i32) : (vector<16x16xf32>, i32) -> (vector<16x16xf32>, i32) {
+ %cond = arith.cmpi slt, %arg3, %c10_i32 : i32
+ scf.condition(%cond) %arg2, %arg3 : vector<16x16xf32>, i32
+ } do {
+ ^bb0(%arg2: vector<16x16xf32>, %arg3: i32):
+ xegpu.store_nd %arg2, %td1[%c0, %c0]
+ {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ : vector<16x16xf32>, !xegpu.tensor_desc<16x16xf32>
+ %next = arith.addi %arg3, %c1_i32 : i32
+ %ld_next = xegpu.load_nd %td0[%c0, %c0]
+ {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ : !xegpu.tensor_desc<16x16xf32> -> vector<16x16xf32>
+ scf.yield %ld_next, %next : vector<16x16xf32>, i32
+ }
+ xegpu.store_nd %r#0, %td1[%c0, %c0]
+ {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ : vector<16x16xf32>, !xegpu.tensor_desc<16x16xf32>
+ gpu.return
+ }
+}
diff --git a/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
index 5c3721630837d..7e621fe9a1526 100644
--- a/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
+++ b/mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp
@@ -282,7 +282,7 @@ struct TestXeGPUSgToLaneDistribute
ConversionTarget target(*ctx);
RewritePatternSet patterns(ctx);
xegpu::populateXeGPUSgToLaneDistributeTypeConversionAndLegality(
- typeConverter, patterns, target);
+ typeConverter, patterns, target, op);
(void)applyPartialConversion(op, target, std::move(patterns));
}
};
>From b4840e065b1a299498c86f7c8a5b274769280427 Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Thu, 28 May 2026 18:40:34 +0000
Subject: [PATCH 8/9] Add pre-compute for scf.for
---
.../mlir/Dialect/XeGPU/Utils/XeGPUUtils.h | 34 +++----
.../XeGPU/Transforms/XeGPUBlocking.cpp | 6 +-
.../Transforms/XeGPUSgToLaneDistribute.cpp | 10 +-
.../Transforms/XeGPUWgToSgDistribute.cpp | 6 +-
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 96 +++++++++++--------
5 files changed, 87 insertions(+), 65 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
index 600c6e6a55f80..aa965af602680 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
@@ -235,29 +235,31 @@ bool matchSplitDimExpansion(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
using SubShapeAndCountFn = std::function<std::pair<SmallVector<int64_t>, int>(
VectorType, DistributeLayoutAttr)>;
-/// Pre-computes VectorType mappings for `scf.while` block arguments under
-/// `topLevelOp` (1:1 shape-changing or 1:N). Pre-computation is needed because
-/// during structural type conversion `scf::WhileOpConversion` detaches the
-/// before/after blocks from their parent region before invoking
-/// `convertSignatureBlock`. At that point, looking up a BlockArgument's layout
-/// via `getDistributeLayoutAttr` walks `v.getParentBlock()->getParent()`,
-/// which trips an LLVM ilist assertion on detached blocks. For scf.for /
-/// scf.if (whose blocks stay attached) the direct lookup works, so they don't
-/// need this pre-computation.
+/// Pre-computes VectorType mappings for SCF loop block arguments under
+/// `topLevelOp` (1:1 shape-changing or 1:N). Covers `scf.while`
+/// before/after-region arguments and `scf.for` iter_args. Pre-computation is
+/// needed because during structural type conversion the SCF converters
+/// transiently detach/replace the loop body before all queries against the
+/// old block arguments have finished, at which point
+/// `getDistributeLayoutAttr(BlockArgument)` returns null (either the parent
+/// op is gone, or the new parent op has not inherited the temporary
+/// `layout_operand_N` attributes set by layout recovery). Caching the
+/// distributed types by `Value` identity sidesteps both failure modes.
+/// `scf.if` has no block arguments and therefore needs no entry here.
DenseMap<Value, SmallVector<Type>>
-precomputeWhileBlockArgTypes(Operation *topLevelOp,
- SubShapeAndCountFn getSubShapeAndCount);
+precomputeLoopBlockArgTypes(Operation *topLevelOp,
+ SubShapeAndCountFn getSubShapeAndCount);
/// Adds a context-aware VectorType conversion to `converter` (1:1
/// shape-changing or 1:N, depending on `getSubShapeAndCount`'s returned
/// count). `getSubShapeAndCount` computes (subShape, count) for a VectorType
-/// and its layout; count <= 0 means no conversion needed. `whileArgTypes`
-/// (typically obtained from `precomputeWhileBlockArgTypes`) provides the
-/// pre-computed types for `scf.while` block arguments; pass an empty map if
-/// the IR has no `scf.while` ops.
+/// and its layout; count <= 0 means no conversion needed. `loopArgTypes`
+/// (typically obtained from `precomputeLoopBlockArgTypes`) provides the
+/// pre-computed types for SCF loop block arguments (`scf.while`,
+/// `scf.for`); pass an empty map if the IR has no such loops.
void addVectorTypeConversion(TypeConverter &converter,
SubShapeAndCountFn getSubShapeAndCount,
- DenseMap<Value, SmallVector<Type>> whileArgTypes);
+ DenseMap<Value, SmallVector<Type>> loopArgTypes);
/// Cleans up UnrealizedConversionCastOps inserted during SCF structural type
/// conversion. Folds cancelling N:1->1:N and 1:N->N:1 cast chains (inserting
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
index 444c0f32222a1..7e1a271aaa984 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
@@ -425,10 +425,10 @@ void XeGPUBlockingPass::runOnOperation() {
-> std::pair<SmallVector<int64_t>, int> {
return getTileShapeAndCount(vecTy.getShape(), layout);
};
- auto whileArgTypes =
- xegpu::precomputeWhileBlockArgTypes(op, getSubShapeAndCount);
+ auto loopArgTypes =
+ xegpu::precomputeLoopBlockArgTypes(op, getSubShapeAndCount);
xegpu::addVectorTypeConversion(converter, getSubShapeAndCount,
- std::move(whileArgTypes));
+ std::move(loopArgTypes));
// Source (N:1) and target (1:1) materializations using
// UnrealizedConversionCastOp.
auto materializeCast = [](OpBuilder &builder, Type type, ValueRange inputs,
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 4dd56a5774f7a..924cc8ac81755 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1748,8 +1748,8 @@ void xegpu::populateXeGPUSgToLaneDistributeTypeConversions(
});
// For VectorType, distribute based on the lane layout (1:1 shape-changing
// conversion). Uses xegpu::addVectorTypeConversion with a pre-computed
- // map for scf.while block args (see precomputeWhileBlockArgTypes for the
- // detached-block rationale).
+ // map for SCF loop block args (see precomputeLoopBlockArgTypes for the
+ // rationale).
auto getSubShapeAndCount = [](VectorType vecTy,
xegpu::DistributeLayoutAttr layout)
-> std::pair<SmallVector<int64_t>, int> {
@@ -1758,10 +1758,10 @@ void xegpu::populateXeGPUSgToLaneDistributeTypeConversions(
return {{}, 0};
return {SmallVector<int64_t>(distTyOrFailure->getShape()), 1};
};
- auto whileArgTypes =
- xegpu::precomputeWhileBlockArgTypes(topLevelOp, getSubShapeAndCount);
+ auto loopArgTypes =
+ xegpu::precomputeLoopBlockArgTypes(topLevelOp, getSubShapeAndCount);
xegpu::addVectorTypeConversion(typeConverter, getSubShapeAndCount,
- std::move(whileArgTypes));
+ std::move(loopArgTypes));
}
void xegpu::populateXeGPUSgToLaneDistributeTypeConversionAndLegality(
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index 2c8e8e77010f5..e1d842fbebf50 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -1510,10 +1510,10 @@ void populateXeGPUWgToSgDistributeTypeConversions(TypeConverter &converter,
return {{}, 0};
return getSgShapeAndCount(vecTy.getShape(), layout);
};
- auto whileArgTypes =
- xegpu::precomputeWhileBlockArgTypes(topLevelOp, getSubShapeAndCount);
+ auto loopArgTypes =
+ xegpu::precomputeLoopBlockArgTypes(topLevelOp, getSubShapeAndCount);
xegpu::addVectorTypeConversion(converter, getSubShapeAndCount,
- std::move(whileArgTypes));
+ std::move(loopArgTypes));
}
void populateXeGPUWgToSgDistributePatterns(RewritePatternSet &patterns) {
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 2eb6a154ae0be..633f089c350e6 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -852,22 +852,30 @@ bool xegpu::matchSplitDimExpansion(
// Context-aware type conversion utilities
//===----------------------------------------------------------------------===//
-// Pre-computes block argument type mappings for scf.while ops.
+// Pre-computes block argument type mappings for SCF loops (scf.while,
+// scf.for).
//
// Block-arg layouts ARE available in the IR (layout recovery propagates
-// them onto region block args). The reason we cannot rely on the regular
-// `getDistributeLayoutAttr(v)` lookup during scf.while conversion is
-// structural, not informational: `scf::WhileOpConversion` detaches the
-// before/after blocks from their parent region before invoking
-// `convertSignatureBlock`. At that point, looking up a BlockArgument's
-// layout walks `v.getParentBlock()->getParent()`, which trips an LLVM
-// ilist assertion on detached blocks. The pre-compute is purely a
-// detached-block workaround; for scf.for / scf.if (whose blocks stay
-// attached during conversion) the direct lookup works fine.
+// them onto the loop op as `layout_operand_N`). The reason we cannot rely
+// on the regular `getDistributeLayoutAttr(v)` lookup during structural
+// conversion is structural, not informational:
+// - For `scf.while`, `scf::WhileOpConversion` detaches the before/after
+// blocks from their parent region before invoking
+// `convertSignatureBlock`. Looking up a detached BlockArgument's layout
+// walks `v.getParentBlock()->getParent()` and trips an LLVM ilist
+// assertion.
+// - For `scf.for`, `scf::ForOpConverter` builds a new `scf.for` and moves
+// the body block into it. The new op does NOT inherit the temporary
+// `layout_operand_N` attributes that layout recovery set on the old
+// op, so any post-move query of a body block argument's layout (e.g.
+// when a pattern that consumes the iter_arg via a non-anchor op like
+// `vector.insert_strided_slice` runs after the move) returns null.
+// Caching the distributed types by `Value` identity sidesteps both failure
+// modes. `scf.if` has no block arguments and is therefore not covered here.
DenseMap<Value, SmallVector<Type>>
-xegpu::precomputeWhileBlockArgTypes(Operation *topLevelOp,
- SubShapeAndCountFn getSubShapeAndCount) {
- DenseMap<Value, SmallVector<Type>> whileArgTypes;
+xegpu::precomputeLoopBlockArgTypes(Operation *topLevelOp,
+ SubShapeAndCountFn getSubShapeAndCount) {
+ DenseMap<Value, SmallVector<Type>> loopArgTypes;
auto recordBlockArgTypes = [&](Value init, BlockArgument arg) {
auto vecTy = dyn_cast<VectorType>(init.getType());
if (!vecTy)
@@ -880,50 +888,62 @@ xegpu::precomputeWhileBlockArgTypes(Operation *topLevelOp,
return;
auto newTy = VectorType::get(subShape, vecTy.getElementType());
SmallVector<Type> types(count, newTy);
- whileArgTypes[arg] = std::move(types);
+ loopArgTypes[arg] = std::move(types);
};
- topLevelOp->walk([&](scf::WhileOp whileOp) {
- // "before" region block arguments correspond to the `inits` operands.
- for (auto [init, arg] :
- llvm::zip(whileOp.getInits(), whileOp.getBeforeArguments()))
- recordBlockArgTypes(init, arg);
- // "after" region block arguments correspond to the operands of the
- // embedded `scf.condition` op (not the `inits`). In general the two
- // type lists may differ.
- scf::ConditionOp condOp = whileOp.getConditionOp();
- for (auto [condArg, arg] :
- llvm::zip(condOp.getArgs(), whileOp.getAfterArguments()))
- recordBlockArgTypes(condArg, arg);
+ topLevelOp->walk([&](Operation *op) {
+ if (auto whileOp = dyn_cast<scf::WhileOp>(op)) {
+ // "before" region block arguments correspond to the `inits` operands.
+ for (auto [init, arg] :
+ llvm::zip(whileOp.getInits(), whileOp.getBeforeArguments()))
+ recordBlockArgTypes(init, arg);
+ // "after" region block arguments correspond to the operands of the
+ // embedded `scf.condition` op (not the `inits`). In general the two
+ // type lists may differ.
+ scf::ConditionOp condOp = whileOp.getConditionOp();
+ for (auto [condArg, arg] :
+ llvm::zip(condOp.getArgs(), whileOp.getAfterArguments()))
+ recordBlockArgTypes(condArg, arg);
+ return;
+ }
+ if (auto forOp = dyn_cast<scf::ForOp>(op)) {
+ // Body block args (excluding the induction variable) correspond to
+ // the `initArgs` operands.
+ for (auto [init, arg] :
+ llvm::zip(forOp.getInitArgs(), forOp.getRegionIterArgs()))
+ recordBlockArgTypes(init, arg);
+ return;
+ }
});
- return whileArgTypes;
+ return loopArgTypes;
}
void xegpu::addVectorTypeConversion(
TypeConverter &converter, SubShapeAndCountFn getSubShapeAndCount,
- DenseMap<Value, SmallVector<Type>> whileArgTypes) {
+ DenseMap<Value, SmallVector<Type>> loopArgTypes) {
// Context-aware VectorType conversion (1:1 shape-changing or 1:N). For
- // scf.while block arguments, uses the pre-computed map. For all other
- // Values, retrieves the layout directly via getDistributeLayoutAttr.
- auto whileArgTypeMap = std::make_shared<DenseMap<Value, SmallVector<Type>>>(
- std::move(whileArgTypes));
+ // SCF loop block arguments (scf.while, scf.for), uses the pre-computed
+ // map. For all other Values, retrieves the layout directly via
+ // getDistributeLayoutAttr.
+ auto loopArgTypeMap = std::make_shared<DenseMap<Value, SmallVector<Type>>>(
+ std::move(loopArgTypes));
converter.addConversion(
- [whileArgTypeMap, getSubShapeAndCount](
+ [loopArgTypeMap, getSubShapeAndCount](
Value v,
SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {
if (!isa<VectorType>(v.getType()))
return std::nullopt;
- // Check pre-computed map first (for scf.while block args).
+ // Check pre-computed map first (for SCF loop block args).
if (isa<BlockArgument>(v)) {
- auto it = whileArgTypeMap->find(v);
- if (it != whileArgTypeMap->end()) {
+ auto it = loopArgTypeMap->find(v);
+ if (it != loopArgTypeMap->end()) {
result.append(it->second.begin(), it->second.end());
return success();
}
}
- // For OpResults and other block arguments (scf.for, scf.if, etc.),
- // retrieve the layout directly.
+ // For OpResults and other block arguments (e.g. region args of
+ // non-loop ops), retrieve the layout directly.
auto layout = xegpu::getDistributeLayoutAttr(v);
if (!layout)
return std::nullopt;
>From e43126180d17fa5a5c8edb43b0e27e2d80d0debb Mon Sep 17 00:00:00 2001
From: nbpatel <nishant.b.patel at intel.com>
Date: Thu, 4 Jun 2026 18:14:22 +0000
Subject: [PATCH 9/9] Address feedback
---
.../mlir/Dialect/XeGPU/Utils/XeGPUUtils.h | 9 ++--
.../XeGPU/Transforms/XeGPUBlocking.cpp | 52 ++-----------------
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 31 +++++++++++
3 files changed, 42 insertions(+), 50 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
index aa965af602680..9ca375044bd08 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
@@ -262,9 +262,12 @@ void addVectorTypeConversion(TypeConverter &converter,
DenseMap<Value, SmallVector<Type>> loopArgTypes);
/// Cleans up UnrealizedConversionCastOps inserted during SCF structural type
-/// conversion. Folds cancelling N:1->1:N and 1:N->N:1 cast chains (inserting
-/// vector.shape_cast when shapes differ but element counts match), and
-/// erases dead casts. Casts in `existingCasts` are preserved.
+/// conversion and/or XeGPU unrolling. Folds cancelling N:1->1:N and 1:N->N:1
+/// cast chains (inserting vector.shape_cast when shapes differ but element
+/// counts match). Unpaired pack (1:N) and unpack (N:1) casts between a single
+/// large VectorType and N identically-typed smaller VectorTypes are lowered
+/// to vector.extract_strided_slice / vector.insert_strided_slice. Dead casts
+/// are erased. Casts in `existingCasts` are preserved.
void cleanupUnrealizedConversionCasts(
Operation *root,
const llvm::SmallSetVector<UnrealizedConversionCastOp, 8> &existingCasts);
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
index 7e1a271aaa984..6ee44d7dd160c 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
@@ -37,49 +37,6 @@ using namespace mlir;
namespace {
-// reslove the unrealized conversion cast ops generated when doing SCF
-// Structural Type Conversion. It will have two formats, N:1 vector
-// cast and 1:N vector cast. vector::insert_strided_slice ops will be
-// used for the first case, and vector::extract_strided_slice ops will be
-// used for the second case.
-static void
-resolveUnrealizedConversionCastOp(UnrealizedConversionCastOp castOp) {
- ValueRange inputs = castOp.getInputs();
- ValueRange outputs = castOp.getOutputs();
-
- auto hasIdenticalVectorTypes = [](ValueRange values) {
- auto types = values.getTypes();
- return llvm::all_of(types, [&](Type type) {
- return isa<VectorType>(type) && type == types.front();
- });
- };
-
- // We only interest in the case where all inputs and outputs have the
- // identical VectorTypes
- if (!hasIdenticalVectorTypes(inputs) || !hasIdenticalVectorTypes(outputs)) {
- LDBG() << "skip unrealized conversion cast op not emulating pack/unpack.";
- return;
- }
-
- VectorType outputTy = dyn_cast<VectorType>(outputs[0].getType());
- OpBuilder builder(castOp);
- if (inputs.size() > 1 && outputs.size() == 1) {
- // the castOp is emulating an unpack op
- ArrayRef<int64_t> shape = outputTy.getShape();
- Value result = xegpu::createVectorWithShapeFromValues(
- builder, castOp.getLoc(), inputs, shape);
- castOp->replaceAllUsesWith(ValueRange(result));
- castOp->erase();
- } else if (castOp.getNumResults() > 1 && castOp.getNumOperands() == 1) {
- // the castOp is emulating a pack op
- ArrayRef<int64_t> tileShape = outputTy.getShape();
- SmallVector<Value> results = xegpu::extractVectorsWithShapeFromValue(
- builder, castOp.getLoc(), inputs[0], tileShape);
- castOp->replaceAllUsesWith(results);
- castOp->erase();
- }
-}
-
//===------------------------------------------------------------------------===//
// The XeGPUBlockingPass leverages the unroll patterns for XeGPU and Vector ops
// to partition operations that process large shapes into multiple operations on
@@ -529,12 +486,13 @@ void XeGPUBlockingPass::runOnOperation() {
SmallVector<NamedAttribute> newAttrs =
xegpu::dropInstDataOnAttrs(op->getAttrs());
op->setAttrs(newAttrs);
-
- // Resolve unrealized conversion cast ops emulating pack/unpack
- if (auto castOp = dyn_cast<UnrealizedConversionCastOp>(op))
- resolveUnrealizedConversionCastOp(castOp);
});
+ // Resolve UnrealizedConversionCastOps generated by SCF structural type
+ // conversion and by XeGPU/Vector unrolling (cancelling cast chains and
+ // unpaired pack/unpack casts).
+ xegpu::cleanupUnrealizedConversionCasts(op, existingCasts);
+
// One more round of folding to clean up the intermediate
// insert/extract strided slice ops.
RewritePatternSet emptyPatterns(ctx);
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 633f089c350e6..bc5575d3e4a31 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -976,6 +976,16 @@ void xegpu::cleanupUnrealizedConversionCasts(
// to vector<16xf32>
// becomes:
// %2 = vector.shape_cast %0 : vector<16x1xf32> to vector<16xf32>
+ //
+ // For unpaired casts that emulate a pack (1:N) or unpack (N:1) between a
+ // single large VectorType and N identically-typed smaller VectorTypes,
+ // lower to vector.extract_strided_slice / vector.insert_strided_slice.
+ auto hasIdenticalVectorTypes = [](ValueRange values) {
+ auto types = values.getTypes();
+ return !types.empty() && llvm::all_of(types, [&](Type type) {
+ return isa<VectorType>(type) && type == types.front();
+ });
+ };
OpBuilder builder(root);
root->walk([&](UnrealizedConversionCastOp op) {
if (existingCasts.contains(op))
@@ -1002,6 +1012,17 @@ void xegpu::cleanupUnrealizedConversionCasts(
} else {
op.replaceAllUsesWith(ValueRange{orig});
}
+ return;
+ }
+ // Unpaired N:1 cast emulating unpack: stitch inputs into the output
+ // shape via vector.insert_strided_slice.
+ auto outputTy = dyn_cast<VectorType>(op.getResult(0).getType());
+ if (op.getNumOperands() > 1 && outputTy &&
+ hasIdenticalVectorTypes(op.getInputs())) {
+ builder.setInsertionPoint(op);
+ Value result = xegpu::createVectorWithShapeFromValues(
+ builder, op.getLoc(), op.getInputs(), outputTy.getShape());
+ op->replaceAllUsesWith(ValueRange(result));
}
return;
}
@@ -1015,6 +1036,16 @@ void xegpu::cleanupUnrealizedConversionCasts(
llvm::equal(ValueRange(defOp.getInputs()).getTypes(),
op->getResultTypes())) {
op.replaceAllUsesWith(defOp.getInputs());
+ return;
+ }
+ // Unpaired 1:N cast emulating pack: split the input into the output
+ // tile shape via vector.extract_strided_slice.
+ auto tileTy = dyn_cast<VectorType>(op.getResult(0).getType());
+ if (tileTy && hasIdenticalVectorTypes(op.getResults())) {
+ builder.setInsertionPoint(op);
+ SmallVector<Value> results = xegpu::extractVectorsWithShapeFromValue(
+ builder, op.getLoc(), op.getInputs()[0], tileTy.getShape());
+ op->replaceAllUsesWith(results);
}
return;
}
More information about the Mlir-commits
mailing list