Skip to content

[mlir][vector] Add linearization pattern for vector.splat #137651

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
merged 5 commits into from
May 1, 2025
Merged

[mlir][vector] Add linearization pattern for vector.splat #137651

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
97a6c57
Add linearization pattern for vector.splat
nbpatel Apr 25, 2025
fef6390
Merge branch 'main' into vector_linearize_splat
nbpatel May 1, 2025
ff82c48
add newline
nbpatel May 1, 2025
9b21851
Remove targetVectorBitWidth
nbpatel May 1, 2025
8dbe3cc
Fix comments
nbpatel May 1, 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
33 changes: 32 additions & 1 deletion mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -293,6 +293,10 @@ struct LinearizeVectorExtract final
LogicalResult
matchAndRewrite(vector::ExtractOp extractOp, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
// Skip if result is not a vector type
if (!isa<VectorType>(extractOp.getType()))
return rewriter.notifyMatchFailure(extractOp,
"scalar extract is not supported.");
Type dstTy = getTypeConverter()->convertType(extractOp.getType());
assert(dstTy && "expected 1-D vector type");

Expand Down Expand Up @@ -415,6 +419,32 @@ struct LinearizeVectorBitCast final
}
};

/// This pattern converts the SplatOp to work on a linearized vector.
/// Following,
/// vector.splat %value : vector<4x4xf32>
/// is converted to:
/// %out_1d = vector.splat %value : vector<16xf32>
/// %out_nd = vector.shape_cast %out_1d : vector<16xf32> to vector<4x4xf32>
struct LinearizeVectorSplat final
: public OpConversionPattern<vector::SplatOp> {
using OpConversionPattern::OpConversionPattern;

LinearizeVectorSplat(const TypeConverter &typeConverter, MLIRContext *context,
PatternBenefit benefit = 1)
: OpConversionPattern(typeConverter, context, benefit) {}

LogicalResult
matchAndRewrite(vector::SplatOp splatOp, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto dstTy = getTypeConverter()->convertType(splatOp.getType());
if (!dstTy)
return rewriter.notifyMatchFailure(splatOp, "cannot convert type.");
rewriter.replaceOpWithNewOp<vector::SplatOp>(splatOp, adaptor.getInput(),
dstTy);
return success();
}
};

} // namespace

/// Return true if the operation `op` does not support scalable vectors and
Expand Down Expand Up @@ -501,7 +531,8 @@ void mlir::vector::populateVectorLinearizeBasePatterns(
const TypeConverter &typeConverter, const ConversionTarget &target,
RewritePatternSet &patterns) {
patterns.add<LinearizeConstantLike, LinearizeVectorizable,
LinearizeVectorBitCast>(typeConverter, patterns.getContext());
LinearizeVectorBitCast, LinearizeVectorSplat>(
typeConverter, patterns.getContext());
}

void mlir::vector::populateVectorLinearizeShuffleLikeOpsPatterns(
Expand Down
34 changes: 34 additions & 0 deletions mlir/test/Dialect/Vector/linearize.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -413,3 +413,37 @@ func.func @test_vector_bitcast(%arg0: vector<[4]x2xf32>) -> vector<[4]x4xf16> {
%1 = vector.bitcast %arg0 : vector<[4]x2xf32> to vector<[4]x4xf16>
return %1 : vector<[4]x4xf16>
}

// -----
// ALL-LABEL: linearize_vector_splat
// ALL-SAME: (%[[ARG:.*]]: i32) -> vector<4x2xi32>
func.func @linearize_vector_splat(%arg0: i32) -> vector<4x2xi32> {
// DEFAULT: %[[SPLAT:.*]] = vector.splat %[[ARG]] : vector<8xi32>
// DEFAULT: %[[CAST:.*]] = vector.shape_cast %[[SPLAT]] : vector<8xi32> to vector<4x2xi32>
// DEFAULT: return %[[CAST]] : vector<4x2xi32>
// BW-128: %[[SPLAT:.*]] = vector.splat %[[ARG]] : vector<8xi32>
// BW-128: %[[CAST:.*]] = vector.shape_cast %[[SPLAT]] : vector<8xi32> to vector<4x2xi32>
// BW-128: return %[[CAST]] : vector<4x2xi32>

// BW-0: %[[SPLAT:.*]] = vector.splat %[[ARG]] : vector<4x2xi32>
// BW-0: return %[[SPLAT]] : vector<4x2xi32>
%0 = vector.splat %arg0 : vector<4x2xi32>
Copy link
Contributor

Choose a reason for hiding this comment

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

This should also work (i.e. scalable vector):

  %0 = vector.splat %arg0 : vector<4x[2]xi32>

Could you try and if it works, add a test? Thanks!

Copy link
Contributor Author

Choose a reason for hiding this comment

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

added the test

return %0 : vector<4x2xi32>
}

// -----
// ALL-LABEL: linearize_scalable_vector_splat
// ALL-SAME: (%[[ARG:.*]]: i32) -> vector<4x[2]xi32>
func.func @linearize_scalable_vector_splat(%arg0: i32) -> vector<4x[2]xi32> {
// DEFAULT: %[[SPLAT:.*]] = vector.splat %[[ARG]] : vector<[8]xi32>
// DEFAULT: %[[CAST:.*]] = vector.shape_cast %[[SPLAT]] : vector<[8]xi32> to vector<4x[2]xi32>
// DEFAULT: return %[[CAST]] : vector<4x[2]xi32>
// BW-128: %[[SPLAT:.*]] = vector.splat %[[ARG]] : vector<[8]xi32>
// BW-128: %[[CAST:.*]] = vector.shape_cast %[[SPLAT]] : vector<[8]xi32> to vector<4x[2]xi32>
// BW-128: return %[[CAST]] : vector<4x[2]xi32>

// BW-0: %[[SPLAT:.*]] = vector.splat %[[ARG]] : vector<4x[2]xi32>
// BW-0: return %[[SPLAT]] : vector<4x[2]xi32>
%0 = vector.splat %arg0 : vector<4x[2]xi32>
return %0 : vector<4x[2]xi32>
}