-
Notifications
You must be signed in to change notification settings - Fork 14.3k
[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.
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
base: main
Are you sure you want to change the base?
[mlir][Vector] add vector.insert canonicalization pattern for vectors created from ub.poison #142944
Conversation
Thank you for submitting a Pull Request (PR) to the LLVM Project! This PR will be automatically labeled and the relevant teams will be notified. If you wish to, you can add reviewers by using the "Reviewers" section on this page. If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers. If you have further questions, they may be answered by the LLVM GitHub User Guide. You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums. |
@llvm/pr-subscribers-mlir-spirv @llvm/pr-subscribers-mlir Author: Yang Bai (yangtetris) ChangesDescriptionThis change introduces a new canonicalization pattern for the MLIR Vector dialect that optimizes chains of constant insertions into vectors initialized with Please be aware that the new pattern doesn't work for poison vectors where only some elements are set, as MLIR doesn't support partial poison vectors for now. New Pattern: InsertConstantToPoison
Refactored Helper Function
Example
It also works for multidimensional vectors.
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Minor comments. Otherwise, it LGTM!
return failure(); | ||
|
||
auto newAttr = DenseElementsAttr::get(destTy, initValues); | ||
rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, destTy, newAttr); |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why not use vector.from_elements here and let it canonicalize to arith.constant if all inputs are constant?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Great idea! That would indeed remove the constant-only restriction. Let me try implementing it. Thanks!
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I just learned that currently vector.from_elements does not have a folder to fold it into an arith.constant op. Do you happen to know if there is any ongoing PR implementing this? If not, I would be interested in creating one.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Feel free to create one, we should have one like that if it doesnt exist already.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Indeed, I think that's is a great idea and something that we would need anyways. Let's implement that first!
// 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. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Calculat the linearized position based on what? I cannot understand what this is trying to say.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sorry for the confusion. I'll updated the function’s documentation to clarify this. Since we adopted a new implementation using from_elements, I also refactored this function. It now focuses solely on calculating the linearized position of the continuous chunk of elements to insert within the destination vector.
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; | ||
}; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Is this safe?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actually, this code originates from #88314 which addressed issues related to mismatches between index types and i64. I think this background is enough to reassures us about the safety. But, it also reminded me that we might need a similar attribute conversion in the from_elements to constant PR. I'll verify later whether that pattern encounters similar issues.
/// 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. |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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
.
if (previousInsertOp.hasDynamicPosition()) | ||
return failure(); | ||
|
||
// The inserted content must be constant. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
if you use from_elements this is not required.
if (!srcAttrs.back()) | ||
return failure(); | ||
|
||
// An insertion at poison index makes the entire chain poisoned. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
What would happen in this case? Can you add a comment if this will get folded by the folder?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The insert operation’s folder will fold this case into a ub.poison, effectively truncating the insert chain’s backward traversal. If the vector is not fully initialized by that point, this pattern will definitively fail. I will add a comment to clarify this behavior in the code.
Update: I also moved this check to another place to give this pattern a better chance of succeeding.
// Currently, MLIR doesn't support partial poison vectors, so we can only | ||
// optimize when the entire vector is completely initialized. |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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
SmallVector<bool> initialized(vectorSize, false); | ||
SmallVector<Attribute> initValues(vectorSize); | ||
|
||
for (auto [insertOp, srcAttr] : llvm::zip(chainInsertOps, srcAttrs)) { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
use zip_equal if they are equal
/// 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. |
There was a problem hiding this comment.
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?
static int64_t calculateInsertPosition(VectorType destTy, | ||
ArrayRef<int64_t> positions) { | ||
llvm::SmallVector<int64_t> completePositions(destTy.getRank(), 0); | ||
copy(positions, completePositions.begin()); |
There was a problem hiding this comment.
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; |
There was a problem hiding this comment.
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
Description
This change introduces a new canonicalization pattern for the MLIR Vector dialect that optimizes chains of insertions into vectors created from
ub.poison
. The optimization identifies when a vector is completely initialized through a series of vector.insert operations and replaces the entire chain with a singlevector.from_elements
operation.Please be aware that the new pattern doesn't work for poison vectors where only some elements are set, as MLIR doesn't support partial poison vectors for now.
New Pattern: InsertConstantToPoison
Refactored Helper Function
calculateInsertPosition
fromfoldDenseElementsAttrDestInsertOp
to avoid code duplication.Example
It also works for multidimensional vectors.
The vector.extract and vector.from_elements ops can be further canonicalized by their own folders or canonicalization patterns.