#include <linux/kernel.h>		// for using printk(), etc.
#include <linux/module.h>		// necessary for creating a kernel module
#include <linux/init.h>			// kernel modules as of 2.4 can use macros (see end of file)
#include <linux/version.h>		// for the KERNEL_VERSION macro
#include <linux/slab.h>
#include <linux/vmalloc.h>	/* for vmalloc()	*/

#ifdef CONFIG_X86_TSC
#include <asm/msr.h>			/* a 64bit precision solution, should go a long way */
#else
#include <linux/timex.h>		/* the platform independent solution, normally resorts to 32bit */
#endif

#if LINUX_VERSION_CODE >= KERNEL_VERSION(2,4,25)
#include <linux/moduleparam.h>	// enable module_param(..) and friends support
#endif
#if LINUX_VERSION_CODE < KERNEL_VERSION(2,6,0) /* only necessary when we're not using KBUILD */
#include <linux/version.h>
#endif

MODULE_AUTHOR("Willem de Bruijn");
MODULE_DESCRIPTION("A benchmark comparing performance of kmalloc and vmalloc");
MODULE_LICENSE("GPL");

#ifdef CONFIG_X86_TSC
typedef unsigned long my_clock_t;
inline my_clock_t my_get_cycles(my_clock_t in){
	register unsigned long unused;
	rdtsc(in,unused);
	return in;
}
#else
typedef unsigned long my_clock_t;
inline my_clock_t my_get_cycles(my_clock_t in){
	return get_cycles();
}
#endif

#define TESTLENGTH PAGE_SIZE * 10

static unsigned long do_memperf_linear_read(char* array, int length){
	my_clock_t start, end;
	int i;
	register char j;
	
	printk(KERN_NOTICE "testing memory by linearly reading %d bytes... ",length);
	start = my_get_cycles(start);
	for(i=0; i < length; i++)
		j = array[i];
	end = my_get_cycles(end);
	printk(KERN_NOTICE "took %lu cycles\n", end - start);
	
	return end - start;
}

static unsigned long do_memperf_linear_write(char* array, int length){
	my_clock_t start, end;
	int i;
	
	printk(KERN_NOTICE "testing memory by linearly writing %d bytes... ",length);
	start = my_get_cycles(start);
	for(i=0; i < length; i++)
		array[i] = 1;
	end = my_get_cycles(end);
	printk(KERN_NOTICE "took %lu cycles\n", end - start);
	
	return end - start;
}

static void do_memperf(char* array, int type, int test){
		
	switch (test){
		/* simple read test */
		case 0	: do_memperf_linear_read(array, TESTLENGTH); break;
		/* simple write test */
		case 1	: do_memperf_linear_write(array, TESTLENGTH); break;
		/* worst case read for virtual memory : a single byte */
		case 2	: do_memperf_linear_read(array, 1); break;
		/* worst case write : a single byte */
		case 3	: do_memperf_linear_write(array, 1); break;
	}
}

static int __init memperf_init(void){
	int test;
	int type;
	char* arrays[8];

	printk(KERN_NOTICE "loading memory benchmarking module\n");
	
	/* setup the structures. All at once, so we can measure page misses for each test */
	for (type=0; type < 2; type++){
		for (test=0; test < 4; test++){
			if (type)
				arrays[(type*4)+test] = vmalloc (TESTLENGTH);
			else
				arrays[(type*4)+test] = kmalloc (TESTLENGTH, GFP_KERNEL);
		
			if (!arrays[(type*4)+test]){
				printk(KERN_NOTICE "unable to allocate the required amount of memory. Aborting (memleak occurred)\n");
				return 0;
			}
		}
	}
		
	/* run the tests */
	for (type=0; type < 2; type++){
		printk(KERN_NOTICE "testing type %d\n", type);
		for (test=0; test < 4; test++)
			do_memperf(arrays[(type*4)+test],type, test);
	}
	
	/* and clean up */
	for (type=0; type < 2; type++){
		for (test=0; test < 4; test++){
			if (type)
				vfree(arrays[(type*4)+test]);
			else
				kfree(arrays[(type*4)+test]);
		}
	}
	return 0;
}

/**
	module destruction.
*/
static void __exit memperf_exit(void){

	printk(KERN_NOTICE "memory benchmarking module cleanup complete\n");
	
}

module_init(memperf_init);
module_exit(memperf_exit);