BLI_task: add TLS support to BLI_task_parallel_mempool

Support thread local storage for BLI_task_parallel_mempool,
as well as support for the reduce and free callbacks.

mempool_iter_threadsafe_* functions have been moved into a private
header thats only shared between task_iterator.c and BLI_mempool.c
so the TLS can be made part of the iterator array without having to
rely on passing in struct offsets.

Add test task.MempoolIterTLS that ensures reduce and free
are working as expected.

Reviewed By: mont29

Ref D11548

1 files changed, 63 insertions, 18 deletions

diff --git a/source/blender/blenlib/intern/task_iterator.c b/source/blender/blenlib/intern/task_iterator.c
index 463d3675cbe..3914a074d11 100644
--- a/source/blender/blenlib/intern/task_iterator.c
+++ b/source/blender/blenlib/intern/task_iterator.c
@@ -29,6 +29,7 @@
 #include "BLI_listbase.h"
 #include "BLI_math.h"
 #include "BLI_mempool.h"
+#include "BLI_mempool_private.h"
 #include "BLI_task.h"
 #include "BLI_threads.h"
 
@@ -348,9 +349,6 @@ void BLI_task_parallel_listbase(ListBase *listbase,
   task_parallel_iterator_do(settings, &state);
 }
 
-#undef MALLOCA
-#undef MALLOCA_FREE
-
 typedef struct ParallelMempoolState {
   void *userdata;
   TaskParallelMempoolFunc func;
@@ -359,11 +357,12 @@ typedef struct ParallelMempoolState {
 static void parallel_mempool_func(TaskPool *__restrict pool, void *taskdata)
 {
   ParallelMempoolState *__restrict state = BLI_task_pool_user_data(pool);
-  BLI_mempool_iter *iter = taskdata;
-  MempoolIterData *item;
+  BLI_mempool_iter *iter = &((ParallelMempoolTaskData *)taskdata)->iter;
+  TaskParallelTLS *tls = &((ParallelMempoolTaskData *)taskdata)->tls;
 
+  MempoolIterData *item;
   while ((item = BLI_mempool_iterstep(iter)) != NULL) {
-    state->func(state->userdata, item);
+    state->func(state->userdata, item, tls);
   }
 }
 
@@ -373,16 +372,14 @@ static void parallel_mempool_func(TaskPool *__restrict pool, void *taskdata)
  * \param mempool: The iterable BLI_mempool to loop over.
  * \param userdata: Common userdata passed to all instances of \a func.
  * \param func: Callback function.
- * \param use_threading: If \a true, actually split-execute loop in threads,
- * else just do a sequential for loop
- * (allows caller to use any kind of test to switch on parallelization or not).
+ * \param settings: See public API doc of TaskParallelSettings for description of all settings.
  *
  * \note There is no static scheduling here.
  */
 void BLI_task_parallel_mempool(BLI_mempool *mempool,
                                void *userdata,
                                TaskParallelMempoolFunc func,
-                               const bool use_threading)
+                               const TaskParallelSettings *settings)
 {
   TaskPool *task_pool;
   ParallelMempoolState state;
@@ -392,14 +389,34 @@ void BLI_task_parallel_mempool(BLI_mempool *mempool,
     return;
   }
 
-  if (!use_threading) {
+  void *userdata_chunk = settings->userdata_chunk;
+  const size_t userdata_chunk_size = settings->userdata_chunk_size;
+  void *userdata_chunk_local = NULL;
+  void *userdata_chunk_array = NULL;
+  const bool use_userdata_chunk = (userdata_chunk_size != 0) && (userdata_chunk != NULL);
+
+  if (!settings->use_threading) {
+    TaskParallelTLS tls = {NULL};
+    if (use_userdata_chunk) {
+      userdata_chunk_local = MALLOCA(userdata_chunk_size);
+      memcpy(userdata_chunk_local, userdata_chunk, userdata_chunk_size);
+      tls.userdata_chunk = userdata_chunk_local;
+    }
+
     BLI_mempool_iter iter;
     BLI_mempool_iternew(mempool, &iter);
 
-    for (void *item = BLI_mempool_iterstep(&iter); item != NULL;
-         item = BLI_mempool_iterstep(&iter)) {
-      func(userdata, item);
+    void *item;
+    while ((item = BLI_mempool_iterstep(&iter))) {
+      func(userdata, item, &tls);
     }
+
+    if (settings->func_free != NULL) {
+      /* `func_free` should only free data that was created during execution of `func`. */
+      settings->func_free(userdata, userdata_chunk_local);
+    }
+
+    MALLOCA_FREE(userdata_chunk_local, userdata_chunk_size);
     return;
   }
 
@@ -415,16 +432,44 @@ void BLI_task_parallel_mempool(BLI_mempool *mempool,
   state.userdata = userdata;
   state.func = func;
 
-  BLI_mempool_iter *mempool_iterators = BLI_mempool_iter_threadsafe_create(mempool,
-                                                                           (size_t)num_tasks);
+  if (use_userdata_chunk) {
+    userdata_chunk_array = MALLOCA(userdata_chunk_size * num_tasks);
+  }
+
+  ParallelMempoolTaskData *mempool_iterator_data = mempool_iter_threadsafe_create(
+      mempool, (size_t)num_tasks);
 
   for (i = 0; i < num_tasks; i++) {
+    if (use_userdata_chunk) {
+      userdata_chunk_local = (char *)userdata_chunk_array + (userdata_chunk_size * i);
+      memcpy(userdata_chunk_local, userdata_chunk, userdata_chunk_size);
+    }
+    mempool_iterator_data[i].tls.userdata_chunk = userdata_chunk_local;
+
     /* Use this pool's pre-allocated tasks. */
-    BLI_task_pool_push(task_pool, parallel_mempool_func, &mempool_iterators[i], false, NULL);
+    BLI_task_pool_push(task_pool, parallel_mempool_func, &mempool_iterator_data[i], false, NULL);
   }
 
   BLI_task_pool_work_and_wait(task_pool);
   BLI_task_pool_free(task_pool);
 
-  BLI_mempool_iter_threadsafe_free(mempool_iterators);
+  if (use_userdata_chunk) {
+    if ((settings->func_free != NULL) || (settings->func_reduce != NULL)) {
+      for (i = 0; i < num_tasks; i++) {
+        if (settings->func_reduce) {
+          settings->func_reduce(
+              userdata, userdata_chunk, mempool_iterator_data[i].tls.userdata_chunk);
+        }
+        if (settings->func_free) {
+          settings->func_free(userdata, mempool_iterator_data[i].tls.userdata_chunk);
+        }
+      }
+    }
+    MALLOCA_FREE(userdata_chunk_array, userdata_chunk_size * num_tasks);
+  }
+
+  mempool_iter_threadsafe_destroy(mempool_iterator_data);
 }
+
+#undef MALLOCA
+#undef MALLOCA_FREE