Utils: Refactor and Test Updates

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

@@ -1,5 +1,4 @@
 #include "compiler/machine_mapping/machine_mapping.h"
-#include "utils/containers.h"
 #include "utils/containers/are_disjoint.h"
 #include "utils/containers/keys.h"
 #include "utils/containers/merge_maps.h"

lib/kernels/include/kernels/accessor.h

@@ -7,7 +7,6 @@
 #include "op-attrs/datatype.h"
 #include "utils/exception.h"
 #include "utils/required.h"
-#include "utils/variant.h"
 
 namespace FlexFlow {
 

lib/kernels/include/kernels/initializer_kernels.h

@@ -4,7 +4,6 @@
 #include "accessor.h"
 #include "kernels/cpu.h"
 #include "op-attrs/datatype_value.dtg.h"
-#include "utils/variant.h"
 
 namespace FlexFlow {
 

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/pcg/include/pcg/optimizer_attrs.h

@@ -3,7 +3,7 @@
 
 #include "pcg/optimizers/adam_optimizer_attrs.h"
 #include "pcg/optimizers/sgd_optimizer_attrs.h"
-#include "utils/variant.h"
+#include <variant>
 
 namespace FlexFlow {
 

lib/pcg/src/pcg/operator_task_space.cc

@@ -5,6 +5,7 @@
 #include "utils/containers/range.h"
 #include "utils/containers/transform.h"
 #include "utils/containers/unordered_set_of.h"
+#include "utils/fmt/unordered_set.h"
 
 namespace FlexFlow {
 
@@ -25,7 +26,7 @@ std::unordered_set<TaskSpaceCoordinate>
 
 TaskSpaceCoordinate
     get_task_space_maximum_coordinate(OperatorTaskSpace const &task) {
-  return maximum(get_task_space_coordinates(task)).value();
+  return maximum(get_task_space_coordinates(task));
 }
 
 size_t num_dims(OperatorTaskSpace const &task) {

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) {