[PROTON] Add proton.state utility (#5110)

`state` is different from `scope` in several ways:

1. State is not recursive; each operation can have only a single state.
Inner most state will overwrite the outer most state.
2. A states is a suffix, meaning that the original call path will append
a state above the name of each kernel.
3. State is compatible with both Python and shadow contexts.

Author: Jokeren

third_party/proton/README.md

@@ -162,6 +162,29 @@ proton-viewer -h
 
 ## Advanced features
 
+### State annotation
+
+In addition to `proton.scope`, we can also customize the call path of each GPU operation using `proton.state`.
+
+`state` is different from `scope` in several ways:
+
+1. State is not recursive; each operation can have only a single state. Inner most state will overwrite the outer most state.
+2. A states is a suffix, meaning that the original call path will append a state above the name of each kernel.
+3. State is compatible with both Python and shadow contexts.
+
+The following example demonstrates a basic use of state:
+
+```python
+with proton.scope("test"):
+    with proton.state("state0"):
+        with proton.scope("test0"):
+            foo0[1,](x, y)
+        with proton.scope("test1"):
+            foo1[1,](x, y)
+```
+
+The call path of `foo1` will be `test->test1->state0`.
+
 ### Instrumentation (experimental)
 
 In addition to profiling, Proton also incorporates MLIR/LLVM based compiler instrumentation passes to get Triton level analysis

third_party/proton/csrc/Proton.cpp

@@ -64,6 +64,13 @@ void initProton(pybind11::module &&m) {
     SessionManager::instance().exitOp(Scope(scopeId, name));
   });
 
+  m.def("enter_state", [](const std::string &state) {
+    SessionManager::instance().setState(state);
+  });
+
+  m.def("exit_state",
+        []() { SessionManager::instance().setState(std::nullopt); });
+
   m.def("add_metrics",
         [](size_t scopeId,
            const std::map<std::string, MetricValueType> &metrics) {

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

@@ -5,6 +5,7 @@
 #include <limits>
 #include <map>
 #include <mutex>
+#include <optional>
 #include <string>
 #include <vector>
 
@@ -31,7 +32,20 @@ class ContextSource {
 public:
   ContextSource() = default;
   virtual ~ContextSource() = default;
-  virtual std::vector<Context> getContexts() = 0;
+
+  std::vector<Context> getContexts() {
+    auto contexts = getContextsImpl();
+    if (state.has_value()) {
+      contexts.push_back(state.value());
+    }
+    return contexts;
+  }
+
+  void setState(std::optional<Context> state) { ContextSource::state = state; }
+
+protected:
+  virtual std::vector<Context> getContextsImpl() = 0;
+  static thread_local std::optional<Context> state;
 };
 
 /// A scope is a context with a unique identifier.

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

@@ -8,7 +8,10 @@ namespace proton {
 /// Unwind the Python stack and early return a list of contexts.
 class PythonContextSource : public ContextSource {
 public:
-  std::vector<Context> getContexts() override;
+  PythonContextSource() = default;
+
+private:
+  std::vector<Context> getContextsImpl() override;
 };
 
 } // namespace proton

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

@@ -12,13 +12,12 @@ class ShadowContextSource : public ContextSource, public ScopeInterface {
 public:
   ShadowContextSource() = default;
 
-  std::vector<Context> getContexts() override { return contextStack; }
-
   void enterScope(const Scope &scope) override;
 
   void exitScope(const Scope &scope) override;
 
 private:
+  std::vector<Context> getContextsImpl() override { return contextStack; }
   std::vector<Context> contextStack;
 };
 

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

@@ -95,6 +95,8 @@ class SessionManager : public Singleton<SessionManager> {
                   const std::map<std::string, MetricValueType> &metrics,
                   bool aggregable);
 
+  void setState(std::optional<Context> context);
+
 private:
   std::unique_ptr<Session> makeSession(size_t id, const std::string &path,
                                        const std::string &profilerName,
@@ -146,6 +148,8 @@ class SessionManager : public Singleton<SessionManager> {
   std::map<ScopeInterface *, size_t> scopeInterfaceCounts;
   // op -> active count
   std::map<OpInterface *, size_t> opInterfaceCounts;
+  // context source -> active count
+  std::map<ContextSource *, size_t> contextSourceCounts;
 };
 
 } // namespace proton

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

@@ -2,6 +2,9 @@
 
 namespace proton {
 
+/*static*/ thread_local std::optional<Context> ContextSource::state =
+    std::nullopt;
+
 std::atomic<size_t> Scope::scopeIdCounter{1};
 
 /*static*/ thread_local std::map<ThreadLocalOpInterface *, bool>

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

@@ -71,7 +71,7 @@ std::string unpackPyobject(PyObject *pyObject) {
 
 } // namespace
 
-std::vector<Context> PythonContextSource::getContexts() {
+std::vector<Context> PythonContextSource::getContextsImpl() {
   pybind11::gil_scoped_acquire gil;
 
   PyFrameObject *frame = PyEval_GetFrame();

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

@@ -111,6 +111,7 @@ void SessionManager::activateSessionImpl(size_t sessionId) {
   sessions[sessionId]->activate();
   registerInterface<ScopeInterface>(sessionId, scopeInterfaceCounts);
   registerInterface<OpInterface>(sessionId, opInterfaceCounts);
+  registerInterface<ContextSource>(sessionId, contextSourceCounts);
 }
 
 void SessionManager::deActivateSessionImpl(size_t sessionId) {
@@ -122,6 +123,7 @@ void SessionManager::deActivateSessionImpl(size_t sessionId) {
   sessions[sessionId]->deactivate();
   unregisterInterface<ScopeInterface>(sessionId, scopeInterfaceCounts);
   unregisterInterface<OpInterface>(sessionId, opInterfaceCounts);
+  unregisterInterface<ContextSource>(sessionId, contextSourceCounts);
 }
 
 void SessionManager::removeSession(size_t sessionId) {
@@ -226,4 +228,14 @@ void SessionManager::addMetrics(
   }
 }
 
+void SessionManager::setState(std::optional<Context> context) {
+  std::shared_lock<std::shared_mutex> lock(mutex);
+  for (auto iter : contextSourceCounts) {
+    auto [contextSource, count] = iter;
+    if (count > 0) {
+      contextSource->setState(context);
+    }
+  }
+}
+
 } // namespace proton

third_party/proton/proton/__init__.py

@@ -1,5 +1,6 @@
 # flake8: noqa
 from .scope import scope, enter_scope, exit_scope
+from .state import state, enter_state, exit_state
 from .profile import (
     start,
     activate,