[Polly] [GPUJIT] Teach GPUJIT to use a pre-existing CUDA context if available.

[bollu]

Summary

On mixing the driver and runtime APIs, it is quite possible that a
context already exists due to runtime API usage. In this case, we should
try to use the same context.

Teach GPUJIT to detect that a context exists and how to pick up this context.

This is a part of (https://reviews.llvm.org/D35991) that was extracted
out.

[grosser]

LGTM.

tools/GPURuntime/GPUJIT.c

29 ↗

(On Diff #109156)

Not needed?

[bollu]

[NFC] remove assert, add error checking.

[bollu]

• [NFC] remove assert & add error checking. For some reason, this did not show up with the previous diff.

[bollu]

• [NFC] ran through update-format

polly/trunk/tools/GPURuntime/GPUJIT.c

 static CuInitFcnTy *CuInitFcnPtr;
 
 typedef CUresult CUDAAPI CuDeviceGetCountFcnTy(int *);
 static CuDeviceGetCountFcnTy *CuDeviceGetCountFcnPtr;
 
 typedef CUresult CUDAAPI CuCtxCreateFcnTy(CUcontext *, unsigned int, CUdevice);
 static CuCtxCreateFcnTy *CuCtxCreateFcnPtr;
 
+typedef CUresult CUDAAPI CuCtxGetCurrentFcnTy(CUcontext *);
+static CuCtxGetCurrentFcnTy *CuCtxGetCurrentFcnPtr;
+
 typedef CUresult CUDAAPI CuDeviceGetFcnTy(CUdevice *, int);
 static CuDeviceGetFcnTy *CuDeviceGetFcnPtr;
 
 typedef CUresult CUDAAPI CuModuleLoadDataExFcnTy(CUmodule *, const void *,
                                                  unsigned int, CUjit_option *,
                                                  void **);
 static CuModuleLoadDataExFcnTy *CuModuleLoadDataExFcnPtr;
 
CuDeviceGetCountFcnPtr =
       (CuDeviceGetCountFcnTy *)getAPIHandleCUDA(HandleCuda, "cuDeviceGetCount");
 
   CuDeviceGetFcnPtr =
       (CuDeviceGetFcnTy *)getAPIHandleCUDA(HandleCuda, "cuDeviceGet");
 
   CuCtxCreateFcnPtr =
       (CuCtxCreateFcnTy *)getAPIHandleCUDA(HandleCuda, "cuCtxCreate_v2");
 
+  CuCtxGetCurrentFcnPtr =
+      (CuCtxGetCurrentFcnTy *)getAPIHandleCUDA(HandleCuda, "cuCtxGetCurrent");
+
   CuModuleLoadDataExFcnPtr = (CuModuleLoadDataExFcnTy *)getAPIHandleCUDA(
       HandleCuda, "cuModuleLoadDataEx");
 
   CuModuleLoadDataFcnPtr =
       (CuModuleLoadDataFcnTy *)getAPIHandleCUDA(HandleCuda, "cuModuleLoadData");
 
   CuModuleGetFunctionFcnPtr = (CuModuleGetFunctionFcnTy *)getAPIHandleCUDA(
       HandleCuda, "cuModuleGetFunction");
if (Context == 0) {
     fprintf(stderr, "Allocate memory for Polly GPU context failed.\n");
     exit(-1);
   }
   Context->Context = malloc(sizeof(CUDAContext));
   if (Context->Context == 0) {
     fprintf(stderr, "Allocate memory for Polly CUDA context failed.\n");
     exit(-1);
   }
-  CuCtxCreateFcnPtr(&(((CUDAContext *)Context->Context)->Cuda), 0, Device);
+
+  // In cases where managed memory is used, it is quite likely that
+  // `cudaMallocManaged` / `polly_mallocManaged` was called before
+  // `polly_initContext` was called.
+  //
+  // If `polly_initContext` calls `CuCtxCreate` when there already was a
+  // pre-existing context created by the runtime API, this causes code running
+  // on P100 to hang. So, we query for a pre-existing context to try and use.
+  // If there is no pre-existing context, we create a new context
+
+  // The possible pre-existing context from previous runtime API calls.
+  CUcontext MaybeRuntimeAPIContext;
+  if (CuCtxGetCurrentFcnPtr(&MaybeRuntimeAPIContext) != CUDA_SUCCESS) {
+    fprintf(stderr, "cuCtxGetCurrent failed.\n");
+    exit(-1);
+  }
+
+  // There was no previous context, initialise it.
+  if (MaybeRuntimeAPIContext == NULL) {
+    if (CuCtxCreateFcnPtr(&(((CUDAContext *)Context->Context)->Cuda), 0,
+                          Device) != CUDA_SUCCESS) {
+      fprintf(stderr, "cuCtxCreateFcnPtr failed.\n");
+      exit(-1);
+    }
+  } else {
+    ((CUDAContext *)Context->Context)->Cuda = MaybeRuntimeAPIContext;
+  }
 
   if (CacheMode)
     CurrentContext = Context;
 
   return Context;
 }
 
 static void freeKernelCUDA(PollyGPUFunction *Kernel) {