Shamrock-code · tdavidcl · May 25, 2025 · May 9, 2025 · May 9, 2025 · May 9, 2025
diff --git a/src/shammodels/ramses/CMakeLists.txt b/src/shammodels/ramses/CMakeLists.txt
@@ -35,6 +35,7 @@ set(Sources
     src/modules/ConsToPrimDust.cpp
     src/modules/ComputeCellAABB.cpp
     src/modules/StencilGenerator.cpp
+    src/modules/NodeBuildTrees.cpp
 )
 
 if(SHAMROCK_USE_SHARED_LIB)

diff --git a/src/shammodels/ramses/include/shammodels/ramses/modules/NodeBuildTrees.hpp b/src/shammodels/ramses/include/shammodels/ramses/modules/NodeBuildTrees.hpp
@@ -0,0 +1,70 @@
+// -------------------------------------------------------//
+//
+// SHAMROCK code for hydrodynamics
+// Copyright (c) 2021-2024 Timothée David--Cléris <[email protected]>
+// 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 NodeBuildTrees.hpp
+ * @author Timothée David--Cléris ([email protected])
+ * @brief
+ *
+ */
+
+#include "shammodels/ramses/solvegraph/TreeEdge.hpp"
+#include "shamrock/solvergraph/FieldRefs.hpp"
+#include "shamrock/solvergraph/INode.hpp"
+#include "shamrock/solvergraph/Indexes.hpp"
+#include "shamtree/RadixTree.hpp"
+
+namespace shammodels::basegodunov::modules {
+
+    template<class Umorton, class TgridVec>
+    class NodeBuildTrees : public shamrock::solvergraph::INode {
+
+        u32 reduction_level = 0;
+
+        using RTree = RadixTree<Umorton, TgridVec>;
+
+        public:
+        NodeBuildTrees() {}
+
+        struct Edges {
+            const shamrock::solvergraph::Indexes<u32> &sizes;
+            const shamrock::solvergraph::FieldRefs<TgridVec> &block_min;
+            const shamrock::solvergraph::FieldRefs<TgridVec> &block_max;
+            solvergraph::TreeEdge<Umorton, TgridVec> &trees;
+        };
+
+        inline void set_edges(
+            std::shared_ptr<shamrock::solvergraph::Indexes<u32>> sizes,
+            std::shared_ptr<shamrock::solvergraph::FieldRefs<TgridVec>> block_min,
+            std::shared_ptr<shamrock::solvergraph::FieldRefs<TgridVec>> block_max,
+            std::shared_ptr<solvergraph::TreeEdge<Umorton, TgridVec>> trees) {
+            __internal_set_ro_edges({sizes, block_min, block_max});
+            __internal_set_rw_edges({trees});
+        }
+
+        inline Edges get_edges() {
+            return Edges{
+                get_ro_edge<shamrock::solvergraph::Indexes<u32>>(0),
+                get_ro_edge<shamrock::solvergraph::FieldRefs<TgridVec>>(1),
+                get_ro_edge<shamrock::solvergraph::FieldRefs<TgridVec>>(2),
+                get_rw_edge<solvergraph::TreeEdge<Umorton, TgridVec>>(0)};
+        }
+
+        void _impl_evaluate_internal();
+
+        void _impl_reset_internal() {};
+
+        inline virtual std::string _impl_get_label() { return "BuildTrees"; };
+
+        virtual std::string _impl_get_tex();
+    };
+
+} // namespace shammodels::basegodunov::modules
diff --git a/src/shammodels/ramses/include/shammodels/ramses/modules/SolverStorage.hpp b/src/shammodels/ramses/include/shammodels/ramses/modules/SolverStorage.hpp
@@ -40,9 +40,10 @@ namespace shammodels::basegodunov {
     template<class T>
     using Component = shambase::StorageComponent<T>;
 
-    template<class Tvec, class TgridVec, class Tmorton>
+    template<class Tvec, class TgridVec, class Tmorton_>
     class SolverStorage {
         public:
+        using Tmorton            = Tmorton_;
         using Tscal              = shambase::VecComponent<Tvec>;
         using Tgridscal          = shambase::VecComponent<TgridVec>;
         static constexpr u32 dim = shambase::VectorProperties<Tvec>::dimension;

diff --git a/src/shammodels/ramses/include/shammodels/ramses/solvegraph/TreeEdge.hpp b/src/shammodels/ramses/include/shammodels/ramses/solvegraph/TreeEdge.hpp
@@ -0,0 +1,35 @@
+// -------------------------------------------------------//
+//
+// SHAMROCK code for hydrodynamics
+// Copyright (c) 2021-2024 Timothée David--Cléris <[email protected]>
+// 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 TreeEdge.hpp
+ * @author Timothée David--Cléris ([email protected])
+ * @brief
+ *
+ */
+
+#include "shambase/DistributedData.hpp"
+#include "shamrock/solvergraph/IDataEdgeNamed.hpp"
+#include "shamtree/RadixTree.hpp"
+
+namespace shammodels::basegodunov::solvergraph {
+    template<class Umorton, class Tvec>
+    class TreeEdge : public shamrock::solvergraph::IDataEdgeNamed {
+        public:
+        using IDataEdgeNamed::IDataEdgeNamed;
+        shambase::DistributedData<RadixTree<Umorton, Tvec>> trees;
+
+        inline auto extract_trees() { return std::move(trees); }
+
+        inline virtual void free_alloc() { trees = {}; }
+    };
+
+} // namespace shammodels::basegodunov::solvergraph
diff --git a/src/shammodels/ramses/src/Solver.cpp b/src/shammodels/ramses/src/Solver.cpp
@@ -160,8 +160,6 @@ void shammodels::basegodunov::Solver<Tvec, TgridVec>::evolve_once() {
     modules::AMRTree amrtree(context, solver_config, storage);
     amrtree.build_trees();
 
-    amrtree.correct_bounding_box();
-
     // modules::StencilGenerator stencil_gen(context,solver_config,storage);
     // stencil_gen.make_stencil();
 

diff --git a/src/shammodels/ramses/src/modules/AMRTree.cpp b/src/shammodels/ramses/src/modules/AMRTree.cpp
@@ -14,67 +14,96 @@
  *
  */
 
+#include "shambase/DistributedData.hpp"
+#include "shambase/memory.hpp"
+#include "shambase/time.hpp"
+#include "shamcomm/logs.hpp"
 #include "shammodels/ramses/modules/AMRTree.hpp"
+#include "shammodels/ramses/modules/NodeBuildTrees.hpp"
+#include "shammodels/ramses/solvegraph/TreeEdge.hpp"
+#include "shamrock/solvergraph/IDataEdgeNamed.hpp"
+#include "shamtree/RadixTree.hpp"
 
 template<class Tvec, class TgridVec>
 void shammodels::basegodunov::modules::AMRTree<Tvec, TgridVec>::build_trees() {
 
     StackEntry stack_loc{};
+    if (true) {
 
-    using MergedPDat = shambase::DistributedData<shamrock::MergedPatchData>;
-    using RTree      = typename Storage::RTree;
+        using Umorton = typename Storage::Tmorton;
 
-    MergedPDat &mpdat = storage.merged_patchdata_ghost.get();
+        std::shared_ptr<solvergraph::TreeEdge<Umorton, TgridVec>> trees
+            = std::make_shared<solvergraph::TreeEdge<Umorton, TgridVec>>("trees", "trees");
 
-    shamrock::patch::PatchDataLayout &mpdl = storage.ghost_layout.get();
+        NodeBuildTrees<Umorton, TgridVec> build_trees;
+        build_trees.set_edges(
+            storage.block_counts_with_ghost, storage.refs_block_min, storage.refs_block_max, trees);
 
-    u32 reduc_level = 0;
+        build_trees.evaluate();
 
-    storage.merge_patch_bounds.set(
-        mpdat.map<shammath::AABB<TgridVec>>([&](u64 id, shamrock::MergedPatchData &merged) {
-            logger::debug_ln("AMR", "compute bound merged patch", id);
+        storage.trees.set(trees->extract_trees());
 
-            TgridVec min_bound = merged.pdat.get_field<TgridVec>(0).compute_min();
-            TgridVec max_bound = merged.pdat.get_field<TgridVec>(1).compute_max();
+    } else {
 
-            return shammath::AABB<TgridVec>{min_bound, max_bound};
-        }));
+        using MergedPDat = shambase::DistributedData<shamrock::MergedPatchData>;
+        using RTree      = typename Storage::RTree;
 
-    shambase::DistributedData<shammath::AABB<TgridVec>> &bounds = storage.merge_patch_bounds.get();
+        MergedPDat &mpdat = storage.merged_patchdata_ghost.get();
 
-    shambase::DistributedData<RTree> trees
-        = mpdat.map<RTree>([&](u64 id, shamrock::MergedPatchData &merged) {
-              logger::debug_ln("AMR", "compute tree for merged patch", id);
+        shamrock::patch::PatchDataLayout &mpdl = storage.ghost_layout.get();
 
-              auto aabb = bounds.get(id);
+        u32 reduc_level = 0;
 
-              TgridVec bmin = aabb.lower;
-              TgridVec bmax = aabb.upper;
+        // TODO split the
 
-              TgridVec diff = bmax - bmin;
-              diff.x()      = shambase::roundup_pow2(diff.x());
-              diff.y()      = shambase::roundup_pow2(diff.y());
-              diff.z()      = shambase::roundup_pow2(diff.z());
-              bmax          = bmin + diff;
+        shambase::DistributedData<shammath::AABB<TgridVec>> bounds
+            = mpdat.map<shammath::AABB<TgridVec>>([&](u64 id, shamrock::MergedPatchData &merged) {
+                  logger::debug_ln("AMR", "compute bound merged patch", id);
 
-              auto &field_pos = merged.pdat.get_field<TgridVec>(0);
+                  TgridVec min_bound = merged.pdat.get_field<TgridVec>(0).compute_min();
+                  TgridVec max_bound = merged.pdat.get_field<TgridVec>(1).compute_max();
 
-              RTree tree(
-                  shamsys::instance::get_compute_scheduler_ptr(),
-                  {bmin, bmax},
-                  field_pos.get_buf(),
-                  field_pos.get_obj_cnt(),
-                  reduc_level);
+                  logger::raw_ln("AABB", id, min_bound, max_bound);
 
-              return tree;
-          });
+                  return shammath::AABB<TgridVec>{min_bound, max_bound};
+              });
 
-    trees.for_each([](u64 id, RTree &tree) {
-        tree.compute_cell_ibounding_box(shamsys::instance::get_compute_queue());
-        tree.convert_bounding_box(shamsys::instance::get_compute_queue());
-    });
+        shambase::DistributedData<RTree> trees
+            = mpdat.map<RTree>([&](u64 id, shamrock::MergedPatchData &merged) {
+                  logger::debug_ln("AMR", "compute tree for merged patch", id);
+
+                  auto aabb = bounds.get(id);
+
+                  TgridVec bmin = aabb.lower;
+                  TgridVec bmax = aabb.upper;
+
+                  TgridVec diff = bmax - bmin;
+                  diff.x()      = shambase::roundup_pow2(diff.x());
+                  diff.y()      = shambase::roundup_pow2(diff.y());
+                  diff.z()      = shambase::roundup_pow2(diff.z());
+                  bmax          = bmin + diff;
+
+                  auto &field_pos = merged.pdat.get_field<TgridVec>(0);
+
+                  RTree tree(
+                      shamsys::instance::get_compute_scheduler_ptr(),
+                      {bmin, bmax},
+                      field_pos.get_buf(),
+                      field_pos.get_obj_cnt(),
+                      reduc_level);
+
+                  return tree;
+              });
+
+        trees.for_each([](u64 id, RTree &tree) {
+            tree.compute_cell_ibounding_box(shamsys::instance::get_compute_queue());
+            tree.convert_bounding_box(shamsys::instance::get_compute_queue());
+        });
+
+        storage.trees.set(std::move(trees));
 
-    storage.trees.set(std::move(trees));
+        correct_bounding_box();
+    }
 }
 
 template<class Tvec, class TgridVec>