-
Notifications
You must be signed in to change notification settings - Fork 160
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Refine colreduction fusion strategy of kStitch. #1257
Conversation
tao_compiler/mlir/disc/transforms/lhlo_legalize_roots_to_loops.cc
Outdated
Show resolved
Hide resolved
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM~ pls clear the tests and merge.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
why remove the pytorch pre version CI task?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The version of PyTorch of pre CI tasks is outdated and the task will fail, and the available versions are always changing.
@@ -1478,13 +1478,27 @@ bool BaseGpuFusionStrategy::tryFuse(ShapeAnalysis& shapeAnalysis, | |||
bool has_row_reduction = llvm::any_of(target.getOpList(), [](Operation* op) { | |||
return isRank2RowReduction(op); | |||
}); | |||
bool has_col_reduciton = llvm::any_of(target.getOpList(), [](Operation* op) { | |||
bool has_col_reduction = llvm::any_of(target.getOpList(), [](Operation* op) { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
maybe has_rank2_col_reduction
is more exactly.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done.
return false; | ||
} | ||
|
||
if (has_col_reduction) { |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
a mirror concern, does disc-fusion pass fused col-reduction ops with different shapes?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
A shape constraint is also set in initFusionPattern
later, here I will delete the same shape constraint.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM!
This PR will add more strict fusion rule for fusionType
kColReduction
: guarantee that the OPs to be fused are both Rank2ColReduction and have same shapes.