SP-ization

CMakeLists.txt

@@ -73,6 +73,7 @@ option(FF_BUILD_ALL_EXAMPLES "build all examples. Overrides others" OFF)
 option(FF_BUILD_UNIT_TESTS "build non-operator unit tests" OFF)
 option(FF_BUILD_SUBSTITUTION_TOOL "build substitution conversion tool" OFF)
 option(FF_BUILD_VISUALIZATION_TOOL "build substitution visualization tool" ON)
+option(FF_BUILD_SP_IZATION_BENCHMARKING "build sp-ization benchmarking" ON)
 option(FF_BUILD_ARG_PARSER "build command line argument parser" OFF)
 option(FF_BUILD_BIN_EXPORT_MODEL_ARCH "build export-model-arch utility" ON)
 

bin/CMakeLists.txt

@@ -10,6 +10,10 @@ if(FF_BUILD_VISUALIZATION_TOOL)
   add_subdirectory(substitution-to-dot)
 endif()
 
+if(FF_BUILD_SP_IZATION_BENCHMARKING)
+  add_subdirectory(sp_ization_benchmarking)
+endif()
+
 if(FF_BUILD_ARG_PARSER)
   add_subdirectory(arg_parser)
 endif()

bin/sp_ization_benchmarking/CMakeLists.txt

@@ -0,0 +1,9 @@
+ff_add_executable(
+  NAME
+    sp-ization-benchmarking
+  SRC_PATTERNS
+    *.cc
+  DEPS
+    utils
+    rapidcheck
+)

bin/sp_ization_benchmarking/distributions.cc

@@ -0,0 +1,55 @@
+#include "distributions.h"
+
+namespace FlexFlow {
+
+Constant::Constant(float val) : val(val) {}
+
+float Constant::operator()() const {
+  return val;
+}
+
+Uniform::Uniform(float a, float b) : a(a), b(b) {}
+
+float Uniform::operator()() const {
+  return a + ((static_cast<double>(std::rand()) / RAND_MAX) * (b - a));
+}
+
+Bernoulli::Bernoulli(float p) : p(p) {}
+
+float Bernoulli::operator()() const {
+  return (Uniform(0, 1)() < p);
+}
+
+Binary::Binary(float a, float b, float p) : a(a), b(b), p(p) {}
+
+float Binary::operator()() const {
+  return (Bernoulli(p)() ? a : b);
+}
+
+Chooser::Chooser(std::vector<float> items) : items(items) {}
+
+float Chooser::operator()() const {
+  return items[std::rand() % items.size()];
+}
+
+UniformNoise::UniformNoise(float lower, float upper)
+    : lower(lower), upper(upper) {}
+
+float UniformNoise::operator()() const {
+  return Uniform(lower, upper)();
+}
+
+float NoNoise::operator()() const {
+  return 1;
+}
+
+GaussianNoise::GaussianNoise(float mean, float stddev)
+    : mean(mean), stddev(stddev) {}
+
+float GaussianNoise::operator()() const {
+  static std::default_random_engine generator;
+  static std::normal_distribution<float> distribution(mean, stddev);
+  return distribution(generator);
+}
+
+} // namespace FlexFlow

bin/sp_ization_benchmarking/distributions.h

@@ -0,0 +1,81 @@
+#ifndef _FLEXFLOW_DISTRIBUTIONS_H
+#define _FLEXFLOW_DISTRIBUTIONS_H
+
+#include "utils/graph/node/node.dtg.h"
+#include <random>
+#include <unordered_map>
+#include <unordered_set>
+
+namespace FlexFlow {
+
+struct Constant {
+  float val;
+  Constant(float val = 1);
+  float operator()() const;
+};
+
+struct Uniform {
+  float a, b;
+  Uniform(float a = 0, float b = 1);
+  float operator()() const;
+};
+
+struct Bernoulli {
+  float p;
+  Bernoulli(float p = 0.5);
+  float operator()() const;
+};
+
+struct Binary {
+  float a, b, p;
+  Binary(float a = 0, float b = 1, float p = 0.5);
+  float operator()() const;
+};
+
+struct Chooser {
+  std::vector<float> items;
+  Chooser(std::vector<float>);
+  float operator()() const;
+};
+
+struct UniformNoise {
+  float lower, upper;
+  UniformNoise(float lower = 0.9, float upper = 1.1);
+  float operator()() const;
+};
+
+struct NoNoise {
+  float operator()() const;
+};
+
+struct GaussianNoise {
+  float mean, stddev;
+  GaussianNoise(float mean = 1, float stddev = .1);
+  float operator()() const;
+};
+
+template <typename Dist>
+std::unordered_map<Node, float>
+    make_cost_map(std::unordered_set<Node> const &nodes,
+                  Dist const &distribution) {
+  std::unordered_map<Node, float> cost_map;
+  for (Node const &node : nodes) {
+    cost_map[node] = distribution();
+  }
+  return cost_map;
+}
+
+template <typename Noise>
+std::unordered_map<Node, float>
+    add_noise_to_cost_map(std::unordered_map<Node, float> cost_map,
+                          Noise const &noise) {
+  std::unordered_map<Node, float> noisy_cost_map;
+  for (auto const &[node, cost] : cost_map) {
+    noisy_cost_map[node] = noise() * cost;
+  }
+  return noisy_cost_map;
+}
+
+} // namespace FlexFlow
+
+#endif

bin/sp_ization_benchmarking/nasnet_bench_graph_generator.h

@@ -0,0 +1,126 @@
+// For context, see https://arxiv.org/abs/1902.09635 &&
+// https://github.com/google-research/nasbench/blob/master/nasbench/api.py
+
+#include "utils/containers.h"
+#include "utils/containers/all_of.h"
+#include "utils/containers/repeat.h"
+#include "utils/containers/transform.h"
+#include "utils/graph/algorithms.h"
+#include "utils/graph/digraph/algorithms.h"
+#include "utils/graph/digraph/algorithms/is_acyclic.h"
+#include "utils/graph/digraph/algorithms/materialize_digraph_view.h"
+#include "utils/graph/digraph/algorithms/transitive_reduction.h"
+#include "utils/graph/instances/adjacency_digraph.h"
+#include "utils/graph/node/algorithms.h"
+#include "utils/graph/series_parallel/digraph_generation.h"
+#include <optional>
+#include <vector>
+
+constexpr size_t MIN_NODES = 6;
+constexpr size_t MAX_NODES = 8;
+constexpr size_t MIN_EDGES = 8;
+constexpr size_t MAX_EDGES = 11;
+constexpr size_t NUM_CELLS = 9;
+
+using AdjacencyMatrix = std::vector<std::vector<bool>>;
+namespace FlexFlow {
+struct NasNetBenchConfig {
+  AdjacencyMatrix adjacency_matrix;
+};
+
+bool is_valid_config(NasNetBenchConfig const &config) {
+  AdjacencyMatrix const &matrix = config.adjacency_matrix;
+  const size_t size = matrix.size();
+
+  auto is_valid_size = [](size_t s) {
+    return s >= MIN_NODES && s <= MAX_NODES;
+  };
+
+  auto is_square_matrix = [&](auto const &m) {
+    return all_of(m, [&](const auto &row) { return row.size() == size; });
+  };
+
+  auto is_upper_triangular = [&](auto const &m) {
+    for (size_t i = 0; i < size; ++i) {
+      for (size_t j = 0; j <= i; ++j) {
+        if (matrix[i][j]) {
+          return false;
+        }
+      }
+    }
+    return true;
+  };
+
+  return is_valid_size(size) && is_square_matrix(matrix) &&
+         is_upper_triangular(matrix);
+}
+
+bool is_valid_cell(DiGraphView const &g) {
+  return (is_acyclic(g)) && (get_sources(g).size() == 1) &&
+         (get_sinks(g).size() == 1) && (num_edges(g) <= MAX_EDGES) &&
+         (num_edges(g) >= MIN_EDGES) && (num_edges(g) <= MAX_NODES) &&
+         (num_edges(g) >= MIN_NODES) &&
+         (num_edges(g) > num_nodes(g)); // filter linear cell and diamond cell
+}
+
+NasNetBenchConfig generate_random_config() {
+  static std::uniform_int_distribution<> size_dist(MIN_NODES, MAX_NODES);
+  Binary bin = Binary(0, 1);
+
+  size_t num_nodes = Uniform(MIN_NODES, MAX_NODES)();
+  std::vector<std::vector<bool>> matrix(num_nodes,
+                                        std::vector<bool>(num_nodes, false));
+
+  for (size_t i = 0; i < num_nodes; ++i) {
+    for (size_t j = i + 1; j < num_nodes; ++j) {
+      matrix[i][j] = bin();
+    }
+  }
+
+  return {matrix};
+}
+
+std::optional<DiGraph>
+    maybe_generate_nasnet_bench_cell(NasNetBenchConfig const &config) {
+  if (!is_valid_config(config)) {
+    return std::nullopt;
+  }
+
+  DiGraph g = DiGraph::create<AdjacencyDiGraph>();
+  std::vector<Node> nodes = add_nodes(g, config.adjacency_matrix.size());
+
+  for (size_t i = 0; i < nodes.size(); ++i) {
+    for (size_t j = i + 1; j < nodes.size(); ++j) {
+      if (config.adjacency_matrix[i][j]) {
+        g.add_edge(DirectedEdge{nodes[i], nodes[j]});
+      }
+    }
+  }
+
+  g = materialize_digraph_view<AdjacencyDiGraph>(transitive_reduction(g));
+
+  if (!is_valid_cell(g)) {
+    return std::nullopt;
+  }
+
+  return g;
+}
+
+DiGraph generate_nasnet_bench_cell() {
+  while (true) {
+    NasNetBenchConfig config = generate_random_config();
+    std::optional<DiGraph> maybe_cell =
+        maybe_generate_nasnet_bench_cell(config);
+    if (maybe_cell) {
+      return maybe_cell.value();
+    }
+  }
+}
+
+DiGraph generate_nasnet_bench_network() {
+  DiGraph g = series_composition(
+      transform(repeat(NUM_CELLS, generate_nasnet_bench_cell),
+                [](auto const cell) -> DiGraphView { return cell; }));
+  return g;
+}
+} // namespace FlexFlow