Thanks to visit codestin.com
Credit goes to github.com

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 1 commit
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
Next Next commit
Add linearization pattern for vector.splat
  • Loading branch information
@nbpatel
nbpatel committed Apr 28, 2025
commit 97a6c57217bab6a815ffcd9bd12905af4d5fca1a
63 changes: 53 additions & 10 deletions mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,9 @@

using namespace mlir;

constexpr unsigned defaultTargetVectorBitWidth =
std::numeric_limits<unsigned>::max();

static bool isLessThanTargetBitWidth(Operation *op, unsigned targetBitWidth) {
auto resultTypes = op->getResultTypes();
for (auto resType : resultTypes) {
Expand Down Expand Up @@ -82,7 +85,7 @@ struct LinearizeConstantLike final

LinearizeConstantLike(
const TypeConverter &typeConverter, MLIRContext *context,
unsigned targetVectBitWidth = std::numeric_limits<unsigned>::max(),
unsigned targetVectBitWidth = defaultTargetVectorBitWidth,
Copy link
Contributor

Choose a reason for hiding this comment

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

Is this needed for the vector.splat pattern?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

no, this will go away once I rebase after #136581 & #138072

PatternBenefit benefit = 1)
: OpTraitConversionPattern(typeConverter, context, benefit),
targetVectorBitWidth(targetVectBitWidth) {}
Expand Down Expand Up @@ -136,7 +139,7 @@ struct LinearizeVectorizable final
public:
LinearizeVectorizable(
const TypeConverter &typeConverter, MLIRContext *context,
unsigned targetVectBitWidth = std::numeric_limits<unsigned>::max(),
unsigned targetVectBitWidth = defaultTargetVectorBitWidth,
PatternBenefit benefit = 1)
: OpTraitConversionPattern(typeConverter, context, benefit),
targetVectorBitWidth(targetVectBitWidth) {}
Expand Down Expand Up @@ -175,7 +178,7 @@ struct LinearizeVectorExtractStridedSlice final
using OpConversionPattern::OpConversionPattern;
LinearizeVectorExtractStridedSlice(
const TypeConverter &typeConverter, MLIRContext *context,
unsigned targetVectBitWidth = std::numeric_limits<unsigned>::max(),
unsigned targetVectBitWidth = defaultTargetVectorBitWidth,
PatternBenefit benefit = 1)
: OpConversionPattern(typeConverter, context, benefit),
targetVectorBitWidth(targetVectBitWidth) {}
Expand Down Expand Up @@ -289,7 +292,7 @@ struct LinearizeVectorShuffle final
using OpConversionPattern::OpConversionPattern;
LinearizeVectorShuffle(
const TypeConverter &typeConverter, MLIRContext *context,
unsigned targetVectBitWidth = std::numeric_limits<unsigned>::max(),
unsigned targetVectBitWidth = defaultTargetVectorBitWidth,
PatternBenefit benefit = 1)
: OpConversionPattern(typeConverter, context, benefit),
targetVectorBitWidth(targetVectBitWidth) {}
Expand Down Expand Up @@ -362,13 +365,17 @@ struct LinearizeVectorExtract final
using OpConversionPattern::OpConversionPattern;
LinearizeVectorExtract(
const TypeConverter &typeConverter, MLIRContext *context,
unsigned targetVectBitWidth = std::numeric_limits<unsigned>::max(),
unsigned targetVectBitWidth = defaultTargetVectorBitWidth,
PatternBenefit benefit = 1)
: OpConversionPattern(typeConverter, context, benefit),
targetVectorBitWidth(targetVectBitWidth) {}
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());
if (!dstTy)
return rewriter.notifyMatchFailure(extractOp,
Expand Down Expand Up @@ -425,7 +432,7 @@ struct LinearizeVectorInsert final
using OpConversionPattern::OpConversionPattern;
LinearizeVectorInsert(
const TypeConverter &typeConverter, MLIRContext *context,
unsigned targetVectBitWidth = std::numeric_limits<unsigned>::max(),
unsigned targetVectBitWidth = defaultTargetVectorBitWidth,
PatternBenefit benefit = 1)
: OpConversionPattern(typeConverter, context, benefit),
targetVectorBitWidth(targetVectBitWidth) {}
Expand Down Expand Up @@ -506,7 +513,7 @@ struct LinearizeVectorBitCast final
using OpConversionPattern::OpConversionPattern;
LinearizeVectorBitCast(
const TypeConverter &typeConverter, MLIRContext *context,
unsigned targetVectBitWidth = std::numeric_limits<unsigned>::max(),
unsigned targetVectBitWidth = defaultTargetVectorBitWidth,
PatternBenefit benefit = 1)
: OpConversionPattern(typeConverter, context, benefit),
targetVectorBitWidth(targetVectBitWidth) {}
Expand All @@ -531,12 +538,48 @@ struct LinearizeVectorBitCast final
unsigned targetVectorBitWidth;
};

/// 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>
/// It ensures that the operation is compatible with the target vector
/// bit width and replaces the original operation with a new SplatOp
/// that operates on the converted type.
struct LinearizeVectorSplat final
: public OpConversionPattern<vector::SplatOp> {
using OpConversionPattern::OpConversionPattern;

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

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();
}

private:
unsigned targetVectorBitWidth;
};

} // namespace

void mlir::vector::populateVectorLinearizeTypeConversionsAndLegality(
TypeConverter &typeConverter, RewritePatternSet &patterns,
ConversionTarget &target, unsigned targetBitWidth) {

typeConverter.addConversion([](Type type) -> Type { return type; });
typeConverter.addConversion([](VectorType type) -> std::optional<Type> {
if (!isLinearizableVector(type))
return type;
Expand All @@ -557,7 +600,7 @@ void mlir::vector::populateVectorLinearizeTypeConversionsAndLegality(
typeConverter.addTargetMaterialization(materializeCast);
target.markUnknownOpDynamicallyLegal(
[=](Operation *op) -> std::optional<bool> {
if ((isa<vector::BitCastOp>(op) ||
if ((isa<vector::BitCastOp, vector::SplatOp>(op) ||
op->hasTrait<OpTrait::ConstantLike>() ||
op->hasTrait<OpTrait::Vectorizable>())) {
return (isLessThanTargetBitWidth(op, targetBitWidth)
Expand All @@ -568,8 +611,8 @@ void mlir::vector::populateVectorLinearizeTypeConversionsAndLegality(
});

patterns.add<LinearizeConstantLike, LinearizeVectorizable,
LinearizeVectorBitCast>(typeConverter, patterns.getContext(),
targetBitWidth);
LinearizeVectorBitCast, LinearizeVectorSplat>(
typeConverter, patterns.getContext(), targetBitWidth);
}

void mlir::vector::populateVectorLinearizeShuffleLikeOpsPatterns(
Expand Down
17 changes: 17 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 @@ -399,3 +399,20 @@ 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>
}