[Polly] [GPUJIT] [Untested] Switch from Runtime API calls for managed memory to Driver API calls.

[bollu]

Summary

Thanks to Philipp Schaad for noticing this inconsitency.
It should now be possible to link on non-cuda systems again.

Note that this is untested. I will test and make sure this works in a
couple hours.

[bollu]

@PhilippSchaad - Can you confirm that this now allows you to link on a machine without CUDA? Sorry for letting this slip by, I mixed up the runtime/driver APIs.

[bollu]

• [NFC] remove blank line

[PhilippSchaad]

Thank you for the quick fix! I can confirm that it works.

polly/trunk/tools/GPURuntime/GPUJIT.c

 /* Dynamic library handles for the CUDA and CUDA runtime library. */
 static void *HandleCuda;
 static void *HandleCudaRT;
 
 /* Type-defines of function pointer to CUDA driver APIs. */
 typedef CUresult CUDAAPI CuMemAllocFcnTy(CUdeviceptr *, size_t);
 static CuMemAllocFcnTy *CuMemAllocFcnPtr;
 
+typedef CUresult CUDAAPI CuMemAllocManagedFcnTy(CUdeviceptr *, size_t,
+                                                unsigned int);
+static CuMemAllocManagedFcnTy *CuMemAllocManagedFcnPtr;
+
 typedef CUresult CUDAAPI CuLaunchKernelFcnTy(
     CUfunction F, unsigned int GridDimX, unsigned int GridDimY,
     unsigned int gridDimZ, unsigned int blockDimX, unsigned int BlockDimY,
     unsigned int BlockDimZ, unsigned int SharedMemBytes, CUstream HStream,
     void **KernelParams, void **Extra);
 static CuLaunchKernelFcnTy *CuLaunchKernelFcnPtr;
 
 typedef CUresult CUDAAPI CuMemcpyDtoHFcnTy(void *, CUdeviceptr, size_t);
if (initialDeviceAPILibrariesCUDA() == 0)
     return 0;
 
   CuLaunchKernelFcnPtr =
       (CuLaunchKernelFcnTy *)getAPIHandleCUDA(HandleCuda, "cuLaunchKernel");
 
   CuMemAllocFcnPtr =
       (CuMemAllocFcnTy *)getAPIHandleCUDA(HandleCuda, "cuMemAlloc_v2");
 
+  CuMemAllocManagedFcnPtr = (CuMemAllocManagedFcnTy *)getAPIHandleCUDA(
+      HandleCuda, "cuMemAllocManaged");
+
   CuMemFreeFcnPtr =
       (CuMemFreeFcnTy *)getAPIHandleCUDA(HandleCuda, "cuMemFree_v2");
 
   CuMemcpyDtoHFcnPtr =
       (CuMemcpyDtoHFcnTy *)getAPIHandleCUDA(HandleCuda, "cuMemcpyDtoH_v2");
 
   CuMemcpyHtoDFcnPtr =
       (CuMemcpyHtoDFcnTy *)getAPIHandleCUDA(HandleCuda, "cuMemcpyHtoD_v2");
void polly_freeManaged(void *mem) {
   // calls in the original source than `malloc` calls. Hence, replacing all
   // `free`s with `cudaFree` does not work, since we would try to free
   // 'illegal' memory.
   // As a quick fix, we keep a free list and check if `mem` is a managed memory
   // pointer. If it is, we call `cudaFree`.
   // If not, we pass it along to the underlying allocator.
   // This is a hack, and can be removed if the underlying issue is fixed.
   if (isManagedPtr(mem)) {
-    if (cudaFree(mem) != cudaSuccess) {
+    if (CuMemFreeFcnPtr((size_t)mem) != CUDA_SUCCESS) {
       fprintf(stderr, "cudaFree failed.\n");
       exit(-1);
     }
     return;
   } else {
     free(mem);
   }
 }
 
 void *polly_mallocManaged(size_t size) {
   // Note: [Size 0 allocations]
   // Sometimes, some runtime computation of size could create a size of 0
   // for an allocation. In these cases, we do not wish to fail.
   // The CUDA API fails on size 0 allocations.
   // So, we allocate size a minimum of size 1.
   if (!size && DebugMode)
     fprintf(stderr, "cudaMallocManaged called with size 0. "
                     "Promoting to size 1");
   size = max(size, 1);
-  polly_initContextCUDA();
-  dump_function();
-  void *a;
-  if (cudaMallocManaged(&a, size, cudaMemAttachGlobal) != cudaSuccess) {
+  PollyGPUContext *_ = polly_initContextCUDA();
+  assert(_ && "polly_initContextCUDA failed");
+
+  void *newMemPtr;
+  const CUresult Res = CuMemAllocManagedFcnPtr((CUdeviceptr *)&newMemPtr, size,
+                                               CU_MEM_ATTACH_GLOBAL);
+  if (Res != CUDA_SUCCESS) {
     fprintf(stderr, "cudaMallocManaged failed for size: %zu\n", size);
     exit(-1);
   }
-  addManagedPtr(a);
-  return a;
+  addManagedPtr(newMemPtr);
+  return newMemPtr;
 }
 
 static void freeDeviceMemoryCUDA(PollyGPUDevicePtr *Allocation) {
   dump_function();
   CUDADevicePtr *DevPtr = (CUDADevicePtr *)Allocation->DevicePtr;
   CuMemFreeFcnPtr((CUdeviceptr)DevPtr->Cuda);
   free(DevPtr);
   free(Allocation);