Beya2019 opened a new pull request #8408:
URL: https://github.com/apache/tvm/pull/8408
TVMScript Parser not support buffer as indices, such as index_buf[0] is the
indices of data_buf[vi, vj, index_buf[0]]
```
@tvm.script.tir
def func(data: ty.handle, index: ty.handle) -> None:
data_buf = tir.match_buffer(data, (1, 5, 6), "float16")
index_buf = tir.match_buffer(index, (1,), "int32")
copy_buf = tir.alloc_buffer((1, 5), "float16")
with tir.block([1, 5], "init") as [vi, vj]:
copy_buf[vi, vj] = data_buf[vi, vj, index_buf[0]]
```
we add the code to support this function, so we can get the result when call
`tvm.script.asscript(func)`
```
primfn(data: handle, index: handle) -> ()
buffers = {index_buf: Buffer(index_buf_1: Pointer(int32), int32, [1], []),
data_buf: Buffer(data_buf_1: Pointer(float16), float16, [1, 5,
6], [])}
buffer_map = {data: data_buf, index: index_buf} {
block([], "root") {
tir.reads([])
tir.writes([])
copy_buf = alloc_buffer(float16[1, 5])
for (i0: int32, 0, 1) {
for (i1: int32, 0, 5) {
block([1, 5], "init") as [vi, vj] {
bind(vi, i0)
bind(vj, i1)
tir.reads([data_buf[vi, vj, 0:6], index_buf[0]])
tir.writes([copy_buf[vi, vj]])
copy_buf[vi, vj] = data_buf[vi, vj, index_buf[0]]
}
}
}
```
@Hzfengsy Would you please have a look at this? Thanks very much.
--
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]
