1 files changed, 56 insertions, 0 deletions
diff --git a/tools/guest-diag/guest_kprobes.sh b/tools/guest-diag/guest_kprobes.sh
new file mode 100644
index 0000000..f676e20
--- /dev/null
+++ b/tools/guest-diag/guest_kprobes.sh
@@ -0,0 +1,56 @@
+#!/bin/sh
+# guest_kprobes.sh - (re)install the UVM trace probes used by semfix_watch.py.
+# Runs in the guest as root; kprobes do not survive a guest boot, so this must
+# run after every VM restart.
+#
+#   uvmt  uvm_gpu_tracking_semaphore_update_completed_value  sem=%di
+#   uvch  __uvm_push_begin_acquire_on_channel_with_info      ch=%di
+#   uvbp  uvm_pushbuffer_begin_push                          pb=%di
+#   uvpbf uvm_pushbuffer_create                              pb=%si
+#
+# Tries the plain symbol first (module symbols usually resolve), then the
+# module:symbol spelling, then without the fetch argument (fetchargs on some
+# symbols are rejected).
+
+TR=/sys/kernel/tracing
+[ -d "$TR" ] || TR=/sys/kernel/debug/tracing
+
+add_probe() {
+    name=$1; sym=$2; fmt=$3
+    echo "-:$name" >> "$TR/kprobe_events" 2>/dev/null
+    if ! grep -qw "$sym" /proc/kallsyms 2>/dev/null; then
+        echo "FAIL $name: symbol $sym not in kallsyms"
+        return 1
+    fi
+    if echo "p:$name $sym $fmt" >> "$TR/kprobe_events" 2>/dev/null; then
+        if grep -q "p:$name" "$TR/kprobe_events" 2>/dev/null; then
+            echo "ok $name $sym"
+            return 0
+        fi
+    fi
+    echo "-:$name" >> "$TR/kprobe_events" 2>/dev/null
+    if echo "p:$name nvidia_uvm:$sym $fmt" >> "$TR/kprobe_events" 2>/dev/null; then
+        if grep -q "p:$name" "$TR/kprobe_events" 2>/dev/null; then
+            echo "ok-module $name $sym"
+            return 0
+        fi
+    fi
+    echo "-:$name" >> "$TR/kprobe_events" 2>/dev/null
+    if echo "p:$name $sym" >> "$TR/kprobe_events" 2>/dev/null; then
+        if grep -q "p:$name" "$TR/kprobe_events" 2>/dev/null; then
+            echo "ok-bare $name $sym"
+            return 0
+        fi
+    fi
+    echo "FAIL $name $sym"
+    return 1
+}
+
+add_probe uvmt uvm_gpu_tracking_semaphore_update_completed_value 'sem=%di val=+16(%di):u32'
+add_probe uvch __uvm_push_begin_acquire_on_channel_with_info 'ch=%di'
+add_probe uvbp uvm_pushbuffer_begin_push 'pb=%di'
+add_probe uvpbf uvm_pushbuffer_create 'pb=%si'
+
+echo 1 > "$TR/events/kprobes/enable" 2>/dev/null || true
+echo "--- installed probes ---"
+grep -E 'uvmt|uvch|uvbp|uvpbf' "$TR/kprobe_events" 2>/dev/null || echo "(none)"