[Codegen][PCF] Add convertForallToGenericNest transform #23332
Conversation
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
Add a transform that converts scf.forall operations into multi-level pcf.generic nests. This is required in cases where a single scf.forall mapping type needs to map to multiple scopes. The immediate case this arises from is converting thread mapped scf.forall ops to combined subgroup + lane scopes. In this case, we can't simply convert to pcf.loop because of the way that automatic redistribution of pcf.loop's lowering works. We need to redistribute the iterations of the scf.forall to *all* workers across both scopes. This means that if we make one of the two scopes (subgroup or lane) a pcf.loop this fails to create the correct IR structure. If we make it loop over subgroups, this fails to predicate the lanes, and if we make it loop over lanes, we end up with more thread divergence than normal (normally only a single subgroup may exhibit divergence). The transform creates an outer loop over workers and an inner scf.forall over the per-worker iteration range. For multi-dimensional cases, affine.linearize_index and affine.delinearize_index are used to flatten/unflatten indices appropriately. Additionally adds a new method getNativeNumProcessorIds to ScopeAttrInterface. This is needed to query for the number of ids to generate.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
Add a transform that converts scf.forall operations into multi-level pcf.generic nests. This is required in cases where a single scf.forall mapping type needs to map to multiple scopes. The immediate case this arises from is converting thread mapped scf.forall ops to combined subgroup + lane scopes. In this case, we can't simply convert to pcf.loop because of the way that automatic redistribution of pcf.loop's lowering works. We need to redistribute the iterations of the scf.forall to all workers across both scopes. This means that if we make one of the two scopes (subgroup or lane) a pcf.loop this fails to create the correct IR structure. If we make it loop over subgroups, this fails to predicate the lanes, and if we make it loop over lanes, we end up with more thread divergence than normal (normally only a single subgroup may exhibit divergence).
The transform creates an outer loop over workers and an inner scf.forall over the per-worker iteration range. For multi-dimensional cases, affine.linearize_index and affine.delinearize_index are used to flatten/unflatten indices appropriately.
Additionally adds a new method getNativeNumProcessorIds to ScopeAttrInterface. This is needed to query for the number of ids to generate.