Skip to content

Conversation

@cyx-6
Copy link
Contributor

@cyx-6 cyx-6 commented Jan 10, 2023

This PR implements the syntax sugar of T.axis.remap for new TVMScript printer. This syntax sugar will synthesize the T.axis.remap when there are more than 2 simple block iterating variable bindings. For example, it will change

for i, j, k in T.grid(128, 128, 128):
  with T.block("update"):
    vi = T.axis.spatial(128, i)
    vj = T.axis.spatial(128, j)
    vk = T.axis.reduce(128, k)

into

for i, j, k in T.grid(128, 128, 128):
  with T.block("update"):
    vi, vj, vk = T.axis.remap("SSR", [i, j, k])

@tvm-bot
Copy link
Collaborator

tvm-bot commented Jan 10, 2023

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

Generated by tvm-bot

Copy link
Member

@junrushao junrushao left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM! Will need to merge #13785 before getting this in

@junrushao
Copy link
Member

@cyx-6 would you please rebase against the latest main given #13785 has been merged? Thanks a lot! I will get this PR merged as soon as I validated it's working in the testsuite

@cyx-6
Copy link
Contributor Author

cyx-6 commented Jan 17, 2023

@cyx-6 would you please rebase against the latest main given #13785 has been merged? Thanks a lot! I will get this PR merged as soon as I validated it's working in the testsuite

rebased

@cyx-6
Copy link
Contributor Author

cyx-6 commented Jan 17, 2023

@tvm-bot rerun

@junrushao
Copy link
Member

@cyx-6 you will have to rebase again :-)

fix outdated code

fix outdated unittest

refactor
Co-authored-by: Junru Shao <junrushao1994@gmail.com>
@junrushao junrushao merged commit 6c5be6f into apache:main Jan 18, 2023
fzi-peccia pushed a commit to fzi-peccia/tvm that referenced this pull request Mar 27, 2023
…#13743)

This PR implements the syntax sugar of `T.axis.remap` for new TVMScript printer. This syntax sugar will synthesize the `T.axis.remap` when there are more than 2 simple block iterating variable bindings. For example, it will change
```python
for i, j, k in T.grid(128, 128, 128):
  with T.block("update"):
    vi = T.axis.spatial(128, i)
    vj = T.axis.spatial(128, j)
    vk = T.axis.reduce(128, k)
```
into
```python
for i, j, k in T.grid(128, 128, 128):
  with T.block("update"):
    vi, vj, vk = T.axis.remap("SSR", [i, j, k])
```

Co-authored-by: Junru Shao <junrushao1994@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants