ekalda commented on code in PR #16782:
URL: https://github.com/apache/tvm/pull/16782#discussion_r1551562489
##########
src/tir/transforms/vectorize_loop.cc:
##########
@@ -433,20 +488,27 @@ class Vectorizer : public StmtMutator, public
ExprFunctor<PrimExpr(const PrimExp
PrimExpr value = this->VisitExpr(op->value);
if (!indices.same_as(op->indices) || !value.same_as(op->value)) {
+ ICHECK(!op->buffer->dtype.is_scalable_vector())
+ << "Vectorizing over scalable buffer elements is not supported in
vectorizer.";
// How many lanes of indexing are present in the index and
- // buffer element type, excluding the last index. T
+ // buffer element type, excluding the last index.
int other_index_lanes = op->buffer->dtype.lanes();
for (size_t i = 0; i < indices.size() - 1; i++) {
other_index_lanes *= indices[i].dtype().lanes();
+ // Only allow the last index to be scalable
+ ICHECK(!indices[i].dtype().is_scalable_vector()) << "Only the last
index can be scalable.";
}
// The total number of lanes of indexing, including the last index.
- int index_lanes = other_index_lanes * indices[indices.size() -
1].dtype().lanes();
+ int lanes_in_last_index = indices[indices.size() -
1].dtype().get_lanes_or_vscale_factor();
+ int index_lanes = other_index_lanes * lanes_in_last_index;
// The total number of lanes in this store operation. Either
// the index or the value will be broadcast out to this number
// of lanes, depending on which has more lanes.
- int total_lanes = std::max(index_lanes, value.dtype().lanes());
+ int value_dtype_lanes = value.dtype().get_lanes_or_vscale_factor();
+ bool is_last_index_scalable = indices[indices.size() -
1].dtype().is_scalable_vector();
Review Comment:
Done
##########
src/tir/transforms/vectorize_loop.cc:
##########
@@ -635,19 +701,22 @@ class Vectorizer : public StmtMutator, public
ExprFunctor<PrimExpr(const PrimExp
if (a.same_as(op->a) && b.same_as(op->b)) {
return GetRef<PrimExpr>(op);
} else {
- int lanes = std::max(a.dtype().lanes(), b.dtype().lanes());
+ int a_lanes = a.dtype().get_lanes_or_vscale_factor();
+ int b_lanes = b.dtype().get_lanes_or_vscale_factor();
+ int lanes = std::max(a_lanes, b_lanes);
if (lanes != 1) {
const RampNode* b_ramp = b.as<RampNode>();
const RampNode* a_ramp = a.as<RampNode>();
- if (a.dtype().lanes() == 1 && b_ramp) {
+ if (!a.dtype().is_scalable_or_fixed_length_vector() && b_ramp) {
Review Comment:
Done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscr...@tvm.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
