lightzhan-intellif opened a new issue, #13559:
URL: https://github.com/apache/tvm/issues/13559
Hi, all
I just encountered a problem with alloc_buffer in TVMScript. Let's look at
the following code:
```python
@T.prim_func
def test_alloc(in_data: T.Buffer[(256, 256), "float32"]):
a = in_data[0, 0]
p = T.alloc_buffer((256, 256), dtype="float32")
for i, j in T.grid(256, 256):
p[i, j] = in_data[i, j]
```
This is just an example, the meaning of it does not matter. When I ran the
code, tvm reported an error:
_ValueError: Block frame or PrimFunc frame not find. Please ensure
'T.alloc_buffer' is called under T.block() or T.prim_func()_
This message makes me very confused. By digging into it deeply, I found that
TVMScript only allows alloc_buffer under the T.block or T.prim_func(just like
what the message said). Actually, from the code text above, alloc_buffer is
exactly under the T.prim_func, the error message does not elaborate on what is
wrong with my code.
Actually, the problem is the stmt `a = in_data[0, 0]` will be interpreted as
a let-bind. That is to say, 'p = T.alloc_buffer((256, 256), dtype="float32")'
is under a let-bind stmt, which is not allowed at present. For non-tvm-native
users, I think it is very unfriendly.
So, I propose to remove this restriction which is very easy to do under the
new parser, just need to change `GetLastFrame` to `FindFrame` in the funciton
`AllocBuffer`:
```c++
Buffer AllocBuffer(Array<PrimExpr> shape, DataType dtype, Optional<Var> data,
Array<PrimExpr> strides, PrimExpr elem_offset, String
storage_scope, int align,
int offset_factor, String buffer_type_str, Array<IntImm>
axis_separators) {
Buffer buffer = BufferDecl(shape, dtype, "", data, strides, elem_offset,
storage_scope, align,
offset_factor, buffer_type_str,
axis_separators);
IRBuilder builder = IRBuilder::Current();
if (Optional<BlockFrame> frame = builder->GetLastFrame<BlockFrame>()) {
// change GetLastFrame->FindFrame
frame.value()->alloc_buffers.push_back(buffer);
} else if (Optional<PrimFuncFrame> frame =
builder->GetLastFrame<PrimFuncFrame>()) { // change GetLastFrame->FindFrame
frame.value()->root_alloc_buffers.push_back(buffer);
} else {
LOG(FATAL) << "ValueError: Block frame or PrimFunc frame not find.
Please ensure "
"'T.alloc_buffer' is called under T.block() or
T.prim_func()";
}
return buffer;
}
```
After these little changes, users can put alloc_buffer under any other
stmts. And it will be lifted to the proper position (under T.block or
T.prim_func) automically rather than emit a confusing message.
What do you make of this idea?
--
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]
