[inductor] Fix ops.scan for non-commutative operators #126633
Conversation
`tl.associative_scan` supports non-commutative combine functions but `tl.reduce` doesn't. This effects non-persistent scans, where we use the reduction from the previous loop iterations as the base for future iterations. Here I work around this by taking the last element of the scan output and using that as the reduced value. This is possibly by using a trick where we create a mask that is 1 at the desired element and 0 elsewhere, then sum over it. ghstack-source-id: 98756771652cefc61e7367dadf795f7e69d24402 Pull Request resolved: #126633
`tl.associative_scan` supports non-commutative combine functions but `tl.reduce` doesn't. This effects non-persistent scans, where we use the reduction from the previous loop iterations as the base for future iterations. Here I work around this by taking the last element of the scan output and using that as the reduced value. This is possibly by using a trick where we create a mask that is 1 at the desired element and 0 elsewhere, then sum over it. ghstack-source-id: 23a6f1b7717633dcf38743aad9d5b6bc937dc935 Pull Request resolved: #126633
There was a problem hiding this comment.
Oh well. We should eventually fix the non-commutativity of tl.reduce, should be quite straightforward. With that fixed, we could pick up the last element reducing with
def last(a, b):
return bOf course, the real fix would be to expose 1-element slices over a dimension in triton.
|
@pytorchbot merge |
Merge failedReason: This PR needs a If not, please add the To add a label, you can comment to pytorchbot, for example For more information, see Details for Dev Infra teamRaised by workflow job |
|
@pytorchbot merge |
Merge startedYour 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 |
|
|
|
Can't merge closed PR #126633 |
|
This newly-added test is failing ROCm. Will attempt to skip it for now. |
Stack from ghstack (oldest at bottom):
tl.associative_scansupports non-commutative combine functions buttl.reducedoesn't. This effects non-persistent scans, where we use the reduction from
the previous loop iterations as the base for future iterations.
Here I work around this by taking the last element of the scan output and using
that as the reduced value. This is done using a trick where we create a
mask that is 1 at the desired element and 0 elsewhere, then sum over it.
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang