[PROTON] Get the context depth of profiling sessions (#6158)

Author: Jokeren

third_party/proton/csrc/Proton.cpp

@@ -78,6 +78,10 @@ void initProton(pybind11::module &&m) {
           SessionManager::instance().addMetrics(scopeId, metrics);
         });
 
+  m.def("get_context_depth", [](size_t sessionId) {
+    return SessionManager::instance().getContextDepth(sessionId);
+  });
+
   pybind11::bind_map<std::map<std::string, MetricValueType>>(m, "MetricMap");
 }
 

third_party/proton/csrc/include/Context/Context.h

@@ -43,6 +43,8 @@ class ContextSource {
 
   void setState(std::optional<Context> state) { ContextSource::state = state; }
 
+  virtual size_t getDepth() = 0;
+
 protected:
   virtual std::vector<Context> getContextsImpl() = 0;
   static thread_local std::optional<Context> state;

third_party/proton/csrc/include/Context/Python.h

@@ -10,6 +10,8 @@ class PythonContextSource : public ContextSource {
 public:
   PythonContextSource() = default;
 
+  size_t getDepth() override;
+
 private:
   std::vector<Context> getContextsImpl() override;
 };

third_party/proton/csrc/include/Context/Shadow.h

@@ -25,6 +25,8 @@ class ShadowContextSource : public ContextSource, public ScopeInterface {
 
   void exitScope(const Scope &scope) override;
 
+  size_t getDepth() override;
+
 private:
   std::vector<Context> getContextsImpl() override;
 

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

@@ -31,6 +31,8 @@ class Session {
 
   void finalize(OutputFormat outputFormat);
 
+  size_t getContextDepth();
+
 private:
   Session(size_t id, const std::string &path, Profiler *profiler,
           std::unique_ptr<ContextSource> contextSource,
@@ -88,6 +90,8 @@ class SessionManager : public Singleton<SessionManager> {
 
   void deactivateAllSessions();
 
+  size_t getContextDepth(size_t sessionId);
+
   void enterScope(const Scope &scope);
 
   void exitScope(const Scope &scope);

third_party/proton/csrc/lib/Context/Python.cpp

@@ -94,4 +94,6 @@ std::vector<Context> PythonContextSource::getContextsImpl() {
   return contexts;
 }
 
+size_t PythonContextSource::getDepth() { return getContextsImpl().size(); }
+
 } // namespace proton

third_party/proton/csrc/lib/Context/Shadow.cpp

@@ -26,6 +26,11 @@ std::vector<Context> ShadowContextSource::getContextsImpl() {
   return threadContextStack;
 }
 
+size_t ShadowContextSource::getDepth() {
+  initializeThreadContext();
+  return threadContextStack.size();
+}
+
 void ShadowContextSource::exitScope(const Scope &scope) {
   if (threadContextStack.empty()) {
     throw std::runtime_error("Context stack is empty");

third_party/proton/csrc/lib/Session/Session.cpp

@@ -69,6 +69,8 @@ void Session::finalize(OutputFormat outputFormat) {
   data->dump(outputFormat);
 }
 
+size_t Session::getContextDepth() { return contextSource->getDepth(); }
+
 std::unique_ptr<Session> SessionManager::makeSession(
     size_t id, const std::string &path, const std::string &profilerName,
     const std::string &profilerPath, const std::string &contextSourceName,
@@ -240,4 +242,10 @@ void SessionManager::setState(std::optional<Context> context) {
   }
 }
 
+size_t SessionManager::getContextDepth(size_t sessionId) {
+  std::lock_guard<std::mutex> lock(mutex);
+  throwIfSessionNotInitialized(sessions, sessionId);
+  return sessions[sessionId]->getContextDepth();
+}
+
 } // namespace proton

third_party/proton/proton/__init__.py

@@ -9,3 +9,4 @@
     profile,
     DEFAULT_PROFILE_NAME,
 )
+from . import context

third_party/proton/proton/context.py

@@ -0,0 +1,18 @@
+from typing import Optional
+from triton._C.libproton import proton as libproton
+from .flags import get_profiling_on
+
+
+def depth(session: Optional[int] = 0) -> Optional[int]:
+    """
+    Get the depth of the context.
+
+    Args:
+        session (int): The session ID of the profiling session. Defaults to 0.
+
+    Returns:
+        depth (int or None): The depth of the context. If profiling is off, returns None.
+    """
+    if not get_profiling_on():
+        return None
+    return libproton.get_context_depth(session)