Restore some lost changes causing verifier errors

Author: gnurizen

go.mod

@@ -13,7 +13,6 @@ require (
 	github.com/elastic/go-freelru v0.16.0
 	github.com/elastic/go-perf v0.0.0-20241029065020-30bec95324b8
 	github.com/google/uuid v1.6.0
-	github.com/ianlancetaylor/demangle v0.0.0-20251118225945-96ee0021ea0f
 	github.com/klauspost/compress v1.18.0
 	github.com/mdlayher/kobject v0.0.0-20200520190114-19ca17470d7d
 	github.com/minio/sha256-simd v1.0.1

go.sum

@@ -116,8 +116,6 @@ github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.2 h1:8Tjv8EJ+pM1xP8mK6egEbD1OgnV
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.2/go.mod h1:pkJQ2tZHJ0aFOVEEot6oZmaVEZcRme73eIFmhiVuRWs=
 github.com/hashicorp/go-version v1.7.0 h1:5tqGy27NaOTB8yJKUZELlFAS/LTKJkrmONwQKeRZfjY=
 github.com/hashicorp/go-version v1.7.0/go.mod h1:fltr4n8CU8Ke44wwGCBoEymUuxUHl09ZGVZPK5anwXA=
-github.com/ianlancetaylor/demangle v0.0.0-20251118225945-96ee0021ea0f h1:Fnl4pzx8SR7k7JuzyW8lEtSFH6EQ8xgcypgIn8pcGIE=
-github.com/ianlancetaylor/demangle v0.0.0-20251118225945-96ee0021ea0f/go.mod h1:gx7rwoVhcfuVKG5uya9Hs3Sxj7EIvldVofAWIUtGouw=
 github.com/josharian/native v1.1.0 h1:uuaP0hAbW7Y4l0ZRQ6C9zfb7Mg1mbFKry/xzDAfmtLA=
 github.com/josharian/native v1.1.0/go.mod h1:7X/raswPFr05uY3HiLlYeyQntB6OO7E/d2Cu7qoaN2w=
 github.com/jsimonetti/rtnetlink v0.0.0-20190606172950-9527aa82566a/go.mod h1:Oz+70psSo5OFh8DBl0Zv2ACw7Esh6pPUphlvZG9x7uw=

interpreter/gpu/cuda.go

@@ -24,8 +24,9 @@ const (
 	// These correspond to the function names in cuda.ebpf.c, not the SEC() paths
 	USDTProgCudaCorrelation   = "cuda_correlation"
 	USDTProgCudaKernel        = "cuda_kernel_exec"
-	USDTProgCudaActivityBatch = "cuda_activity_batch"
-	USDTProgCudaProbe         = "cuda_probe"
+	USDTProgCudaActivityBatch     = "cuda_activity_batch"
+	USDTProgCudaActivityBatchTail = "cuda_activity_batch_tail"
+	USDTProgCudaProbe             = "cuda_probe"
 
 	// BPF attach cookie values — must match CUDA_PROG_* in cuda.ebpf.c.
 	// Used in the low 32 bits of the BPF attach cookie so cuda_probe can
@@ -188,7 +189,7 @@ func (d *data) Attach(ebpf interpreter.EbpfHandler, pid libpf.PID, _ libpf.Addre
 		}
 		cudaTailCallOnce.Do(func() {
 			if err := ebpf.UpdateProgArray(cudaProgsMap, 0,
-				USDTProgCudaActivityBatch); err != nil {
+				USDTProgCudaActivityBatchTail); err != nil {
 				log.Errorf("[cuda] activity_batch tail call failed: %v", err)
 				cudaTailCallFailed = true
 			}

interpreter/gpu/cuda_test.go

@@ -33,6 +33,7 @@ func TestProgramNamesExist(t *testing.T) {
 			gpu.USDTProgCudaCorrelation,
 			gpu.USDTProgCudaKernel,
 			gpu.USDTProgCudaActivityBatch,
+			gpu.USDTProgCudaActivityBatchTail,
 		}
 
 		for _, progName := range progNames {

support/ebpf/cuda.ebpf.c

@@ -39,10 +39,15 @@ struct kernel_timing {
 // Per-CPU scratch space for large structs that exceed the BPF 512-byte stack limit.
 // Used by cuda_kernel_exec and cuda_activity_batch (mutually exclusive on same CPU).
 #define MAX_BATCH_SIZE 128
+#define PTR_BATCH      16
 struct cuda_scratch {
   struct kernel_timing timing;
   struct cupti_activity_kernel5 rec;
-  u64 ptrs[MAX_BATCH_SIZE];
+  // Pre-parsed activity_batch USDT args, set by cuda_probe before tail call.
+  // bpf_get_attach_cookie does not return the correct cookie after bpf_tail_call,
+  // so we parse args in cuda_probe and pass them via scratch.
+  u64 ab_ptrs_base;
+  u32 ab_num_activities;
 };
 
 bpf_map_def SEC("maps") cuda_scratch_heap = {
@@ -100,6 +105,8 @@ int BPF_USDT(
     timing->kernel_name[3] = '\0';
   }
 
+  DEBUG_PRINT("cuda_kernel_exec: pid=%u corr_id=%u dev=%u", pid, correlation_id, device_id);
+
   bpf_perf_event_output(ctx, &cuda_timing_events, BPF_F_CURRENT_CPU, timing, sizeof(*timing));
 
   return 0;
@@ -119,64 +126,100 @@ int BPF_USDT(cuda_activity_batch, u64 ptrs_base, u32 num_activities)
   struct kernel_timing *timing       = &scratch->timing;
   struct cupti_activity_kernel5 *rec = &scratch->rec;
 
+  DEBUG_PRINT("cuda_activity_batch: pid=%u num=%u", pid, (u32)num_activities);
+  DEBUG_PRINT("cuda_activity_batch: ptrs_base=0x%llx", ptrs_base);
+
   if (num_activities > MAX_BATCH_SIZE) {
     num_activities = MAX_BATCH_SIZE;
   }
 
-  // Read all activity pointers into scratch space in one shot.
-  if (bpf_probe_read_user(scratch->ptrs, num_activities * sizeof(u64), (void *)ptrs_base) != 0) {
-    return 0;
-  }
+  // Stack-local pointer batch — small enough for the BPF stack.
+  u64 ptrs[PTR_BATCH] = {};
 
-  for (u32 i = 0; i < num_activities && i < MAX_BATCH_SIZE; i++) {
-    u64 rec_ptr = scratch->ptrs[i & (MAX_BATCH_SIZE - 1)];
-    if (rec_ptr == 0) {
-      continue;
+  // Nested loop: outer iterates over batches of PTR_BATCH pointers,
+  // inner processes each pointer in the batch.  This keeps the verifier's
+  // jump-sequence count well under BPF_COMPLEXITY_LIMIT_JMP_SEQ (8192).
+  for (u32 batch = 0; batch < MAX_BATCH_SIZE / PTR_BATCH; batch++) {
+    u32 base = batch * PTR_BATCH;
+    if (base >= num_activities) {
+      break;
     }
 
-    // Read the full activity record and filter by kind.
-    if (bpf_probe_read_user(rec, sizeof(*rec), (void *)rec_ptr) != 0) {
-      continue;
-    }
-    if (
-      rec->kind != CUPTI_ACTIVITY_KIND_KERNEL &&
-      rec->kind != CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL) {
-      continue;
+    if (bpf_probe_read_user(ptrs, sizeof(ptrs), (void *)(ptrs_base + base * sizeof(u64))) != 0) {
+      break;
     }
 
-    timing->pid            = pid;
-    timing->correlation_id = rec->correlation_id;
-    timing->start          = rec->start;
-    timing->end            = rec->end;
-    timing->graph_node_id  = rec->graph_node_id;
-    timing->device_id      = rec->device_id;
-    timing->stream_id      = rec->stream_id;
-    timing->graph_id       = rec->graph_id;
-
-    const char *name = (const char *)rec->name_ptr;
-    int chars =
-      bpf_probe_read_user_str((char *)&timing->kernel_name, sizeof(timing->kernel_name), name);
-    if (chars <= 0) {
-      timing->kernel_name[0] = 'e';
-      timing->kernel_name[1] = 'r';
-      timing->kernel_name[2] = 'r';
-      timing->kernel_name[3] = '\0';
-    }
+    for (u32 j = 0; j < PTR_BATCH; j++) {
+      if (base + j >= num_activities) {
+        break;
+      }
+
+      u64 rec_ptr = ptrs[j];
+
+      // Read the full activity record and filter by kind.
+      if (bpf_probe_read_user(rec, sizeof(*rec), (void *)rec_ptr) != 0) {
+        continue;
+      }
+      if (
+        rec->kind != CUPTI_ACTIVITY_KIND_KERNEL &&
+        rec->kind != CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL) {
+        continue;
+      }
+
+      timing->pid            = pid;
+      timing->correlation_id = rec->correlation_id;
+      timing->start          = rec->start;
+      timing->end            = rec->end;
+      timing->graph_node_id  = rec->graph_node_id;
+      timing->device_id      = rec->device_id;
+      timing->stream_id      = rec->stream_id;
+      timing->graph_id       = rec->graph_id;
+
+      const char *name = (const char *)rec->name_ptr;
+      int chars =
+        bpf_probe_read_user_str((char *)&timing->kernel_name, sizeof(timing->kernel_name), name);
+      if (chars <= 0) {
+        timing->kernel_name[0] = 'e';
+        timing->kernel_name[1] = 'r';
+        timing->kernel_name[2] = 'r';
+        timing->kernel_name[3] = '\0';
+      }
+
+      DEBUG_PRINT(
+        "cuda_activity_batch: corr_id=%u kind=%u dev=%u",
+        rec->correlation_id,
+        rec->kind,
+        rec->device_id);
 
-    bpf_perf_event_output(ctx, &cuda_timing_events, BPF_F_CURRENT_CPU, timing, sizeof(*timing));
+      bpf_perf_event_output(ctx, &cuda_timing_events, BPF_F_CURRENT_CPU, timing, sizeof(*timing));
+    }
   }
 
   return 0;
 }
 
+// Tail-call entry point for cuda_activity_batch.  Reads pre-parsed USDT args
+// from the scratch map (set by cuda_probe before bpf_tail_call) and forwards
+// them to the inline body generated by BPF_USDT.
+SEC("usdt/cuda_activity_batch_tail")
+int cuda_activity_batch_tail(struct pt_regs *ctx)
+{
+  u32 zero                     = 0;
+  struct cuda_scratch *scratch = bpf_map_lookup_elem(&cuda_scratch_heap, &zero);
+  if (!scratch) {
+    return 0;
+  }
+  return ____cuda_activity_batch(ctx, scratch->ab_ptrs_base, scratch->ab_num_activities);
+}
+
 // Cookie values for the cuda_probe multi-probe dispatcher.
 // Must match the cookie values set in cuda.go.
 #define CUDA_PROG_CORRELATION    0
 #define CUDA_PROG_KERNEL_EXEC    1
 #define CUDA_PROG_ACTIVITY_BATCH 2
 
 // Tail-call prog array for cuda_probe.  Contains a single entry at key 0
-// for cuda_activity_batch, whose batch loop pushes past the BPF verifier's
+// for cuda_activity_batch_tail, whose batch loop pushes past the BPF verifier's
 // BPF_COMPLEXITY_LIMIT_JMP_SEQ (8192) limit.  cuda_correlation and
 // cuda_kernel_exec are inlined directly in cuda_probe.
 bpf_map_def SEC("maps") cuda_progs = {
@@ -205,7 +248,19 @@ int cuda_probe(struct pt_regs *ctx)
       graph_id,
       graph_node_id,
       name);
-  case CUDA_PROG_ACTIVITY_BATCH: bpf_tail_call(ctx, &cuda_progs, 0); break;
+  case CUDA_PROG_ACTIVITY_BATCH: {
+    // Parse USDT args before the tail call — bpf_get_attach_cookie does not
+    // return the correct cookie after bpf_tail_call.
+    u32 zero                     = 0;
+    struct cuda_scratch *scratch = bpf_map_lookup_elem(&cuda_scratch_heap, &zero);
+    if (!scratch) {
+      break;
+    }
+    scratch->ab_ptrs_base      = (u64)bpf_usdt_arg0(ctx);
+    scratch->ab_num_activities = (u32)bpf_usdt_arg1(ctx);
+    bpf_tail_call(ctx, &cuda_progs, 0);
+    break;
+  }
   default: DEBUG_PRINT("cuda_probe: unknown cookie %u", cookie); break;
   }
   return 0;

support/ebpf/tracer.ebpf.amd64

@@ -186,7 +186,7 @@ func TestCUDAVerifierMultiProbe(t *testing.T) {
 	// target — correlation and kernel_exec are inlined in cuda_probe).
 	for _, probe := range probes {
 		if probe.Name == "activity_batch" {
-			err := ebpfHandler.UpdateProgArray("cuda_progs", 0, "cuda_activity_batch")
+			err := ebpfHandler.UpdateProgArray("cuda_progs", 0, "cuda_activity_batch_tail")
 			require.NoError(t, err, "UpdateProgArray failed for cuda_activity_batch")
 			break
 		}