From d02a2669ca80a3e9912dd95abd622a6f0d32fef2 Mon Sep 17 00:00:00 2001 From: leowang1225 <810916296@qq.com> Date: Wed, 30 Dec 2020 16:53:37 +0800 Subject: [PATCH 1/6] [AutoScheduler] Add custom build function Signed-off-by: leowang1225 <810916296@qq.com> --- python/tvm/auto_scheduler/measure.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/python/tvm/auto_scheduler/measure.py b/python/tvm/auto_scheduler/measure.py index 2f177a242835..3b71c0fc763c 100644 --- a/python/tvm/auto_scheduler/measure.py +++ b/python/tvm/auto_scheduler/measure.py @@ -63,6 +63,9 @@ # We use 1e10 instead of sys.float_info.max for better readability in log MAX_FLOAT = 1e10 +class CustomBuildFunc: + """ store custom build_func to class variable. """ + build_func = None @tvm._ffi.register_object("auto_scheduler.MeasureCallback") class MeasureCallback(Object): @@ -303,11 +306,16 @@ class LocalBuilder(ProgramBuilder): This is used in a wrapper of the multiprocessing.Process.join(). n_parallel : int = multiprocessing.cpu_count() Number of threads used to build in parallel. - build_func : str = 'default' - The name of registered build function. + build_func: callable or str + If is 'default', use default build function + If is 'ndk', use function for android ndk + If is callable, use it as custom build function, expect lib_format field. """ def __init__(self, timeout=15, n_parallel=multiprocessing.cpu_count(), build_func="default"): + if not isinstance(build_func, str): + CustomBuildFunc.build_func = build_func + build_func = "custom" self.__init_handle_by_constructor__(_ffi_api.LocalBuilder, timeout, n_parallel, build_func) @@ -628,6 +636,8 @@ def local_build_worker(args): build_func = tar.tar elif build_func == "ndk": build_func = ndk.create_shared + elif build_func == "custom": + build_func = CustomBuildFunc.build_func else: raise ValueError("Invalid build_func" + build_func) From 9e66f14e98918ab65ae0b0342d62f4d27f279634 Mon Sep 17 00:00:00 2001 From: leowang1225 <810916296@qq.com> Date: Mon, 4 Jan 2021 14:07:00 +0800 Subject: [PATCH 2/6] [AutoScheduler] Add custom build function Signed-off-by: leowang1225 <810916296@qq.com> --- python/tvm/auto_scheduler/measure.py | 40 +++++++++++++++++----------- 1 file changed, 24 insertions(+), 16 deletions(-) diff --git a/python/tvm/auto_scheduler/measure.py b/python/tvm/auto_scheduler/measure.py index 3b71c0fc763c..318cebaaa834 100644 --- a/python/tvm/auto_scheduler/measure.py +++ b/python/tvm/auto_scheduler/measure.py @@ -63,9 +63,15 @@ # We use 1e10 instead of sys.float_info.max for better readability in log MAX_FLOAT = 1e10 -class CustomBuildFunc: - """ store custom build_func to class variable. """ - build_func = None +class BuildFunc: + """ store build_func name and callable to class variable. + name: str = "default" + The name of registered build function. + build_func: callable = tar.tar + The callable of registered build function. + """ + name = "default" + build_func = tar.tar @tvm._ffi.register_object("auto_scheduler.MeasureCallback") class MeasureCallback(Object): @@ -306,17 +312,26 @@ class LocalBuilder(ProgramBuilder): This is used in a wrapper of the multiprocessing.Process.join(). n_parallel : int = multiprocessing.cpu_count() Number of threads used to build in parallel. - build_func: callable or str + build_func: callable or str = "default" If is 'default', use default build function If is 'ndk', use function for android ndk If is callable, use it as custom build function, expect lib_format field. """ def __init__(self, timeout=15, n_parallel=multiprocessing.cpu_count(), build_func="default"): - if not isinstance(build_func, str): - CustomBuildFunc.build_func = build_func - build_func = "custom" - self.__init_handle_by_constructor__(_ffi_api.LocalBuilder, timeout, n_parallel, build_func) + if build_func == "default": + BuildFunc.name = "default" + BuildFunc.build_func = tar.tar + elif build_func == "ndk": + BuildFunc.name = "ndk" + BuildFunc.build_func = ndk.create_shared + elif not isinstance(build_func, str): + BuildFunc.name = "custom" + BuildFunc.build_func = build_func + else: + raise ValueError("Invalid build_func" + build_func) + + self.__init_handle_by_constructor__(_ffi_api.LocalBuilder, timeout, n_parallel, BuildFunc.name) @tvm._ffi.register_object("auto_scheduler.LocalRunner") @@ -632,14 +647,7 @@ def local_build_worker(args): The build result of this Builder thread. """ inp, build_func, timeout, verbose = args - if build_func == "default": - build_func = tar.tar - elif build_func == "ndk": - build_func = ndk.create_shared - elif build_func == "custom": - build_func = CustomBuildFunc.build_func - else: - raise ValueError("Invalid build_func" + build_func) + build_func = BuildFunc.build_func res = call_func_with_timeout(timeout, _timed_func, args=(inp, build_func, verbose)) if isinstance(res, TimeoutError): From e3deafb18c4179a68c1daade9ae3675623caec4b Mon Sep 17 00:00:00 2001 From: leowang1225 <810916296@qq.com> Date: Mon, 4 Jan 2021 14:57:27 +0800 Subject: [PATCH 3/6] cheduler] Add custom build function --- python/tvm/auto_scheduler/measure.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/python/tvm/auto_scheduler/measure.py b/python/tvm/auto_scheduler/measure.py index 318cebaaa834..ddc1674f39f0 100644 --- a/python/tvm/auto_scheduler/measure.py +++ b/python/tvm/auto_scheduler/measure.py @@ -63,16 +63,19 @@ # We use 1e10 instead of sys.float_info.max for better readability in log MAX_FLOAT = 1e10 + class BuildFunc: - """ store build_func name and callable to class variable. - name: str = "default" - The name of registered build function. - build_func: callable = tar.tar - The callable of registered build function. + """store build_func name and callable to class variable. + name: str = "default" + The name of registered build function. + build_func: callable = tar.tar + The callable of registered build function. """ + name = "default" build_func = tar.tar + @tvm._ffi.register_object("auto_scheduler.MeasureCallback") class MeasureCallback(Object): """ The base class of measurement callback functions. """ @@ -331,7 +334,9 @@ def __init__(self, timeout=15, n_parallel=multiprocessing.cpu_count(), build_fun else: raise ValueError("Invalid build_func" + build_func) - self.__init_handle_by_constructor__(_ffi_api.LocalBuilder, timeout, n_parallel, BuildFunc.name) + self.__init_handle_by_constructor__( + _ffi_api.LocalBuilder, timeout, n_parallel, BuildFunc.name + ) @tvm._ffi.register_object("auto_scheduler.LocalRunner") @@ -647,6 +652,7 @@ def local_build_worker(args): The build result of this Builder thread. """ inp, build_func, timeout, verbose = args + assert any(build_func == name for name in BuildFunc.name) build_func = BuildFunc.build_func res = call_func_with_timeout(timeout, _timed_func, args=(inp, build_func, verbose)) From 5da7ed7a69e9e45b7077c36be99e96b99ba3fb51 Mon Sep 17 00:00:00 2001 From: leowang1225 <810916296@qq.com> Date: Mon, 4 Jan 2021 15:03:34 +0800 Subject: [PATCH 4/6] [AutoScheduler] Add custom build function Signed-off-by: leowang1225 <810916296@qq.com> --- python/tvm/auto_scheduler/measure.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/tvm/auto_scheduler/measure.py b/python/tvm/auto_scheduler/measure.py index ddc1674f39f0..38fcfd55b013 100644 --- a/python/tvm/auto_scheduler/measure.py +++ b/python/tvm/auto_scheduler/measure.py @@ -652,7 +652,6 @@ def local_build_worker(args): The build result of this Builder thread. """ inp, build_func, timeout, verbose = args - assert any(build_func == name for name in BuildFunc.name) build_func = BuildFunc.build_func res = call_func_with_timeout(timeout, _timed_func, args=(inp, build_func, verbose)) From 7346351075817fd9eeed71db999e292fd7086193 Mon Sep 17 00:00:00 2001 From: leowang1225 <810916296@qq.com> Date: Mon, 4 Jan 2021 15:30:59 +0800 Subject: [PATCH 5/6] [AutoScheduler] Add custom build function Signed-off-by: leowang1225 <810916296@qq.com> --- python/tvm/auto_scheduler/measure.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/tvm/auto_scheduler/measure.py b/python/tvm/auto_scheduler/measure.py index 38fcfd55b013..f1e56258dc6f 100644 --- a/python/tvm/auto_scheduler/measure.py +++ b/python/tvm/auto_scheduler/measure.py @@ -652,6 +652,9 @@ def local_build_worker(args): The build result of this Builder thread. """ inp, build_func, timeout, verbose = args + assert build_func == BuildFunc.name, ( + "BuildFunc.name: " + BuildFunc.name + ", but args is: " + build_func + ) build_func = BuildFunc.build_func res = call_func_with_timeout(timeout, _timed_func, args=(inp, build_func, verbose)) From 976caedffdd24086d36d0ab7caeb9fc065d92700 Mon Sep 17 00:00:00 2001 From: leowang1225 <810916296@qq.com> Date: Tue, 5 Jan 2021 10:25:27 +0800 Subject: [PATCH 6/6] [AutoScheduler] Add custom build function Signed-off-by: leowang1225 <810916296@qq.com> --- python/tvm/auto_scheduler/measure.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/auto_scheduler/measure.py b/python/tvm/auto_scheduler/measure.py index f1e56258dc6f..47ffde4327c4 100644 --- a/python/tvm/auto_scheduler/measure.py +++ b/python/tvm/auto_scheduler/measure.py @@ -328,7 +328,7 @@ def __init__(self, timeout=15, n_parallel=multiprocessing.cpu_count(), build_fun elif build_func == "ndk": BuildFunc.name = "ndk" BuildFunc.build_func = ndk.create_shared - elif not isinstance(build_func, str): + elif callable(build_func): BuildFunc.name = "custom" BuildFunc.build_func = build_func else: