Victor/mcmc cherrypick

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_naive_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,57 @@
+#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 "compiler/mcmc/generic_mcmc_state.h"
+#include "utils/nonnegative_int/nonnegative_range.h"
+#include "utils/random_utils.h"
+#include <optional>
+
+namespace FlexFlow {
+
+template <typename State, typename ScoringFunc>
+void modify_state_for_minimization(
+    Generic_MCMC_state<State, float> &best_state,
+    Generic_MCMC_state<State, float> &current_state,
+    State candidate,
+    ScoringFunc scorer,
+    float temperature) {
+  float best_estimate = best_state.get_score();
+  float new_estimate = scorer(candidate);
+  float delta = new_estimate - best_estimate;
+  if (delta < 0 || (randf() < exp(-delta / temperature))) {
+    current_state = Generic_MCMC_state<State, float>(candidate, new_estimate);
+    if (delta < 0) {
+      best_state = current_state;
+    }
+  }
+}
+
+// GeneratingFunc : State -> nn_int -> std::optional<State>
+// ScoringFunc : State -> float
+
+template <typename State, typename GeneratingFunc, typename ScoringFunc>
+Generic_MCMC_state<State, float>
+    minimize_score(State const &starting_state,
+                   GeneratingFunc const &generator,
+                   ScoringFunc const &scorer,
+                   GenericMCMCConfig const &search_config) {
+  using MCMCState = Generic_MCMC_state<State, float>;
+  MCMCState best_state = MCMCState(starting_state, scorer(starting_state));
+  MCMCState current_state = best_state;
+  for (nonnegative_int i : nonnegative_range(search_config.num_iterations)) {
+    std::optional<State> candidate = generator(current_state.get_state(), i);
+    if (candidate != std::nullopt) {
+      modify_state_for_minimization(best_state,
+                                    current_state,
+                                    candidate.value(),
+                                    scorer,
+                                    search_config.temperature);
+    }
+  }
+  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/generic_mcmc_state.h

@@ -0,0 +1,28 @@
+
+#ifndef _FLEXFLOW_COMPILER_MCMC_GENERIC_MCMC_STATE_H
+#define _FLEXFLOW_COMPILER_MCMC_GENERIC_MCMC_STATE_H
+#include "utils/nonnegative_int/nonnegative_int.h"
+
+namespace FlexFlow {
+
+template <typename State, typename Score>
+struct Generic_MCMC_state {
+public:
+  Generic_MCMC_state(State const &state, Score const &score)
+      : state(state), score(score) {}
+
+  State const &get_state() const {
+    return state;
+  }
+  Score const &get_score() const {
+    return score;
+  }
+
+private:
+  State state;
+  Score score;
+};
+
+} // namespace FlexFlow
+
+#endif

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

@@ -0,0 +1,22 @@
+#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_graph_optimize(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_interval"
+type = "::FlexFlow::nonnegative_int"
+
+[[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});
   };