Make compiler side device-type-agnostic

lib/compiler/include/compiler/cost_estimator/op_cost_estimate_key.h

@@ -3,7 +3,6 @@
 
 #include "compiler/cost_estimator/op_cost_estimate_key.dtg.h"
 #include "compiler/cost_estimator/runtime_only_op_cost_estimate_key.dtg.h"
-#include "pcg/device_id_t.dtg.h"
 #include "pcg/machine_specification.dtg.h"
 #include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h"
 #include "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h"

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

@@ -14,8 +14,7 @@ bool is_valid_machine_view(MachineView const &mv,
 
 std::unordered_set<MachineView>
     get_allowed_machine_views(MachineComputeResourceSlice const &machine_spec,
-                              OperatorTaskSpace const &task,
-                              DeviceType device_type);
+                              OperatorTaskSpace const &task);
 
 } // namespace FlexFlow
 

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

@@ -4,7 +4,6 @@
 #include "compiler/machine_mapping/machine_mapping.dtg.h"
 #include "compiler/machine_mapping/machine_mapping_result.h"
 #include "compiler/series_parallel/pcg/pcg_binary_sp_decomposition.dtg.h"
-#include "pcg/device_id_t.dtg.h"
 #include "pcg/machine_specification.dtg.h"
 #include "pcg/mapped_parallel_computation_graph/mapped_parallel_computation_graph.dtg.h"
 #include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h"

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

@@ -7,13 +7,11 @@
 namespace FlexFlow {
 std::optional<MachineMapping>
     get_random_mapping(ParallelComputationGraph const &pcg,
-                       MachineComputeSpecification const &resources,
-                       DeviceType const &device_type);
+                       MachineComputeSpecification const &resources);
 
 std::optional<MachineMapping>
     get_random_mutation(SearchResult const &mapped_pcg,
-                        MachineComputeSpecification const &resource,
-                        DeviceType const &device_type);
+                        MachineComputeSpecification const &resource);
 } // namespace FlexFlow
 
 #endif

lib/compiler/include/compiler/machine_mapping/machine_view.h

@@ -6,7 +6,6 @@
 #include "op-attrs/operator_task_space.dtg.h"
 #include "op-attrs/parallel_tensor_dim_degrees.dtg.h"
 #include "op-attrs/task_space_coordinate.dtg.h"
-#include "pcg/device_id_t.dtg.h"
 #include "pcg/machine_compute_specification.dtg.h"
 #include "pcg/mapped_parallel_computation_graph/mapped_operator_task_group.h"
 #include "pcg/mapped_parallel_computation_graph/operator_atomic_task_shard_binding.dtg.h"
@@ -20,8 +19,6 @@ namespace FlexFlow {
 
 nonnegative_int mv_get_expected_task_space_num_dims(MachineView const &mv);
 
-DeviceType get_device_type(MachineView const &mv);
-
 std::vector<stride_t> get_strides(MachineView const &mv);
 
 std::vector<MachineSpecificationDimension>
@@ -50,11 +47,6 @@ std::unordered_set<MachineSpaceCoordinate>
     get_machine_space_coordinates(OperatorTaskSpace const &task,
                                   MachineView const &mv);
 
-std::unordered_set<device_id_t>
-    get_device_ids(OperatorTaskSpace const &task,
-                   MachineView const &mv,
-                   MachineComputeSpecification const &ms);
-
 MachineView make_1d_machine_view(MachineSpaceCoordinate const &start,
                                  MachineSpecificationDimension const &dim,
                                  stride_t stride);

lib/compiler/include/compiler/machine_mapping/start_invariant_machine_view.dtg.toml

@@ -12,7 +12,6 @@ features = [
 
 includes = [
   "compiler/machine_mapping/machine_view_dimension.dtg.h",
-  "pcg/device_type.dtg.h"
 ]
 
 src_includes = [
@@ -23,8 +22,3 @@ src_includes = [
 [[fields]]
 name = "dimensions"
 type = "std::vector<::FlexFlow::MachineViewDimension>"
-
-
-[[fields]]
-name = "device_type"
-type = "::FlexFlow::DeviceType"

lib/compiler/include/compiler/machine_mapping/start_invariant_machine_view.h

@@ -19,8 +19,6 @@ StartInvariantMachineView
 
 nonnegative_int num_dims(StartInvariantMachineView const &mv);
 
-DeviceType get_device_type(StartInvariantMachineView const &mv);
-
 std::vector<stride_t> get_strides(StartInvariantMachineView const &mv);
 
 std::vector<MachineSpecificationDimension>

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

@@ -8,7 +8,6 @@ features = [
 ]
 
 includes = [
-  "pcg/device_type.dtg.h",
   "utils/nonnegative_int/nonnegative_int.h"
 ]
 
@@ -23,7 +22,3 @@ type = "::FlexFlow::nonnegative_int"
 [[fields]]
 name = "substitution_frequency"
 type = "float"
-
-[[fields]]
-name = "device_type"
-type = "::FlexFlow::DeviceType"

lib/compiler/include/compiler/task_graph_simulator/pcg_task_graph.dtg.toml

@@ -9,8 +9,8 @@ includes = [
   "utils/graph/digraph/digraph_view.h",
   "utils/bidict/bidict.h",
   "compiler/task_graph_simulator/pcg_task.dtg.h",
-  "pcg/device_id_t.dtg.h",
   "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h",
+  "pcg/machine_space_coordinate.dtg.h",
   "<unordered_set>",
   "<unordered_map>"
 ]
@@ -32,4 +32,4 @@ type = "::FlexFlow::bidict<::FlexFlow::Node, ::FlexFlow::PCGTask>"
 
 [[fields]]
 name = "node_to_devices"
-type = "std::unordered_map<::FlexFlow::Node, std::unordered_set<::FlexFlow::device_id_t>>"
+type = "std::unordered_map<::FlexFlow::Node, std::unordered_set<::FlexFlow::MachineSpaceCoordinate>>"

lib/compiler/src/compiler/cost_estimator/op_cost_estimate_key.cc

@@ -4,7 +4,6 @@
 #include "compiler/machine_mapping/machine_view.dtg.h"
 #include "compiler/machine_mapping/machine_view.h"
 #include "op-attrs/parallel_tensor_shape.dtg.h"
-#include "pcg/device_id_t.dtg.h"
 #include "pcg/machine_specification.dtg.h"
 #include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h"
 #include <unordered_set>

lib/compiler/src/compiler/machine_mapping/allowed_machine_views.cc

@@ -49,8 +49,7 @@ bool is_valid_machine_view(MachineView const &mv,
  */
 static std::unordered_set<MachineView>
     get_candidate_machine_views(MachineComputeResourceSlice const &machine_spec,
-                                OperatorTaskSpace const &task_space,
-                                DeviceType const &device_type) {
+                                OperatorTaskSpace const &task_space) {
 
   auto get_max_stride_upper_bound =
       [](std::vector<positive_int> const &tensor_dims,
@@ -90,17 +89,15 @@ static std::unordered_set<MachineView>
     return strides;
   };
 
-  auto get_candidate_starts = [](MachineComputeResourceSlice const &slice,
-                                 DeviceType const &device_type)
+  auto get_candidate_starts = [](MachineComputeResourceSlice const &slice)
       -> std::unordered_set<MachineSpaceCoordinate> {
-    ASSERT(device_type == DeviceType::GPU);
 
     std::unordered_set<MachineSpaceCoordinate> result;
     for (nonnegative_int node_idx : nonnegative_range(slice.num_nodes)) {
       for (nonnegative_int device_idx :
            nonnegative_range(slice.num_gpus_per_node)) {
         result.insert(
-            MachineSpaceCoordinate{node_idx, device_idx, device_type});
+            MachineSpaceCoordinate{node_idx, device_idx});
       }
     }
     return result;
@@ -127,7 +124,7 @@ static std::unordered_set<MachineView>
   ASSERT(candidate_strides.size() > 0);
 
   std::unordered_set<MachineSpaceCoordinate> candidate_starts =
-      get_candidate_starts(machine_spec, device_type);
+      get_candidate_starts(machine_spec);
   ASSERT(candidate_starts.size() > 0);
 
   std::unordered_multiset<std::vector<MachineSpecificationDimension>>
@@ -151,11 +148,10 @@ static std::unordered_set<MachineView>
 
 std::unordered_set<MachineView>
     get_allowed_machine_views(MachineComputeResourceSlice const &machine_spec,
-                              OperatorTaskSpace const &task_space,
-                              DeviceType device_type) {
+                              OperatorTaskSpace const &task_space) {
 
   std::unordered_set<MachineView> views =
-      get_candidate_machine_views(machine_spec, task_space, device_type);
+      get_candidate_machine_views(machine_spec, task_space);
   return filter(views, [&](MachineView const &mv) {
     return is_valid_machine_view(mv, task_space, machine_spec);
   });

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

@@ -11,15 +11,14 @@ namespace FlexFlow {
 
 std::optional<MachineMapping>
     get_random_mapping(ParallelComputationGraph const &pcg,
-                       MachineComputeSpecification const &resources,
-                       DeviceType const &device_type) {
+                       MachineComputeSpecification const &resources) {
   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(
-            compute_slice_from_specification(resources), task, DeviceType::GPU);
+            compute_slice_from_specification(resources), task);
     if (allowed_machine_views.empty()) {
       return std::nullopt;
     }
@@ -31,8 +30,7 @@ std::optional<MachineMapping>
 
 std::optional<MachineMapping>
     get_random_mutation(SearchResult const &mapped_pcg,
-                        MachineComputeSpecification const &resources,
-                        DeviceType const &device_type) {
+                        MachineComputeSpecification const &resources) {
   ParallelComputationGraph pcg = mapped_pcg.pcg;
   std::vector<parallel_layer_guid_t> layers = topological_ordering(pcg);
   if (layers.size() == 0) {
@@ -46,7 +44,7 @@ std::optional<MachineMapping>
 
   std::vector<MachineView> allowed_machine_views =
       vector_of(get_allowed_machine_views(
-          compute_slice_from_specification(resources), task, device_type));
+          compute_slice_from_specification(resources), task));
   MachineView random_new_machine_view = select_random(allowed_machine_views);
 
   machine_mapping.machine_views.at(random_layer) = random_new_machine_view;

lib/compiler/src/compiler/machine_mapping/machine_resource_split.cc

@@ -88,7 +88,6 @@ MachineSpaceCoordinate
         /*node_idx=*/(coord.node_idx + split.offset)
             .nonnegative_int_from_positive_int(),
         /*device_idx=*/coord.device_idx,
-        /*device_type=*/coord.device_type,
     };
   } else {
     ASSERT(split.dimension == MachineSpecificationDimension::INTRA_NODE);
@@ -97,7 +96,6 @@ MachineSpaceCoordinate
         /*node_idx=*/coord.node_idx,
         /*device_idx=*/
         (coord.device_idx + split.offset).nonnegative_int_from_positive_int(),
-        /*device_type=*/coord.device_type,
     };
   }
 }

lib/compiler/src/compiler/machine_mapping/machine_view.cc

@@ -33,10 +33,6 @@ nonnegative_int mv_get_expected_task_space_num_dims(MachineView const &mv) {
   return num_elements(get_strides(mv));
 }
 
-DeviceType get_device_type(MachineView const &mv) {
-  return mv.start.device_type;
-}
-
 std::vector<stride_t> get_strides(MachineView const &mv) {
   return transform(mv.dimensions,
                    [](MachineViewDimension const &dim) { return dim.stride; });
@@ -130,7 +126,7 @@ MachineSpaceCoordinate
   nonnegative_int device_idx =
       compute_index(machine_view.start.device_idx, intra_dimension_indices);
   MachineSpaceCoordinate ms_coord = MachineSpaceCoordinate{
-      node_idx, device_idx, get_device_type(machine_view)};
+      node_idx, device_idx};
 
   return ms_coord;
 }
@@ -177,19 +173,6 @@ std::unordered_set<MachineSpaceCoordinate>
                    });
 }
 
-std::unordered_set<device_id_t>
-    get_device_ids(OperatorTaskSpace const &task_space,
-                   MachineView const &mv,
-                   MachineComputeSpecification const &ms) {
-  ASSERT(op_task_space_num_dims(task_space) ==
-         mv_get_expected_task_space_num_dims(mv));
-
-  return transform(get_machine_space_coordinates(task_space, mv),
-                   [&](MachineSpaceCoordinate const &coord) {
-                     return get_device_id(ms, coord);
-                   });
-}
-
 MachineView make_1d_machine_view(MachineSpaceCoordinate const &start,
                                  MachineSpecificationDimension const &dim,
                                  stride_t stride) {

lib/compiler/src/compiler/machine_mapping/start_invariant_machine_view.cc

@@ -18,17 +18,13 @@ MachineView machine_view_from_start_invariant(
 
 StartInvariantMachineView
     start_invariant_from_machine_view(MachineView const &mv) {
-  return StartInvariantMachineView{mv.dimensions, get_device_type(mv)};
+  return StartInvariantMachineView{mv.dimensions};
 }
 
 nonnegative_int num_dims(StartInvariantMachineView const &start_inv_mv) {
   return num_elements(start_inv_mv.dimensions);
 }
 
-DeviceType get_device_type(StartInvariantMachineView const &start_inv_mv) {
-  return start_inv_mv.device_type;
-}
-
 std::vector<stride_t>
     get_strides(StartInvariantMachineView const &start_inv_mv) {
   return transform(start_inv_mv.dimensions,
@@ -45,13 +41,12 @@ std::vector<MachineSpecificationDimension>
 StartInvariantMachineView
     start_invariant_machine_view_from_strides_and_machine_spec_dimensions(
         std::vector<stride_t> const &strides,
-        std::vector<MachineSpecificationDimension> const &dims,
-        DeviceType device_type) {
+        std::vector<MachineSpecificationDimension> const &dims) {
   std::vector<MachineViewDimension> dimensions =
       transform(zip(strides, dims), [&](auto const &p) {
         return MachineViewDimension{p.first, p.second};
       });
-  return StartInvariantMachineView{dimensions, device_type};
+  return StartInvariantMachineView{dimensions};
 }
 
 MachineSpaceOffset get_machine_space_offset(
@@ -60,7 +55,7 @@ MachineSpaceOffset get_machine_space_offset(
     TaskSpaceCoordinate const &coord) {
 
   MachineSpaceCoordinate dummy_start =
-      MachineSpaceCoordinate{0_n, 0_n, get_device_type(start_inv_machine_view)};
+      MachineSpaceCoordinate{0_n, 0_n};
 
   MachineView mv =
       machine_view_from_start_invariant(start_inv_machine_view, dummy_start);