• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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_FL_SERVER_KERNEL_ROUND_ROUND_KERNEL_FACTORY_H_
18 #define MINDSPORE_CCSRC_FL_SERVER_KERNEL_ROUND_ROUND_KERNEL_FACTORY_H_
19 
20 #include <memory>
21 #include <string>
22 #include <utility>
23 #include <unordered_map>
24 #include "fl/server/common.h"
25 #include "fl/server/kernel/round/round_kernel.h"
26 
27 namespace mindspore {
28 namespace fl {
29 namespace server {
30 namespace kernel {
31 using RoundKernelCreator = std::function<std::shared_ptr<RoundKernel>()>;
32 // Kernel factory of round kernels.
33 class RoundKernelFactory {
34  public:
35   static RoundKernelFactory &GetInstance();
36   void Register(const std::string &name, RoundKernelCreator &&creator);
37   std::shared_ptr<RoundKernel> Create(const std::string &name);
38 
39  private:
40   RoundKernelFactory() = default;
41   ~RoundKernelFactory() = default;
42   RoundKernelFactory(const RoundKernelFactory &) = delete;
43   RoundKernelFactory &operator=(const RoundKernelFactory &) = delete;
44 
45   std::unordered_map<std::string, RoundKernelCreator> name_to_creator_map_;
46 };
47 
48 class RoundKernelRegister {
49  public:
RoundKernelRegister(const std::string & name,RoundKernelCreator && creator)50   RoundKernelRegister(const std::string &name, RoundKernelCreator &&creator) {
51     RoundKernelFactory::GetInstance().Register(name, std::move(creator));
52   }
53   ~RoundKernelRegister() = default;
54 };
55 
56 #define REG_ROUND_KERNEL(NAME, CLASS)                                                        \
57   static_assert(std::is_base_of<RoundKernel, CLASS>::value, " must be base of RoundKernel"); \
58   static const RoundKernelRegister g_##NAME##_round_kernel_reg(#NAME, []() { return std::make_shared<CLASS>(); });
59 }  // namespace kernel
60 }  // namespace server
61 }  // namespace fl
62 }  // namespace mindspore
63 #endif  // MINDSPORE_CCSRC_FL_SERVER_KERNEL_ROUND_ROUND_KERNEL_FACTORY_H_
64