Bvh.cuh

#ifndef PBAT_GPU_IMPL_GEOMETRY_BVH_CUH
#define PBAT_GPU_IMPL_GEOMETRY_BVH_CUH
 
#include "Aabb.cuh"
#include "Morton.cuh"
#include "pbat/common/Queue.h"
#include "pbat/common/Stack.h"
#include "pbat/geometry/Morton.h"
#include "pbat/geometry/OverlapQueries.h"
#include "pbat/gpu/Aliases.h"
#include "pbat/gpu/impl/common/Buffer.cuh"
#include "pbat/gpu/profiling/Profiling.h"
#include "pbat/math/linalg/mini/Mini.h"
 
#include <exception>
#include <string>
#include <thrust/execution_policy.h>
#include <thrust/for_each.h>
#include <thrust/iterator/counting_iterator.h>
#include <type_traits>
 
namespace pbat {
namespace gpu {
namespace impl {
namespace geometry {

class Bvh
{
  public:
    static auto constexpr kDims = 3;
    using OverlapType           = cuda::std::pair<GpuIndex, GpuIndex>;
    using MortonCodeType        = pbat::geometry::MortonCodeType;
 
    friend class BvhQuery;
 
    static_assert(
        std::is_same_v<GpuIndex, std::int32_t>,
        "gpu::BvhImpl only supported for 32-bit signed integer indices");

    Bvh(GpuIndex nBoxes);

    void Build(Aabb<kDims>& aabbs, Morton::Bound const& min, Morton::Bound const& max);
    void SortByMortonCode(Aabb<kDims>& aabbs, Morton::Bound const& min, Morton::Bound const& max);
    void BuildTree(GpuIndex n);
    void ConstructBoxes(Aabb<kDims>& aabbs);
    template <class FOnOverlapDetected>
    void DetectOverlaps(Aabb<kDims> const& aabbs, FOnOverlapDetected&& fOnOverlapDetected);
    template <class FOnOverlapDetected, auto kStackSize = 64>
    void DetectOverlaps(
        Aabb<kDims> const& leafAabbs,
        Aabb<kDims> const& queryAabbs,
        FOnOverlapDetected fOnOverlapDetected);
    template <
        class FGetQueryObject,
        class FMinDistanceToBox,
        class FDistanceToLeaf,
        class FDistanceUpperBound,
        class FOnNearestNeighbourFound,
        auto kQueueSize = 8,
        auto kStackSize = 64>
    void NearestNeighbours(
        Aabb<kDims> const& aabbs,
        GpuIndex nQueries,
        FGetQueryObject fGetQueryObject,
        FMinDistanceToBox fMinDistanceToBox,
        FDistanceToLeaf fDistanceToleaf,
        FDistanceUpperBound fDistanceUpperBound,
        FOnNearestNeighbourFound fOnNearestNeighbourFound,
        GpuScalar eps = GpuScalar(0));
    template <
        class FGetQueryObject,
        class FMinDistanceToBox,
        class FDistanceToLeaf,
        class FDistanceUpperBound,
        class FOnFound,
        auto kStackSize = 64>
    void RangeSearch(
        Aabb<kDims> const& aabbs,
        GpuIndex nQueries,
        FGetQueryObject fGetQueryObject,
        FMinDistanceToBox fMinDistanceToBox,
        FDistanceToLeaf fDistanceToleaf,
        FDistanceUpperBound fDistanceUpperBound,
        FOnFound fOnFound);
 
    common::Buffer<GpuIndex> inds; 
    Morton morton;                 
    common::Buffer<GpuIndex, 2>
        child; 
    common::Buffer<GpuIndex> parent; 
    common::Buffer<GpuIndex, 2>
        rightmost; 
    Aabb<kDims> iaabbs; 
    common::Buffer<GpuIndex> visits; 
};
 
template <class FOnOverlapDetected>
inline void Bvh::DetectOverlaps(Aabb<kDims> const& aabbs, FOnOverlapDetected&& fOnOverlapDetected)
{
    PBAT_PROFILE_CUDA_NAMED_SCOPE("pbat.gpu.impl.geometry.Bvh.DetectOverlaps");
    auto const nLeafBoxes = aabbs.Size();
    thrust::for_each(
        thrust::device,
        thrust::make_counting_iterator(0),
        thrust::make_counting_iterator(nLeafBoxes),
        [inds      = inds.Raw(),
         child     = child.Raw(),
         rightmost = rightmost.Raw(),
         b         = aabbs.b.Raw(),
         e         = aabbs.e.Raw(),
         ib        = iaabbs.b.Raw(),
         ie        = iaabbs.e.Raw(),
         leafBegin = nLeafBoxes - 1,
         fOnOverlapDetected =
             std::forward<FOnOverlapDetected>(fOnOverlapDetected)] PBAT_DEVICE(GpuIndex s) mutable {
            // Traverse nodes depth-first starting from the root=0 node
            using pbat::common::Stack;
            using namespace pbat::math::linalg;
            using namespace pbat::geometry;
            auto const leaf = leafBegin + s;
            auto Ls         = mini::FromBuffers<3, 1>(b, s);
            auto Us         = mini::FromBuffers<3, 1>(e, s);
            Stack<GpuIndex, 64> stack{};
            stack.Push(0);
            do
            {
                assert(not stack.IsFull());
                GpuIndex const node = stack.Pop();
                // Check each child node for overlap.
                GpuIndex const lc       = child[0][node];
                GpuIndex const rc       = child[1][node];
                bool const bIsLeftLeaf  = lc >= leafBegin;
                bool const bIsRightLeaf = rc >= leafBegin;
                auto Llc                = bIsLeftLeaf ? mini::FromBuffers<3, 1>(b, lc - leafBegin) :
                                                        mini::FromBuffers<3, 1>(ib, lc);
                auto Ulc                = bIsLeftLeaf ? mini::FromBuffers<3, 1>(e, lc - leafBegin) :
                                                        mini::FromBuffers<3, 1>(ie, lc);
                auto Lrc = bIsRightLeaf ? mini::FromBuffers<3, 1>(b, rc - leafBegin) :
                                          mini::FromBuffers<3, 1>(ib, rc);
                auto Urc = bIsRightLeaf ? mini::FromBuffers<3, 1>(e, rc - leafBegin) :
                                          mini::FromBuffers<3, 1>(ie, rc);
                bool const bLeftBoxOverlaps =
                    OverlapQueries::AxisAlignedBoundingBoxes(Ls, Us, Llc, Ulc) and
                    (rightmost[0][node] > leaf);
                bool const bRightBoxOverlaps =
                    OverlapQueries::AxisAlignedBoundingBoxes(Ls, Us, Lrc, Urc) and
                    (rightmost[1][node] > leaf);
 
                // Leaf overlaps another leaf node
                if (bLeftBoxOverlaps and bIsLeftLeaf)
                {
                    GpuIndex const si = inds[s];
                    GpuIndex const sj = inds[lc - leafBegin];
                    fOnOverlapDetected(si, sj);
                }
                if (bRightBoxOverlaps and bIsRightLeaf)
                {
                    GpuIndex const si = inds[s];
                    GpuIndex const sj = inds[rc - leafBegin];
                    fOnOverlapDetected(si, sj);
                }
 
                // Leaf overlaps an internal node -> traverse
                bool const bTraverseLeft  = bLeftBoxOverlaps and not bIsLeftLeaf;
                bool const bTraverseRight = bRightBoxOverlaps and not bIsRightLeaf;
                if (bTraverseLeft)
                    stack.Push(lc);
                if (bTraverseRight)
                    stack.Push(rc);
            } while (not stack.IsEmpty());
        });
}
 
template <class FOnOverlapDetected, auto kStackSize>
inline void Bvh::DetectOverlaps(
    Aabb<kDims> const& leafAabbs,
    Aabb<kDims> const& queryAabbs,
    FOnOverlapDetected fOnOverlapDetected)
{
    PBAT_PROFILE_CUDA_NAMED_SCOPE("pbat.gpu.impl.geometry.Bvh.DetectOverlaps");
 
    static_assert(
        std::is_invocable_v<FOnOverlapDetected, GpuIndex, GpuIndex>,
        "FOnOverlapDetected must be callable with signature void f(GpuIndex,GpuIndex)");
 
    using namespace pbat::math::linalg;
    auto fGetQueryObject = [b = queryAabbs.b.Raw(),
                            e = queryAabbs.e.Raw()] PBAT_DEVICE(GpuIndex q) {
        mini::SMatrix<GpuScalar, 3, 2> LU;
        LU.Col(0) = mini::FromBuffers<3, 1>(b, q);
        LU.Col(1) = mini::FromBuffers<3, 1>(e, q);
        return LU;
    };
    auto fMinDistanceToBox = [] PBAT_DEVICE(
                                 mini::SMatrix<GpuScalar, 3, 2> const& LU1,
                                 mini::SVector<GpuScalar, 3> const& L2,
                                 mini::SVector<GpuScalar, 3> const& U2) {
        bool const bBoxesOverlap = pbat::geometry::OverlapQueries::AxisAlignedBoundingBoxes(
            LU1.Col(0),
            LU1.Col(1),
            L2,
            U2);
        return static_cast<GpuScalar>(not bBoxesOverlap);
    };
    auto fDistanceToLeaf = [] PBAT_DEVICE(
                               [[maybe_unused]] GpuIndex q,
                               [[maybe_unused]] mini::SMatrix<GpuScalar, 3, 2> const& LUQ,
                               [[maybe_unused]] GpuIndex leaf,
                               [[maybe_unused]] GpuIndex i) {
        return GpuScalar(0);
    };
    auto fQueryUpperBound = [] PBAT_DEVICE(GpuIndex q) {
        return GpuScalar(0);
    };
    auto fOnFound = [fOnOverlapDetected] PBAT_DEVICE(
                        GpuIndex q,
                        [[maybe_unused]] GpuIndex leaf,
                        GpuIndex i,
                        [[maybe_unused]] GpuScalar d) mutable {
        fOnOverlapDetected(q, i);
    };
 
    this->RangeSearch<
        decltype(fGetQueryObject),
        decltype(fMinDistanceToBox),
        decltype(fDistanceToLeaf),
        decltype(fQueryUpperBound),
        decltype(fOnFound),
        kStackSize>(
        leafAabbs,
        queryAabbs.Size(),
        fGetQueryObject,
        fMinDistanceToBox,
        fDistanceToLeaf,
        fQueryUpperBound,
        fOnFound);
}
 
template <
    class FGetQueryObject,
    class FMinDistanceToBox,
    class FDistanceToLeaf,
    class FDistanceUpperBound,
    class FOnNearestNeighbourFound,
    auto kQueueSize,
    auto kStackSize>
inline void Bvh::NearestNeighbours(
    Aabb<kDims> const& aabbs,
    GpuIndex nQueries,
    FGetQueryObject fGetQueryObject,
    FMinDistanceToBox fMinDistanceToBox,
    FDistanceToLeaf fDistanceToLeaf,
    FDistanceUpperBound fDistanceUpperBound,
    FOnNearestNeighbourFound fOnNearestNeighbourFound,
    GpuScalar eps)
{
    PBAT_PROFILE_CUDA_NAMED_SCOPE("pbat.gpu.impl.geometry.Bvh.NearestNeighbours");
 
    using TQuery = std::invoke_result_t<FGetQueryObject, GpuIndex>;
    using Point  = pbat::math::linalg::mini::SVector<GpuScalar, kDims>;
 
    static_assert(
        std::is_invocable_v<FGetQueryObject, GpuIndex> and not std::is_same_v<TQuery, void>,
        "FGetQueryObject must be callable with signature TQuery f(GpuIndex)");
    static_assert(
        std::is_invocable_v<FMinDistanceToBox, TQuery, Point, Point> and
            std::is_convertible_v<
                std::invoke_result_t<FMinDistanceToBox, TQuery, Point, Point>,
                GpuScalar>,
        "FMinDistanceToBox must be callable with signature GpuScalar f(TQuery, Point, Point) where "
        "Point=pbat::math::linalg::mini::SVector<GpuScalar, 3>");
    static_assert(
        std::is_invocable_v<FDistanceToLeaf, GpuIndex, TQuery, GpuIndex, GpuIndex> and
            std::is_convertible_v<
                std::invoke_result_t<FDistanceToLeaf, GpuIndex, TQuery, GpuIndex, GpuIndex>,
                GpuScalar>,
        "FDistanceToLeaf must be callable with signature GpuScalar f(GpuIndex, TQuery query, "
        "GpuIndex leaf, "
        "GpuIndex i)");
    static_assert(
        std::is_invocable_v<FDistanceUpperBound, GpuIndex> and
            std::is_convertible_v<std::invoke_result_t<FDistanceUpperBound, GpuIndex>, GpuScalar>,
        "FDistanceUpperBound must be callable with signature GpuScalar f(GpuIndex)");
    static_assert(
        std::is_invocable_v<FOnNearestNeighbourFound, GpuIndex, GpuIndex, GpuScalar, GpuIndex>,
        "FOnNearestNeighbourFound must be callable with signature void f(GpuIndex query, GpuIndex "
        "leaf, GpuScalar dmin, GpuIndex k)");
 
    thrust::for_each(
        thrust::device,
        thrust::make_counting_iterator(0),
        thrust::make_counting_iterator(nQueries),
        [eps,
         fGetQueryObject,
         fMinDistanceToBox,
         fDistanceToLeaf,
         fDistanceUpperBound,
         fOnNearestNeighbourFound,
         leafBegin = aabbs.Size() - 1,
         inds      = inds.Raw(),
         b         = aabbs.b.Raw(),
         e         = aabbs.e.Raw(),
         ib        = iaabbs.b.Raw(),
         ie        = iaabbs.e.Raw(),
         child     = child.Raw()] PBAT_DEVICE(GpuIndex q) mutable {
            using pbat::math::linalg::mini::FromBuffers;
            using Stack = pbat::common::Stack<GpuIndex, kStackSize>;
            using Queue = pbat::common::Queue<GpuIndex, kQueueSize>;
 
            // Depth-first branch and bound distance minimization
            Stack dfs{};
            Queue nn{};
            GpuScalar dmin{fDistanceUpperBound(q)};
            TQuery const query{fGetQueryObject(q)};
 
            dfs.Push(0 /*root*/);
            do
            {
                assert(not dfs.IsFull());
                GpuIndex nodeIdx       = dfs.Pop();
                bool const bIsLeafNode = nodeIdx >= leafBegin;
                auto lo                = dmin - eps;
                auto hi                = dmin + eps;
                if (not bIsLeafNode)
                {
                    auto L  = FromBuffers<3, 1>(ib, nodeIdx);
                    auto U  = FromBuffers<3, 1>(ie, nodeIdx);
                    auto db = fMinDistanceToBox(query, L, U);
                    if (db <= hi)
                    {
                        dfs.Push(child[0][nodeIdx]);
                        dfs.Push(child[1][nodeIdx]);
                    }
                }
                else
                {
                    nodeIdx -= leafBegin;
                    auto L             = FromBuffers<3, 1>(b, nodeIdx);
                    auto U             = FromBuffers<3, 1>(e, nodeIdx);
                    GpuScalar const db = fMinDistanceToBox(query, L, U);
                    if (db <= hi)
                    {
                        GpuIndex const i  = inds[nodeIdx];
                        GpuScalar const d = fDistanceToLeaf(q, query, nodeIdx, i);
                        if (d < lo)
                        {
                            nn.Clear();
                            nn.Push(i);
                            dmin = d;
                        }
                        else if (d <= hi and not nn.IsFull())
                        {
                            nn.Push(i);
                        }
                    }
                }
            } while (not dfs.IsEmpty());
            GpuIndex k{0};
            while (not nn.IsEmpty())
            {
                GpuIndex const i = nn.Top();
                nn.Pop();
                fOnNearestNeighbourFound(q, i, dmin, k++);
            }
        });
}
 
template <
    class FGetQueryObject,
    class FMinDistanceToBox,
    class FDistanceToLeaf,
    class FDistanceUpperBound,
    class FOnFound,
    auto kStackSize>
inline void Bvh::RangeSearch(
    Aabb<kDims> const& aabbs,
    GpuIndex nQueries,
    FGetQueryObject fGetQueryObject,
    FMinDistanceToBox fMinDistanceToBox,
    FDistanceToLeaf fDistanceToleaf,
    FDistanceUpperBound fDistanceUpperBound,
    FOnFound fOnFound)
{
    PBAT_PROFILE_CUDA_NAMED_SCOPE("pbat.gpu.impl.geometry.Bvh.RangeSearch");
 
    using TQuery = std::invoke_result_t<FGetQueryObject, GpuIndex>;
    using Point  = pbat::math::linalg::mini::SVector<GpuScalar, kDims>;
    static_assert(
        std::is_invocable_v<FGetQueryObject, GpuIndex> and not std::is_same_v<TQuery, void>,
        "FGetQueryObject must be callable with signature TQuery f(GpuIndex)");
    static_assert(
        std::is_invocable_v<FMinDistanceToBox, TQuery, Point, Point> and
            std::is_convertible_v<
                std::invoke_result_t<FMinDistanceToBox, TQuery, Point, Point>,
                GpuScalar>,
        "FMinDistanceToBox must be callable with signature GpuScalar f(TQuery, Point, Point) where "
        "Point=pbat::math::linalg::mini::SVector<GpuScalar, 3>");
    static_assert(
        std::is_invocable_v<FDistanceToLeaf, GpuIndex, TQuery, GpuIndex, GpuIndex> and
            std::is_convertible_v<
                std::invoke_result_t<FDistanceToLeaf, GpuIndex, TQuery, GpuIndex, GpuIndex>,
                GpuScalar>,
        "FDistanceToLeaf must be callable with signature GpuScalar f(GpuIndex q, TQuery query, "
        "GpuIndex leaf, "
        "GpuIndex i)");
    static_assert(
        std::is_invocable_v<FDistanceUpperBound, GpuIndex> and
            std::is_convertible_v<std::invoke_result_t<FDistanceUpperBound, GpuIndex>, GpuScalar>,
        "FDistanceUpperBound must be callable with signature GpuScalar f(GpuIndex)");
    static_assert(
        std::is_invocable_v<FOnFound, GpuIndex, GpuIndex, GpuIndex, GpuScalar>,
        "FOnFound must be callable with signature void f(GpuIndex q, GpuIndex "
        "leafIdx, GpuIndex i, GpuScalar d)");
 
    thrust::for_each(
        thrust::device,
        thrust::make_counting_iterator(0),
        thrust::make_counting_iterator(nQueries),
        [leafBegin           = aabbs.Size() - 1,
         inds                = inds.Raw(),
         b                   = aabbs.b.Raw(),
         e                   = aabbs.e.Raw(),
         ib                  = iaabbs.b.Raw(),
         ie                  = iaabbs.e.Raw(),
         child               = child.Raw(),
         fGetQueryObject     = std::forward<FGetQueryObject>(fGetQueryObject),
         fMinDistanceToBox   = std::forward<FMinDistanceToBox>(fMinDistanceToBox),
         fDistanceToLeaf     = std::forward<FDistanceToLeaf>(fDistanceToleaf),
         fDistanceUpperBound = std::forward<FDistanceUpperBound>(fDistanceUpperBound),
         fOnFound            = std::forward<FOnFound>(fOnFound)] PBAT_DEVICE(GpuIndex q) mutable {
            using pbat::math::linalg::mini::FromBuffers;
            using Stack = pbat::common::Stack<GpuIndex, kStackSize>;
 
            // Depth-first branch and bound search
            Stack dfs{};
            GpuScalar const upper{fDistanceUpperBound(q)};
            TQuery const query{fGetQueryObject(q)};
 
            dfs.Push(0 /*root*/);
            do
            {
                assert(not dfs.IsFull());
                GpuIndex nodeIdx       = dfs.Pop();
                bool const bIsLeafNode = nodeIdx >= leafBegin;
                if (not bIsLeafNode)
                {
                    auto L = FromBuffers<3, 1>(ib, nodeIdx);
                    auto U = FromBuffers<3, 1>(ie, nodeIdx);
                    if (fMinDistanceToBox(query, L, U) <= upper)
                    {
                        dfs.Push(child[0][nodeIdx]);
                        dfs.Push(child[1][nodeIdx]);
                    }
                }
                else
                {
                    nodeIdx -= leafBegin;
                    auto L             = FromBuffers<3, 1>(b, nodeIdx);
                    auto U             = FromBuffers<3, 1>(e, nodeIdx);
                    GpuScalar const db = fMinDistanceToBox(query, L, U);
                    if (db <= upper)
                    {
                        GpuIndex const i  = inds[nodeIdx];
                        GpuScalar const d = fDistanceToLeaf(q, query, nodeIdx, i);
                        if (d <= upper)
                        {
                            fOnFound(q, nodeIdx, i, d);
                        }
                    }
                }
            } while (not dfs.IsEmpty());
        });
}
 
} // namespace geometry
} // namespace impl
} // namespace gpu
} // namespace pbat
 
#endif // PBAT_GPU_GEOMETRY_BVHIMPL_H