1 /** 2 * Copyright 2019 Huawei Technologies Co., Ltd 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 #ifndef MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_TBE_TBE_KERNEL_PARALLEL_BUILD_H_ 18 #define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_TBE_TBE_KERNEL_PARALLEL_BUILD_H_ 19 20 #include <utility> 21 #include <string> 22 #include <map> 23 #include <vector> 24 #include <nlohmann/json.hpp> 25 26 #include "backend/kernel_compiler/kernel.h" 27 #include "backend/session/kernel_build_client.h" 28 29 namespace mindspore { 30 namespace kernel { 31 bool TbeOpParallelBuild(const std::vector<AnfNodePtr> &anf_nodes); 32 33 struct KernelBuildTaskInfo { 34 AnfNodePtr node; 35 std::string processor; 36 std::string json_name; 37 std::vector<size_t> input_size_list; 38 std::vector<size_t> output_size_list; 39 int64_t scope_id; 40 }; 41 42 class ParallelBuildManager { 43 public: ParallelBuildManager()44 ParallelBuildManager() { AscendKernelBuildClient::Instance().TbeReset(); } 45 ~ParallelBuildManager(); 46 void SaveTaskInfo(int32_t task_id, const AnfNodePtr &anf_node, const std::string &json_name, 47 const std::vector<size_t> &input_size_list, const std::vector<size_t> &output_size_list, 48 int64_t scope_id = 0); 49 void SavePreBuildTaskInfo(int32_t task_id, const AnfNodePtr &anf_node, const std::string &json_name); 50 void SaveSameOpInfo(const AnfNodePtr &anf_node, const std::string &json_name, 51 const std::vector<size_t> &input_size_list, const std::vector<size_t> &output_size_list); 52 void SaveSameFusionOpInfo(const int64_t scope_id, const std::string &json_name, const std::string &processor, 53 const std::vector<size_t> &input_size_list, const std::vector<size_t> &output_size_list); 54 bool GenSameOpKernelMod() const; 55 bool GenSameFusionOpKernelMod(std::map<int64_t, KernelModPtr> *kernel_mode_ret) const; 56 bool SearchInCache(const std::string &json_name, const std::vector<size_t> &input_size_list, 57 const std::vector<size_t> &output_size_list, AnfNode *node) const; 58 bool IsAllTaskFinish() const; 59 void PreTaskFinishProcess(int32_t task_id, const std::string &pre_build_result); 60 std::pair<int32_t, KernelModPtr> TaskFinishProcess(int32_t task_id, const std::string &build_ret, 61 bool set_kernel_mod = true); 62 KernelModPtr GenKernelMod(const std::vector<size_t> &input_size_list, const std::vector<size_t> &output_size_list, 63 const KernelPackPtr &kernel_pack) const; 64 65 // Interactive with real backend, who could be implemented by Python. 66 static int StartCompileOp(const nlohmann::json &kernel_json); 67 std::string ProcessTbeJob(const nlohmann::json &kernel_json); 68 static bool WaitOne(int *task_id, std::string *task_result, std::string *build_result); 69 void ResetTaskInfo() noexcept; 70 AnfNodePtr GetAnfNodeByTaskID(int32_t task_id); 71 72 private: 73 std::map<int32_t, AnfNodePtr> pre_task_map_; 74 std::map<int32_t, KernelBuildTaskInfo> pre_build_task_map_; 75 std::map<int32_t, KernelBuildTaskInfo> task_map_; 76 std::vector<KernelBuildTaskInfo> same_op_list_; 77 }; 78 } // namespace kernel 79 } // namespace mindspore 80 81 #endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_TBE_TBE_KERNEL_PARALLEL_BUILD_H_ 82