make hugetlb optional

Signed-off-by: Oliver Upton <oliver.upton@linux.dev>

1 files changed, 13 insertions, 2 deletions

diff --git a/main.c b/main.c
index c6f091a..8a4e045 100644
--- a/main.c
+++ b/main.c
@@ -1,4 +1,5 @@
 #include <assert.h>
+#include <stdbool.h>
 #include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -26,6 +27,7 @@ enum mode {
 static size_t test_size = 2UL * GB;
 static size_t test_iterations = 100;
 static enum mode test_mode;
+static bool use_hugetlb;
 
 #define PR_MODE(mode)				\
 	printf("    %d: "#mode"\n", mode);
@@ -46,13 +48,19 @@ static void pr_help(const char *progname)
 	printf("  -i ITERATIONS		(default: %lu)\n", test_iterations);
 	printf("  -m MODE		(default: %d)\n", test_mode);
 	pr_modes();
+	printf("  -H use hugetlb for test buffers\n");
 	printf("  -h prints this message\n");
 }
 
 static void *setup_test_buffer(void)
 {
+	int flags = MAP_PRIVATE | MAP_ANONYMOUS;
+
+	if (use_hugetlb)
+		flags |= MAP_HUGETLB;
+
 	void *buf = mmap(0, test_size, PROT_READ|PROT_WRITE,
-			 MAP_PRIVATE|MAP_ANONYMOUS|MAP_HUGETLB, 0, 0);
+			 flags, 0, 0);
 
 	assert(buf != MAP_FAILED);
 	memset(buf, 0xf, test_size);
@@ -128,7 +136,7 @@ int main(int argc, char **argv)
 {
 	int c;
 
-	while ((c = getopt(argc, argv, "hs:i:m:")) != -1) {
+	while ((c = getopt(argc, argv, "Hhs:i:m:")) != -1) {
 		switch (c) {
 		case 's':
 			test_size = strtoul(optarg, NULL, 0);
@@ -147,6 +155,9 @@ int main(int argc, char **argv)
 				return 1;
 			}
 			break;
+		case 'H':
+			use_hugetlb = true;
+			break;
 		case 'h':
 			pr_help(argv[0]);
 			return 0;