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

Author: Victor Li

lib/substitutions/include/substitutions/unity_substitution_set.h

@@ -1,16 +1,19 @@
 #ifndef _FLEXFLOW_LIB_SUBSTITUTIONS_INCLUDE_SUBSTITUTIONS_UNITY_SUBSTITUTION_SET_H
 #define _FLEXFLOW_LIB_SUBSTITUTIONS_INCLUDE_SUBSTITUTIONS_UNITY_SUBSTITUTION_SET_H
 
+#include "pcg/machine_specification.dtg.h"
 #include "substitutions/substitution.dtg.h"
+#include "utils/fmt/vector.h"
 
 namespace FlexFlow {
 
+std::vector<Substitution>
+    get_substitution_set(MachineSpecification const &resources);
+
 Substitution create_combine_inception(int num_convs, int num_dims, int degree);
 Substitution create_combine_concat(int num_inputs, int num_dims, int degree);
-Substitution create_replicate_linear_combine(int num_dims,
-                                             int degree,
-                                             Activation activation,
-                                             bool use_bias);
+Substitution
+    create_replicate_linear_combine(int num_dims, int degree, bool use_bias);
 Substitution create_partition_linear_combine(int num_dims,
                                              int degree,
                                              Activation activation,
@@ -27,6 +30,7 @@ Substitution create_partition_concat_combine(int num_inputs,
 Substitution create_partition_softmax_combine(ff_dim_t softmax_dim,
                                               ff_dim_t partition_dim,
                                               int degree);
+Substitution create_fuse_linear_activation(Activation activation);
 
 } // namespace FlexFlow
 

lib/substitutions/src/substitutions/unity_substitution_set.cc

@@ -7,6 +7,24 @@
 
 namespace FlexFlow {
 
+std::vector<Substitution>
+    get_substitution_set(MachineSpecification const &resources) {
+  std::vector<Substitution> substitutions;
+  for (int num_dims = 1; num_dims <= MAX_TENSOR_DIM; num_dims++) {
+    for (int degree = 1; degree <= resources.num_nodes; degree *= 2) {
+      substitutions.push_back(
+          create_replicate_linear_combine(num_dims, degree, true));
+      substitutions.push_back(
+          create_replicate_linear_combine(num_dims, degree, false));
+    }
+  }
+  substitutions.push_back(create_fuse_linear_activation(Activation::RELU));
+  substitutions.push_back(create_fuse_linear_activation(Activation::SIGMOID));
+  substitutions.push_back(create_fuse_linear_activation(Activation::TANH));
+  substitutions.push_back(create_fuse_linear_activation(Activation::GELU));
+  return substitutions;
+}
+
 Substitution create_combine_inception(int num_convs, int num_dims, int degree) {
   NOT_IMPLEMENTED();
 }
@@ -122,11 +140,6 @@ Substitution create_partition_linear_combine(int num_dims,
   NOT_IMPLEMENTED();
 }
 
-Substitution create_partition_linear_combine(int degree,
-                                             Activation activation) {
-  NOT_IMPLEMENTED();
-}
-
 Substitution create_partition_conv2d_combine(int num_dims, int degree) {
   NOT_IMPLEMENTED();
 }