[Proton] Fixed pc sampling error (#5787)

Fixed PC sampling error in proton when we have `mode`.

---------

Co-authored-by: Yuanwei Fang <[email protected]>

Author: 2 people

third_party/proton/csrc/Proton.cpp

@@ -16,9 +16,10 @@ void initProton(pybind11::module &&m) {
   m.def("start",
         [](const std::string &path, const std::string &contextSourceName,
            const std::string &dataName, const std::string &profilerName,
-           const std::string &profilerPath) {
+           const std::string &profilerPath, const std::string &mode) {
           auto sessionId = SessionManager::instance().addSession(
-              path, profilerName, profilerPath, contextSourceName, dataName);
+              path, profilerName, profilerPath, contextSourceName, dataName,
+              mode);
           SessionManager::instance().activateSession(sessionId);
           return sessionId;
         });

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

@@ -74,7 +74,7 @@ class SessionManager : public Singleton<SessionManager> {
   size_t addSession(const std::string &path, const std::string &profilerName,
                     const std::string &profilerPath,
                     const std::string &contextSourceName,
-                    const std::string &dataName);
+                    const std::string &dataName, const std::string &mode);
 
   void finalizeSession(size_t sessionId, OutputFormat outputFormat);
 
@@ -106,7 +106,8 @@ class SessionManager : public Singleton<SessionManager> {
                                        const std::string &profilerName,
                                        const std::string &profilerPath,
                                        const std::string &contextSourceName,
-                                       const std::string &dataName);
+                                       const std::string &dataName,
+                                       const std::string &mode);
 
   void activateSessionImpl(size_t sessionId);
 

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

@@ -9,12 +9,14 @@
 namespace proton {
 
 namespace {
-Profiler *getProfiler(const std::string &name, const std::string &path) {
+Profiler *getProfiler(const std::string &name, const std::string &path,
+                      const std::string &mode) {
   if (proton::toLower(name) == "cupti") {
-    return &CuptiProfiler::instance().setLibPath(path);
-  }
-  if (proton::toLower(name) == "cupti_pcsampling") {
-    return &CuptiProfiler::instance().setLibPath(path).enablePCSampling();
+    auto *profiler = &CuptiProfiler::instance();
+    profiler->setLibPath(path);
+    if (proton::toLower(mode) == "pcsampling")
+      profiler->enablePCSampling();
+    return profiler;
   }
   if (proton::toLower(name) == "roctracer") {
     return &RoctracerProfiler::instance();
@@ -72,8 +74,8 @@ void Session::finalize(OutputFormat outputFormat) {
 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,
-    const std::string &dataName) {
-  auto profiler = getProfiler(profilerName, profilerPath);
+    const std::string &dataName, const std::string &mode) {
+  auto profiler = getProfiler(profilerName, profilerPath, mode);
   auto contextSource = makeContextSource(contextSourceName);
   auto data = makeData(dataName, path, contextSource.get());
   auto *session = new Session(id, path, profiler, std::move(contextSource),
@@ -142,7 +144,8 @@ size_t SessionManager::addSession(const std::string &path,
                                   const std::string &profilerName,
                                   const std::string &profilerPath,
                                   const std::string &contextSourceName,
-                                  const std::string &dataName) {
+                                  const std::string &dataName,
+                                  const std::string &mode) {
   std::lock_guard<std::mutex> lock(mutex);
   if (hasSession(path)) {
     auto sessionId = getSessionId(path);
@@ -152,7 +155,7 @@ size_t SessionManager::addSession(const std::string &path,
   auto sessionId = nextSessionId++;
   sessionPaths[path] = sessionId;
   sessions[sessionId] = makeSession(sessionId, path, profilerName, profilerPath,
-                                    contextSourceName, dataName);
+                                    contextSourceName, dataName, mode);
   return sessionId;
 }
 

third_party/proton/proton/profile.py

@@ -113,10 +113,17 @@ def start(
 
     backend_path = _get_backend_default_path(backend)
 
+    _check_mode(backend, mode)
+
+    if mode is None:
+        mode = ""
+
+    backend_path = _get_backend_default_path(backend)
+
     set_profiling_on()
     if hook and hook == "triton":
         register_triton_hook()
-    return libproton.start(name, context, data, backend, backend_path)
+    return libproton.start(name, context, data, backend, backend_path, mode)
 
 
 def activate(session: Optional[int] = None) -> None:

third_party/proton/test/test_lib.py

@@ -32,7 +32,7 @@ def test_op():
 
 def test_session(tmp_path: pathlib.Path):
     temp_file = tmp_path / "test_session.hatchet"
-    session_id = libproton.start(str(temp_file.with_suffix("")), "shadow", "tree", _select_backend(), "")
+    session_id = libproton.start(str(temp_file.with_suffix("")), "shadow", "tree", _select_backend(), "", "")
     libproton.deactivate(session_id)
     libproton.activate(session_id)
     libproton.finalize(session_id, "hatchet")
@@ -42,7 +42,7 @@ def test_session(tmp_path: pathlib.Path):
 
 def test_add_metrics(tmp_path: pathlib.Path):
     temp_file = tmp_path / "test_add_metrics.hatchet"
-    libproton.start(str(temp_file.with_suffix("")), "shadow", "tree", _select_backend(), "")
+    libproton.start(str(temp_file.with_suffix("")), "shadow", "tree", _select_backend(), "", "")
     id1 = libproton.record_scope()
     libproton.enter_scope(id1, "one")
     libproton.add_metrics(id1, {"a": 1.0, "b": 2.0})