[RAMSES][AMR] Histogram of AMR Levels and New Block Index Map

src/shammodels/ramses/include/shammodels/ramses/SolverConfig.hpp

@@ -119,7 +119,8 @@ namespace shammodels::basegodunov {
     template<class Tvec, class TgridVec>
     struct AMRMode {
 
-        using Tscal = shambase::VecComponent<Tvec>;
+        using Tscal      = shambase::VecComponent<Tvec>;
+        bool no_none_cfg = false;
 
         struct None {};
         struct DensityBased {
@@ -129,9 +130,10 @@ namespace shammodels::basegodunov {
         using mode = std::variant<None, DensityBased>;
 
         mode config = None{};
-
         void set_refine_none() { config = None{}; }
         void set_refine_density_based(Tscal crit_mass) { config = DensityBased{crit_mass}; }
+        bool need_level_zero_compute() { return no_none_cfg; }
+        bool need_amr_level_compute() { return no_none_cfg; }
     };
 
     struct BCConfig {

src/shammodels/ramses/include/shammodels/ramses/modules/AMRGridRefinementHandler.hpp

@@ -47,6 +47,7 @@ namespace shammodels::basegodunov::modules {
         using AMRBlock         = typename Config::AMRBlock;
         using BlockCoord       = shamrock::amr::AMRBlockCoord<TgridVec, 3>;
         using OrientedAMRGraph = OrientedAMRGraph<Tvec, TgridVec>;
+        using TgridUint = typename std::make_unsigned<shambase::VecComponent<TgridVec>>::type;
 
         ShamrockCtx &context;
         Config &solver_config;
@@ -58,6 +59,19 @@ namespace shammodels::basegodunov::modules {
         void update_refinement();
 
         private:
+        /** @brief build histrogram for amr block levels
+         *
+         */
+        void amr_block_levels_histogram(
+            const sham::DeviceScheduler_ptr &sched,
+            const sham::DeviceBuffer<TgridUint> &amr_levels,
+            const u32 nb_levels,
+            sham::DeviceBuffer<u32> &block_count_per_level,
+            sham::DeviceBuffer<u32> &block_reordered_indx_map,
+            const u32 len,
+            TgridUint level_min,
+            TgridUint level_max);
+
         /**
          * @brief Generate the list of blocks that need to be refined or derefined.
          *

src/shammodels/ramses/include/shammodels/ramses/modules/ComputeLevel0CellSize.hpp

@@ -0,0 +1,107 @@
+// -------------------------------------------------------//
+//
+// 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
+
+/**
+ * @file ComputeLevel0CellSize.hpp
+ * @author Léodasce Sewanou (leodasce.sewanou@ens-lyon.fr)
+ * @author Timothée David--Cléris (tim.shamrock@proton.me) --no git blame--
+ * @brief
+ *
+ */
+
+#include "shamalgs/primitives/reduction.hpp"
+#include "shammath/AABB.hpp"
+#include "shamrock/solvergraph/IFieldRefs.hpp"
+#include "shamrock/solvergraph/INode.hpp"
+#include "shamrock/solvergraph/IPatchDataLayerRefs.hpp"
+#include "shamrock/solvergraph/ScalarsEdge.hpp"
+#include <memory>
+
+namespace shammodels::basegodunov::modules {
+
+    template<class Tvec, class TgridVec>
+    class ComputeLevel0CellSize : public shamrock::solvergraph::INode {
+        public:
+        ComputeLevel0CellSize() {}
+
+        struct Edges {
+            const shamrock::solvergraph::ScalarsEdge<shammath::AABB<TgridVec>> &patch_boxes;
+            const shamrock::solvergraph::IPatchDataLayerRefs &refs;
+            const shamrock::solvergraph::IFieldRefs<TgridVec> &spans_block_min;
+            const shamrock::solvergraph::IFieldRefs<TgridVec> &spans_block_max;
+            shamrock::solvergraph::ScalarsEdge<TgridVec> &level0_size;
+        };
+
+        void set_edges(
+            std::shared_ptr<shamrock::solvergraph::ScalarsEdge<shammath::AABB<TgridVec>>>
+                patch_boxes,
+            std::shared_ptr<shamrock::solvergraph::IPatchDataLayerRefs> refs,
+            std::shared_ptr<shamrock::solvergraph::IFieldRefs<TgridVec>> spans_block_min,
+            std::shared_ptr<shamrock::solvergraph::IFieldRefs<TgridVec>> spans_block_max,
+            std::shared_ptr<shamrock::solvergraph::ScalarsEdge<TgridVec>> level0_size) {
+            __internal_set_ro_edges({patch_boxes, refs, spans_block_min, spans_block_max});
+            __internal_set_rw_edges({level0_size});
+        }
+
+        Edges get_edges() {
+            return Edges{
+                get_ro_edge<shamrock::solvergraph::ScalarsEdge<shammath::AABB<TgridVec>>>(0),
+                get_ro_edge<shamrock::solvergraph::IPatchDataLayerRefs>(1),
+                get_ro_edge<shamrock::solvergraph::IFieldRefs<TgridVec>>(2),
+                get_ro_edge<shamrock::solvergraph::IFieldRefs<TgridVec>>(3),
+                get_rw_edge<shamrock::solvergraph::ScalarsEdge<TgridVec>>(0)};
+        }
+
+        void _impl_evaluate_internal() {
+            auto edges               = get_edges();
+            edges.level0_size.values = edges.refs.get_const_refs().template map<TgridVec>(
+                [&](u64 id_patch, const shamrock::patch::PatchDataLayer &pdat) {
+                    auto dev_sched       = shamsys::instance::get_compute_scheduler_ptr();
+                    sham::DeviceQueue &q = shamsys::instance::get_compute_scheduler().get_queue();
+                    sham::DeviceBuffer<TgridVec> block_size_buf(pdat.get_obj_cnt(), dev_sched);
+
+                    sham::EventList depends_list;
+                    auto block_min_acc
+                        = edges.spans_block_min.get_field(id_patch).get_buf().get_read_access(
+                            depends_list);
+                    auto block_max_acc
+                        = edges.spans_block_max.get_field(id_patch).get_buf().get_read_access(
+                            depends_list);
+                    auto block_size_acc = block_size_buf.get_write_access(depends_list);
+
+                    auto e = q.submit(depends_list, [&](sycl::handler &cgh) {
+                        cgh.parallel_for(
+                            sycl::range<1>(pdat.get_obj_cnt()), [=](sycl::item<1> gid) {
+                                u32 id             = gid.get_linear_id();
+                                block_size_acc[id] = block_max_acc[id] - block_min_acc[id];
+                            });
+                    });
+                    edges.spans_block_min.get_field(id_patch).get_buf().complete_event_state(e);
+                    edges.spans_block_max.get_field(id_patch).get_buf().complete_event_state(e);
+                    block_size_buf.complete_event_state(e);
+
+                    auto patch_min = shamalgs::primitives::min(
+                        dev_sched, block_size_buf, 0, pdat.get_obj_cnt());
+                    auto patch_max = shamalgs::primitives::max(
+                        dev_sched, block_size_buf, 0, pdat.get_obj_cnt());
+
+                    // shammath::AABB<TgridVec> patch_box = edges.patch_boxes.values.get(id_patch);
+                    // return patch_box.delt();
+                    return patch_max;
+                });
+        }
+
+        inline virtual std::string _impl_get_label() const { return "ComputeLevel0CellSize"; };
+
+        virtual std::string _impl_get_tex() const { return "ComputeLevel0CellSize"; };
+    };
+
+} // namespace shammodels::basegodunov::modules

src/shammodels/ramses/include/shammodels/ramses/modules/SolverStorage.hpp

@@ -11,6 +11,7 @@
 
 /**
  * @file SolverStorage.hpp
+ * @author Léodasce Sewanou (leodasce.sewanou@ens-lyon.fr)
  * @author Timothée David--Cléris (tim.shamrock@proton.me)
  * @brief
  */
@@ -56,9 +57,10 @@ namespace shammodels::basegodunov {
     template<class Tvec, class TgridVec, class Tmorton_>
     class SolverStorage {
         public:
-        using Tmorton            = Tmorton_;
-        using Tscal              = shambase::VecComponent<Tvec>;
-        using Tgridscal          = shambase::VecComponent<TgridVec>;
+        using Tmorton   = Tmorton_;
+        using Tscal     = shambase::VecComponent<Tvec>;
+        using Tgridscal = shambase::VecComponent<TgridVec>;
+        using TgridUint = typename std::make_unsigned<shambase::VecComponent<TgridVec>>::type;
         static constexpr u32 dim = shambase::VectorProperties<Tvec>::dimension;
 
         using RTree = RadixTree<Tmorton, TgridVec>;
@@ -74,6 +76,9 @@ namespace shammodels::basegodunov {
         std::shared_ptr<shamrock::solvergraph::FieldRefs<TgridVec>> refs_block_min;
         std::shared_ptr<shamrock::solvergraph::FieldRefs<TgridVec>> refs_block_max;
 
+        std::shared_ptr<shamrock::solvergraph::ScalarsEdge<TgridVec>> level0_size;
+        std::shared_ptr<shamrock::solvergraph::Field<TgridUint>> amr_block_levels;
+
         std::shared_ptr<shamrock::solvergraph::Indexes<u32>> block_counts;
         std::shared_ptr<shamrock::solvergraph::Indexes<u32>> block_counts_with_ghost;
         std::shared_ptr<shamrock::solvergraph::FieldRefs<Tscal>> refs_rho;

src/shammodels/ramses/src/Solver.cpp

@@ -24,8 +24,10 @@
 #include "shammodels/ramses/SolverConfig.hpp"
 #include "shammodels/ramses/modules/AMRGridRefinementHandler.hpp"
 #include "shammodels/ramses/modules/BlockNeighToCellNeigh.hpp"
+#include "shammodels/ramses/modules/ComputeAMRLevel.hpp"
 #include "shammodels/ramses/modules/ComputeCFL.hpp"
 #include "shammodels/ramses/modules/ComputeCellAABB.hpp"
+#include "shammodels/ramses/modules/ComputeLevel0CellSize.hpp"
 #include "shammodels/ramses/modules/ComputeMass.hpp"
 #include "shammodels/ramses/modules/ComputeSumOverV.hpp"
 #include "shammodels/ramses/modules/ComputeTimeDerivative.hpp"
@@ -408,6 +410,10 @@ void shammodels::basegodunov::Solver<Tvec, TgridVec>::init_solver_graph() {
     storage.cell0block_aabb_lower = std::make_shared<shamrock::solvergraph::Field<Tvec>>(
         1, "cell0block_aabb_lower", "\\mathbf{s}_{\\rm inf,block}");
 
+    // get blocks at level0 sizes for all patches
+    storage.level0_size = std::make_shared<shamrock::solvergraph::ScalarsEdge<TgridVec>>(
+        "level0_amr", "level0_amr");
+
     storage.grad_rho = std::make_shared<shamrock::solvergraph::Field<Tvec>>(
         AMRBlock::block_size, "grad_rho", "\\nabla \\rho");
     storage.dx_v = std::make_shared<shamrock::solvergraph::Field<Tvec>>(
@@ -921,6 +927,31 @@ void shammodels::basegodunov::Solver<Tvec, TgridVec>::init_solver_graph() {
         solver_sequence.push_back(std::make_shared<decltype(node)>(std::move(node)));
     }
 
+    if (solver_config.amr_mode.need_level_zero_compute()) { // compute level0 sizes in patch (to be
+                                                            // enabled later when needed)
+        modules::ComputeLevel0CellSize<Tvec, TgridVec> node_level0_sizes{};
+        node_level0_sizes.set_edges(
+            graph.get_edge_ptr<ScalarsEdge<shammath::AABB<TgridVec>>>("global_patch_boxes"),
+            storage.source_patches,
+            storage.refs_block_min,
+            storage.refs_block_max,
+            storage.level0_size);
+        solver_sequence.push_back(
+            std::make_shared<decltype(node_level0_sizes)>(std::move(node_level0_sizes)));
+    }
+
+    if (solver_config.amr_mode.need_amr_level_compute()) { // compute block amr level in patch
+        modules::ComputeAMRLevel<TgridVec> node_amr_level{};
+        node_amr_level.set_edges(
+            storage.block_counts,
+            storage.level0_size,
+            storage.refs_block_min,
+            storage.refs_block_max,
+            storage.amr_block_levels);
+        solver_sequence.push_back(
+            std::make_shared<decltype(node_amr_level)>(std::move(node_amr_level)));
+    }
+
     if (solver_config.should_compute_rho_mean()) {
         modules::NodeComputeMass<Tvec, TgridVec> node{AMRBlock::block_size};
         node.set_edges(

src/shammodels/ramses/src/modules/AMRGridRefinementHandler.cpp

@@ -20,6 +20,89 @@
 #include "shammodels/ramses/modules/AMRSortBlocks.hpp"
 #include <stdexcept>
 
+/**
+ * @brief build histrogram for amr block levels
+ * @tparam  Tvec
+ * @tparam TgridVec
+ * @param sched pointer to device scheduler
+ * @param amr_levels block's amr level
+ * @param nb_levels number of levels from level_min to level_max
+ * @param block_count_per_level buffer of levels ordered from level_min to level_max
+ * @param block_reordered_indx_map
+ */
+template<class Tvec, class TgridVec>
+void shammodels::basegodunov::modules::AMRGridRefinementHandler<Tvec, TgridVec>::
+    amr_block_levels_histogram(
+        const sham::DeviceScheduler_ptr &sched,
+        const sham::DeviceBuffer<TgridUint> &amr_levels,
+        const u32 nb_levels,
+        sham::DeviceBuffer<u32> &block_count_per_level,
+        sham::DeviceBuffer<u32> &block_reordered_indx_map,
+        const u32 len,
+        TgridUint level_min,
+        TgridUint level_max) {
+
+    SHAM_ASSERT(amr_levels.get_size() == len);
+    SHAM_ASSERT(block_count_per_level.get_size() == nb_levels);
+
+    // compute histogramm to count the number of blocks per level
+    block_count_per_level.fill(0);
+    auto &q1 = shambase::get_check_ref(sched).get_queue();
+    sham::kernel_call(
+        q1,
+        sham::MultiRef{amr_levels},
+        sham::MultiRef{block_count_per_level},
+        len,
+        [](u32 i, const TgridUint *__restrict amr_levels, u32 *__restrict block_count_per_level) {
+            auto cur_level = amr_levels[i];
+            sycl::atomic_ref<
+                u32,
+                sycl::memory_order::relaxed,
+                sycl::memory_scope::system,
+                sycl::access::address_space::global_space>
+                bin_count(block_count_per_level[cur_level]);
+            bin_count++;
+        });
+
+    // // build block index map such that from the left to right
+    // // amr_levl[reordered_blocks] ===> L_min,L_min,L_min...| L1,L1,L1...| .... |L_max L_max L_max
+    block_reordered_indx_map.fill(0);
+    sham::DeviceBuffer<u32> idx_counter_per_lev(nb_levels, sched);
+    idx_counter_per_lev.fill(0);
+    auto &q2 = shambase::get_check_ref(sched).get_queue();
+    sham::kernel_call(
+        q2,
+        sham::MultiRef{amr_levels, block_count_per_level},
+        sham::MultiRef{idx_counter_per_lev, block_reordered_indx_map},
+        len,
+        [level_min](
+            u32 i,
+            const TgridUint *__restrict amr_levels,
+            const u32 *__restrict block_cnt_per_lev,
+            u32 *__restrict idx_counter_per_lev,
+            u32 *__restrict reordered_blocks) {
+            auto amr_lev = amr_levels[i];
+            sycl::atomic_ref<
+                u32,
+                sycl::memory_order::relaxed,
+                sycl::memory_scope::system,
+                sycl::access::address_space::global_space>
+                atomic_cnt_idx_per_lev(idx_counter_per_lev[amr_lev]);
+            u32 old_loc_pos = atomic_cnt_idx_per_lev.fetch_add(1);
+            u32 offset;
+            if (amr_lev == level_min)
+                offset = 0;
+            else {
+                offset = 0;
+                for (auto k = level_min; k < amr_lev; k++) {
+                    offset += block_cnt_per_lev[k];
+                }
+            }
+            u32 glob_pos_idx               = offset + old_loc_pos;
+            reordered_blocks[glob_pos_idx] = i;
+        });
+}
+
 template<class Tvec, class TgridVec>
 template<class UserAcc, class... T>
 void shammodels::basegodunov::modules::AMRGridRefinementHandler<Tvec, TgridVec>::

src/shammodels/ramses/src/pyRamsesModel.cpp

@@ -176,11 +176,13 @@ namespace shammodels::basegodunov {
             .def(
                 "set_amr_mode_none",
                 [](TConfig &self) {
+                    self.amr_mode.no_none_cfg = false;
                     self.amr_mode.set_refine_none();
                 })
             .def(
                 "set_amr_mode_density_based",
                 [](TConfig &self, Tscal crit_mass) {
+                    self.amr_mode.no_none_cfg = true;
                     self.amr_mode.set_refine_density_based(crit_mass);
                 },
                 py::kw_only(),