gemini-code-assist[bot] commented on code in PR #18876:
URL: https://github.com/apache/tvm/pull/18876#discussion_r2891639551
##########
src/s_tir/analysis/calculate_allocated_memory.cc:
##########
@@ -79,8 +79,7 @@ class AllocBufferCalculator : public StmtExprVisitor {
size *= op->buffer->dtype.bytes() * op->buffer->dtype.lanes();
_current_size[storage_scope] += size;
_max_size[storage_scope] = std::max(_current_size[storage_scope],
_max_size[storage_scope]);
- StmtExprVisitor::VisitStmt(op->body);
- _current_size[storage_scope] -= size;
Review Comment:
The logic in `AllocBufferCalculator` appears to be incorrect after this
refactoring. Previously, it would decrement the allocated size after visiting
the `AllocBuffer`'s body, correctly modeling scoped allocation. With the move
to flat `AllocBuffer` nodes, the `_current_size -= size;` line has been
removed, but no alternative mechanism for tracking deallocation at the end of a
scope has been added.
As a result, this visitor now calculates the cumulative size of all
allocations, rather than the peak memory usage. This is likely not the intended
behavior and could lead to incorrect memory analysis.
To fix this, `AllocBufferCalculator` should be made scope-aware. For
example, by overriding `VisitStmt_` for scoping constructs like `ForNode` or
`SBlockNode`, it could track allocations made within that scope and decrement
the size upon exiting the scope.
--
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]
