nonnegative_int additions, code cleanup, etc.

Author: Victor Li

lib/compiler/src/compiler/allowed_machine_views.cc

@@ -60,7 +60,8 @@ static std::unordered_set<MachineView>
         transform(range(1, max_stride_upper_bound + 1),
                   [](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(replicate(nonnegative_int{tensor_dims.size()},
+                                    single_stride_range));
     std::unordered_multiset<MultiDimensionalStride> strides =
         transform(raw_stride_vectors, [](auto const &stride_vec) {
           return MultiDimensionalStride{stride_vec};

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

@@ -7,7 +7,8 @@ namespace FlexFlow {
 
 MachineMapping combine_disjoint_mappings(MachineMapping const &s1,
                                          MachineMapping const &s2) {
-  return MachineMapping{merge_maps(s1.machine_views, s2.machine_views)};
+  return MachineMapping{
+      merge_disjoint_maps(s1.machine_views, s2.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/substitutions/include/substitutions/operator_pattern/operator_attribute_constraint.h

@@ -9,8 +9,8 @@ OperatorAttributeConstraint op_type_equals_constraint(OperatorType);
 
 OperatorAttributeConstraint op_attr_key_equals(OperatorAttributeKey,
                                                OperatorAttributeValue const &);
-OperatorAttributeConstraint op_attr_key_divisible_by(OperatorAttributeKey,
-                                                     int denominator);
+OperatorAttributeConstraint
+    op_attr_key_divisible_by(OperatorAttributeKey, nonnegative_int denominator);
 OperatorAttributeConstraint
     make_equals_constraint(OperatorAttributeExpr const &,
                            OperatorAttributeValue const &);

lib/substitutions/include/substitutions/operator_pattern/operator_attribute_value.variant.toml

@@ -21,6 +21,7 @@ includes = [
   "op-attrs/tensor_shape.dtg.h",
   "op-attrs/datatype.dtg.h",
   "<cstddef>",
+  "utils/nonnegative_int/nonnegative_int.h",
 ]
 
 src_includes = [
@@ -74,3 +75,6 @@ type = "::FlexFlow::TensorDims"
 
 [[values]]
 type = "::FlexFlow::DataType"
+
+[[values]]
+type = "::FlexFlow::nonnegative_int"

lib/substitutions/include/substitutions/substitution_builder.h

@@ -26,7 +26,7 @@ struct SubstitutionBuilder {
   std::vector<OutputGraphExprValue>
       add_output_graph_node(OutputOperatorAttrsAssignment const &node_expr,
                             std::vector<OutputGraphExprValue> const &inputs,
-                            int num_outputs);
+                            nonnegative_int num_outputs);
 
   PatternNode pattern_node_named(std::string const &) const;
   PatternInput pattern_input_named(std::string const &) const;

lib/substitutions/include/substitutions/tensor_pattern/tensor_attribute_pattern.h

@@ -2,11 +2,13 @@
 #define _FLEXFLOW_LIB_SUBSTITUTIONS_INCLUDE_SUBSTITUTIONS_TENSOR_PATTERN_TENSOR_ATTRIBUTE_PATTERN_H
 
 #include "substitutions/tensor_pattern/tensor_attribute_pattern.dtg.h"
+#include "utils/nonnegative_int/nonnegative_int.h"
 
 namespace FlexFlow {
 
 TensorAttributePattern tensor_attribute_pattern_match_all();
-TensorAttributePattern tensor_attr_pattern_require_num_dims(int num_dims);
+TensorAttributePattern
+    tensor_attr_pattern_require_num_dims(nonnegative_int num_dims);
 
 } // namespace FlexFlow
 

lib/substitutions/include/substitutions/tensor_pattern/tensor_attribute_value.variant.toml

@@ -12,10 +12,11 @@ includes = [
   "<vector>",
   "utils/hash/vector.h",
   "utils/fmt/vector.h",
+  "utils/nonnegative_int/nonnegative_int.h",
 ]
 
 [[values]]
-type = "size_t"
+type = "::FlexFlow::nonnegative_int"
 
 [[values]]
-type = "std::vector<size_t>"
+type = "std::vector<::FlexFlow::nonnegative_int>"

lib/substitutions/src/substitutions/apply_substitution/apply_substitution.cc

@@ -46,7 +46,8 @@ SubParallelComputationGraph
         std::unordered_map<parallel_layer_guid_t, ParallelLayerAttrs>
             post_node_data_from_sub = output_graph_data.node_data;
 
-        return merge_maps(post_node_data_from_orig, post_node_data_from_sub);
+        return merge_disjoint_maps(post_node_data_from_orig,
+                                   post_node_data_from_sub);
       }();
 
   std::unordered_set<SubParallelComputationGraphEdge> post_edges = [&] {
@@ -147,7 +148,8 @@ SubParallelComputationGraph
 
         std::unordered_map<open_parallel_tensor_guid_t, ParallelTensorAttrs>
             post_value_data_from_sub = output_graph_data.value_data;
-        return merge_maps(post_value_data_from_orig, post_value_data_from_sub);
+        return merge_disjoint_maps(post_value_data_from_orig,
+                                   post_value_data_from_sub);
       }();
 
   SubParallelComputationGraphData post_data = SubParallelComputationGraphData{

lib/substitutions/src/substitutions/operator_pattern/operator_attribute_constraint.cc

@@ -20,8 +20,9 @@ OperatorAttributeConstraint
   };
 }
 
-OperatorAttributeConstraint op_attr_key_divisible_by(OperatorAttributeKey key,
-                                                     int denominator) {
+OperatorAttributeConstraint
+    op_attr_key_divisible_by(OperatorAttributeKey key,
+                             nonnegative_int denominator) {
   return OperatorAttributeConstraint{
       ConstraintType::DIVISIBLE_BY,
       OperatorAttributeExpr{key},