1 /** 2 * Copyright 2020 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_PS_UTIL_H_ 18 #define MINDSPORE_CCSRC_PS_UTIL_H_ 19 20 #include <map> 21 #include <vector> 22 #include <string> 23 #include <unordered_map> 24 #include "frontend/optimizer/optimizer.h" 25 #include "backend/session/anf_runtime_algorithm.h" 26 #include "backend/kernel_compiler/common_utils.h" 27 #include "backend/kernel_compiler/cpu/sparse_optimizer_cpu_kernel.h" 28 29 namespace mindspore { 30 namespace ps { 31 enum ParamType { kUnKnown = 0, kWeight = 1, kAccumulation = 2 }; 32 33 struct ParamInitInfo { 34 std::string param_name_; 35 ParamType param_type_{kUnKnown}; 36 size_t global_seed_{0}; 37 size_t op_seed_{0}; 38 float init_val_{0}; 39 }; 40 41 constexpr size_t kNodeInputWeightNameOffset = 1; 42 constexpr size_t kNodeInputWeightIndexOffset = 2; 43 44 class Util { 45 public: 46 static bool IsRoleOfPServer(); 47 static bool IsRoleOfScheduler(); 48 static int64_t optimizer_id(const std::string &name); 49 static std::string optimizer_name(int64_t id); 50 static std::string optimizer_node_name(int64_t id); 51 static bool is_optimizer(const std::string &name); 52 static int64_t LocalShard(int64_t first_dim, int64_t rank_id, int64_t server_num); 53 static std::map<int64_t, int64_t> AllRankLocalShard(int64_t first_dim, int64_t rank_id, int64_t server_num); 54 static void ReduceSparseGradient(float *gradients, int *indices, const size_t indices_size, size_t segment_size, 55 const size_t first_dim_size, const size_t outer_dim_size, 56 mindspore::kernel::SparseGradient<int> *unique_sparse_grad); 57 static bool FuseServerCommOps(const pipeline::ResourcePtr &res); 58 59 private: 60 static void DoFusion(const FuncGraphPtr &func_graph, const std::string &cnode_name, 61 const std::string &fused_cnode_name); 62 static kernel::KernelBuildInfoPtr GenerateKernelBuildInfo(const std::vector<AnfNodePtr> &node_list); 63 64 static std::unordered_map<std::string, int64_t> optimizer_to_ids; 65 static std::unordered_map<int64_t, std::string> id_to_optimizers; 66 static std::unordered_map<int64_t, std::string> id_to_optimizer_nodes; 67 static int64_t rank_id_; 68 }; 69 } // namespace ps 70 } // namespace mindspore 71 #endif // MINDSPORE_CCSRC_PS_UTIL_H_ 72