Pass pcg tests

Author: lockshaw

lib/compiler/src/compiler/task_graph_simulator/simulate_task_graph_execution.cc

@@ -8,7 +8,7 @@
 #include "utils/containers/set_of.h"
 #include "utils/containers/sorted.h"
 #include "utils/exception.h"
-#include "utils/graph/digraph/algorithms.h"
+#include "utils/graph/digraph/algorithms/get_initial_nodes.h"
 #include "utils/graph/digraph/algorithms/get_predecessors.h"
 #include "utils/graph/digraph/algorithms/get_successors.h"
 #include "utils/graph/digraph/algorithms/is_acyclic.h"

lib/op-attrs/include/op-attrs/get_operator_task_space.h

@@ -4,12 +4,13 @@
 #include "op-attrs/computation_graph_op_attrs.dtg.h"
 #include "op-attrs/operator_task_space.dtg.h"
 #include "op-attrs/parallel_tensor_dim_degrees.dtg.h"
+#include "op-attrs/tensor_slot_name.dtg.h"
 
 namespace FlexFlow {
 
 OperatorTaskSpace
   get_operator_task_space(ComputationGraphOpAttrs const &attrs,
-                          std::vector<ParallelTensorDimDegrees> const &inputs_degrees);
+                          std::unordered_map<TensorSlotName, ParallelTensorDimDegrees> const &inputs_degrees);
 
 } // namespace FlexFlow
 

lib/op-attrs/src/op-attrs/get_operator_task_space.cc

@@ -1,5 +1,7 @@
 #include "op-attrs/get_operator_task_space.h"
 #include "utils/containers/get_only.h"
+#include "utils/containers/require_only_key.h"
+#include "utils/containers/require_two_keys.h"
 #include "utils/overload.h"
 #include <libassert/assert.hpp>
 #include "op-attrs/ops/element_unary.h"
@@ -13,37 +15,37 @@ namespace FlexFlow {
 
 OperatorTaskSpace
   get_operator_task_space(ComputationGraphOpAttrs const &attrs,
-                          std::vector<ParallelTensorDimDegrees> const &inputs_degrees) {
+                          std::unordered_map<TensorSlotName, ParallelTensorDimDegrees> const &inputs_degrees) {
   return attrs.visit<
     OperatorTaskSpace
   >(overload {
     [&](ElementUnaryAttrs const &attrs) {
-      ASSERT(inputs_degrees.size() == 1);
+      ParallelTensorDimDegrees input = require_only_key(inputs_degrees, TensorSlotName::INPUT);
 
-      return get_operator_task_space(attrs, get_only(inputs_degrees));
+      return get_operator_task_space(attrs, input);
     },
     [&](ElementBinaryAttrs const &attrs) {
-      ASSERT(inputs_degrees.size() == 2);
+      auto [lhs, rhs] = require_two_keys(inputs_degrees, TensorSlotName::LHS_INPUT, TensorSlotName::RHS_INPUT); 
 
       return get_operator_task_space(
         /*attrs=*/attrs,
-        /*lhs_input_degrees=*/inputs_degrees.at(0),
-        /*rhs_input_degrees=*/inputs_degrees.at(1));
+        /*lhs_input_degrees=*/lhs,
+        /*rhs_input_degrees=*/rhs);
     },
     [&](LinearAttrs const &attrs) {
-      ASSERT(inputs_degrees.size() == 1);
+      ParallelTensorDimDegrees input = require_only_key(inputs_degrees, TensorSlotName::INPUT);
 
-      return get_operator_task_space(attrs, get_only(inputs_degrees));
+      return get_operator_task_space(attrs, input);
     },
     [&](InputAttrs const &attrs) {
       ASSERT(inputs_degrees.size() == 0);
 
       return get_operator_task_space(attrs);
     },
     [&](TransposeAttrs const &attrs) {
-      ASSERT(inputs_degrees.size() == 1);
+      ParallelTensorDimDegrees input = require_only_key(inputs_degrees, TensorSlotName::INPUT);
 
-      return get_operator_task_space(attrs, get_only(inputs_degrees));
+      return get_operator_task_space(attrs, input);
     },
     [&](WeightAttrs const &attrs) {
       ASSERT(inputs_degrees.size() == 0);

lib/op-attrs/src/op-attrs/shape_inference.cc

@@ -170,7 +170,7 @@ std::unordered_map<TensorSlotName, TensorShape>
         };
       },
       [&](InputAttrs const &attrs) -> std::unordered_map<TensorSlotName, TensorShape> {
-        TensorShape input = require_only_key(input_shapes, TensorSlotName::INPUT);
+        ASSERT(input_shapes.size() == 0);
 
         return {
           {
@@ -240,7 +240,7 @@ std::unordered_map<TensorSlotName, TensorShape>
         };
       },
       [&](WeightAttrs const &attrs) -> std::unordered_map<TensorSlotName, TensorShape> {
-        TensorShape input = require_only_key(input_shapes, TensorSlotName::INPUT);
+        ASSERT(input_shapes.size() == 0);
 
         return {
           {
@@ -711,7 +711,7 @@ std::unordered_map<TensorSlotName, ParallelTensorShape>
         return {};
       },
       [&](WeightAttrs const &attrs) -> std::unordered_map<TensorSlotName, ParallelTensorShape> {
-        require_only_key(input_shapes, TensorSlotName::INPUT);
+        ASSERT(input_shapes.size() == 0);
 
         return {};
       },

lib/pcg/include/pcg/parallel_computation_graph/parallel_computation_graph_builder.h

@@ -143,10 +143,10 @@ struct ParallelComputationGraphBuilder {
                                  std::string const &name);
 
 private:
-  std::vector<parallel_tensor_guid_t>
+  std::unordered_map<TensorSlotName, parallel_tensor_guid_t>
       add_layer(ParallelLayerAttrs const &layer,
-                std::vector<parallel_tensor_guid_t> const &inputs,
-                std::vector<InitializerAttrs> const &weight_initializers);
+                std::unordered_map<TensorSlotName, parallel_tensor_guid_t> const &inputs,
+                std::unordered_map<TensorSlotName, InitializerAttrs> const &weight_initializers);
 
   parallel_tensor_guid_t
       add_weight(ParallelTensorShape const &weight_tensor_shape,

lib/pcg/include/pcg/parallel_computation_graph/parallel_computation_graph_edge.dtg.toml

@@ -9,9 +9,10 @@ features = [
 ]
 
 includes = [
-  "utils/graph/dataflow_graph/dataflow_edge.dtg.h",
+  "utils/graph/kwarg_dataflow_graph/kwarg_dataflow_edge.dtg.h",
+  "op-attrs/tensor_slot_name.dtg.h",
 ]
 
 [[fields]]
 name = "raw_edge"
-type = "::FlexFlow::DataflowEdge"
+type = "::FlexFlow::KwargDataflowEdge<::FlexFlow::TensorSlotName>"

lib/pcg/include/pcg/parallel_computation_graph/parallel_computation_graph_edge.h

@@ -11,8 +11,8 @@ parallel_tensor_guid_t
     get_parallel_tensor(ParallelComputationGraphEdge const &);
 parallel_layer_guid_t get_src_layer(ParallelComputationGraphEdge const &);
 parallel_layer_guid_t get_dst_layer(ParallelComputationGraphEdge const &);
-nonnegative_int get_src_layer_output_idx(ParallelComputationGraphEdge const &);
-nonnegative_int get_dst_layer_input_idx(ParallelComputationGraphEdge const &);
+TensorSlotName get_src_layer_output_slot_name(ParallelComputationGraphEdge const &);
+TensorSlotName get_dst_layer_input_slot_name(ParallelComputationGraphEdge const &);
 
 } // namespace FlexFlow
 

lib/pcg/include/pcg/parallel_computation_graph/parallel_layer_added_result.dtg.toml

@@ -4,7 +4,6 @@ type = "struct"
 
 features = [
   "eq",
-  "ord",
   "fmt",
 ]
 

lib/pcg/include/pcg/parallel_computation_graph/parallel_tensor_guid_t.dtg.toml

@@ -9,7 +9,7 @@ features = [
 ]
 
 includes = [
-  "utils/graph/dataflow_graph/dataflow_output.dtg.h",
+  "utils/graph/kwarg_dataflow_graph/kwarg_dataflow_output.dtg.h",
   "op-attrs/tensor_slot_name.dtg.h",
 ]
 

lib/pcg/src/pcg/computation_graph.cc

@@ -205,7 +205,7 @@ static std::unordered_map<TensorSlotName, tensor_guid_t>
   std::unordered_map<TensorSlotName, IncomingTensorRole> incoming_slot_roles =
       get_incoming_tensor_roles(attrs);
 
-  assert(incoming_tensors.size() == incoming_slot_roles.size());
+  ASSERT(incoming_tensors.size() == incoming_slot_roles.size());
 
   std::unordered_set<TensorSlotName> slots_with_desired_role = 
     keys(filter_values(incoming_slot_roles,