1 | /* |
---|---|
2 | * Copyright (c) Meta Platforms, Inc. and affiliates. |
3 | * All rights reserved. |
4 | * |
5 | * This source code is licensed under the BSD-style license found in the |
6 | * LICENSE file in the root directory of this source tree. |
7 | */ |
8 | |
9 | // Mediator for initialization and profiler control |
10 | |
11 | #pragma once |
12 | |
13 | #include <atomic> |
14 | #include <chrono> |
15 | #include <functional> |
16 | #include <memory> |
17 | #include <mutex> |
18 | #include <string> |
19 | #include <set> |
20 | #include <thread> |
21 | #include <vector> |
22 | #include <deque> |
23 | |
24 | #include "ActivityProfilerInterface.h" |
25 | #include "ActivityType.h" |
26 | #include "ClientInterface.h" |
27 | #include "GenericTraceActivity.h" |
28 | #include "TraceSpan.h" |
29 | #include "IActivityProfiler.h" |
30 | #include "ActivityTraceInterface.h" |
31 | |
32 | #include "ThreadUtil.h" |
33 | |
34 | extern "C"{ |
35 | void suppressLibkinetoLogMessages(); |
36 | int InitializeInjection(void); |
37 | void libkineto_init(bool cpuOnly, bool logOnError); |
38 | } |
39 | |
40 | namespace libkineto { |
41 | |
42 | class Config; |
43 | class ConfigLoader; |
44 | |
45 | struct CpuTraceBuffer { |
46 | template <class... Args> |
47 | void emplace_activity(Args&&... args) { |
48 | activities.emplace_back( |
49 | std::make_unique<GenericTraceActivity>(std::forward<Args>(args)...)); |
50 | } |
51 | |
52 | static GenericTraceActivity& toRef( |
53 | std::unique_ptr<GenericTraceActivity>& ref) { |
54 | return *ref; |
55 | } |
56 | |
57 | static const GenericTraceActivity& toRef( |
58 | const std::unique_ptr<GenericTraceActivity>& ref) { |
59 | return *ref; |
60 | } |
61 | |
62 | TraceSpan span{0, 0, "none"}; |
63 | int gpuOpCount; |
64 | std::deque<std::unique_ptr<GenericTraceActivity>> activities; |
65 | }; |
66 | |
67 | using ChildActivityProfilerFactory = |
68 | std::function<std::unique_ptr<IActivityProfiler>()>; |
69 | |
70 | class LibkinetoApi { |
71 | public: |
72 | |
73 | explicit LibkinetoApi(ConfigLoader& configLoader) |
74 | : configLoader_(configLoader) { |
75 | } |
76 | |
77 | // Called by client that supports tracing API. |
78 | // libkineto can still function without this. |
79 | void registerClient(ClientInterface* client); |
80 | |
81 | // Called by libkineto on init |
82 | void registerProfiler(std::unique_ptr<ActivityProfilerInterface> profiler) { |
83 | activityProfiler_ = std::move(profiler); |
84 | initClientIfRegistered(); |
85 | } |
86 | |
87 | ActivityProfilerInterface& activityProfiler() { |
88 | return *activityProfiler_; |
89 | } |
90 | |
91 | ClientInterface* client() { |
92 | return client_; |
93 | } |
94 | |
95 | void initProfilerIfRegistered() { |
96 | static std::once_flag once; |
97 | if (activityProfiler_) { |
98 | std::call_once(once, [this] { |
99 | if (!activityProfiler_->isInitialized()) { |
100 | activityProfiler_->init(); |
101 | initChildActivityProfilers(); |
102 | } |
103 | }); |
104 | } |
105 | } |
106 | |
107 | bool isProfilerInitialized() const { |
108 | return activityProfiler_ && activityProfiler_->isInitialized(); |
109 | } |
110 | |
111 | bool isProfilerRegistered() const { |
112 | return activityProfiler_ != nullptr; |
113 | } |
114 | |
115 | void suppressLogMessages() { |
116 | suppressLibkinetoLogMessages(); |
117 | } |
118 | |
119 | // Provides access to profier configuration manaegement |
120 | ConfigLoader& configLoader() { |
121 | return configLoader_; |
122 | } |
123 | |
124 | void registerProfilerFactory( |
125 | ChildActivityProfilerFactory factory) { |
126 | if (isProfilerInitialized()) { |
127 | activityProfiler_->addChildActivityProfiler(factory()); |
128 | } else { |
129 | childProfilerFactories_.push_back(factory); |
130 | } |
131 | } |
132 | |
133 | private: |
134 | |
135 | void initChildActivityProfilers() { |
136 | if (!isProfilerInitialized()) { |
137 | return; |
138 | } |
139 | for (const auto& factory : childProfilerFactories_) { |
140 | activityProfiler_->addChildActivityProfiler(factory()); |
141 | } |
142 | childProfilerFactories_.clear(); |
143 | } |
144 | |
145 | // Client is initialized once both it and libkineto has registered |
146 | void initClientIfRegistered(); |
147 | |
148 | ConfigLoader& configLoader_; |
149 | std::unique_ptr<ActivityProfilerInterface> activityProfiler_{}; |
150 | ClientInterface* client_{}; |
151 | int32_t clientRegisterThread_{0}; |
152 | |
153 | bool isLoaded_{false}; |
154 | std::vector<ChildActivityProfilerFactory> childProfilerFactories_; |
155 | }; |
156 | |
157 | // Singleton |
158 | LibkinetoApi& api(); |
159 | |
160 | } // namespace libkineto |
161 |