/* SPDX-License-Identifier: GPL-2.0 */
/*
 * CUDA uprobe helper.
 *
 * Separate object for uprobes. Shares the pinned map at
 * MLFQ_CUDA_PIN_PATH with the main scheduler so either NVML
 * or CUDA activity is visible to classification.
 *
 * Copyright (c) 2026 Galih Tama <galpt@v.recipes>
 */

#include <vmlinux.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_tracing.h>
#include "intf.h"

char _license[] SEC("license") = "GPL";

/*
 * Pinned map at MLFQ_CUDA_PIN_PATH. Both objects share it.
 */
struct {
	__uint(type, BPF_MAP_TYPE_HASH);
	__uint(max_entries, 1024);
	__type(key, u32);
	__type(value, struct mlfq_gpu_nvml_state);
	__uint(pinning, LIBBPF_PIN_BY_NAME);
} cuda_gpu_map SEC(".maps");

/*
 * mlfq_cuda_uprobe_update - Record a CUDA launch for the current pid.
 *
 * Update the map with the current pid and timestamp.
 */
static __always_inline void mlfq_cuda_uprobe_update(void)
{
	u32 pid = (u32)(bpf_get_current_pid_tgid() >> 32);
	struct mlfq_gpu_nvml_state st;

	st.active = 1;
	st.pad = 0;
	st.last_seen = bpf_ktime_get_ns();
	bpf_map_update_elem(&cuda_gpu_map, &pid, &st, BPF_ANY);
}

SEC("uprobe//usr/lib/libcuda.so.1:cuLaunchKernel")
int mlfq_cuda_launch(void *ctx)
{
	mlfq_cuda_uprobe_update();
	return 0;
}

SEC("uprobe//usr/lib/libcuda.so.1:cuLaunchCooperativeKernel")
int mlfq_cuda_launch_coop(void *ctx)
{
	mlfq_cuda_uprobe_update();
	return 0;
}

SEC("uprobe//usr/lib/x86_64-linux-gnu/libcuda.so.1:cuLaunchKernel")
int mlfq_cuda_launch_x86_64(void *ctx)
{
	mlfq_cuda_uprobe_update();
	return 0;
}

SEC("uprobe//usr/lib/x86_64-linux-gnu/libcuda.so.1:cuLaunchCooperativeKernel")
int mlfq_cuda_launch_coop_x86_64(void *ctx)
{
	mlfq_cuda_uprobe_update();
	return 0;
}

SEC("uprobe//usr/lib64/libcuda.so.1:cuLaunchKernel")
int mlfq_cuda_launch_lib64(void *ctx)
{
	mlfq_cuda_uprobe_update();
	return 0;
}

SEC("uprobe//usr/lib64/libcuda.so.1:cuLaunchCooperativeKernel")
int mlfq_cuda_launch_coop_lib64(void *ctx)
{
	mlfq_cuda_uprobe_update();
	return 0;
}
