DzAvril opened a new pull request #10652:
URL: https://github.com/apache/tvm/pull/10652
When employing double_buffer in tensor core conv2d template, such as:
```python
# conv2d_implicit_gemm_tensorcore.py
def schedule_implicit_gemm_tensorcore(cfg, sch, Conv):
...
# double buffer
sch[AS].double_buffer()
...
```
It results in such assert:
```bash
Check failed: (it != buf_map_.end()) is false: Cannot find allocated buffer
for buffer(im2col_reshape.shared, 0x5619189281f0)
```
Attached is my test script.
[test_tensorcore_conv_nhwc.py.txt](https://github.com/apache/tvm/files/8282512/test_tensorcore_conv_nhwc.py.txt)
As described above, buffer `im2col_reshape.shared` will be doubled. There
are several passes for lowering double_buffer. One of them is storage_flatten.
The lowered tir fed to storage_flatten pass has two types of attributes include
buffer `im2col_reshape.shared`.
- buffer(im2col_reshape.shared, 0x55e3524f0a40), op->attr_key:
double_buffer_scope
- [buffer(buffer, 0x55e3524f44d0), buffer(im2col_reshape.shared,
0x55e3524f0a40)], op->attr_key: buffer_bind_scope
Noticed that at the begin of this pass the pointers of buffer
`im2col_reshape.shared` are same. There is a pass buffer_stride in pass
storage_flatten which has such process for attribute statement.
``` C++
// storage_flatten.cc:BufferStrideLegalize
Stmt VisitStmt_(const AttrStmtNode* op) final {
if (op->attr_key == attr::buffer_dim_align) {
auto buffer = Downcast<tir::Buffer>(op->node);
const CallNode* tuple = op->value.as<CallNode>();
ICHECK(tuple && tuple->op.same_as(builtin::tvm_tuple()));
auto& vinfo = dim_align_[buffer];
int dim = tuple->args[0].as<IntImmNode>()->value;
if (static_cast<size_t>(dim) >= vinfo.size()) {
vinfo.resize(dim + 1);
}
vinfo[dim].align_factor = tuple->args[1].as<IntImmNode>()->value;
vinfo[dim].align_offset = tuple->args[2].as<IntImmNode>()->value;
return this->VisitStmt(op->body);
} else if (op->attr_key == attr::buffer_bind_scope) {
Array<ObjectRef> arr = Downcast<Array<ObjectRef>>(op->node);
ICHECK_EQ(arr.size(), 2U);
Buffer source = Downcast<Buffer>(arr[0]);
Buffer target_with_strides = WithStrides(Downcast<Buffer>(arr[1]));
Buffer source_with_strides = WithStrides(source);
{
BufferEntry entry;
entry.remap_to = source_with_strides;
entry.in_scope = true;
entry.is_external = false;
buf_map_[source] = entry;
}
Stmt body = this->VisitStmt(op->body);
return AttrStmt(Array<ObjectRef>{source_with_strides,
target_with_strides}, op->attr_key,
op->value, body, op->span);
} else {
return StmtExprMutator::VisitStmt_(op);
}
}
```
In branch `op->attr_key == attr::buffer_bind_scope`, buffer
`im2col_reshape.shared` is passed to `WithStrides` and modified in it. Then
buffer `im2col_reshape.shared` changes to `[buffer(buffer, 0x55e3524f44d0),
buffer(im2col_reshape.shared, 0x55e352500860)], op->attr_key:
buffer_bind_scope`. Noticed again the pointer to `im2col_reshape.shared` is
changed afterwards. And because there is none branch for processing attribute
`double_buffer_scope`, this cause mismatch of pointer to buffer
`im2col_reshape.shared`.
So add branch below will solve this issue.
``` C++
else if (op->attr_key == attr::double_buffer_scope) {
auto buffer = Downcast<tir::Buffer>(op->node);
Buffer buffer_with_strides = WithStrides(buffer);
Stmt body = this->VisitStmt(op->body);
return AttrStmt(buffer_with_strides, op->attr_key, op->value, body,
op->span);
}
```
--
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]
