Implicit requirements on value type of AttrStmt

Troubleshooting
#1

Hi all, besides the problem I mentioned in this post. I found another interesting implicit requirement on the value type of AttrStmt.

When I set the attr_key to compute_scope or coproc_uop_scope and the dtype of value bo be tir.IntImm, there will be a crash if I try to build the module.

Reproduce Crashes

Here are two pieces of code to trigger crashes.

(tvm-build) ➜  ~ cat test_compute_scope.py
import tvm
from tvm import ir, tir

a = tir.Var("a", "int32")
b = tir.Var("b", "handle")
iter_var = tir.IterVar(ir.Range(0,1 ), a, 1)
buffer = tir.buffer.decl_buffer((1,))
buffer_map = {b: buffer}
store = tir.Store(buffer.data, tir.const(1), tir.const(1))
attr_stmt = tir.AttrStmt(iter_var, "compute_scope", tir.const(1), store)
f = tir.PrimFunc({a, b}, body=attr_stmt, buffer_map=buffer_map)
mod = tvm.lower(f)
tvm.build(mod)

(tvm-build) ➜  ~ python3 test_compute_scope.py
[1]    8406 segmentation fault (core dumped)  python3 test_compute_scope.py

(tvm-build) ➜  ~ cat test_coproc_uop_scope.py
import tvm
from tvm import ir, tir

a = tir.Var("a", "int32")
b = tir.Var("b", "handle")
iter_var = tir.IterVar(ir.Range(0,1 ), a, 1)
buffer = tir.buffer.decl_buffer((1,))
buffer_map = {b: buffer}
store = tir.Store(buffer.data, tir.const(1), tir.const(1))
attr_stmt = tir.AttrStmt(iter_var, "coproc_uop_scope", tir.const(1), store)
f = tir.PrimFunc({a, b}, body=attr_stmt, buffer_map=buffer_map)
mod = tvm.lower(f)
tvm.build(mod)

(tvm-build) ➜  ~ python test_coproc_uop_scope.py
[1]    9114 segmentation fault (core dumped)  python test_coproc_uop_scope.py

Bug Analysis

To find the root cause of these crashes, I use the gdb to follow the execution and locate the bugs. Then I found that they all crash in src/target/llvm/codegen_cpu.cc.

It seems that there is an implicit requirement on the value type of AttrStmt with specific attr_key. According to the source code, when we set the attr_key to coproc_uop_scope, compute_scope or pragma_import_llvm, TVM keeps the assumption that the value should be String and then directly convert it to StringImmNode, and only in pragma_import_llvm branch, TVM checks whether the pointer is valid. That’s the reason why we crash if we set the attr_key to compute_scope or coproc_uop_scope with the value tir.const(1).

void CodeGenCPU::CreateComputeScope(const AttrStmtNode* op) {
  // ...
  llvm::Function* fcompute = llvm::Function::Create(
      ftype, llvm::Function::InternalLinkage,
      // no checks here
      op->value.as<StringImmNode>()->value.operator llvm::StringRef(), module_.get());
  // ...
}

void CodeGenCPU::VisitStmt_(const AttrStmtNode* op) {
  if (op->attr_key == tir::attr::coproc_uop_scope) {
    // no checks here
    this->CreateStaticInit(op->value.as<StringImmNode>()->value, op->body);
  } else if (op->attr_key == tir::attr::compute_scope) {
    this->CreateComputeScope(op);
  } else if (tir::attr::IsPragmaKey(op->attr_key)) {
    // ...
    } else if (op->attr_key == tir::attr::pragma_import_llvm) {
      const StringImmNode* value = op->value.as<StringImmNode>();
      // check whether the value is valid
      ICHECK(value != nullptr);
      this->HandleImport(value->value);
      this->VisitStmt(op->body);
    }
    // ...
  }
  // ...
}

I think it would be better if we add some checks in src/target/llvm/codegen_cpu.cc to prevent these crashes? And some necessary messages for developers to be aware of the specific requirement on dtype of value will also help a lot. :blush: