Add nvtx context

oneflow/core/auto_parallel/sbp_constructor.h

@@ -17,7 +17,7 @@ limitations under the License.
 #define SBP_CONSTRUCTOR_
 
 #ifndef ENABLE_AUTO_PARALLEL
-#define ENABLE_AUTO_PARALLEL
+//#define ENABLE_AUTO_PARALLEL
 #endif
 
 #include "sbp_graph.h"

oneflow/core/kernel/kernel.cpp

@@ -16,6 +16,7 @@ limitations under the License.
 #include "oneflow/core/kernel/kernel.h"
 #include "oneflow/core/common/cached_caller.h"
 #include "oneflow/core/kernel/runtime_blob_shape_infer_helper.h"
+#include "oneflow/core/profiler/profiler.h"
 
 namespace oneflow {
 
@@ -56,7 +57,9 @@ void Kernel::InitModelAndConstBuf(const KernelCtx& ctx,
 
 void Kernel::Launch(const KernelCtx& ctx,
                     std::function<Blob*(const std::string&)> BnInOp2Blob) const {
+  ::oneflow::profiler::TraceKernelForwardDataContentStart(ctx, this);
   Forward(ctx, BnInOp2Blob);
+  ::oneflow::profiler::TraceKernelForwardDataContentEnd(ctx, this);
 }
 
 const LogicalBlobId& Kernel::BnInOp2Lbi(const std::string& bn_in_op) const {

oneflow/core/profiler/profiler.cpp

+#include "oneflow/core/profiler/profiler.h"
+
+#include <nvtx3/nvToolsExt.h>
+
+namespace oneflow {
+
+namespace profiler {
+
+void RangePush(const std::string& name) { nvtxRangePushA(name.c_str()); }
+
+void RangePop() { nvtxRangePop(); }
+
+void TraceKernelForwardDataContentStart(const KernelCtx& ctx, const Kernel* kernel) {
+  OF_PROFILER_RANGE_PUSH(kernel->op_conf().name());
+}
+
+void TraceKernelForwardDataContentEnd(const KernelCtx& ctx, const Kernel* kernel) {
+  OF_PROFILER_RANGE_POP();
+}
+
+}
+}

oneflow/core/profiler/profiler.h

+#ifndef ONEFLOW_CORE_PROFILER_PROFILER_H_
+#define ONEFLOW_CORE_PROFILER_PROFILER_H_
+
+#include <string>
+#include "oneflow/core/kernel/kernel.h"
+#include "oneflow/core/kernel/kernel_context.h"
+
+namespace oneflow {
+
+class Kernel;
+struct KernelCtx;
+
+namespace profiler {
+
+void TraceKernelForwardDataContentStart(const KernelCtx& ctx, const Kernel* kernel);
+
+void TraceKernelForwardDataContentEnd(const KernelCtx& ctx, const Kernel* kernel);
+
+void RangePush(const std::string& name);
+
+void RangePop();
+
+}  // namespace profiler
+
+}  // namespace oneflow
+
+#define OF_PROFILER_RANGE_PUSH(name) ::oneflow::profiler::RangePush(name)
+#define OF_PROFILER_RANGE_POP() ::oneflow::profiler::RangePop()
+
+#endif
\ No newline at end of file