Sculpt: Initial data-oriented refactor for mask brush

Part of #118145.

Refactor appears to gain a roughly 20% speed increase in the individual
update step function, from an average of 31ms to 25ms.

Pull Request: https://projects.blender.org/blender/blender/pulls/123115

source/blender/editors/sculpt_paint/CMakeLists.txt

@@ -120,6 +120,7 @@ set(SRC
   brushes/fill.cc
   brushes/flatten.cc
   brushes/inflate.cc
+  brushes/mask.cc
   brushes/multires_displacement_eraser.cc
   brushes/scrape.cc
   brushes/smooth.cc

source/blender/editors/sculpt_paint/brushes/mask.cc

@@ -0,0 +1,263 @@
+/* SPDX-FileCopyrightText: 2024 Blender Authors
+ *
+ * SPDX-License-Identifier: GPL-2.0-or-later */
+
+#include "editors/sculpt_paint/brushes/types.hh"
+
+#include "DNA_brush_types.h"
+#include "DNA_mesh_types.h"
+#include "DNA_object_types.h"
+#include "DNA_scene_types.h"
+
+#include "BKE_pbvh.hh"
+#include "BKE_subdiv_ccg.hh"
+
+#include "BLI_array.hh"
+#include "BLI_array_utils.hh"
+#include "BLI_enumerable_thread_specific.hh"
+#include "BLI_math_vector.hh"
+#include "BLI_span.hh"
+#include "BLI_task.h"
+
+#include "editors/sculpt_paint/mesh_brush_common.hh"
+#include "editors/sculpt_paint/sculpt_intern.hh"
+
+namespace blender::ed::sculpt_paint {
+inline namespace mask_cc {
+
+struct LocalData {
+  Vector<float> factors;
+  Vector<float> distances;
+  Vector<float> current_masks;
+  Vector<float> new_masks;
+};
+
+static void invert_mask(const MutableSpan<float> masks)
+{
+  for (float &mask : masks) {
+    mask = 1.0f - mask;
+  }
+}
+
+static void apply_factors(const float strength,
+                          const Span<float> current_masks,
+                          const Span<float> factors,
+                          const MutableSpan<float> masks)
+{
+  BLI_assert(current_masks.size() == masks.size());
+  BLI_assert(factors.size() == masks.size());
+  for (const int i : masks.index_range()) {
+    masks[i] += factors[i] * current_masks[i] * strength;
+  }
+}
+
+static void clamp_mask(const MutableSpan<float> masks)
+{
+  for (float &mask : masks) {
+    mask = std::clamp(mask, 0.0f, 1.0f);
+  }
+}
+
+static void calc_faces(const Brush &brush,
+                       const float strength,
+                       const Span<float3> positions,
+                       const Span<float3> vert_normals,
+                       const PBVHNode &node,
+                       Object &object,
+                       Mesh &mesh,
+                       LocalData &tls,
+                       const MutableSpan<float> mask)
+{
+  SculptSession &ss = *object.sculpt;
+  StrokeCache &cache = *ss.cache;
+
+  const Span<int> verts = bke::pbvh::node_unique_verts(node);
+
+  tls.factors.reinitialize(verts.size());
+  const MutableSpan<float> factors = tls.factors;
+  fill_factor_from_hide(mesh, verts, factors);
+
+  if (brush.flag & BRUSH_FRONTFACE) {
+    calc_front_face(cache.view_normal, vert_normals, verts, factors);
+  }
+
+  tls.distances.reinitialize(verts.size());
+  const MutableSpan<float> distances = tls.distances;
+  calc_distance_falloff(
+      ss, positions, verts, eBrushFalloffShape(brush.falloff_shape), distances, factors);
+  calc_brush_strength_factors(ss, brush, verts, distances, factors);
+
+  if (ss.cache->automasking) {
+    auto_mask::calc_vert_factors(object, *ss.cache->automasking, node, verts, factors);
+  }
+
+  calc_brush_texture_factors(ss, brush, positions, verts, factors);
+
+  tls.new_masks.reinitialize(verts.size());
+  const MutableSpan<float> new_masks = tls.new_masks;
+  array_utils::gather(mask.as_span(), verts, new_masks);
+
+  tls.current_masks = tls.new_masks;
+  const MutableSpan<float> current_masks = tls.current_masks;
+  if (strength > 0.0f) {
+    invert_mask(current_masks);
+  }
+  apply_factors(strength, current_masks, factors, new_masks);
+  clamp_mask(new_masks);
+
+  array_utils::scatter(new_masks.as_span(), verts, mask);
+}
+
+static float calc_new_mask(const float mask, const float factor, const float strength)
+{
+  const float modified_value = strength > 0.0f ? (1.0f - mask) : mask;
+  const float result = mask + factor * strength * modified_value;
+  return std::clamp(result, 0.0f, 1.0f);
+}
+
+static void calc_grids(Object &object, const Brush &brush, const float strength, PBVHNode &node)
+{
+  SculptSession &ss = *object.sculpt;
+
+  SculptBrushTest test;
+  SculptBrushTestFn sculpt_brush_test_sq_fn = SCULPT_brush_test_init_with_falloff_shape(
+      ss, test, brush.falloff_shape);
+  const int thread_id = BLI_task_parallel_thread_id(nullptr);
+  auto_mask::NodeData automask_data = auto_mask::node_begin(
+      object, ss.cache->automasking.get(), node);
+
+  SubdivCCG &subdiv_ccg = *ss.subdiv_ccg;
+  const CCGKey key = *BKE_pbvh_get_grid_key(*ss.pbvh);
+  const Span<CCGElem *> grids = subdiv_ccg.grids;
+  const BitGroupVector<> &grid_hidden = subdiv_ccg.grid_hidden;
+
+  int i = 0;
+  for (const int grid : bke::pbvh::node_grid_indices(node)) {
+    const int grid_verts_start = grid * key.grid_area;
+    CCGElem *elem = grids[grid];
+    for (const int j : IndexRange(key.grid_area)) {
+      if (!grid_hidden.is_empty() && grid_hidden[grid][j]) {
+        i++;
+        continue;
+      }
+      if (!sculpt_brush_test_sq_fn(test, CCG_elem_offset_co(key, elem, j))) {
+        i++;
+        continue;
+      }
+      auto_mask::node_update(automask_data, i);
+      const float fade = SCULPT_brush_strength_factor(ss,
+                                                      brush,
+                                                      CCG_elem_offset_co(key, elem, j),
+                                                      math::sqrt(test.dist),
+                                                      CCG_elem_offset_no(key, elem, j),
+                                                      nullptr,
+                                                      0.0f,
+                                                      BKE_pbvh_make_vref(grid_verts_start + j),
+                                                      thread_id,
+                                                      &automask_data);
+
+      const float current_mask = key.has_mask ? CCG_elem_offset_mask(key, elem, j) : 0.0f;
+      const float new_mask = calc_new_mask(current_mask, fade, strength);
+      CCG_elem_offset_mask(key, elem, j) = new_mask;
+      i++;
+    }
+  }
+}
+
+static void calc_bmesh(Object &object, const Brush &brush, const float strength, PBVHNode &node)
+{
+  SculptSession &ss = *object.sculpt;
+
+  SculptBrushTest test;
+  SculptBrushTestFn sculpt_brush_test_sq_fn = SCULPT_brush_test_init_with_falloff_shape(
+      ss, test, brush.falloff_shape);
+  const int thread_id = BLI_task_parallel_thread_id(nullptr);
+  auto_mask::NodeData automask_data = auto_mask::node_begin(
+      object, ss.cache->automasking.get(), node);
+
+  const int mask_offset = CustomData_get_offset_named(
+      &ss.bm->vdata, CD_PROP_FLOAT, ".sculpt_mask");
+
+  for (BMVert *vert : BKE_pbvh_bmesh_node_unique_verts(&node)) {
+    if (BM_elem_flag_test(vert, BM_ELEM_HIDDEN)) {
+      continue;
+    }
+    if (!sculpt_brush_test_sq_fn(test, vert->co)) {
+      continue;
+    }
+    auto_mask::node_update(automask_data, *vert);
+    const float mask = mask_offset == -1 ? 0.0f : BM_ELEM_CD_GET_FLOAT(vert, mask_offset);
+    const float fade = SCULPT_brush_strength_factor(ss,
+                                                    brush,
+                                                    vert->co,
+                                                    math::sqrt(test.dist),
+                                                    vert->no,
+                                                    nullptr,
+                                                    0.0f,
+                                                    BKE_pbvh_make_vref(intptr_t(vert)),
+                                                    thread_id,
+                                                    &automask_data);
+    const float new_mask = calc_new_mask(mask, fade, strength);
+    BM_ELEM_CD_SET_FLOAT(vert, mask_offset, new_mask);
+  }
+}
+}  // namespace mask_cc
+
+void do_mask_brush(const Sculpt &sd, Object &object, Span<PBVHNode *> nodes)
+{
+  SculptSession &ss = *object.sculpt;
+  const Brush &brush = *BKE_paint_brush_for_read(&sd.paint);
+  const float bstrength = ss.cache->bstrength;
+
+  switch (BKE_pbvh_type(*ss.pbvh)) {
+    case PBVH_FACES: {
+      threading::EnumerableThreadSpecific<LocalData> all_tls;
+      Mesh &mesh = *static_cast<Mesh *>(object.data);
+
+      const PBVH &pbvh = *ss.pbvh;
+      const Span<float3> positions = BKE_pbvh_get_vert_positions(pbvh);
+      const Span<float3> vert_normals = BKE_pbvh_get_vert_normals(pbvh);
+
+      bke::MutableAttributeAccessor attributes = mesh.attributes_for_write();
+
+      bke::SpanAttributeWriter<float> mask = attributes.lookup_or_add_for_write_span<float>(
+          ".sculpt_mask", bke::AttrDomain::Point);
+
+      threading::parallel_for(nodes.index_range(), 1, [&](const IndexRange range) {
+        threading::isolate_task([&]() {
+          LocalData &tls = all_tls.local();
+          for (const int i : range) {
+            calc_faces(brush,
+                       bstrength,
+                       positions,
+                       vert_normals,
+                       *nodes[i],
+                       object,
+                       mesh,
+                       tls,
+                       mask.span);
+          }
+        });
+      });
+      mask.finish();
+      break;
+    }
+    case PBVH_GRIDS: {
+      threading::parallel_for(nodes.index_range(), 1, [&](const IndexRange range) {
+        for (const int i : range) {
+          calc_grids(object, brush, bstrength, *nodes[i]);
+        }
+      });
+      break;
+    }
+    case PBVH_BMESH: {
+      threading::parallel_for(nodes.index_range(), 1, [&](const IndexRange range) {
+        for (const int i : range) {
+          calc_bmesh(object, brush, bstrength, *nodes[i]);
+        }
+      });
+      break;
+    }
+  }
+}
+}  // namespace blender::ed::sculpt_paint

source/blender/editors/sculpt_paint/brushes/types.hh

@@ -20,6 +20,8 @@ void do_draw_vector_displacement_brush(const Sculpt &sd, Object &object, Span<PB
 void do_fill_brush(const Sculpt &sd, Object &object, Span<PBVHNode *> nodes);
 void do_flatten_brush(const Sculpt &sd, Object &ob, Span<PBVHNode *> nodes);
 void do_inflate_brush(const Sculpt &sd, Object &ob, Span<PBVHNode *> nodes);
+/** A brush that modifies mask values instead of position. */
+void do_mask_brush(const Sculpt &sd, Object &object, Span<PBVHNode *> nodes);
 void do_nudge_brush(const Sculpt &sd, Object &object, Span<PBVHNode *> nodes);
 void do_scrape_brush(const Sculpt &sd, Object &object, Span<PBVHNode *> nodes);
 /** Smooth positions with neighboring vertices. */

source/blender/editors/sculpt_paint/mesh_brush_common.hh

@@ -52,6 +52,11 @@ void scale_factors(MutableSpan<float> factors, float strength);
  *   are built for these values, then applied to `positions_orig`.
  */
 
+/**
+ * Calculate initial influence factors based on vertex visibility.
+ */
+void fill_factor_from_hide(const Mesh &mesh, Span<int> vert_indices, MutableSpan<float> r_factors);
+
 /**
  * Calculate initial influence factors based on vertex visibility and masking.
  */

source/blender/editors/sculpt_paint/sculpt.cc

@@ -3842,7 +3842,14 @@ static void do_brush_action(const Sculpt &sd,
       }
       break;
     case SCULPT_TOOL_MASK:
-      SCULPT_do_mask_brush(sd, ob, nodes);
+      switch ((BrushMaskTool)brush.mask_tool) {
+        case BRUSH_MASK_DRAW:
+          do_mask_brush(sd, ob, nodes);
+          break;
+        case BRUSH_MASK_SMOOTH:
+          smooth::do_smooth_mask_brush(sd, ob, nodes, ss.cache->bstrength);
+          break;
+      }
       break;
     case SCULPT_TOOL_POSE:
       pose::do_pose_brush(sd, ob, nodes);
@@ -6510,28 +6517,39 @@ void SCULPT_cube_tip_init(const Sculpt & /*sd*/,
 
 namespace blender::ed::sculpt_paint {
 
+void fill_factor_from_hide(const Mesh &mesh,
+                           const Span<int> verts,
+                           const MutableSpan<float> r_factors)
+{
+  BLI_assert(verts.size() == r_factors.size());
+
+  /* TODO: Avoid overhead of accessing attributes for every PBVH node. */
+  const bke::AttributeAccessor attributes = mesh.attributes();
+  if (const VArray hide_vert = *attributes.lookup<bool>(".hide_vert", bke::AttrDomain::Point)) {
+    const VArraySpan span(hide_vert);
+    for (const int i : verts.index_range()) {
+      r_factors[i] = span[verts[i]] ? 0.0f : 1.0f;
+    }
+  }
+  else {
+    r_factors.fill(1.0f);
+  }
+}
+
 void fill_factor_from_hide_and_mask(const Mesh &mesh,
                                     const Span<int> verts,
                                     const MutableSpan<float> r_factors)
 {
   BLI_assert(verts.size() == r_factors.size());
 
+  fill_factor_from_hide(mesh, verts, r_factors);
+
   /* TODO: Avoid overhead of accessing attributes for every PBVH node. */
   const bke::AttributeAccessor attributes = mesh.attributes();
   if (const VArray mask = *attributes.lookup<float>(".sculpt_mask", bke::AttrDomain::Point)) {
     const VArraySpan span(mask);
     for (const int i : verts.index_range()) {
-      r_factors[i] = 1.0f - span[verts[i]];
-    }
-  }
-  else {
-    r_factors.fill(1.0f);
-  }
-
-  if (const VArray hide_vert = *attributes.lookup<bool>(".hide_vert", bke::AttrDomain::Point)) {
-    const VArraySpan span(hide_vert);
-    for (const int i : verts.index_range()) {
-      r_factors[i] = span[verts[i]] ? 0.0f : r_factors[i];
+      r_factors[i] -= std::max(1.0f - span[verts[i]], 0.0f);
     }
   }
 }

source/blender/editors/sculpt_paint/sculpt_brush_types.cc

@@ -1931,86 +1931,3 @@ void SCULPT_bmesh_topology_rake(const Sculpt &sd,
 }
 
 /** \} */
-
-/* -------------------------------------------------------------------- */
-/** \name Sculpt Mask Brush
- * \{ */
-
-static void do_mask_brush_draw_task(Object &ob,
-                                    const Brush &brush,
-                                    const SculptMaskWriteInfo mask_write,
-                                    PBVHNode *node)
-{
-  using namespace blender::ed::sculpt_paint;
-  SculptSession &ss = *ob.sculpt;
-  const float bstrength = ss.cache->bstrength;
-
-  PBVHVertexIter vd;
-
-  SculptBrushTest test;
-  SculptBrushTestFn sculpt_brush_test_sq_fn = SCULPT_brush_test_init_with_falloff_shape(
-      ss, test, brush.falloff_shape);
-  const int thread_id = BLI_task_parallel_thread_id(nullptr);
-
-  auto_mask::NodeData automask_data = auto_mask::node_begin(
-      ob, ss.cache->automasking.get(), *node);
-
-  BKE_pbvh_vertex_iter_begin (*ss.pbvh, node, vd, PBVH_ITER_UNIQUE) {
-    if (!sculpt_brush_test_sq_fn(test, vd.co)) {
-      continue;
-    }
-
-    auto_mask::node_update(automask_data, vd);
-    const float fade = SCULPT_brush_strength_factor(ss,
-                                                    brush,
-                                                    vd.co,
-                                                    sqrtf(test.dist),
-                                                    vd.no,
-                                                    vd.fno,
-                                                    0.0f,
-                                                    vd.vertex,
-                                                    thread_id,
-                                                    &automask_data);
-
-    float mask = vd.mask;
-    if (bstrength > 0.0f) {
-      mask += fade * bstrength * (1.0f - vd.mask);
-    }
-    else {
-      mask += fade * bstrength * mask;
-    }
-    mask = clamp_f(mask, 0.0f, 1.0f);
-    SCULPT_mask_vert_set(BKE_pbvh_type(*ss.pbvh), mask_write, mask, vd);
-  }
-  BKE_pbvh_vertex_iter_end;
-}
-
-void SCULPT_do_mask_brush_draw(const Sculpt &sd, Object &ob, Span<PBVHNode *> nodes)
-{
-  using namespace blender;
-  const Brush &brush = *BKE_paint_brush_for_read(&sd.paint);
-  const SculptMaskWriteInfo mask_write = SCULPT_mask_get_for_write(*ob.sculpt);
-  threading::parallel_for(nodes.index_range(), 1, [&](const IndexRange range) {
-    for (const int i : range) {
-      do_mask_brush_draw_task(ob, brush, mask_write, nodes[i]);
-    }
-  });
-}
-
-void SCULPT_do_mask_brush(const Sculpt &sd, Object &ob, Span<PBVHNode *> nodes)
-{
-  using namespace blender::ed::sculpt_paint;
-  SculptSession &ss = *ob.sculpt;
-  const Brush &brush = *BKE_paint_brush_for_read(&sd.paint);
-
-  switch ((BrushMaskTool)brush.mask_tool) {
-    case BRUSH_MASK_DRAW:
-      SCULPT_do_mask_brush_draw(sd, ob, nodes);
-      break;
-    case BRUSH_MASK_SMOOTH:
-      smooth::do_smooth_mask_brush(sd, ob, nodes, ss.cache->bstrength);
-      break;
-  }
-}
-
-/** \} */

source/blender/editors/sculpt_paint/sculpt_intern.hh

@@ -2036,8 +2036,6 @@ void SCULPT_do_slide_relax_brush(const Sculpt &sd, Object &ob, blender::Span<PBV
 void SCULPT_do_displacement_smear_brush(const Sculpt &sd,
                                         Object &ob,
                                         blender::Span<PBVHNode *> nodes);
-void SCULPT_do_mask_brush_draw(const Sculpt &sd, Object &ob, blender::Span<PBVHNode *> nodes);
-void SCULPT_do_mask_brush(const Sculpt &sd, Object &ob, blender::Span<PBVHNode *> nodes);
 /** \} */
 
 void SCULPT_bmesh_topology_rake(const Sculpt &sd,