Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[TVMScript] Sugar T.env_thread + T.launch_thread #14217

Merged

Conversation

junrushao
Copy link
Member

This PR introduces a syntactic sugar that combines T.env_thread and T.launch_thread.

Previously, an AttrStmt that specifies thread extent or virtual thread is required to be written in two steps:

bx = T.env_thread("blockIdx.x")  // creates an IterVar
with T.launch_thread(bx, 128):   // specify the iter domain
  ...

With this PR, now this behavior can be merged in a single line:

with T.launch_thread("blockIdx.x", 128) as bx:
  ...

@tvm-bot
Copy link
Collaborator

tvm-bot commented Mar 6, 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

@junrushao junrushao marked this pull request as ready for review March 6, 2023 21:02
@junrushao
Copy link
Member Author

CC: @spectrometerHBH @yzh119 @Hzfengsy

@junrushao junrushao force-pushed the feature/2023-03-05/sugar-launch-env-thread branch from 659fd9f to bbc39df Compare March 6, 2023 21:18
@Hzfengsy
Copy link
Member

Hzfengsy commented Mar 7, 2023

Could you please add a test case for it?

@junrushao
Copy link
Member Author

@Hzfengsy Existing roundtripping tests should suffice already. Do you prefer if I added a new one?

@Hzfengsy
Copy link
Member

Hzfengsy commented Mar 7, 2023

@Hzfengsy Existing roundtripping tests should suffice already. Do you prefer if I added a new one?

It would be great if we can have another explicit one :)

@junrushao junrushao force-pushed the feature/2023-03-05/sugar-launch-env-thread branch from bbc39df to 2cb8a76 Compare March 7, 2023 02:43
@junrushao
Copy link
Member Author

Got it. @Hzfengsy Just added a new testcase

This PR introduces a syntactic sugar that combines T.env_thread and
T.launch_thread.

Previously, an AttrStmt that specifies thread extent or virtual thread
is required to be written in two steps:

```python
bx = T.env_thread("blockIdx.x")  // creates an IterVar
with T.launch_thread(bx, 128):   // specify the iter domain
  ...
```

With this PR, now this behavior can be merged in a single line:

```python
with T.launch_thread("blockIdx.x", 128) as bx:
  ...
```
@junrushao junrushao force-pushed the feature/2023-03-05/sugar-launch-env-thread branch from 2cb8a76 to 7ca1009 Compare March 7, 2023 02:45
@junrushao junrushao merged commit be66a7e into apache:main Mar 7, 2023
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.

4 participants