Don't switch CUDA context prematurely

In c2c53d3, the call to
cuCtxPushCurrent() was moved earlier in the launch sequence, so that we
could call cuContextGetDevice() before we get the kernel from the cache.

The problem is that we run an arbitrary amount of Python code in
compile() in case we have a cache miss, which may switch the context to
something else. So we should be performing our switch after the
compilation is done.

This seems to have caused a CI flake with a random failure in
test_list.py.

Signed-off-by: Greg Bonik <gbonik@nvidia.com>

Author: gbonik

cext/cuda_loader.cpp

@@ -44,14 +44,14 @@ F get_proc_address(cuGetProcAddress_v2_t getter,
 } // anonymous namespace
 
 
-#define DEFINE_CUDA_FUNCTION_GLOBAL(name, _cuda_version) \
+#define DEFINE_CUDA_FUNCTION_GLOBAL(name, _key, _cuda_version) \
     decltype(name)* g_##name;
 
 FOREACH_CUDA_FUNCTION_TO_LOAD(DEFINE_CUDA_FUNCTION_GLOBAL)
 
-#define GET_PROC_ADDRESS(name, cuda_ver) \
+#define GET_PROC_ADDRESS(name, key, cuda_ver) \
         if (!(driver_api.name = \
-                    get_proc_address<decltype(name)*>(_cuGetProcAddress, #name, cuda_ver))) \
+                    get_proc_address<decltype(name)*>(_cuGetProcAddress, key, cuda_ver))) \
             return ErrorRaised;
 
 

cext/cuda_loader.h

@@ -10,64 +10,65 @@
 #include <cuda.h>
 
 #define FOREACH_CUDA_FUNCTION_TO_LOAD(X) \
-    X(cuInit, 2000) \
-    X(cuLibraryLoadData, 12000) \
-    X(cuLibraryUnload, 12000) \
-    X(cuLibraryGetKernel, 12000) \
-    X(cuGetErrorString, 6000) \
-    X(cuLaunchKernel, 7000) \
-    X(cuLaunchKernelEx, 11060) \
-    X(cuPointerGetAttribute, 4000) \
-    X(cuCtxSynchronize, 2000) \
-    X(cuCtxPushCurrent, 4000) \
-    X(cuCtxPopCurrent, 4000) \
-    X(cuCtxGetCurrent, 4000) \
-    X(cuCtxSetCurrent, 4000) \
-    X(cuCtxGetDevice, 2000) \
-    X(cuCtxGetId, 12000) \
-    X(cuDeviceGet, 2000) \
-    X(cuDeviceGetCount, 2000) \
-    X(cuDeviceGetAttribute, 2000) \
-    X(cuDevicePrimaryCtxRetain, 7000) \
-    X(cuDriverGetVersion, 2020) \
-    X(cuEventCreate, 2000) \
-    X(cuEventDestroy, 2000) \
-    X(cuEventQuery, 2000) \
-    X(cuEventRecord, 2000) \
-    X(cuKernelGetFunction, 12000) \
-    X(cuKernelGetAttribute, 12000) \
-    X(cuKernelSetAttribute, 12000) \
-    X(cuMemAlloc, 3020) \
-    X(cuMemAllocHost, 3020) \
-    X(cuMemFree, 3020) \
-    X(cuMemFreeHost, 2000) \
-    X(cuMemGetAddressRange, 3020) \
-    X(cuIpcGetMemHandle, 4010) \
-    X(cuIpcOpenMemHandle, 4010) \
-    X(cuIpcCloseMemHandle, 4010) \
-    X(cuMemcpyHtoDAsync, 3020) \
-    X(cuStreamCreate, 2000) \
-    X(cuStreamDestroy, 4000) \
-    X(cuStreamGetCtx, 9020) \
-    X(cuStreamGetId, 12000) \
-    X(cuStreamIsCapturing, 10000) \
-    X(cuStreamSynchronize, 7000) \
-    X(cuStreamWaitEvent, 7000) \
-    X(cuEventElapsedTime, 12080) \
-    X(cuGraphCreate, 10000) \
-    X(cuGraphDestroy, 10000) \
-    X(cuGraphAddEventRecordNode, 11010) \
-    X(cuGraphAddKernelNode, 12000) \
-    X(cuGraphAddMemsetNode, 10000) \
-    X(cuGraphAddMemAllocNode, 11040) \
-    X(cuGraphAddMemFreeNode, 11040) \
-    X(cuGraphInstantiateWithFlags, 11040) \
-    X(cuGraphExecDestroy, 10000) \
-    X(cuGraphLaunch, 10000) \
-    X(cuTensorMapEncodeTiled, 12000)
-
-
-#define DECLARE_CUDA_FUNC_EXTERN(name, _cuda_version) \
+    X(cuInit, "cuInit", 2000) \
+    X(cuLibraryLoadData, "cuLibraryLoadData", 12000) \
+    X(cuLibraryUnload, "cuLibraryUnload", 12000) \
+    X(cuLibraryGetKernel, "cuLibraryGetKernel", 12000) \
+    X(cuGetErrorString, "cuGetErrorString", 6000) \
+    X(cuLaunchKernel, "cuLaunchKernel", 7000) \
+    X(cuLaunchKernelEx, "cuLaunchKernelEx", 11060) \
+    X(cuPointerGetAttribute, "cuPointerGetAttribute", 4000) \
+    X(cuCtxSynchronize, "cuCtxSynchronize", 2000) \
+    X(cuCtxPushCurrent, "cuCtxPushCurrent", 4000) \
+    X(cuCtxPopCurrent, "cuCtxPopCurrent", 4000) \
+    X(cuCtxGetCurrent, "cuCtxGetCurrent", 4000) \
+    X(cuCtxSetCurrent, "cuCtxSetCurrent", 4000) \
+    X(cuCtxGetDevice, "cuCtxGetDevice", 2000) \
+    X(cuCtxGetDevice_v2, "cuCtxGetDevice", 13000) \
+    X(cuCtxGetId, "cuCtxGetId", 12000) \
+    X(cuDeviceGet, "cuDeviceGet", 2000) \
+    X(cuDeviceGetCount, "cuDeviceGetCount", 2000) \
+    X(cuDeviceGetAttribute, "cuDeviceGetAttribute", 2000) \
+    X(cuDevicePrimaryCtxRetain, "cuDevicePrimaryCtxRetain", 7000) \
+    X(cuDriverGetVersion, "cuDriverGetVersion", 2020) \
+    X(cuEventCreate, "cuEventCreate", 2000) \
+    X(cuEventDestroy, "cuEventDestroy", 2000) \
+    X(cuEventQuery, "cuEventQuery", 2000) \
+    X(cuEventRecord, "cuEventRecord", 2000) \
+    X(cuKernelGetFunction, "cuKernelGetFunction", 12000) \
+    X(cuKernelGetAttribute, "cuKernelGetAttribute", 12000) \
+    X(cuKernelSetAttribute, "cuKernelSetAttribute", 12000) \
+    X(cuMemAlloc, "cuMemAlloc", 3020) \
+    X(cuMemAllocHost, "cuMemAllocHost", 3020) \
+    X(cuMemFree, "cuMemFree", 3020) \
+    X(cuMemFreeHost, "cuMemFreeHost", 2000) \
+    X(cuMemGetAddressRange, "cuMemGetAddressRange", 3020) \
+    X(cuIpcGetMemHandle, "cuIpcGetMemHandle", 4010) \
+    X(cuIpcOpenMemHandle, "cuIpcOpenMemHandle", 4010) \
+    X(cuIpcCloseMemHandle, "cuIpcCloseMemHandle", 4010) \
+    X(cuMemcpyHtoDAsync, "cuMemcpyHtoDAsync", 3020) \
+    X(cuStreamCreate, "cuStreamCreate", 2000) \
+    X(cuStreamDestroy, "cuStreamDestroy", 4000) \
+    X(cuStreamGetCtx, "cuStreamGetCtx", 9020) \
+    X(cuStreamGetId, "cuStreamGetId", 12000) \
+    X(cuStreamIsCapturing, "cuStreamIsCapturing", 10000) \
+    X(cuStreamSynchronize, "cuStreamSynchronize", 7000) \
+    X(cuStreamWaitEvent, "cuStreamWaitEvent", 7000) \
+    X(cuEventElapsedTime, "cuEventElapsedTime", 12080) \
+    X(cuGraphCreate, "cuGraphCreate", 10000) \
+    X(cuGraphDestroy, "cuGraphDestroy", 10000) \
+    X(cuGraphAddEventRecordNode, "cuGraphAddEventRecordNode", 11010) \
+    X(cuGraphAddKernelNode, "cuGraphAddKernelNode", 12000) \
+    X(cuGraphAddMemsetNode, "cuGraphAddMemsetNode", 10000) \
+    X(cuGraphAddMemAllocNode, "cuGraphAddMemAllocNode", 11040) \
+    X(cuGraphAddMemFreeNode, "cuGraphAddMemFreeNode", 11040) \
+    X(cuGraphInstantiateWithFlags, "cuGraphInstantiateWithFlags", 11040) \
+    X(cuGraphExecDestroy, "cuGraphExecDestroy", 10000) \
+    X(cuGraphLaunch, "cuGraphLaunch", 10000) \
+    X(cuTensorMapEncodeTiled, "cuTensorMapEncodeTiled", 12000)
+
+
+#define DECLARE_CUDA_FUNC_EXTERN(name, _key, _cuda_version) \
     decltype(::name)* name;
 
 struct DriverApi {

cext/tile_kernel.cpp

@@ -3433,13 +3433,10 @@ static Result<PreparedLaunch> prepare_launch(
         return ErrorRaised;
     }
 
-    if (!ctx_guard.switch_to(helper->cuda_context))
-        return ErrorRaised;
-
     // Get the compute capability of the device this launch targets.
     // Devices with the same compute capability can share a compiled kernel.
     CUdevice dev;
-    CUresult dev_res = driver->cuCtxGetDevice(&dev);
+    CUresult dev_res = driver->cuCtxGetDevice_v2(&dev, helper->cuda_context);
     if (dev_res != CUDA_SUCCESS) {
         return raise(PyExc_RuntimeError, "Failed to get current CUDA device: %s",
                      get_cuda_error(driver, dev_res));
@@ -3483,6 +3480,9 @@ static Result<PreparedLaunch> prepare_launch(
             kernel_item = kernel_map.insert(std::move(helper->constants), std::move(*res));
     }
 
+    if (!ctx_guard.switch_to(helper->cuda_context))
+        return ErrorRaised;
+
     if (stage_list_args
             && !stage_list_args_on_stream(driver, launch_stream, helper->cuda_context,
                                           helper->arena, helper->list_args,