Implement substitutions

lib/compiler/src/graph_utils.cc

@@ -10,7 +10,7 @@ SerialParallelDecomposition
 std::vector<MultiDiEdge>
     get_sorted_node_input_edges(ParallelComputationGraph const &pcg,
                                 Node const &n) {
-  std::unordered_map<size_t, std::unordered_set<MultiDiEdge>> incoming_edges =
+  std::unordered_map<NodePort, std::unordered_set<MultiDiEdge>> incoming_edges =
       get_incoming_edges_by_idx(pcg, n);
 
   std::vector<MultiDiEdge> result;
@@ -36,11 +36,11 @@ std::unordered_map<MultiDiEdge, ParallelTensorShape>
 
     auto outgoing_edges = get_outgoing_edges_by_idx(pcg, n);
 
-    for (std::size_t i = 0; i < output_tensor_shapes.size(); i++) {
-      if (contains_key(outgoing_edges, i)) {
-        for (MultiDiEdge const &e : outgoing_edges.at(i)) {
-          result.insert({e, output_tensor_shapes[i]});
-        }
+    int i = 0;
+
+    for (auto const &[node_port, edges] : outgoing_edges) {
+      for (MultiDiEdge const &e : edges) {
+        result.insert({e, output_tensor_shapes[i++]});
       }
     }
   }

lib/compiler/src/unity_algorithm.cc

@@ -1,6 +1,6 @@
 #include "compiler/unity_algorithm.h"
 #include "graph_utils.h"
-#include "substitutions_implementation.h"
+#include "substitutions/substitution.h"
 #include "utils/deduplicated_priority_queue.h"
 
 namespace FlexFlow {
@@ -14,7 +14,9 @@ std::unordered_set<Substitution>
 
 std::unordered_set<ParallelComputationGraph>
     apply_substitution(ParallelComputationGraph const &pcg,
-                       Substitution const &);
+                       Substitution const &) {
+  NOT_IMPLEMENTED();
+}
 
 Strategy
     graph_optimize(ComputationGraph &cg,

lib/compiler/test/test_generator.h

@@ -3,6 +3,7 @@
 
 #include "compiler/machine_mapping.h"
 #include "compiler/sub_parallel_computation_graph.h"
+#include "pcg/computation_graph.h"
 #include "rapidcheck.h"
 
 using namespace FlexFlow;
@@ -38,7 +39,7 @@ ParallelComputationGraph
 }
 
 rc::Gen<int> small_integer_generator() {
-  return gen::inRange(1, 4);
+  return rc::gen::inRange(1, 4);
 }
 
 namespace rc {
@@ -51,23 +52,26 @@ Gen<MultiDiGraph> serialParallelMultiDiGraph() {
 template <>
 struct Arbitrary<ComputationGraph> {
   static Gen<ComputationGraph> arbitrary() {
-    return gen::map(serialParallelMultiDiGraph, test_computataion_graph);
+    return gen::map(gen::cast<MultiDiGraphView>(serialParallelMultiDiGraph()),
+                    test_computataion_graph);
   }
 };
 
 template <>
 struct Arbitrary<ParallelComputationGraph> {
   static Gen<ParallelComputationGraph> arbitrary() {
-    return gen::map(serialParallelMultiDiGraph,
-                    test_parallel_computataion_graph);
+    return gen::map(gen::cast<MultiDiGraphView>(serialParallelMultiDiGraph()),
+                    test_parallel_computation_graph);
   }
 };
 
 template <>
 struct Arbitrary<variant<Serial, Node>> {
   static Gen<variant<Serial, Node>> arbitrary() {
     return gen::mapcat(gen::arbitrary<bool>(), [](bool is_node) {
-      return is_node ? gen::arbitrary<Node>() : gen::arbitrary<Serial>();
+      return is_node
+                 ? gen::cast<variant<Serial, Node>>(gen::arbitrary<Node>())
+                 : gen::cast<variant<Serial, Node>>(gen::arbitrary<Serial>());
     });
   }
 };
@@ -76,7 +80,10 @@ template <>
 struct Arbitrary<variant<Parallel, Node>> {
   static Gen<variant<Parallel, Node>> arbitrary() {
     return gen::mapcat(gen::arbitrary<bool>(), [](bool is_node) {
-      return is_node ? gen::arbitrary<Node>() : gen::arbitrary<Parallel>();
+      return is_node
+                 ? gen::cast<variant<Parallel, Node>>(gen::arbitrary<Node>())
+                 : gen::cast<variant<Parallel, Node>>(
+                       gen::arbitrary<Parallel>());
     });
   }
 };

lib/substitutions/CMakeLists.txt

@@ -10,6 +10,7 @@ ff_add_library(
   DEPS
     utils
     op-attrs
+    pcg
 )
 
 add_subdirectory(ffi)