Merge substitution-builder (flexflow#1575)

* Start on pcg builder

* Add tests and some implementation for pcg builder

* Add pcg tests, make dtgen constructors explicit to fix bug

* Add remainder of PCG tests

* Fix build issues in local-execution

* Format

* Address Reyna comments, add topological_order function for PCG

* Pre multidigraph refactor

* Removing visitable from sp code

* Add open dataflow graph, start to replace pcg dataflow graph

* Start refactoring substitutions

* Add utility functions to support pattern matching

* Pre-refactor inputs

* Fix proj url

* Get back to substitutions, now with unordered graph inputs

* Get substitutions building

* substitutions-tests now builds

* Fix bug in filter, pass some initial substitution tests

* Add tests for fmt::to_string, fix some substitutions bugs

* Pass initial unit tests for find_pattern_matches

* Start on unit tests for pcg pattern

* Pass initial test for find_pattern_matches

* Fix small build issue in tests

* Format

* Sync tests in CI with tests in proj

* Fix minor build errors in kernels and local-execution

* Format

* Remove outdated code

* More outdated code removal

* More cleanup, add test for sp decomposition

* Pull apart containers.h

* More sp testing and fixes

* Break up graph algorithms.h

* Pre- full SP algo commit

* Add initial implementation and tests for cbc decomposition and inverse line graph

* Pass test for get_inverse_line_graph

* Add new multidigraph

* Fix get_inverse_line_graph to return a MultiDiGraph instead of a DiGraph

* Add tests for parallel and series reduction finding

* Add really rough implementation of valdez sp decomposition

* Fix local-execution build

* Add implementations and tests for applying series/parallel reductions

* Format

* Clean up sp decomposition interface and tests

* Format

* Add comments for top-level substitutions functions, add proj doxygen support

* Start sketching out substitutions code

* Fix build errors

* Add ability to permute node ids

* Cleanup and start to test new substitutions code

* Add test case for evaluate_substitution_output

* Add naive isomorphism detection code

* Add graph inputs to open dataflow graph isomorphism

* Add input permutation to evaluate_substitution_output

* Fix permute_node_ids

* Add test for permute_input_ids

* Migrate over to mutable implementation of apply_substitution

* Add fast isomorphism checking and an initial implementation of full substitution logic

* Pass initial full substitutions test

* Cleanup old isomorphism checking code

* Fix post-merge bugs

* Fix broken pcg builder test

* Format

* Reorganize code and remove some outdated code pre-code-review

* Format

* Restarting work on this after working on export-model-arch

* Adding in some a simple function to get the currently available substritutions

* nonnegative_int additions, code cleanup, etc.

* A bunch more moving over to nonnegative_int

* Even more nonnegative_int updating

* Fix build

* Fix failing tests

* Format

* Format

---------

Co-authored-by: Colin Unger <[email protected]>
Co-authored-by: Victor Li <[email protected]>

Author: 3 people

bin/export-model-arch/src/export_model_arch.cc

@@ -13,6 +13,7 @@
 #include "utils/cli/cli_parse.h"
 #include "utils/cli/cli_parse_result.h"
 #include "utils/cli/cli_spec.h"
+#include "utils/graph/open_dataflow_graph/algorithms/as_dot.h"
 #include "utils/graph/series_parallel/binary_sp_decomposition_tree/right_associative_binary_sp_tree_from_nary.h"
 #include "utils/graph/series_parallel/get_series_parallel_decomposition.h"
 
@@ -21,11 +22,11 @@ using namespace ::FlexFlow;
 ComputationGraph get_single_operator_computation_graph() {
   ComputationGraphBuilder b;
 
-  size_t batch_size = 8;
-  size_t in_channels = 16;
-  size_t out_channels = 12;
+  nonnegative_int batch_size = 8_n;
+  nonnegative_int in_channels = 16_n;
+  nonnegative_int out_channels = 12_n;
   TensorShape input_shape = TensorShape{
-      TensorDims{FFOrdered<size_t>{
+      TensorDims{FFOrdered<nonnegative_int>{
           batch_size,
           in_channels,
           out_channels,
@@ -69,7 +70,7 @@ tl::expected<ComputationGraph, std::string>
   } else if (model_name == "bert") {
     return get_bert_computation_graph(get_default_bert_config());
   } else if (model_name == "split_test") {
-    int batch_size = 8;
+    nonnegative_int batch_size = 8_n;
     return get_split_test_computation_graph(batch_size);
   } else if (model_name == "single_operator") {
     return get_single_operator_computation_graph();
@@ -100,10 +101,10 @@ tl::expected<JsonSPModelExport, std::string>
     result.value();
   });
 
-  std::pair<V1ComputationGraph, bidict<int, layer_guid_t>> v1_result =
-      to_v1_including_node_numbering(computation_graph);
+  std::pair<V1ComputationGraph, bidict<nonnegative_int, layer_guid_t>>
+      v1_result = to_v1_including_node_numbering(computation_graph);
   V1ComputationGraph v1_cg = v1_result.first;
-  bidict<int, layer_guid_t> layer_numbering = v1_result.second;
+  bidict<nonnegative_int, layer_guid_t> layer_numbering = v1_result.second;
   V1BinarySPDecomposition v1_sp_decomposition =
       to_v1(sp_decomposition, layer_numbering);
 

cmake/flexflow-utils.cmake

@@ -20,6 +20,7 @@ function(define_ff_vars target)
     MAX_TENSOR_DIM=${FF_MAX_DIM}
     MAX_NUM_TASK_REGIONS=${FF_MAX_NUM_TASK_REGIONS}
     MAX_NUM_TASK_ARGUMENTS=${FF_MAX_NUM_TASK_ARGUMENTS}
+    # _FORTIFY_SOURCE=0
     )
 
   if (FF_GPU_BACKEND STREQUAL "cuda")
@@ -39,7 +40,18 @@ function(ff_set_cxx_properties target)
       CXX_EXTENSIONS NO
   )
   target_compile_options(${target}
-    PRIVATE $<$<COMPILE_LANGUAGE:CXX>:> "-ffile-prefix-map=${CMAKE_SOURCE_DIR}=." # add C++ compile flags here
+    PUBLIC 
+    $<$<COMPILE_LANGUAGE:CXX>:> 
+    "-ffile-prefix-map=${CMAKE_SOURCE_DIR}=." 
+    "-fsanitize=undefined" 
+    "-fno-sanitize-recover=all"
+    # add C++ compile flags here
+  )
+  target_link_options(${target}
+    PUBLIC 
+    $<$<COMPILE_LANGUAGE:CXX>:> 
+    "-fsanitize=undefined" 
+    "-fno-sanitize-recover=all"
   )
 endfunction()
 

flake.nix

@@ -38,9 +38,15 @@
       };
       lib = pkgs.lib;
 
-      mkShell = pkgs.mkShell.override {
+      mkShell = attrs: pkgs.mkShell.override {
         stdenv = pkgs.cudaPackages.backendStdenv;
-      };
+      } (attrs // {
+        hardeningDisable = ["all"]; # disable nixpkgs default compiler arguments, otherwise ubsan doesn't catch 
+                                    # signed overflows due to the signedoverflow hardening setting. 
+                                    # for more details, see the following (long-running) nixpkgs github issues: 
+                                    # - https://github.com/NixOS/nixpkgs/issues/18995
+                                    # - https://github.com/NixOS/nixpkgs/issues/60919
+      });
 
       proj = proj-repo.packages.${system}.proj;
     in 
@@ -121,6 +127,8 @@
 
         gpu-ci = mkShell {
           inputsFrom = [ ci ];
+          hardeningDisable = [ "all" ];
+
           buildInputs = builtins.concatLists [
             (with nixGL.packages.${system}; [
               nixGLDefault
@@ -135,6 +143,8 @@
             "${proj-repo.packages.${system}.proj-nvim}"
           ];
 
+          hardeningDisable = [ "all" ];
+
           buildInputs = builtins.concatLists [
             (with pkgs; [
               clang-tools

lib/compiler/include/compiler/series_parallel/computation_graph/computation_graph_binary_sp_decomposition.h

@@ -36,8 +36,9 @@ bool is_right_associative(ComputationGraphBinarySPDecomposition const &);
 std::unordered_multiset<layer_guid_t>
     get_layers(ComputationGraphBinarySPDecomposition const &);
 
-V1BinarySPDecomposition to_v1(ComputationGraphBinarySPDecomposition const &,
-                              bidict<int, layer_guid_t> const &layer_numbering);
+V1BinarySPDecomposition
+    to_v1(ComputationGraphBinarySPDecomposition const &,
+          bidict<nonnegative_int, layer_guid_t> const &layer_numbering);
 
 } // namespace FlexFlow
 

lib/compiler/src/compiler/allowed_machine_views.cc

@@ -11,12 +11,15 @@
 #include "utils/containers/map_from_keys_and_values.h"
 #include "utils/containers/product.h"
 #include "utils/containers/range.h"
-#include "utils/containers/replicate.h"
+#include "utils/containers/repeat_element.h"
 #include "utils/containers/sorted.h"
 #include "utils/containers/transform.h"
 #include "utils/containers/unordered_multiset_of.h"
 #include "utils/containers/unordered_set_of.h"
 #include "utils/containers/zip.h"
+#include "utils/nonnegative_int/ceildiv.h"
+#include "utils/nonnegative_int/nonnegative_range.h"
+#include "utils/nonnegative_int/num_elements.h"
 #include "utils/overload.h"
 
 namespace FlexFlow {
@@ -47,24 +50,29 @@ static std::unordered_set<MachineView>
                                 OperatorTaskSpace const &task,
                                 DeviceType const &device_type) {
 
-  auto get_max_stride_upper_bound = [](std::vector<int> const &tensor_dims,
-                                       int total_devices) -> int {
-    int min_num_devices_with_full_stride_volume = product(transform(
-        tensor_dims, [](int const &num_devices) { return num_devices - 1; }));
-    return std::ceil(total_devices / min_num_devices_with_full_stride_volume);
+  auto get_max_stride_upper_bound =
+      [](std::vector<nonnegative_int> const &tensor_dims,
+         nonnegative_int total_devices) -> nonnegative_int {
+    nonnegative_int min_num_devices_with_full_stride_volume =
+        product(transform(tensor_dims, [](nonnegative_int num_devices) {
+          return nonnegative_int{num_devices.unwrap_nonnegative() - 1};
+        }));
+    return ceildiv(total_devices, min_num_devices_with_full_stride_volume);
   };
 
-  auto candidate_strides = [&](std::vector<int> const &tensor_dims,
-                               int total_devices)
+  auto candidate_strides = [&](std::vector<nonnegative_int> const &tensor_dims,
+                               nonnegative_int total_devices)
       -> std::unordered_multiset<MultiDimensionalStride> {
-    int max_stride_upper_bound =
+    nonnegative_int max_stride_upper_bound =
         get_max_stride_upper_bound(tensor_dims, total_devices);
 
     std::vector<stride_t> single_stride_range =
-        transform(range(1, max_stride_upper_bound + 1),
-                  [](int stride) { return stride_t{stride}; });
+        transform(nonnegative_range(1_n, max_stride_upper_bound + 1_n),
+                  [](nonnegative_int stride) { return stride_t{stride}; });
     std::unordered_multiset<std::vector<stride_t>> raw_stride_vectors =
-        cartesian_product(replicate(tensor_dims.size(), single_stride_range));
+        cartesian_product(
+            repeat_element(/*num_times=*/num_elements(tensor_dims),
+                           /*element=*/single_stride_range));
     std::unordered_multiset<MultiDimensionalStride> strides =
         transform(raw_stride_vectors, [](auto const &stride_vec) {
           return MultiDimensionalStride{stride_vec};
@@ -75,8 +83,9 @@ static std::unordered_set<MachineView>
   auto candidate_starts = [](MachineSpecification const &ms,
                              DeviceType const &device_type) {
     std::unordered_set<MachineSpaceCoordinate> result;
-    for (int node_idx : range(ms.num_nodes)) {
-      for (int device_idx : range(get_num_devices_per_node(ms, device_type))) {
+    for (nonnegative_int node_idx : nonnegative_range(ms.num_nodes)) {
+      for (nonnegative_int device_idx :
+           nonnegative_range(get_num_devices_per_node(ms, device_type))) {
         result.insert(
             MachineSpaceCoordinate{node_idx, device_idx, device_type});
       }
@@ -91,8 +100,8 @@ static std::unordered_set<MachineView>
     return get_all_permutations_with_repetition(options, num_dims(task));
   };
 
-  std::vector<int> tensor_dims = task.degrees;
-  int total_devices = get_num_devices(machine_spec, device_type);
+  std::vector<nonnegative_int> tensor_dims = task.degrees;
+  nonnegative_int total_devices = get_num_devices(machine_spec, device_type);
 
   std::unordered_set<MachineView> machine_views;
 

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

@@ -11,17 +11,19 @@ std::unordered_set<std::pair<MachineSpecification, MachineSpecification>>
   for (int i = 1; i < resource.num_nodes; i *= 2) {
     MachineSpecification sub_resource1 = resource;
     MachineSpecification sub_resource2 = resource;
-    sub_resource1.num_nodes = i;
-    sub_resource2.num_nodes = resource.num_nodes - i;
+    sub_resource1.num_nodes = nonnegative_int{i};
+    sub_resource2.num_nodes =
+        nonnegative_int{resource.num_nodes.unwrap_nonnegative() - i};
     result.insert(std::make_pair(sub_resource1, sub_resource2));
     result.insert(std::make_pair(sub_resource2, sub_resource1));
   }
 
   for (int i = 1; i < resource.num_gpus_per_node; i *= 2) {
     MachineSpecification sub_resource1 = resource;
     MachineSpecification sub_resource2 = resource;
-    sub_resource1.num_gpus_per_node = i;
-    sub_resource2.num_gpus_per_node = resource.num_gpus_per_node - i;
+    sub_resource1.num_gpus_per_node = nonnegative_int{i};
+    sub_resource2.num_gpus_per_node =
+        nonnegative_int{resource.num_gpus_per_node.unwrap_nonnegative() - i};
     result.insert(std::make_pair(sub_resource1, sub_resource2));
     result.insert(std::make_pair(sub_resource2, sub_resource1));
   }

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

@@ -1,20 +1,14 @@
 #include "compiler/machine_mapping/machine_mapping.h"
-#include "pcg/machine_specification.h"
-#include "pcg/machine_view.h"
-#include "pcg/operator_task_space.dtg.h"
-#include "pcg/operator_task_space.h"
-#include "pcg/parallel_computation_graph/parallel_computation_graph.h"
 #include "utils/containers/are_disjoint.h"
-#include "utils/containers/get_one_of.h"
 #include "utils/containers/keys.h"
-#include "utils/containers/map_values.h"
 #include "utils/containers/merge_maps.h"
 
 namespace FlexFlow {
 
 MachineMapping combine_disjoint_mappings(MachineMapping const &m1,
                                          MachineMapping const &m2) {
-  return MachineMapping{merge_maps(m1.machine_views, m2.machine_views)};
+  return MachineMapping{
+      merge_disjoint_maps(m1.machine_views, m2.machine_views)};
 }
 
 bool nodes_are_disjoint(MachineMapping const &m1, MachineMapping const &m2) {

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

@@ -10,8 +10,8 @@ ParallelLayerGuidObliviousMachineMapping binary_combine_mappings(
     ParallelLayerGuidObliviousMachineMapping const &lhs,
     ParallelLayerGuidObliviousMachineMapping const &rhs) {
   return ParallelLayerGuidObliviousMachineMapping{
-      merge_maps(map_keys(lhs.raw_mapping, nest_inside_left_child),
-                 map_keys(rhs.raw_mapping, nest_inside_right_child)),
+      merge_disjoint_maps(map_keys(lhs.raw_mapping, nest_inside_left_child),
+                          map_keys(rhs.raw_mapping, nest_inside_right_child)),
   };
 }
 

lib/compiler/src/compiler/series_parallel/computation_graph/computation_graph_binary_sp_decomposition.cc

@@ -164,7 +164,7 @@ std::unordered_multiset<layer_guid_t>
 
 V1BinarySPDecomposition
     to_v1(ComputationGraphBinarySPDecomposition const &tree,
-          bidict<int, layer_guid_t> const &layer_numbering) {
+          bidict<nonnegative_int, layer_guid_t> const &layer_numbering) {
   return tree.visit<V1BinarySPDecomposition>(
       overload{[&](ComputationGraphBinarySeriesSplit const &series) {
                  return V1BinarySPDecomposition{

lib/compiler/test/src/allowed_machine_views.cc

@@ -15,39 +15,39 @@ TEST_SUITE(FF_TEST_SUITE) {
 
     SUBCASE("1 degree of parallelism") {
       MachineSpecification ms = MachineSpecification{
-          /*num_nodes=*/1,
-          /*num_cpus_per_node=*/5,
-          /*num_gpus_per_node=*/5,
+          /*num_nodes=*/1_n,
+          /*num_cpus_per_node=*/5_n,
+          /*num_gpus_per_node=*/5_n,
           /*inter_node_bandwidth=*/0,
           /*intra_node_bandwidth=*/0,
       };
 
-      OperatorTaskSpace task = OperatorTaskSpace{{3}};
+      OperatorTaskSpace task = OperatorTaskSpace{{3_n}};
 
       std::unordered_set<MachineView> correct = {
           MachineView{
               MachineSpaceCoordinate{
-                  /*node_idx=*/0, /*device_idx=*/0, DeviceType::GPU},
-              {MachineViewDimension{stride_t{1},
+                  /*node_idx=*/0_n, /*device_idx=*/0_n, DeviceType::GPU},
+              {MachineViewDimension{stride_t{1_n},
                                     MachineSpecificationDimension::INTRA_NODE}},
           },
 
           MachineView{
               MachineSpaceCoordinate{
-                  /*node_idx=*/0, /*device_idx=*/1, DeviceType::GPU},
-              {MachineViewDimension{stride_t{1},
+                  /*node_idx=*/0_n, /*device_idx=*/1_n, DeviceType::GPU},
+              {MachineViewDimension{stride_t{1_n},
                                     MachineSpecificationDimension::INTRA_NODE}},
           },
           MachineView{
               MachineSpaceCoordinate{
-                  /*node_idx=*/0, /*device_idx=*/2, DeviceType::GPU},
-              {MachineViewDimension{stride_t{1},
+                  /*node_idx=*/0_n, /*device_idx=*/2_n, DeviceType::GPU},
+              {MachineViewDimension{stride_t{1_n},
                                     MachineSpecificationDimension::INTRA_NODE}},
           },
           MachineView{
               MachineSpaceCoordinate{
-                  /*node_idx=*/0, /*device_idx=*/0, DeviceType::GPU},
-              {MachineViewDimension{stride_t{2},
+                  /*node_idx=*/0_n, /*device_idx=*/0_n, DeviceType::GPU},
+              {MachineViewDimension{stride_t{2_n},
                                     MachineSpecificationDimension::INTRA_NODE}},
           },
       };
@@ -61,18 +61,18 @@ TEST_SUITE(FF_TEST_SUITE) {
     SUBCASE("2 degrees of parallelism") {
 
       MachineSpecification ms = MachineSpecification{
-          /*num_nodes=*/3,
-          /*num_cpus_per_node=*/3,
-          /*num_gpus_per_node=*/3,
+          /*num_nodes=*/3_n,
+          /*num_cpus_per_node=*/3_n,
+          /*num_gpus_per_node=*/3_n,
           /*inter_node_bandwidth=*/0,
           /*intra_node_bandwidth=*/0,
       };
-      OperatorTaskSpace task = OperatorTaskSpace{{2, 3}};
+      OperatorTaskSpace task = OperatorTaskSpace{{2_n, 3_n}};
 
-      auto make_2d_view = [&](int start_node_idx,
-                              int start_device_idx,
-                              int stride1,
-                              int stride2,
+      auto make_2d_view = [&](nonnegative_int start_node_idx,
+                              nonnegative_int start_device_idx,
+                              nonnegative_int stride1,
+                              nonnegative_int stride2,
                               MachineSpecificationDimension m1,
                               MachineSpecificationDimension m2) {
         return MachineView{
@@ -86,13 +86,19 @@ TEST_SUITE(FF_TEST_SUITE) {
       auto intra = MachineSpecificationDimension::INTRA_NODE;
       auto inter = MachineSpecificationDimension::INTER_NODE;
       std::unordered_set<MachineView> correct = {
-          make_2d_view(0, 0, /*stride1=*/1, /*stride2=*/1, inter, intra),
-          make_2d_view(1, 0, /*stride1=*/1, /*stride2=*/1, inter, intra),
-          make_2d_view(0, 0, /*stride1=*/2, /*stride2=*/1, inter, intra),
-
-          make_2d_view(0, 0, /*stride1=*/1, /*stride2=*/1, intra, inter),
-          make_2d_view(0, 1, /*stride1=*/1, /*stride2=*/1, intra, inter),
-          make_2d_view(0, 0, /*stride1=*/2, /*stride2=*/1, intra, inter),
+          make_2d_view(
+              0_n, 0_n, /*stride1=*/1_n, /*stride2=*/1_n, inter, intra),
+          make_2d_view(
+              1_n, 0_n, /*stride1=*/1_n, /*stride2=*/1_n, inter, intra),
+          make_2d_view(
+              0_n, 0_n, /*stride1=*/2_n, /*stride2=*/1_n, inter, intra),
+
+          make_2d_view(
+              0_n, 0_n, /*stride1=*/1_n, /*stride2=*/1_n, intra, inter),
+          make_2d_view(
+              0_n, 1_n, /*stride1=*/1_n, /*stride2=*/1_n, intra, inter),
+          make_2d_view(
+              0_n, 0_n, /*stride1=*/2_n, /*stride2=*/1_n, intra, inter),
       };
 
       std::unordered_set<MachineView> result =