Fix event profiler segfault

libkineto/include/ActivityProfilerInterface.h

@@ -14,6 +14,7 @@
 
 #include "ActivityType.h"
 #include "ActivityTraceInterface.h"
+#include "IActivityProfiler.h"
 
 namespace libkineto {
 
@@ -85,6 +86,11 @@ class ActivityProfilerInterface {
   // Record trace metadata, currently supporting only string key and values,
   // values with the same key are overwritten
   virtual void addMetadata(const std::string& key, const std::string& value) = 0;
+
+  // Add a child activity profiler, this enables frameworks in the application
+  // to enable custom framework events.
+  virtual void addChildActivityProfiler(
+      std::unique_ptr<IActivityProfiler> profiler) {}
 };
 
 } // namespace libkineto

libkineto/include/libkineto.h

@@ -25,6 +25,7 @@
 #include "ClientInterface.h"
 #include "GenericTraceActivity.h"
 #include "TraceSpan.h"
+#include "IActivityProfiler.h"
 
 #include "ThreadUtil.h"
 
@@ -44,6 +45,9 @@ struct CpuTraceBuffer {
   std::vector<GenericTraceActivity> activities;
 };
 
+using ChildActivityProfilerFactory =
+  std::function<std::unique_ptr<IActivityProfiler>()>;
+
 class LibkinetoApi {
  public:
 
@@ -72,6 +76,7 @@ class LibkinetoApi {
   void initProfilerIfRegistered() {
     if (activityProfiler_ && !activityProfiler_->isInitialized()) {
       activityProfiler_->init();
+      initChildActivityProfilers();
     }
   }
 
@@ -87,7 +92,7 @@ class LibkinetoApi {
     netSizeThreshold_ = gpu_ops;
   }
 
-  // Inclide traces with at least this many ops
+  // Include traces with at least this many ops
   // FIXME: Rename and move elsewhere
   int netSizeThreshold() {
     return netSizeThreshold_;
@@ -102,8 +107,27 @@ class LibkinetoApi {
     return configLoader_;
   }
 
+  void registerProfilerFactory(
+      ChildActivityProfilerFactory factory) {
+    if (isProfilerInitialized()) {
+      activityProfiler_->addChildActivityProfiler(factory());
+    } else {
+      childProfilerFactories_.push_back(factory);
+    }
+  }
+
  private:
 
+  void initChildActivityProfilers() {
+    if (!isProfilerInitialized()) {
+      return;
+    }
+    for (const auto& factory : childProfilerFactories_) {
+      activityProfiler_->addChildActivityProfiler(factory());
+    }
+    childProfilerFactories_.clear();
+  }
+
   // Client is initialized once both it and libkineto has registered
   void initClientIfRegistered();
 
@@ -114,6 +138,7 @@ class LibkinetoApi {
 
   bool isLoaded_{false};
   std::atomic_int netSizeThreshold_{};
+  std::vector<ChildActivityProfilerFactory> childProfilerFactories_;
 };
 
 // Singleton

libkineto/src/ActivityProfiler.cpp

@@ -455,7 +455,7 @@ void ActivityProfiler::configureChildProfilers() {
     auto session = profiler->configure(
         start_time_ms,
         config_->activitiesOnDemandDuration().count(),
-        std::set<ActivityType>{ActivityType::CPU_OP} // TODO make configurable
+        config_->selectedActivityTypes()
     );
     if (session) {
       sessions_.push_back(std::move(session));

libkineto/src/ActivityProfiler.h

@@ -118,10 +118,10 @@ class ActivityProfiler {
     metadata_[key] = value;
   }
 
-  void addActivityProfiler(
-      std::shared_ptr<IActivityProfiler> profiler) {
+  void addChildActivityProfiler(
+      std::unique_ptr<IActivityProfiler> profiler) {
     std::lock_guard<std::mutex> guard(mutex_);
-    profilers_.push_back(profiler);
+    profilers_.push_back(std::move(profiler));
   }
 
  protected:
@@ -368,7 +368,7 @@ class ActivityProfiler {
   std::unordered_map<std::string, std::string> metadata_;
 
   // child activity profilers
-  std::vector<std::shared_ptr<IActivityProfiler>> profilers_;
+  std::vector<std::unique_ptr<IActivityProfiler>> profilers_;
 
   // a vector of active profiler plugin sessions
   std::vector<std::unique_ptr<IActivityProfilerSession>> sessions_;

libkineto/src/ActivityProfilerController.h

@@ -66,6 +66,11 @@ class ActivityProfilerController : public ConfigLoader::ConfigHandler {
     profiler_->recordThreadInfo();
   }
 
+  void addChildActivityProfiler(
+      std::unique_ptr<IActivityProfiler> profiler) {
+    profiler_->addChildActivityProfiler(std::move(profiler));
+  }
+
   void addMetadata(const std::string& key, const std::string& value);
 
  private:

libkineto/src/ActivityProfilerProxy.cpp

@@ -98,4 +98,9 @@ void ActivityProfilerProxy::recordThreadInfo() {
   controller_->recordThreadInfo();
 }
 
+void ActivityProfilerProxy::addChildActivityProfiler(
+    std::unique_ptr<IActivityProfiler> profiler) {
+  controller_->addChildActivityProfiler(std::move(profiler));
+}
+
 } // namespace libkineto

libkineto/src/ActivityProfilerProxy.h

@@ -63,6 +63,9 @@ class ActivityProfilerProxy : public ActivityProfilerInterface {
 
   void addMetadata(const std::string& key, const std::string& value) override;
 
+  virtual void addChildActivityProfiler(
+      std::unique_ptr<IActivityProfiler> profiler) override;
+
  private:
   bool cpuOnly_{true};
   ConfigLoader& configLoader_;

libkineto/src/EventProfiler.cpp

@@ -590,7 +590,6 @@ void EventProfiler::dispatchSamples(
 }
 
 void EventProfiler::configure(Config& config, Config* onDemandConfig) {
-  LOG(INFO) << "configure";
   if (!sets_.empty()) {
     sets_[curEnabledSet_].setEnabled(false);
     clearSamples();

libkineto/src/EventProfilerController.cpp

@@ -351,7 +351,9 @@ void EventProfilerController::profilerLoop() {
       now = system_clock::now();
       next_sample_time = now + profiler_->samplePeriod();
       next_report_time = now + profiler_->reportPeriod();
-      next_on_demand_report_time = now + profiler_->onDemandReportPeriod();
+      if (profiler_->isOnDemandActive()) {
+        next_on_demand_report_time = now + profiler_->onDemandReportPeriod();
+      }
       next_multiplex_time = now + profiler_->multiplexPeriod();
       // Collect an initial sample and throw it away
       // The next sample is the first valid one
@@ -386,7 +388,7 @@ void EventProfilerController::profilerLoop() {
       profiler_->reportSamples();
       next_report_time += profiler_->reportPeriod();
     }
-    if (on_demand_report_count && now > next_on_demand_report_time) {
+    if (profiler_->isOnDemandActive() && now > next_on_demand_report_time) {
       VLOG(1) << "OnDemand Report #" << on_demand_report_count++;
       profiler_->reportOnDemandSamples();
       next_on_demand_report_time += profiler_->onDemandReportPeriod();

libkineto/test/ActivityProfilerTest.cpp

@@ -413,11 +413,12 @@ TEST_F(ActivityProfilerTest, SubActivityProfilers) {
   test_activities[2].activityName = "Operator bar";
 
   auto mock_activity_profiler =
-    std::make_shared<MockActivityProfiler>(test_activities);
+    std::make_unique<MockActivityProfiler>(test_activities);
 
   MockCuptiActivities activities;
   ActivityProfiler profiler(activities, /*cpu only*/ true);
-  profiler.addActivityProfiler(mock_activity_profiler);
+  profiler.addChildActivityProfiler(
+      std::move(mock_activity_profiler));
 
   profiler.configure(*cfg_, start_time);
   profiler.startTrace(start_time);