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

[inductor] Stop using x + tl.zeros(...) in generated triton #100163

Closed
wants to merge 8 commits into from

Conversation

peterbell10
Copy link
Collaborator

@peterbell10 peterbell10 commented Apr 27, 2023

Stack from ghstack (oldest at bottom):

For reductions, this changes the accumulator

_tmp2 = tl.zeros([XBLOCK, RBLOCK], tl.int8) + -128

to

_tmp2 = tl.full([XBLOCK, RBLOCK], -128, tl.int32)

which is equivalent since addition does type promotion from int8 to int32

For constant indexing, this changes

tl.store(in_out_ptr0 + (0 + tl.zeros([XBLOCK, 1], tl.int32)), tmp4, None)

to

tl.store(in_out_ptr0 + (tl.full([XBLOCK, 1], 0, tl.int32)), tmp4, None)

For variable indexing, this changes

tl.store(out_ptr0 + (0 + tl.zeros([XBLOCK], tl.int32)), tmp1, None)

to

tl.store(out_ptr0 + (tl.broadcast_to(x0, [XBLOCK])), tmp1, None)

cc @soumith @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @desertfire

For reductions, this changes the accumulator
```python
_tmp2 = tl.zeros([XBLOCK, RBLOCK], tl.int8) + -128
```
to
```python
_tmp2 = tl.full([XBLOCK, RBLOCK], -128, tl.int32)
```
which is equivalent since addition does type promotion from `int8` to `int32`

For constant indexing, this changes
```python
tl.store(in_out_ptr0 + (0 + tl.zeros([XBLOCK, 1], tl.int32)), tmp4, None)
```
to
```python
tl.store(in_out_ptr0 + (tl.full([XBLOCK, 1], 0, tl.int32)), tmp4, None)
```

For variable indexing, this changes
```python
tl.store(out_ptr0 + (0 + tl.zeros([XBLOCK], tl.int32)), tmp1, None)
```
to
```python
tl.store(out_ptr0 + (tl.broadcast_to(x0, [XBLOCK])), tmp1, None)
```

[ghstack-poisoned]
@pytorch-bot
Copy link

pytorch-bot bot commented Apr 27, 2023

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/100163

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

✅ No Failures

As of commit fa372b3:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

…ton"

For reductions, this changes the accumulator
```python
_tmp2 = tl.zeros([XBLOCK, RBLOCK], tl.int8) + -128
```
to
```python
_tmp2 = tl.full([XBLOCK, RBLOCK], -128, tl.int32)
```
which is equivalent since addition does type promotion from `int8` to `int32`

For constant indexing, this changes
```python
tl.store(in_out_ptr0 + (0 + tl.zeros([XBLOCK, 1], tl.int32)), tmp4, None)
```
to
```python
tl.store(in_out_ptr0 + (tl.full([XBLOCK, 1], 0, tl.int32)), tmp4, None)
```

For variable indexing, this changes
```python
tl.store(out_ptr0 + (0 + tl.zeros([XBLOCK], tl.int32)), tmp1, None)
```
to
```python
tl.store(out_ptr0 + (tl.broadcast_to(x0, [XBLOCK])), tmp1, None)
```

cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire

[ghstack-poisoned]
peterbell10 added a commit to peterbell10/pytorch that referenced this pull request Apr 27, 2023
For reductions, this changes the accumulator
```python
_tmp2 = tl.zeros([XBLOCK, RBLOCK], tl.int8) + -128
```
to
```python
_tmp2 = tl.full([XBLOCK, RBLOCK], -128, tl.int32)
```
which is equivalent since addition does type promotion from `int8` to `int32`

For constant indexing, this changes
```python
tl.store(in_out_ptr0 + (0 + tl.zeros([XBLOCK, 1], tl.int32)), tmp4, None)
```
to
```python
tl.store(in_out_ptr0 + (tl.full([XBLOCK, 1], 0, tl.int32)), tmp4, None)
```

For variable indexing, this changes
```python
tl.store(out_ptr0 + (0 + tl.zeros([XBLOCK], tl.int32)), tmp1, None)
```
to
```python
tl.store(out_ptr0 + (tl.broadcast_to(x0, [XBLOCK])), tmp1, None)
```

ghstack-source-id: 32e07c41ea197d555aca29771fee20fc2068e583
Pull Request resolved: pytorch#100163
…ton"

For reductions, this changes the accumulator
```python
_tmp2 = tl.zeros([XBLOCK, RBLOCK], tl.int8) + -128
```
to
```python
_tmp2 = tl.full([XBLOCK, RBLOCK], -128, tl.int32)
```
which is equivalent since addition does type promotion from `int8` to `int32`

For constant indexing, this changes
```python
tl.store(in_out_ptr0 + (0 + tl.zeros([XBLOCK, 1], tl.int32)), tmp4, None)
```
to
```python
tl.store(in_out_ptr0 + (tl.full([XBLOCK, 1], 0, tl.int32)), tmp4, None)
```

For variable indexing, this changes
```python
tl.store(out_ptr0 + (0 + tl.zeros([XBLOCK], tl.int32)), tmp1, None)
```
to
```python
tl.store(out_ptr0 + (tl.broadcast_to(x0, [XBLOCK])), tmp1, None)
```

cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire

[ghstack-poisoned]
peterbell10 added a commit that referenced this pull request Apr 27, 2023
For reductions, this changes the accumulator
```python
_tmp2 = tl.zeros([XBLOCK, RBLOCK], tl.int8) + -128
```
to
```python
_tmp2 = tl.full([XBLOCK, RBLOCK], -128, tl.int32)
```
which is equivalent since addition does type promotion from `int8` to `int32`

For constant indexing, this changes
```python
tl.store(in_out_ptr0 + (0 + tl.zeros([XBLOCK, 1], tl.int32)), tmp4, None)
```
to
```python
tl.store(in_out_ptr0 + (tl.full([XBLOCK, 1], 0, tl.int32)), tmp4, None)
```

For variable indexing, this changes
```python
tl.store(out_ptr0 + (0 + tl.zeros([XBLOCK], tl.int32)), tmp1, None)
```
to
```python
tl.store(out_ptr0 + (tl.broadcast_to(x0, [XBLOCK])), tmp1, None)
```

ghstack-source-id: 933b241a9a3e72a84e09e1cc2569e42de83e504e
Pull Request resolved: #100163
@peterbell10 peterbell10 marked this pull request as ready for review April 27, 2023 17:19
@peterbell10
Copy link
Collaborator Author

@pytorchbot merge

@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Apr 27, 2023
@pytorchmergebot
Copy link
Collaborator

Merge failed

Reason: This PR needs a label
If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Details for Dev Infra team Raised by workflow job

@peterbell10 peterbell10 added the topic: not user facing topic category label Apr 27, 2023
@peterbell10
Copy link
Collaborator Author

@pytorchbot merge

@pytorchmergebot
Copy link
Collaborator

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

@pytorchmergebot
Copy link
Collaborator

Merge failed

Reason: 1 mandatory check(s) failed. The first few are:

Dig deeper by viewing the failures on hud

Details for Dev Infra team Raised by workflow job

Failing merge rule: Core Maintainers

…x + tl.zeros(...)` in generated triton"

For reductions, this changes the accumulator
```python
_tmp2 = tl.zeros([XBLOCK, RBLOCK], tl.int8) + -128
```
to
```python
_tmp2 = tl.full([XBLOCK, RBLOCK], -128, tl.int32)
```
which is equivalent since addition does type promotion from `int8` to `int32`

For constant indexing, this changes
```python
tl.store(in_out_ptr0 + (0 + tl.zeros([XBLOCK, 1], tl.int32)), tmp4, None)
```
to
```python
tl.store(in_out_ptr0 + (tl.full([XBLOCK, 1], 0, tl.int32)), tmp4, None)
```

For variable indexing, this changes
```python
tl.store(out_ptr0 + (0 + tl.zeros([XBLOCK], tl.int32)), tmp1, None)
```
to
```python
tl.store(out_ptr0 + (tl.broadcast_to(x0, [XBLOCK])), tmp1, None)
```

cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire

[ghstack-poisoned]
# tl.reduce doesn't support tl.int1
accumulator = f"{accumulator}.to(tl.int8)"
result_type = triton_compute_type(dtype)
self.suffix.writeline(f"{result_var} = {final_reduction(accumulator)}.to({result_type})")
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

This is only really used for aten.any. It changes the final reduction of a non-persistent reduction from

tmp5 = triton_helpers.max(_tmp5, 1)[:, None]
tmp5 = triton_helpers.max(_tmp5.to(tl.int8), 1)[:, None].to(tl.int1)

Copy link
Collaborator

Choose a reason for hiding this comment

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

This is good to put as a comment in the code!

peterbell10 added a commit to peterbell10/pytorch that referenced this pull request Apr 28, 2023
For reductions, this changes the accumulator
```python
_tmp2 = tl.zeros([XBLOCK, RBLOCK], tl.int8) + -128
```
to
```python
_tmp2 = tl.full([XBLOCK, RBLOCK], -128, tl.int32)
```
which is equivalent since addition does type promotion from `int8` to `int32`

For constant indexing, this changes
```python
tl.store(in_out_ptr0 + (0 + tl.zeros([XBLOCK, 1], tl.int32)), tmp4, None)
```
to
```python
tl.store(in_out_ptr0 + (tl.full([XBLOCK, 1], 0, tl.int32)), tmp4, None)
```

For variable indexing, this changes
```python
tl.store(out_ptr0 + (0 + tl.zeros([XBLOCK], tl.int32)), tmp1, None)
```
to
```python
tl.store(out_ptr0 + (tl.broadcast_to(x0, [XBLOCK])), tmp1, None)
```

ghstack-source-id: d92be321d693cfc9a0faeafd37dfa423b93b4d64
Pull Request resolved: pytorch#100163
…ton"

For reductions, this changes the accumulator
```python
_tmp2 = tl.zeros([XBLOCK, RBLOCK], tl.int8) + -128
```
to
```python
_tmp2 = tl.full([XBLOCK, RBLOCK], -128, tl.int32)
```
which is equivalent since addition does type promotion from `int8` to `int32`

For constant indexing, this changes
```python
tl.store(in_out_ptr0 + (0 + tl.zeros([XBLOCK, 1], tl.int32)), tmp4, None)
```
to
```python
tl.store(in_out_ptr0 + (tl.full([XBLOCK, 1], 0, tl.int32)), tmp4, None)
```

For variable indexing, this changes
```python
tl.store(out_ptr0 + (0 + tl.zeros([XBLOCK], tl.int32)), tmp1, None)
```
to
```python
tl.store(out_ptr0 + (tl.broadcast_to(x0, [XBLOCK])), tmp1, None)
```

cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire

[ghstack-poisoned]
@peterbell10
Copy link
Collaborator Author

@pytorchbot merge -f "Doc build failure is pre-existing"

@pytorchmergebot
Copy link
Collaborator

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

@ezyang ezyang mentioned this pull request May 3, 2023
ezyang added a commit that referenced this pull request May 3, 2023
Revert "[inductor] Stop using `x + tl.zeros(...)` in generated triton (#100163)"

This reverts commit 5b98910.

Revert "[inductor] Fix argmin/max with duplicate values (#99920)"

This reverts commit 659dcc5.

Revert "[inductor] Fix nan-handling of max and min reductions (#99881)"

This reverts commit f9c3fcd.

[ghstack-poisoned]
ezyang added a commit that referenced this pull request May 3, 2023
Revert "[inductor] Stop using `x + tl.zeros(...)` in generated triton (#100163)"

This reverts commit 5b98910.

Revert "[inductor] Fix argmin/max with duplicate values (#99920)"

This reverts commit 659dcc5.

Revert "[inductor] Fix nan-handling of max and min reductions (#99881)"

This reverts commit f9c3fcd.

ghstack-source-id: 85531baedfb245e48512be97c0ed90eba1685664
Pull Request resolved: #100517
peterbell10 added a commit to peterbell10/pytorch that referenced this pull request May 3, 2023
For reductions, this changes the accumulator
```python
_tmp2 = tl.zeros([XBLOCK, RBLOCK], tl.int8) + -128
```
to
```python
_tmp2 = tl.full([XBLOCK, RBLOCK], -128, tl.int32)
```
which is equivalent since addition does type promotion from `int8` to `int32`

For constant indexing, this changes
```python
tl.store(in_out_ptr0 + (0 + tl.zeros([XBLOCK, 1], tl.int32)), tmp4, None)
```
to
```python
tl.store(in_out_ptr0 + (tl.full([XBLOCK, 1], 0, tl.int32)), tmp4, None)
```

For variable indexing, this changes
```python
tl.store(out_ptr0 + (0 + tl.zeros([XBLOCK], tl.int32)), tmp1, None)
```
to
```python
tl.store(out_ptr0 + (tl.broadcast_to(x0, [XBLOCK])), tmp1, None)
```

ghstack-source-id: 3a7b180d4fa58c14afb407ddd453cc29cf819074
Pull Request resolved: pytorch#100163
peterbell10 added a commit to peterbell10/pytorch that referenced this pull request May 3, 2023
For reductions, this changes the accumulator
```python
_tmp2 = tl.zeros([XBLOCK, RBLOCK], tl.int8) + -128
```
to
```python
_tmp2 = tl.full([XBLOCK, RBLOCK], -128, tl.int32)
```
which is equivalent since addition does type promotion from `int8` to `int32`

For constant indexing, this changes
```python
tl.store(in_out_ptr0 + (0 + tl.zeros([XBLOCK, 1], tl.int32)), tmp4, None)
```
to
```python
tl.store(in_out_ptr0 + (tl.full([XBLOCK, 1], 0, tl.int32)), tmp4, None)
```

For variable indexing, this changes
```python
tl.store(out_ptr0 + (0 + tl.zeros([XBLOCK], tl.int32)), tmp1, None)
```
to
```python
tl.store(out_ptr0 + (tl.broadcast_to(x0, [XBLOCK])), tmp1, None)
```

ghstack-source-id: 3a7b180d4fa58c14afb407ddd453cc29cf819074
Pull Request resolved: pytorch#100163
@peterbell10 peterbell10 reopened this May 3, 2023
…ton"

For reductions, this changes the accumulator
```python
_tmp2 = tl.zeros([XBLOCK, RBLOCK], tl.int8) + -128
```
to
```python
_tmp2 = tl.full([XBLOCK, RBLOCK], -128, tl.int32)
```
which is equivalent since addition does type promotion from `int8` to `int32`

For constant indexing, this changes
```python
tl.store(in_out_ptr0 + (0 + tl.zeros([XBLOCK, 1], tl.int32)), tmp4, None)
```
to
```python
tl.store(in_out_ptr0 + (tl.full([XBLOCK, 1], 0, tl.int32)), tmp4, None)
```

For variable indexing, this changes
```python
tl.store(out_ptr0 + (0 + tl.zeros([XBLOCK], tl.int32)), tmp1, None)
```
to
```python
tl.store(out_ptr0 + (tl.broadcast_to(x0, [XBLOCK])), tmp1, None)
```

cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire

[ghstack-poisoned]
…ton"

For reductions, this changes the accumulator
```python
_tmp2 = tl.zeros([XBLOCK, RBLOCK], tl.int8) + -128
```
to
```python
_tmp2 = tl.full([XBLOCK, RBLOCK], -128, tl.int32)
```
which is equivalent since addition does type promotion from `int8` to `int32`

For constant indexing, this changes
```python
tl.store(in_out_ptr0 + (0 + tl.zeros([XBLOCK, 1], tl.int32)), tmp4, None)
```
to
```python
tl.store(in_out_ptr0 + (tl.full([XBLOCK, 1], 0, tl.int32)), tmp4, None)
```

For variable indexing, this changes
```python
tl.store(out_ptr0 + (0 + tl.zeros([XBLOCK], tl.int32)), tmp1, None)
```
to
```python
tl.store(out_ptr0 + (tl.broadcast_to(x0, [XBLOCK])), tmp1, None)
```

cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire

[ghstack-poisoned]
…ton"

For reductions, this changes the accumulator
```python
_tmp2 = tl.zeros([XBLOCK, RBLOCK], tl.int8) + -128
```
to
```python
_tmp2 = tl.full([XBLOCK, RBLOCK], -128, tl.int32)
```
which is equivalent since addition does type promotion from `int8` to `int32`

For constant indexing, this changes
```python
tl.store(in_out_ptr0 + (0 + tl.zeros([XBLOCK, 1], tl.int32)), tmp4, None)
```
to
```python
tl.store(in_out_ptr0 + (tl.full([XBLOCK, 1], 0, tl.int32)), tmp4, None)
```

For variable indexing, this changes
```python
tl.store(out_ptr0 + (0 + tl.zeros([XBLOCK], tl.int32)), tmp1, None)
```
to
```python
tl.store(out_ptr0 + (tl.broadcast_to(x0, [XBLOCK])), tmp1, None)
```

cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire

[ghstack-poisoned]
@facebook-github-bot facebook-github-bot deleted the gh/peterbell10/543/head branch June 8, 2023 18:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

4 participants