Adding support for cudaMalloc* bytes

UO-OACISS · Jul 31, 2020 · 4a64815 · 4a64815
1 parent b9a7d4e
commit 4a64815
Showing 1 changed file with 50 additions and 0 deletions.
diff --git a/src/apex/activity_trace_async.cpp b/src/apex/activity_trace_async.cpp
@@ -475,6 +475,55 @@ bool initialize_first_time() {
     return true;
 }
 
+bool getBytesIfMalloc(CUpti_CallbackId id, const void* params, std::string context) {
+    size_t bytes = 0;
+    switch (id) {
+        case CUPTI_RUNTIME_TRACE_CBID_cudaMalloc_v3020: {
+            bytes = ((cudaMalloc_v3020_params_st*)(params))->size;
+            break;
+        }
+        case CUPTI_RUNTIME_TRACE_CBID_cudaMallocPitch_v3020: {
+            bytes = ((cudaMallocPitch_v3020_params_st*)(params))->width *
+                    ((cudaMallocPitch_v3020_params_st*)(params))->height;
+            break;
+        }
+        case CUPTI_RUNTIME_TRACE_CBID_cudaMallocArray_v3020: {
+            bytes = ((cudaMallocArray_v3020_params_st*)(params))->width *
+                    ((cudaMallocArray_v3020_params_st*)(params))->height;
+            break;
+        }
+        case CUPTI_RUNTIME_TRACE_CBID_cudaMallocHost_v3020: {
+            bytes = ((cudaMallocHost_v3020_params_st*)(params))->size;
+            break;
+        }
+        case CUPTI_RUNTIME_TRACE_CBID_cudaMalloc3D_v3020: {
+            cudaExtent extent = ((cudaMalloc3D_v3020_params_st*)(params))->extent;
+            bytes = extent.depth * extent.height * extent.width;
+            break;
+        }
+        case CUPTI_RUNTIME_TRACE_CBID_cudaMalloc3DArray_v3020: {
+            cudaExtent extent = ((cudaMalloc3DArray_v3020_params_st*)(params))->extent;
+            bytes = extent.depth * extent.height * extent.width;
+            break;
+        }
+        case CUPTI_RUNTIME_TRACE_CBID_cudaMallocMipmappedArray_v5000: {
+            cudaExtent extent = ((cudaMallocMipmappedArray_v5000_params_st*)(params))->extent;
+            bytes = extent.depth * extent.height * extent.width;
+            break;
+        }
+        case CUPTI_RUNTIME_TRACE_CBID_cudaMallocManaged_v6000: {
+            bytes = ((cudaMallocManaged_v6000_params_st*)(params))->size;
+            break;
+        }
+        default: {
+            return false;
+        }
+    }
+    double value = (double)(bytes);
+    store_counter_data("GPU: Bytes Allocated", context, apex::profiler::get_time_ns(), value);
+    return true;
+}
+
 void apex_cupti_callback_dispatch(void *ud, CUpti_CallbackDomain domain,
         CUpti_CallbackId id, const void *params) {
     static bool initialized = initialize_first_time();
@@ -509,6 +558,7 @@ void apex_cupti_callback_dispatch(void *ud, CUpti_CallbackDomain domain,
         map_mutex.lock();
         correlation_map[cbdata->correlationId] = timer;
         map_mutex.unlock();
+        getBytesIfMalloc(id, cbdata->functionParams, tmp);
     } else if (!timer_stack.empty()) {
         auto timer = timer_stack.top();
         apex::stop(timer);