Hzfengsy commented on a change in pull request #8408:
URL: https://github.com/apache/tvm/pull/8408#discussion_r666000311
##########
File path: python/tvm/script/node.py
##########
@@ -108,7 +108,7 @@ def check_index(index: Union[int, PrimExpr]):
span,
)
- slices: List[Slice] = []
+ slices: List[Slice, BufferSlice] = []
Review comment:
```suggestion
slices: List[Union[Slice, BufferSlice]] = []
```
##########
File path: tests/python/unittest/test_tvmscript_complete.py
##########
@@ -190,9 +190,77 @@ def render(e):
assert False
[email protected]
+def func_with_bufferslice_indices(data: ty.handle, index: ty.handle) -> None:
+ data_buf = tir.match_buffer(data, (16, 16), "float32")
+ index_buf = tir.match_buffer(index, (1,), "int32")
+ out_buf = tir.alloc_buffer((16, 16), "float32")
+
+ with tir.block([16, 16]) as [vi, vj]:
+ out_buf[vi, vj] = data_buf[vi, index_buf[0]]
+
+
[email protected]
+def expected_bufferslice_indices(data: ty.handle, index: ty.handle) -> None:
+ index_buf = tir.match_buffer(
Review comment:
Personally recommend writing expected function manually rather than
printing the result out. Two reasons here:
1. make test file clear and short if we can use sugar
2. make sure that the function is EXACTLY what we want. it is more likely to
have bugs in printing functions since we may omit some details when checking it.
--
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]
