[Polly] [HACK] [WIP] Add functionality to GPUJIT to track unified memory.

[bollu]

Summary

• We use this in COSMO to use the unified memory feature of the newer

CUDA APIs (>= 6)

• This is a hack, so please revert this when 2018 rolls in.

[bollu]

Clearly, this entire patch is a hack. I'd like to know how to add tests for this, and how to ensure that this gets reverted at some point in the future :)

[grosser]

Hi Siddharth,

I think it is OK to upstream this code, assuming you write and document it in a generic way and explain how it is used.

tools/GPURuntime/GPUJIT.c

37

Why is this needed. Maybe upstream separately?

1225

There seem to be other situations where a context might already be available. I suggest to upstream this independently.

[bollu]

• Rebase on top of master and take tobias' suggestion into account.

[bollu]

[NFC] upload proper diff.

[grosser]

LGTM in general.

tools/GPURuntime/GPUJIT.c

1437

I don't think it adds any value to discuss COSMO here. Instead, it suggest to just explain what kind of behavior you seeand how this is addressed here.

1439

cudaFree

[bollu]

• [NFC] explain hack better

[bollu]

• [NFC] lint fix.

[singam-sanjay]

This breaks my Julia build.

How are the cudaFree and cudaMallocManaged functions resolved ? Can you add a patch to link GPURuntime to cudart ?

tools/GPURuntime/GPUJIT.h

 /******************************************************************************/
 /*                                                                            */
 /*                     The LLVM Compiler Infrastructure                       */
 /*                                                                            */
 /* This file is dual licensed under the MIT and the University of Illinois    */
 /* Open Source License. See LICENSE.TXT for details.                          */
 /*                                                                            */
 /******************************************************************************/
 /*                                                                            */
 /*  This file defines GPUJIT.                                                 */
 /*                                                                            */
 /******************************************************************************/
 
 #ifndef GPUJIT_H_
 #define GPUJIT_H_
+#include "stddef.h"
 
 /*
  * The following demostrates how we can use the GPURuntime library to
  * execute a GPU kernel.
  *
  * char KernelString[] = "\n\
  *   .version 1.4\n\
  *   .target sm_10, map_f64_to_f32\n\
void polly_copyFromDeviceToHost(PollyGPUDevicePtr *DevData, void *HostData,
                                 long MemSize);
 void polly_synchronizeDevice();
 void polly_launchKernel(PollyGPUFunction *Kernel, unsigned int GridDimX,
                         unsigned int GridDimY, unsigned int BlockSizeX,
                         unsigned int BlockSizeY, unsigned int BlockSizeZ,
                         void **Parameters);
 void polly_freeDeviceMemory(PollyGPUDevicePtr *Allocation);
 void polly_freeContext(PollyGPUContext *Context);
+
+// Note that polly_{malloc/free}Managed are currently not used by Polly.
+// We use them in COSMO by replacing all malloc with polly_mallocManaged and all
+// frees with cudaFree, so we can get managed memory "automatically".
+// Needless to say, this is a hack.
+// Please make sure that this code is not present in Polly when 2018 rolls in.
+// If this is still present, ping Siddharth Bhat <[e-mail]>
+void *polly_mallocManaged(size_t size);
+void polly_freeManaged(void *mem);
 #endif /* GPUJIT_H_ */

tools/GPURuntime/GPUJIT.c

 #include <dlfcn.h>
 #include <stdarg.h>
 #include <stdio.h>
 #include <string.h>
 #include <unistd.h>
 
 static int DebugMode;
 static int CacheMode;
 
 static PollyGPURuntime Runtime = RUNTIME_NONE;
 
 static void debug_print(const char *format, ...) {
   if (!DebugMode)
     return;
 
   va_list args;
   va_start(args, format);

 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) {
Res =
                            GridDimY, GridDimZ, BlockDimX, BlockDimY, BlockDimZ,
                            SharedMemBytes, Stream, Parameters, Extra);
   if (Res != CUDA_SUCCESS) {
     fprintf(stderr, "Launching CUDA kernel failed.\n");
     exit(-1);
   }
 }
 
+// Maximum number of managed memory pointers.
+#define MAX_POINTERS 4000
+// For the rationale behing a list of free pointers, see `polly_freeManaged`.
+void *g_managedptrs[MAX_POINTERS];
+int g_nmanagedptrs = 0;
+
+// Add a pointer as being allocated by cuMallocManaged
+void addManagedPtr(void *mem) {
+  assert(g_nmanagedptrs < MAX_POINTERS && "We have hit the maximum number of "
+                                          "managed pointers allowed. Increase "
+                                          "MAX_POINTERS");
+  g_managedptrs[g_nmanagedptrs++] = mem;
+}
+
+int isManagedPtr(void *mem) {
+  for (int i = 0; i < g_nmanagedptrs; i++) {
+    if (g_managedptrs[i] == mem)
+      return 1;
+  }
+  return 0;
+}
+
+void polly_freeManaged(void *mem) {
+  dump_function();
+
+  // Something extremely weird goes on in the COSMO code: we get calls
+  // to free memory that have not been `malloc`d. Since we replace all free
+  // with `cudaFreee`, sometimes we get calls to free to memory we have not
+  // even mallocd.
+  // So, keep a free list and check if this is one of "our" pointers. If it is,
+  // call `cudaFree`. If not, pass it along to the underlying allocator.
+  if (isManagedPtr(mem)) {
+    if (cudaFree(mem) != cudaSuccess) {
+      fprintf(stderr, "cudaFree failed.\n");
+      exit(-1);
+    }
+    return;
+  } else {
+    free(mem);
+  }
+}
+
+void *polly_mallocManaged(size_t size) {
+  dump_function();
+  void *a;
+  if (cudaMallocManaged(&a, size, cudaMemAttachGlobal) != cudaSuccess) {
+    fprintf(stderr, "cudaMallocManaged failed for size: %zu\n", size);
+    exit(-1);
+  }
+  addManagedPtr(a);
+  return a;
+}
+
 static void freeDeviceMemoryCUDA(PollyGPUDevicePtr *Allocation) {
   dump_function();
   CUDADevicePtr *DevPtr = (CUDADevicePtr *)Allocation->DevicePtr;
   CuMemFreeFcnPtr((CUdeviceptr)DevPtr->Cuda);
   free(DevPtr);
   free(Allocation);
 }