Respond to review

Change-Id: I0569534397a2d0db9587db6424b1674846a76079

Author: ekalda

src/arith/analyzer.cc

@@ -235,17 +235,14 @@ bool Analyzer::CanProve(const PrimExpr& expr, ProofStrength strength) {
   // SVE, we can make some assumptions about the value of vscale and iterate over a
   // space of pre-defined values to attempt to prove the expression.
   if (tir::CheckContains::ExprContains(expr, IsVScaleCall)) {
-    Target curr_target = tvm::Target::Current();
-    if (curr_target.defined() && curr_target->features.defined() &&
-        (curr_target->features.find("has_sve") != curr_target->features.end()) &&
-        curr_target->GetFeature<Bool>("has_sve").value_or(Bool(false)).operator bool()) {
+    if (TargetHasSVE()) {
       return CanProveVscaleExpressionFromKnownValues(this, simplified, kAArch64VScaleValues);
     }
     LOG(WARNING)
         << "The expression contains scalable values. An attempt to prove by substituting "
            "with known values of vscale was not performed. This proof currently only supports "
            "AArch64 SVE targets, but the target was "
-        << curr_target;
+        << Target::Current();
   }
   return false;
 }

src/arith/scalable_expression.cc

@@ -88,5 +88,14 @@ bool CanProveVscaleExpressionFromKnownValues(arith::Analyzer* analyzer, const Pr
   return can_prove_expr;
 }
 
+bool TargetHasSVE() {
+  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));
+  }
+  return has_sve;
+}
+
 }  // namespace arith
 }  // namespace tvm

src/arith/scalable_expression.h

@@ -71,6 +71,12 @@ std::optional<int> ExtractVscaleFactor(const PrimExpr& lanes);
 bool CanProveVscaleExpressionFromKnownValues(arith::Analyzer* analyzer, const PrimExpr& expr,
                                              const std::vector<unsigned int>& vscale_values);
 
+/*!
+ * \brief Check whether the compilation target supports SVE
+ * \return Whether SVE is supported
+ */
+bool TargetHasSVE();
+
 }  // namespace arith
 }  // namespace tvm
 

src/tir/transforms/vectorize_loop.cc

@@ -731,15 +731,12 @@ class LoopVectorizer : public StmtMutator {
   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));
-        }
         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 " << current_target;
+        ICHECK(is_scalable_expr && arith::TargetHasSVE())
+            << "Failed to vectorize loop with extent " << op->extent << " for target "
+            << Target::Current();
       }
       ICHECK(is_zero(op->min));
       return Vectorizer(op->loop_var, op->extent)(op->body);

tests/python/tir-transform/test_tir_transform_vectorize.py

@@ -114,7 +114,6 @@ def test_vectorize_vector_scalable_error4():
     class Module:
         @T.prim_func(private=True)
         def main(A: T.Buffer((25,), "float32")):
-            T.func_attr({"target": sve_target})
             for j in T.vectorized(T.vscale() * 4):
                 A[j * T.vscale() * 4 : j * T.vscale() * 4 + T.vscale() * 4] = T.Broadcast(
                     T.float32(1), T.vscale() * 4