From d629dbbf40e50ca79eb11a5d9424bb123ed6d70d Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Thu, 2 Sep 2021 15:43:45 -0400 Subject: [PATCH 1/3] [AutoScheduler] Propogate global autotvm state to PopenPool workers --- python/tvm/auto_scheduler/measure.py | 3 ++- python/tvm/autotvm/env.py | 6 ++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/python/tvm/auto_scheduler/measure.py b/python/tvm/auto_scheduler/measure.py index b746dbf96f43..138dfbce70f1 100644 --- a/python/tvm/auto_scheduler/measure.py +++ b/python/tvm/auto_scheduler/measure.py @@ -43,6 +43,7 @@ from tvm.driver import build_module from tvm.ir import transform from tvm.autotvm.measure.measure_methods import set_cuda_target_arch +from tvm.autotvm.env import GLOBAL_SCOPE, reset_global_scope from tvm.contrib import tar, ndk from tvm.contrib.popen_pool import PopenWorker, PopenPoolExecutor, StatusKind from tvm.target import Target @@ -692,7 +693,7 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo res : List[BuildResult] The build results of these MeasureInputs. """ - executor = PopenPoolExecutor(n_parallel, timeout) + executor = PopenPoolExecutor(n_parallel, timeout, reset_global_scope, (GLOBAL_SCOPE,)) tuple_res = executor.map_with_error_catching( local_build_worker, [ diff --git a/python/tvm/autotvm/env.py b/python/tvm/autotvm/env.py index ddd510ce55ed..ee7b170739a3 100644 --- a/python/tvm/autotvm/env.py +++ b/python/tvm/autotvm/env.py @@ -30,3 +30,9 @@ def __init__(self): GLOBAL_SCOPE = AutotvmGlobalScope() + + +def reset_global_scope(global_scope): + """Reset global autotvm state. This is needed to initialize PopenPool workers.""" + global GLOBAL_SCOPE + GLOBAL_SCOPE = global_scope From 8384f409bff2fee97595b409a69116a9768197e6 Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Thu, 2 Sep 2021 17:09:44 -0400 Subject: [PATCH 2/3] Fix --- python/tvm/auto_scheduler/measure.py | 4 ++-- python/tvm/autotvm/env.py | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/python/tvm/auto_scheduler/measure.py b/python/tvm/auto_scheduler/measure.py index 138dfbce70f1..c1df9450901e 100644 --- a/python/tvm/auto_scheduler/measure.py +++ b/python/tvm/auto_scheduler/measure.py @@ -43,7 +43,7 @@ from tvm.driver import build_module from tvm.ir import transform from tvm.autotvm.measure.measure_methods import set_cuda_target_arch -from tvm.autotvm.env import GLOBAL_SCOPE, reset_global_scope +from tvm.autotvm.env import AutotvmGlobalScope, reset_global_scope from tvm.contrib import tar, ndk from tvm.contrib.popen_pool import PopenWorker, PopenPoolExecutor, StatusKind from tvm.target import Target @@ -693,7 +693,7 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo res : List[BuildResult] The build results of these MeasureInputs. """ - executor = PopenPoolExecutor(n_parallel, timeout, reset_global_scope, (GLOBAL_SCOPE,)) + executor = PopenPoolExecutor(n_parallel, timeout, reset_global_scope, (AutotvmGlobalScope.current,)) tuple_res = executor.map_with_error_catching( local_build_worker, [ diff --git a/python/tvm/autotvm/env.py b/python/tvm/autotvm/env.py index ee7b170739a3..ae07a852413c 100644 --- a/python/tvm/autotvm/env.py +++ b/python/tvm/autotvm/env.py @@ -36,3 +36,4 @@ def reset_global_scope(global_scope): """Reset global autotvm state. This is needed to initialize PopenPool workers.""" global GLOBAL_SCOPE GLOBAL_SCOPE = global_scope + AutotvmGlobalScope.current = global_scope From 79a1aaf693feddf1d9835779d7eb587ac0eb21df Mon Sep 17 00:00:00 2001 From: Wuwei Lin Date: Thu, 2 Sep 2021 17:27:31 -0400 Subject: [PATCH 3/3] lint --- python/tvm/auto_scheduler/measure.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/tvm/auto_scheduler/measure.py b/python/tvm/auto_scheduler/measure.py index c1df9450901e..72c1126b5035 100644 --- a/python/tvm/auto_scheduler/measure.py +++ b/python/tvm/auto_scheduler/measure.py @@ -693,7 +693,9 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo res : List[BuildResult] The build results of these MeasureInputs. """ - executor = PopenPoolExecutor(n_parallel, timeout, reset_global_scope, (AutotvmGlobalScope.current,)) + executor = PopenPoolExecutor( + n_parallel, timeout, reset_global_scope, (AutotvmGlobalScope.current,) + ) tuple_res = executor.map_with_error_catching( local_build_worker, [