pgjmblvch
/
mindspore2022
forked from huawei/mindspore2022
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

!19251 support virtual computing for tbe op 

Merge pull request !19251 from laiyongqiang/tbe_virtual
This commit is contained in:
i-robot 2021-07-02 09:28:55 +00:00 committed by Gitee
parent d4d3a00e58 ca7973daa5
commit d067d55f0f
2 changed files with 33 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
mindspore/_extends/parallel_compile/tbe_compiler/compiler.py
View File

@ -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):

28
mindspore/_extends/parallel_compile/tbe_compiler/tbe_process.py
View File

@ -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: