[LLVM][Codegen] Enable SVE/VLA for RISCV targets

Author: cbalint13

src/arith/analyzer.cc

@@ -231,17 +231,18 @@ bool Analyzer::CanProve(const PrimExpr& expr, ProofStrength strength) {
   // Current analysis may not be powerful enough to prove expressions containing
   // the same symbolic value multiple times. However, when the symbolic values are
   // "T.vscale" and the compile target uses a scalable architecture extension like
-  // SVE, we can make some assumptions about the value of vscale and iterate over a
+  // VLA, 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.
   Target curr_target = Target::Current();
   if (ContainsVscaleCall(simplified)) {
-    if (TargetHasSVE(curr_target)) {
-      return CanProveVscaleExpressionFromKnownValues(this, simplified, kAArch64VScaleValues);
+    if (TargetHasVLA(curr_target)) {
+      auto kVScaleValues = GetVScaleValues(curr_target);
+      return CanProveVscaleExpressionFromKnownValues(this, simplified, kVScaleValues);
     }
     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 "
+           "VLA targets, but the target was "
         << curr_target;
   }
   return false;

src/arith/const_int_bound.cc

@@ -364,15 +364,16 @@ class ConstIntBoundAnalyzer::Impl
     // only special handle >> and & which can be
     // used for index calculation.
 
+    auto curr_target = Target::Current();
     if (op->op.same_as(tir::builtin::shift_right())) {
       return VisitRightShift(op);
     } else if (op->op.same_as(tir::builtin::shift_left())) {
       return VisitLeftShift(op);
     } else if (op->op.same_as(tir::builtin::bitwise_and())) {
       return VisitBitwiseAnd(op);
-    } else if (op->op.same_as(tir::builtin::vscale()) && TargetHasSVE(Target::Current())) {
-      unsigned int max_val =
-          *std::max_element(kAArch64VScaleValues.begin(), kAArch64VScaleValues.end());
+    } else if (op->op.same_as(tir::builtin::vscale()) && TargetHasVLA(curr_target)) {
+      auto kVScaleValues = GetVScaleValues(curr_target);
+      unsigned int max_val = *std::max_element(kVScaleValues.begin(), kVScaleValues.end());
       return MakeBound(1, max_val);
     } else {
       return Everything(op->dtype);

src/arith/scalable_expression.cc

@@ -86,14 +86,45 @@ bool CanProveVscaleExpressionFromKnownValues(arith::Analyzer* analyzer, const Pr
   return can_prove_expr;
 }
 
-bool TargetHasSVE(Optional<Target> target) {
+bool TargetHasVLA(Optional<Target> target) {
   if (!target.defined()) {
     target = Target::Current();
   }
+  bool has_vla{false};
   if (target.defined()) {
-    return Downcast<Target>(target)->GetFeature<Bool>("has_sve").value_or(Bool(false));
+    // aarch64
+    has_vla = Downcast<Target>(target)->GetFeature<Bool>("has_sve").value_or(Bool(false));
+    // riscv{32,64}
+    static const PackedFunc* target_has_feature_fn_ptr =
+        runtime::Registry::Get("target.target_has_feature");
+    ICHECK(target_has_feature_fn_ptr != nullptr)
+        << "The `target.target_has_feature` func is not in tvm registry.";
+    has_vla |= static_cast<bool>((*target_has_feature_fn_ptr)("v", target));
   }
-  return false;
+  return has_vla;
+}
+
+const std::vector<unsigned int> GetVScaleValues(Optional<Target> target) {
+  unsigned int vector_width = 0;
+  std::vector<unsigned int> kVScaleValues;
+  if (!target.defined()) {
+    target = Target::Current();
+  }
+  if (target.defined()) {
+    static const PackedFunc* llvm_get_vector_width_fn_ptr =
+        runtime::Registry::Get("target.llvm_get_vector_width");
+    ICHECK(llvm_get_vector_width_fn_ptr != nullptr)
+        << "The `target.llvm_get_vector_width` func is not in tvm registry.";
+    vector_width = static_cast<int>((*llvm_get_vector_width_fn_ptr)(target));
+  }
+  // scale list with powers of two
+  for (unsigned int i = 0;; ++i) {
+    auto power = static_cast<unsigned int>(std::pow(2, i));
+    if (power > (vector_width / 8)) break;
+    kVScaleValues.push_back(power);
+  }
+
+  return kVScaleValues;
 }
 
 }  // namespace arith

src/arith/scalable_expression.h

@@ -35,9 +35,6 @@
 namespace tvm {
 namespace arith {
 
-/*! \brief A list of known vscale values to try for an AArch64 SVE target. */
-static const std::vector<unsigned int> kAArch64VScaleValues = {1, 2, 4, 8, 16};
-
 /*!
  * \brief Check if an expr is a call to the vscale intrinsic.
  * \param expr The expr to check
@@ -80,10 +77,18 @@ bool CanProveVscaleExpressionFromKnownValues(arith::Analyzer* analyzer, const Pr
 
 /*!
  * \brief Check whether the compilation target supports SVE
+ * \brief Check whether the compilation target supports VLA
+ * \param target The target to check.
+ * \return Whether VLA is supported
+ */
+bool TargetHasVLA(Optional<Target> target = NullOpt);
+
+/*!
+ * \brief Get a list of known vscale values to try for an VLA target.
  * \param target The target to check.
- * \return Whether SVE is supported
+ * \return A list of vscale values as std::vector<usigned int>
  */
-bool TargetHasSVE(Optional<Target> target = std::nullopt);
+const std::vector<unsigned int> GetVScaleValues(Optional<Target> target = std::nullopt);
 
 }  // namespace arith
 }  // namespace tvm

src/target/llvm/codegen_aarch64.cc

@@ -57,8 +57,8 @@ void CodeGenAArch64::SetTargetAttributes(llvm::Function* func) {
 #if TVM_LLVM_VERSION >= 130
   // Add vscale_range() function attribute when appropriate.
   if (llvm_target_->TargetHasCPUFeature("sve") || llvm_target_->TargetHasCPUFeature("sme")) {
-    unsigned int max_val =
-        *std::max_element(arith::kAArch64VScaleValues.begin(), arith::kAArch64VScaleValues.end());
+    auto kVScaleValues = arith::GetVScaleValues(Target::Current());
+    unsigned int max_val = *std::max_element(kVScaleValues.begin(), kVScaleValues.end());
     func->addFnAttr(
         llvm::Attribute::getWithVScaleRangeArgs(*llvm_target_->GetContext(), 1, max_val));
   }

src/tir/transforms/vectorize_loop.cc

@@ -80,8 +80,8 @@ bool EnableBufferLevelPredication(Target target) {
     return enable_buffer_predication.value();
   }
 
-  // Use buffer-level predication by default for AArch64 SVE targets
-  return arith::TargetHasSVE(target);
+  // Use buffer-level predication by default for VLA targets
+  return arith::TargetHasVLA(target);
 }
 
 /*!
@@ -972,7 +972,7 @@ class LoopVectorizer : public StmtMutator {
 
       if (!extent_as_int || extent_as_int->value < 1) {
         bool is_scalable_expr = CheckContains::ExprContains(op->extent, arith::IsVScaleCall);
-        ICHECK(is_scalable_expr && arith::TargetHasSVE(target_))
+        ICHECK(is_scalable_expr && arith::TargetHasVLA(target_))
             << "Failed to vectorize loop with extent " << op->extent << " for target " << target_;
       }
       ICHECK(is_zero(op->min));

tests/python/arith/test_arith_simplify.py

@@ -113,7 +113,7 @@ def test_simplify_vscale_comparison_without_sve_target(capfd):
     warning_msg = (
         "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 llvm -keys=arm_cpu,cpu -mtriple=aarch64-linux-gnu"
+        "VLA targets, but the target was llvm -keys=arm_cpu,cpu -mtriple=aarch64-linux-gnu"
     )
     capture = capfd.readouterr().err
     assert warning_msg in capture