Geometry Nodes: refactor array devirtualization

Goals:
* Better high level control over where devirtualization occurs. There is always
  a trade-off between performance and compile-time/binary-size.
* Simplify using array devirtualization.
* Better performance for cases where devirtualization wasn't used before.

Many geometry nodes accept fields as inputs. Internally, that means that the
execution functions have to accept so called "virtual arrays" as inputs. Those
 can be e.g. actual arrays, just single values, or lazily computed arrays.
Due to these different possible virtual arrays implementations, access to
individual elements is slower than it would be if everything was just a normal
array (access does through a virtual function call). For more complex execution
functions, this overhead does not matter, but for small functions (like a simple
addition) it very much does. The virtual function call also prevents the compiler
from doing some optimizations (e.g. loop unrolling and inserting simd instructions).

The solution is to "devirtualize" the virtual arrays for small functions where the
overhead is measurable. Essentially, the function is generated many times with
different array types as input. Then there is a run-time dispatch that calls the
best implementation. We have been doing devirtualization in e.g. math nodes
for a long time already. This patch just generalizes the concept and makes it
easier to control. It also makes it easier to investigate the different trade-offs
when it comes to devirtualization.

Nodes that we've optimized using devirtualization before didn't get a speedup.
However, a couple of nodes are using devirtualization now, that didn't before.
Those got a 2-4x speedup in common cases.
* Map Range
* Random Value
* Switch
* Combine XYZ

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

5 files changed, 31 insertions, 42 deletions

diff --git a/source/blender/nodes/geometry/nodes/node_geo_curve_resample.cc b/source/blender/nodes/geometry/nodes/node_geo_curve_resample.cc
index 121ee015d2d..e0348f27e51 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_curve_resample.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_curve_resample.cc
@@ -551,8 +551,10 @@ static Field<int> get_curve_count_field(GeoNodeExecParams params,
                                         const GeometryNodeCurveResampleMode mode)
 {
   if (mode == GEO_NODE_CURVE_RESAMPLE_COUNT) {
-    static fn::CustomMF_SI_SO<int, int> max_one_fn("Clamp Above One",
-                                                   [](int value) { return std::max(1, value); });
+    static fn::CustomMF_SI_SO<int, int> max_one_fn(
+        "Clamp Above One",
+        [](int value) { return std::max(1, value); },
+        fn::CustomMF_presets::AllSpanOrSingle());
     auto clamp_op = std::make_shared<FieldOperation>(
         FieldOperation(max_one_fn, {Field<int>(params.extract_input<Field<int>>("Count"))}));
 
@@ -561,12 +563,14 @@ static Field<int> get_curve_count_field(GeoNodeExecParams params,
 
   if (mode == GEO_NODE_CURVE_RESAMPLE_LENGTH) {
     static fn::CustomMF_SI_SI_SO<float, float, int> get_count_fn(
-        "Length Input to Count", [](const float curve_length, const float sample_length) {
+        "Length Input to Count",
+        [](const float curve_length, const float sample_length) {
           /* Find the number of sampled segments by dividing the total length by
            * the sample length. Then there is one more sampled point than segment. */
           const int count = int(curve_length / sample_length) + 1;
           return std::max(1, count);
-        });
+        },
+        fn::CustomMF_presets::AllSpanOrSingle());
 
     auto get_count_op = std::make_shared<FieldOperation>(
         FieldOperation(get_count_fn,
@@ -588,9 +592,11 @@ static Field<int> get_curve_count_field(GeoNodeExecParams params,
 static Field<bool> get_selection_field(GeoNodeExecParams params)
 {
   static fn::CustomMF_SI_SI_SO<bool, int, bool> get_selection_fn(
-      "Create Curve Selection", [](const bool orig_selection, const int evaluated_points_num) {
+      "Create Curve Selection",
+      [](const bool orig_selection, const int evaluated_points_num) {
         return orig_selection && evaluated_points_num > 1;
-      });
+      },
+      fn::CustomMF_presets::AllSpanOrSingle());
 
   auto selection_op = std::make_shared<FieldOperation>(
       FieldOperation(get_selection_fn,
diff --git a/source/blender/nodes/geometry/nodes/node_geo_curve_sample.cc b/source/blender/nodes/geometry/nodes/node_geo_curve_sample.cc
index 6661d03a851..9077b59254c 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_curve_sample.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_curve_sample.cc
@@ -203,9 +203,11 @@ static Field<float> get_length_input_field(const GeoNodeExecParams &params,
     /* Just make sure the length is in bounds of the curve. */
     Field<float> length_field = params.get_input<Field<float>>("Length");
     auto clamp_fn = std::make_unique<fn::CustomMF_SI_SO<float, float>>(
-        __func__, [curve_total_length](float length) {
+        __func__,
+        [curve_total_length](float length) {
           return std::clamp(length, 0.0f, curve_total_length);
-        });
+        },
+        fn::CustomMF_presets::AllSpanOrSingle());
     auto clamp_op = std::make_shared<FieldOperation>(
         FieldOperation(std::move(clamp_fn), {std::move(length_field)}));
 
@@ -215,10 +217,12 @@ static Field<float> get_length_input_field(const GeoNodeExecParams &params,
   /* Convert the factor to a length and clamp it to the bounds of the curve. */
   Field<float> factor_field = params.get_input<Field<float>>("Factor");
   auto clamp_fn = std::make_unique<fn::CustomMF_SI_SO<float, float>>(
-      __func__, [curve_total_length](float factor) {
+      __func__,
+      [curve_total_length](float factor) {
         const float length = factor * curve_total_length;
         return std::clamp(length, 0.0f, curve_total_length);
-      });
+      },
+      fn::CustomMF_presets::AllSpanOrSingle());
   auto process_op = std::make_shared<FieldOperation>(
       FieldOperation(std::move(clamp_fn), {std::move(factor_field)}));
 
diff --git a/source/blender/nodes/geometry/nodes/node_geo_extrude_mesh.cc b/source/blender/nodes/geometry/nodes/node_geo_extrude_mesh.cc
index 0072fbcde93..4591cfa1da2 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_extrude_mesh.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_extrude_mesh.cc
@@ -1274,7 +1274,9 @@ static void node_geo_exec(GeoNodeExecParams params)
   /* Create a combined field from the offset and the scale so the field evaluator
    * can take care of the multiplication and to simplify each extrude function. */
   static fn::CustomMF_SI_SI_SO<float3, float, float3> multiply_fn{
-      "Scale", [](const float3 &offset, const float scale) { return offset * scale; }};
+      "Scale",
+      [](const float3 &offset, const float scale) { return offset * scale; },
+      fn::CustomMF_presets::AllSpanOrSingle()};
   std::shared_ptr<FieldOperation> multiply_op = std::make_shared<FieldOperation>(
       FieldOperation(multiply_fn, {std::move(offset_field), std::move(scale_field)}));
   const Field<float3> final_offset{std::move(multiply_op)};
diff --git a/source/blender/nodes/geometry/nodes/node_geo_mesh_to_points.cc b/source/blender/nodes/geometry/nodes/node_geo_mesh_to_points.cc
index 6a7aab1cea7..1a0cc53cb6c 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_mesh_to_points.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_mesh_to_points.cc
@@ -124,7 +124,9 @@ static void node_geo_exec(GeoNodeExecParams params)
   /* Use another multi-function operation to make sure the input radius is greater than zero.
    * TODO: Use mutable multi-function once that is supported. */
   static fn::CustomMF_SI_SO<float, float> max_zero_fn(
-      __func__, [](float value) { return std::max(0.0f, value); });
+      __func__,
+      [](float value) { return std::max(0.0f, value); },
+      fn::CustomMF_presets::AllSpanOrSingle());
   auto max_zero_op = std::make_shared<FieldOperation>(
       FieldOperation(max_zero_fn, {std::move(radius)}));
   Field<float> positive_radius(std::move(max_zero_op), 0);
diff --git a/source/blender/nodes/geometry/nodes/node_geo_switch.cc b/source/blender/nodes/geometry/nodes/node_geo_switch.cc
index 6e27b346407..ddc87e3dac4 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_switch.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_switch.cc
@@ -139,35 +139,6 @@ static void node_gather_link_searches(GatherLinkSearchOpParams &params)
   }
 }
 
-template<typename T> class SwitchFieldsFunction : public fn::MultiFunction {
- public:
-  SwitchFieldsFunction()
-  {
-    static fn::MFSignature signature = create_signature();
-    this->set_signature(&signature);
-  }
-  static fn::MFSignature create_signature()
-  {
-    fn::MFSignatureBuilder signature{"Switch"};
-    signature.single_input<bool>("Switch");
-    signature.single_input<T>("False");
-    signature.single_input<T>("True");
-    signature.single_output<T>("Output");
-    return signature.build();
-  }
-
-  void call(IndexMask mask, fn::MFParams params, fn::MFContext UNUSED(context)) const override
-  {
-    const VArray<bool> &switches = params.readonly_single_input<bool>(0, "Switch");
-    const VArray<T> &falses = params.readonly_single_input<T>(1, "False");
-    const VArray<T> &trues = params.readonly_single_input<T>(2, "True");
-    MutableSpan<T> values = params.uninitialized_single_output_if_required<T>(3, "Output");
-    for (int64_t i : mask) {
-      new (&values[i]) T(switches[i] ? trues[i] : falses[i]);
-    }
-  }
-};
-
 template<typename T> void switch_fields(GeoNodeExecParams &params, const StringRef suffix)
 {
   if (params.lazy_require_input("Switch")) {
@@ -190,7 +161,11 @@ template<typename T> void switch_fields(GeoNodeExecParams &params, const StringR
     Field<T> falses_field = params.extract_input<Field<T>>(name_false);
     Field<T> trues_field = params.extract_input<Field<T>>(name_true);
 
-    auto switch_fn = std::make_unique<SwitchFieldsFunction<T>>();
+    static fn::CustomMF_SI_SI_SI_SO<bool, T, T, T> switch_fn{
+        "Switch", [](bool condition, const T &false_value, const T &true_value) {
+          return condition ? true_value : false_value;
+        }};
+
     auto switch_op = std::make_shared<FieldOperation>(FieldOperation(
         std::move(switch_fn),
         {std::move(switches_field), std::move(falses_field), std::move(trues_field)}));