From b61868b8bfb8dca1513186fd40c056d8f5d29084 Mon Sep 17 00:00:00 2001 From: LaiYongqiang Date: Thu, 8 Jul 2021 18:56:07 +0800 Subject: [PATCH] add reset op info for tune --- .../_extends/parallel_compile/tbe_compiler/tbe_process.py | 4 ++++ mindspore/_extends/parallel_compile/tbe_compiler/tuner.py | 5 +++-- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/mindspore/_extends/parallel_compile/tbe_compiler/tbe_process.py b/mindspore/_extends/parallel_compile/tbe_compiler/tbe_process.py index 36d7e9581de..335e40e3975 100644 --- a/mindspore/_extends/parallel_compile/tbe_compiler/tbe_process.py +++ b/mindspore/_extends/parallel_compile/tbe_compiler/tbe_process.py @@ -29,6 +29,7 @@ from mindspore import log from .tbe_common import check_kernel_info, TBEException from .helper import _op_select_format, _check_supported + # tune type NO_TUNE = "NO_TUNE" GA_TUNE = "GA" @@ -355,7 +356,10 @@ class TbeProcess: log.error("Auto tune init failed, place check your hardware config or go back to normal compile!") self.tune_init = False return error_id + self.__reset_op_info = self.get_reset_op_info() self.__tuner.tune_init = True + json_info["reset_op_info"] = self.__reset_op_info + op_json = json.dumps(json_info) self.__all_tune_tasks.append(task_id) self.__running_tune_tasks.append(task_id) diff --git a/mindspore/_extends/parallel_compile/tbe_compiler/tuner.py b/mindspore/_extends/parallel_compile/tbe_compiler/tuner.py index 886bb7eb0d5..22a905972c8 100644 --- a/mindspore/_extends/parallel_compile/tbe_compiler/tuner.py +++ b/mindspore/_extends/parallel_compile/tbe_compiler/tuner.py @@ -18,8 +18,8 @@ import datetime import json import sys import traceback +from tbe.common.rl_bank.bank_manager import set_current_op_name from te.platform.cce_conf import te_set_version -from te_fusion.fusion_manager import set_current_op_name from te_fusion.fusion_util import fusion_op, dump_fusion_json from te_fusion.parallel_compilation import init_multi_process_env, get_finished_compilation_task, \ deinit_multi_process_env, start_ga_multi_process @@ -331,13 +331,14 @@ class TbeTuner: raise ValueError("Json string Errors, key:fusion_op not found.") kernel_name = json_info["fusion_op"]["fusion_op_name"] full_name = json_info["fusion_op"]["full_name"] + reset_op_info = json_info["reset_op_info"] set_current_op_name(kernel_name) converted_json = fusion_to_fusion(json.dumps(json_info), tune_mode="RL") job_type = RL_COMPILE base_kernel = './kernel_meta/' + kernel_name + '.o' compile_info = None try: - fusion_op(converted_json) + fusion_op(converted_json, reset_op_info=reset_op_info) # pylint: disable=broad-except except Exception: exc_type, exc_value, _ = sys.exc_info()