more than 64 threads

Signed-off-by: Pantelis Antoniou <[email protected]>

src/blake3/blake3_thread.c

@@ -170,6 +170,9 @@ int blake3_thread_submit_work(blake3_thread *t, const blake3_thread_work *work)
 
 	/* atomically update the work */
 
+	assert(t);
+	assert(work);
+
 	pthread_mutex_lock(&t->lock);
 	exp_work = NULL;
 	if (!atomic_compare_exchange_strong(&t->work, &exp_work, work)) {
@@ -224,29 +227,36 @@ blake3_thread *blake3_thread_pool_reserve(blake3_thread_pool *tp)
 {
 	blake3_thread *t;
 	unsigned int slot;
-	uint64_t v, exp;
+	uint64_t *free, v, exp;
+	unsigned int i;
 
 	t = NULL;
-	v = atomic_load(&tp->free);
-	while (v) {
-		slot = lowest_one(v);
-		assert(v & BIT64(slot));
-		exp = v;		/* expecting the previous value */
-		v &= ~BIT64(slot);	/* clear this bit */
-		if (atomic_compare_exchange_strong(&tp->free, &exp, v)) {
-			t = tp->threads + slot;
-			assert(slot == t->id);
-			return t;
+	for (i = 0, free = tp->freep; i < tp->free_count; i++, free++) {
+		v = atomic_load(free);
+		while (v) {
+			slot = lowest_one(v);
+			assert(v & BIT64(slot));
+			exp = v;		/* expecting the previous value */
+			v &= ~BIT64(slot);	/* clear this bit */
+			if (atomic_compare_exchange_strong(free, &exp, v)) {
+				slot += i * 64;
+				t = tp->threads + slot;
+				assert(slot == t->id);
+				return t;
+			}
+			v = exp;
 		}
-		v = exp;
 	}
 
 	return NULL;
 }
 
 void blake3_thread_pool_unreserve(blake3_thread_pool *tp, blake3_thread *t)
 {
-	atomic_fetch_or(&tp->free, BIT64(t->id));
+	uint64_t *free;
+
+	free = tp->freep + (unsigned int)(t->id / 64);
+	atomic_fetch_or(free, BIT64(t->id & 63));
 }
 
 void blake3_thread_pool_cleanup(blake3_thread_pool *tp)
@@ -267,6 +277,9 @@ void blake3_thread_pool_cleanup(blake3_thread_pool *tp)
 		free(tp->threads);
 	}
 
+	if (tp->freep)
+		free(tp->freep);
+
 	memset(tp, 0, sizeof(*tp));
 }
 
@@ -279,19 +292,31 @@ int blake3_thread_pool_init(blake3_thread_pool *tp, unsigned int num_threads)
 	(void)rc;
 
 	assert(tp);
-
+
+	if (!num_threads)
+		return -1;
+
 	memset(tp, 0, sizeof(*tp));
 
-	tp->num_threads = num_threads >= 64 ? 64 : num_threads;
+	tp->num_threads = num_threads;
+
+	tp->free_count = (tp->num_threads / 64) + ((tp->num_threads & 63) ? 1 : 0);
+
+	tp->freep = malloc(tp->free_count * sizeof(uint64_t));
+	if (!tp->freep)
+		goto err_out;
+
+	for (i = 0; i < tp->free_count; i++)
+		tp->freep[i] = (uint64_t)-1;
+	if (tp->num_threads & 63)
+		tp->freep[tp->free_count - 1] = BIT64(tp->num_threads & 63) - 1;
 
 	tp->threads = malloc(sizeof(*tp->threads) * tp->num_threads);
 	if (!tp->threads)
 		goto err_out;
 
 	memset(tp->threads, 0, sizeof(*tp->threads) * tp->num_threads);
 
-	tp->free = BIT64(tp->num_threads & 63) - 1;
-
 	for (i = 0, t = tp->threads; i < tp->num_threads; i++, t++) {
 
 		t->tp = tp;

src/blake3/blake3_thread.h

@@ -20,7 +20,7 @@ typedef struct _blake3_thread_work {
 	void *arg;
 } blake3_thread_work;
 
-//#define BLAKE3_THREAD_PORTABLE
+#define BLAKE3_THREAD_PORTABLE
 typedef struct _blake3_thread {
 	struct _blake3_thread_pool *tp;
 	unsigned int id;
@@ -41,7 +41,9 @@ typedef struct _blake3_thread {
 typedef struct _blake3_thread_pool {
 	unsigned int num_threads;
 	struct _blake3_thread *threads;
+	uint64_t *freep;
 	uint64_t free;
+	unsigned int free_count;
 } blake3_thread_pool;
 
 blake3_thread_pool *blake3_thread_pool_create(unsigned int num_threads);

src/internal/fy-b3sum.c

@@ -44,7 +44,7 @@ static void test_worker_thread_fn(void *arg)
 	}
 }
 
-void test_worker_threads(void)
+void test_worker_threads(unsigned int num_threads)
 {
 	blake3_thread_pool *tp;
 	blake3_thread **threads, *t;
@@ -55,9 +55,12 @@ void test_worker_threads(void)
 
 	(void)rc;
 
-	scval = sysconf(_SC_NPROCESSORS_ONLN);
-	assert(scval > 0);
-	num_cpus = (unsigned int)scval;
+	if (num_threads == 0) {
+		scval = sysconf(_SC_NPROCESSORS_ONLN);
+		assert(scval > 0);
+		num_cpus = (unsigned int)scval;
+	} else
+		num_cpus = num_threads;
 
 	fprintf(stderr, "calling: blake3_thread_pool_init()\n");
 	tp = blake3_thread_pool_create(num_cpus);
@@ -106,7 +109,7 @@ void test_worker_threads(void)
 	blake3_thread_pool_destroy(tp);
 }
 
-void test_thread_join(void)
+void test_thread_join(unsigned int num_threads)
 {
 	blake3_thread_pool *tp;
 	void **args;
@@ -116,9 +119,12 @@ void test_thread_join(void)
 
 	(void)rc;
 
-	scval = sysconf(_SC_NPROCESSORS_ONLN);
-	assert(scval > 0);
-	num_cpus = (unsigned int)scval;
+	if (num_threads == 0) {
+		scval = sysconf(_SC_NPROCESSORS_ONLN);
+		assert(scval > 0);
+		num_cpus = (unsigned int)scval;
+	} else
+		num_cpus = num_threads;
 
 	fprintf(stderr, "calling: blake3_thread_pool_create()\n");
 	tp = blake3_thread_pool_create(num_cpus);
@@ -229,8 +235,8 @@ int main(int argc, char *argv[])
 			break;
 		case OPT_NUM_THREADS:
 			num_threads = atoi(optarg);
-			if (num_threads < 0 || num_threads > 64) {
-				fprintf(stderr, "bad num_threads=%d (must be >= 0 and < 64)\n", num_threads);
+			if (num_threads < 0) {
+				fprintf(stderr, "bad num_threads=%d (must be >= 0)\n", num_threads);
 				display_usage(opt == 'h' ? stdout : stderr, argv[0]);
 				return EXIT_FAILURE;
 			}
@@ -275,8 +281,8 @@ int main(int argc, char *argv[])
 	}
 
 	if (do_test) {
-		test_worker_threads();
-		test_thread_join();
+		test_worker_threads(num_threads);
+		test_thread_join(num_threads);
 		return EXIT_SUCCESS;
 	}