Caffe2 - C++ API
A deep learning, cross platform ML framework
onnxifi_graph_info.cc
1 #include "caffe2/onnx/onnxifi_graph_info.h"
2 #include "caffe2/core/logging.h"
3 
4 namespace caffe2 {
5 namespace onnx {
6 
7 SharedPtrBackendGraphInfo OnnxBackendGraphMap::lookup(const std::string& key) {
8  std::lock_guard<std::mutex> guard(backend_graph_map_lock_);
9  auto it = backend_graph_map_.find(key);
10  if (it != backend_graph_map_.end()) {
11  return it->second;
12  }
13  return nullptr;
14 }
15 
16 SharedPtrBackendGraphInfo OnnxBackendGraphMap::insert(
17  const std::string& key,
18  std::function<SharedPtrBackendGraphInfo()> creator) {
19  // First acquire lock.
20  std::lock_guard<std::mutex> guard(backend_graph_map_lock_);
21  // Then check if the backend_graph_info already exists in the map.
22  if (backend_graph_map_.find(key) != backend_graph_map_.end()) {
23  LOG(INFO) << "Reusing onnxifi backend for: " << key;
24  // If it already exists, return it.
25  // The onus is on the caller to release onnxGraph pointed by
26  // backend_graph_info
27  return backend_graph_map_[key];
28  }
29  LOG(INFO) << "Creating onnxifi backend for: " << key;
30  const auto ret_pair = backend_graph_map_.emplace(key, creator());
31  return ret_pair.first->second;
32 }
33 
34 void OnnxBackendGraphMap::remove(const std::string& key) {
35  SharedPtrBackendGraphInfo tmp;
36  {
37  std::lock_guard<std::mutex> guard(backend_graph_map_lock_);
38  auto it = backend_graph_map_.find(key);
39  if (it != backend_graph_map_.end()) {
40  if (it->second.unique()) {
41  LOG(INFO) << "Removing onnxifi backend for " << key;
42  tmp = it->second;
43  backend_graph_map_.erase(it);
44  }
45  }
46  }
47 }
48 
49 OnnxBackendGraphMap* getOnnxBackendGraphMap() {
50  static OnnxBackendGraphMap onnx_backend_graph_map;
51  return &onnx_backend_graph_map;
52 }
53 
54 } // namespace onnx
55 } // namespace caffe2
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13