Refactor 4/n: Simplify activity logger step 2/3

libkineto/include/GenericTraceActivity.h

@@ -20,6 +20,7 @@ namespace libkineto {
 
 // Link type, used in GenericTraceActivity.flow.type
 constexpr unsigned int kLinkFwdBwd = 1;
+constexpr unsigned int kLinkAsyncCpuGpu = 2;
 
 // @lint-ignore-every CLANGTIDY cppcoreguidelines-non-private-member-variables-in-classes
 // @lint-ignore-every CLANGTIDY cppcoreguidelines-pro-type-member-init
@@ -57,6 +58,10 @@ class GenericTraceActivity : public ITraceActivity {
     return activityType;
   }
 
+  const ITraceActivity* linkedActivity() const override {
+    return nullptr;
+  }
+
   int flowType() const override {
     return flow.type;
   }
@@ -65,12 +70,12 @@ class GenericTraceActivity : public ITraceActivity {
     return flow.id;
   }
 
-  const std::string name() const override {
-    return activityName;
+  bool flowStart() const override {
+    return flow.start;
   }
 
-  const ITraceActivity* linkedActivity() const override {
-    return flow.linkedActivity;
+  const std::string name() const override {
+    return activityName;
   }
 
   const TraceSpan* traceSpan() const override {
@@ -103,13 +108,13 @@ class GenericTraceActivity : public ITraceActivity {
   ActivityType activityType;
   std::string activityName;
   struct Flow {
-    Flow(): linkedActivity(nullptr), id(0), type(0) {}
-    ITraceActivity* linkedActivity; // Only set in destination side.
+    Flow(): id(0), type(0), start(0) {}
     // Ids must be unique within each type
-    uint32_t id : 28;
+    uint32_t id : 27;
     // Type will be used to connect flows between profilers, as
     // well as look up flow information (name etc)
     uint32_t type : 4;
+    uint32_t start : 1;
   } flow;
 
  private:

libkineto/include/ITraceActivity.h

@@ -33,6 +33,7 @@ struct ITraceActivity {
   // Part of a flow, identified by flow id and type
   virtual int flowType() const = 0;
   virtual int flowId() const = 0;
+  virtual bool flowStart() const = 0;
   virtual ActivityType type() const = 0;
   virtual const std::string name() const = 0;
   // Optional linked activity

libkineto/include/TraceSpan.h

@@ -36,8 +36,6 @@ struct TraceSpan {
   std::string name;
   // 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/src/ActivityBuffers.h

@@ -19,6 +19,16 @@ namespace KINETO_NAMESPACE {
 struct ActivityBuffers {
   std::list<std::unique_ptr<libkineto::CpuTraceBuffer>> cpu;
   std::unique_ptr<CuptiActivityBufferMap> gpu;
+
+  // Add a wrapper object to the underlying struct stored in the buffer
+  template<class T>
+  const ITraceActivity& addActivityWrapper(const T& act) {
+    wrappers_.push_back(std::make_unique<T>(act));
+    return *wrappers_.back().get();
+  }
+
+ private:
+  std::vector<std::unique_ptr<const ITraceActivity>> wrappers_;
 };
 
 } // namespace KINETO_NAMESPACE

libkineto/src/CuptiActivity.h

@@ -30,7 +30,7 @@ struct TraceSpan;
 // Abstract base class, templated on Cupti activity type
 template<class T>
 struct CuptiActivity : public ITraceActivity {
-  explicit CuptiActivity(const T* activity, const ITraceActivity& linked)
+  explicit CuptiActivity(const T* activity, const ITraceActivity* linked)
       : activity_(*activity), linked_(linked) {}
   int64_t timestamp() const override {
     return nsToUs(unixEpochTimestamp(activity_.start));
@@ -39,27 +39,28 @@ struct CuptiActivity : public ITraceActivity {
     return nsToUs(activity_.end - activity_.start);
   }
   int64_t correlationId() const override {return activity_.correlationId;}
-  int flowType() const override {return 0;}
-  int flowId() const override {return 0;}
+  const ITraceActivity* linkedActivity() const override {return linked_;}
+  int flowType() const override {return kLinkAsyncCpuGpu;}
+  int flowId() const override {return correlationId();}
   const T& raw() const {return activity_;}
-  const ITraceActivity* linkedActivity() const override {return &linked_;}
   const TraceSpan* traceSpan() const override {return nullptr;}
 
  protected:
   const T& activity_;
-  const ITraceActivity& linked_;
+  const ITraceActivity* linked_{nullptr};
 };
 
 // CUpti_ActivityAPI - CUDA runtime activities
 struct RuntimeActivity : public CuptiActivity<CUpti_ActivityAPI> {
   explicit RuntimeActivity(
       const CUpti_ActivityAPI* activity,
-      const ITraceActivity& linked,
+      const ITraceActivity* linked,
       int32_t threadId)
       : CuptiActivity(activity, linked), threadId_(threadId) {}
   int64_t deviceId() const override {return processId();}
   int64_t resourceId() const override {return threadId_;}
   ActivityType type() const override {return ActivityType::CUDA_RUNTIME;}
+  bool flowStart() const override;
   const std::string name() const override {return runtimeCbidName(activity_.cbid);}
   void log(ActivityLogger& logger) const override;
   const std::string metadataJson() const override;
@@ -72,11 +73,12 @@ struct RuntimeActivity : public CuptiActivity<CUpti_ActivityAPI> {
 // Can also be instantiated directly.
 template<class T>
 struct GpuActivity : public CuptiActivity<T> {
-  explicit GpuActivity(const T* activity, const ITraceActivity& linked)
+  explicit GpuActivity(const T* activity, const ITraceActivity* linked)
       : CuptiActivity<T>(activity, linked) {}
   int64_t deviceId() const override {return raw().deviceId;}
   int64_t resourceId() const override {return raw().streamId;}
   ActivityType type() const override;
+  bool flowStart() const override {return false;}
   const std::string name() const override;
   void log(ActivityLogger& logger) const override;
   const std::string metadataJson() const override;

libkineto/src/CuptiActivity.tpp

@@ -67,16 +67,28 @@ inline ActivityType GpuActivity<CUpti_ActivityMemset>::type() const {
 }
 
 inline void RuntimeActivity::log(ActivityLogger& logger) const {
-  logger.handleRuntimeActivity(*this);
+  logger.handleGenericActivity(*this);
 }
 
 template<class T>
 inline void GpuActivity<T>::log(ActivityLogger& logger) const {
   logger.handleGpuActivity(*this);
 }
 
+inline bool RuntimeActivity::flowStart() const {
+  return activity_.cbid == CUPTI_RUNTIME_TRACE_CBID_cudaLaunchKernel_v7000 ||
+      (activity_.cbid >= CUPTI_RUNTIME_TRACE_CBID_cudaMemcpy_v3020 &&
+       activity_.cbid <= CUPTI_RUNTIME_TRACE_CBID_cudaMemset2DAsync_v3020) ||
+      activity_.cbid ==
+          CUPTI_RUNTIME_TRACE_CBID_cudaLaunchCooperativeKernel_v9000 ||
+      activity_.cbid ==
+          CUPTI_RUNTIME_TRACE_CBID_cudaLaunchCooperativeKernelMultiDevice_v9000;
+}
+
 inline const std::string RuntimeActivity::metadataJson() const {
-  return "";
+  return fmt::format(R"JSON(
+      "cbid": {}, "correlation": {})JSON",
+      activity_.cbid, activity_.correlationId);
 }
 
 template<class T>

libkineto/src/CuptiActivityBuffer.h

@@ -15,6 +15,8 @@
 #include <sys/types.h>
 #include <vector>
 
+#include "ITraceActivity.h"
+
 namespace KINETO_NAMESPACE {
 
 class CuptiActivityBuffer {
@@ -44,6 +46,8 @@ class CuptiActivityBuffer {
 
   std::vector<uint8_t> buf_;
   size_t size_;
+
+  std::vector<std::unique_ptr<const ITraceActivity>> wrappers_;
 };
 
 using CuptiActivityBufferMap =