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
34extern "C" {
35 void suppressLibkinetoLogMessages();
36 int InitializeInjection(void);
37 void libkineto_init(bool cpuOnly, bool logOnError);
38}
39
40namespace libkineto {
41
42class Config;
43class ConfigLoader;
44
45struct 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
67using ChildActivityProfilerFactory =
68 std::function<std::unique_ptr<IActivityProfiler>()>;
69
70class 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
158LibkinetoApi& api();
159
160} // namespace libkineto
161