1 /** 2 * Copyright 2021 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_AKG_AKG_KERNEL_BUILD_H_ 18 #define MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_AKG_AKG_KERNEL_BUILD_H_ 19 20 #include <sys/shm.h> 21 22 #include <string> 23 #include <utility> 24 #include <vector> 25 #include <map> 26 #include <set> 27 #include "ir/anf.h" 28 #include "backend/kernel_compiler/kernel.h" 29 #include "backend/session/kernel_build_client.h" 30 #include "backend/kernel_compiler/akg/akg_kernel_json_generator.h" 31 32 namespace mindspore { 33 namespace kernel { 34 using JsonNodePair = std::pair<AkgKernelJsonGenerator, AnfNodePtr>; 35 36 class AkgKernelBuilder { 37 public: 38 AkgKernelBuilder() = default; 39 virtual ~AkgKernelBuilder() = default; 40 41 virtual KernelBuildClient *GetClient() = 0; 42 virtual KernelPackPtr AkgSearchCache(const std::string &kernel_name) = 0; 43 virtual KernelPackPtr AkgInsertCache(const std::string &kernel_name) = 0; 44 virtual void AkgSetKernelMod(const KernelPackPtr &kernel_pack, const AkgKernelJsonGenerator &json_generator, 45 const AnfNodePtr &anf_node) = 0; 46 virtual void AkgSaveJsonInfo(const string &kernel_name, const string &kernel_json) = 0; 47 bool AkgKernelParallelBuild(const std::vector<AnfNodePtr> &anf_nodes); 48 49 private: 50 std::vector<JsonNodePair> GetNotCachedKernels(const std::vector<JsonNodePair> &build_args); 51 std::vector<std::string> GetKernelJsonsByHashId(const std::vector<JsonNodePair> &build_args, 52 std::set<size_t> fetched_ids); 53 bool InsertToCache(const std::vector<JsonNodePair> &build_args); 54 bool HandleRepeatNodes(); 55 bool AkgOpParallelBuild(const std::vector<JsonNodePair> &build_args); 56 std::vector<JsonNodePair> repeat_nodes_; 57 std::string CollectBuildAttrs(); 58 }; 59 60 class AkgKernelPool { 61 public: 62 class LockMng { 63 public: LockMng(int32_t fd)64 explicit LockMng(int32_t fd) { 65 fd_ = fd; 66 locked_ = TryLock(); 67 } 68 ~LockMng()69 virtual ~LockMng() { 70 if (locked_) { 71 Unlock(); 72 } 73 } 74 75 bool locked_{false}; 76 77 private: 78 bool TryLock() const; 79 void Unlock() const; 80 81 int32_t fd_{-1}; 82 }; 83 84 public: 85 AkgKernelPool() = default; ~AkgKernelPool()86 virtual ~AkgKernelPool() { 87 // Close key file 88 if (fd_ != -1) { 89 (void)close(fd_); 90 } 91 } 92 93 int32_t Init(const std::vector<JsonNodePair> &build_args); 94 int32_t Release() const; 95 int32_t FetchKernels(std::set<size_t> *out); 96 int32_t UpdateAndWait(const std::set<size_t> &ids); 97 98 constexpr inline static size_t kMaxKernelNum_{1000}; 99 100 // allocate memory for todo_list, doing_list, done_list 101 constexpr inline static size_t kListNum_{3}; 102 103 constexpr inline static auto kKeyName_ = "./akg_build_tmp.key"; 104 105 constexpr inline static int32_t kToDoIdx_ = 0; 106 constexpr inline static int32_t kDoingIdx_ = 1; 107 constexpr inline static int32_t kDoneIdx_ = 2; 108 109 private: ListBegin(int32_t list_idx)110 inline size_t *ListBegin(int32_t list_idx) const { return kernel_lists_[list_idx]; } 111 ListEnd(int32_t list_idx)112 inline size_t *ListEnd(int32_t list_idx) const { 113 return kernel_lists_[list_idx] + kernel_lists_[list_idx][kMaxKernelNum_]; 114 } 115 ResetListSize(int32_t list_idx,size_t val)116 inline void ResetListSize(int32_t list_idx, size_t val) { kernel_lists_[list_idx][kMaxKernelNum_] = val; } 117 IncListSize(int32_t list_idx,size_t val)118 inline void IncListSize(int32_t list_idx, size_t val) { kernel_lists_[list_idx][kMaxKernelNum_] += val; } 119 120 void *CreateSharedMem(const std::string &path); 121 std::string GetCurrentPath() const; 122 InitKernelLists(void * addr)123 inline void InitKernelLists(void *addr) { 124 kernel_lists_[kToDoIdx_] = reinterpret_cast<size_t *>(addr); 125 kernel_lists_[kDoingIdx_] = kernel_lists_[kToDoIdx_] + kMaxKernelNum_ + 1; 126 kernel_lists_[kDoneIdx_] = kernel_lists_[kDoingIdx_] + kMaxKernelNum_ + 1; 127 } 128 129 int32_t AddKernels(const std::vector<JsonNodePair> &kernel_jsons); 130 int32_t Wait() const; 131 132 int32_t shm_id_{-1}; 133 bool is_creator_{false}; 134 int32_t fd_{-1}; 135 136 // includes 3 lists: todo_list, doing_list, done_list. 137 // each list has kMaxKernelNum_ + 1 elements and, the count of elements in each list 138 // is stored in kernel_lists_[xx][kMaxKernelNum_] 139 size_t *kernel_lists_[kListNum_]{nullptr, nullptr, nullptr}; 140 141 std::set<size_t> self_kernel_ids_; 142 }; 143 } // namespace kernel 144 } // namespace mindspore 145 146 #endif // MINDSPORE_CCSRC_BACKEND_KERNEL_COMPILER_AKG_AKG_KERNEL_BUILD_H_ 147