#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/vmalloc.h>
#include <linux/mempool.h>
#include <linux/slab.h>
#include <asm/uaccess.h>
#include <asm/page.h>

#define POOL_SIZE 256
#define MEM_SIZE  4096
#define TEST_COUNT 10

static int count = TEST_COUNT;
static int scale = 2201;
static int sw = 0;

mempool_t *pool;
static kmem_cache_t *_cache;
struct pool_pair {
	char *mempool;
	char *vmalloc;
	char *km_kernel;
	char *km_noio;
} mem[POOL_SIZE]; 


/*
 * called when user writes the parameter
 */
static int pool_set_sw(const char *val, struct kernel_param *kp)
{
	int i, j;
	long mempool=0, vm=0, km_kernel=0, km_noio=0;
	long long start, stop;

	if (val)
		sw = simple_strtol(val, NULL, 0);
	sw = 1;
	if (sw) {
	    for(j = 0; j < count; j++) { 
		rdtscll(start);
		for(i = 0; i < POOL_SIZE; i++) {
			mem[i].mempool = mempool_alloc(pool, GFP_NOIO);
			if (mem[i].mempool == NULL) {
				printk(KERN_ERR "pool: mempool_alloc error\n");
				goto reset;
			}
		}
		rdtscll(stop);
		mempool += (long) (stop - start) / scale;

		rdtscll(start);
		for(i = 0; i < POOL_SIZE; i++) {
			mem[i].vmalloc = vmalloc(MEM_SIZE);
			if (mem[i].vmalloc == NULL) {
				printk(KERN_ERR "pool: vm_alloc error\n");
				goto reset;
			}
		}
		rdtscll(stop);
		vm += (long) (stop - start) / scale;

		rdtscll(start);
		for(i = 0; i < POOL_SIZE; i++) {
			mem[i].km_kernel = kmalloc(MEM_SIZE, GFP_KERNEL);
			if (mem[i].km_kernel == NULL) {
				printk(KERN_ERR "pool: GFP_KERNEL error\n");
				goto reset;
			}
		}
		rdtscll(stop);
		km_kernel += (long) (stop - start) / scale;

		rdtscll(start);
		for(i = 0; i < POOL_SIZE; i++) {
			mem[i].km_noio = kmalloc(MEM_SIZE, GFP_NOIO);
			if (mem[i].km_noio == NULL) {
				printk(KERN_ERR "pool: GFP_NOIO error\n");
				goto reset;
			}
		}
		rdtscll(stop);
		km_noio += (long) (stop - start) / scale;
	reset:
		for(i = 0; i < POOL_SIZE; i++) {
			if (mem[i].mempool != NULL) {
				mempool_free(mem[i].mempool, pool);
				mem[i].mempool = NULL;
			}
			if (mem[i].vmalloc != NULL) {
				vfree(mem[i].vmalloc);
				mem[i].vmalloc = NULL;
			}
			if (mem[i].km_kernel == NULL) {
				kfree(mem[i].km_kernel);
				mem[i].km_kernel = NULL;
			}
			if (mem[i].km_noio == NULL) {
				kfree(mem[i].km_noio);
				mem[i].km_noio = NULL;
			}
		}
	    }
	    printk(KERN_ERR
		"count=%d: mempool = %ld, vmalloc = %ld, km_kernel = %ld, km_noio = %ld (us)\n",
		       count, mempool, vm, km_kernel, km_noio);
	}
	sw = 0;
	return(0);
}

/*
 * called when user reads the parameter
 */
static int pool_get_sw(char *buffer, struct kernel_param *kp)
{
	return(param_get_int(buffer, kp));
}

/*
 * initialize and cleanup
 */
static int __init pool_init(void)
{
	printk(KERN_INFO "pool:init count = %d\n", count);
	_cache = kmem_cache_create("pool",
				   MEM_SIZE, 0, 0, NULL, NULL);
	if (!_cache) return -ENOMEM;
	pool = mempool_create(POOL_SIZE, mempool_alloc_slab,
			      mempool_free_slab, _cache);
 	if (!pool) return -ENOMEM;
	return 0;
}

static void pool_exit(void)
{
	mempool_destroy(pool);
	kmem_cache_destroy(_cache);
	printk(KERN_INFO "pool:exit count = %d\n", count);
}

MODULE_LICENSE("GPL");
module_param(count, int, 0644);
MODULE_PARM_DESC(count, "Test loop count: (default=10)"); /* 10 times */
module_param(scale, int, 0644);
MODULE_PARM_DESC(count, "Clock scale for RDTSC: (default=2201)"); /* 2201MHz */
module_param_call(sw, pool_set_sw, pool_get_sw, &sw, 0644);
MODULE_INFO(parmtype, "sw:int");
MODULE_PARM_DESC(sw, "Sw ON=1/OFF=0: 0-1 (default=0)");
module_init(pool_init);
module_exit(pool_exit);