github/ROCm
1
1
Fork 0
mirror of https://github.com/ROCm/ROCm.git synced 2026-04-05 03:01:17 -04:00
Code Issues Packages Projects Releases Wiki Activity

[BACKEND] make sure we always bind to primary context in loadBinary (#1912)

Browse Source
This commit is contained in:
Natalia Gimelshein
2023-07-07 14:28:03 -07:00
committed by GitHub
parent 38d767ea93
commit 778ed64a66
2 changed files with 11 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/triton/runtime/backends/cuda.c
View File

@@ -69,6 +69,13 @@ static PyObject *loadBinary(PyObject *self, PyObject *args) {
int32_t n_regs = 0;
int32_t n_spills = 0;
// create driver handles
CUcontext pctx = 0;
CUDA_CHECK(cuCtxGetCurrent(&pctx));
if (!pctx) {
CUDA_CHECK(cuDevicePrimaryCtxRetain(&pctx, device));
CUDA_CHECK(cuCtxSetCurrent(pctx));
}
CUDA_CHECK(cuModuleLoadData(&mod, data));
CUDA_CHECK(cuModuleGetFunction(&fun, mod, name));
// get allocated registers and spilled registers from the function