Add support for multiple protocols in client API

libkineto/src/AbstractConfig.h

@@ -52,6 +52,7 @@ class AbstractConfig {
     return timestamp_;
   }
 
+  // Source config string that this was parsed from
   const std::string& source() const {
     return source_;
   }

libkineto/src/ActivityLoggerFactory.h

@@ -0,0 +1,64 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ * All rights reserved.
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <algorithm>
+#include <fmt/format.h>
+#include <functional>
+#include <map>
+#include <string>
+
+namespace KINETO_NAMESPACE {
+
+class ActivityLogger;
+
+class ActivityLoggerFactory {
+
+ public:
+  using FactoryFunc =
+    std::function<std::unique_ptr<ActivityLogger>(const std::string& url)>;
+
+  // Add logger factory for a protocol prefix
+  void addProtocol(const std::string& protocol, FactoryFunc f) {
+    factories_[tolower(protocol)] = f;
+  }
+
+  // Create a logger, invoking the factory for the protocol specified in url
+  std::unique_ptr<ActivityLogger> makeLogger(const std::string& url) const {
+    std::string protocol = extractProtocol(url);
+    auto it = factories_.find(tolower(protocol));
+    if (it != factories_.end()) {
+      return it->second(stripProtocol(url));
+    }
+    throw std::invalid_argument(fmt::format(
+        "No logger registered for the {} protocol prefix",
+        protocol));
+    return nullptr;
+  }
+
+ private:
+  static std::string tolower(std::string s) {
+    std::transform(s.begin(), s.end(), s.begin(),
+        [](unsigned char c) { return std::tolower(c); }
+    );
+    return s;
+  }
+
+  static std::string extractProtocol(std::string url) {
+    return url.substr(0, url.find("://"));
+  }
+
+  static std::string stripProtocol(std::string url) {
+    size_t pos = url.find("://");
+    return pos == url.npos ? url : url.substr(pos + 3);
+  }
+
+  std::map<std::string, FactoryFunc> factories_;
+};
+
+} // namespace KINETO_NAMESPACE

libkineto/src/ActivityProfilerController.cpp

@@ -10,6 +10,7 @@
 #include <chrono>
 #include <thread>
 
+#include "ActivityLoggerFactory.h"
 #include "ActivityTrace.h"
 #include "CuptiActivityInterface.h"
 #include "ThreadUtil.h"
@@ -39,26 +40,29 @@ ActivityProfilerController::~ActivityProfilerController() {
   VLOG(0) << "Stopped activity profiler";
 }
 
+static ActivityLoggerFactory initLoggerFactory() {
+  ActivityLoggerFactory factory;
+  factory.addProtocol("file", [](const std::string& url) {
+      return std::unique_ptr<ActivityLogger>(new ChromeTraceLogger(url));
+  });
+  return factory;
+}
+
 static ActivityLoggerFactory& loggerFactory() {
-  static ActivityLoggerFactory factory{nullptr};
+  static ActivityLoggerFactory factory = initLoggerFactory();
   return factory;
 }
 
-void ActivityProfilerController::setLoggerFactory(
-    const ActivityLoggerFactory& factory) {
-  loggerFactory() = factory;
+void ActivityProfilerController::addLoggerFactory(
+    const std::string& protocol, ActivityLoggerFactory::FactoryFunc factory) {
+  loggerFactory().addProtocol(protocol, factory);
 }
 
 static std::unique_ptr<ActivityLogger> makeLogger(const Config& config) {
   if (config.activitiesLogToMemory()) {
     return std::make_unique<MemoryTraceLogger>(config);
   }
-  if (loggerFactory()) {
-    return loggerFactory()(config);
-  }
-  return std::make_unique<ChromeTraceLogger>(
-      config.activitiesLogFile(),
-      CuptiActivityInterface::singleton().smCount());
+  return loggerFactory().makeLogger(config.activitiesLogUrl());
 }
 
 void ActivityProfilerController::profilerLoop() {
@@ -138,7 +142,7 @@ std::unique_ptr<ActivityTraceInterface> ActivityProfilerController::stopTrace()
   auto logger = std::make_unique<MemoryTraceLogger>(profiler_->config());
   profiler_->processTrace(*logger);
   profiler_->reset();
-  return std::make_unique<ActivityTrace>(std::move(logger), CuptiActivityInterface::singleton());
+  return std::make_unique<ActivityTrace>(std::move(logger), loggerFactory());
 }
 
 void ActivityProfilerController::addMetadata(

libkineto/src/ActivityProfilerController.h

@@ -11,6 +11,7 @@
 #include <memory>
 #include <thread>
 
+#include "ActivityLoggerFactory.h"
 #include "ActivityProfiler.h"
 #include "ActivityProfilerInterface.h"
 #include "ActivityTraceInterface.h"
@@ -20,9 +21,6 @@ namespace KINETO_NAMESPACE {
 
 class Config;
 
-using ActivityLoggerFactory =
-    std::function<std::unique_ptr<ActivityLogger>(const Config&)>;
-
 class ActivityProfilerController {
  public:
   explicit ActivityProfilerController(bool cpuOnly);
@@ -32,7 +30,9 @@ class ActivityProfilerController {
 
   ~ActivityProfilerController();
 
-  static void setLoggerFactory(const ActivityLoggerFactory& factory);
+  static void addLoggerFactory(
+      const std::string& protocol,
+      ActivityLoggerFactory::FactoryFunc factory);
 
   void scheduleTrace(const Config& config);
 

libkineto/src/ActivityTrace.h

@@ -10,8 +10,8 @@
 #include <memory>
 #include <string>
 
+#include "ActivityLoggerFactory.h"
 #include "ActivityTraceInterface.h"
-#include "CuptiActivityInterface.h"
 #include "output_json.h"
 #include "output_membuf.h"
 
@@ -20,22 +20,31 @@ namespace libkineto {
 class ActivityTrace : public ActivityTraceInterface {
  public:
   ActivityTrace(
-      std::unique_ptr<MemoryTraceLogger> logger,
-      CuptiActivityInterface& cuptiActivities)
-    : logger_(std::move(logger)), cuptiActivities_(cuptiActivities) {}
+      std::unique_ptr<MemoryTraceLogger> tmpLogger,
+      const ActivityLoggerFactory& factory)
+    : memLogger_(std::move(tmpLogger)),
+      loggerFactory_(factory) {
+  }
 
   const std::vector<std::unique_ptr<TraceActivity>>* activities() override {
-    return logger_->traceActivities();
+    return memLogger_->traceActivities();
   };
 
-  void save(const std::string& path) override {
-    ChromeTraceLogger chrome_logger(path, cuptiActivities_.smCount());
-    logger_->log(chrome_logger);
+  void save(const std::string& url) override {
+    std::string prefix;
+    // if no protocol is specified, default to file
+    if (url.find("://") == url.npos) {
+      prefix = "file://";
+    }
+    memLogger_->log(*loggerFactory_.makeLogger(prefix + url));
   };
 
  private:
-  std::unique_ptr<MemoryTraceLogger> logger_;
-  CuptiActivityInterface& cuptiActivities_;
+  // Activities are logged into a buffer
+  std::unique_ptr<MemoryTraceLogger> memLogger_;
+
+  // Alternative logger used by save() if protocol prefix is specified
+  const ActivityLoggerFactory& loggerFactory_;
 };
 
 } // namespace libkineto

libkineto/src/Config.cpp

@@ -61,6 +61,7 @@ const string kHeartbeatMonitorPeriodKey =
 const string kActivitiesEnabledKey = "ACTIVITIES_ENABLED";
 const string kActivityTypesKey = "ACTIVITY_TYPES";
 const string kActivitiesLogFileKey = "ACTIVITIES_LOG_FILE";
+const string kActivitiesLogUrlKey = "ACTIVITIES_LOG_URL";
 const string kActivitiesDurationKey = "ACTIVITIES_DURATION_SECS";
 const string kActivitiesDurationMsecsKey = "ACTIVITIES_DURATION_MSECS";
 const string kActivitiesIterationsKey = "ACTIVITIES_ITERATIONS";
@@ -118,16 +119,16 @@ const string kConfigFile = "/etc/libkineto.conf";
 // Max devices supported on any system
 constexpr uint8_t kMaxDevices = 8;
 
-static std::map<std::string, std::function<AbstractConfig*(const Config&)>>&
+static std::map<std::string, std::function<AbstractConfig*(Config&)>>&
 configFactories() {
-  static std::map<std::string, std::function<AbstractConfig*(const Config&)>>
+  static std::map<std::string, std::function<AbstractConfig*(Config&)>>
       factories;
   return factories;
 }
 
 void Config::addConfigFactory(
     std::string name,
-    std::function<AbstractConfig*(const Config&)> factory) {
+    std::function<AbstractConfig*(Config&)> factory) {
   configFactories()[name] = factory;
 }
 
@@ -286,6 +287,7 @@ bool Config::handleOption(const std::string& name, std::string& val) {
     activityProfilerEnabled_ = toBool(val);
   } else if (name == kActivitiesLogFileKey) {
     activitiesLogFile_ = val;
+    activitiesLogUrl_ = fmt::format("file://{}", val);
     activitiesOnDemandTimestamp_ = timestamp();
   } else if (name == kActivitiesMaxGpuBufferSizeKey) {
     activitiesMaxGpuBufferSize_ = toInt32(val) * 1024 * 1024;

libkineto/src/Config.h

@@ -53,6 +53,15 @@ class Config : public AbstractConfig {
     return activitiesLogFile_;
   }
 
+  // Log activitiy trace to this url
+  const std::string& activitiesLogUrl() const {
+    return activitiesLogUrl_;
+  }
+
+  void setActivitiesLogUrl(const std::string& url) {
+    activitiesLogUrl_ = url;
+  }
+
   bool activitiesLogToMemory() const {
     return activitiesLogToMemory_;
   }
@@ -272,7 +281,7 @@ class Config : public AbstractConfig {
 
   static void addConfigFactory(
       std::string name,
-      std::function<AbstractConfig*(const Config&)> factory);
+      std::function<AbstractConfig*(Config&)> factory);
 
   void print(std::ostream& s) const;
 
@@ -338,6 +347,8 @@ class Config : public AbstractConfig {
   // The activity profiler settings are all on-demand
   std::string activitiesLogFile_;
 
+  std::string activitiesLogUrl_;
+
   // Log activities to memory buffer
   bool activitiesLogToMemory_{false};
 

libkineto/src/CudaDeviceProperties.cpp

@@ -23,13 +23,16 @@ static const std::vector<cudaDeviceProp> createDeviceProps() {
   cudaError_t error_id = cudaGetDeviceCount(&device_count);
   // Return empty vector if error.
   if (error_id != cudaSuccess) {
+    LOG(ERROR) << "cudaGetDeviceCount failed with code " << error_id;
     return {};
   }
+  VLOG(0) << "Device count is " << device_count;
   for (size_t i = 0; i < device_count; ++i) {
     cudaDeviceProp prop;
     error_id = cudaGetDeviceProperties(&prop, i);
     // Return empty vector if any device property fail to get.
     if (error_id != cudaSuccess) {
+      LOG(ERROR) << "cudaGetDeviceProperties failed with " << error_id;
       return {};
     }
     props.push_back(prop);
@@ -75,6 +78,12 @@ const std::string& devicePropertiesJson() {
   return devicePropsJson;
 }
 
+int smCount(uint32_t deviceId) {
+  const std::vector<cudaDeviceProp> &props = deviceProps();
+  return deviceId >= props.size() ? 0 :
+     props[deviceId].multiProcessorCount;
+}
+
 float kernelOccupancy(
     uint32_t deviceId,
     uint16_t registersPerThread,

libkineto/src/CudaDeviceProperties.h

@@ -12,6 +12,8 @@
 
 namespace KINETO_NAMESPACE {
 
+int smCount(uint32_t deviceId);
+
 // Return estimated achieved occupancy for a kernel
 float kernelOccupancy(
     uint32_t deviceId,

libkineto/src/init.cpp

@@ -146,7 +146,8 @@ int InitializeInjection(void) {
 }
 
 void suppressLibkinetoLogMessages() {
-  SET_LOG_SEVERITY_LEVEL(ERROR);
+  //SET_LOG_SEVERITY_LEVEL(ERROR);
+  SET_LOG_VERBOSITY_LEVEL(2, {});
 }
 
 } // extern C

libkineto/src/output_json.cpp

@@ -30,6 +30,8 @@ using namespace libkineto;
 namespace KINETO_NAMESPACE {
 
 static constexpr int kSchemaVersion = 1;
+static const std::string kDefaultLogFileFmt =
+    "/tmp/libkineto_activities_{}.json";
 
 void ChromeTraceLogger::handleTraceStart(
     const std::unordered_map<std::string, std::string>& metadata) {
@@ -52,6 +54,10 @@ void ChromeTraceLogger::handleTraceStart(
   "traceEvents": [)JSON";
 }
 
+static std::string defaultFileName() {
+  return fmt::format(kDefaultLogFileFmt, processId());
+}
+
 void ChromeTraceLogger::openTraceFile() {
   traceOf_.open(fileName_, std::ofstream::out | std::ofstream::trunc);
   if (!traceOf_) {
@@ -61,13 +67,10 @@ void ChromeTraceLogger::openTraceFile() {
   }
 }
 
-ChromeTraceLogger::ChromeTraceLogger(const std::string& traceFileName, int smCount)
-    : fileName_(traceFileName) {
+ChromeTraceLogger::ChromeTraceLogger(const std::string& traceFileName) {
+  fileName_ = traceFileName.empty() ? defaultFileName() : traceFileName;
   traceOf_.clear(std::ios_base::badbit);
   openTraceFile();
-#ifdef HAS_CUPTI
-  smCount_ = CuptiActivityInterface::singleton().smCount();
-#endif
 }
 
 static int64_t us(int64_t timestamp) {
@@ -332,9 +335,10 @@ void ChromeTraceLogger::handleGpuActivity(
   constexpr int threads_per_warp = 32;
   float blocks_per_sm = -1.0;
   float warps_per_sm = -1.0;
-  if (smCount_) {
+  int sm_count = smCount(kernel->deviceId);
+  if (sm_count) {
     blocks_per_sm =
-        (kernel->gridX * kernel->gridY * kernel->gridZ) / (float) smCount_;
+        (kernel->gridX * kernel->gridY * kernel->gridZ) / (float) sm_count;
     warps_per_sm =
         blocks_per_sm * (kernel->blockX * kernel->blockY * kernel->blockZ)
         / threads_per_warp;