gemini-code-assist[bot] commented on code in PR #19776:
URL: https://github.com/apache/tvm/pull/19776#discussion_r3411039536
##########
src/tirx/transform/vectorize_loop.cc:
##########
@@ -202,15 +220,25 @@ class TryPredicateBufferAccesses : public StmtExprMutator
{
num_accesses_rewritten_ += 1;
auto writer = node.CopyOnWrite();
- writer->predicate = lane_mask;
+ if (node->predicate.defined() && allow_offset_predication_) {
+ writer->predicate = node->predicate.value() & lane_mask;
+ } else {
+ writer->predicate = lane_mask;
+ }
Review Comment:
Using the bitwise AND operator `&` on `PrimExpr` creates a bitwise AND call
(`bitwise_and` intrinsic) rather than a logical `And` node. In TVM TIR,
predicates and logical conditions should be combined using logical AND (`&&` or
`And(...)`) to allow the arithmetic analyzer and simplification passes to
optimize them effectively. Using `&` can prevent these simplifications and
potentially lead to suboptimal code generation.
```suggestion
if (node->predicate.defined() && allow_offset_predication_) {
writer->predicate = node->predicate.value() && lane_mask;
} else {
writer->predicate = lane_mask;
}
```
##########
src/tirx/transform/vectorize_loop.cc:
##########
@@ -999,6 +1031,29 @@ class LoopVectorizer : public StmtMutator {
}
private:
+ Stmt VectorizeFixedLoopForRVV(const ForNode* op, int64_t extent) {
+ // Match the existing TIRx scalable-vector convention. LLVM/RVV still
+ // selects the runtime vector length with vsetvli.
+ static constexpr int kDefaultVScaleFactor = 4;
+ DataType index_dtype = op->loop_var->dtype;
+ PrimExpr zero = make_const(index_dtype, 0);
+ PrimExpr fixed_extent = make_const(index_dtype, extent);
+ PrimExpr scalable_lanes = CreateNewLanes(/*is_scalable=*/true,
kDefaultVScaleFactor);
+ PrimExpr num_chunks = ceildiv(fixed_extent, scalable_lanes);
Review Comment:
If the loop variable `op->loop_var` has a 64-bit integer type (`int64_t`),
`fixed_extent` will be `Int(64)`. However, `CreateNewLanes` hardcodes the
returned expression to `Int(32)`. Mixing `Int(64)` and `Int(32)` in binary
operations like `ceildiv` or loop bounds will cause a type mismatch error in
TVM. To prevent compiler crashes on 64-bit loop variables, we should explicitly
cast `scalable_lanes` to `index_dtype` if they do not match.
```suggestion
PrimExpr scalable_lanes = CreateNewLanes(/*is_scalable=*/true,
kDefaultVScaleFactor);
if (scalable_lanes.dtype() != index_dtype) {
scalable_lanes = Cast(index_dtype, scalable_lanes);
}
PrimExpr num_chunks = ceildiv(fixed_extent, scalable_lanes);
```
##########
src/tirx/transform/vectorize_loop.cc:
##########
@@ -190,10 +199,19 @@ class TryPredicateBufferAccesses : public StmtExprMutator
{
}
Ramp ramp = Downcast<Ramp>(node->indices[0]);
- // The vectorized access pattern must match the base of the predicate
- if (!ffi::StructuralEqual()(ramp->base, base_)) {
+ bool same_base = ffi::StructuralEqual()(ramp->base, base_);
+ if (!same_base && !allow_offset_predication_) {
return node;
}
+ if (!same_base) {
+ // The lane mask describes which lanes are active, independent of the
+ // memory base. This covers accesses such as A[offset + i] guarded by
+ // a predicate over i.
+ if (!ffi::StructuralEqual()(ramp->stride, stride_) ||
+ !ffi::StructuralEqual()(ramp->lanes, lanes_)) {
+ return node;
+ }
+ }
Review Comment:
Even when `same_base` is true, we should still verify that the stride and
lanes of the buffer access match the predicate's stride and lanes. If they
differ, applying the same lane mask would be incorrect. Simplifying this check
to always validate stride and lanes makes the predication logic much more
robust and cleaner.
```c
bool same_base = ffi::StructuralEqual()(ramp->base, base_);
if (!same_base && !allow_offset_predication_) {
return node;
}
if (!ffi::StructuralEqual()(ramp->stride, stride_) ||
!ffi::StructuralEqual()(ramp->lanes, lanes_)) {
return node;
}
```
--
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]
