fix some errors introduced in merge

lib/compiler/src/compiler/compiler.cc

@@ -1,27 +1,27 @@
 #include "compiler/compiler.h"
 #include "compiler/unity_algorithm/unity_algorithm.h"
+#include "pcg/pcg_from_computation_graph.h"
 #include "utils/overload.h"
 
 namespace FlexFlow {
 
 SearchResult optimize(ComputationGraph const &computation_graph,
                       MachineSpecification const &machine_specification,
                       CostEstimator const &cost_estimator,
                       AlgorithmConfig const &search_config) {
   return search_config.visit<SearchResult>(overload{
       [&](DataParallelismConfig const &config) -> SearchResult {
         throw std::runtime_error(
             "Data parallel search algorithm is not implemented yet");
       },
       [&](UnitySearchConfig const &config) {
         ParallelComputationGraph pcg =
-            parallel_computation_graph_from_computation_graph(
-                computation_graph);
+            pcg_from_computation_graph(computation_graph);
         std::vector<Substitution> substitutions; // TODO: Implement this
         return graph_optimize(
             pcg, cost_estimator, machine_specification, substitutions, config);
       },
   });
 }
 
 } // namespace FlexFlow

lib/compiler/src/compiler/unity_algorithm/unity_algorithm.cc

@@ -12,6 +12,7 @@
 #include "compiler/unity_algorithm/graph_optimize_state.h"
 #include "pcg/machine_specification.dtg.h"
 #include "pcg/operator_task_space.h"
+#include "substitutions/apply_substitution/apply_substitution.h"
 #include "substitutions/pcg_pattern.h"
 #include "substitutions/sub_parallel_computation_graph.h"
 #include "substitutions/substitution.h"
@@ -26,113 +27,113 @@
  * Applies a substitution to all possible positions in PCG
  */
 std::vector<ParallelComputationGraph>
     all_pcgs_obtained_by_applying_a_substitution(
         ParallelComputationGraph const &pcg,
         std::vector<Substitution> const &substitutions) {
   std::vector<ParallelComputationGraph> results;
   SubParallelComputationGraph subpcg = sub_pcg_from_full_pcg(pcg);
   for (Substitution const &substitution : substitutions) {
     for (PCGPatternMatch const &pattern_match :
          find_pattern_matches(substitution.pcg_pattern, subpcg)) {
       SubParallelComputationGraph subpcg_from_substitution =
           apply_substitution(subpcg, substitution, pattern_match);
       results.push_back(
           pcg_from_sub_pcg_by_dropping_inputs(subpcg_from_substitution));
     }
   }
   return results;
 }
 
 SearchResult graph_optimize(ParallelComputationGraph &pcg,
                             CostEstimator const &cost_estimator,
                             MachineSpecification const &resources,
                             std::vector<Substitution> const &substitutions,
                             UnitySearchConfig const &search_config) {
 
   MachineMappingCache cached_subgraph_costs = empty_machine_mapping_cache();
   DeduplicatedPriorityQueue<GraphOptimizeState> candidates;
 
   MachineMappingContext context = MachineMappingContext{
       /*cost_estimator=*/cost_estimator,
       /*allowed_machine_views=*/
       [&](UnmappedOpCostEstimateKey const &key,
           MachineSpecification const &resources)
           -> std::unordered_set<MachineView> {
         return get_allowed_machine_views(
             resources, key.op_task_space, DeviceType::GPU);
       },
   };
 
   auto optimize_pcg = [&](ParallelComputationGraph const &pcg)
       -> std::pair<GraphOptimizeState, std::optional<MachineMapping>> {
     PCGBinarySPDecomposition sp_decomp =
         expect(get_pcg_balanced_binary_sp_decomposition(pcg),
                "Failed to get SP decomposition of PCG");
 
     MachineMappingProblemTree problem_tree =
         get_machine_mapping_problem_tree(pcg, sp_decomp);
     MachineMappingConstraints constraints =
         get_unconstrained_solution_for_layers(get_all_leaf_paths(problem_tree));
 
     MachineMappingResult mm_result = get_optimal_machine_mapping(
         cached_subgraph_costs,
         context,
         get_machine_mapping_problem_tree(pcg, sp_decomp),
         resources,
         constraints);
 
     return {
         GraphOptimizeState{
             /*pcg=*/pcg,
             /*runtime_with_optimal_mm=*/get_runtime_cost(mm_result),
         },
         get_machine_mapping_from_machine_mapping_result(sp_decomp, mm_result),
     };
   };
 
   GraphOptimizeState best_state = optimize_pcg(pcg).first;
   candidates.push(best_state);
 
   for (int iteration = 0;
        !candidates.empty() && iteration < search_config.budget;
        ++iteration) {
     GraphOptimizeState current_state = candidates.top();
     candidates.pop();
 
     if (current_state < best_state) {
       best_state = current_state;
     } else if (current_state.runtime_with_optimal_mm >
                best_state.runtime_with_optimal_mm * search_config.alpha) {
       continue;
     }
 
     for (ParallelComputationGraph const &new_pcg :
          all_pcgs_obtained_by_applying_a_substitution(current_state.pcg,
                                                       substitutions)) {
       std::optional<GraphOptimizeState> new_pcg_optimize_result =
           optimize_pcg(new_pcg).first;
       if (new_pcg_optimize_result == std::nullopt) {
         continue;
       }
       GraphOptimizeState new_state = new_pcg_optimize_result.value();
       if (new_state.runtime_with_optimal_mm <= search_config.threshold &&
           get_nodes(new_pcg.raw_graph).size() <= search_config.max_num_ops) {
         candidates.push(new_state);
       }
     }
   }
 
   std::optional<MachineMapping> best_mapping =
       optimize_pcg(best_state.pcg).second;
 
   if (best_mapping == std::nullopt) {
     throw std::runtime_error("Failed to find any solutions");
   }
 
   return SearchResult{
       /*pcg=*/best_state.pcg,
       /*machine_mapping=*/best_mapping.value(),
   };
 }
 
 } // namespace FlexFlow

lib/compiler/test/src/compiler/machine_mapping/machine_mapping_problem_tree/get_machine_mapping_problem_tree.cc

@@ -98,7 +98,7 @@ TEST_SUITE(FF_TEST_SUITE) {
     PCGOperatorAttrs input_attrs = PCGOperatorAttrs{InputAttrs{}};
 
     auto make_operator_task_space = [&](ParallelTensorShape const &shape) {
-      std::vector<int> degrees;
+      std::vector<nonnegative_int> degrees;
       extend(degrees, vector_of(ff_ordered_shard_degrees(shape)));
       degrees.push_back(get_sum_degree(shape));
       degrees.push_back(get_discard_copy_degree(shape));

lib/compiler/test/src/compiler/unity_algorithm/graph_optimize_state.cc

@@ -10,12 +10,13 @@ TEST_SUITE(FF_TEST_SUITE) {
 
     ParallelTensorShape input_shape = ParallelTensorShape{
         ParallelTensorDims{
-            FFOrdered<ShardParallelDim>{
-                ShardParallelDim{10, 1},
-            },
+            FFOrdered<ShardParallelDim>{ShardParallelDim{
+                nonnegative_int{10},
+                nonnegative_int{1},
+            }},
             ReplicaParallelDimSet{
-                SumDegree{1},
-                DiscardCopyDegree{1},
+                SumDegree{nonnegative_int{1}},
+                DiscardCopyDegree{nonnegative_int{1}},
             },
         },
         DataType::FLOAT,

lib/compiler/test/src/compiler/unity_algorithm/unity_algorithm.cc

@@ -7,81 +7,43 @@
 #include "op-attrs/shard_parallel_dim.h"
 #include "pcg/computation_graph_builder.h"
 #include "pcg/parallel_computation_graph/parallel_computation_graph_builder.h"
+#include "pcg/pcg_from_computation_graph.h"
 #include "utils/integer_conversions.h"
 
 using namespace FlexFlow;
 
 TEST_SUITE(FF_TEST_SUITE) {
   TEST_CASE("graph_optimize") {
-    // TODO: recover this by implementing
-    // parallel_computation_graph_from_computation_graph ComputationGraph cg =
-    // [&] {
-    //   ComputationGraphBuilder b;
-    //   TensorShape input_tensor_shape = TensorShape{
-    //     TensorDims{
-    //       FFOrdered<size_t> {32, 64},
-    //     },
-    //     DataType::FLOAT,
-    //   };
-    //   tensor_guid_t t = b.create_input(input_tensor_shape, CreateGrad::YES);
-    //   t = b.dense(t,
-    //               /*outDim=*/16,
-    //               /*activation=*/std::nullopt);
-    //   t = b.gelu(t);
-    //   t = b.dense(t,
-    //               /*outDim=*/12,
-    //               /*activation=*/std::nullopt,
-    //               /*use_bias=*/false,
-    //               /*data_type=*/DataType::FLOAT,
-    //               /*kernel_initializer=*/std::nullopt,
-    //               /*bias_initializer=*/std::nullopt);
-    //   t = b.relu(t);
-    //   t = b.dense(t,
-    //               /*outDim=*/8,
-    //               /*activation=*/Activation::RELU);
-    //   return b.computation_graph;
-    // }();
-
-    // ParallelComputationGraph pcg =
-    // parallel_computation_graph_from_computation_graph(cg);
-
-    ParallelComputationGraph pcg = [&] {
-      ParallelComputationGraphBuilder b;
-      int in_channels = 24;
-      int batch_size = 4;
-      int batch_degree = 2;
-      parallel_tensor_guid_t t = b.create_input_tensor(ParallelTensorShape{
-          ParallelTensorDims{
-              FFOrdered<ShardParallelDim>{
-                  ShardParallelDim{size_t_from_int(batch_size), batch_degree},
-                  ShardParallelDim{size_t_from_int(in_channels), 1},
-              },
-              ReplicaParallelDimSet{
-                  SumDegree{1},
-                  DiscardCopyDegree{1},
-              },
+    ComputationGraph cg = [&] {
+      ComputationGraphBuilder b;
+      TensorShape input_tensor_shape = TensorShape{
+          TensorDims{
+              FFOrdered<nonnegative_int>{nonnegative_int{32},
+                                         nonnegative_int{64}},
           },
           DataType::FLOAT,
-      });
+      };
+      tensor_guid_t t = b.create_input(input_tensor_shape, CreateGrad::YES);
       t = b.dense(t,
-                  /*outDim=*/16,
+                  /*outDim=*/nonnegative_int{16},
                   /*activation=*/std::nullopt);
       t = b.gelu(t);
       t = b.dense(t,
-                  /*outDim=*/12,
+                  /*outDim=*/nonnegative_int{12},
                   /*activation=*/std::nullopt,
                   /*use_bias=*/false,
                   /*data_type=*/DataType::FLOAT,
                   /*kernel_initializer=*/std::nullopt,
                   /*bias_initializer=*/std::nullopt);
       t = b.relu(t);
       t = b.dense(t,
-                  /*outDim=*/8,
+                  /*outDim=*/nonnegative_int{8},
                   /*activation=*/Activation::RELU);
-
-      return b.pcg;
+      return b.computation_graph;
     }();
 
+    ParallelComputationGraph pcg = pcg_from_computation_graph(cg);
+
     CostEstimator cost_estimator = make_fake_cost_estimator(
         [](OpCostEstimateKey const &k) {
           return OpCostMetrics{
@@ -93,9 +55,9 @@ TEST_SUITE(FF_TEST_SUITE) {
         [](TensorSetMovement const &) { return 1.0; });
 
     MachineSpecification full_machine_spec = MachineSpecification{
-        /*num_nodes=*/2,
-        /*num_cpus_per_node=*/1,
-        /*num_gpus_per_node=*/1,
+        /*num_nodes=*/nonnegative_int{2},
+        /*num_cpus_per_node=*/nonnegative_int{1},
+        /*num_gpus_per_node=*/nonnegative_int{1},
         /*inter_node_bandwidth=*/1,
         /*intra_node_bandwidth=*/1,
     };

lib/pcg/include/pcg/parallel_computation_graph/parallel_computation_graph.h

@@ -85,9 +85,6 @@ ParallelComputationGraph without_layer_names(ParallelComputationGraph const &);
 bool pcgs_are_isomorphic(ParallelComputationGraph const &,
                          ParallelComputationGraph const &);
 
-ParallelComputationGraph
-    parallel_computation_graph_from_computation_graph(ComputationGraph const &);
-
 } // namespace FlexFlow
 
 #endif

lib/pcg/src/pcg/parallel_computation_graph/parallel_computation_graph.cc

@@ -247,9 +247,4 @@ bool pcgs_are_isomorphic(ParallelComputationGraph const &lhs,
       .has_value();
 }
 
-ParallelComputationGraph parallel_computation_graph_from_computation_graph(
-    ComputationGraph const &) {
-  NOT_IMPLEMENTED();
-}
-
 } // namespace FlexFlow