Geometry Nodes: optimize Set Position node

This implements four optimizations in the Set Position node:
* Check whether the position input is the current position and ignore
  it if it is. This results in a speedup when only the Offset input is used.
* Use multi-threading when copying to computed values to the
  position attribute. All geometry types benefit from this.
* Use devirtualization for the offset and position input. This optimizes
  the common case that they are either single values or computed
  in the fly in a span.
* Write to `Mesh->mvert` directly instead of creating a temporary span.
  This makes setting mesh vertex positions even more efficient.

In my simple benchmark I'm using a White Noise node to offset the
position of 1,000,000 vertices. The speed is `20 ms -> 4.5 ms` in the
multi-threaded case and `32 ms -> 22 ms` in the single-threaded case.

1 files changed, 77 insertions, 12 deletions

diff --git a/source/blender/nodes/geometry/nodes/node_geo_set_position.cc b/source/blender/nodes/geometry/nodes/node_geo_set_position.cc
index a8e59a807e7..218e4d03464 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_set_position.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_set_position.cc
@@ -16,6 +16,11 @@
 
 #include "DEG_depsgraph_query.h"
 
+#include "BLI_task.hh"
+
+#include "DNA_mesh_types.h"
+#include "DNA_meshdata_types.h"
+
 #include "node_geometry_util.hh"
 
 namespace blender::nodes::node_geo_set_position_cc {
@@ -29,6 +34,77 @@ static void node_declare(NodeDeclarationBuilder &b)
   b.add_output<decl::Geometry>(N_("Geometry"));
 }
 
+static void set_computed_position_and_offset(GeometryComponent &component,
+                                             const VArray<float3> &in_positions,
+                                             const VArray<float3> &in_offsets,
+                                             const AttributeDomain domain,
+                                             const IndexMask selection)
+{
+
+  OutputAttribute_Typed<float3> positions = component.attribute_try_get_for_output<float3>(
+      "position", domain, {0, 0, 0});
+
+  const int grain_size = 10000;
+
+  switch (component.type()) {
+    case GEO_COMPONENT_TYPE_MESH: {
+      Mesh *mesh = static_cast<MeshComponent &>(component).get_for_write();
+      MutableSpan<MVert> mverts{mesh->mvert, mesh->totvert};
+      if (in_positions.is_same(positions.varray())) {
+        devirtualize_varray(in_offsets, [&](const auto in_offsets) {
+          threading::parallel_for(
+              selection.index_range(), grain_size, [&](const IndexRange range) {
+                for (const int i : selection.slice(range)) {
+                  const float3 offset = in_offsets[i];
+                  add_v3_v3(mverts[i].co, offset);
+                }
+              });
+        });
+      }
+      else {
+        devirtualize_varray2(
+            in_positions, in_offsets, [&](const auto in_positions, const auto in_offsets) {
+              threading::parallel_for(
+                  selection.index_range(), grain_size, [&](const IndexRange range) {
+                    for (const int i : selection.slice(range)) {
+                      const float3 new_position = in_positions[i] + in_offsets[i];
+                      copy_v3_v3(mverts[i].co, new_position);
+                    }
+                  });
+            });
+      }
+      break;
+    }
+    default: {
+      MutableSpan<float3> out_positions_span = positions.as_span();
+      if (in_positions.is_same(positions.varray())) {
+        devirtualize_varray(in_offsets, [&](const auto in_offsets) {
+          threading::parallel_for(
+              selection.index_range(), grain_size, [&](const IndexRange range) {
+                for (const int i : selection.slice(range)) {
+                  out_positions_span[i] += in_offsets[i];
+                }
+              });
+        });
+      }
+      else {
+        devirtualize_varray2(
+            in_positions, in_offsets, [&](const auto in_positions, const auto in_offsets) {
+              threading::parallel_for(
+                  selection.index_range(), grain_size, [&](const IndexRange range) {
+                    for (const int i : selection.slice(range)) {
+                      out_positions_span[i] = in_positions[i] + in_offsets[i];
+                    }
+                  });
+            });
+      }
+      break;
+    }
+  }
+
+  positions.save();
+}
+
 static void set_position_in_component(GeometryComponent &component,
                                       const Field<bool> &selection_field,
                                       const Field<float3> &position_field,
@@ -53,20 +129,9 @@ static void set_position_in_component(GeometryComponent &component,
   position_evaluator.add(offset_field);
   position_evaluator.evaluate();
 
-  /* TODO: We could have different code paths depending on whether the offset input is a single
-   * value or not */
-
   const VArray<float3> &positions_input = position_evaluator.get_evaluated<float3>(0);
   const VArray<float3> &offsets_input = position_evaluator.get_evaluated<float3>(1);
-
-  OutputAttribute_Typed<float3> positions = component.attribute_try_get_for_output<float3>(
-      "position", domain, {0, 0, 0});
-  MutableSpan<float3> position_mutable = positions.as_span();
-
-  for (int i : selection) {
-    position_mutable[i] = positions_input[i] + offsets_input[i];
-  }
-  positions.save();
+  set_computed_position_and_offset(component, positions_input, offsets_input, domain, selection);
 }
 
 static void node_geo_exec(GeoNodeExecParams params)