Finish kwarg fixes for ComputationGraph

Author: lockshaw

lib/op-attrs/include/op-attrs/ops/embedding.h

@@ -5,6 +5,7 @@
 #include "op-attrs/ops/embedding_attrs.dtg.h"
 #include "op-attrs/parallel_tensor_shape.h"
 #include "op-attrs/tensor_shape.h"
+#include "op-attrs/tensor_slot_name.dtg.h"
 #include "utils/record_formatter.h"
 #include <tl/expected.hpp>
 
@@ -28,7 +29,7 @@ tl::expected<ParallelTensorShape, std::string>
  * see
  * https://github.com/pytorch/pytorch/blob/1eba9b3aa3c43f86f4a2c807ac8e12c4a7767340/torch/nn/modules/sparse.py#L180-L182
  */
-std::vector<InitializerAttrs> get_initializers(
+std::unordered_map<TensorSlotName, InitializerAttrs> get_initializers(
     EmbeddingAttrs const &,
     std::optional<InitializerAttrs> const &initializer_attrs = std::nullopt);
 

lib/op-attrs/src/op-attrs/ops/embedding.cc

@@ -130,7 +130,7 @@ tl::expected<ParallelTensorShape, std::string>
       unpar, sum_degree, discard_copy_degree, shard_degrees);
 }
 
-std::vector<InitializerAttrs> get_initializers(
+std::unordered_map<TensorSlotName, InitializerAttrs> get_initializers(
     EmbeddingAttrs const &,
     std::optional<InitializerAttrs> const &maybe_initializer_attrs) {
   InitializerAttrs default_initializer_attrs = InitializerAttrs{
@@ -141,7 +141,12 @@ std::vector<InitializerAttrs> get_initializers(
       },
   };
 
-  return {maybe_initializer_attrs.value_or(default_initializer_attrs)};
+  return {
+    {
+      TensorSlotName::WEIGHT,
+      maybe_initializer_attrs.value_or(default_initializer_attrs),
+    },
+  };
 }
 
 } // namespace FlexFlow

lib/pcg/include/pcg/computation_graph/layer_added_result.dtg.toml

@@ -9,10 +9,13 @@ features = [
 includes = [
   "pcg/layer_guid_t.dtg.h",
   "pcg/tensor_guid_t.dtg.h",
-  "utils/fmt/vector.h",
   "op-attrs/tensor_slot_name.dtg.h",
 ]
 
+src_includes = [
+  "utils/fmt/unordered_map.h",
+]
+
 [[fields]]
 name = "layer"
 type = "::FlexFlow::layer_guid_t"

lib/pcg/include/pcg/computation_graph_builder.h

@@ -255,11 +255,11 @@ struct ComputationGraphBuilder {
   TensorShape get_shape(tensor_guid_t const &) const;
 
 private:
-  std::vector<tensor_guid_t> add_layer(
+  std::unordered_map<TensorSlotName, tensor_guid_t> add_layer(
       LayerAttrs const &layer,
-      std::vector<tensor_guid_t> const &inputs,
-      std::vector<InitializerAttrs> const &weights,
-      std::optional<std::vector<CreateGrad>> const &outputs = std::nullopt);
+      std::unordered_map<TensorSlotName, tensor_guid_t> const &inputs,
+      std::unordered_map<TensorSlotName, InitializerAttrs> const &weights,
+      std::optional<std::unordered_map<TensorSlotName, CreateGrad>> const &outputs = std::nullopt);
 
   tensor_guid_t
       broadcast(tensor_guid_t const &, TensorDims const &, std::string const &);

lib/pcg/include/pcg/file_format/v1/graphs/v1_dataflow_graph.h

@@ -10,6 +10,10 @@ features = [
   "fmt",
 ]
 
+template_params = [
+  "SlotName",
+]
+
 includes = [
   "utils/nonnegative_int/nonnegative_int.h",
 ]
@@ -19,13 +23,13 @@ name = "srcNode"
 type = "::FlexFlow::nonnegative_int"
 
 [[fields]]
-name = "srcIdx"
-type = "::FlexFlow::nonnegative_int"
+name = "srcSlot"
+type = "SlotName"
 
 [[fields]]
 name = "dstNode"
 type = "::FlexFlow::nonnegative_int"
 
 [[fields]]
-name = "dstIdx"
-type = "::FlexFlow::nonnegative_int"
+name = "dstSlot"
+type = "SlotName"

lib/pcg/include/pcg/file_format/v1/graphs/v1_graph_edge.dtg.toml

@@ -1,5 +1,5 @@
 namespace = "FlexFlow"
-name = "V1DataflowGraph"
+name = "V1KwargDataflowGraph"
 type = "struct"
 features = [
   "eq",
@@ -10,6 +10,10 @@ features = [
   "fmt",
 ]
 
+template_params = [
+  "SlotName",
+]
+
 includes = [
   "<vector>",
   "<unordered_set>",
@@ -30,4 +34,4 @@ type = "std::vector<::FlexFlow::nonnegative_int>"
 
 [[fields]]
 name = "edges"
-type = "std::unordered_set<::FlexFlow::V1GraphEdge>"
+type = "std::unordered_set<::FlexFlow::V1GraphEdge<SlotName>>"

…mat/v1/graphs/v1_dataflow_graph.dtg.toml …/graphs/v1_kwarg_dataflow_graph.dtg.tomllib/pcg/include/pcg/file_format/v1/graphs/v1_dataflow_graph.dtg.toml renamed to lib/pcg/include/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.dtg.toml lib/pcg/include/pcg/file_format/v1/graphs/v1_dataflow_graph.dtg.toml renamed to lib/pcg/include/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.dtg.toml

@@ -0,0 +1,42 @@
+#ifndef _FLEXFLOW_LIB_PCG_INCLUDE_PCG_FILE_FORMAT_V1_GRAPHS_V1_KWARG_DATAFLOW_GRAPH_H
+#define _FLEXFLOW_LIB_PCG_INCLUDE_PCG_FILE_FORMAT_V1_GRAPHS_V1_KWARG_DATAFLOW_GRAPH_H
+
+#include "pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.dtg.h"
+#include "utils/graph/kwarg_dataflow_graph/kwarg_dataflow_graph_view.h"
+#include "utils/bidict/algorithms/bidict_from_enumerating.h"
+#include "utils/containers/enumerate.h"
+#include "utils/containers/sorted.h"
+#include "utils/containers/values.h"
+#include "utils/graph/node/algorithms.h"
+#include "utils/integer_conversions.h"
+#include "utils/graph/kwarg_dataflow_graph/algorithms/get_all_kwarg_dataflow_edges.h"
+
+namespace FlexFlow {
+
+template <typename SlotName>
+V1KwargDataflowGraph<SlotName> to_v1(KwargDataflowGraphView<SlotName> const &g) {
+  bidict<nonnegative_int, Node> node_enumeration_bidict =
+      bidict_from_enumerating(get_nodes(g));
+  std::unordered_map<Node, nonnegative_int> node_enumeration =
+      node_enumeration_bidict.reversed().as_unordered_map();
+  return to_v1(g, node_enumeration);
+}
+
+template <typename SlotName>
+V1KwargDataflowGraph<SlotName> to_v1(KwargDataflowGraphView<SlotName> const &g,
+                      std::unordered_map<Node, nonnegative_int> const &nodes) {
+  std::unordered_set<V1GraphEdge<SlotName>> edges;
+  for (KwargDataflowEdge<SlotName> const &e : get_all_kwarg_dataflow_edges(g)) {
+    edges.insert(V1GraphEdge{
+        nodes.at(e.src.node), e.src.slot_name, nodes.at(e.dst.node), e.dst.slot_name});
+  }
+
+  return V1KwargDataflowGraph<SlotName>{
+      sorted(values(nodes)),
+      edges,
+  };
+}
+
+} // namespace FlexFlow
+
+#endif

lib/pcg/include/pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.h

@@ -1,5 +1,5 @@
 namespace = "FlexFlow"
-name = "V1LabelledDataflowGraph"
+name = "V1LabelledKwargDataflowGraph"
 type = "struct"
 features = [
   "eq",
@@ -13,11 +13,12 @@ features = [
 template_params = [
   "NodeLabel",
   "OutputLabel",
+  "SlotName",
 ]
 
 includes = [
   "<unordered_map>",
-  "pcg/file_format/v1/graphs/v1_dataflow_graph.dtg.h",
+  "pcg/file_format/v1/graphs/v1_kwarg_dataflow_graph.dtg.h",
   "pcg/file_format/v1/graphs/v1_graph_output.dtg.h",
   "utils/nonnegative_int/nonnegative_int.h",
 ]
@@ -35,8 +36,8 @@ type = "std::unordered_map<::FlexFlow::nonnegative_int, NodeLabel>"
 
 [[fields]]
 name = "output_labels"
-type = "std::unordered_map<::FlexFlow::nonnegative_int, std::vector<OutputLabel>>"
+type = "std::unordered_map<::FlexFlow::nonnegative_int, std::unordered_map<SlotName, OutputLabel>>"
 
 [[fields]]
 name = "graph"
-type = "::FlexFlow::V1DataflowGraph"
+type = "::FlexFlow::V1KwargDataflowGraph<SlotName>"