tqchen commented on code in PR #19926:
URL: https://github.com/apache/tvm/pull/19926#discussion_r3510406193


##########
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:
   Fixed in d6e16196f4. Non-primitive Var/SizeVar cannot safely use the legacy 
VisitExpr_ hooks because those return PrimExpr, so the generic path now 
dispatches through new virtual exact-Expr VisitExpr(const Var&) / 
VisitExpr(const SizeVar&) hooks. Scalar variables still dispatch through the 
existing PrimExpr vtable, while pointer-aware mutators (including 
PrimFuncSpecializer) override the exact Var hook. The pointer-specialization 
regression passes.



##########
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:
   Fixed in d6e16196f4. The Let codegen path now recognizes both exact 
PointerType and primitive handle views before using handle_data_type_. 
RegisterHandleTypeFromPointer also preserves legacy GetType inference for 
primitive ptr_byte_offset calls and exact-type lookup for generic Expr calls. A 
C-source regression verifies float* output.



##########
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:
   Fixed in d6e16196f4. The Bind codegen path now recognizes both exact 
PointerType and primitive handle views before using handle_data_type_. The same 
regression verifies that both Bind and Let pointer-offset locals are emitted as 
float*.



##########
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:
   Fixed in d6e16196f4. IRMutatorWithAnalyzer now mutates every Bind value as 
Expr first, then conditionally performs side-effect analysis and analyzer 
binding only when the mutated result is a PrimExpr. Added a regression that 
simplifies a scalar argument nested inside a pointer-valued Bind.



##########
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:
   Fixed in d6e16196f4. IRVisitorWithAnalyzer now visits every Bind value as 
Expr before conditionally narrowing to PrimExpr for analyzer binding. Added an 
OOBChecker regression where an out-of-bounds BufferLoad is nested inside a 
pointer-valued Bind.



-- 
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]

Reply via email to