1 #pragma once 2 #include <c10/macros/Export.h> 3 #include <torch/csrc/jit/ir/ir.h> 4 #include <mutex> 5 #include <string> 6 #include <unordered_map> 7 8 namespace torch::jit { 9 10 class UpgradersMap { 11 public: 12 void set_content( 13 std::unordered_map<std::string, std::shared_ptr<Graph>>&& content); 14 int count(); 15 const std::unordered_map<std::string, std::shared_ptr<Graph>>& get_content(); 16 bool is_populated(); 17 // THESE METHODS ARE ONLY USED FOR TESTING PURPOSES 18 void test_only_set_content( 19 const std::unordered_map<std::string, std::string>& content); 20 void test_only_remove_content( 21 const std::unordered_map<std::string, std::string>& content); 22 23 private: 24 std::unordered_map<std::string, std::shared_ptr<Graph>> content_; 25 std::mutex lock; 26 bool isPopulated = false; 27 }; 28 29 TORCH_API void populate_upgraders_map( 30 std::unordered_map<std::string, std::shared_ptr<Graph>>&& content); 31 32 TORCH_API int get_upgraders_map_size(); 33 34 TORCH_API bool is_upgraders_map_populated(); 35 36 TORCH_API const std::unordered_map<std::string, std::shared_ptr<Graph>>& 37 dump_upgraders_map(); 38 39 // THESE TWO METHODS BELOW ARE ONLY USED FOR TESTING 40 TORCH_API void test_only_populate_upgraders( 41 const std::unordered_map<std::string, std::string>& content); 42 43 TORCH_API void test_only_remove_upgraders( 44 const std::unordered_map<std::string, std::string>& content); 45 46 } // namespace torch::jit 47