Geometry Nodes: multi threaded field evaluation

This adds a new `ParallelMultiFunction` which wraps another multi-function
and evaluates it with multiple threads. The speeds up field evaluation
quite a bit (the effect is most noticeable when the number of evaluations
and the field is large).

There are still other single-threaded performance bottlenecks in field
evaluation that will need to be solved separately. Most notably here
is the process of copying the computed data into the position attribute
in the Set Position node.

Differential Revision: https://developer.blender.org/D12457

3 files changed, 147 insertions, 2 deletions

diff --git a/source/blender/functions/intern/field.cc b/source/blender/functions/intern/field.cc
index 43f28efd002..6a4518ad4a6 100644
--- a/source/blender/functions/intern/field.cc
+++ b/source/blender/functions/intern/field.cc
@@ -21,6 +21,7 @@
 #include "BLI_vector_set.hh"
 
 #include "FN_field.hh"
+#include "FN_multi_function_parallel.hh"
 
 namespace blender::fn {
 
@@ -360,7 +361,13 @@ Vector<const GVArray *> evaluate_fields(ResourceScope &scope,
     build_multi_function_procedure_for_fields(
         procedure, scope, field_tree_info, varying_fields_to_evaluate);
     MFProcedureExecutor procedure_executor{"Procedure", procedure};
-    MFParamsBuilder mf_params{procedure_executor, &mask};
+    /* Add multi threading capabilities to the field evaluation. */
+    const int grain_size = 10000;
+    fn::ParallelMultiFunction parallel_procedure_executor{procedure_executor, grain_size};
+    /* Utility variable to make easy to switch the executor. */
+    const MultiFunction &executor_fn = parallel_procedure_executor;
+
+    MFParamsBuilder mf_params{executor_fn, &mask};
     MFContextBuilder mf_context;
 
     /* Provide inputs to the procedure executor. */
@@ -401,7 +408,7 @@ Vector<const GVArray *> evaluate_fields(ResourceScope &scope,
       mf_params.add_uninitialized_single_output(span);
     }
 
-    procedure_executor.call(mask, mf_params, mf_context);
+    executor_fn.call(mask, mf_params, mf_context);
   }
 
   /* Evaluate constant fields if necessary. */
diff --git a/source/blender/functions/intern/generic_virtual_array.cc b/source/blender/functions/intern/generic_virtual_array.cc
index bd033a429de..9a83d8cd497 100644
--- a/source/blender/functions/intern/generic_virtual_array.cc
+++ b/source/blender/functions/intern/generic_virtual_array.cc
@@ -387,4 +387,47 @@ void GVMutableArray_GSpan::disable_not_applied_warning()
   show_not_saved_warning_ = false;
 }
 
+/* --------------------------------------------------------------------
+ * GVArray_For_SlicedGVArray.
+ */
+
+void GVArray_For_SlicedGVArray::get_impl(const int64_t index, void *r_value) const
+{
+  varray_.get(index + offset_, r_value);
+}
+
+void GVArray_For_SlicedGVArray::get_to_uninitialized_impl(const int64_t index, void *r_value) const
+{
+  varray_.get_to_uninitialized(index + offset_, r_value);
+}
+
+/* --------------------------------------------------------------------
+ * GVArray_Slice.
+ */
+
+GVArray_Slice::GVArray_Slice(const GVArray &varray, const IndexRange slice)
+{
+  if (varray.is_span()) {
+    /* Create a new virtual for the sliced span. */
+    const GSpan span = varray.get_internal_span();
+    const GSpan sliced_span = span.slice(slice.start(), slice.size());
+    varray_span_.emplace(sliced_span);
+    varray_ = &*varray_span_;
+  }
+  else if (varray.is_single()) {
+    /* Can just use the existing virtual array, because it's the same value for the indices in the
+     * slice anyway. */
+    varray_ = &varray;
+  }
+  else {
+    /* Generic version when none of the above method works.
+     * We don't necessarily want to materialize the input varray because there might be
+     * large distances between the required indices. Then we would materialize many elements that
+     * are not accessed later on.
+     */
+    varray_any_.emplace(varray, slice);
+    varray_ = &*varray_any_;
+  }
+}
+
 }  // namespace blender::fn
diff --git a/source/blender/functions/intern/multi_function_parallel.cc b/source/blender/functions/intern/multi_function_parallel.cc
new file mode 100644
index 00000000000..5a8c621f0b3
--- /dev/null
+++ b/source/blender/functions/intern/multi_function_parallel.cc
@@ -0,0 +1,95 @@
+/*
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+ */
+
+#include "FN_multi_function_parallel.hh"
+
+#include "BLI_task.hh"
+
+namespace blender::fn {
+
+ParallelMultiFunction::ParallelMultiFunction(const MultiFunction &fn, const int64_t grain_size)
+    : fn_(fn), grain_size_(grain_size)
+{
+  this->set_signature(&fn.signature());
+
+  threading_supported_ = true;
+  for (const int param_index : fn.param_indices()) {
+    const MFParamType param_type = fn.param_type(param_index);
+    if (param_type.data_type().category() == MFDataType::Vector) {
+      /* Vector parameters do not support threading yet. */
+      threading_supported_ = false;
+      break;
+    }
+  }
+}
+
+void ParallelMultiFunction::call(IndexMask full_mask, MFParams params, MFContext context) const
+{
+  if (full_mask.size() <= grain_size_ || !threading_supported_) {
+    fn_.call(full_mask, params, context);
+    return;
+  }
+
+  threading::parallel_for(full_mask.index_range(), grain_size_, [&](const IndexRange mask_slice) {
+    Vector<int64_t> sub_mask_indices;
+    const IndexMask sub_mask = full_mask.slice_and_offset(mask_slice, sub_mask_indices);
+    if (sub_mask.is_empty()) {
+      return;
+    }
+    const int64_t input_slice_start = full_mask[mask_slice.first()];
+    const int64_t input_slice_size = full_mask[mask_slice.last()] - input_slice_start + 1;
+    const IndexRange input_slice_range{input_slice_start, input_slice_size};
+
+    MFParamsBuilder sub_params{fn_, sub_mask.min_array_size()};
+    ResourceScope &scope = sub_params.resource_scope();
+
+    /* All parameters are sliced so that the wrapped multi-function does not have to take care of
+     * the index offset. */
+    for (const int param_index : fn_.param_indices()) {
+      const MFParamType param_type = fn_.param_type(param_index);
+      switch (param_type.category()) {
+        case MFParamType::SingleInput: {
+          const GVArray &varray = params.readonly_single_input(param_index);
+          const GVArray &sliced_varray = scope.construct<GVArray_Slice>(varray, input_slice_range);
+          sub_params.add_readonly_single_input(sliced_varray);
+          break;
+        }
+        case MFParamType::SingleMutable: {
+          const GMutableSpan span = params.single_mutable(param_index);
+          const GMutableSpan sliced_span = span.slice(input_slice_start, input_slice_size);
+          sub_params.add_single_mutable(sliced_span);
+          break;
+        }
+        case MFParamType::SingleOutput: {
+          const GMutableSpan span = params.uninitialized_single_output(param_index);
+          const GMutableSpan sliced_span = span.slice(input_slice_start, input_slice_size);
+          sub_params.add_uninitialized_single_output(sliced_span);
+          break;
+        }
+        case MFParamType::VectorInput:
+        case MFParamType::VectorMutable:
+        case MFParamType::VectorOutput: {
+          BLI_assert_unreachable();
+          break;
+        }
+      }
+    }
+
+    fn_.call(sub_mask, sub_params, context);
+  });
+}
+
+}  // namespace blender::fn