This is an automated email from the ASF dual-hosted git repository.
syfeng pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new d0650bad66 [Bugfix][MetaSchedule] Filter out dynamic extents (#11747)
d0650bad66 is described below
commit d0650bad66d0ff89a01347537021bc442a98c223
Author: Junru Shao <[email protected]>
AuthorDate: Thu Jun 16 04:19:48 2022 -0700
[Bugfix][MetaSchedule] Filter out dynamic extents (#11747)
Previously only static shape computation is allowed in our tuning
system. However, one special case is overlooked: the reduction iter vars
could still have dynamic iteration domains which depend on other data
parallel vars. This PR rules out this case by carefully checking all the
loop extents during task extraction.
Related issue: https://github.com/apache/tvm/issues/11746.
---
python/tvm/meta_schedule/runner/rpc_runner.py | 11 +++++++----
src/meta_schedule/extracted_task.cc | 14 +++++++++++++-
tests/python/unittest/test_meta_schedule_integration.py | 9 +++++++++
3 files changed, 29 insertions(+), 5 deletions(-)
diff --git a/python/tvm/meta_schedule/runner/rpc_runner.py
b/python/tvm/meta_schedule/runner/rpc_runner.py
index 16e422cc60..9ff2489f8e 100644
--- a/python/tvm/meta_schedule/runner/rpc_runner.py
+++ b/python/tvm/meta_schedule/runner/rpc_runner.py
@@ -15,8 +15,8 @@
# specific language governing permissions and limitations
# under the License.
"""RPC Runner"""
-import logging
import concurrent.futures
+import logging
import os.path as osp
from contextlib import contextmanager
from typing import Callable, List, Optional, Union
@@ -26,6 +26,7 @@ from tvm.rpc import RPCSession
from tvm.runtime import Device, Module
from ..utils import (
+ cpu_count,
derived_object,
get_global_func_on_rpc_session,
get_global_func_with_default_on_worker,
@@ -242,7 +243,7 @@ class RPCRunner(PyRunner):
f_alloc_argument: Union[T_ALLOC_ARGUMENT, str, None] = None,
f_run_evaluator: Union[T_RUN_EVALUATOR, str, None] = None,
f_cleanup: Union[T_CLEANUP, str, None] = None,
- max_workers: int = 1,
+ max_workers: Optional[int] = 1,
initializer: Optional[Callable[[], None]] = None,
) -> None:
"""Constructor
@@ -267,8 +268,8 @@ class RPCRunner(PyRunner):
The function name to run the evaluator or the function itself.
f_cleanup: Union[T_CLEANUP, str, None]
The function name to cleanup the session or the function itself.
- max_workers: int = 1
- The maximum number of connections. Defaults to 1.
+ max_workers: Optional[int] = None
+ The maximum number of connections. Defaults to number of logical
CPU cores.
initializer: Optional[Callable[[], None]]
The initializer function.
"""
@@ -282,6 +283,8 @@ class RPCRunner(PyRunner):
self.f_alloc_argument = f_alloc_argument
self.f_run_evaluator = f_run_evaluator
self.f_cleanup = f_cleanup
+ if max_workers is None:
+ max_workers = cpu_count()
logger.info("RPCRunner: max_workers = %d", max_workers)
self.pool = PopenPoolExecutor(
max_workers=max_workers,
diff --git a/src/meta_schedule/extracted_task.cc
b/src/meta_schedule/extracted_task.cc
index abd7235acb..358f56efab 100644
--- a/src/meta_schedule/extracted_task.cc
+++ b/src/meta_schedule/extracted_task.cc
@@ -72,7 +72,19 @@ Optional<tir::PrimFunc> DefaultTaskFilterImpl(const
Array<te::Tensor>& args, boo
return NullOpt;
}
}
- return te::CreatePrimFunc(args);
+ PrimFunc func = te::CreatePrimFunc(args);
+ bool dynamic_loop_extent = false;
+ PostOrderVisit(func->body, [&dynamic_loop_extent](const ObjectRef& obj) ->
void {
+ if (const auto* loop = obj.as<tir::ForNode>()) {
+ if (!loop->extent->IsInstance<IntImmNode>()) {
+ dynamic_loop_extent = true;
+ }
+ }
+ });
+ if (dynamic_loop_extent) {
+ return NullOpt;
+ }
+ return func;
}
Optional<tir::PrimFunc> DefaultTaskFilter(const Array<te::Tensor>& args) {
diff --git a/tests/python/unittest/test_meta_schedule_integration.py
b/tests/python/unittest/test_meta_schedule_integration.py
index 6d5016cd81..4868640ade 100644
--- a/tests/python/unittest/test_meta_schedule_integration.py
+++ b/tests/python/unittest/test_meta_schedule_integration.py
@@ -19,6 +19,7 @@ import numpy as np
import pytest
import tvm
import tvm.testing
+from tvm import IRModule
from tvm import meta_schedule as ms
from tvm import relay, te, tir
from tvm.meta_schedule.testing.relay_workload import get_network
@@ -60,6 +61,14 @@ def test_meta_schedule_apply_history_best_no_current():
assert ms.ApplyHistoryBest.current() is None
+def test_meta_schedule_dynamic_loop_extent():
+ a = relay.var("a", shape=(1, 8, 8, 512), dtype="float32")
+ b = relay.nn.adaptive_avg_pool2d(a, (7, 7), "NHWC")
+ mod = IRModule({"main": relay.Function([a], b)})
+ extracted_tasks = ms.extract_task_from_relay(mod, target="llvm", params={})
+ assert not extracted_tasks
+
+
@requires_torch
def test_meta_schedule_integration_extract_from_resnet():
mod, params, _ = get_network(name="resnet_18", input_shape=[1, 3, 224,
224])
