gemini-code-assist[bot] commented on code in PR #19638:
URL: https://github.com/apache/tvm/pull/19638#discussion_r3318659880
##########
src/target/llvm/codegen_aarch64.cc:
##########
@@ -58,9 +57,18 @@ void CodeGenAArch64::AddFunction(const GlobalVar& gvar,
const PrimFunc& f) {
void CodeGenAArch64::SetTargetAttributes(llvm::Function* func) {
// Add vscale_range() function attribute when appropriate.
if (llvm_target_->TargetHasCPUFeature("sve") ||
llvm_target_->TargetHasCPUFeature("sme")) {
- auto kVScaleValues = arith::GetVScaleValues(Target::Current());
- if (!kVScaleValues.empty()) {
- unsigned int max_val = *std::max_element(kVScaleValues.begin(),
kVScaleValues.end());
+ // Compute max_val = largest power-of-two <= vector_width/8.
+ static auto llvm_get_vector_width_fn =
+ tvm::ffi::Function::GetGlobalRequired("target.llvm_get_vector_width");
+ unsigned int vector_width =
+ static_cast<unsigned
int>(llvm_get_vector_width_fn(Target::Current()).cast<int>());
+ unsigned int max_val = 0;
+ for (unsigned int i = 0;; ++i) {
+ unsigned int power = 1u << i;
+ if (power > (vector_width / 8)) break;
+ max_val = power;
+ }
Review Comment:

To prevent any potential shift-out-of-bounds undefined behavior (if
`vector_width / 8` is extremely large), we should limit the loop index `i` to
be less than 32. This ensures the shift operation `1u << i` is always
well-defined.
```c
unsigned int limit = vector_width / 8;
for (unsigned int i = 0; i < 32; ++i) {
unsigned int power = 1u << i;
if (power > limit) break;
max_val = power;
}
```
##########
src/arith/rewrite_simplify.cc:
##########
@@ -34,15 +34,41 @@
#include <utility>
#include "../target/datatype/registry.h"
+#include "../tirx/analysis/check_contains.h"
#include "conjunctive_normal_form.h"
#include "const_fold.h"
#include "constraint_extract.h"
#include "pattern_match.h"
-#include "scalable_expression.h"
namespace tvm {
namespace arith {
+namespace {
+// File-local helper: true if `expr` is a call to tirx::builtin::vscale().
+bool IsVScaleCall(const PrimExpr& expr) {
+ if (const auto* call = expr.as<tirx::CallNode>()) {
+ return call->op.same_as(tirx::builtin::vscale());
+ }
+ return false;
+}
+
+// File-local helper: true if `expr` contains a call to
tirx::builtin::vscale().
+bool ContainsVscaleCall(const PrimExpr& expr) {
+ return tirx::CheckContains::ExprContains(expr, IsVScaleCall);
+}
+
+// File-local helper: returns the vscale multiplier if `lanes` is of the form
+// `multiplier * vscale()` or `vscale() * multiplier`, nullopt otherwise.
+std::optional<int> ExtractVscaleFactor(const PrimExpr& lanes) {
+ PVar<IntImm> multiplier;
+ PCallExpr<PVscaleOp> vscale;
+ if (PMatchesOneOf(multiplier * vscale, vscale * multiplier).Match(lanes)) {
+ return multiplier.Eval()->value;
+ }
Review Comment:

The function returns `std::optional<int>`, but `multiplier.Eval()->value` is
of type `int64_t`. To prevent implicit narrowing conversion warnings on
stricter compilers, use an explicit `static_cast<int>`.
```c
if (PMatchesOneOf(multiplier * vscale, vscale * multiplier).Match(lanes)) {
return static_cast<int>(multiplier.Eval()->value);
}
```
##########
src/tirx/ir/expr.cc:
##########
@@ -29,13 +29,34 @@
#include <optional>
-#include "../../arith/scalable_expression.h"
#include "../../support/str_escape.h"
#include "buffer_common.h"
namespace tvm {
namespace tirx {
+namespace {
+// File-local helper: returns the vscale multiplier if `lanes` is of the form
+// `multiplier * vscale()` or `vscale() * multiplier`, nullopt otherwise.
+std::optional<int> ExtractVscaleFactor(const PrimExpr& lanes) {
+ auto is_vscale = [](const PrimExpr& e) -> bool {
+ if (const auto* call = e.as<CallNode>()) {
+ return call->op.same_as(tirx::builtin::vscale());
+ }
+ return false;
+ };
Review Comment:

Since we are already inside `namespace tvm::tirx`, the `tirx::` namespace
qualification is redundant and can be simplified to `builtin::vscale()`.
```suggestion
auto is_vscale = [](const PrimExpr& e) -> bool {
if (const auto* call = e.as<CallNode>()) {
return call->op.same_as(builtin::vscale());
}
return false;
};
```
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]