feat: allow explicit distribution as prefix sum via --explicit-distri…

…bution-is-prefix-sum
KarlsruheGraphGeneration · Nov 13, 2024 · bda1e97 · bda1e97
1 parent d3e2b9c
commit bda1e97
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 11 deletions.
diff --git a/app/KaGen.cpp b/app/KaGen.cpp
@@ -447,6 +447,9 @@ This is mostly useful for experimental graph generators or when using KaGen to l
             "--explicit-distribution", config.input_graph.explicit_distribution_filename,
             "A text file containing the number of vertices on each PE, one line per PE. Only used when "
             "--distribution=explicit.");
+        cmd->add_flag(
+            "--explicit-distribution-is-prefix-sum", config.input_graph.explicit_distribution_is_prefix_sum,
+            "Interpret the explicit distribution as prefix sum instead of vertices on each PE.");
         cmd->add_option("--input-format", config.input_graph.format)
             ->transform(CLI::CheckedTransformer(GetInputFormatMap()).description(""))
             ->description(R"(The following file formats are supported:

diff --git a/kagen/context.cpp b/kagen/context.cpp
@@ -392,6 +392,7 @@ PGeneratorConfig CreateConfigFromString(const std::string& options_str, PGenerat
         config.input_graph.distribution = distribution_it->second;
 
         config.input_graph.explicit_distribution_filename = get_string_or_default("explicit_distribution");
+        config.input_graph.explicit_distribution_is_prefix_sum = get_bool_or_default("explicit_distribution_is_prefix_sum");
 
         const auto        formats     = GetInputFormatMap();
         const std::string format_name = get_string_or_default("input_format", StringifyEnum(config.input_graph.format));

diff --git a/kagen/context.h b/kagen/context.h
@@ -40,11 +40,12 @@ struct ImageMeshConfig {
 };
 
 struct InputGraphConfig {
-    std::string       filename                       = "";
-    FileFormat        format                         = FileFormat::EXTENSION;
-    GraphDistribution distribution                   = GraphDistribution::BALANCE_VERTICES;
-    std::string       explicit_distribution_filename = "";
-    int               width                          = 64;
+    std::string       filename                            = "";
+    FileFormat        format                              = FileFormat::EXTENSION;
+    GraphDistribution distribution                        = GraphDistribution::BALANCE_VERTICES;
+    std::string       explicit_distribution_filename      = "";
+    bool              explicit_distribution_is_prefix_sum = false;
+    int               width                               = 64;
 
     int vtx_width    = 64;
     int adjncy_width = 64;

diff --git a/kagen/io.cpp b/kagen/io.cpp
@@ -124,7 +124,7 @@ CreateGraphReader(const FileFormat format, const InputGraphConfig& config, const
 
 namespace {
 
-std::vector<SInt> ReadExplicitVertexDistribution(const std::string& filename) {
+std::vector<SInt> ReadExplicitVertexDistribution(const std::string& filename, const bool is_prefix_sum) {
     MappedFileToker toker(filename);
     toker.SkipSpaces();
 
@@ -134,9 +134,11 @@ std::vector<SInt> ReadExplicitVertexDistribution(const std::string& filename) {
         number_of_vertices.push_back(toker.ScanUnsigned());
         toker.SkipLine();
     }
-    number_of_vertices.push_back(0);
 
-    std::exclusive_scan(number_of_vertices.begin(), number_of_vertices.end(), number_of_vertices.begin(), 0);
+    if (!is_prefix_sum) {
+        number_of_vertices.push_back(0);
+        std::exclusive_scan(number_of_vertices.begin(), number_of_vertices.end(), number_of_vertices.begin(), 0);
+    }
 
     return number_of_vertices;
 }
@@ -186,9 +188,10 @@ GraphFragment ReadGraphFragment(
         }
 
         case GraphDistribution::EXPLICIT: {
-            auto distribution = ReadExplicitVertexDistribution(config.explicit_distribution_filename);
-            from              = distribution[rank];
-            to_node           = distribution[rank + 1];
+            auto distribution = ReadExplicitVertexDistribution(
+                config.explicit_distribution_filename, config.explicit_distribution_is_prefix_sum);
+            from    = distribution[rank];
+            to_node = distribution[rank + 1];
             break;
         }
     }