gemini-code-assist[bot] commented on code in PR #19644:
URL: https://github.com/apache/tvm/pull/19644#discussion_r3328848544
##########
src/relax/transform/dataflow_inplace.cc:
##########
@@ -39,6 +39,67 @@
namespace tvm {
namespace relax {
+// Ops that may return a tensor sharing storage with the first argument.
+bool IsViewMemoryOp(const OpNode* op_node) {
+ static const std::unordered_set<std::string> kViewOps = {
+ "relax.expand_dims", "relax.squeeze", "relax.reshape",
+ "relax.permute_dims", "relax.memory.view",
"relax.memory.ensure_zero_offset",
+ };
+ return kViewOps.count(op_node->name);
+}
+
+// Look up alias ids for a call argument (only Var args are expected in
dataflow blocks).
+std::unordered_set<int> GetVarAliasSetFromExpr(
+ const Expr& arg, const std::unordered_map<Var, std::unordered_set<int>>&
alias_sets) {
+ if (auto* var_node = arg.as<VarNode>()) {
+ Var var = ffi::GetRef<Var>(var_node);
+ if (!alias_sets.count(var)) {
+ return {-1};
+ }
+ return alias_sets.at(var);
+ }
+ return {-1};
+}
+
+// In-place on arg `candidate` is invalid if another distinct operand may
alias the same
+// storage (e.g. two expand_dims views of x bound to different vars). We only
reject when
+// alias analysis gives a definite overlap; unknown (-1) other operands are
handled elsewhere.
+// Repeated uses of the same var (e.g. add(z, z)) are allowed.
+bool InplaceArgDisjointFromOtherCallArgs(
+ const CallNode* call_node, int candidate,
+ const std::unordered_map<Var, std::unordered_set<int>>& alias_sets) {
+ const auto* cand_var_node = call_node->args[candidate].as<VarNode>();
+ if (!cand_var_node) {
+ return false;
+ }
+ Var cand_var = ffi::GetRef<Var>(cand_var_node);
+ auto cand_set = GetVarAliasSetFromExpr(call_node->args[candidate],
alias_sets);
+ if (cand_set.count(-1)) {
+ return false;
+ }
+ for (size_t j = 0; j < call_node->args.size(); j++) {
+ if (static_cast<int>(j) == candidate) {
+ continue;
+ }
+ const Expr& other_arg = call_node->args[j];
+ if (const auto* other_var_node = other_arg.as<VarNode>()) {
+ if (other_var_node == cand_var_node ||
ffi::GetRef<Var>(other_var_node).same_as(cand_var)) {
+ continue;
+ }
+ }
+ auto other_set = GetVarAliasSetFromExpr(other_arg, alias_sets);
+ if (other_set.count(-1)) {
+ continue;
+ }
+ for (int alias_idx : other_set) {
+ if (cand_set.count(alias_idx)) {
+ return false;
+ }
+ }
+ }
+ return true;
+}
Review Comment:
### Correctness Bug & Simplification Opportunity
1. **Correctness Bug**: The check `if (other_set.count(-1)) { continue; }`
will skip checking any other alias indices for `other_arg` if its alias set
contains `-1`. However, an alias set can contain both `-1` (representing an
unknown/external alias) and other valid, known alias indices (e.g., if it's the
result of a mystery call that took the candidate as an argument). Skipping the
check in this case means we fail to detect a definite overlap, which can lead
to incorrect in-place rewrites and silent runtime correctness bugs.
2. **Simplification**: The pointer comparison `other_var_node ==
cand_var_node` and `same_as` check are redundant and can be simplified to a
single `other_arg.same_as(call_node->args[candidate])` check, which is cleaner
and avoids unnecessary casting.
Since `cand_set` is already checked to not contain `-1` (at line 77), we can
safely remove the `other_set.count(-1)` check entirely. Any `-1` in `other_set`
will not match any index in `cand_set` anyway.
```c
bool InplaceArgDisjointFromOtherCallArgs(
const CallNode* call_node, int candidate,
const std::unordered_map<Var, std::unordered_set<int>>& alias_sets) {
const auto* cand_var_node = call_node->args[candidate].as<VarNode>();
if (!cand_var_node) {
return false;
}
auto cand_set = GetVarAliasSetFromExpr(call_node->args[candidate],
alias_sets);
if (cand_set.count(-1)) {
return false;
}
for (size_t j = 0; j < call_node->args.size(); j++) {
if (static_cast<int>(j) == candidate) {
continue;
}
const Expr& other_arg = call_node->args[j];
if (other_arg.same_as(call_node->args[candidate])) {
continue;
}
auto other_set = GetVarAliasSetFromExpr(other_arg, alias_sets);
for (int alias_idx : other_set) {
if (cand_set.count(alias_idx)) {
return false;
}
}
}
return true;
}
```
--
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]
