Skip to content

[MLIR][XeGPU] Add distribution for vector.create_mask from Wg to Sg #169571

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
+74 −20
Open

[MLIR][XeGPU] Add distribution for vector.create_mask from Wg to Sg #169571

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3df7560
Add distribution for vector.create_mask op
nbpatel Nov 25, 2025
f9c930d
Templatize mask ops
nbpatel Nov 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 29 additions & 20 deletions mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1270,15 +1270,15 @@ struct WgToSgVectorTransposeOp
}
};

// This pattern distributes the vector.constant_mask ops to work at subgroup
// level.
struct WgToSgVectorConstantMaskOp
: public OpConversionPattern<vector::ConstantMaskOp> {
using OpConversionPattern<vector::ConstantMaskOp>::OpConversionPattern;

LogicalResult
matchAndRewrite(vector::ConstantMaskOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
// Distribute vector mask ops to work at subgroup level.
template <typename MaskOpType>
struct WgToSgVectorMaskOp : public OpConversionPattern<MaskOpType> {
using OpConversionPattern<MaskOpType>::OpConversionPattern;

LogicalResult matchAndRewrite(
MaskOpType op,
typename OpConversionPattern<MaskOpType>::OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
xegpu::DistributeLayoutAttr layout =
xegpu::getDistributeLayoutAttr(op.getResult());
if (!layout || !layout.isForWorkgroup())
Expand All @@ -1288,9 +1288,16 @@ struct WgToSgVectorConstantMaskOp
VectorType type = op.getResult().getType();
auto wgShape = type.getShape();

ArrayRef<int64_t> wgMaskDimSizes = op.getMaskDimSizes();
SmallVector<Value> wgMaskDimSizes;
if constexpr (std::is_same_v<MaskOpType, vector::ConstantMaskOp>) {
for (int64_t maskSize : op.getMaskDimSizes()) {
wgMaskDimSizes.push_back(
arith::ConstantIndexOp::create(rewriter, loc, maskSize));
}
} else if constexpr (std::is_same_v<MaskOpType, vector::CreateMaskOp>) {
wgMaskDimSizes = llvm::to_vector(op.getOperands());
}

// Get subgroup ID.
Value sgId =
gpu::SubgroupIdOp::create(rewriter, loc, /*upper_bound=*/nullptr);
auto sgOffsets =
Expand All @@ -1302,19 +1309,17 @@ struct WgToSgVectorConstantMaskOp
VectorType resultType = VectorType::get(sgShape, type.getElementType());

// In each dimension, each subgroup computes its local mask size as:
// min(max(wgMaskSize[d] - offset[d], 0), sgDimSize[d])
// min(max(wgMaskDimSize[d] - offset[d], 0), sgDimSize[d])
SmallVector<Value> newCreateMaskOps;
for (auto offsetSet : *sgOffsets) {
SmallVector<Value> maskOperands;

for (auto [i, wgMaskSize] : llvm::enumerate(wgMaskDimSizes)) {
Value wgMaskSizeVal =
arith::ConstantIndexOp::create(rewriter, loc, wgMaskSize);
for (auto [i, wgMaskDimSize] : llvm::enumerate(wgMaskDimSizes)) {
Value dimSizeVal =
arith::ConstantIndexOp::create(rewriter, loc, sgShape[i]);
Value offset = offsetSet[i];
Value adjustedMaskSize =
arith::SubIOp::create(rewriter, loc, wgMaskSizeVal, offset);
arith::SubIOp::create(rewriter, loc, wgMaskDimSize, offset);
Value zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
Value nonNegative =
arith::MaxSIOp::create(rewriter, loc, adjustedMaskSize, zero);
Expand All @@ -1335,6 +1340,8 @@ struct WgToSgVectorConstantMaskOp
}
};

using WgToSgVectorConstantMaskOp = WgToSgVectorMaskOp<vector::ConstantMaskOp>;
using WgToSgVectorCreateMaskOp = WgToSgVectorMaskOp<vector::CreateMaskOp>;
} // namespace

namespace mlir {
Expand All @@ -1350,7 +1357,8 @@ void populateXeGPUWgToSgDistributePatterns(RewritePatternSet &patterns) {
WgToSgStoreScatterOpWithOffset, WgToSgLoadMatrixOp,
WgToSgStoreMatrixOp, WgToSgVectorStepOp, WgToSgVectorShapeCastOp,
WgToSgMultiDimReductionOp, WgToSgVectorTransposeOp,
WgToSgVectorConstantMaskOp>(patterns.getContext());
WgToSgVectorConstantMaskOp, WgToSgVectorCreateMaskOp>(
patterns.getContext());
}
} // namespace xegpu
} // namespace mlir
Expand Down Expand Up @@ -1477,9 +1485,10 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
return isLegal(layout);
});

target.addDynamicallyLegalOp<
vector::ShapeCastOp, vector::StepOp, vector::TransposeOp,
vector::BroadcastOp, vector::MultiDimReductionOp, vector::ConstantMaskOp>(
target.addDynamicallyLegalOp<vector::ShapeCastOp, vector::StepOp,
vector::TransposeOp, vector::BroadcastOp,
vector::MultiDimReductionOp,
vector::ConstantMaskOp, vector::CreateMaskOp>(
[=](Operation *op) -> bool {
// Check for either a SliceAttr or LayoutAttr on the result.
auto layout = xegpu::getDistributeLayoutAttr(op->getResult(0));
Expand Down
8 changes: 8 additions & 0 deletions mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops-rr.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -138,5 +138,13 @@ gpu.module @test_distribution {
%constant_mask = vector.constant_mask [16, 16] {layout_result_0 = #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16]>} : vector<256x128xi1>
gpu.return
}

gpu.func @vector_create_mask_2D() {
// CHECK-COUNT-4: vector.create_mask {{.*}}, {{.*}} : vector<16x16xi1>
// CHECK-NOT: vector.create_mask
%cst16 = arith.constant 16 : index
%constant_mask = vector.create_mask %cst16, %cst16 {layout_result_0 = #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16]>} : vector<256x128xi1>
gpu.return
}
}

37 changes: 37 additions & 0 deletions mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -583,6 +583,43 @@ gpu.module @test_distribution {
gpu.return
}

// CHECK-LABEL: vector_create_mask_1D
gpu.func @vector_create_mask_1D() {
// CHECK-DAG: %[[SGID:.*]] = gpu.subgroup_id : index
// CHECK-DAG: %[[REMU:.*]] = index.remu %[[SGID]], %[[C2:.*]]
// CHECK-DAG: %[[MUL:.*]] = index.mul %[[REMU]], %[[C16:.*]]
// CHECK-DAG: %[[REMU2:.*]] = index.remu %[[MUL]], %[[C32:.*]]
// CHECK-DAG: %[[SUB:.*]] = arith.subi %[[C8:.*]], %[[REMU2]] : index
// CHECK-DAG: %[[MAX:.*]] = arith.maxsi %[[SUB]], %[[C0:.*]] : index
// CHECK-DAG: %[[MIN:.*]] = arith.minsi %[[MAX]], %[[C16:.*]] : index
// CHECK-DAG: %[[MASK:.*]] = vector.create_mask %[[MIN]] : vector<16xi1>
%cst8 = arith.constant 8 : index
%constant_mask = vector.create_mask %cst8 {layout_result_0 = #xegpu.layout<sg_layout = [2], sg_data = [16]>} : vector<32xi1>
gpu.return
}

// CHECK-LABEL: vector_create_mask_2D
gpu.func @vector_create_mask_2D() {
// CHECK-DAG: %[[SGID:.*]] = gpu.subgroup_id : index
// CHECK-DAG: %[[SGIDX:.*]] = index.remu %[[SGID]], %[[C4:.*]]
// CHECK-DAG: %[[SGIDY_TMP:.*]] = index.divu %[[SGID]], %[[C4:.*]]
// CHECK-DAG: %[[SGIDY:.*]] = index.remu %[[SGIDY_TMP]], %[[C8:.*]]
// CHECK-DAG: %[[ROW:.*]] = index.mul %[[SGIDY]], %[[C32:.*]]
// CHECK-DAG: %[[COL:.*]] = index.mul %[[SGIDX]], %[[C32:.*]]
// CHECK-DAG: %[[MODROW:.*]] = index.remu %[[ROW]], %[[C256:.*]]
// CHECK-DAG: %[[MODCOL:.*]] = index.remu %[[COL]], %[[C128:.*]]
// CHECK-DAG: %[[SUBROW:.*]] = arith.subi %[[C16:.*]], %[[MODROW]] : index
// CHECK-DAG: %[[MAXROW:.*]] = arith.maxsi %[[SUBROW]], %[[C0:.*]] : index
// CHECK-DAG: %[[MINROW:.*]] = arith.minsi %[[MAXROW]], %[[C32:.*]] : index
// CHECK-DAG: %[[SUBCOL:.*]] = arith.subi %[[C16:.*]], %[[MODCOL]] : index
// CHECK-DAG: %[[MAXCOL:.*]] = arith.maxsi %[[SUBCOL]], %[[C0:.*]] : index
// CHECK-DAG: %[[MINCOL:.*]] = arith.minsi %[[MAXCOL]], %[[C32:.*]] : index
// CHECK-DAG: %[[MASK:.*]] = vector.create_mask %[[MINROW]], %[[MINCOL]] : vector<32x32xi1>
%cst16 = arith.constant 16 : index
%constant_mask = vector.create_mask %cst16, %cst16 {layout_result_0 = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>} : vector<256x128xi1>
gpu.return
}

// CHECK-LABEL: distribute_load_slice_attr
gpu.func @distribute_load_slice_attr() {
%2 = memref.alloca() {alignment = 1024} : memref<4096xf32>
Expand Down