extracts input broadcasting

ginkgo-project · MarcelKoch · Jan 30, 2023 · Apr 28, 2022 · Oct 18, 2022 · Oct 19, 2022
commit b3aaa68458d1dd33c2b2e0f96044d690378abadd
diff --git a/benchmark/blas/distributed/multi_vector.cpp b/benchmark/blas/distributed/multi_vector.cpp
@@ -71,20 +71,7 @@ Parameters for a benchmark case are:
 
     auto exec = executor_factory_mpi.at(FLAGS_executor)(comm.get());
 
-    std::string json_input;
-    if (rank == 0) {
-        std::string line;
-        while (std::cin >> line) {
-            json_input += line;
-        }
-    }
-    auto input_size = json_input.size();
-    comm.broadcast(exec->get_master(), &input_size, 1, 0);
-    json_input.resize(input_size);
-    comm.broadcast(exec->get_master(), &json_input[0],
-                   static_cast<int>(input_size), 0);
-
-
+    std::string json_input = broadcast_json_input(std::cin, comm);
     rapidjson::Document test_cases;
     test_cases.Parse(json_input.c_str());
     if (!test_cases.IsArray()) {

diff --git a/benchmark/solver/distributed/solver.cpp b/benchmark/solver/distributed/solver.cpp
@@ -135,27 +135,12 @@ int main(int argc, char* argv[])
     ss_rel_res_goal << std::scientific << FLAGS_rel_res_goal;
 
 
-    std::string json_input;
-    if (rank == 0) {
-        if (FLAGS_overhead) {
-            // Fake test case to run once
-            json_input += R"(
+    std::string json_input = FLAGS_overhead
+                                 ? R"(
 [{"filename": "overhead.mtx",
   "optimal": {"spmv": "csr-csr"}]
-)";
-        } else {
-            std::string line;
-            while (std::cin >> line) {
-                json_input += line;
-            }
-        }
-    }
-    auto input_size = json_input.size();
-    comm.broadcast(exec->get_master(), &input_size, 1, 0);
-    json_input.resize(input_size);
-    comm.broadcast(exec->get_master(), &json_input[0],
-                   static_cast<int>(input_size), 0);
-
+)"
+                                 : broadcast_json_input(std::cin, comm);
     rapidjson::Document test_cases;
     test_cases.Parse(json_input.c_str());
 

diff --git a/benchmark/spmv/distributed/spmv.cpp b/benchmark/spmv/distributed/spmv.cpp
@@ -124,19 +124,7 @@ int main(int argc, char* argv[])
         }
     }
 
-    std::string json_input;
-    if (rank == 0) {
-        std::string line;
-        while (std::cin >> line) {
-            json_input += line;
-        }
-    }
-    auto input_size = json_input.size();
-    comm.broadcast(exec->get_master(), &input_size, 1, 0);
-    json_input.resize(input_size);
-    comm.broadcast(exec->get_master(), &json_input[0],
-                   static_cast<int>(input_size), 0);
-
+    std::string json_input = broadcast_json_input(std::cin, comm);
     rapidjson::Document test_cases;
     test_cases.Parse(json_input.c_str());
     if (!test_cases.IsArray()) {

diff --git a/benchmark/utils/distributed_helpers.hpp b/benchmark/utils/distributed_helpers.hpp
@@ -48,6 +48,29 @@ using dist_mtx =
                                            GlobalIndexType>;
 
 
+std::string broadcast_json_input(std::istream& is,
+                                 gko::experimental::mpi::communicator comm)
+{
+    auto exec = gko::ReferenceExecutor::create();
+
+    std::string json_input;
+    if (comm.rank() == 0) {
+        std::string line;
+        while (std::cin >> line) {
+            json_input += line;
+        }
+    }
+
+    auto input_size = json_input.size();
+    comm.broadcast(exec->get_master(), &input_size, 1, 0);
+    json_input.resize(input_size);
+    comm.broadcast(exec->get_master(), &json_input[0],
+                   static_cast<int>(input_size), 0);
+
+    return json_input;
+}
+
+
 template <typename ValueType, typename IndexType>
 gko::matrix_data<ValueType, IndexType> generate_matrix_data(
     rapidjson::Value& test_case, gko::experimental::mpi::communicator comm)