Skip to content

Commit a5862a5

Browse files
authored
[SVE] Use only powers of two as possible vscale values (#17001)
When analyzing scalable expressions, the analyzer will iterate over a series of known vscale values in the range 1-16. However, we can tighten this range to only values that are a power of two, as stated in the [LLVM lang ref](https://llvm.org/docs/LangRef.html#llvm-vscale-intrinsic:~:text=This%20function%20attribute%20indicates%20vscale%20is%20a%20power%2Dof%2Dtwo%20within%20a%20specified%20range) and more generally the [reference manual](https://developer.arm.com/documentation/ddi0487/latest/). This comes from a discussion in #16921 (comment)
1 parent 2e56421 commit a5862a5

File tree

3 files changed

+8
-5
lines changed

3 files changed

+8
-5
lines changed

src/arith/const_int_bound.cc

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -371,7 +371,9 @@ class ConstIntBoundAnalyzer::Impl
371371
} else if (op->op.same_as(tir::builtin::bitwise_and())) {
372372
return VisitBitwiseAnd(op);
373373
} else if (op->op.same_as(tir::builtin::vscale()) && TargetHasSVE()) {
374-
return MakeBound(1, kAArch64VScaleValues.size());
374+
unsigned int max_val =
375+
*std::max_element(kAArch64VScaleValues.begin(), kAArch64VScaleValues.end());
376+
return MakeBound(1, max_val);
375377
} else {
376378
return Everything(op->dtype);
377379
}

src/arith/scalable_expression.h

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,8 +35,7 @@ namespace tvm {
3535
namespace arith {
3636

3737
/*! \brief A list of known vscale values to try for an AArch64 SVE target. */
38-
static const std::vector<unsigned int> kAArch64VScaleValues = {1, 2, 3, 4, 5, 6, 7, 8,
39-
9, 10, 11, 12, 13, 14, 15, 16};
38+
static const std::vector<unsigned int> kAArch64VScaleValues = {1, 2, 4, 8, 16};
4039

4140
/*!
4241
* \brief Check if an expr is a call to the vscale intrinsic.

src/target/llvm/codegen_aarch64.cc

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -57,8 +57,10 @@ void CodeGenAArch64::SetTargetAttributes(llvm::Function* func) {
5757
#if TVM_LLVM_VERSION >= 130
5858
// Add vscale_range() function attribute when appropriate.
5959
if (llvm_target_->TargetHasCPUFeature("sve") || llvm_target_->TargetHasCPUFeature("sme")) {
60-
func->addFnAttr(llvm::Attribute::getWithVScaleRangeArgs(
61-
*llvm_target_->GetContext(), 1, tvm::arith::kAArch64VScaleValues.size()));
60+
unsigned int max_val =
61+
*std::max_element(arith::kAArch64VScaleValues.begin(), arith::kAArch64VScaleValues.end());
62+
func->addFnAttr(
63+
llvm::Attribute::getWithVScaleRangeArgs(*llvm_target_->GetContext(), 1, max_val));
6264
}
6365
#endif
6466
CodeGenCPU::SetTargetAttributes(func);

0 commit comments

Comments
 (0)