From ad578fb012fe5b71530349b120b2ae1b3f74ad0e Mon Sep 17 00:00:00 2001
From: Juncheng <liujuncheng1022@gmail.com>
Date: Tue, 20 Jul 2021 22:22:41 +0800
Subject: [PATCH] Add ProfilerStart/ProfilerStop API (#5542)
Co-authored-by: oneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>
---
oneflow/api/python/profiler.cpp | 4 ++++
oneflow/core/profiler/profiler.cpp | 14 ++++++++++++++
oneflow/core/profiler/profiler.h | 4 ++++
oneflow/python/framework/profiler.py | 10 ++++++++++
4 files changed, 32 insertions(+)
diff --git a/oneflow/api/python/profiler.cpp b/oneflow/api/python/profiler.cpp
index d8e0e83cf..5848dc8a9 100644
--- a/oneflow/api/python/profiler.cpp
+++ b/oneflow/api/python/profiler.cpp
@@ -26,6 +26,10 @@ ONEFLOW_API_PYBIND11_MODULE("profiler", m) {
m.def("RangePush", [](const std::string& str) { OF_PROFILER_RANGE_PUSH(str); });
m.def("RangePop", []() { OF_PROFILER_RANGE_POP(); });
+
+ m.def("ProfilerStart", []() { profiler::ProfilerStart(); });
+
+ m.def("ProfilerStop", []() { profiler::ProfilerStop(); });
}
} // namespace oneflow
diff --git a/oneflow/core/profiler/profiler.cpp b/oneflow/core/profiler/profiler.cpp
index 5ea3d84ac..b3502d340 100644
--- a/oneflow/core/profiler/profiler.cpp
+++ b/oneflow/core/profiler/profiler.cpp
@@ -19,6 +19,8 @@ limitations under the License.
#include <nvtx3/nvToolsExt.h>
#include <sys/syscall.h>
#include <iostream>
+#include <cuda_profiler_api.h>
+#include "oneflow/core/device/cuda_util.h"
#endif // OF_ENABLE_PROFILER
namespace oneflow {
@@ -107,6 +109,18 @@ void LogHostMemoryUsage(const std::string& name) {
#endif // OF_ENABLE_PROFILER
}
+void ProfilerStart() {
+#ifdef OF_ENABLE_PROFILER
+ OF_CUDA_CHECK(cudaProfilerStart());
+#endif // OF_ENABLE_PROFILER
+}
+
+void ProfilerStop() {
+#ifdef OF_ENABLE_PROFILER
+ OF_CUDA_CHECK(cudaProfilerStop());
+#endif // OF_ENABLE_PROFILER
+}
+
} // namespace profiler
} // namespace oneflow
diff --git a/oneflow/core/profiler/profiler.h b/oneflow/core/profiler/profiler.h
index f6c26131a..7a90fa8c5 100644
--- a/oneflow/core/profiler/profiler.h
+++ b/oneflow/core/profiler/profiler.h
@@ -32,6 +32,10 @@ void RangePop();
void LogHostMemoryUsage(const std::string& name);
+void ProfilerStart();
+
+void ProfilerStop();
+
class RangeGuardCtx;
class RangeGuard final {
diff --git a/oneflow/python/framework/profiler.py b/oneflow/python/framework/profiler.py
index 0fe5c2a14..6ba67601d 100644
--- a/oneflow/python/framework/profiler.py
+++ b/oneflow/python/framework/profiler.py
@@ -27,3 +27,13 @@ def RangePush(range_name):
@oneflow_export("profiler.range_pop")
def RangePop():
oneflow._oneflow_internal.profiler.RangePop()
+
+
+@oneflow_export("profiler.profiler_start")
+def ProfilerStart():
+ oneflow._oneflow_internal.profiler.ProfilerStart()
+
+
+@oneflow_export("profiler.profiler_stop")
+def ProfilerStop():
+ oneflow._oneflow_internal.profiler.ProfilerStop()
--
GitLab