Skip to content
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

[AutoScheduler] Querying and sampling in task extraction #7571

Merged
merged 2 commits into from
Mar 4, 2021

Conversation

comaniac
Copy link
Contributor

@comaniac comaniac commented Mar 2, 2021

This PR deals with the case that we may want to extract tuning tasks from a JIT model that compiles each subgraph when first time executing the model graph. In this case, we need at least one valid schedule (for GPU) to make sure we can visit the whole graph and collect all tasks, but this may not be the case especially for the backward model.

The above puzzle forms a chicken egg problem -- we don't have valid schedules so we need to extract tasks to tune, but we need at least one valid schedule in order to extract tasks. This PR solves this puzzle by sampling one valid schedule during task extraction:

    env_sample = auto_scheduler.ApplyHistoryBestOrSample("temp.log", num_measure=2)
    env_tracing_task = auto_scheduler.relay_integration.TracingEnvironment(
        auto_scheduler.relay_integration.TracingMode.EXTRACT_COMPLEX_TASK_ONLY
    )
    with env_sample:
        with env_tracing_task:
            # Run the JIT model.

The idea is once we found that there's no valid schedule of the task being collected, we immediately sample a valid schedule and use it for the current subgraph. In this way, we guarantee the graph can be visited from beginning to the end.

cc @merrymercy

@comaniac comaniac changed the title [AutoScheduler] Query in task extraction [AutoScheduler] Querying and sampling in task extraction Mar 2, 2021
@comaniac comaniac merged commit 3f5f84d into apache:main Mar 4, 2021
@comaniac
Copy link
Contributor Author

comaniac commented Mar 4, 2021

Thanks @merrymercy @jcf94

@comaniac comaniac deleted the ansor_extract_n_sample branch March 4, 2021 00:09
trevor-m pushed a commit to trevor-m/tvm that referenced this pull request May 6, 2021
* [AutoScheduler] Query in task extraction

* trigger ci
trevor-m pushed a commit to neo-ai/tvm that referenced this pull request May 11, 2021
* [AutoScheduler] Query in task extraction

* trigger ci
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.

3 participants