flexflow · lockshaw · Oct 21, 2024 · Oct 21, 2024 · Oct 21, 2024 · Oct 23, 2024
diff --git a/.proj.toml b/.proj.toml
@@ -6,6 +6,21 @@ cuda_launch_cmd = [
   "nixGL",
   "--",
 ]
+layout_ignore_paths = [
+  "lib/runtime",
+  "lib/kernels",
+  "lib/utils/test/common",
+  "lib/compiler/ffi",
+  "lib/op-attrs/ffi",
+  "lib/pcg/ffi",
+  "lib/substitutions/ffi",
+  "lib/utils/ffi",
+  "lib/ffi",
+  "lib/utils/include/utils/graph/docs",
+  "lib/compiler/test/src/internal",
+  "lib/local-execution/test/src/internal",
+  "bin/protobuf-to-json",
+]
 
 [targets.utils]
 type = "lib"
@@ -70,6 +85,13 @@ has-cpu-only-benchmarks = false
 has-cuda-tests = true
 has-cuda-benchmarks = false
 
+[targets.local-pcg-execution]
+type = "lib"
+has-cpu-only-tests = true
+has-cpu-only-benchmarks = false
+has-cuda-tests = false
+has-cuda-benchmarks = false
+
 [targets.models]
 type = "lib"
 has-cpu-only-tests = true

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -102,7 +102,6 @@ include(spdlog)
 include(doctestlib) # named doctestlib to avoid a name collision with doctest.cmake in rapidcheck
 include(gbenchmark)
 include(libassert)
-include(visit_struct)
 include(CTest)
 include(fmt)
 include(legion)

diff --git a/bin/export-model-arch/include/export-model-arch/json_sp_model_export.dtg.toml b/bin/export-model-arch/include/export-model-arch/json_sp_model_export.dtg.toml
@@ -0,0 +1,26 @@
+namespace = "FlexFlow"
+name = "JsonSPModelExport"
+type = "struct"
+features = [
+  "eq",
+  "hash",
+  "fmt",
+  "json", 
+]
+
+includes = [
+  "pcg/file_format/v1/v1_computation_graph.dtg.h",
+  "pcg/file_format/v1/v1_binary_sp_decomposition/v1_binary_sp_decomposition.dtg.h",
+]
+
+src_includes = [
+  "pcg/file_format/v1/v1_binary_sp_decomposition/json.h",
+]
+
+[[fields]]
+name = "sp_decomposition"
+type = "::FlexFlow::V1BinarySPDecomposition"
+
+[[fields]]
+name = "computation_graph"
+type = "::FlexFlow::V1ComputationGraph"
diff --git a/bin/export-model-arch/include/export_model_arch/json_sp_model_export.struct.toml b/bin/export-model-arch/include/export_model_arch/json_sp_model_export.struct.toml
diff --git a/...xport-model-arch/src/export_model_arch.cc → ...-model-arch/src/export-model-arch/main.cc b/...xport-model-arch/src/export_model_arch.cc → ...-model-arch/src/export-model-arch/main.cc
diff --git a/bin/substitution-to-dot/CMakeLists.txt b/bin/substitution-to-dot/CMakeLists.txt
@@ -2,7 +2,9 @@ ff_add_executable(
   NAME
     substitution-to-dot
   SRC_PATTERNS
-    *.cc
+    src/*.cc
+  PRIVATE_INCLUDE
+    include/
   DEPS
     substitution-generator
 )
diff --git a/...ubstitution-to-dot/substitution_to_dot.cc → ...on-to-dot/src/substitution-to-dot/main.cc b/...ubstitution-to-dot/substitution_to_dot.cc → ...on-to-dot/src/substitution-to-dot/main.cc
diff --git a/cmake/visit_struct.cmake b/cmake/visit_struct.cmake
diff --git a/flake.lock b/flake.lock
diff --git a/flake.nix b/flake.nix
@@ -18,7 +18,7 @@
     flake-utils.url = "github:numtide/flake-utils";
 
     proj-repo = {
-      url = "git+https://git.sr.ht/~lockshaw/proj";
+      url = "git+https://git.sr.ht/~lockshaw/proj?ref=emulated-fs";
       inputs.nixpkgs.follows = "nixpkgs";
       inputs.flake-utils.follows = "flake-utils";
     };
@@ -93,9 +93,6 @@
                                 -DFF_USE_EXTERNAL_EXPECTED=ON \
                                 -DFF_USE_EXTERNAL_GBENCHMARK=ON \
                                 -DFF_USE_EXTERNAL_LIBASSERT=ON \
-                                -DFF_USE_EXTERNAL_RANGEV3=ON \
-                                -DFF_USE_EXTERNAL_BOOST_PREPROCESSOR=ON \
-                                -DFF_USE_EXTERNAL_TYPE_INDEX=ON"
           '';
 
           buildInputs = builtins.concatLists [

diff --git a/lib/CMakeLists.txt b/lib/CMakeLists.txt
@@ -4,6 +4,7 @@ add_subdirectory(runtime)
 add_subdirectory(op-attrs)
 add_subdirectory(kernels)
 add_subdirectory(local-execution)
+add_subdirectory(local-pcg-execution)
 add_subdirectory(task-spec)
 add_subdirectory(utils)
 add_subdirectory(ffi)

diff --git a/lib/compiler/include/compiler/allowed_machine_views.h b/lib/compiler/include/compiler/allowed_machine_views.h
@@ -1,18 +1,18 @@
 #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"
+#include "op-attrs/operator_task_space.dtg.h"
+#include "pcg/machine_compute_specification.dtg.h"
+#include "compiler/machine_mapping/machine_view.dtg.h"
 
 namespace FlexFlow {
 
 bool is_valid_machine_view(MachineView const &mv,
                            OperatorTaskSpace const &task,
-                           MachineSpecification const &ms);
+                           MachineComputeSpecification const &ms);
 
 std::unordered_set<MachineView>
-    get_allowed_machine_views(MachineSpecification const &machine_spec,
+    get_allowed_machine_views(MachineComputeSpecification const &machine_spec,
                               OperatorTaskSpace const &task,
                               DeviceType device_type);
 

diff --git a/lib/compiler/include/compiler/compiler.h b/lib/compiler/include/compiler/compiler.h
diff --git a/lib/compiler/include/compiler/cost_estimator/communication_edge.h b/lib/compiler/include/compiler/cost_estimator/communication_edge.h
@@ -0,0 +1,50 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_COMMUNICATION_EDGE_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_COMMUNICATION_EDGE_H
+
+#include "pcg/machine_space_coordinate.dtg.h"
+
+namespace FlexFlow {
+
+struct CommunicationEdge {
+  CommunicationEdge() = delete;
+
+  CommunicationEdge(MachineSpaceCoordinate const &src,
+                    MachineSpaceCoordinate const &dst);
+
+  bool operator==(CommunicationEdge const &) const;
+  bool operator!=(CommunicationEdge const &) const;
+
+  bool operator<(CommunicationEdge const &) const;
+  bool operator>(CommunicationEdge const &) const;
+  bool operator<=(CommunicationEdge const &) const;
+  bool operator>=(CommunicationEdge const &) const;
+
+  MachineSpaceCoordinate const &get_src() const;
+  MachineSpaceCoordinate const &get_dst() const;
+private:
+  MachineSpaceCoordinate src;
+  MachineSpaceCoordinate dst;
+private:
+  std::tuple<
+    decltype(src) const &,
+    decltype(dst) const &
+  > tie() const;
+
+  friend struct ::std::hash<CommunicationEdge>;
+};
+
+std::string format_as(CommunicationEdge const &);
+std::ostream &operator<<(std::ostream &, CommunicationEdge const &);
+
+} // namespace FlexFlow
+
+namespace std {
+
+template <>
+struct hash<::FlexFlow::CommunicationEdge> {
+  size_t operator()(::FlexFlow::CommunicationEdge const &) const;
+};
+
+}
+
+#endif
diff --git a/lib/compiler/include/compiler/cost_estimator/cost_estimator.h b/lib/compiler/include/compiler/cost_estimator/cost_estimator.h
@@ -6,7 +6,7 @@
 #include "compiler/cost_estimator/tensor_set_movement.dtg.h"
 #include "op-attrs/parallel_tensor_shape.dtg.h"
 #include "op-attrs/pcg_operator_attrs.dtg.h"
-#include "pcg/machine_view.dtg.h"
+#include "compiler/machine_mapping/machine_view.dtg.h"
 #include <vector>
 
 namespace FlexFlow {

diff --git a/lib/compiler/include/compiler/cost_estimator/network_cost_model.h b/lib/compiler/include/compiler/cost_estimator/network_cost_model.h
@@ -0,0 +1,14 @@
+#ifndef _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_NETWORK_COST_MODEL_H
+#define _FLEXFLOW_LIB_COMPILER_INCLUDE_COMPILER_COST_ESTIMATOR_NETWORK_COST_MODEL_H
+
+#include "compiler/cost_estimator/tensor_set_movement.dtg.h"
+#include "pcg/machine_specification.dtg.h"
+
+namespace FlexFlow {
+
+float estimate_communication_cost(MachineSpecification const &,
+                                  TensorSetMovement const &);
+
+} // namespace FlexFlow
+
+#endif
diff --git a/lib/compiler/include/compiler/cost_estimator/op_cost_estimate_key.dtg.toml b/lib/compiler/include/compiler/cost_estimator/op_cost_estimate_key.dtg.toml
@@ -0,0 +1,46 @@
+namespace = "FlexFlow"
+name = "OpCostEstimateKey"
+type = "struct"
+features = [
+  "eq",
+  "ord",
+  "fmt",
+  "hash",
+]
+
+includes = [
+  "op-attrs/pcg_operator_attrs.dtg.h",
+  "op-attrs/parallel_tensor_shape.dtg.h",
+  "<vector>",
+  "compiler/machine_mapping/machine_view.dtg.h",
+  "pcg/optimizer_attrs.dtg.h",
+]
+
+src_includes = [
+  "utils/hash/vector.h",
+  "utils/fmt/vector.h",
+]
+
+[[fields]]
+name = "op_attrs"
+type = "::FlexFlow::PCGOperatorAttrs"
+
+[[fields]]
+name = "input_shapes"
+type = "std::vector<::FlexFlow::ParallelTensorShape>"
+
+[[fields]]
+name = "weight_shapes"
+type = "std::vector<::FlexFlow::ParallelTensorShape>"
+
+[[fields]]
+name = "output_shapes"
+type = "std::vector<::FlexFlow::ParallelTensorShape>"
+
+[[fields]]
+name = "optimizer_attrs"
+type = "::FlexFlow::OptimizerAttrs"
+
+[[fields]]
+name = "machine_view"
+type = "::FlexFlow::MachineView"