bvh/v2/reinsertion_optimizer.h

0001 #ifndef BVH_V2_REINSERTION_OPTIMIZER_H
0002 #define BVH_V2_REINSERTION_OPTIMIZER_H
0003
0004 #include "bvh/v2/bvh.h"
0005 #include "bvh/v2/thread_pool.h"
0006 #include "bvh/v2/executor.h"
0007
0008 #include <vector>
0009 #include <algorithm>
0010
0011 namespace bvh::v2 {
0012
0013 template <typename Node>
0014 class ReinsertionOptimizer {
0015     using Scalar = typename Node::Scalar;
0016     using BBox = bvh::v2::BBox<Scalar, Node::dimension>;
0017
0018 public:
0019     struct Config {
0020         /// Fraction of the number of nodes to optimize per iteration.
0021         Scalar batch_size_ratio = static_cast<Scalar>(0.05);
0022
0023         /// Maximum number of iterations.
0024         size_t max_iter_count = 3;
0025     };
0026
0027     static void optimize(ThreadPool& thread_pool, Bvh<Node>& bvh, const Config& config = {}) {
0028         ParallelExecutor executor(thread_pool);
0029         optimize(executor, bvh, config);
0030     }
0031
0032     static void optimize(Bvh<Node>& bvh, const Config& config = {}) {
0033         SequentialExecutor executor;
0034         optimize(executor, bvh, config);
0035     }
0036
0037 private:
0038     struct Candidate {
0039         size_t node_id = 0;
0040         Scalar cost = -std::numeric_limits<Scalar>::max();
0041
0042         BVH_ALWAYS_INLINE bool operator > (const Candidate& other) const {
0043             return cost > other.cost;
0044         }
0045     };
0046
0047     struct Reinsertion {
0048         size_t from = 0;
0049         size_t to = 0;
0050         Scalar area_diff = static_cast<Scalar>(0);
0051
0052         BVH_ALWAYS_INLINE bool operator > (const Reinsertion& other) const {
0053             return area_diff > other.area_diff;
0054         }
0055     };
0056
0057     Bvh<Node>& bvh_;
0058     std::vector<size_t> parents_;
0059
0060     ReinsertionOptimizer(Bvh<Node>& bvh, std::vector<size_t>&& parents)
0061         : bvh_(bvh)
0062         , parents_(std::move(parents))
0063     {}
0064
0065     template <typename Derived>
0066     static void optimize(Executor<Derived>& executor, Bvh<Node>& bvh, const Config& config) {
0067         auto parents = compute_parents(executor, bvh);
0068         ReinsertionOptimizer<Node>(bvh, std::move(parents)).optimize(executor, config);
0069     }
0070
0071     template <typename Derived>
0072     static std::vector<size_t> compute_parents(Executor<Derived>& executor, const Bvh<Node>& bvh) {
0073         std::vector<size_t> parents(bvh.nodes.size());
0074         parents[0] = 0;
0075         executor.for_each(0, bvh.nodes.size(),
0076             [&] (size_t begin, size_t end) {
0077                 for (size_t i = begin; i < end; ++i) {
0078                     auto& node = bvh.nodes[i];
0079                     if (!node.is_leaf()) {
0080                         parents[node.index.first_id() + 0] = i;
0081                         parents[node.index.first_id() + 1] = i;
0082                     }
0083                 }
0084             });
0085         return parents;
0086     }
0087
0088     BVH_ALWAYS_INLINE std::vector<Candidate> find_candidates(size_t target_count) {
0089         // Gather the `target_count` nodes that have the highest cost.
0090         // Note that this may produce fewer nodes if the BVH has fewer than `target_count` nodes.
0091         const auto node_count = std::min(bvh_.nodes.size(), target_count + 1);
0092         std::vector<Candidate> candidates;
0093         for (size_t i = 1; i < node_count; ++i)
0094             candidates.push_back(Candidate { i, bvh_.nodes[i].get_bbox().get_half_area() });
0095         std::make_heap(candidates.begin(), candidates.end(), std::greater<>{});
0096         for (size_t i = node_count; i < bvh_.nodes.size(); ++i) {
0097             auto cost = bvh_.nodes[i].get_bbox().get_half_area();
0098             if (candidates.front().cost < cost) {
0099                 std::pop_heap(candidates.begin(), candidates.end(), std::greater<>{});
0100                 candidates.back() = Candidate { i, cost };
0101                 std::push_heap(candidates.begin(), candidates.end(), std::greater<>{});
0102             }
0103         }
0104         return candidates;
0105     }
0106
0107     Reinsertion find_reinsertion(size_t node_id) {
0108         assert(node_id != 0);
0109
0110         /*
0111          * Here is an example that explains how the cost of a reinsertion is computed. For the
0112          * reinsertion from A to C, in the figure below, we need to remove P1, replace it by B,
0113          * and create a node that holds A and C and place it where C was.
0114          *
0115          *             R
0116          *            / \
0117          *          Pn   Q1
0118          *          /     \
0119          *        ...     ...
0120          *        /         \
0121          *       P1          C
0122          *      / \
0123          *     A   B
0124          *
0125          * The resulting area *decrease* is (SA(x) means the surface area of x):
0126          *
0127          *     SA(P1) +                                                : P1 was removed
0128          *     SA(P2) - SA(B) +                                        : P2 now only contains B
0129          *     SA(P3) - SA(B U sibling(P2)) +                          : Same but for P3
0130          *     ... +
0131          *     SA(Pn) - SA(B U sibling(P2) U ... U sibling(P(n - 1)) + : Same but for Pn
0132          *     0 +                                                     : R does not change
0133          *     SA(Q1) - SA(Q1 U A) +                                   : Q1 now contains A
0134          *     SA(Q2) - SA(Q2 U A) +                                   : Q2 now contains A
0135          *     ... +
0136          *     -SA(A U C)                                              : For the parent of A and C
0137          */
0138
0139         Reinsertion best_reinsertion { /*.from */ node_id, 0, 0 };
0140         auto node_area   = bvh_.nodes[node_id].get_bbox().get_half_area();
0141         auto parent_area = bvh_.nodes[parents_[node_id]].get_bbox().get_half_area();
0142         auto area_diff = parent_area;
0143         auto sibling_id = Bvh<Node>::get_sibling_id(node_id);
0144         auto pivot_bbox = bvh_.nodes[sibling_id].get_bbox();
0145         auto parent_id = parents_[node_id];
0146         auto pivot_id = parent_id;
0147
0148         std::vector<std::pair<Scalar, size_t>> stack;
0149         do {
0150             // Try to find a reinsertion in the sibling at the current level of the tree
0151             stack.emplace_back(area_diff, sibling_id);
0152             while (!stack.empty()) {
0153                 auto top = stack.back();
0154                 stack.pop_back();
0155                 if (top.first - node_area <= best_reinsertion.area_diff)
0156                     continue;
0157
0158                 auto& dst_node = bvh_.nodes[top.second];
0159                 auto merged_area = dst_node.get_bbox().extend(bvh_.nodes[node_id].get_bbox()).get_half_area();
0160                 auto reinsert_area = top.first - merged_area;
0161                 if (reinsert_area > best_reinsertion.area_diff) {
0162                     best_reinsertion.to = top.second;
0163                     best_reinsertion.area_diff = reinsert_area;
0164                 }
0165
0166                 if (!dst_node.is_leaf()) {
0167                     auto child_area = reinsert_area + dst_node.get_bbox().get_half_area();
0168                     stack.emplace_back(child_area, dst_node.index.first_id() + 0);
0169                     stack.emplace_back(child_area, dst_node.index.first_id() + 1);
0170                 }
0171             }
0172
0173             // Compute the bounding box on the path from the node to the root, and record the
0174             // corresponding decrease in area.
0175             if (pivot_id != parent_id) {
0176                 pivot_bbox.extend(bvh_.nodes[sibling_id].get_bbox());
0177                 area_diff += bvh_.nodes[pivot_id].get_bbox().get_half_area() - pivot_bbox.get_half_area();
0178             }
0179
0180             sibling_id = Bvh<Node>::get_sibling_id(pivot_id);
0181             pivot_id = parents_[pivot_id];
0182         } while (pivot_id != 0);
0183
0184         if (best_reinsertion.to == Bvh<Node>::get_sibling_id(best_reinsertion.from) ||
0185             best_reinsertion.to == parents_[best_reinsertion.from])
0186             best_reinsertion = {};
0187         return best_reinsertion;
0188     }
0189
0190     BVH_ALWAYS_INLINE void reinsert_node(size_t from, size_t to) {
0191         auto sibling_id = Bvh<Node>::get_sibling_id(from);
0192         auto parent_id  = parents_[from];
0193         auto sibling_node = bvh_.nodes[sibling_id];
0194         auto dst_node     = bvh_.nodes[to];
0195
0196         bvh_.nodes[to].index = Node::Index::make_inner(Bvh<Node>::get_left_sibling_id(from));
0197         bvh_.nodes[sibling_id] = dst_node;
0198         bvh_.nodes[parent_id] = sibling_node;
0199
0200         if (!dst_node.is_leaf()) {
0201             parents_[dst_node.index.first_id() + 0] = sibling_id;
0202             parents_[dst_node.index.first_id() + 1] = sibling_id;
0203         }
0204         if (!sibling_node.is_leaf()) {
0205             parents_[sibling_node.index.first_id() + 0] = parent_id;
0206             parents_[sibling_node.index.first_id() + 1] = parent_id;
0207         }
0208
0209         parents_[sibling_id] = to;
0210         parents_[from] = to;
0211         refit_from(to);
0212         refit_from(parent_id);
0213     }
0214
0215     BVH_ALWAYS_INLINE void refit_from(size_t index) {
0216         do {
0217             auto& node = bvh_.nodes[index];
0218             if (!node.is_leaf()) {
0219                 node.set_bbox(
0220                     bvh_.nodes[node.index.first_id() + 0].get_bbox().extend(
0221                     bvh_.nodes[node.index.first_id() + 1].get_bbox()));
0222             }
0223             index = parents_[index];
0224         } while (index != 0);
0225     }
0226
0227     BVH_ALWAYS_INLINE auto get_conflicts(size_t from, size_t to) {
0228         return std::array<size_t, 5> {
0229             to, from,
0230             Bvh<Node>::get_sibling_id(from),
0231             parents_[to],
0232             parents_[from]
0233         };
0234     }
0235
0236     template <typename Derived>
0237     void optimize(Executor<Derived>& executor, const Config& config) {
0238         auto batch_size = std::max(size_t{1},
0239             static_cast<size_t>(static_cast<Scalar>(bvh_.nodes.size()) * config.batch_size_ratio));
0240         std::vector<Reinsertion> reinsertions;
0241         std::vector<bool> touched(bvh_.nodes.size());
0242
0243         for (size_t iter = 0; iter < config.max_iter_count; ++iter) {
0244             auto candidates = find_candidates(batch_size);
0245
0246             std::fill(touched.begin(), touched.end(), false);
0247             reinsertions.resize(candidates.size());
0248             executor.for_each(0, candidates.size(),
0249                 [&] (size_t begin, size_t end) {
0250                     for (size_t i = begin; i < end; ++i)
0251                         reinsertions[i] = find_reinsertion(candidates[i].node_id);
0252                 });
0253
0254             reinsertions.erase(std::remove_if(reinsertions.begin(), reinsertions.end(),
0255                 [] (auto& r) { return r.area_diff <= 0; }), reinsertions.end());
0256             std::sort(reinsertions.begin(), reinsertions.end(), std::greater<>{});
0257
0258             for (auto& reinsertion : reinsertions) {
0259                 auto conflicts = get_conflicts(reinsertion.from, reinsertion.to);
0260                 if (std::any_of(conflicts.begin(), conflicts.end(), [&] (size_t i) { return touched[i]; }))
0261                     continue;
0262                 for (auto conflict : conflicts)
0263                     touched[conflict] = true;
0264                 reinsert_node(reinsertion.from, reinsertion.to);
0265             }
0266         }
0267     }
0268 };
0269
0270 } // namespace bvh::v2
0271
0272 #endif