Implement MCMC mapping/optimization algorithm

lib/compiler/include/compiler/machine_mapping/allowed_machine_views.h

@@ -0,0 +1,21 @@
+#ifndef _FLEXFLOW_COMPILER_ALLOWED_MACHINE_VIEWS_H
+#define _FLEXFLOW_COMPILER_ALLOWED_MACHINE_VIEWS_H
+
+#include "pcg/machine_specification.dtg.h"
+#include "pcg/machine_view.dtg.h"
+#include "pcg/operator_task_space.dtg.h"
+
+namespace FlexFlow {
+
+bool is_valid_machine_view(MachineView const &mv,
+                           OperatorTaskSpace const &task,
+                           MachineSpecification const &ms);
+
+std::unordered_set<MachineView>
+    get_allowed_machine_views(MachineSpecification const &machine_spec,
+                              OperatorTaskSpace const &task,
+                              DeviceType device_type);
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/machine_mapping/apply_substitution_and_update_machine_mapping.h

@@ -0,0 +1,32 @@
+#ifndef _FLEXFLOW_LIB_SUBSTITUTIONS_INCLUDE_SUBSTITUTIONS_APPLY_SUBSTITUTION_APPLY_SUBSTITUTION_AND_UPDATE_MACHINE_MAPPING_H
+#define _FLEXFLOW_LIB_SUBSTITUTIONS_INCLUDE_SUBSTITUTIONS_APPLY_SUBSTITUTION_APPLY_SUBSTITUTION_AND_UPDATE_MACHINE_MAPPING_H
+
+#include "compiler/search_result.dtg.h"
+#include "substitutions/pcg_pattern_match.dtg.h"
+#include "substitutions/sub_parallel_computation_graph.dtg.h"
+#include "substitutions/substitution.dtg.h"
+
+namespace FlexFlow {
+/**
+ * @brief Applies \p substitution to \p mapped_pcg at the location specified by
+ * \p match, returning the resulting SearchResult (mapped pcg)
+ *
+ * @param mapped_pcg
+ * @param substitution
+ * @param match The location at which to apply substitution. This location in
+ * sub_pcg should match substitution's PCGPattern. Likely created by running
+ * FlexFlow::find_pattern_matches(PCGPattern const &,
+ * SubParallelComputationGraph const &).
+ * @return SearchResult A mapped pcg similar to mapped_pcg, but with
+ * the subgraph of the pcg specified by match replaced with the result of the
+ * output expression of substitution and the machine mapping updated to account
+ * for the new output
+ */
+SearchResult apply_substitution_and_update_machine_mapping(
+    SearchResult const &mapped_pcg,
+    Substitution const &sub,
+    PCGPatternMatch const &match);
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/machine_mapping/machine_mapping_mutation_set.h

@@ -0,0 +1,19 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MCMC_MACHINE_MAPPING_MUTATION_SET_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_MCMC_MACHINE_MAPPING_MUTATION_SET_H
+
+#include "compiler/machine_mapping/machine_mapping.h"
+#include "compiler/search_result.dtg.h"
+
+namespace FlexFlow {
+std::optional<MachineMapping>
+    get_random_mapping(ParallelComputationGraph &pcg,
+                       MachineSpecification const &resources,
+                       DeviceType const &device_type);
+
+std::optional<MachineMapping>
+    get_random_mutation(SearchResult mapped_pcg,
+                        MachineSpecification const &resource,
+                        DeviceType const &device_type);
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/mcmc/generic_mcmc_algorithm.h

@@ -0,0 +1,41 @@
+#ifndef _FLEXFLOW_COMPILER_MCMC_GENERIC_MCMC_ALGORITHM_H
+#define _FLEXFLOW_COMPILER_MCMC_GENERIC_MCMC_ALGORITHM_H
+
+#include "compiler/mcmc/generic_mcmc_config.dtg.h"
+#include "utils/containers/transform.h"
+#include "utils/nonnegative_int/nonnegative_range.h"
+#include "utils/optional.h"
+#include "utils/random_utils.h"
+
+namespace FlexFlow {
+
+// SamplingFn : State -> std::optional<State>
+// CostFn : State -> float
+
+template <typename State, typename SamplingFn, typename CostFn>
+State run_mcmc(State const &starting_state,
+               SamplingFn const &sampler,
+               CostFn const &cost,
+               GenericMCMCConfig const &search_config) {
+  State best_state = starting_state;
+  State current_state = best_state;
+  for (nonnegative_int i : nonnegative_range(search_config.num_iterations)) {
+    std::optional<State> maybe_new_state =
+        transform(sampler(current_state), [&](State const &s) {
+          float delta = cost(s) - cost(best_state);
+          if (randf() < exp(-delta / search_config.temperature)) {
+            if (delta < 0) {
+              best_state = s;
+            }
+            return s;
+          }
+          return current_state;
+        });
+    current_state = or_else(maybe_new_state, [&]() { return current_state; });
+  }
+  return best_state;
+}
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/mcmc/generic_mcmc_config.struct.toml

@@ -0,0 +1,19 @@
+namespace = "FlexFlow"
+name = "GenericMCMCConfig"
+features = [
+  "eq",
+  "hash",
+  "fmt",
+]
+
+includes = [
+  "utils/nonnegative_int/nonnegative_int.h"
+]
+
+[[fields]]
+name = "temperature"
+type = "float"
+
+[[fields]]
+name = "num_iterations"
+type = "::FlexFlow::nonnegative_int"

lib/compiler/include/compiler/mcmc/mcmc_over_mapped_pcg.h

@@ -0,0 +1,23 @@
+#ifndef _FLEXFLOW_COMPILER_MCMC_OVER_MAPPED_PCG_H
+#define _FLEXFLOW_COMPILER_MCMC_OVER_MAPPED_PCG_H
+
+#include "compiler/cost_estimator/runtime_only_cost_estimator.h"
+#include "compiler/mcmc/mcmc_over_mapped_pcg_config.dtg.h"
+#include "compiler/search_result.dtg.h"
+#include "pcg/computation_graph.h"
+#include "pcg/machine_specification.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h"
+#include "substitutions/sub_parallel_computation_graph.h"
+#include "substitutions/substitution.h"
+
+namespace FlexFlow {
+
+SearchResult
+    mcmc_over_mapped_pcg(ParallelComputationGraph &pcg,
+                         RuntimeOnlyCostEstimator const &cost_estimator,
+                         MachineSpecification const &resources,
+                         MCMCOverMappedPCGConfig const &search_config);
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/mcmc/mcmc_over_mapped_pcg_config.struct.toml

@@ -0,0 +1,28 @@
+namespace = "FlexFlow"
+name = "MCMCOverMappedPCGConfig"
+features = [
+  "eq",
+  "hash",
+  "fmt",
+]
+
+includes = [
+  "pcg/device_type.dtg.h",
+  "utils/nonnegative_int/nonnegative_int.h"
+]
+
+[[fields]]
+name = "temperature"
+type = "float"
+
+[[fields]]
+name = "num_iterations"
+type = "::FlexFlow::nonnegative_int"
+
+[[fields]]
+name = "substitution_frequency"
+type = "float"
+
+[[fields]]
+name = "device_type"
+type = "::FlexFlow::DeviceType"

lib/compiler/include/compiler/search_result.h

@@ -0,0 +1,13 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_GRAPH_OPTIMIZE_RESULT_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_GRAPH_OPTIMIZE_RESULT_H
+
+#include "compiler/search_result.dtg.h"
+
+namespace FlexFlow {
+
+std::string format_as(SearchResult const &);
+std::ostream &operator<<(std::ostream &, SearchResult const &);
+
+} // namespace FlexFlow
+
+#endif

lib/compiler/include/compiler/search_result.struct.toml

@@ -0,0 +1,17 @@
+namespace = "FlexFlow"
+name = "SearchResult"
+features = [
+]
+
+includes = [
+  "pcg/parallel_computation_graph/parallel_computation_graph.h",
+  "compiler/machine_mapping/machine_mapping.h",
+]
+
+[[fields]]
+name = "pcg"
+type = "::FlexFlow::ParallelComputationGraph"
+
+[[fields]]
+name = "machine_mapping"
+type = "::FlexFlow::MachineMapping"

lib/compiler/src/compiler/allowed_machine_views.cc

@@ -57,6 +57,8 @@ static std::unordered_set<MachineView>
         product(transform(tensor_dims, [](positive_int num_devices) {
           return nonnegative_int{num_devices.int_from_positive_int() - 1};
         }));
+    min_num_devices_with_full_stride_volume =
+        std::max(min_num_devices_with_full_stride_volume, 1_n);
     return ceildiv(total_devices,
                    positive_int{min_num_devices_with_full_stride_volume});
   };

lib/compiler/src/compiler/machine_mapping/apply_substitution_and_update_machine_mapping.cc

@@ -0,0 +1,78 @@
+#include "compiler/machine_mapping/apply_substitution_and_update_machine_mapping.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph_edge.h"
+#include "pcg/parallel_computation_graph/parallel_tensor_guid_t.h"
+#include "substitutions/apply_substitution/apply_substitution.h"
+#include "substitutions/apply_substitution/evaluate_substitution_output.h"
+#include "substitutions/apply_substitution/output_expr_to_result_sub_pcg_mapping.h"
+#include "substitutions/open_parallel_tensor_guid_t.h"
+#include "substitutions/pcg_pattern_match.h"
+#include "substitutions/sub_parallel_computation_graph.h"
+#include "substitutions/sub_parallel_computation_graph_data.dtg.h"
+#include "substitutions/sub_parallel_computation_graph_edge.h"
+#include "utils/containers/filter.h"
+#include "utils/containers/is_subseteq_of.h"
+#include "utils/containers/keys.h"
+#include "utils/containers/merge_maps.h"
+#include "utils/containers/restrict_keys.h"
+#include "utils/containers/set_minus.h"
+#include "utils/containers/values.h"
+#include "utils/containers/vector_of.h"
+#include "utils/random_utils.h"
+#include <libassert/assert.hpp>
+
+namespace FlexFlow {
+
+SearchResult apply_substitution_and_update_machine_mapping(
+    SearchResult const &mapped_pcg,
+    Substitution const &sub,
+    PCGPatternMatch const &match) {
+  SubParallelComputationGraph spcg = sub_pcg_from_full_pcg(mapped_pcg.pcg);
+
+  std::pair<SubParallelComputationGraph, OutputExprToResultSubPCGMapping>
+      substitution_output_result =
+          evaluate_substitution_output(spcg, sub, match);
+
+  SubParallelComputationGraph post_substitution_graph =
+      apply_substitution_from_output_result(
+          substitution_output_result, spcg, sub, match);
+
+  std::unordered_map<parallel_layer_guid_t, ParallelLayerAttrs> post_node_data =
+      get_sub_pcg_data(post_substitution_graph).node_data;
+
+  std::unordered_set<parallel_layer_guid_t>
+      substitution_output_parallel_layers =
+          get_parallel_layers(substitution_output_result.first);
+
+  std::unordered_map<parallel_layer_guid_t, MachineView> machine_views =
+      mapped_pcg.machine_mapping.machine_views;
+
+  std::unordered_set<parallel_layer_guid_t> matched_nodes =
+      unordered_set_of(values(match.node_assignment));
+
+  std::vector<MachineView> substituted_machine_views = vector_of(
+      transform(matched_nodes, [&](parallel_layer_guid_t const &node) {
+        return machine_views.at(node);
+      }));
+
+  for (parallel_layer_guid_t layer : substitution_output_parallel_layers) {
+    machine_views.insert_or_assign(layer,
+                                   select_random(substituted_machine_views));
+  }
+
+  ASSERT(is_subseteq_of(keys(post_node_data), keys(machine_views)));
+
+  std::unordered_map<parallel_layer_guid_t, MachineView>
+      post_node_machine_views =
+          filter(machine_views,
+                 [&](std::pair<parallel_layer_guid_t, MachineView> const &p) {
+                   return post_node_data.count(p.first);
+                 });
+
+  ASSERT(keys(post_node_data) == keys(post_node_machine_views));
+
+  return SearchResult{
+      pcg_from_sub_pcg_by_dropping_inputs(post_substitution_graph),
+      MachineMapping{post_node_machine_views}};
+}
+
+} // namespace FlexFlow

lib/compiler/src/compiler/machine_mapping/machine_mapping_mutation_set.cc

@@ -0,0 +1,52 @@
+#include "compiler/machine_mapping/machine_mapping_mutation_set.h"
+#include "compiler/allowed_machine_views.h"
+#include "pcg/machine_view.h"
+#include "pcg/operator_task_space.h"
+#include "utils/containers/vector_of.h"
+#include "utils/nonnegative_int/nonnegative_range.h"
+#include "utils/random_utils.h"
+
+namespace FlexFlow {
+
+std::optional<MachineMapping>
+    get_random_mapping(ParallelComputationGraph &pcg,
+                       MachineSpecification const &resources,
+                       DeviceType const &device_type) {
+  std::vector<parallel_layer_guid_t> layers = topological_ordering(pcg);
+  std::unordered_map<parallel_layer_guid_t, MachineView> machine_views;
+  for (parallel_layer_guid_t layer : layers) {
+    OperatorTaskSpace task = get_operator_task_space(pcg, layer);
+    std::unordered_set<MachineView> allowed_machine_views =
+        get_allowed_machine_views(resources, task, DeviceType::GPU);
+    if (allowed_machine_views.empty()) {
+      return std::nullopt;
+    }
+    machine_views.insert(
+        {layer, select_random(vector_of(allowed_machine_views))});
+  }
+  return MachineMapping{machine_views};
+}
+
+std::optional<MachineMapping>
+    get_random_mutation(SearchResult mapped_pcg,
+                        MachineSpecification const &resources,
+                        DeviceType const &device_type) {
+  ParallelComputationGraph pcg = mapped_pcg.pcg;
+  std::vector<parallel_layer_guid_t> layers = topological_ordering(pcg);
+  if (layers.size() == 0) {
+    return std::nullopt;
+  }
+  parallel_layer_guid_t random_layer = select_random(layers);
+
+  MachineMapping machine_mapping = mapped_pcg.machine_mapping;
+  MachineView machine_view = machine_mapping.machine_views.at(random_layer);
+  OperatorTaskSpace task = get_operator_task_space(pcg, random_layer);
+
+  std::vector<MachineView> allowed_machine_views =
+      vector_of(get_allowed_machine_views(resources, task, device_type));
+  MachineView random_new_machine_view = select_random(allowed_machine_views);
+
+  machine_mapping.machine_views.at(random_layer) = random_new_machine_view;
+  return machine_mapping;
+}
+} // namespace FlexFlow

lib/compiler/src/compiler/mcmc/generic_mcmc_algorithm.cc

@@ -0,0 +1,15 @@
+#include "compiler/mcmc/generic_mcmc_algorithm.h"
+#include "utils/archetypes/value_type.h"
+
+namespace FlexFlow {
+
+using State = value_type<0>;
+using SamplingFn = std::function<std::optional<State>(State)>;
+using CostFn = std::function<float(State)>;
+
+template State run_mcmc(State const &starting_state,
+                        SamplingFn const &sampler,
+                        CostFn const &cost,
+                        GenericMCMCConfig const &search_config);
+
+} // namespace FlexFlow