Skip to content

Commit 55722ae

Browse files
gnurizengnurizen
committed
Support new activity_batch parcagpucupti probe
In order to reduce bpf overhead send through up to 128 kernel launch timing activities to the usdt probe. The old single shot kernel_executed probe is still supported. Inline correlation and kernel_exec into cuda_probe, tail-call only activity_batch The unwinder is sensitive to tail calls, so minimize them: inline cuda_correlation and cuda_kernel_exec directly into cuda_probe's switch statement using bpf_usdt_arg() for USDT arg reading.
1 parent ede3a25 commit 55722ae

12 files changed

Lines changed: 459 additions & 127 deletions

File tree

‎interpreter/gpu/cuda.go‎

Lines changed: 70 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -22,14 +22,28 @@ import (
2222
const (
2323
// eBPF program names for USDT probes
2424
// These correspond to the function names in cuda.ebpf.c, not the SEC() paths
25-
USDTProgCudaCorrelation = "cuda_correlation"
26-
USDTProgCudaKernel = "cuda_kernel_exec"
27-
USDTProgCudaProbe = "cuda_probe"
25+
USDTProgCudaCorrelation = "cuda_correlation"
26+
USDTProgCudaKernel = "cuda_kernel_exec"
27+
USDTProgCudaActivityBatch = "cuda_activity_batch"
28+
USDTProgCudaProbe = "cuda_probe"
29+
30+
// Tail-call prog array indices — must match CUDA_PROG_* in cuda.ebpf.c.
31+
// These are used as both the BPF attach cookie (low 32 bits) and the
32+
// cuda_progs prog array key for tail-call dispatch.
33+
CudaProgCorrelation = 0
34+
CudaProgKernelExec = 1
35+
CudaProgActivityBatch = 2
2836
)
2937

38+
const cudaProgsMap = "cuda_progs"
39+
3040
var (
3141
// gpuFixers maps PID to gpuTraceFixer
3242
gpuFixers sync.Map
43+
44+
// cudaTailCallOnce ensures the cuda_progs prog array is populated exactly
45+
// once. The tail-call targets must be in place before cuda_probe fires.
46+
cudaTailCallOnce sync.Once
3347
)
3448

3549
// SymbolizedCudaTrace holds a symbolized trace awaiting GPU timing information.
@@ -110,15 +124,36 @@ func Loader(ebpf interpreter.EbpfHandler, info *interpreter.LoaderInfo) (interpr
110124
return nil, nil
111125
}
112126

113-
// Filter to only the probes we need
114-
var requiredProbes []pfelf.USDTProbe
115-
for _, probe := range parcagpuProbes {
116-
if probe.Name == "cuda_correlation" || probe.Name == "kernel_executed" {
117-
requiredProbes = append(requiredProbes, probe)
127+
// Filter to only the probes we need.
128+
// Always require cuda_correlation. Prefer activity_batch over kernel_executed.
129+
var correlationProbe *pfelf.USDTProbe
130+
var kernelProbe *pfelf.USDTProbe
131+
var batchProbe *pfelf.USDTProbe
132+
for i := range parcagpuProbes {
133+
switch parcagpuProbes[i].Name {
134+
case "cuda_correlation":
135+
correlationProbe = &parcagpuProbes[i]
136+
case "kernel_executed":
137+
kernelProbe = &parcagpuProbes[i]
138+
case "activity_batch":
139+
batchProbe = &parcagpuProbes[i]
118140
}
119141
}
120-
if len(requiredProbes) != 2 {
121-
log.Warnf("parcagpu USDT probes in %s missing required probes (need cuda_correlation and kernel_executed): %v", info.FileName(), parcagpuProbes)
142+
if correlationProbe == nil {
143+
log.Warnf("parcagpu USDT probes in %s missing cuda_correlation: %v", info.FileName(), parcagpuProbes)
144+
return nil, nil
145+
}
146+
147+
var requiredProbes []pfelf.USDTProbe
148+
requiredProbes = append(requiredProbes, *correlationProbe)
149+
if batchProbe != nil {
150+
requiredProbes = append(requiredProbes, *batchProbe)
151+
log.Debugf("parcagpu: using activity_batch mode for %s", info.FileName())
152+
} else if kernelProbe != nil {
153+
requiredProbes = append(requiredProbes, *kernelProbe)
154+
log.Debugf("parcagpu: using kernel_executed mode for %s", info.FileName())
155+
} else {
156+
log.Warnf("parcagpu USDT probes in %s missing kernel probe (need activity_batch or kernel_executed): %v", info.FileName(), parcagpuProbes)
122157
return nil, nil
123158
}
124159
parcagpuProbes = requiredProbes
@@ -137,25 +172,44 @@ func Loader(ebpf interpreter.EbpfHandler, info *interpreter.LoaderInfo) (interpr
137172

138173
func (d *data) Attach(ebpf interpreter.EbpfHandler, pid libpf.PID, _ libpf.Address,
139174
_ remotememory.RemoteMemory) (interpreter.Instance, error) {
140-
// Maps usdt probe name to ebpf program name.
141-
// Use the first character of the probe name as a cookie.
142-
// 'c' -> cuda_correlation
143-
// 'k' -> cuda_kernel_exec
175+
// Map USDT probe names to eBPF program names and tail-call indices.
176+
// The cookie doubles as the cuda_progs prog array key for tail-call dispatch.
144177
cookies := make([]uint64, len(d.probes))
145178
progNames := make([]string, len(d.probes))
146179
for i, probe := range d.probes {
147-
cookies[i] = uint64(probe.Name[0])
148-
// Map probe names to specific program names for single-shot mode
149180
switch probe.Name {
150181
case "cuda_correlation":
182+
cookies[i] = CudaProgCorrelation
151183
progNames[i] = USDTProgCudaCorrelation
152184
case "kernel_executed":
185+
cookies[i] = CudaProgKernelExec
153186
progNames[i] = USDTProgCudaKernel
187+
case "activity_batch":
188+
cookies[i] = CudaProgActivityBatch
189+
progNames[i] = USDTProgCudaActivityBatch
154190
default:
155191
log.Debugf("unknown parcagpu USDT probe name: %s", probe.Name)
156192
}
157193
}
158194

195+
// Populate the tail-call prog array for activity_batch (the only tail-call
196+
// target — correlation and kernel_exec are inlined in cuda_probe).
197+
hasActivityBatch := false
198+
for _, probe := range d.probes {
199+
if probe.Name == "activity_batch" {
200+
hasActivityBatch = true
201+
break
202+
}
203+
}
204+
if hasActivityBatch {
205+
cudaTailCallOnce.Do(func() {
206+
if err := ebpf.UpdateProgArray(cudaProgsMap, 0,
207+
USDTProgCudaActivityBatch); err != nil {
208+
log.Warnf("[cuda] failed to populate %s: %v", cudaProgsMap, err)
209+
}
210+
})
211+
}
212+
159213
var lc interpreter.LinkCloser
160214
if d.link == nil {
161215
var err error

‎interpreter/gpu/cuda_test.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ func TestProgramNamesExist(t *testing.T) {
3232
progNames := []string{
3333
gpu.USDTProgCudaCorrelation,
3434
gpu.USDTProgCudaKernel,
35+
gpu.USDTProgCudaActivityBatch,
3536
}
3637

3738
for _, progName := range progNames {

‎interpreter/instancestubs.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,10 @@ func (mockup *EbpfHandlerStubs) AttachUSDTProbes(libpf.PID, string, string, []pf
7777
return nil, nil
7878
}
7979

80+
func (mockup *EbpfHandlerStubs) UpdateProgArray(string, uint32, string) error {
81+
return nil
82+
}
83+
8084
func (mockup *EbpfHandlerStubs) AttachUprobe(
8185
libpf.PID, string, uint64, string) (LinkCloser, error) {
8286
return nil, nil

‎interpreter/types.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,11 @@ type EbpfHandler interface {
136136

137137
// AttachUprobe attaches an eBPF uprobe to a function at a specific offset in a binary
138138
AttachUprobe(pid libpf.PID, path string, offset uint64, progName string) (LinkCloser, error)
139+
140+
// UpdateProgArray loads an eBPF program by name and inserts it into the
141+
// named BPF_MAP_TYPE_PROG_ARRAY at the given key. The program is loaded
142+
// once and cached; subsequent calls with the same progName reuse it.
143+
UpdateProgArray(mapName string, key uint32, progName string) error
139144
}
140145

141146
type LinkCloser interface {

‎processmanager/ebpf/ebpf.go‎

Lines changed: 108 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,10 @@ type ebpfMapsImpl struct {
8080
usdtSpecsMap *cebpf.Map
8181
nextSpecID uint32
8282
specIDLock sync.Mutex
83+
84+
// allMaps holds all loaded BPF maps by name for generic lookups
85+
// (e.g. UpdateProgArray).
86+
allMaps map[string]*cebpf.Map
8387
}
8488

8589
// Compile time check to make sure ebpfMapsImpl satisfies the interface .
@@ -98,6 +102,7 @@ func LoadMaps(ctx context.Context, maps map[string]*cebpf.Map,
98102
probeProgsMap: maps["kprobe_progs"],
99103
usdtSpecsMap: maps["__bpf_usdt_specs"],
100104
nextSpecID: 0,
105+
allMaps: maps,
101106
}
102107
impl.errCounter = make(map[metrics.MetricID]int64)
103108

@@ -212,12 +217,13 @@ func (lc *linkCloser) Unload() error {
212217
}
213218

214219
// AttachUSDTProbes allows interpreters to attach to usdt probes.
220+
//
221+
// When both multiProgName and singleProgNames are provided, multi-probe
222+
// attachment is attempted first (if the kernel supports it). On failure the
223+
// function automatically falls back to single-shot per-probe attachment.
215224
func (impl *ebpfMapsImpl) AttachUSDTProbes(pid libpf.PID, path, multiProgName string,
216225
probes []pfelf.USDTProbe, cookies []uint64, singleProgNames []string) (interpreter.LinkCloser, error) {
217226
useMulti := util.HasMultiUprobeSupport()
218-
if !useMulti && len(probes) > 1 && multiProgName != "" && len(singleProgNames) == 0 {
219-
return nil, errors.New("uprobe multi attach requires kernel support (kernel 6.6+)")
220-
}
221227

222228
containerPath := fmt.Sprintf("/proc/%d/root/%s", pid, path)
223229
exe, err := link.OpenExecutable(containerPath)
@@ -280,77 +286,88 @@ func (impl *ebpfMapsImpl) AttachUSDTProbes(pid libpf.PID, path, multiProgName st
280286
}
281287
}
282288

283-
// If multiProgName is empty or multi-probe not supported, use individual programs (one per probe)
284-
if multiProgName == "" || !useMulti {
285-
if singleProgNames == nil {
286-
return nil, fmt.Errorf("singleProgNames required when multiProgName is empty or multi-probe not supported")
287-
}
288-
if len(singleProgNames) != len(probes) {
289-
return nil, fmt.Errorf(
290-
"number of single program names %d does not match number of probes %d",
291-
len(singleProgNames), len(probes))
289+
// Try multi-probe first if the kernel supports it and a dispatcher was given.
290+
if multiProgName != "" && useMulti {
291+
lc, multiErr := impl.attachMultiProbe(exe, path, pid, multiProgName,
292+
names, addresses, offsets, finalCookies, specIDs)
293+
if multiErr == nil {
294+
return lc, nil
292295
}
293-
// Single-shot mode with multiple programs
294-
// Load all programs first
295-
progs := make([]*cebpf.Program, len(probes))
296-
for i := range probes {
297-
prog := impl.userProgs[singleProgNames[i]]
298-
if prog == nil {
299-
if err := impl.loadUSDTProgram(singleProgNames[i], false); err != nil {
300-
return nil, err
301-
}
302-
prog = impl.userProgs[singleProgNames[i]]
303-
}
304-
progs[i] = prog
296+
if singleProgNames == nil {
297+
return nil, multiErr
305298
}
299+
log.Warnf("multi-probe attach failed (%v), falling back to single-shot", multiErr)
300+
}
306301

307-
// Attach individual probes with their own programs
308-
var links []link.Link
309-
for i, probe := range probes {
310-
prog := progs[i]
311-
if prog == nil {
312-
// Clean up already attached probes
313-
for _, l := range links {
314-
l.Close()
315-
}
316-
return nil, fmt.Errorf("program %d is nil for probe %s", i, probe.Name)
317-
}
302+
// Single-shot mode: one program per probe (direct path or fallback).
303+
if singleProgNames == nil {
304+
return nil, fmt.Errorf("singleProgNames required when multi-probe not available")
305+
}
306+
if len(singleProgNames) != len(probes) {
307+
return nil, fmt.Errorf(
308+
"number of single program names %d does not match number of probes %d",
309+
len(singleProgNames), len(probes))
310+
}
318311

319-
// Prepare uprobe options
320-
uprobeOpts := &link.UprobeOptions{
321-
Address: probe.Location,
322-
RefCtrOffset: probe.SemaphoreOffset,
312+
progs := make([]*cebpf.Program, len(probes))
313+
for i := range probes {
314+
prog := impl.userProgs[singleProgNames[i]]
315+
if prog == nil {
316+
if err := impl.loadUSDTProgram(singleProgNames[i], false); err != nil {
317+
return nil, err
323318
}
319+
prog = impl.userProgs[singleProgNames[i]]
320+
}
321+
progs[i] = prog
322+
}
324323

325-
// Set cookie if provided
326-
if finalCookies != nil && i < len(finalCookies) {
327-
uprobeOpts.Cookie = finalCookies[i]
324+
var links []link.Link
325+
for i, probe := range probes {
326+
prog := progs[i]
327+
if prog == nil {
328+
for _, l := range links {
329+
l.Close()
328330
}
331+
return nil, fmt.Errorf("program %d is nil for probe %s", i, probe.Name)
332+
}
329333

330-
// Attach uprobe at the probe location
331-
l, err := exe.Uprobe(probe.Name, prog, uprobeOpts)
332-
if err != nil {
333-
// Clean up already attached probes
334-
for _, lnk := range links {
335-
lnk.Close()
336-
}
337-
return nil, fmt.Errorf("failed to attach USDT probe %s at location 0x%x: %w",
338-
probe.Name, probe.Location, err)
339-
}
340-
links = append(links, l)
334+
uprobeOpts := &link.UprobeOptions{
335+
Address: probe.Location,
336+
RefCtrOffset: probe.SemaphoreOffset,
337+
}
338+
if finalCookies != nil && i < len(finalCookies) {
339+
uprobeOpts.Cookie = finalCookies[i]
341340
}
342341

343-
log.Infof("Attached %d individual probes to %s in PID %d", len(links), path, pid)
344-
return &linkCloser{
345-
unloadLink: links,
346-
unloadSpecIDs: specIDs,
347-
specMap: impl.usdtSpecsMap,
348-
}, nil
342+
l, err := exe.Uprobe(probe.Name, prog, uprobeOpts)
343+
if err != nil {
344+
for _, lnk := range links {
345+
lnk.Close()
346+
}
347+
return nil, fmt.Errorf("failed to attach USDT probe %s at location 0x%x: %w",
348+
probe.Name, probe.Location, err)
349+
}
350+
links = append(links, l)
349351
}
350352

353+
log.Infof("Attached %d individual probes to %s in PID %d", len(links), path, pid)
354+
return &linkCloser{
355+
unloadLink: links,
356+
unloadSpecIDs: specIDs,
357+
specMap: impl.usdtSpecsMap,
358+
}, nil
359+
}
360+
361+
// attachMultiProbe loads the multi-probe dispatcher and attaches all probes
362+
// via UprobeMulti.
363+
func (impl *ebpfMapsImpl) attachMultiProbe(
364+
exe *link.Executable, path string, pid libpf.PID, multiProgName string,
365+
names []string, addresses, offsets, finalCookies []uint64,
366+
specIDs []uint32,
367+
) (interpreter.LinkCloser, error) {
351368
prog := impl.userProgs[multiProgName]
352369
if prog == nil {
353-
if err := impl.loadUSDTProgram(multiProgName, useMulti); err != nil {
370+
if err := impl.loadUSDTProgram(multiProgName, true); err != nil {
354371
return nil, err
355372
}
356373
prog = impl.userProgs[multiProgName]
@@ -374,14 +391,43 @@ func (impl *ebpfMapsImpl) AttachUSDTProbes(pid libpf.PID, path, multiProgName st
374391
}, nil
375392
}
376393

377-
// loadProgram loads an eBPF program from progSpec and populates the related maps.
394+
// UpdateProgArray loads a USDT eBPF program by name and inserts it into
395+
// the named BPF_MAP_TYPE_PROG_ARRAY at the given key.
396+
func (impl *ebpfMapsImpl) UpdateProgArray(mapName string, key uint32, progName string) error {
397+
m := impl.allMaps[mapName]
398+
if m == nil {
399+
return fmt.Errorf("map %s not found", mapName)
400+
}
401+
if impl.userProgs == nil {
402+
impl.userProgs = make(map[string]*cebpf.Program)
403+
}
404+
prog := impl.userProgs[progName]
405+
if prog == nil {
406+
if err := impl.loadUSDTProgram(progName, false); err != nil {
407+
return fmt.Errorf("failed to load program %s: %w", progName, err)
408+
}
409+
prog = impl.userProgs[progName]
410+
}
411+
fd := uint32(prog.FD())
412+
if err := m.Update(unsafe.Pointer(&key), unsafe.Pointer(&fd),
413+
cebpf.UpdateAny); err != nil {
414+
return fmt.Errorf("failed to insert %s into %s[%d]: %w",
415+
progName, mapName, key, err)
416+
}
417+
log.Infof("Inserted program %s into %s[%d]", progName, mapName, key)
418+
return nil
419+
}
420+
421+
// loadUSDTProgram loads an eBPF program from progSpec and populates the related maps.
378422
func (impl *ebpfMapsImpl) loadUSDTProgram(progName string, useMulti bool) error {
379423
progSpec := impl.coll.Programs[progName]
380424
if progSpec == nil {
381425
return fmt.Errorf("eBPF program %s not found in collection", progName)
382426
}
383427
programOptions := cebpf.ProgramOptions{
384-
// TODO: wire in debug level
428+
// LogSizeStart is the initial buffer for the verifier log on load failure.
429+
// Complex USDT programs (e.g. cuda_probe) produce large verifier output.
430+
LogSizeStart: 1 << 24, // 16MB
385431
}
386432
restoreRlimit, err := rlimit.MaximizeMemlock()
387433
if err != nil {

0 commit comments

Comments
 (0)