Add correlation ID filter to track traced kernel launches

Implement a thread-safe correlation ID filter to ensure activity records
are only emitted for kernel launches that have been traced via callbacks.
This prevents emitting activity records for kernels that were filtered out
by the rate limiter or other filtering logic.

The filter uses a C++ unordered_set wrapped with mutex protection and
exposed via a C API. When a callback is fired for a kernel launch, the
correlation ID is inserted into the filter. When the corresponding activity
record arrives, it's checked against the filter - only activities whose
correlation IDs are in the filter will fire the kernel_executed probe.

Changes:
- Add correlation_filter.cpp/h with thread-safe insert/check_and_remove API
- Update cupti-prof.c to insert correlation IDs on callback and check on activity
- Add debug output to trace callback and activity matching flow
- Update CMakeLists.txt to build the correlation filter

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

Author: gnurizen

cupti/CMakeLists.txt

@@ -1,5 +1,5 @@
 cmake_minimum_required(VERSION 3.18)
-project(parcagpucupti C)
+project(parcagpucupti C CXX)
 
 # CUDA root directory (can be overridden via -DCUDA_ROOT=...)
 if(NOT DEFINED CUDA_ROOT)
@@ -20,18 +20,20 @@ set(CUDAToolkit_INCLUDE_DIRS
 # Add CUDA library directory to link directories
 link_directories(${CUDA_LIBDIR})
 
-# Create shared library
-add_library(parcagpucupti SHARED cupti-prof.c)
+# Create shared library with both C and C++ sources
+add_library(parcagpucupti SHARED cupti-prof.c correlation_filter.cpp)
 
 # Set properties
 set_target_properties(parcagpucupti PROPERTIES
     C_STANDARD 11
     C_STANDARD_REQUIRED ON
+    CXX_STANDARD 17
+    CXX_STANDARD_REQUIRED ON
     POSITION_INDEPENDENT_CODE ON
 )
 
-# Add debug symbols
-target_compile_options(parcagpucupti PRIVATE -g)
+# Add debug symbols and disable C++ exceptions
+target_compile_options(parcagpucupti PRIVATE -g -fno-exceptions)
 
 # Include directories
 target_include_directories(parcagpucupti PRIVATE

cupti/correlation_filter.cpp

@@ -0,0 +1,77 @@
+#include "correlation_filter.h"
+#include <unordered_set>
+#include <mutex>
+#include <memory>
+
+// CorrelationFilter implementation using std::unordered_set with mutex protection
+// This provides thread-safe access with minimal overhead for our use case
+class CorrelationFilter {
+public:
+    CorrelationFilter() = default;
+
+    // Insert a correlation ID into the filter
+    // Thread-safe
+    void insert(uint32_t correlation_id) {
+        std::lock_guard<std::mutex> lock(mutex_);
+        set_.insert(correlation_id);
+    }
+
+    // Check if correlation ID exists and remove it atomically
+    // Returns true if found and removed, false if not found
+    // Thread-safe
+    bool check_and_remove(uint32_t correlation_id) {
+        std::lock_guard<std::mutex> lock(mutex_);
+        auto it = set_.find(correlation_id);
+        if (it != set_.end()) {
+            set_.erase(it);
+            return true;
+        }
+        return false;
+    }
+
+    // Get current size
+    // Thread-safe
+    size_t size() const {
+        std::lock_guard<std::mutex> lock(mutex_);
+        return set_.size();
+    }
+
+private:
+    std::unordered_set<uint32_t> set_;
+    mutable std::mutex mutex_;
+};
+
+// C API implementation
+extern "C" {
+
+CorrelationFilterHandle correlation_filter_create(void) {
+    return new CorrelationFilter();
+}
+
+void correlation_filter_destroy(CorrelationFilterHandle filter) {
+    if (filter) {
+        delete static_cast<CorrelationFilter*>(filter);
+    }
+}
+
+void correlation_filter_insert(CorrelationFilterHandle filter, uint32_t correlation_id) {
+    if (filter) {
+        static_cast<CorrelationFilter*>(filter)->insert(correlation_id);
+    }
+}
+
+bool correlation_filter_check_and_remove(CorrelationFilterHandle filter, uint32_t correlation_id) {
+    if (filter) {
+        return static_cast<CorrelationFilter*>(filter)->check_and_remove(correlation_id);
+    }
+    return false;
+}
+
+size_t correlation_filter_size(CorrelationFilterHandle filter) {
+    if (filter) {
+        return static_cast<CorrelationFilter*>(filter)->size();
+    }
+    return 0;
+}
+
+} // extern "C"

cupti/correlation_filter.h

@@ -0,0 +1,35 @@
+#pragma once
+
+#include <stdint.h>
+#include <stdbool.h>
+#include <stddef.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// Opaque handle to correlation filter
+typedef void* CorrelationFilterHandle;
+
+// Create a new correlation filter
+CorrelationFilterHandle correlation_filter_create(void);
+
+// Destroy the correlation filter
+void correlation_filter_destroy(CorrelationFilterHandle filter);
+
+// Insert a correlation ID into the filter
+// Thread-safe: can be called from multiple threads concurrently
+void correlation_filter_insert(CorrelationFilterHandle filter, uint32_t correlation_id);
+
+// Check if a correlation ID exists and remove it if found
+// Returns true if the correlation ID was found and removed, false otherwise
+// Thread-safe: safe to call concurrently with inserts
+bool correlation_filter_check_and_remove(CorrelationFilterHandle filter, uint32_t correlation_id);
+
+// Get the current size of the filter (number of tracked correlation IDs)
+// Note: This is an approximate count in concurrent scenarios
+size_t correlation_filter_size(CorrelationFilterHandle filter);
+
+#ifdef __cplusplus
+}
+#endif

cupti/cupti-prof.c

@@ -11,6 +11,8 @@
 
 #include <cupti.h>
 
+#include "correlation_filter.h"
+
 // Debug logging control
 static bool debug_enabled = false;
 
@@ -61,6 +63,9 @@ static bool rateLimiterTryAcquire(uint64_t nowNs) {
   return false;
 }
 
+// Correlation ID filter
+static CorrelationFilterHandle correlationFilter = NULL;
+
 static void init_debug(void) {
   static bool initialized = false;
   if (!initialized) {
@@ -200,6 +205,14 @@ int InitializeInjection(void) {
     DEBUG_PRINTF("[CUPTI] Enabled CONCURRENT_KERNEL activity\n");
   }
 
+  // Create correlation filter
+  correlationFilter = correlation_filter_create();
+  if (correlationFilter) {
+    DEBUG_PRINTF("[CUPTI] Correlation filter created and enabled\n");
+  } else {
+    fprintf(stderr, "[CUPTI] Warning: Failed to create correlation filter\n");
+  }
+
   atexit(cleanup);
 
   DEBUG_PRINTF("[CUPTI] Successfully initialized CUPTI callbacks with external "
@@ -236,13 +249,17 @@ static void parcagpuCuptiCallback(void *userdata, CUpti_CallbackDomain domain,
   if (domain == CUPTI_CB_DOMAIN_RUNTIME_API &&
       cbdata->callbackSite == CUPTI_API_ENTER) {
     runtimeEnterCorrelationId = correlationId;
+    DEBUG_PRINTF("[CUPTI] Runtime API ENTER: correlationId=%u\n", correlationId);
     return;
   }
 
   // We hook on EXIT because that makes our probe overhead not add to GPU
   // launch latency and hopefully covers some of the overhead in the shadow of
   // GPU async work.
   if (cbdata->callbackSite != CUPTI_API_EXIT) {
+    if (cbdata->callbackSite == CUPTI_API_ENTER && domain == CUPTI_CB_DOMAIN_DRIVER_API) {
+      DEBUG_PRINTF("[CUPTI] Driver API ENTER: correlationId=%u (will check on EXIT)\n", correlationId);
+    }
     return;
   }
 
@@ -253,22 +270,22 @@ static void parcagpuCuptiCallback(void *userdata, CUpti_CallbackDomain domain,
   if (domain == CUPTI_CB_DOMAIN_DRIVER_API) {
     // Skip if this driver call is under a runtime call (same correlation ID)
     if (correlationId == runtimeEnterCorrelationId) {
-      DEBUG_PRINTF("[CUPTI] Skipping driver EXIT correlationId=%u - runtime "
+      DEBUG_PRINTF("[CUPTI] Skipping driver EXIT correlationId=%u (runtimeEnter=%u) - runtime "
                    "will handle\n",
-                   correlationId);
+                   correlationId, runtimeEnterCorrelationId);
       return;
     }
     // Pure driver call (no runtime wrapper) - use negative cbid
     signedCbid = -(int)cbid;
     DEBUG_PRINTF(
-        "[CUPTI] Driver API callback: cbid=%d, correlationId=%u, func=%s\n",
-        cbid, correlationId, name);
+        "[CUPTI] Driver API EXIT callback: cbid=%d, correlationId=%u, runtimeEnter=%u, func=%s\n",
+        cbid, correlationId, runtimeEnterCorrelationId, name);
   } else if (domain == CUPTI_CB_DOMAIN_RUNTIME_API) {
     signedCbid = (int)cbid;
-    runtimeEnterCorrelationId = 0; // Clear after use
     DEBUG_PRINTF(
-        "[CUPTI] Runtime API callback: cbid=%d, correlationId=%u, func=%s\n",
-        cbid, correlationId, name);
+        "[CUPTI] Runtime API EXIT callback: cbid=%d, correlationId=%u, runtimeEnter=%u, func=%s\n",
+        cbid, correlationId, runtimeEnterCorrelationId, name);
+    runtimeEnterCorrelationId = 0; // Clear after use
   } else {
     return;
   }
@@ -286,6 +303,14 @@ static void parcagpuCuptiCallback(void *userdata, CUpti_CallbackDomain domain,
 
   outstandingEvents++;
   DTRACE_PROBE3(parcagpu, cuda_correlation, correlationId, signedCbid, name);
+
+  // Insert correlation ID into filter so we can match it later in activity buffer
+  if (correlationFilter) {
+    correlation_filter_insert(correlationFilter, correlationId);
+    DEBUG_PRINTF("[CUPTI] Inserted correlationId=%u into filter (size=%zu)\n",
+                 correlationId, correlation_filter_size(correlationFilter));
+  }
+
   // If we let too many events pile up it overwhelms the perf_event buffers,
   // just another reason to explore just passing the activity buffer through to
   // eBPF.
@@ -342,9 +367,26 @@ static void parcagpuBufferCompleted(CUcontext ctx, uint32_t streamId,
                    k->graphId, k->graphNodeId, k->name, k->correlationId,
                    k->deviceId, k->streamId, k->start, k->end,
                    k->end - k->start);
-      DTRACE_PROBE8(parcagpu, kernel_executed, k->start, k->end,
-                    k->correlationId, k->deviceId, k->streamId, k->graphId,
-                    k->graphNodeId, k->name);
+
+      // Check if this correlation ID is in our filter
+      bool should_fire = true;
+      if (correlationFilter) {
+        should_fire = correlation_filter_check_and_remove(correlationFilter, k->correlationId);
+        if (!should_fire) {
+          DEBUG_PRINTF("[CUPTI] Filtered out correlationId=%u (not tracked)\n",
+                       k->correlationId);
+        } else {
+          DEBUG_PRINTF("[CUPTI] Matched correlationId=%u - firing kernel_executed (filter size=%zu)\n",
+                       k->correlationId, correlation_filter_size(correlationFilter));
+        }
+      }
+
+      // Only fire probe if correlation ID was tracked (or filter disabled)
+      if (should_fire) {
+        DTRACE_PROBE8(parcagpu, kernel_executed, k->start, k->end,
+                      k->correlationId, k->deviceId, k->streamId, k->graphId,
+                      k->graphNodeId, k->name);
+      }
     }
   }
 
@@ -388,5 +430,15 @@ void cleanup(void) {
     subscriber = 0;
   }
 
+  // Destroy correlation filter
+  if (correlationFilter) {
+    size_t remaining = correlation_filter_size(correlationFilter);
+    if (remaining > 0) {
+      DEBUG_PRINTF("[CUPTI] Warning: %zu correlation IDs still in filter at cleanup\n", remaining);
+    }
+    correlation_filter_destroy(correlationFilter);
+    correlationFilter = NULL;
+  }
+
   DEBUG_PRINTF("[CUPTI] Cleanup completed\n");
 }