[PROTON] Reworked the mechanism for finding libraries for profiling backends. (#5751)

Jokeren · web-flow · commit f47cc3eaaa11 · 2025-01-29T18:24:47.000-08:00
We should be able to use the default CUPTI path when Triton is
installed, whether in development mode or within site packages.
Additionally, we resolved incorrect architecture-to-package name
mappings in setup.py.
diff --git a/python/setup.py b/python/setup.py
@@ -315,7 +315,8 @@ def download_and_copy(name, src_func, dst_path, variable, version, url_func):
     base_dir = os.path.dirname(__file__)
     system = platform.system()
     arch = platform.machine()
-    arch = {"arm64": "aarch64"}.get(arch, arch)
+    # NOTE: This might be wrong for jetson if both grace chips and jetson chips return aarch64
+    arch = {"arm64": "sbsa", "aarch64": "sbsa"}.get(arch, arch)
     supported = {"Linux": "linux", "Darwin": "linux"}
     url = url_func(supported[system], arch, version)
     src_path = src_func(supported[system], arch, version)
@@ -407,11 +408,7 @@ def get_proton_cmake_args(self):
         if cupti_include_dir == "":
             cupti_include_dir = os.path.join(get_base_dir(), "third_party", "nvidia", "backend", "include")
         cmake_args += ["-DCUPTI_INCLUDE_DIR=" + cupti_include_dir]
-        cupti_lib_dir = get_env_with_keys(["TRITON_CUPTI_LIB_PATH"])
-        if cupti_lib_dir == "":
-            cupti_lib_dir = os.path.join(get_base_dir(), "third_party", "nvidia", "backend", "lib", "cupti")
-        cmake_args += ["-DCUPTI_LIB_DIR=" + cupti_lib_dir]
-        roctracer_include_dir = get_env_with_keys(["ROCTRACER_INCLUDE_PATH"])
+        roctracer_include_dir = get_env_with_keys(["TRITON_ROCTRACER_INCLUDE_PATH"])
         if roctracer_include_dir == "":
             roctracer_include_dir = os.path.join(get_base_dir(), "third_party", "amd", "backend", "include")
         cmake_args += ["-DROCTRACER_INCLUDE_DIR=" + roctracer_include_dir]
diff --git a/third_party/proton/CMakeLists.txt b/third_party/proton/CMakeLists.txt
@@ -29,11 +29,6 @@ if(APPLE)
   set(PROTON_PYTHON_LDFLAGS "-undefined dynamic_lookup")
 endif()
 
-if(DEFINED CUPTI_LIB_DIR)
-  message(STATUS "CUPTI lib directory: ${CUPTI_LIB_DIR}")
-  add_compile_definitions(CUPTI_LIB_DIR=${CUPTI_LIB_DIR})
-endif()
-
 include_directories(${CUPTI_INCLUDE_DIR})
 include_directories(SYSTEM ${ROCTRACER_INCLUDE_DIR})
 target_compile_definitions(proton PRIVATE __HIP_PLATFORM_AMD__)
diff --git a/third_party/proton/csrc/Proton.cpp b/third_party/proton/csrc/Proton.cpp
@@ -15,9 +15,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 &dataName, const std::string &profilerName,
+           const std::string &profilerPath) {
           auto sessionId = SessionManager::instance().addSession(
-              path, profilerName, contextSourceName, dataName);
+              path, profilerName, profilerPath, contextSourceName, dataName);
           SessionManager::instance().activateSession(sessionId);
           return sessionId;
         });
diff --git a/third_party/proton/csrc/include/Driver/Dispatch.h b/third_party/proton/csrc/include/Driver/Dispatch.h
@@ -44,13 +44,13 @@ namespace proton {
 
 struct ExternLibBase {
   using RetType = int; // Generic type, can be overridden in derived structs
-  static constexpr const char *name = "";       // Placeholder
-  static constexpr const char *defaultDir = ""; // Placeholder
-  static constexpr RetType success = 0;         // Placeholder
+  static constexpr const char *name = ""; // Placeholder
+  static constexpr RetType success = 0;   // Placeholder
   ExternLibBase() = delete;
   ExternLibBase(const ExternLibBase &) = delete;
   ExternLibBase &operator=(const ExternLibBase &) = delete;
   static inline void *lib{nullptr};
+  static inline std::string defaultDir{""};
 };
 
 template <typename ExternLib> class Dispatch {
@@ -59,25 +59,24 @@ template <typename ExternLib> class Dispatch {
 
   static void init(const char *name, void **lib) {
     if (*lib == nullptr) {
-      // First reuse the existing handle
-      *lib = dlopen(name, RTLD_NOLOAD);
-    }
-    if (*lib == nullptr) {
-      // If not found, try to load it from LD_LIBRARY_PATH
-      *lib = dlopen(name, RTLD_LOCAL | RTLD_LAZY);
-    }
-    if (*lib == nullptr) {
-      // If still not found, try to load it from the default path
+      // If not found, try to load it from the default path
       auto dir = std::string(ExternLib::defaultDir);
       if (dir.length() > 0) {
         auto fullPath = dir + "/" + name;
         *lib = dlopen(fullPath.c_str(), RTLD_LOCAL | RTLD_LAZY);
+      } else {
+        // Only if the default path is not set, we try to load it from the
+        // system.
+        // First reuse the existing handle
+        *lib = dlopen(name, RTLD_NOLOAD);
+        if (*lib == nullptr) {
+          // If not found, try to load it from LD_LIBRARY_PATH
+          *lib = dlopen(name, RTLD_LOCAL | RTLD_LAZY);
+        }
       }
     }
     if (*lib == nullptr) {
-      throw std::runtime_error("Could not find `" + std::string(name) +
-                               "`. Make sure it is in your "
-                               "LD_LIBRARY_PATH.");
+      throw std::runtime_error("Could not load `" + std::string(name) + "`");
     }
   }
 
diff --git a/third_party/proton/csrc/include/Driver/GPU/CuptiApi.h b/third_party/proton/csrc/include/Driver/GPU/CuptiApi.h
@@ -3,6 +3,7 @@
 
 #include "cupti.h"
 #include "cupti_pcsampling.h"
+#include <string>
 
 namespace proton {
 
@@ -106,6 +107,8 @@ CUptiResult pcSamplingStart(CUpti_PCSamplingStartParams *pParams);
 template <bool CheckSuccess>
 CUptiResult pcSamplingStop(CUpti_PCSamplingStopParams *pParams);
 
+void setLibPath(const std::string &path);
+
 } // namespace cupti
 
 } // namespace proton
diff --git a/third_party/proton/csrc/include/Profiler/GPUProfiler.h b/third_party/proton/csrc/include/Profiler/GPUProfiler.h
@@ -41,6 +41,11 @@ class GPUProfiler : public Profiler,
   }
   bool isPCSamplingEnabled() const { return pcSamplingEnabled; }
 
+  ConcreteProfilerT &setLibPath(const std::string &libPath) {
+    pImpl->setLibPath(libPath);
+    return dynamic_cast<ConcreteProfilerT &>(*this);
+  }
+
 protected:
   // OpInterface
   void startOp(const Scope &scope) override {
@@ -136,6 +141,7 @@ class GPUProfiler : public Profiler,
         : profiler(profiler) {}
     virtual ~GPUProfilerPimplInterface() = default;
 
+    virtual void setLibPath(const std::string &libPath) = 0;
     virtual void doStart() = 0;
     virtual void doFlush() = 0;
     virtual void doStop() = 0;
diff --git a/third_party/proton/csrc/include/Session/Session.h b/third_party/proton/csrc/include/Session/Session.h
@@ -72,6 +72,7 @@ class SessionManager : public Singleton<SessionManager> {
   ~SessionManager() = default;
 
   size_t addSession(const std::string &path, const std::string &profilerName,
+                    const std::string &profilerPath,
                     const std::string &contextSourceName,
                     const std::string &dataName);
 
@@ -103,6 +104,7 @@ class SessionManager : public Singleton<SessionManager> {
 private:
   std::unique_ptr<Session> makeSession(size_t id, const std::string &path,
                                        const std::string &profilerName,
+                                       const std::string &profilerPath,
                                        const std::string &contextSourceName,
                                        const std::string &dataName);
 
diff --git a/third_party/proton/csrc/lib/Driver/GPU/CuptiApi.cpp b/third_party/proton/csrc/lib/Driver/GPU/CuptiApi.cpp
@@ -6,16 +6,10 @@ namespace proton {
 
 namespace cupti {
 
-#define STRINGIFY(x) #x
-#define TOSTRING(x) STRINGIFY(x)
 struct ExternLibCupti : public ExternLibBase {
   using RetType = CUptiResult;
   static constexpr const char *name = "libcupti.so";
-#ifdef CUPTI_LIB_DIR
-  static constexpr const char *defaultDir = TOSTRING(CUPTI_LIB_DIR);
-#else
-  static constexpr const char *defaultDir = "";
-#endif
+  static inline std::string defaultDir = "";
   static constexpr RetType success = CUPTI_SUCCESS;
   static void *lib;
 };
@@ -116,6 +110,8 @@ DEFINE_DISPATCH(ExternLibCupti, pcSamplingStart, cuptiPCSamplingStart,
 DEFINE_DISPATCH(ExternLibCupti, pcSamplingStop, cuptiPCSamplingStop,
                 CUpti_PCSamplingStopParams *);
 
+void setLibPath(const std::string &path) { ExternLibCupti::defaultDir = path; }
+
 } // namespace cupti
 
 } // namespace proton
diff --git a/third_party/proton/csrc/lib/Profiler/Cupti/CuptiProfiler.cpp b/third_party/proton/csrc/lib/Profiler/Cupti/CuptiProfiler.cpp
@@ -198,6 +198,9 @@ struct CuptiProfiler::CuptiProfilerPimpl
       : GPUProfiler<CuptiProfiler>::GPUProfilerPimplInterface(profiler) {}
   virtual ~CuptiProfilerPimpl() = default;
 
+  void setLibPath(const std::string &libPath) override {
+    cupti::setLibPath(libPath);
+  }
   void doStart() override;
   void doFlush() override;
   void doStop() override;
diff --git a/third_party/proton/csrc/lib/Profiler/RocTracer/RoctracerProfiler.cpp b/third_party/proton/csrc/lib/Profiler/RocTracer/RoctracerProfiler.cpp
@@ -189,6 +189,7 @@ struct RoctracerProfiler::RoctracerProfilerPimpl
       : GPUProfiler<RoctracerProfiler>::GPUProfilerPimplInterface(profiler) {}
   virtual ~RoctracerProfilerPimpl() = default;
 
+  void setLibPath(const std::string &libPath) override {}
   void doStart() override;
   void doFlush() override;
   void doStop() override;
diff --git a/third_party/proton/csrc/lib/Session/Session.cpp b/third_party/proton/csrc/lib/Session/Session.cpp
@@ -9,17 +9,17 @@
 namespace proton {
 
 namespace {
-Profiler *getProfiler(const std::string &profilerName) {
-  if (proton::toLower(profilerName) == "cupti") {
-    return &CuptiProfiler::instance();
+Profiler *getProfiler(const std::string &name, const std::string &path) {
+  if (proton::toLower(name) == "cupti") {
+    return &CuptiProfiler::instance().setLibPath(path);
   }
-  if (proton::toLower(profilerName) == "cupti_pcsampling") {
-    return &CuptiProfiler::instance().enablePCSampling();
+  if (proton::toLower(name) == "cupti_pcsampling") {
+    return &CuptiProfiler::instance().setLibPath(path).enablePCSampling();
   }
-  if (proton::toLower(profilerName) == "roctracer") {
+  if (proton::toLower(name) == "roctracer") {
     return &RoctracerProfiler::instance();
   }
-  throw std::runtime_error("Unknown profiler: " + profilerName);
+  throw std::runtime_error("Unknown profiler: " + name);
 }
 
 std::unique_ptr<Data> makeData(const std::string &dataName,
@@ -71,8 +71,9 @@ 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 &contextSourceName, const std::string &dataName) {
-  auto profiler = getProfiler(profilerName);
+    const std::string &profilerPath, const std::string &contextSourceName,
+    const std::string &dataName) {
+  auto profiler = getProfiler(profilerName, profilerPath);
   auto contextSource = makeContextSource(contextSourceName);
   auto data = makeData(dataName, path, contextSource.get());
   auto *session = new Session(id, path, profiler, std::move(contextSource),
@@ -139,6 +140,7 @@ void SessionManager::removeSession(size_t sessionId) {
 
 size_t SessionManager::addSession(const std::string &path,
                                   const std::string &profilerName,
+                                  const std::string &profilerPath,
                                   const std::string &contextSourceName,
                                   const std::string &dataName) {
   std::lock_guard<std::mutex> lock(mutex);
@@ -149,8 +151,8 @@ size_t SessionManager::addSession(const std::string &path,
   }
   auto sessionId = nextSessionId++;
   sessionPaths[path] = sessionId;
-  sessions[sessionId] =
-      makeSession(sessionId, path, profilerName, contextSourceName, dataName);
+  sessions[sessionId] = makeSession(sessionId, path, profilerName, profilerPath,
+                                    contextSourceName, dataName);
   return sessionId;
 }
 
diff --git a/third_party/proton/proton/profile.py b/third_party/proton/proton/profile.py
@@ -1,6 +1,7 @@
 import functools
 import triton
 import os
+import pathlib
 
 from triton._C.libproton import proton as libproton
 from .hook import register_triton_hook, unregister_triton_hook
@@ -20,6 +21,18 @@ def _select_backend() -> str:
         raise ValueError("No backend is available for the current target.")
 
 
+def _get_backend_default_path(backend: str) -> str:
+    lib_path = ""
+    if backend == "cupti":
+        # First try to get the path from the environment variable that overrides the default path
+        lib_path = os.getenv("TRITON_CUPTI_LIB_PATH", None)
+        if lib_path is None:
+            # Get the default path for the cupti backend,
+            # which is the most compatible with the current CUPTI header file triton is compiled with
+            lib_path = str(pathlib.Path(__file__).parent.parent.absolute() / "backends" / "nvidia" / "lib" / "cupti")
+    return lib_path
+
+
 def _check_env(backend: str) -> None:
     if backend == "roctracer":
         hip_device_envs = ["HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"]
@@ -79,10 +92,12 @@ def start(
 
     _check_env(backend)
 
+    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)
+    return libproton.start(name, context, data, backend, backend_path)
 
 
 def activate(session: Optional[int] = None) -> None:
diff --git a/third_party/proton/test/test_lib.py b/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})
