Refactor trace activities (#59360)

Summary:
Pull Request resolved: pytorch/pytorch#59360

Pull Request resolved: #206

Replace ClientTraceActivity with GenericActivity.
In addition:
* Add a couple of new activity types for user annotations
* Simplify code for GPU-side user annotations
* Add accessor to containing trace span object in activities. Later we can replace this with a trace context / trace session object.
* Simplified MemoryTraceLogger
* Added early exit for cupti push/pop correlation ID

Reviewed By: ilia-cher

Differential Revision: D28231675

fbshipit-source-id: 7129f2493016efb4d3697094f24475e2c39e6e65

Author: gdankel

libkineto/include/ActivityType.h

@@ -7,16 +7,29 @@
 
 #pragma once
 
+#include <array>
+#include <string>
+
 namespace libkineto {
 
 enum class ActivityType {
     CPU_OP = 0, // cpu side ops
+    USER_ANNOTATION,
+    GPU_USER_ANNOTATION,
     GPU_MEMCPY,
     GPU_MEMSET,
     CONCURRENT_KERNEL, // on-device kernels
     EXTERNAL_CORRELATION,
     CUDA_RUNTIME, // host side cuda runtime events
     CPU_INSTANT_EVENT, // host side point-like events
+    ENUM_COUNT
 };
 
+const char* toString(ActivityType t);
+ActivityType toActivityType(const std::string& str);
+
+// Return an array of all activity types except COUNT
+constexpr int activityTypeCount = (int)ActivityType::ENUM_COUNT;
+const std::array<ActivityType, activityTypeCount> activityTypes();
+
 } // namespace libkineto

libkineto/include/GenericTraceActivity.h

@@ -14,19 +14,28 @@
 
 #include "ThreadUtil.h"
 #include "TraceActivity.h"
+#include "TraceSpan.h"
 
 namespace libkineto {
 
 // @lint-ignore-every CLANGTIDY cppcoreguidelines-non-private-member-variables-in-classes
 // @lint-ignore-every CLANGTIDY cppcoreguidelines-pro-type-member-init
-struct GenericTraceActivity : TraceActivity {
+class GenericTraceActivity : public TraceActivity {
+
+ public:
+  GenericTraceActivity() = delete;
+
+  GenericTraceActivity(
+      const TraceSpan& trace, ActivityType type, const std::string& name)
+      : activityType(type), activityName(name), traceSpan_(&trace) {
+  }
 
   int64_t deviceId() const override {
     return device;
   }
 
   int64_t resourceId() const override {
-    return sysThreadId;
+    return resource;
   }
 
   int64_t timestamp() const override {
@@ -38,7 +47,7 @@ struct GenericTraceActivity : TraceActivity {
   }
 
   int64_t correlationId() const override {
-    return correlation;
+    return id;
   }
 
   ActivityType type() const override {
@@ -53,6 +62,10 @@ struct GenericTraceActivity : TraceActivity {
     return nullptr;
   }
 
+  const TraceSpan* traceSpan() const override {
+    return traceSpan_;
+  }
+
   void log(ActivityLogger& logger) const override;
 
   //Encode client side metadata as a key/value string.
@@ -68,13 +81,14 @@ struct GenericTraceActivity : TraceActivity {
 
   int64_t startTime{0};
   int64_t endTime{0};
-  int64_t correlation{0};
-  int device{-1};
-  int32_t sysThreadId{0};
-  std::string activityName;
+  int32_t id{0};
+  int32_t device{0};
+  int32_t resource{0};
   ActivityType activityType;
+  std::string activityName;
 
  private:
+  const TraceSpan* traceSpan_;
   std::vector<std::string> metadata_;
 };
 

libkineto/include/TraceActivity.h

@@ -14,6 +14,7 @@
 namespace libkineto {
 
 class ActivityLogger;
+struct TraceSpan;
 
 // Generic activity interface is borrowed from tensorboard protobuf format.
 struct TraceActivity {
@@ -33,6 +34,8 @@ struct TraceActivity {
   virtual const std::string name() const = 0;
   // Optional linked activity
   virtual const TraceActivity* linkedActivity() const = 0;
+  // Optional containing trace object
+  virtual const TraceSpan* traceSpan() const = 0;
   // Log activity
   virtual void log(ActivityLogger& logger) const = 0;
 

libkineto/include/TraceSpan.h

@@ -7,21 +7,37 @@
 
 #pragma once
 
+#include <atomic>
 #include <string>
 #include <thread>
 
 namespace libkineto {
 
 struct TraceSpan {
+  TraceSpan() = delete;
+  TraceSpan(
+      int64_t startTime, int64_t endTime, std::string name)
+      : startTime(startTime), endTime(endTime), name(std::move(name)) {
+  }
+  TraceSpan(
+      int opCount, int it, std::string name, std::string prefix)
+      : opCount(opCount),
+        iteration(it),
+        name(std::move(name)),
+        prefix(std::move(prefix)) {
+  }
+
   // FIXME: change to duration?
   int64_t startTime{0};
   int64_t endTime{0};
   int opCount{0};
   int iteration{-1};
   // Name is used to identify timeline
   std::string name;
-  // Prefix used to distinguish sub-nets on the same timeline
+  // Prefix used to distinguish trace spans on the same timeline
   std::string prefix;
+  // Tracked by profiler for iteration trigger
+  bool tracked{false};
 };
 
 } // namespace libkineto

libkineto/include/libkineto.h

@@ -38,7 +38,7 @@ namespace libkineto {
 class Config;
 
 struct CpuTraceBuffer {
-  TraceSpan span;
+  TraceSpan span{0, 0, "none"};
   int gpuOpCount;
   std::vector<GenericTraceActivity> activities;
 };