diff --git a/src/tir/analysis/verify_well_formed.cc b/src/tir/analysis/verify_well_formed.cc index 878618fbe6fdf..e0318e14080f4 100644 --- a/src/tir/analysis/verify_well_formed.cc +++ b/src/tir/analysis/verify_well_formed.cc @@ -60,24 +60,39 @@ class BlockVarAccessVerifier : public StmtExprVisitor { void VisitExpr_(const VarNode* op) final { auto it = loop_vars_.find(op); - if (it != loop_vars_.end() && it->second < cur_block_level_) { + if (it != loop_vars_.end() && it->second < block_stack_.size()) { has_error_ = true; if (assert_mode_) { - report_error(op); + if (it->second == 0) { + LOG(FATAL) << "Well-formedness check failed: " + << "Loop iterator var " << op->name_hint + << " is defined outside of any block, " + << "but is used inside the non-opaque current block \"" + << block_stack_.back()->name_hint << "\"."; + } else { + LOG(FATAL) << "Well-formedness check failed: " + << "Loop iterator var " << op->name_hint << " is defined in block \"" + << block_stack_[it->second - 1]->name_hint << "\", " + << "but is used inside the non-opaque current block \"" + << block_stack_.back()->name_hint << "\"."; + } } } } void VisitStmt_(const ForNode* op) final { ICHECK(loop_vars_.find(op->loop_var.get()) == loop_vars_.end()); - loop_vars_[op->loop_var.get()] = cur_block_level_; + loop_vars_[op->loop_var.get()] = block_stack_.size(); StmtExprVisitor::VisitStmt_(op); loop_vars_.erase(op->loop_var.get()); } void VisitStmt_(const BlockNode* op) final { // Do not check boundary if it's a opaque block. - cur_block_level_ += !op->iter_vars.empty(); + bool is_non_opaque = op->iter_vars.size(); + if (is_non_opaque) { + block_stack_.push_back(op); + } // Step 0. Skip block iter var's domain @@ -103,22 +118,18 @@ class BlockVarAccessVerifier : public StmtExprVisitor { } this->VisitStmt(op->body); - cur_block_level_ -= !op->iter_vars.empty(); + if (is_non_opaque) { + block_stack_.pop_back(); + } } private: - void report_error(const VarNode* var) { - // TODO(siyuan): use the error message from the parser. - LOG(FATAL) << "Well-formedness check failed: outside defined var " << var->name_hint - << " is used inside the current block."; - } - /*! \brief The map from outside loop vars to its corresponding block level. */ std::unordered_map loop_vars_; /*! \brief Whether it's in assert mode. */ bool assert_mode_; /*! \brief Current nested block stack level. */ - size_t cur_block_level_{0}; + std::vector block_stack_; /*! \brief Whether there is error. */ bool has_error_{false}; };