From ca7973daa56b9ccd98ec749048a77f690593bc6e Mon Sep 17 00:00:00 2001 From: LaiYongqiang Date: Thu, 1 Jul 2021 20:54:08 +0800 Subject: [PATCH] adapt virtual computing --- .../parallel_compile/tbe_compiler/compiler.py | 12 ++++++-- .../tbe_compiler/tbe_process.py | 28 ++++++++++++++++--- 2 files changed, 33 insertions(+), 7 deletions(-) diff --git a/mindspore/_extends/parallel_compile/tbe_compiler/compiler.py b/mindspore/_extends/parallel_compile/tbe_compiler/compiler.py index a6f2b1aaf2d..472377ff808 100755 --- a/mindspore/_extends/parallel_compile/tbe_compiler/compiler.py +++ b/mindspore/_extends/parallel_compile/tbe_compiler/compiler.py @@ -18,6 +18,7 @@ import os import sys from te.platform.cce_conf import te_set_version from te_fusion.fusion_util import fusion_op +from te_fusion.fusion_manager import set_context_parameter import tbe.common.context.op_info as operator_info sys.path.append(os.path.abspath(os.path.dirname(__file__))) # pylint: disable=wrong-import-position @@ -71,6 +72,7 @@ def build_op(build_type, json_str, tune_mode=None): op_type = kernel_info['op_info']['Type'] rl_tune_switch = kernel_info['op_info']['rl_tune_switch'] rl_tune_list = kernel_info['op_info']['rl_tune_list'] + reset_op_info = kernel_info["reset_op_info"] try: custom_flag = False @@ -121,7 +123,9 @@ def build_op(build_type, json_str, tune_mode=None): import tbe.common.context.op_context as op_context with op_context.OpContext("dynamic"): op_info = operator_info.OpInfo(op_type, op_type) - op_context.get_context().add_op_info(op_info) + context = op_context.get_context() + context.add_op_info(op_info) + set_context_parameter(context, None, None, reset_op_info) op_func(*inputs_args, *outputs_args, *attrs_args, kernel_name=kernel_name) compile_info = op_context.get_context().get_compile_info() if tune_mode is not None: @@ -141,7 +145,7 @@ def build_op(build_type, json_str, tune_mode=None): auto_tiling_mode=None, device_id=None, fuzz_build_info=None, - reset_op_info=None, + reset_op_info=reset_op_info, switch_str=rl_tune_switch, lic_opt_list=rl_tune_list) if tune_mode is not None: @@ -165,6 +169,7 @@ def compile_fusion_op(json_str): Exception: If specific keyword is not found. """ args = json.loads(json_str) + reset_op_info = args["reset_op_info"] te_set_version(args['fusion_op']["socVersion"]) if 'fusion_op' not in args or not args['fusion_op']: raise ValueError("Json string Errors, key:fusion_op not found.") @@ -172,7 +177,8 @@ def compile_fusion_op(json_str): fusion_op_arg = args['fusion_op'] rl_tune_switch = args['fusion_op']['rl_tune_switch'] rl_tune_list = args['fusion_op']['rl_tune_list'] - return fusion_op(json.dumps(fusion_op_arg), switch_str=rl_tune_switch, lic_opt_list=rl_tune_list) + return fusion_op(json.dumps(fusion_op_arg), reset_op_info=reset_op_info, switch_str=rl_tune_switch, + lic_opt_list=rl_tune_list) def compile_with_json(json_str): diff --git a/mindspore/_extends/parallel_compile/tbe_compiler/tbe_process.py b/mindspore/_extends/parallel_compile/tbe_compiler/tbe_process.py index b304e0b9ff2..36d7e9581de 100644 --- a/mindspore/_extends/parallel_compile/tbe_compiler/tbe_process.py +++ b/mindspore/_extends/parallel_compile/tbe_compiler/tbe_process.py @@ -21,6 +21,10 @@ import sys import os import time import json +import tbe.common.context.op_context as op_context +from tbe.common.buildcfg import build_config +import te.platform.vector_random_buff as vector_random_buff +import te.platform.cube_random_buff as cube_random_buff from mindspore import log from .tbe_common import check_kernel_info, TBEException from .helper import _op_select_format, _check_supported @@ -127,6 +131,7 @@ class TbeProcess: self.__process_num = multiprocessing.cpu_count() self.compile_process_num = 24 self.__pool = None + self.__reset_op_info = None self.__next_task_id = 1 self.__running_tasks = [] self.__all_tune_tasks = [] @@ -261,13 +266,12 @@ class TbeProcess: full_name = op_json["op_info"]["full_name"] return full_name in self.selected_tune_ops - def select_tune_mode(self, op_json): + def select_tune_mode(self, json_info): """ Select the corresponding tune mode from op json and env info for the op - :param op_json: ori json + :param json_info: ori json :return: NO_TUNE RL_TUNE or GA_TUNE """ - json_info = json.loads(op_json) tune_mode = json_info["SocInfo"]["autoTilingMode"] kernel_names = self.get_kernel_names(json_info) if self.offline_tune: @@ -300,6 +304,18 @@ class TbeProcess: kernel_names.append(json_info['op_info']['name']) return kernel_names + + def get_reset_op_info(self): + """ get reset op info """ + with op_context.OpContext("static"): + context = op_context.get_context() + with build_config(kernel_meta_parent_dir=".", compatible=True): + vector_random_buff.vector_random_buff() + cube_random_buff.cube_random_buff() + reset_op_info = context.get_addition("reset_op_info") + return reset_op_info + + def start_compile_op(self, op_json): """ start compile op async. @@ -315,13 +331,17 @@ class TbeProcess: if not self.tune_init: return error_id self.__next_task_id = self.__next_task_id + 1 - tune_mode = self.select_tune_mode(op_json) + json_info = json.loads(op_json) + tune_mode = self.select_tune_mode(json_info) self.__task_info[task_id] = op_json if tune_mode == NO_TUNE: if self.__process_num > self.compile_process_num: self.__process_num = self.compile_process_num if self.__pool is None: self.__pool = multiprocessing.Pool(processes=self.__process_num) + self.__reset_op_info = self.get_reset_op_info() + json_info["reset_op_info"] = self.__reset_op_info + op_json = json.dumps(json_info) task_future = self.__pool.apply_async(func=run_compiler, args=(op_json,)) self.__running_tasks.append((task_id, task_future)) else: