[Backend] Fix predicates for device assert inside reduction/scan region #5033
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
davidberard98
commented
Nov 1, 2024
| @@ -5917,7 +5917,7 @@ def test_side_effectful_reduction(device): | |||
| if device != "cuda": | |||
| pytest.skip() | |||
|
|
|||
| @triton.jit(debug=True) | |||
| @triton.jit | |||
There was a problem hiding this comment.
I think the debug=True needs to be added as a kwarg to the invocation of the triton kernel. Previously I wasn't seeing any asserts in the ttgir
There was a problem hiding this comment.
Thanks for spotting this, I've opened #5037 to fix it
Comment on lines
166
to
174
| // Predicate to ensure we don't read from invalid memory. | ||
| // definitions: | ||
| // "Lane": the strip of values that are being reduced along. | ||
| // relevant variables: | ||
| // interleave: for two consecutive elements in a lane, the difference | ||
| // between their thread ids is the interleave. | ||
| // numLanesToReduce: how many lanes we're reducing across. | ||
| // totalNumLanes: how many lanes exist in total. If the reduction | ||
| // skips some threads, totalNumLanes might not equal numLanesToReduce. |
There was a problem hiding this comment.
@peterbell10 is this accurate? tbh I didn't quite understand what scenario requires a predicate - I verified that this fixes my scenario, but I don't know if it regresses the scenario you were initially targeting.
peterbell10
reviewed
Nov 1, 2024
Comment on lines
178
to
180
| Value laneId = | ||
| urem(udiv(threadId, i32_val(interleave)), i32_val(totalNumLanes)); | ||
| Value lanePred = icmp_slt(laneId, i32_val(numLaneToReduce)); |
There was a problem hiding this comment.
The definition of lane is the position of a thread within its warp, so this is a bit confusing. Would it work to do this?
Suggested change
| Value laneId = | |
| urem(udiv(threadId, i32_val(interleave)), i32_val(totalNumLanes)); | |
| Value lanePred = icmp_slt(laneId, i32_val(numLaneToReduce)); | |
| Value laneId = urem(threadId, warpSize); | |
| Value lanePred = icmp_slt(laneId, i32_val(totalNumLanes * interleave)); |
There was a problem hiding this comment.
@peterbell10 thanks for the suggestion!
Instead I'm using
Value lanePred = icmp_slt(laneId, i32_val(numLaneToReduce * interleave));since I presume that the reason for predicating is due to the difference between numLaneToReduce vs. totalNumLanes?
davidberard98
force-pushed
the
fix-reduction-assert
branch
from
fe5c1fc
to
d430559
Compare
|
note: the other test_side_effectful_reduction and side_effectful_scan tests are failing after #5035, but somehow not failing on the test_side_effectful_reduction_2d test added by this PR. |
davidberard98
force-pushed
the
fix-reduction-assert
branch
from
eb32075
to
2ea66e9
Compare
davidberard98
force-pushed
the
fix-reduction-assert
branch
from
2614adc
to
2ea66e9
Compare
pytorchmergebot
pushed a commit
to pytorch/pytorch
that referenced
this pull request
Nov 1, 2024
In upstream triton, triton-lang/triton#4589 introduces overflow checks. However, overflow checks likely add some overhead, and have some correctness bugs at the moment (e.g. triton-lang/triton#5033). Let's set `sanitize_overflow=False` but keep `debug=True` so that we can keep using device_assert but without the additional asserts added by `sanitize_overflow`. Pull Request resolved: #139502 Approved by: https://github.com/bertmaher
davidberard98
force-pushed
the
fix-reduction-assert
branch
from
2ea66e9
to
40986be
Compare
peterbell10
force-pushed
the
fix-reduction-assert
branch
from
40986be
to
4e5ba83
Compare
Reductions have special handling for side effectful "combine ops" (e.g. "add" for a sum reduction). In the presence of side effects, a predicate is computed to determine whether a thread should participate in the reduction, to ensure that invalid/uninitialized data is not operated on. See [] for more details. Previously, the predicate logic was incorrect for 2D reductions. This PR fixes the logic and adds a python test.
davidberard98
force-pushed
the
fix-reduction-assert
branch
from
4e5ba83
to
1f3198c
Compare
rahulsingh-intel
pushed a commit
to rahulsingh-intel/pytorch
that referenced
this pull request
Nov 5, 2024
…9502) In upstream triton, triton-lang/triton#4589 introduces overflow checks. However, overflow checks likely add some overhead, and have some correctness bugs at the moment (e.g. triton-lang/triton#5033). Let's set `sanitize_overflow=False` but keep `debug=True` so that we can keep using device_assert but without the additional asserts added by `sanitize_overflow`. Pull Request resolved: pytorch#139502 Approved by: https://github.com/bertmaher
peterbell10
approved these changes
Nov 5, 2024
bertmaher
pushed a commit
that referenced
this pull request
Nov 5, 2024
…on (#5033) Reductions have special handling for side effectful "combine ops" (e.g. "add" for a sum reduction). In the presence of side effects, a predicate is computed to determine whether a thread should participate in the reduction, to ensure that invalid/uninitialized data is not operated on. See #4811 for more details. ~Previously, the predicate logic was incorrect for 2D reductions. This PR fixes the logic and adds a python test.~ Edit: after additional discussion with @peterbell10, we removed the lanePred logic. Here's our thinking on why this is valid: * lanePred info is computed based entirely on the blocked layout info and properties of the reduction * the blocked layout won't tell you which threads do or don't have uninitialized data Instead, it sounds like the motivation for #4811 is based on uninitialized values that can be indicated by the `pred` variable passed into `warpReduce()`.
antiagainst
added a commit
to antiagainst/triton
that referenced
this pull request
Nov 5, 2024
…can region (triton-lang#5033)" This reverts commit 732aee7.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Reductions have special handling for side effectful "combine ops" (e.g. "add" for a sum reduction). In the presence of side effects, a predicate is computed to determine whether a thread should participate in the reduction, to ensure that invalid/uninitialized data is not operated on. See #4811 for more details.
Previously, the predicate logic was incorrect for 2D reductions. This PR fixes the logic and adds a python test.Edit: after additional discussion with @peterbell10, we removed the lanePred logic. Here's our thinking on why this is valid:
Instead, it sounds like the motivation for #4811 is based on uninitialized values that can be indicated by the
predvariable passed intowarpReduce().