CLBVHObjectIterator.hpp

Go to the documentation of this file.

// -------------------------------------------------------//
//
// SHAMROCK code for hydrodynamics
// Copyright (c) 2021-2026 Timothée David--Cléris <tim.shamrock@proton.me>
// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1
// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information
//
// -------------------------------------------------------//
 
#pragma once

 
#include "shambackends/DeviceBuffer.hpp"
#include "shammath/AABB.hpp"
#include "shammath/sfc/morton.hpp"
#include "shamtree/CellIterator.hpp"
#include "shamtree/KarrasRadixTree.hpp"
#include "shamtree/KarrasRadixTreeAABB.hpp"
#include "shamtree/LeafCellIterator.hpp"
#include "shamtree/MortonReducedSet.hpp"
 
namespace shamtree {

    template<class Tmorton, class Tvec, u32 dim>
    struct CLBVHTraverserAccessed;

    template<class Tmorton, class Tvec, u32 dim>
    struct CLBVHTraverser;

    template<class Tmorton, class Tvec, u32 dim>
    struct CLBVHTraverserHost;

    template<class Tmorton, class Tvec, u32 dim>
    struct CLBVHObjectIteratorAccessed;

    template<class Tmorton, class Tvec, u32 dim>
    struct CLBVHObjectIterator;

    template<class Tmorton, class Tvec, u32 dim>
    struct CLBVHObjectIteratorHost;
 
} // namespace shamtree
 
template<class Tmorton, class Tvec, u32 dim>
struct shamtree::CLBVHTraverserAccessed {

    static constexpr u32 tree_depth_max
        = shamrock::sfc::MortonCodes<Tmorton, 3>::significant_bits + 1;
 
    KarrasTreeTraverserAccessed tree_traverser; 
    const Tvec *aabb_min;                       
    const Tvec *aabb_max;                       

    template<class Functor1, class Functor2, class Functor3>
    inline void traverse_tree_base(
        Functor1 &&traverse_condition,
        Functor2 &&on_found_leaf,
        Functor3 &&on_excluded_node) const {
 
        tree_traverser.template stack_based_traversal<tree_depth_max>(
            std::forward<Functor1>(traverse_condition),
            std::forward<Functor2>(on_found_leaf),
            std::forward<Functor3>(on_excluded_node));
    }

    template<class Functor1, class Functor2, class Functor3>
    inline void traverse_tree_base(
        u32 root_node,
        Functor1 &&traverse_condition,
        Functor2 &&on_found_leaf,
        Functor3 &&on_excluded_node) const {
 
        tree_traverser.template stack_based_traversal<tree_depth_max>(
            root_node,
            std::forward<Functor1>(traverse_condition),
            std::forward<Functor2>(on_found_leaf),
            std::forward<Functor3>(on_excluded_node));
    }
 
    template<class Functor1, class Functor2>
    inline void rtree_for(Functor1 &&traverse_condition_with_aabb, Functor2 &&on_found_leaf) const {
 
        traverse_tree_base(
            [&](u32 node_id) { // interaction crit
                return traverse_condition_with_aabb(
                    node_id, shammath::AABB<Tvec>{aabb_min[node_id], aabb_max[node_id]});
            },
            [&](u32 node_id) { // on leaf found
                on_found_leaf(node_id);
            },
            [&](u32) {});
    }
};
 
template<class Tmorton, class Tvec, u32 dim>
struct shamtree::CLBVHObjectIteratorAccessed {

    static constexpr u32 tree_depth_max
        = shamrock::sfc::MortonCodes<Tmorton, 3>::significant_bits + 1;
 
    LeafCellIterator::acc cell_iterator;                       
    CLBVHTraverserAccessed<Tmorton, Tvec, dim> tree_traverser; 

    template<class Functor1, class Functor2>
    inline void rtree_for(
        Functor1 &&traverse_condition_with_aabb, Functor2 &&on_found_object) const {
 
        tree_traverser.rtree_for(
            std::forward<Functor1>(traverse_condition_with_aabb),
            [&](u32 node_id) { // on leaf found
                u32 leaf_id = node_id - tree_traverser.tree_traverser.offset_leaf;
                cell_iterator.for_each_in_leaf_cell(leaf_id, on_found_object);
            });
    }
};
 
template<class Tmorton, class Tvec, u32 dim>
struct shamtree::CLBVHTraverser {
    KarrasTreeTraverser tree_traverser;       
    const sham::DeviceBuffer<Tvec> &aabb_min; 
    const sham::DeviceBuffer<Tvec> &aabb_max; 

    using acc = CLBVHTraverserAccessed<Tmorton, Tvec, dim>;

    inline acc get_read_access(sham::EventList &deps) const {
        return acc{
            tree_traverser.get_read_access(deps),
            aabb_min.get_read_access(deps),
            aabb_max.get_read_access(deps)};
    }

    inline void complete_event_state(sycl::event e) const {
        tree_traverser.complete_event_state(e);
        aabb_min.complete_event_state(e);
        aabb_max.complete_event_state(e);
    }
};
 
template<class Tmorton, class Tvec, u32 dim>
struct shamtree::CLBVHTraverserHost {
    KarrasTreeTraverserHost tree_traverser; 
    std::vector<Tvec> aabb_min;             
    std::vector<Tvec> aabb_max;             

    using acc = CLBVHTraverserAccessed<Tmorton, Tvec, dim>;

    inline acc get_read_access() const {
        return acc{tree_traverser.get_read_access(), aabb_min.data(), aabb_max.data()};
    }
};
 
template<class Tmorton, class Tvec, u32 dim>
struct shamtree::CLBVHObjectIterator {
    LeafCellIterator cell_iterator;                    
    CLBVHTraverser<Tmorton, Tvec, dim> tree_traverser; 

    using acc = CLBVHObjectIteratorAccessed<Tmorton, Tvec, dim>;

    inline acc get_read_access(sham::EventList &deps) const {
        return acc{cell_iterator.get_read_access(deps), tree_traverser.get_read_access(deps)};
    }

    inline void complete_event_state(sycl::event e) const {
        cell_iterator.complete_event_state(e);
        tree_traverser.complete_event_state(e);
    }
};
 
template<class Tmorton, class Tvec, u32 dim>
struct shamtree::CLBVHObjectIteratorHost {
    LeafCellIteratorHost cell_iterator;                    
    CLBVHTraverserHost<Tmorton, Tvec, dim> tree_traverser; 

    using acc = CLBVHObjectIteratorAccessed<Tmorton, Tvec, dim>;

    inline acc get_read_access() const {
        return acc{cell_iterator.get_read_access(), tree_traverser.get_read_access()};
    }
};