-
Notifications
You must be signed in to change notification settings - Fork 3.7k
[AUTO_SCHEDULER] Add feature extraction directly from PrimFunc #10455
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
Conversation
|
cc @comaniac |
|
@tkonolige can you kick another job now that eric PR has landed. |
Allow users to directly extract features from a PrimFunc. Extracted features can be used to get an estimate of flops, memory load size, or arithmetic intensity from a PrimFunc. Also fix feature extraction to correctly measure the number of arithmetic operations width vector datatypes.
7f6deb3 to
9970eb9
Compare
src/auto_scheduler/feature.cc
Outdated
| } | ||
|
|
||
| void VisitStmt_(const StoreNode* node) final { | ||
| MathOpCounter math_op_counter; |
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.
Do we hit this function now? Note that Load / Store node are deprecated after eric PR cc @Lunderberg
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.
No, we don't hit this. I've removed it.
| + T.load("float32x128", A.data, T.ramp(k, 128, 128), T.broadcast(True, 128)) | ||
| * T.broadcast(T.load("float32", B.data, y * 128 + k), 128), | ||
| T.broadcast(True, 128), | ||
| ) |
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.
cc @Lunderberg @vinx13 @junrushao1994 on recommendation on using T.load / store in new tests.
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.
I rewrote this to use buffer store.
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.
👍
comaniac
left a comment
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.
Otherwise LGTM.
Also cc @merrymercy @junrushao1994 to check the changes in feature extractions.
…e#10455) * [AUTO_SCHEDULER] Add feature extraction directly from PrimFunc Allow users to directly extract features from a PrimFunc. Extracted features can be used to get an estimate of flops, memory load size, or arithmetic intensity from a PrimFunc. Also fix feature extraction to correctly measure the number of arithmetic operations width vector datatypes. * fix param name * log scale in cc instead of python * rename functions, remove load/store * forgot rename in tests * forgot to commit rename
Allow users to directly extract features from a PrimFunc. Extracted features can be used to get an estimate of flops, memory load size, or arithmetic intensity of a PrimFunc. I've improved support to handle PrimFuncs created from a full lowering pass in addition to the limited lowering that auto_scheduler uses.
Also fix feature extraction to correctly measure the number of arithmetic operations width vector datatypes.
This change will not change features used by auto_scheduler's cost models.
@junrushao1994 @merrymercy @masahi