From 505dcbca4e9cce6602a54e4598c807198c098416 Mon Sep 17 00:00:00 2001
From: Gonzalo Brito Gadeschi <gonzalob@nvidia.com>
Date: Sun, 23 Jun 2024 11:53:07 -0700
Subject: [PATCH] Clean all tree nodes in first iteration

---
 src/barnes_hut.h | 8 ++++----
 src/kernels.h    | 4 ++--
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/barnes_hut.h b/src/barnes_hut.h
index 6a3bfa8..fed852b 100644
--- a/src/barnes_hut.h
+++ b/src/barnes_hut.h
@@ -14,10 +14,10 @@
 using clock_timer = std::chrono::steady_clock;
 
 template<typename T, typename Index_t>
-void barnes_hut_step(System<T>& system, Arguments arguments, AtomicQuadTree<T, Index_t> tree) {
+void barnes_hut_step(System<T>& system, Arguments arguments, AtomicQuadTree<T, Index_t> tree, bool first) {
     auto start_timer = clock_timer::now();
 
-    clear_tree(system, tree);
+    clear_tree(system, tree, first? tree.capacity : tree.bump_allocator->load(memory_order_relaxed));
     compute_bounded_atomic_quad_tree(system, tree);
     build_atomic_tree(system, tree);
     auto built_tree_timer = clock_timer::now();
@@ -57,8 +57,8 @@ void run_barnes_hut(System<T>& system, Arguments arguments) {
         std::cout << "Tree init complete\n";
     }
     for (size_t step = 0; step < arguments.steps; step++) {
-        barnes_hut_step<T, Index_t>(system, arguments, tree);
-	saver.save_points(system);
+        barnes_hut_step<T, Index_t>(system, arguments, tree, step == 0);
+        saver.save_points(system);
     }
 }
 
diff --git a/src/kernels.h b/src/kernels.h
index d56f655..ffe744a 100644
--- a/src/kernels.h
+++ b/src/kernels.h
@@ -213,12 +213,12 @@ auto compute_bounded_atomic_quad_tree(System<T>& system, AtomicQuadTree<T, Index
 }
 
 template<typename T, typename Index_t>
-auto clear_tree(System<T>& system, AtomicQuadTree<T, Index_t> tree) {
+auto clear_tree(System<T>& system, AtomicQuadTree<T, Index_t> tree, Index_t last_node) {
     // clear the tree, ready for next iteration
     auto r = system.body_indices();
     std::for_each_n(
         std::execution::par_unseq,
-        r.begin(), tree.bump_allocator->load(memory_order_acquire),
+        r.begin(), last_node,
         [tree] (auto tree_index) mutable { tree.clear(tree_index); });
 }