task-spec reorganization

Author: lockshaw

.flake/pkgs/ffdb/ffdb.py

@@ -1,8 +1,8 @@
-from proj.config_file import get_config_root
+from proj import get_repo_root
 from pathlib import Path
 import gdb
 
-gdb.execute(f'directory {get_config_root(Path.cwd())}')
+gdb.execute(f'directory {get_repo_root(Path.cwd())}')
 gdb.prompt_hook = lambda x: '(ffdb) '
 gdb.execute('set history save on')
 gdb.execute('catch throw')

flake.lock

@@ -18,7 +18,8 @@
     flake-utils.url = "github:numtide/flake-utils";
 
     proj-repo = {
-      url = "git+https://git.sr.ht/~lockshaw/proj?ref=emulated-fs";
+      # url = "git+https://git.sr.ht/~lockshaw/proj?ref=emulated-fs";
+      url = "git+file:///home/lockshaw/x/proj/proj?ref=emulated-fs";
       inputs.nixpkgs.follows = "nixpkgs";
       inputs.flake-utils.follows = "flake-utils";
     };

flake.nix

@@ -17,9 +17,6 @@ struct GraphOptimizeState {
   bool operator<(GraphOptimizeState const &other) const;
 };
 
-std::string format_as(GraphOptimizeState const &);
-std::ostream &operator<<(std::ostream &, GraphOptimizeState const &);
-
 } // namespace FlexFlow
 
 namespace std {

lib/compiler/include/compiler/graph_optimize_state.h

@@ -59,6 +59,8 @@ MachineView make_1d_machine_view(MachineSpaceCoordinate const &start,
                                  MachineSpecificationDimension const &dim,
                                  stride_t stride);
 
+MachineView make_single_device_machine_view(MachineSpaceCoordinate const &);
+
 OperatorAtomicTaskShardBinding
   operator_atomic_task_shard_binding_from_machine_view(ComputationGraphOpAttrs const &,
                                                        std::vector<ParallelTensorDimDegrees> const &,

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

@@ -12,6 +12,10 @@ MappedParallelComputationGraph
     ParallelComputationGraph const &pcg,
     MachineMapping const &mapping) {
 
+  std::unordered_set<parallel_layer_guid_t> pcg_layers = get_parallel_layers(pcg);
+  std::unordered_set<parallel_layer_guid_t> mapped_layers = keys(mapping.machine_views);
+  ASSERT(pcg_layers == mapped_layers);
+
   return MappedParallelComputationGraph{
     /*pcg=*/pcg,
     /*mapped_tasks=*/
@@ -24,6 +28,7 @@ MappedParallelComputationGraph
           std::vector<ParallelTensorDimDegrees> inputs_dim_degrees = 
             get_incoming_input_degrees(pcg, l);
 
+          ASSERT(contains_key(mapping.machine_views, l));
           MachineView machine_view = mapping.machine_views.at(l);
 
           return mapped_operator_task_group_from_machine_view(

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

@@ -193,6 +193,11 @@ MachineView make_1d_machine_view(MachineSpaceCoordinate const &start,
       start, {stride}, {dim});
 }
 
+MachineView make_single_device_machine_view(MachineSpaceCoordinate const &coord) {
+  return machine_view_from_strides_and_machine_spec_dimensions(
+      coord, {}, {});
+}
+
 OperatorAtomicTaskShardBinding
   operator_atomic_task_shard_binding_from_machine_view(ComputationGraphOpAttrs const &op_attrs,
                                                        std::vector<ParallelTensorDimDegrees> const &inputs_dim_degrees,

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

@@ -1,14 +1,17 @@
 #include "compiler/graph_optimize_state.h"
 #include "compiler/machine_mapping/machine_mapping.dtg.h"
-#include "compiler/mapped_parallel_computation_graph.h"
+#include "compiler/machine_mapping/machine_mapping.h"
+#include "compiler/machine_mapping/machine_view.h"
+#include "pcg/mapped_parallel_computation_graph/mapped_parallel_computation_graph.h"
 #include "compiler/machine_mapping/machine_view.dtg.h"
 #include "pcg/parallel_computation_graph/parallel_computation_graph_builder.h"
 #include <doctest/doctest.h>
+#include "test/utils/doctest/check_without_stringify.h"
 
 using namespace FlexFlow;
 
 TEST_SUITE(FF_TEST_SUITE) {
-  TEST_CASE("GraphOptimizeState::operator==") {
+  TEST_CASE("GraphOptimizeState operator==") {
     TensorShape input_shape = TensorShape{
         TensorDims{
             FFOrdered{
@@ -19,13 +22,6 @@ TEST_SUITE(FF_TEST_SUITE) {
         DataType::FLOAT,
     };
 
-    // `machine_mapping` is determined by the PCG and the device mapping
-    // algorithm, and `runtime` is determined by the PCG and the device mapping,
-    // so their values here do not matter.
-    MachineMapping empty_machine_mapping = MachineMapping{
-      std::unordered_map<parallel_layer_guid_t, MachineView>{},
-    };
-
     InitializerAttrs zero_init = InitializerAttrs{ZeroInitializerAttrs{}};
 
     auto create_pcg = [&]() -> ParallelComputationGraph {
@@ -56,26 +52,45 @@ TEST_SUITE(FF_TEST_SUITE) {
       return builder.pcg;
     };
 
+    auto create_machine_mapping_for_pcg = [](ParallelComputationGraph const &pcg) -> MachineMapping {
+      MachineSpaceCoordinate device = MachineSpaceCoordinate{
+        /*node_idx=*/0_n,
+        /*device_idx=*/0_n,
+        /*device_type=*/DeviceType::GPU,
+      };
+
+      MachineView machine_view = make_single_device_machine_view(device);
+
+      return MachineMapping{
+        generate_map(get_parallel_layers(pcg),
+                     [&](parallel_layer_guid_t) {
+                        return machine_view; 
+                     }),
+      };
+    };
+
     ParallelComputationGraph pcg1 = create_pcg();
+    MachineMapping machine_mapping_1 = create_machine_mapping_for_pcg(pcg1);
 
     SUBCASE("returns true if the PCGs are isomorphic") {
       ParallelComputationGraph pcg2 = create_pcg();
+      MachineMapping machine_mapping_2 = create_machine_mapping_for_pcg(pcg2);
 
       GraphOptimizeState state1 = GraphOptimizeState{
           GraphOptimizeResult{
-            mapped_pcg_from_pcg_and_mapping(pcg1, empty_machine_mapping),
+            mapped_pcg_from_pcg_and_mapping(pcg1, machine_mapping_1),
           },
           0,
       };
 
       GraphOptimizeState state2 = GraphOptimizeState{
           GraphOptimizeResult{
-            mapped_pcg_from_pcg_and_mapping(pcg2, empty_machine_mapping),
+            mapped_pcg_from_pcg_and_mapping(pcg2, machine_mapping_2),
           },
           0,
       };
 
-      CHECK(state1 == state2);
+      CHECK_WITHOUT_STRINGIFY(state1 == state2);
     }
 
     SUBCASE("returns false it the PCGs are not isomorphic") {
@@ -93,23 +108,25 @@ TEST_SUITE(FF_TEST_SUITE) {
                          /*bias_initializer=*/zero_init,
                          /*name=*/"dense0");
 
-      ParallelComputationGraph pcg_ = builder_.pcg;
+      ParallelComputationGraph other_pcg = builder_.pcg;
+
+      MachineMapping other_machine_mapping = create_machine_mapping_for_pcg(other_pcg);
 
       GraphOptimizeState state1 = GraphOptimizeState{
           GraphOptimizeResult{
-            mapped_pcg_from_pcg_and_mapping(pcg1, empty_machine_mapping),
+            mapped_pcg_from_pcg_and_mapping(pcg1, machine_mapping_1),
           },
           0,
       };
 
       GraphOptimizeState state_ = GraphOptimizeState{
           GraphOptimizeResult{
-            mapped_pcg_from_pcg_and_mapping(pcg_, empty_machine_mapping),
+            mapped_pcg_from_pcg_and_mapping(other_pcg, other_machine_mapping),
           },
           0,
       };
 
-      CHECK_FALSE(state1 == state_);
+      CHECK_FALSE_WITHOUT_STRINGIFY(state1 == state_);
     }
   }
 }

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

@@ -15,12 +15,12 @@
 #include "pcg/machine_specification_dimension.dtg.h"
 #include "compiler/machine_mapping/machine_view.dtg.h"
 #include "compiler/machine_mapping/machine_view.h"
-#include "pcg/machine_view_dimension.dtg.h"
+#include "compiler/machine_mapping/machine_view_dimension.dtg.h"
 #include "pcg/parallel_computation_graph/parallel_computation_graph.h"
 #include "pcg/parallel_computation_graph/parallel_computation_graph_builder.h"
 #include "pcg/parallel_computation_graph/parallel_layer_guid_t.dtg.h"
 #include "pcg/parallel_computation_graph/parallel_tensor_guid_t.h"
-#include "pcg/stride_t.dtg.h"
+#include "compiler/machine_mapping/stride_t.dtg.h"
 #include "substitutions/sub_parallel_computation_graph.dtg.h"
 #include "substitutions/sub_parallel_computation_graph.h"
 #include "utils/containers/get_only.h"

lib/compiler/test/src/compiler/task_graph_simulator/task_simulator.cc

@@ -67,8 +67,7 @@ class GenericTensorAccessorR {
              decltype(device_type) const &>
       tie() const;
 
-  // TODO(@lockshaw)(#pr): delete
-  // friend ::std::hash<GenericTensorAccessorR>;
+  friend ::std::hash<GenericTensorAccessorR>;
 };
 
 std::string format_as(GenericTensorAccessorR const &);
@@ -137,8 +136,7 @@ class GenericTensorAccessorW {
              decltype(device_type) const &>
       tie() const;
 
-  // TODO(@lockshaw)(#pr): delete
-  // friend ::std::hash<GenericTensorAccessorW>;
+  friend ::std::hash<GenericTensorAccessorW>;
 };
 
 std::string format_as(GenericTensorAccessorW const &);
@@ -228,19 +226,18 @@ real_type_t<DT> accessor_get_only_value(GenericTensorAccessorR const &acc) {
 
 } // namespace FlexFlow
 
-// TODO(@lockshaw)(#pr): delete
-// namespace std {
-//
-// template <>
-// struct hash<::FlexFlow::GenericTensorAccessorR> {
-//   size_t operator()(::FlexFlow::GenericTensorAccessorR const &) const;
-// };
-//
-// template <>
-// struct hash<::FlexFlow::GenericTensorAccessorW> {
-//   size_t operator()(::FlexFlow::GenericTensorAccessorW const &) const;
-// };
-//
-// }
-//
+namespace std {
+
+template <>
+struct hash<::FlexFlow::GenericTensorAccessorR> {
+  size_t operator()(::FlexFlow::GenericTensorAccessorR const &) const;
+};
+
+template <>
+struct hash<::FlexFlow::GenericTensorAccessorW> {
+  size_t operator()(::FlexFlow::GenericTensorAccessorW const &) const;
+};
+
+}
+
 #endif