• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /**
2  * Copyright 2022 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 #ifndef MINDSPORE_CCSRC_BACKEND_OPTIMIZER_GRAPH_KERNEL_SPLIT_MODEL_SPLIT_MODEL_FACTORY_H_
17 #define MINDSPORE_CCSRC_BACKEND_OPTIMIZER_GRAPH_KERNEL_SPLIT_MODEL_SPLIT_MODEL_FACTORY_H_
18 
19 #include <memory>
20 #include <string>
21 #include "backend/common/graph_kernel/split_model/split_model.h"
22 #include "utils/hash_map.h"
23 #include "include/backend/visible.h"
24 
25 namespace mindspore::graphkernel::inner {
26 class BACKEND_EXPORT SplitModelFactory {
27  public:
Instance()28   static SplitModelFactory &Instance() {
29     static SplitModelFactory instance = SplitModelFactory();
30     return instance;
31   }
32   SplitModelPtr CreateSplitModel(const std::string &processor);
33   using RegFunc = std::function<std::shared_ptr<SplitModel>()>;
Register(const std::string & processor,const RegFunc & func)34   void Register(const std::string &processor, const RegFunc &func) { creators[processor] = func; }
35 
36  private:
37   mindspore::HashMap<std::string, RegFunc> creators;
38 };
39 
40 class SplitModelRegister {
41  public:
SplitModelRegister(const std::string & processor,const SplitModelFactory::RegFunc & func)42   SplitModelRegister(const std::string &processor, const SplitModelFactory::RegFunc &func) : func_(func) {
43     SplitModelFactory::Instance().Register(processor, func);
44   }
45   ~SplitModelRegister() = default;
46 
47  protected:
48   // for pclint-plus
49   SplitModelFactory::RegFunc func_;
50 };
51 
52 #define SPLIT_MODEL_JOIN(x, y) x##y
53 #define SPLIT_MODEL_UNIQUE_NAME(prefix, cnt) SPLIT_MODEL_JOIN(prefix, cnt)
54 #define SPLIT_MODEL_REGISTER(processor, cls, ...)                                                                     \
55   const mindspore::graphkernel::inner::SplitModelRegister SPLIT_MODEL_UNIQUE_NAME(split_model_, __COUNTER__)(         \
56     processor, [__VA_ARGS__]() noexcept {                                                                             \
57       return std::static_pointer_cast<mindspore::graphkernel::inner::SplitModel>(std::make_shared<cls>(__VA_ARGS__)); \
58     })
59 }  // namespace mindspore::graphkernel::inner
60 #endif  // MINDSPORE_CCSRC_BACKEND_OPTIMIZER_GRAPH_KERNEL_SPLIT_MODEL_SPLIT_MODEL_FACTORY_H_
61