|
|
@ -40,6 +40,7 @@ constexpr auto kParallelCompileModule = "mindspore._extends.parallel_compile.tbe
|
|
|
|
constexpr auto kCreateParallelCompiler = "create_tbe_parallel_compiler";
|
|
|
|
constexpr auto kCreateParallelCompiler = "create_tbe_parallel_compiler";
|
|
|
|
constexpr auto kStartCompileOp = "start_compile_op";
|
|
|
|
constexpr auto kStartCompileOp = "start_compile_op";
|
|
|
|
constexpr auto kWaitOne = "wait_one";
|
|
|
|
constexpr auto kWaitOne = "wait_one";
|
|
|
|
|
|
|
|
constexpr auto kResetTaskInfo = "reset_task_info";
|
|
|
|
|
|
|
|
|
|
|
|
bool TbeOpParallelBuild(std::vector<AnfNodePtr> anf_nodes) {
|
|
|
|
bool TbeOpParallelBuild(std::vector<AnfNodePtr> anf_nodes) {
|
|
|
|
auto build_manger = std::make_shared<ParallelBuildManager>();
|
|
|
|
auto build_manger = std::make_shared<ParallelBuildManager>();
|
|
|
@ -96,6 +97,8 @@ bool TbeOpParallelBuild(std::vector<AnfNodePtr> anf_nodes) {
|
|
|
|
|
|
|
|
|
|
|
|
ParallelBuildManager::ParallelBuildManager() { tbe_parallel_compiler_ = TbePythonFuncs::TbeParallelCompiler(); }
|
|
|
|
ParallelBuildManager::ParallelBuildManager() { tbe_parallel_compiler_ = TbePythonFuncs::TbeParallelCompiler(); }
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
ParallelBuildManager::~ParallelBuildManager() { ResetTaskInfo(); }
|
|
|
|
|
|
|
|
|
|
|
|
int32_t ParallelBuildManager::StartCompileOp(const nlohmann::json &kernel_json) const {
|
|
|
|
int32_t ParallelBuildManager::StartCompileOp(const nlohmann::json &kernel_json) const {
|
|
|
|
PyObject *pRes = nullptr;
|
|
|
|
PyObject *pRes = nullptr;
|
|
|
|
PyObject *pArgs = PyTuple_New(1);
|
|
|
|
PyObject *pArgs = PyTuple_New(1);
|
|
|
@ -234,5 +237,16 @@ KernelModPtr ParallelBuildManager::GenKernelMod(const string &json_name, const s
|
|
|
|
kernel_mod_ptr->SetWorkspaceSizeList(kernel_json_info.workspaces);
|
|
|
|
kernel_mod_ptr->SetWorkspaceSizeList(kernel_json_info.workspaces);
|
|
|
|
return kernel_mod_ptr;
|
|
|
|
return kernel_mod_ptr;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
void ParallelBuildManager::ResetTaskInfo() {
|
|
|
|
|
|
|
|
if (task_map_.empty()) {
|
|
|
|
|
|
|
|
MS_LOG(INFO) << "All tasks are compiled success.";
|
|
|
|
|
|
|
|
return;
|
|
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
task_map_.clear();
|
|
|
|
|
|
|
|
same_op_list_.clear();
|
|
|
|
|
|
|
|
PyObject *pArg = Py_BuildValue("()");
|
|
|
|
|
|
|
|
(void)PyObject_CallMethod(tbe_parallel_compiler_, kResetTaskInfo, "O", pArg);
|
|
|
|
|
|
|
|
}
|
|
|
|
} // namespace kernel
|
|
|
|
} // namespace kernel
|
|
|
|
} // namespace mindspore
|
|
|
|
} // namespace mindspore
|
|
|
|