1 files changed, 12 insertions, 7 deletions

diff --git a/copy_experiments/mon_cross_ctx_copies.cu b/copy_experiments/mon_cross_ctx_copies.cu
index d9749e0..93d0e4c 100644
--- a/copy_experiments/mon_cross_ctx_copies.cu
+++ b/copy_experiments/mon_cross_ctx_copies.cu
@@ -92,7 +92,6 @@ void cpu_copy_mon(int loops, char* cpu_mem) {
 void* copy_thread(void* args_raw) {
         CUcontext ctx;
         int dev, i;
-        cudaError_t err;
         int *barrier;
         int *barrier_dev_ptr; // On Pascal+ this will be the same as *barrier
         char *pinned_hostmem, *devmem;
@@ -106,8 +105,8 @@ void* copy_thread(void* args_raw) {
         // Explictly create a context (avoids creating a primary context implictly)
         // This has been verified on CUDA 11.1 to give each thread a different context
         // handle
-        cudaGetDevice(&dev);
-        SAFE(cuCtxCreate(&ctx, 0, dev));
+        SAFE(cudaGetDevice(&dev));
+        SAFE_D(cuCtxCreate(&ctx, 0, dev));
 
         uint64_t dev_ns, dev_ns2;
         double host_s, host_s2;
@@ -249,21 +248,27 @@ static error_t arg_parser(int key, char* arg, struct argp_state *state) {
 }
 
 int main(int argc, char**argv) {
-        int tmp, dev;
+        int tmp, dev = 0;
+        CUdevice dev_itrl;
         uint64_t *ctx_times[MAX_THREADS] = {0};
         pthread_t t[MAX_THREADS];
         global_args_t g_args = {0};
 
+        // Temporarially initialize CUDA to query device attributes
+        SAFE_D(cuInit(0));
+        SAFE_D(cuDeviceGet(&dev_itrl, dev));
         // Due to some laziness in how we handle barriers, this flag needs to be true
         /// XXX: Still seems to work fine if it isn't???
-        cudaGetDevice(&dev);
-        cudaDeviceGetAttribute(&tmp, cudaDevAttrCanMapHostMemory, dev);
+        SAFE_D(cuDeviceGetAttribute(&tmp, CU_DEVICE_ATTRIBUTE_CAN_MAP_HOST_MEMORY, dev_itrl));
         if (!tmp) {
                 fprintf(stderr, "Unsupported platform. It must be possible to map host (CPU)"
-                                "DRAM into the GPU virtual address space for accurate clock"
+                                " DRAM into the GPU virtual address space for accurate clock "
                                 "synchronization. Exiting...\n");
                 return 1;
         }
+        // Terminate the context used for attrib check so it's not accidentially
+        // reused in subprocesses
+        SAFE_D(cuDevicePrimaryCtxRelease(dev_itrl));
 
         struct argp argp = {opts, arg_parser, 0, desc};
         argp_parse(&argp, argc, argv, 0, 0, &g_args);