gemini-code-assist[bot] commented on code in PR #19926:
URL: https://github.com/apache/tvm/pull/19926#discussion_r3510319126
##########
src/tirx/ir/expr_functor.cc:
##########
@@ -27,6 +27,31 @@
namespace tvm {
namespace tirx {
+void ExprVisitor::VisitExpr(const Expr& expr) {
+ if (auto prim = expr.as<PrimExpr>()) {
+ return this->VisitExpr(prim.value());
+ }
+ if (const auto* op = expr.as<SizeVarNode>()) return VisitExpr_(op);
+ if (const auto* op = expr.as<VarNode>()) return VisitExpr_(op);
+ if (const auto* op = expr.as<CallNode>()) return VisitExpr_(op);
+ if (const auto* op = expr.as<StringImmNode>()) return VisitExpr_(op);
+ TVM_FFI_THROW(TypeError) << "Unsupported non-primitive expression: " <<
expr->GetTypeKey();
+}
+
+Expr ExprMutator::VisitExpr(const Expr& expr) {
+ if (auto prim = expr.as<PrimExpr>()) {
+ return this->VisitExpr(prim.value());
+ }
+ if (expr.as<VarNode>() || expr.as<StringImmNode>()) return expr;
Review Comment:
By returning `expr` directly for `VarNode` without calling `VisitExpr_`, any
custom mutator that overrides `VisitExpr_(const VarNode*)` or `VisitExpr_(const
SizeVarNode*)` will fail to mutate/remap variables when they are encountered as
generic `Expr`s (such as in `BindNode::value` or `CallNode::args`).\n\nWe
should dispatch `VarNode` and `SizeVarNode` to their respective `VisitExpr_`
methods, similar to how `ExprVisitor::VisitExpr` does.
```c
if (const auto* op = expr.as<SizeVarNode>()) return
this->VisitExpr_(op);\n if (const auto* op = expr.as<VarNode>()) return
this->VisitExpr_(op);\n if (expr.as<StringImmNode>()) return expr;
```
##########
src/target/source/codegen_c.cc:
##########
@@ -996,13 +1010,13 @@ void CodeGenC::VisitExpr_(const LetNode* op,
std::ostream& os) { // NOLINT(*)
var_idmap_[op->var.get()] = value;
} else {
PrintIndent();
- if (op->var.ty().IsHandle() && handle_data_type_.count(op->var.get())) {
+ if (op->var->ty.as<PointerTypeNode>() &&
handle_data_type_.count(op->var.get())) {
Review Comment:
Checking only `op->var->ty.as<PointerTypeNode>()` misses the case where the
variable's type is a `PrimType` of handle type. This causes the code to bypass
the `handle_data_type_` lookup and print generic handle types instead of the
specific pointer types.\n\nWe should use the same robust `is_pointer` check
that was implemented in `PrintFunctionSignature`.
```c
bool is_pointer = op->var->ty.as<PointerTypeNode>() ||\n
(op->var->ty.as<PrimTypeNode>() && op->var.ty().IsHandle());\n if
(is_pointer && handle_data_type_.count(op->var.get())) {
```
##########
src/arith/ir_mutator_with_analyzer.cc:
##########
@@ -203,17 +203,18 @@ Stmt IRMutatorWithAnalyzer::VisitStmt_(const SBlockNode*
op) {
}
Stmt IRMutatorWithAnalyzer::VisitStmt_(const BindNode* op) {
- PrimExpr value = this->VisitExpr(op->value);
+ ffi::Optional<PrimExpr> prim_value = op->value.as<PrimExpr>();
+ if (!prim_value.defined()) {
+ return ffi::GetRef<Stmt>(op);
+ }
+ PrimExpr value = this->VisitExpr(prim_value.value());
if (SideEffect(value) <= CallEffectKind::kPure) {
analyzer_->Bind(op->var, value);
}
- if (value.same_as(op->value)) {
- return ffi::GetRef<Stmt>(op);
- } else {
- auto n = this->CopyOnWrite(op);
- n->value = std::move(value);
- return Stmt(n);
- }
+ if (value.same_as(op->value)) return ffi::GetRef<Stmt>(op);
+ auto n = this->CopyOnWrite(op);
+ n->value = std::move(value);
+ return Stmt(n);
Review Comment:
Returning early when `op->value` is not a `PrimExpr` completely skips
mutating non-primitive bindings (such as pointer-typed bindings). This means
any nested variables or expressions inside these bindings will not be
mutated/remapped.\n\nWe should mutate `op->value` as an `Expr` first, and then
conditionally bind to the analyzer if the mutated value is a `PrimExpr`.
```c
Expr value = this->VisitExpr(op->value);\n if (auto prim_value =
value.as<PrimExpr>()) {\n if (SideEffect(prim_value.value()) <=
CallEffectKind::kPure) {\n analyzer_->Bind(op->var, prim_value.value());\n
}\n }\n if (value.same_as(op->value)) return ffi::GetRef<Stmt>(op);\n
auto n = this->CopyOnWrite(op);\n n->value = std::move(value);\n return
Stmt(n);
```
##########
src/arith/ir_visitor_with_analyzer.cc:
##########
@@ -50,8 +50,10 @@ void IRVisitorWithAnalyzer::VisitStmt_(const SBlockNode* op)
{
}
void IRVisitorWithAnalyzer::VisitStmt_(const BindNode* op) {
- this->VisitExpr(op->value);
- analyzer_->Bind(op->var, op->value);
+ if (ffi::Optional<PrimExpr> value = op->value.as<PrimExpr>()) {
+ this->VisitExpr(value.value());
+ analyzer_->Bind(op->var, value.value());
+ }
Review Comment:
Skipping the visit of `op->value` when it is not a `PrimExpr` means that any
variables or sub-expressions inside non-primitive bindings will not be visited.
This can break analyses that rely on visiting all expressions in the AST.\n\nWe
should visit `op->value` as an `Expr` first, and then conditionally bind to the
analyzer if it is a `PrimExpr`.
```c
this->VisitExpr(op->value);\n if (ffi::Optional<PrimExpr> value =
op->value.as<PrimExpr>()) {\n analyzer_->Bind(op->var, value.value());\n }
```
##########
src/target/source/codegen_c.cc:
##########
@@ -1121,13 +1135,13 @@ void CodeGenC::VisitStmt_(const BindNode* op) {
var_idmap_[op->var.get()] = value;
} else {
PrintIndent();
- if (op->var.ty().IsHandle() && handle_data_type_.count(op->var.get())) {
+ if (op->var->ty.as<PointerTypeNode>() &&
handle_data_type_.count(op->var.get())) {
Review Comment:
Checking only `op->var->ty.as<PointerTypeNode>()` misses the case where the
variable's type is a `PrimType` of handle type. This causes the code to bypass
the `handle_data_type_` lookup and print generic handle types instead of the
specific pointer types.\n\nWe should use the same robust `is_pointer` check
that was implemented in `PrintFunctionSignature`.
```c
bool is_pointer = op->var->ty.as<PointerTypeNode>() ||\n
(op->var->ty.as<PrimTypeNode>() && op->var.ty().IsHandle());\n if
(is_pointer && handle_data_type_.count(op->var.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]
