CUDAPointHashGridSearcher3.hpp

Go to the documentation of this file.

// This code is based on Jet framework.
// Copyright (c) 2018 Doyub Kim
// CubbyFlow is voxel-based fluid simulation engine for computer games.
// Copyright (c) 2020 CubbyFlow Team
// Core Part: Chris Ohk, Junwoo Hwang, Jihong Sin, Seungwoo Yoo
// AI Part: Dongheon Cho, Minseo Kim
// We are making my contributions/submissions to this project solely in our
// personal capacity and are not conveying any rights to any intellectual
// property of any third parties.

#ifndef CUBBYFLOW_CUDA_POINT_HASH_GRID_SEARCHER3_HPP
#define CUBBYFLOW_CUDA_POINT_HASH_GRID_SEARCHER3_HPP

#ifdef CUBBYFLOW_USE_CUDA

#include <Core/CUDA/CUDAArrayView.hpp>
#include <Core/Matrix/Matrix.hpp>

#include <cuda_runtime.h>

namespace CubbyFlow
{
class CUDAPointHashGridSearcher3 final
{
 public:
    class Builder;

    class HashUtils
    {
     public:
        HashUtils() = default;

        inline CUBBYFLOW_CUDA_HOST_DEVICE HashUtils(float gridSpacing,
                                                    uint3 resolution);

        inline CUBBYFLOW_CUDA_HOST_DEVICE void GetNearbyKeys(
            float4 position, uint32_t* nearbyKeys) const;

        inline CUBBYFLOW_CUDA_HOST_DEVICE int3
        GetBucketIndex(float4 position) const;

        inline CUBBYFLOW_CUDA_HOST_DEVICE uint32_t
        GetHashKeyFromBucketIndex(int3 bucketIndex) const;

        inline CUBBYFLOW_CUDA_HOST_DEVICE uint32_t
        GetHashKeyFromPosition(float4 position) const;

     private:
        float m_gridSpacing;
        uint3 m_resolution;
    };

    template <typename Callback>
    class ForEachNearbyPointFunc
    {
     public:
        inline CUBBYFLOW_CUDA_HOST_DEVICE ForEachNearbyPointFunc(
            float r, float gridSpacing, uint3 resolution, const uint32_t* sit,
            const uint32_t* eit, const uint32_t* si, const float4* p,
            const float4* o, Callback cb);

        template <typename Index>
        inline CUBBYFLOW_CUDA_HOST_DEVICE void operator()(Index idx);

     private:
        HashUtils m_hashUtils;
        float m_radius;
        const uint32_t* m_startIndexTable;
        const uint32_t* m_endIndexTable;
        const uint32_t* m_sortedIndices;
        const float4* m_points;
        const float4* m_origins;
        Callback m_callback;
    };

    CUDAPointHashGridSearcher3(const uint3& resolution, float gridSpacing);

    CUDAPointHashGridSearcher3(uint32_t resolutionX, uint32_t resolutionY,
                               uint32_t resolutionZ, float gridSpacing);

    CUDAPointHashGridSearcher3(const CUDAPointHashGridSearcher3& other);

    CUDAPointHashGridSearcher3(CUDAPointHashGridSearcher3&& other) noexcept;

    ~CUDAPointHashGridSearcher3() = default;

    CUDAPointHashGridSearcher3& operator=(
        const CUDAPointHashGridSearcher3& other);

    CUDAPointHashGridSearcher3& operator=(
        CUDAPointHashGridSearcher3&& other) noexcept;

    void Build(const ConstCUDAArrayView1<float4>& points);

    template <typename Callback>
    void ForEachNearbyPoint(const ConstCUDAArrayView1<float4>& origins,
                            float radius, Callback callback) const;

    float GridSpacing() const;

    Vector3UZ Resolution() const;

    ConstCUDAArrayView1<float4> SortedPoints() const;

    ConstCUDAArrayView1<uint32_t> Keys() const;

    ConstCUDAArrayView1<uint32_t> StartIndexTable() const;

    ConstCUDAArrayView1<uint32_t> EndIndexTable() const;

    ConstCUDAArrayView1<uint32_t> SortedIndices() const;

    void Set(const CUDAPointHashGridSearcher3& other);

    static Builder GetBuilder();

 private:
    float m_gridSpacing = 1.0f;
    uint3 m_resolution = make_uint3(1, 1, 1);
    CUDAArray1<float4> m_points;
    CUDAArray1<uint32_t> m_keys;
    CUDAArray1<uint32_t> m_startIndexTable;
    CUDAArray1<uint32_t> m_endIndexTable;
    CUDAArray1<uint32_t> m_sortedIndices;
};

using CUDAPointHashGridSearcher3Ptr =
    std::shared_ptr<CUDAPointHashGridSearcher3>;

class CUDAPointHashGridSearcher3::Builder final
{
 public:
    Builder& WithResolution(const Vector3UZ& resolution);

    Builder& WithGridSpacing(float gridSpacing);

    CUDAPointHashGridSearcher3 Build() const;

    CUDAPointHashGridSearcher3Ptr MakeShared() const;

 private:
    float m_gridSpacing = 1.0f;
    Vector3UZ m_resolution{ 64, 64, 64 };
};
}  // namespace CubbyFlow

#include <Core/CUDA/CUDAPointHashGridSearcher3-Impl.hpp>

#endif

#endif