• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /**
2  * Copyright 2024 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 #include "ir/func_graph.h"
18 #include "gtest/gtest.h"
19 
20 #ifndef MINDSPORE_UT_COMMON_RESOURCE_H
21 #define MINDSPORE_UT_COMMON_RESOURCE_H
22 namespace UT {
23 using UTKeyInfo = std::pair<std::string, std::string>;
24 
25 class UTResourceManager {
26  public:
27   UTResourceManager() = default;
28   ~UTResourceManager() {
29     for (const auto &it : all_func_graphs_) {
30       auto key_info = it.first;
31       std::cout << "Unexpected unreleased func graph resource of case:" << key_info.first << "." << key_info.second
32                 << std::endl;
33     }
34     if (!all_func_graphs_.empty()) {
35       std::cout << "Please check `TearDown` function of testcase, and make sure all func graphs can be dropped after "
36                    "case executed, otherwise core dumped might occur."
37                 << std::endl;
38     }
39   }
40 
41   void HoldFuncGraph(const mindspore::FuncGraphPtr &fg) {
42     const char *suite_name = testing::UnitTest::GetInstance()->current_test_suite()->name();
43     const char *test_name = testing::UnitTest::GetInstance()->current_test_info()->name();
44     auto new_fg = std::make_shared<mindspore::FuncGraph>();
45     std::cout << "Hold func graph of case:" << suite_name << "." << test_name << std::endl;
46     (void)all_func_graphs_[UTKeyInfo{suite_name, test_name}].insert(fg);
47   }
48 
49   mindspore::FuncGraphPtr MakeAndHoldFuncGraph() {
50     auto func_graph = std::make_shared<mindspore::FuncGraph>();
51     HoldFuncGraph(func_graph);
52     return func_graph;
53   }
54 
55   void DropFuncGraph(const UTKeyInfo &ut_info) {
56     if (all_func_graphs_.find(ut_info) == all_func_graphs_.cend()) {
57       return;
58     }
59     std::cout << "Drop func graph of case:" << ut_info.first << "." << ut_info.second << std::endl;
60     (void)all_func_graphs_.erase(ut_info);
61   }
62 
63   void DropAllFuncGraphs() { all_func_graphs_.clear(); }
64 
65   static std::shared_ptr<UTResourceManager> GetInstance();
66 
67  private:
68   static std::shared_ptr<UTResourceManager> inst_resource_manager_;
69   std::map<UTKeyInfo, std::set<mindspore::FuncGraphPtr>> all_func_graphs_;
70 };
71 
72 }  // namespace UT
73 
74 #endif  // MINDSPORE_UT_COMMON_RESOURCE_H
75