Skip to content

Commit 4336650

Browse files
Jay Guhaijieg
Jay Gu
authored and
haijieg
committed
[NFC] Share CUDA loading and context management
Move CUDA context switching, CUlibrary ownership, and CUkernel lookup out of tile_kernel.cpp and into cuda_loader. Signed-off-by: Jay Gu <jagu@.nvidia.com>
1 parent 52da7cf commit 4336650

3 files changed

Lines changed: 142 additions & 113 deletions

File tree

‎cext/cuda_loader.cpp‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,3 +92,88 @@ Result<const DriverApi*> get_driver_api() {
9292
}
9393
return &instance;
9494
}
95+
96+
97+
CudaLibrary::CudaLibrary(const DriverApi* driver, CUlibrary lib)
98+
: driver_(driver), lib_(lib) {}
99+
100+
101+
CudaLibrary::CudaLibrary(CudaLibrary&& other)
102+
: driver_(other.driver_), lib_(other.lib_) {
103+
other.lib_ = nullptr;
104+
}
105+
106+
107+
CudaLibrary::~CudaLibrary() {
108+
if (lib_) {
109+
CUresult res = driver_->cuLibraryUnload(lib_);
110+
CHECK(res == CUDA_SUCCESS);
111+
}
112+
}
113+
114+
115+
const CUlibrary& CudaLibrary::get() const {
116+
return lib_;
117+
}
118+
119+
120+
static Result<CudaLibrary> load_cuda_library(const DriverApi* driver, const void* code) {
121+
CUlibrary lib;
122+
CUresult res = driver->cuLibraryLoadData(&lib, code, nullptr, nullptr, 0,
123+
nullptr, nullptr, 0);
124+
if (res == CUDA_SUCCESS)
125+
return CudaLibrary(driver, lib);
126+
127+
return raise(PyExc_RuntimeError, "Failed to load CUDA library: %s",
128+
get_cuda_error(driver, res));
129+
}
130+
131+
132+
Result<CudaKernel> load_cuda_kernel(
133+
const DriverApi* driver,
134+
const char* cubin_data,
135+
size_t cubin_size,
136+
const char* func_name) {
137+
(void) cubin_size;
138+
139+
Result<CudaLibrary> lib = load_cuda_library(driver, cubin_data);
140+
if (!lib.is_ok()) return ErrorRaised;
141+
142+
CUkernel kernel;
143+
CUresult res = driver->cuLibraryGetKernel(&kernel, lib->get(), func_name);
144+
if (res == CUDA_SUCCESS)
145+
return CudaKernel{std::move(*lib), kernel};
146+
147+
return raise(PyExc_RuntimeError, "Failed to get kernel %s from library: %s",
148+
func_name, get_cuda_error(driver, res));
149+
}
150+
151+
152+
Status CudaContextGuard::switch_to(CUcontext target) {
153+
if (!target) return OK;
154+
155+
CUcontext current;
156+
CUresult res = driver->cuCtxGetCurrent(&current);
157+
if (res != CUDA_SUCCESS) {
158+
return raise(PyExc_RuntimeError, "Failed to get current CUDA context: %s",
159+
get_cuda_error(driver, res));
160+
}
161+
if (current == target) return OK;
162+
163+
res = driver->cuCtxPushCurrent(target);
164+
if (res != CUDA_SUCCESS) {
165+
return raise(PyExc_RuntimeError, "Failed to switch CUDA context: %s",
166+
get_cuda_error(driver, res));
167+
}
168+
need_to_pop = true;
169+
return OK;
170+
}
171+
172+
173+
CudaContextGuard::~CudaContextGuard() {
174+
if (need_to_pop) {
175+
CUcontext old;
176+
CUresult res = driver->cuCtxPopCurrent(&old);
177+
CHECK(res == CUDA_SUCCESS);
178+
}
179+
}

‎cext/cuda_loader.h‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,51 @@ struct DriverApi {
7777
Result<const DriverApi*> get_driver_api();
7878

7979

80+
class CudaContextGuard {
81+
const DriverApi* driver;
82+
bool need_to_pop = false;
83+
public:
84+
CudaContextGuard(const CudaContextGuard&) = delete;
85+
void operator=(const CudaContextGuard&) = delete;
86+
87+
explicit CudaContextGuard(const DriverApi* driver) : driver(driver) {}
88+
89+
Status switch_to(CUcontext target);
90+
91+
~CudaContextGuard();
92+
};
93+
94+
95+
// RAII wrapper around a CUDA library loaded into a context.
96+
class CudaLibrary {
97+
public:
98+
explicit CudaLibrary(const DriverApi* driver, CUlibrary lib);
99+
CudaLibrary(CudaLibrary&& other);
100+
CudaLibrary(const CudaLibrary&) = delete;
101+
void operator=(const CudaLibrary&) = delete;
102+
~CudaLibrary();
103+
104+
const CUlibrary& get() const;
105+
106+
private:
107+
const DriverApi* driver_;
108+
CUlibrary lib_;
109+
};
110+
111+
112+
struct CudaKernel {
113+
CudaLibrary lib;
114+
CUkernel kernel;
115+
};
116+
117+
118+
Result<CudaKernel> load_cuda_kernel(
119+
const DriverApi* driver,
120+
const char* cubin_data,
121+
size_t cubin_size,
122+
const char* func_name);
123+
124+
80125
class CudaGraph {
81126
const DriverApi* d;
82127
CUgraph graph;

‎cext/tile_kernel.cpp‎

Lines changed: 12 additions & 113 deletions
Original file line numberDiff line numberDiff line change
@@ -236,50 +236,6 @@ PyTypeObject CallingConvention::pytype = {
236236
};
237237

238238

239-
// RAII wrapper around CUlibrary
240-
namespace { class CudaLibrary {
241-
public:
242-
explicit CudaLibrary(const DriverApi* driver, CUlibrary lib) : driver_(driver), lib_(lib) {}
243-
244-
CudaLibrary(CudaLibrary&& other) : driver_(other.driver_), lib_(other.lib_) {
245-
other.lib_ = nullptr;
246-
}
247-
248-
CudaLibrary(const CudaLibrary&) = delete;
249-
void operator=(const CudaLibrary&) = delete;
250-
251-
~CudaLibrary() {
252-
if (lib_) {
253-
CUresult res = driver_->cuLibraryUnload(lib_);
254-
CHECK(res == CUDA_SUCCESS);
255-
}
256-
}
257-
258-
const CUlibrary& get() const {
259-
return lib_;
260-
}
261-
262-
private:
263-
const DriverApi* driver_;
264-
CUlibrary lib_;
265-
}; }
266-
267-
static Result<CudaLibrary> load_cuda_library(const DriverApi* driver, const void* code) {
268-
CUlibrary lib;
269-
CUresult res = driver->cuLibraryLoadData(&lib, code, nullptr, nullptr, 0,
270-
nullptr, nullptr, 0);
271-
if (res == CUDA_SUCCESS)
272-
return CudaLibrary(driver, lib);
273-
274-
return raise(PyExc_RuntimeError, "Failed to load CUDA library: %s",
275-
get_cuda_error(driver, res));
276-
}
277-
278-
struct CudaKernel {
279-
CudaLibrary lib;
280-
CUkernel kernel;
281-
};
282-
283239
static Status enable_maximum_dynamic_shared_memory(const DriverApi *driver,
284240
const CUkernel kernel,
285241
const char *func_name) {
@@ -340,26 +296,6 @@ static Status enable_maximum_dynamic_shared_memory(const DriverApi *driver,
340296
return OK;
341297
}
342298

343-
static Result<CudaKernel> load_cuda_kernel(const DriverApi* driver,
344-
const char* cubin_data,
345-
size_t cubin_size,
346-
const char* func_name) {
347-
(void) cubin_size;
348-
349-
Result<CudaLibrary> lib = load_cuda_library(driver, cubin_data);
350-
if (!lib.is_ok()) return ErrorRaised;
351-
352-
CUkernel kernel;
353-
CUresult res = driver->cuLibraryGetKernel(&kernel, lib->get(), func_name);
354-
if (res != CUDA_SUCCESS) {
355-
return raise(PyExc_RuntimeError, "Failed to get kernel %s from library: %s",
356-
func_name, get_cuda_error(driver, res));
357-
}
358-
359-
return CudaKernel{std::move(*lib), kernel};
360-
}
361-
362-
363299
// X(Name, #Attrs, MinStack, StackEffect)
364300
#define FOREACH_SIZE_OPCODE(X) \
365301
X(Const, 1, 0, 1) \
@@ -2308,6 +2244,7 @@ static Result<CUstream> parse_stream(PyObject* py_stream) {
23082244
return stream;
23092245
}
23102246

2247+
23112248
using StreamBufferPoolMap = HashMap<unsigned long long, StreamBufferPool*>;
23122249

23132250
// Protected by GIL or g_launch_mutex.
@@ -2340,47 +2277,6 @@ static Result<StreamBufferPool*> get_stream_buffer_pool(const DriverApi* driver,
23402277
}
23412278
}
23422279

2343-
namespace { struct ContextGuard {
2344-
bool need_to_pop;
2345-
const DriverApi* driver_;
2346-
2347-
ContextGuard(const DriverApi* driver) : need_to_pop(false), driver_(driver) {}
2348-
2349-
ContextGuard() = delete;
2350-
ContextGuard(const ContextGuard&) = delete;
2351-
void operator=(const ContextGuard&) = delete;
2352-
2353-
~ContextGuard() {
2354-
if (need_to_pop) {
2355-
CUcontext old;
2356-
CUresult res = driver_->cuCtxPopCurrent(&old);
2357-
CHECK(res == CUDA_SUCCESS);
2358-
}
2359-
}
2360-
}; }
2361-
2362-
static Status maybe_switch_context(const DriverApi* driver, CUcontext target, ContextGuard& guard) {
2363-
if (!target) return OK;
2364-
2365-
CUcontext current;
2366-
CUresult res = driver->cuCtxGetCurrent(&current);
2367-
if (res != CUDA_SUCCESS) {
2368-
return raise(PyExc_RuntimeError, "Failed to get current CUDA context: %s",
2369-
get_cuda_error(driver, res));
2370-
}
2371-
2372-
if (current == target) return OK;
2373-
2374-
res = driver->cuCtxPushCurrent(target);
2375-
if (res != CUDA_SUCCESS) {
2376-
return raise(PyExc_RuntimeError, "Failed to switch CUDA context: %s",
2377-
get_cuda_error(driver, res));
2378-
}
2379-
2380-
guard.need_to_pop = true;
2381-
return OK;
2382-
}
2383-
23842280
struct Grid {
23852281
enum { Len = 3 };
23862282
unsigned dims[Len];
@@ -2541,7 +2437,7 @@ static Result<PreparedLaunch> prepare_launch(
25412437
bool capture_kernel_image,
25422438
bool stage_list_args,
25432439
StreamBufferTransaction& tx,
2544-
ContextGuard& ctx_guard) {
2440+
CudaContextGuard& ctx_guard) {
25452441

25462442
LaunchHelperPtr helper = launch_helper_get();
25472443

@@ -2623,7 +2519,7 @@ static Result<PreparedLaunch> prepare_launch(
26232519
kernel_item = kernel_map.insert(std::move(helper->constants), std::move(*res));
26242520
}
26252521

2626-
if (!maybe_switch_context(driver, helper->cuda_context, ctx_guard))
2522+
if (!ctx_guard.switch_to(helper->cuda_context))
26272523
return ErrorRaised;
26282524

26292525
if (stage_list_args
@@ -2661,7 +2557,7 @@ static Status launch(const DriverApi* driver,
26612557
PyObject* const* pyargs,
26622558
Py_ssize_t num_pyargs
26632559
) {
2664-
ContextGuard ctx_guard(driver);
2560+
CudaContextGuard ctx_guard(driver);
26652561
StreamBufferTransaction tx;
26662562
Result<PreparedLaunch> prep = prepare_launch(
26672563
driver, dispatcher_pyobj, launch_stream, pyargs, num_pyargs,
@@ -2974,6 +2870,9 @@ parse_parameter_annotation_nodes_seq(PyObject* nodes_seq) {
29742870
return result;
29752871
}
29762872

2873+
2874+
2875+
29772876
static int TileContext_init(PyObject* self, PyObject* args, PyObject* kwargs) {
29782877
const char* keywords[] = {"config", nullptr};
29792878
PyObject* config = nullptr;
@@ -3330,7 +3229,7 @@ static PyObject* cuda_tile_benchmark(PyObject* mod, PyObject* const* args, Py_ss
33303229
Result<const DriverApi*> driver = get_driver_api();
33313230
if (!driver.is_ok()) return nullptr;
33323231

3333-
ContextGuard ctx_guard(*driver);
3232+
CudaContextGuard ctx_guard(*driver);
33343233
StreamBufferTransaction tx;
33353234
Result<PreparedLaunch> prep = prepare_launch(
33363235
*driver, launch_args.dispatcher, launch_args.stream,
@@ -3363,7 +3262,7 @@ static PyObject* cuda_tile_export_ipc_benchmark_payload(PyObject*, PyObject* con
33633262
Result<const DriverApi*> driver = get_driver_api();
33643263
if (!driver.is_ok()) return nullptr;
33653264

3366-
ContextGuard ctx_guard(*driver);
3265+
CudaContextGuard ctx_guard(*driver);
33673266
StreamBufferTransaction tx;
33683267
Result<PreparedLaunch> prep = prepare_launch(
33693268
*driver, launch_args.dispatcher, launch_args.stream,
@@ -3390,6 +3289,7 @@ static PyObject* cuda_tile_export_ipc_benchmark_payload(PyObject*, PyObject* con
33903289
const char* symbol = PyUnicode_AsUTF8AndSize(kernel_image.symbol.get(), &symbol_size);
33913290
if (!symbol) return nullptr;
33923291

3292+
33933293
CUdevice device;
33943294
CUresult res = (*driver)->cuCtxGetDevice(&device);
33953295
if (res != CUDA_SUCCESS) {
@@ -3512,8 +3412,8 @@ static PyObject* cuda_tile_benchmark_with_ipc_payload(PyObject*, PyObject* const
35123412
return nullptr;
35133413
}
35143414

3515-
ContextGuard ctx_guard(*driver);
3516-
if (!maybe_switch_context(*driver, ctx, ctx_guard))
3415+
CudaContextGuard ctx_guard(*driver);
3416+
if (!ctx_guard.switch_to(ctx))
35173417
return nullptr;
35183418

35193419
IpcHandleCreator ipc_mem_handle(*driver);
@@ -3797,4 +3697,3 @@ Status tile_kernel_init(PyObject* m) {
37973697

37983698
return OK;
37993699
}
3800-

0 commit comments

Comments
 (0)