Add multi-thread version of normal update to mesh (#92)

Summary:

Inspired by D63747584, adding a multi-threaded version of normal update in mesh.

Additional Changes:

- Clear `facePerVertex_` before updating
- Utilize Eigen::Vector::normalize() to potentially optimize SIMD (if available) while maintaining the same behavior (i.e., leaving the vector unchanged when the norm is too small)
- Use .noalias() for improved performance.
- Prefer using member functions over assignment (e.g., setZero() instead of m = Eigen::Vector3f::Zero()) for better performance.

Differential Revision: D63765102

CMakeLists.txt

@@ -156,6 +156,7 @@ mt_library(
     common
     fmt_eigen
     online_qr
+    Dispenso::dispenso
     Eigen3::Eigen
     Microsoft.GSL::GSL
 )

momentum/math/mesh.cpp

@@ -7,8 +7,11 @@
 
 #include "momentum/math/mesh.h"
 
+#include "momentum/common/profile.h"
 #include "momentum/math/utility.h"
 
+#include <dispenso/parallel_for.h>
+
 #include <algorithm>
 
 namespace momentum {
@@ -35,13 +38,73 @@ void MeshT<T>::updateNormals() {
     }
   }
   // re-normalize normals
-  for (size_t i = 0; i < normals.size(); i++) {
-    const T len = normals[i].norm();
-    // avoid divide-by-zero:
-    if (len != 0) {
-      normals[i] /= len;
+  for (auto& normal : normals) {
+    normal.normalize(); // if the input vector is too small, the this is left unchanged
+  }
+}
+
+template <typename T>
+void MeshT<T>::updateNormalsMt(size_t maxThreads) {
+  // fill vector for faces per vertex
+  facePerVertex_.clear(); // need to clear previous data
+  facePerVertex_.resize(vertices.size());
+  for (size_t i = 0; i < faces.size(); i++) {
+    const auto& face = faces[i];
+    for (size_t j = 0; j < 3; j++) {
+      facePerVertex_[face[j]].push_back(i);
     }
   }
+
+  // resize face normals
+  faceNormals_.resize(faces.size());
+
+  // set options for parallel for
+  auto dispensoOptions = dispenso::ParForOptions();
+  dispensoOptions.maxThreads = maxThreads;
+
+  // compute normals per face
+  MT_PROFILE_PUSH("Compute face normals");
+  const auto verticesNum = static_cast<int>(vertices.size());
+  dispenso::parallel_for(
+      0,
+      faces.size(),
+      [&](const size_t i) {
+        // Skip faces with out-of-boundaries indexes
+        const auto& face = faces[i];
+        if (std::any_of(face.begin(), face.end(), [verticesNum](int idx) {
+              return idx < 0 || idx >= verticesNum;
+            })) {
+          faceNormals_[i].setZero();
+          return;
+        }
+
+        // calculate normal
+        faceNormals_[i].noalias() =
+            (vertices[face[1]] - vertices[face[0]]).cross(vertices[face[2]] - vertices[face[0]]);
+      },
+      dispensoOptions);
+  MT_PROFILE_POP();
+
+  // for each vertex, add the face normals
+  MT_PROFILE_PUSH("add face normals");
+  normals.resize(vertices.size());
+  std::fill(normals.begin(), normals.end(), Eigen::Vector3<T>::Zero());
+  dispenso::parallel_for(
+      0,
+      facePerVertex_.size(),
+      [&](const size_t i) {
+        for (const auto& faceIdx : facePerVertex_[i]) {
+          if (IsNanNoOpt(faceNormals_[faceIdx][0])) {
+            continue;
+          }
+
+          normals[i].noalias() += faceNormals_[faceIdx];
+        }
+
+        normals[i].normalize(); // if the input vector is too small, the this is left unchanged
+      },
+      dispensoOptions);
+  MT_PROFILE_POP();
 }
 
 template <typename T, typename T2, int N>
@@ -81,6 +144,8 @@ void MeshT<T>::reset() {
   texcoords.clear();
   texcoord_faces.clear();
   texcoord_lines.clear();
+  facePerVertex_.clear();
+  faceNormals_.clear();
 }
 
 template MeshT<float> MeshT<float>::cast<float>() const;

momentum/math/mesh.h

@@ -36,15 +36,45 @@ struct MeshT {
   std::vector<Eigen::Vector3i> texcoord_faces; // list of texture coordinate indices per face
   std::vector<std::vector<int32_t>> texcoord_lines; // list of texture coordinate indices per line
 
+  MeshT(
+      const std::vector<Eigen::Vector3<T>>& vertices = {},
+      const std::vector<Eigen::Vector3<T>>& normals = {},
+      const std::vector<Eigen::Vector3i>& faces = {},
+      const std::vector<std::vector<int32_t>>& lines = {},
+      const std::vector<Eigen::Vector3b>& colors = {},
+      const std::vector<T>& confidence = {},
+      const std::vector<Eigen::Vector2f>& texcoords = {},
+      const std::vector<Eigen::Vector3i>& texcoord_faces = {},
+      const std::vector<std::vector<int32_t>>& texcoord_lines = {})
+      : vertices(vertices),
+        normals(normals),
+        faces(faces),
+        lines(lines),
+        colors(colors),
+        confidence(confidence),
+        texcoords(texcoords),
+        texcoord_faces(texcoord_faces),
+        texcoord_lines(texcoord_lines) {
+    // Empty
+  }
+
   /// Compute vertex normals (by averaging connected face normals)
   void updateNormals();
 
+  /// Compute vertex normals (by averaging connected face normals) in multi-threaded fashion
+  void updateNormalsMt(size_t maxThreads = std::numeric_limits<uint32_t>::max());
+
   /// Cast data type of mesh vertices, normals and confidence
   template <typename T2>
   MeshT<T2> cast() const;
 
   /// Reset mesh
   void reset();
+
+ private:
+  // Cache for updateNormalsMt()
+  std::vector<std::vector<size_t>> facePerVertex_; // for each vertex, list of faces it belongs to
+  std::vector<Vector3<T>> faceNormals_;
 };
 
 } // namespace momentum

momentum/test/math/mesh_test.cpp

@@ -43,6 +43,25 @@ TEST(Momentum_Mesh, UpdateNormals) {
   }
 }
 
+TEST(Momentum_Mesh, UpdateNormalsMt) {
+  // Construct simple mesh with 3 vertices and 1 face
+  Mesh m;
+  m.vertices = {{0.0, 0.0, 0.0}, {1.0, 0.0, 0.0}, {1.0, 1.0, 0.0}};
+  m.faces = {{0, 1, 2}};
+
+  // Compute mesh normals
+  m.updateNormalsMt();
+  EXPECT_EQ(m.normals.size(), 3);
+
+  // Ensure all normals are (0, 0, 1)
+  const float eps = 1e-5;
+  for (int i = 0; i < m.normals.size(); ++i) {
+    EXPECT_NEAR(m.normals[i][0], 0.f, eps);
+    EXPECT_NEAR(m.normals[i][1], 0.f, eps);
+    EXPECT_NEAR(m.normals[i][2], 1.f, eps);
+  }
+}
+
 TEST(Momentum_Mesh, Reset) {
   // Construct non-empty mesh
   Mesh m = {