Merge branch 'fix_eigenvalue_output_in_parallel_with_spinors' into 'm…

…aster' Bugfix: avoid segmentation fault in eigenvalue output in parallel with spinors See merge request npneq/inq!1139
LLNL · Sep 30, 2024 · 7f70c8c · 7f70c8c
2 parents 2292a58 + 428b49a
commit 7f70c8c
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 6 deletions.
diff --git a/src/ground_state/eigenvalue_output.hpp b/src/ground_state/eigenvalue_output.hpp
@@ -61,13 +61,17 @@ class eigenvalues_output {
 			}
 			iphi++;
 		}
+
+		all_kpoint_index = parallel::gather(+kpoint_index.flatted(), el.kpin_spinor_states_part(), el.kpin_states_comm(), 0);
+		if(el.kpin_states_comm().root()){
+			for([[maybe_unused]] auto index : all_kpoint_index) assert(index >= 0 and index < electrons_.brillouin_zone().size());
+		}
 
-		all_kpoint_index = parallel::gather(+kpoint_index.flatted(), el.kpin_states_part(), el.kpin_states_comm(), 0);
-		all_spin_index = parallel::gather(+spin_index.flatted(), el.kpin_states_part(), el.kpin_states_comm(), 0);
-		all_states_index = parallel::gather(+state_index.flatted(), el.kpin_states_part(), el.kpin_states_comm(), 0);
-		all_eigenvalues = parallel::gather(+el.eigenvalues().flatted(), el.kpin_states_part(), el.kpin_states_comm(), 0);
-		all_occupations = parallel::gather(+occs.flatted(), el.kpin_states_part(), el.kpin_states_comm(), 0);
-		all_normres = parallel::gather(+normres.flatted(), el.kpin_states_part(), el.kpin_states_comm(), 0);
+		all_spin_index = parallel::gather(+spin_index.flatted(), el.kpin_spinor_states_part(), el.kpin_states_comm(), 0);
+		all_states_index = parallel::gather(+state_index.flatted(), el.kpin_spinor_states_part(), el.kpin_states_comm(), 0);
+		all_eigenvalues = parallel::gather(+el.eigenvalues().flatted(), el.kpin_spinor_states_part(), el.kpin_states_comm(), 0);
+		all_occupations = parallel::gather(+occs.flatted(), el.kpin_spinor_states_part(), el.kpin_states_comm(), 0);
+		all_normres = parallel::gather(+normres.flatted(), el.kpin_spinor_states_part(), el.kpin_states_comm(), 0);
 	}
 
 	static std::string spin_string(int index){

diff --git a/src/systems/electrons.hpp b/src/systems/electrons.hpp
@@ -78,6 +78,7 @@ class electrons {
 	std::shared_ptr<spdlog::logger> logger_;
 	parallel::partition kpin_part_;
 	parallel::arbitrary_partition kpin_states_part_;
+	parallel::arbitrary_partition kpin_spinor_states_part_;
 
 public:
 
@@ -180,6 +181,7 @@ class electrons {
 		}
 
 		kpin_states_part_ = parallel::arbitrary_partition(kpin_part_.local_size()*max_local_set_size_, kpin_states_comm_);
+		kpin_spinor_states_part_ = parallel::arbitrary_partition(kpin_part_.local_size()*max_local_spinor_set_size_, kpin_states_comm_);
 
 		assert(long(kpin_.size()) == kpin_part_.local_size());
 		assert(max_local_set_size_ > 0);
@@ -406,6 +408,10 @@ class electrons {
 	auto & kpin_states_part() const {
 		return kpin_states_part_;
 	}
+
+	auto & kpin_spinor_states_part() const {
+		return kpin_spinor_states_part_;
+	}
 
 	auto & max_local_set_size() const {
 		return max_local_set_size_;