[TKW] Fix types, shapes and propagate resolved indexing #177
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
raikonenfnu
force-pushed
the
fixTypesUpstream
branch
3 times, most recently
from
25c6d23
to
e03e3ce
Compare
|
Currently broken on igemm because ExtractSlice used in iGEMM semantic needs to follow upstream. We need to implement an extract/indexing op for semantic required in local thread reduction |
raikonenfnu
force-pushed
the
fixTypesUpstream
branch
3 times, most recently
from
d3f6edb
to
ab70a52
Compare
raikonenfnu
changed the title
harsh-nod
reviewed
Oct 7, 2024
harsh-nod
reviewed
Oct 7, 2024
harsh-nod
reviewed
Oct 7, 2024
harsh-nod
reviewed
Oct 7, 2024
Signed-off-by: Stanley Winata <stanley.winata@amd.com>
raikonenfnu
force-pushed
the
fixTypesUpstream
branch
from
ab70a52
to
08d43e0
Compare
Signed-off-by: Stanley Winata <stanley.winata@amd.com>
harsh-nod
reviewed
Oct 17, 2024
harsh-nod
reviewed
Oct 17, 2024
harsh-nod
reviewed
Oct 17, 2024
harsh-nod
reviewed
Oct 17, 2024
harsh-nod
reviewed
Oct 17, 2024
iree/turbine/kernel/ops/wave_ops.py
Outdated
|
|
||
| # Typically only fastest dim has non-unit dim, | ||
| # but if all unit-dim get fastest/last one. | ||
| all_unit_dims = lambda index: all(x.size == 1 for x in index.values()) |
There was a problem hiding this comment.
Do we need this check? Can we merge this with the non_unit_dim check? if len(non_unit_dim) == 0, use fastest dim.
There was a problem hiding this comment.
done! I think :)
harsh-nod
requested changes
Oct 17, 2024
harsh-nod
approved these changes
Oct 17, 2024
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.
This PR add supports for doing accumulate on non-induction variables. For most part our stack already supports this, but we'd need to fix reduction symbolic shape which used to be just the input shape to input shape - reduction dim. Without this our thread shape analysis wouldn't be able to handle broadcasting of reduced op properly.
To support the above, we also introduce
Extractop in addition to the existingExtractSliceto fix the shape types. Specifically, ExtractSlice follow upstream's semantic where we just slice but do not reduce any dimensions. For the ReduceOp case, specifically the local reduction, we'd want it to have a reducing semantic on the fastest dimension.Additionally, we also add a propagation of our resolution for thread shape. This is helpful in the test we added which is a broadcast-sub followed by an exp2.