[AMD] Added instr.sched guards for the FA-like kernels #5163
+159
−82
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.
Extended AMDGPU instruction scheduling for the Flash Attention like kernels. The introduced source code changes adds
sched.barriers
at the beginning and at the end of eachscf.For
op (calledguards
) which contains at least 2tt.Dot
,tt.reduce
and at least onemath::Exp2Op
ops. The guards prevent moves of instructions from basic block adjacent to the bodies forfor-loops
. According to test results, it results in increase performance for the FA kernels due to a reduction of VGPRs spilling.I am not making a trivial change, such as fixing a typo in a comment.
I have written a PR description following these
rules.
I have run
pre-commit run --from-ref origin/main --to-ref HEAD
.Select one of the following.
/test
forlit
tests/unittest
for C++ tests/python/test
for end-to-end testsSelect one of the following.
lit
tests.lit
tests I have added follow these best practices,including the "tests should be minimal" section. (Usually running Python code
and using the instructions it generates is not minimal.)