Skip to content

[mlir][Vector] add vector.insert canonicalization pattern for vectors created from ub.poison #142944

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 1 commit into
base: main
Choose a base branch
from
Open

[mlir][Vector] add vector.insert canonicalization pattern for vectors created from ub.poison #142944

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
174 changes: 145 additions & 29 deletions mlir/lib/Dialect/Vector/IR/VectorOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3149,6 +3149,42 @@ LogicalResult InsertOp::verify() {
return success();
}

// Calculate the linearized position for inserting elements and extract values
// from the source attribute. Returns the starting position in the destination
// vector where elements should be inserted.
static int64_t calculateInsertPositionAndExtractValues(
VectorType destTy, const ArrayRef<int64_t> &positions, Attribute srcAttr,
SmallVector<Attribute> &valueToInsert) {
llvm::SmallVector<int64_t> completePositions(destTy.getRank(), 0);
copy(positions, completePositions.begin());
int64_t insertBeginPosition =
linearize(completePositions, computeStrides(destTy.getShape()));

Type destEltType = destTy.getElementType();

/// Converts the expected type to an IntegerAttr if there's
/// a mismatch.
auto convertIntegerAttr = [](Attribute attr, Type expectedType) -> Attribute {
if (auto intAttr = mlir::dyn_cast<IntegerAttr>(attr)) {
if (intAttr.getType() != expectedType)
return IntegerAttr::get(expectedType, intAttr.getInt());
}
return attr;
};

// The `convertIntegerAttr` method specifically handles the case
// for `llvm.mlir.constant` which can hold an attribute with a
// different type than the return type.
if (auto denseSource = llvm::dyn_cast<DenseElementsAttr>(srcAttr)) {
for (auto value : denseSource.getValues<Attribute>())
valueToInsert.push_back(convertIntegerAttr(value, destEltType));
} else {
valueToInsert.push_back(convertIntegerAttr(srcAttr, destEltType));
}

return insertBeginPosition;
}

namespace {

// If insertOp is only inserting unit dimensions it can be transformed to a
Expand Down Expand Up @@ -3191,6 +3227,109 @@ class InsertSplatToSplat final : public OpRewritePattern<InsertOp> {
}
};

// Pattern to optimize a chain of constant insertions into a poison vector.
//
// This pattern identifies chains of vector.insert operations that:
// 1. Start from an ub.poison operation.
// 2. Insert only constant values at static positions.
// 3. Completely initialize all elements in the resulting vector.
//
// When these conditions are met, the entire chain can be replaced with a
// single arith.constant operation containing a dense elements attribute.
//
// Example transformation:
// %poison = ub.poison : vector<2xi32>
// %0 = vector.insert %c1, %poison[0] : i32 into vector<2xi32>
// %1 = vector.insert %c2, %0[1] : i32 into vector<2xi32>
// ->
// %result = arith.constant dense<[1, 2]> : vector<2xi32>

// TODO: Support the case where only some elements of the poison vector are set.
// Currently, MLIR doesn't support partial poison vectors.

class InsertConstantToPoison final : public OpRewritePattern<InsertOp> {
public:
using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(InsertOp op,
PatternRewriter &rewriter) const override {

VectorType destTy = op.getDestVectorType();
if (destTy.isScalable())
return failure();
// Check if the result is used as the dest operand of another vector.insert
// Only care about the last op in a chain of insertions.
for (Operation *user : op.getResult().getUsers())
if (auto insertOp = dyn_cast<InsertOp>(user))
if (insertOp.getDest() == op.getResult())
return failure();

InsertOp firstInsertOp;
InsertOp previousInsertOp = op;
SmallVector<InsertOp> chainInsertOps;
SmallVector<Attribute> srcAttrs;
while (previousInsertOp) {
// Dynamic position is not supported.
if (previousInsertOp.hasDynamicPosition())
return failure();

// The inserted content must be constant.
chainInsertOps.push_back(previousInsertOp);
srcAttrs.push_back(Attribute());
matchPattern(previousInsertOp.getValueToStore(),
m_Constant(&srcAttrs.back()));
if (!srcAttrs.back())
return failure();

// An insertion at poison index makes the entire chain poisoned.
if (is_contained(previousInsertOp.getStaticPosition(),
InsertOp::kPoisonIndex))
return failure();

firstInsertOp = previousInsertOp;
previousInsertOp = previousInsertOp.getDest().getDefiningOp<InsertOp>();
}

if (!firstInsertOp.getDest().getDefiningOp<ub::PoisonOp>())
return failure();

// Need to make sure all elements are initialized.
int64_t vectorSize = destTy.getNumElements();
int64_t initializedCount = 0;
SmallVector<bool> initialized(vectorSize, false);
SmallVector<Attribute> initValues(vectorSize);

for (auto [insertOp, srcAttr] : llvm::zip(chainInsertOps, srcAttrs)) {
// Calculate the linearized position for inserting elements, as well as
// convert the source attribute to the proper type.
SmallVector<Attribute> valueToInsert;
int64_t insertBeginPosition = calculateInsertPositionAndExtractValues(
destTy, insertOp.getStaticPosition(), srcAttr, valueToInsert);
for (auto index :
llvm::seq<int64_t>(insertBeginPosition,
insertBeginPosition + valueToInsert.size())) {
if (initialized[index])
continue;

initialized[index] = true;
++initializedCount;
initValues[index] = valueToInsert[index - insertBeginPosition];
}
// If all elements in the vector have been initialized, we can stop
// processing the remaining insert operations in the chain.
if (initializedCount == vectorSize)
break;
}

// some positions are not initialized.
if (initializedCount != vectorSize)
return failure();

auto newAttr = DenseElementsAttr::get(destTy, initValues);
rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, destTy, newAttr);
return success();
}
};

} // namespace

static Attribute
Expand All @@ -3217,35 +3356,11 @@ foldDenseElementsAttrDestInsertOp(InsertOp insertOp, Attribute srcAttr,
!insertOp->hasOneUse())
return {};

// Calculate the linearized position of the continuous chunk of elements to
// insert.
llvm::SmallVector<int64_t> completePositions(destTy.getRank(), 0);
copy(insertOp.getStaticPosition(), completePositions.begin());
int64_t insertBeginPosition =
linearize(completePositions, computeStrides(destTy.getShape()));

// Calculate the linearized position for inserting elements, as well as
// convert the source attribute to the proper type.
SmallVector<Attribute> insertedValues;
Type destEltType = destTy.getElementType();

/// Converts the expected type to an IntegerAttr if there's
/// a mismatch.
auto convertIntegerAttr = [](Attribute attr, Type expectedType) -> Attribute {
if (auto intAttr = mlir::dyn_cast<IntegerAttr>(attr)) {
if (intAttr.getType() != expectedType)
return IntegerAttr::get(expectedType, intAttr.getInt());
}
return attr;
};

// The `convertIntegerAttr` method specifically handles the case
// for `llvm.mlir.constant` which can hold an attribute with a
// different type than the return type.
if (auto denseSource = llvm::dyn_cast<DenseElementsAttr>(srcAttr)) {
for (auto value : denseSource.getValues<Attribute>())
insertedValues.push_back(convertIntegerAttr(value, destEltType));
} else {
insertedValues.push_back(convertIntegerAttr(srcAttr, destEltType));
}
int64_t insertBeginPosition = calculateInsertPositionAndExtractValues(
destTy, insertOp.getStaticPosition(), srcAttr, insertedValues);

auto allValues = llvm::to_vector(denseDst.getValues<Attribute>());
copy(insertedValues, allValues.begin() + insertBeginPosition);
Expand All @@ -3256,7 +3371,8 @@ foldDenseElementsAttrDestInsertOp(InsertOp insertOp, Attribute srcAttr,

void InsertOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
results.add<InsertToBroadcast, BroadcastFolder, InsertSplatToSplat>(context);
results.add<InsertToBroadcast, BroadcastFolder, InsertSplatToSplat,
InsertConstantToPoison>(context);
}

OpFoldResult vector::InsertOp::fold(FoldAdaptor adaptor) {
Expand Down
32 changes: 32 additions & 0 deletions mlir/test/Dialect/Vector/canonicalize.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2320,6 +2320,38 @@ func.func @insert_2d_constant() -> (vector<2x3xi32>, vector<2x3xi32>, vector<2x3

// -----

// CHECK-LABEL: func.func @fully_insert_scalar_constant_to_poison_vector
// CHECK: %[[VAL0:.+]] = arith.constant dense<[10, 20]> : vector<2xi64>
// CHECK-NEXT: return %[[VAL0]]
func.func @fully_insert_scalar_constant_to_poison_vector() -> vector<2xi64> {
%poison = ub.poison : vector<2xi64>
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%e0 = arith.constant 10 : i64
%e1 = arith.constant 20 : i64
%v1 = vector.insert %e0, %poison[%c0] : i64 into vector<2xi64>
%v2 = vector.insert %e1, %v1[%c1] : i64 into vector<2xi64>
return %v2 : vector<2xi64>
}

// -----

// CHECK-LABEL: func.func @fully_insert_vector_constant_to_poison_vector
// CHECK: %[[VAL0:.+]] = arith.constant dense<{{\[\[1, 2, 3\], \[4, 5, 6\]\]}}> : vector<2x3xi64>
// CHECK-NEXT: return %[[VAL0]]
func.func @fully_insert_vector_constant_to_poison_vector() -> vector<2x3xi64> {
%poison = ub.poison : vector<2x3xi64>
%cv0 = arith.constant dense<[1, 2, 3]> : vector<3xi64>
%cv1 = arith.constant dense<[4, 5, 6]> : vector<3xi64>
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%v1 = vector.insert %cv0, %poison[%c0] : vector<3xi64> into vector<2x3xi64>
%v2 = vector.insert %cv1, %v1[%c1] : vector<3xi64> into vector<2x3xi64>
return %v2 : vector<2x3xi64>
}

// -----

// CHECK-LABEL: func.func @insert_2d_splat_constant
// CHECK-DAG: %[[ACST:.*]] = arith.constant dense<0> : vector<2x3xi32>
// CHECK-DAG: %[[BCST:.*]] = arith.constant dense<{{\[\[99, 0, 0\], \[0, 0, 0\]\]}}> : vector<2x3xi32>
Expand Down