-
Notifications
You must be signed in to change notification settings - Fork 608
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[GPU] Support multiple contraction dims in MmaSchedules (#18720)
This adds support for multiple M, N, and K dims in problems when deducing a GPUMMASchedule. The new heuristic is similar to the old one, but works on pairs of M and N dims. For example: ``` tensor<M1xM0xK1xK0> * tensor<N1xN0xK1xK0> -> tensor<M1xN1xM0xN0> ``` This will try to distribute the seeded tile counts to `M0` and `N0` (first attempting to distribute evenly, and then distributing to N followed by N), and then distribute the residual counts to `M1` and `N1`. The K tile counts will be partitioned to `K0` first, and then the residual tile counts will be partitioned to `K1`. This PR also updates the config selection logic for the TileAndFuse pipeline to make use of the multiple contraction dimensions in mma schedules. --------- Signed-off-by: Max Dawkins <max.dawkins@gmail.com>
- Loading branch information
Showing
8 changed files
with
492 additions
and
232 deletions.
There are no files selected for viewing
361 changes: 239 additions & 122 deletions
361
compiler/src/iree/compiler/Codegen/Common/GPU/GPUHeuristics.cpp
Large diffs are not rendered by default.
Oops, something went wrong.
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
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
Oops, something went wrong.