Skip to content

[MLIR][Vector] Add distribution pattern for vector::ConstantMaskOp #172268

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

Merged
akroviakov merged 3 commits into from
Dec 16, 2025
Merged

[MLIR][Vector] Add distribution pattern for vector::ConstantMaskOp #172268

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6d81029
[MLIR][Vector] Add distribution pattern for `vector::ConstantMaskOp`
akroviakov Dec 15, 2025
cbf6a94
Use templates
akroviakov Dec 16, 2025
498608b
Update pattern description
akroviakov Dec 16, 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
45 changes: 31 additions & 14 deletions mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1100,12 +1100,14 @@ struct WarpOpShapeCast : public WarpDistributionPattern {
}
};

/// Sink out vector.create_mask op feeding into a warp op yield.
/// Sink out vector.create_mask / vector.constant_mask op feeding into a warp op
/// yield.
/// ```
/// %0 = ...
/// %1 = gpu.warp_execute_on_lane_0(%arg0) -> (vector<1xf32>) {
/// ...
/// %mask = vector.create_mask %0 : vector<32xi1>
/// // or %mask = vector.constant_mask[2] : vector<32xi1>
/// gpu.yield %mask : vector<32xi1>
/// }
/// ```
Expand All @@ -1118,31 +1120,45 @@ struct WarpOpShapeCast : public WarpDistributionPattern {
/// %cmp = arith.cmpi ult, %laneid, %0
/// %ub = arith.select %cmp, %c0, %c1
/// %1 = vector.create_mask %ub : vector<1xi1>
template <typename OpType,
typename = std::enable_if_t<llvm::is_one_of<
OpType, vector::CreateMaskOp, vector::ConstantMaskOp>::value>>
struct WarpOpCreateMask : public WarpDistributionPattern {
using Base::Base;
LogicalResult matchAndRewrite(WarpExecuteOnLane0Op warpOp,
Copy link

@Garra1980 Garra1980 Dec 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Probably function comments need to be tweaked a bit

PatternRewriter &rewriter) const override {
OpOperand *yieldOperand =
getWarpResult(warpOp, llvm::IsaPred<vector::CreateMaskOp>);
OpOperand *yieldOperand = getWarpResult(warpOp, (llvm::IsaPred<OpType>));
if (!yieldOperand)
return failure();

auto mask = yieldOperand->get().getDefiningOp<vector::CreateMaskOp>();
Operation *mask = yieldOperand->get().getDefiningOp<OpType>();

// Early exit if any values needed for calculating the new mask indices
// are defined inside the warp op.
if (!llvm::all_of(mask->getOperands(), [&](Value value) {
if (mask->getOperands().size() &&
!llvm::all_of(mask->getOperands(), [&](Value value) {
return warpOp.isDefinedOutsideOfRegion(value);
}))
return failure();

Location loc = mask.getLoc();
Location loc = mask->getLoc();
unsigned operandIndex = yieldOperand->getOperandNumber();

auto distType = cast<VectorType>(warpOp.getResult(operandIndex).getType());
VectorType seqType = mask.getVectorType();
VectorType seqType = cast<VectorType>(mask->getResult(0).getType());
ArrayRef<int64_t> seqShape = seqType.getShape();
ArrayRef<int64_t> distShape = distType.getShape();
SmallVector<Value> materializedOperands;
if constexpr (std::is_same_v<OpType, vector::CreateMaskOp>) {
materializedOperands.append(mask->getOperands().begin(),
mask->getOperands().end());
} else {
auto constantMaskOp = cast<vector::ConstantMaskOp>(mask);
auto dimSizes = constantMaskOp.getMaskDimSizesAttr().asArrayRef();
for (auto dimSize : dimSizes)
materializedOperands.push_back(
arith::ConstantIndexOp::create(rewriter, loc, dimSize).getResult());
}

rewriter.setInsertionPointAfter(warpOp);

Expand Down Expand Up @@ -1170,7 +1186,7 @@ struct WarpOpCreateMask : public WarpDistributionPattern {
// mask sizes are always in the range [0, mask_vector_size[i]).
Value maskDimIdx = affine::makeComposedAffineApply(
rewriter, loc, s1 - s0 * distShape[i],
{delinearizedIds[i], mask.getOperand(i)});
{delinearizedIds[i], materializedOperands[i]});
newOperands.push_back(maskDimIdx);
}

Expand Down Expand Up @@ -2282,12 +2298,13 @@ void mlir::vector::populatePropagateWarpVectorDistributionPatterns(
const WarpShuffleFromIdxFn &warpShuffleFromIdxFn, PatternBenefit benefit,
PatternBenefit readBenefit) {
patterns.add<WarpOpTransferRead>(patterns.getContext(), readBenefit);
patterns
.add<WarpOpElementwise, WarpOpDeadResult, WarpOpBroadcast,
WarpOpShapeCast, WarpOpExtract, WarpOpForwardOperand, WarpOpConstant,
WarpOpInsertScalar, WarpOpInsert, WarpOpCreateMask,
WarpOpExtractStridedSlice, WarpOpInsertStridedSlice, WarpOpStep>(
patterns.getContext(), benefit);
patterns.add<WarpOpElementwise, WarpOpDeadResult, WarpOpBroadcast,
WarpOpShapeCast, WarpOpExtract, WarpOpForwardOperand,
WarpOpConstant, WarpOpInsertScalar, WarpOpInsert,
WarpOpCreateMask<vector::CreateMaskOp>,
WarpOpCreateMask<vector::ConstantMaskOp>,
WarpOpExtractStridedSlice, WarpOpInsertStridedSlice, WarpOpStep>(
patterns.getContext(), benefit);
patterns.add<WarpOpExtractScalar>(patterns.getContext(), warpShuffleFromIdxFn,
benefit);
patterns.add<WarpOpScfForOp>(patterns.getContext(), distributionMapFn,
Expand Down
33 changes: 33 additions & 0 deletions mlir/test/Dialect/Vector/vector-warp-distribute.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1779,6 +1779,21 @@ func.func @warp_propagate_unconnected_read_write(%laneid: index, %buffer: memref
// CHECK-DIST-AND-PROP: %[[CST:.+]] = arith.constant dense<2.000000e+00> : vector<4xf32>
// CHECK-DIST-AND-PROP: vector.transfer_read {{.*}} : memref<128xf32>, vector<4xf32>
// CHECK-DIST-AND-PROP: vector.transfer_write %[[CST]], {{.*}} : vector<4xf32>, memref<128xf32>
// -----

func.func @warp_propagate_constant_mask(%laneid: index) -> vector<1xi1> {
%r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<1xi1>) {
%1 = vector.constant_mask [1] : vector<32xi1>
gpu.yield %1 : vector<32xi1>
}
return %r : vector<1xi1>
}

// CHECK-PROP-DAG: #[[$SUB:.*]] = affine_map<()[s0] -> (-s0 + 1)>
// CHECK-PROP-LABEL: func @warp_propagate_constant_mask
// CHECK-PROP-SAME: %[[LANEID:.+]]: index
// CHECK-PROP: %[[MDIST:.+]] = affine.apply #[[$SUB]]()[%[[LANEID]]]
// CHECK-PROP: vector.create_mask %[[MDIST]] : vector<1xi1>

// -----

Expand Down Expand Up @@ -1813,6 +1828,24 @@ func.func @warp_propagate_multi_dim_create_mask(%laneid: index, %m0: index, %m1:
// CHECK-PROP: %[[DISTM0:.+]] = affine.apply #[[$SUBM0]]()[%[[M0]], %[[LANEID]]]
// CHECK-PROP: %[[DISTM1:.+]] = affine.apply #[[$SUBM1]]()[%[[M1]], %[[LANEID]]]
// CHECK-PROP: vector.create_mask %[[DISTM0]], %[[DISTM1]], %[[M2]] : vector<1x2x4xi1>
// -----

func.func @warp_propagate_multi_dim_constant_mask(%laneid: index) -> vector<1x2x4xi1> {
%r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<1x2x4xi1>) {
%1 = vector.constant_mask [1, 1, 2]: vector<16x4x4xi1>
gpu.yield %1 : vector<16x4x4xi1>
}
return %r : vector<1x2x4xi1>
}

// CHECK-PROP-DAG: #[[$SUBM0:.*]] = affine_map<()[s0] -> (-(s0 floordiv 2) + 1)>
// CHECK-PROP-DAG: #[[$SUBM1:.*]] = affine_map<()[s0] -> (s0 * -2 + (s0 floordiv 2) * 4 + 1)>
// CHECK-PROP-LABEL: func @warp_propagate_multi_dim_constant_mask
// CHECK-PROP-SAME: %[[LANEID:.+]]: index
// CHECK-PROP: %[[CST2:.+]] = arith.constant 2 : index
// CHECK-PROP: %[[DISTM0:.+]] = affine.apply #[[$SUBM0]]()[%[[LANEID]]]
// CHECK-PROP: %[[DISTM1:.+]] = affine.apply #[[$SUBM1]]()[%[[LANEID]]]
// CHECK-PROP: vector.create_mask %[[DISTM0]], %[[DISTM1]], %[[CST2]] : vector<1x2x4xi1>

// -----

Expand Down