From 58479912d26938b75cb6e6dea9c054d3ef93d120 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Wed, 8 Jun 2022 14:05:50 -0700 Subject: [PATCH 1/7] init commit --- python/tvm/auto_scheduler/measure.py | 41 ++++++++----------- python/tvm/autotvm/measure/measure_methods.py | 32 +++++++++++---- 2 files changed, 42 insertions(+), 31 deletions(-) diff --git a/python/tvm/auto_scheduler/measure.py b/python/tvm/auto_scheduler/measure.py index 2a4a03bbe8e7..939f5260e814 100644 --- a/python/tvm/auto_scheduler/measure.py +++ b/python/tvm/auto_scheduler/measure.py @@ -31,37 +31,28 @@ We implement these in python to utilize python's multiprocessing and error handling. """ +import logging +import multiprocessing import os -import time import shutil import tempfile -import multiprocessing -import logging +import time import tvm._ffi -from tvm.runtime import Object, module, ndarray +from tvm.autotvm.env import AutotvmGlobalScope, reset_global_scope +from tvm.contrib import ndk, tar +from tvm.contrib.popen_pool import PopenPoolExecutor, PopenWorker, StatusKind from tvm.driver import build_module from tvm.ir import transform -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.runtime import Object, module, ndarray from tvm.target import Target - from . import _ffi_api from .loop_state import StateObject -from .utils import ( - call_func_with_timeout, - check_remote, - get_const_tuple, - get_func_name, - make_traceback_info, - request_remote, -) -from .workload_registry import ( - serialize_workload_registry_entry, - deserialize_workload_registry_entry, -) +from .utils import (call_func_with_timeout, check_remote, get_const_tuple, + get_func_name, make_traceback_info, request_remote) +from .workload_registry import (deserialize_workload_registry_entry, + serialize_workload_registry_entry) # pylint: disable=invalid-name logger = logging.getLogger("auto_scheduler") @@ -221,7 +212,8 @@ def recover_measure_input(inp, rebuild_state=False): The fully recovered MeasureInput with all fields rebuilt. """ # pylint: disable=import-outside-toplevel - from .search_task import SearchTask # lazily import to avoid recursive dependency + from .search_task import \ + SearchTask # lazily import to avoid recursive dependency task = inp.task task.target, task.target_host = Target.canon_target_and_host(task.target, task.target_host) @@ -555,8 +547,8 @@ def __init__( device=0, ): # pylint: disable=import-outside-toplevel - from tvm.rpc.tracker import Tracker from tvm.rpc.server import Server + from tvm.rpc.tracker import Tracker self.tracker = Tracker(port=9000, port_end=10000, silent=True) device_key = "$local$device$%d" % self.tracker.port @@ -630,7 +622,7 @@ def _local_build_worker(inp_serialized, build_func, verbose): filename = os.path.join(dirname, "tmp_func." + build_func.output_format) try: - with transform.PassContext(): + with transform.PassContext().current(): func = build_module.build(sch, args, target=task.target) func.export_library(filename, build_func) # pylint: disable=broad-except @@ -841,7 +833,8 @@ def prepare_runner_args(inp, build_res): """ # pylint: disable=import-outside-toplevel - from .search_task import get_task_input_buffer # lazily import to avoid recursive dependency + from .search_task import \ + get_task_input_buffer # lazily import to avoid recursive dependency task_input_names = inp.task.task_input_names tensor_input_map = prepare_input_map(build_res.args) diff --git a/python/tvm/autotvm/measure/measure_methods.py b/python/tvm/autotvm/measure/measure_methods.py index f582bd1974aa..e76566229fa7 100644 --- a/python/tvm/autotvm/measure/measure_methods.py +++ b/python/tvm/autotvm/measure/measure_methods.py @@ -31,9 +31,9 @@ import time import traceback import typing +import warnings from collections import namedtuple from random import getrandbits -import warnings import tvm._ffi import tvm.ir.transform @@ -505,10 +505,6 @@ def _build_func_common(measure_input, runtime=None, check_gpu=None, build_option if not config.valid(): raise InstantiationError(config.errors) - opts = build_option or {} - if check_gpu: # Add verify pass to filter out invalid configs in advance. - opts["tir.add_lower_pass"] = [(2, gpu_verify_pass(**check_gpu))] - # if target is vta, we need to use vta build if ( hasattr(measure_input.target, "device_name") @@ -519,7 +515,28 @@ def _build_func_common(measure_input, runtime=None, check_gpu=None, build_option func = vta.build(s, args, target_host=task.target_host) else: - with tvm.ir.transform.PassContext(config=opts): + current_pass_context: tvm.ir.transform.PassContext = ( + tvm.ir.transform.PassContext.current() + ) + current_config = dict(current_pass_context.config) + if build_option is not None: + current_config.update(build_option) + + if "tir.add_lower_pass" in current_config: + current_add_lower_pass = list(current_config["tir.add_lower_pass"]) + else: + current_add_lower_pass = [] + if check_gpu: + current_add_lower_pass.append((2, gpu_verify_pass(**check_gpu))) + current_config["tir.add_lower_pass"] = current_add_lower_pass + + with tvm.ir.transform.PassContext( + opt_level=current_pass_context.opt_level, + required_pass=current_pass_context.required_pass, + disabled_pass=current_pass_context.disabled_pass, + instruments=current_pass_context.instruments, + config=current_config, + ): func = build(s, args, target_host=task.target_host, runtime=runtime) return func, tuple((get_const_tuple(x.shape), x.dtype) for x in args) @@ -572,7 +589,8 @@ def __call__(self, measure_input, tmp_dir, **kwargs): if self.build_func.output_format == ".model-library-format": # Late import to preserve autoTVM with USE_MICRO OFF try: - from tvm import micro # pylint: disable=import-outside-toplevel + from tvm import \ + micro # pylint: disable=import-outside-toplevel except ImportError: raise ImportError("Requires USE_MICRO") micro.export_model_library_format(func, filename) From fd03d32d013737bea60b592dc1bdbb0b458ae473 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Wed, 8 Jun 2022 16:19:51 -0700 Subject: [PATCH 2/7] lint --- python/tvm/auto_scheduler/measure.py | 22 ++++++++++++------- python/tvm/autotvm/measure/measure_methods.py | 3 +-- 2 files changed, 15 insertions(+), 10 deletions(-) diff --git a/python/tvm/auto_scheduler/measure.py b/python/tvm/auto_scheduler/measure.py index 939f5260e814..6f331499b042 100644 --- a/python/tvm/auto_scheduler/measure.py +++ b/python/tvm/auto_scheduler/measure.py @@ -49,10 +49,18 @@ from . import _ffi_api from .loop_state import StateObject -from .utils import (call_func_with_timeout, check_remote, get_const_tuple, - get_func_name, make_traceback_info, request_remote) -from .workload_registry import (deserialize_workload_registry_entry, - serialize_workload_registry_entry) +from .utils import ( + call_func_with_timeout, + check_remote, + get_const_tuple, + get_func_name, + make_traceback_info, + request_remote, +) +from .workload_registry import ( + deserialize_workload_registry_entry, + serialize_workload_registry_entry, +) # pylint: disable=invalid-name logger = logging.getLogger("auto_scheduler") @@ -212,8 +220,7 @@ def recover_measure_input(inp, rebuild_state=False): The fully recovered MeasureInput with all fields rebuilt. """ # pylint: disable=import-outside-toplevel - from .search_task import \ - SearchTask # lazily import to avoid recursive dependency + from .search_task import SearchTask # lazily import to avoid recursive dependency task = inp.task task.target, task.target_host = Target.canon_target_and_host(task.target, task.target_host) @@ -833,8 +840,7 @@ def prepare_runner_args(inp, build_res): """ # pylint: disable=import-outside-toplevel - from .search_task import \ - get_task_input_buffer # lazily import to avoid recursive dependency + from .search_task import get_task_input_buffer # lazily import to avoid recursive dependency task_input_names = inp.task.task_input_names tensor_input_map = prepare_input_map(build_res.args) diff --git a/python/tvm/autotvm/measure/measure_methods.py b/python/tvm/autotvm/measure/measure_methods.py index e76566229fa7..8fc0da89c4c6 100644 --- a/python/tvm/autotvm/measure/measure_methods.py +++ b/python/tvm/autotvm/measure/measure_methods.py @@ -589,8 +589,7 @@ def __call__(self, measure_input, tmp_dir, **kwargs): if self.build_func.output_format == ".model-library-format": # Late import to preserve autoTVM with USE_MICRO OFF try: - from tvm import \ - micro # pylint: disable=import-outside-toplevel + from tvm import micro # pylint: disable=import-outside-toplevel except ImportError: raise ImportError("Requires USE_MICRO") micro.export_model_library_format(func, filename) From e84bcf92dd10158ce20c543abcab9b8ea4048355 Mon Sep 17 00:00:00 2001 From: Andrew Zhao Luo Date: Thu, 9 Jun 2022 09:29:11 -0700 Subject: [PATCH 3/7] empty commit From 1eae767082f4ec2daf2bc49003344ad119739ac9 Mon Sep 17 00:00:00 2001 From: Andrew Zhao Luo Date: Thu, 9 Jun 2022 18:17:53 -0700 Subject: [PATCH 4/7] test results --- tests/python/integration/test_tuning.py | 123 +++++++++++++++++++++++- 1 file changed, 121 insertions(+), 2 deletions(-) diff --git a/tests/python/integration/test_tuning.py b/tests/python/integration/test_tuning.py index 03f38aa9cc9e..d1f6b7120887 100644 --- a/tests/python/integration/test_tuning.py +++ b/tests/python/integration/test_tuning.py @@ -25,11 +25,14 @@ import pytest import tvm +from tvm.autotvm.measure import measure_methods +from tvm.ir.transform import PassContext import tvm.relay import tvm.testing from tvm import autotvm, te from tvm.autotvm.tuner import RandomTuner from tvm.target import Target +from tvm.contrib import tar def setup_module(): @@ -153,7 +156,6 @@ def run_test_with_all_multiprocessing(func, *args, **kwargs): mp.set_start_method(old_start_method, force=True) -@pytest.mark.xfail(strict=False, reason="See https://github.com/apache/tvm/issues/10489") @tvm.testing.parametrize_targets("cuda", "opencl") def test_tuning_gpu(target, dev): def runner(target, dev): @@ -174,12 +176,129 @@ def runner(target, dev): assert len(results) == 20 - successful_results = [r for r in results if r.error_no == autotvm.MeasureErrorNo.NO_ERROR] + successful_results = [ + r + for r in results + if r.error_no == autotvm.MeasureErrorNo.NO_ERROR + # Autotvm can filter some records before building if we know they won't work ahead of time. + or r.error_no == autotvm.MeasureErrorNo.INSTANTIATION_ERROR + ] assert len(successful_results) > 0, f"No successful tuning runs: {results!r}" run_test_with_all_multiprocessing(runner, target, dev) +from tvm.ir.instrument import pass_instrument + + +@tvm.testing.parametrize_targets("cuda", "opencl") +def test_tuning_gpu_inherits_pass_context(target, dev): + """Autotvm tuner inherits PassContexts but also adds a gpu verification pass by default. + + Test that using PassContext inherits passes properly but also runs gpu verification pass. + """ + from tvm.tir.analysis import _ffi_api as _analysis_ffi_api + + @pass_instrument + class PassInstrumentChecker: + """Pass Instrument that simply sees if it's been run.""" + + def __init__(self): + self.has_been_run = False + + def run_after_pass(self, mod, info): + self.has_been_run = True + + class GPUVerifyPassMocked: + """Context manager that mocks tir.analysis.verify_gpu_code meant + to verify the pass has been run. This is done by patching the ffi func handles.""" + + FFI_FUNC_HANDLE = "tir.analysis.verify_gpu_code" + FUNC_NAME = "verify_gpu_code" + + def __init__(self) -> None: + self.old_impl = tvm._ffi.get_global_func(self.FFI_FUNC) + self.has_been_run = False + + def gpu_verify_pass_mocked(self): + """Get the replacement for the gpu verification pass.""" + + def _gpu_verify_pass_mocked(*args, **kwargs): + self.has_been_run = True + return self.old_impl(*args, **kwargs) + + return _gpu_verify_pass_mocked + + def __enter__(self): + tvm._ffi.register_func( + self.FFI_FUNC_HANDLE, self.gpu_verify_pass_mocked(), override=True + ) + + # Also overwrite the python bindings + setattr( + _analysis_ffi_api, self.FUNC_NAME, tvm._ffi.get_global_func(self.FFI_FUNC_HANDLE) + ) + + def __exit__(self, *args, **kwargs): + # Restore FFI status back to normal + tvm._ffi.register_func(self.FFI_FUNC_HANDLE, self.old_impl, override=True) + setattr(_analysis_ffi_api, self.FUNC_NAME, self.old_impl) + + class OverwrittenBuildFunc(measure_methods._WrappedBuildFunc): + """BuildFunc that mocks and patches as necessary to test proper passes are run.""" + + def __call__(self, measure_input, tmp_dir, **kwargs): + instrument = PassInstrumentChecker() + mocked_pass_checker = GPUVerifyPassMocked() + with mocked_pass_checker: + with PassContext(instruments=[instrument]): + regular_result = super().__call__(measure_input, tmp_dir, **kwargs) + + # Check instrument has been run, meaning context was inherited by builder + assert instrument.has_been_run + + # But also check the gpu verification pass has been run + # (which was not in the inherited ctx) + assert mocked_pass_checker.has_been_run + + return regular_result + + class MockedLocalBuilder(measure_methods.LocalBuilder): + """As measure_methods.LocalBuilder but overwrites the PassContext for testing.""" + + def __init__( + self, + timeout=10, + n_parallel=None, + build_kwargs=None, + build_func="default", + do_fork=False, + runtime=None, + ): + super().__init__(timeout, n_parallel, build_kwargs, build_func, do_fork, runtime) + self.build_func = OverwrittenBuildFunc(tar.tar, runtime) + + def runner(target, dev): + task, target = get_sample_task(target, None) + logging.info("task config space: %s", task.config_space) + + # Note: we use the MockedLocalBuilder here instead of autotvm.LocalBuilder() + measure_option = autotvm.measure_option(MockedLocalBuilder(), autotvm.LocalRunner()) + + results = [] + + tuner = RandomTuner(task) + tuner.tune( + n_trial=1, + measure_option=measure_option, + callbacks=(lambda _tuner, _inputs, rs: results.extend(rs),), + ) + + assert len(results) == 1 + + run_test_with_all_multiprocessing(runner, target, dev) + + def test_tuning_cpu(): def runner(): ir_mod = tvm.parser.fromtext( From e7b0cd49ba7803931d5550dba6ac3eafd0753f5c Mon Sep 17 00:00:00 2001 From: Andrew Zhao Luo Date: Thu, 9 Jun 2022 18:19:54 -0700 Subject: [PATCH 5/7] reset progress --- tests/python/integration/test_tuning.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/tests/python/integration/test_tuning.py b/tests/python/integration/test_tuning.py index d1f6b7120887..275a36fdf3c5 100644 --- a/tests/python/integration/test_tuning.py +++ b/tests/python/integration/test_tuning.py @@ -156,6 +156,7 @@ def run_test_with_all_multiprocessing(func, *args, **kwargs): mp.set_start_method(old_start_method, force=True) +@pytest.mark.xfail(strict=False, reason="See https://github.com/apache/tvm/issues/10489") @tvm.testing.parametrize_targets("cuda", "opencl") def test_tuning_gpu(target, dev): def runner(target, dev): @@ -176,13 +177,7 @@ def runner(target, dev): assert len(results) == 20 - successful_results = [ - r - for r in results - if r.error_no == autotvm.MeasureErrorNo.NO_ERROR - # Autotvm can filter some records before building if we know they won't work ahead of time. - or r.error_no == autotvm.MeasureErrorNo.INSTANTIATION_ERROR - ] + successful_results = [r for r in results if r.error_no == autotvm.MeasureErrorNo.NO_ERROR] assert len(successful_results) > 0, f"No successful tuning runs: {results!r}" run_test_with_all_multiprocessing(runner, target, dev) From 235ab6247b6deb3e62fa003d0bd91581d80c8b6b Mon Sep 17 00:00:00 2001 From: Andrew Zhao Luo Date: Thu, 9 Jun 2022 18:20:09 -0700 Subject: [PATCH 6/7] lint --- tests/python/integration/test_tuning.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/python/integration/test_tuning.py b/tests/python/integration/test_tuning.py index 275a36fdf3c5..bfc086395bfd 100644 --- a/tests/python/integration/test_tuning.py +++ b/tests/python/integration/test_tuning.py @@ -33,6 +33,7 @@ from tvm.autotvm.tuner import RandomTuner from tvm.target import Target from tvm.contrib import tar +from tvm.ir.instrument import pass_instrument def setup_module(): @@ -183,9 +184,6 @@ def runner(target, dev): run_test_with_all_multiprocessing(runner, target, dev) -from tvm.ir.instrument import pass_instrument - - @tvm.testing.parametrize_targets("cuda", "opencl") def test_tuning_gpu_inherits_pass_context(target, dev): """Autotvm tuner inherits PassContexts but also adds a gpu verification pass by default. From 16a347056b535dfbbdf3dd2a7eb15d45025e38b2 Mon Sep 17 00:00:00 2001 From: Andrew Luo Date: Fri, 10 Jun 2022 08:25:00 -0700 Subject: [PATCH 7/7] fix --- tests/python/integration/test_tuning.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/python/integration/test_tuning.py b/tests/python/integration/test_tuning.py index bfc086395bfd..a3dca33e71ee 100644 --- a/tests/python/integration/test_tuning.py +++ b/tests/python/integration/test_tuning.py @@ -25,15 +25,15 @@ import pytest import tvm -from tvm.autotvm.measure import measure_methods -from tvm.ir.transform import PassContext import tvm.relay import tvm.testing from tvm import autotvm, te +from tvm.autotvm.measure import measure_methods from tvm.autotvm.tuner import RandomTuner -from tvm.target import Target from tvm.contrib import tar from tvm.ir.instrument import pass_instrument +from tvm.ir.transform import PassContext +from tvm.target import Target def setup_module(): @@ -210,7 +210,7 @@ class GPUVerifyPassMocked: FUNC_NAME = "verify_gpu_code" def __init__(self) -> None: - self.old_impl = tvm._ffi.get_global_func(self.FFI_FUNC) + self.old_impl = tvm._ffi.get_global_func(self.FFI_FUNC_HANDLE) self.has_been_run = False def gpu_verify_pass_mocked(self):