diff --git a/mindspore/ccsrc/backend/common/somas/somas.cc b/mindspore/ccsrc/backend/common/somas/somas.cc index f0600b3161e..92998334b28 100644 --- a/mindspore/ccsrc/backend/common/somas/somas.cc +++ b/mindspore/ccsrc/backend/common/somas/somas.cc @@ -503,10 +503,19 @@ bool Somas::InitSomasTensors(const session::KernelGraph *graph) { return true; } -/** - * The InitSomasStreamAndNode function initializes Somas streams and nodes using the given KernelGraph. - * - * @param graph: Pointer to the KernelGraph used for initializing Somas streams and nodes. +/** + * The InitSomasStreamAndNode function is responsible for initializing SomasStream and SomasNode instances. + * A SomasStream represents a stream of execution on the device, and a SomasNode is classified into either + * kCommonNode or kCommunicationNode and added to the corresponding stream list. + * + * The detailed initialization process involves: + * 1. Iterating through the KernelGraph to identify and classify nodes as either kCommonNode or kCommunicationNode. + * 2. Adding the classified nodes to the appropriate SomasStream lists, thereby establishing the relationship + * between SomasStreams and SomasNodes. + * 3. The initialized SomasStreams and SomasNodes are used in further processing and computation within the Somas framework. + * + * @param graph: Pointer to the KernelGraph object, which contains the graph information and is used as the basis + * for initializing Somas streams and nodes. */ void Somas::InitSomasStreamAndNode(const session::KernelGraph *graph) { // Logging the start of the initialization process @@ -999,9 +1008,20 @@ void Somas::InitBasicInfo(const session::KernelGraph *graph) { } /** - * The GetNextOutputProcess function processes the output of GetNext operations in the graph. + * The GetNextOutputProcess function processes the output tensors of GetNext operations in the given graph. + * It iterates through each node in the execution order of the graph, identifies nodes corresponding to GetNext operations, + * and processes their output tensors. During this process, the function calculates the total aligned size + * of the tensors associated with these GetNext operations and assigns specific lifelong and type values to them. * - * @param graph: A pointer to the kernel graph containing GetNext operations. + * Detailed steps: + * 1. Iterate through each node (kernel) in the graph's execution order. + * 2. Check if the node corresponds to a GetNext operation. + * 3. If the node is a GetNext operation, find the node in the nodes_map_. + * 4. For each output tensor of the GetNext node, calculate the aligned size and add it to the total size. + * 5. Assign the lifelong value of the tensor to kLifeLongGraphAll and set the tensor type to kGetNextOutput. + * + * @param graph: A pointer to the KernelGraph containing GetNext operations. The graph provides the execution + * order of nodes and is used to identify and process GetNext nodes and their output tensors. */ void Somas::GetNextOutputProcess(const session::KernelGraph *graph) { MS_EXCEPTION_IF_NULL(graph); @@ -1029,9 +1049,21 @@ void Somas::GetNextOutputProcess(const session::KernelGraph *graph) { } /** - * The IndependentNodeOutputProcess function processes the output of independent nodes in the graph. + * The IndependentNodeOutputProcess function processes the output of independent nodes in the given graph. + * It iterates through each node in the execution order of the graph, identifies independent nodes, + * and processes their output tensors. During this process, the function calculates the total size + * of the tensors associated with these independent nodes and assigns a specific lifelong value to them, + * indicating their lifespan until the end of the graph. * - * @param graph: A pointer to the kernel graph containing independent nodes. + * Detailed steps: + * 1. Iterate through each node (kernel) in the graph's execution order. + * 2. Check if the node is independent. + * 3. If the node is independent, find the node in the nodes_map_. + * 4. For each output tensor of the independent node, calculate the aligned size and add it to the total size. + * 5. Assign the lifelong value of the tensor to kLifeLongGraphEnd. + * + * @param graph: A pointer to the KernelGraph containing independent nodes. The graph provides the execution + * order of nodes and is used to identify and process independent nodes and their output tensors. */ void Somas::IndependentNodeOutputProcess(const session::KernelGraph *graph) { MS_EXCEPTION_IF_NULL(graph); @@ -1104,8 +1136,24 @@ void Somas::SummaryInputProcess(const session::KernelGraph *graph) { } #endif -// This function processes reference nodes in the given KernelGraph to determine constraints and relationships -// between input and output tensors. +/** + * The RefNodeProcess function handles the processing of reference nodes within the given graph. + * It iterates through each kernel node, identifies reference nodes, and processes their input and output tensors. + * The function also calculates the total sizes of the input and output tensors of the reference nodes. + * + * Detailed Steps: + * 1. Iterate through each kernel node in the graph's execution order. + * 2. For each kernel node, retrieve the associated kernel module and its list of output sizes. + * 3. Iterate through each output size, check if it belongs to a reference node, and process it. + * 4. If the output is of a reference node, retrieve the corresponding original node and output tensor. + * 5. Set the type of the output tensor to kRefNodeOutput and update the total output size. + * 6. If the original node is a real kernel node, retrieve the corresponding input tensor and set its type to kRefNodeInput. + * 7. Update the total input size and add the input and output tensor IDs to ref_node_constraints_. + * + * @param graph: A pointer to the KernelGraph containing the nodes to be processed. + * The graph provides the execution order of nodes and is used to identify and process reference nodes + * and their associated tensors. + */ void Somas::RefNodeProcess(const session::KernelGraph *graph) { // Ensure the input graph is not null. MS_EXCEPTION_IF_NULL(graph); @@ -1179,12 +1227,28 @@ void Somas::RefNodeProcess(const session::KernelGraph *graph) { } } } - - // Log the total sizes of special tensors (RefNode inputs and outputs). - MS_LOG(INFO) << "Special Tensor total size: RefNode: input " << total_input_size << " output " << total_output_size; } -// The NonTaskSplitProcess function is responsible for processing the non-task operations in the given KernelGraph. +/** + * The NonTaskSplitProcess function processes non-task operations within the given kernel graph. + * Non-task operations are special operations in the computation graph that do not participate in task scheduling, + * and therefore need to be handled separately. This function identifies such operations, processes their input + * and output tensors, and adds constraints for these tensors in the SOMAS solver. + * + * Detailed Steps: + * 1. Iterate through each kernel node in the graph's execution order. + * 2. For each kernel node, check if it represents a non-task operation. + * 3. If it is a non-task operation, initialize a vector to store the IDs of input and output tensors of the operation. + * 4. Retrieve the corresponding SOMAS node for the non-task operation. + * 5. Check if the operation has at least one input tensor. If not, log an exception. + * 6. Set the type of the first input tensor to kRefNodeInput and add its ID to the vector. + * 7. Iterate through the output tensors of the operation, set their type to kRefNodeOutput, and add their IDs to the vector. + * 8. Add the vector containing the IDs of input and output tensors to the ref_node_constraints_ list for further processing. + * + * @param graph: A pointer to the KernelGraph containing the nodes to be processed. + * The graph provides the execution order of nodes and is used to identify and process non-task operations + * and their associated tensors. + */ void Somas::NonTaskSplitProcess(const session::KernelGraph *graph) { // Check if the input graph is not null. MS_EXCEPTION_IF_NULL(graph); @@ -1233,10 +1297,25 @@ void Somas::NonTaskSplitProcess(const session::KernelGraph *graph) { } } - -// The UnReuseNodeProcess function processes nodes that are not reusable. For each node in full_name_list, -// it sets the lifelong_value_ attribute of its input, output, and workspace tensors to kLifeLongGraphAll, -// indicating that these tensors are active throughout the graph's lifecycle and are not reusable. +/** + * The UnReuseNodeProcess function iterates through the given kernel graph and processes nodes + * that are marked as "UnReuse". For these nodes, it sets the lifelong value of their input, + * output, and workspace tensors to kLifeLongGraphAll, indicating that the memory of these + * tensors should not be reused throughout the entire graph execution. + * + * Detailed Steps: + * 1. Define a list of full names of nodes that should be processed as "UnReuse" nodes. + * 2. If there are no nodes to process, the function returns immediately. + * 3. Iterate through each kernel node in the graph's execution order. + * 4. Check if the current node's full name matches any in the defined list of "UnReuse" nodes. + * 5. If a match is found, log the information and retrieve the corresponding SOMAS node. + * 6. Process the input, output, and workspace tensors of the SOMAS node by setting their + * lifelong value to kLifeLongGraphAll, thereby marking them as non-reusable. + * + * @param graph: A pointer to the KernelGraph containing the nodes to be processed. The graph provides + * the execution order of nodes and is used to identify and process the "UnReuse" nodes + * and their associated tensors. + */ void Somas::UnReuseNodeProcess(const session::KernelGraph *graph) { MS_EXCEPTION_IF_NULL(graph); @@ -1287,9 +1366,21 @@ void Somas::UnReuseNodeProcess(const session::KernelGraph *graph) { } } -// The GenContiguousList function generates lists of contiguous tensors. For each node of type kCommunicationNode, -// it checks whether its input and output tensors are contiguous. If not, it updates the aligned_size_ of these tensors, -// marks them as contiguous, and adds them to the contiguous_tensors_list_. +/** + * GenContiguousList Function + * -------------------------- + * The GenContiguousList function is a crucial component in the SOMAS (Solver for Memory Assignment + * Scheduling) framework. It goes through every node in the nodes_list_ and identifies nodes of type + * kCommunicationNode to ensure that the input and output tensors for these nodes are contiguous. + * + * 1. For every kCommunicationNode type node, the function checks the contiguity of its input and output tensors. + * 2. If these tensors are not contiguous, it updates the aligned_size_ of the tensors, marks them as contiguous, + * and then adds them to the contiguous_tensors_list_. + * 3. It also ensures there are no duplicate tensor IDs in the input or output tensors of the node, throwing an exception if any are found. + * + * @param graph: A pointer to the kernel graph. Represents the computational graph of the session. + * + */ void Somas::GenContiguousList(const session::KernelGraph *graph) { MS_EXCEPTION_IF_NULL(graph); @@ -1350,10 +1441,23 @@ void Somas::GenContiguousList(const session::KernelGraph *graph) { } -/** - * The ComputeConflictPairs function computes the conflict pairs among tensors in the SOMAS - * Scheduler to resolve memory conflicts. It utilizes a bitset model for efficient conflict - * computation, and it can operate in multi-thread mode for large-scale tensor sets. +/** + * ComputeConflictPairs Function + * ----------------------------- + * The ComputeConflictPairs function is responsible for computing conflict pairs among tensors + * in the SOMAS Scheduler to avoid and resolve memory conflicts during tensor allocation. + * + * The function implements the following workflow: + * 1. It checks if there are tensors available for conflict computing, logging a message and returning if none are found. + * 2. The nodes_list_ is sorted for further computations. + * 3. It then updates tensor destinations using the UpdateTensorDestinations function. + * 4. Initialize nodes_dependency vector with bitsets for each node in the nodes_list_. + * 5. Compute ancestor paths via bitset for time dependence. + * 6. Initialize the reuse_matrix_ to store tensor conflicts information. + * 7. Depending on the number of tensors, it decides whether to compute conflicts in single-thread or multi-thread mode. + * 8. If multi-threading is used, tasks are created and executed in parallel using the ThreadPool. + * 9. Finally, it logs the time taken for conflict computation. + * */ void Somas::ComputeConflictPairs() { // Check if there are no tensors for conflict computing @@ -1534,12 +1638,25 @@ void Somas::ComputeMultiTensorConflicts(const std::vector &calc_ } /** - * The ComputeOneTensorConflicts function computes the conflicts for a single tensor with all other tensors in the graph. + * ComputeOneTensorConflicts Function + * ---------------------------------- + * The ComputeOneTensorConflicts function is vital in identifying memory conflicts within + * the graph. It computes the conflicts between a single specified tensor and all other tensors in the graph. * - * @param calc_tensor: A shared pointer to the SomasTensor object for which conflicts need to be computed. - * @param all_tensors_list: A list of all tensors in the graph. - * @param nodes_dependency: A vector representing the dependencies between nodes in the graph. - * @param tensor_relation: A pointer to a vector representing the relations between tensors in the graph. + * 1. The function iterates over each tensor in the all_tensors_list. + * 2. It skips checking conflicts for tensors that are the same as the calc_tensor, lifelong tensors, + * semi-lifelong start tensors, tensors with reference overlap, and tensors with aligned size zero. + * 3. If calc_tensor and the target tensor share the same source node ID, they are not checked for conflicts. + * 4. If a conflict has already been identified between calc_tensor and the target tensor, it is not checked again. + * 5. The function then checks whether all consumers of calc_tensor are dependencies of the source node of the target tensor. + * 6. If they are, or if the target tensor's source node ID is the same as one of the destination node IDs of calc_tensor, + * the tensors can’t be reused, and the function continues to the next iteration. + * 7. Otherwise, the tensor pair is marked as having dependencies and therefore can be reused. + * + * @param calc_tensor: A shared pointer to the SomasTensor object representing the tensor for which conflicts need to be computed. + * @param all_tensors_list: A vector containing shared pointers to all SomasTensor objects in the graph. + * @param nodes_dependency: A vector of DynamicBitSets representing the dependencies between nodes in the graph. + * @param tensor_relation: A pointer to a vector of DynamicBitSets representing the relations between tensors in the graph. */ void Somas::ComputeOneTensorConflicts(const std::shared_ptr &calc_tensor, const std::vector &all_tensors_list, @@ -1593,6 +1710,34 @@ void Somas::ComputeOneTensorConflicts(const std::shared_ptr &calc_t bool Somas::NodeSort(const SomasNodePtr &node1, const SomasNodePtr &node2) { return node1->GetId() < node2->GetId(); } +/** + * The Assign function is responsible for assigning memory to each tensor in the computation graph. + * The overall process can be described in several key steps as follows: + * + * 1. Preprocess Reference Nodes: + * - Invokes UpdateRefTensorsConflict() to compute and update conflicts between reference tensors. + * - Identifies and records contiguous tensors that contain reference tensors. + * - Filters out and removes tensors and contiguous lists that do not require memory assignments. + * + * 2. Prepare Solver Information: + * - Iterates through the list of tensors (tensors_list_) and extracts the SomasSolverTensorDesc information. + * - Populates the solver_tensor_desc_map_ with the tensor descriptors for the solver to use. + * + * 3. Solving Process: + * - A new SomasSolverPre instance is created, and the Solving method is invoked with the prepared information. + * - The solver uses the information provided, along with the constraint matrix (reuse_matrix_), + * to allocate memory for each tensor while avoiding conflicts and optimizing memory usage. + * - If the solving process fails, it logs the error and the function returns false. + * + * 4. Update Tensor Offsets: + * - Based on the results of the solving process, the offsets of each tensor in the tensors_list_ are updated. + * - The offsets of reference tensors and contiguous tensors are further adjusted with UpdateRefTensorsOffset() + * and UpdateContiguousTensorsOffset(). + * - The overall memory offset (mem_offset_) is set based on the maximum offset value obtained from the solver. + * + * @param graph: Pointer to the KernelGraph object representing the computation graph of the session. + * @return bool: Returns true if the memory assignment is successful, logs the error, and returns false if failed. + */ bool Somas::Assign(const session::KernelGraph *graph) { MS_LOG(DEBUG) << "Somas Assign start..."; if (tensors_list_.empty()) { @@ -1688,6 +1833,34 @@ bool Somas::Assign(const session::KernelGraph *graph) { return true; } +/** + * The GetContiguousListContainRefTensor function is responsible for identifying and mapping contiguous lists + * of tensors that contain reference tensors. Reference tensors share the same memory space and thus have constraints + * on memory assignment. + * + * The function performs the following key steps: + * + * 1. Initialize the Map: + * - Initializes the contiguous_list_with_ref_index_map map, which will hold the mappings. + * - Retrieves a map of reference tensors in contiguous lists using GetRefTensorsInContiguousList(). + * + * 2. Identify and Map Contiguous Lists with Reference Tensors: + * - Iterates through each pair of reference tensors in the ref_tensors_in_contiguous_map. + * - For each reference tensor pair, it searches the contiguous_tensors_list_ to find the lists that contain them. + * - If both reference tensors are found, their corresponding contiguous list indices and positions within the list + * are recorded. + * - The found indices are then used to update the contiguous_list_with_ref_index_map. + * - Performs error checking to identify any inconsistencies or anomalies in the mapping, logging warnings if any issues + * are found. + * + * 3. Additional Error Checking: + * - Goes through the generated map to further check for inconsistencies, such as mismatched list sizes or unconsidered + * reference pairs, and logs warnings if necessary. + * + * @return std::map: Returns a map where the key represents the index of the contiguous list containing + * the first reference tensor, and the value is the index of the contiguous list + * containing the second reference tensor. + */ std::map Somas::GetContiguousListContainRefTensor() { // key: contiguous list index with ref node input; value: contiguous list index with ref node output std::map contiguous_list_with_ref_index_map; @@ -1765,6 +1938,15 @@ std::map Somas::GetContiguousListContainRefTensor() { return contiguous_list_with_ref_index_map; } +/** + * @brief Identifies and returns a map containing reference tensors within contiguous lists. + * + * The function iterates through each list of reference node constraints and counts the number of contiguous tensors + * in each list. It logs warnings for any detected irregularities in the list sizes and the number of contiguous + * tensors. If a list has exactly two contiguous tensors, they are added to the map. + * + * @return std::map: A map where each entry represents a pair of reference tensors within contiguous lists. + */ std::map Somas::GetRefTensorsInContiguousList() { // key: refnode input value: refnode output std::map ref_tensors_in_contiguous_map; @@ -1788,6 +1970,13 @@ std::map Somas::GetRefTensorsInContiguousList() { return ref_tensors_in_contiguous_map; } +/** + * @brief Updates the offset of contiguous tensors based on the reference list map. + * + * The function uses the provided map to update the offset of tensors in each contiguous list. It ensures that + * tensors in the same index positions across two linked contiguous lists have the same offset. Additionally, + * it performs postprocessing to adjust the gaps between contiguous tensors. + */ void Somas::UpdateContiguousTensorsOffset(const std::map &contiguous_ref_list_map) { // Handle contiguous ref node for (auto ref_list_pair : contiguous_ref_list_map) { @@ -1805,6 +1994,12 @@ void Somas::UpdateContiguousTensorsOffset(const std::map &contig } } +/** + * @brief Performs postprocessing to update the offset of reference tensors. + * + * This function iterates through each reference node constraint list and updates the offset of all tensors + * in the list to match the offset of the first tensor in the list. + */ void Somas::UpdateRefTensorsOffset() { // Ref Node Postprocessing MS_LOG(INFO) << "\nStart Solving Postprocessing for Ref Node"; @@ -1832,23 +2027,37 @@ void Somas::UpdateRefOverlapTensorsConflicts() { MS_LOG(INFO) << "End Solving Preprocessing for Ref Overlap"; } +/** + * The UpdateRefTensorsConflict function updates the conflicts between reference tensors in the graph. + * It iterates through each list of reference node constraints, examines the reusability of tensors, + * and updates the reuse_matrix_ accordingly. Additionally, it modifies the aligned_size_ of non-contiguous + * tensors in the reference node list, ensuring that they are ignored by the solver during the subsequent processing. + */ void Somas::UpdateRefTensorsConflict() { - // Keep all constraints for first tensor in list + // Iterate over each list of reference node constraints. for (auto ref_node_list : ref_node_constraints_) { - size_t tid_0 = ref_node_list[0]; + size_t tid_0 = ref_node_list[0]; // Store the ID of the first tensor in the current list. + + // Loop through all tensors in the tensor list. for (SomasTensorPtr tensor : tensors_list_) { + // Check if the first tensor (tid_0) can be reused with the current tensor. if (reuse_matrix_[tid_0].IsBitTrue(tensor->GetId()) == false) { - continue; + continue; // Skip to the next tensor if they cannot be reused. } + + // Iterate over all tensor IDs in the current reference node list. for (size_t tid : ref_node_list) { + // If the current tensor ID (tid) cannot be reused with the tensor, update the reuse_matrix_ accordingly. if (reuse_matrix_[tid].IsBitTrue(tensor->GetId()) == false) { reuse_matrix_[tid_0].SetBitFalse(tensor->GetId()); reuse_matrix_[tensor->GetId()].SetBitFalse(tid_0); - break; + break; // Break out of the loop as one non-reusable tensor ID is found. } } } - // Set rest to size 0, so that solver ignores them (if not contiguous) + + // Update the aligned_size_ for the rest of the tensors in ref_node_list to 0 if they are not contiguous. + // This ensures that the solver ignores them. for (size_t i = 1; i < ref_node_list.size(); ++i) { if (!tensors_map_[ref_node_list[i]]->contiguous_) { tensors_map_[ref_node_list[i]]->aligned_size_ = 0; @@ -1856,7 +2065,6 @@ void Somas::UpdateRefTensorsConflict() { } } } - std::string Somas::GetSplitName(const std::string &scope_name) const { auto index = scope_name.rfind('/'); if (index == std::string::npos) {