module.c

#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/fs.h>
#include <linux/device.h>
#include <linux/cdev.h>
#include <linux/kdebug.h>
#include <asm/nmi.h>
#include <linux/smp.h>
#include <linux/sched.h>
#include <linux/sched.h>
#include <asm/processor.h>
#include <asm/io.h>
#include <asm/ptrace.h>
#include <../../kernel/sched/sched.h>

const char TAG[] = "Perf Driver :";

// Model-Specific Registers
static const u32 IA32_PERF_GLOBAL_CTRL = 0x038F;
static const u32 IA32_PERF_GLOBAL_OVF_CTRL = 0x0390;
static const u32 IA32_PERF_FIXED_CTR_CTRL = 0x038D;
static const u32 IA32_FIXED_CTR0 = 0x0309;
static const u32 IA32_FIXED_CTR1 = 0x030A;
static const u32 IA32_FIXED_CTR2 = 0x030B;
static const u32 IA32_FIXED_CTR_CTRL = 0x038D;

static const u32 IA32_PMC0 = 0x00C1;
static const u32 IA32_PMC1 = 0x00C2;
static const u32 IA32_PMC2 = 0x00C3;
static const u32 IA32_PMC3 = 0x00C4;
static const u32 IA32_PERFEVTSEL0 = 0x0186;
static const u32 IA32_PERFEVTSEL1 = 0x0187;
static const u32 IA32_PERFEVTSEL2 = 0x0188;
static const u32 IA32_PERFEVTSEL3 = 0x0189;

static const u64 INST_UNHALTED = 0x003c;

// Valid entries for IA32_PERF_GLOBAL_OVF_CTRL
static const u64 CLR_OVF_PMC0 = (u64) 0x1 << 0;
static const u64 CLR_OVF_PMC1 = (u64) 0x1 << 1;

static const u64 CLR_OVF_BUFFER = (u64) 0x1 << 62;

// Valid entries for IA32_PERFEVTSELn
static const u64 INT_ENABLE = (u64) 0x1 << 20;
static const u64 COUNTER_ENABLE = (u64) 0x1 << 22;
static const u64 USR_MODE = (u64) 0x1 << 16;
static const u64 OS_MODE = (u64) 0x1 << 17;

static void wrmsr64_safe_on_cpu(int cpu, u32 reg, u64 val) {
	wrmsr_safe_on_cpu(cpu, reg, (u32) ((u64) (unsigned long) val), (u32) ((u64) (unsigned long) val >> 32));
}

static void clear_msrs(void) {
	int cpu;
	for_each_online_cpu(cpu)
	{
		wrmsr64_safe_on_cpu(cpu, IA32_PERF_GLOBAL_CTRL, 0x0);
		wrmsr64_safe_on_cpu(cpu, IA32_PERF_GLOBAL_OVF_CTRL, 0x0);
		wrmsr64_safe_on_cpu(cpu, IA32_PERF_FIXED_CTR_CTRL, 0x0);

		wrmsr64_safe_on_cpu(cpu, IA32_FIXED_CTR0, 0x0);
		wrmsr64_safe_on_cpu(cpu, IA32_FIXED_CTR1, 0x0);
		wrmsr64_safe_on_cpu(cpu, IA32_FIXED_CTR2, 0x0);

		wrmsr64_safe_on_cpu(cpu, IA32_PMC0, 0x0);
		wrmsr64_safe_on_cpu(cpu, IA32_PMC1, 0x0);
		wrmsr64_safe_on_cpu(cpu, IA32_PMC2, 0x0);
		wrmsr64_safe_on_cpu(cpu, IA32_PMC3, 0x0);

		wrmsr64_safe_on_cpu(cpu, IA32_PERFEVTSEL0, 0x0);
		wrmsr64_safe_on_cpu(cpu, IA32_PERFEVTSEL1, 0x0);
		wrmsr64_safe_on_cpu(cpu, IA32_PERFEVTSEL2, 0x0);
		wrmsr64_safe_on_cpu(cpu, IA32_PERFEVTSEL3, 0x0);
	}
}

static void read_msrs(void) {
	int cpu;
	for_each_online_cpu(cpu)
	{
		u32 lo, hi;
		u64 val;

		rdmsr_safe_on_cpu(cpu, IA32_PMC0, &lo, &hi);
		val = ((u64) hi << 32) | (u64) lo;
		printk(KERN_INFO "%s IA32_PMC0 value : %llx\n", TAG, val);

		rdmsr_safe_on_cpu(cpu, IA32_FIXED_CTR2, &lo, &hi);
		val = ((u64) hi << 32) | (u64) lo;
		printk(KERN_INFO "%s IA32_FIXED_CTR2 value : %llx\n", TAG, val);
	}
}

unsigned long read_msrs_on_cpu(int cpu) {
	u32 lo, hi;
	u64 val;

	rdmsr_safe_on_cpu(cpu, IA32_PMC0, &lo, &hi);
	val = ((u64) hi << 32) | (u64) lo;

	printk(KERN_INFO "%s IA32_PMC0 value : %llx\n", TAG, val);
	return val;
}

static void init_msr_counters(void) {
	int cpu;
	for_each_online_cpu(cpu)
	{
		u64 counterVal = (u64) (-999);

		// Clear overflow by setting bit in GLOBAL_OVF_CTRL
		wrmsr64_safe_on_cpu(cpu, IA32_PERF_GLOBAL_OVF_CTRL, CLR_OVF_PMC0 | CLR_OVF_BUFFER);

		// Program programmable counter 0
		wrmsr64_safe_on_cpu(cpu, IA32_PERFEVTSEL0, (u64) INST_UNHALTED | INT_ENABLE | COUNTER_ENABLE | USR_MODE);

		// Reset PMC0 counter
		wrmsr64_safe_on_cpu(cpu, IA32_PMC0, counterVal);
	}
}

static void begin_sampling(void) {
	int cpu;
	for_each_online_cpu(cpu)
	{
		wrmsr64_safe_on_cpu(cpu, IA32_PERF_GLOBAL_CTRL, 0x1);
	}
}

static void end_sampling(void) {
	int cpu;
	for_each_online_cpu(cpu)
	{
		wrmsr64_safe_on_cpu(cpu, IA32_PERF_GLOBAL_CTRL, 0x0);
	}
}

int ipid = 1;
module_param(ipid, int, 0);

unsigned long long counter = 0;
struct task_struct *ts;

u64 totalCount = 0;

static int __kprobes PMC_handler(unsigned int cmd, struct pt_regs *regs) {
	struct task_struct *t1;
	int cpu;
	u64 counterVal = (u64) (0xffffffffff00);

	t1 = ts;

	if (t1 == NULL) {
		return NOTIFY_STOP;
	}

	cpu = smp_processor_id();

	// Disable the PMU's temporarly
	wrmsr64_safe_on_cpu(cpu, IA32_PERF_GLOBAL_CTRL, 0x0);

	printk(KERN_INFO "Perf Sample: CPU %d PID %d PC %lx\n", cpu, current->pid, task_pt_regs(ts)->ip);
	totalCount += read_msrs_on_cpu(cpu);
	wrmsr64_safe_on_cpu(cpu, IA32_PERF_GLOBAL_OVF_CTRL, CLR_OVF_PMC0 | CLR_OVF_BUFFER);
	wrmsr64_safe_on_cpu(cpu, IA32_PMC0, counterVal);		// Write the new counter value
	// Enable interrupts
	apic_write(APIC_LVTPC, APIC_DM_NMI);
	// Start counters again
	wrmsr64_safe_on_cpu(cpu, IA32_PERF_GLOBAL_CTRL, 0x1);

	return NOTIFY_STOP;
}

static int init_PMC_handler(void) {
	int ret;
	apic_write(APIC_LVTPC, APIC_DM_NMI);
	ret = register_nmi_handler(NMI_LOCAL, PMC_handler, 0, "perf_handler");
	return 0;
}

static void killall(void) {
	end_sampling();
	clear_msrs();
	unregister_nmi_handler(NMI_LOCAL, "perf_handler");
	printk(KERN_INFO "%s device unregistered\n", TAG);
}

/*****************************************************
 *       Context Switch Hook handler
 *****************************************************/

extern void (*contextSwitchHook)(struct rq*, struct task_struct*, struct task_struct*, int);
extern void (*taskExitHook)(struct task_struct*);

struct task_struct *ttemp;

void context_switch_handler(struct rq* rq, struct task_struct *prev, struct task_struct *next, int cpu) {
	// Our process is pushed out of the cpu
	if (prev->pid == ipid) {
		// Stop the PMU counter
		wrmsr64_safe_on_cpu(cpu, IA32_PERF_GLOBAL_CTRL, 0x0);
	}

	// If our process is getting scheduled next
	if (next->pid == ipid) {
		// Start the PMU counter
		wrmsr64_safe_on_cpu(cpu, IA32_PERF_GLOBAL_CTRL, 0x1);
	}
}

void task_exit_handler(struct task_struct *task) {
	if (task->pid == ipid) {
		printk(KERN_INFO "Total Counter = %ld \n", totalCount);
		read_msrs();
		ts = NULL;
		end_sampling();
		unregister_nmi_handler(NMI_LOCAL, "perf_handler");
	}
}

static int __init perf_init(void) /* Constructor */
{
	if (init_PMC_handler() != 0) {
		printk(KERN_ERR "%s failed to register NMI notifier!\n", TAG);
		killall();
		return -1;
	}

	printk(KERN_INFO "Starting Perf Module for pid %d\n", ipid);

	clear_msrs();
	init_msr_counters();
	read_msrs();
	/*begin_sampling();*/

	ts = pid_task(find_vpid(ipid), PIDTYPE_PID);
	if (ts == NULL) {
		printk(KERN_INFO "Invalid PID\n");
	}
	printk(KERN_INFO "%s initialization successful\n", TAG);

	// Register the context switch hook
	contextSwitchHook = context_switch_handler;
	taskExitHook = task_exit_handler;

	return 0;
}

static void __exit perf_exit(void) /* Destructor */
{
	// Unregister the context switch hook
	contextSwitchHook = NULL;
	taskExitHook = NULL;
	killall();
}

module_init(perf_init);
module_exit(perf_exit);

MODULE_LICENSE("GPL");
MODULE_DESCRIPTION("perf counter driver");