mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
ac7acfb894
Summary: This PR adds Auto-Trace implementation for Trace ID. By default, the python side will generate a uuid in the same format as the one set in the backend by kineto. Upon running an auto-trace, the python generated trace id will overwrite the one set in kineto using the Config variable. Since we don't expect users to generate on-demand traces after an auto-trace we can simply keep overwriting the backend trace id whenever autotrace is ran. If we one day want to eventually do something like this, we simply have to add a call in kineto on the backend to generate a new ID upon start of profiling. We also implement a custom callback in the frontend such that users can generate their own trace ids if they wish to. This works similarly as the default, only difference being that they have to manually set this callback after a profiler is generated. We use a specific call to set this rather then putting it in the frontend initializer in case users want to change the trace_id for different repeats. Test Plan: Tested both default and custom callbacks using the verbose prints added. Trace ids on the frontend and the prints on the backend for the manifold upload matched. Differential Revision: D65178308 Pull Request resolved: https://github.com/pytorch/pytorch/pull/139310 Approved by: https://github.com/shengfukevin
149 lines
3.8 KiB
C++
149 lines
3.8 KiB
C++
#pragma once
|
|
|
|
#include <memory>
|
|
#include <string>
|
|
|
|
// Skip Kineto dependency on mobile unless explicitly asked for.
|
|
// When is it explicitly asked for?
|
|
// KinetoEdgeCPUProfiler uses KinetoProfiler for cpu
|
|
// event profiling. This has a dependency on cpu only libkineto
|
|
#if defined(USE_KINETO) && defined(C10_MOBILE) && \
|
|
!defined(EDGE_PROFILER_USE_KINETO)
|
|
#undef USE_KINETO
|
|
#endif
|
|
|
|
#include <ActivityType.h>
|
|
|
|
#include <torch/csrc/Export.h>
|
|
#include <torch/csrc/profiler/api.h>
|
|
|
|
#ifdef USE_KINETO
|
|
// Forward declarations so we don't have to include `libkineto.h` in a header.
|
|
namespace libkineto {
|
|
class GenericTraceActivity;
|
|
struct CpuTraceBuffer;
|
|
class ActivityTraceInterface;
|
|
} // namespace libkineto
|
|
#endif
|
|
|
|
namespace torch {
|
|
namespace profiler {
|
|
|
|
#ifdef USE_KINETO
|
|
constexpr bool kKinetoAvailable{true};
|
|
#else
|
|
constexpr bool kKinetoAvailable{false};
|
|
#endif
|
|
|
|
namespace impl::kineto {
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// -- Interface (Does not require Kineto) -------------------------------------
|
|
// ----------------------------------------------------------------------------
|
|
struct DeviceAndResource {
|
|
int32_t device;
|
|
int32_t resource;
|
|
};
|
|
const DeviceAndResource kineto_ids();
|
|
|
|
#ifdef USE_KINETO
|
|
using trace_t = libkineto::CpuTraceBuffer;
|
|
using interface_trace_t = libkineto::ActivityTraceInterface;
|
|
using activity_t = libkineto::GenericTraceActivity;
|
|
#else
|
|
struct DummyTraceBuffer {};
|
|
struct DummyTraceInterface {};
|
|
|
|
using trace_t = DummyTraceBuffer;
|
|
using interface_trace_t = DummyTraceBuffer;
|
|
struct activity_t;
|
|
#endif // USE_KINETO
|
|
|
|
void addMetadata(
|
|
activity_t* activity,
|
|
const std::string& key,
|
|
const std::string& value);
|
|
|
|
// Wraps: libkineto::CpuTraceBuffer
|
|
struct TraceWrapper {
|
|
TraceWrapper(const int64_t start_time, const std::string& name);
|
|
|
|
// The caller is expected to hold a mutex when calling `addCPUActivity`.
|
|
activity_t* addCPUActivity(
|
|
const std::string& name,
|
|
const libkineto::ActivityType type,
|
|
const DeviceAndResource device_and_resource,
|
|
const uint64_t correlation_id,
|
|
const int64_t start_time,
|
|
const int64_t end_time);
|
|
|
|
void transferCpuTrace(int64_t end_time);
|
|
|
|
explicit operator bool() const;
|
|
|
|
std::unique_ptr<trace_t>& get() {
|
|
return cpu_trace_;
|
|
}
|
|
|
|
private:
|
|
std::unique_ptr<trace_t> cpu_trace_;
|
|
};
|
|
|
|
// Wraps libkineto::ActivityTraceInterface
|
|
struct ActivityTraceWrapper {
|
|
explicit ActivityTraceWrapper(std::unique_ptr<interface_trace_t>&& trace);
|
|
ActivityTraceWrapper() = default;
|
|
explicit operator bool() const;
|
|
void save(const std::string& path);
|
|
|
|
const std::unique_ptr<interface_trace_t>& get() {
|
|
return trace_;
|
|
}
|
|
|
|
private:
|
|
std::unique_ptr<interface_trace_t> trace_;
|
|
#ifdef USE_KINETO
|
|
bool saved_ = false; // Kineto's save is destructive
|
|
#endif
|
|
};
|
|
|
|
using ActivitySet = std::set<torch::autograd::profiler::ActivityType>;
|
|
void prepareTrace(
|
|
const bool cpuOnly,
|
|
const ActivitySet& activities,
|
|
const torch::profiler::impl::ExperimentalConfig& config,
|
|
const std::string& trace_id = "");
|
|
|
|
void toggleCollectionDynamic(const bool enable);
|
|
void startTrace();
|
|
ActivityTraceWrapper stopTrace();
|
|
void pushCorrelationId(uint64_t correlation_id);
|
|
void pushUserCorrelationId(uint64_t correlation_id);
|
|
void popCorrelationId();
|
|
void popUserCorrelationId();
|
|
void recordThreadInfo();
|
|
bool collectivesProfilerExists();
|
|
|
|
void logInvariantViolation(
|
|
const std::string& assertion,
|
|
const std::string& error,
|
|
const std::string& profile_id,
|
|
const std::string& group_profile_id);
|
|
|
|
} // namespace impl::kineto
|
|
|
|
} // namespace profiler
|
|
|
|
namespace autograd::profiler {
|
|
c10::DeviceType deviceTypeFromActivity(libkineto::ActivityType activity_type);
|
|
|
|
TORCH_API void addMetadataJson(
|
|
const std::string& key,
|
|
const std::string& value);
|
|
|
|
TORCH_API void profilerStep();
|
|
|
|
} // namespace autograd::profiler
|
|
|
|
} // namespace torch
|