diff --git a/mindspore/_extends/parallel_compile/tbe_compiler/compiler.py b/mindspore/_extends/parallel_compile/tbe_compiler/compiler.py index c385f7dee06..68dbebd5b67 100755 --- a/mindspore/_extends/parallel_compile/tbe_compiler/compiler.py +++ b/mindspore/_extends/parallel_compile/tbe_compiler/compiler.py @@ -28,7 +28,8 @@ build_in_impl_path = get_build_in_impl_path() # op function list op_build = "compile" op_pre_build = "pre_build" - +fusion_type_map = {'Convolution': 0, 'ElemWise': 1, 'CommReduce': 2, + 'Segment': 3, 'Opaque': 4} def _initialize(impl_path): """Initialize""" @@ -108,7 +109,7 @@ def build_op(build_type, json_str): # pre build if build_type == op_pre_build: - op_func(*inputs_args, *outputs_args, *attrs_args, kernel_name) + op_func(*inputs_args, *outputs_args, *attrs_args, kernel_name=kernel_name) # disable only pattern configuration op_build_cfg_en() return get_op_pattern() @@ -159,11 +160,16 @@ def compile_with_json(json_str): json_info = json.loads(json_str) if "fusion_op" in json_info: ret = compile_fusion_op(json_str) + elif "compile_type" in json_info: + ret = build_op(op_pre_build, json_str) else: ret = build_op(op_build, json_str) return ret - if __name__ == "__main__": in_args = sys.stdin.readline() - compile_with_json(in_args) + result = compile_with_json(in_args) + if result in fusion_type_map: + exit(fusion_type_map[result]) + else: + exit(100) diff --git a/mindspore/_extends/parallel_compile/tbe_compiler/tbe_process.py b/mindspore/_extends/parallel_compile/tbe_compiler/tbe_process.py index 9a3846c4f93..6dabc6121df 100644 --- a/mindspore/_extends/parallel_compile/tbe_compiler/tbe_process.py +++ b/mindspore/_extends/parallel_compile/tbe_compiler/tbe_process.py @@ -75,7 +75,6 @@ def check_supported(op_json: str): return ret - def run_compiler(op_json): """ run compiler to compile op with subprocess @@ -88,15 +87,16 @@ def run_compiler(op_json): """ try: tbe_compiler = os.path.join(os.path.split(os.path.realpath(__file__))[0], "compiler.py") - subprocess.run([sys.executable, tbe_compiler], input=op_json, timeout=300, - text=True, capture_output=True, check=True) - return "Success", "Success" + completed_object = subprocess.run([sys.executable, tbe_compiler], input=op_json, timeout=300, + text=True, capture_output=True, check=False) + if completed_object: + code = completed_object.returncode + return "Success", str(code) except subprocess.TimeoutExpired: tb = traceback.format_exc() - return "TBEException", "CompileTimeOut: " + tb + "\ninput_args: " + op_json + return "TBEException", "PreCompileTimeOut: " + tb + "\ninput_args: " + op_json except subprocess.CalledProcessError as e: - return "TBEException", "CompileProcessFailed:\n" + e.stdout + "\n" + e.stderr + "\ninput_args: " + op_json - + return "TBEException", "PreCompileProcessFailed:\n" + e.stdout + "\n" + e.stderr + "\ninput_args: " + op_json class CompilerPool: """compiler pool""" @@ -154,11 +154,11 @@ class CompilerPool: task_id, task_future = self.__running_tasks.pop(0) ret_type, result = task_future.get(330) if ret_type == "Success": - ret = task_id, "Success" + ret = task_id, "Success", result elif ret_type in ("Exception", "TBEException"): - ret = task_id, ret_type + ":" + result + ret = task_id, ret_type + ":" + result, "_" else: - ret = task_id, "Exception: Not support return type:" + str(ret_type) + ret = task_id, "Exception: Not support return type:" + str(ret_type), "_" return ret def reset_task_info(self): diff --git a/mindspore/ccsrc/device/ascend/kernel_build_ascend.cc b/mindspore/ccsrc/device/ascend/kernel_build_ascend.cc index afce5f3607d..68be42ab6dc 100644 --- a/mindspore/ccsrc/device/ascend/kernel_build_ascend.cc +++ b/mindspore/ccsrc/device/ascend/kernel_build_ascend.cc @@ -62,6 +62,31 @@ static kernel::KernelModPtr SerialCompileImpl(const AnfNodePtr &anf_node) { return kernel_mod_ptr; } +static bool KernelPreBuildParallelCompile(const mindspore::session::KernelGraph *kernel_graph_ptr) { + MS_EXCEPTION_IF_NULL(kernel_graph_ptr); + std::vector tbe_nodes; + for (const auto &anf_node : kernel_graph_ptr->execution_order()) { + MS_EXCEPTION_IF_NULL(anf_node); + if (!AnfAlgo::IsRealKernel(anf_node)) { + continue; + } + KernelType kernel_type = AnfAlgo::GetKernelType(anf_node); + switch (kernel_type) { + case KernelType::TBE_KERNEL: { + if (AnfAlgo::GetKernelMod(anf_node) == nullptr) { + tbe_nodes.push_back(anf_node); + } + break; + } + default: { + break; + } + } + } + bool ret = kernel::TbeOpParallelPreBuild(tbe_nodes); + return ret; +} + static bool KernelBuildParallelCompile(const mindspore::session::KernelGraph *kernel_graph_ptr) { MS_EXCEPTION_IF_NULL(kernel_graph_ptr); std::vector tbe_nodes; @@ -188,6 +213,12 @@ bool IsAtomicNode(const CNodePtr &kernel_node) { return atomic_flag; } +bool KernelPreBuild(const mindspore::session::KernelGraph *kernel_graph_ptr) { + MS_EXCEPTION_IF_NULL(kernel_graph_ptr); + bool ret = device::ascend::KernelPreBuildParallelCompile(kernel_graph_ptr); + return ret; +} + bool KernelBuild(const mindspore::session::KernelGraph *kernel_graph_ptr) { MS_EXCEPTION_IF_NULL(kernel_graph_ptr); TbeUtils::LoadCache(); diff --git a/mindspore/ccsrc/device/ascend/kernel_build_ascend.h b/mindspore/ccsrc/device/ascend/kernel_build_ascend.h index 5dea36a1835..d987b6ce7a8 100644 --- a/mindspore/ccsrc/device/ascend/kernel_build_ascend.h +++ b/mindspore/ccsrc/device/ascend/kernel_build_ascend.h @@ -22,6 +22,10 @@ namespace mindspore { namespace device { namespace ascend { +/** + * @brief kernel pre build for ascend. + */ +bool KernelPreBuild(const mindspore::session::KernelGraph *kernel_graph_ptr); /** * @brief kernel build for ascend. */ diff --git a/mindspore/ccsrc/kernel/kernel_fusion.cc b/mindspore/ccsrc/kernel/kernel_fusion.cc index 4e1ad97e230..be79eca15a8 100644 --- a/mindspore/ccsrc/kernel/kernel_fusion.cc +++ b/mindspore/ccsrc/kernel/kernel_fusion.cc @@ -102,7 +102,8 @@ std::map KernelFusion(const std::vector while (!build_manger->IsAllTaskFinish()) { int task_id = -1; char *task_result = nullptr; - auto ret = build_manger->WaitOne(&task_id, &task_result); + char *pre_build_result = nullptr; + auto ret = build_manger->WaitOne(&task_id, &task_result, &pre_build_result); if (!ret) { MS_EXCEPTION(ArgumentError) << "Build Failed. wait one ret:" << ret << ", task id:" << task_id; } diff --git a/mindspore/ccsrc/kernel/tbe/tbe_adapter.h b/mindspore/ccsrc/kernel/tbe/tbe_adapter.h index 0208d6c6a63..51c4cfd777f 100644 --- a/mindspore/ccsrc/kernel/tbe/tbe_adapter.h +++ b/mindspore/ccsrc/kernel/tbe/tbe_adapter.h @@ -27,7 +27,7 @@ // the TBE back-end operator implementation difference namespace mindspore { namespace kernel { -enum kCreaterType : int { SINGLE_BUILD = 0, PREBUILD, OP_SELECT_FORMAT, CHECK_SUPPORTED }; +enum kCreaterType : int { SINGLE_BUILD = 0, PREBUILD, OP_SELECT_FORMAT, CHECK_SUPPORTED, OP_PRE_COMPILE }; namespace tbe { using FAttrsPass = void (*)(const AnfNodePtr &anf_node, const std::vector> &op_info_attrs, nlohmann::json *attrs_json); diff --git a/mindspore/ccsrc/kernel/tbe/tbe_kernel_parallel_build.cc b/mindspore/ccsrc/kernel/tbe/tbe_kernel_parallel_build.cc index 577af45d595..e3973d2403a 100644 --- a/mindspore/ccsrc/kernel/tbe/tbe_kernel_parallel_build.cc +++ b/mindspore/ccsrc/kernel/tbe/tbe_kernel_parallel_build.cc @@ -42,6 +42,40 @@ constexpr auto kStartCompileOp = "start_compile_op"; constexpr auto kWaitOne = "wait_one"; constexpr auto kResetTaskInfo = "reset_task_info"; +bool TbeOpParallelPreBuild(const std::vector &anf_nodes) { + auto build_manger = std::make_shared(); + MS_EXCEPTION_IF_NULL(build_manger); + for (const auto &anf_node : anf_nodes) { + // gen kernel json + nlohmann::json kernel_json; + TbeKernelJsonCreator creator(OP_PRE_COMPILE); + if (!creator.GenTbeSingleKernelJson(anf_node, &kernel_json)) { + MS_LOG(ERROR) << "GenTbeSingleKernelJson failed"; + return false; + } + kernel_json["compile_type"] = "pre_build"; + // op build + auto task_id = build_manger->StartCompileOp(kernel_json); + build_manger->SavePreTaskInfo(task_id, anf_node); + } + while (!build_manger->IsAllPreTaskFinish()) { + int task_id = -1; + char *task_result = nullptr; + char *pre_build_result = nullptr; + auto ret = build_manger->WaitOne(&task_id, &task_result, &pre_build_result); + if (!ret) { + MS_EXCEPTION(ArgumentError) << "Pre Build Failed. wait one ret:" << ret << ", task id:" << task_id; + } + + if ((task_result != nullptr) && (strcmp(task_result, "Success") != 0)) { + MS_EXCEPTION(ArgumentError) << "task pre compile Failed, task id:" << task_id << ", cause:" << task_result; + } + + build_manger->PreTaskFinishProcess(task_id, pre_build_result); + } + return true; +} + bool TbeOpParallelBuild(std::vector anf_nodes) { auto build_manger = std::make_shared(); MS_EXCEPTION_IF_NULL(build_manger); @@ -82,7 +116,8 @@ bool TbeOpParallelBuild(std::vector anf_nodes) { while (!build_manger->IsAllTaskFinish()) { int task_id = -1; char *task_result = nullptr; - auto ret = build_manger->WaitOne(&task_id, &task_result); + char *pre_build_result = nullptr; + auto ret = build_manger->WaitOne(&task_id, &task_result, &pre_build_result); if (!ret) { MS_EXCEPTION(ArgumentError) << "Build Failed. wait one ret:" << ret << ", task id:" << task_id; } @@ -116,7 +151,7 @@ int32_t ParallelBuildManager::StartCompileOp(const nlohmann::json &kernel_json) return task_id; } -bool ParallelBuildManager::WaitOne(int *task_id, char **task_result) const { +bool ParallelBuildManager::WaitOne(int *task_id, char **task_result, char **pre_build_result) const { MS_LOG(INFO) << "wait task start."; MS_EXCEPTION_IF_NULL(task_id); MS_EXCEPTION_IF_NULL(task_result); @@ -128,10 +163,15 @@ bool ParallelBuildManager::WaitOne(int *task_id, char **task_result) const { MS_EXCEPTION(ArgumentError) << "Failed to call function wait_one"; return false; } - (void)PyArg_ParseTuple(pRes, "is", task_id, task_result); + (void)PyArg_ParseTuple(pRes, "iss", task_id, task_result, pre_build_result); return true; } +void ParallelBuildManager::SavePreTaskInfo(int32_t task_id, const mindspore::AnfNodePtr &anf_node) { + MS_LOG(INFO) << "SavePreTaskInfo, task id: " << task_id; + pre_task_map_[task_id] = anf_node; +} + void ParallelBuildManager::SaveTaskInfo(int32_t task_id, const mindspore::AnfNodePtr &anf_node, const std::string &json_name, const std::vector &input_size_list, const std::vector &output_size_list, int32_t scope_id) { @@ -150,11 +190,24 @@ void ParallelBuildManager::SaveTaskInfo(int32_t task_id, const mindspore::AnfNod task_map_[task_id] = task_info; } +bool ParallelBuildManager::IsAllPreTaskFinish() const { + MS_LOG(INFO) << "wait pre build process task_num: " << pre_task_map_.size(); + return pre_task_map_.empty(); +} + bool ParallelBuildManager::IsAllTaskFinish() const { MS_LOG(INFO) << "wait process task_num: " << task_map_.size(); return task_map_.empty(); } +void ParallelBuildManager::PreTaskFinishProcess(int32_t task_id, const std::string &pre_build_result) { + auto task_iter = pre_task_map_.find(task_id); + if (task_iter == pre_task_map_.end()) { + MS_EXCEPTION(ArgumentError) << "can find pre task_id:" << task_id; + } + (void)pre_task_map_.erase(task_iter); +} + std::pair ParallelBuildManager::TaskFinishProcess(int32_t task_id, bool set_kernel_mod) { auto task_iter = task_map_.find(task_id); if (task_iter == task_map_.end()) { diff --git a/mindspore/ccsrc/kernel/tbe/tbe_kernel_parallel_build.h b/mindspore/ccsrc/kernel/tbe/tbe_kernel_parallel_build.h index 776aa0b1fcf..c900baf036e 100644 --- a/mindspore/ccsrc/kernel/tbe/tbe_kernel_parallel_build.h +++ b/mindspore/ccsrc/kernel/tbe/tbe_kernel_parallel_build.h @@ -26,6 +26,7 @@ #include namespace mindspore { namespace kernel { +bool TbeOpParallelPreBuild(const std::vector &anf_nodes); bool TbeOpParallelBuild(std::vector anf_nodes); struct KernelBuildTaskInfo { @@ -42,6 +43,7 @@ class ParallelBuildManager { ParallelBuildManager(); ~ParallelBuildManager(); int32_t StartCompileOp(const nlohmann::json &kernel_json) const; + void SavePreTaskInfo(int32_t task_id, const AnfNodePtr &anf_node); void SaveTaskInfo(int32_t task_id, const AnfNodePtr &anf_node, const std::string &json_name, const std::vector &input_size_list, const std::vector &output_size_list, int32_t scope_id = 0); @@ -52,8 +54,10 @@ class ParallelBuildManager { const std::vector &input_size_list, const std::vector &output_size_list, AnfNode *node) const; - bool WaitOne(int *task_id, char **task_result) const; + bool WaitOne(int *task_id, char **task_result, char **pre_build_result) const; + bool IsAllPreTaskFinish() const; bool IsAllTaskFinish() const; + void PreTaskFinishProcess(int32_t task_id, const std::string &pre_build_result); std::pair TaskFinishProcess(int32_t task_id, bool set_kernel_mod = true); KernelModPtr GenKernelMod(const string &json_name, const string &processor, const std::vector &input_size_list, const std::vector &output_size_list, @@ -62,6 +66,7 @@ class ParallelBuildManager { private: PyObject *tbe_parallel_compiler_; + std::map pre_task_map_; std::map task_map_; std::vector same_op_list_; }; diff --git a/mindspore/ccsrc/session/ascend_session.cc b/mindspore/ccsrc/session/ascend_session.cc index 0665a6f76c7..bd497f6d112 100644 --- a/mindspore/ccsrc/session/ascend_session.cc +++ b/mindspore/ccsrc/session/ascend_session.cc @@ -546,6 +546,7 @@ void AscendSession::InitRuntimeResource() { void AscendSession::HardwareOptimize(const std::shared_ptr &kernel_graph) const { MS_LOG(INFO) << "HardwareOptimize start!"; + device::ascend::KernelPreBuild(kernel_graph.get()); opt::AscendBackendOptimization(kernel_graph); MS_EXCEPTION_IF_NULL(kernel_graph); kernel_graph->SetExecOrderByDefault();