kernel hash driver benchmark

hashbench.c

#include <crypto/hash.h>
#include <linux/crc32c.h>
#include <linux/kernel.h>
#include <linux/ktime.h>
#include <linux/module.h>
#include <linux/proc_fs.h>
#include <linux/random.h>
#include <linux/slab.h>
#include <linux/types.h>
#include <linux/vmalloc.h>
#include <linux/xxhash.h>

static struct proc_dir_entry *proc_file;
#define PROC_FILENAME     "hashbench"
#define hb_err(fmt, ...)  pr_info(PROC_FILENAME ": " fmt, ##__VA_ARGS__)
#define hb_info(fmt, ...) pr_info(PROC_FILENAME ": " fmt, ##__VA_ARGS__)

static void benchmark(char *csum_driver, size_t pages, size_t iterations) {
	// set up shash
	struct crypto_shash *csum_shash = crypto_alloc_shash(csum_driver, 0, 0);
	if (IS_ERR(csum_shash)) {
		hb_err("failed to allocate %s shash for checksum\n", csum_driver);
		return;
	}
	const char *driver_name = crypto_shash_driver_name(csum_shash);
	hb_info("driver=%s pages=%zu iters=%zu\n", driver_name, pages, iterations);

	SHASH_DESC_ON_STACK(shash, csum_shash);
	shash->tfm = csum_shash;

	// generate random data
	size_t vmem_size = pages * PAGE_SIZE;
	u8 *vmem = vmalloc(vmem_size);
	if (!vmem) {
		hb_err("failed to allocate memory\n");
		goto free_shash;
	}
	u8 *end = vmem + pages * PAGE_SIZE;
	get_random_bytes(vmem, vmem_size);

	// benchmark
	ktime_t start = ktime_get();

	u8 hash[512];
	for (size_t i = 0; i < iterations; i++) {
		if (crypto_shash_init(shash)) {
			hb_err("failed to init shash\n");
			goto free_vmem;
		}
		for (u8 *page = vmem; page != end; page += PAGE_SIZE) {
			if (crypto_shash_update(shash, page, PAGE_SIZE)) {
				hb_err("failed to update shash");
				goto free_vmem;
			}
		}
		if (crypto_shash_final(shash, hash)) {
			hb_err("failed to finalize shash");
			goto free_vmem;
		}
	}
	ktime_t delta = ktime_sub(ktime_get(), start);
	hb_info("total: %6lld ms, %6lld ns/page \n", ktime_to_ms(delta),
	        ktime_to_ns(delta) / pages / iterations);

free_vmem:
	vfree(vmem);
free_shash:
	crypto_free_shash(csum_shash);
}

static ssize_t hashbench_write(struct file *file, const char __user *buffer,
                               size_t count, loff_t *pos) {
	char input[256];
	char driver[64];
	size_t pages;
	size_t iterations;

	if (count > 255)
		return -EINVAL;

	if (copy_from_user(input, buffer, count))
		return -EFAULT;

	input[count] = '\0';

	if (sscanf(input, "%s %zu %zu", driver, &pages, &iterations) != 3) {
		return -EINVAL;
	}
	benchmark(driver, pages, iterations);

	return count;
}

static const struct proc_ops procops = {
    .proc_write = hashbench_write,
};

static int __init hash_benchmark_init(void) {
	// Create /proc/hashbench file
	proc_file = proc_create(PROC_FILENAME, 0666, NULL, &procops);
	if (NULL == proc_file) {
		hb_err("Failed to create /proc/hashbench\n");
		return -ENOMEM;
	}
	return 0;
}

static void __exit hash_benchmark_exit(void) {
	proc_remove(proc_file);
}

module_init(hash_benchmark_init);
module_exit(hash_benchmark_exit);

MODULE_LICENSE("GPL");
MODULE_AUTHOR("memchr");
MODULE_DESCRIPTION("Benchmark shash drivers");