[Bugfix] Add a nullptr check to tir.Buffer to fix the illegal memory access

[syang-ng]

Bug Description

I found there will be a segfault when we try to load a buffer whose data_ points to nullptr.

Here is the code to trigger the segmentation fault error.

import tvm
from tvm import tir

var = tir.Var('a',dtype='int32')
buf = tir.decl_buffer((1,), name='buf')
buf_load = tir.expr.BufferLoad(buffer=buf, indices=tvm.runtime.convert([0]))
buf_load_stmt = tir.stmt.Evaluate(buf_load)
for_loop = tir.stmt.For(loop_var=var, kind=1, min_val=1, extent=buf_load, body=buf_load_stmt)
buf_func = tir.PrimFunc(params={}, body=for_loop)

tvm.lower(buf_func)

(tvm-build) ➜  ~ vim buffer_nullptr.py
(tvm-build) ➜  ~ python3 buffer_nullptr.py
[1]    28338 segmentation fault (core dumped)  python3 buffer_nullptr.py

Bug Analysis

If we call the function tvm.build(buf_func), it will execute Buffer::vload() in src/tir/ir/buffer.cc:

PrimExpr Buffer::vload(Array<PrimExpr> begin, DataType dtype) const {
  // specially handle bool, stored as DataType::Int(8)
  const BufferNode* n = operator->();
  ICHECK(dtype.element_of() == n->dtype.element_of() && dtype.lanes() % n->dtype.lanes() == 0)
      << "Cannot load " << dtype << " from buffer of " << n->dtype;
  if (dtype == DataType::Bool()) {
    return tir::Cast(DataType::Bool(),
                     tir::Load(DataType::Int(8), n->data, BufferOffset(n, begin, DataType::Int(8)),
                               const_true()));
  } else {
    return tir::Load(dtype, n->data, BufferOffset(n, begin, dtype), const_true(dtype.lanes()));
  }
}

Since n might be nullptr, n->dtype.element_of() will actually access the illegal memory address, then a segfault occurs.

So to prevent illegal memory access, it's necessary to add a checker to check whether n points to a valid object. And I also found the similar checkers had been added into TVM such as line 40 in include/tvm/ir/env_func.h:

Thus, I add the checker in src/tir/ir/buffer.cc to prevent illegal memory access in this commit.

As tvm::tir::Buffer extends from tvm::runtime::Object, it can also solve the problem by adding the checkers to tvm::runtime::Object.

[junrushao]

LGTM

[tqchen]

Thanks @syang-ng . It would be great to add a unit test case to cover this bug.

[syang-ng]

I have added a test case to cover this problem. Could you help review my new commit? Thanks! @tqchen @junrushao1994

[Hzfengsy]

LGTM. Just curious why the data_ will be nullptr

[junrushao]

@Hzfengsy in the example below:

buf = tir.decl_buffer((1,), name='buf')

which actually indicates:

buf = tir.decl_buffer(data=None, shape=(1,), name='buf')

I suppose it won't happen in our common workflow, but is an adversarial case when using our APIs in an arbitrary way

[junrushao]

Thanks @syang-ng for the enhancement!