bvh/v2/binned_sah_builder.h

0001 #ifndef BVH_V2_BINNED_SAH_BUILDER_H
0002 #define BVH_V2_BINNED_SAH_BUILDER_H
0003
0004 #include "bvh/v2/top_down_sah_builder.h"
0005
0006 #include <stack>
0007 #include <tuple>
0008 #include <algorithm>
0009 #include <optional>
0010 #include <numeric>
0011 #include <cassert>
0012
0013 namespace bvh::v2 {
0014
0015 /// Single-threaded top-down builder that partitions primitives based on a binned approximation of
0016 /// the Surface Area Heuristic (SAH). This builder is inspired by
0017 /// "On Fast Construction of SAH-based Bounding Volume Hierarchies", by I. Wald.
0018 template <typename Node, size_t BinCount = 8>
0019 class BinnedSahBuilder : public TopDownSahBuilder<Node> {
0020     using typename TopDownSahBuilder<Node>::Scalar;
0021     using typename TopDownSahBuilder<Node>::Vec;
0022     using typename TopDownSahBuilder<Node>::BBox;
0023
0024     using TopDownSahBuilder<Node>::build;
0025     using TopDownSahBuilder<Node>::config_;
0026     using TopDownSahBuilder<Node>::bboxes_;
0027     using TopDownSahBuilder<Node>::centers_;
0028
0029 public:
0030     using typename TopDownSahBuilder<Node>::Config;
0031
0032     BVH_ALWAYS_INLINE static Bvh<Node> build(
0033         std::span<const BBox> bboxes,
0034         std::span<const Vec> centers,
0035         const Config& config = {})
0036     {
0037         return BinnedSahBuilder(bboxes, centers, config).build();
0038     }
0039
0040 protected:
0041     struct Split {
0042         size_t bin_id;
0043         Scalar cost;
0044         size_t axis;
0045     };
0046
0047     struct Bin {
0048         BBox bbox = BBox::make_empty();
0049         size_t prim_count = 0;
0050
0051         Bin() = default;
0052
0053         BVH_ALWAYS_INLINE Scalar get_cost(const SplitHeuristic<Scalar>& sah) const {
0054             return sah.get_leaf_cost(0, prim_count, bbox);
0055         }
0056
0057         BVH_ALWAYS_INLINE void add(const BBox& bbox, size_t prim_count = 1) {
0058             this->bbox.extend(bbox);
0059             this->prim_count += prim_count;
0060         }
0061
0062         BVH_ALWAYS_INLINE void add(const Bin& bin) { add(bin.bbox, bin.prim_count); }
0063     };
0064
0065     using Bins = std::array<Bin, BinCount>;
0066     using PerAxisBins = std::array<Bins, Node::dimension>;
0067
0068     std::vector<size_t> prim_ids_;
0069
0070     BVH_ALWAYS_INLINE BinnedSahBuilder(
0071         std::span<const BBox> bboxes,
0072         std::span<const Vec> centers,
0073         const Config& config)
0074         : TopDownSahBuilder<Node>(bboxes, centers, config)
0075         , prim_ids_(bboxes.size())
0076     {
0077         std::iota(prim_ids_.begin(), prim_ids_.end(), 0);
0078     }
0079
0080     std::vector<size_t>& get_prim_ids() override { return prim_ids_; }
0081
0082     BVH_ALWAYS_INLINE void fill_bins(
0083         PerAxisBins& per_axis_bins,
0084         const BBox& bbox,
0085         size_t begin,
0086         size_t end)
0087     {
0088         auto bin_scale = Vec(BinCount) / bbox.get_diagonal();
0089         auto bin_offset = -bbox.min * bin_scale;
0090
0091         for (size_t i = begin; i < end; ++i) {
0092             auto pos = fast_mul_add(centers_[prim_ids_[i]], bin_scale, bin_offset);
0093             static_for<0, Node::dimension>([&] (size_t axis) {
0094                 size_t index = std::min(BinCount - 1,
0095                     static_cast<size_t>(robust_max(pos[axis], static_cast<Scalar>(0.))));
0096                 per_axis_bins[axis][index].add(bboxes_[prim_ids_[i]]);
0097             });
0098         }
0099     }
0100
0101     void find_best_split(size_t axis, const Bins& bins, Split& best_split) {
0102         Bin right_accum;
0103         std::array<Scalar, BinCount> right_costs;
0104         for (size_t i = BinCount - 1; i > 0; --i) {
0105             right_accum.add(bins[i]);
0106             right_costs[i] = right_accum.get_cost(config_.sah);
0107         }
0108
0109         Bin left_accum;
0110         for (size_t i = 0; i < BinCount - 1; ++i) {
0111             left_accum.add(bins[i]);
0112             auto cost = left_accum.get_cost(config_.sah) + right_costs[i + 1];
0113             if (cost < best_split.cost)
0114                 best_split = Split { i + 1, cost, axis };
0115         }
0116     }
0117
0118     size_t fallback_split(size_t axis, size_t begin, size_t end) {
0119         size_t mid = (begin + end + 1) / 2;
0120         std::partial_sort(
0121             prim_ids_.begin() + begin,
0122             prim_ids_.begin() + mid,
0123             prim_ids_.begin() + end,
0124             [&] (size_t i, size_t j) { return centers_[i][axis] < centers_[j][axis]; });
0125         return mid;
0126     }
0127
0128     std::optional<size_t> try_split(const BBox& bbox, size_t begin, size_t end) override {
0129         PerAxisBins per_axis_bins;
0130         fill_bins(per_axis_bins, bbox, begin, end);
0131
0132         auto largest_axis = bbox.get_diagonal().get_largest_axis();
0133         auto best_split = Split { BinCount / 2, std::numeric_limits<Scalar>::max(), largest_axis };
0134         for (size_t axis = 0; axis < Node::dimension; ++axis)
0135             find_best_split(axis, per_axis_bins[axis], best_split);
0136
0137         // Make sure that the split is good before proceeding with it
0138         auto leaf_cost = config_.sah.get_non_split_cost(begin, end, bbox);
0139         if (best_split.cost >= leaf_cost) {
0140             if (end - begin <= config_.max_leaf_size)
0141                 return std::nullopt;
0142             return fallback_split(largest_axis, begin, end);
0143         }
0144
0145         auto split_pos = fast_mul_add(
0146             bbox.get_diagonal()[best_split.axis] / static_cast<Scalar>(BinCount),
0147             static_cast<Scalar>(best_split.bin_id),
0148             bbox.min[best_split.axis]);
0149
0150         size_t index = std::partition(prim_ids_.begin() + begin, prim_ids_.begin() + end,
0151             [&] (size_t i) { return centers_[i][best_split.axis] < split_pos; }) - prim_ids_.begin();
0152         if (index == begin || index == end)
0153             return fallback_split(largest_axis, begin, end);
0154
0155         return std::make_optional(index);
0156     }
0157 };
0158
0159 } // namespace bvh::v2
0160
0161 #endif