BLI_task: add a callback to initialize TLS

Useful when TLS requires it's own allocated structures.

2 files changed, 19 insertions, 0 deletions

diff --git a/source/blender/blenlib/BLI_task.h b/source/blender/blenlib/BLI_task.h
index dbe8ec3dcc0..418db14e2f3 100644
--- a/source/blender/blenlib/BLI_task.h
+++ b/source/blender/blenlib/BLI_task.h
@@ -129,6 +129,9 @@ typedef struct TaskParallelTLS {
 typedef void (*TaskParallelRangeFunc)(void *__restrict userdata,
                                       const int iter,
                                       const TaskParallelTLS *__restrict tls);
+
+typedef void (*TaskParallelInitFunc)(const void *__restrict userdata, void *__restrict chunk);
+
 typedef void (*TaskParallelReduceFunc)(const void *__restrict userdata,
                                        void *__restrict chunk_join,
                                        void *__restrict chunk);
@@ -151,6 +154,10 @@ typedef struct TaskParallelSettings {
   /* Function called from calling thread once whole range have been
    * processed.
    */
+  /* Function called to initialize user data chunk,
+   * typically to allocate data, freed by `func_free`.
+   */
+  TaskParallelInitFunc func_init;
   /* Function called to join user data chunk into another, to reduce
    * the result to the original userdata_chunk memory.
    * The reduce functions should have no side effects, so that they
diff --git a/source/blender/blenlib/intern/task_iterator.c b/source/blender/blenlib/intern/task_iterator.c
index 0ff408ddb0a..33af4894b48 100644
--- a/source/blender/blenlib/intern/task_iterator.c
+++ b/source/blender/blenlib/intern/task_iterator.c
@@ -186,6 +186,9 @@ static void task_parallel_iterator_no_threads(const TaskParallelSettings *settin
   if (use_userdata_chunk) {
     userdata_chunk_local = MALLOCA(userdata_chunk_size);
     memcpy(userdata_chunk_local, userdata_chunk, userdata_chunk_size);
+    if (settings->func_init != NULL) {
+      settings->func_init(state->userdata, userdata_chunk_local);
+    }
   }
 
   /* Also marking it as non-threaded for the iterator callback. */
@@ -247,6 +250,9 @@ static void task_parallel_iterator_do(const TaskParallelSettings *settings,
     if (use_userdata_chunk) {
       userdata_chunk_local = (char *)userdata_chunk_array + (userdata_chunk_size * i);
       memcpy(userdata_chunk_local, userdata_chunk, userdata_chunk_size);
+      if (settings->func_init != NULL) {
+        settings->func_init(state->userdata, userdata_chunk_local);
+      }
     }
     /* Use this pool's pre-allocated tasks. */
     BLI_task_pool_push(task_pool, parallel_iterator_func, userdata_chunk_local, false, NULL);
@@ -422,6 +428,9 @@ void BLI_task_parallel_mempool(BLI_mempool *mempool,
     if (use_userdata_chunk) {
       userdata_chunk_local = MALLOCA(userdata_chunk_size);
       memcpy(userdata_chunk_local, userdata_chunk, userdata_chunk_size);
+      if (settings->func_init != NULL) {
+        settings->func_init(state.userdata, userdata_chunk_local);
+      }
       tls.userdata_chunk = userdata_chunk_local;
     }
 
@@ -465,6 +474,9 @@ void BLI_task_parallel_mempool(BLI_mempool *mempool,
     if (use_userdata_chunk) {
       userdata_chunk_local = (char *)userdata_chunk_array + (userdata_chunk_size * i);
       memcpy(userdata_chunk_local, userdata_chunk, userdata_chunk_size);
+      if (settings->func_init != NULL) {
+        settings->func_init(userdata, userdata_chunk_local);
+      }
     }
     mempool_iterator_data[i].tls.userdata_chunk = userdata_chunk_local;