Utils: Refactor and Test Updates

lib/compiler/src/machine_mapping.cc

@@ -6,7 +6,6 @@
 #include "pcg/machine_view.dtg.h"
 #include "pcg/machine_view.h"
 #include "pcg/parallel_computation_graph/parallel_computation_graph.h"
-#include "utils/containers.h"
 #include "utils/containers/are_disjoint.h"
 #include "utils/containers/as_vector.h"
 #include "utils/containers/contains_key.h"

lib/compiler/test/src/test_machine_mapping.cc

@@ -10,8 +10,8 @@ TEST_SUITE(FF_TEST_SUITE) {
 
   //     RC_ASSERT(comb.machine_views.size() ==
   //               m0.machine_views.size() + m1.machine_views.size());
-  //     RC_ASSERT(is_submap(comb.machine_views, m0.machine_views));
-  //     RC_ASSERT(is_submap(comb.machine_views, m1.machine_views));
+  //     RC_ASSERT(is_submapeq_of(comb.machine_views, m0.machine_views));
+  //     RC_ASSERT(is_submapeq_of(comb.machine_views, m1.machine_views));
   //   });
   // }
 

lib/kernels/test/src/test_dropout.cc

@@ -1,7 +1,6 @@
 #include "doctest/doctest.h"
 #include "kernels/dropout_kernels.h"
 #include "test_utils.h"
-#include "utils/containers.h"
 
 using namespace ::FlexFlow;
 TEST_SUITE(FF_TEST_SUITE) {

lib/runtime/src/model.cc

@@ -27,7 +27,6 @@
 #include "parallel_tensor_mapping.h"
 #include "task_spec/task_argument_accessor.h"
 #include "test_utils.h"
-#include "utils/containers.h"
 #include "utils/random_utils.h"
 #include <dirent.h>
 #include <queue>

lib/runtime/src/task_invocation_compilation.cc

@@ -1,5 +1,4 @@
 #include "task_invocation_compilation.h"
-#include "utils/containers.h"
 
 namespace FlexFlow {
 

lib/runtime/src/task_spec/task_invocation_compilation.cc

@@ -1,5 +1,4 @@
 #include "task_invocation_compilation.h"
-#include "utils/containers.h"
 
 namespace FlexFlow {
 

lib/runtime/src/task_spec/tensor_args_format.cc

@@ -1,5 +1,5 @@
 #include "tensor_args_format.h"
-#include "utils/containers.h"
+#include "utils/containers/flatmap.h"
 
 namespace FlexFlow {
 

lib/substitutions/src/substitutions/pcg_pattern_match.cc

@@ -2,7 +2,7 @@
 #include "substitutions/pcg_pattern.h"
 #include "substitutions/sub_parallel_computation_graph.h"
 #include "utils/bidict/algorithms/bidict_from_keys_and_values.h"
-#include "utils/bidict/algorithms/merge_bidicts.h"
+#include "utils/bidict/algorithms/merge_disjoint_bidicts.h"
 #include "utils/containers/map_values.h"
 #include "utils/containers/zip.h"
 
@@ -27,7 +27,7 @@ bidict<PatternNodeOutput, parallel_tensor_guid_t>
         bidict_from_keys_and_values(pattern_node_outputs,
                                     matched_layer_output_tensors);
 
-    result = merge_bidicts(result, mapping);
+    result = merge_disjoint_bidicts(result, mapping);
   }
 
   return result;

lib/substitutions/src/substitutions/substitution_internal/output_expr_to_result_sub_pcg_mapping.cc

@@ -2,7 +2,7 @@
 #include "substitutions/output_graph/output_graph_expr.h"
 #include "substitutions/sub_parallel_computation_graph.h"
 #include "utils/bidict/algorithms/bidict_from_keys_and_values.h"
-#include "utils/bidict/algorithms/merge_bidicts.h"
+#include "utils/bidict/algorithms/merge_disjoint_bidicts.h"
 
 namespace FlexFlow {
 
@@ -23,7 +23,7 @@ bidict<parallel_tensor_guid_t, OutputGraphExprNodeOutput>
         mapping_for_layer = bidict_from_keys_and_values(
             layer_outputs, output_graph_expr_outputs);
 
-    result = merge_bidicts(result, mapping_for_layer);
+    result = merge_disjoint_bidicts(result, mapping_for_layer);
   }
 
   return result;

lib/utils/include/utils/bidict/algorithms/merge_bidicts.h → lib/utils/include/utils/bidict/algorithms/merge_disjoint_bidicts.h

@@ -1,17 +1,25 @@
-#ifndef _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_BIDICT_ALGORITHMS_MERGE_BIDICTS_H
-#define _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_BIDICT_ALGORITHMS_MERGE_BIDICTS_H
+#ifndef _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_BIDICT_ALGORITHMS_MERGE_DISJOINT_BIDICTS_H
+#define _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_BIDICT_ALGORITHMS_MERGE_DISJOINT_BIDICTS_H
 
 #include "utils/bidict/algorithms/left_entries.h"
 #include "utils/bidict/algorithms/right_entries.h"
 #include "utils/bidict/bidict.h"
 #include "utils/containers/are_disjoint.h"
+#include "utils/exception.h"
 
 namespace FlexFlow {
 
 template <typename K, typename V>
-bidict<K, V> merge_bidicts(bidict<K, V> const &lhs, bidict<K, V> const &rhs) {
-  assert(are_disjoint(left_entries(lhs), left_entries(rhs)));
-  assert(are_disjoint(right_entries(lhs), right_entries(rhs)));
+bidict<K, V> merge_disjoint_bidicts(bidict<K, V> const &lhs,
+                                    bidict<K, V> const &rhs) {
+  if (!are_disjoint(left_entries(lhs), left_entries(rhs))) {
+    throw mk_runtime_error(
+        fmt::format("Left entries of {} and {} are non-disjoint", lhs, rhs));
+  }
+  if (!are_disjoint(right_entries(lhs), right_entries(rhs))) {
+    throw mk_runtime_error(
+        fmt::format("Right entries of {} and {} are non-disjoint", lhs, rhs));
+  }
 
   bidict<K, V> result;
   for (auto const &kv : lhs) {