AMD specific scheduling pass for TTGIR instructions #483
Conversation
oplavsic
force-pushed
the
scheduling_pass_checkpoint
branch
4 times, most recently
from
3b8581f to
171a67e
Compare
oplavsic
force-pushed
the
scheduling_pass_checkpoint
branch
2 times, most recently
from
3ae1238 to
2276fac
Compare
oplavsic
changed the title
oplavsic
requested review from
jayfurmanek,
zhanglx13,
alefimov-amd and
jtang10
python/triton/compiler/compiler.py
Outdated
| @@ -164,8 +164,11 @@ def optimize_ttgir(mod, num_stages, num_warps, num_ctas, target, cluster_info, e | |||
| pm.add_tritongpu_remove_layout_conversions_pass() | |||
| pm.add_tritongpu_decompose_conversions_pass() | |||
| pm.add_tritongpu_ws_fixup_missing_attrs_pass() | |||
| if num_stages != 0: | |||
| if is_hip() and num_stages != 0: | |||
There was a problem hiding this comment.
The check against num_stages makes it not work with the stream pipeliner. Is it intentional?
There was a problem hiding this comment.
This scheduler is designed to do scheduling for FA kernel with dot slicing enabled. So it's kind of replacing the stream-pipeline pass.
| if (!operandsSorted.empty() && | ||
| operandsSorted[operandsSorted.size() - 1].getDefiningOp()) { | ||
|
|
||
| moveAfter(op, operandsSorted[operandsSorted.size() - 1].getDefiningOp()); |
There was a problem hiding this comment.
What if last operand in a list is a block argument?
I afraid it could crash of spoil IR
There was a problem hiding this comment.
Wouldn't in that case getDefiningOp return false, so this code wouldn't even step into "if" statement?
| return; | ||
| } | ||
|
|
||
| std::sort(operandsSorted.begin(), operandsSorted.end(), |
There was a problem hiding this comment.
I am kind of suspicious about this sort...
std::sort assumes you have fully ordered objects, which is not true for dominance.
There was a problem hiding this comment.
Hmm, let's see. Here is some documentation:
/// Return true if operation A properly dominates operation B, i.e. if A and B
/// are in the same block and A properly dominates B within the block, or if
/// the block that contains A properly dominates the block that contains B. In
/// an SSACFG region, Operation A dominates Operation B in the same block if A
/// preceeds B. In a Graph region, all operations in a block dominate all
/// other operations in the same block.
In SSACFG regions (and here we are dealing with these), blocks are ordered in
containing region (https://mlir.llvm.org/docs/LangRef/#high-level-structure), So if two operands are in different blocks, since they are ordered, one will dominate the other, so instruction in the dominant block will dominate instruction in the other block. If they are in the same block, operations are ordered, and by definition (from the comment above), the one that precedes the other is the dominant.
I could be missing something since I'm a bit rusty on MLIR SSA specifics. Can you think of some specific case that could cause problems?
| @@ -0,0 +1,397 @@ | |||
| #include "mlir/Analysis/SliceAnalysis.h" | |||
There was a problem hiding this comment.
nitpic, add license:
| #include "mlir/Analysis/SliceAnalysis.h" | |
| /* | |
| * Copyright (c) 2023, Advanced Micro Devices, Inc. All rights reserved. | |
| * | |
| * Permission is hereby granted, free of charge, to any person obtaining | |
| * a copy of this software and associated documentation files | |
| * (the "Software"), to deal in the Software without restriction, | |
| * including without limitation the rights to use, copy, modify, merge, | |
| * publish, distribute, sublicense, and/or sell copies of the Software, | |
| * and to permit persons to whom the Software is furnished to do so, | |
| * subject to the following conditions: | |
| * | |
| * The above copyright notice and this permission notice shall be | |
| * included in all copies or substantial portions of the Software. | |
| * | |
| * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, | |
| * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF | |
| * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. | |
| * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY | |
| * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, | |
| * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE | |
| * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. | |
| */ | |
| #include "mlir/Analysis/SliceAnalysis.h" |
There was a problem hiding this comment.
Done! Btw, do we need this in every amd specific .cpp file?
oplavsic
force-pushed
the
scheduling_pass_checkpoint
branch
from
66214bd to
0320a62
Compare
This PR introduces AMD specific scheduling pass. Main purpose it has for now is to hoist Q tensor out of the loop in FA fwd pass, and to schedule instructions produced by dot slicing pass.