163 lines
3.8 KiB
C++
163 lines
3.8 KiB
C++
/*
|
|
* Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
* All rights reserved.
|
|
*
|
|
* This source code is licensed under the BSD-style license found in the
|
|
* LICENSE file in the root directory of this source tree.
|
|
*/
|
|
|
|
// Mediator for initialization and profiler control
|
|
|
|
#pragma once
|
|
|
|
#include <atomic>
|
|
#include <chrono>
|
|
#include <functional>
|
|
#include <memory>
|
|
#include <mutex>
|
|
#include <string>
|
|
#include <set>
|
|
#include <thread>
|
|
#include <vector>
|
|
#include <deque>
|
|
|
|
#include "ActivityProfilerInterface.h"
|
|
#include "ActivityType.h"
|
|
#include "ClientInterface.h"
|
|
#include "GenericTraceActivity.h"
|
|
#include "TraceSpan.h"
|
|
#include "IActivityProfiler.h"
|
|
#include "ActivityTraceInterface.h"
|
|
#include "ILoggerObserver.h"
|
|
#include "LoggingAPI.h"
|
|
|
|
#include "ThreadUtil.h"
|
|
|
|
extern "C" {
|
|
void suppressLibkinetoLogMessages();
|
|
int InitializeInjection(void);
|
|
void libkineto_init(bool cpuOnly, bool logOnError);
|
|
bool hasTestEnvVar();
|
|
}
|
|
|
|
namespace libkineto {
|
|
|
|
class Config;
|
|
class ConfigLoader;
|
|
|
|
struct CpuTraceBuffer {
|
|
template <class... Args>
|
|
void emplace_activity(Args&&... args) {
|
|
activities.emplace_back(
|
|
std::make_unique<GenericTraceActivity>(std::forward<Args>(args)...));
|
|
}
|
|
|
|
static GenericTraceActivity& toRef(
|
|
std::unique_ptr<GenericTraceActivity>& ref) {
|
|
return *ref;
|
|
}
|
|
|
|
static const GenericTraceActivity& toRef(
|
|
const std::unique_ptr<GenericTraceActivity>& ref) {
|
|
return *ref;
|
|
}
|
|
|
|
TraceSpan span{0, 0, "none"};
|
|
int gpuOpCount;
|
|
std::deque<std::unique_ptr<GenericTraceActivity>> activities;
|
|
};
|
|
|
|
using ChildActivityProfilerFactory =
|
|
std::function<std::unique_ptr<IActivityProfiler>()>;
|
|
|
|
class LibkinetoApi {
|
|
public:
|
|
|
|
explicit LibkinetoApi(ConfigLoader& configLoader)
|
|
: configLoader_(configLoader) {
|
|
}
|
|
|
|
// Called by client that supports tracing API.
|
|
// libkineto can still function without this.
|
|
void registerClient(ClientInterface* client);
|
|
|
|
// Called by libkineto on init
|
|
void registerProfiler(std::unique_ptr<ActivityProfilerInterface> profiler) {
|
|
activityProfiler_ = std::move(profiler);
|
|
initClientIfRegistered();
|
|
}
|
|
|
|
ActivityProfilerInterface& activityProfiler() {
|
|
return *activityProfiler_;
|
|
}
|
|
|
|
ClientInterface* client() {
|
|
return client_;
|
|
}
|
|
|
|
void initProfilerIfRegistered() {
|
|
static std::once_flag once;
|
|
if (activityProfiler_) {
|
|
std::call_once(once, [this] {
|
|
if (!activityProfiler_->isInitialized()) {
|
|
activityProfiler_->init();
|
|
initChildActivityProfilers();
|
|
}
|
|
});
|
|
}
|
|
}
|
|
|
|
bool isProfilerInitialized() const {
|
|
return activityProfiler_ && activityProfiler_->isInitialized();
|
|
}
|
|
|
|
bool isProfilerRegistered() const {
|
|
return activityProfiler_ != nullptr;
|
|
}
|
|
|
|
void suppressLogMessages() {
|
|
suppressLibkinetoLogMessages();
|
|
}
|
|
|
|
// Provides access to profier configuration manaegement
|
|
ConfigLoader& configLoader() {
|
|
return configLoader_;
|
|
}
|
|
|
|
void registerProfilerFactory(
|
|
ChildActivityProfilerFactory factory) {
|
|
if (isProfilerInitialized()) {
|
|
activityProfiler_->addChildActivityProfiler(factory());
|
|
} else {
|
|
childProfilerFactories_.push_back(factory);
|
|
}
|
|
}
|
|
|
|
private:
|
|
|
|
void initChildActivityProfilers() {
|
|
if (!isProfilerInitialized()) {
|
|
return;
|
|
}
|
|
for (const auto& factory : childProfilerFactories_) {
|
|
activityProfiler_->addChildActivityProfiler(factory());
|
|
}
|
|
childProfilerFactories_.clear();
|
|
}
|
|
|
|
// Client is initialized once both it and libkineto has registered
|
|
void initClientIfRegistered();
|
|
|
|
ConfigLoader& configLoader_;
|
|
std::unique_ptr<ActivityProfilerInterface> activityProfiler_{};
|
|
ClientInterface* client_{};
|
|
int32_t clientRegisterThread_{0};
|
|
|
|
std::vector<ChildActivityProfilerFactory> childProfilerFactories_;
|
|
};
|
|
|
|
// Singleton
|
|
LibkinetoApi& api();
|
|
|
|
} // namespace libkineto
|