Crash triggered by 8 lines of TIR code

Troubleshooting
5

Hi all! I think I found the root cause of this segfault when I read related C++ backend implementation.

We can found that, function tvm::tir::IRConvertSSA::VisitExpr_(tvm::tir::LoadNode const*) implemented in tir/transforms/ir_utils.cc(Line 113) forgets to check the size of scope_[op->buffer_var.get()]:

PrimExpr VisitExpr_(const LoadNode* op) final {
  PrimExpr expr = StmtExprMutator::VisitExpr_(op);
  op = expr.as<LoadNode>();
  if (scope_.count(op->buffer_var.get())) {
    return Load(op->dtype, scope_[op->buffer_var.get()].back(), op->index, op->predicate);
  } else {
    return expr;
  }
}

As document said, calling std::vector::back on an empty container causes undefined behavior. So if scope_[op->buffer_var.get()] is actually empty, we will get undefined behavior, like this segfault.

To avoid undefined behavior, it’s necessary to check the size of vector before calling back(). And we can find another VisitExpr_ about AttrStmtNode actually uses scope_.count(v) and scope_[v].size() != 0 to check vector before calling scope_[v].back():

Stmt VisitStmt_(const AttrStmtNode* op) final {
  if (const VarNode* v = op->node.as<VarNode>()) {
    Stmt stmt = StmtExprMutator::VisitStmt_(op);
    op = stmt.as<AttrStmtNode>();
    if (scope_.count(v) && scope_[v].size() != 0) {
      return AttrStmt(scope_[v].back(), op->attr_key, op->value, op->body);
    } else {
      return stmt;
    }
  } else {
    return StmtExprMutator::VisitStmt_(op);
  }
}