xref: /aosp_15_r20/external/pytorch/c10/util/Gauge.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <memory>
4 #include <string_view>
5 
6 #include <c10/macros/Macros.h>
7 #include <c10/util/SmallVector.h>
8 
9 namespace c10::monitor {
10 namespace detail {
11 
12 class GaugeImpl;
13 
14 class GaugeBackendIf {
15  public:
16   virtual ~GaugeBackendIf() = default;
17   virtual void record(int64_t value) noexcept = 0;
18 };
19 
20 class GaugeBackendFactoryIf {
21  public:
22   virtual ~GaugeBackendFactoryIf() = default;
23 
24   // May return nullptr if the gauge will be ignored by the given backend.
25   virtual std::unique_ptr<GaugeBackendIf> create(
26       std::string_view key) noexcept = 0;
27 };
28 
29 void C10_API registerGaugeBackend(std::unique_ptr<GaugeBackendFactoryIf>);
30 } // namespace detail
31 
32 // A handle to a Gauge.
33 class C10_API GaugeHandle {
34  public:
35   explicit GaugeHandle(std::string_view key);
36   void record(int64_t value);
37 
38  private:
39   detail::GaugeImpl& impl_;
40 };
41 
42 } // namespace c10::monitor
43 
44 #define STATIC_GAUGE(_key)                            \
45   []() -> ::c10::monitor::GaugeHandle& {              \
46     static ::c10::monitor::GaugeHandle handle(#_key); \
47     return handle;                                    \
48   }()
49