gemini-code-assist[bot] commented on code in PR #19733:
URL: https://github.com/apache/tvm/pull/19733#discussion_r3396554431
##########
src/s_tir/transform/remove_store_undef.cc:
##########
@@ -61,11 +69,15 @@ class StoreUndefLocator : public StmtExprVisitor {
<< "Error: T.undef() used in BufferStore expressions "
<< "must not have other side effects";
undef_stores_.insert(op);
+ // Record the destination buffer as undef-carrying so that any later
+ // BufferLoad of it propagates undef (this is the flattened form of an
+ // old let-bound undef variable, e.g. ``val = T.undef(); A[0] = val``).
+ undef_buffers_.insert(op->buffer->data.get());
Review Comment:
Unconditionally treating any buffer that has an undef store as
undef-carrying is unsafe for non-scalar buffers (e.g., arrays) or buffers with
mixed stores. If a multi-dimensional or non-scalar buffer has a single element
stored with `T.undef()`, this implementation will:
1. Treat all subsequent loads from that buffer (at any index) as undef.
2. Completely remove the `AllocBuffer` statement for that buffer, leaving
other valid stores/loads to an unallocated buffer, causing compilation or
runtime errors.
To prevent this, we should only track undef propagation for **scalar
buffers** (constant size of 1) and ensure they are **only** ever stored with
undef values (no mixed stores).
```c
undef_stores_.insert(op);
// Record the destination buffer as undef-carrying so that any later
// BufferLoad of it propagates undef (this is the flattened form of an
// old let-bound undef variable, e.g. ``val = T.undef(); A[0] = val``).
bool is_scalar = true;
int64_t size = 1;
for (const auto& dim : op->buffer->shape) {
if (const auto* int_dim = dim.as<IntImmNode>()) {
size *= int_dim->value;
} else {
is_scalar = false;
break;
}
}
if (is_scalar && size == 1) {
undef_buffers_.insert(op->buffer->data.get());
}
```
--
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]
