Merge branch 'master' into dlrm-pcg

flexflow · Feb 10, 2025 · ac950a7 · ac950a7
2 parents 55d56c3 + ccef91d
commit ac950a7
Show file tree

Hide file tree

Showing 17 changed files with 313 additions and 14 deletions.
diff --git a/lib/op-attrs/include/op-attrs/pcg_operator_attrs.h b/lib/op-attrs/include/op-attrs/pcg_operator_attrs.h
@@ -1,6 +1,7 @@
 #ifndef _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_PCG_OPERATOR_ATTRS_H
 #define _FLEXFLOW_LIB_OP_ATTRS_INCLUDE_OP_ATTRS_PCG_OPERATOR_ATTRS_H
 
+#include "op-attrs/computation_graph_op_attrs.dtg.h"
 #include "op-attrs/computation_graph_op_attrs.h"
 #include "op-attrs/pcg_operator_attrs.dtg.h"
 

diff --git a/lib/pcg/include/pcg/parallel_computation_graph/parallel_layer_attrs.h b/lib/pcg/include/pcg/parallel_computation_graph/parallel_layer_attrs.h
@@ -1,12 +1,15 @@
 #ifndef _FLEXFLOW_LIB_PCG_INCLUDE_PCG_PARALLEL_LAYER_ATTRS_H
 #define _FLEXFLOW_LIB_PCG_INCLUDE_PCG_PARALLEL_LAYER_ATTRS_H
 
+#include "pcg/layer_attrs.dtg.h"
 #include "pcg/parallel_computation_graph/parallel_layer_attrs.dtg.h"
 
 namespace FlexFlow {
 
 OperatorType get_op_type(ParallelLayerAttrs const &);
 
+ParallelLayerAttrs parallel_layer_attrs_from_layer_attrs(LayerAttrs const &);
+
 } // namespace FlexFlow
 
 #endif
diff --git a/lib/pcg/include/pcg/parallel_tensor_attrs.h b/lib/pcg/include/pcg/parallel_tensor_attrs.h
@@ -9,6 +9,9 @@ namespace FlexFlow {
 
 TensorAttrs get_piece_attrs(ParallelTensorAttrs const &);
 
+ParallelTensorAttrs
+    parallel_tensor_attrs_from_tensor_attrs(TensorAttrs const &);
+
 } // namespace FlexFlow
 
 #endif
diff --git a/lib/pcg/include/pcg/pcg_from_computation_graph.h b/lib/pcg/include/pcg/pcg_from_computation_graph.h
@@ -0,0 +1,13 @@
+#ifndef _FLEXFLOW_PCG_INCLUDE_PCG_CG_TO_PCG_H
+#define _FLEXFLOW_PCG_INCLUDE_PCG_CG_TO_PCG_H
+
+#include "pcg/computation_graph.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h"
+
+namespace FlexFlow {
+
+ParallelComputationGraph pcg_from_computation_graph(ComputationGraph const &cg);
+
+} // namespace FlexFlow
+
+#endif
diff --git a/lib/pcg/include/pcg/tensor_attrs.struct.toml b/lib/pcg/include/pcg/tensor_attrs.struct.toml
@@ -26,14 +26,14 @@ src_includes = [
 name = "shape"
 type = "::FlexFlow::TensorShape"
 
-[[fields]]
-name = "initializer"
-type = "std::optional<::FlexFlow::InitializerAttrs>"
-
 [[fields]]
 name = "sync_type"
 type = "std::optional<::FlexFlow::ParamSync>"
 
+[[fields]]
+name = "initializer"
+type = "std::optional<::FlexFlow::InitializerAttrs>"
+
 [[fields]]
 name = "create_gradients"
 type = "::FlexFlow::CreateGrad"
diff --git a/lib/pcg/src/pcg/computation_graph_builder.cc b/lib/pcg/src/pcg/computation_graph_builder.cc
@@ -37,7 +37,7 @@ namespace FlexFlow {
 static TensorAttrs make_weight_attrs(
     TensorShape const &shape,
     std::optional<InitializerAttrs> const &initializer_attrs) {
-  return TensorAttrs{shape, initializer_attrs, std::nullopt, CreateGrad::YES};
+  return TensorAttrs{shape, std::nullopt, initializer_attrs, CreateGrad::YES};
 }
 
 static TensorAttrs make_output_attrs(TensorShape const &shape) {
@@ -86,7 +86,7 @@ tensor_guid_t ComputationGraphBuilder::create_weight(
     std::optional<ParamSync> param_sync,
     std::optional<std::string> const &maybe_name) {
   TensorAttrs tensor_attrs =
-      TensorAttrs{shape, initializer, param_sync, create_grad};
+      TensorAttrs{shape, param_sync, initializer, create_grad};
 
   return this->create_weight(tensor_attrs, maybe_name);
 }

diff --git a/lib/pcg/src/pcg/parallel_computation_graph/parallel_layer_attrs.cc b/lib/pcg/src/pcg/parallel_computation_graph/parallel_layer_attrs.cc
@@ -1,10 +1,18 @@
 #include "pcg/parallel_computation_graph/parallel_layer_attrs.h"
 #include "op-attrs/pcg_operator_attrs.h"
+#include "pcg/parallel_computation_graph/parallel_layer_attrs.dtg.h"
 
 namespace FlexFlow {
 
 OperatorType get_op_type(ParallelLayerAttrs const &a) {
   return get_op_type(a.op_attrs);
 }
 
+ParallelLayerAttrs
+    parallel_layer_attrs_from_layer_attrs(LayerAttrs const &layer_attrs) {
+  return ParallelLayerAttrs{
+      pcg_op_attrs_from_compgraph_op_attrs(layer_attrs.attrs),
+      layer_attrs.name};
+}
+
 } // namespace FlexFlow
diff --git a/lib/pcg/src/pcg/parallel_tensor_attrs.cc b/lib/pcg/src/pcg/parallel_tensor_attrs.cc
@@ -1,12 +1,22 @@
 #include "pcg/parallel_tensor_attrs.h"
+#include "op-attrs/parallel_tensor_shape.h"
+#include "pcg/parallel_computation_graph/parallel_tensor_attrs.dtg.h"
 
 namespace FlexFlow {
 
 TensorAttrs get_piece_attrs(ParallelTensorAttrs const &parallel_attrs) {
   return TensorAttrs{get_piece_shape(parallel_attrs.shape),
-                     parallel_attrs.initializer,
                      parallel_attrs.sync_type,
+                     parallel_attrs.initializer,
                      parallel_attrs.create_gradients};
 }
 
+ParallelTensorAttrs
+    parallel_tensor_attrs_from_tensor_attrs(TensorAttrs const &tensor_attrs) {
+  return ParallelTensorAttrs{lift_to_parallel(tensor_attrs.shape),
+                             tensor_attrs.sync_type,
+                             tensor_attrs.initializer,
+                             tensor_attrs.create_gradients};
+}
+
 } // namespace FlexFlow
diff --git a/lib/pcg/src/pcg/pcg_from_computation_graph.cc b/lib/pcg/src/pcg/pcg_from_computation_graph.cc
@@ -0,0 +1,36 @@
+#include "pcg/pcg_from_computation_graph.h"
+#include "op-attrs/pcg_operator_attrs.h"
+#include "pcg/computation_graph.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_layer_attrs.dtg.h"
+#include "pcg/parallel_computation_graph/parallel_layer_attrs.h"
+#include "pcg/parallel_tensor_attrs.h"
+#include "pcg/tensor_attrs.dtg.h"
+#include "utils/graph/instances/unordered_set_labelled_open_dataflow_graph.h"
+#include "utils/graph/labelled_dataflow_graph/algorithms/rewrite_node_labels.h"
+#include "utils/graph/labelled_dataflow_graph/algorithms/rewrite_value_labels.h"
+#include "utils/graph/labelled_dataflow_graph/labelled_dataflow_graph.h"
+#include "utils/graph/labelled_dataflow_graph/labelled_dataflow_graph_view.h"
+#include "utils/graph/open_dataflow_graph/open_dataflow_value.dtg.h"
+
+namespace FlexFlow {
+
+ParallelComputationGraph
+    pcg_from_computation_graph(ComputationGraph const &cg) {
+  auto layer_map = [&](Node const &_, LayerAttrs const &layer) {
+    return parallel_layer_attrs_from_layer_attrs(layer);
+  };
+  auto tensor_map = [&](OpenDataflowValue const &_, TensorAttrs const &tensor) {
+    return parallel_tensor_attrs_from_tensor_attrs(tensor);
+  };
+  auto graph_view = rewrite_value_labels(
+      rewrite_node_labels(cg.raw_graph, layer_map), tensor_map);
+  return ParallelComputationGraph{
+      LabelledDataflowGraph<ParallelLayerAttrs, ParallelTensorAttrs>::
+          create_copy_of<
+              UnorderedSetLabelledOpenDataflowGraph<ParallelLayerAttrs,
+                                                    ParallelTensorAttrs>>(
+              graph_view)};
+}
+
+} // namespace FlexFlow
diff --git a/lib/pcg/test/src/pcg/pcg_from_computation_graph.cc b/lib/pcg/test/src/pcg/pcg_from_computation_graph.cc
@@ -0,0 +1,137 @@
+#include "pcg/pcg_from_computation_graph.h"
+#include "pcg/computation_graph.h"
+#include "pcg/parallel_computation_graph/parallel_computation_graph.h"
+#include "utils/containers/get_only.h"
+#include <doctest/doctest.h>
+
+using namespace ::FlexFlow;
+
+TEST_SUITE(FF_TEST_SUITE) {
+  TEST_CASE("pcg_from_computation_graph") {
+    std::string input_name = "input";
+    std::string dense_name = "dense";
+    std::string relu_name = "relu";
+
+    ComputationGraph cg = [&] {
+      ComputationGraph cg = make_empty_computation_graph();
+
+      TensorShape input_shape = TensorShape{
+          TensorDims{FFOrdered<nonnegative_int>{10_n, 12_n}}, DataType::FLOAT};
+      TensorAttrs input_attrs = TensorAttrs{input_shape,
+                                            /*sync_type=*/std::nullopt,
+                                            /*initializer=*/std::nullopt,
+                                            CreateGrad::YES};
+      LayerAttrs input_layer_attrs =
+          LayerAttrs{ComputationGraphOpAttrs{InputAttrs{}}, input_name};
+      LayerAddedResult input_added = add_layer(
+          cg, input_layer_attrs, /*inputs=*/{}, /*outputs=*/{input_attrs});
+      tensor_guid_t input_tensor = get_only(input_added.outputs);
+
+      LinearAttrs linear_attrs = LinearAttrs{/*out_channels=*/8_n,
+                                             /*use_bias=*/true,
+                                             /*data_type=*/DataType::FLOAT,
+                                             /*activation=*/Activation::RELU,
+                                             /*regularizer=*/std::nullopt};
+      TensorShape dense_output_shape = TensorShape{
+          TensorDims{FFOrdered<nonnegative_int>{10_n, 8_n}}, DataType::FLOAT};
+      LayerAttrs dense_layer_attrs =
+          LayerAttrs{ComputationGraphOpAttrs{linear_attrs}, dense_name};
+      LayerAddedResult dense_added =
+          add_layer(cg,
+                    /*attrs=*/dense_layer_attrs,
+                    /*inputs=*/{input_tensor},
+                    /*outputs=*/
+                    {TensorAttrs{dense_output_shape,
+                                 /*sync_type=*/std::nullopt,
+                                 /*initializer=*/std::nullopt,
+                                 CreateGrad::YES}});
+      tensor_guid_t dense_output = get_only(dense_added.outputs);
+
+      ElementUnaryAttrs relu_attrs =
+          ElementUnaryAttrs{OperatorType::RELU, /*scalar=*/std::nullopt};
+      LayerAttrs relu_layer_attrs =
+          LayerAttrs{ComputationGraphOpAttrs{relu_attrs}, relu_name};
+      add_layer(cg,
+                /*attrs=*/relu_layer_attrs,
+                /*inputs=*/{dense_output},
+                /*outputs=*/
+                {TensorAttrs{dense_output_shape,
+                             /*sync_type=*/std::nullopt,
+                             /*initializer=*/std::nullopt,
+                             CreateGrad::YES}});
+
+      return cg;
+    }();
+
+    ParallelComputationGraph correct = [&] {
+      ParallelComputationGraph pcg = empty_parallel_computation_graph();
+
+      ParallelTensorShape input_shape = ParallelTensorShape{
+          ParallelTensorDims{
+              FFOrdered<ShardParallelDim>{ShardParallelDim{10_n, 1_n},
+                                          ShardParallelDim{12_n, 1_n}},
+              ReplicaParallelDimSet{SumDegree{1_n}, DiscardCopyDegree{1_n}}},
+          DataType::FLOAT};
+
+      ParallelLayerAttrs input_layer_attrs =
+          ParallelLayerAttrs{PCGOperatorAttrs{InputAttrs{}}, input_name};
+
+      ParallelLayerAddedResult input_added = add_parallel_layer(
+          pcg,
+          /*attrs=*/input_layer_attrs,
+          /*inputs=*/{},
+          /*outputs=*/
+          {ParallelTensorAttrs{
+              input_shape, std::nullopt, std::nullopt, CreateGrad::YES}});
+
+      parallel_tensor_guid_t input_tensor = get_only(input_added.outputs);
+
+      LinearAttrs linear_attrs = LinearAttrs{/*out_channels=*/8_n,
+                                             /*use_bias=*/true,
+                                             /*data_type=*/DataType::FLOAT,
+                                             /*activation=*/Activation::RELU,
+                                             /*regularizer=*/std::nullopt};
+
+      ParallelLayerAttrs dense_layer_attrs =
+          ParallelLayerAttrs{PCGOperatorAttrs{linear_attrs}, dense_name};
+
+      ParallelTensorShape dense_output_shape = ParallelTensorShape{
+          ParallelTensorDims{
+              FFOrdered<ShardParallelDim>{ShardParallelDim{10_n, 1_n},
+                                          ShardParallelDim{8_n, 1_n}},
+              ReplicaParallelDimSet{SumDegree{1_n}, DiscardCopyDegree{1_n}}},
+          DataType::FLOAT};
+
+      ParallelLayerAddedResult dense_added =
+          add_parallel_layer(pcg,
+                             /*attrs=*/dense_layer_attrs,
+                             /*inputs=*/{input_tensor},
+                             /*outputs=*/
+                             {ParallelTensorAttrs{dense_output_shape,
+                                                  /*sync_type=*/std::nullopt,
+                                                  /*initializer=*/std::nullopt,
+                                                  CreateGrad::YES}});
+
+      parallel_tensor_guid_t dense_output = get_only(dense_added.outputs);
+
+      ElementUnaryAttrs relu_attrs =
+          ElementUnaryAttrs{OperatorType::RELU, std::nullopt};
+      ParallelLayerAttrs relu_layer_attrs =
+          ParallelLayerAttrs{PCGOperatorAttrs{relu_attrs}, relu_name};
+
+      add_parallel_layer(pcg,
+                         /*attrs=*/relu_layer_attrs,
+                         /*inputs=*/{dense_output},
+                         /*outputs=*/
+                         {ParallelTensorAttrs{dense_output_shape,
+                                              /*sync_type=*/std::nullopt,
+                                              /*initializer=*/std::nullopt,
+                                              CreateGrad::YES}});
+      return pcg;
+    }();
+
+    ParallelComputationGraph result = pcg_from_computation_graph(cg);
+
+    CHECK(pcgs_are_isomorphic(result, correct));
+  }
+}
diff --git a/lib/utils/include/utils/graph/labelled_dataflow_graph/algorithms/rewrite_node_labels.h b/lib/utils/include/utils/graph/labelled_dataflow_graph/algorithms/rewrite_node_labels.h
@@ -1,6 +1,7 @@
 #ifndef _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_GRAPH_LABELLED_DATAFLOW_GRAPH_ALGORITHMS_REWRITE_NODE_LABELS_H
 #define _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_GRAPH_LABELLED_DATAFLOW_GRAPH_ALGORITHMS_REWRITE_NODE_LABELS_H
 
+#include "utils/graph/labelled_dataflow_graph/algorithms/view_as_labelled_open_dataflow_graph.h"
 #include "utils/graph/labelled_open_dataflow_graph/algorithms/rewrite_node_labels.h"
 
 namespace FlexFlow {

diff --git a/lib/utils/include/utils/graph/labelled_dataflow_graph/algorithms/rewrite_value_labels.h b/lib/utils/include/utils/graph/labelled_dataflow_graph/algorithms/rewrite_value_labels.h
@@ -0,0 +1,23 @@
+#ifndef _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_GRAPH_LABELLED_DATAFLOW_GRAPH_ALGORITHMS_REWRITE_VALUE_LABELS_H
+#define _FLEXFLOW_LIB_UTILS_INCLUDE_UTILS_GRAPH_LABELLED_DATAFLOW_GRAPH_ALGORITHMS_REWRITE_VALUE_LABELS_H
+
+#include "utils/graph/labelled_dataflow_graph/algorithms/view_as_labelled_open_dataflow_graph.h"
+#include "utils/graph/labelled_open_dataflow_graph/algorithms/rewrite_value_labels.h"
+
+namespace FlexFlow {
+
+template <
+    typename NodeLabel,
+    typename ValueLabel,
+    typename F,
+    typename NewValueLabel =
+        std::invoke_result_t<F, OpenDataflowValue const &, ValueLabel const &>>
+LabelledDataflowGraphView<NodeLabel, NewValueLabel> rewrite_value_labels(
+    LabelledDataflowGraphView<NodeLabel, ValueLabel> const &g, F f) {
+  return rewrite_value_labels<NodeLabel, ValueLabel, F, NewValueLabel>(
+      view_as_labelled_open_dataflow_graph(g), f);
+}
+
+} // namespace FlexFlow
+
+#endif
diff --git a/lib/utils/src/utils/graph/labelled_dataflow_graph/algorithms/rewrite_node_labels.cc b/lib/utils/src/utils/graph/labelled_dataflow_graph/algorithms/rewrite_node_labels.cc
@@ -1 +1,15 @@
 #include "utils/graph/labelled_dataflow_graph/algorithms/rewrite_node_labels.h"
+#include "utils/archetypes/value_type.h"
+
+namespace FlexFlow {
+
+using NodeLabel = value_type<0>;
+using ValueLabel = value_type<1>;
+using NewNodeLabel = value_type<2>;
+using F = std::function<NewNodeLabel(Node const &, NodeLabel const &)>;
+
+template LabelledDataflowGraphView<NewNodeLabel, ValueLabel>
+    rewrite_node_labels(
+        LabelledDataflowGraphView<NodeLabel, ValueLabel> const &, F);
+
+} // namespace FlexFlow
diff --git a/lib/utils/src/utils/graph/labelled_dataflow_graph/algorithms/rewrite_value_labels.cc b/lib/utils/src/utils/graph/labelled_dataflow_graph/algorithms/rewrite_value_labels.cc
@@ -0,0 +1,16 @@
+#include "utils/graph/labelled_dataflow_graph/algorithms/rewrite_value_labels.h"
+#include "utils/archetypes/value_type.h"
+
+namespace FlexFlow {
+
+using NodeLabel = value_type<0>;
+using ValueLabel = value_type<1>;
+using NewValueLabel = value_type<2>;
+using F =
+    std::function<NewValueLabel(OpenDataflowValue const &, ValueLabel const &)>;
+
+template LabelledDataflowGraphView<NodeLabel, NewValueLabel>
+    rewrite_value_labels(
+        LabelledDataflowGraphView<NodeLabel, ValueLabel> const &, F);
+
+} // namespace FlexFlow