From 3f594b40ee72ac9a21421b37551536f4e9fcec17 Mon Sep 17 00:00:00 2001 From: Fan Date: Fri, 25 Oct 2019 12:13:43 +0800 Subject: [PATCH] fix rebase --- CMakeLists.txt | 2 +- contrib/tvmop/compile.py | 52 +++++++++------------------------- contrib/tvmop/core/__init__.py | 7 +---- contrib/tvmop/opdef.py | 39 +++++++++++++++++++++---- 4 files changed, 50 insertions(+), 50 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index c86b57a04b8d..a06aa9dba485 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -763,7 +763,7 @@ if(USE_TVM_OP) endif() endif() - set(TVM_OP_COMPILE_OPTIONS "-o${CMAKE_CURRENT_BINARY_DIR}/libtvmop.so --config ${CMAKE_CURRENT_BINARY_DIR}/tvmop.conf") + set(TVM_OP_COMPILE_OPTIONS "-o${CMAKE_CURRENT_BINARY_DIR}/libtvmop.so" "--config" "${CMAKE_CURRENT_BINARY_DIR}/tvmop.conf") if(CUDA_ARCH_BIN) set(TVM_OP_COMPILE_OPTIONS "${TVM_OP_COMPILE_OPTIONS}" "--cuda-arch" "${CUDA_ARCH_BIN}") endif() diff --git a/contrib/tvmop/compile.py b/contrib/tvmop/compile.py index 888b504504e9..b0254218077a 100644 --- a/contrib/tvmop/compile.py +++ b/contrib/tvmop/compile.py @@ -79,46 +79,17 @@ def get_cuda_arch(arch): func_list_llvm = [] func_list_cuda = [] - config_spaces = ConfigSpaces() # TODO: attach instruction features to the library, e.g., avx-512, etc. - for op in __OP_DEF__: - if tvm.module.enabled(get_target(op.target)): - func_list = func_list_llvm if op.target == "cpu" else func_list_cuda - for each_kwargs in op.arg_combination: - if (op.attrs_valid(**each_kwargs)): - name = op.name \ - + ''.join(["{}_{}".format(key, each_kwargs[key]) for key in op.attrs]) - if op.dispatch is True: - config_space = autotvm.ConfigSpace() - with autotvm.task.ApplyConfig(config_space): - sch, args = op.func(fallback=False, **each_kwargs) - # register dispatch schedules - for i in range(len(config_space)): - config_entity = config_space.get(i) - with autotvm.task.ApplyConfig(config_entity): - sch, args = op.func(fallback=False, **each_kwargs) - subname = name + "index_" + str(i) + \ - ''.join(["%s_%d" % (arg.dtype, len(arg.shape)) for arg in args]) - func_lower = tvm.lower(sch, args, - name=subname, - binds=op.get_binds(args)) - func_list.append(func_lower) - # register config space - config_spaces[name] = ConfigSpace.from_tvm(config_space) - # register fallback schedule - config_space = autotvm.ConfigSpace() - with autotvm.task.ApplyConfig(config_space): - sch, args = op.func(fallback=True, **each_kwargs) - subname = name + "fallback" + \ - ''.join(["%s_%d" % (arg.dtype, len(arg.shape)) for arg in args]) - func_lower = tvm.lower(sch, args, name=subname, binds=op.get_binds(args)) - func_list.append(func_lower) - else: - sch, args = op.func(**each_kwargs) - subname = name + ''.join(["%s_%d" % (arg.dtype, len(arg.shape)) for arg in args]) - func_lower = tvm.lower(sch, args, name=subname, binds=op.get_binds(args)) - func_list.append(func_lower) + for operator_def in __OP_DEF__: + for sch, args, name in operator_def.invoke_all(): + name = operator_def.get_op_name(name, args) + if tvm.module.enabled(get_target(operator_def.target)): + func_list = func_list_llvm if operator_def.target == "cpu" else func_list_cuda + func_lower = tvm.lower(sch, args, + name=name, + binds=operator_def.get_binds(args)) + func_list.append(func_lower) lowered_funcs = {get_target("cpu"): func_list_llvm} if len(func_list_cuda) > 0: @@ -131,5 +102,10 @@ def get_cuda_arch(arch): set_cuda_target_arch(cuda_arch) func_binary = tvm.build(lowered_funcs, name="tvmop") func_binary.export_library(arguments.target_path) + + config_spaces = ConfigSpaces() + for operator_def in __OP_DEF__: + for config_space, name in operator_def.get_config_spaces(): + config_spaces[name] = ConfigSpace.from_tvm(config_space) with open(arguments.config_path, "w") as f: json.dump(config_spaces.to_json_dict(), f) diff --git a/contrib/tvmop/core/__init__.py b/contrib/tvmop/core/__init__.py index de9e6b38182b..e309f237df05 100644 --- a/contrib/tvmop/core/__init__.py +++ b/contrib/tvmop/core/__init__.py @@ -15,9 +15,4 @@ # specific language governing permissions and limitations # under the License. -<<<<<<< HEAD -from . import umath, fromnumeric -======= -# coding: utf-8 -from . import multiarray ->>>>>>> infra for dispatch tvm op +from . import umath, fromnumeric, multiarray diff --git a/contrib/tvmop/opdef.py b/contrib/tvmop/opdef.py index 3159e5183ce8..57537bac648e 100644 --- a/contrib/tvmop/opdef.py +++ b/contrib/tvmop/opdef.py @@ -77,12 +77,41 @@ def __call__(self, *args, **kwargs): def invoke_all(self): for each_kwargs in self.arg_combination: if self.attrs_valid(**each_kwargs): - sch, args = self.func(**each_kwargs) name = self.name \ - + ''.join(["{}_{}".format(key, each_kwargs[key]) for key in self.attrs]) \ - + ''.join(["%s_%d" % (arg.dtype, len(arg.shape)) - for arg in args if hasattr(arg, 'shape')]) - yield sch, args, name + + ''.join(["{}_{}".format(key, each_kwargs[key]) for key in self.attrs]) + if self.dispatch is False: + sch, args = self.func(**each_kwargs) + yield sch, args, name + else: + # register dispatch schedules + config_space = autotvm.ConfigSpace() + with autotvm.task.ApplyConfig(config_space): + sch, args = self.func(fallback=False, **each_kwargs) + for i in range(len(config_space)): + config_entity = config_space.get(i) + with autotvm.task.ApplyConfig(config_entity): + sch, args = self.func(fallback=False, **each_kwargs) + subname = name + "index_" + str(i) + yield sch, args, subname + # register fallback schedule + config_space = autotvm.ConfigSpace() + with autotvm.task.ApplyConfig(config_space): + sch, args = self.func(fallback=True, **each_kwargs) + subname = name + "fallback" + yield sch, args, subname + + def get_op_name(self, name, args): + return name + ''.join(["%s_%d" % (arg.dtype, len(arg.shape)) for arg in args if hasattr(arg, 'shape')]) + + def get_config_spaces(self): + for each_kwargs in self.arg_combination: + if self.attrs_valid(**each_kwargs) and self.dispatch is True: + name = self.name \ + + ''.join(["{}_{}".format(key, each_kwargs[key]) for key in self.attrs]) + config_space = autotvm.ConfigSpace() + with autotvm.task.ApplyConfig(config_space): + self.func(fallback=False, **each_kwargs) + yield config_space, name def get_binds(self, args): if self.auto_broadcast: