[AUTO_SCHEDULER] Supported CSE (variable definitions) in feature extraction by tkonolige · Pull Request #10686 · apache/tvm

tkonolige · 2022-03-18T23:36:28Z

Add supported for LetStmts in feature extraction. A stack of variable definitions is maintained and added to the arithmetic analyzer at the appropriate points. The buffer access analysis now creates a new arithmetic analysis context per set of loops to avoid redefining variables which is unsafe in the presence of let statements.

@tqchen @junrushao1994 @merrymercy

…action Add supported for LetStmts in feature extraction. A stack of variable definitions is maintained and added to the arithmetic analyzer at the appropriate points. The buffer access analysis now creates a new arithmetic analysis context per set of loops to avoid redefining variables which is unsafe in the presence of let statements.

tqchen · 2022-03-19T17:26:12Z

-    return 1;
-  }
+// Return the maximum extent of a for loop
+int64_t GetLoopExtent(const ForNode* node, const Analyzer& ana) {


NOTE: this is a change of behavior that might worth checking, const_int bound may return INT_MAX for cases that it cannot analyze

I've added a check that return 1 if const_int_bound returns inf (it way of saying it cannot analyze).

This LGTM. although I do not expect regression in this case, would be good to run some tests to confirm, not a blocker to merge though.

tqchen · 2022-03-19T17:33:00Z



 def test_primfunc():
-    features = auto_scheduler.feature.named_features_from_primfunc(tir_matmul)


Would be good to keep the original test because that is the original intended usage, since the original function was designed for auto-scheduler with limited passes (to enable fast feature extraction). The test should cover the original intended use.

The additional support enables possible explorations for new use cases. Consider add another testcase test_primfunc_lowered_with_let to get the fully lowered version.

It would be helpful to construct a simplified TVM script version of the lowered version to see exactly where let bound may results in a change of values.

+1 for TQ's idea

I wrote the original test and its intended usage was for analysis after lowering.

I specifically did not want to construct a test from the explicitly lowered tvmscript so future large changes like CSE would cause the test to fail.

I see, explicitly constructed script would help to clarify the regression that the change intend to cover.

To rephrase what you are saying (correct if that was not what you meant):

In the future, if CSE changes its behavior to generate other cases, the feature extractor from an explicit tvm-script may not capture that changes (as new CSE will generate new IRs that may have other properties) -- the test won't fail but I guess you mean it fails to capture the possible problem in CSE changes.

This is indeed a fair pt. Although running things through is a bit like integration tests(that tests the combined effect of CSE, lowering and feature extraction).

The case of tvmscript inputs is closer to unit-test to test the intended behavior of the specific pass. Such unit-tests usually have a clear intend, and easier to debug when it triggers a regression. Although its coverage won't move as the integration tests do.

I still would suggest keep test_primfunc given the original intended behavior of the pass was limited. It does not hurt to have a case to cover the original intended case.

Adding another test_primfunc_lowered_with_let would give quite clear context, and provide some safety net when we start to explore the ability of using the extractor with CSE passes.

As for the unit-test case, it is merely a recommendation(as per above discussion on unit-test and integration tests)

To rephrase what you are saying (correct if that was not what you meant):

In the future, if CSE changes its behavior to generate other cases, the feature extractor from an explicit tvm-script may not capture that changes (as new CSE will generate new IRs that may have other properties) -- the test won't fail but I guess you mean it fails to capture the possible problem in CSE changes.

What I mean is that I want a test that ensures that changes to the lowering pipeline do not break feature extraction. CSE is an example of a test that broke it. I agree that it could be viewed as an integration test and I can move the test if you'd like.

Great, seems our interpretations aligns with each other. I agree that having a test that thread through lowered is useful if we bring in future usecases that applies feature extractor on lowered code.

To summarize the recommendation:

Keep the original test_primfunc(perhaps rename to test_prim_func_before_lowering to clairfy the intend)

This was the original intended behavior of the extractor and auto-scheduler only relied on this behavior.

Add a new test test_primfunc_lowered_with_let

I understand that there is an exploration effort to try out this and possibly followup features that depends on this. Although auto-scheduler do not yet rely on this behavior, so it would be good to keep it separate.

Like you said it is more like an integration tests(but i agree in this case it is OK to keep it here)

(Optional): add a unit-test testcase with explicitly constructed tvmscript that covers the lowered_let. The main rationale is that the intend of the unit-test case would be clear and it becomes easier to debug in facing regression, but I feel in this case it can be optional.

junrushao · 2022-03-19T17:39:40Z

-  }
+// Return the maximum extent of a for loop
+int64_t GetLoopExtent(const ForNode* node, const Analyzer& ana) {
+  return ana.const_int_bound(node->extent)->max_value;


The reason that we didn't use const-int-bound in feature extraction is that it becomes less helpful when encountering TIR of symbolic shapes.

auto scheduler doesn't work with symbolic shapes, so why would this be a problem? Seems like this is strictly an improvement over what we currently have.

junrushao · 2022-03-19T17:41:13Z



 def test_primfunc():
-    features = auto_scheduler.feature.named_features_from_primfunc(tir_matmul)


+1 for TQ's idea

tkonolige · 2022-03-28T16:43:06Z

@junrushao1994 Could you review again.

junrushao

Thanks!

junrushao · 2022-03-28T20:40:19Z

CC @tqchen for a second look

…ction (apache#10686) Add supported for LetStmts in feature extraction. A stack of variable definitions is maintained and added to the arithmetic analyzer at the appropriate points. The buffer access analysis now creates a new arithmetic analysis context per set of loops to avoid redefining variables which is unsafe in the presence of let statements.

tkonolige force-pushed the cse_feature_extraction branch from b82e065 to 189e1bd Compare March 18, 2022 23:39

tqchen assigned junrushao Mar 19, 2022

tqchen reviewed Mar 19, 2022

View reviewed changes

junrushao requested changes Mar 19, 2022

View reviewed changes

Tristan Konolige added 3 commits March 21, 2022 10:10

add check for analyzer failure

86248ff

split tests into lowered and unlowered

05039ab

cast to int in loop comparison

53f40b1

junrushao approved these changes Mar 28, 2022

View reviewed changes

tqchen approved these changes Mar 28, 2022

View reviewed changes

junrushao merged commit ab8e7c8 into apache:main Mar 28, 2022



		def test_primfunc():
		features = auto_scheduler.feature.named_features_from_primfunc(tir_matmul)

Conversation

tkonolige commented Mar 18, 2022

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tqchen Mar 21, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tqchen Mar 19, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tqchen Mar 21, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tkonolige Mar 21, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tkonolige commented Mar 28, 2022

Uh oh!

junrushao left a comment

Choose a reason for hiding this comment

Uh oh!

junrushao commented Mar 28, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

tqchen Mar 21, 2022 •

edited

Loading

tqchen Mar 19, 2022 •

edited

Loading

tqchen Mar 21, 2022 •

edited

Loading

tkonolige Mar 21, 2022 •

edited

Loading