diff --git a/python/tvm/auto_scheduler/measure.py b/python/tvm/auto_scheduler/measure.py index 8a8b92201d15..9f592550cda8 100644 --- a/python/tvm/auto_scheduler/measure.py +++ b/python/tvm/auto_scheduler/measure.py @@ -51,19 +51,51 @@ from .loop_state import StateObject from .utils import ( get_const_tuple, - NoDaemonPool, call_func_with_timeout, request_remote, check_remote, ) +from .compute_dag import ComputeDAG +from .search_task import SearchTask +from .workload_registry import workload_func_name, get_workload_func # The maximum length of error message MAX_ERROR_MSG_LEN = 512 -# We use fork and a global variable to copy arguments between processes. -# This can avoid expensive serialization of TVM IR when using multiprocessing.Pool -GLOBAL_BUILD_ARGUMENTS = None -GLOBAL_RUN_ARGUMENTS = None + +def recover_measure_input(inp, rebuild_state=False): + """ + Recover a deserialized MeasureInput by rebuilding the missing fields. + 1. Rebuid the compute_dag in inp.task + 2. (Optional) Rebuild the stages in inp.state + + Parameters + ---------- + inp: MeasureInput + The deserialized MeasureInput + rebuild_state: bool = False + Whether rebuild the stages in MeasureInput.State + + Returns + ------- + new_input: MeasureInput + The fully recovered MeasureInput with all fields rebuilt. + """ + task = inp.task + new_task = SearchTask( + ComputeDAG(task.workload_key), + task.workload_key, + task.target, + task.target_host, + task.hardware_params, + ) + + if rebuild_state: + new_state = new_task.compute_dag.infer_bound_from_state(inp.state) + else: + new_state = inp.state + + return MeasureInput(new_task, new_state) @tvm._ffi.register_object("auto_scheduler.MeasureCallback") @@ -87,6 +119,31 @@ def __init__(self, task, state): state = state if isinstance(state, StateObject) else state.state_object self.__init_handle_by_constructor__(_ffi_api.MeasureInput, task, state) + def serialize(self): + """Custom serialization to workaround MeasureInput not exposing all its + members to the TVM ffi interface. + + Note that we do not implement __getstate__ as it does not seem to work + with initialization of the workload registry (maybe because of + initialization order?). + """ + serialize = tvm.get_global_func("auto_scheduler.SerializeMeasureInput", True) + assert serialize + # We serialize the workload function so that it can be used on the deserialized side. + return { + "measureinput": serialize(self), + "name": workload_func_name(self.task.workload_key), + "func": get_workload_func(self.task), + } + + @staticmethod + def deserialize(state): + deserialize = tvm.get_global_func("auto_scheduler.DeserializeMeasureInput", True) + assert deserialize + tvm.auto_scheduler.workload_registry.WORKLOAD_FUNC_REGISTRY[state["name"]] = state["func"] + x = deserialize(state["measureinput"]) + return recover_measure_input(x) + @tvm._ffi.register_object("auto_scheduler.BuildResult") class BuildResult(Object): @@ -486,29 +543,63 @@ def make_error_msg(): return error_msg -def local_build_worker(index): +def _timed_func(inp_serialized, build_func, verbose): + tic = time.time() + inp = MeasureInput.deserialize(inp_serialized) + task = inp.task + + error_no = MeasureErrorNo.NO_ERROR + error_msg = None + args = [] + + try: + sch, args = task.compute_dag.apply_steps_from_state(inp.state, layout_rewrite=True) + # pylint: disable=broad-except + except Exception: + error_no = MeasureErrorNo.INSTANTIATION_ERROR + error_msg = make_error_msg() + + if error_no == 0: + dirname = tempfile.mkdtemp() + filename = os.path.join(dirname, "tmp_func." + build_func.output_format) + + try: + # TODO(merrymercy): Port the unroll pass. + with transform.PassContext(): + func = build_module.build( + sch, args, target=task.target, target_host=task.target_host + ) + func.export_library(filename, build_func) + # pylint: disable=broad-except + except Exception: + error_no = MeasureErrorNo.COMPILE_HOST + error_msg = make_error_msg() + else: + filename = "" + + if verbose >= 1: + if error_no == MeasureErrorNo.NO_ERROR: + print(".", end="") + else: + print(".E", end="") # Build error + return filename, args, error_no, error_msg, time.time() - tic + + +def local_build_worker(args): """ Build function of LocalBuilder to be ran in the Builder thread pool. Parameters ---------- - index : int - The MeasureInput index to be processed by the current Builder thread. + args: Tuple[MeasureInput, str, int, int] + inputs, build-func, time, verbose args passed to local_builder_build Returns ------- res : BuildResult The build result of this Builder thread. """ - global GLOBAL_BUILD_ARGUMENTS - - # We use fork and a global variable to copy arguments between processes. - # This can avoid expensive serialization of TVM IR when using multiprocessing.Pool - if not GLOBAL_BUILD_ARGUMENTS: - raise ValueError("GLOBAL_BUILD_ARGUMENTS not found") - measure_inputs, build_func, timeout, verbose = GLOBAL_BUILD_ARGUMENTS - assert isinstance(build_func, str) - + inp, build_func, timeout, verbose = args if build_func == "default": build_func = tar.tar elif build_func == "ndk": @@ -516,48 +607,7 @@ def local_build_worker(index): else: raise ValueError("Invalid build_func" + build_func) - def timed_func(): - tic = time.time() - inp = measure_inputs[index] - task = inp.task - - error_no = MeasureErrorNo.NO_ERROR - error_msg = None - args = [] - - try: - sch, args = task.compute_dag.apply_steps_from_state(inp.state, layout_rewrite=True) - # pylint: disable=broad-except - except Exception: - error_no = MeasureErrorNo.INSTANTIATION_ERROR - error_msg = make_error_msg() - - if error_no == 0: - dirname = tempfile.mkdtemp() - filename = os.path.join(dirname, "tmp_func." + build_func.output_format) - - try: - # TODO(merrymercy): Port the unroll pass. - with transform.PassContext(): - func = build_module.build( - sch, args, target=task.target, target_host=task.target_host - ) - func.export_library(filename, build_func) - # pylint: disable=broad-except - except Exception: - error_no = MeasureErrorNo.COMPILE_HOST - error_msg = make_error_msg() - else: - filename = "" - - if verbose >= 1: - if error_no == MeasureErrorNo.NO_ERROR: - print(".", end="") - else: - print(".E", end="") # Build error - return filename, args, error_no, error_msg, time.time() - tic - - res = call_func_with_timeout(timeout, timed_func) + res = call_func_with_timeout(timeout, _timed_func, args=(inp, build_func, verbose)) if isinstance(res, TimeoutError): if verbose >= 1: print(".T", end="") # Build timeout @@ -590,14 +640,20 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo res : List[BuildResult] The build results of these MeasureInputs. """ - # We use fork and a global variable to copy arguments between processes. - # This can avoid expensive serialization of TVM IR when using multiprocessing.Pool - global GLOBAL_BUILD_ARGUMENTS - - GLOBAL_BUILD_ARGUMENTS = (inputs, build_func, timeout, verbose) - - pool = NoDaemonPool(n_parallel) - tuple_res = pool.map(local_build_worker, range(len(inputs))) + # This pool is not doing computationally intensive work, so we can use threads + pool = multiprocessing.pool.ThreadPool(n_parallel) + tuple_res = pool.map( + local_build_worker, + [ + ( + i.serialize(), + build_func, + timeout, + verbose, + ) + for i in inputs + ], + ) pool.terminate() pool.join() del pool @@ -609,6 +665,70 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo return results +def _timed_eval_func( + inp_serialized, + build_res, + number, + repeat, + min_repeat_ms, + cooldown_interval, + enable_cpu_cache_flush, + verbose, +): + inp = MeasureInput.deserialize(inp_serialized) + tic = time.time() + error_no = 0 + error_msg = None + try: + func = module.load_module(build_res.filename) + ctx = ndarray.context(str(inp.task.target), 0) + # Limitation: + # We can not get PackFunction directly in the remote mode as it is wrapped + # under the std::function. We could lift the restriction later once we fold + # the PackedFunc as an object. Currently, we pass function name to work + # around it. + f_prepare = "cache_flush_cpu_non_first_arg" if enable_cpu_cache_flush else "" + time_f = func.time_evaluator( + func.entry_name, + ctx, + number=number, + repeat=repeat, + min_repeat_ms=min_repeat_ms, + f_preproc=f_prepare, + ) + # pylint: disable=broad-except + except Exception: + costs = (max_float,) + error_no = MeasureErrorNo.COMPILE_DEVICE + error_msg = make_error_msg() + + if error_no == 0: + try: + args = [ndarray.empty(get_const_tuple(x.shape), x.dtype, ctx) for x in build_res.args] + random_fill = tvm.get_global_func("tvm.contrib.random.random_fill", True) + assert random_fill, "Please make sure USE_RANDOM is ON in the config.cmake" + for arg in args: + random_fill(arg) + ctx.sync() + costs = time_f(*args).results + # pylint: disable=broad-except + except Exception: + costs = (max_float,) + error_no = MeasureErrorNo.RUNTIME_DEVICE + error_msg = make_error_msg() + + shutil.rmtree(os.path.dirname(build_res.filename)) + toc = time.time() + time.sleep(cooldown_interval) + + if verbose >= 1: + if error_no == MeasureErrorNo.NO_ERROR: + print("*", end="") + else: + print("*E", end="") # Run error + return costs, error_no, error_msg, toc - tic + build_res.time_cost, toc + + @tvm._ffi.register_func("auto_scheduler.local_runner.run") def local_run( inputs, @@ -667,61 +787,6 @@ def local_run( """ max_float = 1e10 # We use 1e10 instead of sys.float_info.max for better readability in log - def timed_func(inp, build_res): - tic = time.time() - error_no = 0 - error_msg = None - try: - func = module.load_module(build_res.filename) - ctx = ndarray.context(str(inp.task.target), 0) - # Limitation: - # We can not get PackFunction directly in the remote mode as it is wrapped - # under the std::function. We could lift the restriction later once we fold - # the PackedFunc as an object. Currently, we pass function name to work - # around it. - f_prepare = "cache_flush_cpu_non_first_arg" if enable_cpu_cache_flush else "" - time_f = func.time_evaluator( - func.entry_name, - ctx, - number=number, - repeat=repeat, - min_repeat_ms=min_repeat_ms, - f_preproc=f_prepare, - ) - # pylint: disable=broad-except - except Exception: - costs = (max_float,) - error_no = MeasureErrorNo.COMPILE_DEVICE - error_msg = make_error_msg() - - if error_no == 0: - try: - args = [ - ndarray.empty(get_const_tuple(x.shape), x.dtype, ctx) for x in build_res.args - ] - random_fill = tvm.get_global_func("tvm.contrib.random.random_fill", True) - assert random_fill, "Please make sure USE_RANDOM is ON in the config.cmake" - for arg in args: - random_fill(arg) - ctx.sync() - costs = time_f(*args).results - # pylint: disable=broad-except - except Exception: - costs = (max_float,) - error_no = MeasureErrorNo.RUNTIME_DEVICE - error_msg = make_error_msg() - - shutil.rmtree(os.path.dirname(build_res.filename)) - toc = time.time() - time.sleep(cooldown_interval) - - if verbose >= 1: - if error_no == MeasureErrorNo.NO_ERROR: - print("*", end="") - else: - print("*E", end="") # Run error - return costs, error_no, error_msg, toc - tic + build_res.time_cost, toc - measure_results = [] assert len(inputs) == len(build_results), "Measure input size should be equal to build results" for inp, build_res in zip(inputs, build_results): @@ -734,7 +799,20 @@ def timed_func(inp, build_res): time.time(), ) else: - res = call_func_with_timeout(timeout, timed_func, args=(inp, build_res)) + res = call_func_with_timeout( + timeout, + _timed_eval_func, + args=( + inp.serialize(), + build_res, + number, + repeat, + min_repeat_ms, + cooldown_interval, + enable_cpu_cache_flush, + verbose, + ), + ) if isinstance(res, TimeoutError): if verbose >= 1: print("*T", end="") # Run timeout @@ -753,40 +831,104 @@ def timed_func(inp, build_res): return measure_results -def rpc_run_worker(index): +def _timed_rpc_run( + inp_serialized, + build_res, + key, + host, + port, + priority, + timeout, + number, + repeat, + min_repeat_ms, + cooldown_interval, + enable_cpu_cache_flush, + verbose, +): + inp = MeasureInput.deserialize(inp_serialized) + tic = time.time() + error_no = 0 + error_msg = None + try: + # upload built module + remote = request_remote(key, host, port, priority, timeout) + remote.upload(build_res.filename) + func = remote.load_module(os.path.split(build_res.filename)[1]) + ctx = remote.context(str(inp.task.target), 0) + # Limitation: + # We can not get PackFunction directly in the remote mode as it is wrapped + # under the std::function. We could lift the restriction later once we fold + # the PackedFunc as an object. Currently, we pass function name to work + # around it. + f_prepare = "cache_flush_cpu_non_first_arg" if enable_cpu_cache_flush else "" + time_f = func.time_evaluator( + func.entry_name, + ctx, + number=number, + repeat=repeat, + min_repeat_ms=min_repeat_ms, + f_preproc=f_prepare, + ) + # pylint: disable=broad-except + except Exception: + costs = (max_float,) + error_no = MeasureErrorNo.COMPILE_DEVICE + error_msg = make_error_msg() + + if error_no == 0: + try: + args = [ndarray.empty(get_const_tuple(x.shape), x.dtype, ctx) for x in build_res.args] + try: + random_fill = remote.get_function("tvm.contrib.random.random_fill") + except AttributeError: + raise AttributeError( + "Please make sure USE_RANDOM is ON in the config.cmake " "on the remote devices" + ) + for arg in args: + random_fill(arg) + ctx.sync() + + costs = time_f(*args).results + # clean up remote files + remote.remove(build_res.filename) + remote.remove(os.path.splitext(build_res.filename)[0] + ".so") + remote.remove("") + # pylint: disable=broad-except + except Exception: + costs = (max_float,) + error_no = MeasureErrorNo.RUNTIME_DEVICE + error_msg = make_error_msg() + + shutil.rmtree(os.path.dirname(build_res.filename)) + toc = time.time() + + time.sleep(cooldown_interval) + if verbose >= 1: + if error_no == MeasureErrorNo.NO_ERROR: + print("*", end="") + else: + print("*E", end="") # Run error + + return costs, error_no, error_msg, toc - tic + build_res.time_cost, toc + + +def _rpc_run_worker(args): """Function to be ran in the RPCRunner thread pool. Parameters ---------- - index : int - The MeasureInput and BuildResult index to be processed by the current Runner thread. + args : Tuple[MeasureInput, BuildResult, ...] + Single input and build result plus the rest of the arguments to `rpc_runner_run`. Returns ------- res : MeasureResult The measure result of this Runner thread. """ - global GLOBAL_RUN_ARGUMENTS - ( - inputs, - build_results, - key, - host, - port, - priority, - timeout, - number, - repeat, - min_repeat_ms, - cooldown_interval, - enable_cpu_cache_flush, - verbose, - ) = GLOBAL_RUN_ARGUMENTS - max_float = 1e10 # We use 1e10 instead of sys.float_info.max for better readability in log - inp = inputs[index] - build_res = build_results[index] + _, build_res, _, _, _, _, timeout, _, _, _, _, _, verbose = args if build_res.error_no != MeasureErrorNo.NO_ERROR: return ( (max_float,), @@ -796,76 +938,7 @@ def rpc_run_worker(index): time.time(), ) - def timed_func(): - tic = time.time() - error_no = 0 - error_msg = None - try: - # upload built module - remote = request_remote(key, host, port, priority, timeout) - remote.upload(build_res.filename) - func = remote.load_module(os.path.split(build_res.filename)[1]) - ctx = remote.context(str(inp.task.target), 0) - # Limitation: - # We can not get PackFunction directly in the remote mode as it is wrapped - # under the std::function. We could lift the restriction later once we fold - # the PackedFunc as an object. Currently, we pass function name to work - # around it. - f_prepare = "cache_flush_cpu_non_first_arg" if enable_cpu_cache_flush else "" - time_f = func.time_evaluator( - func.entry_name, - ctx, - number=number, - repeat=repeat, - min_repeat_ms=min_repeat_ms, - f_preproc=f_prepare, - ) - # pylint: disable=broad-except - except Exception: - costs = (max_float,) - error_no = MeasureErrorNo.COMPILE_DEVICE - error_msg = make_error_msg() - - if error_no == 0: - try: - args = [ - ndarray.empty(get_const_tuple(x.shape), x.dtype, ctx) for x in build_res.args - ] - try: - random_fill = remote.get_function("tvm.contrib.random.random_fill") - except AttributeError: - raise AttributeError( - "Please make sure USE_RANDOM is ON in the config.cmake " - "on the remote devices" - ) - for arg in args: - random_fill(arg) - ctx.sync() - - costs = time_f(*args).results - # clean up remote files - remote.remove(build_res.filename) - remote.remove(os.path.splitext(build_res.filename)[0] + ".so") - remote.remove("") - # pylint: disable=broad-except - except Exception: - costs = (max_float,) - error_no = MeasureErrorNo.RUNTIME_DEVICE - error_msg = make_error_msg() - - shutil.rmtree(os.path.dirname(build_res.filename)) - toc = time.time() - - time.sleep(cooldown_interval) - if verbose >= 1: - if error_no == MeasureErrorNo.NO_ERROR: - print("*", end="") - else: - print("*E", end="") # Run error - - return costs, error_no, error_msg, toc - tic + build_res.time_cost, toc - - res = call_func_with_timeout(timeout, timed_func) + res = call_func_with_timeout(timeout, _timed_rpc_run, args=args) if isinstance(res, TimeoutError): if verbose >= 1: @@ -950,26 +1023,30 @@ def rpc_runner_run( res : List[MeasureResult] The measure results of these MeasureInputs. """ - global GLOBAL_RUN_ARGUMENTS - GLOBAL_RUN_ARGUMENTS = ( - inputs, - build_results, - key, - host, - port, - priority, - timeout, - number, - repeat, - min_repeat_ms, - cooldown_interval, - enable_cpu_cache_flush, - verbose, - ) - assert len(inputs) == len(build_results), "Measure input size should be equal to build results" - pool = NoDaemonPool(n_parallel) - tuple_res = pool.map(rpc_run_worker, range(len(build_results))) + # This pool is not doing computationally intensive work, so we can use threads + pool = multiprocessing.pool.ThreadPool(n_parallel) + tuple_res = pool.map( + _rpc_run_worker, + [ + ( + inp.serialize(), + build_res, + key, + host, + port, + priority, + timeout, + number, + repeat, + min_repeat_ms, + cooldown_interval, + enable_cpu_cache_flush, + verbose, + ) + for inp, build_res in zip(inputs, build_results) + ], + ) pool.terminate() pool.join() del pool diff --git a/python/tvm/auto_scheduler/measure_record.py b/python/tvm/auto_scheduler/measure_record.py index 1d0d7650a0f6..f0d930e3257e 100644 --- a/python/tvm/auto_scheduler/measure_record.py +++ b/python/tvm/auto_scheduler/measure_record.py @@ -21,9 +21,7 @@ import tvm._ffi from tvm.runtime import Object -from .compute_dag import ComputeDAG -from .measure import MeasureErrorNo, MeasureInput, MeasureCallback -from .search_task import SearchTask +from .measure import MeasureErrorNo, MeasureCallback from . import _ffi_api @@ -175,38 +173,3 @@ def load_best(filename, workload_key=None, target=None): best_res = res return best_inp, best_res - - -def recover_measure_input(inp, rebuild_state=False): - """ - Recover a deserialized MeasureInput by rebuilding the missing fields. - 1. Rebuid the compute_dag in inp.task - 2. (Optional) Rebuild the stages in inp.state - - Parameters - ---------- - inp: MeasureInput - The deserialized MeasureInput - rebuild_state: bool = False - Whether rebuild the stages in MeasureInput.State - - Returns - ------- - new_input: MeasureInput - The fully recovered MeasureInput with all fields rebuilt. - """ - task = inp.task - new_task = SearchTask( - ComputeDAG(task.workload_key), - task.workload_key, - task.target, - task.target_host, - task.hardware_params, - ) - - if rebuild_state: - new_state = new_task.compute_dag.infer_bound_from_state(inp.state) - else: - new_state = inp.state - - return MeasureInput(new_task, new_state) diff --git a/python/tvm/auto_scheduler/utils.py b/python/tvm/auto_scheduler/utils.py index 75fec9c891e8..2d0ec3efd75d 100644 --- a/python/tvm/auto_scheduler/utils.py +++ b/python/tvm/auto_scheduler/utils.py @@ -129,32 +129,6 @@ def deserialize_args(args): return ret -class NoDaemonProcess(multiprocessing.Process): - @property - def daemon(self): - return False - - @daemon.setter - def daemon(self, value): - pass - - -class NoDaemonContext(type(multiprocessing.get_context())): - Process = NoDaemonProcess - - -class NoDaemonPool(multiprocessing.pool.Pool): - """A no daemon pool version of multiprocessing.Pool. - This allows us to start new processes inside the worker function""" - - def __init__(self, *args, **kwargs): - kwargs["context"] = NoDaemonContext() - super().__init__(*args, **kwargs) - - def __reduce__(self): - pass - - def kill_child_processes(parent_pid, sig=signal.SIGTERM): """kill all child processes recursively""" try: @@ -169,17 +143,19 @@ def kill_child_processes(parent_pid, sig=signal.SIGTERM): return +def _func_wrapper(que, func, args, kwargs): + """Call function and return the result over the queue.""" + if kwargs: + que.put(func(*args, **kwargs)) + else: + que.put(func(*args)) + + def call_func_with_timeout(timeout, func, args=(), kwargs=None): """Call a function with timeout""" - def func_wrapper(que): - if kwargs: - que.put(func(*args, **kwargs)) - else: - que.put(func(*args)) - que = multiprocessing.Queue(2) - process = multiprocessing.Process(target=func_wrapper, args=(que,)) + process = multiprocessing.Process(target=_func_wrapper, args=(que, func, args, kwargs)) process.start() process.join(timeout) diff --git a/python/tvm/auto_scheduler/workload_registry.py b/python/tvm/auto_scheduler/workload_registry.py index 1d9ee6da4f7a..c2d7f90771e3 100644 --- a/python/tvm/auto_scheduler/workload_registry.py +++ b/python/tvm/auto_scheduler/workload_registry.py @@ -175,6 +175,41 @@ def workload_key_to_tensors(workload_key): return lookup(*args) +def get_workload_func(task): + """Get the workload function for a given task + + Parameters + ---------- + task : SearchTask + Task to get workload of. + + Returns + ------- + workload : callable + The registered workload function. + """ + name = workload_func_name(task.workload_key) + lookup = WORKLOAD_FUNC_REGISTRY[name] + assert callable(lookup) + return lookup + + +def workload_func_name(workload_key): + """Decode a workload key to the registered function name. + + Parameters + ---------- + workload_key : str + The input workload key. + + Returns + ------- + name : str + The function name of this workload key. + """ + return decode_workload_key_to_func_args(workload_key)[0] + + def save_workload_func_registry(filename): """Dump workload function registry to a pickle binary file. diff --git a/python/tvm/testing.py b/python/tvm/testing.py index 51fa2d0d7def..e5b17f3d7b53 100644 --- a/python/tvm/testing.py +++ b/python/tvm/testing.py @@ -634,6 +634,23 @@ def requires_micro(*args): return _compose(args, _requires_micro) +def requires_rpc(*args): + """Mark a test as requiring rpc to run. + + Parameters + ---------- + f : function + Function to mark + """ + _requires_rpc = [ + pytest.mark.skipif( + tvm.support.libinfo().get("USE_RPC", "OFF") != "ON", + reason="RPC support not enabled. Set USE_RPC=ON in config.cmake to enable.", + ) + ] + return _compose(args, _requires_rpc) + + def _target_to_requirement(target): # mapping from target to decorator if target.startswith("cuda"): diff --git a/src/auto_scheduler/measure_record.cc b/src/auto_scheduler/measure_record.cc old mode 100755 new mode 100644 index 66f521e17e80..1bc2c78a99f0 --- a/src/auto_scheduler/measure_record.cc +++ b/src/auto_scheduler/measure_record.cc @@ -107,6 +107,54 @@ struct Handler<::tvm::auto_scheduler::StateNode> { } }; +template <> +struct Handler<::tvm::auto_scheduler::HardwareParamsNode> { + inline static void Write(dmlc::JSONWriter* writer, + const ::tvm::auto_scheduler::HardwareParamsNode& data) { + writer->BeginArray(false); + writer->WriteArrayItem(data.num_cores); + writer->WriteArrayItem(data.vector_unit_bytes); + writer->WriteArrayItem(data.cache_line_bytes); + writer->WriteArrayItem(data.max_shared_memory_per_block); + writer->WriteArrayItem(data.max_registers_per_block); + writer->WriteArrayItem(data.max_threads_per_block); + writer->WriteArrayItem(data.max_vthread_extent); + writer->WriteArrayItem(data.warp_size); + writer->EndArray(); + } + inline static void Read(dmlc::JSONReader* reader, + ::tvm::auto_scheduler::HardwareParamsNode* data) { + bool s; + reader->BeginArray(); + s = reader->NextArrayItem(); + CHECK(s); + reader->Read(&data->num_cores); + s = reader->NextArrayItem(); + CHECK(s); + reader->Read(&data->vector_unit_bytes); + s = reader->NextArrayItem(); + CHECK(s); + reader->Read(&data->cache_line_bytes); + s = reader->NextArrayItem(); + CHECK(s); + reader->Read(&data->max_shared_memory_per_block); + s = reader->NextArrayItem(); + CHECK(s); + reader->Read(&data->max_registers_per_block); + s = reader->NextArrayItem(); + CHECK(s); + reader->Read(&data->max_threads_per_block); + s = reader->NextArrayItem(); + CHECK(s); + reader->Read(&data->max_vthread_extent); + s = reader->NextArrayItem(); + CHECK(s); + reader->Read(&data->warp_size); + s = reader->NextArrayItem(); + CHECK(!s); + } +}; + template <> struct Handler<::tvm::auto_scheduler::SearchTaskNode> { inline static void Write(dmlc::JSONWriter* writer, @@ -114,11 +162,13 @@ struct Handler<::tvm::auto_scheduler::SearchTaskNode> { writer->BeginArray(false); writer->WriteArrayItem(std::string(data.workload_key)); writer->WriteArrayItem(data.target->str()); + writer->WriteArrayItem(*data.hardware_params.get()); writer->EndArray(); } inline static void Read(dmlc::JSONReader* reader, ::tvm::auto_scheduler::SearchTaskNode* data) { bool s; std::string str_value; + auto hardware_params_node = ::tvm::make_object<::tvm::auto_scheduler::HardwareParamsNode>(); reader->BeginArray(); s = reader->NextArrayItem(); ICHECK(s); @@ -129,7 +179,12 @@ struct Handler<::tvm::auto_scheduler::SearchTaskNode> { reader->Read(&str_value); data->target = ::tvm::Target(str_value); s = reader->NextArrayItem(); - ICHECK(!s); + if (s) { + reader->Read(hardware_params_node.get()); + s = reader->NextArrayItem(); + data->hardware_params = ::tvm::auto_scheduler::HardwareParams(hardware_params_node); + ICHECK(!s); + } } }; @@ -216,7 +271,7 @@ namespace auto_scheduler { TVM_REGISTER_OBJECT_TYPE(RecordToFileNode); TVM_REGISTER_OBJECT_TYPE(RecordReaderNode); -const std::string AUTO_SCHEDULER_LOG_VERSION = "v0.2"; // NOLINT(*) +const std::string AUTO_SCHEDULER_LOG_VERSION = "v0.3"; // NOLINT(*) RecordToFile::RecordToFile(String filename) { auto node = make_object(); @@ -340,5 +395,21 @@ TVM_REGISTER_GLOBAL("auto_scheduler.SaveRecords") std::ofstream ofs(filename, std::ofstream::app); WriteMeasureRecords(&ofs, in, res); }); + +TVM_REGISTER_GLOBAL("auto_scheduler.SerializeMeasureInput") + .set_body_typed([](const MeasureInput& input) { + std::ostringstream os; + dmlc::JSONWriter writer(&os); + writer.Write(*input.get()); + return os.str(); + }); + +TVM_REGISTER_GLOBAL("auto_scheduler.DeserializeMeasureInput").set_body_typed([](String json) { + std::istringstream ss(json); + dmlc::JSONReader reader(&ss); + auto inp = make_object(); + reader.Read(inp.get()); + return ObjectRef(inp); +}); } // namespace auto_scheduler } // namespace tvm diff --git a/tests/python/unittest/test_auto_scheduler_measure.py b/tests/python/unittest/test_auto_scheduler_measure.py index 4369d203b476..80ce98d0b1c1 100644 --- a/tests/python/unittest/test_auto_scheduler_measure.py +++ b/tests/python/unittest/test_auto_scheduler_measure.py @@ -17,6 +17,7 @@ """ Test measurement and log serialization. """ +import multiprocessing import tvm from tvm import topi from tvm import te, auto_scheduler @@ -182,12 +183,10 @@ def test_recover_measure_input(): raw_inp = inputs[0] - correct_inp = auto_scheduler.measure_record.recover_measure_input(raw_inp) + correct_inp = auto_scheduler.measure.recover_measure_input(raw_inp) assert str(correct_inp.task.compute_dag) == str(inp.task.compute_dag) - correct_inp = auto_scheduler.measure_record.recover_measure_input( - raw_inp, rebuild_state=True - ) + correct_inp = auto_scheduler.measure.recover_measure_input(raw_inp, rebuild_state=True) assert str(correct_inp.state) == str(inp.state) @@ -232,6 +231,19 @@ def test_measure_local_builder_rpc_runner(): del measure_ctx +def measure_local_builder_rpc_runner_spawn(): + assert multiprocessing.get_start_method(False) == "spawn" + test_measure_local_builder_rpc_runner() + + +@tvm.testing.requires_llvm +def test_measure_local_builder_rpc_runner_spawn(): + ctx = multiprocessing.get_context("spawn") + p = ctx.Process(target=measure_local_builder_rpc_runner_spawn) + p.start() + p.join() + + if __name__ == "__main__": test_record_split_reorder_fuse_annotation() test_record_compute_at_root_inline_cache_read_write() @@ -239,4 +251,5 @@ def test_measure_local_builder_rpc_runner(): test_record_pragma_storage_align_rfactor() test_recover_measure_input() test_measure_local_builder_runner() + test_measure_local_builder_runner_spawn() test_measure_local_builder_rpc_runner() diff --git a/tests/python/unittest/test_auto_scheduler_search_policy.py b/tests/python/unittest/test_auto_scheduler_search_policy.py index 07cf4c8141a0..5329f3d50685 100644 --- a/tests/python/unittest/test_auto_scheduler_search_policy.py +++ b/tests/python/unittest/test_auto_scheduler_search_policy.py @@ -18,6 +18,7 @@ """Test search policy""" import random +import multiprocessing import numpy as np import tempfile @@ -26,6 +27,7 @@ from tvm import auto_scheduler from test_auto_scheduler_common import matmul_auto_scheduler_test, PropagatingThread +import multiprocessing def search_common( @@ -122,6 +124,19 @@ def test_sketch_search_policy_basic(): t.join() +def sketch_search_policy_basic_spawn(): + assert multiprocessing.get_start_method(False) == "spawn" + test_sketch_search_policy_basic() + + +@tvm.testing.requires_llvm +def test_sketch_search_policy_basic_spawn(): + ctx = multiprocessing.get_context("spawn") + p = ctx.Process(target=sketch_search_policy_basic_spawn) + p.start() + p.join() + + @tvm.testing.requires_llvm def test_sketch_search_policy_xgbmodel(): # wrap the search in a new thread to avoid the conflict @@ -156,9 +171,8 @@ def test_sketch_search_policy_cuda_rpc_runner(): t.join() +@tvm.testing.requires_cuda def test_sketch_search_policy_cuda_xgbmodel_rpc_runner(): - if not tvm.runtime.enabled("cuda"): - return measure_ctx = auto_scheduler.LocalRPCMeasureContext() # wrap the search in a new thread to avoid the conflict # between python's multiprocessing and tvm's thread pool @@ -179,6 +193,7 @@ def test_sketch_search_policy_cuda_xgbmodel_rpc_runner(): if __name__ == "__main__": test_workload_registry_search_basic() test_sketch_search_policy_basic() + test_sketch_search_policy_basic_spawn() test_sketch_search_policy_xgbmodel() test_sketch_search_policy_cuda_rpc_runner() test_sketch_search_policy_cuda_xgbmodel_rpc_runner() diff --git a/tests/python/unittest/test_auto_scheduler_task_scheduler.py b/tests/python/unittest/test_auto_scheduler_task_scheduler.py index 72b998a5a38a..7851d922013d 100644 --- a/tests/python/unittest/test_auto_scheduler_task_scheduler.py +++ b/tests/python/unittest/test_auto_scheduler_task_scheduler.py @@ -18,6 +18,7 @@ import tempfile +import multiprocessing import numpy as np from tvm import auto_scheduler @@ -68,6 +69,18 @@ def objective_func(costs): task_scheduler.tune(tune_option, search_policy="sketch.random") +def task_scheduler_round_robin_spawn(): + assert multiprocessing.get_start_method(False) == "spawn" + test_task_scheduler_round_robin() + + +def test_task_scheduler_round_robin_spawn(): + ctx = multiprocessing.get_context("spawn") + p = ctx.Process(target=task_scheduler_round_robin_spawn) + p.start() + p.join() + + def test_task_scheduler_gradient(): tasks = [] for n in [2, 4]: @@ -109,4 +122,5 @@ def objective_func(costs): if __name__ == "__main__": test_task_scheduler_round_robin() + test_task_scheduler_round_robin_spawn() test_task_scheduler_gradient() diff --git a/tests/python/unittest/test_autotvm_dispatch_context.py b/tests/python/unittest/test_autotvm_dispatch_context.py index 4064ede3cc06..6ca062047fd7 100644 --- a/tests/python/unittest/test_autotvm_dispatch_context.py +++ b/tests/python/unittest/test_autotvm_dispatch_context.py @@ -21,12 +21,13 @@ from tvm import autotvm -def test_fallback(): - @autotvm.template("testing/dispatch_fallback") - def simple_template(a, b): - cfg = autotvm.get_config() - assert cfg.is_fallback +@autotvm.template("testing/dispatch_fallback") +def simple_template(a, b): + cfg = autotvm.get_config() + assert cfg.is_fallback + +def test_fallback(): simple_template(2, 3) diff --git a/tests/python/unittest/test_runtime_rpc.py b/tests/python/unittest/test_runtime_rpc.py index d25eff23ae76..2739f5637967 100644 --- a/tests/python/unittest/test_runtime_rpc.py +++ b/tests/python/unittest/test_runtime_rpc.py @@ -30,6 +30,7 @@ from tvm.rpc.tracker import Tracker +@tvm.testing.requires_rpc def test_bigendian_rpc(): """Test big endian rpc when there is a PowerPC RPC server available""" host = os.environ.get("TVM_POWERPC_TEST_HOST", None) @@ -61,22 +62,23 @@ def verify_rpc(remote, target, shape, dtype): verify_rpc(remote, target, (10,), dtype) -def test_rpc_simple(): - if not tvm.runtime.enabled("rpc"): - return +@tvm.register_func("rpc.test.addone") +def addone(x): + return x + 1 + + +@tvm.register_func("rpc.test.strcat") +def strcat(name, x): + return "%s:%d" % (name, x) - @tvm.register_func("rpc.test.addone") - def addone(x): - return x + 1 - @tvm.register_func("rpc.test.strcat") - def strcat(name, x): - return "%s:%d" % (name, x) +@tvm.register_func("rpc.test.except") +def remotethrow(name): + raise ValueError("%s" % name) - @tvm.register_func("rpc.test.except") - def remotethrow(name): - raise ValueError("%s" % name) +@tvm.testing.requires_rpc +def test_rpc_simple(): server = rpc.Server("localhost", key="x1") client = rpc.connect(server.host, server.port, key="x1") f1 = client.get_function("rpc.test.addone") @@ -90,14 +92,13 @@ def remotethrow(name): assert f2("abc", 11) == "abc:11" -def test_rpc_runtime_string(): - if not tvm.runtime.enabled("rpc"): - return +@tvm.register_func("rpc.test.runtime_str_concat") +def strcat(x, y): + return x + y - @tvm.register_func("rpc.test.runtime_str_concat") - def strcat(x, y): - return x + y +@tvm.testing.requires_rpc +def test_rpc_runtime_string(): server = rpc.Server("localhost", key="x1") client = rpc.connect(server.host, server.port, key="x1") func = client.get_function("rpc.test.runtime_str_concat") @@ -106,14 +107,15 @@ def strcat(x, y): assert str(func(x, y)) == "abcdef" -def test_rpc_array(): - if not tvm.runtime.enabled("rpc"): - return - x = np.random.randint(0, 10, size=(3, 4)) +@tvm.register_func("rpc.test.remote_array_func") +def remote_array_func(y): + x = np.ones((3, 4)) + np.testing.assert_equal(y.asnumpy(), x) - @tvm.register_func("rpc.test.remote_array_func") - def remote_array_func(y): - np.testing.assert_equal(y.asnumpy(), x) + +@tvm.testing.requires_rpc +def test_rpc_array(): + x = np.ones((3, 4)) server = rpc.Server("localhost") remote = rpc.connect(server.host, server.port) @@ -124,6 +126,7 @@ def remote_array_func(y): fremote(r_cpu) +@tvm.testing.requires_rpc def test_rpc_large_array(): # testcase of large array creation server = rpc.Server("localhost") @@ -137,6 +140,7 @@ def test_rpc_large_array(): np.testing.assert_equal(b.asnumpy(), b_np) +@tvm.testing.requires_rpc def test_rpc_echo(): def check(remote): fecho = remote.get_function("testing.echo") @@ -180,9 +184,8 @@ def check_minrpc(): check_minrpc() +@tvm.testing.requires_rpc def test_rpc_file_exchange(): - if not tvm.runtime.enabled("rpc"): - return server = rpc.Server("localhost") remote = rpc.connect(server.host, server.port) blob = bytearray(np.random.randint(0, 10, size=(10))) @@ -191,10 +194,9 @@ def test_rpc_file_exchange(): assert rev == blob +@tvm.testing.requires_rpc @tvm.testing.requires_llvm def test_rpc_remote_module(): - if not tvm.runtime.enabled("rpc"): - return # graph n = tvm.runtime.convert(102) A = te.placeholder((n,), name="A") @@ -317,11 +319,13 @@ def check_remote_link_cl(remote): check_minrpc() -def test_rpc_return_func(): - @tvm.register_func("rpc.test.remote_func") - def addone(x): - return lambda y: x + y +@tvm.register_func("rpc.test.remote_func") +def addone(x): + return lambda y: x + y + +@tvm.testing.requires_rpc +def test_rpc_return_func(): server = rpc.Server("localhost", key="x1") client = rpc.connect(server.host, server.port, key="x1") f1 = client.get_function("rpc.test.remote_func") @@ -329,6 +333,7 @@ def addone(x): assert fadd(12) == 22 +@tvm.testing.requires_rpc def test_rpc_session_constructor_args(): # start server server0 = rpc.Server("localhost", key="x0") @@ -365,21 +370,23 @@ def check_error_handling(): check_error_handling() -def test_rpc_return_ndarray(): +@tvm.register_func("rpc.test.remote_return_nd") +def my_module(name): # Use closure to check the ref counter correctness nd = tvm.nd.array(np.zeros(10).astype("float32")) - @tvm.register_func("rpc.test.remote_return_nd") - def my_module(name): - if name == "get_arr": - return lambda: nd - elif name == "ref_count": - return lambda: tvm.testing.object_use_count(nd) - elif name == "get_elem": - return lambda idx: nd.asnumpy()[idx] - elif name == "get_arr_elem": - return lambda arr, idx: arr.asnumpy()[idx] + if name == "get_arr": + return lambda: nd + elif name == "ref_count": + return lambda: tvm.testing.object_use_count(nd) + elif name == "get_elem": + return lambda idx: nd.asnumpy()[idx] + elif name == "get_arr_elem": + return lambda arr, idx: arr.asnumpy()[idx] + +@tvm.testing.requires_rpc +def test_rpc_return_ndarray(): # start server server = rpc.Server("localhost", key="x1") client = rpc.connect(server.host, server.port, key="x1") @@ -392,26 +399,19 @@ def my_module(name): # array test def run_arr_test(): arr = get_arr() - assert ref_count() == 2 - arr2 = get_arr() - assert ref_count() == 3 - assert arr.context == client.cpu(0) - arr.copyfrom(np.ones(10).astype(arr.dtype)) - assert arr2.asnumpy()[0] == 1.0 - assert get_elem(0) == 1.0 - assert get_arr_elem(arr2, 0) == 1.0 - - assert ref_count() == 1 + assert get_elem(0) == 0.0 + assert get_arr_elem(arr, 0) == 0.0 + run_arr_test() - # check recycle correctness - assert ref_count() == 1 -def test_local_func(): - @tvm.register_func("rpc.test.remote_func2") - def addone(x): - return lambda y: x + y +@tvm.register_func("rpc.test.remote_func2") +def addone(x): + return lambda y: x + y + +@tvm.testing.requires_rpc +def test_local_func(): client = rpc.LocalSession() f1 = client.get_function("rpc.test.remote_func2") fadd = f1(10) @@ -423,6 +423,7 @@ def addone(x): assert rev == blob +@tvm.testing.requires_rpc def test_rpc_tracker_register(): # test registration tracker = Tracker("localhost", port=9000, port_end=10000) @@ -459,6 +460,15 @@ def test_rpc_tracker_register(): tracker.terminate() +def _target(host, port, device_key, timeout): + client = rpc.connect_tracker(host, port) + remote = client.request(device_key, session_timeout=timeout) + while True: + pass + remote.cpu() + + +@tvm.testing.requires_rpc def test_rpc_tracker_request(): # test concurrent request tracker = Tracker("localhost", port=9000, port_end=10000) @@ -472,16 +482,11 @@ def test_rpc_tracker_request(): ) client = rpc.connect_tracker(tracker.host, tracker.port) - def target(host, port, device_key, timeout): - client = rpc.connect_tracker(host, port) - remote = client.request(device_key, session_timeout=timeout) - while True: - pass - remote.cpu() - - proc1 = multiprocessing.Process(target=target, args=(tracker.host, tracker.port, device_key, 4)) + proc1 = multiprocessing.Process( + target=_target, args=(tracker.host, tracker.port, device_key, 4) + ) proc2 = multiprocessing.Process( - target=target, args=(tracker.host, tracker.port, device_key, 200) + target=_target, args=(tracker.host, tracker.port, device_key, 200) ) proc1.start() time.sleep(0.5)