Add ProfilerStart/ProfilerStop API (#5542)

Co-authored-by: oneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>

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

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

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 {

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()