#!/bin/sh
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
# uvm_rmmem_probe.sh - guest-side kprobe that dumps every UVM rm_mem mapping
# as it is used: {type, CPU VA, GPU0 VA, GPU1 VA, size} for each call to
# uvm_rm_mem_get_cpu_va. This is the ground truth for VA -> physical
# translations that the host-side scan has to rediscover (the CPU VA is a
# direct-map address for sysmem objects that are mapped page by page, and a
# vmap address for objects mapped through nv_vm_map_pages; only the former can
# be turned into a physical address by subtracting the direct-map base).
#
# The pushbuffer/semaphore-pool content can also be read straight out of the
# kernel with nested fetchargs, e.g. a0=+0(+48(%di)):x64 reads the first u64
# of the object's CPU mapping.
#
# Run as root in the guest, after nvidia-uvm is loaded and before starting the
# client (the first calls happen during bring-up):
#   sudo sh uvm_rmmem_probe.sh
#   ... start the CUDA workload ...
#   sudo grep rmcpu /sys/kernel/tracing/trace
#
# Note: kprobe writes must be plain "echo ... >> file" redirections. Piping
# through tee returns EINVAL on this kernel.
set -e
T=/sys/kernel/tracing

echo 0 >> $T/events/kprobes/rmcpu/enable 2>/dev/null || true
echo 0 >> $T/events/kprobes/pb/enable 2>/dev/null || true
echo > $T/kprobe_events
printf 'p:rmcpu uvm_rm_mem_get_cpu_va type=+0(%%di):u32 cpu_va=+48(%%di):x64 gpu0=+56(%%di):x64 gpu1=+64(%%di):x64 size=+2120(%%di):x64\n' >> $T/kprobe_events
printf 'p:pb uvm_rm_mem_get_cpu_va gpu0=+56(%%di):x64 a0=+0(+48(%%di)):x64 a1=+8(+48(%%di)):x64 a2=+16(+48(%%di)):x64 a3=+24(+48(%%di)):x64 a4=+32(+48(%%di)):x64 a5=+40(+48(%%di)):x64 a6=+48(+48(%%di)):x64 a7=+56(+48(%%di)):x64\n' >> $T/kprobe_events
cat > $T/events/kprobes/rmcpu/filter <<'EOF'
gpu0 != 0x121000000
EOF
cat > $T/events/kprobes/pb/filter <<'EOF'
gpu0 == 0x120000000
EOF
echo > $T/trace
echo 1 >> $T/events/kprobes/rmcpu/enable
echo 1 >> $T/events/kprobes/pb/enable
cat $T/kprobe_events
echo "enabled; traces in $T/trace"