mindspore/ccsrc/frontend/parallel/step_parallel.h

@@ -38,146 +38,198 @@ using OperatorInfoPtr = std::shared_ptr<mindspore::parallel::OperatorInfo>;
 
 namespace mindspore {
 namespace parallel {
-const uint64_t kUSecondInSecond = 1000000;
-const int32_t RECURSION_LIMIT = 3;
+// Constants
+const uint64_t kUSecondInSecond = 1000000;  // Microseconds in a second
+const int32_t RECURSION_LIMIT = 3;  // Recursion limit for some functions
 
+// Struct to store information about a loss node
 struct LossNodeInfo {
-  bool has_tuple_getitem = false;
-  int64_t dout_index = 0;  // now don't support the sens is a tuple
-  CNodePtr loss_node = nullptr;
+  bool has_tuple_getitem = false;  // Indicates if the node has a tuple_getitem operation
+  int64_t dout_index = 0;  // Index of the "dout" tensor in the node (currently doesn't support tuple "sens")
+  CNodePtr loss_node = nullptr;  // Pointer to the loss node
 };
 
+// Struct to store communication information
 struct CommInfo {
-  int64_t device_num = 1;
-  int64_t global_rank = 0;
-  std::string world_group;
-  std::string communication_backend;
+  int64_t device_num = 1;  // Number of devices
+  int64_t global_rank = 0;  // Global rank of the current device
+  std::string world_group;  // World group name
+  std::string communication_backend;  // Communication backend used
 };
 
+// Function to create input nodes for an operator
 std::vector<AnfNodePtr> CreateInput(const Operator &op, const AnfNodePtr &node, const std::string &instance_name);
+
+// Function to perform forward communication for a list of operators
 void ForwardCommunication(OperatorVector forward_op, const CNodePtr &node);
 
+// Function to insert redistribution operations into the graph
 void InsertRedistribution(const RedistributionOpListPtr &redistribution_oplist_ptr, const CNodePtr &node,
                           const FuncGraphPtr &func_graph, int64_t pos, const CNodePtr &pre_node);
 
+// Function to get the input tensor layout for a given node
 TensorLayout GetTensorInLayout(const CNodePtr &pre_node, const PrimitivePtr &pre_prim,
                                const OperatorInfoPtr &distribute_operator_pre);
 
+// Function to get the distribute operator associated with a CNode
 OperatorInfoPtr GetDistributeOperator(const CNodePtr &node);
 
+// Function to perform tensor redistribution
 void Redistribution(const std::pair<AnfNodePtr, int64_t> &node_pair, const OperatorInfoPtr &distribute_operator,
                     const CNodePtr &middle_node, int64_t index, TensorRedistribution tensor_redistribution,
                     const CNodePtr &pre_node);
 
+// Function to check if a strategy is found in attributes
 bool StrategyFound(const mindspore::HashMap<std::string, ValuePtr> &attrs);
 
+// Function to check if a specific attribute is found in attributes
 bool AttrFound(const mindspore::HashMap<std::string, ValuePtr> &attrs, const std::string &target);
 
+// Function to get the accumulated gradient node for a parameter
 AnfNodePtr GetAccuGrad(const std::vector<AnfNodePtr> &parameters, const std::string &weight_name);
 
+// Function to mark forward CNodes in the computation graph
 void MarkForwardCNode(const FuncGraphPtr &root);
 
+// Function to check if there are communication ops in the graph
 bool FindCommunicationOp(const std::vector<AnfNodePtr> &all_nodes);
 
+// Function to perform step redistribution
 void StepRedistribution(const CNodePtr &node, const OperatorInfoPtr &distribute_operator, const CNodePtr &insert_node,
                         const TensorRedistribution &tensor_redistribution, const CNodePtr &pre_node);
 
+// Function to replace operators in a step
 void StepReplaceOp(OperatorVector replace_op, const CNodePtr &node);
 
+// Function to insert virtual div operators
 void InsertVirtualDivOp(const VirtualDivOp &virtual_div_op, const CNodePtr &node);
 
+// Function to find a CNode in a function graph
 std::pair<bool, CNodePtr> FindCNode(const AnfNodePtr &anode, const std::string &name, const FuncGraphPtr &func_graph,
                                     size_t max_depth);
 
-// Generate and init parallel operator
+// Function to generate and initialize a parallel operator
 OperatorInfoPtr OperatorInstance(const PrimitivePtr &prim, const PrimitiveAttrs &attrs,
                                  const std::vector<Shapes> &shape_list);
 
-// Generate without initing parallel operator
+// Function to generate a parallel operator without initialization
 OperatorInfoPtr NewOperatorInstance(const PrimitivePtr &prim, const PrimitiveAttrs &attrs,
                                     std::vector<Shapes> shape_list);
 
-// Extract strategy from attr
+// Function to extract strategy from an attribute
 StrategyPtr ExtractStrategy(const ValuePtr &strategy);
 
-// Extract shape from anfnode
+// Function to extract shapes from a CNode
 std::vector<Shapes> ExtractShape(const CNodePtr &node);
 
-// Find finally sub graph
+// Function to find a sub-graph within a function graph
 std::pair<AnfNodePtr, int64_t> FindSubGraph(const FuncGraphPtr &func_graph, const AnfNodePtr &parameter);
 
-// Set distribute shape for parameters abstract
+// Function to set parallel shapes for parameter abstracts
 std::string SetParallelShape(const AnfNodePtr &parameter, const std::pair<AnfNodePtr, int64_t> &res);
 
-// change parameters'shape in resource
+// Function to change parameters' shapes in resource
 void CoverSliceShape(const FuncGraphPtr &root);
 
+// Function to label batch size split
 void LableBatchSizeSplit(const CNodePtr &node);
 
+// Function to set virtual dataset strategy
 void SetVirtualDatasetStrategy(const CNodePtr &node);
+
+// Function to check if virtual output should be inserted
 bool IsInsertVirtualOutput(const FuncGraphPtr &root);
 
+// Function to set strided slice split strategy
 void SetStridedSliceSplitStrategy(const std::vector<AnfNodePtr> &all_nodes);
 
-// Create parallel operator for primitive node(has strategy)
+// Function to create parallel operators and extract information from the computation graph
 void ExtractInformation(const std::vector<AnfNodePtr> &all_nodes);
 
+// Function to get input layout from a CNode
 TensorLayout GetInputLayoutFromCNode(const std::pair<AnfNodePtr, int64_t> &node_pair);
 
+// Function to find the next layout from a CNode
 std::shared_ptr<TensorLayout> FindNextLayout(const CNodePtr &node);
 
+// Function to get output layout from a CNode
 std::shared_ptr<TensorLayout> GetOutputLayoutFromCNode(const CNodePtr &cnode, size_t output_index);
 
+// Function to find the previous layout from an AnfNode
 std::shared_ptr<TensorLayout> FindPrevParallelCareNodeLayout(const AnfNodePtr &node, size_t output_index);
 
+// Function to find the previous layout from an AnfNode
 std::shared_ptr<TensorLayout> FindPrevLayout(const AnfNodePtr &node);
 
+// Function to initialize reshaping
 void ReshapeInit(const std::vector<AnfNodePtr> &all_nodes);
 
+// Function to generate batch parallel strategy for an operator
 StrategyPtr GenerateBatchParallelStrategy(const OperatorInfoPtr operator_, const PrimitivePtr prim);
 
-// Add node for whole graph
+// Function to add parallel communication operations to the computation graph
 void ParallelCommunication(const FuncGraphPtr &root, const std::vector<AnfNodePtr> &all_nodes,
                            const FuncGraphManagerPtr &manager);
 
+// Function to create a mapping of parameter names to parameter nodes
 ParameterMap NodeParameterName(const CNodePtr &node, int64_t index, size_t curr_depth);
 
+// Function to checkpoint the strategy for multi-train
 void CheckpointStrategy(const std::vector<AnfNodePtr> &all_nodes, const FuncGraphPtr &root);
 
-// main step of Parallel
+// Main function for performing step parallel optimization
 bool StepParallel(const FuncGraphPtr &func_graph, const opt::OptimizerPtr &optimizer);
 
+// Function to get the index of a TupleGetItem operation
 int64_t GetTupleGetItemIndex(const CNodePtr &cnode);
 
+// Function to initialize parallel context
 Status ParallelInit();
 
+// Function to forward the graph and return a set of forward graphs
 std::set<FuncGraphPtr> ForwardGraph(const FuncGraphPtr &root);
 
+// Function to extract input tensor names from a CNode
 std::vector<std::string> ExtractInputsTensorName(const CNodePtr &node);
 
+// Function to find the next layout for a parameter node
 std::shared_ptr<TensorLayout> FindParameterNextLayout(const AnfNodePtr &node);
 
+// Function to check if a parameter node is used in a function graph
 bool IsUsedParameter(const FuncGraphPtr &graph, const AnfNodePtr &parameter);
 
+// Function to apply parallel optimization on a parameter
 void ApplyParallelOptOnParam(TensorLayout *tensor_layout, const OperatorInfoPtr &distribute_operator,
                              const CNodePtr &cnode, const AnfNodePtr &parameter, size_t index);
 
+// Function to set the strategy for the last node
 void SetLastNodeStrategy(const StrategyPtr strategyPtr);
 
+// Function to create groups based on a checkpoint file
 bool CreateGroupsByCkptFile(const std::string &file);
 
+// Function to find the unique IDs of the last nodes in a function graph
 void FindLastNodesUniqueId(const FuncGraphPtr &root, std::vector<std::string> *unique_ids,
                            std::vector<size_t> *indexes);
 
+// Function to insert virtual output for the computation graph
 void InsertVirtualOutput(const FuncGraphPtr &root, const std::vector<AnfNodePtr> &all_nodes);
 
+// Function to get the name of the Mirror operator
 std::string MirrorOpName();
 
+// Function to get communication information
 CommInfo GetCommInfo();
 
+// Function to get the name of a primitive operator in a CNode
 std::string GetPrimName(const CNodePtr &node);
 
+// Function to reorder nodes for pipeline splitting
 void ReorderForPipelineSplit(const FuncGraphPtr &root, const FuncGraphManagerPtr &manager, int64_t pipeline_stages);
+
+}  // namespace parallel
+}  // namespace mindspore
+
 }  // namespace parallel
 }  // namespace mindspore