[lang] Keep custom reduce and scan callback bodies self-contained - #103
Closed
azazhu wants to merge 9 commits into
Closed
[lang] Keep custom reduce and scan callback bodies self-contained#103azazhu wants to merge 9 commits into
azazhu wants to merge 9 commits into
Conversation
Signed-off-by: feiwen zhu <mzhu@nvidia.com>
Signed-off-by: feiwen zhu <mzhu@nvidia.com>
Signed-off-by: feiwen zhu <mzhu@nvidia.com>
Signed-off-by: feiwen zhu <mzhu@nvidia.com>
Signed-off-by: feiwen zhu <mzhu@nvidia.com>
Return the source-level names of captured values from the pre-pass and hand them to the legalization pass, instead of stashing them on the IRContext and importing a private helper from _compile.py. Key them by the value that survives Assign elimination. Also document why reduce/scan bodies are LICM barriers and why capture legalization must run after constant materialization, and name the callback's parameters explicitly in the runtime-capture error. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: feiwen zhu <mzhu@nvidia.com>
Check, right after a callback body is built, that every value inside it is a 0-d tile. This gives one consistent diagnostic for non-scalar constants whether they are captured from the enclosing scope or created inside the callback, and it replaces the capture-only shape check in the legalization pass. Document the restriction in the ct.reduce()/ct.scan() docstrings. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: feiwen zhu <mzhu@nvidia.com>
The kernels checking that a whole reduce/scan can be hoisted stored tiles
whose rank did not match the array; give them matching ranks and launch
them so the hoisted result is verified. Cover scan in the LICM barrier
test as well. The scan constant-capture test used identity 0 together with
'(a + b) % 2', which is not an identity for arbitrary inputs; restrict the
inputs to {0, 1} so the reference does not depend on where the identity is
combined.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: feiwen zhu <mzhu@nvidia.com>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: feiwen zhu <mzhu@nvidia.com>
azazhu
changed the title
Collaborator
|
@azazhu i will close this PR since this issue should be fixed in tileiras. |
This file contains hidden or 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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
Classification: bug fix with a documented behavior change.
Description
Custom
ct.reduce()/ct.scan()callbacks are pure combine functions over their block arguments. Until now the body emitted for a callback could depend on the enclosing scope: loop-invariant code motion hoisted constants out of the body, and constants captured from the kernel were referenced from outside the body. This change keeps callback bodies self-contained, so they rely on nothing but their block arguments and constants materialized inside the body.Testing
python -m pytest -q test/test_reduction.py test/test_scan.py test/test_code_motion.py450 passed, 72 skipped, 4 xfailedpython -m flake8on all changed Python filesgit diff --check main..HEADChecklist