mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
ad8aef0f98
Mostly in torch/csrc/jit/runtime and in `ATen/cuda/` Pull Request resolved: https://github.com/pytorch/pytorch/pull/110314 Approved by: https://github.com/seemethere
47 lines
1.4 KiB
C++
47 lines
1.4 KiB
C++
#pragma once
|
|
#include <c10/macros/Export.h>
|
|
#include <torch/csrc/jit/ir/ir.h>
|
|
#include <mutex>
|
|
#include <string>
|
|
#include <unordered_map>
|
|
|
|
namespace torch::jit {
|
|
|
|
class UpgradersMap {
|
|
public:
|
|
void set_content(
|
|
std::unordered_map<std::string, std::shared_ptr<Graph>>&& content);
|
|
int count();
|
|
const std::unordered_map<std::string, std::shared_ptr<Graph>>& get_content();
|
|
bool is_populated();
|
|
// THESE METHODS ARE ONLY USED FOR TESTING PURPOSES
|
|
void test_only_set_content(
|
|
const std::unordered_map<std::string, std::string>& content);
|
|
void test_only_remove_content(
|
|
const std::unordered_map<std::string, std::string>& content);
|
|
|
|
private:
|
|
std::unordered_map<std::string, std::shared_ptr<Graph>> content_;
|
|
std::mutex lock;
|
|
bool isPopulated = false;
|
|
};
|
|
|
|
TORCH_API void populate_upgraders_map(
|
|
std::unordered_map<std::string, std::shared_ptr<Graph>>&& content);
|
|
|
|
TORCH_API int get_upgraders_map_size();
|
|
|
|
TORCH_API bool is_upgraders_map_populated();
|
|
|
|
TORCH_API const std::unordered_map<std::string, std::shared_ptr<Graph>>&
|
|
dump_upgraders_map();
|
|
|
|
// THESE TWO METHODS BELOW ARE ONLY USED FOR TESTING
|
|
TORCH_API void test_only_populate_upgraders(
|
|
const std::unordered_map<std::string, std::string>& content);
|
|
|
|
TORCH_API void test_only_remove_upgraders(
|
|
const std::unordered_map<std::string, std::string>& content);
|
|
|
|
} // namespace torch::jit
|