[LLVM][Codegen] Enable SVE for RISCV targets

Author: cbalint13

src/arith/analyzer.cc

@@ -236,12 +236,13 @@ bool Analyzer::CanProve(const PrimExpr& expr, ProofStrength strength) {
   Target curr_target = Target::Current();
   if (ContainsVscaleCall(simplified)) {
     if (TargetHasSVE(curr_target)) {
-      return CanProveVscaleExpressionFromKnownValues(this, simplified, kAArch64VScaleValues);
+      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 "
+           "SVE 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()) && TargetHasSVE(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

@@ -90,10 +90,41 @@ bool TargetHasSVE(Optional<Target> target) {
   if (!target.defined()) {
     target = Target::Current();
   }
+  bool has_sve{false};
   if (target.defined()) {
-    return Downcast<Target>(target)->GetFeature<Bool>("has_sve").value_or(Bool(false));
+    // aarch64
+    has_sve = 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_sve |= static_cast<bool>((*target_has_feature_fn_ptr)("v", target));
   }
-  return false;
+  return has_sve;
+}
+
+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
@@ -85,6 +82,13 @@ bool CanProveVscaleExpressionFromKnownValues(arith::Analyzer* analyzer, const Pr
  */
 bool TargetHasSVE(Optional<Target> target = NullOpt);
 
+/*!
+ * \brief Get a list of known vscale values to try for an SVE target.
+ * \param target The target to check.
+ * \return A list of vscale values as std::vector<usigned int>
+ */
+const std::vector<unsigned int> GetVScaleValues(Optional<Target> target = 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/target/llvm/llvm_instance.cc

@@ -53,18 +53,6 @@
 #include <llvm/Support/raw_ostream.h>
 #include <llvm/Target/TargetMachine.h>
 #include <llvm/Target/TargetOptions.h>
-#if TVM_LLVM_VERSION >= 190
-#include <llvm/TargetParser/RISCVISAInfo.h>
-#else
-#if TVM_LLVM_VERSION >= 140
-#include <llvm/Support/RISCVISAInfo.h>
-#endif
-#endif
-#if TVM_LLVM_VERSION >= 160
-#include <llvm/TargetParser/RISCVTargetParser.h>
-#else
-#include <llvm/Support/TargetParser.h>
-#endif
 #include <tvm/runtime/container/array.h>
 #include <tvm/runtime/container/map.h>
 #include <tvm/runtime/container/optional.h>
@@ -299,34 +287,25 @@ LLVMTargetInfo::LLVMTargetInfo(LLVMInstance& instance, const TargetJSON& target)
     // code model
     code_model_ = llvm::CodeModel::Medium;
 #if TVM_LLVM_VERSION >= 140
-    // VLEN inference
-    const auto cpu_name = GetOrCreateTargetMachine(false)->getMCSubtargetInfo()->getCPU();
-    const auto canon_arch = llvm::RISCV::getMArchFromMcpu(cpu_name);
-    auto ISAInfo =
-        llvm::RISCVISAInfo::parseArchString(canon_arch, /*EnableExperimentalExtensions=*/true);
-    // infer VLEN from LLVM RISCVInfo parser
-    if (!llvm::errorToBool(ISAInfo.takeError()) && (vector_width_ == 0)) {
-      vector_width_ = (*ISAInfo)->getMinVLen();
-    }
-    // infer VLEN from LLVM options (zvlXXXb override)
-    for (const auto& attr : attrs_) {
-      if (attr.find("zvl") != std::string::npos) {
-        std::string vec;
-        for (char c : attr) {
-          if (std::isdigit(c)) vec += c;
+    // get VLEN from the LLVM backend (zvlXXXb)
+    Map<String, String> features = GetAllLLVMCpuFeatures();
+    // check vector ISA
+    if (features.count("v") > 0) {
+      vector_width_ = 0;
+      int zvlbits = 0;
+      for (const auto& [attr, val] : features) {
+        if (std::string(attr).find("zvl") != std::string::npos) {
+          std::string vec;
+          for (char c : std::string(attr)) {
+            if (std::isdigit(c)) vec += c;
+          }
+          zvlbits = std::stoi(vec);
+          // max of the multiple zvlXXXb
+          if (vector_width_ < zvlbits) vector_width_ = zvlbits;
         }
-        vector_width_ = std::stoi(vec);
       }
     }
 #endif
-    if (vector_width_ > 0) {
-      // push cl-opt to LLVM
-      llvm_options_.push_back(
-          ParseOptionString("-riscv-v-vector-bits-min:int=" + std::to_string(vector_width_)));
-    } else {
-      // fallback default (codegen will warn)
-      llvm_options_.push_back(ParseOptionString("-riscv-v-vector-bits-min:int=256"));
-    }
   }
 
   // Target options
@@ -943,9 +922,7 @@ const int LLVMTargetInfo::GetVectorWidth() {
     } else if (arch == llvm::Triple::arm || arch == llvm::Triple::aarch64) {
       vector_width_ = 128;
     } else if (arch == llvm::Triple::riscv32 || arch == llvm::Triple::riscv64) {
-      vector_width_ = 256;
-      LOG(WARNING) << "LLVM RVV VLEN inference failed, "
-                   << "using 256 bits, set -vector-width=XXX to override";
+      vector_width_ = 128;
     } else {
       // fallback default
       vector_width_ = 128;

src/tir/transforms/vectorize_loop.cc

@@ -80,7 +80,7 @@ bool EnableBufferLevelPredication(Target target) {
     return enable_buffer_predication.value();
   }
 
-  // Use buffer-level predication by default for AArch64 SVE targets
+  // Use buffer-level predication by default for SVE targets
   return arith::TargetHasSVE(target);
 }