#!/bin/bash
# kret.sh - kretprobe the nvGpuOps* call chain in the guest's nvidia.ko and
# capture return values while /tmp/cdt runs. Run with sudo.
set -u

TR=/sys/kernel/tracing
SYMS="nvGpuOpsGetGpuInfo nvGpuOpsDeviceCreate nvGpuOpsRmDeviceCreate \
nvGpuOpsRmSubDeviceCreate queryFbInfo findDeviceClasses queryVirtMode \
getSysmemLinkInfo getSystemMemoryWindow getNvswitchInfo getEgmInfo getAtsInfo \
nvGpuOpsQueryGpuConfidentialComputeCaps allocNvlinkStatus"

echo > $TR/kprobe_events
for f in $SYMS; do
    if ! grep -q " [tT] $f\b" /proc/kallsyms; then
        echo "skip $f (not in kallsyms)"
        continue
    fi
    if ! echo "r:ret_$f $f \$retval" >> $TR/kprobe_events 2>/dev/null; then
        echo "FAILED to register $f"
    fi
done
echo "--- registered:"
grep ret_ $TR/kprobe_events 2>/dev/null | head -20

if [ -d $TR/events/kprobes ]; then
    echo 1 > $TR/events/kprobes/enable
    dmesg -C
    /tmp/cdt >/dev/null 2>&1
    echo "--- results (retval per probe):"
    grep -E "ret_" $TR/trace | awk '{n=$4; sub(/^ret_/,"",n); printf "%s %s\n", n, $NF}' | sort | uniq -c | sort -rn | head -30
    echo 0 > $TR/events/kprobes/enable
    echo > $TR/trace
fi