xref: /aosp_15_r20/external/pytorch/c10/util/DynamicCounter.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #include <c10/util/DynamicCounter.h>
2 
3 #include <c10/util/Synchronized.h>
4 
5 #include <stdexcept>
6 #include <string>
7 #include <unordered_set>
8 #include <vector>
9 
10 namespace c10::monitor {
11 
12 namespace {
13 using DynamicCounterBackends =
14     std::vector<std::shared_ptr<detail::DynamicCounterBackendIf>>;
15 
dynamicCounterBackends()16 Synchronized<DynamicCounterBackends>& dynamicCounterBackends() {
17   static auto instance = new Synchronized<DynamicCounterBackends>();
18   return *instance;
19 }
20 
registeredCounters()21 Synchronized<std::unordered_set<std::string>>& registeredCounters() {
22   static auto instance = new Synchronized<std::unordered_set<std::string>>();
23   return *instance;
24 }
25 } // namespace
26 
27 namespace detail {
registerDynamicCounterBackend(std::unique_ptr<DynamicCounterBackendIf> backend)28 void registerDynamicCounterBackend(
29     std::unique_ptr<DynamicCounterBackendIf> backend) {
30   dynamicCounterBackends().withLock(
31       [&](auto& backends) { backends.push_back(std::move(backend)); });
32 }
33 } // namespace detail
34 
35 struct DynamicCounter::Guard {
Guardc10::monitor::DynamicCounter::Guard36   Guard(std::string_view key, Callback&& getCounterCallback)
37       : key_{key},
38         getCounterCallback_(std::move(getCounterCallback)),
39         backends_{dynamicCounterBackends().withLock(
__anon5cf5b6150302c10::monitor::DynamicCounter::Guard40             [](auto& backends) { return backends; })} {
__anon5cf5b6150402(auto& registeredCounters) 41     registeredCounters().withLock([&](auto& registeredCounters) {
42       if (!registeredCounters.insert(std::string(key)).second) {
43         throw std::logic_error(
44             "Counter " + std::string(key) + " already registered");
45       }
46     });
47 
48     for (const auto& backend : backends_) {
49       // Avoid copying the user-provided callback to avoid unexpected behavior
50       // changes when more than one backend is registered.
__anon5cf5b6150502() 51       backend->registerCounter(key, [&]() { return getCounterCallback_(); });
52     }
53   }
54 
~Guardc10::monitor::DynamicCounter::Guard55   ~Guard() {
56     for (const auto& backend : backends_) {
57       backend->unregisterCounter(key_);
58     }
59 
60     registeredCounters().withLock(
61         [&](auto& registeredCounters) { registeredCounters.erase(key_); });
62   }
63 
64  private:
65   std::string key_;
66   Callback getCounterCallback_;
67   DynamicCounterBackends backends_;
68 };
69 
DynamicCounter(std::string_view key,Callback getCounterCallback)70 DynamicCounter::DynamicCounter(
71     std::string_view key,
72     Callback getCounterCallback)
73     : guard_{std::make_unique<Guard>(key, std::move(getCounterCallback))} {}
74 
75 DynamicCounter::~DynamicCounter() = default;
76 
77 } // namespace c10::monitor
78