8000 Add a custom profiler configuration option by fwenguang · Pull Request #151656 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

Add a custom profiler configuration option #151656

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 2 commits into from
+31 −3
Closed

Add a custom profiler configuration option #151656

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9068ef8
Add a custom profiler configuration option
fwenguang Apr 21, 2025
4d47924
fix ci
fwenguang May 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 16 additions & 1 deletion torch/csrc/profiler/kineto_shim.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -206,6 +206,7 @@ class ExperimentalConfigWrapper {
configss << "\nCUPTI_PROFILER_ENABLE_PER_KERNEL="
<< (config_.profiler_measure_per_kernel ? "true" : "false")
<< "\n";
configss << "CUSTOM_CONFIG=" << config_.custom_profiler_config << "\n";
LOG(INFO) << "Generated config = " << configss.str();

libkineto::api().activityProfiler().prepareTrace(
Expand Down Expand Up @@ -238,6 +239,18 @@ static const std::string setTraceID(const std::string& trace_id) {
configss << "REQUEST_GROUP_TRACE_ID=" << trace_id << "\n";
return configss.str();
}

static const std::string appendCustomConfig(
const std::string& config,
const std::string& custom_profiler_config) {
if (custom_profiler_config.empty()) {
return config;
}
std::stringstream configss;
configss << config;
configss << "CUSTOM_CONFIG=" << custom_profiler_config << "\n";
return configss.str();
}
#endif

void prepareTrace(
Expand Down Expand Up @@ -294,7 +307,9 @@ void prepareTrace(
return;
}

const std::string configStr = setTraceID(trace_id);
const std::string traceIdStr = setTraceID(trace_id);
const std::string configStr =
appendCustomConfig(traceIdStr, config.custom_profiler_config);

libkineto::api().activityProfiler().prepareTrace(k_activities, configStr);
#endif // USE_KINETO
Expand Down
2 changes: 2 additions & 0 deletions torch/csrc/profiler/orchestration/observer.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ ExperimentalConfig::ExperimentalConfig(
bool disable_external_correlation,
bool profile_all_threads,
bool capture_overload_names,
std::string custom_profiler_config,
bool adjust_timestamps)
: profiler_metrics{std::move(profiler_metrics)},
profiler_measure_per_kernel{profiler_measure_per_kernel},
Expand All @@ -31,6 +32,7 @@ ExperimentalConfig::ExperimentalConfig(
disable_external_correlation{disable_external_correlation},
profile_all_threads{profile_all_threads},
capture_overload_names{capture_overload_names},
custom_profiler_config(std::move(custom_profiler_config)),
adjust_timestamps{adjust_timestamps} {}

/*explicit*/ ExperimentalConfig::operator bool() const {
Expand Down
7 changes: 7 additions & 0 deletions torch/csrc/profiler/orchestration/observer.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ struct TORCH_API ExperimentalConfig {
bool disable_external_correlation = false,
bool profile_all_threads = false,
bool capture_overload_names = false,
std::string custom_profiler_config = "",
bool adjust_timestamps = false);
explicit operator bool() const;

Expand Down Expand Up @@ -101,6 +102,12 @@ struct TORCH_API ExperimentalConfig {
* function schema and stored in the profile */
bool capture_overload_names;

/*
* A custom_profiler_config option is introduced to allow custom backends
* to apply custom configurations as needed.
*/
std::string custom_profiler_config;

/*
* Controls whether or not timestamp adjustment occurs after profiling.
* The purpose of this is to adjust Vulkan event timelines to align with those
Expand Down
8 changes: 6 additions & 2 deletions torch/csrc/profiler/python/init.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -340,7 +340,8 @@ void initPythonBindings(PyObject* module) {
bool /* adjust_profiler_step */,
bool /* disable_external_correlation*/,
bool /* profile_all_threads */,
bool /* capture_overload_names */
bool /* capture_overload_names */,
std::string /* custom_profiler_config*/
>(),
"An experimental config for Kineto features. Please note that"
"backward compatibility is not guaranteed.\n"
Expand All @@ -359,6 +360,7 @@ void initPythonBindings(PyObject* module) {
" disable_external_correlation (bool) : whether to disable external correlation\n",
" profile_all_threads (bool) : whether to profile all threads\n",
" capture_overload_names (bool) : whether to include ATen overload names in the profile\n",
" custom_profiler_config (string) : Used to pass some configurations to the custom profiler backend.\n",
py::arg("profiler_metrics") = std::vector<std::string>(),
py::arg("profiler_measure_per_kernel") = false,
py::arg("verbose") = false,
Expand All @@ -367,7 +369,8 @@ void initPythonBindings(PyObject* module) {
py::arg("adjust_profiler_step") = false,
py::arg("disable_external_correlation") = false,
py::arg("profile_all_threads") = false,
py::arg("capture_overload_names") = false)
py::arg("capture_overload_names") = false,
py::arg("custom_profiler_config") = "")
.def(py::pickle(
[](const ExperimentalConfig& p) { // __getstate__
py::list py_metrics;
Expand All @@ -390,6 +393,7 @@ void initPythonBindings(PyObject* module) {
p.disable_external_correlation,
p.profile_all_threads,
p.capture_overload_names,
p.custom_profiler_config,
p.performance_events);
},
[](const py::tuple& t) { // __setstate__
Expand Down
Loading
0