Skip to content

Commit dbc85fc

Browse files
JokerenJokeren
authored
[PROTON] Simplify backend lib settings (#8246)
1 parent 83683fc commit dbc85fc

File tree

16 files changed

+93
-98
lines changed

16 files changed

+93
-98
lines changed

‎python/triton/knobs.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import re
77
import subprocess
88
import sysconfig
9+
import pathlib
910

1011
from dataclasses import dataclass
1112
from contextlib import contextmanager
@@ -513,7 +514,9 @@ class amd_knobs(base_knobs):
513514

514515

515516
class proton_knobs(base_knobs):
516-
cupti_lib_dir: env_opt_str = env_opt_str("TRITON_CUPTI_LIB_PATH")
517+
cupti_lib_dir: env_str = env_str(
518+
"TRITON_CUPTI_LIB_PATH",
519+
str(pathlib.Path(__file__).parent.absolute() / "backends" / "nvidia" / "lib" / "cupti"))
517520
enable_nvtx: env_bool = env_bool("TRITON_ENABLE_NVTX", True)
518521

519522

‎third_party/proton/csrc/Proton.cpp‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,16 +17,15 @@ static void initProton(pybind11::module &&m) {
1717
"start",
1818
[](const std::string &path, const std::string &contextSourceName,
1919
const std::string &dataName, const std::string &profilerName,
20-
const std::string &profilerPath, const std::string &mode) {
20+
const std::string &mode) {
2121
auto sessionId = SessionManager::instance().addSession(
22-
path, profilerName, profilerPath, contextSourceName, dataName,
23-
mode);
22+
path, profilerName, contextSourceName, dataName, mode);
2423
SessionManager::instance().activateSession(sessionId);
2524
return sessionId;
2625
},
2726
pybind11::arg("path"), pybind11::arg("contextSourceName"),
2827
pybind11::arg("dataName"), pybind11::arg("profilerName"),
29-
pybind11::arg("profilerPath") = "", pybind11::arg("mode") = "");
28+
pybind11::arg("mode") = "");
3029

3130
m.def("activate", [](size_t sessionId) {
3231
SessionManager::instance().activateSession(sessionId);

‎third_party/proton/csrc/include/Driver/Dispatch.h‎

Lines changed: 27 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
#include <dlfcn.h>
55

6+
#include "Utility/Env.h"
67
#include <stdexcept>
78
#include <string>
89

@@ -44,13 +45,14 @@ namespace proton {
4445

4546
struct ExternLibBase {
4647
using RetType = int; // Generic type, can be overridden in derived structs
47-
static constexpr const char *name = ""; // Placeholder
48-
static constexpr RetType success = 0; // Placeholder
48+
static constexpr const char *name = ""; // Placeholder
49+
static constexpr const char *symbolName{}; // Placeholder
50+
static constexpr const char *pathEnv{}; // Placeholder
51+
static constexpr RetType success = 0; // Placeholder
4952
ExternLibBase() = delete;
5053
ExternLibBase(const ExternLibBase &) = delete;
5154
ExternLibBase &operator=(const ExternLibBase &) = delete;
5255
static inline void *lib{nullptr};
53-
static inline std::string defaultDir{""};
5456
};
5557

5658
template <typename ExternLib> class Dispatch {
@@ -60,8 +62,9 @@ template <typename ExternLib> class Dispatch {
6062
static void init(const char *name, void **lib) {
6163
if (*lib == nullptr) {
6264
// If not found, try to load it from the default path
63-
auto dir = std::string(ExternLib::defaultDir);
64-
if (dir.length() > 0) {
65+
auto dir =
66+
ExternLib::pathEnv == nullptr ? "" : getStrEnv(ExternLib::pathEnv);
67+
if (!dir.empty()) {
6568
auto fullPath = dir + "/" + name;
6669
*lib = dlopen(fullPath.c_str(), RTLD_LOCAL | RTLD_LAZY);
6770
} else {
@@ -105,6 +108,25 @@ template <typename ExternLib> class Dispatch {
105108
}
106109
return ret;
107110
}
111+
112+
static std::string getLibPath() {
113+
if (ExternLib::lib == nullptr) {
114+
// Force initialization
115+
Dispatch<ExternLib>::init(ExternLib::name, &ExternLib::lib);
116+
if (ExternLib::lib == nullptr) {
117+
return "";
118+
}
119+
}
120+
if (ExternLib::lib != nullptr) {
121+
void *sym = dlsym(ExternLib::lib,
122+
ExternLib::symbolName); // pick any known symbol
123+
Dl_info info;
124+
if (dladdr(sym, &info)) {
125+
return info.dli_fname;
126+
}
127+
}
128+
return "";
129+
}
108130
};
109131

110132
} // namespace proton

‎third_party/proton/csrc/include/Driver/GPU/CuptiApi.h‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#ifndef PROTON_DRIVER_GPU_CUPTI_API_H_
22
#define PROTON_DRIVER_GPU_CUPTI_API_H_
33

4+
#include "Driver/Dispatch.h"
45
#include "cupti.h"
56
#include "cupti_pcsampling.h"
67
#include <string>
@@ -9,6 +10,15 @@ namespace proton {
910

1011
namespace cupti {
1112

13+
struct ExternLibCupti : public ExternLibBase {
14+
using RetType = CUptiResult;
15+
static constexpr const char *name = "libcupti.so";
16+
static constexpr const char *symbolName = "cuptiUnsubscribe";
17+
static constexpr const char *pathEnv = "TRITON_CUPTI_LIB_PATH";
18+
static constexpr RetType success = CUPTI_SUCCESS;
19+
static inline void *lib = nullptr;
20+
};
21+
1222
template <bool CheckSuccess> CUptiResult getVersion(uint32_t *version);
1323

1424
template <bool CheckSuccess>
@@ -107,10 +117,6 @@ CUptiResult pcSamplingStart(CUpti_PCSamplingStartParams *pParams);
107117
template <bool CheckSuccess>
108118
CUptiResult pcSamplingStop(CUpti_PCSamplingStopParams *pParams);
109119

110-
void setLibPath(const std::string &path);
111-
112-
const std::string getLibPath();
113-
114120
} // namespace cupti
115121

116122
} // namespace proton

‎third_party/proton/csrc/include/Driver/GPU/RoctracerApi.h‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,22 @@
11
#ifndef PROTON_DRIVER_GPU_ROCTRACER_API_H_
22
#define PROTON_DRIVER_GPU_ROCTRACER_API_H_
33

4+
#include "Driver/Dispatch.h"
45
#include "roctracer/roctracer.h"
56

67
namespace proton {
78

89
namespace roctracer {
910

11+
struct ExternLibRoctracer : public ExternLibBase {
12+
using RetType = roctracer_status_t;
13+
static constexpr const char *name = "libroctracer64.so";
14+
static constexpr const char *symbolName = "roctracer_start";
15+
static constexpr const char *pathEnv{};
16+
static constexpr RetType success = ROCTRACER_STATUS_SUCCESS;
17+
static inline void *lib = nullptr;
18+
};
19+
1020
template <bool CheckSuccess>
1121
roctracer_status_t setProperties(roctracer_domain_t domain, void *properties);
1222

‎third_party/proton/csrc/include/Profiler/GPUProfiler.h‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -32,11 +32,6 @@ class GPUProfiler : public Profiler,
3232
std::unordered_map<uint64_t, std::pair<size_t, size_t>>>;
3333
using ApiExternIdSet = ThreadSafeSet<size_t, std::unordered_set<size_t>>;
3434

35-
ConcreteProfilerT &setLibPath(const std::string &libPath) {
36-
pImpl->setLibPath(libPath);
37-
return dynamic_cast<ConcreteProfilerT &>(*this);
38-
}
39-
4035
protected:
4136
// OpInterface
4237
void startOp(const Scope &scope) override {
@@ -148,7 +143,6 @@ class GPUProfiler : public Profiler,
148143
: profiler(profiler) {}
149144
virtual ~GPUProfilerPimplInterface() = default;
150145

151-
virtual void setLibPath(const std::string &libPath) = 0;
152146
virtual void doStart() = 0;
153147
virtual void doFlush() = 0;
154148
virtual void doStop() = 0;

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,6 @@ class SessionManager : public Singleton<SessionManager> {
7676
~SessionManager() = default;
7777

7878
size_t addSession(const std::string &path, const std::string &profilerName,
79-
const std::string &profilerPath,
8079
const std::string &contextSourceName,
8180
const std::string &dataName, const std::string &mode);
8281

@@ -125,7 +124,6 @@ class SessionManager : public Singleton<SessionManager> {
125124

126125
std::unique_ptr<Session> makeSession(size_t id, const std::string &path,
127126
const std::string &profilerName,
128-
const std::string &profilerPath,
129127
const std::string &contextSourceName,
130128
const std::string &dataName,
131129
const std::string &mode);

‎third_party/proton/csrc/include/Utility/Env.h‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,13 @@
1+
#ifndef PROTON_UTILITY_ENV_H_
2+
#define PROTON_UTILITY_ENV_H_
3+
14
#include <algorithm>
25
#include <cstdlib>
36
#include <mutex>
47
#include <string>
58

9+
namespace proton {
10+
611
static std::mutex getenv_mutex;
712

813
inline bool getBoolEnv(const std::string &env, bool defaultValue) {
@@ -15,3 +20,13 @@ inline bool getBoolEnv(const std::string &env, bool defaultValue) {
1520
[](unsigned char c) { return std::tolower(c); });
1621
return str == "on" || str == "true" || str == "1";
1722
}
23+
24+
inline std::string getStrEnv(const std::string &env) {
25+
std::lock_guard<std::mutex> lock(getenv_mutex);
26+
const char *s = std::getenv(env.c_str());
27+
return std::string(s ? s : "");
28+
}
29+
30+
} // namespace proton
31+
32+
#endif // PROTON_UTILITY_ENV_H_

‎third_party/proton/csrc/lib/Driver/GPU/CuptiApi.cpp‎

Lines changed: 0 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -6,16 +6,6 @@ namespace proton {
66

77
namespace cupti {
88

9-
struct ExternLibCupti : public ExternLibBase {
10-
using RetType = CUptiResult;
11-
static constexpr const char *name = "libcupti.so";
12-
static inline std::string defaultDir = "";
13-
static constexpr RetType success = CUPTI_SUCCESS;
14-
static void *lib;
15-
};
16-
17-
void *ExternLibCupti::lib = nullptr;
18-
199
DEFINE_DISPATCH(ExternLibCupti, getVersion, cuptiGetVersion, uint32_t *);
2010

2111
DEFINE_DISPATCH(ExternLibCupti, getContextId, cuptiGetContextId, CUcontext,
@@ -110,26 +100,6 @@ DEFINE_DISPATCH(ExternLibCupti, pcSamplingStart, cuptiPCSamplingStart,
110100
DEFINE_DISPATCH(ExternLibCupti, pcSamplingStop, cuptiPCSamplingStop,
111101
CUpti_PCSamplingStopParams *);
112102

113-
void setLibPath(const std::string &path) { ExternLibCupti::defaultDir = path; }
114-
115-
// TODO(Keren): generalize to AMD
116-
const std::string getLibPath() {
117-
if (ExternLibCupti::lib == nullptr) {
118-
// Force initialization
119-
Dispatch<ExternLibCupti>::init(ExternLibCupti::name, &ExternLibCupti::lib);
120-
if (ExternLibCupti::lib == nullptr) {
121-
return "";
122-
}
123-
}
124-
void *sym =
125-
dlsym(ExternLibCupti::lib, "cuptiUnsubscribe"); // pick any known symbol
126-
Dl_info info;
127-
if (dladdr(sym, &info)) {
128-
return info.dli_fname;
129-
}
130-
return "";
131-
}
132-
133103
} // namespace cupti
134104

135105
} // namespace proton

‎third_party/proton/csrc/lib/Driver/GPU/NvtxApi.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,8 @@ namespace nvtx {
1919

2020
void enable() {
2121
// Get cupti lib path and append it to NVTX_INJECTION64_PATH
22-
const std::string cuptiLibPath = cupti::getLibPath();
22+
const std::string cuptiLibPath =
23+
Dispatch<cupti::ExternLibCupti>::getLibPath();
2324
if (!cuptiLibPath.empty()) {
2425
setenv("NVTX_INJECTION64_PATH", cuptiLibPath.c_str(), 1);
2526
}

0 commit comments

Comments
 (0)