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