|
| 1 | +#include "task-spec/dynamic_graph/make_dynamic_open_dataflow_graph_from_cg.h" |
| 2 | +#include "op-attrs/pcg_operator_attrs.h" |
| 3 | +#include "pcg/computation_graph.h" |
| 4 | +#include "task-spec/dynamic_graph/dynamic_layer_guid_t.dtg.h" |
| 5 | +#include "task-spec/dynamic_graph/dynamic_open_dataflow_graph.h" |
| 6 | +#include "task-spec/dynamic_graph/dynamic_tensor_role.h" |
| 7 | +#include "utils/containers/generate_map.h" |
| 8 | +#include <optional> |
| 9 | +#include <unordered_map> |
| 10 | +#include <utility> |
| 11 | + |
| 12 | +namespace FlexFlow { |
| 13 | + |
| 14 | +DynamicOpenDataflowGraph |
| 15 | + make_dynamic_open_dataflow_graph_from_cg(ComputationGraph const &cg) { |
| 16 | + DynamicOpenDataflowGraph result = make_empty_dynamic_open_dataflow_graph(); |
| 17 | + |
| 18 | + for (auto const &[layer, attrs] : get_layer_attrs_mapping(cg)) { |
| 19 | + DynamicNodeAttrs result_attrs{ |
| 20 | + /*task_type=*/std::nullopt, |
| 21 | + /*device_coord=*/std::nullopt, |
| 22 | + /*mapping=*/std::nullopt, |
| 23 | + /*op_attrs=*/pcg_op_attrs_from_compgraph_op_attrs(attrs.op_attrs), |
| 24 | + /*pcg_layer_guid=*/dynamic_layer_guid_t{layer}, |
| 25 | + /*per_device_op_state=*/std::nullopt, |
| 26 | + }; |
| 27 | + |
| 28 | + std::unordered_map<DynamicTensorSlot, DynamicValueAttrs> result_inputs = |
| 29 | + transform( |
| 30 | + get_incoming_tensors(cg, layer), |
| 31 | + [](TensorSlotName const &slot_name, tensor_guid_t const &tensor) { |
| 32 | + return std::pair<DynamicTensorSlot, DynamicValueAttrs>{ |
| 33 | + DynamicTensorSlot{/*slot_name=*/slot_name, |
| 34 | + /*slot_tensor_role=*/std::nullopt}, |
| 35 | + DynamicValueAttrs{ |
| 36 | + /*tensor_guid=*/dynamic_tensor_guid_t{tensor}, |
| 37 | + /*parallel_tensor_shape=*/std::nullopt, |
| 38 | + /*shard_coord=*/std::nullopt, |
| 39 | + /*accessor=*/std::nullopt, |
| 40 | + /*role=*/std::nullopt}, |
| 41 | + }; |
| 42 | + }); |
| 43 | + std::unordered_map<DynamicTensorSlot, DynamicValueAttrs> result_outputs = |
| 44 | + transform( |
| 45 | + get_outgoing_tensors(cg, layer), |
| 46 | + [](TensorSlotName const &slot_name, tensor_guid_t const &tensor) { |
| 47 | + return std::pair<DynamicTensorSlot, DynamicValueAttrs>{ |
| 48 | + DynamicTensorSlot{/*slot_name=*/slot_name, |
| 49 | + /*slot_tensor_role=*/std::nullopt}, |
| 50 | + DynamicValueAttrs{ |
| 51 | + /*tensor_guid=*/dynamic_tensor_guid_t{tensor}, |
| 52 | + /*parallel_tensor_shape=*/std::nullopt, |
| 53 | + /*shard_coord=*/std::nullopt, |
| 54 | + /*accessor=*/std::nullopt, |
| 55 | + /*role=*/std::nullopt}, |
| 56 | + }; |
| 57 | + }); |
| 58 | + |
| 59 | + result.invocations.emplace(result_outputs, result_attrs, result_outputs); |
| 60 | + } |
| 61 | + |
| 62 | + return result; |
| 63 | +} |
| 64 | + |
| 65 | +} // namespace FlexFlow |
0 commit comments