lhutton1 commented on code in PR #16893:
URL: https://github.com/apache/tvm/pull/16893#discussion_r1572052392
##########
src/tir/transforms/vectorize_loop.cc:
##########
@@ -725,17 +728,33 @@ class Vectorizer : public StmtMutator, public
ExprFunctor<PrimExpr(const PrimExp
class LoopVectorizer : public StmtMutator {
public:
+ LoopVectorizer(PrimFunc f) {
+ auto target = f->attrs.GetAttr<tvm::Target>(tvm::attr::kTarget);
+ if (target.defined()) {
+ target_ = Downcast<Target>(target);
+ has_sve_ = target_->GetFeature<Bool>("has_sve").value_or(Bool(false));
+ }
+ }
+
Stmt VisitStmt_(const ForNode* op) final {
if (op->kind == ForKind::kVectorized) {
+ auto* extent_as_int = op->extent.as<IntImmNode>();
+ if (!extent_as_int || extent_as_int->value < 1) {
+ bool is_scalable_expr = CheckContains::ExprContains(op->extent,
arith::IsVScaleCall);
+ ICHECK(is_scalable_expr && has_sve_)
+ << "Failed to vectorize loop with extent " << op->extent << " for
target " << target_;
+ }
ICHECK(is_zero(op->min));
return Vectorizer(op->loop_var, op->extent)(op->body);
} else {
return StmtMutator::VisitStmt_(op);
}
}
-};
-Stmt VectorizeLoop(Stmt stmt) { return LoopVectorizer()(std::move(stmt)); }
Review Comment:
I don't see it being used in the codebase :)
##########
tests/python/tir-transform/test_tir_transform_vectorize.py:
##########
@@ -99,26 +105,29 @@ def main(A: T.Buffer((25,), "float32")):
error_msg = f"Vectorizing over existing scalable vectors is not supported."
with pytest.raises(tvm.error.InternalError, match=error_msg):
- tvm.tir.transform.VectorizeLoop()(Module)
+ with tvm.target.Target(sve_target):
+ tvm.tir.transform.VectorizeLoop()(Module)
def test_vectorize_vector_scalable_error4():
@I.ir_module
class Module:
@T.prim_func(private=True)
def main(A: T.Buffer((25,), "float32")):
+ T.func_attr({"target": sve_target})
Review Comment:
Is this still needed?
##########
src/tir/transforms/vectorize_loop.cc:
##########
@@ -727,6 +730,17 @@ class LoopVectorizer : public StmtMutator {
public:
Stmt VisitStmt_(const ForNode* op) final {
if (op->kind == ForKind::kVectorized) {
+ auto* extent_as_int = op->extent.as<IntImmNode>();
+ if (!extent_as_int || extent_as_int->value < 1) {
+ Target current_target = Target::Current();
+ bool has_sve{false};
+ if (current_target.defined()) {
+ has_sve =
current_target->GetFeature<Bool>("has_sve").value_or(Bool(false));
Review Comment:
I did something similar here:
https://github.com/apache/tvm/blob/d4056ca79571d4265a12beeedd1b1565953df936/src/arith/analyzer.cc#L241,
but your approach is much neater :). Perhaps we could create a utility
function to keep them consistent?
--
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
