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 5 commits into
base: main
Choose a base branch
from
+174 −15
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
5 commits
Select commit Hold shift + click to select a range
c511a4e
add vector.insert canonicalization pattern for vectors created from u…
Jun 5, 2025
bba3d6c
refine comments & add hasOneUse check
Jun 25, 2025
9701de8
use from_elements to replace arith.constant
Jun 26, 2025
27bb005
update doc
Jun 26, 2025
b3e9850
fix typo and update test
Jun 26, 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
145 changes: 138 additions & 7 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,19 @@ LogicalResult InsertOp::verify() {
return success();
}

// Calculate the linearized position of the continuous chunk of elements to
// insert, based on the shape of the value to insert and the positions to insert
// at.
static int64_t calculateInsertPosition(VectorType destTy,
ArrayRef<int64_t> positions) {
llvm::SmallVector<int64_t> completePositions(destTy.getRank(), 0);
copy(positions, completePositions.begin());
Copy link
Member

Choose a reason for hiding this comment

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

This should probably assert we are not copying past the end of the buffer.

int64_t insertBeginPosition =
linearize(completePositions, computeStrides(destTy.getShape()));

return insertBeginPosition;
Comment on lines +3159 to +3162
Copy link
Member

Choose a reason for hiding this comment

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

No need for a local variable IMO

}

namespace {

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

/// Pattern to optimize a chain of insertions into a poison vector.
///
/// This pattern identifies chains of vector.insert operations that:
/// 1. Start from an ub.poison operation.
/// 2. Only insert values at static positions.
/// 3. Completely initialize all elements in the resulting vector.
Copy link
Member

Choose a reason for hiding this comment

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

If all elements are initialized, why does it matter if the operation started from a ub.poison?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

If the vector didn't start from ub.poison, the existing foldDenseElementsAttrDestInsertOp in InsertOp::fold should already be able to handle the folding. However, it does not guarantee all elements are initialized. That’s why I want to add this new pattern explicitly targets ub.poison.

Copy link
Member

Choose a reason for hiding this comment

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

+1, this feels like an awkward special case. Why can't the existing folder handle this case instead?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

The reason I chose not to use the folder is that this method requires traversing the entire chain of insert operations, which results in linear time complexity for each traversal. If someone uses createOrFold to create a sequence of insert operations, the access pattern would look like this:

insert op 1 -> ub.poison  
insert op 2 -> insert op 1 -> ub.poison  
insert op 3 -> insert op 2 -> insert op 1 -> ub.poison  
...
insert op n -> insert op n - 1 -> ... -> inset op 1 -> ub.poison

The condition for being fully initialized is only satisfied at the very last line, but the overall complexity becomes O(n²) due to repeated traversals.

However, I also admit that this is an awkward case. I’m completely okay if you think it’s not worth introducing a new canonicalization pattern. Perhaps it would be much easier to address this once partial poisoned vectors can be easily represented.

Copy link
Member

Choose a reason for hiding this comment

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

It shouldn't matter at all if the start is a ub.poison. If you think that the pattern isn't a canonicalization if the start is something else, then it shouldn't be a canonicalization at all. Personally, this seems like a cleanup and we could live with this in cleanup before lowering to backends.

But the initial start shouldn't matter at all whatever the case.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

It shouldn't matter at all if the start is a ub.poison.

After a second thought, I think it indeed makes sense. We don’t need to worry about whether a canonicalization pattern may overlap with a folder in some cases. I'll remove the check for ub.poison.

/// 4. All intermediate insert operations have only one use.
///
/// When these conditions are met, the entire chain can be replaced with a
/// single vector.from_elements operation.
///
/// 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 = vector.from_elements %c1, %c2 : 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 InsertToPoison 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;
while (previousInsertOp) {
// Dynamic position is not supported.
if (previousInsertOp.hasDynamicPosition())
return failure();

chainInsertOps.push_back(previousInsertOp);

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

// Check that intermediate inserts have only one use to avoid an explosion
// of vectors.
if (previousInsertOp && !previousInsertOp->hasOneUse())
return failure();
}

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

// Currently, MLIR doesn't support partial poison vectors, so we can only
// optimize when the entire vector is completely initialized.
Comment on lines +3264 to +3265
Copy link
Member

Choose a reason for hiding this comment

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

Not exactly, you can always do a shufflevector instruction to get these things in the right order, but we can ignore this for now.

Copy link
Contributor

Choose a reason for hiding this comment

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

My understanding is that this is about representing a constant vector with poison and non-poison values. AFAIK, we can't represent that in MLIR right now

int64_t vectorSize = destTy.getNumElements();
int64_t initializedCount = 0;
SmallVector<bool> initialized(vectorSize, false);
SmallVector<Value> elements(vectorSize);

for (auto insertOp : chainInsertOps) {
// The insert op folder will fold an insert at poison index into a
// ub.poison, which truncates the insert chain's backward traversal.
if (is_contained(insertOp.getStaticPosition(), InsertOp::kPoisonIndex))
return failure();

// Calculate the linearized position for inserting elements.
int64_t insertBeginPosition =
calculateInsertPosition(destTy, insertOp.getStaticPosition());

// The valueToStore operand may be a vector or a scalar. Need to handle
// both cases.
SmallVector<Value> elementsToInsert;
int64_t elementsToInsertSize = 1;
if (auto srcVectorType =
llvm::dyn_cast<VectorType>(insertOp.getValueToStoreType())) {

elementsToInsertSize = srcVectorType.getNumElements();
elementsToInsert.reserve(elementsToInsertSize);
SmallVector<int64_t> strides = computeStrides(srcVectorType.getShape());
// Get all elements from the vector in row-major order.
for (int64_t linearIdx = 0; linearIdx < elementsToInsertSize;
linearIdx++) {
SmallVector<int64_t> position = delinearize(linearIdx, strides);
Value extractedElement = rewriter.create<vector::ExtractOp>(
insertOp.getLoc(), insertOp.getValueToStore(), position);
elementsToInsert.push_back(extractedElement);
}
} else {
elementsToInsert.push_back(insertOp.getValueToStore());
}

for (auto index :
llvm::seq<int64_t>(insertBeginPosition,
insertBeginPosition + elementsToInsertSize)) {
if (initialized[index])
continue;

initialized[index] = true;
++initializedCount;
elements[index] = elementsToInsert[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();

rewriter.replaceOpWithNewOp<vector::FromElementsOp>(op, destTy, elements);
return success();
}
};

} // namespace

static Attribute
Expand All @@ -3217,13 +3351,9 @@ 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());
// Calculate the linearized position for inserting elements.
int64_t insertBeginPosition =
linearize(completePositions, computeStrides(destTy.getShape()));

calculateInsertPosition(destTy, insertOp.getStaticPosition());
SmallVector<Attribute> insertedValues;
Type destEltType = destTy.getElementType();

Expand Down Expand Up @@ -3256,7 +3386,8 @@ foldDenseElementsAttrDestInsertOp(InsertOp insertOp, Attribute srcAttr,

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

OpFoldResult vector::InsertOp::fold(FoldAdaptor adaptor) {
Expand Down
12 changes: 4 additions & 8 deletions mlir/test/Conversion/ConvertToSPIRV/vector-unroll.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,20 +83,16 @@ func.func @vaddi_reduction(%arg0 : vector<8xi32>, %arg1 : vector<8xi32>) -> (i32
// CHECK-LABEL: @transpose
// CHECK-SAME: (%[[ARG0:.+]]: vector<3xi32>, %[[ARG1:.+]]: vector<3xi32>)
func.func @transpose(%arg0 : vector<2x3xi32>) -> (vector<3x2xi32>) {
// CHECK: %[[UB:.*]] = ub.poison : vector<2xi32>
// CHECK: %[[EXTRACT0:.*]] = vector.extract %[[ARG0]][0] : i32 from vector<3xi32>
// CHECK: %[[INSERT0:.*]]= vector.insert %[[EXTRACT0]], %[[UB]] [0] : i32 into vector<2xi32>
// CHECK: %[[EXTRACT1:.*]] = vector.extract %[[ARG1]][0] : i32 from vector<3xi32>
// CHECK: %[[INSERT1:.*]] = vector.insert %[[EXTRACT1]], %[[INSERT0]][1] : i32 into vector<2xi32>
// CHECK: %[[FROM_ELEMENTS0:.*]] = vector.from_elements %[[EXTRACT0]], %[[EXTRACT1]] : vector<2xi32>
// CHECK: %[[EXTRACT2:.*]] = vector.extract %[[ARG0]][1] : i32 from vector<3xi32>
// CHECK: %[[INSERT2:.*]] = vector.insert %[[EXTRACT2]], %[[UB]] [0] : i32 into vector<2xi32>
// CHECK: %[[EXTRACT3:.*]] = vector.extract %[[ARG1]][1] : i32 from vector<3xi32>
// CHECK: %[[INSERT3:.*]] = vector.insert %[[EXTRACT3]], %[[INSERT2]] [1] : i32 into vector<2xi32>
// CHECK: %[[FROM_ELEMENTS1:.*]] = vector.from_elements %[[EXTRACT2]], %[[EXTRACT3]] : vector<2xi32>
// CHECK: %[[EXTRACT4:.*]] = vector.extract %[[ARG0]][2] : i32 from vector<3xi32>
// CHECK: %[[INSERT4:.*]] = vector.insert %[[EXTRACT4]], %[[UB]] [0] : i32 into vector<2xi32>
// CHECK: %[[EXTRACT5:.*]] = vector.extract %[[ARG1]][2] : i32 from vector<3xi32>
// CHECK: %[[INSERT5:.*]] = vector.insert %[[EXTRACT5]], %[[INSERT4]] [1] : i32 into vector<2xi32>
// CHECK: return %[[INSERT1]], %[[INSERT3]], %[[INSERT5]] : vector<2xi32>, vector<2xi32>, vector<2xi32>
// CHECK: %[[FROM_ELEMENTS2:.*]] = vector.from_elements %[[EXTRACT4]], %[[EXTRACT5]] : vector<2xi32>
// CHECK: return %[[FROM_ELEMENTS0]], %[[FROM_ELEMENTS1]], %[[FROM_ELEMENTS2]] : vector<2xi32>, vector<2xi32>, vector<2xi32>
%0 = vector.transpose %arg0, [1, 0] : vector<2x3xi32> to vector<3x2xi32>
return %0 : vector<3x2xi32>
}
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
Loading