parallel uv unwrap (#5950)

- Added function PCAPartition() to the Tensor PointCloud and TriangleMesh for fast partitioning
- Added function SelectFacesByMask() to Tensor TriangleMesh to create a new mesh with the selected faces
- Added new arguments to ComputeUVAtlas() to automatically partition the mesh and unwrap it in parallel

15 files changed, 700 insertions, 64 deletions

diff --git a/cpp/open3d/t/geometry/PointCloud.cpp b/cpp/open3d/t/geometry/PointCloud.cpp
index f24ad21b..6ec721ff 100644
--- a/cpp/open3d/t/geometry/PointCloud.cpp
+++ b/cpp/open3d/t/geometry/PointCloud.cpp
@@ -34,6 +34,7 @@
 #include "open3d/t/geometry/TriangleMesh.h"
 #include "open3d/t/geometry/VtkUtils.h"
 #include "open3d/t/geometry/kernel/GeometryMacros.h"
+#include "open3d/t/geometry/kernel/PCAPartition.h"
 #include "open3d/t/geometry/kernel/PointCloud.h"
 #include "open3d/t/geometry/kernel/Transform.h"
 #include "open3d/t/pipelines/registration/Registration.h"
@@ -1244,6 +1245,15 @@ PointCloud PointCloud::Crop(const OrientedBoundingBox &obb, bool invert) const {
             obb.GetPointIndicesWithinBoundingBox(GetPointPositions()), invert);
 }
 
+int PointCloud::PCAPartition(int max_points) {
+    int num_partitions;
+    core::Tensor partition_ids;
+    std::tie(num_partitions, partition_ids) =
+            kernel::pcapartition::PCAPartition(GetPointPositions(), max_points);
+    SetPointAttr("partition_ids", partition_ids.To(GetDevice()));
+    return num_partitions;
+}
+
 }  // namespace geometry
 }  // namespace t
 }  // namespace open3d
diff --git a/cpp/open3d/t/geometry/PointCloud.h b/cpp/open3d/t/geometry/PointCloud.h
index eda27467..68827fa4 100644
--- a/cpp/open3d/t/geometry/PointCloud.h
+++ b/cpp/open3d/t/geometry/PointCloud.h
@@ -682,6 +682,13 @@ public:
                           double scale = 1.0,
                           bool capping = true) const;
 
+    /// Partition the point cloud by recursively doing PCA.
+    /// This function creates a new point attribute with the name
+    /// "partition_ids".
+    /// \param max_points The maximum allowed number of points in a partition.
+    /// \return The number of partitions.
+    int PCAPartition(int max_points);
+
 protected:
     core::Device device_ = core::Device("CPU:0");
     TensorMap point_attr_;
diff --git a/cpp/open3d/t/geometry/TriangleMesh.cpp b/cpp/open3d/t/geometry/TriangleMesh.cpp
index 1d0224da..4134ee6f 100644
--- a/cpp/open3d/t/geometry/TriangleMesh.cpp
+++ b/cpp/open3d/t/geometry/TriangleMesh.cpp
@@ -28,10 +28,12 @@
 #include "open3d/t/geometry/PointCloud.h"
 #include "open3d/t/geometry/RaycastingScene.h"
 #include "open3d/t/geometry/VtkUtils.h"
+#include "open3d/t/geometry/kernel/PCAPartition.h"
 #include "open3d/t/geometry/kernel/PointCloud.h"
 #include "open3d/t/geometry/kernel/Transform.h"
 #include "open3d/t/geometry/kernel/TriangleMesh.h"
 #include "open3d/t/geometry/kernel/UVUnwrapping.h"
+#include "open3d/utility/ParallelScan.h"
 
 namespace open3d {
 namespace t {
@@ -666,11 +668,15 @@ TriangleMesh TriangleMesh::FillHoles(double hole_size) const {
     return CreateTriangleMeshFromVtkPolyData(result);
 }
 
-void TriangleMesh::ComputeUVAtlas(size_t size,
-                                  float gutter,
-                                  float max_stretch) {
-    kernel::uvunwrapping::ComputeUVAtlas(*this, size, size, gutter,
-                                         max_stretch);
+std::tuple<float, int, int> TriangleMesh::ComputeUVAtlas(
+        size_t size,
+        float gutter,
+        float max_stretch,
+        int parallel_partitions,
+        int nthreads) {
+    return kernel::uvunwrapping::ComputeUVAtlas(*this, size, size, gutter,
+                                                max_stretch,
+                                                parallel_partitions, nthreads);
 }
 
 namespace {
@@ -971,6 +977,90 @@ TriangleMesh TriangleMesh::ExtrudeLinear(const core::Tensor &vector,
     return ExtrudeLinearTriangleMesh(*this, vector, scale, capping);
 }
 
+int TriangleMesh::PCAPartition(int max_faces) {
+    core::Tensor verts = GetVertexPositions();
+    core::Tensor tris = GetTriangleIndices();
+    if (!tris.GetLength()) {
+        utility::LogError("Mesh must have at least one face.");
+    }
+    core::Tensor tris_centers = verts.IndexGet({tris}).Mean({1});
+
+    int num_parititions;
+    core::Tensor partition_ids;
+    std::tie(num_parititions, partition_ids) =
+            kernel::pcapartition::PCAPartition(tris_centers, max_faces);
+    SetTriangleAttr("partition_ids", partition_ids.To(GetDevice()));
+    return num_parititions;
+}
+
+TriangleMesh TriangleMesh::SelectFacesByMask(const core::Tensor &mask) const {
+    core::AssertTensorShape(mask, {GetTriangleIndices().GetLength()});
+    core::AssertTensorDtype(mask, core::Bool);
+    GetTriangleAttr().AssertSizeSynchronized();
+    GetVertexAttr().AssertSizeSynchronized();
+
+    // select triangles
+    core::Tensor tris = GetTriangleIndices().IndexGet({mask});
+    core::Tensor tris_cpu = tris.To(core::Device()).Contiguous();
+    const int64_t num_tris = tris_cpu.GetLength();
+
+    // create mask for vertices that are part of the selected faces
+    const int64_t num_verts = GetVertexPositions().GetLength();
+    core::Tensor vertex_mask = core::Tensor::Zeros({num_verts}, core::Int32);
+    std::vector<int64_t> prefix_sum(num_verts + 1, 0);
+    {
+        int32_t *vertex_mask_ptr = vertex_mask.GetDataPtr<int32_t>();
+        if (tris_cpu.GetDtype() == core::Int32) {
+            int32_t *vert_idx_ptr = tris_cpu.GetDataPtr<int32_t>();
+            for (int64_t i = 0; i < tris_cpu.GetLength() * 3; ++i) {
+                vertex_mask_ptr[vert_idx_ptr[i]] = 1;
+            }
+        } else {
+            int64_t *vert_idx_ptr = tris_cpu.GetDataPtr<int64_t>();
+            for (int64_t i = 0; i < tris_cpu.GetLength() * 3; ++i) {
+                vertex_mask_ptr[vert_idx_ptr[i]] = 1;
+            }
+        }
+        utility::InclusivePrefixSum(
+                vertex_mask_ptr, vertex_mask_ptr + num_verts, &prefix_sum[1]);
+    }
+
+    // update triangle indices
+    if (tris_cpu.GetDtype() == core::Int32) {
+        int32_t *vert_idx_ptr = tris_cpu.GetDataPtr<int32_t>();
+        for (int64_t i = 0; i < num_tris * 3; ++i) {
+            int64_t new_idx = prefix_sum[vert_idx_ptr[i]];
+            vert_idx_ptr[i] = int32_t(new_idx);
+        }
+    } else {
+        int64_t *vert_idx_ptr = tris_cpu.GetDataPtr<int64_t>();
+        for (int64_t i = 0; i < num_tris * 3; ++i) {
+            int64_t new_idx = prefix_sum[vert_idx_ptr[i]];
+            vert_idx_ptr[i] = new_idx;
+        }
+    }
+
+    tris = tris_cpu.To(GetDevice());
+    vertex_mask = vertex_mask.To(GetDevice(), core::Bool);
+    core::Tensor verts = GetVertexPositions().IndexGet({vertex_mask});
+    TriangleMesh result(verts, tris);
+
+    // copy attributes
+    for (auto item : GetVertexAttr()) {
+        if (!result.HasVertexAttr(item.first)) {
+            result.SetVertexAttr(item.first,
+                                 item.second.IndexGet({vertex_mask}));
+        }
+    }
+    for (auto item : GetTriangleAttr()) {
+        if (!result.HasTriangleAttr(item.first)) {
+            result.SetTriangleAttr(item.first, item.second.IndexGet({mask}));
+        }
+    }
+
+    return result;
+}
+
 }  // namespace geometry
 }  // namespace t
 }  // namespace open3d
diff --git a/cpp/open3d/t/geometry/TriangleMesh.h b/cpp/open3d/t/geometry/TriangleMesh.h
index 8d5c8c20..7828ac16 100644
--- a/cpp/open3d/t/geometry/TriangleMesh.h
+++ b/cpp/open3d/t/geometry/TriangleMesh.h
@@ -813,9 +813,23 @@ public:
     /// \param gutter This is the space around the uv islands in pixels.
     /// \param max_stretch The maximum amount of stretching allowed. The
     /// parameter range is [0..1] with 0 meaning no stretch allowed.
-    void ComputeUVAtlas(size_t size = 512,
-                        float gutter = 1.0f,
-                        float max_stretch = 1.f / 6);
+    /// \param parallel_partitions The approximate number of partitions created
+    /// before computing the UV atlas for parallelizing the computation.
+    /// Parallelization can be enabled with values > 1. Note that
+    /// parallelization increases the number of UV islands and can lead to
+    /// results with lower quality.
+    /// \param nthreads The number of threads used
+    /// when parallel_partitions is > 1. Set to 0 for automatic number of thread
+    /// detection.
+    ///
+    /// \return Tuple with (max stretch, num_charts, num_partitions) storing the
+    /// actual amount of stretch, the number of created charts, and the number
+    /// of parallel partitions created.
+    std::tuple<float, int, int> ComputeUVAtlas(size_t size = 512,
+                                               float gutter = 1.0f,
+                                               float max_stretch = 1.f / 6,
+                                               int parallel_partitions = 1,
+                                               int nthreads = 0);
 
     /// Bake vertex attributes into textures.
     ///
@@ -903,6 +917,19 @@ public:
                                double scale = 1.0,
                                bool capping = true) const;
 
+    /// Partition the mesh by recursively doing PCA.
+    /// This function creates a new triangle attribute with the name
+    /// "partition_ids".
+    /// \param max_faces The maximum allowed number of faces in a partition.
+    /// \return The number of partitions.
+    int PCAPartition(int max_faces);
+
+    /// Returns a new mesh with the faces selected by a boolean mask.
+    /// \param mask A boolean mask with the shape (N) with N as the number of
+    /// faces in the mesh.
+    /// \return A new mesh with the selected faces.
+    TriangleMesh SelectFacesByMask(const core::Tensor &mask) const;
+
 protected:
     core::Device device_ = core::Device("CPU:0");
     TensorMap vertex_attr_;
diff --git a/cpp/open3d/t/geometry/VtkUtils.cpp b/cpp/open3d/t/geometry/VtkUtils.cpp
index a9ed11fe..e6b7245d 100644
--- a/cpp/open3d/t/geometry/VtkUtils.cpp
+++ b/cpp/open3d/t/geometry/VtkUtils.cpp
@@ -310,17 +310,19 @@ static void AddTensorMapToVtkFieldData(
             continue;
         }
         // we only support 2D tensors
-        if (key_tensor.second.NumDims() != 2) {
-            utility::LogWarning(
-                    "Ignoring attribute '{}' for TensorMap with primary key "
-                    "'{}' because of incompatible ndim={}",
-                    key_tensor.first, tmap.GetPrimaryKey(),
-                    key_tensor.second.NumDims());
-            continue;
-        }
 
         if (include.count(key_tensor.first) &&
             !exclude.count(key_tensor.first)) {
+            if (key_tensor.second.NumDims() != 2) {
+                utility::LogWarning(
+                        "Ignoring attribute '{}' for TensorMap with primary "
+                        "key "
+                        "'{}' because of incompatible ndim={}",
+                        key_tensor.first, tmap.GetPrimaryKey(),
+                        key_tensor.second.NumDims());
+                continue;
+            }
+
             auto array = CreateVtkDataArrayFromTensor(key_tensor.second, copy);
             array->SetName(key_tensor.first.c_str());
             field_data->AddArray(array);
diff --git a/cpp/open3d/t/geometry/VtkUtils.h b/cpp/open3d/t/geometry/VtkUtils.h
index 1c108d41..85aeabe6 100644
--- a/cpp/open3d/t/geometry/VtkUtils.h
+++ b/cpp/open3d/t/geometry/VtkUtils.h
@@ -38,9 +38,8 @@ int DtypeToVtkType(const core::Dtype& dtype);
 ///  will not be added to the vtkPolyData. The exclusion set has precedence over
 ///  the included keys.
 /// \param face_attr_exclude A set of keys for which face attributes will not be
-/// added
-///  to the vtkPolyData. The exclusion set has precedence over the included
-///  keys.
+/// added to the vtkPolyData. The exclusion set has precedence over the included
+/// keys.
 vtkSmartPointer<vtkPolyData> CreateVtkPolyDataFromGeometry(
         const Geometry& geometry,
         const std::unordered_set<std::string>& point_attr_include,
diff --git a/cpp/open3d/t/geometry/kernel/CMakeLists.txt b/cpp/open3d/t/geometry/kernel/CMakeLists.txt
index e0327954..081d24a6 100644
--- a/cpp/open3d/t/geometry/kernel/CMakeLists.txt
+++ b/cpp/open3d/t/geometry/kernel/CMakeLists.txt
@@ -3,6 +3,7 @@ open3d_ispc_add_library(tgeometry_kernel OBJECT)
 target_sources(tgeometry_kernel PRIVATE
     Image.cpp
     ImageCPU.cpp
+    PCAPartition.cpp
     PointCloud.cpp
     PointCloudCPU.cpp
     TriangleMesh.cpp
diff --git a/cpp/open3d/t/geometry/kernel/PCAPartition.cpp b/cpp/open3d/t/geometry/kernel/PCAPartition.cpp
new file mode 100644
index 00000000..e1cfedc1
--- /dev/null
+++ b/cpp/open3d/t/geometry/kernel/PCAPartition.cpp
@@ -0,0 +1,143 @@
+// ----------------------------------------------------------------------------
+// -                        Open3D: www.open3d.org                            -
+// ----------------------------------------------------------------------------
+// Copyright (c) 2018-2023 www.open3d.org
+// SPDX-License-Identifier: MIT
+// ----------------------------------------------------------------------------
+
+#include "open3d/t/geometry/kernel/PCAPartition.h"
+
+#include <Eigen/Eigenvalues>
+#include <numeric>
+
+#include "open3d/core/TensorCheck.h"
+#include "open3d/utility/Eigen.h"
+
+namespace open3d {
+namespace t {
+namespace geometry {
+namespace kernel {
+namespace pcapartition {
+
+namespace {
+typedef std::vector<size_t> IdxVec;
+
+/// Split a partition using PCA.
+/// \tparam TReal The data type for the points. Either float or double.
+/// \param points Array of 3D points with shape (N,3) as contiguous memory.
+/// \param num_points Number of points N.
+/// \param indices The point indices of the input partition.
+/// \param output Output vector for storing the 2 newly created partitions after
+/// splitting.
+template <class TReal>
+void Split(const TReal* const points,
+           const size_t num_points,
+           const IdxVec& indices,
+           std::vector<IdxVec>& output) {
+    Eigen::Vector3d mean;
+    Eigen::Matrix3d cov;
+    std::tie(mean, cov) = utility::ComputeMeanAndCovariance(points, indices);
+
+    Eigen::SelfAdjointEigenSolver<Eigen::Matrix3d> solver;
+    solver.computeDirect(cov, Eigen::ComputeEigenvectors);
+    Eigen::Matrix<TReal, 3, 1> ax = solver.eigenvectors().col(2).cast<TReal>();
+
+    Eigen::Matrix<TReal, 3, 1> mu = mean.cast<TReal>();
+    Eigen::Map<const Eigen::Matrix<TReal, 3, Eigen::Dynamic>> pts(points, 3,
+                                                                  num_points);
+    TReal mindot = ax.dot(pts.col(indices.front()) - mu);
+    TReal maxdot = mindot;
+    for (auto idx : indices) {
+        TReal dot = ax.dot(pts.col(idx) - mu);
+        mindot = std::min(dot, mindot);
+        maxdot = std::max(dot, maxdot);
+    }
+    TReal center = TReal(0.5) * (mindot + maxdot);
+    IdxVec part1, part2;
+    for (auto idx : indices) {
+        TReal dot = ax.dot(pts.col(idx) - mu);
+        if (dot < center) {
+            part1.push_back(idx);
+        } else {
+            part2.push_back(idx);
+        }
+    }
+
+    if (part1.empty() || part2.empty()) {
+        // this should not happen make sure that we return two partitions that
+        // are both smaller than the input indices
+        part1.clear();
+        part2.clear();
+        part1.insert(part1.begin(), indices.begin(),
+                     indices.begin() + indices.size() / 2);
+        part2.insert(part2.begin(), indices.begin() + indices.size() / 2,
+                     indices.end());
+    }
+
+    output.emplace_back(IdxVec(std::move(part1)));
+    output.emplace_back(IdxVec(std::move(part2)));
+}
+}  // namespace
+
+std::tuple<int, core::Tensor> PCAPartition(core::Tensor& points,
+                                           int max_points) {
+    if (max_points <= 0) {
+        utility::LogError("max_points must be > 0 but is {}", max_points);
+    }
+    core::AssertTensorDtypes(points, {core::Float32, core::Float64});
+    core::AssertTensorShape(points, {utility::nullopt, 3});
+    const size_t max_points_(max_points);
+    const size_t num_points = points.GetLength();
+    if (num_points == 0) {
+        utility::LogError("number of points must be greater zero");
+    }
+
+    auto points_cpu = points.To(core::Device()).Contiguous();
+    core::AssertTensorDtypes(points_cpu, {core::Float32, core::Float64});
+
+    std::vector<IdxVec> partitions_to_process;
+    std::vector<IdxVec> partitions_result;
+    {
+        partitions_to_process.emplace_back(IdxVec(num_points));
+        IdxVec& indices = partitions_to_process.back();
+        std::iota(std::begin(indices), std::end(indices), 0);
+    }
+
+    while (partitions_to_process.size()) {
+        std::vector<IdxVec> tmp;
+        for (auto& indices : partitions_to_process) {
+            if (indices.size() <= max_points_) {
+                partitions_result.emplace_back(IdxVec(std::move(indices)));
+            } else {
+                if (points.GetDtype() == core::Float32) {
+                    Split(points_cpu.GetDataPtr<float>(), num_points, indices,
+                          tmp);
+                } else {
+                    Split(points_cpu.GetDataPtr<double>(), num_points, indices,
+                          tmp);
+                }
+            }
+        }
+        // tmp contains the partitions for the next iteration
+        partitions_to_process.clear();
+        partitions_to_process.swap(tmp);
+    }
+
+    auto partition_id = core::Tensor::Empty({int64_t(num_points)}, core::Int32);
+    int32_t* pid_ptr = partition_id.GetDataPtr<int32_t>();
+    for (size_t i = 0; i < partitions_result.size(); ++i) {
+        IdxVec& indices = partitions_result[i];
+        for (const auto idx : indices) {
+            pid_ptr[idx] = int32_t(i);
+        }
+    }
+
+    int num_partitions = partitions_result.size();
+    return std::tie(num_partitions, partition_id);
+}
+
+}  // namespace pcapartition
+}  // namespace kernel
+}  // namespace geometry
+}  // namespace t
+}  // namespace open3d
\ No newline at end of file
diff --git a/cpp/open3d/t/geometry/kernel/PCAPartition.h b/cpp/open3d/t/geometry/kernel/PCAPartition.h
new file mode 100644
index 00000000..4127fa64
--- /dev/null
+++ b/cpp/open3d/t/geometry/kernel/PCAPartition.h
@@ -0,0 +1,30 @@
+// ----------------------------------------------------------------------------
+// -                        Open3D: www.open3d.org                            -
+// ----------------------------------------------------------------------------
+// Copyright (c) 2018-2023 www.open3d.org
+// SPDX-License-Identifier: MIT
+// ----------------------------------------------------------------------------
+
+#pragma once
+
+#include "open3d/core/Tensor.h"
+
+namespace open3d {
+namespace t {
+namespace geometry {
+namespace kernel {
+namespace pcapartition {
+
+/// Partition the point cloud by recursively doing PCA.
+/// \param points Points tensor with shape (N,3).
+/// \param max_points The maximum allowed number of points in a partition.
+/// \return The number of partitions and an int32 tensor with the partition id
+/// for each point. The output tensor uses always the CPU device.
+std::tuple<int, core::Tensor> PCAPartition(core::Tensor& points,
+                                           int max_points);
+
+}  // namespace pcapartition
+}  // namespace kernel
+}  // namespace geometry
+}  // namespace t
+}  // namespace open3d
diff --git a/cpp/open3d/t/geometry/kernel/UVUnwrapping.cpp b/cpp/open3d/t/geometry/kernel/UVUnwrapping.cpp
index d6bcb93e..866e620d 100644
--- a/cpp/open3d/t/geometry/kernel/UVUnwrapping.cpp
+++ b/cpp/open3d/t/geometry/kernel/UVUnwrapping.cpp
@@ -7,7 +7,11 @@
 
 #include "open3d/t/geometry/kernel/UVUnwrapping.h"
 
+// clang-format off
+// include tbb before uvatlas
+#include <tbb/parallel_for.h>
 #include <UVAtlas.h>
+// clang-format on
 
 namespace open3d {
 namespace t {
@@ -15,15 +19,25 @@ namespace geometry {
 namespace kernel {
 namespace uvunwrapping {
 
-void ComputeUVAtlas(TriangleMesh& mesh,
-                    const size_t width,
-                    const size_t height,
-                    const float gutter,
-                    const float max_stretch,
-                    float* max_stretch_out,
-                    size_t* num_charts_out) {
-    using namespace DirectX;
+namespace {
+using namespace DirectX;
+
+struct UVAtlasPartitionOutput {
+    std::vector<int64_t> original_face_idx;
+    std::vector<UVAtlasVertex> vb;
+    // raw index buffer. Elements are uint32_t for DXGI_FORMAT_R32_UINT.
+    std::vector<uint8_t> ib;
+    std::vector<uint32_t> partition_result_adjacency;
+    float max_stretch_out;
+    size_t num_charts_out;
+};
+
+void ComputeUVAtlasPartition(TriangleMesh mesh,
+                             const float max_stretch,
+                             bool fast_mode,
+                             UVAtlasPartitionOutput& output) {
     const int64_t num_verts = mesh.GetVertexPositions().GetLength();
+
     std::unique_ptr<XMFLOAT3[]> pos(new XMFLOAT3[num_verts]);
     {
         core::Tensor vertices = mesh.GetVertexPositions()
@@ -94,26 +108,166 @@ void ComputeUVAtlas(TriangleMesh& mesh,
     // UVAtlas will create new vertices. remap stores the original vertex id for
     // all created vertices.
     std::vector<uint32_t> remap;
+    std::vector<uint32_t> face_partitioning;
+    std::vector<uint32_t> partition_result_adjacency;
 
-    HRESULT hr = UVAtlasCreate(
+    HRESULT hr = UVAtlasPartition(
             pos.get(), num_verts, triangles_ptr, DXGI_FORMAT_R32_UINT,
-            num_triangles, 0, max_stretch, width, height, gutter, adj.data(),
-            nullptr, nullptr, nullptr, UVATLAS_DEFAULT_CALLBACK_FREQUENCY,
-            UVATLAS_DEFAULT, vb, ib, nullptr, &remap, max_stretch_out,
-            num_charts_out);
+            num_triangles, 0, max_stretch, adj.data(), nullptr, nullptr,
+            nullptr, UVATLAS_DEFAULT_CALLBACK_FREQUENCY,
+            fast_mode ? UVATLAS_GEODESIC_FAST : UVATLAS_DEFAULT, vb, ib,
+            &face_partitioning, &remap, partition_result_adjacency,
+            &output.max_stretch_out, &output.num_charts_out);
+
     if (FAILED(hr)) {
         if (hr == static_cast<HRESULT>(0x8007000DL)) {
-            utility::LogError("UVAtlasCreate: Non-manifold mesh");
+            utility::LogError("UVAtlasPartition: Non-manifold mesh");
         } else if (hr == static_cast<HRESULT>(0x80070216L)) {
-            utility::LogError("UVAtlasCreate: Arithmetic overflow");
+            utility::LogError("UVAtlasPartition: Arithmetic overflow");
         } else if (hr == static_cast<HRESULT>(0x80070032L)) {
-            utility::LogError("UVAtlasCreate: Not supported");
+            utility::LogError("UVAtlasPartition: Not supported");
         }
-        utility::LogError("UVAtlasCreate failed with code 0x{:X}",
+        utility::LogError("UVAtlasPartition failed with code 0x{:X}",
                           static_cast<uint32_t>(hr));
     }
 
-    const uint32_t* out_triangles_ptr = reinterpret_cast<uint32_t*>(ib.data());
+    output.original_face_idx =
+            mesh.GetTriangleAttr("original_idx").ToFlatVector<int64_t>();
+    output.ib = std::move(ib);
+    output.vb = std::move(vb);
+    output.partition_result_adjacency = std::move(partition_result_adjacency);
+}
+}  // namespace
+
+std::tuple<float, int, int> ComputeUVAtlas(TriangleMesh& mesh,
+                                           const size_t width,
+                                           const size_t height,
+                                           const float gutter,
+                                           const float max_stretch,
+                                           int parallel_partitions,
+                                           int nthreads) {
+    const int64_t num_verts = mesh.GetVertexPositions().GetLength();
+
+    // create temporary mesh for partitioning
+    TriangleMesh mesh_tmp(
+            mesh.GetVertexPositions().To(core::Device()).Contiguous(),
+            mesh.GetTriangleIndices().To(core::Device()).Contiguous());
+    if (parallel_partitions > 1) {
+        const int max_points_per_partition =
+                (num_verts - 1) / (parallel_partitions - 1);
+        parallel_partitions = mesh_tmp.PCAPartition(max_points_per_partition);
+    }
+    utility::LogInfo("actual parallel_partitions {}", parallel_partitions);
+    mesh_tmp.SetTriangleAttr(
+            "original_idx",
+            core::Tensor::Arange(0, mesh_tmp.GetTriangleIndices().GetLength(),
+                                 1, core::Int64));
+
+    std::vector<TriangleMesh> mesh_partitions;
+    if (parallel_partitions > 1) {
+        for (int i = 0; i < parallel_partitions; ++i) {
+            core::Tensor mask = mesh_tmp.GetTriangleAttr("partition_ids").Eq(i);
+            mesh_partitions.emplace_back(mesh_tmp.SelectFacesByMask(mask));
+        }
+    } else {
+        mesh_partitions.emplace_back(mesh_tmp);
+    }
+
+    std::vector<UVAtlasPartitionOutput> uvatlas_partitions(parallel_partitions);
+
+    // By default UVAtlas uses a fast mode if there are more than 25k faces.
+    // This makes sure that we always use fast mode if we parallelize.
+    const bool fast_mode = parallel_partitions > 1;
+    auto LoopFn = [&](const tbb::blocked_range<size_t>& range) {
+        for (size_t i = range.begin(); i < range.end(); ++i) {
+            auto& output = uvatlas_partitions[i];
+            ComputeUVAtlasPartition(mesh_partitions[i], max_stretch, fast_mode,
+                                    output);
+        }
+    };
+
+    if (parallel_partitions > 1) {
+        if (nthreads > 0) {
+            tbb::task_arena arena(nthreads);
+            arena.execute([&]() {
+                tbb::parallel_for(
+                        tbb::blocked_range<size_t>(0, parallel_partitions),
+                        LoopFn);
+            });
+        } else {
+            tbb::parallel_for(
+                    tbb::blocked_range<size_t>(0, parallel_partitions), LoopFn);
+        }
+    } else {
+        LoopFn(tbb::blocked_range<size_t>(0, parallel_partitions));
+    }
+
+    // merge outputs for the packing algorithm
+    UVAtlasPartitionOutput& combined_output = uvatlas_partitions.front();
+    for (int i = 1; i < parallel_partitions; ++i) {
+        auto& output = uvatlas_partitions[i];
+
+        // append vectors and update indices
+        const uint32_t vidx_offset = combined_output.vb.size();
+        uint32_t* indices_ptr = reinterpret_cast<uint32_t*>(output.ib.data());
+        const int64_t num_indices = output.ib.size() / sizeof(uint32_t);
+        for (int64_t indices_i = 0; indices_i < num_indices; ++indices_i) {
+            indices_ptr[indices_i] += vidx_offset;
+        }
+        combined_output.vb.insert(combined_output.vb.end(), output.vb.begin(),
+                                  output.vb.end());
+
+        const uint32_t fidx_offset =
+                combined_output.ib.size() / (sizeof(uint32_t) * 3);
+        const uint32_t invalid = std::numeric_limits<uint32_t>::max();
+        for (auto& x : output.partition_result_adjacency) {
+            if (x != invalid) {
+                x += fidx_offset;
+            }
+        }
+        combined_output.ib.insert(combined_output.ib.end(), output.ib.begin(),
+                                  output.ib.end());
+        combined_output.partition_result_adjacency.insert(
+                combined_output.partition_result_adjacency.end(),
+                output.partition_result_adjacency.begin(),
+                output.partition_result_adjacency.end());
+
+        combined_output.original_face_idx.insert(
+                combined_output.original_face_idx.end(),
+                output.original_face_idx.begin(),
+                output.original_face_idx.end());
+
+        // update stats
+        combined_output.max_stretch_out = std::max(
+                combined_output.max_stretch_out, output.max_stretch_out);
+        combined_output.num_charts_out += output.num_charts_out;
+
+        // free memory
+        output = UVAtlasPartitionOutput();
+    }
+
+    HRESULT hr = UVAtlasPack(combined_output.vb, combined_output.ib,
+                             DXGI_FORMAT_R32_UINT, width, height, gutter,
+                             combined_output.partition_result_adjacency,
+                             nullptr, UVATLAS_DEFAULT_CALLBACK_FREQUENCY);
+
+    if (FAILED(hr)) {
+        if (hr == static_cast<HRESULT>(0x8007000DL)) {
+            utility::LogError("UVAtlasPack: Non-manifold mesh");
+        } else if (hr == static_cast<HRESULT>(0x80070216L)) {
+            utility::LogError("UVAtlasPack: Arithmetic overflow");
+        } else if (hr == static_cast<HRESULT>(0x80070032L)) {
+            utility::LogError("UVAtlasPack: Not supported");
+        }
+        utility::LogError("UVAtlasPack failed with code 0x{:X}",
+                          static_cast<uint32_t>(hr));
+    }
+
+    auto& ib = combined_output.ib;
+    auto& vb = combined_output.vb;
+    auto& original_face_idx = combined_output.original_face_idx;
+    const int64_t num_triangles = mesh.GetTriangleIndices().GetLength();
+    const uint32_t* indices_ptr = reinterpret_cast<uint32_t*>(ib.data());
     if (ib.size() != sizeof(uint32_t) * 3 * num_triangles) {
         utility::LogError(
                 "Unexpected output index buffer size. Got {} expected {}.",
@@ -124,23 +278,21 @@ void ComputeUVAtlas(TriangleMesh& mesh,
         // copy uv coords
         float* texture_uvs_ptr = texture_uvs.GetDataPtr<float>();
         for (int64_t i = 0; i < num_triangles; ++i) {
+            int64_t original_i = original_face_idx[i];
             for (int j = 0; j < 3; ++j) {
-                const uint32_t& vidx = out_triangles_ptr[i * 3 + j];
+                const uint32_t& vidx = indices_ptr[i * 3 + j];
                 const auto& vert = vb[vidx];
-                texture_uvs_ptr[i * 6 + j * 2 + 0] = vert.uv.x;
-                texture_uvs_ptr[i * 6 + j * 2 + 1] = vert.uv.y;
-
-                if (remap[vidx] != triangles_ptr[i * 3 + j]) {
-                    utility::LogError(
-                            "Output index buffer differs from input index "
-                            "buffer.");
-                }
+                texture_uvs_ptr[original_i * 6 + j * 2 + 0] = vert.uv.x;
+                texture_uvs_ptr[original_i * 6 + j * 2 + 1] = vert.uv.y;
             }
         }
     }
 
     texture_uvs = texture_uvs.To(mesh.GetDevice());
     mesh.SetTriangleAttr("texture_uvs", texture_uvs);
+
+    return std::tie(combined_output.max_stretch_out,
+                    combined_output.num_charts_out, parallel_partitions);
 }
 
 }  // namespace uvunwrapping
diff --git a/cpp/open3d/t/geometry/kernel/UVUnwrapping.h b/cpp/open3d/t/geometry/kernel/UVUnwrapping.h
index 972ba565..e2262b15 100644
--- a/cpp/open3d/t/geometry/kernel/UVUnwrapping.h
+++ b/cpp/open3d/t/geometry/kernel/UVUnwrapping.h
@@ -40,13 +40,26 @@ namespace uvunwrapping {
 /// of stretch.
 ///
 /// \param num_charts_out Output parameter with the number of charts created.
-void ComputeUVAtlas(TriangleMesh& mesh,
-                    const size_t width = 512,
-                    const size_t height = 512,
-                    const float gutter = 1.0f,
-                    const float max_stretch = 1.f / 6,
-                    float* max_stretch_out = nullptr,
-                    size_t* num_charts_out = nullptr);
+///
+/// \param parallel_partitions The approximate number of partitions created
+/// before computing the UV atlas for parallelizing the computation.
+/// Parallelization can be enabled with values > 1. Note that
+/// parallelization increases the number of UV islands and can lead to results
+/// with lower quality.
+///
+/// \param nthreads The number of threads used when parallel_partitions
+/// is > 1. Set to 0 for automatic number of thread detection.
+///
+/// \return Tuple with (max stretch, num_charts, num_partitions) storing the
+/// actual amount of stretch, the number of created charts, and the number of
+/// parallel partitions created.
+std::tuple<float, int, int> ComputeUVAtlas(TriangleMesh& mesh,
+                                           const size_t width = 512,
+                                           const size_t height = 512,
+                                           const float gutter = 1.0f,
+                                           const float max_stretch = 1.f / 6,
+                                           int parallel_partitions = 1,
+                                           int nthreads = 0);
 
 }  // namespace uvunwrapping
 }  // namespace kernel
diff --git a/cpp/open3d/utility/Eigen.cpp b/cpp/open3d/utility/Eigen.cpp
index 35396309..b3f6fb34 100644
--- a/cpp/open3d/utility/Eigen.cpp
+++ b/cpp/open3d/utility/Eigen.cpp
@@ -365,6 +365,51 @@ template std::tuple<Eigen::Vector3d, Eigen::Matrix3d> ComputeMeanAndCovariance(
         const std::vector<Eigen::Vector3d> &points,
         const std::vector<int> &indices);
 
+template <typename RealType, typename IdxType>
+std::tuple<Eigen::Vector3d, Eigen::Matrix3d> ComputeMeanAndCovariance(
+        const RealType *const points, const std::vector<IdxType> &indices) {
+    Eigen::Vector3d mean;
+    Eigen::Matrix3d covariance;
+    Eigen::Matrix<double, 9, 1> cumulants;
+    cumulants.setZero();
+    for (const auto &idx : indices) {
+        const Eigen::Map<const Eigen::Matrix<RealType, 3, 1>> point(points +
+                                                                    3 * idx);
+        cumulants(0) += point(0);
+        cumulants(1) += point(1);
+        cumulants(2) += point(2);
+        cumulants(3) += point(0) * point(0);
+        cumulants(4) += point(0) * point(1);
+        cumulants(5) += point(0) * point(2);
+        cumulants(6) += point(1) * point(1);
+        cumulants(7) += point(1) * point(2);
+        cumulants(8) += point(2) * point(2);
+    }
+    cumulants /= (double)indices.size();
+    mean(0) = cumulants(0);
+    mean(1) = cumulants(1);
+    mean(2) = cumulants(2);
+    covariance(0, 0) = cumulants(3) - cumulants(0) * cumulants(0);
+    covariance(1, 1) = cumulants(6) - cumulants(1) * cumulants(1);
+    covariance(2, 2) = cumulants(8) - cumulants(2) * cumulants(2);
+    covariance(0, 1) = cumulants(4) - cumulants(0) * cumulants(1);
+    covariance(1, 0) = covariance(0, 1);
+    covariance(0, 2) = cumulants(5) - cumulants(0) * cumulants(2);
+    covariance(2, 0) = covariance(0, 2);
+    covariance(1, 2) = cumulants(7) - cumulants(1) * cumulants(2);
+    covariance(2, 1) = covariance(1, 2);
+    return std::make_tuple(mean, covariance);
+}
+
+template std::tuple<Eigen::Vector3d, Eigen::Matrix3d> ComputeMeanAndCovariance(
+        const float *const points, const std::vector<size_t> &indices);
+template std::tuple<Eigen::Vector3d, Eigen::Matrix3d> ComputeMeanAndCovariance(
+        const double *const points, const std::vector<size_t> &indices);
+template std::tuple<Eigen::Vector3d, Eigen::Matrix3d> ComputeMeanAndCovariance(
+        const float *const points, const std::vector<int> &indices);
+template std::tuple<Eigen::Vector3d, Eigen::Matrix3d> ComputeMeanAndCovariance(
+        const double *const points, const std::vector<int> &indices);
+
 Eigen::Matrix3d SkewMatrix(const Eigen::Vector3d &vec) {
     Eigen::Matrix3d skew;
     // clang-format off
diff --git a/cpp/open3d/utility/Eigen.h b/cpp/open3d/utility/Eigen.h
index 07e53d58..d3438d6d 100644
--- a/cpp/open3d/utility/Eigen.h
+++ b/cpp/open3d/utility/Eigen.h
@@ -123,5 +123,16 @@ template <typename IdxType>
 std::tuple<Eigen::Vector3d, Eigen::Matrix3d> ComputeMeanAndCovariance(
         const std::vector<Eigen::Vector3d> &points,
         const std::vector<IdxType> &indices);
+
+/// Function to compute the mean and covariance matrix of a set of points.
+/// \tparam RealType Either float or double.
+/// \tparam IdxType Either size_t or int.
+/// \param points Contiguous memory with the 3D points.
+/// \param indices The indices for which the mean and covariance will be
+/// computed. \return The mean and covariance matrix.
+template <typename RealType, typename IdxType>
+std::tuple<Eigen::Vector3d, Eigen::Matrix3d> ComputeMeanAndCovariance(
+        const RealType *const points, const std::vector<IdxType> &indices);
+
 }  // namespace utility
 }  // namespace open3d
diff --git a/cpp/pybind/t/geometry/pointcloud.cpp b/cpp/pybind/t/geometry/pointcloud.cpp
index 0e072bc0..cc7f98ec 100644
--- a/cpp/pybind/t/geometry/pointcloud.cpp
+++ b/cpp/pybind/t/geometry/pointcloud.cpp
@@ -661,6 +661,30 @@ Example:
         lines = pcd.extrude_linear([0,1,0])
         o3d.visualization.draw([{'name': 'lines', 'geometry': lines}])
 
+)");
+
+    pointcloud.def("pca_partition", &PointCloud::PCAPartition, "max_points"_a,
+                   R"(Partition the point cloud by recursively doing PCA.
+
+This function creates a new point attribute with the name "partition_ids" storing 
+the partition id for each point.
+
+Args:
+    max_points (int): The maximum allowed number of points in a partition.
+
+
+Example:
+
+    This code computes parititions a point cloud such that each partition
+    contains at most 20 points::
+
+        import open3d as o3d
+        import numpy as np
+        pcd = o3d.t.geometry.PointCloud(np.random.rand(100,3))
+        num_partitions = pcd.pca_partition(max_points=20)
+
+        # print the partition ids and the number of points for each of them.
+        print(np.unique(pcd.point.partition_ids.numpy(), return_counts=True))
 
 )");
 }
diff --git a/cpp/pybind/t/geometry/trianglemesh.cpp b/cpp/pybind/t/geometry/trianglemesh.cpp
index 582bb89e..2954495f 100644
--- a/cpp/pybind/t/geometry/trianglemesh.cpp
+++ b/cpp/pybind/t/geometry/trianglemesh.cpp
@@ -197,12 +197,12 @@ The attributes of the triangle mesh have different levels::
             R"(Compute the convex hull of a point cloud using qhull. This runs on the CPU.
 
 Args:
-    joggle_inputs (default False). Handle precision problems by
-    randomly perturbing the input data. Set to True if perturbing the input
-    iis acceptable but you need convex simplicial output. If False,
-    neighboring facets may be merged in case of precision problems. See
-    `QHull docs <http://www.qhull.org/html/qh-impre.htm#joggle`__ for more
-    details.
+    joggle_inputs (bool with default False): Handle precision problems by
+        randomly perturbing the input data. Set to True if perturbing the input
+        iis acceptable but you need convex simplicial output. If False,
+        neighboring facets may be merged in case of precision problems. See
+        `QHull docs <http://www.qhull.org/html/qh-impre.htm#joggle`__ for more
+        details.
 
 Returns:
     TriangleMesh representing the convexh hull. This contains an
@@ -663,6 +663,7 @@ Example:
     triangle_mesh.def(
             "compute_uvatlas", &TriangleMesh::ComputeUVAtlas, "size"_a = 512,
             "gutter"_a = 1.f, "max_stretch"_a = 1.f / 6,
+            "parallel_partitions"_a = 1, "nthreads"_a = 0,
             R"(Creates an UV atlas and adds it as triangle attr 'texture_uvs' to the mesh.
 
 Input meshes must be manifold for this method to work.
@@ -671,6 +672,7 @@ Zhou et al, "Iso-charts: Stretch-driven Mesh Parameterization using Spectral
              Analysis", Eurographics Symposium on Geometry Processing (2004)
 Sander et al. "Signal-Specialized Parametrization" Europgraphics 2002
 This function always uses the CPU device.
+
 Args:
     size (int): The target size of the texture (size x size). The uv coordinates
         will still be in the range [0..1] but parameters like gutter use pixels
@@ -678,10 +680,25 @@ Args:
     gutter (float): This is the space around the uv islands in pixels.
     max_stretch (float): The maximum amount of stretching allowed. The parameter
         range is [0..1] with 0 meaning no stretch allowed.
+
+    parallel_partitions (int): The approximate number of partitions created
+        before computing the UV atlas for parallelizing the computation.
+        Parallelization can be enabled with values > 1. Note that
+        parallelization increases the number of UV islands and can lead to results
+        with lower quality.
+
+    nthreads (int): The number of threads used when parallel_partitions
+        is > 1. Set to 0 for automatic number of thread detection.
+
 Returns:
-    None. This function modifies the mesh in-place.
+    This function creates a face attribute "texture_uvs" and returns a tuple
+    with (max stretch, num_charts, num_partitions) storing the 
+    actual amount of stretch, the number of created charts, and the number of
+    parallel partitions created.
+
 Example:
     This code creates a uv map for the Stanford Bunny mesh::
+
         import open3d as o3d
         bunny = o3d.data.BunnyMesh()
         mesh = o3d.t.geometry.TriangleMesh.from_legacy(o3d.io.read_triangle_mesh(bunny.path))
@@ -727,6 +744,7 @@ Returns:
 
 Example:
     We generate a texture storing the xyz coordinates for each texel::
+
         import open3d as o3d
         from matplotlib import pyplot as plt
 
@@ -781,6 +799,7 @@ Returns:
 Example:
     We generate a texture visualizing the index of the triangle to which the
     texel belongs to::
+
         import open3d as o3d
         from matplotlib import pyplot as plt
 
@@ -809,16 +828,17 @@ Example:
                       R"(Sweeps the triangle mesh rotationally about an axis.
 Args:
     angle (float): The rotation angle in degree.
-
     axis (open3d.core.Tensor): The rotation axis.
-
     resolution (int): The resolution defines the number of intermediate sweeps
         about the rotation axis.
     translation (float): The translation along the rotation axis.
+
 Returns:
     A triangle mesh with the result of the sweep operation.
+
 Example:
     This code generates a spring with a triangle cross-section::
+
         import open3d as o3d
 
         mesh = o3d.t.geometry.TriangleMesh([[1,1,0], [0.7,1,0], [1,0.7,0]], [[0,1,2]])
@@ -830,19 +850,81 @@ Example:
                       "vector"_a, "scale"_a = 1.0, "capping"_a = true,
                       R"(Sweeps the line set along a direction vector.
 Args:
-
     vector (open3d.core.Tensor): The direction vector.
-
     scale (float): Scalar factor which essentially scales the direction vector.
+
 Returns:
     A triangle mesh with the result of the sweep operation.
+
 Example:
     This code generates a wedge from a triangle::
+
         import open3d as o3d
         triangle = o3d.t.geometry.TriangleMesh([[1.0,1.0,0.0], [0,1,0], [1,0,0]], [[0,1,2]])
         wedge = triangle.extrude_linear([0,0,1])
         o3d.visualization.draw([{'name': 'wedge', 'geometry': wedge}])
 )");
+
+    triangle_mesh.def("pca_partition", &TriangleMesh::PCAPartition,
+                      "max_faces"_a,
+                      R"(Partition the mesh by recursively doing PCA.
+
+This function creates a new face attribute with the name "partition_ids" storing 
+the partition id for each face.
+
+Args:
+    max_faces (int): The maximum allowed number of faces in a partition.
+
+
+Example:
+
+    This code partitions a mesh such that each partition contains at most 20k 
+    faces::
+
+        import open3d as o3d
+        import numpy as np
+        bunny = o3d.data.BunnyMesh()
+        mesh = o3d.t.geometry.TriangleMesh.from_legacy(o3d.io.read_triangle_mesh(bunny.path))
+        num_partitions = mesh.pca_partition(max_faces=20000)
+
+        # print the partition ids and the number of faces for each of them.
+        print(np.unique(mesh.triangle.partition_ids.numpy(), return_counts=True))
+
+)");
+
+    triangle_mesh.def(
+            "select_faces_by_mask", &TriangleMesh::SelectFacesByMask, "mask"_a,
+            R"(Returns a new mesh with the faces selected by a boolean mask.
+
+Args:
+    mask ()
+    mask (open3d.core.Tensor): A boolean mask with the shape (N) with N as the 
+        number of faces in the mesh.
+    
+Returns:
+    A new mesh with the selected faces.
+
+Example:
+
+    This code partitions the mesh using PCA and then visualized the individual 
+    parts::
+
+        import open3d as o3d
+        import numpy as np
+        bunny = o3d.data.BunnyMesh()
+        mesh = o3d.t.geometry.TriangleMesh.from_legacy(o3d.io.read_triangle_mesh(bunny.path))
+        num_partitions = mesh.pca_partition(max_faces=20000)
+
+        parts = []
+        for i in range(num_partitions):
+            mask = mesh.triangle.partition_ids == i
+            part = mesh.select_faces_by_mask(mask)
+            part.vertex.colors = np.tile(np.random.rand(3), (part.vertex.positions.shape[0],1))
+            parts.append(part)
+
+        o3d.visualization.draw(parts)
+
+)");
 }
 
 }  // namespace geometry