From 4f745e9d30ed234eb16c7a70bdd11222d2e27316 Mon Sep 17 00:00:00 2001
From: zhubr <n-a-zhubr@yandex.ru>
Date: Sun, 6 Jun 2021 16:13:15 +0300
Subject: [PATCH 1/2] Handle memory allocation failure properly, allow to
 override SIZE (mostly needed for smaller devices), fix declaration syntax for
 older compilers.

---
 main.c | 26 +++++++++++++++++++++-----
 util.c |  1 +
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/main.c b/main.c
index 6205555..fc57db2 100644
--- a/main.c
+++ b/main.c
@@ -40,7 +40,7 @@
 #include "asm-opt.h"
 #include "version.h"
 
-#define SIZE             (32 * 1024 * 1024)
+#define SIZE_DEFAULT     (32 * 1024 * 1024)
 #define BLOCKSIZE        2048
 #ifndef MAXREPEATS
 # define MAXREPEATS      10
@@ -480,12 +480,14 @@ int latency_bench(int size, int count, int use_hugepage)
     return 1;
 }
 
-int main(void)
+int main(int argc, char *argv[])
 {
-    int latbench_size = SIZE * 2, latbench_count = LATBENCH_COUNT;
+    int latbench_size, latbench_count = LATBENCH_COUNT;
+    int custom_size = SIZE_DEFAULT;
     int64_t *srcbuf, *dstbuf, *tmpbuf;
     void *poolbuf;
-    size_t bufsize = SIZE;
+    size_t bufsize;
+    bench_info *bi;
 #ifdef __linux__
     size_t fbsize = 0;
     int64_t *fbbuf = mmap_framebuffer(&fbsize);
@@ -494,11 +496,25 @@ int main(void)
 
     printf("tinymembench v" VERSION " (simple benchmark for memory throughput and latency)\n");
 
+    if (argc > 1) {
+	custom_size = atoi (argv[1]);
+	if ((custom_size < 4096) || (custom_size > 16*SIZE_DEFAULT)) {
+	    printf("WARNING: unexpected size specified, using default instead.\n");
+	    custom_size = SIZE_DEFAULT;
+	}
+    }
+    latbench_size = custom_size * 2;
+    bufsize = custom_size;
 
     poolbuf = alloc_four_nonaliased_buffers((void **)&srcbuf, bufsize,
                                             (void **)&dstbuf, bufsize,
                                             (void **)&tmpbuf, BLOCKSIZE,
                                             NULL, 0);
+    if (!poolbuf) {
+	printf("FATAL: memory allocation failed. Please override SIZE_DEFAULT (%ld bytes) on command line.\n", (long)SIZE_DEFAULT);
+	return 1;
+    }
+
     printf("\n");
     printf("==========================================================================\n");
     printf("== Memory bandwidth tests                                               ==\n");
@@ -517,7 +533,7 @@ int main(void)
     bandwidth_bench(dstbuf, srcbuf, tmpbuf, bufsize, BLOCKSIZE, " ", c_benchmarks);
     printf(" ---\n");
     bandwidth_bench(dstbuf, srcbuf, tmpbuf, bufsize, BLOCKSIZE, " ", libc_benchmarks);
-    bench_info *bi = get_asm_benchmarks();
+    bi = get_asm_benchmarks();
     if (bi->f) {
         printf(" ---\n");
         bandwidth_bench(dstbuf, srcbuf, tmpbuf, bufsize, BLOCKSIZE, " ", bi);
diff --git a/util.c b/util.c
index a0d562c..10d8c87 100644
--- a/util.c
+++ b/util.c
@@ -337,6 +337,7 @@ void *alloc_four_nonaliased_buffers(void **buf1_, int size1,
 
     ptr = buf = 
         (char *)malloc(size1 + size2 + size3 + size4 + 9 * ALIGN_PADDING);
+    if (!ptr) return ptr;
     memset(buf, 0xCC, size1 + size2 + size3 + size4 + 9 * ALIGN_PADDING);
 
     ptr = align_up(ptr, ALIGN_PADDING);

From ad962c18602368d32b4382d47f31b7429dd00899 Mon Sep 17 00:00:00 2001
From: zhubr <n-a-zhubr@yandex.ru>
Date: Sun, 6 Jun 2021 18:21:14 +0300
Subject: [PATCH 2/2] Allow to override LATBENCH_COUNT on command line (mostly
 needed for slow devices).

---
 main.c | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/main.c b/main.c
index fc57db2..3285e8e 100644
--- a/main.c
+++ b/main.c
@@ -503,6 +503,13 @@ int main(int argc, char *argv[])
 	    custom_size = SIZE_DEFAULT;
 	}
     }
+    if (argc > 2) {
+       latbench_count = atoi (argv[2]);
+       if ((latbench_count < 100) || (latbench_count > 100*LATBENCH_COUNT)) {
+           printf("WARNING: unexpected count specified, using default instead.\n");
+           latbench_count = LATBENCH_COUNT;
+       }
+    }
     latbench_size = custom_size * 2;
     bufsize = custom_size;