[PROTON] Implement the get_data api to export profile directly in Python

third_party/proton/csrc/Proton.cpp

@@ -165,6 +165,13 @@ static void initProton(pybind11::module &&m) {
   m.def("get_context_depth", [](size_t sessionId) {
     return SessionManager::instance().getContextDepth(sessionId);
   });
+
+  m.def(
+      "get_data",
+      [](size_t sessionId) {
+        return SessionManager::instance().getData(sessionId);
+      },
+      pybind11::arg("sessionId"));
 }
 
 PYBIND11_MODULE(libproton, m) {

third_party/proton/csrc/include/Data/Data.h

@@ -43,6 +43,9 @@ class Data : public ScopeInterface {
   /// Clear all caching data.
   virtual void clear() = 0;
 
+  /// To Json
+  virtual std::string toJsonString() const = 0;
+
   /// Dump the data to the given output format.
   void dump(const std::string &outputFormat);
 

third_party/proton/csrc/include/Data/TraceData.h

@@ -22,6 +22,8 @@ class TraceData : public Data {
   addMetrics(size_t scopeId,
              const std::map<std::string, MetricValueType> &metrics) override;
 
+  std::string toJsonString() const override;
+
   void clear() override;
 
   class Trace;

third_party/proton/csrc/include/Data/TreeData.h

@@ -3,9 +3,13 @@
 
 #include "Context/Context.h"
 #include "Data.h"
+#include "nlohmann/json.hpp"
 #include <stdexcept>
+#include <string>
 #include <unordered_map>
 
+using json = nlohmann::json;
+
 namespace proton {
 
 class TreeData : public Data {
@@ -25,6 +29,8 @@ class TreeData : public Data {
   addMetrics(size_t scopeId,
              const std::map<std::string, MetricValueType> &metrics) override;
 
+  std::string toJsonString() const override;
+
   void clear() override;
 
 protected:
@@ -34,6 +40,12 @@ class TreeData : public Data {
   void exitScope(const Scope &scope) override;
 
 private:
+  // `tree` and `scopeIdToContextId` can be accessed by both the user thread and
+  // the background threads concurrently, so methods that access them should be
+  // protected by a (shared) mutex.
+  class Tree;
+  json buildHatchetJson(TreeData::Tree *tree) const;
+
   void dumpHatchet(std::ostream &os) const;
 
   void doDump(std::ostream &os, OutputFormat outputFormat) const override;
@@ -42,10 +54,6 @@ class TreeData : public Data {
     return OutputFormat::Hatchet;
   }
 
-  // `tree` and `scopeIdToContextId` can be accessed by both the user thread and
-  // the background threads concurrently, so methods that access them should be
-  // protected by a (shared) mutex.
-  class Tree;
   std::unique_ptr<Tree> tree;
   // ScopeId -> ContextId
   std::unordered_map<size_t, size_t> scopeIdToContextId;

third_party/proton/csrc/include/Session/Session.h

@@ -93,6 +93,8 @@ class SessionManager : public Singleton<SessionManager> {
 
   size_t getContextDepth(size_t sessionId);
 
+  std::string getData(size_t sessionId);
+
   void enterScope(const Scope &scope);
 
   void exitScope(const Scope &scope);

third_party/proton/csrc/lib/Data/TraceData.cpp

@@ -229,6 +229,8 @@ void TraceData::addMetrics(
   }
 }
 
+std::string TraceData::toJsonString() const { throw NotImplemented(); }
+
 void TraceData::clear() {
   std::unique_lock<std::shared_mutex> lock(mutex);
   scopeIdToContextId.clear();

third_party/proton/csrc/lib/Data/TreeData.cpp

@@ -2,16 +2,13 @@
 #include "Context/Context.h"
 #include "Data/Metric.h"
 #include "Device.h"
-#include "nlohmann/json.hpp"
 
 #include <limits>
 #include <map>
 #include <mutex>
 #include <set>
 #include <stdexcept>
 
-using json = nlohmann::json;
-
 namespace proton {
 
 class TreeData::Tree {
@@ -106,6 +103,134 @@ class TreeData::Tree {
   std::map<size_t, TreeNode> treeNodeMap;
 };
 
+json TreeData::buildHatchetJson(TreeData::Tree *tree) const {
+  std::map<size_t, json *> jsonNodes;
+  json output = json::array();
+  output.push_back(json::object());
+  jsonNodes[TreeData::Tree::TreeNode::RootId] = &(output.back());
+  std::set<std::string> inclusiveValueNames;
+  std::map<uint64_t, std::set<uint64_t>> deviceIds;
+  tree->template walk<TreeData::Tree::WalkPolicy::PreOrder>(
+      [&](TreeData::Tree::TreeNode &treeNode) {
+        const auto contextName = treeNode.name;
+        auto contextId = treeNode.id;
+        json *jsonNode = jsonNodes[contextId];
+        (*jsonNode)["frame"] = {{"name", contextName}, {"type", "function"}};
+        (*jsonNode)["metrics"] = json::object();
+        for (auto [metricKind, metric] : treeNode.metrics) {
+          if (metricKind == MetricKind::Kernel) {
+            std::shared_ptr<KernelMetric> kernelMetric =
+                std::dynamic_pointer_cast<KernelMetric>(metric);
+            uint64_t duration = std::get<uint64_t>(
+                kernelMetric->getValue(KernelMetric::Duration));
+            uint64_t invocations = std::get<uint64_t>(
+                kernelMetric->getValue(KernelMetric::Invocations));
+            uint64_t deviceId = std::get<uint64_t>(
+                kernelMetric->getValue(KernelMetric::DeviceId));
+            uint64_t deviceType = std::get<uint64_t>(
+                kernelMetric->getValue(KernelMetric::DeviceType));
+            std::string deviceTypeName =
+                getDeviceTypeString(static_cast<DeviceType>(deviceType));
+            (*jsonNode)["metrics"]
+                       [kernelMetric->getValueName(KernelMetric::Duration)] =
+                           duration;
+            (*jsonNode)["metrics"]
+                       [kernelMetric->getValueName(KernelMetric::Invocations)] =
+                           invocations;
+            (*jsonNode)["metrics"]
+                       [kernelMetric->getValueName(KernelMetric::DeviceId)] =
+                           std::to_string(deviceId);
+            (*jsonNode)["metrics"]
+                       [kernelMetric->getValueName(KernelMetric::DeviceType)] =
+                           deviceTypeName;
+            inclusiveValueNames.insert(
+                kernelMetric->getValueName(KernelMetric::Duration));
+            inclusiveValueNames.insert(
+                kernelMetric->getValueName(KernelMetric::Invocations));
+            deviceIds[deviceType].insert(deviceId);
+          } else if (metricKind == MetricKind::PCSampling) {
+            auto pcSamplingMetric =
+                std::dynamic_pointer_cast<PCSamplingMetric>(metric);
+            for (size_t i = 0; i < PCSamplingMetric::Count; i++) {
+              auto valueName = pcSamplingMetric->getValueName(i);
+              inclusiveValueNames.insert(valueName);
+              std::visit(
+                  [&](auto &&value) {
+                    (*jsonNode)["metrics"][valueName] = value;
+                  },
+                  pcSamplingMetric->getValues()[i]);
+            }
+          } else if (metricKind == MetricKind::Cycle) {
+            auto cycleMetric = std::dynamic_pointer_cast<CycleMetric>(metric);
+            uint64_t duration = std::get<uint64_t>(
+                cycleMetric->getValue(CycleMetric::Duration));
+            double normalizedDuration = std::get<double>(
+                cycleMetric->getValue(CycleMetric::NormalizedDuration));
+            uint64_t deviceId = std::get<uint64_t>(
+                cycleMetric->getValue(CycleMetric::DeviceId));
+            uint64_t deviceType = std::get<uint64_t>(
+                cycleMetric->getValue(CycleMetric::DeviceType));
+            (*jsonNode)["metrics"]
+                       [cycleMetric->getValueName(CycleMetric::Duration)] =
+                           duration;
+            (*jsonNode)["metrics"][cycleMetric->getValueName(
+                CycleMetric::NormalizedDuration)] = normalizedDuration;
+            (*jsonNode)["metrics"]
+                       [cycleMetric->getValueName(CycleMetric::DeviceId)] =
+                           std::to_string(deviceId);
+            (*jsonNode)["metrics"]
+                       [cycleMetric->getValueName(CycleMetric::DeviceType)] =
+                           std::to_string(deviceType);
+            deviceIds[deviceType].insert(deviceId);
+          } else if (metricKind == MetricKind::Flexible) {
+            // Flexible metrics are handled in a different way
+          } else {
+            throw std::runtime_error("MetricKind not supported");
+          }
+        }
+        for (auto [_, flexibleMetric] : treeNode.flexibleMetrics) {
+          auto valueName = flexibleMetric.getValueName(0);
+          if (!flexibleMetric.isExclusive(0))
+            inclusiveValueNames.insert(valueName);
+          std::visit(
+              [&](auto &&value) { (*jsonNode)["metrics"][valueName] = value; },
+              flexibleMetric.getValues()[0]);
+        }
+        (*jsonNode)["children"] = json::array();
+        auto children = treeNode.children;
+        for (auto _ : children) {
+          (*jsonNode)["children"].push_back(json::object());
+        }
+        auto idx = 0;
+        for (auto child : children) {
+          auto [index, childId] = child;
+          jsonNodes[childId] = &(*jsonNode)["children"][idx];
+          idx++;
+        }
+      });
+  for (auto valueName : inclusiveValueNames) {
+    output[TreeData::Tree::TreeNode::RootId]["metrics"][valueName] = 0;
+  }
+  output.push_back(json::object());
+  auto &deviceJson = output.back();
+  for (auto [deviceType, deviceIdSet] : deviceIds) {
+    auto deviceTypeName =
+        getDeviceTypeString(static_cast<DeviceType>(deviceType));
+    if (!deviceJson.contains(deviceTypeName))
+      deviceJson[deviceTypeName] = json::object();
+    for (auto deviceId : deviceIdSet) {
+      Device device = getDevice(static_cast<DeviceType>(deviceType), deviceId);
+      deviceJson[deviceTypeName][std::to_string(deviceId)] = {
+          {"clock_rate", device.clockRate},
+          {"memory_clock_rate", device.memoryClockRate},
+          {"bus_width", device.busWidth},
+          {"arch", device.arch},
+          {"num_sms", device.numSms}};
+    }
+  }
+  return output;
+}
+
 void TreeData::enterScope(const Scope &scope) {
   // enterOp and addMetric maybe called from different threads
   std::unique_lock<std::shared_mutex> lock(mutex);
@@ -201,136 +326,16 @@ void TreeData::clear() {
 }
 
 void TreeData::dumpHatchet(std::ostream &os) const {
-  std::map<size_t, json *> jsonNodes;
-  json output = json::array();
-  output.push_back(json::object());
-  jsonNodes[Tree::TreeNode::RootId] = &(output.back());
-  std::set<std::string> inclusiveValueNames;
-  std::map<uint64_t, std::set<uint64_t>> deviceIds;
-  this->tree->template walk<Tree::WalkPolicy::PreOrder>([&](Tree::TreeNode
-                                                                &treeNode) {
-    const auto contextName = treeNode.name;
-    auto contextId = treeNode.id;
-    json *jsonNode = jsonNodes[contextId];
-    (*jsonNode)["frame"] = {{"name", contextName}, {"type", "function"}};
-    (*jsonNode)["metrics"] = json::object();
-    for (auto [metricKind, metric] : treeNode.metrics) {
-      if (metricKind == MetricKind::Kernel) {
-        std::shared_ptr<KernelMetric> kernelMetric =
-            std::dynamic_pointer_cast<KernelMetric>(metric);
-        uint64_t duration =
-            std::get<uint64_t>(kernelMetric->getValue(KernelMetric::Duration));
-        uint64_t invocations = std::get<uint64_t>(
-            kernelMetric->getValue(KernelMetric::Invocations));
-        uint64_t deviceId =
-            std::get<uint64_t>(kernelMetric->getValue(KernelMetric::DeviceId));
-        uint64_t deviceType = std::get<uint64_t>(
-            kernelMetric->getValue(KernelMetric::DeviceType));
-        std::string deviceTypeName =
-            getDeviceTypeString(static_cast<DeviceType>(deviceType));
-        (*jsonNode)["metrics"]
-                   [kernelMetric->getValueName(KernelMetric::Duration)] =
-                       duration;
-        (*jsonNode)["metrics"]
-                   [kernelMetric->getValueName(KernelMetric::Invocations)] =
-                       invocations;
-        (*jsonNode)["metrics"]
-                   [kernelMetric->getValueName(KernelMetric::DeviceId)] =
-                       std::to_string(deviceId);
-        (*jsonNode)["metrics"]
-                   [kernelMetric->getValueName(KernelMetric::DeviceType)] =
-                       deviceTypeName;
-        inclusiveValueNames.insert(
-            kernelMetric->getValueName(KernelMetric::Duration));
-        inclusiveValueNames.insert(
-            kernelMetric->getValueName(KernelMetric::Invocations));
-        deviceIds[deviceType].insert(deviceId);
-      } else if (metricKind == MetricKind::PCSampling) {
-        auto pcSamplingMetric =
-            std::dynamic_pointer_cast<PCSamplingMetric>(metric);
-        for (size_t i = 0; i < PCSamplingMetric::Count; i++) {
-          auto valueName = pcSamplingMetric->getValueName(i);
-          inclusiveValueNames.insert(valueName);
-          std::visit(
-              [&](auto &&value) { (*jsonNode)["metrics"][valueName] = value; },
-              pcSamplingMetric->getValues()[i]);
-        }
-      } else if (metricKind == MetricKind::Cycle) {
-        auto cycleMetric = std::dynamic_pointer_cast<CycleMetric>(metric);
-        uint64_t duration =
-            std::get<uint64_t>(cycleMetric->getValue(CycleMetric::Duration));
-        double normalizedDuration = std::get<double>(
-            cycleMetric->getValue(CycleMetric::NormalizedDuration));
-        uint64_t deviceId =
-            std::get<uint64_t>(cycleMetric->getValue(CycleMetric::DeviceId));
-        uint64_t deviceType =
-            std::get<uint64_t>(cycleMetric->getValue(CycleMetric::DeviceType));
-        (*jsonNode)["metrics"]
-                   [cycleMetric->getValueName(CycleMetric::Duration)] =
-                       duration;
-        (*jsonNode)["metrics"][cycleMetric->getValueName(
-            CycleMetric::NormalizedDuration)] = normalizedDuration;
-        (*jsonNode)["metrics"]
-                   [cycleMetric->getValueName(CycleMetric::DeviceId)] =
-                       std::to_string(deviceId);
-        (*jsonNode)["metrics"]
-                   [cycleMetric->getValueName(CycleMetric::DeviceType)] =
-                       std::to_string(deviceType);
-        deviceIds[deviceType].insert(deviceId);
-      } else if (metricKind == MetricKind::Flexible) {
-        // Flexible metrics are handled in a different way
-      } else {
-        throw std::runtime_error("MetricKind not supported");
-      }
-    }
-    for (auto [_, flexibleMetric] : treeNode.flexibleMetrics) {
-      auto valueName = flexibleMetric.getValueName(0);
-      if (!flexibleMetric.isExclusive(0))
-        inclusiveValueNames.insert(valueName);
-      std::visit(
-          [&](auto &&value) { (*jsonNode)["metrics"][valueName] = value; },
-          flexibleMetric.getValues()[0]);
-    }
-    (*jsonNode)["children"] = json::array();
-    auto children = treeNode.children;
-    for (auto _ : children) {
-      (*jsonNode)["children"].push_back(json::object());
-    }
-    auto idx = 0;
-    for (auto child : children) {
-      auto [index, childId] = child;
-      jsonNodes[childId] = &(*jsonNode)["children"][idx];
-      idx++;
-    }
-  });
-  // Hints for all inclusive metrics
-  for (auto valueName : inclusiveValueNames) {
-    output[Tree::TreeNode::RootId]["metrics"][valueName] = 0;
-  }
-  // Prepare the device information
-  // Note that this is done from the application thread,
-  // query device information from the tool thread (e.g., CUPTI) will have
-  // problems
-  output.push_back(json::object());
-  auto &deviceJson = output.back();
-  for (auto [deviceType, deviceIdSet] : deviceIds) {
-    auto deviceTypeName =
-        getDeviceTypeString(static_cast<DeviceType>(deviceType));
-    if (!deviceJson.contains(deviceTypeName))
-      deviceJson[deviceTypeName] = json::object();
-    for (auto deviceId : deviceIdSet) {
-      Device device = getDevice(static_cast<DeviceType>(deviceType), deviceId);
-      deviceJson[deviceTypeName][std::to_string(deviceId)] = {
-          {"clock_rate", device.clockRate},
-          {"memory_clock_rate", device.memoryClockRate},
-          {"bus_width", device.busWidth},
-          {"arch", device.arch},
-          {"num_sms", device.numSms}};
-    }
-  }
+  auto output = buildHatchetJson(tree.get());
   os << std::endl << output.dump(4) << std::endl;
 }
 
+std::string TreeData::toJsonString() const {
+  std::shared_lock<std::shared_mutex> lock(mutex);
+  auto output = buildHatchetJson(tree.get());
+  return output.dump();
+}
+
 void TreeData::doDump(std::ostream &os, OutputFormat outputFormat) const {
   if (outputFormat == OutputFormat::Hatchet) {
     dumpHatchet(os);