source/blender/nodes/geometry/nodes/node_geo_index_of_nearest.cc

@@ -1,7 +1,5 @@
 /* SPDX-License-Identifier: GPL-2.0-or-later */
 
-#include "BKE_attribute_math.hh"
-
 #include "BLI_kdtree.h"
 #include "BLI_multi_value_map.hh"
 #include "BLI_task.hh"
@@ -19,6 +17,55 @@ static void node_declare(NodeDeclarationBuilder &b)
   b.add_output<decl::Bool>(N_("Has Neighbor")).field_source();
 }
 
+static KDTree_3d *build_kdtree(const Span<float3> &positions, const IndexMask mask)
+{
+  KDTree_3d *tree = BLI_kdtree_3d_new(mask.size());
+  for (const int i : mask) {
+    BLI_kdtree_3d_insert(tree, i, positions[i]);
+  }
+  BLI_kdtree_3d_balance(tree);
+  return tree;
+}
+
+static int find_nearest_non_self(const KDTree_3d &tree, const float3 &position, const int index)
+{
+  return BLI_kdtree_3d_find_nearest_cb_cpp(
+      &tree, position, 0, [index](const int other, const float * /*co*/, const float /*dist_sq*/) {
+        return index == other ? 0 : 1;
+      });
+}
+
+static void find_neighbors(const KDTree_3d &tree,
+                           const Span<float3> positions,
+                           const IndexMask mask,
+                           MutableSpan<int> indices)
+{
+  threading::parallel_for(mask.index_range(), 1024, [&](const IndexRange range) {
+    for (const int i : mask.slice(range)) {
+      indices[i] = find_nearest_non_self(tree, positions[i], i);
+    }
+  });
+}
+
+static Vector<IndexMask> masks_from_group_ids(const Span<int> group_ids,
+                                              const IndexMask mask,
+                                              MultiValueMap<int, int64_t> &storage)
+{
+  mask.foreach_index([&](const int i) { storage.add(group_ids[i], i); });
+  Vector<IndexMask> masks;
+  masks.reserve(storage.size());
+  for (const Span<int64_t> indices : storage.values()) {
+    masks.append(indices);
+  }
+  return masks;
+}
+
+static Vector<IndexMask> masks_from_group_ids(const Span<int> group_ids,
+                                              MultiValueMap<int, int64_t> &storage)
+{
+  return masks_from_group_ids(group_ids, group_ids.index_range(), storage);
+}
+
 class IndexOfNearestFieldInput final : public bke::GeometryFieldInput {
  private:
   const Field<float3> positions_field_;
@@ -26,70 +73,70 @@ class IndexOfNearestFieldInput final : public bke::GeometryFieldInput {
 
  public:
   IndexOfNearestFieldInput(Field<float3> positions_field, Field<int> group_field)
-      : bke::GeometryFieldInput(CPPType::get<int>(), "Nearest to"),
+      : bke::GeometryFieldInput(CPPType::get<int>(), "Index of Nearest"),
         positions_field_(std::move(positions_field)),
         group_field_(std::move(group_field))
   {
   }
 
   GVArray get_varray_for_context(const bke::GeometryFieldContext &context,
-                                 IndexMask mask) const final
+                                 const IndexMask mask) const final
   {
-    fn::FieldEvaluator evaluator{context, &mask};
+    if (!context.attributes()) {
+      return {};
+    }
+    const int domain_size = context.attributes()->domain_size(context.domain());
+    fn::FieldEvaluator evaluator{context, domain_size};
     evaluator.add(positions_field_);
     evaluator.add(group_field_);
     evaluator.evaluate();
+    const VArraySpan<float3> positions = evaluator.get_evaluated<float3>(0);
+    const VArray<int> group = evaluator.get_evaluated<int>(1);
 
-    const VArray<float3> &positions = evaluator.get_evaluated<float3>(0);
-    const VArray<int> &group = evaluator.get_evaluated<int>(1);
+    Array<int> result(mask.min_array_size());
 
-    MultiValueMap<int, int64_t> group_masks;
-    mask.foreach_index([&](const int index) { group_masks.add(group[index], index); });
-
-    Array<int> indices(mask.min_array_size());
-
-    const auto nearest_for = [this, &positions](const IndexMask mask, MutableSpan<int> r_indices) {
-      devirtualize_varray(positions, [mask, r_indices, this](const auto positions) {
-        KDTree_3d *tree = BLI_kdtree_3d_new(mask.size());
-        mask.foreach_index([tree, positions](const int index) {
-          BLI_kdtree_3d_insert(tree, index, positions[index]);
-        });
-
-        BLI_kdtree_3d_balance(tree);
-
-        threading::parallel_for(mask.index_range(), 512, [&](const IndexRange range) {
-          mask.slice(range).foreach_index([&](const auto index) {
-            r_indices[index] = this->kdtree_find_neighboard(tree, positions[index], index);
-          });
-        });
-
-        BLI_kdtree_3d_free(tree);
-      });
-    };
-
-    for (const Span<int64_t> mask_span : group_masks.values()) {
-      if (mask_span.size() == 1) {
-        indices[mask_span.first()] = -1;
-      }
-      nearest_for(mask_span, indices);
+    if (group.is_single()) {
+      const IndexMask full_mask = positions.index_range();
+      KDTree_3d *tree = build_kdtree(positions, full_mask);
+      find_neighbors(*tree, positions, mask, result);
+      BLI_kdtree_3d_free(tree);
     }
+    else {
+      /* The goal is to build each tree and use it immediately, rather than building all trees and
+       * sampling them later. That should help to keep the tree in caches before balancing and when
+       * sampling many points. */
+      const VArraySpan<int> group_ids(group);
+      MultiValueMap<int, int64_t> group_mask_storage;
+      const Vector<IndexMask> tree_masks = masks_from_group_ids(group_ids, group_mask_storage);
 
-    return VArray<int>::ForContainer(std::move(indices));
-  }
+      MultiValueMap<int, int64_t> evaluate_masks_storage;
+      Vector<IndexMask> evaluate_masks;
+      if (mask.size() < domain_size) {
+        /* Separate masks for evaluation are only necessary if the mask mask
+         * for field input evaluation doesn't have every element selected. */
+        evaluate_masks = masks_from_group_ids(group_ids, mask, evaluate_masks_storage);
+      }
 
- protected:
-  static int kdtree_find_neighboard(KDTree_3d *tree, const float3 &position, const int &index)
-  {
-    return BLI_kdtree_3d_find_nearest_cb_cpp(
-        tree,
-        position,
-        0,
-        [index](const int other_new_i, const float * /*co*/, const float /*dist_sq*/) {
-          if (index == other_new_i) {
-            return 0;
+      /* The grain size should be larger as each tree gets smaller. */
+      const int avg_tree_size = group_ids.size() / group_mask_storage.size();
+      const int grain_size = std::max(8192 / avg_tree_size, 1);
+      threading::parallel_for(tree_masks.index_range(), grain_size, [&](const IndexRange range) {
+        for (const int i : range) {
+          const IndexMask tree_mask = tree_masks[i];
+          const IndexMask evaluate_mask = evaluate_masks.is_empty() ? tree_mask :
+                                                                      evaluate_masks[i];
+          if (tree_masks[i].size() < 2) {
+            result.as_mutable_span().fill_indices(evaluate_mask.indices(), 0);
           }
-          return 1;
-        });
+          else {
+            KDTree_3d *tree = build_kdtree(positions, tree_mask);
+            find_neighbors(*tree, positions, evaluate_mask, result);
+            BLI_kdtree_3d_free(tree);
+          }
+        }
+      });
+    }
+    return VArray<int>::ForContainer(std::move(result));
   }
 
  public:
@@ -99,53 +146,105 @@ class IndexOfNearestFieldInput final : public bke::GeometryFieldInput {
     group_field_.node().for_each_field_input_recursive(fn);
   }
 
-  uint64_t hash() const override
+  uint64_t hash() const final
   {
     return get_default_hash_2(positions_field_, group_field_);
   }
 
-  bool is_equal_to(const fn::FieldNode &other) const override
+  bool is_equal_to(const fn::FieldNode &other) const final
   {
-    if (const IndexOfNearestFieldInput *other_field =
-            dynamic_cast<const IndexOfNearestFieldInput *>(&other)) {
+    if (const auto *other_field = dynamic_cast<const IndexOfNearestFieldInput *>(&other)) {
       return positions_field_ == other_field->positions_field_ &&
              group_field_ == other_field->group_field_;
     }
     return false;
   }
 
-  std::optional<eAttrDomain> preferred_domain(const GeometryComponent &component) const override
+  std::optional<eAttrDomain> preferred_domain(const GeometryComponent &component) const final
   {
     return bke::try_detect_field_domain(component, positions_field_);
   }
 };
 
+class HasNeighborFieldInput final : public bke::GeometryFieldInput {
+ private:
+  const Field<int> group_field_;
+
+ public:
+  HasNeighborFieldInput(Field<int> group_field)
+      : bke::GeometryFieldInput(CPPType::get<bool>(), "Has Neighbor"),
+        group_field_(std::move(group_field))
+  {
+  }
+
+  GVArray get_varray_for_context(const bke::GeometryFieldContext &context,
+                                 const IndexMask mask) const final
+  {
+    if (!context.attributes()) {
+      return {};
+    }
+    const int domain_size = context.attributes()->domain_size(context.domain());
+    fn::FieldEvaluator evaluator{context, domain_size};
+    evaluator.add(group_field_);
+    evaluator.evaluate();
+    const VArray<int> group = evaluator.get_evaluated<int>(0);
+
+    if (group.is_single()) {
+      return VArray<bool>::ForSingle(true, mask.min_array_size());
+    }
+
+    /* When a group ID is contained in the set, it means there is only one element with that ID. */
+    Map<int, int> counts;
+    const VArraySpan<int> group_span(group);
+    mask.foreach_index([&](const int i) {
+      counts.add_or_modify(
+          group_span[i], [](int *count) { *count = 0; }, [](int *count) { (*count)++; });
+    });
+    Array<bool> result(mask.min_array_size());
+    mask.foreach_index([&](const int i) { result[i] = counts.lookup(group_span[i]) > 1; });
+    return VArray<bool>::ForContainer(std::move(result));
+  }
+
+ public:
+  void for_each_field_input_recursive(FunctionRef<void(const FieldInput &)> fn) const
+  {
+    group_field_.node().for_each_field_input_recursive(fn);
+  }
+
+  uint64_t hash() const final
+  {
+    return get_default_hash_2(3984756934876, group_field_);
+  }
+
+  bool is_equal_to(const fn::FieldNode &other) const final
+  {
+    if (const auto *other_field = dynamic_cast<const HasNeighborFieldInput *>(&other)) {
+      return group_field_ == other_field->group_field_;
+    }
+    return false;
+  }
+
+  std::optional<eAttrDomain> preferred_domain(const GeometryComponent &component) const final
+  {
+    return bke::try_detect_field_domain(component, group_field_);
+  }
+};
+
 static void node_geo_exec(GeoNodeExecParams params)
 {
   Field<float3> position_field = params.extract_input<Field<float3>>("Position");
   Field<int> group_field = params.extract_input<Field<int>>("Group ID");
 
-  Field<int> index_of_nearest_field(std::make_shared<IndexOfNearestFieldInput>(
-      std::move(position_field), std::move(group_field)));
-
   if (params.output_is_required("Index")) {
-    static auto clamp_fn = mf::build::SI1_SO<int, int>(
-        "Index Clamping",
-        [](const int index) { return math::max(0, index); },
-        mf::build::exec_presets::Materialized());
-    auto clamp_op = std::make_shared<FieldOperation>(
-        FieldOperation(std::move(clamp_fn), {index_of_nearest_field}));
-    params.set_output("Index", Field<int>(clamp_op, 0));
+    params.set_output("Index",
+                      Field<int>(std::make_shared<IndexOfNearestFieldInput>(
+                          std::move(position_field), group_field)));
   }
 
   if (params.output_is_required("Has Neighbor")) {
-    static auto valid_fn = mf::build::SI1_SO<int, bool>(
-        "Index Validating",
-        [](const int index) { return index != -1; },
-        mf::build::exec_presets::Materialized());
-    auto valid_op = std::make_shared<FieldOperation>(
-        FieldOperation(std::move(valid_fn), {std::move(index_of_nearest_field)}));
-    params.set_output("Has Neighbor", Field<bool>(valid_op, 0));
+    params.set_output(
+        "Has Neighbor",
+        Field<bool>(std::make_shared<HasNeighborFieldInput>(std::move(group_field))));
   }
 }