Skip to content

[AUTO_SCHEDULER] Add feature extraction directly from PrimFunc#10455

Merged
masahi merged 6 commits into
apache:mainfrom
tkonolige:features_from_stores
Mar 10, 2022
Merged

[AUTO_SCHEDULER] Add feature extraction directly from PrimFunc#10455
masahi merged 6 commits into
apache:mainfrom
tkonolige:features_from_stores

Conversation

@tkonolige

Copy link
Copy Markdown
Contributor

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

@masahi

masahi commented Mar 3, 2022

Copy link
Copy Markdown
Member

cc @comaniac

@masahi

masahi commented Mar 9, 2022

Copy link
Copy Markdown
Member

@tkonolige can you kick another job now that eric PR has landed.

Tristan Konolige added 3 commits March 8, 2022 16:21
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.
@tkonolige tkonolige force-pushed the features_from_stores branch from 7f6deb3 to 9970eb9 Compare March 9, 2022 00:22
Comment thread src/auto_scheduler/feature.cc Outdated
}

void VisitStmt_(const StoreNode* node) final {
MathOpCounter math_op_counter;

@masahi masahi Mar 9, 2022

Copy link
Copy Markdown
Member

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

Copy link
Copy Markdown
Contributor Author

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),
)

Copy link
Copy Markdown
Member

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.

Copy link
Copy Markdown
Contributor Author

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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍

@comaniac comaniac left a comment

Copy link
Copy Markdown
Contributor

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.

Comment thread python/tvm/auto_scheduler/feature.py Outdated
Comment thread python/tvm/auto_scheduler/feature.py Outdated
@masahi masahi merged commit 7e49f53 into apache:main Mar 10, 2022
pfk-beta pushed a commit to pfk-beta/tvm that referenced this pull request Apr 11, 2022
…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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants