Skip to content

Commit e348357

Browse files
nicolaloinicolaloi
committed
Cache info to optimize computations when ratio of selected indices is high.
1 parent 3fa7f80 commit e348357

File tree

4 files changed

+173
-69
lines changed

4 files changed

+173
-69
lines changed

‎cpp/benchmarks/t/pipelines/registration/Feature.cpp

-2
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,6 @@
55
// SPDX-License-Identifier: MIT
66
// ----------------------------------------------------------------------------
77

8-
#include <map>
9-
108
#include "open3d/t/pipelines/registration/Feature.h"
119

1210
#include <benchmark/benchmark.h>

‎cpp/open3d/t/pipelines/kernel/Feature.cpp

+16-15
Original file line numberDiff line numberDiff line change
@@ -15,15 +15,15 @@ namespace t {
1515
namespace pipelines {
1616
namespace kernel {
1717

18-
void ComputeFPFHFeature(const core::Tensor &points,
19-
const core::Tensor &normals,
20-
const core::Tensor &indices,
21-
const core::Tensor &distance2,
22-
const core::Tensor &counts,
23-
core::Tensor &fpfhs,
24-
const utility::optional<core::Tensor> &mask,
25-
const utility::optional<core::Tensor>
26-
&map_batch_info_idx_to_point_idx) {
18+
void ComputeFPFHFeature(
19+
const core::Tensor &points,
20+
const core::Tensor &normals,
21+
const core::Tensor &indices,
22+
const core::Tensor &distance2,
23+
const core::Tensor &counts,
24+
core::Tensor &fpfhs,
25+
const utility::optional<core::Tensor> &mask,
26+
const utility::optional<core::Tensor> &map_info_idx_to_point_idx) {
2727
if (mask.has_value()) {
2828
const int64_t size =
2929
mask.value().To(core::Int64).Sum({0}).Item<int64_t>();
@@ -32,21 +32,22 @@ void ComputeFPFHFeature(const core::Tensor &points,
3232
} else {
3333
core::AssertTensorShape(fpfhs, {points.GetLength(), 33});
3434
}
35-
if (map_batch_info_idx_to_point_idx.has_value()) {
36-
core::AssertTensorShape(map_batch_info_idx_to_point_idx.value(),
37-
{counts.GetLength()});
35+
if (map_info_idx_to_point_idx.has_value()) {
36+
const bool is_radius_search = indices.GetShape().size() == 1;
37+
core::AssertTensorShape(
38+
map_info_idx_to_point_idx.value(),
39+
{counts.GetLength() - (is_radius_search ? 1 : 0)});
3840
}
3941
const core::Tensor points_d = points.Contiguous();
4042
const core::Tensor normals_d = normals.Contiguous();
4143
const core::Tensor counts_d = counts.To(core::Int32);
4244
if (points_d.IsCPU()) {
4345
ComputeFPFHFeatureCPU(points_d, normals_d, indices, distance2, counts_d,
44-
fpfhs, mask, map_batch_info_idx_to_point_idx);
46+
fpfhs, mask, map_info_idx_to_point_idx);
4547
} else {
4648
core::CUDAScopedDevice scoped_device(points.GetDevice());
4749
CUDA_CALL(ComputeFPFHFeatureCUDA, points_d, normals_d, indices,
48-
distance2, counts_d, fpfhs, mask,
49-
map_batch_info_idx_to_point_idx);
50+
distance2, counts_d, fpfhs, mask, map_info_idx_to_point_idx);
5051
}
5152
utility::LogDebug(
5253
"[ComputeFPFHFeature] Computed {:d} features from "

‎cpp/open3d/t/pipelines/kernel/FeatureImpl.h

+10-12
Original file line numberDiff line numberDiff line change
@@ -116,19 +116,17 @@ void ComputeFPFHFeatureCPU
116116
const core::Tensor &counts,
117117
core::Tensor &fpfhs,
118118
const utility::optional<core::Tensor> &mask,
119-
const utility::optional<core::Tensor>
120-
&map_batch_info_idx_to_point_idx) {
119+
const utility::optional<core::Tensor> &map_info_idx_to_point_idx) {
121120
const core::Dtype dtype = points.GetDtype();
122121
const core::Device device = points.GetDevice();
123122
const int64_t n_points = points.GetLength();
124123

125124
const bool filter_fpfh =
126-
mask.has_value() && map_batch_info_idx_to_point_idx.has_value();
127-
if (mask.has_value() ^ map_batch_info_idx_to_point_idx.has_value()) {
125+
mask.has_value() && map_info_idx_to_point_idx.has_value();
126+
if (mask.has_value() ^ map_info_idx_to_point_idx.has_value()) {
128127
utility::LogError(
129-
"Parameters mask and map_batch_info_idx_to_point_idx must be "
130-
"both "
131-
"provided or both not provided.");
128+
"Parameters mask and map_info_idx_to_point_idx must "
129+
"either be both provided or both not provided.");
132130
}
133131
if (filter_fpfh) {
134132
if (mask.value().GetShape()[0] != n_points) {
@@ -137,19 +135,19 @@ void ComputeFPFHFeatureCPU
137135
"be equal to the number of points {:d}.",
138136
(int)mask.value().GetShape()[0], n_points);
139137
}
140-
if (map_batch_info_idx_to_point_idx.value().GetShape()[0] !=
141-
counts.GetShape()[0]) {
138+
if (map_info_idx_to_point_idx.value().GetShape()[0] !=
139+
counts.GetShape()[0] - (indices.GetShape().size() == 1 ? 1 : 0)) {
142140
utility::LogError(
143-
"Parameter map_batch_info_idx_to_point_idx was provided, "
141+
"Parameter map_info_idx_to_point_idx was provided, "
144142
"but its size"
145143
"{:d} should be equal to the size of counts {:d}.",
146-
(int)map_batch_info_idx_to_point_idx.value().GetShape()[0],
144+
(int)map_info_idx_to_point_idx.value().GetShape()[0],
147145
(int)counts.GetShape()[0]);
148146
}
149147
}
150148

151149
core::Tensor map_spfh_info_idx_to_point_idx =
152-
map_batch_info_idx_to_point_idx.value_or(
150+
map_info_idx_to_point_idx.value_or(
153151
core::Tensor::Empty({0}, core::Int64, device));
154152

155153
const core::Tensor map_fpfh_idx_to_point_idx =

‎cpp/open3d/t/pipelines/registration/Feature.cpp

+147-40
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
#include "open3d/t/pipelines/registration/Feature.h"
99

10+
#include "open3d/core/ParallelFor.h"
1011
#include "open3d/core/nns/NearestNeighborSearch.h"
1112
#include "open3d/t/geometry/PointCloud.h"
1213
#include "open3d/t/pipelines/kernel/Feature.h"
@@ -43,58 +44,138 @@ core::Tensor ComputeFPFHFeature(
4344
core::Int32);
4445
bool tree_set = false;
4546

47+
const bool filter_fpfh = indices.has_value();
48+
// If we are computing a subset of the FPFH feature,
49+
// cache some information to speed up the computation
50+
// if the ratio of the indices to the total number of points is high.
51+
const double cache_info_indices_ratio_thresh = 0.01;
52+
bool cache_fpfh_info = true;
53+
4654
core::Tensor mask_fpfh_points;
47-
core::Tensor mask_required_points;
48-
if (indices.has_value()) {
55+
core::Tensor indices_fpfh_points;
56+
core::Tensor map_point_idx_to_required_point_idx;
57+
core::Tensor map_required_point_idx_to_point_idx;
58+
core::Tensor save_p_indices, save_p_distance2, save_p_counts;
59+
core::Tensor mask_spfh_points;
60+
61+
// If we are computing a subset of the FPFH feature, we need to find
62+
// the subset of points (neighbors) required to compute the FPFH features.
63+
if (filter_fpfh) {
64+
if (indices.value().GetLength() == 0) {
65+
return core::Tensor::Zeros({0, 33}, dtype, device);
66+
}
4967
mask_fpfh_points =
5068
core::Tensor::Zeros({num_points}, core::Bool, device);
51-
mask_required_points =
52-
core::Tensor::Zeros({num_points}, core::Bool, device);
53-
core::Tensor indices_tmp, distance2_tmp, counts_tmp;
5469
mask_fpfh_points.IndexSet({indices.value()},
5570
core::Tensor::Ones({1}, core::Bool, device));
5671
const core::Tensor query_point_positions =
57-
input.GetPointPositions().IndexGet({indices.value()});
72+
input.GetPointPositions().IndexGet({mask_fpfh_points});
73+
core::Tensor p_indices, p_distance2, p_counts;
5874
if (radius.has_value() && max_nn.has_value()) {
5975
tree_set = tree.HybridIndex(radius.value());
6076
if (!tree_set) {
6177
utility::LogError("Building HybridIndex failed.");
6278
}
63-
std::tie(indices_tmp, distance2_tmp, counts_tmp) =
64-
tree.HybridSearch(query_point_positions, radius.value(),
65-
max_nn.value());
79+
std::tie(p_indices, p_distance2, p_counts) = tree.HybridSearch(
80+
query_point_positions, radius.value(), max_nn.value());
6681
} else if (!radius.has_value() && max_nn.has_value()) {
6782
tree_set = tree.KnnIndex();
6883
if (!tree_set) {
6984
utility::LogError("Building KnnIndex failed.");
7085
}
71-
std::tie(indices_tmp, distance2_tmp) =
86+
std::tie(p_indices, p_distance2) =
7287
tree.KnnSearch(query_point_positions, max_nn.value());
88+
89+
// Make counts full with min(max_nn, num_points).
90+
const int fill_value =
91+
max_nn.value() > num_points ? num_points : max_nn.value();
92+
p_counts = core::Tensor::Full({query_point_positions.GetLength()},
93+
fill_value, core::Int32, device);
7394
} else if (radius.has_value() && !max_nn.has_value()) {
7495
tree_set = tree.FixedRadiusIndex(radius.value());
7596
if (!tree_set) {
7697
utility::LogError("Building RadiusIndex failed.");
7798
}
78-
std::tie(indices_tmp, distance2_tmp, counts_tmp) =
79-
tree.FixedRadiusSearch(query_point_positions,
80-
radius.value());
99+
std::tie(p_indices, p_distance2, p_counts) = tree.FixedRadiusSearch(
100+
query_point_positions, radius.value());
81101
} else {
82102
utility::LogError("Both max_nn and radius are none.");
83103
}
84104

85-
indices_tmp = indices_tmp.To(core::Int64).View({-1});
105+
core::Tensor mask_required_points =
106+
core::Tensor::Zeros({num_points}, core::Bool, device);
86107
mask_required_points.IndexSet(
87-
{indices_tmp}, core::Tensor::Ones({1}, core::Bool, device));
108+
{p_indices.To(core::Int64).View({-1})},
109+
core::Tensor::Ones({1}, core::Bool, device));
110+
map_required_point_idx_to_point_idx =
111+
mask_required_points.NonZero().GetItem(
112+
{core::TensorKey::Index(0)});
113+
indices_fpfh_points =
114+
mask_fpfh_points.NonZero().GetItem({core::TensorKey::Index(0)});
88115

89-
} else {
90-
mask_fpfh_points = core::Tensor::Zeros({0}, core::Bool, device);
91-
mask_required_points = core::Tensor::Zeros({0}, core::Bool, device);
116+
const bool is_radius_search = p_indices.GetShape().size() == 1;
117+
118+
// Cache the info if the ratio of the indices to the total number of
119+
// points is high and we are not doing a radius search. Radius search
120+
// requires a different pipeline since tensor output p_counts is a
121+
// prefix sum.
122+
cache_fpfh_info = !is_radius_search &&
123+
(indices_fpfh_points.GetLength() >=
124+
cache_info_indices_ratio_thresh * num_points);
125+
126+
if (cache_fpfh_info) {
127+
map_point_idx_to_required_point_idx =
128+
core::Tensor::Full({num_points}, -1, core::Int32, device);
129+
map_point_idx_to_required_point_idx.IndexSet(
130+
{map_required_point_idx_to_point_idx},
131+
core::Tensor::Arange(
132+
0, map_required_point_idx_to_point_idx.GetLength(),
133+
1, core::Int32, device));
134+
135+
core::SizeVector save_p_indices_shape = p_indices.GetShape();
136+
save_p_indices_shape[0] =
137+
map_required_point_idx_to_point_idx.GetLength();
138+
save_p_indices = core::Tensor::Zeros(save_p_indices_shape,
139+
core::Int32, device);
140+
save_p_distance2 = core::Tensor::Zeros(save_p_indices.GetShape(),
141+
dtype, device);
142+
save_p_counts = core::Tensor::Zeros(
143+
{map_required_point_idx_to_point_idx.GetLength() +
144+
(is_radius_search ? 1 : 0)},
145+
core::Int32, device);
146+
147+
core::Tensor map_fpfh_point_idx_to_required_point_idx =
148+
map_point_idx_to_required_point_idx
149+
.IndexGet({indices_fpfh_points})
150+
.To(core::Int64);
151+
152+
save_p_indices.IndexSet({map_fpfh_point_idx_to_required_point_idx},
153+
p_indices);
154+
save_p_distance2.IndexSet(
155+
{map_fpfh_point_idx_to_required_point_idx}, p_distance2);
156+
save_p_counts.IndexSet({map_fpfh_point_idx_to_required_point_idx},
157+
p_counts);
158+
159+
// If we are filtering FPFH features, we have already computed some
160+
// info about the FPFH points' neighbors. Now we just need to
161+
// compute the info for the remaining required points, so skip the
162+
// computation for the already computed info.
163+
mask_spfh_points =
164+
core::Tensor::Zeros({num_points}, core::Bool, device);
165+
mask_spfh_points.IndexSet(
166+
{map_required_point_idx_to_point_idx},
167+
core::Tensor::Ones({1}, core::Bool, device));
168+
mask_spfh_points.IndexSet(
169+
{indices_fpfh_points},
170+
core::Tensor::Zeros({1}, core::Bool, device));
171+
} else {
172+
mask_spfh_points = mask_required_points;
173+
}
92174
}
93175

94176
const core::Tensor query_point_positions =
95-
mask_required_points.GetShape()[0] > 0
96-
? input.GetPointPositions().IndexGet({mask_required_points})
97-
: input.GetPointPositions();
177+
filter_fpfh ? input.GetPointPositions().IndexGet({mask_spfh_points})
178+
: input.GetPointPositions();
98179

99180
// Compute nearest neighbors and squared distances.
100181
core::Tensor p_indices, p_distance2, p_counts;
@@ -119,14 +200,25 @@ core::Tensor ComputeFPFHFeature(
119200
utility::LogError("Building KnnIndex failed.");
120201
}
121202
}
122-
std::tie(p_indices, p_distance2) =
123-
tree.KnnSearch(query_point_positions, max_nn.value());
124-
125-
// Make counts full with min(max_nn, num_points).
126-
const int fill_value =
127-
max_nn.value() > num_points ? num_points : max_nn.value();
128-
p_counts = core::Tensor::Full({query_point_positions.GetLength()},
129-
fill_value, core::Int32, device);
203+
204+
// tree.KnnSearch complains if the query point cloud is empty.
205+
if (query_point_positions.GetLength() > 0) {
206+
std::tie(p_indices, p_distance2) =
207+
tree.KnnSearch(query_point_positions, max_nn.value());
208+
209+
const int fill_value =
210+
max_nn.value() > num_points ? num_points : max_nn.value();
211+
212+
p_counts = core::Tensor::Full({query_point_positions.GetLength()},
213+
fill_value, core::Int32, device);
214+
} else {
215+
p_indices = core::Tensor::Zeros({0, max_nn.value()}, core::Int32,
216+
device);
217+
p_distance2 =
218+
core::Tensor::Zeros({0, max_nn.value()}, dtype, device);
219+
p_counts = core::Tensor::Zeros({0}, core::Int32, device);
220+
}
221+
130222
utility::LogDebug(
131223
"Use KNNSearch [max_nn: {}] for computing FPFH feature.",
132224
max_nn.value());
@@ -147,18 +239,33 @@ core::Tensor ComputeFPFHFeature(
147239
}
148240

149241
core::Tensor fpfh;
150-
if (indices.has_value()) {
151-
const auto mask_fpfh_points_indices =
152-
mask_fpfh_points.NonZero().GetItem({core::TensorKey::Index(0)});
153-
const auto map_batch_info_idx_to_point_idx =
154-
mask_required_points.NonZero().GetItem(
155-
{core::TensorKey::Index(0)});
156-
fpfh = core::Tensor::Zeros({mask_fpfh_points_indices.GetLength(), 33},
157-
dtype, device);
242+
if (filter_fpfh) {
243+
const int64_t size = indices_fpfh_points.GetLength();
244+
fpfh = core::Tensor::Zeros({size, 33}, dtype, device);
245+
core::Tensor final_p_indices, final_p_distance2, final_p_counts;
246+
if (cache_fpfh_info) {
247+
core::Tensor map_spfh_idx_to_required_point_idx =
248+
map_point_idx_to_required_point_idx
249+
.IndexGet({mask_spfh_points})
250+
.To(core::Int64);
251+
save_p_indices.IndexSet({map_spfh_idx_to_required_point_idx},
252+
p_indices);
253+
save_p_distance2.IndexSet({map_spfh_idx_to_required_point_idx},
254+
p_distance2);
255+
save_p_counts.IndexSet({map_spfh_idx_to_required_point_idx},
256+
p_counts);
257+
final_p_indices = save_p_indices;
258+
final_p_distance2 = save_p_distance2;
259+
final_p_counts = save_p_counts;
260+
} else {
261+
final_p_indices = p_indices;
262+
final_p_distance2 = p_distance2;
263+
final_p_counts = p_counts;
264+
}
158265
pipelines::kernel::ComputeFPFHFeature(
159-
input.GetPointPositions(), input.GetPointNormals(), p_indices,
160-
p_distance2, p_counts, fpfh, mask_fpfh_points,
161-
map_batch_info_idx_to_point_idx);
266+
input.GetPointPositions(), input.GetPointNormals(),
267+
final_p_indices, final_p_distance2, final_p_counts, fpfh,
268+
mask_fpfh_points, map_required_point_idx_to_point_idx);
162269
} else {
163270
const int64_t size = input.GetPointPositions().GetLength();
164271
fpfh = core::Tensor::Zeros({size, 33}, dtype, device);

0 commit comments

Comments
 (0)