diff --git a/python/tvm/auto_scheduler/measure.py b/python/tvm/auto_scheduler/measure.py index b746dbf96f43..72c1126b5035 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 AutotvmGlobalScope, 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,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) + 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 ddd510ce55ed..ae07a852413c 100644 --- a/python/tvm/autotvm/env.py +++ b/python/tvm/autotvm/env.py @@ -30,3 +30,10 @@ 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 + AutotvmGlobalScope.current = global_scope