添加了pass的注释

erge branch 'zbtrs' of https://gitlink.org.cn/ssk015/mindspore2022 into zbtrs

mindspore/ccsrc/backend/common/optimizer/common_backend_optimization.cc

@@ -80,18 +80,50 @@ void BackendCommonOptimization(const std::shared_ptr<session::KernelGraph> &kern
   MS_LOG(INFO) << "Status record: end common optimization. graph id: " << kernel_graph->graph_id();
 }
 
+/**
+ * @brief Performs the final common optimization on a given kernel graph.
+ * 
+ * This function conducts a series of optimization passes on the kernel graph. The
+ * passes are bundled together in a PassManager which is then used by the GraphOptimizer
+ * to optimize the graph. After optimizations, the kernel graph's execution order is set 
+ * by default. If the dump IR flag is enabled, this function will also dump the intermediate
+ * representation of the optimized kernel graph.
+ * 
+ * @param kernel_graph A shared pointer to the kernel graph to be optimized. It is 
+ *                     a data structure representing the computation graph at the kernel level 
+ *                     in the MindSpore computational framework.
+ * 
+ * @return void This function does not return any value. It modifies the provided 
+ *              kernel_graph in-place by applying the specified optimizations.
+ * 
+ * @note The function will throw exceptions if kernel_graph is null or if there are issues 
+ *       accessing the global context (MsContext).
+ * 
+ * @see GraphOptimizer, PassManager, MsContext
+ */
 void CommonFinalOptimization(const std::shared_ptr<session::KernelGraph> &kernel_graph) {
+  // Check if kernel_graph is null and throw an exception if true.
   MS_EXCEPTION_IF_NULL(kernel_graph);
-  // Run optimizer passes.
+  
+  // Initialize the graph optimizer and the pass manager for final optimization.
   auto optimizer = std::make_shared<GraphOptimizer>();
   auto pm = std::make_shared<PassManager>("final_opt");
+  
+  // Add optimization passes to the pass manager.
   pm->AddPass(std::make_shared<OptimizeUpdateState>());
   pm->AddPass(std::make_shared<AddAkgKernelAttrs>());
+  
+  // Attach the pass manager to the graph optimizer.
   optimizer->AddPassManager(pm);
+  
+  // Apply optimization on the kernel graph using the graph optimizer.
   (void)optimizer->Optimize(kernel_graph);
+  
+  // Set the default execution order for the kernel graph.
   kernel_graph->SetExecOrderByDefault();
+
 #ifdef ENABLE_DUMP_IR
-  // Dump IR if save_graphs is set.
+  // Check and dump the kernel graph's intermediate representation if the save_graphs flag is enabled.
   auto context = MsContext::GetInstance();
   MS_EXCEPTION_IF_NULL(context);
   const bool save_graphs = context->get_param<bool>(MS_CTX_SAVE_GRAPHS_FLAG);
@@ -131,17 +163,54 @@ void CommonUnifyMindIR(const std::shared_ptr<session::KernelGraph> &kernel_graph
 #endif
 }
 
+/**
+ * @brief Applies a pass to add dynamic shape attributes to the given kernel graph.
+ * 
+ * This function sets up a GraphOptimizer and PassManager, then applies the
+ * AddDynamicShapeAttr pass on the kernel graph to add necessary dynamic shape
+ * attributes to nodes in the graph.
+ * 
+ * @param kernel_graph A shared pointer to the kernel graph to be processed.
+ * 
+ * @return void Modifies the provided kernel_graph in-place.
+ */
 void AddDynamicShapeAttrPass(const std::shared_ptr<session::KernelGraph> &kernel_graph) {
+  // Initialize the graph optimizer.
   auto opt = std::make_shared<GraphOptimizer>();
+
+  // Setup a pass manager for adding dynamic shape attributes.
   auto pm = std::make_shared<PassManager>("add_dynamic_shape_attr");
+
+  // Add the specific pass for adding dynamic shape attributes.
   pm->AddPass(std::make_shared<AddDynamicShapeAttr>());
+
+  // Add the pass manager to the graph optimizer.
   opt->AddPassManager(pm);
+
+  // Execute the optimization on the kernel graph.
   (void)opt->Optimize(kernel_graph);
 }
 
+/**
+ * @brief Eliminates illegal data types from the provided kernel graph.
+ * 
+ * This function initializes a GraphOptimizer and PassManager, then applies the
+ * EliminateFuncDataType pass on the kernel graph to remove or replace any
+ * nodes with illegal data types. The kernel graph's execution order is 
+ * then set by default. If the dump IR flag is enabled, this function will 
+ * also dump the intermediate representation of the kernel graph both before 
+ * and after the optimization.
+ * 
+ * @param kernel_graph A shared pointer to the kernel graph to be processed.
+ * 
+ * @return void Modifies the provided kernel_graph in-place.
+ */
 void EliminateIllegalDataTypePass(const std::shared_ptr<session::KernelGraph> &kernel_graph) {
+  // Check if kernel_graph is null and log the start of this pass.
   MS_EXCEPTION_IF_NULL(kernel_graph);
   MS_LOG(INFO) << "Start eliminate illegal data type for kernel graph id:" << kernel_graph->graph_id();
+
+  // Check and dump the kernel graph's intermediate representation if the save_graphs flag is enabled.
 #ifdef ENABLE_DUMP_IR
   auto context_ptr = MsContext::GetInstance();
   MS_EXCEPTION_IF_NULL(context_ptr);
@@ -152,12 +221,26 @@ void EliminateIllegalDataTypePass(const std::shared_ptr<session::KernelGraph> &k
     DumpIR(file_name, kernel_graph);
   }
 #endif
+
+  // Initialize the graph optimizer.
   auto opt = std::make_shared<GraphOptimizer>();
+
+  // Setup a pass manager for eliminating illegal data types.
   auto pm = std::make_shared<PassManager>("common_eliminate_illegal_data_type_pm");
+
+  // Add the specific pass for eliminating illegal data types.
   pm->AddPass(std::make_shared<EliminateFuncDataType>());
+
+  // Add the pass manager to the graph optimizer.
   opt->AddPassManager(pm);
+
+  // Execute the optimization on the kernel graph.
   (void)opt->Optimize(kernel_graph);
+
+  // Set the default execution order for the kernel graph.
   kernel_graph->SetExecOrderByDefault();
+
+  // Check and dump the kernel graph's intermediate representation post optimization if save_graphs flag is enabled.
 #ifdef ENABLE_DUMP_IR
   if (save_graphs) {
     std::string file_name =
@@ -167,9 +250,26 @@ void EliminateIllegalDataTypePass(const std::shared_ptr<session::KernelGraph> &k
 #endif
 }
 
+
+/**
+ * @brief Converts custom ops related to dynamic shapes in the given kernel graph.
+ * 
+ * This function sets up a GraphOptimizer and PassManager, then applies a series
+ * of passes (ConvertCustomOp and LinkCustomOp) on the kernel graph to handle
+ * custom operations related to dynamic shapes. If the dump IR flag is enabled,
+ * this function will also dump the intermediate representation of the kernel 
+ * graph both before and after the optimization.
+ * 
+ * @param kernel_graph A shared pointer to the kernel graph to be processed.
+ * 
+ * @return void Modifies the provided kernel_graph in-place.
+ */
 void DynamicShapeConvertPass(const std::shared_ptr<session::KernelGraph> &kernel_graph) {
+  // Check if kernel_graph is null and log the start of this pass.
   MS_EXCEPTION_IF_NULL(kernel_graph);
   MS_LOG(INFO) << "Start dynamic shape convert for kernel graph id:" << kernel_graph->graph_id();
+
+  // Check and dump the kernel graph's intermediate representation if save_graphs flag is enabled.
 #ifdef ENABLE_DUMP_IR
   auto context_ptr = MsContext::GetInstance();
   MS_EXCEPTION_IF_NULL(context_ptr);
@@ -180,12 +280,24 @@ void DynamicShapeConvertPass(const std::shared_ptr<session::KernelGraph> &kernel
     DumpIR(file_name, kernel_graph);
   }
 #endif
+
+  // Initialize the graph optimizer.
   auto optimizer = std::make_shared<opt::GraphOptimizer>();
+
+  // Setup a pass manager for dynamic shape conversion.
   auto dynamic_shape_convert_pm = std::make_shared<opt::PassManager>("dynamic_shape_convert_pm");
+
+  // Add specific passes related to dynamic shape conversion.
   dynamic_shape_convert_pm->AddPass(std::make_shared<opt::dynamic_shape::ConvertCustomOp>());
   dynamic_shape_convert_pm->AddPass(std::make_shared<opt::dynamic_shape::LinkCustomOp>());
+
+  // Add the pass manager to the graph optimizer.
   optimizer->AddPassManager(dynamic_shape_convert_pm);
+
+  // Execute the optimization on the kernel graph.
   (void)optimizer->Optimize(kernel_graph);
+
+  // Check and dump the kernel graph's intermediate representation post optimization if save_graphs flag is enabled.
 #ifdef ENABLE_DUMP_IR
   if (save_graphs) {
     std::string file_name =

mindspore/ccsrc/backend/common/optimizer/const_input_to_attr.cc

@@ -126,6 +126,10 @@ bool ConstInputToAttrInfoRegistry::GetRegisterByOpName(const std::string &op_nam
   return false;
 }
 
+/*
+ * @brief If the input of cnode is a const tensor, and the index of input is in the input_attrs,
+ * then set the const tensor to attr of cnode.
+ */
 void ConstInputToAttr(const CNodePtr &cnode, const mindspore::HashSet<size_t> &input_attrs) {
   MS_EXCEPTION_IF_NULL(cnode);
   std::vector<AnfNodePtr> new_inputs;
@@ -148,6 +152,7 @@ void ConstInputToAttr(const CNodePtr &cnode, const mindspore::HashSet<size_t> &i
       input_node = AnfUtils::VisitKernel(input_node, 0).first;
     }
     if (input_attrs.find(i) != input_attrs.end() && input_node->isa<ValueNode>() && !HasAbstractMonad(input_node)) {
+      // set const input to primitive attr and erase original const input
       auto value_node = input_node->cast<ValueNodePtr>();
       MS_EXCEPTION_IF_NULL(value_node);
       MS_LOG(DEBUG) << "start erase input[" << i << "] of cnode[" + cnode->DebugString() + "]";

mindspore/ccsrc/backend/common/optimizer/node_pass.cc

@@ -31,10 +31,22 @@ const size_t kSwitchBranchIndex = 2;
 const size_t kCallArgsIndex = 1;
 const size_t kPartialArgsIndex = 1;
 
+/**
+ * @brief Maps the output of a subgraph to its caller node in the main graph.
+ * 
+ * This function checks if the given node (`cnode`) is of type `switch` or `call` and then maps
+ * the output of the corresponding subgraph to the `cnode`. This is useful for tracing back 
+ * from a subgraph output to its caller in the main graph.
+ * 
+ * @param cnode The computation node being checked.
+ * @param out_caller_map Pointer to a map that will hold the relationship between subgraph outputs and their caller nodes.
+ */
 void AddOutputAndCallerToMap(const CNodePtr &cnode, mindspore::HashMap<AnfNodePtr, AnfNodePtr> *out_caller_map) {
+  // Ensure the provided node and map are not null.
   MS_EXCEPTION_IF_NULL(cnode);
   MS_EXCEPTION_IF_NULL(out_caller_map);
   auto inputs = cnode->inputs();
+  // Check if the node is a 'switch' type.
   if (common::AnfAlgo::CheckPrimitiveType(cnode, prim::kPrimSwitch)) {
     auto partial_node = dyn_cast<CNode>(inputs.at(kSwitchBranchIndex));
     MS_EXCEPTION_IF_NULL(partial_node);
@@ -45,6 +57,7 @@ void AddOutputAndCallerToMap(const CNodePtr &cnode, mindspore::HashMap<AnfNodePt
     auto switch_subgraph = GetValueNode<FuncGraphPtr>(partial_inputs.at(kPartialArgsIndex));
     MS_EXCEPTION_IF_NULL(switch_subgraph);
     (*out_caller_map)[switch_subgraph->output()] = cnode;
+  // Check if the node is a 'call' type.
   } else if (common::AnfAlgo::CheckPrimitiveType(cnode, prim::kPrimCall)) {
     auto call_subgraph = GetValueNode<FuncGraphPtr>(inputs.at(kCallArgsIndex));
     MS_EXCEPTION_IF_NULL(call_subgraph);
@@ -52,16 +65,32 @@ void AddOutputAndCallerToMap(const CNodePtr &cnode, mindspore::HashMap<AnfNodePt
   }
 }
 
+/**
+ * @brief Executes the node pass over the given function graph.
+ * 
+ * This function iteratively processes nodes in the `func_graph` and applies transformations, 
+ * possibly optimizing or modifying the graph. The changes are tracked, and the function 
+ * returns whether the graph was modified or not.
+ * 
+ * In the context of MindSpore, this function aids in intermediate optimizations and transformations 
+ * during the model compilation phase.
+ * 
+ * @param func_graph The function graph (or computational graph) to be processed.
+ * @return A boolean indicating whether the graph was changed during processing.
+ */
 bool NodePass::Run(const FuncGraphPtr &func_graph) {
+  // Ensure the function graph is not null.
   MS_EXCEPTION_IF_NULL(func_graph);
   FuncGraphManagerPtr manager = func_graph->manager();
   MS_EXCEPTION_IF_NULL(manager);
   manager->AddFuncGraph(func_graph);
 
+  // Initializations
   mindspore::HashMap<AnfNodePtr, AnfNodePtr> subgraph_out_caller_map = {};
   mindspore::HashSet<AnfNodePtr> seen_node;
   std::deque<std::pair<AnfNodePtr, FuncGraphPtr>> todo{{func_graph->output(), func_graph}};
   bool changes = false;
+  // Main loop to process each node in the graph
   while (!todo.empty()) {
     AnfNodePtr node = todo.front().first;
     auto fg = todo.front().second;

mindspore/ccsrc/backend/common/optimizer/optimizer.cc

@@ -44,12 +44,24 @@ void PatternProcessPass::Build() {
   pattern_ = SexpToNode(DefinePattern(), fg, primitive_vars_.get(), multigraph_);
 }
 
+/**
+ * @brief Matches a node in the computation graph against a predefined pattern and processes it accordingly.
+ * 
+ * This method attempts to match a given node against a specific pattern. If a match is found, it processes 
+ * the node and returns the result.
+ *
+ * @param func_graph The computation graph containing the node.
+ * @param node The node in the computation graph to match against the pattern.
+ * @return Processed node if a match is found; otherwise, nullptr.
+ */
 AnfNodePtr PatternProcessPass::Run(const FuncGraphPtr &func_graph, const AnfNodePtr &node) {
+  // Ensure the pattern is built before matching.
   if (pattern_ == nullptr) {
     Build();
   }
 
   auto primitive = GetCNodePrimitive(pattern_);
+  // Check if the node matches the primitive pattern.
   if (IsPrimitiveCNode(node, primitive)) {
     MS_EXCEPTION_IF_NULL(primitive_vars_);
     MS_EXCEPTION_IF_NULL(equiv_);
@@ -62,6 +74,11 @@ AnfNodePtr PatternProcessPass::Run(const FuncGraphPtr &func_graph, const AnfNode
   return nullptr;
 }
 
+/**
+ * @brief Retrieves the original nodes that match the primitive variables.
+ * 
+ * @return A vector containing the original nodes matched against the primitive variables.
+ */
 std::vector<AnfNodePtr> PatternProcessPass::GetOrigNodes() const {
   std::vector<AnfNodePtr> orig_nodes;
   for (auto &prim_var : *primitive_vars_) {
@@ -78,18 +95,41 @@ std::vector<AnfNodePtr> PatternProcessPass::GetOrigNodes() const {
   return orig_nodes;
 }
 
+/**
+ * @brief Creates a new CNode with the given inputs and attaches it to the provided computation graph.
+ *
+ * @param inputs The input nodes for the new CNode.
+ * @param fg The computation graph to which the new CNode should be added.
+ * @return The newly created CNode.
+ */
 CNodePtr PatternProcessPass::NewCNode(const std::vector<AnfNodePtr> &inputs, const FuncGraphPtr &fg) const {
   MS_EXCEPTION_IF_NULL(fg);
   auto orig_nodes = GetOrigNodes();
   return opt::NewCNode(inputs, fg, orig_nodes);
 }
 
+/**
+ * @brief Creates a new CNode by copying an existing one and attaches it to the provided kernel graph.
+ *
+ * @param cnode The CNode to copy.
+ * @param fg The kernel graph to which the new CNode should be added.
+ * @return The newly created CNode.
+ */
 CNodePtr PatternProcessPass::NewCNode(const CNodePtr &cnode, const KernelGraphPtr &fg) const {
   MS_EXCEPTION_IF_NULL(fg);
   auto orig_nodes = GetOrigNodes();
   return opt::NewCNode(cnode, fg, orig_nodes);
 }
 
+/**
+ * @brief Matches a node against another predefined pattern.
+ *
+ * This method checks if a node matches another specific pattern, distinct from the primary pattern.
+ *
+ * @param node The node in the computation graph to match against the secondary pattern.
+ * @param equiv The equivalence mapping of nodes to variables in the primary pattern.
+ * @return True if the node matches the secondary pattern; otherwise, false.
+ */
 bool MultipleOutputPatternProcessPass::MatchAnotherPattern(const AnfNodePtr &node, const EquivPtr &equiv) const {
   MS_EXCEPTION_IF_NULL(node);
   MS_EXCEPTION_IF_NULL(equiv);
@@ -105,6 +145,11 @@ bool MultipleOutputPatternProcessPass::MatchAnotherPattern(const AnfNodePtr &nod
   return false;
 }
 
+/**
+ * @brief Retrieves the original nodes that match both the primary and secondary primitive variables.
+ * 
+ * @return A vector containing the original nodes matched against both sets of primitive variables.
+ */
 std::vector<AnfNodePtr> MultipleOutputPatternProcessPass::GetOrigNodes() const {
   std::vector<AnfNodePtr> orig_nodes = PatternProcessPass::GetOrigNodes();
   for (auto &prim_var : *child_primitive_vars_) {
@@ -118,35 +163,75 @@ std::vector<AnfNodePtr> MultipleOutputPatternProcessPass::GetOrigNodes() const {
   return orig_nodes;
 }
 
+/**
+ * @brief Adds a pass manager to the graph optimizer's list of pass managers.
+ *
+ * A pass manager contains a set of transformation passes to be applied to the computation graph.
+ *
+ * @param pass_manager The pass manager to be added.
+ */
 void GraphOptimizer::AddPassManager(const PassManagerPtr &pass_manager) {
   if (pass_manager != nullptr) {
     pass_managers_.push_back(pass_manager);
   }
 }
 
+/**
+ * @brief Optimizes a given computation graph by applying various transformation passes.
+ *
+ * This method applies a set of transformation passes to the given computation graph, seeking to 
+ * optimize its structure and operations. These transformations can include tasks like constant 
+ * folding, operation fusion, or simplification. The goal is to generate a more efficient graph 
+ * without changing the original semantics. This function can either apply all the transformation 
+ * passes once or multiple times based on the `run_only_once` parameter.
+ *
+ * @param func_graph The computation graph that needs to be optimized. Represented as a pointer 
+ *                   to a FuncGraph.
+ * @param run_only_once If true, each transformation pass is applied only once. Otherwise, the 
+ *                      passes may be applied multiple times until no further changes are observed 
+ *                      in the graph.
+ * @return Returns the optimized computation graph.
+ */
 FuncGraphPtr GraphOptimizer::Optimize(const FuncGraphPtr &func_graph, bool run_only_once) {
+  // Check if the input computation graph is null. If so, raise an exception.
   MS_EXCEPTION_IF_NULL(func_graph);
+
+  // Determine if we should run the optimization only once based on the number of pass managers.
   run_only_once_ = (pass_managers_.size() == 1) ? true : run_only_once;
-  // cppcheck-suppress *
+
+  // Create or retrieve a manager for the computation graph.
+  // The manager is responsible for handling various tasks like graph traversal, node replacement,
+  // and overall graph management.
   auto manager = Manage(func_graph, true);
 
   bool changed = true;
+  // Loop until no further changes are observed in the graph.
   while (changed) {
     changed = false;
+    // Iterate through all the pass managers and run their associated transformation passes 
+    // on the computation graph.
     for (size_t i = 0; i < pass_managers_.size(); ++i) {
       const PassManagerPtr &pm = pass_managers_[i];
+      // If a transformation results in a change to the graph, set the `changed` flag to true.
       if (pm != nullptr && pm->Run(func_graph)) {
         changed = true;
       }
     }
+    // If we're only supposed to run the optimization once, break out of the loop.
     if (run_only_once_) {
       break;
     }
   }
 
+  // List to hold the collection of computation graphs.
   std::vector<FuncGraphPtr> func_graphs;
   func_graphs.push_back(func_graph);
+
+  // Perform a topological sort on the computation graph nodes starting from the return node.
+  // This ensures that the graph nodes are arranged in a specific order based on their dependencies.
   (void)TopoSort(func_graph->get_return());
+
+  // Return the optimized computation graph.
   return func_graph;
 }
 }  // namespace opt

mindspore/ccsrc/backend/common/pass/add_akg_kernel_attrs.cc

@@ -25,57 +25,104 @@
 namespace mindspore {
 namespace opt {
 namespace {
+/**
+ * @brief Clones the primitive of the given node.
+ * 
+ * In scenarios where multiple CNodes may share a primitive pointer, 
+ * this function ensures that each CNode gets its own cloned primitive.
+ * 
+ * @param node The AnfNode which has the primitive to be cloned.
+ */
 void ClonePrimitive(const AnfNodePtr &node) {
-  // Several CNode may share a primitive pointer, so we clone the primitive before setting attr.
   auto cnode = node->cast<CNodePtr>();
   if (cnode == nullptr) return;
+  
+  // Clone the primitive and set it back to the CNode.
   auto prim_node = NewValueNode(common::AnfAlgo::GetCNodePrimitive(cnode)->Clone());
   cnode->set_input(kAnfPrimitiveIndex, prim_node);
 }
-}  // namespace
 
+namespace {
+
+/**
+ * @brief Processes the Cast operation node to set appropriate AKG kernel attributes.
+ * 
+ * @param node The AnfNode representing the Cast operation.
+ */
 void ProcessCast(const AnfNodePtr &node) {
-  // The x and output are akg op input and output param.
+  ClonePrimitive(node);
   std::vector<std::string> input_names = {"x", kAttrDstType};
   std::vector<std::string> output_names = {"output"};
-  ClonePrimitive(node);
+  
+  // Set the input and output names as attributes.
   common::AnfAlgo::SetNodeAttr(kAttrInputNames, MakeValue(input_names), node);
   common::AnfAlgo::SetNodeAttr(kAttrOutputNames, MakeValue(output_names), node);
+  
+  // Set the destination type attribute.
   TypeId output_type = AnfAlgo::GetOutputDeviceDataType(node, 0);
   common::AnfAlgo::SetNodeAttr(kAttrDstType, TypeIdToType(output_type), node);
 }
 
+/**
+ * @brief Processes the MatMul operation node to set appropriate AKG kernel attributes.
+ * 
+ * @param node The AnfNode representing the MatMul operation.
+ */
 void ProcessMatMul(const AnfNodePtr &node) {
   ClonePrimitive(node);
+  
+  // Set the destination type attribute.
   TypeId output_type = AnfAlgo::GetOutputDeviceDataType(node, 0);
   common::AnfAlgo::SetNodeAttr(kAttrDstType, TypeIdToType(output_type), node);
+  
+  // Set the left and right format attributes.
   auto left_format = AnfAlgo::GetInputFormat(node, 0);
   auto right_format = AnfAlgo::GetInputFormat(node, 1);
   common::AnfAlgo::SetNodeAttr("left_format", MakeValue(left_format), node);
   common::AnfAlgo::SetNodeAttr("right_format", MakeValue(right_format), node);
 }
 
+/**
+ * @brief Process nodes in the FuncGraph to set appropriate AKG kernel attributes.
+ * 
+ * @param func_graph The function graph being processed.
+ * @param node The node in the graph to be processed.
+ * @param equiv Not used in the current implementation.
+ * @return AnfNodePtr Modified node if changes were made, otherwise nullptr.
+ */
 const AnfNodePtr AddAkgKernelAttrs::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
                                             const EquivPtr &) const {
   MS_EXCEPTION_IF_NULL(func_graph);
   MS_EXCEPTION_IF_NULL(node);
   auto shape = node->Shape();
-  // dynamic shape nodes will re-infer the shape and dtype
+  
+  // Nodes with dynamic shape will re-infer the shape and dtype.
   if (shape == nullptr || shape->IsDynamic()) {
     return nullptr;
   }
+
+  // Process Cast and MatMul nodes.
   if (IsPrimitiveCNode(node, prim::kPrimCast)) {
     ProcessCast(node);
   } else if (IsPrimitiveCNode(node, prim::kPrimMatMul) || IsPrimitiveCNode(node, prim::kPrimBatchMatMul)) {
     ProcessMatMul(node);
   }
+
   return nullptr;
 }
 
+/**
+ * @brief Define the pattern to be matched in the graph.
+ * 
+ * @return BaseRef Pattern to be matched.
+ */
 const BaseRef AddAkgKernelAttrs::DefinePattern() const {
   VarPtr X = std::make_shared<Var>();
   VarPtr Xs = std::make_shared<SeqVar>();
   return VectorRef({X, Xs});
 }
+
+}  // namespace
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/add_akg_kernel_attrs_old.cc

@@ -0,0 +1,128 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "backend/common/pass/add_akg_kernel_attrs.h"
+#include <memory>
+#include <vector>
+#include <string>
+#include "base/core_ops.h"
+#include "backend/common/session/anf_runtime_algorithm.h"
+#include "include/common/utils/anfalgo.h"
+
+namespace mindspore {
+namespace opt {
+namespace {
+/**
+ * @brief Clones the primitive of the given node.
+ * 
+ * In scenarios where multiple CNodes may share a primitive pointer, 
+ * this function ensures that each CNode gets its own cloned primitive.
+ * 
+ * @param node The AnfNode which has the primitive to be cloned.
+ */
+void ClonePrimitive(const AnfNodePtr &node) {
+  auto cnode = node->cast<CNodePtr>();
+  if (cnode == nullptr) return;
+  
+  // Clone the primitive and set it back to the CNode.
+  auto prim_node = NewValueNode(common::AnfAlgo::GetCNodePrimitive(cnode)->Clone());
+  cnode->set_input(kAnfPrimitiveIndex, prim_node);
+}
+
+namespace {
+
+/**
+ * @brief Processes the Cast operation node to set appropriate AKG kernel attributes.
+ * 
+ * @param node The AnfNode representing the Cast operation.
+ */
+void ProcessCast(const AnfNodePtr &node) {
+  ClonePrimitive(node);
+  std::vector<std::string> input_names = {"x", kAttrDstType};
+  std::vector<std::string> output_names = {"output"};
+  
+  // Set the input and output names as attributes.
+  common::AnfAlgo::SetNodeAttr(kAttrInputNames, MakeValue(input_names), node);
+  common::AnfAlgo::SetNodeAttr(kAttrOutputNames, MakeValue(output_names), node);
+  
+  // Set the destination type attribute.
+  TypeId output_type = AnfAlgo::GetOutputDeviceDataType(node, 0);
+  common::AnfAlgo::SetNodeAttr(kAttrDstType, TypeIdToType(output_type), node);
+}
+
+/**
+ * @brief Processes the MatMul operation node to set appropriate AKG kernel attributes.
+ * 
+ * @param node The AnfNode representing the MatMul operation.
+ */
+void ProcessMatMul(const AnfNodePtr &node) {
+  ClonePrimitive(node);
+  
+  // Set the destination type attribute.
+  TypeId output_type = AnfAlgo::GetOutputDeviceDataType(node, 0);
+  common::AnfAlgo::SetNodeAttr(kAttrDstType, TypeIdToType(output_type), node);
+  
+  // Set the left and right format attributes.
+  auto left_format = AnfAlgo::GetInputFormat(node, 0);
+  auto right_format = AnfAlgo::GetInputFormat(node, 1);
+  common::AnfAlgo::SetNodeAttr("left_format", MakeValue(left_format), node);
+  common::AnfAlgo::SetNodeAttr("right_format", MakeValue(right_format), node);
+}
+
+/**
+ * @brief Process nodes in the FuncGraph to set appropriate AKG kernel attributes.
+ * 
+ * @param func_graph The function graph being processed.
+ * @param node The node in the graph to be processed.
+ * @param equiv Not used in the current implementation.
+ * @return AnfNodePtr Modified node if changes were made, otherwise nullptr.
+ */
+const AnfNodePtr AddAkgKernelAttrs::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
+                                            const EquivPtr &) const {
+  MS_EXCEPTION_IF_NULL(func_graph);
+  MS_EXCEPTION_IF_NULL(node);
+  auto shape = node->Shape();
+  
+  // Nodes with dynamic shape will re-infer the shape and dtype.
+  if (shape == nullptr || shape->IsDynamic()) {
+    return nullptr;
+  }
+
+  // Process Cast and MatMul nodes.
+  if (IsPrimitiveCNode(node, prim::kPrimCast)) {
+    ProcessCast(node);
+  } else if (IsPrimitiveCNode(node, prim::kPrimMatMul) || IsPrimitiveCNode(node, prim::kPrimBatchMatMul)) {
+    ProcessMatMul(node);
+  }
+
+  return nullptr;
+}
+
+/**
+ * @brief Define the pattern to be matched in the graph.
+ * 
+ * @return BaseRef Pattern to be matched.
+ */
+const BaseRef AddAkgKernelAttrs::DefinePattern() const {
+  VarPtr X = std::make_shared<Var>();
+  VarPtr Xs = std::make_shared<SeqVar>();
+  return VectorRef({X, Xs});
+}
+
+}  // namespace
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/add_dynamic_shape_attr.cc

@@ -21,17 +21,33 @@
 
 namespace mindspore {
 namespace opt {
+/**
+ * @brief Process nodes in the FuncGraph to set appropriate dynamic shape attributes.
+ * 
+ * If the given node has dynamic shape attributes, the function will log 
+ * the node and set the graph's dynamic attribute to true.
+ * 
+ * @param func_graph The function graph being processed.
+ * @param node The node in the graph to be processed.
+ * @param equiv Not used in the current implementation.
+ * @return AnfNodePtr The input node as it's not modified within this method.
+ */
 const AnfNodePtr AddDynamicShapeAttr::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
                                               const EquivPtr &) const {
   MS_EXCEPTION_IF_NULL(func_graph);
   MS_EXCEPTION_IF_NULL(node);
+  
+  // If the node has dynamic shape, set the dynamic attribute of the graph to true.
   if (common::AnfAlgo::IsDynamicShape(node)) {
     MS_LOG(DEBUG) << "Set Dynamic Shape Attr to Node:" << node->fullname_with_scope();
+    
     auto kernel_graph = func_graph->cast<KernelGraphPtr>();
     MS_EXCEPTION_IF_NULL(kernel_graph);
     kernel_graph->SetGraphDynamicAttr(true);
   }
+
   return node;
 }
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/add_dynamic_shape_attr_old.cc

@@ -0,0 +1,53 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "backend/common/pass/add_dynamic_shape_attr.h"
+#include "ir/anf.h"
+#include "backend/common/optimizer/optimizer.h"
+#include "include/common/utils/anfalgo.h"
+
+namespace mindspore {
+namespace opt {
+/**
+ * @brief Process nodes in the FuncGraph to set appropriate dynamic shape attributes.
+ * 
+ * If the given node has dynamic shape attributes, the function will log 
+ * the node and set the graph's dynamic attribute to true.
+ * 
+ * @param func_graph The function graph being processed.
+ * @param node The node in the graph to be processed.
+ * @param equiv Not used in the current implementation.
+ * @return AnfNodePtr The input node as it's not modified within this method.
+ */
+const AnfNodePtr AddDynamicShapeAttr::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
+                                              const EquivPtr &) const {
+  MS_EXCEPTION_IF_NULL(func_graph);
+  MS_EXCEPTION_IF_NULL(node);
+  
+  // If the node has dynamic shape, set the dynamic attribute of the graph to true.
+  if (common::AnfAlgo::IsDynamicShape(node)) {
+    MS_LOG(DEBUG) << "Set Dynamic Shape Attr to Node:" << node->fullname_with_scope();
+    
+    auto kernel_graph = func_graph->cast<KernelGraphPtr>();
+    MS_EXCEPTION_IF_NULL(kernel_graph);
+    kernel_graph->SetGraphDynamicAttr(true);
+  }
+
+  return node;
+}
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/add_training_attr.cc

@@ -29,19 +29,35 @@
 namespace mindspore {
 namespace opt {
 namespace {
+// Define the operations and their respective sets for marking.
 mindspore::HashMap<std::string, mindspore::HashSet<std::string>> MarkOp{
-  {"LSTM", {"LSTMGradWeight", "LSTMGrad", "LSTMGradData"}}};
+  {"LSTM", {"LSTMGradWeight", "LSTMGrad", "LSTMGradData"}}
+};
 
+/**
+ * @brief Checks if any user of the given node is within the provided set of operations.
+ * 
+ * This method recursively traverses the users of the node and checks if any of them belong to 
+ * the provided set. This helps in identifying certain patterns or connections in the graph.
+ * 
+ * @param manager The FuncGraphManager for managing nodes in the graph.
+ * @param cnode Node to check for its users.
+ * @param set The set of operation names to check against.
+ * @return bool Indicates whether any user of the node is in the set.
+ */
 bool CheckOP(const FuncGraphManagerPtr &manager, const AnfNodePtr &cnode, const mindspore::HashSet<std::string> &set) {
   for (const auto &node_index : manager->node_users()[cnode]) {
     auto output = node_index.first;
     MS_EXCEPTION_IF_NULL(output);
+
+    // If the user is a TupleGetItem, recursively check its users.
     if (common::AnfAlgo::CheckPrimitiveType(output, prim::kPrimTupleGetItem)) {
       if (CheckOP(manager, output, set)) {
         return true;
       }
     } else if (output->isa<CNode>()) {
       auto name = common::AnfAlgo::GetCNodeName(output);
+      // If the user's name is in the set, return true.
       if (set.find(name) != set.end()) {
         return true;
       }
@@ -49,23 +65,47 @@ bool CheckOP(const FuncGraphManagerPtr &manager, const AnfNodePtr &cnode, const
   }
   return false;
 }
+
+/**
+ * @brief Adds a training attribute to the given CNode based on its connections.
+ * 
+ * The method checks if the CNode is connected to certain operations and based on the check 
+ * result, it sets the 'kAttrIsTraining' attribute on the node.
+ * 
+ * @param func_graph The function graph containing the node.
+ * @param cnode The CNode to which the attribute should be added.
+ */
 void AddAttrTraining(const FuncGraphPtr &func_graph, const CNodePtr &cnode) {
   MS_EXCEPTION_IF_NULL(func_graph);
   MS_EXCEPTION_IF_NULL(cnode);
   auto manager = func_graph->manager();
   MS_EXCEPTION_IF_NULL(manager);
+
+  // If the node has no users, exit.
   if (manager->node_users().find(cnode) == manager->node_users().end()) {
     return;
   }
+
   auto set = MarkOp[common::AnfAlgo::GetCNodeName(cnode)];
   if (CheckOP(manager, cnode, set)) {
+    // Add 'IsTraining' attribute with a value of 'true'.
     cnode->AddAttr(kAttrIsTraining, MakeValue(true));
   } else {
+    // Add 'IsTraining' attribute with a value of 'false'.
     cnode->AddAttr(kAttrIsTraining, MakeValue(false));
   }
 }
-}  // namespace
 
+namespace {
+
+/**
+ * @brief Processes the given node in the function graph and adds training attributes if necessary.
+ * 
+ * @param func_graph The function graph being processed.
+ * @param node The AnfNode being examined.
+ * @param equiv Not used in the current implementation.
+ * @return AnfNodePtr Returns the modified node or nullptr if no changes are made.
+ */
 const AnfNodePtr AddTrainingAttr::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
                                           const EquivPtr &) const {
   if (node == nullptr || func_graph == nullptr || common::AnfAlgo::CheckPrimitiveType(node, prim::kPrimTupleGetItem) ||
@@ -75,14 +115,23 @@ const AnfNodePtr AddTrainingAttr::Process(const FuncGraphPtr &func_graph, const
   if (!node->isa<CNode>()) {
     return nullptr;
   }
+
+  // Get the name of the CNode.
   auto name = common::AnfAlgo::GetCNodeName(node);
+
+  // If the name isn't in the MarkOp, exit.
   auto iter = MarkOp.find(name);
   if (iter == MarkOp.end()) {
     return nullptr;
   }
+
+  // Add training attributes and return the modified node.
   auto cnode = node->cast<CNodePtr>();
   AddAttrTraining(func_graph, cnode);
   return cnode;
 }
+
+}  // namespace
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/add_training_attr_old.cc

@@ -0,0 +1,137 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "backend/common/pass/add_training_attr.h"
+
+#include <vector>
+#include <memory>
+#include <utility>
+
+#include "utils/hash_map.h"
+#include "utils/hash_set.h"
+#include "ir/graph_utils.h"
+#include "backend/common/optimizer/helper.h"
+#include "include/common/utils/anfalgo.h"
+
+namespace mindspore {
+namespace opt {
+namespace {
+// Define the operations and their respective sets for marking.
+mindspore::HashMap<std::string, mindspore::HashSet<std::string>> MarkOp{
+  {"LSTM", {"LSTMGradWeight", "LSTMGrad", "LSTMGradData"}}
+};
+
+/**
+ * @brief Checks if any user of the given node is within the provided set of operations.
+ * 
+ * This method recursively traverses the users of the node and checks if any of them belong to 
+ * the provided set. This helps in identifying certain patterns or connections in the graph.
+ * 
+ * @param manager The FuncGraphManager for managing nodes in the graph.
+ * @param cnode Node to check for its users.
+ * @param set The set of operation names to check against.
+ * @return bool Indicates whether any user of the node is in the set.
+ */
+bool CheckOP(const FuncGraphManagerPtr &manager, const AnfNodePtr &cnode, const mindspore::HashSet<std::string> &set) {
+  for (const auto &node_index : manager->node_users()[cnode]) {
+    auto output = node_index.first;
+    MS_EXCEPTION_IF_NULL(output);
+
+    // If the user is a TupleGetItem, recursively check its users.
+    if (common::AnfAlgo::CheckPrimitiveType(output, prim::kPrimTupleGetItem)) {
+      if (CheckOP(manager, output, set)) {
+        return true;
+      }
+    } else if (output->isa<CNode>()) {
+      auto name = common::AnfAlgo::GetCNodeName(output);
+      // If the user's name is in the set, return true.
+      if (set.find(name) != set.end()) {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
+/**
+ * @brief Adds a training attribute to the given CNode based on its connections.
+ * 
+ * The method checks if the CNode is connected to certain operations and based on the check 
+ * result, it sets the 'kAttrIsTraining' attribute on the node.
+ * 
+ * @param func_graph The function graph containing the node.
+ * @param cnode The CNode to which the attribute should be added.
+ */
+void AddAttrTraining(const FuncGraphPtr &func_graph, const CNodePtr &cnode) {
+  MS_EXCEPTION_IF_NULL(func_graph);
+  MS_EXCEPTION_IF_NULL(cnode);
+  auto manager = func_graph->manager();
+  MS_EXCEPTION_IF_NULL(manager);
+
+  // If the node has no users, exit.
+  if (manager->node_users().find(cnode) == manager->node_users().end()) {
+    return;
+  }
+
+  auto set = MarkOp[common::AnfAlgo::GetCNodeName(cnode)];
+  if (CheckOP(manager, cnode, set)) {
+    // Add 'IsTraining' attribute with a value of 'true'.
+    cnode->AddAttr(kAttrIsTraining, MakeValue(true));
+  } else {
+    // Add 'IsTraining' attribute with a value of 'false'.
+    cnode->AddAttr(kAttrIsTraining, MakeValue(false));
+  }
+}
+
+namespace {
+
+/**
+ * @brief Processes the given node in the function graph and adds training attributes if necessary.
+ * 
+ * @param func_graph The function graph being processed.
+ * @param node The AnfNode being examined.
+ * @param equiv Not used in the current implementation.
+ * @return AnfNodePtr Returns the modified node or nullptr if no changes are made.
+ */
+const AnfNodePtr AddTrainingAttr::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
+                                          const EquivPtr &) const {
+  if (node == nullptr || func_graph == nullptr || common::AnfAlgo::CheckPrimitiveType(node, prim::kPrimTupleGetItem) ||
+      common::AnfAlgo::CheckPrimitiveType(node, prim::kPrimMakeTuple)) {
+    return nullptr;
+  }
+  if (!node->isa<CNode>()) {
+    return nullptr;
+  }
+
+  // Get the name of the CNode.
+  auto name = common::AnfAlgo::GetCNodeName(node);
+
+  // If the name isn't in the MarkOp, exit.
+  auto iter = MarkOp.find(name);
+  if (iter == MarkOp.end()) {
+    return nullptr;
+  }
+
+  // Add training attributes and return the modified node.
+  auto cnode = node->cast<CNodePtr>();
+  AddAttrTraining(func_graph, cnode);
+  return cnode;
+}
+
+}  // namespace
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/adjust_depend_for_parallel_optimizer_recompute_all_gather.cc

@@ -21,36 +21,63 @@
 namespace mindspore {
 namespace opt {
 constexpr const int64_t kFusionGap = 2;
+/**
+ * @brief Adjusts the dependencies for parallel optimizer recompute AllGather operations within the function graph.
+ * 
+ * This method processes the function graph, identifies the AllGather operations that are subject 
+ * to recompute optimization and adjusts their dependencies to ensure proper computation sequence.
+ * The dependencies of AllGather operations might need adjustment due to parallel optimization strategies.
+ * 
+ * @param graph The function graph to process.
+ * @return bool Indicates whether any adjustments were made to the function graph.
+ */
 bool AdjustDependForParallelOptimizerRecomputeAllGather::Run(const FuncGraphPtr &graph) {
+  // Validate the input function graph.
   MS_EXCEPTION_IF_NULL(graph);
+  
+  // A mapping to determine if an AllGather operation within a fusion group is set for recompute.
   mindspore::HashMap<int64_t, bool> forward_allgather_recompute_value_in_fusion_group;
+  
+  // Get the nodes of the function graph in topological order.
   std::vector<AnfNodePtr> node_list = TopoSort(graph->get_return());
+
+  // Variables to keep track of the AllGather operations and their attributes.
   std::vector<int64_t> parallel_optimizer_recompute_allgather_fusion_ids;
   std::vector<AnfNodePtr> parallel_optimizer_recompute_allgathers;
   std::vector<AnfNodePtr> parallel_optimizer_recompute_first_fusion_allgathers;
   int64_t unrecompute_max_fusion_id = -1;
   int64_t recompute_min_fusion_id = 0;
+
+  // Process each node to identify the AllGather operations subject to recompute.
   for (auto &node : node_list) {
     MS_EXCEPTION_IF_NULL(node);
+    
+    // Filter out non-kernel nodes.
     if (!node->cast<CNodePtr>() || !AnfUtils::IsRealKernel(node)) {
       continue;
     }
     auto cnode = node->cast<CNodePtr>();
+
+    // Filter out nodes that aren't AllGather or associated with fusion and parallel optimizer.
     if (!common::AnfAlgo::IsAllgather(cnode) || !common::AnfAlgo::IsFusion(cnode) ||
         !common::AnfAlgo::IsFromParallelOptimizer(cnode)) {
       continue;
     }
+
+    // If the node is set for recompute, update the fusion IDs and track it.
     if (common::AnfAlgo::IsRecompute(cnode)) {
       int64_t fusion_id = common::AnfAlgo::GetNodeAttr<int64_t>(cnode, kAttrFusion);
+      // Check and add unique fusion_ids.
       if (std::find(parallel_optimizer_recompute_allgather_fusion_ids.begin(),
                     parallel_optimizer_recompute_allgather_fusion_ids.end(),
                     fusion_id) == parallel_optimizer_recompute_allgather_fusion_ids.end()) {
+        // If the fusion id is not in the vector, it means that the allgather node is the first allgather node in the
+        // fusion group. Leave the fusion group alone till another allgather node in the same fusion group is found.
         parallel_optimizer_recompute_allgather_fusion_ids.push_back(fusion_id);
-        if (recompute_min_fusion_id == 0 || fusion_id < recompute_min_fusion_id) {
-          recompute_min_fusion_id = fusion_id;
-        }
+        recompute_min_fusion_id = recompute_min_fusion_id == 0 ? fusion_id : std::min(fusion_id, recompute_min_fusion_id);
         parallel_optimizer_recompute_first_fusion_allgathers.push_back(node);
       } else {
+        // Now here's another allgather node in the same fusion group. Handle it.
         parallel_optimizer_recompute_allgathers.push_back(node);
       }
     } else {
@@ -60,57 +87,105 @@ bool AdjustDependForParallelOptimizerRecomputeAllGather::Run(const FuncGraphPtr
                                  common::AnfAlgo::GetNodeAttr<bool>(cnode, kAttrRecompute);
       auto [iter, inserted] =
         forward_allgather_recompute_value_in_fusion_group.emplace(unrecompute_fusion_id, would_be_recomputed);
+      
+      // Ensure consistency within the fusion group.
       if (!inserted && iter->second != would_be_recomputed) {
-        MS_LOG(EXCEPTION) << "In same fusion group, the allgather recompute attribute should be equal. "
+        MS_LOG(EXCEPTION) << "In the same fusion group, the AllGather recompute attribute should be equal. "
                              "The normal node is:"
                           << cnode->fullname_with_scope();
       }
     }
   }
+
+  // Adjust the fusion IDs for the identified AllGather operations.
   IncreaseAllgatherFusionId(parallel_optimizer_recompute_allgathers,
                             parallel_optimizer_recompute_first_fusion_allgathers, unrecompute_max_fusion_id,
                             recompute_min_fusion_id);
+
+  // Adjust the dependencies for the AllGather operations.
   return AdjustAllgatherDepend(graph, parallel_optimizer_recompute_allgathers);
 }
 
+
+/**
+ * @brief Increases the fusion ID of AllGather operations to adjust for recompute optimization.
+ * 
+ * The purpose of this method is to handle cases where duplicated AllGather operations would be fused with forward 
+ * AllGather operations. In such cases, fusion IDs need adjustment to ensure proper execution.
+ * 
+ * @param parallel_optimizer_recompute_allgathers All the AllGather operations marked for recompute.
+ * @param parallel_optimizer_recompute_first_fusion_allgathers First AllGather operations in the recomputed segments.
+ * @param unrecompute_max_fusion_id Maximum fusion ID of unrecomputed AllGather operations.
+ * @param recompute_min_fusion_id Minimum fusion ID of recomputed AllGather operations.
+ */
 void AdjustDependForParallelOptimizerRecomputeAllGather::IncreaseAllgatherFusionId(
   const std::vector<AnfNodePtr> &parallel_optimizer_recompute_allgathers,
   const std::vector<AnfNodePtr> &parallel_optimizer_recompute_first_fusion_allgathers,
   int64_t unrecompute_max_fusion_id, int64_t recompute_min_fusion_id) {
-  // means that there may some forward allgather and duplicated allgather would be fused.
+  
+  // If the condition holds, some forward AllGather and duplicated AllGather may be fused.
   if (recompute_min_fusion_id <= unrecompute_max_fusion_id) {
     MS_LOG(WARNING) << "Increase the duplicated allgather fusion id";
+    
+    // Adjust fusion ID for the first AllGather operations in recomputed segments.
     for (auto &adjust_node : parallel_optimizer_recompute_first_fusion_allgathers) {
+      // Calculate the new fusion ID.
       int64_t current_fusion_id = common::AnfAlgo::GetNodeAttr<int64_t>(adjust_node, kAttrFusion);
       int64_t destination_fusion_id =
         (kFusionGap + current_fusion_id + unrecompute_max_fusion_id) - recompute_min_fusion_id;
+      // Set the new fusion ID.
       common::AnfAlgo::SetNodeAttr(kAttrFusion, MakeValue(destination_fusion_id), adjust_node);
     }
+
+    // Adjust fusion ID for the remaining AllGather operations marked for recompute.
     for (auto &adjust_node : parallel_optimizer_recompute_allgathers) {
+      // Calculate the new fusion ID.
       int64_t current_fusion_id = common::AnfAlgo::GetNodeAttr<int64_t>(adjust_node, kAttrFusion);
       int64_t destination_fusion_id =
         (kFusionGap + current_fusion_id + unrecompute_max_fusion_id) - recompute_min_fusion_id;
+      // Set the new fusion ID.
       common::AnfAlgo::SetNodeAttr(kAttrFusion, MakeValue(destination_fusion_id), adjust_node);
     }
   }
 }
 
+/**
+ * @brief Adjusts the dependencies of AllGather operations subject to recompute within the function graph.
+ * 
+ * This method ensures that the execution dependencies of AllGather operations are maintained as expected, 
+ * especially after introducing the recompute optimization.
+ * 
+ * @param graph The function graph to process.
+ * @param parallel_optimizer_recompute_allgathers All the AllGather operations marked for recompute.
+ * @return bool Indicates whether any adjustments were made to the function graph.
+ */
 bool AdjustDependForParallelOptimizerRecomputeAllGather::AdjustAllgatherDepend(
   const FuncGraphPtr &graph, const std::vector<AnfNodePtr> &parallel_optimizer_recompute_allgathers) {
+  
+  // Get the manager for the function graph.
   FuncGraphManagerPtr manager = graph->manager();
   bool changed = false;
+
+  // Process each AllGather operation marked for recompute.
   for (auto &node : parallel_optimizer_recompute_allgathers) {
     auto cnode = node->cast<CNodePtr>();
     auto depend_node = common::AnfAlgo::GetInputNode(cnode, 0);
+
+    // If the dependency is a Depend operation, adjust the dependency.
     if (IsPrimitiveCNode(depend_node, prim::kPrimDepend)) {
+       // depend node is a "depend" primitive
       auto depend_cnode = depend_node->cast<CNodePtr>();
       AnfNodeIndexSet allgather_node_set = manager->node_users()[cnode];
       for (auto &node_pair : allgather_node_set) {
         auto allgather_next_node = node_pair.first;
         CNodePtr allgather_next_cnode = node_pair.first->cast<CNodePtr>();
+
+        // Continue if the node isn't a valid CNode or a primitive.
         if (allgather_next_cnode == nullptr || !IsValueNode<Primitive>(allgather_next_cnode->input(0))) {
           continue;
         }
+
+        // Construct a new Depend operation and adjust the connections.
         std::vector<AnfNodePtr> inputs = {NewValueNode(std::make_shared<Primitive>(prim::kPrimDepend->name())),
                                           allgather_next_node, common::AnfAlgo::GetInputNode(depend_cnode, 1)};
         auto new_depend = graph->NewCNode(inputs);
@@ -119,8 +194,11 @@ bool AdjustDependForParallelOptimizerRecomputeAllGather::AdjustAllgatherDepend(
         (void)manager->Replace(allgather_next_node, new_depend);
         changed = true;
       }
-    } else if (IsPrimitiveCNode(depend_node, prim::kPrimCast) &&
+    } 
+    // Handle the case where the dependency is a Cast operation with an underlying Depend operation.
+    else if (IsPrimitiveCNode(depend_node, prim::kPrimCast) &&
                IsPrimitiveCNode(common::AnfAlgo::GetInputNode(depend_node->cast<CNodePtr>(), 0), prim::kPrimDepend)) {
+      // The logic here mirrors the above block for Depend operations, with an additional layer of handling for Cast.
       auto cast_cnode = depend_node->cast<CNodePtr>();
       auto cast_depend_node = common::AnfAlgo::GetInputNode(cast_cnode, 0);
       auto cast_depend_cnode = cast_depend_node->cast<CNodePtr>();
@@ -145,5 +223,6 @@ bool AdjustDependForParallelOptimizerRecomputeAllGather::AdjustAllgatherDepend(
   }
   return changed;
 }
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/adjust_depend_for_parallel_optimizer_recompute_all_gather_old.cc

@@ -0,0 +1,228 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "backend/common/pass/adjust_depend_for_parallel_optimizer_recompute_all_gather.h"
+#include <algorithm>
+#include "include/common/utils/anfalgo.h"
+
+namespace mindspore {
+namespace opt {
+constexpr const int64_t kFusionGap = 2;
+/**
+ * @brief Adjusts the dependencies for parallel optimizer recompute AllGather operations within the function graph.
+ * 
+ * This method processes the function graph, identifies the AllGather operations that are subject 
+ * to recompute optimization and adjusts their dependencies to ensure proper computation sequence.
+ * The dependencies of AllGather operations might need adjustment due to parallel optimization strategies.
+ * 
+ * @param graph The function graph to process.
+ * @return bool Indicates whether any adjustments were made to the function graph.
+ */
+bool AdjustDependForParallelOptimizerRecomputeAllGather::Run(const FuncGraphPtr &graph) {
+  // Validate the input function graph.
+  MS_EXCEPTION_IF_NULL(graph);
+  
+  // A mapping to determine if an AllGather operation within a fusion group is set for recompute.
+  mindspore::HashMap<int64_t, bool> forward_allgather_recompute_value_in_fusion_group;
+  
+  // Get the nodes of the function graph in topological order.
+  std::vector<AnfNodePtr> node_list = TopoSort(graph->get_return());
+
+  // Variables to keep track of the AllGather operations and their attributes.
+  std::vector<int64_t> parallel_optimizer_recompute_allgather_fusion_ids;
+  std::vector<AnfNodePtr> parallel_optimizer_recompute_allgathers;
+  std::vector<AnfNodePtr> parallel_optimizer_recompute_first_fusion_allgathers;
+  int64_t unrecompute_max_fusion_id = -1;
+  int64_t recompute_min_fusion_id = 0;
+
+  // Process each node to identify the AllGather operations subject to recompute.
+  for (auto &node : node_list) {
+    MS_EXCEPTION_IF_NULL(node);
+    
+    // Filter out non-kernel nodes.
+    if (!node->cast<CNodePtr>() || !AnfUtils::IsRealKernel(node)) {
+      continue;
+    }
+    auto cnode = node->cast<CNodePtr>();
+
+    // Filter out nodes that aren't AllGather or associated with fusion and parallel optimizer.
+    if (!common::AnfAlgo::IsAllgather(cnode) || !common::AnfAlgo::IsFusion(cnode) ||
+        !common::AnfAlgo::IsFromParallelOptimizer(cnode)) {
+      continue;
+    }
+
+    // If the node is set for recompute, update the fusion IDs and track it.
+    if (common::AnfAlgo::IsRecompute(cnode)) {
+      int64_t fusion_id = common::AnfAlgo::GetNodeAttr<int64_t>(cnode, kAttrFusion);
+      // Check and add unique fusion_ids.
+      if (std::find(parallel_optimizer_recompute_allgather_fusion_ids.begin(),
+                    parallel_optimizer_recompute_allgather_fusion_ids.end(),
+                    fusion_id) == parallel_optimizer_recompute_allgather_fusion_ids.end()) {
+        // If the fusion id is not in the vector, it means that the allgather node is the first allgather node in the
+        // fusion group. Leave the fusion group alone till another allgather node in the same fusion group is found.
+        parallel_optimizer_recompute_allgather_fusion_ids.push_back(fusion_id);
+        recompute_min_fusion_id = recompute_min_fusion_id == 0 ? fusion_id : std::min(fusion_id, recompute_min_fusion_id);
+        parallel_optimizer_recompute_first_fusion_allgathers.push_back(node);
+      } else {
+        // Now here's another allgather node in the same fusion group. Handle it.
+        parallel_optimizer_recompute_allgathers.push_back(node);
+      }
+    } else {
+      int64_t unrecompute_fusion_id = common::AnfAlgo::GetNodeAttr<int64_t>(cnode, kAttrFusion);
+      unrecompute_max_fusion_id = std::max(unrecompute_fusion_id, unrecompute_max_fusion_id);
+      bool would_be_recomputed = common::AnfAlgo::HasNodeAttr(kAttrRecompute, cnode) &&
+                                 common::AnfAlgo::GetNodeAttr<bool>(cnode, kAttrRecompute);
+      auto [iter, inserted] =
+        forward_allgather_recompute_value_in_fusion_group.emplace(unrecompute_fusion_id, would_be_recomputed);
+      
+      // Ensure consistency within the fusion group.
+      if (!inserted && iter->second != would_be_recomputed) {
+        MS_LOG(EXCEPTION) << "In the same fusion group, the AllGather recompute attribute should be equal. "
+                             "The normal node is:"
+                          << cnode->fullname_with_scope();
+      }
+    }
+  }
+
+  // Adjust the fusion IDs for the identified AllGather operations.
+  IncreaseAllgatherFusionId(parallel_optimizer_recompute_allgathers,
+                            parallel_optimizer_recompute_first_fusion_allgathers, unrecompute_max_fusion_id,
+                            recompute_min_fusion_id);
+
+  // Adjust the dependencies for the AllGather operations.
+  return AdjustAllgatherDepend(graph, parallel_optimizer_recompute_allgathers);
+}
+
+
+/**
+ * @brief Increases the fusion ID of AllGather operations to adjust for recompute optimization.
+ * 
+ * The purpose of this method is to handle cases where duplicated AllGather operations would be fused with forward 
+ * AllGather operations. In such cases, fusion IDs need adjustment to ensure proper execution.
+ * 
+ * @param parallel_optimizer_recompute_allgathers All the AllGather operations marked for recompute.
+ * @param parallel_optimizer_recompute_first_fusion_allgathers First AllGather operations in the recomputed segments.
+ * @param unrecompute_max_fusion_id Maximum fusion ID of unrecomputed AllGather operations.
+ * @param recompute_min_fusion_id Minimum fusion ID of recomputed AllGather operations.
+ */
+void AdjustDependForParallelOptimizerRecomputeAllGather::IncreaseAllgatherFusionId(
+  const std::vector<AnfNodePtr> &parallel_optimizer_recompute_allgathers,
+  const std::vector<AnfNodePtr> &parallel_optimizer_recompute_first_fusion_allgathers,
+  int64_t unrecompute_max_fusion_id, int64_t recompute_min_fusion_id) {
+  
+  // If the condition holds, some forward AllGather and duplicated AllGather may be fused.
+  if (recompute_min_fusion_id <= unrecompute_max_fusion_id) {
+    MS_LOG(WARNING) << "Increase the duplicated allgather fusion id";
+    
+    // Adjust fusion ID for the first AllGather operations in recomputed segments.
+    for (auto &adjust_node : parallel_optimizer_recompute_first_fusion_allgathers) {
+      // Calculate the new fusion ID.
+      int64_t current_fusion_id = common::AnfAlgo::GetNodeAttr<int64_t>(adjust_node, kAttrFusion);
+      int64_t destination_fusion_id =
+        (kFusionGap + current_fusion_id + unrecompute_max_fusion_id) - recompute_min_fusion_id;
+      // Set the new fusion ID.
+      common::AnfAlgo::SetNodeAttr(kAttrFusion, MakeValue(destination_fusion_id), adjust_node);
+    }
+
+    // Adjust fusion ID for the remaining AllGather operations marked for recompute.
+    for (auto &adjust_node : parallel_optimizer_recompute_allgathers) {
+      // Calculate the new fusion ID.
+      int64_t current_fusion_id = common::AnfAlgo::GetNodeAttr<int64_t>(adjust_node, kAttrFusion);
+      int64_t destination_fusion_id =
+        (kFusionGap + current_fusion_id + unrecompute_max_fusion_id) - recompute_min_fusion_id;
+      // Set the new fusion ID.
+      common::AnfAlgo::SetNodeAttr(kAttrFusion, MakeValue(destination_fusion_id), adjust_node);
+    }
+  }
+}
+
+/**
+ * @brief Adjusts the dependencies of AllGather operations subject to recompute within the function graph.
+ * 
+ * This method ensures that the execution dependencies of AllGather operations are maintained as expected, 
+ * especially after introducing the recompute optimization.
+ * 
+ * @param graph The function graph to process.
+ * @param parallel_optimizer_recompute_allgathers All the AllGather operations marked for recompute.
+ * @return bool Indicates whether any adjustments were made to the function graph.
+ */
+bool AdjustDependForParallelOptimizerRecomputeAllGather::AdjustAllgatherDepend(
+  const FuncGraphPtr &graph, const std::vector<AnfNodePtr> &parallel_optimizer_recompute_allgathers) {
+  
+  // Get the manager for the function graph.
+  FuncGraphManagerPtr manager = graph->manager();
+  bool changed = false;
+
+  // Process each AllGather operation marked for recompute.
+  for (auto &node : parallel_optimizer_recompute_allgathers) {
+    auto cnode = node->cast<CNodePtr>();
+    auto depend_node = common::AnfAlgo::GetInputNode(cnode, 0);
+
+    // If the dependency is a Depend operation, adjust the dependency.
+    if (IsPrimitiveCNode(depend_node, prim::kPrimDepend)) {
+       // depend node is a "depend" primitive
+      auto depend_cnode = depend_node->cast<CNodePtr>();
+      AnfNodeIndexSet allgather_node_set = manager->node_users()[cnode];
+      for (auto &node_pair : allgather_node_set) {
+        auto allgather_next_node = node_pair.first;
+        CNodePtr allgather_next_cnode = node_pair.first->cast<CNodePtr>();
+
+        // Continue if the node isn't a valid CNode or a primitive.
+        if (allgather_next_cnode == nullptr || !IsValueNode<Primitive>(allgather_next_cnode->input(0))) {
+          continue;
+        }
+
+        // Construct a new Depend operation and adjust the connections.
+        std::vector<AnfNodePtr> inputs = {NewValueNode(std::make_shared<Primitive>(prim::kPrimDepend->name())),
+                                          allgather_next_node, common::AnfAlgo::GetInputNode(depend_cnode, 1)};
+        auto new_depend = graph->NewCNode(inputs);
+        new_depend->set_abstract(depend_node->abstract());
+        manager->SetEdge(node, 1, common::AnfAlgo::GetInputNode(depend_cnode, 0));
+        (void)manager->Replace(allgather_next_node, new_depend);
+        changed = true;
+      }
+    } 
+    // Handle the case where the dependency is a Cast operation with an underlying Depend operation.
+    else if (IsPrimitiveCNode(depend_node, prim::kPrimCast) &&
+               IsPrimitiveCNode(common::AnfAlgo::GetInputNode(depend_node->cast<CNodePtr>(), 0), prim::kPrimDepend)) {
+      // The logic here mirrors the above block for Depend operations, with an additional layer of handling for Cast.
+      auto cast_cnode = depend_node->cast<CNodePtr>();
+      auto cast_depend_node = common::AnfAlgo::GetInputNode(cast_cnode, 0);
+      auto cast_depend_cnode = cast_depend_node->cast<CNodePtr>();
+      AnfNodeIndexSet allgather_node_set = manager->node_users()[cnode];
+      for (auto &node_pair : allgather_node_set) {
+        auto allgather_next_node = node_pair.first;
+        CNodePtr allgather_next_cnode = node_pair.first->cast<CNodePtr>();
+        if (allgather_next_cnode == nullptr || !IsValueNode<Primitive>(allgather_next_cnode->input(0))) {
+          continue;
+        }
+        std::vector<AnfNodePtr> inputs = {NewValueNode(std::make_shared<Primitive>(prim::kPrimDepend->name())),
+                                          allgather_next_node, common::AnfAlgo::GetInputNode(cast_depend_cnode, 1)};
+        auto new_depend = graph->NewCNode(inputs);
+        new_depend->set_abstract(cast_depend_node->abstract());
+        manager->SetEdge(depend_node, 1, common::AnfAlgo::GetInputNode(cast_depend_cnode, 0));
+        (void)manager->Replace(allgather_next_node, new_depend);
+        changed = true;
+      }
+    } else {
+      MS_LOG(WARNING) << "The parallel optimizer recompute allgather has no depend edge";
+    }
+  }
+  return changed;
+}
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/common_subexpression_elimination.cc

@@ -32,135 +32,242 @@ namespace opt {
 namespace {
 using KernelWithIndex = std::pair<AnfNodePtr, int64_t>;
 
+/**
+ * @brief Check whether to ignore the current node during CSE process.
+ *
+ * @param node The node to check.
+ * @return true if the node should be ignored, false otherwise.
+ */
 bool CheckIgnoreCase(const AnfNodePtr &node) {
-  MS_EXCEPTION_IF_NULL(node);
-  if (common::AnfAlgo::GetCNodeName(node) != kTransDataOpName) {
-    return false;
-  }
-  auto cnode = node->cast<CNodePtr>();
-  MS_EXCEPTION_IF_NULL(cnode);
-  bool need_ignore = true;
-  auto input_size = cnode->inputs().size() - 1;
-  for (size_t k = 0; k < input_size; ++k) {
-    auto input = common::AnfAlgo::VisitKernelWithReturnType(common::AnfAlgo::GetInputNode(cnode, k), 0).first;
-    if (input != nullptr && input->isa<CNode>()) {
-      need_ignore = false;
-      break;
+    // Ensure the node is not null.
+    MS_EXCEPTION_IF_NULL(node);
+
+    // Check if the node's operation type is `kTransDataOpName`.
+    if (common::AnfAlgo::GetCNodeName(node) != kTransDataOpName) {
+        return false;
     }
-  }
-  return need_ignore;
+
+    // Cast to a computation node for further inspection.
+    auto cnode = node->cast<CNodePtr>();
+    MS_EXCEPTION_IF_NULL(cnode);
+
+    // By default, we consider ignoring the node.
+    bool need_ignore = true;
+
+    // Deduct 1 to ignore self in inputs.
+    auto input_size = cnode->inputs().size() - 1;
+
+    // Loop over all inputs to the computation node.
+    for (size_t k = 0; k < input_size; ++k) {
+        auto input = common::AnfAlgo::VisitKernelWithReturnType(common::AnfAlgo::GetInputNode(cnode, k), 0).first;
+        
+        // If any of the inputs are computation nodes, the node shouldn't be ignored.
+        if (input != nullptr && input->isa<CNode>()) {
+            need_ignore = false;
+            break;
+        }
+    }
+
+    // Return whether or not this node should be ignored.
+    return need_ignore;
 }
 
+/**
+ * @brief Eliminates duplicate 'TupleGetItem' operations in the computational graph.
+ *
+ * This function identifies duplicate 'TupleGetItem' operations in the graph, which have 
+ * the same input and index, and replaces duplicates with the original.
+ *
+ * @param graph The target computational graph.
+ * @param manager Manager for handling graph operations.
+ */
 void EliminateDuplicatedTupleGetItem(const FuncGraphPtr &graph, const FuncGraphManagerPtr &manager) {
-  MS_EXCEPTION_IF_NULL(graph);
-  MS_EXCEPTION_IF_NULL(manager);
+    // Ensure provided graph and manager are not null.
+    MS_EXCEPTION_IF_NULL(graph);
+    MS_EXCEPTION_IF_NULL(manager);
 
-  // key: (getitem_input, getitem_index), value: getitem_list
-  std::map<KernelWithIndex, std::vector<AnfNodePtr>> getitem_dup_map;
-  const auto &node_list = TopoSort(graph->get_return());
-  for (auto &node : node_list) {
-    if (!node->isa<CNode>() || !IsPrimitiveCNode(node, prim::kPrimTupleGetItem)) {
-      continue;
-    }
-    auto getitem_cnode = node->cast<CNodePtr>();
-    KernelWithIndex input_with_index{getitem_cnode->input(kRealInputNodeIndexInTupleGetItem),
-                                     GetGetitemIndex(getitem_cnode)};
-    if (getitem_dup_map.count(input_with_index) == 0) {
-      getitem_dup_map.emplace(input_with_index, std::vector<AnfNodePtr>{node});
-    } else {
-      getitem_dup_map[input_with_index].push_back(node);
-    }
-  }
+    // This map will store each unique 'TupleGetItem' operation and its duplicates.
+    std::map<KernelWithIndex, std::vector<AnfNodePtr>> getitem_dup_map;
 
-  // remove duplicated
-  for (auto &item : getitem_dup_map) {
-    auto &getitem_list = item.second;
-    if (getitem_list.size() > 1) {
-      auto first_getitem = getitem_list[0];
-      std::for_each(getitem_list.begin() + 1, getitem_list.end(), [first_getitem, manager](const AnfNodePtr &getitem) {
-        (void)manager->Replace(getitem, first_getitem);
-      });
+    // Sort nodes to ensure we process them in topological order.
+    const auto &node_list = TopoSort(graph->get_return());
+    
+    // Process each node in the graph.
+    for (auto &node : node_list) {
+        if (!node->isa<CNode>() || !IsPrimitiveCNode(node, prim::kPrimTupleGetItem)) {
+            continue;
+        }
+        auto getitem_cnode = node->cast<CNodePtr>();
+
+        // Create a key for this 'TupleGetItem' operation based on its input and index.
+        KernelWithIndex input_with_index{getitem_cnode->input(kRealInputNodeIndexInTupleGetItem),
+                                         GetGetitemIndex(getitem_cnode)};
+
+        // Store the operation in the map.
+        if (getitem_dup_map.count(input_with_index) == 0) {
+            getitem_dup_map.emplace(input_with_index, std::vector<AnfNodePtr>{node});
+        } else {
+            getitem_dup_map[input_with_index].push_back(node);
+        }
+    }
+
+    // For each 'TupleGetItem' operation in the map, if there are duplicates, replace them with the original.
+    for (auto &item : getitem_dup_map) {
+        auto &getitem_list = item.second;
+        if (getitem_list.size() > 1) {
+            auto first_getitem = getitem_list[0];
+            std::for_each(getitem_list.begin() + 1, getitem_list.end(), [first_getitem, manager](const AnfNodePtr &getitem) {
+                (void)manager->Replace(getitem, first_getitem);
+            });
+        }
     }
-  }
 }
 }  // namespace
 
+/**
+ * @brief Compare KernelBuildInfo between two nodes.
+ *
+ * @param main The primary node for comparison.
+ * @param node The node to compare with the primary node.
+ * @return true if the KernelBuildInfo are equal, false otherwise.
+ */
 bool BackendCSE::CheckEqualKernelBuildInfo(const AnfNodePtr &main, const AnfNodePtr &node) const {
-  MS_EXCEPTION_IF_NULL(main);
-  MS_EXCEPTION_IF_NULL(node);
-  if (main->isa<CNode>()) {
-    auto main_name = common::AnfAlgo::GetCNodeName(main);
-    if (main_name == prim::kPrimTensorMove->name() || main_name == prim::kPrimMemCpyAsync->name()) {
-      return false;
+    // Ensure the main and secondary nodes are not null.
+    MS_EXCEPTION_IF_NULL(main);
+    MS_EXCEPTION_IF_NULL(node);
+
+    // If main is a computation node and has specific operations, return false immediately.
+    if (main->isa<CNode>()) {
+        auto main_name = common::AnfAlgo::GetCNodeName(main);
+        if (main_name == prim::kPrimTensorMove->name() || main_name == prim::kPrimMemCpyAsync->name()) {
+            return false;
+        }
     }
-  }
-  auto main_kernel_info = dynamic_cast<device::KernelInfo *>(main->kernel_info());
-  auto node_kernel_info = dynamic_cast<device::KernelInfo *>(node->kernel_info());
-  if (main_kernel_info == nullptr && node_kernel_info == nullptr) {
-    return true;
-  }
-  if (main_kernel_info != nullptr && node_kernel_info != nullptr) {
-    return *main_kernel_info == *node_kernel_info;
-  }
-  return false;
+
+    // Cast the kernel info of both nodes to specific kernel info type.
+    auto main_kernel_info = dynamic_cast<device::KernelInfo *>(main->kernel_info());
+    auto node_kernel_info = dynamic_cast<device::KernelInfo *>(node->kernel_info());
+
+    // Check if both kernel infos are null. If yes, return true.
+    if (main_kernel_info == nullptr && node_kernel_info == nullptr) {
+        return true;
+    }
+    // If both kernel infos are valid, compare them.
+    if (main_kernel_info != nullptr && node_kernel_info != nullptr) {
+        return *main_kernel_info == *node_kernel_info;
+    }
+
+    // If one of the kernel infos is null but the other isn't, return false.
+    return false;
 }
 
+/**
+ * @brief Compare if the inputs of two CNodes are equal.
+ *
+ * @param main The primary CNode for comparison.
+ * @param node The CNode to compare with the primary CNode.
+ * @return true if inputs are equal, false otherwise.
+ */
 bool BackendCSE::CheckEqualCnodeInputs(const AnfNodePtr &main, const AnfNodePtr &node) const {
-  auto c_main = main->cast<CNodePtr>();
-  MS_EXCEPTION_IF_NULL(c_main);
-  auto c_node = node->cast<CNodePtr>();
-  MS_EXCEPTION_IF_NULL(c_node);
-  const auto &inp1 = c_main->inputs();
-  const auto &inp2 = c_node->inputs();
-  if (inp1.size() != inp2.size()) {
-    return false;
-  }
-  for (size_t j = 0; j < inp1.size(); j++) {
-    auto inp1_j = inp1[j];
-    auto inp2_j = inp2[j];
-    MS_EXCEPTION_IF_NULL(inp1_j);
-    MS_EXCEPTION_IF_NULL(inp2_j);
-    if (!(*inp1_j == *inp2_j)) {
-      return false;
+    // Cast nodes to specific computation node type.
+    auto c_main = main->cast<CNodePtr>();
+    MS_EXCEPTION_IF_NULL(c_main);
+    auto c_node = node->cast<CNodePtr>();
+    MS_EXCEPTION_IF_NULL(c_node);
+
+    // Retrieve the inputs of the computation nodes.
+    const auto &inp1 = c_main->inputs();
+    const auto &inp2 = c_node->inputs();
+
+    // If the number of inputs differ, the nodes are not identical.
+    if (inp1.size() != inp2.size()) {
+        return false;
     }
-  }
-  return true;
+
+    // Compare each input of both computation nodes.
+    for (size_t j = 0; j < inp1.size(); j++) {
+        auto inp1_j = inp1[j];
+        auto inp2_j = inp2[j];
+        MS_EXCEPTION_IF_NULL(inp1_j);
+        MS_EXCEPTION_IF_NULL(inp2_j);
+        if (!(*inp1_j == *inp2_j)) {
+            return false;
+        }
+    }
+
+    // If all inputs match, the nodes are identical.
+    return true;
 }
 
+/**
+ * @brief Compares two value nodes to determine if they hold the same value.
+ *
+ * This function is used to ascertain whether the main and secondary value nodes 
+ * encapsulate the same underlying value.
+ *
+ * @param main The main value node for comparison.
+ * @param node The secondary value node for comparison.
+ * @return True if the values in the nodes are the same, otherwise false.
+ */
 bool BackendCSE::CheckValueNode(const ValueNodePtr &main, const ValueNodePtr &node) const {
-  MS_EXCEPTION_IF_NULL(main);
-  MS_EXCEPTION_IF_NULL(node);
+    MS_EXCEPTION_IF_NULL(main);
+    MS_EXCEPTION_IF_NULL(node);
 
-  auto main_value = main->value();
-  MS_EXCEPTION_IF_NULL(main_value);
-  auto node_value = node->value();
-  MS_EXCEPTION_IF_NULL(node_value);
-  if (main_value->isa<Primitive>() && node_value->isa<Primitive>()) {
-    return false;
-  } else if (main_value->isa<tensor::Tensor>() && node_value->isa<tensor::Tensor>()) {
-    return (AbsOf(main) == AbsOf(node)) && CheckEqualKernelBuildInfo(main, node);
-  }
-  return (AbsOf(main) == AbsOf(node)) && (*main_value == *node_value);
+    auto main_value = main->value();
+    MS_EXCEPTION_IF_NULL(main_value);
+    auto node_value = node->value();
+    MS_EXCEPTION_IF_NULL(node_value);
+
+    // If both values are primitive types, they are not considered equal.
+    if (main_value->isa<Primitive>() && node_value->isa<Primitive>()) {
+        return false;
+    } 
+    // Special handling for tensor values.
+    else if (main_value->isa<tensor::Tensor>() && node_value->isa<tensor::Tensor>()) {
+        return (AbsOf(main) == AbsOf(node)) && CheckEqualKernelBuildInfo(main, node);
+    }
+
+    // For other value types, directly compare their values.
+    return (AbsOf(main) == AbsOf(node)) && (*main_value == *node_value);
 }
 
+/**
+ * @brief Checks if two computation nodes are functionally identical.
+ *
+ * This method performs several checks to determine if the main and secondary computation nodes 
+ * can be considered the same in terms of their function within a computational graph.
+ *
+ * @param main The main computation node to compare.
+ * @param node The secondary computation node to compare.
+ * @return True if the computation nodes are functionally identical, otherwise false.
+ */
 bool BackendCSE::CheckCNode(const CNodePtr &main, const CNodePtr &node) const {
-  MS_EXCEPTION_IF_NULL(main);
-  MS_EXCEPTION_IF_NULL(node);
+    MS_EXCEPTION_IF_NULL(main);
+    MS_EXCEPTION_IF_NULL(node);
 
-  auto context_ptr = MsContext::GetInstance();
-  MS_EXCEPTION_IF_NULL(context_ptr);
-  if (!context_ptr->get_param<bool>(MS_CTX_ENABLE_LOOP_SINK) && CheckIgnoreCase(main)) {
-    return false;
-  }
-  if (HasHiddenSideEffect(main) || HasHiddenSideEffect(node)) {
-    return false;
-  }
-  if (!CheckEqualKernelBuildInfo(main, node)) {
-    return false;
-  }
-  return CheckEqualCnodeInputs(main, node);
+    auto context_ptr = MsContext::GetInstance();
+    MS_EXCEPTION_IF_NULL(context_ptr);
+
+    // Check if loop sink is enabled and if the main node should be ignored.
+    if (!context_ptr->get_param<bool>(MS_CTX_ENABLE_LOOP_SINK) && CheckIgnoreCase(main)) {
+        return false;
+    }
+    // Check for hidden side effects in the main and secondary computation nodes.
+    if (HasHiddenSideEffect(main) || HasHiddenSideEffect(node)) {
+        return false;
+    }
+
+    // Compare kernel build info and the inputs of the two nodes.
+    return CheckEqualKernelBuildInfo(main, node) && CheckEqualCnodeInputs(main, node);
 }
 
+/**
+ * @brief Verify if one node can be replaced by another in the CSE process.
+ *
+ * @param main The node which might replace the other node.
+ * @param node The node which might be replaced.
+ * @return true if the nodes can replace each other, false otherwise.
+ */
 bool BackendCSE::CheckReplace(const AnfNodePtr &main, const AnfNodePtr &node) const {
   MS_EXCEPTION_IF_NULL(main);
   MS_EXCEPTION_IF_NULL(node);
@@ -180,6 +287,12 @@ bool BackendCSE::CheckReplace(const AnfNodePtr &main, const AnfNodePtr &node) co
   return false;
 }
 
+/*
+ * @brief Perform CSE and tuple_getitem elimination for one graph.
+ *
+ * @param graph The graph to be processed.
+ * @param manager The manager of the graph.
+ */
 bool BackendCSE::Cse(const FuncGraphPtr graph, const FuncGraphManagerPtr manager) const {
   MS_EXCEPTION_IF_NULL(manager);
   auto ret = BuildOrderGroupAndDoReplaceForOneGraph(graph, manager);
@@ -189,6 +302,13 @@ bool BackendCSE::Cse(const FuncGraphPtr graph, const FuncGraphManagerPtr manager
   return ret;
 }
 
+/**
+ * @brief Execute common subexpression elimination (CSE) on a functional graph.
+ * Implementations are in BackendCSE.
+ *
+ * @param func_graph The functional graph to process.
+ * @return true if the CSE operation was successful, false otherwise.
+ */
 bool CommonSubexpressionElimination::Run(const FuncGraphPtr &func_graph) {
   MS_EXCEPTION_IF_NULL(func_graph);
   auto backend_cse = std::make_shared<BackendCSE>();

mindspore/ccsrc/backend/common/pass/common_subexpression_elimination_old.cc

@@ -0,0 +1,319 @@
+/**
+ * Copyright 2019-2022 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "backend/common/pass/common_subexpression_elimination.h"
+
+#include <memory>
+#include <vector>
+#include <map>
+#include <utility>
+#include "runtime/device/kernel_info.h"
+#include "base/core_ops.h"
+#include "utils/flags.h"
+#include "utils/ms_context.h"
+#include "include/common/utils/utils.h"
+#include "include/common/utils/anfalgo.h"
+#include "backend/common/optimizer/helper.h"
+
+namespace mindspore {
+namespace opt {
+namespace {
+using KernelWithIndex = std::pair<AnfNodePtr, int64_t>;
+
+/**
+ * @brief Check whether to ignore the current node during CSE process.
+ *
+ * @param node The node to check.
+ * @return true if the node should be ignored, false otherwise.
+ */
+bool CheckIgnoreCase(const AnfNodePtr &node) {
+    // Ensure the node is not null.
+    MS_EXCEPTION_IF_NULL(node);
+
+    // Check if the node's operation type is `kTransDataOpName`.
+    if (common::AnfAlgo::GetCNodeName(node) != kTransDataOpName) {
+        return false;
+    }
+
+    // Cast to a computation node for further inspection.
+    auto cnode = node->cast<CNodePtr>();
+    MS_EXCEPTION_IF_NULL(cnode);
+
+    // By default, we consider ignoring the node.
+    bool need_ignore = true;
+
+    // Deduct 1 to ignore self in inputs.
+    auto input_size = cnode->inputs().size() - 1;
+
+    // Loop over all inputs to the computation node.
+    for (size_t k = 0; k < input_size; ++k) {
+        auto input = common::AnfAlgo::VisitKernelWithReturnType(common::AnfAlgo::GetInputNode(cnode, k), 0).first;
+        
+        // If any of the inputs are computation nodes, the node shouldn't be ignored.
+        if (input != nullptr && input->isa<CNode>()) {
+            need_ignore = false;
+            break;
+        }
+    }
+
+    // Return whether or not this node should be ignored.
+    return need_ignore;
+}
+
+/**
+ * @brief Eliminates duplicate 'TupleGetItem' operations in the computational graph.
+ *
+ * This function identifies duplicate 'TupleGetItem' operations in the graph, which have 
+ * the same input and index, and replaces duplicates with the original.
+ *
+ * @param graph The target computational graph.
+ * @param manager Manager for handling graph operations.
+ */
+void EliminateDuplicatedTupleGetItem(const FuncGraphPtr &graph, const FuncGraphManagerPtr &manager) {
+    // Ensure provided graph and manager are not null.
+    MS_EXCEPTION_IF_NULL(graph);
+    MS_EXCEPTION_IF_NULL(manager);
+
+    // This map will store each unique 'TupleGetItem' operation and its duplicates.
+    std::map<KernelWithIndex, std::vector<AnfNodePtr>> getitem_dup_map;
+
+    // Sort nodes to ensure we process them in topological order.
+    const auto &node_list = TopoSort(graph->get_return());
+    
+    // Process each node in the graph.
+    for (auto &node : node_list) {
+        if (!node->isa<CNode>() || !IsPrimitiveCNode(node, prim::kPrimTupleGetItem)) {
+            continue;
+        }
+        auto getitem_cnode = node->cast<CNodePtr>();
+
+        // Create a key for this 'TupleGetItem' operation based on its input and index.
+        KernelWithIndex input_with_index{getitem_cnode->input(kRealInputNodeIndexInTupleGetItem),
+                                         GetGetitemIndex(getitem_cnode)};
+
+        // Store the operation in the map.
+        if (getitem_dup_map.count(input_with_index) == 0) {
+            getitem_dup_map.emplace(input_with_index, std::vector<AnfNodePtr>{node});
+        } else {
+            getitem_dup_map[input_with_index].push_back(node);
+        }
+    }
+
+    // For each 'TupleGetItem' operation in the map, if there are duplicates, replace them with the original.
+    for (auto &item : getitem_dup_map) {
+        auto &getitem_list = item.second;
+        if (getitem_list.size() > 1) {
+            auto first_getitem = getitem_list[0];
+            std::for_each(getitem_list.begin() + 1, getitem_list.end(), [first_getitem, manager](const AnfNodePtr &getitem) {
+                (void)manager->Replace(getitem, first_getitem);
+            });
+        }
+    }
+}
+}  // namespace
+
+/**
+ * @brief Compare KernelBuildInfo between two nodes.
+ *
+ * @param main The primary node for comparison.
+ * @param node The node to compare with the primary node.
+ * @return true if the KernelBuildInfo are equal, false otherwise.
+ */
+bool BackendCSE::CheckEqualKernelBuildInfo(const AnfNodePtr &main, const AnfNodePtr &node) const {
+    // Ensure the main and secondary nodes are not null.
+    MS_EXCEPTION_IF_NULL(main);
+    MS_EXCEPTION_IF_NULL(node);
+
+    // If main is a computation node and has specific operations, return false immediately.
+    if (main->isa<CNode>()) {
+        auto main_name = common::AnfAlgo::GetCNodeName(main);
+        if (main_name == prim::kPrimTensorMove->name() || main_name == prim::kPrimMemCpyAsync->name()) {
+            return false;
+        }
+    }
+
+    // Cast the kernel info of both nodes to specific kernel info type.
+    auto main_kernel_info = dynamic_cast<device::KernelInfo *>(main->kernel_info());
+    auto node_kernel_info = dynamic_cast<device::KernelInfo *>(node->kernel_info());
+
+    // Check if both kernel infos are null. If yes, return true.
+    if (main_kernel_info == nullptr && node_kernel_info == nullptr) {
+        return true;
+    }
+    // If both kernel infos are valid, compare them.
+    if (main_kernel_info != nullptr && node_kernel_info != nullptr) {
+        return *main_kernel_info == *node_kernel_info;
+    }
+
+    // If one of the kernel infos is null but the other isn't, return false.
+    return false;
+}
+
+/**
+ * @brief Compare if the inputs of two CNodes are equal.
+ *
+ * @param main The primary CNode for comparison.
+ * @param node The CNode to compare with the primary CNode.
+ * @return true if inputs are equal, false otherwise.
+ */
+bool BackendCSE::CheckEqualCnodeInputs(const AnfNodePtr &main, const AnfNodePtr &node) const {
+    // Cast nodes to specific computation node type.
+    auto c_main = main->cast<CNodePtr>();
+    MS_EXCEPTION_IF_NULL(c_main);
+    auto c_node = node->cast<CNodePtr>();
+    MS_EXCEPTION_IF_NULL(c_node);
+
+    // Retrieve the inputs of the computation nodes.
+    const auto &inp1 = c_main->inputs();
+    const auto &inp2 = c_node->inputs();
+
+    // If the number of inputs differ, the nodes are not identical.
+    if (inp1.size() != inp2.size()) {
+        return false;
+    }
+
+    // Compare each input of both computation nodes.
+    for (size_t j = 0; j < inp1.size(); j++) {
+        auto inp1_j = inp1[j];
+        auto inp2_j = inp2[j];
+        MS_EXCEPTION_IF_NULL(inp1_j);
+        MS_EXCEPTION_IF_NULL(inp2_j);
+        if (!(*inp1_j == *inp2_j)) {
+            return false;
+        }
+    }
+
+    // If all inputs match, the nodes are identical.
+    return true;
+}
+
+/**
+ * @brief Compares two value nodes to determine if they hold the same value.
+ *
+ * This function is used to ascertain whether the main and secondary value nodes 
+ * encapsulate the same underlying value.
+ *
+ * @param main The main value node for comparison.
+ * @param node The secondary value node for comparison.
+ * @return True if the values in the nodes are the same, otherwise false.
+ */
+bool BackendCSE::CheckValueNode(const ValueNodePtr &main, const ValueNodePtr &node) const {
+    MS_EXCEPTION_IF_NULL(main);
+    MS_EXCEPTION_IF_NULL(node);
+
+    auto main_value = main->value();
+    MS_EXCEPTION_IF_NULL(main_value);
+    auto node_value = node->value();
+    MS_EXCEPTION_IF_NULL(node_value);
+
+    // If both values are primitive types, they are not considered equal.
+    if (main_value->isa<Primitive>() && node_value->isa<Primitive>()) {
+        return false;
+    } 
+    // Special handling for tensor values.
+    else if (main_value->isa<tensor::Tensor>() && node_value->isa<tensor::Tensor>()) {
+        return (AbsOf(main) == AbsOf(node)) && CheckEqualKernelBuildInfo(main, node);
+    }
+
+    // For other value types, directly compare their values.
+    return (AbsOf(main) == AbsOf(node)) && (*main_value == *node_value);
+}
+
+/**
+ * @brief Checks if two computation nodes are functionally identical.
+ *
+ * This method performs several checks to determine if the main and secondary computation nodes 
+ * can be considered the same in terms of their function within a computational graph.
+ *
+ * @param main The main computation node to compare.
+ * @param node The secondary computation node to compare.
+ * @return True if the computation nodes are functionally identical, otherwise false.
+ */
+bool BackendCSE::CheckCNode(const CNodePtr &main, const CNodePtr &node) const {
+    MS_EXCEPTION_IF_NULL(main);
+    MS_EXCEPTION_IF_NULL(node);
+
+    auto context_ptr = MsContext::GetInstance();
+    MS_EXCEPTION_IF_NULL(context_ptr);
+
+    // Check if loop sink is enabled and if the main node should be ignored.
+    if (!context_ptr->get_param<bool>(MS_CTX_ENABLE_LOOP_SINK) && CheckIgnoreCase(main)) {
+        return false;
+    }
+    // Check for hidden side effects in the main and secondary computation nodes.
+    if (HasHiddenSideEffect(main) || HasHiddenSideEffect(node)) {
+        return false;
+    }
+
+    // Compare kernel build info and the inputs of the two nodes.
+    return CheckEqualKernelBuildInfo(main, node) && CheckEqualCnodeInputs(main, node);
+}
+
+/**
+ * @brief Verify if one node can be replaced by another in the CSE process.
+ *
+ * @param main The node which might replace the other node.
+ * @param node The node which might be replaced.
+ * @return true if the nodes can replace each other, false otherwise.
+ */
+bool BackendCSE::CheckReplace(const AnfNodePtr &main, const AnfNodePtr &node) const {
+  MS_EXCEPTION_IF_NULL(main);
+  MS_EXCEPTION_IF_NULL(node);
+
+  // attrs of nop node inserted by backend maybe omitted, so two nodes have same inputs will have different outputs
+  auto main_abs = main->abstract();
+  auto node_abs = node->abstract();
+  if (main_abs != nullptr && node_abs != nullptr && !(*main_abs == *node_abs)) {
+    return false;
+  }
+
+  if (main->isa<ValueNode>() && node->isa<ValueNode>()) {
+    return CheckValueNode(main->cast<ValueNodePtr>(), node->cast<ValueNodePtr>());
+  } else if (main->isa<CNode>() && node->isa<CNode>()) {
+    return CheckCNode(main->cast<CNodePtr>(), node->cast<CNodePtr>());
+  }
+  return false;
+}
+
+/*
+ * @brief Perform CSE and tuple_getitem elimination for one graph.
+ *
+ * @param graph The graph to be processed.
+ * @param manager The manager of the graph.
+ */
+bool BackendCSE::Cse(const FuncGraphPtr graph, const FuncGraphManagerPtr manager) const {
+  MS_EXCEPTION_IF_NULL(manager);
+  auto ret = BuildOrderGroupAndDoReplaceForOneGraph(graph, manager);
+  if (ret) {
+    EliminateDuplicatedTupleGetItem(graph, manager);
+  }
+  return ret;
+}
+
+/**
+ * @brief Execute common subexpression elimination (CSE) on a functional graph.
+ * Implementations are in BackendCSE.
+ *
+ * @param func_graph The functional graph to process.
+ * @return true if the CSE operation was successful, false otherwise.
+ */
+bool CommonSubexpressionElimination::Run(const FuncGraphPtr &func_graph) {
+  MS_EXCEPTION_IF_NULL(func_graph);
+  auto backend_cse = std::make_shared<BackendCSE>();
+  MS_EXCEPTION_IF_NULL(backend_cse);
+  return backend_cse->Cse(func_graph, func_graph->manager());
+}
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/communication_op_fusion.cc

@@ -37,19 +37,39 @@ constexpr auto kAttrDefaultOp = "default_op";
 constexpr size_t kAlignSize = 2 << 9;
 constexpr int64_t kDefaultThresholdMb2Byte = 262144;
 
+/**
+ * @brief Generate kernel build information for a range of communication operations.
+ * 
+ * This function processes a range of communication operation nodes (from `start_index` to `end_index`) 
+ * within `communication_op_info` and extracts their device format, device type, and other properties to 
+ * create a kernel build information object.
+ *
+ * @param communication_op_info The structure containing all communication operation nodes.
+ * @param start_index The starting index within the `communication_op_info` structure.
+ * @param end_index The ending index within the `communication_op_info` structure.
+ * @return The constructed kernel build information object.
+ */
 kernel::KernelBuildInfoPtr GenerateKernelBuildInfo(const CommunicationOpInfo &communication_op_info, size_t start_index,
                                                    size_t end_index) {
+  // Check the validity of the given range.
   if (end_index >= communication_op_info.communication_op_nodes.size()) {
     MS_LOG(EXCEPTION) << "end index out of communication_op_nodes size";
   }
+  
+  // Initialize containers for device information.
   std::vector<std::string> inputs_device_format;
   std::vector<std::string> outputs_device_format;
   std::vector<TypeId> inputs_device_type;
   std::vector<TypeId> outputs_device_type;
   std::vector<std::vector<size_t>> outputs_shape;
+
   kernel::KernelBuildInfo::KernelBuildInfoBuilder builder;
+
+  // Process each communication operation node in the given range.
   for (size_t idx = start_index; idx <= end_index; ++idx) {
     auto cnode = communication_op_info.communication_op_nodes[idx];
+
+    // Extract rank size for certain operations.
     int64_t rank_size = 1;
     if (common::AnfAlgo::HasNodeAttr(kAttrRankSize, cnode) &&
         common::AnfAlgo::GetCNodeName(cnode) == kAllGatherOpName) {
@@ -59,12 +79,15 @@ kernel::KernelBuildInfoPtr GenerateKernelBuildInfo(const CommunicationOpInfo &co
     if (rank_size_t == 0) {
       MS_LOG(EXCEPTION) << "Rank size should not be zero.";
     }
-    MS_EXCEPTION_IF_NULL(cnode);
+
+    // Extract device format and type for inputs of the communication operation.
     size_t input_num = common::AnfAlgo::GetInputTensorNum(cnode);
     for (size_t input_index = 0; input_index < input_num; ++input_index) {
       inputs_device_format.push_back(AnfAlgo::GetInputFormat(cnode, input_index));
       inputs_device_type.push_back(AnfAlgo::GetInputDeviceDataType(cnode, input_index));
     }
+
+    // Extract device format, type, and shape for outputs of the communication operation.
     for (size_t rank_index = 0; rank_index < rank_size_t; ++rank_index) {
       size_t output_num = common::AnfAlgo::GetOutputTensorNum(cnode);
       for (size_t output_index = 0; output_index < output_num; ++output_index) {
@@ -77,20 +100,39 @@ kernel::KernelBuildInfoPtr GenerateKernelBuildInfo(const CommunicationOpInfo &co
         outputs_shape.push_back(common::AnfAlgo::GetOutputInferShape(cnode, output_index));
       }
     }
+
+    // Add more general attributes to the builder.
     builder.SetFusionType(AnfAlgo::GetFusionType(cnode));
     builder.SetProcessor(AnfAlgo::GetProcessor(cnode));
     builder.SetKernelType(AnfAlgo::GetKernelType(cnode));
   }
+
+  // Finalize the builder with extracted information.
   builder.SetInputsFormat(inputs_device_format);
   builder.SetOutputsFormat(outputs_device_format);
   builder.SetInputsDeviceType(inputs_device_type);
   builder.SetOutputsDeviceType(outputs_device_type);
+
   return builder.Build();
 }
 
+
+/**
+ * @brief Generate a unique key for a fusion group based on a given node.
+ * 
+ * This function uses various attributes such as `fusion`, `group`, `op`, and data type from 
+ * the given node to generate a unique key which can be used to identify or group 
+ * related fusion operations in MindSpore.
+ *
+ * @param node The AnfNode for which the fusion group key is to be generated.
+ * @return A string representing the fusion group key.
+ */
 std::string GetFusionGroupKey(const AnfNodePtr &node) {
+  // Extract the primitive from the given node.
   auto primitive = common::AnfAlgo::GetCNodePrimitive(node);
   MS_EXCEPTION_IF_NULL(primitive);
+
+  // Extract and check fusion attribute.
   ValuePtr attr_fusion = primitive->GetAttr(kAttrFusion);
   if (attr_fusion == nullptr) {
     return "";
@@ -99,6 +141,8 @@ std::string GetFusionGroupKey(const AnfNodePtr &node) {
   if (fusion == 0) {
     return "";
   }
+
+  // Extract or use default values for group and op attributes.
   std::string group = kAttrDefaultGroup;
   ValuePtr attr_group = primitive->GetAttr(kAttrGroup);
   if (attr_group != nullptr) {
@@ -109,23 +153,51 @@ std::string GetFusionGroupKey(const AnfNodePtr &node) {
   if (attr_op != nullptr) {
     op = GetValue<std::string>(attr_op);
   }
+
+  // Extract the data type for the node.
   auto dtype = common::AnfAlgo::GetPrevNodeOutputInferDataType(node, 0);
+
+  // Construct and return the fusion group key.
   return group + op + std::to_string(fusion) + TypeIdLabel(dtype);
 }
 
+/**
+ * @brief Checks if multiple fusion inputs have duplicates.
+ * 
+ * In the context of fusion operations within MindSpore, different communication operations 
+ * in one segment should not share the same input. This function checks if this condition is met.
+ *
+ * @param fusion_inputs A vector of AnfNodePtr representing the fusion inputs to check.
+ */
 void CheckInputs(const std::vector<AnfNodePtr> &fusion_inputs) {
   std::set<AnfNodePtr> inputs_set(fusion_inputs.begin(), fusion_inputs.end());
+  
+  // If the set size is smaller than the vector, there are duplicates.
   if (inputs_set.size() < fusion_inputs.size()) {
     MS_LOG(EXCEPTION) << "Different communication op in one segment cannot share the same input";
   }
 }
 
+/**
+ * @brief Validate the segments for a given communication operation node.
+ * 
+ * This function checks the validity of segments for a communication operation node. 
+ * It ensures the correct ordering and range of segment indices.
+ *
+ * @param communication_op_node_size The size of the communication operation node.
+ * @param segment_index The segment index to be validated.
+ * @return true if the segments are valid, otherwise an exception is thrown.
+ */
 bool CheckSegments(size_t communication_op_node_size, const std::vector<size_t> *segment_index) {
   MS_EXCEPTION_IF_NULL(segment_index);
+
+  // Check the last segment index for validity.
   auto segments = segment_index->size();
   if (segment_index->at(segments - 1) != communication_op_node_size - 1) {
     MS_LOG(EXCEPTION) << "the last segment index is invalid.";
   }
+
+  // Validate the ordering of segment indices.
   for (size_t i = 0; i < segments - 1; ++i) {
     if (segment_index->at(i) > segment_index->at(i + 1)) {
       MS_LOG(EXCEPTION) << "illegal split: segment_index[" << i << "]=" << segment_index->at(i) << ", segment_index[ "
@@ -134,14 +206,30 @@ bool CheckSegments(size_t communication_op_node_size, const std::vector<size_t>
   }
   return true;
 }
+
 }  // namespace
 
+/**
+ * @brief Get the segmentation indices for fusion in communication operations.
+ * 
+ * This function determines the segments' boundaries for communication operations fusion
+ * based on the provided CommunicationOpInfo and a specific group. The segmentation might
+ * vary based on the operation name (e.g., HcomSendOp, ReceiveOp) and other factors.
+ * 
+ * @param communication_op_info The CommunicationOpInfo containing communication operation nodes.
+ * @param segment_index A vector to be filled with the segmentation indices.
+ * @param group The specific group for which the segmentation is required.
+ * @return true if the segmentation was successful, otherwise false.
+ */
 bool CommunicationOpFusion::GetSplitSegments(const CommunicationOpInfo &communication_op_info,
                                              std::vector<size_t> *segment_index, const std::string &group) const {
   MS_EXCEPTION_IF_NULL(segment_index);
+
+  // Get the size of communication operation nodes.
   size_t communication_op_node_size = communication_op_info.communication_op_nodes.size();
   MS_LOG(INFO) << "graph " << op_name_ << " node size " << communication_op_node_size;
 
+  // Special handling for send and receive operations.
   if (op_name_ == kHcomSendOpName || op_name_ == kReceiveOpName) {
     if (communication_op_node_size == 0) {
       return false;
@@ -152,15 +240,20 @@ bool CommunicationOpFusion::GetSplitSegments(const CommunicationOpInfo &communic
 
   auto parallel_context = parallel::ParallelContext::GetInstance();
   MS_EXCEPTION_IF_NULL(parallel_context);
+
+  // Determine split indices.
   std::vector<uint32_t> split_indices;
   if (!parallel_context->enable_parallel_optimizer()) {
     split_indices = parallel_context->GetAllReduceFusionSplitIndices(group);
   }
 
+  // Handling based on split indices.
   if (!split_indices.empty()) {
     uint32_t last_index = 0;
     for (size_t i = 0; i < split_indices.size(); ++i) {
       uint32_t index = split_indices[i];
+
+      // Check validity of split index.
       if (index <= last_index && i != 0) {
         MS_LOG(EXCEPTION) << "invalid " << op_name_ << " split index " << i << " " << index;
       }
@@ -172,15 +265,19 @@ bool CommunicationOpFusion::GetSplitSegments(const CommunicationOpInfo &communic
       segment_index->push_back(index);
       last_index = index;
     }
+    // Ensure the inclusion of the last index.
     if (last_index != communication_op_node_size - 1) {
       segment_index->push_back(communication_op_node_size - 1);
     }
   } else {
+    // Default segmentation strategy.
     for (size_t i = 0; i < groups_ - 1; ++i) {
       segment_index->push_back((i + 1) * (communication_op_node_size / groups_) - 1);
     }
     segment_index->push_back(communication_op_node_size - 1);
   }
+
+  // Further segmentation strategy for data parallelism mode with AllReduce operations.
   auto parallel_mode = parallel_context->parallel_mode();
   if (parallel_mode == parallel::kDataParallel && op_name_ == kAllReduceOpName) {
     auto threshold = parallel_context->dp_fusion_threshold_mb();
@@ -188,19 +285,36 @@ bool CommunicationOpFusion::GetSplitSegments(const CommunicationOpInfo &communic
     MS_LOG(INFO) << "The split threshold for AllReduce is " << threshold << ", the segment num is "
                  << segment_index->size();
   }
+
   return CheckSegments(communication_op_node_size, segment_index);
 }
 
+/**
+ * @brief Determine the segmentation indices for AllReduce operations based on a threshold.
+ * 
+ * This function refines the segmentation indices for AllReduce communication operations
+ * based on a given memory threshold. If the cumulative memory size exceeds the threshold,
+ * a new segment is started.
+ *
+ * @param nodes A vector of CNodePtr representing the nodes being considered for segmentation.
+ * @param threshold The memory threshold in MB for the segmentation.
+ * @param segment_index A vector to be filled with the refined segmentation indices.
+ */
 void CommunicationOpFusion::GetAllReduceSplitSegment(const std::vector<CNodePtr> &nodes, int64_t threshold,
                                                      std::vector<size_t> *segment_index) const {
   MS_EXCEPTION_IF_NULL(segment_index);
+
+  // Check the provided threshold.
   if (threshold <= 0) {
     MS_LOG(WARNING) << "Split threshold is " << threshold << ". AllReduce nodes will take default fusion strategy.";
     return;
   }
-  threshold *= kDefaultThresholdMb2Byte;
+  threshold *= kDefaultThresholdMb2Byte; // Convert MB to Byte.
+
   std::vector<size_t> real_segment_index;
   size_t start_index = 0;
+
+  // Iterate through the segments to refine them based on the threshold.
   for (auto index : *segment_index) {
     if (index >= nodes.size()) {
       MS_LOG(WARNING) << "split index is greater than or equal to total gradient's number " << nodes.size();
@@ -216,14 +330,17 @@ void CommunicationOpFusion::GetAllReduceSplitSegment(const std::vector<CNodePtr>
         accumulate += tensor_size;
       }
     }
+    // Add segments with remaining accumulated size.
     if (accumulate != 0) {
       real_segment_index.push_back(index);
     }
     start_index = index + 1;
   }
+  
   *segment_index = std::move(real_segment_index);
 }
 
+
 // Hard coded Load(%paraxxx, cnode()) to Load(%paraxxx, U) to prevent
 // cycle after AllReduce fused. It's a workaround.
 // case 1:
@@ -265,29 +382,50 @@ void CommunicationOpFusion::GetAllReduceSplitSegment(const std::vector<CNodePtr>
 // ...
 // %109 = AssignAdd(%para485, Tensor(34), cnode_u)
 // %110 = UpdateState(cnode_u, xxx)
+/**
+ * @brief Adjust the input of an AllReduce node if it has a Load dependency.
+ * 
+ * This function identifies whether an AllReduce operation has a Load operation dependency.
+ * If found, it performs several adjustments, mainly replacing certain nodes' input to UMonads.
+ * 
+ * @param cnode The CNodePtr representing the AllReduce node to be checked and adjusted.
+ */
 static void AdjustAllReduceInputWithLoad(const CNodePtr &cnode) {
+  // Constant definitions for indexing and size checks.
   const size_t monad_index = 2;
   const size_t tuple_inputs_size = 2;
   const size_t load_inputs_size = 3;
+
+  // Search for a Load operation dependency.
   auto cnode_load = BroadFirstSearchFirstOf({cnode}, [](const CNodePtr &search_cnode) {
     if (!IsPrimitiveCNode(search_cnode, prim::kPrimLoad)) {
       return false;
     }
+    // Ensuring the Load CNode has the expected number of inputs.
     if (search_cnode->inputs().size() != load_inputs_size) {
       MS_LOG(EXCEPTION) << "Load CNode should have 3 inputs, but: " << search_cnode->DebugString();
     }
     return search_cnode->input(monad_index)->isa<CNode>();
   });
+
+  // If a Load operation is found, perform adjustments.
   if (cnode_load != nullptr) {
+    // Create a UMonad ValueNode.
     auto const_u_monad = NewValueNode(kUMonad);
     const_u_monad->set_abstract(kUMonad->ToAbstract());
     const auto &cnode_u = cnode_load->input(monad_index);
+
     MS_LOG(DEBUG) << "Replace Load with CNode U to constant U for cnode: " << cnode_load->DebugString();
+
+    // Ensure the cnode belongs to a valid FuncGraph.
     MS_EXCEPTION_IF_NULL(cnode->func_graph());
     MS_EXCEPTION_IF_NULL(cnode->func_graph()->manager());
+
+    // Replace the UMonad input of Load CNode.
     auto manager = cnode->func_graph()->manager();
     manager->SetEdge(cnode_load, monad_index, const_u_monad);
-    // Update the u_monad input of UpdateState from CNode U same as Load to constant U.
+
+    // Identify UpdateState dependencies and adjust them.
     CNodePtr cnode_update_state = nullptr;
     CNodePtr cnode_make_tuple = nullptr;
     const auto &cnode_load_users = manager->node_users()[cnode_load];
@@ -316,14 +454,14 @@ static void AdjustAllReduceInputWithLoad(const CNodePtr &cnode) {
         }
       }
     }
+
+    // Perform adjustments based on identified UpdateState dependencies.
     if (cnode_update_state != nullptr) {
       if (cnode_make_tuple == nullptr || cnode_make_tuple->inputs().size() == tuple_inputs_size) {
-        // case 1 and case 3: Replace cnode_update_state to cnode_u;
         MS_LOG(DEBUG) << "Replace UpdateState with CNode U: " << cnode_update_state->DebugString()
                       << " ::TO:: " << cnode_u->DebugString();
         manager->Replace(cnode_update_state, cnode_u);
       } else if (cnode_make_tuple->inputs().size() > tuple_inputs_size) {
-        // case 2: remove cnode_load from cnode_make_tuple;
         MS_LOG(DEBUG) << "Drop " << cnode_load->DebugString() << " from " << cnode_make_tuple->DebugString();
         const auto &make_tuple_inputs = cnode_make_tuple->inputs();
         AnfNodePtrList new_tuple_inputs(make_tuple_inputs.size() - 1);
@@ -339,9 +477,22 @@ static void AdjustAllReduceInputWithLoad(const CNodePtr &cnode) {
   }
 }
 
+/**
+ * @brief Creates a fused communication operation.
+ * 
+ * This function creates a new CNode representing a fused communication operation based
+ * on the specified communication operation information.
+ * 
+ * @param func_graph The function graph in which the new node will be created.
+ * @param communication_op_info The structure containing information about communication operations.
+ * @param start_index The starting index of the communication operation to be fused.
+ * @param end_index The ending index of the communication operation to be fused.
+ * @return AnfNodePtr The newly created CNode representing the fused operation.
+ */
 AnfNodePtr CommunicationOpFusion::CreateFusedCommunicationOp(const FuncGraphPtr &func_graph,
                                                              const CommunicationOpInfo &communication_op_info,
                                                              size_t start_index, size_t end_index) const {
+  // Validate the input function graph.
   MS_EXCEPTION_IF_NULL(func_graph);
   auto prim = std::make_shared<Primitive>(op_name_);
   MS_EXCEPTION_IF_NULL(prim);
@@ -430,64 +581,110 @@ AnfNodePtr CommunicationOpFusion::CreateFusedCommunicationOp(const FuncGraphPtr
   return fused_node;
 }
 
-bool CommunicationOpFusion::DoFusion(const FuncGraphPtr &func_graph, const CommunicationOpInfo &communication_op_info,
+/**
+ * @brief Performs fusion on communication operations within a function graph.
+ * 
+ * This method attempts to fuse a sequence of communication operations within a given function graph
+ * into a single composite operation. This is done in order to optimize performance and reduce
+ * overhead of executing multiple small operations. The method uses segment indices to determine
+ * which operations should be fused together.
+ * 
+ * @param func_graph The function graph in which communication operations are to be fused.
+ * @param communication_op_info The structure containing information about communication operations.
+ * @param segment_index A list of indices that denote segments of operations to be fused.
+ * @return bool Indicates whether any fusion has taken place.
+ */
+bool CommunicationOpFusion::DoFusion(const FuncGraphPtr &func_graph, 
+                                     const CommunicationOpInfo &communication_op_info,
                                      const std::vector<size_t> &segment_index) const {
+  // Validate the input function graph.
   MS_EXCEPTION_IF_NULL(func_graph);
   auto manager = func_graph->manager();
   MS_EXCEPTION_IF_NULL(manager);
+
   bool changed = false;
   size_t start_index = 0;
+
+  // Iterate through segment indices to determine operations to fuse.
   for (size_t segment_idx = 0; segment_idx < segment_index.size(); ++segment_idx) {
     size_t end_index = segment_index.at(segment_idx);
+    
+    // Skip if the segment is too small to fuse.
     if (end_index - start_index < 1) {
       start_index = end_index + 1;
       continue;
     }
+
+    // Cast the function graph to a KernelGraph.
     auto kernel_graph = func_graph->cast<KernelGraphPtr>();
     MS_EXCEPTION_IF_NULL(kernel_graph);
     auto graph_id = kernel_graph->graph_id();
-    AnfNodePtr new_communication_op =
+    
+    // Create a new fused communication operation for the segment.
+    AnfNodePtr new_communication_op = 
       CreateFusedCommunicationOp(func_graph, communication_op_info, start_index, end_index);
     AnfAlgo::SetGraphId(graph_id, new_communication_op.get());
-    // replace old communication op with new communication op
+
+    // Replace the old communication operations in the segment with the new fused operation.
     for (auto idx = start_index; idx <= end_index; ++idx) {
-      std::vector<AnfNodePtr> tuple_getitem_input;
-      tuple_getitem_input.push_back(NewValueNode(prim::kPrimTupleGetItem));
-      tuple_getitem_input.push_back(new_communication_op);
-      auto offset = SizeToLong(idx - start_index);
-      auto index = NewValueNode(offset);
-      MS_EXCEPTION_IF_NULL(index);
-      auto imm = std::make_shared<Int64Imm>(idx - start_index);
-      MS_EXCEPTION_IF_NULL(imm);
-      auto abstract_scalar = std::make_shared<abstract::AbstractScalar>();
-      MS_EXCEPTION_IF_NULL(abstract_scalar);
-      index->set_abstract(abstract_scalar);
-      tuple_getitem_input.push_back(index);
+      // Create a new TupleGetItem node to extract outputs from the fused operation.
+      std::vector<AnfNodePtr> tuple_getitem_input = {
+        NewValueNode(prim::kPrimTupleGetItem),
+        new_communication_op,
+        NewValueNode(SizeToLong(idx - start_index))
+      };
       AnfNodePtr tuple_getitem = func_graph->NewCNode(tuple_getitem_input);
       MS_EXCEPTION_IF_NULL(tuple_getitem);
+
+      // Update abstract and replace the old operation with TupleGetItem node.
       auto communication_op_node_item = communication_op_info.communication_op_nodes.at(idx);
       MS_EXCEPTION_IF_NULL(communication_op_node_item);
       tuple_getitem->set_abstract(communication_op_node_item->abstract());
+
+      // Handle internal outputs if the operation is an internal output of the kernel graph.
       if (kernel_graph->IsInternalOutput(communication_op_node_item, 0)) {
-        kernel_graph->ReplaceInternalOutput(communication_op_node_item, new_communication_op, 0, LongToSize(offset));
+        kernel_graph->ReplaceInternalOutput(communication_op_node_item, new_communication_op, 0, LongToSize(idx - start_index));
       }
+
+      // Perform the replacement in the graph manager.
       if (!manager->Replace(communication_op_node_item, tuple_getitem)) {
         MS_LOG(EXCEPTION) << "Manager replace node failed";
       }
     }
+
+    // Move to the next segment.
     start_index = end_index + 1;
     changed = true;
   }
+
   return changed;
 }
 
+
+/**
+ * @brief Fuses communication operations within the provided function graph.
+ * 
+ * This method processes the given function graph, identifies and groups the communication
+ * operations that are candidates for fusion, and then fuses them. Fusion can lead to 
+ * improved performance by reducing the overhead of executing multiple smaller operations.
+ * 
+ * @param func_graph The function graph to process.
+ * @return bool Indicates whether any fusion operations were applied to the function graph.
+ */
 bool CommunicationOpFusion::Run(const FuncGraphPtr &func_graph) {
+  // Validate the input function graph.
   MS_EXCEPTION_IF_NULL(func_graph);
   const float input_grad_size_num = 0.0;
   const float input_grad_time_num = 0.0;
-  // divide candidate fusion groups with same (group,op,fusion,dtype) attrs, fusion==0 means not fusion
+  
+  // Initialize a mapping to keep track of candidate operations that can be fused.
+  // Operations are grouped by certain attributes such as group, operation, fusion, and dtype.
   mindspore::HashMap<std::string, CommunicationOpInfo> candidate_groups;
+  
+  // Retrieve the nodes of the function graph in topological order.
   std::vector<AnfNodePtr> node_list = TopoSort(func_graph->get_return());
+  
+  // Iterate through the nodes and identify communication operations for potential fusion.
   for (auto &node : node_list) {
     if (node != nullptr && node->isa<CNode>() && common::AnfAlgo::GetCNodeName(node) == op_name_) {
       std::string key = GetFusionGroupKey(node);
@@ -503,14 +700,20 @@ bool CommunicationOpFusion::Run(const FuncGraphPtr &func_graph) {
       candidate_groups[key].input_grad_time.push_back(input_grad_time_num);
     }
   }
-  // split candidate group to segments according to _group class member
+  
   bool changed = false;
+
+  // Process each group of candidate operations.
   for (auto &it : candidate_groups) {
+    // Skip groups with only one operation since they cannot be fused.
     if (it.second.communication_op_nodes.size() <= 1) {
       continue;
     }
+
     auto first_node = it.second.communication_op_nodes[0];
     TraceGuard guard(std::make_shared<TraceOpt>(first_node->debug_info()));
+
+    // If the nodes have the "index" attribute, sort them based on this attribute.
     if (common::AnfAlgo::HasNodeAttr(kAttrIndex, first_node) &&
         common::AnfAlgo::GetNodeAttr<int64_t>(first_node, kAttrIndex) > 0) {
       std::stable_sort(it.second.communication_op_nodes.begin(), it.second.communication_op_nodes.end(),
@@ -519,14 +722,19 @@ bool CommunicationOpFusion::Run(const FuncGraphPtr &func_graph) {
                                 common::AnfAlgo::GetNodeAttr<int64_t>(b, kAttrIndex);
                        });
     }
+
+    // Determine segments of the group to be fused.
     std::vector<size_t> segment_index;
     if (GetSplitSegments(it.second, &segment_index, it.first)) {
+      // Apply the fusion on the identified segments.
       if (DoFusion(func_graph, it.second, segment_index)) {
         changed = true;
       }
     }
   }
+
   return changed;
 }
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/communication_op_fusion_old.cc

@@ -0,0 +1,740 @@
+/**
+ * Copyright 2019-2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "backend/common/pass/communication_op_fusion.h"
+
+#include <vector>
+#include <set>
+#include <memory>
+
+#include "utils/hash_map.h"
+#include "ir/graph_utils.h"
+#include "base/core_ops.h"
+#include "runtime/device/kernel_info.h"
+#include "backend/common/session/anf_runtime_algorithm.h"
+#include "include/common/utils/anfalgo.h"
+#include "kernel/kernel_build_info.h"
+#include "backend/common/optimizer/helper.h"
+#include "include/common/utils/parallel_context.h"
+
+namespace mindspore {
+namespace opt {
+namespace {
+constexpr auto kAttrDefaultGroup = "default_group";
+constexpr auto kAttrDefaultOp = "default_op";
+constexpr size_t kAlignSize = 2 << 9;
+constexpr int64_t kDefaultThresholdMb2Byte = 262144;
+
+/**
+ * @brief Generate kernel build information for a range of communication operations.
+ * 
+ * This function processes a range of communication operation nodes (from `start_index` to `end_index`) 
+ * within `communication_op_info` and extracts their device format, device type, and other properties to 
+ * create a kernel build information object.
+ *
+ * @param communication_op_info The structure containing all communication operation nodes.
+ * @param start_index The starting index within the `communication_op_info` structure.
+ * @param end_index The ending index within the `communication_op_info` structure.
+ * @return The constructed kernel build information object.
+ */
+kernel::KernelBuildInfoPtr GenerateKernelBuildInfo(const CommunicationOpInfo &communication_op_info, size_t start_index,
+                                                   size_t end_index) {
+  // Check the validity of the given range.
+  if (end_index >= communication_op_info.communication_op_nodes.size()) {
+    MS_LOG(EXCEPTION) << "end index out of communication_op_nodes size";
+  }
+  
+  // Initialize containers for device information.
+  std::vector<std::string> inputs_device_format;
+  std::vector<std::string> outputs_device_format;
+  std::vector<TypeId> inputs_device_type;
+  std::vector<TypeId> outputs_device_type;
+  std::vector<std::vector<size_t>> outputs_shape;
+
+  kernel::KernelBuildInfo::KernelBuildInfoBuilder builder;
+
+  // Process each communication operation node in the given range.
+  for (size_t idx = start_index; idx <= end_index; ++idx) {
+    auto cnode = communication_op_info.communication_op_nodes[idx];
+
+    // Extract rank size for certain operations.
+    int64_t rank_size = 1;
+    if (common::AnfAlgo::HasNodeAttr(kAttrRankSize, cnode) &&
+        common::AnfAlgo::GetCNodeName(cnode) == kAllGatherOpName) {
+      rank_size = common::AnfAlgo::GetNodeAttr<int64_t>(cnode, kAttrRankSize);
+    }
+    size_t rank_size_t = LongToSize(rank_size);
+    if (rank_size_t == 0) {
+      MS_LOG(EXCEPTION) << "Rank size should not be zero.";
+    }
+
+    // Extract device format and type for inputs of the communication operation.
+    size_t input_num = common::AnfAlgo::GetInputTensorNum(cnode);
+    for (size_t input_index = 0; input_index < input_num; ++input_index) {
+      inputs_device_format.push_back(AnfAlgo::GetInputFormat(cnode, input_index));
+      inputs_device_type.push_back(AnfAlgo::GetInputDeviceDataType(cnode, input_index));
+    }
+
+    // Extract device format, type, and shape for outputs of the communication operation.
+    for (size_t rank_index = 0; rank_index < rank_size_t; ++rank_index) {
+      size_t output_num = common::AnfAlgo::GetOutputTensorNum(cnode);
+      for (size_t output_index = 0; output_index < output_num; ++output_index) {
+        outputs_device_format.push_back(AnfAlgo::GetOutputFormat(cnode, output_index));
+        outputs_device_type.push_back(AnfAlgo::GetOutputDeviceDataType(cnode, output_index));
+        std::vector<size_t> shape = common::AnfAlgo::GetOutputInferShape(cnode, output_index);
+        if (!shape.empty()) {
+          shape[0] /= rank_size_t;
+        }
+        outputs_shape.push_back(common::AnfAlgo::GetOutputInferShape(cnode, output_index));
+      }
+    }
+
+    // Add more general attributes to the builder.
+    builder.SetFusionType(AnfAlgo::GetFusionType(cnode));
+    builder.SetProcessor(AnfAlgo::GetProcessor(cnode));
+    builder.SetKernelType(AnfAlgo::GetKernelType(cnode));
+  }
+
+  // Finalize the builder with extracted information.
+  builder.SetInputsFormat(inputs_device_format);
+  builder.SetOutputsFormat(outputs_device_format);
+  builder.SetInputsDeviceType(inputs_device_type);
+  builder.SetOutputsDeviceType(outputs_device_type);
+
+  return builder.Build();
+}
+
+
+/**
+ * @brief Generate a unique key for a fusion group based on a given node.
+ * 
+ * This function uses various attributes such as `fusion`, `group`, `op`, and data type from 
+ * the given node to generate a unique key which can be used to identify or group 
+ * related fusion operations in MindSpore.
+ *
+ * @param node The AnfNode for which the fusion group key is to be generated.
+ * @return A string representing the fusion group key.
+ */
+std::string GetFusionGroupKey(const AnfNodePtr &node) {
+  // Extract the primitive from the given node.
+  auto primitive = common::AnfAlgo::GetCNodePrimitive(node);
+  MS_EXCEPTION_IF_NULL(primitive);
+
+  // Extract and check fusion attribute.
+  ValuePtr attr_fusion = primitive->GetAttr(kAttrFusion);
+  if (attr_fusion == nullptr) {
+    return "";
+  }
+  auto fusion = GetValue<int64_t>(attr_fusion);
+  if (fusion == 0) {
+    return "";
+  }
+
+  // Extract or use default values for group and op attributes.
+  std::string group = kAttrDefaultGroup;
+  ValuePtr attr_group = primitive->GetAttr(kAttrGroup);
+  if (attr_group != nullptr) {
+    group = GetValue<std::string>(attr_group);
+  }
+  std::string op = kAttrDefaultOp;
+  ValuePtr attr_op = primitive->GetAttr(kAttrOp);
+  if (attr_op != nullptr) {
+    op = GetValue<std::string>(attr_op);
+  }
+
+  // Extract the data type for the node.
+  auto dtype = common::AnfAlgo::GetPrevNodeOutputInferDataType(node, 0);
+
+  // Construct and return the fusion group key.
+  return group + op + std::to_string(fusion) + TypeIdLabel(dtype);
+}
+
+/**
+ * @brief Checks if multiple fusion inputs have duplicates.
+ * 
+ * In the context of fusion operations within MindSpore, different communication operations 
+ * in one segment should not share the same input. This function checks if this condition is met.
+ *
+ * @param fusion_inputs A vector of AnfNodePtr representing the fusion inputs to check.
+ */
+void CheckInputs(const std::vector<AnfNodePtr> &fusion_inputs) {
+  std::set<AnfNodePtr> inputs_set(fusion_inputs.begin(), fusion_inputs.end());
+  
+  // If the set size is smaller than the vector, there are duplicates.
+  if (inputs_set.size() < fusion_inputs.size()) {
+    MS_LOG(EXCEPTION) << "Different communication op in one segment cannot share the same input";
+  }
+}
+
+/**
+ * @brief Validate the segments for a given communication operation node.
+ * 
+ * This function checks the validity of segments for a communication operation node. 
+ * It ensures the correct ordering and range of segment indices.
+ *
+ * @param communication_op_node_size The size of the communication operation node.
+ * @param segment_index The segment index to be validated.
+ * @return true if the segments are valid, otherwise an exception is thrown.
+ */
+bool CheckSegments(size_t communication_op_node_size, const std::vector<size_t> *segment_index) {
+  MS_EXCEPTION_IF_NULL(segment_index);
+
+  // Check the last segment index for validity.
+  auto segments = segment_index->size();
+  if (segment_index->at(segments - 1) != communication_op_node_size - 1) {
+    MS_LOG(EXCEPTION) << "the last segment index is invalid.";
+  }
+
+  // Validate the ordering of segment indices.
+  for (size_t i = 0; i < segments - 1; ++i) {
+    if (segment_index->at(i) > segment_index->at(i + 1)) {
+      MS_LOG(EXCEPTION) << "illegal split: segment_index[" << i << "]=" << segment_index->at(i) << ", segment_index[ "
+                        << (i + 1) << "]=" << segment_index->at(i + 1);
+    }
+  }
+  return true;
+}
+
+}  // namespace
+
+/**
+ * @brief Get the segmentation indices for fusion in communication operations.
+ * 
+ * This function determines the segments' boundaries for communication operations fusion
+ * based on the provided CommunicationOpInfo and a specific group. The segmentation might
+ * vary based on the operation name (e.g., HcomSendOp, ReceiveOp) and other factors.
+ * 
+ * @param communication_op_info The CommunicationOpInfo containing communication operation nodes.
+ * @param segment_index A vector to be filled with the segmentation indices.
+ * @param group The specific group for which the segmentation is required.
+ * @return true if the segmentation was successful, otherwise false.
+ */
+bool CommunicationOpFusion::GetSplitSegments(const CommunicationOpInfo &communication_op_info,
+                                             std::vector<size_t> *segment_index, const std::string &group) const {
+  MS_EXCEPTION_IF_NULL(segment_index);
+
+  // Get the size of communication operation nodes.
+  size_t communication_op_node_size = communication_op_info.communication_op_nodes.size();
+  MS_LOG(INFO) << "graph " << op_name_ << " node size " << communication_op_node_size;
+
+  // Special handling for send and receive operations.
+  if (op_name_ == kHcomSendOpName || op_name_ == kReceiveOpName) {
+    if (communication_op_node_size == 0) {
+      return false;
+    }
+    (void)segment_index->emplace_back(communication_op_node_size - 1);
+    return true;
+  }
+
+  auto parallel_context = parallel::ParallelContext::GetInstance();
+  MS_EXCEPTION_IF_NULL(parallel_context);
+
+  // Determine split indices.
+  std::vector<uint32_t> split_indices;
+  if (!parallel_context->enable_parallel_optimizer()) {
+    split_indices = parallel_context->GetAllReduceFusionSplitIndices(group);
+  }
+
+  // Handling based on split indices.
+  if (!split_indices.empty()) {
+    uint32_t last_index = 0;
+    for (size_t i = 0; i < split_indices.size(); ++i) {
+      uint32_t index = split_indices[i];
+
+      // Check validity of split index.
+      if (index <= last_index && i != 0) {
+        MS_LOG(EXCEPTION) << "invalid " << op_name_ << " split index " << i << " " << index;
+      }
+      if (index >= communication_op_node_size) {
+        MS_LOG(WARNING) << op_name_ << "'s split index " << index
+                        << " is Greater than or equal to total gradient's number " << communication_op_node_size;
+        continue;
+      }
+      segment_index->push_back(index);
+      last_index = index;
+    }
+    // Ensure the inclusion of the last index.
+    if (last_index != communication_op_node_size - 1) {
+      segment_index->push_back(communication_op_node_size - 1);
+    }
+  } else {
+    // Default segmentation strategy.
+    for (size_t i = 0; i < groups_ - 1; ++i) {
+      segment_index->push_back((i + 1) * (communication_op_node_size / groups_) - 1);
+    }
+    segment_index->push_back(communication_op_node_size - 1);
+  }
+
+  // Further segmentation strategy for data parallelism mode with AllReduce operations.
+  auto parallel_mode = parallel_context->parallel_mode();
+  if (parallel_mode == parallel::kDataParallel && op_name_ == kAllReduceOpName) {
+    auto threshold = parallel_context->dp_fusion_threshold_mb();
+    GetAllReduceSplitSegment(communication_op_info.communication_op_nodes, threshold, segment_index);
+    MS_LOG(INFO) << "The split threshold for AllReduce is " << threshold << ", the segment num is "
+                 << segment_index->size();
+  }
+
+  return CheckSegments(communication_op_node_size, segment_index);
+}
+
+/**
+ * @brief Determine the segmentation indices for AllReduce operations based on a threshold.
+ * 
+ * This function refines the segmentation indices for AllReduce communication operations
+ * based on a given memory threshold. If the cumulative memory size exceeds the threshold,
+ * a new segment is started.
+ *
+ * @param nodes A vector of CNodePtr representing the nodes being considered for segmentation.
+ * @param threshold The memory threshold in MB for the segmentation.
+ * @param segment_index A vector to be filled with the refined segmentation indices.
+ */
+void CommunicationOpFusion::GetAllReduceSplitSegment(const std::vector<CNodePtr> &nodes, int64_t threshold,
+                                                     std::vector<size_t> *segment_index) const {
+  MS_EXCEPTION_IF_NULL(segment_index);
+
+  // Check the provided threshold.
+  if (threshold <= 0) {
+    MS_LOG(WARNING) << "Split threshold is " << threshold << ". AllReduce nodes will take default fusion strategy.";
+    return;
+  }
+  threshold *= kDefaultThresholdMb2Byte; // Convert MB to Byte.
+
+  std::vector<size_t> real_segment_index;
+  size_t start_index = 0;
+
+  // Iterate through the segments to refine them based on the threshold.
+  for (auto index : *segment_index) {
+    if (index >= nodes.size()) {
+      MS_LOG(WARNING) << "split index is greater than or equal to total gradient's number " << nodes.size();
+      continue;
+    }
+    size_t accumulate = 0;
+    for (size_t j = start_index; j <= index; ++j) {
+      auto tensor_size = AnfAlgo::GetOutputTensorMemSize(nodes[j], 0);
+      if (accumulate + tensor_size > LongToSize(threshold)) {
+        real_segment_index.push_back(j);
+        accumulate = 0;
+      } else {
+        accumulate += tensor_size;
+      }
+    }
+    // Add segments with remaining accumulated size.
+    if (accumulate != 0) {
+      real_segment_index.push_back(index);
+    }
+    start_index = index + 1;
+  }
+  
+  *segment_index = std::move(real_segment_index);
+}
+
+
+// Hard coded Load(%paraxxx, cnode()) to Load(%paraxxx, U) to prevent
+// cycle after AllReduce fused. It's a workaround.
+// case 1:
+// cnode_load = Load(%para2, cnode_u)
+// %100 = UpdateState(cnode_u, cnode_load)
+// ...
+// %109 = AssignAdd(%para485, Tensor(34), %100)
+// %110 = UpdateState(%100, xxx)
+// will convert to:
+// cnode_load = Load(%para2, U)
+// ...
+// %109 = AssignAdd(%para485, Tensor(34), cnode_u)
+// %110 = UpdateState(cnode_u, xxx)
+//
+// case 2:
+// cnode_load = Load(%para2, cnode_u)
+// %99 = make_tuple(yyy, ..., cnode_load, ...)
+// %100 = UpdateState(cnode_u, %99)
+// ...
+// %109 = AssignAdd(%para485, Tensor(34), %100)
+// %110 = UpdateState(%100, xxx)
+// will convert to:
+// cnode_load = Load(%para2, U)
+// %99 = make_tuple(yyy, ...)
+// %100 = UpdateState(cnode_u, %99)
+// ...
+// %109 = AssignAdd(%para485, Tensor(34), %100)
+// %110 = UpdateState(%100, xxx)
+//
+// case 3:
+// cnode_load = Load(%para2, cnode_u)
+// %99 = make_tuple(cnode_load)
+// %100 = UpdateState(cnode_u, %99)
+// ...
+// %109 = AssignAdd(%para485, Tensor(34), %100)
+// %110 = UpdateState(%100, xxx)
+// will convert to:
+// cnode_load = Load(%para2, U)
+// ...
+// %109 = AssignAdd(%para485, Tensor(34), cnode_u)
+// %110 = UpdateState(cnode_u, xxx)
+/**
+ * @brief Adjust the input of an AllReduce node if it has a Load dependency.
+ * 
+ * This function identifies whether an AllReduce operation has a Load operation dependency.
+ * If found, it performs several adjustments, mainly replacing certain nodes' input to UMonads.
+ * 
+ * @param cnode The CNodePtr representing the AllReduce node to be checked and adjusted.
+ */
+static void AdjustAllReduceInputWithLoad(const CNodePtr &cnode) {
+  // Constant definitions for indexing and size checks.
+  const size_t monad_index = 2;
+  const size_t tuple_inputs_size = 2;
+  const size_t load_inputs_size = 3;
+
+  // Search for a Load operation dependency.
+  auto cnode_load = BroadFirstSearchFirstOf({cnode}, [](const CNodePtr &search_cnode) {
+    if (!IsPrimitiveCNode(search_cnode, prim::kPrimLoad)) {
+      return false;
+    }
+    // Ensuring the Load CNode has the expected number of inputs.
+    if (search_cnode->inputs().size() != load_inputs_size) {
+      MS_LOG(EXCEPTION) << "Load CNode should have 3 inputs, but: " << search_cnode->DebugString();
+    }
+    return search_cnode->input(monad_index)->isa<CNode>();
+  });
+
+  // If a Load operation is found, perform adjustments.
+  if (cnode_load != nullptr) {
+    // Create a UMonad ValueNode.
+    auto const_u_monad = NewValueNode(kUMonad);
+    const_u_monad->set_abstract(kUMonad->ToAbstract());
+    const auto &cnode_u = cnode_load->input(monad_index);
+
+    MS_LOG(DEBUG) << "Replace Load with CNode U to constant U for cnode: " << cnode_load->DebugString();
+
+    // Ensure the cnode belongs to a valid FuncGraph.
+    MS_EXCEPTION_IF_NULL(cnode->func_graph());
+    MS_EXCEPTION_IF_NULL(cnode->func_graph()->manager());
+
+    // Replace the UMonad input of Load CNode.
+    auto manager = cnode->func_graph()->manager();
+    manager->SetEdge(cnode_load, monad_index, const_u_monad);
+
+    // Identify UpdateState dependencies and adjust them.
+    CNodePtr cnode_update_state = nullptr;
+    CNodePtr cnode_make_tuple = nullptr;
+    const auto &cnode_load_users = manager->node_users()[cnode_load];
+    for (auto &load_user : cnode_load_users) {
+      if (IsPrimitiveCNode(load_user.first, prim::kPrimMakeTuple)) {
+        const auto &cnode_make_tuple_users = manager->node_users()[load_user.first];
+        for (auto &make_tuple_user : cnode_make_tuple_users) {
+          if (IsPrimitiveCNode(make_tuple_user.first, prim::kPrimUpdateState)) {
+            const auto &cnode_user = make_tuple_user.first->cast<CNodePtr>();
+            if (cnode_user->input(1) == cnode_u) {
+              cnode_update_state = cnode_user;
+              cnode_make_tuple = load_user.first->cast<CNodePtr>();
+              break;
+            }
+          }
+        }
+        if (cnode_update_state != nullptr) {
+          break;
+        }
+      }
+      if (IsPrimitiveCNode(load_user.first, prim::kPrimUpdateState)) {
+        const auto &cnode_user = load_user.first->cast<CNodePtr>();
+        if (cnode_user->input(1) == cnode_u) {
+          cnode_update_state = cnode_user;
+          break;
+        }
+      }
+    }
+
+    // Perform adjustments based on identified UpdateState dependencies.
+    if (cnode_update_state != nullptr) {
+      if (cnode_make_tuple == nullptr || cnode_make_tuple->inputs().size() == tuple_inputs_size) {
+        MS_LOG(DEBUG) << "Replace UpdateState with CNode U: " << cnode_update_state->DebugString()
+                      << " ::TO:: " << cnode_u->DebugString();
+        manager->Replace(cnode_update_state, cnode_u);
+      } else if (cnode_make_tuple->inputs().size() > tuple_inputs_size) {
+        MS_LOG(DEBUG) << "Drop " << cnode_load->DebugString() << " from " << cnode_make_tuple->DebugString();
+        const auto &make_tuple_inputs = cnode_make_tuple->inputs();
+        AnfNodePtrList new_tuple_inputs(make_tuple_inputs.size() - 1);
+        std::copy_if(make_tuple_inputs.cbegin(), make_tuple_inputs.cend(), new_tuple_inputs.begin(),
+                     [cnode_load](const auto &inp) { return inp != cnode_load; });
+        auto new_cnode_make_tuple = cnode_make_tuple->func_graph()->NewCNode(new_tuple_inputs);
+        manager->Replace(cnode_make_tuple, new_cnode_make_tuple);
+      } else {
+        MS_LOG(EXCEPTION) << "Cannot replace UpdateState with CNode U: " << cnode_update_state->DebugString()
+                          << " as make_tuple CNode cannot match " << cnode_make_tuple->DebugString();
+      }
+    }
+  }
+}
+
+/**
+ * @brief Creates a fused communication operation.
+ * 
+ * This function creates a new CNode representing a fused communication operation based
+ * on the specified communication operation information.
+ * 
+ * @param func_graph The function graph in which the new node will be created.
+ * @param communication_op_info The structure containing information about communication operations.
+ * @param start_index The starting index of the communication operation to be fused.
+ * @param end_index The ending index of the communication operation to be fused.
+ * @return AnfNodePtr The newly created CNode representing the fused operation.
+ */
+AnfNodePtr CommunicationOpFusion::CreateFusedCommunicationOp(const FuncGraphPtr &func_graph,
+                                                             const CommunicationOpInfo &communication_op_info,
+                                                             size_t start_index, size_t end_index) const {
+  // Validate the input function graph.
+  MS_EXCEPTION_IF_NULL(func_graph);
+  auto prim = std::make_shared<Primitive>(op_name_);
+  MS_EXCEPTION_IF_NULL(prim);
+  std::vector<AnfNodePtr> fusion_inputs = {NewValueNode(prim)};
+  // get all inputs of current segment
+  if (end_index >= communication_op_info.communication_op_nodes.size()) {
+    MS_LOG(EXCEPTION) << "End index is out of communication_op_nodes size";
+  }
+  std::vector<AnfNodePtr> orig_nodes;
+  for (size_t idx = start_index; idx <= end_index; ++idx) {
+    auto cnode = communication_op_info.communication_op_nodes[idx];
+    MS_EXCEPTION_IF_NULL(cnode);
+    if (idx != start_index) {
+      AdjustAllReduceInputWithLoad(cnode);
+    }
+    (void)fusion_inputs.insert(fusion_inputs.end(), cnode->inputs().begin() + 1, cnode->inputs().end());
+    (void)orig_nodes.emplace_back(cnode);
+  }
+  CheckInputs(fusion_inputs);
+  AnfNodePtr fused_node = NewCNode(fusion_inputs, func_graph, orig_nodes);
+  MS_EXCEPTION_IF_NULL(fused_node);
+  auto kernel_info = std::make_shared<device::KernelInfo>();
+  MS_EXCEPTION_IF_NULL(kernel_info);
+  fused_node->set_kernel_info(kernel_info);
+  auto final_node = communication_op_info.communication_op_nodes[end_index];
+  size_t node_num = end_index - start_index + 1;
+  int64_t rank_size = 1;
+  if (common::AnfAlgo::HasNodeAttr(kAttrRankSize, final_node) &&
+      common::AnfAlgo::GetCNodeName(final_node) == kAllGatherOpName) {
+    rank_size = common::AnfAlgo::GetNodeAttr<int64_t>(final_node, kAttrRankSize);
+  }
+  size_t rank_size_t = LongToSize(rank_size);
+  if (rank_size_t == 0) {
+    MS_LOG(EXCEPTION) << "Rank size should not be zero.";
+  }
+  size_t output_num = node_num * rank_size_t;
+  std::vector<TypeId> dtypes(output_num, common::AnfAlgo::GetOutputInferDataType(final_node, 0));
+  std::vector<std::vector<size_t>> shapes;
+  int64_t fusion_total_size = 0;
+  for (size_t i = 0; i < rank_size_t; ++i) {
+    for (size_t idx = start_index; idx <= end_index; ++idx) {
+      auto input_node = communication_op_info.communication_op_nodes[idx];
+      MS_EXCEPTION_IF_NULL(input_node);
+      std::vector<size_t> shape = common::AnfAlgo::GetOutputInferShape(input_node, 0);
+      if (!shape.empty()) {
+        shape[0] /= rank_size_t;
+      }
+      shapes.push_back(shape);
+      size_t tensor_size = AnfAlgo::GetOutputTensorMemSize(input_node, 0);
+      TypeId output_type = AnfAlgo::GetOutputDeviceDataType(input_node, 0);
+      size_t type_size = GetTypeByte(TypeIdToType(output_type));
+      if (type_size == 0) {
+        MS_LOG(EXCEPTION) << "Divisor 'type_size' should not be 0.";
+      }
+      tensor_size = (tensor_size / kAlignSize + 1) * kAlignSize / type_size;
+      fusion_total_size += static_cast<int64_t>(tensor_size);
+    }
+  }
+  common::AnfAlgo::SetOutputInferTypeAndShape(dtypes, shapes, fused_node.get());
+  auto kernel_build_info = GenerateKernelBuildInfo(communication_op_info, start_index, end_index);
+  AnfAlgo::SetSelectKernelBuildInfo(kernel_build_info, fused_node.get());
+  const std::vector<std::string> kHcclFusionAttrs = {
+    kAttrFusion, kAttrGroup, kAttrGroupBack, kAttrSrTag,        kAttrDestRank,          kAttrSrcRank,
+    kAttrDType,  kAttrOp,    kAttrRankSize,  kAttrGroupRankIds, kAttrReuseCommunication};
+  for (const auto &attr : kHcclFusionAttrs) {
+    if (common::AnfAlgo::HasNodeAttr(attr, final_node)) {
+      common::AnfAlgo::CopyNodeAttr(attr, final_node, fused_node);
+    }
+  }
+  if (common::AnfAlgo::HasNodeAttr(kAttrShape, final_node)) {
+    std::vector<int64_t> fusion_total_shape{fusion_total_size};
+    common::AnfAlgo::SetNodeAttr(kAttrShape, MakeValue(fusion_total_shape), fused_node);
+  }
+  bool is_recompute =
+    final_node->GetAttr(kAttrDuplicated) != nullptr && GetValue<bool>(final_node->GetAttr(kAttrDuplicated));
+  if (common::AnfAlgo::GetCNodeName(final_node) == kAllGatherOpName && is_recompute) {
+    auto fused_cnode = fused_node->cast<CNodePtr>();
+    fused_cnode->AddAttr("duplicated", MakeValue(true));
+    auto fused_prim = GetCNodePrimitive(fused_cnode);
+    auto final_node_prim = GetCNodePrimitive(final_node);
+    fused_prim->set_instance_name(final_node_prim->instance_name());
+  }
+  if (common::AnfAlgo::HasNodeAttr(kAttrNotDelayFusion, final_node)) {
+    common::AnfAlgo::CopyNodeAttr(kAttrNotDelayFusion, final_node, fused_node);
+  }
+  return fused_node;
+}
+
+/**
+ * @brief Performs fusion on communication operations within a function graph.
+ * 
+ * This method attempts to fuse a sequence of communication operations within a given function graph
+ * into a single composite operation. This is done in order to optimize performance and reduce
+ * overhead of executing multiple small operations. The method uses segment indices to determine
+ * which operations should be fused together.
+ * 
+ * @param func_graph The function graph in which communication operations are to be fused.
+ * @param communication_op_info The structure containing information about communication operations.
+ * @param segment_index A list of indices that denote segments of operations to be fused.
+ * @return bool Indicates whether any fusion has taken place.
+ */
+bool CommunicationOpFusion::DoFusion(const FuncGraphPtr &func_graph, 
+                                     const CommunicationOpInfo &communication_op_info,
+                                     const std::vector<size_t> &segment_index) const {
+  // Validate the input function graph.
+  MS_EXCEPTION_IF_NULL(func_graph);
+  auto manager = func_graph->manager();
+  MS_EXCEPTION_IF_NULL(manager);
+
+  bool changed = false;
+  size_t start_index = 0;
+
+  // Iterate through segment indices to determine operations to fuse.
+  for (size_t segment_idx = 0; segment_idx < segment_index.size(); ++segment_idx) {
+    size_t end_index = segment_index.at(segment_idx);
+    
+    // Skip if the segment is too small to fuse.
+    if (end_index - start_index < 1) {
+      start_index = end_index + 1;
+      continue;
+    }
+
+    // Cast the function graph to a KernelGraph.
+    auto kernel_graph = func_graph->cast<KernelGraphPtr>();
+    MS_EXCEPTION_IF_NULL(kernel_graph);
+    auto graph_id = kernel_graph->graph_id();
+    
+    // Create a new fused communication operation for the segment.
+    AnfNodePtr new_communication_op = 
+      CreateFusedCommunicationOp(func_graph, communication_op_info, start_index, end_index);
+    AnfAlgo::SetGraphId(graph_id, new_communication_op.get());
+
+    // Replace the old communication operations in the segment with the new fused operation.
+    for (auto idx = start_index; idx <= end_index; ++idx) {
+      // Create a new TupleGetItem node to extract outputs from the fused operation.
+      std::vector<AnfNodePtr> tuple_getitem_input = {
+        NewValueNode(prim::kPrimTupleGetItem),
+        new_communication_op,
+        NewValueNode(SizeToLong(idx - start_index))
+      };
+      AnfNodePtr tuple_getitem = func_graph->NewCNode(tuple_getitem_input);
+      MS_EXCEPTION_IF_NULL(tuple_getitem);
+
+      // Update abstract and replace the old operation with TupleGetItem node.
+      auto communication_op_node_item = communication_op_info.communication_op_nodes.at(idx);
+      MS_EXCEPTION_IF_NULL(communication_op_node_item);
+      tuple_getitem->set_abstract(communication_op_node_item->abstract());
+
+      // Handle internal outputs if the operation is an internal output of the kernel graph.
+      if (kernel_graph->IsInternalOutput(communication_op_node_item, 0)) {
+        kernel_graph->ReplaceInternalOutput(communication_op_node_item, new_communication_op, 0, LongToSize(idx - start_index));
+      }
+
+      // Perform the replacement in the graph manager.
+      if (!manager->Replace(communication_op_node_item, tuple_getitem)) {
+        MS_LOG(EXCEPTION) << "Manager replace node failed";
+      }
+    }
+
+    // Move to the next segment.
+    start_index = end_index + 1;
+    changed = true;
+  }
+
+  return changed;
+}
+
+
+/**
+ * @brief Fuses communication operations within the provided function graph.
+ * 
+ * This method processes the given function graph, identifies and groups the communication
+ * operations that are candidates for fusion, and then fuses them. Fusion can lead to 
+ * improved performance by reducing the overhead of executing multiple smaller operations.
+ * 
+ * @param func_graph The function graph to process.
+ * @return bool Indicates whether any fusion operations were applied to the function graph.
+ */
+bool CommunicationOpFusion::Run(const FuncGraphPtr &func_graph) {
+  // Validate the input function graph.
+  MS_EXCEPTION_IF_NULL(func_graph);
+  const float input_grad_size_num = 0.0;
+  const float input_grad_time_num = 0.0;
+  
+  // Initialize a mapping to keep track of candidate operations that can be fused.
+  // Operations are grouped by certain attributes such as group, operation, fusion, and dtype.
+  mindspore::HashMap<std::string, CommunicationOpInfo> candidate_groups;
+  
+  // Retrieve the nodes of the function graph in topological order.
+  std::vector<AnfNodePtr> node_list = TopoSort(func_graph->get_return());
+  
+  // Iterate through the nodes and identify communication operations for potential fusion.
+  for (auto &node : node_list) {
+    if (node != nullptr && node->isa<CNode>() && common::AnfAlgo::GetCNodeName(node) == op_name_) {
+      std::string key = GetFusionGroupKey(node);
+      if (key.empty()) {
+        continue;
+      }
+      if (candidate_groups.find(key) == candidate_groups.end()) {
+        CommunicationOpInfo communication_op_info;
+        candidate_groups[key] = communication_op_info;
+      }
+      candidate_groups[key].communication_op_nodes.push_back(node->cast<CNodePtr>());
+      candidate_groups[key].input_grad_size.push_back(input_grad_size_num);
+      candidate_groups[key].input_grad_time.push_back(input_grad_time_num);
+    }
+  }
+  
+  bool changed = false;
+
+  // Process each group of candidate operations.
+  for (auto &it : candidate_groups) {
+    // Skip groups with only one operation since they cannot be fused.
+    if (it.second.communication_op_nodes.size() <= 1) {
+      continue;
+    }
+
+    auto first_node = it.second.communication_op_nodes[0];
+    TraceGuard guard(std::make_shared<TraceOpt>(first_node->debug_info()));
+
+    // If the nodes have the "index" attribute, sort them based on this attribute.
+    if (common::AnfAlgo::HasNodeAttr(kAttrIndex, first_node) &&
+        common::AnfAlgo::GetNodeAttr<int64_t>(first_node, kAttrIndex) > 0) {
+      std::stable_sort(it.second.communication_op_nodes.begin(), it.second.communication_op_nodes.end(),
+                       [](const CNodePtr &a, const CNodePtr &b) {
+                         return common::AnfAlgo::GetNodeAttr<int64_t>(a, kAttrIndex) <
+                                common::AnfAlgo::GetNodeAttr<int64_t>(b, kAttrIndex);
+                       });
+    }
+
+    // Determine segments of the group to be fused.
+    std::vector<size_t> segment_index;
+    if (GetSplitSegments(it.second, &segment_index, it.first)) {
+      // Apply the fusion on the identified segments.
+      if (DoFusion(func_graph, it.second, segment_index)) {
+        changed = true;
+      }
+    }
+  }
+
+  return changed;
+}
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/const_to_attr_strided_slice_grad.cc

@@ -29,45 +29,84 @@ namespace opt {
 namespace {
 const size_t strides_index = 5;
 
+/**
+ * @brief Retrieves the stride values from the StridedSliceGrad node.
+ *
+ * This function is designed to extract the stride values from the given StridedSliceGrad node in MindSpore.
+ * It primarily checks the validity of the node, extracts the desired stride values, and returns whether the operation was successful.
+ *
+ * @param strided_slice_grad The given StridedSliceGrad node.
+ * @param strides_values Pointer to the list where the stride values will be stored.
+ * @return True if the stride values were successfully retrieved, otherwise False.
+ */
 bool GetStridesValues(const CNodePtr &strided_slice_grad, ValuePtrList *strides_values) {
+  // Check for null inputs.
   MS_EXCEPTION_IF_NULL(strided_slice_grad);
   MS_EXCEPTION_IF_NULL(strides_values);
+
   constexpr size_t kSizeChange = 6;
+
+  // Ensure the strided_slice_grad node has the expected size.
   if (strided_slice_grad->size() < kSizeChange) {
     MS_LOG(DEBUG) << "Op strided_slice_grad's inputs size less than 6, graph not changed";
     return false;
   }
+
+  // Extract the strides input from the strided_slice_grad node.
   auto strides_input = strided_slice_grad->input(strides_index);
   MS_EXCEPTION_IF_NULL(strides_input);
+
+  // Cast the strides input to a value node.
   auto strides_value_node = strides_input->cast<ValueNodePtr>();
   if (strides_value_node == nullptr) {
     MS_LOG(DEBUG) << "strides is not a value node.";
     return false;
   }
+
+  // Retrieve the value of the strides_value_node.
   auto value = strides_value_node->value();
   if (value == nullptr) {
     MS_LOG(DEBUG) << "strides has no value.";
     return false;
   }
+
+  // Cast the retrieved value to a tuple.
   auto value_tuple = value->cast<ValueTuplePtr>();
   if (value_tuple == nullptr) {
     MS_LOG(DEBUG) << "strides is not a value tuple.";
     return false;
   }
+
+  // Assign the retrieved tuple value to the strides_values pointer.
   *strides_values = value_tuple->value();
+
   return true;
 }
 
+/**
+ * @brief Checks the values of the given strides.
+ *
+ * This function checks whether the provided stride values meet specific conditions.
+ * It validates if each stride value is a scalar of integer type and equal to 1, which seems to be a requirement for StridedSliceGrad in MindSpore.
+ *
+ * @param strides_values List containing the stride values.
+ * @return True if all stride values meet the required conditions, otherwise False.
+ */
 bool CheckValues(const ValuePtrList &strides_values) {
+  // Check for an empty strides_values list.
   if (strides_values.empty()) {
     MS_LOG(DEBUG) << "strides_values is empty";
     return false;
   }
+
+  // Iterate over the stride values and check each value.
   for (auto &value : strides_values) {
     MS_EXCEPTION_IF_NULL(value);
     if (value->isa<Scalar>()) {
       auto scalar = value->cast<ScalarPtr>();
       MS_EXCEPTION_IF_NULL(scalar);
+
+      // Check if the scalar value is an integer and equals 1.
       if (scalar->isa<Int32Imm>()) {
         if (GetValue<int>(scalar) != 1) {
           MS_LOG(DEBUG) << "StridedSliceGrad has no 1 value";
@@ -87,61 +126,112 @@ bool CheckValues(const ValuePtrList &strides_values) {
       return false;
     }
   }
+
   return true;
 }
 
+
+/**
+ * @brief Checks the attributes of the StridedSliceGrad node.
+ * 
+ * This function verifies the existence and values of the `new_axis_mask` and `shrink_axis_mask` attributes
+ * in the given StridedSliceGrad node.
+ *
+ * @param strided_slice_grad The StridedSliceGrad node to check.
+ * @return True if both attributes exist and their values are 0; otherwise, False.
+ */
 bool CheckAttrs(const CNodePtr &strided_slice_grad) {
+  // Ensure strided_slice_grad is not null.
   MS_EXCEPTION_IF_NULL(strided_slice_grad);
+
+  // Check for the existence of the required attributes.
   if (!common::AnfAlgo::HasNodeAttr(kAttrNewAxisMask, strided_slice_grad) ||
       !common::AnfAlgo::HasNodeAttr(kAttrShrinkAxisMask, strided_slice_grad)) {
     MS_LOG(INFO) << "new_axis_mask or shrink_axis_mask not exist in cnode[" + strided_slice_grad->DebugString() + "]";
     return false;
   }
+
+  // Extract the values of the required attributes.
   auto new_axis_mask = common::AnfAlgo::GetNodeAttr<int64_t>(strided_slice_grad, kAttrNewAxisMask);
   auto shrink_axis_mask = common::AnfAlgo::GetNodeAttr<int64_t>(strided_slice_grad, kAttrShrinkAxisMask);
+
+  // Check the attribute values.
   if (new_axis_mask != 0 || shrink_axis_mask != 0) {
     MS_LOG(INFO) << "new_axis_mask or shrink_axis_mask not equal 0";
     return false;
   }
+
   return true;
 }
-}  // namespace
 
+// Namespace wrapping might be in the original code. Assuming it's still the case.
+namespace {
+
+/**
+ * @brief Defines the pattern for the ConstToAttrStridedSliceGradPass.
+ *
+ * This function sets up the pattern to identify the StridedSliceGrad operations that should be 
+ * processed by the ConstToAttrStridedSliceGradPass.
+ * 
+ * @return The pattern for identifying the appropriate nodes.
+ */
 const BaseRef ConstToAttrStridedSliceGradPass::DefinePattern() const {
   VarPtr Xs = std::make_shared<SeqVar>();
   auto strided_slice_grad_prim = std::make_shared<Primitive>(kStridedSliceGradOpName);
   return VectorRef({strided_slice_grad_prim, Xs});
 }
 
+/**
+ * @brief Processes the identified StridedSliceGrad node.
+ * 
+ * For the StridedSliceGrad nodes that match the previously defined pattern, this function checks 
+ * attributes and strides' values. If they satisfy certain conditions (based on the device and other factors), 
+ * the node's inputs are then converted to attributes.
+ *
+ * @param graph The functional graph containing the node.
+ * @param node The StridedSliceGrad node to process.
+ * @param [unused] The equivalence class of the node. (currently unused in this function)
+ * @return nullptr, indicating the node was processed (or not modified).
+ */
 const AnfNodePtr ConstToAttrStridedSliceGradPass::Process(const FuncGraphPtr &graph, const AnfNodePtr &node,
                                                           const EquivPtr &) const {
+  // Ensure inputs are not null.
   MS_EXCEPTION_IF_NULL(graph);
   MS_EXCEPTION_IF_NULL(node);
+
   auto strided_slice_grad = node->cast<CNodePtr>();
   MS_EXCEPTION_IF_NULL(strided_slice_grad);
 
   auto ms_context = MsContext::GetInstance();
   MS_EXCEPTION_IF_NULL(ms_context);
 
+  // Check if the processing is specific to Ascend device.
   if (ms_context->get_param<std::string>(MS_CTX_DEVICE_TARGET) == kAscendDevice) {
+    // Validate the node's attributes.
     if (!CheckAttrs(strided_slice_grad)) {
       MS_LOG(INFO) << "Check strided_slice_grad's attrs failed, graph not changed";
       return nullptr;
     }
 
     ValuePtrList strides_values;
+    // Extract the stride values.
     if (!GetStridesValues(strided_slice_grad, &strides_values)) {
       return nullptr;
     }
 
+    // Validate the extracted stride values.
     if (!CheckValues(strides_values)) {
       MS_LOG(INFO) << "Check strides' values failed, graph not changed";
       return nullptr;
     }
   }
 
+  // Convert node's constant inputs to attributes.
   ConstInputToAttr(strided_slice_grad, {1, 2, 3, 4});
   return nullptr;
 }
+
+}  // end of the namespace
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/const_to_attr_strided_slice_grad_old.cc

@@ -0,0 +1,237 @@
+/**
+ * Copyright 2020-2022 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "backend/common/pass/const_to_attr_strided_slice_grad.h"
+#include <memory>
+#include <vector>
+#include <string>
+#include "ir/primitive.h"
+#include "utils/ms_context.h"
+#include "include/common/utils/utils.h"
+#include "abstract/abstract_value.h"
+#include "backend/common/optimizer/const_input_to_attr.h"
+#include "include/common/utils/anfalgo.h"
+
+namespace mindspore {
+namespace opt {
+namespace {
+const size_t strides_index = 5;
+
+/**
+ * @brief Retrieves the stride values from the StridedSliceGrad node.
+ *
+ * This function is designed to extract the stride values from the given StridedSliceGrad node in MindSpore.
+ * It primarily checks the validity of the node, extracts the desired stride values, and returns whether the operation was successful.
+ *
+ * @param strided_slice_grad The given StridedSliceGrad node.
+ * @param strides_values Pointer to the list where the stride values will be stored.
+ * @return True if the stride values were successfully retrieved, otherwise False.
+ */
+bool GetStridesValues(const CNodePtr &strided_slice_grad, ValuePtrList *strides_values) {
+  // Check for null inputs.
+  MS_EXCEPTION_IF_NULL(strided_slice_grad);
+  MS_EXCEPTION_IF_NULL(strides_values);
+
+  constexpr size_t kSizeChange = 6;
+
+  // Ensure the strided_slice_grad node has the expected size.
+  if (strided_slice_grad->size() < kSizeChange) {
+    MS_LOG(DEBUG) << "Op strided_slice_grad's inputs size less than 6, graph not changed";
+    return false;
+  }
+
+  // Extract the strides input from the strided_slice_grad node.
+  auto strides_input = strided_slice_grad->input(strides_index);
+  MS_EXCEPTION_IF_NULL(strides_input);
+
+  // Cast the strides input to a value node.
+  auto strides_value_node = strides_input->cast<ValueNodePtr>();
+  if (strides_value_node == nullptr) {
+    MS_LOG(DEBUG) << "strides is not a value node.";
+    return false;
+  }
+
+  // Retrieve the value of the strides_value_node.
+  auto value = strides_value_node->value();
+  if (value == nullptr) {
+    MS_LOG(DEBUG) << "strides has no value.";
+    return false;
+  }
+
+  // Cast the retrieved value to a tuple.
+  auto value_tuple = value->cast<ValueTuplePtr>();
+  if (value_tuple == nullptr) {
+    MS_LOG(DEBUG) << "strides is not a value tuple.";
+    return false;
+  }
+
+  // Assign the retrieved tuple value to the strides_values pointer.
+  *strides_values = value_tuple->value();
+
+  return true;
+}
+
+/**
+ * @brief Checks the values of the given strides.
+ *
+ * This function checks whether the provided stride values meet specific conditions.
+ * It validates if each stride value is a scalar of integer type and equal to 1, which seems to be a requirement for StridedSliceGrad in MindSpore.
+ *
+ * @param strides_values List containing the stride values.
+ * @return True if all stride values meet the required conditions, otherwise False.
+ */
+bool CheckValues(const ValuePtrList &strides_values) {
+  // Check for an empty strides_values list.
+  if (strides_values.empty()) {
+    MS_LOG(DEBUG) << "strides_values is empty";
+    return false;
+  }
+
+  // Iterate over the stride values and check each value.
+  for (auto &value : strides_values) {
+    MS_EXCEPTION_IF_NULL(value);
+    if (value->isa<Scalar>()) {
+      auto scalar = value->cast<ScalarPtr>();
+      MS_EXCEPTION_IF_NULL(scalar);
+
+      // Check if the scalar value is an integer and equals 1.
+      if (scalar->isa<Int32Imm>()) {
+        if (GetValue<int>(scalar) != 1) {
+          MS_LOG(DEBUG) << "StridedSliceGrad has no 1 value";
+          return false;
+        }
+      } else if (scalar->isa<Int64Imm>()) {
+        if (GetValue<int64_t>(scalar) != 1) {
+          MS_LOG(DEBUG) << "StridedSliceGrad has no 1 value";
+          return false;
+        }
+      } else {
+        MS_LOG(DEBUG) << "Strides value is not an integer";
+        return false;
+      }
+    } else {
+      MS_LOG(DEBUG) << "The value " << value << "of tuple is not a scalar";
+      return false;
+    }
+  }
+
+  return true;
+}
+
+
+/**
+ * @brief Checks the attributes of the StridedSliceGrad node.
+ * 
+ * This function verifies the existence and values of the `new_axis_mask` and `shrink_axis_mask` attributes
+ * in the given StridedSliceGrad node.
+ *
+ * @param strided_slice_grad The StridedSliceGrad node to check.
+ * @return True if both attributes exist and their values are 0; otherwise, False.
+ */
+bool CheckAttrs(const CNodePtr &strided_slice_grad) {
+  // Ensure strided_slice_grad is not null.
+  MS_EXCEPTION_IF_NULL(strided_slice_grad);
+
+  // Check for the existence of the required attributes.
+  if (!common::AnfAlgo::HasNodeAttr(kAttrNewAxisMask, strided_slice_grad) ||
+      !common::AnfAlgo::HasNodeAttr(kAttrShrinkAxisMask, strided_slice_grad)) {
+    MS_LOG(INFO) << "new_axis_mask or shrink_axis_mask not exist in cnode[" + strided_slice_grad->DebugString() + "]";
+    return false;
+  }
+
+  // Extract the values of the required attributes.
+  auto new_axis_mask = common::AnfAlgo::GetNodeAttr<int64_t>(strided_slice_grad, kAttrNewAxisMask);
+  auto shrink_axis_mask = common::AnfAlgo::GetNodeAttr<int64_t>(strided_slice_grad, kAttrShrinkAxisMask);
+
+  // Check the attribute values.
+  if (new_axis_mask != 0 || shrink_axis_mask != 0) {
+    MS_LOG(INFO) << "new_axis_mask or shrink_axis_mask not equal 0";
+    return false;
+  }
+
+  return true;
+}
+
+// Namespace wrapping might be in the original code. Assuming it's still the case.
+namespace {
+
+/**
+ * @brief Defines the pattern for the ConstToAttrStridedSliceGradPass.
+ *
+ * This function sets up the pattern to identify the StridedSliceGrad operations that should be 
+ * processed by the ConstToAttrStridedSliceGradPass.
+ * 
+ * @return The pattern for identifying the appropriate nodes.
+ */
+const BaseRef ConstToAttrStridedSliceGradPass::DefinePattern() const {
+  VarPtr Xs = std::make_shared<SeqVar>();
+  auto strided_slice_grad_prim = std::make_shared<Primitive>(kStridedSliceGradOpName);
+  return VectorRef({strided_slice_grad_prim, Xs});
+}
+
+/**
+ * @brief Processes the identified StridedSliceGrad node.
+ * 
+ * For the StridedSliceGrad nodes that match the previously defined pattern, this function checks 
+ * attributes and strides' values. If they satisfy certain conditions (based on the device and other factors), 
+ * the node's inputs are then converted to attributes.
+ *
+ * @param graph The functional graph containing the node.
+ * @param node The StridedSliceGrad node to process.
+ * @param [unused] The equivalence class of the node. (currently unused in this function)
+ * @return nullptr, indicating the node was processed (or not modified).
+ */
+const AnfNodePtr ConstToAttrStridedSliceGradPass::Process(const FuncGraphPtr &graph, const AnfNodePtr &node,
+                                                          const EquivPtr &) const {
+  // Ensure inputs are not null.
+  MS_EXCEPTION_IF_NULL(graph);
+  MS_EXCEPTION_IF_NULL(node);
+
+  auto strided_slice_grad = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(strided_slice_grad);
+
+  auto ms_context = MsContext::GetInstance();
+  MS_EXCEPTION_IF_NULL(ms_context);
+
+  // Check if the processing is specific to Ascend device.
+  if (ms_context->get_param<std::string>(MS_CTX_DEVICE_TARGET) == kAscendDevice) {
+    // Validate the node's attributes.
+    if (!CheckAttrs(strided_slice_grad)) {
+      MS_LOG(INFO) << "Check strided_slice_grad's attrs failed, graph not changed";
+      return nullptr;
+    }
+
+    ValuePtrList strides_values;
+    // Extract the stride values.
+    if (!GetStridesValues(strided_slice_grad, &strides_values)) {
+      return nullptr;
+    }
+
+    // Validate the extracted stride values.
+    if (!CheckValues(strides_values)) {
+      MS_LOG(INFO) << "Check strides' values failed, graph not changed";
+      return nullptr;
+    }
+  }
+
+  // Convert node's constant inputs to attributes.
+  ConstInputToAttr(strided_slice_grad, {1, 2, 3, 4});
+  return nullptr;
+}
+
+}  // end of the namespace
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/conv_transpose_to_conv_bp.cc

@@ -27,29 +27,67 @@ namespace {
 constexpr size_t kCNodePrimitiveIdx = 0;
 }
 
+/**
+ * @brief This class aims to transform the convolution transpose (Deconvolution) operator 
+ * into a convolution backpropagation to the input operator.
+ */
+
+/**
+ * @brief Defines the pattern to be matched for this transformation.
+ *
+ * The pattern looks for a sequence of nodes where the main operation is a Convolution Transpose operation.
+ * 
+ * @return A VectorRef containing the pattern to be matched.
+ */
 const BaseRef ConvTransposeToConvBackpropInputPass::DefinePattern() const {
+  // A variable representing a sequence of nodes.
   VarPtr Xs = std::make_shared<SeqVar>();
+
+  // Represents the Conv2DTranspose operation.
   auto conv_transpose = std::make_shared<Primitive>(kConv2DTransposeOpName);
+
+  // Define the pattern: The Conv2DTranspose operation followed by any sequence of nodes.
   return VectorRef({conv_transpose, Xs});
 }
 
+/**
+ * @brief Processes the matched node and transforms it as required.
+ *
+ * If the provided node matches the pattern (i.e., it is a Convolution Transpose operation), 
+ * this function renames it to Convolution Backpropagation to the input operation.
+ * 
+ * @param graph The function graph that the node belongs to.
+ * @param node The matched node to be processed.
+ * @param An equivalence mapping (not used in this function but retained for consistency).
+ * 
+ * @return The processed node (either transformed or untouched).
+ */
 const AnfNodePtr ConvTransposeToConvBackpropInputPass::Process(const FuncGraphPtr &graph, const AnfNodePtr &node,
                                                                const EquivPtr &) const {
+  // Basic null checks.
   MS_EXCEPTION_IF_NULL(graph);
   MS_EXCEPTION_IF_NULL(node);
+
+  // Cast the node to a computational node.
   auto conv_transpose = node->cast<CNodePtr>();
   MS_EXCEPTION_IF_NULL(conv_transpose);
 
+  // Ensure the computational node has inputs.
   if (conv_transpose->inputs().empty()) {
     MS_LOG(EXCEPTION) << "Cnode inputs should not be empty, cnode: " << node->DebugString()
                       << trace::DumpSourceLines(conv_transpose);
   }
 
+  // Extract the main operation of the computational node.
   auto prim = GetValueNode<PrimitivePtr>(conv_transpose->input(kCNodePrimitiveIdx));
   MS_EXCEPTION_IF_NULL(prim);
+
+  // Rename the operation to Conv2DBackpropInput (Convolution Backpropagation to the input).
   prim->Named::operator=(Named(kConv2DBackpropInputOpName));
 
+  // Return the modified node.
   return node;
 }
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/conv_transpose_to_conv_bp_old.cc

@@ -0,0 +1,93 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "backend/common/pass/conv_transpose_to_conv_bp.h"
+#include <memory>
+#include <vector>
+#include "ir/primitive.h"
+#include "include/common/utils/utils.h"
+#include "utils/trace_base.h"
+#include "backend/common/optimizer/helper.h"
+
+namespace mindspore {
+namespace opt {
+namespace {
+constexpr size_t kCNodePrimitiveIdx = 0;
+}
+
+/**
+ * @brief This class aims to transform the convolution transpose (Deconvolution) operator 
+ * into a convolution backpropagation to the input operator.
+ */
+
+/**
+ * @brief Defines the pattern to be matched for this transformation.
+ *
+ * The pattern looks for a sequence of nodes where the main operation is a Convolution Transpose operation.
+ * 
+ * @return A VectorRef containing the pattern to be matched.
+ */
+const BaseRef ConvTransposeToConvBackpropInputPass::DefinePattern() const {
+  // A variable representing a sequence of nodes.
+  VarPtr Xs = std::make_shared<SeqVar>();
+
+  // Represents the Conv2DTranspose operation.
+  auto conv_transpose = std::make_shared<Primitive>(kConv2DTransposeOpName);
+
+  // Define the pattern: The Conv2DTranspose operation followed by any sequence of nodes.
+  return VectorRef({conv_transpose, Xs});
+}
+
+/**
+ * @brief Processes the matched node and transforms it as required.
+ *
+ * If the provided node matches the pattern (i.e., it is a Convolution Transpose operation), 
+ * this function renames it to Convolution Backpropagation to the input operation.
+ * 
+ * @param graph The function graph that the node belongs to.
+ * @param node The matched node to be processed.
+ * @param An equivalence mapping (not used in this function but retained for consistency).
+ * 
+ * @return The processed node (either transformed or untouched).
+ */
+const AnfNodePtr ConvTransposeToConvBackpropInputPass::Process(const FuncGraphPtr &graph, const AnfNodePtr &node,
+                                                               const EquivPtr &) const {
+  // Basic null checks.
+  MS_EXCEPTION_IF_NULL(graph);
+  MS_EXCEPTION_IF_NULL(node);
+
+  // Cast the node to a computational node.
+  auto conv_transpose = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(conv_transpose);
+
+  // Ensure the computational node has inputs.
+  if (conv_transpose->inputs().empty()) {
+    MS_LOG(EXCEPTION) << "Cnode inputs should not be empty, cnode: " << node->DebugString()
+                      << trace::DumpSourceLines(conv_transpose);
+  }
+
+  // Extract the main operation of the computational node.
+  auto prim = GetValueNode<PrimitivePtr>(conv_transpose->input(kCNodePrimitiveIdx));
+  MS_EXCEPTION_IF_NULL(prim);
+
+  // Rename the operation to Conv2DBackpropInput (Convolution Backpropagation to the input).
+  prim->Named::operator=(Named(kConv2DBackpropInputOpName));
+
+  // Return the modified node.
+  return node;
+}
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/convert_attr_to_unify_mindir.cc

@@ -23,27 +23,49 @@
 
 namespace mindspore {
 namespace opt {
+/**
+ * @brief Converts attributes of a node to a unified format for MindIR.
+ *
+ * This function checks a node's attributes and, if necessary, converts them into a format 
+ * that is suitable for MindIR representation. This standardization aids in serialization
+ * and further processing.
+ *
+ * @param node The node whose attributes are to be processed and converted.
+ * 
+ * @return The processed CNode with converted attributes or nullptr if the node isn't suitable.
+ */
 const AnfNodePtr ConvertAttrToUnifyMindIR::Process(const FuncGraphPtr &, const AnfNodePtr &node,
                                                    const EquivPtr &) const {
+  // Check if the node is valid and is a real computational node.
   if (node == nullptr || !AnfUtils::IsRealCNodeKernel(node)) {
     return nullptr;
   }
 
   auto cnode = node->cast<CNodePtr>();
   MS_EXCEPTION_IF_NULL(cnode);
+  
+  // Extract the operator from the computational node.
   auto inputs = cnode->inputs();
   AnfNodePtr op = inputs[0];
   MS_EXCEPTION_IF_NULL(op);
+  
+  // If the operator is a primitive value, process its attributes.
   if (IsValueNode<Primitive>(op)) {
     auto prim = GetValueNode<PrimitivePtr>(op);
     MS_EXCEPTION_IF_NULL(prim);
+    
     auto attrs = prim->attrs();
     std::string type_name = prim->name();
+    
+    // Iterate over the attributes of the primitive.
     for (auto attr : attrs) {
+      // Attempt to convert the attribute value to a string representation.
       bool converted = CheckAndConvertUtils::ConvertAttrValueToString(type_name, attr.first, &attr.second);
       if (converted) {
         prim->set_attr(attr.first, attr.second);
       }
+
+      // Attempt to convert intermediate representation attributes to operator attributes.
       bool converted_ir_attr = CheckAndConvertUtils::CheckIrAttrtoOpAttr(type_name, attr.first, &attr.second);
       if (converted_ir_attr) {
         prim->set_attr(attr.first, attr.second);
@@ -53,5 +75,6 @@ const AnfNodePtr ConvertAttrToUnifyMindIR::Process(const FuncGraphPtr &, const A
 
   return node;
 }
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/convert_attr_to_unify_mindir_old.cc

@@ -0,0 +1,80 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "backend/common/pass/convert_attr_to_unify_mindir.h"
+
+#include <string>
+#include "utils/check_convert_utils.h"
+#include "backend/common/session/anf_runtime_algorithm.h"
+#include "include/common/utils/anfalgo.h"
+#include "kernel/common_utils.h"
+
+namespace mindspore {
+namespace opt {
+/**
+ * @brief Converts attributes of a node to a unified format for MindIR.
+ *
+ * This function checks a node's attributes and, if necessary, converts them into a format 
+ * that is suitable for MindIR representation. This standardization aids in serialization
+ * and further processing.
+ *
+ * @param node The node whose attributes are to be processed and converted.
+ * 
+ * @return The processed CNode with converted attributes or nullptr if the node isn't suitable.
+ */
+const AnfNodePtr ConvertAttrToUnifyMindIR::Process(const FuncGraphPtr &, const AnfNodePtr &node,
+                                                   const EquivPtr &) const {
+  // Check if the node is valid and is a real computational node.
+  if (node == nullptr || !AnfUtils::IsRealCNodeKernel(node)) {
+    return nullptr;
+  }
+
+  auto cnode = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(cnode);
+  
+  // Extract the operator from the computational node.
+  auto inputs = cnode->inputs();
+  AnfNodePtr op = inputs[0];
+  MS_EXCEPTION_IF_NULL(op);
+  
+  // If the operator is a primitive value, process its attributes.
+  if (IsValueNode<Primitive>(op)) {
+    auto prim = GetValueNode<PrimitivePtr>(op);
+    MS_EXCEPTION_IF_NULL(prim);
+    
+    auto attrs = prim->attrs();
+    std::string type_name = prim->name();
+    
+    // Iterate over the attributes of the primitive.
+    for (auto attr : attrs) {
+      // Attempt to convert the attribute value to a string representation.
+      bool converted = CheckAndConvertUtils::ConvertAttrValueToString(type_name, attr.first, &attr.second);
+      if (converted) {
+        prim->set_attr(attr.first, attr.second);
+      }
+
+      // Attempt to convert intermediate representation attributes to operator attributes.
+      bool converted_ir_attr = CheckAndConvertUtils::CheckIrAttrtoOpAttr(type_name, attr.first, &attr.second);
+      if (converted_ir_attr) {
+        prim->set_attr(attr.first, attr.second);
+      }
+    }
+  }
+
+  return node;
+}
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/convert_const_input_to_attr.cc

@@ -24,53 +24,61 @@
 
 namespace mindspore {
 namespace opt {
+/**
+ * @brief Converts certain constant inputs of a computational node to node attributes.
+ *
+ * For specific operators, this function checks if they have constant inputs that need to be 
+ * converted to attributes. Such conversions are essential for some backend compilations.
+ *
+ * @param node The node to be processed.
+ * 
+ * @return Processed CNode or nullptr if no conversion occurred.
+ */
 const AnfNodePtr ConvertConstInputToAttr::Process(const FuncGraphPtr &, const AnfNodePtr &node,
                                                   const EquivPtr &) const {
+  // Check if the node is valid and is a real computational node.
   if (node == nullptr || !AnfUtils::IsRealCNodeKernel(node)) {
     return nullptr;
   }
 
   auto cnode = node->cast<CNodePtr>();
   MS_EXCEPTION_IF_NULL(cnode);
+
+  // Obtain the relevant input-to-attribute conversion registry for the node's operator.
   ConstInputToAttrInfoRegister reg;
   if (!ConstInputToAttrInfoRegistry::Instance().GetRegisterByOpName(common::AnfAlgo::GetCNodeName(cnode), &reg)) {
     return nullptr;
   }
+
+  // Specific checks for certain node types.
   if (common::AnfAlgo::GetCNodeName(cnode) == prim::kPrimEmbeddingLookup->name() ||
       common::AnfAlgo::GetCNodeName(cnode) == prim::kPrimEmbeddingLookupCommGrad->name()) {
     if (!common::AnfAlgo::HasNodeAttr(kAttrPrimitiveTarget, cnode)) {
       return nullptr;
     }
   }
+
   auto ms_context = MsContext::GetInstance();
   MS_EXCEPTION_IF_NULL(ms_context);
+
+  // Get the device target from context.
   auto device = ms_context->get_param<std::string>(MS_CTX_DEVICE_TARGET);
   if (common::AnfAlgo::GetCNodeName(cnode) == prim::kPrimGatherD->name()) {
     if (device != kGPUDevice) {
       return nullptr;
     }
   }
+
+  // Check if the node handles dynamic shapes.
   if (common::AnfAlgo::IsDynamicShape(cnode)) {
-    if (device == kGPUDevice) {
-      if (DynamicShapeConstInputToAttrGPU.find(common::AnfAlgo::GetCNodeName(cnode)) ==
-          DynamicShapeConstInputToAttrGPU.end()) {
-        MS_LOG(INFO) << "current node is dynamic shape " << cnode->fullname_with_scope();
-        return nullptr;
-      }
-    } else if (device == kCPUDevice) {
-      if (DynamicShapeConstInputToAttrCPU.find(common::AnfAlgo::GetCNodeName(cnode)) ==
-          DynamicShapeConstInputToAttrCPU.end()) {
-        MS_LOG(INFO) << "current node is dynamic shape " << cnode->fullname_with_scope();
-        return nullptr;
-      }
-    } else {
-      if (DynamicShapeConstInputToAttr.find(common::AnfAlgo::GetCNodeName(cnode)) ==
-          DynamicShapeConstInputToAttr.end()) {
-        MS_LOG(INFO) << "current node is dynamic shape " << cnode->fullname_with_scope();
-        return nullptr;
-      }
+    // Ensure the operator can handle dynamic shape based on the device.
+    if (!IsSupportedDynamicShapeOp(device, common::AnfAlgo::GetCNodeName(cnode))) {
+      MS_LOG(INFO) << "current node is dynamic shape " << cnode->fullname_with_scope();
+      return nullptr;
     }
   }
+
+  // Specific checks for Ascend device.
   if (device == kAscendDevice &&
       NeedConvertToValueNodeSet.find(common::AnfAlgo::GetCNodeName(cnode)) != NeedConvertToValueNodeSet.end() &&
       !common::AnfAlgo::HasNodeAttr(kAttrNeedConvertToValueNode, cnode)) {
@@ -82,9 +90,22 @@ const AnfNodePtr ConvertConstInputToAttr::Process(const FuncGraphPtr &, const An
     return nullptr;
   }
 
+  // Convert the designated constant inputs to attributes.
   ConstInputToAttr(cnode, reg.GetConstInputAttrInfo());
 
   return node;
 }
+
+// Helper function: Check if an operator can handle dynamic shape for a specific device.
+bool IsSupportedDynamicShapeOp(const std::string &device, const std::string &op_name) {
+  if (device == kGPUDevice) {
+    return DynamicShapeConstInputToAttrGPU.find(op_name) != DynamicShapeConstInputToAttrGPU.end();
+  } else if (device == kCPUDevice) {
+    return DynamicShapeConstInputToAttrCPU.find(op_name) != DynamicShapeConstInputToAttrCPU.end();
+  } else {
+    return DynamicShapeConstInputToAttr.find(op_name) != DynamicShapeConstInputToAttr.end();
+  }
+}
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/convert_const_input_to_attr_old.cc

@@ -0,0 +1,111 @@
+/**
+ * Copyright 2020-2022 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "backend/common/pass/convert_const_input_to_attr.h"
+#include <vector>
+#include <algorithm>
+#include "backend/common/optimizer/const_input_to_attr.h"
+#include "include/common/utils/utils.h"
+#include "utils/ms_context.h"
+#include "base/core_ops.h"
+#include "include/common/utils/anfalgo.h"
+
+namespace mindspore {
+namespace opt {
+/**
+ * @brief Converts certain constant inputs of a computational node to node attributes.
+ *
+ * For specific operators, this function checks if they have constant inputs that need to be 
+ * converted to attributes. Such conversions are essential for some backend compilations.
+ *
+ * @param node The node to be processed.
+ * 
+ * @return Processed CNode or nullptr if no conversion occurred.
+ */
+const AnfNodePtr ConvertConstInputToAttr::Process(const FuncGraphPtr &, const AnfNodePtr &node,
+                                                  const EquivPtr &) const {
+  // Check if the node is valid and is a real computational node.
+  if (node == nullptr || !AnfUtils::IsRealCNodeKernel(node)) {
+    return nullptr;
+  }
+
+  auto cnode = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(cnode);
+
+  // Obtain the relevant input-to-attribute conversion registry for the node's operator.
+  ConstInputToAttrInfoRegister reg;
+  if (!ConstInputToAttrInfoRegistry::Instance().GetRegisterByOpName(common::AnfAlgo::GetCNodeName(cnode), &reg)) {
+    return nullptr;
+  }
+
+  // Specific checks for certain node types.
+  if (common::AnfAlgo::GetCNodeName(cnode) == prim::kPrimEmbeddingLookup->name() ||
+      common::AnfAlgo::GetCNodeName(cnode) == prim::kPrimEmbeddingLookupCommGrad->name()) {
+    if (!common::AnfAlgo::HasNodeAttr(kAttrPrimitiveTarget, cnode)) {
+      return nullptr;
+    }
+  }
+
+  auto ms_context = MsContext::GetInstance();
+  MS_EXCEPTION_IF_NULL(ms_context);
+
+  // Get the device target from context.
+  auto device = ms_context->get_param<std::string>(MS_CTX_DEVICE_TARGET);
+  if (common::AnfAlgo::GetCNodeName(cnode) == prim::kPrimGatherD->name()) {
+    if (device != kGPUDevice) {
+      return nullptr;
+    }
+  }
+
+  // Check if the node handles dynamic shapes.
+  if (common::AnfAlgo::IsDynamicShape(cnode)) {
+    // Ensure the operator can handle dynamic shape based on the device.
+    if (!IsSupportedDynamicShapeOp(device, common::AnfAlgo::GetCNodeName(cnode))) {
+      MS_LOG(INFO) << "current node is dynamic shape " << cnode->fullname_with_scope();
+      return nullptr;
+    }
+  }
+
+  // Specific checks for Ascend device.
+  if (device == kAscendDevice &&
+      NeedConvertToValueNodeSet.find(common::AnfAlgo::GetCNodeName(cnode)) != NeedConvertToValueNodeSet.end() &&
+      !common::AnfAlgo::HasNodeAttr(kAttrNeedConvertToValueNode, cnode)) {
+    auto input_attrs = reg.GetConstInputAttrInfo();
+    std::vector<size_t> need_convert_to_constant;
+    std::transform(input_attrs.begin(), input_attrs.end(), std::back_inserter(need_convert_to_constant),
+                   [](size_t i) { return i + 1; });
+    common::AnfAlgo::SetNodeAttr(kAttrNeedConvertToValueNode, MakeValue(need_convert_to_constant), cnode);
+    return nullptr;
+  }
+
+  // Convert the designated constant inputs to attributes.
+  ConstInputToAttr(cnode, reg.GetConstInputAttrInfo());
+
+  return node;
+}
+
+// Helper function: Check if an operator can handle dynamic shape for a specific device.
+bool IsSupportedDynamicShapeOp(const std::string &device, const std::string &op_name) {
+  if (device == kGPUDevice) {
+    return DynamicShapeConstInputToAttrGPU.find(op_name) != DynamicShapeConstInputToAttrGPU.end();
+  } else if (device == kCPUDevice) {
+    return DynamicShapeConstInputToAttrCPU.find(op_name) != DynamicShapeConstInputToAttrCPU.end();
+  } else {
+    return DynamicShapeConstInputToAttr.find(op_name) != DynamicShapeConstInputToAttr.end();
+  }
+}
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/convert_const_input_to_tensor_input.cc

@@ -29,55 +29,109 @@
 namespace mindspore {
 namespace opt {
 namespace {
+/**
+ * @brief Converts a ValueNode containing a Scalar or ValueTuple to a tensor node.
+ *
+ * If the provided node contains a Scalar or ValueTuple value, this function 
+ * will convert that value into a Tensor and encapsulate it in a new ValueNode.
+ *
+ * @param kernel_graph The computational graph the node belongs to.
+ * @param input_node The node containing Scalar or ValueTuple to be converted.
+ *
+ * @return A ValueNode containing the Tensor representation or nullptr if conversion fails.
+ */
 AnfNodePtr CreateTensorInput(const KernelGraphPtr &kernel_graph, const AnfNodePtr &input_node) {
+  // Validate if the provided node is not null.
   MS_EXCEPTION_IF_NULL(input_node);
+  
   auto value_node = input_node->cast<ValueNodePtr>();
   MS_EXCEPTION_IF_NULL(value_node);
+
+  // Extract the value held by the ValueNode.
   auto value = value_node->value();
   MS_EXCEPTION_IF_NULL(value);
+
   tensor::TensorPtr tensor_ptr = nullptr;
+
+  // Convert the value to a tensor.
+  // If it's a Scalar, use ScalarToTensor function.
   if (value->isa<Scalar>()) {
     tensor_ptr = ScalarToTensor(value->cast<ScalarPtr>());
-  } else if (value->isa<ValueTuple>()) {
+  } 
+  // If it's a ValueTuple, use CreateTupleTensor function.
+  else if (value->isa<ValueTuple>()) {
     tensor_ptr = CreateTupleTensor(value->cast<ValueTuplePtr>());
-  } else {
+  } 
+  // Throw an exception if the value is neither Scalar nor ValueTuple.
+  else {
     MS_LOG(EXCEPTION) << "The value should be a scalar or value tuple";
   }
+
+  // If tensor conversion fails, return nullptr.
   if (tensor_ptr == nullptr) {
     MS_LOG(DEBUG) << "Create tensor failed";
     return nullptr;
   }
+
+  // Wrap the tensor in a new ValueNode.
   auto tensor_input = std::make_shared<ValueNode>(tensor_ptr);
   MS_EXCEPTION_IF_NULL(tensor_input);
   tensor_input->set_abstract(tensor_ptr->ToAbstract());
+
+  // If a KernelGraph is provided, add the new ValueNode to the graph.
   if (kernel_graph != nullptr) {
     tensor_input = kernel_graph->NewValueNode(tensor_input);
     kernel_graph->AddValueNodeToGraph(tensor_input);
   } else {
     tensor_input = MakeValueNode(tensor_input);
   }
+
+  // Assign the original node's scope to the new ValueNode.
   tensor_input->set_scope(input_node->scope());
+
   return tensor_input;
 }
+
 }  // namespace
 
+/**
+ * @brief Convert const inputs of a CNode to tensor inputs.
+ *
+ * If a given computational node has Scalar or ValueTuple as inputs, 
+ * this function ensures those inputs are converted into Tensors.
+ *
+ * @param func_graph The computational graph the node belongs to.
+ * @param cnode The node whose const inputs need to be converted.
+ *
+ * @return A new CNode with converted tensor inputs or nullptr if no conversion occurred.
+ */
 AnfNodePtr ConvertConstInputToTensorInput::ConstInputToTensorInput(const FuncGraphPtr &func_graph,
                                                                    const CNodePtr &cnode) const {
   MS_EXCEPTION_IF_NULL(func_graph);
   MS_EXCEPTION_IF_NULL(cnode);
+
+  // A set of node types that shouldn't be converted.
   const std::set<std::string> no_need_to_convert_nodes = {kStackOpName};
+
   auto node_type = common::AnfAlgo::GetCNodeName(cnode);
   if (no_need_to_convert_nodes.find(node_type) != no_need_to_convert_nodes.end()) {
     return nullptr;
   }
+
   std::vector<AnfNodePtr> new_inputs;
   auto kernel_graph = func_graph->cast<std::shared_ptr<session::KernelGraph>>();
   auto inputs = cnode->inputs();
+
+  // Keep the first input (typically the operator node) unchanged.
   new_inputs.push_back(inputs[0]);
+
   bool need_update = false;
-  // the first input is primitive node which is not the real input
+
+  // Check each input for Scalar or ValueTuple.
   for (size_t i = 0; i < inputs.size() - 1; ++i) {
     auto input_node = inputs[i + 1];
+
+    // Convert Scalar or ValueTuple inputs to Tensor inputs.
     if (IsValueNode<Scalar>(input_node) || IsValueNode<ValueTuple>(input_node)) {
       auto tensor_input = CreateTensorInput(kernel_graph, input_node);
       if (tensor_input == nullptr) {
@@ -90,8 +144,9 @@ AnfNodePtr ConvertConstInputToTensorInput::ConstInputToTensorInput(const FuncGra
       new_inputs.push_back(input_node);
     }
   }
+
+  // If any inputs were converted, create a new CNode with updated inputs.
   if (need_update) {
-    MS_EXCEPTION_IF_NULL(func_graph);
     auto new_cnode = NewCNode(new_inputs, func_graph);
     MS_EXCEPTION_IF_NULL(new_cnode);
     if (common::AnfAlgo::CheckPrimitiveType(cnode, prim::kPrimDepend)) {
@@ -109,17 +164,29 @@ AnfNodePtr ConvertConstInputToTensorInput::ConstInputToTensorInput(const FuncGra
   return nullptr;
 }
 
+/**
+ * @brief Entry function to initiate the conversion of const inputs to tensor inputs.
+ *
+ * @param func_graph The computational graph the node belongs to.
+ * @param node The node to be processed.
+ * @param equiv Placeholder for equivalence classes in pattern matching (unused in this function).
+ *
+ * @return Processed CNode or nullptr if no conversion occurred.
+ */
 const AnfNodePtr ConvertConstInputToTensorInput::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
                                                          const EquivPtr &) const {
-  if (node == nullptr || func_graph == nullptr || common::AnfAlgo::CheckPrimitiveType(node, prim::kPrimTupleGetItem) ||
+  if (node == nullptr || func_graph == nullptr ||
+      common::AnfAlgo::CheckPrimitiveType(node, prim::kPrimTupleGetItem) ||
       common::AnfAlgo::CheckPrimitiveType(node, prim::kPrimMakeTuple)) {
     return nullptr;
   }
+
   if (!node->isa<CNode>()) {
     return nullptr;
   }
 
   return ConstInputToTensorInput(func_graph, node->cast<CNodePtr>());
 }
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/convert_const_input_to_tensor_input_old.cc

@@ -0,0 +1,192 @@
+/**
+ * Copyright 2020-2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "backend/common/pass/convert_const_input_to_tensor_input.h"
+
+#include <vector>
+#include <memory>
+#include <set>
+
+#include "ir/graph_utils.h"
+#include "backend/common/optimizer/helper.h"
+#include "backend/common/session/anf_runtime_algorithm.h"
+#include "include/common/utils/anfalgo.h"
+#include "backend/common/session/kernel_graph.h"
+#include "kernel/common_utils.h"
+
+namespace mindspore {
+namespace opt {
+namespace {
+/**
+ * @brief Converts a ValueNode containing a Scalar or ValueTuple to a tensor node.
+ *
+ * If the provided node contains a Scalar or ValueTuple value, this function 
+ * will convert that value into a Tensor and encapsulate it in a new ValueNode.
+ *
+ * @param kernel_graph The computational graph the node belongs to.
+ * @param input_node The node containing Scalar or ValueTuple to be converted.
+ *
+ * @return A ValueNode containing the Tensor representation or nullptr if conversion fails.
+ */
+AnfNodePtr CreateTensorInput(const KernelGraphPtr &kernel_graph, const AnfNodePtr &input_node) {
+  // Validate if the provided node is not null.
+  MS_EXCEPTION_IF_NULL(input_node);
+  
+  auto value_node = input_node->cast<ValueNodePtr>();
+  MS_EXCEPTION_IF_NULL(value_node);
+
+  // Extract the value held by the ValueNode.
+  auto value = value_node->value();
+  MS_EXCEPTION_IF_NULL(value);
+
+  tensor::TensorPtr tensor_ptr = nullptr;
+
+  // Convert the value to a tensor.
+  // If it's a Scalar, use ScalarToTensor function.
+  if (value->isa<Scalar>()) {
+    tensor_ptr = ScalarToTensor(value->cast<ScalarPtr>());
+  } 
+  // If it's a ValueTuple, use CreateTupleTensor function.
+  else if (value->isa<ValueTuple>()) {
+    tensor_ptr = CreateTupleTensor(value->cast<ValueTuplePtr>());
+  } 
+  // Throw an exception if the value is neither Scalar nor ValueTuple.
+  else {
+    MS_LOG(EXCEPTION) << "The value should be a scalar or value tuple";
+  }
+
+  // If tensor conversion fails, return nullptr.
+  if (tensor_ptr == nullptr) {
+    MS_LOG(DEBUG) << "Create tensor failed";
+    return nullptr;
+  }
+
+  // Wrap the tensor in a new ValueNode.
+  auto tensor_input = std::make_shared<ValueNode>(tensor_ptr);
+  MS_EXCEPTION_IF_NULL(tensor_input);
+  tensor_input->set_abstract(tensor_ptr->ToAbstract());
+
+  // If a KernelGraph is provided, add the new ValueNode to the graph.
+  if (kernel_graph != nullptr) {
+    tensor_input = kernel_graph->NewValueNode(tensor_input);
+    kernel_graph->AddValueNodeToGraph(tensor_input);
+  } else {
+    tensor_input = MakeValueNode(tensor_input);
+  }
+
+  // Assign the original node's scope to the new ValueNode.
+  tensor_input->set_scope(input_node->scope());
+
+  return tensor_input;
+}
+
+}  // namespace
+
+/**
+ * @brief Convert const inputs of a CNode to tensor inputs.
+ *
+ * If a given computational node has Scalar or ValueTuple as inputs, 
+ * this function ensures those inputs are converted into Tensors.
+ *
+ * @param func_graph The computational graph the node belongs to.
+ * @param cnode The node whose const inputs need to be converted.
+ *
+ * @return A new CNode with converted tensor inputs or nullptr if no conversion occurred.
+ */
+AnfNodePtr ConvertConstInputToTensorInput::ConstInputToTensorInput(const FuncGraphPtr &func_graph,
+                                                                   const CNodePtr &cnode) const {
+  MS_EXCEPTION_IF_NULL(func_graph);
+  MS_EXCEPTION_IF_NULL(cnode);
+
+  // A set of node types that shouldn't be converted.
+  const std::set<std::string> no_need_to_convert_nodes = {kStackOpName};
+
+  auto node_type = common::AnfAlgo::GetCNodeName(cnode);
+  if (no_need_to_convert_nodes.find(node_type) != no_need_to_convert_nodes.end()) {
+    return nullptr;
+  }
+
+  std::vector<AnfNodePtr> new_inputs;
+  auto kernel_graph = func_graph->cast<std::shared_ptr<session::KernelGraph>>();
+  auto inputs = cnode->inputs();
+
+  // Keep the first input (typically the operator node) unchanged.
+  new_inputs.push_back(inputs[0]);
+
+  bool need_update = false;
+
+  // Check each input for Scalar or ValueTuple.
+  for (size_t i = 0; i < inputs.size() - 1; ++i) {
+    auto input_node = inputs[i + 1];
+
+    // Convert Scalar or ValueTuple inputs to Tensor inputs.
+    if (IsValueNode<Scalar>(input_node) || IsValueNode<ValueTuple>(input_node)) {
+      auto tensor_input = CreateTensorInput(kernel_graph, input_node);
+      if (tensor_input == nullptr) {
+        new_inputs.push_back(input_node);
+        continue;
+      }
+      new_inputs.push_back(tensor_input);
+      need_update = true;
+    } else {
+      new_inputs.push_back(input_node);
+    }
+  }
+
+  // If any inputs were converted, create a new CNode with updated inputs.
+  if (need_update) {
+    auto new_cnode = NewCNode(new_inputs, func_graph);
+    MS_EXCEPTION_IF_NULL(new_cnode);
+    if (common::AnfAlgo::CheckPrimitiveType(cnode, prim::kPrimDepend)) {
+      new_cnode->set_abstract(new_inputs[1]->abstract());
+    } else {
+      new_cnode->set_abstract(cnode->abstract());
+    }
+    new_cnode->set_scope(cnode->scope());
+    common::AnfAlgo::CopyNodeAttrs(cnode, new_cnode);
+    if (kernel_graph != nullptr) {
+      kernel_graph->FrontBackendlMapUpdate(cnode, new_cnode);
+    }
+    return new_cnode;
+  }
+  return nullptr;
+}
+
+/**
+ * @brief Entry function to initiate the conversion of const inputs to tensor inputs.
+ *
+ * @param func_graph The computational graph the node belongs to.
+ * @param node The node to be processed.
+ * @param equiv Placeholder for equivalence classes in pattern matching (unused in this function).
+ *
+ * @return Processed CNode or nullptr if no conversion occurred.
+ */
+const AnfNodePtr ConvertConstInputToTensorInput::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
+                                                         const EquivPtr &) const {
+  if (node == nullptr || func_graph == nullptr ||
+      common::AnfAlgo::CheckPrimitiveType(node, prim::kPrimTupleGetItem) ||
+      common::AnfAlgo::CheckPrimitiveType(node, prim::kPrimMakeTuple)) {
+    return nullptr;
+  }
+
+  if (!node->isa<CNode>()) {
+    return nullptr;
+  }
+
+  return ConstInputToTensorInput(func_graph, node->cast<CNodePtr>());
+}
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/convert_const_scalar_to_tensor.cc

@@ -25,64 +25,119 @@
 namespace mindspore {
 namespace opt {
 namespace {
+/**
+ * @brief Create a tensor input from a given scalar node.
+ *
+ * If the provided node is a ValueNode containing a Scalar, this function 
+ * converts that Scalar into a Tensor and wraps it in a new ValueNode.
+ *
+ * @param kernel_graph The computational graph the node belongs to.
+ * @param input_node The scalar node to be converted.
+ *
+ * @return A ValueNode containing the Tensor representation or nullptr for non-scalars.
+ */
 AnfNodePtr CreateTensorInput(const KernelGraphPtr &kernel_graph, const AnfNodePtr &input_node) {
+  // Ensure the provided node is not null.
   MS_EXCEPTION_IF_NULL(input_node);
+
+  // If the input_node is not a ValueNode, return nullptr.
   if (!input_node->isa<ValueNode>()) {
     return nullptr;
   }
+
   auto value_node = input_node->cast<ValueNodePtr>();
   MS_EXCEPTION_IF_NULL(value_node);
+
+  // Extract the value held by the ValueNode.
   auto value = value_node->value();
   MS_EXCEPTION_IF_NULL(value);
+
+  // If the value is not a Scalar, return nullptr.
   if (!value->isa<Scalar>()) {
     return nullptr;
   }
+
+  // Convert the scalar value to a tensor.
   tensor::TensorPtr tensor_ptr = ScalarToTensor(value->cast<ScalarPtr>());
   if (tensor_ptr == nullptr) {
     MS_LOG(WARNING) << "Create tensor of" << input_node->DebugString() << "failed";
     return nullptr;
   }
+
+  // Wrap the tensor in a new ValueNode.
   auto tensor_input = std::make_shared<ValueNode>(tensor_ptr);
   MS_EXCEPTION_IF_NULL(tensor_input);
   tensor_input->set_abstract(tensor_ptr->ToAbstract());
+
+  // If a KernelGraph is provided, add the new ValueNode to the graph.
   if (kernel_graph != nullptr) {
     tensor_input = kernel_graph->NewValueNode(tensor_input);
     kernel_graph->AddValueNodeToGraph(tensor_input);
   } else {
     tensor_input = MakeValueNode(tensor_input);
   }
+
+  // Assign the original node's scope to the new ValueNode.
   tensor_input->set_scope(input_node->scope());
+
   return tensor_input;
 }
+
 }  // namespace
 
+/**
+ * @brief Convert nodes with scalar values to tensor nodes.
+ *
+ * For a given computational node, if the node has scalar values as inputs,
+ * this function will convert these scalars into tensors.
+ *
+ * @param func_graph The computational graph the node belongs to.
+ * @param node The node whose scalar inputs need to be converted.
+ * @param equiv Placeholder for equivalence classes in pattern matching (unused in this function).
+ *
+ * @return A new CNode with converted tensor inputs or nullptr if no conversion occurred.
+ */
 const AnfNodePtr ConvertConstScalarToTensor::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
                                                      const EquivPtr &) const {
+  // Check if the node is null, the function graph is null or if the node is of type TupleGetItem.
   if (node == nullptr || func_graph == nullptr || common::AnfAlgo::CheckPrimitiveType(node, prim::kPrimTupleGetItem)) {
     return nullptr;
   }
-  // input is scalar, and link to graph return
+
+  // If the node is the graph's output and is a scalar, convert it to a tensor.
   if (node->isa<ValueNode>() && node == func_graph->output()) {
     return CreateTensorInput(func_graph->cast<KernelGraphPtr>(), node);
   }
+
+  // If the node is not a CNode, return nullptr.
   if (!node->isa<CNode>()) {
     return nullptr;
   }
+
   auto cnode = node->cast<CNodePtr>();
   MS_EXCEPTION_IF_NULL(cnode);
+
   bool input_changed = false;
+
+  // Iterate over all inputs of the CNode.
   for (size_t i = 0; i < cnode->inputs().size(); ++i) {
+    // If the input is a scalar, convert it to a tensor.
     auto new_input = CreateTensorInput(func_graph->cast<KernelGraphPtr>(), cnode->inputs()[i]);
     if (new_input != nullptr) {
       cnode->set_input(i, new_input);
       input_changed = true;
     }
   }
+
+  // If no inputs were converted or the graph isn't of type KernelGraph, return nullptr.
   auto kernel_graph = func_graph->cast<KernelGraphPtr>();
   if (kernel_graph == nullptr || !input_changed) {
     return nullptr;
   }
+
+  // Return a new CNode with converted tensor inputs.
   return NewCNode(cnode, kernel_graph);
 }
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/convert_const_scalar_to_tensor_old.cc

@@ -0,0 +1,143 @@
+/**
+ * Copyright 2020-2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "backend/common/pass/convert_const_scalar_to_tensor.h"
+#include <memory>
+#include <utility>
+#include "include/common/utils/convert_utils.h"
+#include "backend/common/optimizer/helper.h"
+#include "backend/common/session/anf_runtime_algorithm.h"
+#include "include/common/utils/anfalgo.h"
+#include "backend/common/session/kernel_graph.h"
+
+namespace mindspore {
+namespace opt {
+namespace {
+/**
+ * @brief Create a tensor input from a given scalar node.
+ *
+ * If the provided node is a ValueNode containing a Scalar, this function 
+ * converts that Scalar into a Tensor and wraps it in a new ValueNode.
+ *
+ * @param kernel_graph The computational graph the node belongs to.
+ * @param input_node The scalar node to be converted.
+ *
+ * @return A ValueNode containing the Tensor representation or nullptr for non-scalars.
+ */
+AnfNodePtr CreateTensorInput(const KernelGraphPtr &kernel_graph, const AnfNodePtr &input_node) {
+  // Ensure the provided node is not null.
+  MS_EXCEPTION_IF_NULL(input_node);
+
+  // If the input_node is not a ValueNode, return nullptr.
+  if (!input_node->isa<ValueNode>()) {
+    return nullptr;
+  }
+
+  auto value_node = input_node->cast<ValueNodePtr>();
+  MS_EXCEPTION_IF_NULL(value_node);
+
+  // Extract the value held by the ValueNode.
+  auto value = value_node->value();
+  MS_EXCEPTION_IF_NULL(value);
+
+  // If the value is not a Scalar, return nullptr.
+  if (!value->isa<Scalar>()) {
+    return nullptr;
+  }
+
+  // Convert the scalar value to a tensor.
+  tensor::TensorPtr tensor_ptr = ScalarToTensor(value->cast<ScalarPtr>());
+  if (tensor_ptr == nullptr) {
+    MS_LOG(WARNING) << "Create tensor of" << input_node->DebugString() << "failed";
+    return nullptr;
+  }
+
+  // Wrap the tensor in a new ValueNode.
+  auto tensor_input = std::make_shared<ValueNode>(tensor_ptr);
+  MS_EXCEPTION_IF_NULL(tensor_input);
+  tensor_input->set_abstract(tensor_ptr->ToAbstract());
+
+  // If a KernelGraph is provided, add the new ValueNode to the graph.
+  if (kernel_graph != nullptr) {
+    tensor_input = kernel_graph->NewValueNode(tensor_input);
+    kernel_graph->AddValueNodeToGraph(tensor_input);
+  } else {
+    tensor_input = MakeValueNode(tensor_input);
+  }
+
+  // Assign the original node's scope to the new ValueNode.
+  tensor_input->set_scope(input_node->scope());
+
+  return tensor_input;
+}
+
+}  // namespace
+
+/**
+ * @brief Convert nodes with scalar values to tensor nodes.
+ *
+ * For a given computational node, if the node has scalar values as inputs,
+ * this function will convert these scalars into tensors.
+ *
+ * @param func_graph The computational graph the node belongs to.
+ * @param node The node whose scalar inputs need to be converted.
+ * @param equiv Placeholder for equivalence classes in pattern matching (unused in this function).
+ *
+ * @return A new CNode with converted tensor inputs or nullptr if no conversion occurred.
+ */
+const AnfNodePtr ConvertConstScalarToTensor::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
+                                                     const EquivPtr &) const {
+  // Check if the node is null, the function graph is null or if the node is of type TupleGetItem.
+  if (node == nullptr || func_graph == nullptr || common::AnfAlgo::CheckPrimitiveType(node, prim::kPrimTupleGetItem)) {
+    return nullptr;
+  }
+
+  // If the node is the graph's output and is a scalar, convert it to a tensor.
+  if (node->isa<ValueNode>() && node == func_graph->output()) {
+    return CreateTensorInput(func_graph->cast<KernelGraphPtr>(), node);
+  }
+
+  // If the node is not a CNode, return nullptr.
+  if (!node->isa<CNode>()) {
+    return nullptr;
+  }
+
+  auto cnode = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(cnode);
+
+  bool input_changed = false;
+
+  // Iterate over all inputs of the CNode.
+  for (size_t i = 0; i < cnode->inputs().size(); ++i) {
+    // If the input is a scalar, convert it to a tensor.
+    auto new_input = CreateTensorInput(func_graph->cast<KernelGraphPtr>(), cnode->inputs()[i]);
+    if (new_input != nullptr) {
+      cnode->set_input(i, new_input);
+      input_changed = true;
+    }
+  }
+
+  // If no inputs were converted or the graph isn't of type KernelGraph, return nullptr.
+  auto kernel_graph = func_graph->cast<KernelGraphPtr>();
+  if (kernel_graph == nullptr || !input_changed) {
+    return nullptr;
+  }
+
+  // Return a new CNode with converted tensor inputs.
+  return NewCNode(cnode, kernel_graph);
+}
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/convert_tuple_input_to_dynamic_input.cc

@@ -24,16 +24,32 @@
 namespace mindspore {
 namespace opt {
 namespace {
+/**
+ * @brief Split the inputs of a tuple into separate inputs.
+ *
+ * If the provided node outputs a tuple, this function extracts each element of the tuple 
+ * and appends them to the provided plant_inputs vector.
+ *
+ * @param graph The computational graph the node belongs to.
+ * @param tuple_input The tuple node whose elements need to be extracted.
+ * @param plant_inputs A vector where the extracted tuple elements will be appended.
+ *
+ * @return Number of extracted tuple elements or -1 if not a tuple output.
+ */
 int64_t SplitTupleInputs(const FuncGraphPtr &graph, const AnfNodePtr &tuple_input,
                          std::vector<AnfNodePtr> *plant_inputs) {
+  // Ensure provided node is a tuple output.
   if (!common::AnfAlgo::IsTupleOutput(tuple_input)) {
     auto abs = tuple_input->abstract();
     MS_EXCEPTION_IF_NULL(abs);
     MS_LOG(WARNING) << "The Function only split the output type is tuple type but got" << abs->ToString();
     return -1;
   }
+
   MS_EXCEPTION_IF_NULL(plant_inputs);
   auto input_size = common::AnfAlgo::GetOutputTensorNum(tuple_input);
+
+  // If tuple_input is a MakeTuple node, extract its inputs directly.
   if (tuple_input->isa<CNode>() && common::AnfAlgo::CheckPrimitiveType(tuple_input, prim::kPrimMakeTuple)) {
     auto make_tuple = tuple_input->cast<CNodePtr>();
     MS_EXCEPTION_IF_NULL(make_tuple);
@@ -46,27 +62,47 @@ int64_t SplitTupleInputs(const FuncGraphPtr &graph, const AnfNodePtr &tuple_inpu
     }
     return input_size;
   }
+
+  // For general tuple inputs, create TupleGetItem nodes for each element.
   for (size_t index = 0; index < input_size; ++index) {
     auto dynamic_input_node = CreatTupleGetItemNode(graph, tuple_input, index);
     (void)plant_inputs->emplace_back(dynamic_input_node);
   }
+
   return input_size;
 }
 
+/**
+ * @brief Convert MakeTuple inputs of a CNode to individual inputs.
+ *
+ * For a given CNode, if any of its inputs are of type MakeTuple, this function will 
+ * split the tuple and replace the MakeTuple input with its individual elements.
+ *
+ * @param graph The computational graph the CNode belongs to.
+ * @param cnode_ptr The CNode whose inputs need processing.
+ */
 void ConvertMakeTupleInputToPlantInputs(const FuncGraphPtr &graph, const CNodePtr &cnode_ptr) {
   MS_EXCEPTION_IF_NULL(cnode_ptr);
   MS_EXCEPTION_IF_NULL(graph);
+
+  // If the CNode is either a Call or Partial node, no processing is required.
   if (common::AnfAlgo::CheckPrimitiveType(cnode_ptr, prim::kPrimCall) ||
       common::AnfAlgo::CheckPrimitiveType(cnode_ptr, prim::kPrimPartial)) {
     return;
   }
+
   std::vector<AnfNodePtr> plant_inputs;
   std::vector<int64_t> dyn_input_sizes;
+
+  // Always add the CNode's primitive to the new input list.
   plant_inputs.push_back(common::AnfAlgo::GetCNodePrimitiveNode(cnode_ptr));
+
   size_t input_num = cnode_ptr->inputs().size() - 1;
   for (size_t i = 0; i < input_num; ++i) {
     auto input_node = common::AnfAlgo::GetInputNode(cnode_ptr, i);
     MS_EXCEPTION_IF_NULL(input_node);
+
+    // If the input is a tuple output, split and append its elements to plant_inputs.
     if (common::AnfAlgo::IsTupleOutput(input_node)) {
       (void)dyn_input_sizes.emplace_back(SplitTupleInputs(graph, input_node, &plant_inputs));
     } else {
@@ -74,27 +110,54 @@ void ConvertMakeTupleInputToPlantInputs(const FuncGraphPtr &graph, const CNodePt
       plant_inputs.push_back(input_node);
     }
   }
-  // If there is dynamic input, set the dyn_input_sizes as an attribute and update the inputs.
+
+  // If there are dynamic inputs, update the CNode's attributes and inputs.
   if (std::any_of(dyn_input_sizes.begin(), dyn_input_sizes.end(), [](int64_t s) { return s >= 0; })) {
     common::AnfAlgo::SetNodeAttr(kAttrDynInputSizes, MakeValue(dyn_input_sizes), cnode_ptr);
     cnode_ptr->set_inputs(plant_inputs);
   }
 }
+
 }  // namespace
 
+/**
+ * @brief Define the pattern to be matched.
+ *
+ * This function returns a vector pattern where the first item is a variable 
+ * and the subsequent items are a sequence of variables.
+ *
+ * @return A BaseRef representing the defined pattern.
+ */
 const BaseRef ConvertTupleInputToDynamicInput::DefinePattern() const {
   VarPtr V = std::make_shared<Var>();
   VarPtr Xs = std::make_shared<SeqVar>();
   return VectorRef({V, Xs});
 }
 
+/**
+ * @brief Converts MakeTuple inputs of a CNode to individual dynamic inputs.
+ *
+ * If a node in the computational graph has MakeTuple type inputs, this function ensures 
+ * that those inputs are converted to individual dynamic inputs.
+ *
+ * @param func_graph The computational graph the node belongs to.
+ * @param node The node to be processed.
+ * @param equiv Placeholder for equivalence classes in pattern matching (unused in this function).
+ *
+ * @return The original node after potential modification.
+ */
 const AnfNodePtr ConvertTupleInputToDynamicInput::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
                                                           const EquivPtr &) const {
+  // Validate if the given node is not null, is of type CNode, and is a real kernel.
   if (node == nullptr || !node->isa<CNode>() || !AnfUtils::IsRealKernel(node)) {
     return nullptr;
   }
+
+  // Convert MakeTuple inputs of the CNode to individual dynamic inputs.
   ConvertMakeTupleInputToPlantInputs(func_graph, node->cast<CNodePtr>());
+
   return node;
 }
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/convert_tuple_input_to_dynamic_input_old.cc

@@ -0,0 +1,163 @@
+/**
+ * Copyright 2020 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "backend/common/pass/convert_tuple_input_to_dynamic_input.h"
+
+#include <algorithm>
+#include <memory>
+
+#include "backend/common/optimizer/helper.h"
+#include "include/common/utils/anfalgo.h"
+
+namespace mindspore {
+namespace opt {
+namespace {
+/**
+ * @brief Split the inputs of a tuple into separate inputs.
+ *
+ * If the provided node outputs a tuple, this function extracts each element of the tuple 
+ * and appends them to the provided plant_inputs vector.
+ *
+ * @param graph The computational graph the node belongs to.
+ * @param tuple_input The tuple node whose elements need to be extracted.
+ * @param plant_inputs A vector where the extracted tuple elements will be appended.
+ *
+ * @return Number of extracted tuple elements or -1 if not a tuple output.
+ */
+int64_t SplitTupleInputs(const FuncGraphPtr &graph, const AnfNodePtr &tuple_input,
+                         std::vector<AnfNodePtr> *plant_inputs) {
+  // Ensure provided node is a tuple output.
+  if (!common::AnfAlgo::IsTupleOutput(tuple_input)) {
+    auto abs = tuple_input->abstract();
+    MS_EXCEPTION_IF_NULL(abs);
+    MS_LOG(WARNING) << "The Function only split the output type is tuple type but got" << abs->ToString();
+    return -1;
+  }
+
+  MS_EXCEPTION_IF_NULL(plant_inputs);
+  auto input_size = common::AnfAlgo::GetOutputTensorNum(tuple_input);
+
+  // If tuple_input is a MakeTuple node, extract its inputs directly.
+  if (tuple_input->isa<CNode>() && common::AnfAlgo::CheckPrimitiveType(tuple_input, prim::kPrimMakeTuple)) {
+    auto make_tuple = tuple_input->cast<CNodePtr>();
+    MS_EXCEPTION_IF_NULL(make_tuple);
+    size_t tuple_input_num = common::AnfAlgo::GetInputTensorNum(make_tuple);
+    for (size_t j = 0; j < tuple_input_num; ++j) {
+      // using for graph kernel
+      auto dyn_input_node = common::AnfAlgo::GetInputNode(make_tuple, j);
+      MS_EXCEPTION_IF_NULL(dyn_input_node);
+      (void)plant_inputs->emplace_back(dyn_input_node);
+    }
+    return input_size;
+  }
+
+  // For general tuple inputs, create TupleGetItem nodes for each element.
+  for (size_t index = 0; index < input_size; ++index) {
+    auto dynamic_input_node = CreatTupleGetItemNode(graph, tuple_input, index);
+    (void)plant_inputs->emplace_back(dynamic_input_node);
+  }
+
+  return input_size;
+}
+
+/**
+ * @brief Convert MakeTuple inputs of a CNode to individual inputs.
+ *
+ * For a given CNode, if any of its inputs are of type MakeTuple, this function will 
+ * split the tuple and replace the MakeTuple input with its individual elements.
+ *
+ * @param graph The computational graph the CNode belongs to.
+ * @param cnode_ptr The CNode whose inputs need processing.
+ */
+void ConvertMakeTupleInputToPlantInputs(const FuncGraphPtr &graph, const CNodePtr &cnode_ptr) {
+  MS_EXCEPTION_IF_NULL(cnode_ptr);
+  MS_EXCEPTION_IF_NULL(graph);
+
+  // If the CNode is either a Call or Partial node, no processing is required.
+  if (common::AnfAlgo::CheckPrimitiveType(cnode_ptr, prim::kPrimCall) ||
+      common::AnfAlgo::CheckPrimitiveType(cnode_ptr, prim::kPrimPartial)) {
+    return;
+  }
+
+  std::vector<AnfNodePtr> plant_inputs;
+  std::vector<int64_t> dyn_input_sizes;
+
+  // Always add the CNode's primitive to the new input list.
+  plant_inputs.push_back(common::AnfAlgo::GetCNodePrimitiveNode(cnode_ptr));
+
+  size_t input_num = cnode_ptr->inputs().size() - 1;
+  for (size_t i = 0; i < input_num; ++i) {
+    auto input_node = common::AnfAlgo::GetInputNode(cnode_ptr, i);
+    MS_EXCEPTION_IF_NULL(input_node);
+
+    // If the input is a tuple output, split and append its elements to plant_inputs.
+    if (common::AnfAlgo::IsTupleOutput(input_node)) {
+      (void)dyn_input_sizes.emplace_back(SplitTupleInputs(graph, input_node, &plant_inputs));
+    } else {
+      dyn_input_sizes.push_back(-1);
+      plant_inputs.push_back(input_node);
+    }
+  }
+
+  // If there are dynamic inputs, update the CNode's attributes and inputs.
+  if (std::any_of(dyn_input_sizes.begin(), dyn_input_sizes.end(), [](int64_t s) { return s >= 0; })) {
+    common::AnfAlgo::SetNodeAttr(kAttrDynInputSizes, MakeValue(dyn_input_sizes), cnode_ptr);
+    cnode_ptr->set_inputs(plant_inputs);
+  }
+}
+
+}  // namespace
+
+/**
+ * @brief Define the pattern to be matched.
+ *
+ * This function returns a vector pattern where the first item is a variable 
+ * and the subsequent items are a sequence of variables.
+ *
+ * @return A BaseRef representing the defined pattern.
+ */
+const BaseRef ConvertTupleInputToDynamicInput::DefinePattern() const {
+  VarPtr V = std::make_shared<Var>();
+  VarPtr Xs = std::make_shared<SeqVar>();
+  return VectorRef({V, Xs});
+}
+
+/**
+ * @brief Converts MakeTuple inputs of a CNode to individual dynamic inputs.
+ *
+ * If a node in the computational graph has MakeTuple type inputs, this function ensures 
+ * that those inputs are converted to individual dynamic inputs.
+ *
+ * @param func_graph The computational graph the node belongs to.
+ * @param node The node to be processed.
+ * @param equiv Placeholder for equivalence classes in pattern matching (unused in this function).
+ *
+ * @return The original node after potential modification.
+ */
+const AnfNodePtr ConvertTupleInputToDynamicInput::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
+                                                          const EquivPtr &) const {
+  // Validate if the given node is not null, is of type CNode, and is a real kernel.
+  if (node == nullptr || !node->isa<CNode>() || !AnfUtils::IsRealKernel(node)) {
+    return nullptr;
+  }
+
+  // Convert MakeTuple inputs of the CNode to individual dynamic inputs.
+  ConvertMakeTupleInputToPlantInputs(func_graph, node->cast<CNodePtr>());
+
+  return node;
+}
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/convert_tuple_output_to_maketuple.cc

@@ -25,69 +25,139 @@
 namespace mindspore {
 namespace opt {
 namespace {
+/**
+ * @brief Converts a tuple input to a MakeTuple node.
+ *
+ * If a given node has a tuple output, this function ensures its representation 
+ * is a MakeTuple node. This is especially useful for operations that expect 
+ * their inputs to be explicitly represented as MakeTuple nodes, instead of 
+ * implicit tuple structures.
+ *
+ * @param graph The computational graph the node belongs to.
+ * @param tuple_anf The node that potentially has tuple output.
+ *
+ * @return An AnfNode pointer which is either the original node or a MakeTuple representation.
+ */
 AnfNodePtr ConvertTupleInputToMakeTuple(const FuncGraphPtr &graph, const AnfNodePtr &tuple_anf) {
+  // Ensure provided nodes and graph are not null.
   MS_EXCEPTION_IF_NULL(tuple_anf);
   MS_EXCEPTION_IF_NULL(graph);
 
+  // If the node does not produce a tuple output, return it as-is.
   if (!common::AnfAlgo::IsTupleOutput(tuple_anf)) {
     return tuple_anf;
   }
+
+  // Initially consider the graph as a KernelGraph. 
   auto kernel_graph = graph->cast<KernelGraphPtr>();
   FuncGraphPtr anf_graph = tuple_anf->func_graph();
+
+  // If the tuple node is associated with a functional graph, 
+  // attempt to cast it to a KernelGraph.
   if (anf_graph != nullptr) {
     kernel_graph = anf_graph->cast<KernelGraphPtr>();
   }
+
   MS_EXCEPTION_IF_NULL(kernel_graph);
+
+  // Check if the tuple node is already mapped to a MakeTuple node.
   if (kernel_graph->FindTupleParameterToMakeTupleMap(tuple_anf)) {
     return kernel_graph->FindTupleParameterToMakeTupleMap(tuple_anf);
   }
+
+  // Convert the tuple node to a MakeTuple node.
   auto make_tuple = kernel_graph->TransTupleToMakeTuple(tuple_anf);
   MS_EXCEPTION_IF_NULL(make_tuple);
+
+  // Store the mapping from the original tuple node to the new MakeTuple node.
   kernel_graph->InsertTupleParameterToMakeTupleMap(tuple_anf, make_tuple);
-  // replace graph inputs if input is a parameter
+
+  // If the tuple node was an input to the graph, replace it with the MakeTuple node.
   kernel_graph->ReplaceGraphInput(tuple_anf, make_tuple);
+
   return make_tuple;
 }
+
 }  // namespace
 
+/**
+ * @brief Define the pattern to be matched.
+ *
+ * This function returns a vector pattern where the first item is a variable 
+ * and the subsequent items are a sequence of variables.
+ *
+ * @return A BaseRef representing the defined pattern.
+ */
 const BaseRef ConvertTupleOutputToMaketuple::DefinePattern() const {
   VarPtr V = std::make_shared<Var>();
   VarPtr Xs = std::make_shared<SeqVar>();
   return VectorRef({V, Xs});
 }
 
+/**
+ * @brief Converts tuple outputs of nodes to MakeTuple nodes.
+ *
+ * If a node in the computational graph outputs a tuple, this function ensures 
+ * its representation is a MakeTuple node. If the node's input is a tuple output, 
+ * this function replaces that input with a MakeTuple node.
+ *
+ * @param func_graph The computational graph the node belongs to.
+ * @param node The node to be processed.
+ * @param equiv Placeholder for equivalence classes in pattern matching (unused in this function).
+ *
+ * @return A new CNode with replaced inputs if changes were made; otherwise, nullptr.
+ */
 const AnfNodePtr ConvertTupleOutputToMaketuple::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
                                                         const EquivPtr &) const {
+  // Validate if the given node is not null and is of type CNode.
   if (node == nullptr || !node->isa<CNode>()) {
     return nullptr;
   }
+
   auto cnode = node->cast<CNodePtr>();
   MS_EXCEPTION_IF_NULL(cnode);
+
+  // If the node fetches an item from a tuple, check its real input.
   if (IsPrimitiveCNode(cnode, prim::kPrimTupleGetItem)) {
     auto real_input = common::AnfAlgo::GetTupleGetItemRealInput(cnode);
     MS_EXCEPTION_IF_NULL(real_input);
+
+    // If the real input isn't a parameter or a value node, return nullptr.
     if (!real_input->isa<Parameter>() && !real_input->isa<ValueNode>()) {
       return nullptr;
     }
   }
+
+  // If the node updates the state, no need to proceed further.
   if (IsPrimitiveCNode(cnode, prim::kPrimUpdateState)) {
     return nullptr;
   }
+
   bool cnode_input_changed = false;
+
+  // Iterate over all inputs of the node.
   for (size_t i = 0; i < cnode->inputs().size(); ++i) {
     const auto &input = cnode->inputs()[i];
+
+    // If the input produces a tuple output and is not of type kPrimCall,
+    // Convert that tuple input to a MakeTuple node.
     if (input->Type() != nullptr && AnfUtils::IsRealKernel(input) && common::AnfAlgo::IsTupleOutput(input) &&
         !common::AnfAlgo::CheckPrimitiveType(input, prim::kPrimCall)) {
       cnode->set_input(i, ConvertTupleInputToMakeTuple(func_graph, input));
       cnode_input_changed = true;
     }
   }
+
+  // If no inputs were changed or the graph isn't of type KernelGraph, return nullptr.
   FuncGraphPtr graph = node->func_graph();
   auto kernel_graph = graph->cast<KernelGraphPtr>();
   if (kernel_graph == nullptr || !cnode_input_changed) {
     return nullptr;
   }
+
+  // Return a new CNode with replaced inputs.
   return NewCNode(cnode, kernel_graph);
 }
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/convert_tuple_output_to_maketuple_old.cc

@@ -0,0 +1,163 @@
+/**
+ * Copyright 2020-2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "backend/common/pass/convert_tuple_output_to_maketuple.h"
+
+#include <algorithm>
+#include <memory>
+
+#include "backend/common/optimizer/helper.h"
+#include "backend/common/session/kernel_graph.h"
+#include "include/common/utils/anfalgo.h"
+
+namespace mindspore {
+namespace opt {
+namespace {
+/**
+ * @brief Converts a tuple input to a MakeTuple node.
+ *
+ * If a given node has a tuple output, this function ensures its representation 
+ * is a MakeTuple node. This is especially useful for operations that expect 
+ * their inputs to be explicitly represented as MakeTuple nodes, instead of 
+ * implicit tuple structures.
+ *
+ * @param graph The computational graph the node belongs to.
+ * @param tuple_anf The node that potentially has tuple output.
+ *
+ * @return An AnfNode pointer which is either the original node or a MakeTuple representation.
+ */
+AnfNodePtr ConvertTupleInputToMakeTuple(const FuncGraphPtr &graph, const AnfNodePtr &tuple_anf) {
+  // Ensure provided nodes and graph are not null.
+  MS_EXCEPTION_IF_NULL(tuple_anf);
+  MS_EXCEPTION_IF_NULL(graph);
+
+  // If the node does not produce a tuple output, return it as-is.
+  if (!common::AnfAlgo::IsTupleOutput(tuple_anf)) {
+    return tuple_anf;
+  }
+
+  // Initially consider the graph as a KernelGraph. 
+  auto kernel_graph = graph->cast<KernelGraphPtr>();
+  FuncGraphPtr anf_graph = tuple_anf->func_graph();
+
+  // If the tuple node is associated with a functional graph, 
+  // attempt to cast it to a KernelGraph.
+  if (anf_graph != nullptr) {
+    kernel_graph = anf_graph->cast<KernelGraphPtr>();
+  }
+
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+
+  // Check if the tuple node is already mapped to a MakeTuple node.
+  if (kernel_graph->FindTupleParameterToMakeTupleMap(tuple_anf)) {
+    return kernel_graph->FindTupleParameterToMakeTupleMap(tuple_anf);
+  }
+
+  // Convert the tuple node to a MakeTuple node.
+  auto make_tuple = kernel_graph->TransTupleToMakeTuple(tuple_anf);
+  MS_EXCEPTION_IF_NULL(make_tuple);
+
+  // Store the mapping from the original tuple node to the new MakeTuple node.
+  kernel_graph->InsertTupleParameterToMakeTupleMap(tuple_anf, make_tuple);
+
+  // If the tuple node was an input to the graph, replace it with the MakeTuple node.
+  kernel_graph->ReplaceGraphInput(tuple_anf, make_tuple);
+
+  return make_tuple;
+}
+
+}  // namespace
+
+/**
+ * @brief Define the pattern to be matched.
+ *
+ * This function returns a vector pattern where the first item is a variable 
+ * and the subsequent items are a sequence of variables.
+ *
+ * @return A BaseRef representing the defined pattern.
+ */
+const BaseRef ConvertTupleOutputToMaketuple::DefinePattern() const {
+  VarPtr V = std::make_shared<Var>();
+  VarPtr Xs = std::make_shared<SeqVar>();
+  return VectorRef({V, Xs});
+}
+
+/**
+ * @brief Converts tuple outputs of nodes to MakeTuple nodes.
+ *
+ * If a node in the computational graph outputs a tuple, this function ensures 
+ * its representation is a MakeTuple node. If the node's input is a tuple output, 
+ * this function replaces that input with a MakeTuple node.
+ *
+ * @param func_graph The computational graph the node belongs to.
+ * @param node The node to be processed.
+ * @param equiv Placeholder for equivalence classes in pattern matching (unused in this function).
+ *
+ * @return A new CNode with replaced inputs if changes were made; otherwise, nullptr.
+ */
+const AnfNodePtr ConvertTupleOutputToMaketuple::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
+                                                        const EquivPtr &) const {
+  // Validate if the given node is not null and is of type CNode.
+  if (node == nullptr || !node->isa<CNode>()) {
+    return nullptr;
+  }
+
+  auto cnode = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(cnode);
+
+  // If the node fetches an item from a tuple, check its real input.
+  if (IsPrimitiveCNode(cnode, prim::kPrimTupleGetItem)) {
+    auto real_input = common::AnfAlgo::GetTupleGetItemRealInput(cnode);
+    MS_EXCEPTION_IF_NULL(real_input);
+
+    // If the real input isn't a parameter or a value node, return nullptr.
+    if (!real_input->isa<Parameter>() && !real_input->isa<ValueNode>()) {
+      return nullptr;
+    }
+  }
+
+  // If the node updates the state, no need to proceed further.
+  if (IsPrimitiveCNode(cnode, prim::kPrimUpdateState)) {
+    return nullptr;
+  }
+
+  bool cnode_input_changed = false;
+
+  // Iterate over all inputs of the node.
+  for (size_t i = 0; i < cnode->inputs().size(); ++i) {
+    const auto &input = cnode->inputs()[i];
+
+    // If the input produces a tuple output and is not of type kPrimCall,
+    // Convert that tuple input to a MakeTuple node.
+    if (input->Type() != nullptr && AnfUtils::IsRealKernel(input) && common::AnfAlgo::IsTupleOutput(input) &&
+        !common::AnfAlgo::CheckPrimitiveType(input, prim::kPrimCall)) {
+      cnode->set_input(i, ConvertTupleInputToMakeTuple(func_graph, input));
+      cnode_input_changed = true;
+    }
+  }
+
+  // If no inputs were changed or the graph isn't of type KernelGraph, return nullptr.
+  FuncGraphPtr graph = node->func_graph();
+  auto kernel_graph = graph->cast<KernelGraphPtr>();
+  if (kernel_graph == nullptr || !cnode_input_changed) {
+    return nullptr;
+  }
+
+  // Return a new CNode with replaced inputs.
+  return NewCNode(cnode, kernel_graph);
+}
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/custom_op_const_input_to_attr.cc

@@ -23,26 +23,48 @@
 
 namespace mindspore {
 namespace opt {
-const AnfNodePtr CustomOpConstInputToAttr::Process(const FuncGraphPtr &, const AnfNodePtr &node,
-                                                   const EquivPtr &) const {
+/**
+ * @brief Convert constant inputs to attributes for a custom operation node.
+ *
+ * This function is responsible for identifying constant inputs in a custom operation node and 
+ * converting them into operation attributes. This is useful in optimizing the representation 
+ * of custom operations where some inputs could be better represented as attributes.
+ *
+ * @param func_graph The computational graph the node belongs to (unused in this snippet).
+ * @param node The node to be processed.
+ * @param equiv Placeholder for equivalence classes in pattern matching (unused in this snippet).
+ *
+ * @return Processed AnfNode pointer.
+ */
+const AnfNodePtr CustomOpConstInputToAttr::Process(const FuncGraphPtr &, const AnfNodePtr &node, const EquivPtr &) const {
+  // Validate if the given node is not null and if it's a real CNode kernel.
   if (node == nullptr || !AnfUtils::IsRealCNodeKernel(node)) {
     return nullptr;
   }
+
+  // Cast the node to a CNode pointer.
   auto cnode = node->cast<CNodePtr>();
   MS_EXCEPTION_IF_NULL(cnode);
+
+  // Ensure that the given CNode is a custom operation node.
   if (!IsPrimitiveCNode(cnode, prim::kPrimCustom)) {
     return nullptr;
   }
 
+  // Get the indices in the CNode which correspond to custom operation attributes.
   mindspore::HashSet<size_t> attr_indices;
   GetCustomOpAttrIndex(common::AnfAlgo::GetCNodePrimitive(cnode), &attr_indices);
+
+  // If no custom attribute indices are identified, return early.
   if (attr_indices.empty()) {
     return nullptr;
   }
 
+  // Convert constant inputs at the identified indices to attributes.
   ConstInputToAttr(cnode, attr_indices);
 
   return node;
 }
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/custom_op_const_input_to_attr_old.cc

@@ -0,0 +1,70 @@
+/**
+ * Copyright 2021-2022 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "backend/common/pass/custom_op_const_input_to_attr.h"
+
+#include <memory>
+
+#include "utils/hash_set.h"
+#include "backend/common/optimizer/const_input_to_attr.h"
+#include "include/common/utils/anfalgo.h"
+
+namespace mindspore {
+namespace opt {
+/**
+ * @brief Convert constant inputs to attributes for a custom operation node.
+ *
+ * This function is responsible for identifying constant inputs in a custom operation node and 
+ * converting them into operation attributes. This is useful in optimizing the representation 
+ * of custom operations where some inputs could be better represented as attributes.
+ *
+ * @param func_graph The computational graph the node belongs to (unused in this snippet).
+ * @param node The node to be processed.
+ * @param equiv Placeholder for equivalence classes in pattern matching (unused in this snippet).
+ *
+ * @return Processed AnfNode pointer.
+ */
+const AnfNodePtr CustomOpConstInputToAttr::Process(const FuncGraphPtr &, const AnfNodePtr &node, const EquivPtr &) const {
+  // Validate if the given node is not null and if it's a real CNode kernel.
+  if (node == nullptr || !AnfUtils::IsRealCNodeKernel(node)) {
+    return nullptr;
+  }
+
+  // Cast the node to a CNode pointer.
+  auto cnode = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(cnode);
+
+  // Ensure that the given CNode is a custom operation node.
+  if (!IsPrimitiveCNode(cnode, prim::kPrimCustom)) {
+    return nullptr;
+  }
+
+  // Get the indices in the CNode which correspond to custom operation attributes.
+  mindspore::HashSet<size_t> attr_indices;
+  GetCustomOpAttrIndex(common::AnfAlgo::GetCNodePrimitive(cnode), &attr_indices);
+
+  // If no custom attribute indices are identified, return early.
+  if (attr_indices.empty()) {
+    return nullptr;
+  }
+
+  // Convert constant inputs at the identified indices to attributes.
+  ConstInputToAttr(cnode, attr_indices);
+
+  return node;
+}
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/custom_op_reg_info_to_attr.cc

@@ -27,10 +27,28 @@
 namespace mindspore {
 namespace opt {
 namespace {
+/**
+ * @brief Parses and sets default attribute values for a given primitive based on attribute type.
+ *
+ * This function is responsible for parsing different types of attributes from string representation
+ * and then setting them to the provided primitive. Supported attribute types include:
+ * "int", "str", "bool", "float", "listInt", "listStr", "listBool", and "listFloat".
+ *
+ * @param op_name Name of the operation.
+ * @param attr_name Name of the attribute.
+ * @param attr_value String representation of the attribute value.
+ * @param attr_type Type of the attribute.
+ * @param prim Pointer to the primitive where the parsed attribute should be set.
+ * 
+ * @throws Exception when prim is null or when encountering unsupported attribute type.
+ */
 void ParseAttrDefaultValue(const std::string &op_name, const std::string &attr_name, const std::string &attr_value,
                            const std::string &attr_type, const PrimitivePtr &prim) {
+  // Ensure that the provided primitive is not null.
   MS_EXCEPTION_IF_NULL(prim);
+
   try {
+    // Parse and set the attribute based on its type.
     if (attr_type == "int") {
       prim->set_attr(attr_name, std::make_shared<Int64Imm>(std::stoi(attr_value)));
     } else if (attr_type == "str") {
@@ -45,103 +63,163 @@ void ParseAttrDefaultValue(const std::string &op_name, const std::string &attr_n
       std::stringstream ss(attr_value);
       std::string elem;
       std::vector<ValuePtr> value;
+
+      // Parse comma-separated integers and add to the list.
       while (std::getline(ss, elem, ',')) {
         value.push_back(std::make_shared<Int64Imm>(std::stoi(elem)));
       }
+
       prim->set_attr(attr_name, std::make_shared<ValueList>(value));
     } else if (attr_type == "listStr") {
       std::stringstream ss(attr_value);
       std::string elem;
       std::vector<ValuePtr> value;
+
+      // Parse comma-separated strings and add to the list.
       while (std::getline(ss, elem, ',')) {
         value.push_back(std::make_shared<StringImm>(elem));
       }
+
       prim->set_attr(attr_name, std::make_shared<ValueList>(value));
     } else if (attr_type == "listBool") {
       std::stringstream ss(attr_value);
       std::string elem;
       std::vector<ValuePtr> value;
+
+      // Parse comma-separated booleans and add to the list.
       while (std::getline(ss, elem, ',')) {
         bool cur_value = false;
         std::istringstream(elem) >> std::boolalpha >> cur_value;
         value.push_back(std::make_shared<BoolImm>(cur_value));
       }
+
       prim->set_attr(attr_name, std::make_shared<ValueList>(value));
     } else if (attr_type == "listFloat") {
       std::stringstream ss(attr_value);
       std::string elem;
       std::vector<ValuePtr> value;
+
+      // Parse comma-separated floats and add to the list.
       while (std::getline(ss, elem, ',')) {
         value.push_back(std::make_shared<FP32Imm>(std::stof(elem)));
       }
+
       prim->set_attr(attr_name, std::make_shared<ValueList>(value));
     } else {
+      // Unsupported attribute type.
       MS_LOG(EXCEPTION) << "Unsupported attr type: " << attr_type;
     }
   } catch (const std::exception &e) {
+    // Handle exceptions during parsing and attribute setting.
     MS_LOG(EXCEPTION) << "Parse attr [" << attr_name << "] of op [" << op_name << "] failed! attr type: " << attr_type
                       << ", default value: " << attr_value << ", error message: " << e.what();
   }
 }
 
+
+namespace {
+
+/**
+ * @brief Add missing attributes with their default values to a given CNode based on Op registration info.
+ *
+ * @param cnode Pointer to the CNode.
+ * @param imply_type Type of kernel operation.
+ * @param missing_attrs Set of missing attribute names.
+ *
+ * @throws Exception if cnode or its primitive is null.
+ */
 void AddMissingAttrs(const CNodePtr &cnode, kernel::OpImplyType imply_type,
                      const std::unordered_set<std::string> &missing_attrs) {
+  // Ensure that the CNode is not null.
   MS_EXCEPTION_IF_NULL(cnode);
+
+  // Retrieve the primitive associated with the CNode.
   auto primitive = common::AnfAlgo::GetCNodePrimitive(cnode);
   MS_EXCEPTION_IF_NULL(primitive);
+
+  // Clone the primitive for modifications.
   primitive = primitive->Clone();
+
+  // Retrieve the operation name from the CNode.
   auto op_name = common::AnfAlgo::GetCNodeName(cnode);
+
+  // Get the operation registration information based on the operation name and its type.
   auto op_info_ptr = mindspore::kernel::OpLib::FindOp(op_name, imply_type);
   MS_EXCEPTION_IF_NULL(op_info_ptr);
+
+  // Retrieve all the attributes associated with the operation.
   auto all_attrs = op_info_ptr->attrs_ptr();
+
   bool need_update = false;
   for (const auto &attr : all_attrs) {
     auto attr_name = attr->name();
+
+    // Skip attributes that are not missing.
     if (missing_attrs.find(attr_name) == missing_attrs.end()) {
       continue;
     }
-    // If attr's param_type is required, it should have default value.
-    // If attr have default value, we should parse it no matter whether its param_type is required or not.
+
+    // Retrieve the default value for the attribute.
     auto default_value = attr->default_value();
+
+    // If the attribute type isn't required and doesn't have a default value, continue to the next attribute.
     if (default_value.empty() && attr->param_type() != "required") {
       continue;
     }
+
+    // If the attribute doesn't have a default value, raise an exception.
     if (default_value.empty()) {
       MS_LOG(EXCEPTION) << "attr [" << attr_name << "] in the registration information of op [" << op_name
                         << "] does not have a value." << trace::DumpSourceLines(cnode);
     }
+
+    // Parse and set the default attribute value to the primitive.
     ParseAttrDefaultValue(op_name, attr_name, default_value, attr->type(), primitive);
+
+    // Indicate that an update to the CNode is needed.
     need_update = true;
   }
+
+  // If there were changes, update the primitive in the CNode.
   if (need_update) {
     cnode->set_input(kAnfPrimitiveIndex, NewValueNode(primitive));
   }
 }
-}  // namespace
 
 const AnfNodePtr CustomOpRegInfoToAttr::Process(const FuncGraphPtr &, const AnfNodePtr &node, const EquivPtr &) const {
+  // Check for null node or non-real CNode kernels.
   if (node == nullptr || !AnfUtils::IsRealCNodeKernel(node)) {
     return nullptr;
   }
+
+  // Convert the node to a CNode pointer.
   auto cnode = node->cast<CNodePtr>();
   MS_EXCEPTION_IF_NULL(cnode);
+
+  // Check if the CNode is of type kPrimCustom.
   if (!IsPrimitiveCNode(cnode, prim::kPrimCustom)) {
     return nullptr;
   }
 
+  // Retrieve the primitive associated with the CNode.
   auto primitive = common::AnfAlgo::GetCNodePrimitive(cnode);
   MS_EXCEPTION_IF_NULL(primitive);
+
+  // Determine the function type of the CNode.
   auto func_type = common::AnfAlgo::GetNodeAttr<std::string>(cnode, kAttrFuncType);
-  // AKG/AICPU need to process attr, TBE will process later in the json creating phase.
+
+  // If the node's function type is neither AKG nor AICPU, return early.
   if (kCustomTypeAkg.find(func_type) == kCustomTypeAkg.end() || func_type == kCustomTypeAICPU) {
     return nullptr;
   }
-  // Early return if current node does not have attr
+
+  // Check if the CNode has any attributes. If not, return early.
   auto attr_names = primitive->GetAttr(kAttrAttrNames);
   if (attr_names == nullptr) {
     return nullptr;
   }
-  // Early return if all attr in reg info exist in the node's attr
+
+  // Check if all attributes in the registration info are present in the CNode's attributes.
   std::unordered_set<std::string> missing_attrs;
   auto attr_names_vec = GetValue<std::vector<std::string>>(attr_names);
   for (const auto &name : attr_names_vec) {
@@ -149,14 +227,23 @@ const AnfNodePtr CustomOpRegInfoToAttr::Process(const FuncGraphPtr &, const AnfN
       (void)missing_attrs.insert(name);
     }
   }
+
+  // If no attributes are missing, return early.
   if (missing_attrs.empty()) {
     return nullptr;
   }
+
+  // Determine the implication type for the operation.
   kernel::OpImplyType imply_type =
     func_type == kCustomTypeAICPU ? kernel::OpImplyType::kAICPU : kernel::OpImplyType::kAKG;
+
+  // Add the missing attributes to the CNode.
   AddMissingAttrs(cnode, imply_type, missing_attrs);
 
   return node;
 }
+
+}  // namespace
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/custom_op_reg_info_to_attr_old.cc

@@ -0,0 +1,249 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "backend/common/pass/custom_op_reg_info_to_attr.h"
+
+#include <memory>
+#include <string>
+#include <vector>
+#include <unordered_set>
+
+#include "kernel/oplib/opinfo.h"
+#include "kernel/oplib/oplib.h"
+#include "include/common/utils/anfalgo.h"
+
+namespace mindspore {
+namespace opt {
+namespace {
+/**
+ * @brief Parses and sets default attribute values for a given primitive based on attribute type.
+ *
+ * This function is responsible for parsing different types of attributes from string representation
+ * and then setting them to the provided primitive. Supported attribute types include:
+ * "int", "str", "bool", "float", "listInt", "listStr", "listBool", and "listFloat".
+ *
+ * @param op_name Name of the operation.
+ * @param attr_name Name of the attribute.
+ * @param attr_value String representation of the attribute value.
+ * @param attr_type Type of the attribute.
+ * @param prim Pointer to the primitive where the parsed attribute should be set.
+ * 
+ * @throws Exception when prim is null or when encountering unsupported attribute type.
+ */
+void ParseAttrDefaultValue(const std::string &op_name, const std::string &attr_name, const std::string &attr_value,
+                           const std::string &attr_type, const PrimitivePtr &prim) {
+  // Ensure that the provided primitive is not null.
+  MS_EXCEPTION_IF_NULL(prim);
+
+  try {
+    // Parse and set the attribute based on its type.
+    if (attr_type == "int") {
+      prim->set_attr(attr_name, std::make_shared<Int64Imm>(std::stoi(attr_value)));
+    } else if (attr_type == "str") {
+      prim->set_attr(attr_name, std::make_shared<StringImm>(attr_value));
+    } else if (attr_type == "bool") {
+      bool value = false;
+      std::istringstream(attr_value) >> std::boolalpha >> value;
+      prim->set_attr(attr_name, std::make_shared<BoolImm>(value));
+    } else if (attr_type == "float") {
+      prim->set_attr(attr_name, std::make_shared<FP32Imm>(std::stof(attr_value)));
+    } else if (attr_type == "listInt") {
+      std::stringstream ss(attr_value);
+      std::string elem;
+      std::vector<ValuePtr> value;
+
+      // Parse comma-separated integers and add to the list.
+      while (std::getline(ss, elem, ',')) {
+        value.push_back(std::make_shared<Int64Imm>(std::stoi(elem)));
+      }
+
+      prim->set_attr(attr_name, std::make_shared<ValueList>(value));
+    } else if (attr_type == "listStr") {
+      std::stringstream ss(attr_value);
+      std::string elem;
+      std::vector<ValuePtr> value;
+
+      // Parse comma-separated strings and add to the list.
+      while (std::getline(ss, elem, ',')) {
+        value.push_back(std::make_shared<StringImm>(elem));
+      }
+
+      prim->set_attr(attr_name, std::make_shared<ValueList>(value));
+    } else if (attr_type == "listBool") {
+      std::stringstream ss(attr_value);
+      std::string elem;
+      std::vector<ValuePtr> value;
+
+      // Parse comma-separated booleans and add to the list.
+      while (std::getline(ss, elem, ',')) {
+        bool cur_value = false;
+        std::istringstream(elem) >> std::boolalpha >> cur_value;
+        value.push_back(std::make_shared<BoolImm>(cur_value));
+      }
+
+      prim->set_attr(attr_name, std::make_shared<ValueList>(value));
+    } else if (attr_type == "listFloat") {
+      std::stringstream ss(attr_value);
+      std::string elem;
+      std::vector<ValuePtr> value;
+
+      // Parse comma-separated floats and add to the list.
+      while (std::getline(ss, elem, ',')) {
+        value.push_back(std::make_shared<FP32Imm>(std::stof(elem)));
+      }
+
+      prim->set_attr(attr_name, std::make_shared<ValueList>(value));
+    } else {
+      // Unsupported attribute type.
+      MS_LOG(EXCEPTION) << "Unsupported attr type: " << attr_type;
+    }
+  } catch (const std::exception &e) {
+    // Handle exceptions during parsing and attribute setting.
+    MS_LOG(EXCEPTION) << "Parse attr [" << attr_name << "] of op [" << op_name << "] failed! attr type: " << attr_type
+                      << ", default value: " << attr_value << ", error message: " << e.what();
+  }
+}
+
+
+namespace {
+
+/**
+ * @brief Add missing attributes with their default values to a given CNode based on Op registration info.
+ *
+ * @param cnode Pointer to the CNode.
+ * @param imply_type Type of kernel operation.
+ * @param missing_attrs Set of missing attribute names.
+ *
+ * @throws Exception if cnode or its primitive is null.
+ */
+void AddMissingAttrs(const CNodePtr &cnode, kernel::OpImplyType imply_type,
+                     const std::unordered_set<std::string> &missing_attrs) {
+  // Ensure that the CNode is not null.
+  MS_EXCEPTION_IF_NULL(cnode);
+
+  // Retrieve the primitive associated with the CNode.
+  auto primitive = common::AnfAlgo::GetCNodePrimitive(cnode);
+  MS_EXCEPTION_IF_NULL(primitive);
+
+  // Clone the primitive for modifications.
+  primitive = primitive->Clone();
+
+  // Retrieve the operation name from the CNode.
+  auto op_name = common::AnfAlgo::GetCNodeName(cnode);
+
+  // Get the operation registration information based on the operation name and its type.
+  auto op_info_ptr = mindspore::kernel::OpLib::FindOp(op_name, imply_type);
+  MS_EXCEPTION_IF_NULL(op_info_ptr);
+
+  // Retrieve all the attributes associated with the operation.
+  auto all_attrs = op_info_ptr->attrs_ptr();
+
+  bool need_update = false;
+  for (const auto &attr : all_attrs) {
+    auto attr_name = attr->name();
+
+    // Skip attributes that are not missing.
+    if (missing_attrs.find(attr_name) == missing_attrs.end()) {
+      continue;
+    }
+
+    // Retrieve the default value for the attribute.
+    auto default_value = attr->default_value();
+
+    // If the attribute type isn't required and doesn't have a default value, continue to the next attribute.
+    if (default_value.empty() && attr->param_type() != "required") {
+      continue;
+    }
+
+    // If the attribute doesn't have a default value, raise an exception.
+    if (default_value.empty()) {
+      MS_LOG(EXCEPTION) << "attr [" << attr_name << "] in the registration information of op [" << op_name
+                        << "] does not have a value." << trace::DumpSourceLines(cnode);
+    }
+
+    // Parse and set the default attribute value to the primitive.
+    ParseAttrDefaultValue(op_name, attr_name, default_value, attr->type(), primitive);
+
+    // Indicate that an update to the CNode is needed.
+    need_update = true;
+  }
+
+  // If there were changes, update the primitive in the CNode.
+  if (need_update) {
+    cnode->set_input(kAnfPrimitiveIndex, NewValueNode(primitive));
+  }
+}
+
+const AnfNodePtr CustomOpRegInfoToAttr::Process(const FuncGraphPtr &, const AnfNodePtr &node, const EquivPtr &) const {
+  // Check for null node or non-real CNode kernels.
+  if (node == nullptr || !AnfUtils::IsRealCNodeKernel(node)) {
+    return nullptr;
+  }
+
+  // Convert the node to a CNode pointer.
+  auto cnode = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(cnode);
+
+  // Check if the CNode is of type kPrimCustom.
+  if (!IsPrimitiveCNode(cnode, prim::kPrimCustom)) {
+    return nullptr;
+  }
+
+  // Retrieve the primitive associated with the CNode.
+  auto primitive = common::AnfAlgo::GetCNodePrimitive(cnode);
+  MS_EXCEPTION_IF_NULL(primitive);
+
+  // Determine the function type of the CNode.
+  auto func_type = common::AnfAlgo::GetNodeAttr<std::string>(cnode, kAttrFuncType);
+
+  // If the node's function type is neither AKG nor AICPU, return early.
+  if (kCustomTypeAkg.find(func_type) == kCustomTypeAkg.end() || func_type == kCustomTypeAICPU) {
+    return nullptr;
+  }
+
+  // Check if the CNode has any attributes. If not, return early.
+  auto attr_names = primitive->GetAttr(kAttrAttrNames);
+  if (attr_names == nullptr) {
+    return nullptr;
+  }
+
+  // Check if all attributes in the registration info are present in the CNode's attributes.
+  std::unordered_set<std::string> missing_attrs;
+  auto attr_names_vec = GetValue<std::vector<std::string>>(attr_names);
+  for (const auto &name : attr_names_vec) {
+    if (!primitive->HasAttr(name)) {
+      (void)missing_attrs.insert(name);
+    }
+  }
+
+  // If no attributes are missing, return early.
+  if (missing_attrs.empty()) {
+    return nullptr;
+  }
+
+  // Determine the implication type for the operation.
+  kernel::OpImplyType imply_type =
+    func_type == kCustomTypeAICPU ? kernel::OpImplyType::kAICPU : kernel::OpImplyType::kAKG;
+
+  // Add the missing attributes to the CNode.
+  AddMissingAttrs(cnode, imply_type, missing_attrs);
+
+  return node;
+}
+
+}  // namespace
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/eliminate_func_data_type.cc

@@ -24,17 +24,27 @@
 
 namespace mindspore::opt {
 namespace {
+/**
+ * @brief Checks if a function data type exists within an abstract tuple.
+ * 
+ * Recursively searches within nested abstract tuples to determine if any element has a type of function.
+ * 
+ * @param node_abs The abstract base pointer pointing to the abstract tuple.
+ * @return bool Returns true if a function data type is found, false otherwise.
+ */
 bool FuncDataTypeExistsInAbstractTuple(const AbstractBasePtr &node_abs) {
   MS_EXCEPTION_IF_NULL(node_abs);
   auto abs_tuple = dyn_cast<abstract::AbstractTuple>(node_abs);
   MS_EXCEPTION_IF_NULL(abs_tuple);
   for (const auto &abs : abs_tuple->elements()) {
     MS_EXCEPTION_IF_NULL(abs);
+    // If the element is another abstract tuple, recursively check it.
     if (abs->isa<abstract::AbstractTuple>() && FuncDataTypeExistsInAbstractTuple(abs)) {
       return true;
     }
     auto type = abs->BuildType();
     MS_EXCEPTION_IF_NULL(type);
+    // If the type of the current element is a function, return true.
     if (type->type_id() == kObjectTypeFunction) {
       return true;
     }
@@ -42,12 +52,27 @@ bool FuncDataTypeExistsInAbstractTuple(const AbstractBasePtr &node_abs) {
   return false;
 }
 
+/**
+ * @brief Removes the input function node from the kernel graph's input list.
+ * 
+ * This function searches through the list of inputs in a kernel graph and removes the specified input function node.
+ * 
+ * @param kernel_graph The kernel graph where the input function node is to be removed.
+ * @param func_input_node The input function node to be removed.
+ */
 void RemoveInputFuncNodeForKernelGraph(const KernelGraphPtr &kernel_graph, const AnfNodePtr &func_input_node) {
   MS_EXCEPTION_IF_NULL(kernel_graph);
   MS_EXCEPTION_IF_NULL(func_input_node);
+  
+  // If the kernel graph has mutable inputs, process them.
   if (kernel_graph->MutableInputs() != nullptr) {
+    // Make a copy of the original inputs.
     std::vector<AnfNodePtr> original_inputs = *(kernel_graph->MutableInputs());
+    
+    // Clear the original list of inputs.
     kernel_graph->MutableInputs()->clear();
+
+    // Iterate over the original inputs and re-add them, excluding the specified function input node.
     std::for_each(original_inputs.begin(), original_inputs.end(),
                   [&kernel_graph, &func_input_node](const AnfNodePtr &node) {
                     MS_EXCEPTION_IF_NULL(node);
@@ -55,15 +80,28 @@ void RemoveInputFuncNodeForKernelGraph(const KernelGraphPtr &kernel_graph, const
                       kernel_graph->MutableInputs()->emplace_back(node);
                     }
                   });
+    
+    // Reset the input nodes of the kernel graph.
     kernel_graph->SetInputNodes();
   }
 }
 
+
+/**
+ * @brief Eliminates function data types within an abstract tuple.
+ * 
+ * Replaces function data type with a constant abstract scalar value (currently Int32Imm(1)). For nested abstract
+ * tuples, the function works recursively.
+ * 
+ * @param abs_tuple The abstract tuple pointer whose elements need to be checked.
+ * @return abstract::AbstractBasePtrList The modified list of abstract base pointers after the elimination.
+ */
 abstract::AbstractBasePtrList EliminateFuncDataTypeForAbstractTuple(const abstract::AbstractTuplePtr &abs_tuple) {
   MS_EXCEPTION_IF_NULL(abs_tuple);
   AbstractBasePtrList new_abs;
   const auto &elements = abs_tuple->elements();
   for (const auto &abs : elements) {
+    // Handle nested abstract tuple by recursive call.
     if (abs->isa<abstract::AbstractTuple>()) {
       new_abs.emplace_back(std::make_shared<abstract::AbstractTuple>(
         EliminateFuncDataTypeForAbstractTuple(dyn_cast<abstract::AbstractTuple>(abs))));
@@ -71,6 +109,7 @@ abstract::AbstractBasePtrList EliminateFuncDataTypeForAbstractTuple(const abstra
     }
     auto type = abs->BuildType();
     MS_EXCEPTION_IF_NULL(type);
+    // Replace the function data type with a constant.
     if (type->type_id() == kObjectTypeFunction) {
       new_abs.emplace_back(std::make_shared<abstract::AbstractScalar>(std::make_shared<Int32Imm>(1)));
     } else {
@@ -79,21 +118,37 @@ abstract::AbstractBasePtrList EliminateFuncDataTypeForAbstractTuple(const abstra
   }
   return new_abs;
 }
-}  // namespace
 
+/**
+ * @brief Initializes the EliminateFuncDataType class.
+ * 
+ * Creates and initializes a constant value and its associated abstract representation.
+ */
 void EliminateFuncDataType::Init() {
   constant_ = NewValueNode(MakeValue<int32_t>(1));
   constant_abs_ = std::make_shared<abstract::AbstractScalar>(std::make_shared<Int32Imm>(1));
   constant_->set_abstract(constant_abs_);
 }
 
+/**
+ * @brief Processes a node to eliminate function data types.
+ * 
+ * Replaces nodes with function data type or abstract tuples containing function data type. For parameters, they are 
+ * replaced with a constant. For non-parameter nodes, their abstract is replaced with a constant.
+ * 
+ * @param func_graph The functional graph that the node belongs to.
+ * @param node The AnfNode to be processed.
+ * @param equiv Not used currently but retained for interface consistency.
+ * @return AnfNodePtr Returns nullptr after processing.
+ */
 const AnfNodePtr EliminateFuncDataType::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
                                                 const EquivPtr &) const {
   static uint32_t parameter_already_processed_graph_id = UINT32_MAX;
-  // Case 1: for parameter node which has func data type, replace it with constant.
   MS_EXCEPTION_IF_NULL(func_graph);
   auto kernel_graph = func_graph->cast<KernelGraphPtr>();
   MS_EXCEPTION_IF_NULL(kernel_graph);
+
+  // Case 1: Replace parameters with function data type or abstract tuples containing function data type with a constant.
   if (kernel_graph->graph_id() != parameter_already_processed_graph_id) {
     auto manage = kernel_graph->manager();
     MS_EXCEPTION_IF_NULL(manage);
@@ -104,10 +159,9 @@ const AnfNodePtr EliminateFuncDataType::Process(const FuncGraphPtr &func_graph,
       MS_EXCEPTION_IF_NULL(param);
       auto abs = param->abstract();
       MS_EXCEPTION_IF_NULL(abs);
-      if (abs->isa<abstract::AbstractTuple>() && FuncDataTypeExistsInAbstractTuple(abs)) {
-        RemoveInputFuncNodeForKernelGraph(kernel_graph, param);
-        (void)tr.Replace(param, constant_);
-      } else if (common::AnfAlgo::GetOutputInferDataType(param, 0) == kObjectTypeFunction) {
+      // If the parameter has an abstract tuple with a function data type or is of function data type, replace it.
+      if ((abs->isa<abstract::AbstractTuple>() && FuncDataTypeExistsInAbstractTuple(abs)) ||
+          (common::AnfAlgo::GetOutputInferDataType(param, 0) == kObjectTypeFunction)) {
         RemoveInputFuncNodeForKernelGraph(kernel_graph, param);
         (void)tr.Replace(param, constant_);
       } else {
@@ -118,7 +172,8 @@ const AnfNodePtr EliminateFuncDataType::Process(const FuncGraphPtr &func_graph,
     kernel_graph->set_parameters(std::move(new_params));
     parameter_already_processed_graph_id = kernel_graph->graph_id();
   }
-  // Case 2: for non-parameter node which has func data type, replace its abstract with constant.
+
+  // Case 2: For non-parameter nodes, replace their abstract with a constant if they have function data type.
   MS_EXCEPTION_IF_NULL(node);
   const auto &abs = node->abstract();
   if (abs != nullptr) {
@@ -131,4 +186,5 @@ const AnfNodePtr EliminateFuncDataType::Process(const FuncGraphPtr &func_graph,
   }
   return nullptr;
 }
+
 }  // namespace mindspore::opt

mindspore/ccsrc/backend/common/pass/eliminate_func_data_type_old.cc

@@ -0,0 +1,190 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "backend/common/pass/eliminate_func_data_type.h"
+#include <memory>
+#include <utility>
+#include "ir/anf.h"
+#include "backend/common/session/kernel_graph.h"
+#include "backend/common/optimizer/helper.h"
+#include "include/common/utils/anfalgo.h"
+
+namespace mindspore::opt {
+namespace {
+/**
+ * @brief Checks if a function data type exists within an abstract tuple.
+ * 
+ * Recursively searches within nested abstract tuples to determine if any element has a type of function.
+ * 
+ * @param node_abs The abstract base pointer pointing to the abstract tuple.
+ * @return bool Returns true if a function data type is found, false otherwise.
+ */
+bool FuncDataTypeExistsInAbstractTuple(const AbstractBasePtr &node_abs) {
+  MS_EXCEPTION_IF_NULL(node_abs);
+  auto abs_tuple = dyn_cast<abstract::AbstractTuple>(node_abs);
+  MS_EXCEPTION_IF_NULL(abs_tuple);
+  for (const auto &abs : abs_tuple->elements()) {
+    MS_EXCEPTION_IF_NULL(abs);
+    // If the element is another abstract tuple, recursively check it.
+    if (abs->isa<abstract::AbstractTuple>() && FuncDataTypeExistsInAbstractTuple(abs)) {
+      return true;
+    }
+    auto type = abs->BuildType();
+    MS_EXCEPTION_IF_NULL(type);
+    // If the type of the current element is a function, return true.
+    if (type->type_id() == kObjectTypeFunction) {
+      return true;
+    }
+  }
+  return false;
+}
+
+/**
+ * @brief Removes the input function node from the kernel graph's input list.
+ * 
+ * This function searches through the list of inputs in a kernel graph and removes the specified input function node.
+ * 
+ * @param kernel_graph The kernel graph where the input function node is to be removed.
+ * @param func_input_node The input function node to be removed.
+ */
+void RemoveInputFuncNodeForKernelGraph(const KernelGraphPtr &kernel_graph, const AnfNodePtr &func_input_node) {
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  MS_EXCEPTION_IF_NULL(func_input_node);
+  
+  // If the kernel graph has mutable inputs, process them.
+  if (kernel_graph->MutableInputs() != nullptr) {
+    // Make a copy of the original inputs.
+    std::vector<AnfNodePtr> original_inputs = *(kernel_graph->MutableInputs());
+    
+    // Clear the original list of inputs.
+    kernel_graph->MutableInputs()->clear();
+
+    // Iterate over the original inputs and re-add them, excluding the specified function input node.
+    std::for_each(original_inputs.begin(), original_inputs.end(),
+                  [&kernel_graph, &func_input_node](const AnfNodePtr &node) {
+                    MS_EXCEPTION_IF_NULL(node);
+                    if (node != func_input_node) {
+                      kernel_graph->MutableInputs()->emplace_back(node);
+                    }
+                  });
+    
+    // Reset the input nodes of the kernel graph.
+    kernel_graph->SetInputNodes();
+  }
+}
+
+
+/**
+ * @brief Eliminates function data types within an abstract tuple.
+ * 
+ * Replaces function data type with a constant abstract scalar value (currently Int32Imm(1)). For nested abstract
+ * tuples, the function works recursively.
+ * 
+ * @param abs_tuple The abstract tuple pointer whose elements need to be checked.
+ * @return abstract::AbstractBasePtrList The modified list of abstract base pointers after the elimination.
+ */
+abstract::AbstractBasePtrList EliminateFuncDataTypeForAbstractTuple(const abstract::AbstractTuplePtr &abs_tuple) {
+  MS_EXCEPTION_IF_NULL(abs_tuple);
+  AbstractBasePtrList new_abs;
+  const auto &elements = abs_tuple->elements();
+  for (const auto &abs : elements) {
+    // Handle nested abstract tuple by recursive call.
+    if (abs->isa<abstract::AbstractTuple>()) {
+      new_abs.emplace_back(std::make_shared<abstract::AbstractTuple>(
+        EliminateFuncDataTypeForAbstractTuple(dyn_cast<abstract::AbstractTuple>(abs))));
+      continue;
+    }
+    auto type = abs->BuildType();
+    MS_EXCEPTION_IF_NULL(type);
+    // Replace the function data type with a constant.
+    if (type->type_id() == kObjectTypeFunction) {
+      new_abs.emplace_back(std::make_shared<abstract::AbstractScalar>(std::make_shared<Int32Imm>(1)));
+    } else {
+      new_abs.emplace_back(abs);
+    }
+  }
+  return new_abs;
+}
+
+/**
+ * @brief Initializes the EliminateFuncDataType class.
+ * 
+ * Creates and initializes a constant value and its associated abstract representation.
+ */
+void EliminateFuncDataType::Init() {
+  constant_ = NewValueNode(MakeValue<int32_t>(1));
+  constant_abs_ = std::make_shared<abstract::AbstractScalar>(std::make_shared<Int32Imm>(1));
+  constant_->set_abstract(constant_abs_);
+}
+
+/**
+ * @brief Processes a node to eliminate function data types.
+ * 
+ * Replaces nodes with function data type or abstract tuples containing function data type. For parameters, they are 
+ * replaced with a constant. For non-parameter nodes, their abstract is replaced with a constant.
+ * 
+ * @param func_graph The functional graph that the node belongs to.
+ * @param node The AnfNode to be processed.
+ * @param equiv Not used currently but retained for interface consistency.
+ * @return AnfNodePtr Returns nullptr after processing.
+ */
+const AnfNodePtr EliminateFuncDataType::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
+                                                const EquivPtr &) const {
+  static uint32_t parameter_already_processed_graph_id = UINT32_MAX;
+  MS_EXCEPTION_IF_NULL(func_graph);
+  auto kernel_graph = func_graph->cast<KernelGraphPtr>();
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+
+  // Case 1: Replace parameters with function data type or abstract tuples containing function data type with a constant.
+  if (kernel_graph->graph_id() != parameter_already_processed_graph_id) {
+    auto manage = kernel_graph->manager();
+    MS_EXCEPTION_IF_NULL(manage);
+    auto tr = manage->Transact();
+    std::vector<AnfNodePtr> new_params;
+    const auto &original_params = kernel_graph->parameters();
+    for (const auto &param : original_params) {
+      MS_EXCEPTION_IF_NULL(param);
+      auto abs = param->abstract();
+      MS_EXCEPTION_IF_NULL(abs);
+      // If the parameter has an abstract tuple with a function data type or is of function data type, replace it.
+      if ((abs->isa<abstract::AbstractTuple>() && FuncDataTypeExistsInAbstractTuple(abs)) ||
+          (common::AnfAlgo::GetOutputInferDataType(param, 0) == kObjectTypeFunction)) {
+        RemoveInputFuncNodeForKernelGraph(kernel_graph, param);
+        (void)tr.Replace(param, constant_);
+      } else {
+        new_params.emplace_back(param);
+      }
+    }
+    tr.Commit();
+    kernel_graph->set_parameters(std::move(new_params));
+    parameter_already_processed_graph_id = kernel_graph->graph_id();
+  }
+
+  // Case 2: For non-parameter nodes, replace their abstract with a constant if they have function data type.
+  MS_EXCEPTION_IF_NULL(node);
+  const auto &abs = node->abstract();
+  if (abs != nullptr) {
+    if (abs->isa<abstract::AbstractTuple>()) {
+      auto abs_tuple = dyn_cast<abstract::AbstractTuple>(abs);
+      node->set_abstract(std::make_shared<abstract::AbstractTuple>(EliminateFuncDataTypeForAbstractTuple(abs_tuple)));
+    } else if (common::AnfAlgo::GetOutputInferDataType(node, 0) == kObjectTypeFunction) {
+      node->set_abstract(constant_abs_);
+    }
+  }
+  return nullptr;
+}
+
+}  // namespace mindspore::opt

mindspore/ccsrc/backend/common/pass/eliminate_redundant_op.cc

@@ -27,20 +27,40 @@
 namespace mindspore {
 namespace opt {
 namespace {
+/**
+ * @brief Traverse the computation graph to find the real preceding CNode.
+ * 
+ * This function traces back through specific operations like MakeTuple, TupleGetItem, Depend, and UpdateState 
+ * to find the real preceding CNode in the graph.
+ * 
+ * @param node The node to start tracing from.
+ * @param index The index of the input to consider for the given node.
+ * @param pass_vector A vector to record nodes and indices that have been passed during tracing.
+ * @return CNodePtr The real preceding CNode, or nullptr if not found.
+ */
+
+// get the real previous cnode; skipping virtual nodes
 CNodePtr GetRealPrevCNode(const AnfNodePtr &node, size_t index, std::vector<KernelWithIndex> *pass_vector) {
+  // Ensure pass_vector is not null.
   MS_EXCEPTION_IF_NULL(pass_vector);
+
+  // Return nullptr if the node is null or not a CNode.
   if (node == nullptr || !node->isa<CNode>()) {
     return nullptr;
   }
   auto cnode = node->cast<CNodePtr>();
   MS_EXCEPTION_IF_NULL(cnode);
+
+  // If cnode is a real kernel, record it and return.
   if (AnfUtils::IsRealCNodeKernel(cnode)) {
     pass_vector->push_back(make_pair(cnode, IntToSize(1)));
     return cnode;
   }
 
+  // Handle different cases based on the primitive of the CNode.
   auto input0 = cnode->input(0);
   MS_EXCEPTION_IF_NULL(input0);
+
   constexpr size_t kInput2 = 2;
   if (IsPrimitive(input0, prim::kPrimMakeTuple)) {
     auto temp_node = cnode->input(index + IntToSize(1));
@@ -66,19 +86,59 @@ CNodePtr GetRealPrevCNode(const AnfNodePtr &node, size_t index, std::vector<Kern
   }
 }
 
+/**
+ * @brief A placeholder function that always returns true.
+ * 
+ * This function is typically used as a condition check where no actual check is required.
+ * 
+ * @param node1 The first CNode.
+ * @param node2 The second CNode.
+ * @return true Always true.
+ */
 bool TransOpEliminateCondition(const CNodePtr &, const CNodePtr &) { return true; }
 
+/**
+ * @brief Check if two CNodes have symmetrical kernel information.
+ * 
+ * This function determines if two CNodes can be considered equivalent based on their kernel information.
+ * 
+ * @param node1 The first CNode to compare.
+ * @param node2 The second CNode to compare.
+ * @return true If both nodes have symmetrical kernel information.
+ * @return false Otherwise.
+ */
 bool CastEliminateCondition(const CNodePtr &node1, const CNodePtr &node2) {
   return HasSymmetricalKernelInfo(node1, node2);
 }
 
+
+/**
+ * @brief Checks if two TransData operation nodes can be eliminated based on their input/output formats.
+ * 
+ * @param node1 The first CNode.
+ * @param node2 The second CNode.
+ * @return true If the nodes satisfy the conditions to be eliminated.
+ * @return false Otherwise.
+ */
 bool TransDataOpEliminateCondition(const CNodePtr &node1, const CNodePtr &node2) {
   return AnfAlgo::GetInputFormat(node1, 0) == AnfAlgo::GetOutputFormat(node2, 0) &&
          AnfAlgo::GetOutputFormat(node1, 0) == AnfAlgo::GetInputFormat(node2, 0) &&
          kernel::IsSameShape(AnfAlgo::GetInputDeviceShape(node2, 0), AnfAlgo::GetOutputDeviceShape(node1, 0));
 }
-}  // namespace
 
+namespace {
+
+/**
+ * @brief Process nodes that match a redundant pattern.
+ * 
+ * If two subsequent nodes are identified as redundant, the redundant node (or nodes) are eliminated or replaced.
+ * 
+ * @param func_graph The function graph containing the nodes.
+ * @param cnode The CNode being checked.
+ * @param prev_cnode The preceding CNode to the cnode being checked.
+ * @param pass_vector A vector of nodes that have been traversed in search of redundancy.
+ * @return AnfNodePtr The processed node.
+ */
 const AnfNodePtr EliminateRedundantOp::ProcessMatchedNodes(const FuncGraphPtr &func_graph, const CNodePtr &cnode,
                                                            const CNodePtr &prev_cnode,
                                                            std::vector<KernelWithIndex> *pass_vector) const {
@@ -92,6 +152,7 @@ const AnfNodePtr EliminateRedundantOp::ProcessMatchedNodes(const FuncGraphPtr &f
   auto &users = manager->node_users();
 
   auto pass_size = pass_vector->size();
+  // Check if any traversed node is of type Depend or has more than one user.
   for (size_t idx = 1; idx <= pass_size - 1; ++idx) {
     auto nd = (*pass_vector)[idx].first;
     if (common::AnfAlgo::CheckPrimitiveType(nd, prim::kPrimDepend)) {
@@ -102,14 +163,16 @@ const AnfNodePtr EliminateRedundantOp::ProcessMatchedNodes(const FuncGraphPtr &f
     }
   }
 
-  // when no depend node and no node used more than once, no need to rebuild the pass nodes
   constexpr size_t kOffset = 2;
   if (!has_depend_node) {
+    // No depend node found, so directly return the input of prev_cnode.
     return prev_cnode->input(1);
   } else if (!has_node_used_more_than_once) {
+    // If a node is only used once and it's a depend node, replace it with its input.
     (void)manager->Replace(prev_cnode, prev_cnode->input(1));
     return cnode->input(1);
-  } else {  // rebuild the pass nodes
+  } else {
+    // Rebuild the nodes that were passed over.
     if (pass_size < kOffset) {
       MS_LOG(ERROR) << "pass_size should >= 2";
     }
@@ -127,6 +190,9 @@ const AnfNodePtr EliminateRedundantOp::ProcessMatchedNodes(const FuncGraphPtr &f
   }
 }
 
+/**
+ * @brief Initializes the map of redundant operation patterns and their corresponding elimination conditions.
+ */
 void EliminateRedundantOp::Init() {
   (void)redundant_process_map_.emplace(std::pair<std::string, RedundantOpPair>(
     kFour2FiveOpName, std::pair<std::string, ConditionFunc>(kFive2FourOpName, TransOpEliminateCondition)));
@@ -140,44 +206,76 @@ void EliminateRedundantOp::Init() {
     kTransDataRNNOpName, std::pair<std::string, ConditionFunc>(kTransDataRNNOpName, TransDataOpEliminateCondition)));
 }
 
+
+/**
+ * @brief Attempts to eliminate a CNode if it matches certain redundancy criteria.
+ * 
+ * Checks the node's name against a known set of redundant operations. If there's a match, 
+ * it searches for the previous node to see if both nodes can be combined or eliminated.
+ * 
+ * @param func_graph The function graph containing the CNode.
+ * @param cnode The CNode being checked.
+ * @return AnfNodePtr The processed node or nullptr if elimination is not possible.
+ */
 const AnfNodePtr EliminateRedundantOp::DoEliminate(const FuncGraphPtr &func_graph, const CNodePtr &cnode) const {
-  // match the first name
+  // Check if the current node's name matches a known redundant operation.
   auto name1 = common::AnfAlgo::GetCNodeName(cnode);
   auto it = redundant_process_map_.find(name1);
   if (it == redundant_process_map_.end()) {
     return nullptr;
   }
+
+  // Collect nodes that are traversed while searching for redundancy.
   std::vector<KernelWithIndex> pass_vector;
   pass_vector.push_back(make_pair(cnode, 1));
+  
+  // Find the preceding CNode.
   auto prev_cnode = GetRealPrevCNode(cnode->input(1), 0, &pass_vector);
   if (prev_cnode == nullptr) {
     return nullptr;
   }
-  // match the second name
+  
+  // Check if the preceding CNode's name matches the expected redundant operation's name.
   auto name2 = common::AnfAlgo::GetCNodeName(prev_cnode);
   if (name2 != it->second.first) {
     return nullptr;
   }
-  // match condition
+
+  // Check if the two nodes meet the elimination condition.
   auto condition_func = it->second.second;
-  if (condition_func == nullptr) {
-    return nullptr;
-  }
-  if (!condition_func(cnode, prev_cnode)) {
+  if (condition_func == nullptr || !condition_func(cnode, prev_cnode)) {
     return nullptr;
   }
 
+  // Process the matched nodes and perform the elimination.
   return ProcessMatchedNodes(func_graph, cnode, prev_cnode, &pass_vector);
 }
 
+/**
+ * @brief The main process function that checks a node for possible elimination.
+ * 
+ * If the node is a CNode and is part of a function graph, it delegates to the DoEliminate method.
+ * 
+ * @param func_graph The function graph containing the node.
+ * @param node The AnfNode being checked.
+ * @param EquivPtr Unused parameter.
+ * @return AnfNodePtr The processed node or nullptr if elimination is not possible.
+ */
 const AnfNodePtr EliminateRedundantOp::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
                                                const EquivPtr &) const {
   MS_EXCEPTION_IF_NULL(node);
+  
+  // Convert the AnfNode to CNode for further checks.
   auto cnode = node->cast<CNodePtr>();
+  
+  // Return nullptr if the node is not a CNode or the function graph is not given.
   if (cnode == nullptr || func_graph == nullptr) {
     return nullptr;
   }
+
+  // Delegate to the DoEliminate method.
   return DoEliminate(func_graph, cnode);
 }
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/eliminate_redundant_op_old.cc

@@ -0,0 +1,281 @@
+/**
+ * Copyright 2020-2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "backend/common/pass/eliminate_redundant_op.h"
+#include <memory>
+#include <utility>
+#include "utils/hash_map.h"
+#include "backend/common/session/anf_runtime_algorithm.h"
+#include "include/common/utils/anfalgo.h"
+#include "include/common/utils/utils.h"
+#include "backend/common/optimizer/helper.h"
+#include "base/core_ops.h"
+#include "kernel/common_utils.h"
+
+namespace mindspore {
+namespace opt {
+namespace {
+/**
+ * @brief Traverse the computation graph to find the real preceding CNode.
+ * 
+ * This function traces back through specific operations like MakeTuple, TupleGetItem, Depend, and UpdateState 
+ * to find the real preceding CNode in the graph.
+ * 
+ * @param node The node to start tracing from.
+ * @param index The index of the input to consider for the given node.
+ * @param pass_vector A vector to record nodes and indices that have been passed during tracing.
+ * @return CNodePtr The real preceding CNode, or nullptr if not found.
+ */
+
+// get the real previous cnode; skipping virtual nodes
+CNodePtr GetRealPrevCNode(const AnfNodePtr &node, size_t index, std::vector<KernelWithIndex> *pass_vector) {
+  // Ensure pass_vector is not null.
+  MS_EXCEPTION_IF_NULL(pass_vector);
+
+  // Return nullptr if the node is null or not a CNode.
+  if (node == nullptr || !node->isa<CNode>()) {
+    return nullptr;
+  }
+  auto cnode = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(cnode);
+
+  // If cnode is a real kernel, record it and return.
+  if (AnfUtils::IsRealCNodeKernel(cnode)) {
+    pass_vector->push_back(make_pair(cnode, IntToSize(1)));
+    return cnode;
+  }
+
+  // Handle different cases based on the primitive of the CNode.
+  auto input0 = cnode->input(0);
+  MS_EXCEPTION_IF_NULL(input0);
+
+  constexpr size_t kInput2 = 2;
+  if (IsPrimitive(input0, prim::kPrimMakeTuple)) {
+    auto temp_node = cnode->input(index + IntToSize(1));
+    MS_EXCEPTION_IF_NULL(temp_node);
+    pass_vector->push_back(make_pair(cnode, index + IntToSize(1)));
+    return GetRealPrevCNode(temp_node, 0, pass_vector);
+  } else if (IsPrimitive(input0, prim::kPrimTupleGetItem)) {
+    auto input2 = cnode->input(kInput2);
+    MS_EXCEPTION_IF_NULL(input2);
+    auto value_node = input2->cast<ValueNodePtr>();
+    MS_EXCEPTION_IF_NULL(value_node);
+    auto item_idx = GetValue<int64_t>(value_node->value());
+    pass_vector->push_back(make_pair(cnode, IntToSize(1)));
+    return GetRealPrevCNode(cnode->input(1), LongToSize(item_idx), pass_vector);
+  } else if (IsPrimitive(input0, prim::kPrimDepend)) {
+    pass_vector->push_back(make_pair(cnode, IntToSize(1)));
+    return GetRealPrevCNode(cnode->input(1), 0, pass_vector);
+  } else if (IsPrimitive(input0, prim::kPrimUpdateState)) {
+    pass_vector->push_back(make_pair(cnode, IntToSize(kUpdateStateRealInput)));
+    return GetRealPrevCNode(cnode->input(kUpdateStateRealInput), 0, pass_vector);
+  } else {
+    return nullptr;
+  }
+}
+
+/**
+ * @brief A placeholder function that always returns true.
+ * 
+ * This function is typically used as a condition check where no actual check is required.
+ * 
+ * @param node1 The first CNode.
+ * @param node2 The second CNode.
+ * @return true Always true.
+ */
+bool TransOpEliminateCondition(const CNodePtr &, const CNodePtr &) { return true; }
+
+/**
+ * @brief Check if two CNodes have symmetrical kernel information.
+ * 
+ * This function determines if two CNodes can be considered equivalent based on their kernel information.
+ * 
+ * @param node1 The first CNode to compare.
+ * @param node2 The second CNode to compare.
+ * @return true If both nodes have symmetrical kernel information.
+ * @return false Otherwise.
+ */
+bool CastEliminateCondition(const CNodePtr &node1, const CNodePtr &node2) {
+  return HasSymmetricalKernelInfo(node1, node2);
+}
+
+
+/**
+ * @brief Checks if two TransData operation nodes can be eliminated based on their input/output formats.
+ * 
+ * @param node1 The first CNode.
+ * @param node2 The second CNode.
+ * @return true If the nodes satisfy the conditions to be eliminated.
+ * @return false Otherwise.
+ */
+bool TransDataOpEliminateCondition(const CNodePtr &node1, const CNodePtr &node2) {
+  return AnfAlgo::GetInputFormat(node1, 0) == AnfAlgo::GetOutputFormat(node2, 0) &&
+         AnfAlgo::GetOutputFormat(node1, 0) == AnfAlgo::GetInputFormat(node2, 0) &&
+         kernel::IsSameShape(AnfAlgo::GetInputDeviceShape(node2, 0), AnfAlgo::GetOutputDeviceShape(node1, 0));
+}
+
+namespace {
+
+/**
+ * @brief Process nodes that match a redundant pattern.
+ * 
+ * If two subsequent nodes are identified as redundant, the redundant node (or nodes) are eliminated or replaced.
+ * 
+ * @param func_graph The function graph containing the nodes.
+ * @param cnode The CNode being checked.
+ * @param prev_cnode The preceding CNode to the cnode being checked.
+ * @param pass_vector A vector of nodes that have been traversed in search of redundancy.
+ * @return AnfNodePtr The processed node.
+ */
+const AnfNodePtr EliminateRedundantOp::ProcessMatchedNodes(const FuncGraphPtr &func_graph, const CNodePtr &cnode,
+                                                           const CNodePtr &prev_cnode,
+                                                           std::vector<KernelWithIndex> *pass_vector) const {
+  MS_EXCEPTION_IF_NULL(func_graph);
+  MS_EXCEPTION_IF_NULL(pass_vector);
+  FuncGraphManagerPtr manager = func_graph->manager();
+  MS_EXCEPTION_IF_NULL(manager);
+
+  bool has_depend_node = false;
+  bool has_node_used_more_than_once = false;
+  auto &users = manager->node_users();
+
+  auto pass_size = pass_vector->size();
+  // Check if any traversed node is of type Depend or has more than one user.
+  for (size_t idx = 1; idx <= pass_size - 1; ++idx) {
+    auto nd = (*pass_vector)[idx].first;
+    if (common::AnfAlgo::CheckPrimitiveType(nd, prim::kPrimDepend)) {
+      has_depend_node = true;
+    }
+    if (users[nd].size() > 1) {
+      has_node_used_more_than_once = true;
+    }
+  }
+
+  constexpr size_t kOffset = 2;
+  if (!has_depend_node) {
+    // No depend node found, so directly return the input of prev_cnode.
+    return prev_cnode->input(1);
+  } else if (!has_node_used_more_than_once) {
+    // If a node is only used once and it's a depend node, replace it with its input.
+    (void)manager->Replace(prev_cnode, prev_cnode->input(1));
+    return cnode->input(1);
+  } else {
+    // Rebuild the nodes that were passed over.
+    if (pass_size < kOffset) {
+      MS_LOG(ERROR) << "pass_size should >= 2";
+    }
+    for (size_t idx = pass_size - kOffset; idx > 0; --idx) {
+      auto new_node = NewCNode((*pass_vector)[idx].first->inputs(), func_graph);
+      if (idx == pass_size - kOffset) {
+        new_node->set_input((*pass_vector)[idx].second,
+                            (*pass_vector)[idx + 1].first->input((*pass_vector)[idx + 1].second));
+      } else {
+        new_node->set_input((*pass_vector)[idx].second, (*pass_vector)[idx + 1].first);
+      }
+      (*pass_vector)[idx].first = new_node;
+    }
+    return (*pass_vector)[1].first;
+  }
+}
+
+/**
+ * @brief Initializes the map of redundant operation patterns and their corresponding elimination conditions.
+ */
+void EliminateRedundantOp::Init() {
+  (void)redundant_process_map_.emplace(std::pair<std::string, RedundantOpPair>(
+    kFour2FiveOpName, std::pair<std::string, ConditionFunc>(kFive2FourOpName, TransOpEliminateCondition)));
+  (void)redundant_process_map_.emplace(std::pair<std::string, RedundantOpPair>(
+    kFive2FourOpName, std::pair<std::string, ConditionFunc>(kFour2FiveOpName, TransOpEliminateCondition)));
+  (void)redundant_process_map_.emplace(std::pair<std::string, RedundantOpPair>(
+    prim::kPrimCast->name(), std::pair<std::string, ConditionFunc>(prim::kPrimCast->name(), CastEliminateCondition)));
+  (void)redundant_process_map_.emplace(std::pair<std::string, RedundantOpPair>(
+    kTransDataOpName, std::pair<std::string, ConditionFunc>(kTransDataOpName, TransDataOpEliminateCondition)));
+  (void)redundant_process_map_.emplace(std::pair<std::string, RedundantOpPair>(
+    kTransDataRNNOpName, std::pair<std::string, ConditionFunc>(kTransDataRNNOpName, TransDataOpEliminateCondition)));
+}
+
+
+/**
+ * @brief Attempts to eliminate a CNode if it matches certain redundancy criteria.
+ * 
+ * Checks the node's name against a known set of redundant operations. If there's a match, 
+ * it searches for the previous node to see if both nodes can be combined or eliminated.
+ * 
+ * @param func_graph The function graph containing the CNode.
+ * @param cnode The CNode being checked.
+ * @return AnfNodePtr The processed node or nullptr if elimination is not possible.
+ */
+const AnfNodePtr EliminateRedundantOp::DoEliminate(const FuncGraphPtr &func_graph, const CNodePtr &cnode) const {
+  // Check if the current node's name matches a known redundant operation.
+  auto name1 = common::AnfAlgo::GetCNodeName(cnode);
+  auto it = redundant_process_map_.find(name1);
+  if (it == redundant_process_map_.end()) {
+    return nullptr;
+  }
+
+  // Collect nodes that are traversed while searching for redundancy.
+  std::vector<KernelWithIndex> pass_vector;
+  pass_vector.push_back(make_pair(cnode, 1));
+  
+  // Find the preceding CNode.
+  auto prev_cnode = GetRealPrevCNode(cnode->input(1), 0, &pass_vector);
+  if (prev_cnode == nullptr) {
+    return nullptr;
+  }
+  
+  // Check if the preceding CNode's name matches the expected redundant operation's name.
+  auto name2 = common::AnfAlgo::GetCNodeName(prev_cnode);
+  if (name2 != it->second.first) {
+    return nullptr;
+  }
+
+  // Check if the two nodes meet the elimination condition.
+  auto condition_func = it->second.second;
+  if (condition_func == nullptr || !condition_func(cnode, prev_cnode)) {
+    return nullptr;
+  }
+
+  // Process the matched nodes and perform the elimination.
+  return ProcessMatchedNodes(func_graph, cnode, prev_cnode, &pass_vector);
+}
+
+/**
+ * @brief The main process function that checks a node for possible elimination.
+ * 
+ * If the node is a CNode and is part of a function graph, it delegates to the DoEliminate method.
+ * 
+ * @param func_graph The function graph containing the node.
+ * @param node The AnfNode being checked.
+ * @param EquivPtr Unused parameter.
+ * @return AnfNodePtr The processed node or nullptr if elimination is not possible.
+ */
+const AnfNodePtr EliminateRedundantOp::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
+                                               const EquivPtr &) const {
+  MS_EXCEPTION_IF_NULL(node);
+  
+  // Convert the AnfNode to CNode for further checks.
+  auto cnode = node->cast<CNodePtr>();
+  
+  // Return nullptr if the node is not a CNode or the function graph is not given.
+  if (cnode == nullptr || func_graph == nullptr) {
+    return nullptr;
+  }
+
+  // Delegate to the DoEliminate method.
+  return DoEliminate(func_graph, cnode);
+}
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/erase_visit_attr_old.cc

@@ -0,0 +1,36 @@
+/**
+ * Copyright 2019 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "backend/common/pass/erase_visit_attr.h"
+#include <memory>
+#include "kernel/common_utils.h"
+#include "backend/common/session/anf_runtime_algorithm.h"
+#include "include/common/utils/anfalgo.h"
+#include "backend/common/optimizer/helper.h"
+
+namespace mindspore {
+namespace opt {
+const BaseRef EraseVisitAttr::DefinePattern() const {
+  std::shared_ptr<Var> V = std::make_shared<CondVar>(Visited);
+  std::shared_ptr<Var> Xs = std::make_shared<SeqVar>();
+  return VectorRef({V, Xs});
+}
+
+const AnfNodePtr EraseVisitAttr::Process(const FuncGraphPtr &, const AnfNodePtr &node, const EquivPtr &) const {
+  common::AnfAlgo::EraseNodeAttr(kAttrVisited, node);
+  return nullptr;
+}
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/getitem_tuple.cc

@@ -22,45 +22,99 @@
 namespace mindspore {
 namespace opt {
 namespace {
+namespace {
+
+/**
+ * @brief Check if a given node is a ValueNode.
+ * 
+ * @param n The BaseRef node to check.
+ * @return true If n is an instance of AnfNode and is a ValueNode.
+ * @return false Otherwise.
+ */
 bool IsC(const BaseRef &n) {
+  // Ensure the node is not null.
   MS_EXCEPTION_IF_NULL(n);
+
+  // Check if the node is an instance of AnfNode.
   if (utils::isa<AnfNodePtr>(n)) {
     auto in = utils::cast<AnfNodePtr>(n);
     MS_EXCEPTION_IF_NULL(in);
+
+    // Check if the AnfNode is a ValueNode.
     return in->isa<ValueNode>();
   } else {
     return false;
   }
 }
-}  // namespace
 
+/**
+ * @brief Define a pattern to match a sequence of nodes in MindSpore.
+ * 
+ * The goal is to capture a pattern where a TupleGetItem operation 
+ * retrieves an item from a MakeTuple operation.
+ * 
+ * @return BaseRef The pattern defined as a sequence of operations.
+ */
 const BaseRef GetitemTuple::DefinePattern() const {
+  // Define a sequence variable to capture multiple nodes.
   VarPtr Xs = std::make_shared<SeqVar>();
+  
+  // Define a condition variable to check if a node meets the IsC condition.
   VarPtr C = std::make_shared<CondVar>(IsC);
+  
+  // Return the pattern sequence.
   return VectorRef({prim::kPrimTupleGetItem, VectorRef({prim::kPrimMakeTuple, Xs}), C});
 }
 
+/**
+ * @brief Process a matched pattern and retrieve a specific item from a tuple.
+ * 
+ * If the pattern matches, this method returns the corresponding item from the MakeTuple.
+ * Otherwise, it returns nullptr.
+ * 
+ * @param func_graph The function graph (not used in this function).
+ * @param node The AnfNodePtr node being processed.
+ * @param equiv The equivalence class (not used in this function).
+ * @return AnfNodePtr The matched item node or nullptr if not found.
+ */
 const AnfNodePtr GetitemTuple::Process(const FuncGraphPtr &, const AnfNodePtr &node, const EquivPtr &) const {
+  // Ensure the node is not null.
   MS_EXCEPTION_IF_NULL(node);
+
+  // Cast node to CNode.
   auto tuple_getitem = node->cast<CNodePtr>();
   MS_EXCEPTION_IF_NULL(tuple_getitem);
+
+  // Get the MakeTuple operation and index node from TupleGetItem.
   AnfNodePtr make_tuple_anf = tuple_getitem->input(kRealInputNodeIndexInTupleGetItem);
   MS_EXCEPTION_IF_NULL(make_tuple_anf);
   AnfNodePtr index_node = tuple_getitem->input(kInputNodeOutputIndexInTupleGetItem);
   MS_EXCEPTION_IF_NULL(index_node);
+
+  // Check if the index node is a 64-bit integer.
   if (IsValueNode<Int64Imm>(index_node)) {
     auto value_node = index_node->cast<ValueNodePtr>();
     MS_EXCEPTION_IF_NULL(value_node);
+    
+    // Retrieve the index value.
     auto index = GetValue<int64_t>(value_node->value());
+    
     auto make_tuple = make_tuple_anf->cast<CNodePtr>();
     MS_EXCEPTION_IF_NULL(make_tuple);
+    
+    // If the index is valid, retrieve and return the corresponding node from MakeTuple.
     if (make_tuple->inputs().size() > LongToSize(index + 1)) {
       auto ret = make_tuple->input(LongToSize(index + 1));
       MS_EXCEPTION_IF_NULL(ret);
       return ret;
     }
   }
+
+  // Return nullptr if no matching node is found.
   return nullptr;
 }
+
+}  // namespace
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/getitem_tuple_old.cc

@@ -0,0 +1,120 @@
+/**
+ * Copyright 2019 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "backend/common/pass/getitem_tuple.h"
+
+#include <memory>
+#include "base/core_ops.h"
+#include "include/common/utils/utils.h"
+
+namespace mindspore {
+namespace opt {
+namespace {
+namespace {
+
+/**
+ * @brief Check if a given node is a ValueNode.
+ * 
+ * @param n The BaseRef node to check.
+ * @return true If n is an instance of AnfNode and is a ValueNode.
+ * @return false Otherwise.
+ */
+bool IsC(const BaseRef &n) {
+  // Ensure the node is not null.
+  MS_EXCEPTION_IF_NULL(n);
+
+  // Check if the node is an instance of AnfNode.
+  if (utils::isa<AnfNodePtr>(n)) {
+    auto in = utils::cast<AnfNodePtr>(n);
+    MS_EXCEPTION_IF_NULL(in);
+
+    // Check if the AnfNode is a ValueNode.
+    return in->isa<ValueNode>();
+  } else {
+    return false;
+  }
+}
+
+/**
+ * @brief Define a pattern to match a sequence of nodes in MindSpore.
+ * 
+ * The goal is to capture a pattern where a TupleGetItem operation 
+ * retrieves an item from a MakeTuple operation.
+ * 
+ * @return BaseRef The pattern defined as a sequence of operations.
+ */
+const BaseRef GetitemTuple::DefinePattern() const {
+  // Define a sequence variable to capture multiple nodes.
+  VarPtr Xs = std::make_shared<SeqVar>();
+  
+  // Define a condition variable to check if a node meets the IsC condition.
+  VarPtr C = std::make_shared<CondVar>(IsC);
+  
+  // Return the pattern sequence.
+  return VectorRef({prim::kPrimTupleGetItem, VectorRef({prim::kPrimMakeTuple, Xs}), C});
+}
+
+/**
+ * @brief Process a matched pattern and retrieve a specific item from a tuple.
+ * 
+ * If the pattern matches, this method returns the corresponding item from the MakeTuple.
+ * Otherwise, it returns nullptr.
+ * 
+ * @param func_graph The function graph (not used in this function).
+ * @param node The AnfNodePtr node being processed.
+ * @param equiv The equivalence class (not used in this function).
+ * @return AnfNodePtr The matched item node or nullptr if not found.
+ */
+const AnfNodePtr GetitemTuple::Process(const FuncGraphPtr &, const AnfNodePtr &node, const EquivPtr &) const {
+  // Ensure the node is not null.
+  MS_EXCEPTION_IF_NULL(node);
+
+  // Cast node to CNode.
+  auto tuple_getitem = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(tuple_getitem);
+
+  // Get the MakeTuple operation and index node from TupleGetItem.
+  AnfNodePtr make_tuple_anf = tuple_getitem->input(kRealInputNodeIndexInTupleGetItem);
+  MS_EXCEPTION_IF_NULL(make_tuple_anf);
+  AnfNodePtr index_node = tuple_getitem->input(kInputNodeOutputIndexInTupleGetItem);
+  MS_EXCEPTION_IF_NULL(index_node);
+
+  // Check if the index node is a 64-bit integer.
+  if (IsValueNode<Int64Imm>(index_node)) {
+    auto value_node = index_node->cast<ValueNodePtr>();
+    MS_EXCEPTION_IF_NULL(value_node);
+    
+    // Retrieve the index value.
+    auto index = GetValue<int64_t>(value_node->value());
+    
+    auto make_tuple = make_tuple_anf->cast<CNodePtr>();
+    MS_EXCEPTION_IF_NULL(make_tuple);
+    
+    // If the index is valid, retrieve and return the corresponding node from MakeTuple.
+    if (make_tuple->inputs().size() > LongToSize(index + 1)) {
+      auto ret = make_tuple->input(LongToSize(index + 1));
+      MS_EXCEPTION_IF_NULL(ret);
+      return ret;
+    }
+  }
+
+  // Return nullptr if no matching node is found.
+  return nullptr;
+}
+
+}  // namespace
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/insert_assign_for_custom_op.cc

@@ -28,20 +28,35 @@ constexpr auto kCustomOutput = 0;
 constexpr auto kCustomInput = 1;
 constexpr auto kCustomAttrInplaceAssignOutput = "inplace_assign_output";
 
-// Used to find Custom op outputs' inplace assign index
+/**
+ * @brief Fetch inplace index for Custom op outputs.
+ * 
+ * This function retrieves the inplace index information for Custom ops with hybrid type.
+ * 
+ * @param cnode The node whose inplace index is to be fetched.
+ * @return A 2D vector containing pairs of inplace indices.
+ */
 std::vector<std::vector<int64_t>> GetHybridInplaceIndex(const CNodePtr &cnode) {
+  // Check if the node's function type is "hybrid".
   if (common::AnfAlgo::GetNodeAttr<std::string>(cnode, kAttrFuncType) != kCustomTypeHybrid) {
     return {};
   }
 
+  // Verify if the node has an inplace attribute.
   if (!common::AnfAlgo::HasNodeAttr(kCustomAttrInplaceAssignOutput, cnode)) {
     return {};
   }
+
+  // Extract the inplace index string.
   auto inplace_index_str = common::AnfAlgo::GetNodeAttr<std::string>(cnode, kCustomAttrInplaceAssignOutput);
   std::regex delimiters(" ");
+  
+  // Tokenize the inplace index string to individual indices.
   std::vector<std::string> index(
     std::sregex_token_iterator(inplace_index_str.begin(), inplace_index_str.end(), delimiters, -1),
     std::sregex_token_iterator());
+  
+  // Convert tokenized strings to integer pairs.
   std::vector<std::vector<int64_t>> inplace_index;
   std::vector<int64_t> tmp;
   for (size_t i = 0; i < index.size(); i++) {
@@ -54,22 +69,32 @@ std::vector<std::vector<int64_t>> GetHybridInplaceIndex(const CNodePtr &cnode) {
   return inplace_index;
 }
 
+/**
+ * @brief Insert an Assign operation into the graph.
+ * 
+ * This function adds Assign, UpdateState, and Load operations after a given node in the graph.
+ * 
+ * @param func_graph The function graph being processed.
+ * @param src The source node for the Assign operation.
+ * @param dst The destination node for the Assign operation.
+ * @return A CNodePtr pointing to the newly created Load node.
+ */
 CNodePtr InsertAssign(const FuncGraphPtr &func_graph, const AnfNodePtr &src, const CNodePtr &dst) {
-  // Insert UpdateState, Load and Assign, need mount a UMonad node.
+  // Create a UMonad node.
   auto u = NewValueNode(kUMonad);
   u->set_abstract(kUMonad->ToAbstract());
 
-  // Insert Assign
+  // Construct the Assign node with the given source and destination.
   AnfNodePtrList assign_inputs = {NewValueNode(prim::kPrimAssign), dst, src, u};
   auto assign_cnode = func_graph->NewCNode(assign_inputs);
   assign_cnode->set_abstract(dst->abstract());
 
-  // Insert UpdateState
+  // Insert an UpdateState node after the Assign operation.
   AnfNodePtrList update_state_inputs = {NewValueNode(prim::kPrimUpdateState), u, assign_cnode};
   auto update_state_cnode = func_graph->NewCNode(update_state_inputs);
   update_state_cnode->set_abstract(kUMonad->ToAbstract());
 
-  // Insert Load
+  // Insert a Load node after the UpdateState operation.
   AnfNodePtrList load_inputs = {NewValueNode(prim::kPrimLoad), dst, update_state_cnode};
   auto load_cnode = func_graph->NewCNode(load_inputs);
   load_cnode->set_abstract(dst->abstract());
@@ -77,30 +102,60 @@ CNodePtr InsertAssign(const FuncGraphPtr &func_graph, const AnfNodePtr &src, con
   return load_cnode;
 }
 
+/**
+ * @brief Insert Assign operation after a Custom node.
+ * 
+ * @param func_graph The function graph being processed.
+ * @param cnode The Custom node after which the Assign operation will be added.
+ * @return A CNodePtr pointing to the new node if insertion is successful, otherwise nullptr.
+ */
 CNodePtr InsertAssignAfterCustom(const FuncGraphPtr &func_graph, const CNodePtr &cnode) {
+  // Fetch inplace info for the given node.
   auto inplace_info = GetHybridInplaceIndex(cnode);
+  
+  // Currently, only support a single inplace pair.
   if (inplace_info.size() != 1) return nullptr;
+
   auto input_size = common::AnfAlgo::GetInputTensorNum(cnode);
   if (auto i = LongToSize(inplace_info[0][kCustomInput]); i < input_size) {
+    // Insert the Assign operation after the custom node.
     return InsertAssign(func_graph, cnode->input(i + 1), cnode);
   } else {
     return nullptr;
   }
 }
 
+/**
+ * @brief Insert Assign operation after a TupleGetItem node.
+ * 
+ * If the TupleGetItem's input is a Custom node with matching inplace info, this function
+ * will insert an Assign operation after the TupleGetItem.
+ * 
+ * @param func_graph The function graph being processed.
+ * @param cnode The TupleGetItem node after which the Assign operation will be added.
+ * @return A CNodePtr pointing to the new node if insertion is successful, otherwise nullptr.
+ */
 CNodePtr InsertAssignAfterTupleGetItem(const FuncGraphPtr &func_graph, const CNodePtr &cnode) {
+  // Fetch the actual input node to the TupleGetItem.
   auto input_node = cnode->input(kRealInputNodeIndexInTupleGetItem);
   MS_EXCEPTION_IF_NULL(input_node);
   auto real_input = dyn_cast<CNode>(input_node);
+  
   if (real_input == nullptr) {
     return nullptr;
   }
+  
+  // Get the tuple item index.
   auto value_ptr = GetValueNode(cnode->input(kInputNodeOutputIndexInTupleGetItem));
   MS_EXCEPTION_IF_NULL(value_ptr);
   auto gt_idx = GetValue<int64_t>(value_ptr);
+  
   if (IsPrimitiveCNode(real_input, prim::kPrimCustom)) {
+    // Fetch inplace info for the real input node.
     auto inplace_info = GetHybridInplaceIndex(real_input);
+    
     for (auto index : inplace_info) {
+      // If there's a matching inplace info, insert Assign after the TupleGetItem.
       if (index[kCustomOutput] == gt_idx && index[kCustomInput] >= 0) {
         auto custom_input_size = common::AnfAlgo::GetInputTensorNum(real_input);
         if (auto i = LongToSize(index[kCustomInput]); i < custom_input_size) {
@@ -112,15 +167,24 @@ CNodePtr InsertAssignAfterTupleGetItem(const FuncGraphPtr &func_graph, const CNo
   return nullptr;
 }
 
+/**
+ * @brief Process nodes to insert Assign operations for Custom ops.
+ * 
+ * This function checks for nodes of type Custom or TupleGetItem and attempts to insert 
+ * Assign operations after them, if needed.
+ * 
+ * @param func_graph The function graph being processed.
+ * @param node The current node being examined.
+ * @param Equiv equivalence class pointer, unused in this function.
+ * @return A new AnfNodePtr if the node was replaced, otherwise nullptr.
+ */
 const AnfNodePtr InsertAssignForCustomOp::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
                                                   const EquivPtr &) const {
-  MS_EXCEPTION_IF_NULL(func_graph);
-  MS_EXCEPTION_IF_NULL(node);
-  auto cnode = node->cast<CNodePtr>();
+  // Check if the given node is a Custom or TupleGetItem node and has not been visited.
   if (cnode == nullptr) {
     return nullptr;
   }
-
+  
   if (IsPrimitiveCNode(cnode, prim::kPrimCustom) && visited_.find(cnode) == visited_.end()) {
     visited_.insert(cnode);
     return InsertAssignAfterCustom(func_graph, cnode);

mindspore/ccsrc/backend/common/pass/insert_assign_for_custom_op_old.cc

@@ -0,0 +1,199 @@
+/**
+ * Copyright 2022 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "backend/common/pass/insert_assign_for_custom_op.h"
+
+#include <memory>
+#include <vector>
+#include <string>
+#include <regex>
+#include "backend/common/optimizer/helper.h"
+#include "include/common/utils/anfalgo.h"
+
+namespace mindspore {
+namespace opt {
+constexpr auto kCustomOutput = 0;
+constexpr auto kCustomInput = 1;
+constexpr auto kCustomAttrInplaceAssignOutput = "inplace_assign_output";
+
+/**
+ * @brief Fetch inplace index for Custom op outputs.
+ * 
+ * This function retrieves the inplace index information for Custom ops with hybrid type.
+ * 
+ * @param cnode The node whose inplace index is to be fetched.
+ * @return A 2D vector containing pairs of inplace indices.
+ */
+std::vector<std::vector<int64_t>> GetHybridInplaceIndex(const CNodePtr &cnode) {
+  // Check if the node's function type is "hybrid".
+  if (common::AnfAlgo::GetNodeAttr<std::string>(cnode, kAttrFuncType) != kCustomTypeHybrid) {
+    return {};
+  }
+
+  // Verify if the node has an inplace attribute.
+  if (!common::AnfAlgo::HasNodeAttr(kCustomAttrInplaceAssignOutput, cnode)) {
+    return {};
+  }
+
+  // Extract the inplace index string.
+  auto inplace_index_str = common::AnfAlgo::GetNodeAttr<std::string>(cnode, kCustomAttrInplaceAssignOutput);
+  std::regex delimiters(" ");
+  
+  // Tokenize the inplace index string to individual indices.
+  std::vector<std::string> index(
+    std::sregex_token_iterator(inplace_index_str.begin(), inplace_index_str.end(), delimiters, -1),
+    std::sregex_token_iterator());
+  
+  // Convert tokenized strings to integer pairs.
+  std::vector<std::vector<int64_t>> inplace_index;
+  std::vector<int64_t> tmp;
+  for (size_t i = 0; i < index.size(); i++) {
+    tmp.push_back(std::stol(index[i]));
+    if (i & 1) {
+      inplace_index.push_back(tmp);
+      tmp.clear();
+    }
+  }
+  return inplace_index;
+}
+
+/**
+ * @brief Insert an Assign operation into the graph.
+ * 
+ * This function adds Assign, UpdateState, and Load operations after a given node in the graph.
+ * 
+ * @param func_graph The function graph being processed.
+ * @param src The source node for the Assign operation.
+ * @param dst The destination node for the Assign operation.
+ * @return A CNodePtr pointing to the newly created Load node.
+ */
+CNodePtr InsertAssign(const FuncGraphPtr &func_graph, const AnfNodePtr &src, const CNodePtr &dst) {
+  // Create a UMonad node.
+  auto u = NewValueNode(kUMonad);
+  u->set_abstract(kUMonad->ToAbstract());
+
+  // Construct the Assign node with the given source and destination.
+  AnfNodePtrList assign_inputs = {NewValueNode(prim::kPrimAssign), dst, src, u};
+  auto assign_cnode = func_graph->NewCNode(assign_inputs);
+  assign_cnode->set_abstract(dst->abstract());
+
+  // Insert an UpdateState node after the Assign operation.
+  AnfNodePtrList update_state_inputs = {NewValueNode(prim::kPrimUpdateState), u, assign_cnode};
+  auto update_state_cnode = func_graph->NewCNode(update_state_inputs);
+  update_state_cnode->set_abstract(kUMonad->ToAbstract());
+
+  // Insert a Load node after the UpdateState operation.
+  AnfNodePtrList load_inputs = {NewValueNode(prim::kPrimLoad), dst, update_state_cnode};
+  auto load_cnode = func_graph->NewCNode(load_inputs);
+  load_cnode->set_abstract(dst->abstract());
+
+  return load_cnode;
+}
+
+/**
+ * @brief Insert Assign operation after a Custom node.
+ * 
+ * @param func_graph The function graph being processed.
+ * @param cnode The Custom node after which the Assign operation will be added.
+ * @return A CNodePtr pointing to the new node if insertion is successful, otherwise nullptr.
+ */
+CNodePtr InsertAssignAfterCustom(const FuncGraphPtr &func_graph, const CNodePtr &cnode) {
+  // Fetch inplace info for the given node.
+  auto inplace_info = GetHybridInplaceIndex(cnode);
+  
+  // Currently, only support a single inplace pair.
+  if (inplace_info.size() != 1) return nullptr;
+
+  auto input_size = common::AnfAlgo::GetInputTensorNum(cnode);
+  if (auto i = LongToSize(inplace_info[0][kCustomInput]); i < input_size) {
+    // Insert the Assign operation after the custom node.
+    return InsertAssign(func_graph, cnode->input(i + 1), cnode);
+  } else {
+    return nullptr;
+  }
+}
+
+/**
+ * @brief Insert Assign operation after a TupleGetItem node.
+ * 
+ * If the TupleGetItem's input is a Custom node with matching inplace info, this function
+ * will insert an Assign operation after the TupleGetItem.
+ * 
+ * @param func_graph The function graph being processed.
+ * @param cnode The TupleGetItem node after which the Assign operation will be added.
+ * @return A CNodePtr pointing to the new node if insertion is successful, otherwise nullptr.
+ */
+CNodePtr InsertAssignAfterTupleGetItem(const FuncGraphPtr &func_graph, const CNodePtr &cnode) {
+  // Fetch the actual input node to the TupleGetItem.
+  auto input_node = cnode->input(kRealInputNodeIndexInTupleGetItem);
+  MS_EXCEPTION_IF_NULL(input_node);
+  auto real_input = dyn_cast<CNode>(input_node);
+  
+  if (real_input == nullptr) {
+    return nullptr;
+  }
+  
+  // Get the tuple item index.
+  auto value_ptr = GetValueNode(cnode->input(kInputNodeOutputIndexInTupleGetItem));
+  MS_EXCEPTION_IF_NULL(value_ptr);
+  auto gt_idx = GetValue<int64_t>(value_ptr);
+  
+  if (IsPrimitiveCNode(real_input, prim::kPrimCustom)) {
+    // Fetch inplace info for the real input node.
+    auto inplace_info = GetHybridInplaceIndex(real_input);
+    
+    for (auto index : inplace_info) {
+      // If there's a matching inplace info, insert Assign after the TupleGetItem.
+      if (index[kCustomOutput] == gt_idx && index[kCustomInput] >= 0) {
+        auto custom_input_size = common::AnfAlgo::GetInputTensorNum(real_input);
+        if (auto i = LongToSize(index[kCustomInput]); i < custom_input_size) {
+          return InsertAssign(func_graph, real_input->input(i + 1), cnode);
+        }
+      }
+    }
+  }
+  return nullptr;
+}
+
+/**
+ * @brief Process nodes to insert Assign operations for Custom ops.
+ * 
+ * This function checks for nodes of type Custom or TupleGetItem and attempts to insert 
+ * Assign operations after them, if needed.
+ * 
+ * @param func_graph The function graph being processed.
+ * @param node The current node being examined.
+ * @param Equiv equivalence class pointer, unused in this function.
+ * @return A new AnfNodePtr if the node was replaced, otherwise nullptr.
+ */
+const AnfNodePtr InsertAssignForCustomOp::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
+                                                  const EquivPtr &) const {
+  // Check if the given node is a Custom or TupleGetItem node and has not been visited.
+  if (cnode == nullptr) {
+    return nullptr;
+  }
+  
+  if (IsPrimitiveCNode(cnode, prim::kPrimCustom) && visited_.find(cnode) == visited_.end()) {
+    visited_.insert(cnode);
+    return InsertAssignAfterCustom(func_graph, cnode);
+  } else if (IsPrimitiveCNode(cnode, prim::kPrimTupleGetItem) && visited_.find(cnode) == visited_.end()) {
+    visited_.insert(cnode);
+    return InsertAssignAfterTupleGetItem(func_graph, cnode);
+  }
+
+  return nullptr;
+}
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/optimize_dependence.cc

@@ -31,10 +31,23 @@ constexpr auto kSingleInputIndex = 1;
 constexpr auto kIsolatedDependRealInputIndex = 0;
 constexpr auto kIsolatedDependVirtualInputIndex = 1;
 namespace {
+/**
+ * @brief Create a new depend node with the given inputs.
+ * 
+ * This function creates a new depend node, copying the attributes from the original node and updating the inputs.
+ * If the function graph is a KernelGraph, it creates a CNode via the KernelGraph, otherwise it uses the FuncGraph.
+ * 
+ * @param func_graph The current function graph.
+ * @param cnode The original node from which the new depend node will be created.
+ * @param new_depend_inputs The inputs for the new depend node.
+ * @return A CNodePtr pointing to the new depend node.
+ */
 CNodePtr CreateNewDependNode(const FuncGraphPtr &func_graph, const CNodePtr &cnode,
                              const std::vector<AnfNodePtr> &new_depend_inputs) {
   MS_EXCEPTION_IF_NULL(func_graph);
   MS_EXCEPTION_IF_NULL(cnode);
+
+  // Check if the function graph is a KernelGraph.
   auto kernel_graph = func_graph->cast<KernelGraphPtr>();
   if (kernel_graph == nullptr) {
     auto new_depend = func_graph->NewCNode(new_depend_inputs);
@@ -43,22 +56,40 @@ CNodePtr CreateNewDependNode(const FuncGraphPtr &func_graph, const CNodePtr &cno
     new_depend->set_scope(cnode->scope());
     return new_depend;
   }
+
+  // If it's a KernelGraph, create a new CNode using the KernelGraph.
   auto new_depend = kernel_graph->NewCNode(cnode);
   MS_EXCEPTION_IF_NULL(new_depend);
   new_depend->set_inputs(new_depend_inputs);
   return new_depend;
 }
 
+/**
+ * @brief Check if the given node is an isolated virtual node.
+ * 
+ * This function checks if the provided CNode is either a 'Depend' or 'Load' node. If it is, it further checks
+ * if the node is isolated and associated with a virtual operation. If all these conditions are met, 
+ * it returns the real input of the node.
+ * 
+ * @param cnode The CNode to check.
+ * @return A CNodePtr pointing to the real input of the node if the node is an isolated virtual node, otherwise nullptr.
+ */
 CNodePtr CheckIsolatedVirtualNode(const CNodePtr &cnode) {
   MS_EXCEPTION_IF_NULL(cnode);
+
+  // Check if the CNode is a 'Depend' or 'Load' node.
   if (common::AnfAlgo::GetCNodeName(cnode) != prim::kPrimDepend->name() &&
       common::AnfAlgo::GetCNodeName(cnode) != prim::kPrimLoad->name()) {
     return nullptr;
   }
+
+  // Check if the input of the node has a monadic abstract.
   auto virtual_input_op = common::AnfAlgo::GetInputNode(cnode, kIsolatedDependVirtualInputIndex);
   if (!HasAbstractMonad(virtual_input_op)) {
     return nullptr;
   }
+
+  // Return the real input of the node if it is a CNode.
   auto real_input_op = common::AnfAlgo::GetInputNode(cnode, kIsolatedDependRealInputIndex);
   MS_EXCEPTION_IF_NULL(real_input_op);
   if (!real_input_op->isa<CNode>()) {
@@ -68,62 +99,112 @@ CNodePtr CheckIsolatedVirtualNode(const CNodePtr &cnode) {
   return real_input_cnode;
 }
 
+
+/**
+ * @brief Eliminate the isolated virtual node input and replace it.
+ * 
+ * This function removes the isolated virtual node input from the node and replaces it with another node.
+ * 
+ * @param func_graph The current function graph.
+ * @param cnode The node whose isolated virtual input should be eliminated.
+ * @param eliminate_node The node that will replace the isolated virtual node.
+ * @return A AnfNodePtr pointing to the new node created after elimination.
+ */
 AnfNodePtr EliminateIsolatedVirtualNodeInput(const FuncGraphPtr &func_graph, const CNodePtr &cnode,
                                              const CNodePtr &eliminate_node) {
+  // Basic checks.
   MS_EXCEPTION_IF_NULL(func_graph);
   MS_EXCEPTION_IF_NULL(cnode);
   MS_EXCEPTION_IF_NULL(eliminate_node);
+
+  // Create the new node.
   auto replace_node = eliminate_node->input(kSingleInputIndex);
   std::vector<AnfNodePtr> new_depend_inputs = cnode->inputs();
   new_depend_inputs[kIsolatedDependRealInputIndex + 1] = replace_node;
   auto new_depend = CreateNewDependNode(func_graph, cnode, new_depend_inputs);
+
+  // Replace the old node with the new node in the graph.
   (void)func_graph->manager()->Replace(cnode, new_depend);
   return new_depend;
 }
 
+/**
+ * @brief Get the replace node for a given node.
+ * 
+ * This function determines if the provided node needs to be replaced and retrieves the appropriate node.
+ * Nodes like transdata or cast nodes which are not really used by others may need to be replaced.
+ * 
+ * @param func_graph The current function graph.
+ * @param node The node to check for replacement.
+ * @return A AnfNodePtr pointing to the replacement node if any, otherwise nullptr.
+ */
 AnfNodePtr GetReplaceNode(const FuncGraphPtr &func_graph, const AnfNodePtr &node) {
+  // Basic checks.
   MS_EXCEPTION_IF_NULL(func_graph);
   MS_EXCEPTION_IF_NULL(node);
   if (!node->isa<CNode>()) {
     return nullptr;
   }
+  
   auto cnode = node->cast<CNodePtr>();
   MS_EXCEPTION_IF_NULL(cnode);
+
   auto replace_cnode = cnode;
-  // Process updatestate and depend as isolated node env.
+
+  // Check if it's an isolated node.
   auto isolated_cnode = CheckIsolatedVirtualNode(replace_cnode);
   if (isolated_cnode != nullptr) {
     replace_cnode = isolated_cnode;
   }
+
   string op_name = common::AnfAlgo::GetCNodeName(replace_cnode);
-  // Currently we only eliminate transdata or cast nodes.
+
+  // Filter out nodes based on op name.
   if (op_name != kTransDataOpName && op_name != prim::kPrimCast->name()) {
     return nullptr;
   }
+
+  // More checks.
   if (!IsNotRealUsedByOthers(func_graph, replace_cnode)) {
     return nullptr;
   }
+  
   CheckCNodeInputSize(replace_cnode, kSingleInputIndex);
+  
   if (isolated_cnode != nullptr) {
     auto new_depend_node = EliminateIsolatedVirtualNodeInput(func_graph, cnode, replace_cnode);
     return new_depend_node;
   }
+
   return cnode->input(kSingleInputIndex);
 }
 
+/**
+ * @brief Replace MakeTuple node inputs if required.
+ * 
+ * This function replaces the MakeTuple node inputs with appropriate nodes if necessary.
+ * 
+ * @param func_graph The current function graph.
+ * @param cnode The MakeTuple node to be checked.
+ * @return A AnfNodePtr pointing to the new MakeTuple node if any changes were made, otherwise nullptr.
+ */
 AnfNodePtr ReplaceMakeTuple(const FuncGraphPtr &func_graph, const CNodePtr &cnode) {
+  // Basic checks.
   MS_EXCEPTION_IF_NULL(func_graph);
   MS_EXCEPTION_IF_NULL(cnode);
+
   if (common::AnfAlgo::GetCNodeName(cnode) != prim::kPrimMakeTuple->name()) {
     return nullptr;
   }
+  
   std::vector<AnfNodePtr> new_make_tuple_inputs = {common::AnfAlgo::GetCNodePrimitiveNode(cnode)};
   bool need_update = false;
   size_t input_num = common::AnfAlgo::GetInputTensorNum(cnode);
+  
+  // Check and possibly replace each input.
   for (size_t index = 0; index < input_num; ++index) {
     auto input = common::AnfAlgo::GetInputNode(cnode, index);
     AnfNodePtr replace_input = GetReplaceNode(func_graph, input);
-    // If replace input is not null, it will be the input of the TransData or Cast.
     if (replace_input == nullptr) {
       new_make_tuple_inputs.push_back(input);
       continue;
@@ -131,6 +212,8 @@ AnfNodePtr ReplaceMakeTuple(const FuncGraphPtr &func_graph, const CNodePtr &cnod
     new_make_tuple_inputs.push_back(replace_input);
     need_update = true;
   }
+  
+  // If updates were made, replace the old MakeTuple node with a new one.
   if (need_update) {
     auto kernel_graph = func_graph->cast<std::shared_ptr<session::KernelGraph>>();
     CNodePtr new_make_tuple = nullptr;
@@ -139,6 +222,7 @@ AnfNodePtr ReplaceMakeTuple(const FuncGraphPtr &func_graph, const CNodePtr &cnod
     } else {
       new_make_tuple = kernel_graph->NewCNode(cnode);
     }
+
     MS_EXCEPTION_IF_NULL(new_make_tuple);
     new_make_tuple->set_inputs(new_make_tuple_inputs);
     auto manager = func_graph->manager();
@@ -146,23 +230,40 @@ AnfNodePtr ReplaceMakeTuple(const FuncGraphPtr &func_graph, const CNodePtr &cnod
     manager->Replace(cnode, new_make_tuple);
     return new_make_tuple;
   }
+
   return nullptr;
 }
 }  // namespace
 
+/**
+ * @brief Define a pattern for the OptimizeDependence pass.
+ * 
+ * This function defines a pattern to be matched in the graph for the subsequent optimization.
+ * 
+ * @return BaseRef that represents the defined pattern.
+ */
 const BaseRef OptimizeDependence::DefinePattern() const {
   VarPtr X = std::make_shared<Var>();
   VarPtr Xs = std::make_shared<SeqVar>();
   return VectorRef({X, Xs});
 }
 
+/**
+ * @brief Search for nodes of type TransData and Cast within the inputs of a given node.
+ * 
+ * This function identifies inputs which are of type TransData or Cast.
+ * 
+ * @param cnode Node whose inputs are to be examined.
+ * @return A vector of indices pointing to TransData and Cast inputs.
+ */
 std::vector<size_t> SearchTransDataAndCast(const CNodePtr &cnode) {
-  // Search Depend and UpdateState only.
+  // Only search within Depend and UpdateState nodes.
   if (!cnode->IsApply(prim::kPrimDepend) && !cnode->IsApply(prim::kPrimUpdateState)) {
     return {};
   }
-  // Find inputs which is Cast or TransData.
+
   std::vector<size_t> result;
+  // Identify indices of inputs which are Cast, TransData or MakeTuple.
   for (size_t i = 1; i < cnode->size(); ++i) {
     auto &input = cnode->input(i);
     if (common::AnfAlgo::CheckPrimitiveType(input, prim::kPrimCast) ||
@@ -174,25 +275,40 @@ std::vector<size_t> SearchTransDataAndCast(const CNodePtr &cnode) {
   return result;
 }
 
+/**
+ * @brief Process nodes to optimize dependence in the graph.
+ * 
+ * This function replaces specific inputs of the Depend or UpdateState nodes
+ * that match certain conditions for optimization.
+ * 
+ * @param func_graph The function graph being processed.
+ * @param node The current node being examined.
+ * @param Equiv equivalence class pointer, unused in this function.
+ * @return A new AnfNodePtr if the node was replaced, otherwise nullptr.
+ */
 const AnfNodePtr OptimizeDependence::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
                                              const EquivPtr &) const {
+  // Basic validations.
   MS_EXCEPTION_IF_NULL(func_graph);
   MS_EXCEPTION_IF_NULL(node);
   auto cnode = dyn_cast<CNode>(node);
   if (cnode == nullptr) {
     return nullptr;
   }
-  // Search inputs to be replaced.
+  
+  // Identify which inputs might need to be replaced.
   auto candidate_inputs = SearchTransDataAndCast(cnode);
   if (candidate_inputs.empty()) {
     return nullptr;
   }
-  // Get new nodes which will act as new inputs of Depend or UpdateState.
+
   std::vector<AnfNodePtr> new_inputs = cnode->inputs();
   bool inputs_changed = false;
+  
+  // Attempt to replace each identified candidate input.
   for (auto index : candidate_inputs) {
     if (index >= new_inputs.size()) {
-      MS_LOG(EXCEPTION) << "Index is out of the size of " << cnode->DebugString() << trace::DumpSourceLines(cnode);
+      MS_LOG(EXCEPTION) << "Index out of bounds for " << cnode->DebugString() << trace::DumpSourceLines(cnode);
     }
     auto replace_node = GetConvertNode(func_graph, cnode, index);
     if (replace_node != nullptr) {
@@ -200,35 +316,53 @@ const AnfNodePtr OptimizeDependence::Process(const FuncGraphPtr &func_graph, con
       inputs_changed = true;
     }
   }
+
+  // If any inputs were changed, replace the original node with the new inputs.
   if (!inputs_changed) {
     return nullptr;
   }
-  // Create a new Depend node to replace the old one if inputs changed.
   auto new_depend = CreateNewDependNode(func_graph, cnode, new_inputs);
   (void)func_graph->manager()->Replace(cnode, new_depend);
   return nullptr;
 }
 
+/**
+ * @brief Obtain the node to replace the given node in the graph.
+ * 
+ * This function identifies the suitable replacement for a given node in the graph, focusing on optimizations.
+ * 
+ * @param graph The current function graph.
+ * @param node The node to be replaced.
+ * @param index The index pointing to the node's input that might be replaced.
+ * @return AnfNodePtr pointing to the replacement node if found, otherwise nullptr.
+ */
 const AnfNodePtr OptimizeDependence::GetConvertNode(const FuncGraphPtr &graph, const AnfNodePtr &node,
                                                     const size_t index) const {
+  // Basic validations.
   MS_EXCEPTION_IF_NULL(graph);
   MS_EXCEPTION_IF_NULL(node);
   auto depend_cnode = node->cast<CNodePtr>();
   MS_EXCEPTION_IF_NULL(depend_cnode);
+  
   auto replacing_node = depend_cnode->input(index);
   MS_EXCEPTION_IF_NULL(replacing_node);
   if (!replacing_node->isa<CNode>()) {
     return nullptr;
   }
+  
   auto replacing_cnode = replacing_node->cast<CNodePtr>();
   MS_EXCEPTION_IF_NULL(replacing_cnode);
-  // Deal with the make_tuple with TransData or Cast inputs.
+
+  // Handle the case of MakeTuple nodes with TransData or Cast inputs.
   auto make_tuple_replace_node = ReplaceMakeTuple(graph, replacing_cnode);
   if (make_tuple_replace_node != nullptr) {
     return make_tuple_replace_node;
   }
+  
+  // Get the node to replace the given cnode.
   AnfNodePtr replace_node = GetReplaceNode(graph, replacing_cnode);
   return replace_node;
 }
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/optimize_dependence_old.cc

@@ -0,0 +1,368 @@
+/**
+ * Copyright 2019 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "backend/common/pass/optimize_dependence.h"
+#include <memory>
+#include <vector>
+#include <string>
+#include "backend/common/optimizer/helper.h"
+#include "base/core_ops.h"
+#include "include/common/utils/utils.h"
+#include "include/common/utils/anfalgo.h"
+#include "utils/trace_base.h"
+#include "backend/common/session/kernel_graph.h"
+
+namespace mindspore {
+namespace opt {
+constexpr auto kSingleInputIndex = 1;
+constexpr auto kIsolatedDependRealInputIndex = 0;
+constexpr auto kIsolatedDependVirtualInputIndex = 1;
+namespace {
+/**
+ * @brief Create a new depend node with the given inputs.
+ * 
+ * This function creates a new depend node, copying the attributes from the original node and updating the inputs.
+ * If the function graph is a KernelGraph, it creates a CNode via the KernelGraph, otherwise it uses the FuncGraph.
+ * 
+ * @param func_graph The current function graph.
+ * @param cnode The original node from which the new depend node will be created.
+ * @param new_depend_inputs The inputs for the new depend node.
+ * @return A CNodePtr pointing to the new depend node.
+ */
+CNodePtr CreateNewDependNode(const FuncGraphPtr &func_graph, const CNodePtr &cnode,
+                             const std::vector<AnfNodePtr> &new_depend_inputs) {
+  MS_EXCEPTION_IF_NULL(func_graph);
+  MS_EXCEPTION_IF_NULL(cnode);
+
+  // Check if the function graph is a KernelGraph.
+  auto kernel_graph = func_graph->cast<KernelGraphPtr>();
+  if (kernel_graph == nullptr) {
+    auto new_depend = func_graph->NewCNode(new_depend_inputs);
+    MS_EXCEPTION_IF_NULL(new_depend);
+    new_depend->set_abstract(cnode->abstract());
+    new_depend->set_scope(cnode->scope());
+    return new_depend;
+  }
+
+  // If it's a KernelGraph, create a new CNode using the KernelGraph.
+  auto new_depend = kernel_graph->NewCNode(cnode);
+  MS_EXCEPTION_IF_NULL(new_depend);
+  new_depend->set_inputs(new_depend_inputs);
+  return new_depend;
+}
+
+/**
+ * @brief Check if the given node is an isolated virtual node.
+ * 
+ * This function checks if the provided CNode is either a 'Depend' or 'Load' node. If it is, it further checks
+ * if the node is isolated and associated with a virtual operation. If all these conditions are met, 
+ * it returns the real input of the node.
+ * 
+ * @param cnode The CNode to check.
+ * @return A CNodePtr pointing to the real input of the node if the node is an isolated virtual node, otherwise nullptr.
+ */
+CNodePtr CheckIsolatedVirtualNode(const CNodePtr &cnode) {
+  MS_EXCEPTION_IF_NULL(cnode);
+
+  // Check if the CNode is a 'Depend' or 'Load' node.
+  if (common::AnfAlgo::GetCNodeName(cnode) != prim::kPrimDepend->name() &&
+      common::AnfAlgo::GetCNodeName(cnode) != prim::kPrimLoad->name()) {
+    return nullptr;
+  }
+
+  // Check if the input of the node has a monadic abstract.
+  auto virtual_input_op = common::AnfAlgo::GetInputNode(cnode, kIsolatedDependVirtualInputIndex);
+  if (!HasAbstractMonad(virtual_input_op)) {
+    return nullptr;
+  }
+
+  // Return the real input of the node if it is a CNode.
+  auto real_input_op = common::AnfAlgo::GetInputNode(cnode, kIsolatedDependRealInputIndex);
+  MS_EXCEPTION_IF_NULL(real_input_op);
+  if (!real_input_op->isa<CNode>()) {
+    return nullptr;
+  }
+  auto real_input_cnode = real_input_op->cast<CNodePtr>();
+  return real_input_cnode;
+}
+
+
+/**
+ * @brief Eliminate the isolated virtual node input and replace it.
+ * 
+ * This function removes the isolated virtual node input from the node and replaces it with another node.
+ * 
+ * @param func_graph The current function graph.
+ * @param cnode The node whose isolated virtual input should be eliminated.
+ * @param eliminate_node The node that will replace the isolated virtual node.
+ * @return A AnfNodePtr pointing to the new node created after elimination.
+ */
+AnfNodePtr EliminateIsolatedVirtualNodeInput(const FuncGraphPtr &func_graph, const CNodePtr &cnode,
+                                             const CNodePtr &eliminate_node) {
+  // Basic checks.
+  MS_EXCEPTION_IF_NULL(func_graph);
+  MS_EXCEPTION_IF_NULL(cnode);
+  MS_EXCEPTION_IF_NULL(eliminate_node);
+
+  // Create the new node.
+  auto replace_node = eliminate_node->input(kSingleInputIndex);
+  std::vector<AnfNodePtr> new_depend_inputs = cnode->inputs();
+  new_depend_inputs[kIsolatedDependRealInputIndex + 1] = replace_node;
+  auto new_depend = CreateNewDependNode(func_graph, cnode, new_depend_inputs);
+
+  // Replace the old node with the new node in the graph.
+  (void)func_graph->manager()->Replace(cnode, new_depend);
+  return new_depend;
+}
+
+/**
+ * @brief Get the replace node for a given node.
+ * 
+ * This function determines if the provided node needs to be replaced and retrieves the appropriate node.
+ * Nodes like transdata or cast nodes which are not really used by others may need to be replaced.
+ * 
+ * @param func_graph The current function graph.
+ * @param node The node to check for replacement.
+ * @return A AnfNodePtr pointing to the replacement node if any, otherwise nullptr.
+ */
+AnfNodePtr GetReplaceNode(const FuncGraphPtr &func_graph, const AnfNodePtr &node) {
+  // Basic checks.
+  MS_EXCEPTION_IF_NULL(func_graph);
+  MS_EXCEPTION_IF_NULL(node);
+  if (!node->isa<CNode>()) {
+    return nullptr;
+  }
+  
+  auto cnode = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(cnode);
+
+  auto replace_cnode = cnode;
+
+  // Check if it's an isolated node.
+  auto isolated_cnode = CheckIsolatedVirtualNode(replace_cnode);
+  if (isolated_cnode != nullptr) {
+    replace_cnode = isolated_cnode;
+  }
+
+  string op_name = common::AnfAlgo::GetCNodeName(replace_cnode);
+
+  // Filter out nodes based on op name.
+  if (op_name != kTransDataOpName && op_name != prim::kPrimCast->name()) {
+    return nullptr;
+  }
+
+  // More checks.
+  if (!IsNotRealUsedByOthers(func_graph, replace_cnode)) {
+    return nullptr;
+  }
+  
+  CheckCNodeInputSize(replace_cnode, kSingleInputIndex);
+  
+  if (isolated_cnode != nullptr) {
+    auto new_depend_node = EliminateIsolatedVirtualNodeInput(func_graph, cnode, replace_cnode);
+    return new_depend_node;
+  }
+
+  return cnode->input(kSingleInputIndex);
+}
+
+/**
+ * @brief Replace MakeTuple node inputs if required.
+ * 
+ * This function replaces the MakeTuple node inputs with appropriate nodes if necessary.
+ * 
+ * @param func_graph The current function graph.
+ * @param cnode The MakeTuple node to be checked.
+ * @return A AnfNodePtr pointing to the new MakeTuple node if any changes were made, otherwise nullptr.
+ */
+AnfNodePtr ReplaceMakeTuple(const FuncGraphPtr &func_graph, const CNodePtr &cnode) {
+  // Basic checks.
+  MS_EXCEPTION_IF_NULL(func_graph);
+  MS_EXCEPTION_IF_NULL(cnode);
+
+  if (common::AnfAlgo::GetCNodeName(cnode) != prim::kPrimMakeTuple->name()) {
+    return nullptr;
+  }
+  
+  std::vector<AnfNodePtr> new_make_tuple_inputs = {common::AnfAlgo::GetCNodePrimitiveNode(cnode)};
+  bool need_update = false;
+  size_t input_num = common::AnfAlgo::GetInputTensorNum(cnode);
+  
+  // Check and possibly replace each input.
+  for (size_t index = 0; index < input_num; ++index) {
+    auto input = common::AnfAlgo::GetInputNode(cnode, index);
+    AnfNodePtr replace_input = GetReplaceNode(func_graph, input);
+    if (replace_input == nullptr) {
+      new_make_tuple_inputs.push_back(input);
+      continue;
+    }
+    new_make_tuple_inputs.push_back(replace_input);
+    need_update = true;
+  }
+  
+  // If updates were made, replace the old MakeTuple node with a new one.
+  if (need_update) {
+    auto kernel_graph = func_graph->cast<std::shared_ptr<session::KernelGraph>>();
+    CNodePtr new_make_tuple = nullptr;
+    if (kernel_graph == nullptr) {
+      new_make_tuple = func_graph->NewCNode(new_make_tuple_inputs);
+    } else {
+      new_make_tuple = kernel_graph->NewCNode(cnode);
+    }
+
+    MS_EXCEPTION_IF_NULL(new_make_tuple);
+    new_make_tuple->set_inputs(new_make_tuple_inputs);
+    auto manager = func_graph->manager();
+    MS_EXCEPTION_IF_NULL(manager);
+    manager->Replace(cnode, new_make_tuple);
+    return new_make_tuple;
+  }
+
+  return nullptr;
+}
+}  // namespace
+
+/**
+ * @brief Define a pattern for the OptimizeDependence pass.
+ * 
+ * This function defines a pattern to be matched in the graph for the subsequent optimization.
+ * 
+ * @return BaseRef that represents the defined pattern.
+ */
+const BaseRef OptimizeDependence::DefinePattern() const {
+  VarPtr X = std::make_shared<Var>();
+  VarPtr Xs = std::make_shared<SeqVar>();
+  return VectorRef({X, Xs});
+}
+
+/**
+ * @brief Search for nodes of type TransData and Cast within the inputs of a given node.
+ * 
+ * This function identifies inputs which are of type TransData or Cast.
+ * 
+ * @param cnode Node whose inputs are to be examined.
+ * @return A vector of indices pointing to TransData and Cast inputs.
+ */
+std::vector<size_t> SearchTransDataAndCast(const CNodePtr &cnode) {
+  // Only search within Depend and UpdateState nodes.
+  if (!cnode->IsApply(prim::kPrimDepend) && !cnode->IsApply(prim::kPrimUpdateState)) {
+    return {};
+  }
+
+  std::vector<size_t> result;
+  // Identify indices of inputs which are Cast, TransData or MakeTuple.
+  for (size_t i = 1; i < cnode->size(); ++i) {
+    auto &input = cnode->input(i);
+    if (common::AnfAlgo::CheckPrimitiveType(input, prim::kPrimCast) ||
+        common::AnfAlgo::CheckPrimitiveType(input, prim::kPrimTransData) ||
+        common::AnfAlgo::CheckPrimitiveType(input, prim::kPrimMakeTuple)) {
+      (void)result.emplace_back(i);
+    }
+  }
+  return result;
+}
+
+/**
+ * @brief Process nodes to optimize dependence in the graph.
+ * 
+ * This function replaces specific inputs of the Depend or UpdateState nodes
+ * that match certain conditions for optimization.
+ * 
+ * @param func_graph The function graph being processed.
+ * @param node The current node being examined.
+ * @param Equiv equivalence class pointer, unused in this function.
+ * @return A new AnfNodePtr if the node was replaced, otherwise nullptr.
+ */
+const AnfNodePtr OptimizeDependence::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
+                                             const EquivPtr &) const {
+  // Basic validations.
+  MS_EXCEPTION_IF_NULL(func_graph);
+  MS_EXCEPTION_IF_NULL(node);
+  auto cnode = dyn_cast<CNode>(node);
+  if (cnode == nullptr) {
+    return nullptr;
+  }
+  
+  // Identify which inputs might need to be replaced.
+  auto candidate_inputs = SearchTransDataAndCast(cnode);
+  if (candidate_inputs.empty()) {
+    return nullptr;
+  }
+
+  std::vector<AnfNodePtr> new_inputs = cnode->inputs();
+  bool inputs_changed = false;
+  
+  // Attempt to replace each identified candidate input.
+  for (auto index : candidate_inputs) {
+    if (index >= new_inputs.size()) {
+      MS_LOG(EXCEPTION) << "Index out of bounds for " << cnode->DebugString() << trace::DumpSourceLines(cnode);
+    }
+    auto replace_node = GetConvertNode(func_graph, cnode, index);
+    if (replace_node != nullptr) {
+      new_inputs[index] = replace_node;
+      inputs_changed = true;
+    }
+  }
+
+  // If any inputs were changed, replace the original node with the new inputs.
+  if (!inputs_changed) {
+    return nullptr;
+  }
+  auto new_depend = CreateNewDependNode(func_graph, cnode, new_inputs);
+  (void)func_graph->manager()->Replace(cnode, new_depend);
+  return nullptr;
+}
+
+/**
+ * @brief Obtain the node to replace the given node in the graph.
+ * 
+ * This function identifies the suitable replacement for a given node in the graph, focusing on optimizations.
+ * 
+ * @param graph The current function graph.
+ * @param node The node to be replaced.
+ * @param index The index pointing to the node's input that might be replaced.
+ * @return AnfNodePtr pointing to the replacement node if found, otherwise nullptr.
+ */
+const AnfNodePtr OptimizeDependence::GetConvertNode(const FuncGraphPtr &graph, const AnfNodePtr &node,
+                                                    const size_t index) const {
+  // Basic validations.
+  MS_EXCEPTION_IF_NULL(graph);
+  MS_EXCEPTION_IF_NULL(node);
+  auto depend_cnode = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(depend_cnode);
+  
+  auto replacing_node = depend_cnode->input(index);
+  MS_EXCEPTION_IF_NULL(replacing_node);
+  if (!replacing_node->isa<CNode>()) {
+    return nullptr;
+  }
+  
+  auto replacing_cnode = replacing_node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(replacing_cnode);
+
+  // Handle the case of MakeTuple nodes with TransData or Cast inputs.
+  auto make_tuple_replace_node = ReplaceMakeTuple(graph, replacing_cnode);
+  if (make_tuple_replace_node != nullptr) {
+    return make_tuple_replace_node;
+  }
+  
+  // Get the node to replace the given cnode.
+  AnfNodePtr replace_node = GetReplaceNode(graph, replacing_cnode);
+  return replace_node;
+}
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/optimize_updatestate.cc

@@ -27,47 +27,77 @@ constexpr size_t kInputIndex = 1;
 constexpr size_t kAttachIndex = 2;
 constexpr size_t kAdditionalAttachIndex = 3;
 
+/**
+ * @brief Define a pattern for the OptimizeUpdateState.
+ * 
+ * This function defines a pattern with the prim::kPrimUpdateState primitive followed by a sequence variable.
+ * 
+ * @return A VectorRef representing the defined pattern.
+ */
 const BaseRef OptimizeUpdateState::DefinePattern() const {
+  // Create a sequence variable.
   VarPtr Xs = std::make_shared<SeqVar>();
   return VectorRef({prim::kPrimUpdateState, Xs});
 }
 
+/**
+ * @brief Processes the given node to optimize UpdateState operations.
+ * 
+ * This function optimizes UpdateState operations by reducing unnecessary attach nodes. If the attach node is only used
+ * by the UpdateState operation and is not a parameter, it will be dropped. If there are changes in the attaches, a new
+ * UpdateState node will be created with the optimized attaches.
+ * 
+ * @param func_graph The current function graph.
+ * @param node The node to be processed.
+ * @param EquivPtr Placeholder for equivalence. Not used in this context.
+ * @return An optimized AnfNodePtr if applicable, or nullptr if no optimization was applied.
+ */
 const AnfNodePtr OptimizeUpdateState::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
                                               const EquivPtr &) const {
+  // Ensure func_graph and node are not null.
   MS_EXCEPTION_IF_NULL(func_graph);
   MS_EXCEPTION_IF_NULL(node);
+
+  // Cast the node to a CNode.
   auto update_state = dyn_cast<CNode>(node);
   MS_EXCEPTION_IF_NULL(update_state);
+
+  // If the UpdateState node doesn't have additional attaches, skip processing.
   if (update_state->size() <= kAdditionalAttachIndex) {
-    // Skip UpdateState nodes with no additional attaches.
     return nullptr;
   }
+
   auto manager = func_graph->manager();
   MS_EXCEPTION_IF_NULL(manager);
+  
   auto &node_users = manager->node_users();
   std::vector<AnfNodePtr> new_inputs;
   (void)new_inputs.emplace_back(update_state->input(0));
   (void)new_inputs.emplace_back(update_state->input(kInputIndex));
   (void)new_inputs.emplace_back(update_state->input(kAttachIndex));
+
+  // Process each attach and decide whether to keep it.
   for (size_t i = kAdditionalAttachIndex; i < update_state->size(); ++i) {
     auto &attach = update_state->input(i);
     auto &users = node_users[attach];
-    // In heterogeneous, parameters in subgraphs may only be used by UpdateState and should not be eliminated.
     if ((users.size() == 1) && (users.front().first == update_state) && !attach->isa<Parameter>()) {
-      // If the only user of attach is the UpdateState node, drop the attach node.
+      // If the only user of attach is the UpdateState node and it's not a parameter, skip adding this attach.
       continue;
     }
     (void)new_inputs.emplace_back(attach);
   }
+
+  // If there were no changes in the attaches, return nullptr.
   if (new_inputs.size() == update_state->size()) {
-    // Attaches not changed.
     return nullptr;
   }
-  // Attaches changed, make a new UpdateState.
+
+  // If the attaches changed, create a new UpdateState node.
   auto new_update_state = func_graph->NewCNode(new_inputs);
   new_update_state->set_abstract(update_state->abstract());
   new_update_state->set_scope(update_state->scope());
   return new_update_state;
 }
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/optimize_updatestate_old.cc

@@ -0,0 +1,103 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "backend/common/pass/optimize_updatestate.h"
+#include <memory>
+#include <vector>
+#include <string>
+#include "base/core_ops.h"
+#include "include/common/utils/utils.h"
+
+namespace mindspore {
+namespace opt {
+constexpr size_t kInputIndex = 1;
+constexpr size_t kAttachIndex = 2;
+constexpr size_t kAdditionalAttachIndex = 3;
+
+/**
+ * @brief Define a pattern for the OptimizeUpdateState.
+ * 
+ * This function defines a pattern with the prim::kPrimUpdateState primitive followed by a sequence variable.
+ * 
+ * @return A VectorRef representing the defined pattern.
+ */
+const BaseRef OptimizeUpdateState::DefinePattern() const {
+  // Create a sequence variable.
+  VarPtr Xs = std::make_shared<SeqVar>();
+  return VectorRef({prim::kPrimUpdateState, Xs});
+}
+
+/**
+ * @brief Processes the given node to optimize UpdateState operations.
+ * 
+ * This function optimizes UpdateState operations by reducing unnecessary attach nodes. If the attach node is only used
+ * by the UpdateState operation and is not a parameter, it will be dropped. If there are changes in the attaches, a new
+ * UpdateState node will be created with the optimized attaches.
+ * 
+ * @param func_graph The current function graph.
+ * @param node The node to be processed.
+ * @param EquivPtr Placeholder for equivalence. Not used in this context.
+ * @return An optimized AnfNodePtr if applicable, or nullptr if no optimization was applied.
+ */
+const AnfNodePtr OptimizeUpdateState::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
+                                              const EquivPtr &) const {
+  // Ensure func_graph and node are not null.
+  MS_EXCEPTION_IF_NULL(func_graph);
+  MS_EXCEPTION_IF_NULL(node);
+
+  // Cast the node to a CNode.
+  auto update_state = dyn_cast<CNode>(node);
+  MS_EXCEPTION_IF_NULL(update_state);
+
+  // If the UpdateState node doesn't have additional attaches, skip processing.
+  if (update_state->size() <= kAdditionalAttachIndex) {
+    return nullptr;
+  }
+
+  auto manager = func_graph->manager();
+  MS_EXCEPTION_IF_NULL(manager);
+  
+  auto &node_users = manager->node_users();
+  std::vector<AnfNodePtr> new_inputs;
+  (void)new_inputs.emplace_back(update_state->input(0));
+  (void)new_inputs.emplace_back(update_state->input(kInputIndex));
+  (void)new_inputs.emplace_back(update_state->input(kAttachIndex));
+
+  // Process each attach and decide whether to keep it.
+  for (size_t i = kAdditionalAttachIndex; i < update_state->size(); ++i) {
+    auto &attach = update_state->input(i);
+    auto &users = node_users[attach];
+    if ((users.size() == 1) && (users.front().first == update_state) && !attach->isa<Parameter>()) {
+      // If the only user of attach is the UpdateState node and it's not a parameter, skip adding this attach.
+      continue;
+    }
+    (void)new_inputs.emplace_back(attach);
+  }
+
+  // If there were no changes in the attaches, return nullptr.
+  if (new_inputs.size() == update_state->size()) {
+    return nullptr;
+  }
+
+  // If the attaches changed, create a new UpdateState node.
+  auto new_update_state = func_graph->NewCNode(new_inputs);
+  new_update_state->set_abstract(update_state->abstract());
+  new_update_state->set_scope(update_state->scope());
+  return new_update_state;
+}
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/reduce_sum_optimizer.cc

@@ -23,23 +23,43 @@ namespace opt {
 namespace {
 const int axis_input_index = 2;
 
+
+/**
+ * @brief Check if rank computation is needed for a given CNode.
+ *
+ * This function examines the axis input of the CNode. If the axis input is a tuple that
+ * is either empty or contains negative values, the rank computation is deemed necessary.
+ * 
+ * @param cnode The CNode to be checked.
+ * @return true if rank computation is needed, false otherwise.
+ */
 bool IsNeedComputeRank(const CNodePtr &cnode) {
+  // Ensure the CNode is not null.
   MS_EXCEPTION_IF_NULL(cnode);
+
+  // Retrieve the axis input from the CNode.
   auto axis_input = cnode->input(axis_input_index);
   MS_EXCEPTION_IF_NULL(axis_input);
+
+  // If the axis input is not a tuple, rank computation is not needed.
   if (!IsValueNode<ValueTuple>(axis_input)) {
     return false;
   }
+
+  // Cast the axis input to a ValueNode and retrieve its value.
   auto value_node = axis_input->cast<ValueNodePtr>();
   MS_EXCEPTION_IF_NULL(value_node);
   auto value = value_node->value();
   MS_EXCEPTION_IF_NULL(value);
+
   if (value->isa<ValueTuple>()) {
     auto value_tuple = value->cast<ValueTuplePtr>();
     MS_EXCEPTION_IF_NULL(value_tuple);
+    // If the tuple is empty, rank computation is needed.
     if (value_tuple->value().empty()) {
       return true;
     }
+    // Check each item in the tuple. If any item is negative, rank computation is needed.
     for (auto &iter : value_tuple->value()) {
       auto item = GetValue<int64_t>(iter->cast<ScalarPtr>());
       if (item < 0) {
@@ -49,44 +69,74 @@ bool IsNeedComputeRank(const CNodePtr &cnode) {
   }
   return false;
 }
-}  // namespace
 
+/**
+ * @brief Create a new Rank operation node.
+ *
+ * The Rank operation computes the rank (number of dimensions) of a tensor.
+ * This function creates a new CNode that represents a Rank operation in the MindSpore graph.
+ * 
+ * @param cnode The original CNode for which the Rank operation is being created.
+ * @param kernel_graph The current kernel graph.
+ * @return Pointer to the newly created Rank operation node.
+ */
 AnfNodePtr ReduceSumOptimizer::NewRankOp(const AnfNodePtr &cnode, const KernelGraphPtr &kernel_graph) const {
+  // Initialize inputs for the Rank operation.
   std::vector<AnfNodePtr> rank_inputs;
   auto prim = std::make_shared<Primitive>(prim::kPrimRank->name());
   rank_inputs.push_back(NewValueNode(prim));
+  
+  // Get the previous node's output.
   auto prev_node = common::AnfAlgo::GetPrevNodeOutput(cnode, 1);
   rank_inputs.push_back(prev_node.first);
+  
+  // Create the Rank operation node.
   auto rank_op = NewCNode(rank_inputs, kernel_graph);
   MS_EXCEPTION_IF_NULL(rank_op);
   rank_op->set_abstract(prev_node.first->abstract());
+
   return rank_op;
 }
 
+/**
+ * @brief Create a new Range operation node.
+ *
+ * The Range operation generates a sequence of numbers. In this context, it is used to
+ * generate a sequence based on the rank of a tensor.
+ * This function creates a new CNode that represents a Range operation in the MindSpore graph.
+ * 
+ * @param rank_op The Rank operation node that provides the limit for the Range operation.
+ * @param kernel_graph The current kernel graph.
+ * @return Pointer to the newly created Range operation node.
+ */
 AnfNodePtr ReduceSumOptimizer::NewRangeOp(const AnfNodePtr &rank_op, const KernelGraphPtr &kernel_graph) const {
+  // Initialize inputs for the Range operation.
   std::vector<AnfNodePtr> range_inputs;
   auto prim = std::make_shared<Primitive>(prim::kPrimRange->name());
   range_inputs.push_back(NewValueNode(prim));
-  // "start"
+
+  // Set the "start" of the range.
   auto start_ = NewValueNode(SizeToLong(0));
   MS_EXCEPTION_IF_NULL(start_);
   auto imm_start = std::make_shared<Int64Imm>(SizeToLong(0));
   start_->set_abstract(std::make_shared<abstract::AbstractScalar>(imm_start));
   range_inputs.push_back(start_);
 
-  // "limit"
+  // Set the "limit" of the range.
   range_inputs.push_back(rank_op);
 
-  // "delta"
+  // Set the "delta" of the range.
   auto delta_ = NewValueNode(SizeToLong(1));
   MS_EXCEPTION_IF_NULL(delta_);
   auto imm_delta = std::make_shared<Int64Imm>(SizeToLong(1));
   delta_->set_abstract(std::make_shared<abstract::AbstractScalar>(imm_delta));
   range_inputs.push_back(delta_);
-  // new range op
+
+  // Create the Range operation node.
   auto range_op = NewCNode(range_inputs, kernel_graph);
   MS_EXCEPTION_IF_NULL(range_op);
   range_op->set_abstract(rank_op->abstract());
+
   return range_op;
 }
 
@@ -150,34 +200,74 @@ AnfNodePtr ReduceSumOptimizer::NewAssistValueNode(const CNodePtr &cnode, const K
   return nullptr;
 }
 
+/**
+ * @brief Processes the given node for optimization in the context of a ReduceSum operation.
+ * 
+ * This function optimizes ReduceSum operations under certain conditions. If the current node is not a ReduceSum
+ * operation or if it does not have a dynamic shape, it will be skipped. If the node has unknown dimensions
+ * and needs rank computation, an assist node will be inserted. Otherwise, a new assist value node will be created.
+ *
+ * @param func_graph The current function graph.
+ * @param node The node to be processed.
+ * @param EquivPtr Placeholder for equivalence. Not used in this context.
+ * @return An optimized AnfNodePtr if applicable, or nullptr if no optimization was applied.
+ */
 const AnfNodePtr ReduceSumOptimizer::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
                                              const EquivPtr &) const {
+  // Ensure func_graph and node are not null.
   MS_EXCEPTION_IF_NULL(func_graph);
   MS_EXCEPTION_IF_NULL(node);
+  
+  // Cast the node to a CNode.
   auto cnode = node->cast<CNodePtr>();
   MS_EXCEPTION_IF_NULL(cnode);
+
+  // Get the name of the CNode operation.
   auto op_name = common::AnfAlgo::GetCNodeName(cnode);
+  
+  // If the current node is not a ReduceSum operation, skip processing.
   if (op_name != kReduceSumOpName) {
     MS_LOG(DEBUG) << "Current node is not: " << kReduceSumOpName << ", skip!";
     return nullptr;
   }
+  
+  // If the current node does not have a dynamic shape, skip processing.
   if (!common::AnfAlgo::IsDynamicShape(cnode)) {
     MS_LOG(DEBUG) << "Current node is not dynamic shape, skip!";
     return nullptr;
   }
+
+  // Mark the node as visited.
   common::AnfAlgo::SetNodeAttr(kAttrVisited, MakeValue(true), node);
+
+  // Cast the func_graph to a KernelGraph.
   auto kernel_graph = func_graph->cast<std::shared_ptr<session::KernelGraph>>();
   MS_EXCEPTION_IF_NULL(kernel_graph);
+
+  // If the CNode has unknown dimensions and requires rank computation, insert an assist node.
   if (AnfUtils::IsDimUnknown(cnode) && IsNeedComputeRank(cnode)) {
     return InsertAssistNode(cnode, kernel_graph);
   }
+  
+  // Create a new assist value node.
   return NewAssistValueNode(cnode, kernel_graph);
 }
 
+/**
+ * @brief Define a pattern for the ReduceSumOptimizer.
+ * 
+ * This function defines a pattern with a condition variable followed by a sequence variable. The condition variable
+ * ensures that the pattern is not visited more than once.
+ * 
+ * @return A VectorRef representing the defined pattern.
+ */
 const BaseRef ReduceSumOptimizer::DefinePattern() const {
+  // Create a condition variable to ensure the pattern is not visited more than once.
   std::shared_ptr<Var> V = std::make_shared<CondVar>(UnVisited);
+  // Create a sequence variable.
   std::shared_ptr<Var> Xs = std::make_shared<SeqVar>();
   return VectorRef({V, Xs});
 }
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/reduce_sum_optimizer_old.cc

@@ -0,0 +1,273 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "backend/common/pass/reduce_sum_optimizer.h"
+#include <vector>
+#include "include/common/utils/anfalgo.h"
+
+namespace mindspore {
+namespace opt {
+namespace {
+const int axis_input_index = 2;
+
+
+/**
+ * @brief Check if rank computation is needed for a given CNode.
+ *
+ * This function examines the axis input of the CNode. If the axis input is a tuple that
+ * is either empty or contains negative values, the rank computation is deemed necessary.
+ * 
+ * @param cnode The CNode to be checked.
+ * @return true if rank computation is needed, false otherwise.
+ */
+bool IsNeedComputeRank(const CNodePtr &cnode) {
+  // Ensure the CNode is not null.
+  MS_EXCEPTION_IF_NULL(cnode);
+
+  // Retrieve the axis input from the CNode.
+  auto axis_input = cnode->input(axis_input_index);
+  MS_EXCEPTION_IF_NULL(axis_input);
+
+  // If the axis input is not a tuple, rank computation is not needed.
+  if (!IsValueNode<ValueTuple>(axis_input)) {
+    return false;
+  }
+
+  // Cast the axis input to a ValueNode and retrieve its value.
+  auto value_node = axis_input->cast<ValueNodePtr>();
+  MS_EXCEPTION_IF_NULL(value_node);
+  auto value = value_node->value();
+  MS_EXCEPTION_IF_NULL(value);
+
+  if (value->isa<ValueTuple>()) {
+    auto value_tuple = value->cast<ValueTuplePtr>();
+    MS_EXCEPTION_IF_NULL(value_tuple);
+    // If the tuple is empty, rank computation is needed.
+    if (value_tuple->value().empty()) {
+      return true;
+    }
+    // Check each item in the tuple. If any item is negative, rank computation is needed.
+    for (auto &iter : value_tuple->value()) {
+      auto item = GetValue<int64_t>(iter->cast<ScalarPtr>());
+      if (item < 0) {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
+/**
+ * @brief Create a new Rank operation node.
+ *
+ * The Rank operation computes the rank (number of dimensions) of a tensor.
+ * This function creates a new CNode that represents a Rank operation in the MindSpore graph.
+ * 
+ * @param cnode The original CNode for which the Rank operation is being created.
+ * @param kernel_graph The current kernel graph.
+ * @return Pointer to the newly created Rank operation node.
+ */
+AnfNodePtr ReduceSumOptimizer::NewRankOp(const AnfNodePtr &cnode, const KernelGraphPtr &kernel_graph) const {
+  // Initialize inputs for the Rank operation.
+  std::vector<AnfNodePtr> rank_inputs;
+  auto prim = std::make_shared<Primitive>(prim::kPrimRank->name());
+  rank_inputs.push_back(NewValueNode(prim));
+  
+  // Get the previous node's output.
+  auto prev_node = common::AnfAlgo::GetPrevNodeOutput(cnode, 1);
+  rank_inputs.push_back(prev_node.first);
+  
+  // Create the Rank operation node.
+  auto rank_op = NewCNode(rank_inputs, kernel_graph);
+  MS_EXCEPTION_IF_NULL(rank_op);
+  rank_op->set_abstract(prev_node.first->abstract());
+
+  return rank_op;
+}
+
+/**
+ * @brief Create a new Range operation node.
+ *
+ * The Range operation generates a sequence of numbers. In this context, it is used to
+ * generate a sequence based on the rank of a tensor.
+ * This function creates a new CNode that represents a Range operation in the MindSpore graph.
+ * 
+ * @param rank_op The Rank operation node that provides the limit for the Range operation.
+ * @param kernel_graph The current kernel graph.
+ * @return Pointer to the newly created Range operation node.
+ */
+AnfNodePtr ReduceSumOptimizer::NewRangeOp(const AnfNodePtr &rank_op, const KernelGraphPtr &kernel_graph) const {
+  // Initialize inputs for the Range operation.
+  std::vector<AnfNodePtr> range_inputs;
+  auto prim = std::make_shared<Primitive>(prim::kPrimRange->name());
+  range_inputs.push_back(NewValueNode(prim));
+
+  // Set the "start" of the range.
+  auto start_ = NewValueNode(SizeToLong(0));
+  MS_EXCEPTION_IF_NULL(start_);
+  auto imm_start = std::make_shared<Int64Imm>(SizeToLong(0));
+  start_->set_abstract(std::make_shared<abstract::AbstractScalar>(imm_start));
+  range_inputs.push_back(start_);
+
+  // Set the "limit" of the range.
+  range_inputs.push_back(rank_op);
+
+  // Set the "delta" of the range.
+  auto delta_ = NewValueNode(SizeToLong(1));
+  MS_EXCEPTION_IF_NULL(delta_);
+  auto imm_delta = std::make_shared<Int64Imm>(SizeToLong(1));
+  delta_->set_abstract(std::make_shared<abstract::AbstractScalar>(imm_delta));
+  range_inputs.push_back(delta_);
+
+  // Create the Range operation node.
+  auto range_op = NewCNode(range_inputs, kernel_graph);
+  MS_EXCEPTION_IF_NULL(range_op);
+  range_op->set_abstract(rank_op->abstract());
+
+  return range_op;
+}
+
+AnfNodePtr ReduceSumOptimizer::InsertAssistNode(const CNodePtr &cnode, const KernelGraphPtr &kernel_graph) const {
+  // the input dim is unknown, need rank + range, don't supported now;
+  MS_LOG(EXCEPTION)
+    << "Can not support the case that input is dim unknown and axis is empty or axis contain value less 0. node: "
+    << trace::DumpSourceLines(cnode);
+}
+
+// create a new assist value node to deal with the following two case
+// 1: the axis_input is empty, the new tensor of the new value node should be 'range(shape.size())',
+// the shape is the first input'shape of ReduceSum;
+// 2: the value of axis_input contain the value less 0,
+// the new tensor of the new value node should be "shape.size() + the_old_value_less_0",
+// the shape is the first input'shape of ReduceSum;
+AnfNodePtr ReduceSumOptimizer::NewAssistValueNode(const CNodePtr &cnode, const KernelGraphPtr &kernel_graph) const {
+  // axis is a tuple ,maybe empty or contain a value less 0;
+  auto axis_input = cnode->input(axis_input_index);
+  if (IsValueNode<ValueTuple>(axis_input)) {
+    std::vector<AnfNodePtr> new_inputs = {common::AnfAlgo::GetCNodePrimitiveNode(cnode)};
+    new_inputs.push_back(cnode->input(1));
+    auto value_node = axis_input->cast<ValueNodePtr>();
+    MS_EXCEPTION_IF_NULL(value_node);
+    auto value = value_node->value();
+    MS_EXCEPTION_IF_NULL(value);
+    if (value->isa<ValueTuple>()) {
+      auto value_tuple = value->cast<ValueTuplePtr>();
+      MS_EXCEPTION_IF_NULL(value_tuple);
+      auto x_shape = dyn_cast<abstract::Shape>(cnode->input(1)->Shape());
+      MS_EXCEPTION_IF_NULL(x_shape);
+      std::vector<int64_t> axes_value;
+      ValuePtr valuePtr = nullptr;
+      if (value_tuple->value().empty()) {
+        // case 1: tensor is empty;
+        for (size_t i = 0; i < x_shape->shape().size(); i++) {
+          axes_value.emplace_back(SizeToLong(i));
+        }
+      } else {
+        // case 2: contain value less 0;
+        for (auto &iter : value_tuple->value()) {
+          auto item = GetValue<int64_t>(iter->cast<ScalarPtr>());
+          if (item < 0) {
+            (void)axes_value.emplace_back(item + static_cast<int64_t>(x_shape->shape().size()));
+          } else {
+            (void)axes_value.emplace_back(item);
+          }
+        }
+      }
+      valuePtr = MakeValue<std::vector<int64_t>>(axes_value);
+      auto assist_node = NewValueNode(valuePtr);
+      assist_node->set_abstract(std::make_shared<abstract::AbstractTensor>(kInt64, axes_value));
+      auto assist_value_node = kernel_graph->NewValueNode(assist_node);
+      new_inputs.push_back(assist_value_node);
+      auto new_node = NewCNode(cnode, kernel_graph);
+      MS_EXCEPTION_IF_NULL(new_node);
+      new_node->set_inputs(new_inputs);
+      return new_node;
+    }
+  }
+  return nullptr;
+}
+
+/**
+ * @brief Processes the given node for optimization in the context of a ReduceSum operation.
+ * 
+ * This function optimizes ReduceSum operations under certain conditions. If the current node is not a ReduceSum
+ * operation or if it does not have a dynamic shape, it will be skipped. If the node has unknown dimensions
+ * and needs rank computation, an assist node will be inserted. Otherwise, a new assist value node will be created.
+ *
+ * @param func_graph The current function graph.
+ * @param node The node to be processed.
+ * @param EquivPtr Placeholder for equivalence. Not used in this context.
+ * @return An optimized AnfNodePtr if applicable, or nullptr if no optimization was applied.
+ */
+const AnfNodePtr ReduceSumOptimizer::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
+                                             const EquivPtr &) const {
+  // Ensure func_graph and node are not null.
+  MS_EXCEPTION_IF_NULL(func_graph);
+  MS_EXCEPTION_IF_NULL(node);
+  
+  // Cast the node to a CNode.
+  auto cnode = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(cnode);
+
+  // Get the name of the CNode operation.
+  auto op_name = common::AnfAlgo::GetCNodeName(cnode);
+  
+  // If the current node is not a ReduceSum operation, skip processing.
+  if (op_name != kReduceSumOpName) {
+    MS_LOG(DEBUG) << "Current node is not: " << kReduceSumOpName << ", skip!";
+    return nullptr;
+  }
+  
+  // If the current node does not have a dynamic shape, skip processing.
+  if (!common::AnfAlgo::IsDynamicShape(cnode)) {
+    MS_LOG(DEBUG) << "Current node is not dynamic shape, skip!";
+    return nullptr;
+  }
+
+  // Mark the node as visited.
+  common::AnfAlgo::SetNodeAttr(kAttrVisited, MakeValue(true), node);
+
+  // Cast the func_graph to a KernelGraph.
+  auto kernel_graph = func_graph->cast<std::shared_ptr<session::KernelGraph>>();
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+
+  // If the CNode has unknown dimensions and requires rank computation, insert an assist node.
+  if (AnfUtils::IsDimUnknown(cnode) && IsNeedComputeRank(cnode)) {
+    return InsertAssistNode(cnode, kernel_graph);
+  }
+  
+  // Create a new assist value node.
+  return NewAssistValueNode(cnode, kernel_graph);
+}
+
+/**
+ * @brief Define a pattern for the ReduceSumOptimizer.
+ * 
+ * This function defines a pattern with a condition variable followed by a sequence variable. The condition variable
+ * ensures that the pattern is not visited more than once.
+ * 
+ * @return A VectorRef representing the defined pattern.
+ */
+const BaseRef ReduceSumOptimizer::DefinePattern() const {
+  // Create a condition variable to ensure the pattern is not visited more than once.
+  std::shared_ptr<Var> V = std::make_shared<CondVar>(UnVisited);
+  // Create a sequence variable.
+  std::shared_ptr<Var> Xs = std::make_shared<SeqVar>();
+  return VectorRef({V, Xs});
+}
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/replace_node_by_proxy.cc

@@ -23,74 +23,110 @@
 
 namespace mindspore {
 namespace opt {
+/**
+ * \brief Generates the kernel build information for a given CNode.
+ * 
+ * The function extracts the required information about input and output formats, device types, etc., 
+ * from the given CNode and generates the corresponding kernel build information.
+ * 
+ * \param cnode The CNode for which the kernel build information is to be generated.
+ * \return Returns a pointer to the KernelBuildInfo object.
+ */
 kernel::KernelBuildInfoPtr ReplaceNodeByProxy::GenerateKernelBuildInfo(const CNodePtr &cnode) {
   MS_EXCEPTION_IF_NULL(cnode);
+  
+  // Lists to store device formats and types of inputs and outputs.
   std::vector<std::string> inputs_device_format;
   std::vector<std::string> outputs_device_format;
   std::vector<TypeId> inputs_device_type;
   std::vector<TypeId> outputs_device_type;
   std::vector<std::vector<size_t>> outputs_shape;
   kernel::KernelBuildInfo::KernelBuildInfoBuilder builder;
+
+  // Populate the input information lists.
   size_t input_num = common::AnfAlgo::GetInputTensorNum(cnode);
   for (size_t input_index = 0; input_index < input_num; ++input_index) {
     inputs_device_format.push_back(AnfAlgo::GetInputFormat(cnode, input_index));
     inputs_device_type.push_back(AnfAlgo::GetInputDeviceDataType(cnode, input_index));
   }
+  
+  // Populate the output information lists.
   size_t output_num = common::AnfAlgo::GetOutputTensorNum(cnode);
   for (size_t output_index = 0; output_index < output_num; ++output_index) {
     outputs_device_format.push_back(AnfAlgo::GetOutputFormat(cnode, output_index));
     outputs_device_type.push_back(AnfAlgo::GetOutputDeviceDataType(cnode, output_index));
     outputs_shape.push_back(common::AnfAlgo::GetOutputInferShape(cnode, output_index));
   }
+
+  // Set attributes for the builder object.
   builder.SetFusionType(AnfAlgo::GetFusionType(cnode));
   builder.SetProcessor(AnfAlgo::GetProcessor(cnode));
   builder.SetKernelType(AnfAlgo::GetKernelType(cnode));
-
   builder.SetInputsFormat(inputs_device_format);
   builder.SetOutputsFormat(outputs_device_format);
   builder.SetInputsDeviceType(inputs_device_type);
   builder.SetOutputsDeviceType(outputs_device_type);
+  
   return builder.Build();
 }
 
+/**
+ * \brief Replaces specific nodes in the computational graph with their proxy versions.
+ * 
+ * The function identifies nodes in the graph representing the 'kEmbeddingLookupOpName' operation and 
+ * replaces them with proxy versions ('kEmbeddingLookupProxyOpName') while preserving their attributes 
+ * and kernel build information.
+ * 
+ * \param func_graph Pointer to the computational graph.
+ * \return Returns true if the process completes successfully.
+ */
 bool ReplaceNodeByProxy::Run(const FuncGraphPtr &func_graph) {
   MS_EXCEPTION_IF_NULL(func_graph);
+  
+  // Get the manager for the graph and its nodes.
   auto manager = func_graph->manager();
   MS_EXCEPTION_IF_NULL(manager);
+  
+  // Obtain a topological sorted list of nodes in the graph.
   std::vector<AnfNodePtr> node_list = TopoSort(func_graph->get_return());
+
+  // Iterate through nodes and identify ones representing the 'kEmbeddingLookupOpName' operation.
   for (auto node : node_list) {
     if (node != nullptr && node->isa<CNode>() && common::AnfAlgo::GetCNodeName(node) == kEmbeddingLookupOpName) {
       TraceGuard guard(std::make_shared<TraceOpt>(node->debug_info()));
+
       auto cnode = node->cast<CNodePtr>();
       MS_EXCEPTION_IF_NULL(cnode);
+      
+      // Create a new proxy node with the same inputs as the original node.
       auto prim = std::make_shared<Primitive>(kEmbeddingLookupProxyOpName);
-      MS_EXCEPTION_IF_NULL(prim);
       std::vector<AnfNodePtr> proxy_inputs = {NewValueNode(prim)};
       proxy_inputs.insert(proxy_inputs.end(), cnode->inputs().begin() + 1, cnode->inputs().end());
       AnfNodePtr proxy_node = func_graph->NewCNode(proxy_inputs);
-      MS_EXCEPTION_IF_NULL(proxy_node);
-
+      
+      // Set kernel information for the proxy node.
       auto kernel_info = std::make_shared<device::KernelInfo>();
-      MS_EXCEPTION_IF_NULL(kernel_info);
       proxy_node->set_kernel_info(kernel_info);
-
-      AbstractBasePtrList abstract_list;
       common::AnfAlgo::CopyNodeAttr(kAttrPsKey, cnode, proxy_node);
       common::AnfAlgo::CopyNodeAttr("offset", cnode, proxy_node);
+
+      AbstractBasePtrList abstract_list;
       abstract_list.push_back(cnode->abstract());
       auto abstract_tuple = std::make_shared<abstract::AbstractTuple>(abstract_list);
-      MS_EXCEPTION_IF_NULL(abstract_tuple);
       proxy_node->set_abstract(abstract_tuple);
 
       auto kernel_build_info = GenerateKernelBuildInfo(cnode);
       AnfAlgo::SetSelectKernelBuildInfo(kernel_build_info, proxy_node.get());
-
+      
+      // Replace the original node with the proxy node in the graph.
       if (!manager->Replace(cnode, proxy_node)) {
         MS_LOG(EXCEPTION) << "Replace node by proxy node failed.";
       }
     }
   }
+  
   return true;
 }
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/replace_node_by_proxy_old.cc

@@ -0,0 +1,132 @@
+/**
+ * Copyright 2020 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "backend/common/pass/replace_node_by_proxy.h"
+#include <vector>
+#include <memory>
+#include "runtime/device/kernel_info.h"
+#include "backend/common/session/anf_runtime_algorithm.h"
+#include "include/common/utils/anfalgo.h"
+#include "kernel/kernel_build_info.h"
+
+namespace mindspore {
+namespace opt {
+/**
+ * \brief Generates the kernel build information for a given CNode.
+ * 
+ * The function extracts the required information about input and output formats, device types, etc., 
+ * from the given CNode and generates the corresponding kernel build information.
+ * 
+ * \param cnode The CNode for which the kernel build information is to be generated.
+ * \return Returns a pointer to the KernelBuildInfo object.
+ */
+kernel::KernelBuildInfoPtr ReplaceNodeByProxy::GenerateKernelBuildInfo(const CNodePtr &cnode) {
+  MS_EXCEPTION_IF_NULL(cnode);
+  
+  // Lists to store device formats and types of inputs and outputs.
+  std::vector<std::string> inputs_device_format;
+  std::vector<std::string> outputs_device_format;
+  std::vector<TypeId> inputs_device_type;
+  std::vector<TypeId> outputs_device_type;
+  std::vector<std::vector<size_t>> outputs_shape;
+  kernel::KernelBuildInfo::KernelBuildInfoBuilder builder;
+
+  // Populate the input information lists.
+  size_t input_num = common::AnfAlgo::GetInputTensorNum(cnode);
+  for (size_t input_index = 0; input_index < input_num; ++input_index) {
+    inputs_device_format.push_back(AnfAlgo::GetInputFormat(cnode, input_index));
+    inputs_device_type.push_back(AnfAlgo::GetInputDeviceDataType(cnode, input_index));
+  }
+  
+  // Populate the output information lists.
+  size_t output_num = common::AnfAlgo::GetOutputTensorNum(cnode);
+  for (size_t output_index = 0; output_index < output_num; ++output_index) {
+    outputs_device_format.push_back(AnfAlgo::GetOutputFormat(cnode, output_index));
+    outputs_device_type.push_back(AnfAlgo::GetOutputDeviceDataType(cnode, output_index));
+    outputs_shape.push_back(common::AnfAlgo::GetOutputInferShape(cnode, output_index));
+  }
+
+  // Set attributes for the builder object.
+  builder.SetFusionType(AnfAlgo::GetFusionType(cnode));
+  builder.SetProcessor(AnfAlgo::GetProcessor(cnode));
+  builder.SetKernelType(AnfAlgo::GetKernelType(cnode));
+  builder.SetInputsFormat(inputs_device_format);
+  builder.SetOutputsFormat(outputs_device_format);
+  builder.SetInputsDeviceType(inputs_device_type);
+  builder.SetOutputsDeviceType(outputs_device_type);
+  
+  return builder.Build();
+}
+
+/**
+ * \brief Replaces specific nodes in the computational graph with their proxy versions.
+ * 
+ * The function identifies nodes in the graph representing the 'kEmbeddingLookupOpName' operation and 
+ * replaces them with proxy versions ('kEmbeddingLookupProxyOpName') while preserving their attributes 
+ * and kernel build information.
+ * 
+ * \param func_graph Pointer to the computational graph.
+ * \return Returns true if the process completes successfully.
+ */
+bool ReplaceNodeByProxy::Run(const FuncGraphPtr &func_graph) {
+  MS_EXCEPTION_IF_NULL(func_graph);
+  
+  // Get the manager for the graph and its nodes.
+  auto manager = func_graph->manager();
+  MS_EXCEPTION_IF_NULL(manager);
+  
+  // Obtain a topological sorted list of nodes in the graph.
+  std::vector<AnfNodePtr> node_list = TopoSort(func_graph->get_return());
+
+  // Iterate through nodes and identify ones representing the 'kEmbeddingLookupOpName' operation.
+  for (auto node : node_list) {
+    if (node != nullptr && node->isa<CNode>() && common::AnfAlgo::GetCNodeName(node) == kEmbeddingLookupOpName) {
+      TraceGuard guard(std::make_shared<TraceOpt>(node->debug_info()));
+
+      auto cnode = node->cast<CNodePtr>();
+      MS_EXCEPTION_IF_NULL(cnode);
+      
+      // Create a new proxy node with the same inputs as the original node.
+      auto prim = std::make_shared<Primitive>(kEmbeddingLookupProxyOpName);
+      std::vector<AnfNodePtr> proxy_inputs = {NewValueNode(prim)};
+      proxy_inputs.insert(proxy_inputs.end(), cnode->inputs().begin() + 1, cnode->inputs().end());
+      AnfNodePtr proxy_node = func_graph->NewCNode(proxy_inputs);
+      
+      // Set kernel information for the proxy node.
+      auto kernel_info = std::make_shared<device::KernelInfo>();
+      proxy_node->set_kernel_info(kernel_info);
+      common::AnfAlgo::CopyNodeAttr(kAttrPsKey, cnode, proxy_node);
+      common::AnfAlgo::CopyNodeAttr("offset", cnode, proxy_node);
+
+      AbstractBasePtrList abstract_list;
+      abstract_list.push_back(cnode->abstract());
+      auto abstract_tuple = std::make_shared<abstract::AbstractTuple>(abstract_list);
+      proxy_node->set_abstract(abstract_tuple);
+
+      auto kernel_build_info = GenerateKernelBuildInfo(cnode);
+      AnfAlgo::SetSelectKernelBuildInfo(kernel_build_info, proxy_node.get());
+      
+      // Replace the original node with the proxy node in the graph.
+      if (!manager->Replace(cnode, proxy_node)) {
+        MS_LOG(EXCEPTION) << "Replace node by proxy node failed.";
+      }
+    }
+  }
+  
+  return true;
+}
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/pass/sparse_process.cc

@@ -31,83 +31,159 @@ using CSRTensorPtr = mindspore::tensor::CSRTensorPtr;
 constexpr auto kCSRValueNodeNum = 2;
 constexpr auto kSparseAttrIndex = 1;
 
-// Convert CSRTensor Parameter or ValueNode to Tuple by setting its abstract.
+/**
+ * \brief Convert a CSRTensor Parameter or ValueNode to a Tuple by setting its abstract.
+ * 
+ * This function transforms a given CSRTensor's abstract to an AbstractTuple.
+ * 
+ * \param sparse AnfNode pointer representing the CSRTensor.
+ */
 void AbstractCSRToAbstractTuple(const AnfNodePtr &sparse) {
+  // Check for null pointer
   MS_EXCEPTION_IF_NULL(sparse);
+
+  // Validate if 'sparse' is either a Parameter or a ValueNode
   if (!(sparse->isa<Parameter>() || sparse->isa<ValueNode>())) {
     return;
   }
+
   auto param_abs = sparse->abstract();
+  // Check for null abstract
   MS_EXCEPTION_IF_NULL(param_abs);
+
+  // Check if the abstract is a CSRTensor type
   if (param_abs->isa<abstract::AbstractCSRTensor>()) {
     auto abs_sparse = param_abs->cast<abstract::AbstractCSRTensorPtr>();
+
+    // Construct a list from the CSRTensor abstract components
     std::vector<AbstractBasePtr> abstract_list{abs_sparse->indptr(), abs_sparse->indices(), abs_sparse->values(),
                                                abs_sparse->dense_shape()};
+
+    // Convert the abstract list to an AbstractTuple
     auto abs_tuple = std::make_shared<abstract::AbstractTuple>(abstract_list);
     abs_tuple->set_type(abs_tuple->BuildType());
     sparse->set_abstract(abs_tuple);
   }
 }
 
+/**
+ * \brief Create and add a new ValueNode to the graph.
+ * 
+ * \param val ValueNode pointer to be added.
+ * \param abs Abstract associated with the ValueNode.
+ * \param kernel_graph KernelGraph pointer to which the node will be added.
+ * 
+ * \return Returns the newly created ValueNode pointer.
+ */
 ValueNodePtr MakeNewValueNodeToGraph(const ValueNodePtr &val, const AbstractBasePtr &abs,
                                      const KernelGraphPtr &kernel_graph) {
+  // Ensure kernel graph is not null
   MS_EXCEPTION_IF_NULL(kernel_graph);
+
+  // Create a new ValueNode
   auto node = kernel_graph->NewValueNode(val);
   MS_EXCEPTION_IF_NULL(node);
   node->set_abstract(abs);
+
+  // Add the ValueNode to the graph
   kernel_graph->AddValueNodeToGraph(node);
   return node;
 }
 
+/**
+ * \brief Split a ValueNode and add its components to the input list.
+ * 
+ * This function extracts components from a CSRTensor ValueNode and adds them to a list.
+ * 
+ * \param node AnfNode pointer to be split.
+ * \param new_inputs Vector to store the new AnfNode pointers.
+ * \param kernel_graph KernelGraph pointer used for node creation.
+ * 
+ * \return Returns true if the ValueNode was split, otherwise returns false.
+ */
 bool SplitValueNode(const AnfNodePtr &node, std::vector<AnfNodePtr> *new_inputs, const KernelGraphPtr &kernel_graph) {
   ValuePtr value = node->cast<ValueNodePtr>()->value();
+
+  // Ensure 'value' is not null
   MS_EXCEPTION_IF_NULL(value);
+
+  // Check if the 'value' is a CSRTensor
   if (!value->isa<CSRTensor>()) return false;
+
   auto csr_tensor = value->cast<CSRTensorPtr>();
   MS_EXCEPTION_IF_NULL(csr_tensor);
   auto csr_abs = node->abstract()->cast<abstract::AbstractCSRTensorPtr>();
   MS_EXCEPTION_IF_NULL(csr_abs);
+
+  // Split the CSRTensor into its components and add to 'new_inputs'
   auto new_indptr = MakeNewValueNodeToGraph(NewValueNode(csr_tensor->GetIndptr()), csr_abs->indptr(), kernel_graph);
   new_inputs->push_back(new_indptr);
+
   auto new_indices = MakeNewValueNodeToGraph(NewValueNode(csr_tensor->GetIndices()), csr_abs->indices(), kernel_graph);
   new_inputs->push_back(new_indices);
+
   auto new_values = MakeNewValueNodeToGraph(NewValueNode(csr_tensor->GetValues()), csr_abs->values(), kernel_graph);
   new_inputs->push_back(new_values);
+
   return true;
 }
 
+
+/**
+ * \brief Split a Parameter node representing CSRTensor into its components and add to the input list.
+ * 
+ * If the node represents a CSRTensor, this function splits its components (indptr, indices, values) and 
+ * updates their abstract information. It also maintains a map (`csr_params_map`) to cache previously processed nodes.
+ * 
+ * \param node AnfNode pointer representing the Parameter to be split.
+ * \param new_inputs Vector to store the newly split AnfNode pointers.
+ * \param kernel_graph KernelGraph pointer used for node creation.
+ * 
+ * \return Returns true if the Parameter node was split, otherwise returns false.
+ */
 bool SplitParameter(const AnfNodePtr &node, std::vector<AnfNodePtr> *new_inputs, const KernelGraphPtr &kernel_graph) {
+  // Validate the provided node
   MS_EXCEPTION_IF_NULL(node);
+  
   auto node_abs = node->abstract();
   MS_EXCEPTION_IF_NULL(node_abs);
+  
+  // Map to store the nodes and their split components
   static HashMap<AnfNodePtr, std::vector<AnfNodePtr>> csr_params_map;
   auto param = node->cast<ParameterPtr>();
   MS_EXCEPTION_IF_NULL(param);
+
+  // Process the node if its abstract type is CSRTensor
   if (node_abs->isa<abstract::AbstractCSRTensor>()) {
     auto param_abs = node_abs->cast<abstract::AbstractCSRTensorPtr>();
+    // Validate the components of CSRTensor
     MS_EXCEPTION_IF_NULL(param_abs);
     MS_EXCEPTION_IF_NULL(param_abs->indptr());
     MS_EXCEPTION_IF_NULL(param_abs->indices());
     MS_EXCEPTION_IF_NULL(param_abs->values());
-    auto new_indptr =
-      MakeNewValueNodeToGraph(NewValueNode(param_abs->indptr()->BuildValue()), param_abs->indptr(), kernel_graph);
+
+    // Split the components and add to new_inputs
+    auto new_indptr = MakeNewValueNodeToGraph(NewValueNode(param_abs->indptr()->BuildValue()), param_abs->indptr(), kernel_graph);
     MS_EXCEPTION_IF_NULL(new_indptr);
     new_inputs->push_back(new_indptr);
-    auto new_indices =
-      MakeNewValueNodeToGraph(NewValueNode(param_abs->indices()->BuildValue()), param_abs->indices(), kernel_graph);
+
+    auto new_indices = MakeNewValueNodeToGraph(NewValueNode(param_abs->indices()->BuildValue()), param_abs->indices(), kernel_graph);
     MS_EXCEPTION_IF_NULL(new_indices);
     new_inputs->push_back(new_indices);
-    // Set CSRTensor Parameter abstract to Tensor by its values.
+
+    // Update the node's abstract to Tensor using its values
     node->set_abstract(param_abs->values()->Broaden());
     new_inputs->push_back(node);
-    // Set csr_params_map
+
+    // Cache the processed node in csr_params_map
     if (csr_params_map.find(node) == csr_params_map.end()) {
       csr_params_map[node].emplace_back(new_indptr);
       csr_params_map[node].emplace_back(new_indices);
     }
     return true;
-    // If the cnode has a csr_tensor_param which has been split, use the map to find its indptr and indices.
-  } else if (node_abs->isa<abstract::AbstractTensor>() && csr_params_map.find(node) != csr_params_map.end()) {
+  } 
+  // If node was previously processed and cached, retrieve its components
+  else if (node_abs->isa<abstract::AbstractTensor>() && csr_params_map.find(node) != csr_params_map.end()) {
     if (csr_params_map[node].size() != kCSRValueNodeNum) {
       MS_LOG(ERROR) << "csr_params_map[" << node->DebugString() << "] has " << csr_params_map[node].size()
                     << " inputs, but expect two inputs! They are all added in new_inputs.";
@@ -119,27 +195,46 @@ bool SplitParameter(const AnfNodePtr &node, std::vector<AnfNodePtr> *new_inputs,
   return false;
 }
 
+/**
+ * \brief Split a CNode representing MakeCSR/MakeCOO or MakeTuple into its components and add to the input list.
+ * 
+ * \param node AnfNode pointer representing the CNode to be split.
+ * \param new_inputs Vector to store the CNode components.
+ * 
+ * \return Returns true if the CNode was split, otherwise returns false.
+ */
 bool SplitCNode(const AnfNodePtr &node, std::vector<AnfNodePtr> *new_inputs) {
   auto cnode = node->cast<CNodePtr>();
   auto sparse_prim = common::AnfAlgo::GetCNodePrimitive(cnode);
   MS_EXCEPTION_IF_NULL(sparse_prim);
-  // Currently, only MakeCSR/MakeCOO and MakeTuple nodes can be split.
+
+  // Only process specific CNodes: MakeCSR/MakeCOO and MakeTuple
   if (make_sparse_set.count(sparse_prim->name()) <= 0 && sparse_prim->name().compare(prim::kPrimMakeTuple->name()) != 0)
     return false;
 
   auto sparse_inputs = cnode->inputs();
-  // skip the last input, as it always represents shape, and has already been
-  // registered as primitive attribute.
+  // Exclude the last input, since it represents shape and is registered as primitive attribute
   for (size_t j = 1; j < sparse_inputs.size() - 1; ++j) {
     new_inputs->push_back(sparse_inputs[j]);
   }
   return true;
 }
 
+/**
+ * \brief Retrieve abstract list from a node based on the provided primitive type.
+ * 
+ * This function processes nodes with CSRTensor and COOTensor abstracts and retrieves their components as a list.
+ * 
+ * \param node AnfNode pointer to be processed.
+ * \param prim Primitive pointer indicating the type of node (MakeCSRTensor or MakeCOOTensor).
+ * 
+ * \return Returns a vector of abstracts for the provided node and primitive.
+ */
 std::vector<AbstractBasePtr> GetAbstractList(const AnfNodePtr &node, const PrimitivePtr &prim) {
   MS_EXCEPTION_IF_NULL(node);
   MS_EXCEPTION_IF_NULL(prim);
   std::string prim_name = prim->name();
+
   if (prim_name == prim::kPrimMakeCSRTensor->name()) {
     auto abs_sparse = dyn_cast<abstract::AbstractCSRTensor>(node->abstract());
     MS_EXCEPTION_IF_NULL(abs_sparse);
@@ -152,9 +247,21 @@ std::vector<AbstractBasePtr> GetAbstractList(const AnfNodePtr &node, const Primi
   return {};
 }
 
+/**
+ * \brief Converts a node representing a MakeSparse operation into a MakeTuple operation.
+ * 
+ * This function modifies the given CNode representing a MakeSparse operation and changes it to a MakeTuple operation.
+ * 
+ * \param node AnfNode pointer representing the node to be converted.
+ * \param kernel_graph KernelGraph pointer used for node manipulation.
+ * 
+ * \return Returns a new CNodePtr representing the converted MakeTuple operation.
+ */
 CNodePtr ConvertMakeSparseToMakeTuple(const AnfNodePtr &node, const KernelGraphPtr &kernel_graph) {
+  // Validation
   MS_EXCEPTION_IF_NULL(node);
   MS_EXCEPTION_IF_NULL(kernel_graph);
+
   CNodePtr cnode = node->cast<CNodePtr>();
   MS_EXCEPTION_IF_NULL(cnode);
 
@@ -162,6 +269,7 @@ CNodePtr ConvertMakeSparseToMakeTuple(const AnfNodePtr &node, const KernelGraphP
   inputs.emplace_back(NewValueNode(prim::kPrimMakeTuple));
   (void)inputs.insert(inputs.end(), cnode->inputs().begin() + 1, cnode->inputs().end());
 
+  // Create new node with MakeTuple operation and set its attributes
   auto new_node = NewCNode(inputs, cnode->func_graph());
   std::vector<AbstractBasePtr> abstract_list = GetAbstractList(node, common::AnfAlgo::GetCNodePrimitive(cnode));
   auto abs_res = std::make_shared<abstract::AbstractTuple>(abstract_list);
@@ -173,9 +281,22 @@ CNodePtr ConvertMakeSparseToMakeTuple(const AnfNodePtr &node, const KernelGraphP
   return new_node;
 }
 
+/**
+ * \brief Converts a node representing a SparseGetAttr operation into a TupleGetItem operation.
+ * 
+ * This function modifies the CNode representing a SparseGetAttr operation and changes it to a TupleGetItem operation.
+ * 
+ * \param index The index of the item in the tuple to retrieve.
+ * \param node AnfNode pointer representing the node to be converted.
+ * \param kernel_graph KernelGraph pointer used for node creation.
+ * 
+ * \return Returns a new CNodePtr representing the converted TupleGetItem operation.
+ */
 CNodePtr ConvertSparseGetAttrToTupleGetItem(int64_t index, const AnfNodePtr &node, const KernelGraphPtr &kernel_graph) {
+  // Validation
   MS_EXCEPTION_IF_NULL(node);
   MS_EXCEPTION_IF_NULL(kernel_graph);
+
   CNodePtr cnode = node->cast<CNodePtr>();
   MS_EXCEPTION_IF_NULL(cnode);
 
@@ -187,6 +308,8 @@ CNodePtr ConvertSparseGetAttrToTupleGetItem(int64_t index, const AnfNodePtr &nod
   auto index_node = NewValueNode(index);
   AbstractBasePtr index_abs = std::make_shared<abstract::AbstractScalar>(std::make_shared<Int64Imm>(index));
   index_node->set_abstract(index_abs);
+
+  // Create new node with TupleGetItem operation and set its attributes
   auto new_node =
     NewCNode({NewValueNode(prim::kPrimTupleGetItem), inputs[kSparseAttrIndex], index_node}, cnode->func_graph());
   new_node->set_abstract(node->abstract());
@@ -196,9 +319,21 @@ CNodePtr ConvertSparseGetAttrToTupleGetItem(int64_t index, const AnfNodePtr &nod
   return new_node;
 }
 
+/**
+ * \brief Processes the inputs of a CNode representing a sparse operation.
+ * 
+ * For nodes representing specific sparse operations, this function splits their components and adjusts their attributes.
+ * 
+ * \param node AnfNode pointer representing the node to be processed.
+ * \param kernel_graph KernelGraph pointer used for node creation and manipulation.
+ * 
+ * \return Returns a new CNodePtr with adjusted inputs and attributes.
+ */
 CNodePtr FetchInputsForSparseOP(const AnfNodePtr &node, const KernelGraphPtr &kernel_graph) {
+  // Validation
   MS_EXCEPTION_IF_NULL(node);
   MS_EXCEPTION_IF_NULL(kernel_graph);
+
   CNodePtr cnode = node->cast<CNodePtr>();
   MS_EXCEPTION_IF_NULL(cnode);
 
@@ -206,6 +341,8 @@ CNodePtr FetchInputsForSparseOP(const AnfNodePtr &node, const KernelGraphPtr &ke
     MS_LOG(INFO) << "Do not process CNode " << cnode << " (" << cnode->DebugString() << "), because it has been split.";
     return nullptr;
   }
+
+  // Process the inputs based on their type and whether they represent a sparse operation
   const auto &inputs = cnode->inputs();
   std::vector<AnfNodePtr> new_inputs;
   new_inputs.push_back(inputs[0]);
@@ -215,15 +352,14 @@ CNodePtr FetchInputsForSparseOP(const AnfNodePtr &node, const KernelGraphPtr &ke
     } else if (inputs[i]->isa<ValueNode>()) {
       if (SplitValueNode(inputs[i], &new_inputs, kernel_graph)) continue;
     } else if (inputs[i]->isa<Parameter>()) {
-      // 1. Split CSRTensor param to multiple tensors.
-      // 2. Set CSRTensor abstract to AbstractTensor that is related its values.
       if (SplitParameter(inputs[i], &new_inputs, kernel_graph)) continue;
     }
     new_inputs.push_back(inputs[i]);
   }
+
+  // Create new node with the processed inputs and set its attributes
   auto new_node = NewCNode(new_inputs, cnode->func_graph());
   new_node->set_abstract(node->abstract());
-  // Set attr "has_been_split" to prevent the node is split more than once.
   new_node->AddAttr("has_been_split", MakeValue(true));
   if (kernel_graph != nullptr) {
     kernel_graph->FrontBackendlMapUpdate(cnode, new_node);
@@ -231,18 +367,34 @@ CNodePtr FetchInputsForSparseOP(const AnfNodePtr &node, const KernelGraphPtr &ke
   return new_node;
 }
 
+/**
+ * \brief Main processing function for nodes representing sparse operations.
+ * 
+ * This function examines the type of sparse operation represented by the node and dispatches the appropriate processing
+ * function.
+ * 
+ * \param func_graph FuncGraph pointer for the current computational graph.
+ * \param node AnfNode pointer representing the node to be processed.
+ * \param EquivPtr Placeholder for future functionality, currently not used.
+ * 
+ * \return Returns a processed AnfNodePtr if the node represents a recognized sparse operation, or nullptr otherwise.
+ */
 const AnfNodePtr SparseProcess::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
                                         const EquivPtr &) const {
+  // Validation
   MS_EXCEPTION_IF_NULL(func_graph);
   MS_EXCEPTION_IF_NULL(node);
   if (node == nullptr || !node->isa<CNode>() || !AnfUtils::IsRealKernel(node)) {
     return nullptr;
   }
+
   auto cnode = node->cast<CNodePtr>();
   MS_EXCEPTION_IF_NULL(cnode);
   auto prim = common::AnfAlgo::GetCNodePrimitive(cnode);
   MS_EXCEPTION_IF_NULL(prim);
   std::string prim_name = prim->name();
+
+  // Dispatch based on the type of sparse operation represented by the node
   auto kernel_graph = func_graph->cast<std::shared_ptr<session::KernelGraph>>();
   if (make_sparse_set.find(prim_name) != make_sparse_set.end()) {
     return ConvertMakeSparseToMakeTuple(node, kernel_graph);
@@ -253,5 +405,6 @@ const AnfNodePtr SparseProcess::Process(const FuncGraphPtr &func_graph, const An
   }
   return nullptr;
 }
+
 }  // namespace opt
 }  // namespace mindspore

mindspore/ccsrc/backend/common/pass/sparse_process_old.cc

@@ -0,0 +1,410 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <vector>
+#include <memory>
+#include "backend/common/pass/sparse_process.h"
+#include "ir/anf.h"
+#include "include/common/utils/convert_utils.h"
+#include "utils/anf_utils.h"
+#include "backend/common/optimizer/helper.h"
+#include "include/common/utils/anfalgo.h"
+
+namespace mindspore {
+namespace opt {
+using CSRTensor = mindspore::tensor::CSRTensor;
+using CSRTensorPtr = mindspore::tensor::CSRTensorPtr;
+
+constexpr auto kCSRValueNodeNum = 2;
+constexpr auto kSparseAttrIndex = 1;
+
+/**
+ * \brief Convert a CSRTensor Parameter or ValueNode to a Tuple by setting its abstract.
+ * 
+ * This function transforms a given CSRTensor's abstract to an AbstractTuple.
+ * 
+ * \param sparse AnfNode pointer representing the CSRTensor.
+ */
+void AbstractCSRToAbstractTuple(const AnfNodePtr &sparse) {
+  // Check for null pointer
+  MS_EXCEPTION_IF_NULL(sparse);
+
+  // Validate if 'sparse' is either a Parameter or a ValueNode
+  if (!(sparse->isa<Parameter>() || sparse->isa<ValueNode>())) {
+    return;
+  }
+
+  auto param_abs = sparse->abstract();
+  // Check for null abstract
+  MS_EXCEPTION_IF_NULL(param_abs);
+
+  // Check if the abstract is a CSRTensor type
+  if (param_abs->isa<abstract::AbstractCSRTensor>()) {
+    auto abs_sparse = param_abs->cast<abstract::AbstractCSRTensorPtr>();
+
+    // Construct a list from the CSRTensor abstract components
+    std::vector<AbstractBasePtr> abstract_list{abs_sparse->indptr(), abs_sparse->indices(), abs_sparse->values(),
+                                               abs_sparse->dense_shape()};
+
+    // Convert the abstract list to an AbstractTuple
+    auto abs_tuple = std::make_shared<abstract::AbstractTuple>(abstract_list);
+    abs_tuple->set_type(abs_tuple->BuildType());
+    sparse->set_abstract(abs_tuple);
+  }
+}
+
+/**
+ * \brief Create and add a new ValueNode to the graph.
+ * 
+ * \param val ValueNode pointer to be added.
+ * \param abs Abstract associated with the ValueNode.
+ * \param kernel_graph KernelGraph pointer to which the node will be added.
+ * 
+ * \return Returns the newly created ValueNode pointer.
+ */
+ValueNodePtr MakeNewValueNodeToGraph(const ValueNodePtr &val, const AbstractBasePtr &abs,
+                                     const KernelGraphPtr &kernel_graph) {
+  // Ensure kernel graph is not null
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+
+  // Create a new ValueNode
+  auto node = kernel_graph->NewValueNode(val);
+  MS_EXCEPTION_IF_NULL(node);
+  node->set_abstract(abs);
+
+  // Add the ValueNode to the graph
+  kernel_graph->AddValueNodeToGraph(node);
+  return node;
+}
+
+/**
+ * \brief Split a ValueNode and add its components to the input list.
+ * 
+ * This function extracts components from a CSRTensor ValueNode and adds them to a list.
+ * 
+ * \param node AnfNode pointer to be split.
+ * \param new_inputs Vector to store the new AnfNode pointers.
+ * \param kernel_graph KernelGraph pointer used for node creation.
+ * 
+ * \return Returns true if the ValueNode was split, otherwise returns false.
+ */
+bool SplitValueNode(const AnfNodePtr &node, std::vector<AnfNodePtr> *new_inputs, const KernelGraphPtr &kernel_graph) {
+  ValuePtr value = node->cast<ValueNodePtr>()->value();
+
+  // Ensure 'value' is not null
+  MS_EXCEPTION_IF_NULL(value);
+
+  // Check if the 'value' is a CSRTensor
+  if (!value->isa<CSRTensor>()) return false;
+
+  auto csr_tensor = value->cast<CSRTensorPtr>();
+  MS_EXCEPTION_IF_NULL(csr_tensor);
+  auto csr_abs = node->abstract()->cast<abstract::AbstractCSRTensorPtr>();
+  MS_EXCEPTION_IF_NULL(csr_abs);
+
+  // Split the CSRTensor into its components and add to 'new_inputs'
+  auto new_indptr = MakeNewValueNodeToGraph(NewValueNode(csr_tensor->GetIndptr()), csr_abs->indptr(), kernel_graph);
+  new_inputs->push_back(new_indptr);
+
+  auto new_indices = MakeNewValueNodeToGraph(NewValueNode(csr_tensor->GetIndices()), csr_abs->indices(), kernel_graph);
+  new_inputs->push_back(new_indices);
+
+  auto new_values = MakeNewValueNodeToGraph(NewValueNode(csr_tensor->GetValues()), csr_abs->values(), kernel_graph);
+  new_inputs->push_back(new_values);
+
+  return true;
+}
+
+
+/**
+ * \brief Split a Parameter node representing CSRTensor into its components and add to the input list.
+ * 
+ * If the node represents a CSRTensor, this function splits its components (indptr, indices, values) and 
+ * updates their abstract information. It also maintains a map (`csr_params_map`) to cache previously processed nodes.
+ * 
+ * \param node AnfNode pointer representing the Parameter to be split.
+ * \param new_inputs Vector to store the newly split AnfNode pointers.
+ * \param kernel_graph KernelGraph pointer used for node creation.
+ * 
+ * \return Returns true if the Parameter node was split, otherwise returns false.
+ */
+bool SplitParameter(const AnfNodePtr &node, std::vector<AnfNodePtr> *new_inputs, const KernelGraphPtr &kernel_graph) {
+  // Validate the provided node
+  MS_EXCEPTION_IF_NULL(node);
+  
+  auto node_abs = node->abstract();
+  MS_EXCEPTION_IF_NULL(node_abs);
+  
+  // Map to store the nodes and their split components
+  static HashMap<AnfNodePtr, std::vector<AnfNodePtr>> csr_params_map;
+  auto param = node->cast<ParameterPtr>();
+  MS_EXCEPTION_IF_NULL(param);
+
+  // Process the node if its abstract type is CSRTensor
+  if (node_abs->isa<abstract::AbstractCSRTensor>()) {
+    auto param_abs = node_abs->cast<abstract::AbstractCSRTensorPtr>();
+    // Validate the components of CSRTensor
+    MS_EXCEPTION_IF_NULL(param_abs);
+    MS_EXCEPTION_IF_NULL(param_abs->indptr());
+    MS_EXCEPTION_IF_NULL(param_abs->indices());
+    MS_EXCEPTION_IF_NULL(param_abs->values());
+
+    // Split the components and add to new_inputs
+    auto new_indptr = MakeNewValueNodeToGraph(NewValueNode(param_abs->indptr()->BuildValue()), param_abs->indptr(), kernel_graph);
+    MS_EXCEPTION_IF_NULL(new_indptr);
+    new_inputs->push_back(new_indptr);
+
+    auto new_indices = MakeNewValueNodeToGraph(NewValueNode(param_abs->indices()->BuildValue()), param_abs->indices(), kernel_graph);
+    MS_EXCEPTION_IF_NULL(new_indices);
+    new_inputs->push_back(new_indices);
+
+    // Update the node's abstract to Tensor using its values
+    node->set_abstract(param_abs->values()->Broaden());
+    new_inputs->push_back(node);
+
+    // Cache the processed node in csr_params_map
+    if (csr_params_map.find(node) == csr_params_map.end()) {
+      csr_params_map[node].emplace_back(new_indptr);
+      csr_params_map[node].emplace_back(new_indices);
+    }
+    return true;
+  } 
+  // If node was previously processed and cached, retrieve its components
+  else if (node_abs->isa<abstract::AbstractTensor>() && csr_params_map.find(node) != csr_params_map.end()) {
+    if (csr_params_map[node].size() != kCSRValueNodeNum) {
+      MS_LOG(ERROR) << "csr_params_map[" << node->DebugString() << "] has " << csr_params_map[node].size()
+                    << " inputs, but expect two inputs! They are all added in new_inputs.";
+    }
+    new_inputs->insert(new_inputs->end(), csr_params_map[node].begin(), csr_params_map[node].end());
+    new_inputs->push_back(node);
+    return true;
+  }
+  return false;
+}
+
+/**
+ * \brief Split a CNode representing MakeCSR/MakeCOO or MakeTuple into its components and add to the input list.
+ * 
+ * \param node AnfNode pointer representing the CNode to be split.
+ * \param new_inputs Vector to store the CNode components.
+ * 
+ * \return Returns true if the CNode was split, otherwise returns false.
+ */
+bool SplitCNode(const AnfNodePtr &node, std::vector<AnfNodePtr> *new_inputs) {
+  auto cnode = node->cast<CNodePtr>();
+  auto sparse_prim = common::AnfAlgo::GetCNodePrimitive(cnode);
+  MS_EXCEPTION_IF_NULL(sparse_prim);
+
+  // Only process specific CNodes: MakeCSR/MakeCOO and MakeTuple
+  if (make_sparse_set.count(sparse_prim->name()) <= 0 && sparse_prim->name().compare(prim::kPrimMakeTuple->name()) != 0)
+    return false;
+
+  auto sparse_inputs = cnode->inputs();
+  // Exclude the last input, since it represents shape and is registered as primitive attribute
+  for (size_t j = 1; j < sparse_inputs.size() - 1; ++j) {
+    new_inputs->push_back(sparse_inputs[j]);
+  }
+  return true;
+}
+
+/**
+ * \brief Retrieve abstract list from a node based on the provided primitive type.
+ * 
+ * This function processes nodes with CSRTensor and COOTensor abstracts and retrieves their components as a list.
+ * 
+ * \param node AnfNode pointer to be processed.
+ * \param prim Primitive pointer indicating the type of node (MakeCSRTensor or MakeCOOTensor).
+ * 
+ * \return Returns a vector of abstracts for the provided node and primitive.
+ */
+std::vector<AbstractBasePtr> GetAbstractList(const AnfNodePtr &node, const PrimitivePtr &prim) {
+  MS_EXCEPTION_IF_NULL(node);
+  MS_EXCEPTION_IF_NULL(prim);
+  std::string prim_name = prim->name();
+
+  if (prim_name == prim::kPrimMakeCSRTensor->name()) {
+    auto abs_sparse = dyn_cast<abstract::AbstractCSRTensor>(node->abstract());
+    MS_EXCEPTION_IF_NULL(abs_sparse);
+    return {abs_sparse->indptr(), abs_sparse->indices(), abs_sparse->values(), abs_sparse->dense_shape()};
+  } else if (prim_name == prim::kPrimMakeCOOTensor->name()) {
+    auto abs_sparse = dyn_cast<abstract::AbstractCOOTensor>(node->abstract());
+    MS_EXCEPTION_IF_NULL(abs_sparse);
+    return {abs_sparse->indices(), abs_sparse->values(), abs_sparse->dense_shape()};
+  }
+  return {};
+}
+
+/**
+ * \brief Converts a node representing a MakeSparse operation into a MakeTuple operation.
+ * 
+ * This function modifies the given CNode representing a MakeSparse operation and changes it to a MakeTuple operation.
+ * 
+ * \param node AnfNode pointer representing the node to be converted.
+ * \param kernel_graph KernelGraph pointer used for node manipulation.
+ * 
+ * \return Returns a new CNodePtr representing the converted MakeTuple operation.
+ */
+CNodePtr ConvertMakeSparseToMakeTuple(const AnfNodePtr &node, const KernelGraphPtr &kernel_graph) {
+  // Validation
+  MS_EXCEPTION_IF_NULL(node);
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+
+  CNodePtr cnode = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(cnode);
+
+  std::vector<AnfNodePtr> inputs;
+  inputs.emplace_back(NewValueNode(prim::kPrimMakeTuple));
+  (void)inputs.insert(inputs.end(), cnode->inputs().begin() + 1, cnode->inputs().end());
+
+  // Create new node with MakeTuple operation and set its attributes
+  auto new_node = NewCNode(inputs, cnode->func_graph());
+  std::vector<AbstractBasePtr> abstract_list = GetAbstractList(node, common::AnfAlgo::GetCNodePrimitive(cnode));
+  auto abs_res = std::make_shared<abstract::AbstractTuple>(abstract_list);
+  new_node->set_abstract(abs_res);
+  new_node->set_scope(cnode->scope());
+  if (kernel_graph != nullptr) {
+    kernel_graph->FrontBackendlMapUpdate(cnode, new_node);
+  }
+  return new_node;
+}
+
+/**
+ * \brief Converts a node representing a SparseGetAttr operation into a TupleGetItem operation.
+ * 
+ * This function modifies the CNode representing a SparseGetAttr operation and changes it to a TupleGetItem operation.
+ * 
+ * \param index The index of the item in the tuple to retrieve.
+ * \param node AnfNode pointer representing the node to be converted.
+ * \param kernel_graph KernelGraph pointer used for node creation.
+ * 
+ * \return Returns a new CNodePtr representing the converted TupleGetItem operation.
+ */
+CNodePtr ConvertSparseGetAttrToTupleGetItem(int64_t index, const AnfNodePtr &node, const KernelGraphPtr &kernel_graph) {
+  // Validation
+  MS_EXCEPTION_IF_NULL(node);
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+
+  CNodePtr cnode = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(cnode);
+
+  const auto &inputs = cnode->inputs();
+  if (inputs.size() <= kSparseAttrIndex) {
+    MS_LOG(EXCEPTION) << "For SparseGetAttr, CNode must have 2 inputs (Prim, Sparse)";
+  }
+  AbstractCSRToAbstractTuple(inputs[kSparseAttrIndex]);
+  auto index_node = NewValueNode(index);
+  AbstractBasePtr index_abs = std::make_shared<abstract::AbstractScalar>(std::make_shared<Int64Imm>(index));
+  index_node->set_abstract(index_abs);
+
+  // Create new node with TupleGetItem operation and set its attributes
+  auto new_node =
+    NewCNode({NewValueNode(prim::kPrimTupleGetItem), inputs[kSparseAttrIndex], index_node}, cnode->func_graph());
+  new_node->set_abstract(node->abstract());
+  if (kernel_graph != nullptr) {
+    kernel_graph->FrontBackendlMapUpdate(cnode, new_node);
+  }
+  return new_node;
+}
+
+/**
+ * \brief Processes the inputs of a CNode representing a sparse operation.
+ * 
+ * For nodes representing specific sparse operations, this function splits their components and adjusts their attributes.
+ * 
+ * \param node AnfNode pointer representing the node to be processed.
+ * \param kernel_graph KernelGraph pointer used for node creation and manipulation.
+ * 
+ * \return Returns a new CNodePtr with adjusted inputs and attributes.
+ */
+CNodePtr FetchInputsForSparseOP(const AnfNodePtr &node, const KernelGraphPtr &kernel_graph) {
+  // Validation
+  MS_EXCEPTION_IF_NULL(node);
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+
+  CNodePtr cnode = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(cnode);
+
+  if (cnode->GetAttr("has_been_split") != nullptr) {
+    MS_LOG(INFO) << "Do not process CNode " << cnode << " (" << cnode->DebugString() << "), because it has been split.";
+    return nullptr;
+  }
+
+  // Process the inputs based on their type and whether they represent a sparse operation
+  const auto &inputs = cnode->inputs();
+  std::vector<AnfNodePtr> new_inputs;
+  new_inputs.push_back(inputs[0]);
+  for (size_t i = 1; i < inputs.size(); ++i) {
+    if (inputs[i]->isa<CNode>()) {
+      if (SplitCNode(inputs[i], &new_inputs)) continue;
+    } else if (inputs[i]->isa<ValueNode>()) {
+      if (SplitValueNode(inputs[i], &new_inputs, kernel_graph)) continue;
+    } else if (inputs[i]->isa<Parameter>()) {
+      if (SplitParameter(inputs[i], &new_inputs, kernel_graph)) continue;
+    }
+    new_inputs.push_back(inputs[i]);
+  }
+
+  // Create new node with the processed inputs and set its attributes
+  auto new_node = NewCNode(new_inputs, cnode->func_graph());
+  new_node->set_abstract(node->abstract());
+  new_node->AddAttr("has_been_split", MakeValue(true));
+  if (kernel_graph != nullptr) {
+    kernel_graph->FrontBackendlMapUpdate(cnode, new_node);
+  }
+  return new_node;
+}
+
+/**
+ * \brief Main processing function for nodes representing sparse operations.
+ * 
+ * This function examines the type of sparse operation represented by the node and dispatches the appropriate processing
+ * function.
+ * 
+ * \param func_graph FuncGraph pointer for the current computational graph.
+ * \param node AnfNode pointer representing the node to be processed.
+ * \param EquivPtr Placeholder for future functionality, currently not used.
+ * 
+ * \return Returns a processed AnfNodePtr if the node represents a recognized sparse operation, or nullptr otherwise.
+ */
+const AnfNodePtr SparseProcess::Process(const FuncGraphPtr &func_graph, const AnfNodePtr &node,
+                                        const EquivPtr &) const {
+  // Validation
+  MS_EXCEPTION_IF_NULL(func_graph);
+  MS_EXCEPTION_IF_NULL(node);
+  if (node == nullptr || !node->isa<CNode>() || !AnfUtils::IsRealKernel(node)) {
+    return nullptr;
+  }
+
+  auto cnode = node->cast<CNodePtr>();
+  MS_EXCEPTION_IF_NULL(cnode);
+  auto prim = common::AnfAlgo::GetCNodePrimitive(cnode);
+  MS_EXCEPTION_IF_NULL(prim);
+  std::string prim_name = prim->name();
+
+  // Dispatch based on the type of sparse operation represented by the node
+  auto kernel_graph = func_graph->cast<std::shared_ptr<session::KernelGraph>>();
+  if (make_sparse_set.find(prim_name) != make_sparse_set.end()) {
+    return ConvertMakeSparseToMakeTuple(node, kernel_graph);
+  } else if (sparse_attr_map.find(prim_name) != sparse_attr_map.end()) {
+    return ConvertSparseGetAttrToTupleGetItem(sparse_attr_map.at(prim_name), node, kernel_graph);
+  } else if (sparse_op_set.find(prim_name) != sparse_op_set.end()) {
+    return FetchInputsForSparseOP(node, kernel_graph);
+  }
+  return nullptr;
+}
+
+}  // namespace opt
+}  // namespace mindspore

mindspore/ccsrc/backend/common/session/ascend_auto_monad.cc

@@ -60,12 +60,29 @@ const char OUTPUT[] = "output";
 const char ITEREND[] = "PROFILING_ITER_END";
 
 #ifdef ENABLE_DUMP_IR
+
+/**
+ * @brief Check if graph saving is enabled.
+ * 
+ * This function checks if the flag for saving graphs is enabled in the current context.
+ * 
+ * @return True if graph saving is enabled, false otherwise.
+ */
 bool IsSaveGraph() {
   auto context_ptr = MsContext::GetInstance();
   MS_EXCEPTION_IF_NULL(context_ptr);
   return context_ptr->get_param<bool>(MS_CTX_SAVE_GRAPHS_FLAG);
 }
 
+/**
+ * @brief Recursively dump all graphs starting from a given kernel graph.
+ * 
+ * This function recursively traverses and dumps all child graphs starting from a given kernel graph.
+ * It uses a memoization set to avoid processing the same graph multiple times.
+ * 
+ * @param kg The kernel graph to start the traversal.
+ * @param memo A set to keep track of processed graphs.
+ */
 void DumpAllGraphs(NotNull<KernelGraphPtr> kg, std::set<KernelGraphPtr> *memo) {
   if (memo->find(kg) != memo->end()) {
     return;
@@ -81,15 +98,34 @@ void DumpAllGraphs(NotNull<KernelGraphPtr> kg, std::set<KernelGraphPtr> *memo) {
   }
 }
 
+/**
+ * @brief Dump the entire graph hierarchy for debugging purposes.
+ * 
+ * This function is used to dump the entire graph hierarchy, including all child graphs,
+ * for debugging purposes. It checks if graph saving is enabled and then initiates the dumping process.
+ * 
+ * @param kg The kernel graph to start the dumping process.
+ */
 void DumpGraphForDebug(const NotNull<KernelGraphPtr> kg) {
   if (IsSaveGraph()) {
     std::set<KernelGraphPtr> memo;
     DumpAllGraphs(kg, &memo);
   }
 }
+
 #endif
 
+
 #ifndef ENABLE_SECURITY
+
+/**
+ * @brief Dump the execution order of nodes in the kernel graph.
+ * 
+ * This function dumps the execution order of nodes in the kernel graph to a file. It checks if graph saving is enabled,
+ * constructs the file path, and writes the execution order along with labels and other relevant information to the file.
+ * 
+ * @param kg The kernel graph whose execution order needs to be dumped.
+ */
 void DumpExecuteOrder(const NotNull<KernelGraphPtr> kg) {
   if (!IsSaveGraph()) {
     return;
@@ -140,9 +176,18 @@ void DumpExecuteOrder(const NotNull<KernelGraphPtr> kg) {
   }
   fout.close();
 }
+
 #endif
 
-// Return kNoLabel when label id attribute not set for the graph.
+/**
+ * @brief Get the label ID associated with the given kernel graph.
+ * 
+ * This function retrieves the label ID associated with a kernel graph. If the label ID attribute is not set for the graph,
+ * it returns kNoLabel.
+ * 
+ * @param kg The kernel graph for which to get the label ID.
+ * @return The label ID or kNoLabel if not set.
+ */
 uint32_t GetGraphLabel(const KernelGraphPtr &kg) {
   auto value = kg->get_attr(kAttrLabelIndex);
   if (value == nullptr) {
@@ -151,8 +196,30 @@ uint32_t GetGraphLabel(const KernelGraphPtr &kg) {
   return GetValue<uint32_t>(value);
 }
 
+/**
+ * @brief Check if two abstract values are compatible.
+ * 
+ * This function checks if two abstract values are compatible with each other. The compatibility criteria may vary based
+ * on the specific use case.
+ * 
+ * @param a1 The first abstract value.
+ * @param a2 The second abstract value.
+ * @return True if the abstract values are compatible, false otherwise.
+ */
 bool IsCompatible(const abstract::AbstractBasePtr &a1, const abstract::AbstractBasePtr &a2);
 
+
+/**
+ * @brief Check if two AbstractTuple abstract values are compatible.
+ * 
+ * This function checks if two AbstractTuple abstract values are compatible by comparing the elements of the tuples.
+ * If the number of elements in the two tuples is different, they are considered incompatible.
+ * Otherwise, it recursively checks the compatibility of each element in the tuples.
+ * 
+ * @param a1 The first AbstractTuple abstract value.
+ * @param a2 The second AbstractTuple abstract value.
+ * @return True if the AbstractTuple abstract values are compatible, false otherwise.
+ */
 bool CheckAbstractTupleIsCompatible(const abstract::AbstractBasePtr &a1, const abstract::AbstractBasePtr &a2) {
   auto &a1_tuple = static_cast<abstract::AbstractTuple &>(*a1);
   auto &a2_tuple = static_cast<abstract::AbstractTuple &>(*a2);
@@ -171,6 +238,17 @@ bool CheckAbstractTupleIsCompatible(const abstract::AbstractBasePtr &a1, const a
   return true;
 }
 
+/**
+ * @brief Check if one abstract tensor is compatible with an abstract scalar.
+ * 
+ * This function checks if one abstract tensor is compatible with an abstract scalar. It considers the following two cases as compatible:
+ *   1. a1: AbstractTensor (shape: (), element: AbstractScalar)
+ *   2. a2: AbstractScalar
+ * 
+ * @param a1 The first abstract value to check.
+ * @param a2 The second abstract value to check.
+ * @return True if the abstract values are compatible, false otherwise.
+ */
 bool CheckTensorAndScalar(const abstract::AbstractBasePtr &a1, const abstract::AbstractBasePtr &a2) {
   if (a1->isa<abstract::AbstractTensor>() && a2->isa<abstract::AbstractScalar>()) {
     auto a1_element = static_cast<abstract::AbstractUndetermined &>(*a1).element();
@@ -181,7 +259,17 @@ bool CheckTensorAndScalar(const abstract::AbstractBasePtr &a1, const abstract::A
   return false;
 }
 
-// Check if one abstract is compatible with another abstract.
+/**
+ * @brief Check if two abstract values are compatible.
+ * 
+ * This function checks if two abstract values are compatible with each other. It considers various cases for compatibility,
+ * including AbstractTuple, AbstractTensor, AbstractRef, and others. If the abstract values match any of these cases,
+ * they are considered compatible.
+ * 
+ * @param a1 The first abstract value.
+ * @param a2 The second abstract value.
+ * @return True if the abstract values are compatible, false otherwise.
+ */
 bool IsCompatible(const abstract::AbstractBasePtr &a1, const abstract::AbstractBasePtr &a2) {
   if (a1 == nullptr || a2 == nullptr) {
     return false;
@@ -194,8 +282,8 @@ bool IsCompatible(const abstract::AbstractBasePtr &a1, const abstract::AbstractB
     return CheckAbstractTupleIsCompatible(a1, a2);
   }
   // Consider the following two cases as compatible：
-  // a1: AbstractScalar(Type: Bool, Value: AnyValue, Shape: NoShape)
-  // a2: AbstractTensor(shape: (), element: AbstractScalar(Type: Bool, Value: AnyValue, Shape: NoShape), value:...)
+  //   1. a1: AbstractScalar(Type: Bool, Value: AnyValue, Shape: NoShape)
+  //   2. a2: AbstractTensor(shape: (), element: AbstractScalar(Type: Bool, Value: AnyValue, Shape: NoShape), value:...)
   if (CheckTensorAndScalar(a1, a2) || CheckTensorAndScalar(a2, a1)) {
     return true;
   }
@@ -221,6 +309,7 @@ bool IsCompatible(const abstract::AbstractBasePtr &a1, const abstract::AbstractB
   return *shape1 == *shape2;
 }
 
+
 struct CallBranch {
   KernelGraphPtr graph;
   std::vector<AnfNodePtr> args;

mindspore/ccsrc/backend/common/session/ascend_inference_session.cc

@@ -29,6 +29,15 @@
 
 namespace mindspore {
 namespace session {
+/**
+ * @brief Load input data into the device for inference.
+ * 
+ * This function loads input data into the device for inference. It synchronizes input tensors from the host
+ * to the device for all non-weight input parameters in the given kernel graph.
+ * 
+ * @param kernel_graph The kernel graph for which input data needs to be loaded.
+ * @param inputs_const A vector of input tensors to be loaded.
+ */
 void AscendInferenceSession::LoadInputData(const std::shared_ptr<KernelGraph> &kernel_graph,
                                            const std::vector<tensor::TensorPtr> &inputs_const) const {
   MS_EXCEPTION_IF_NULL(kernel_graph);
@@ -39,6 +48,7 @@ void AscendInferenceSession::LoadInputData(const std::shared_ptr<KernelGraph> &k
   for (size_t i = 0; i < input_nodes.size(); ++i) {
     tensor::TensorPtr tensor = nullptr;
     if (!input_nodes[i]->isa<Parameter>() || !AnfAlgo::OutputAddrExist(input_nodes[i], 0)) {
+      // Skip inputs that are not Parameters or do not have output addresses.
       MS_LOG(INFO) << "Kernel graph inputs have anfnode which is not Parameter or without output addr.";
       continue;
     }
@@ -47,6 +57,7 @@ void AscendInferenceSession::LoadInputData(const std::shared_ptr<KernelGraph> &k
     auto device_address = AnfAlgo::GetMutableOutputAddr(pk_node, 0);
     MS_EXCEPTION_IF_NULL(device_address);
     if (!common::AnfAlgo::IsParameterWeight(pk_node)) {
+      // Load input data for non-weight parameters.
       tensor = inputs[no_weight_input++];
       if (!device_address->SyncHostToDevice(trans::GetRuntimePaddingShape(pk_node, 0),
                                             LongToSize(tensor->data().nbytes()), tensor->data_type(), tensor->data_c(),
@@ -57,14 +68,24 @@ void AscendInferenceSession::LoadInputData(const std::shared_ptr<KernelGraph> &k
   }
 }
 
+/**
+ * @brief Compile a graph and load weight data to the device.
+ * 
+ * This function compiles the given graph and loads weight data to the device for weight parameters.
+ * It calls the base class's CompileGraphImpl and then loads weight data for the graph's input parameters.
+ * 
+ * @param func_graph The function graph to be compiled.
+ * @return The unique graph identifier after compilation.
+ */
 GraphId AscendInferenceSession::CompileGraphImpl(NotNull<FuncGraphPtr> func_graph) {
   auto graph_id = AscendSession::CompileGraphImpl(func_graph);
   auto kernel_graph = GetGraph(graph_id);
   MS_EXCEPTION_IF_NULL(kernel_graph);
-  // load weight data to device
+  // Load weight data to device for input parameters
   auto input_nodes = kernel_graph->inputs();
   for (size_t i = 0; i < input_nodes.size(); ++i) {
     if (!input_nodes[i]->isa<Parameter>() || !AnfAlgo::OutputAddrExist(input_nodes[i], 0)) {
+      // Skip inputs that are not Parameters or do not have output addresses.
       MS_LOG(INFO) << "Kernel graph inputs have anfnode which is not Parameter or without output addr.";
       continue;
     }
@@ -73,6 +94,7 @@ GraphId AscendInferenceSession::CompileGraphImpl(NotNull<FuncGraphPtr> func_grap
     auto device_address = AnfAlgo::GetMutableOutputAddr(pk_node, 0);
     MS_EXCEPTION_IF_NULL(device_address);
     if (common::AnfAlgo::IsParameterWeight(pk_node)) {
+      // Load weight data to the device for weight parameters.
       const auto &param_value = pk_node->default_param();
       MS_EXCEPTION_IF_NULL(param_value);
       auto tensor = std::dynamic_pointer_cast<tensor::Tensor>(param_value);
@@ -87,6 +109,18 @@ GraphId AscendInferenceSession::CompileGraphImpl(NotNull<FuncGraphPtr> func_grap
   return graph_id;
 }
 
+
+/**
+ * @brief Check if the provided model inputs are compatible with the graph inputs.
+ * 
+ * This function checks whether the provided model inputs are compatible with the graph inputs
+ * for the given graph. It compares the number of inputs, their shapes, and data types.
+ * 
+ * @param graph_id The unique identifier of the graph to be checked.
+ * @param inputs The vector of input tensors to be checked.
+ * @param error_msg A pointer to a string that will contain an error message in case of failure.
+ * @return True if the inputs are compatible; false otherwise.
+ */
 bool AscendInferenceSession::CheckModelInputs(uint32_t graph_id, const std::vector<tensor::TensorPtr> &inputs,
                                               std::string *error_msg) const {
   MS_LOG(INFO) << "Start check client inputs, graph id : " << graph_id;
@@ -95,7 +129,8 @@ bool AscendInferenceSession::CheckModelInputs(uint32_t graph_id, const std::vect
   auto kernel_graph_inputs = kernel_graph->inputs();
   size_t no_weight_input = 0;
   vector<ParameterPtr> paras;
-  // find parameters of graph inputs
+  
+  // Find parameters of graph inputs (excluding weight parameters).
   for (size_t i = 0; i < kernel_graph_inputs.size(); ++i) {
     if (!kernel_graph_inputs[i]->isa<Parameter>()) {
       MS_LOG(ERROR) << "Kernel graph inputs have anfnode which is not Parameter.";
@@ -107,9 +142,9 @@ bool AscendInferenceSession::CheckModelInputs(uint32_t graph_id, const std::vect
     }
   }
 
-  // check inputs
+  // Check inputs
   for (size_t i = 0; i < paras.size(); ++i) {
-    // compare input number
+    // Compare input tensor with the corresponding graph parameter.
     if (paras.size() != inputs.size()) {
       MS_LOG(ERROR) << "Input number is inconsistent. The actual input number [" << inputs.size()
                     << "] but the graph input number is [" << paras.size() << "]";
@@ -139,13 +174,21 @@ bool AscendInferenceSession::CheckModelInputs(uint32_t graph_id, const std::vect
   return true;
 }
 
+/**
+ * @brief Compare an input tensor with a graph parameter.
+ * 
+ * This function compares an input tensor with a graph parameter in terms of shape and data type.
+ * 
+ * @param input The input tensor to be compared.
+ * @param parameter The graph parameter to be compared.
+ * @return True if the input and parameter are compatible; false otherwise.
+ */
 bool AscendInferenceSession::CompareInput(const tensor::TensorPtr &input, const ParameterPtr &parameter) const {
   MS_EXCEPTION_IF_NULL(input);
   MS_EXCEPTION_IF_NULL(parameter);
-  // compare dims
-  auto parameter_shape = AnfAlgo::GetOutputDeviceShape(parameter, 0);
 
-  // compare shape
+  // Compare shapes
+  auto parameter_shape = AnfAlgo::GetOutputDeviceShape(parameter, 0);
   auto input_shape = input->shape();
   vector<size_t> trans_input;
   (void)std::transform(input_shape.begin(), input_shape.end(), std::back_inserter(trans_input),
@@ -160,7 +203,7 @@ bool AscendInferenceSession::CompareInput(const tensor::TensorPtr &input, const
     return false;
   }
 
-  // compare data type
+  // Compare data types
   auto kernel_build_info = AnfAlgo::GetSelectKernelBuildInfo(parameter);
   if (input->data_type() != kernel_build_info->GetOutputDeviceType(0)) {
     MS_LOG(ERROR) << "Input data type is inconsistent. The actual data type is " << input->data_type()
@@ -171,6 +214,7 @@ bool AscendInferenceSession::CompareInput(const tensor::TensorPtr &input, const
   return true;
 }
 
+
 template <typename T>
 std::string AscendInferenceSession::PrintInputShape(std::vector<T> shape) const {
   string res = "[";

mindspore/ccsrc/backend/common/session/cpu_session_old.cc

@@ -0,0 +1,411 @@
+/**
+ * Copyright 2019-2022 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "backend/common/session/cpu_session.h"
+#include <algorithm>
+#include <sstream>
+#include <exception>
+#include "ir/anf.h"
+#include "utils/ms_utils.h"
+#include "utils/trace_base.h"
+#include "common/graph_kernel/graph_kernel_flags.h"
+#include "backend/common/session/anf_runtime_algorithm.h"
+#include "include/common/utils/anfalgo.h"
+#include "plugin/factory/ms_factory.h"
+#include "runtime/device/kernel_runtime.h"
+#include "plugin/device/cpu/kernel/cpu_kernel.h"
+#include "plugin/device/cpu/kernel/akg/akg_cpu_kernel_build.h"
+#include "plugin/device/cpu/hal/device/kernel_select_cpu.h"
+#include "backend/common/optimizer/optimizer.h"
+#include "backend/common/optimizer/pass_manager.h"
+#include "plugin/device/cpu/optimizer/insert_cast_cpu.h"
+#include "plugin/device/cpu/optimizer/insert_format_transform_op.h"
+#include "common/graph_kernel/adapter/graph_kernel_optimization.h"
+#include "backend/common/pass/replace_node_by_proxy.h"
+#include "backend/common/pass/erase_visit_attr.h"
+#include "include/common/debug/anf_ir_dump.h"
+#include "backend/common/optimizer/common_backend_optimization.h"
+#include "include/common/debug/dump_proto.h"
+#ifndef ENABLE_SECURITY
+#include "debug/data_dump/dump_json_parser.h"
+#endif
+#if ((defined ENABLE_CPU) && (!defined _WIN32))
+#include "ps/util.h"
+#include "ps/ps_context.h"
+#endif
+#ifdef ENABLE_DUMP_IR
+#include "debug/rdr/graph_recorder.h"
+#endif
+
+namespace mindspore {
+namespace session {
+void CPUSession::Init(uint32_t device_id) {
+#ifndef ENABLE_SECURITY
+  // Dump json config file if dump is enabled
+  auto &json_parser = DumpJsonParser::GetInstance();
+  json_parser.Parse();
+  json_parser.CopyMSCfgJsonToDir(rank_id_);
+#endif
+  InitExecutor(kCPUDevice, device_id);
+}
+
+ParameterPtr CPUSession::CreateNewParameterFromParameter(const AnfNodePtr &anf, KernelGraph *graph) {
+  MS_EXCEPTION_IF_NULL(anf);
+  MS_EXCEPTION_IF_NULL(graph);
+  if (!anf->isa<Parameter>()) {
+    MS_LOG(EXCEPTION) << "anf[" << anf->DebugString() << "] is not a parameter";
+  }
+  auto valid_inputs = graph->MutableValidInputs();
+  auto graph_inputs = graph->MutableInputs();
+  MS_EXCEPTION_IF_NULL(graph_inputs);
+  TraceManager::DebugTrace(std::make_shared<TraceCopy>(anf->debug_info()));
+  ParameterPtr new_parameter = graph->NewParameter(anf->cast<ParameterPtr>());
+  TraceManager::EndTrace();
+  graph_inputs->push_back(new_parameter);
+  valid_inputs->push_back(true);
+  return new_parameter;
+}
+
+// Remove after PS feature finish adapting push/pull in auto_monad.
+void CPUSession::Reorder(std::vector<CNodePtr> *node_list) {
+  common::AnfAlgo::ReorderPosteriorExecList(NOT_NULL(node_list));
+}
+
+void CPUSession::Optimize(const std::shared_ptr<KernelGraph> &kernel_graph) {
+  auto optimizer = std::make_shared<opt::GraphOptimizer>();
+  auto pm = std::make_shared<opt::PassManager>();
+#if ((defined ENABLE_CPU) && (!defined _WIN32) && !defined(__APPLE__))
+  auto ms_context = MsContext::GetInstance();
+  MS_EXCEPTION_IF_NULL(ms_context);
+  if (ms_context->get_param<int>(MS_CTX_EXECUTION_MODE) != kPynativeMode && ps::PSContext::instance()->is_ps_mode()) {
+    AssignParamKey(kernel_graph);
+    if (ps::PSContext::instance()->is_worker()) {
+      std::string pass_name = "replace_node_by_proxy";
+      pass_name.append(std::to_string(graph_sum_));
+      pm->AddPass(std::make_shared<opt::ReplaceNodeByProxy>(pass_name));
+    }
+  }
+#endif
+  pm->AddPass(std::make_shared<opt::InsertFormatTransformOpCPU>("insert_format_transform_op_cpu"));
+  pm->AddPass(std::make_shared<opt::InsertCastCPU>("insert_cast"));
+  pm->AddPass(std::make_shared<opt::EraseVisitAttr>());
+  optimizer->AddPassManager(pm);
+  (void)optimizer->Optimize(kernel_graph);
+  kernel_graph->SetExecOrderByDefault();
+}
+
+void CPUSession::GraphKernelOptimize(const std::shared_ptr<KernelGraph> &kernel_graph) {
+#ifdef ENABLE_AKG
+  if (!graphkernel::GraphKernelFlags::GetInstance().IsEnableGraphKernel()) {
+    return;
+  }
+  graphkernel::GraphKernelOptimize(kernel_graph);
+  kernel_graph->SetExecOrderByDefault();
+#endif
+}
+
+GraphId CPUSession::CompileGraphImpl(const AnfNodePtrList &lst, const AnfNodePtrList &outputs) {
+  auto graph_id = graph_sum_;
+  auto graph = ConstructKernelGraph(lst, outputs, DeviceAddressType::kCPU);
+  MS_EXCEPTION_IF_NULL(graph);
+  opt::AddDynamicShapeAttrPass(graph);
+  MS_LOG(INFO) << "Set kernel info";
+  SetKernelInfo(graph.get());
+  MS_LOG(INFO) << "Set kernel info end";
+  Optimize(graph);
+  FinalOptimize(graph);
+  GraphKernelOptimize(graph);
+  MS_LOG(INFO) << "Build kernel";
+  BuildKernel(graph.get());
+  // Remove reorder after PS feature finish adapting push/pull in auto_monad.
+  auto execution_order = graph->execution_order();
+  Reorder(&execution_order);
+  graph->set_execution_order(execution_order);
+  // runtime init
+  if (!runtime_.Init()) {
+    MS_LOG(EXCEPTION) << "Kernel runtime init error.";
+  }
+  MS_LOG(INFO) << "Assign kernel graph address";
+  runtime_.AssignKernelGraphAddress(graph.get());
+  // set summary node
+#ifndef ENABLE_SECURITY
+  SetSummaryNodes(graph.get());
+#endif
+  runtime_.IncreaseSummaryRefCount(graph->summary_nodes());
+  DumpGraphs({graph});
+  return graph_id;
+}
+
+void CPUSession::CreateOutputTensors(const GraphId &graph_id, const std::vector<tensor::TensorPtr> &input_tensors,
+                                     VectorRef *outputs,
+                                     std::map<tensor::TensorPtr, session::KernelWithIndex> *tensor_to_node,
+                                     KernelMapTensor *) {
+  auto kernel_graph = GetGraph(graph_id);
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  runtime_.CreateOutputTensors(kernel_graph.get(), input_tensors, outputs, tensor_to_node);
+}
+
+void CPUSession::LoadInputData(const std::shared_ptr<KernelGraph> &kernel_graph,
+                               const std::vector<tensor::TensorPtr> &inputs_const) const {
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  auto &input_nodes = kernel_graph->input_nodes();
+  if (input_nodes.size() != inputs_const.size()) {
+    MS_LOG(EXCEPTION) << "Input size " << inputs_const.size() << " is not equal to input node size "
+                      << input_nodes.size();
+  }
+  for (size_t input_idx = 0; input_idx < input_nodes.size(); ++input_idx) {
+    auto &input_node = input_nodes[input_idx];
+    MS_EXCEPTION_IF_NULL(input_node);
+    if (!input_node->isa<Parameter>() || HasAbstractMonad(input_node)) {
+      continue;
+    }
+    auto address = AnfAlgo::GetMutableOutputAddr(input_node, 0);
+    auto tensor = inputs_const[input_idx];
+    auto tensor_address = tensor->device_address();
+    MS_EXCEPTION_IF_NULL(address);
+    MS_EXCEPTION_IF_NULL(tensor);
+    if (tensor_address == nullptr || tensor_address == address) {
+      continue;
+    }
+    auto input_param = input_node->cast<ParameterPtr>();
+    if (common::AnfAlgo::IsParameterWeight(input_param) && !tensor->IsUpdatedByDevice()) {
+      continue;
+    }
+    if (std::dynamic_pointer_cast<device::DeviceAddress>(tensor_address)->DeviceType() !=
+        device::DeviceAddressType::kCPU) {
+      tensor->data_sync(false);
+    }
+  }
+}
+
+void CPUSession::PreExecuteGraph(const std::shared_ptr<KernelGraph> &kernel_graph,
+                                 const std::vector<tensor::TensorPtr> &inputs, VectorRef *const outputs) {
+  MS_LOG(INFO) << "Bind input output address";
+  runtime_.BindInputOutput(kernel_graph.get(), inputs, outputs);
+
+#if ((defined ENABLE_CPU) && (!defined _WIN32) && !defined(__APPLE__))
+  InitPSParamAndOptim(kernel_graph, inputs);
+#endif
+}
+
+void CPUSession::PostExecuteGraph(const std::shared_ptr<KernelGraph> &kernel_graph,
+                                  const std::vector<tensor::TensorPtr> &, VectorRef *const) {
+#ifndef ENABLE_SECURITY
+  Summary(kernel_graph.get());
+#endif
+}
+
+void CPUSession::ExecuteGraph(const std::shared_ptr<KernelGraph> &kernel_graph) {
+  bool ret = runtime_.Run(*kernel_graph, false);
+  if (!ret) {
+    MS_LOG(EXCEPTION) << "Run graph failed";
+  }
+}
+
+KernelGraphPtr CPUSession::BuildOpImpl(const OpRunInfo &op_run_info, const GraphInfo &graph_info,
+                                       const std::vector<tensor::TensorPtr> &input_tensors,
+                                       const std::vector<int64_t> &tensors_mask) {
+  // Check if the graph cache exists.
+  auto it = run_op_graphs_.find(graph_info);
+  if (it != run_op_graphs_.end()) {
+    return it->second;
+  }
+
+  // Prepare the graph
+  const auto &kernel_graph = ConstructSingleOpGraph(op_run_info, input_tensors, tensors_mask);
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  SetKernelInfo(kernel_graph.get());
+  Optimize(kernel_graph);
+  BuildKernel(kernel_graph.get());
+  auto enable_op_graph_cache = MsContext::GetInstance()->get_param<bool>(MS_CTX_ENABLE_PYNATIVE_OP_GRAPH_CACHE);
+  if (enable_op_graph_cache) {
+    run_op_graphs_[graph_info] = kernel_graph;
+  }
+  return kernel_graph;
+}
+
+void CPUSession::SetOutputFlags(const VectorRef &base_ref) {
+  for (size_t i = 0; i < base_ref.size(); ++i) {
+    if (utils::isa<VectorRef>(base_ref[i])) {
+      auto ref_iter = utils::cast<VectorRef>(base_ref[i]);
+      SetOutputFlags(ref_iter);
+    } else if (utils::isa<tensor::TensorPtr>(base_ref[i])) {
+      auto tensor_ptr = utils::cast<std::shared_ptr<tensor::Tensor>>(base_ref[i]);
+      tensor_ptr->SetNeedWait(false);
+      tensor_ptr->data_sync(false);
+    }
+  }
+}
+
+void CPUSession::UpdateDynamicOutputShape(const std::map<tensor::TensorPtr, KernelWithIndex> &tensor_to_node) {
+  for (const auto &tensor_node : tensor_to_node) {
+    if (common::AnfAlgo::IsDynamicShape(tensor_node.second.first)) {
+      const auto &kernel = tensor_node.second.first;
+      const auto &output_index = tensor_node.second.second;
+      const auto &shape = common::AnfAlgo::GetOutputInferShape(kernel, output_index);
+      std::vector<int64_t> refresh_shape;
+      (void)std::copy(shape.begin(), shape.end(), std::back_inserter(refresh_shape));
+      MS_EXCEPTION_IF_NULL(tensor_node.first);
+      tensor_node.first->set_shape(refresh_shape);
+    }
+  }
+}
+
+void CPUSession::RunOpImplOrigin(const GraphInfo &graph_info, OpRunInfo *op_run_info,
+                                 std::vector<tensor::TensorPtr> *input_tensors, VectorRef *outputs,
+                                 const std::vector<int64_t> &tensors_mask) {
+  RunOpImpl(graph_info, op_run_info, input_tensors, outputs, tensors_mask);
+}
+
+void CPUSession::RunOpImpl(const GraphInfo &graph_info, OpRunInfo *op_run_info,
+                           std::vector<tensor::TensorPtr> *input_tensors, VectorRef *outputs,
+                           const std::vector<int64_t> &tensors_mask) {
+  MS_EXCEPTION_IF_NULL(input_tensors);
+  MS_EXCEPTION_IF_NULL(op_run_info);
+  ProcessInputTensorsForHeterogeneous("CPU", *input_tensors);
+  const auto &kernel_graph = BuildOpImpl(*op_run_info, graph_info, *input_tensors, tensors_mask);
+  EraseValueNodeTensor(tensors_mask, input_tensors);
+
+  // Remove reorder after PS feature finish adapting push/pull in auto_monad.
+  auto execution_order = kernel_graph->execution_order();
+  Reorder(&execution_order);
+  kernel_graph->set_execution_order(execution_order);
+
+  // runtime init
+  if (!runtime_.Init()) {
+    MS_LOG(EXCEPTION) << "Kernel runtime init error.";
+  }
+  runtime_.AssignKernelGraphAddress(kernel_graph.get());
+  std::map<tensor::TensorPtr, session::KernelWithIndex> tensor_to_node;
+  runtime_.CreateOutputTensors(kernel_graph.get(), *input_tensors, outputs, &tensor_to_node);
+  runtime_.BindInputOutput(kernel_graph.get(), *input_tensors, outputs);
+
+  bool ret = runtime_.Run(*kernel_graph, false);
+  if (!ret) {
+    MS_LOG(EXCEPTION) << "Run Op failed";
+  }
+  UpdateDynamicOutputShape(tensor_to_node);
+  // update output abstract of dynamic op to op_run_info
+  if (op_run_info->is_dynamic_shape) {
+    UpdateOutputAbstract(kernel_graph, op_run_info);
+  }
+  SetOutputFlags(*outputs);
+  runtime_.RunOpClearMemory(*kernel_graph);
+}
+
+void CPUSession::SetKernelInfo(const KernelGraph *kernel_graph) {
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  auto &kernel_nodes = kernel_graph->execution_order();
+  for (const auto &kernel_node : kernel_nodes) {
+    MS_EXCEPTION_IF_NULL(kernel_node);
+    device::cpu::SetKernelInfo(kernel_node);
+  }
+}
+
+namespace {
+void KernelNotSupportException(const AnfNodePtr &kernel_node) {
+  std::string kernel_name = common::AnfAlgo::GetCNodeName(kernel_node);
+  std::stringstream operator_info;
+  operator_info << "Operator[" << kernel_name << "] ";
+  auto kernel_info = dynamic_cast<device::KernelInfo *>(kernel_node->kernel_info());
+  if (kernel_info == nullptr) {
+    operator_info << "is not support.";
+    MS_LOG(EXCEPTION) << operator_info.str();
+  }
+  auto kernel_build_Info = kernel_info->select_kernel_build_info();
+  if (kernel_build_Info == nullptr) {
+    operator_info << "is not support.";
+    MS_LOG(EXCEPTION) << operator_info.str();
+  }
+  size_t input_num = kernel_build_Info->GetInputNum();
+  if (input_num > 0) {
+    operator_info << " input(";
+    for (size_t i = 0; i < input_num; ++i) {
+      operator_info << TypeIdLabel(kernel_build_Info->GetInputDeviceType(i));
+      if (i != input_num - 1) {
+        operator_info << ",";
+      }
+    }
+    operator_info << ") ";
+  }
+  size_t output_num = kernel_build_Info->GetOutputNum();
+  if (output_num > 0) {
+    operator_info << "output(";
+    for (size_t i = 0; i < output_num; ++i) {
+      operator_info << TypeIdLabel(kernel_build_Info->GetOutputDeviceType(i));
+      if (i != kernel_build_Info->GetOutputNum() - 1) {
+        operator_info << ",";
+      }
+    }
+    operator_info << ") ";
+  }
+  operator_info << "is not support.";
+  MS_LOG(EXCEPTION) << operator_info.str() << trace::DumpSourceLines(kernel_node);
+}
+}  // namespace
+
+void CPUSession::BuildKernel(const KernelGraph *kernel_graph) {
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  auto &kernel_nodes = kernel_graph->execution_order();
+  kernel::KernelMeta *bin_map = kernel::KernelMeta::GetInstance();
+  MS_EXCEPTION_IF_NULL(bin_map);
+  std::vector<AnfNodePtr> akg_nodes;
+  for (const auto &kernel_node : kernel_nodes) {
+    MS_EXCEPTION_IF_NULL(kernel_node);
+    std::string kernel_name = common::AnfAlgo::GetCNodeName(kernel_node);
+    MS_LOG(INFO) << "Cpu building operator[" << kernel_name << "].";
+    if (session::AnfRuntimeAlgorithm::GetKernelType(kernel_node) == KernelType::AKG_KERNEL) {
+      if (!bin_map->initialized()) {
+        bin_map->Initialize();
+      }
+      akg_nodes.push_back(kernel_node);
+      continue;
+    }
+    std::shared_ptr<kernel::NativeCpuKernelMod> cpu_kernel_mod =
+      kernel::Factory<kernel::NativeCpuKernelMod>::Instance().Create(kernel_name);
+    if (cpu_kernel_mod == nullptr) {
+      KernelNotSupportException(kernel_node);
+    }
+    // This branch would be removed When KernelMode rectification is complete
+    auto discard_cpu_kernel_mod = std::dynamic_pointer_cast<kernel::DeprecatedNativeCpuKernelMod>(cpu_kernel_mod);
+    if (discard_cpu_kernel_mod) {
+      try {
+        discard_cpu_kernel_mod->SetCpuRefMapToKernelInfo(kernel_node);
+        discard_cpu_kernel_mod->Init(kernel_node);
+      } catch (std::exception &e) {
+        MS_LOG(EXCEPTION) << e.what() << trace::DumpSourceLines(kernel_node);
+      }
+      AnfAlgo::SetKernelMod(discard_cpu_kernel_mod, kernel_node.get());
+      MS_LOG(INFO) << "Cpu build success operator[" << kernel_name << "].";
+    } else {
+      auto kernel_attrs = cpu_kernel_mod->GetOpSupport();
+      SetCpuRefMapToKernelInfo(kernel_node, kernel_attrs);
+      auto [base_operator, input_tensors, output_tensors] = kernel::GetArgsFromCNode(kernel_node);
+      auto ret = cpu_kernel_mod->Init(base_operator, input_tensors, output_tensors);
+      if (!ret) {
+        MS_LOG(EXCEPTION) << trace::DumpSourceLines(kernel_node);
+      }
+      AnfAlgo::SetKernelMod(cpu_kernel_mod, kernel_node.get());
+      MS_LOG(INFO) << "Cpu build success operator[" << kernel_name << "].";
+    }
+  }
+#ifdef ENABLE_AKG
+  kernel::AkgCpuKernelBuilder akg_cpu_kernel_builder;
+  (void)akg_cpu_kernel_builder.AkgKernelParallelBuild(akg_nodes);
+#endif
+}
+}  // namespace session
+}  // namespace mindspore

mindspore/ccsrc/backend/common/session/gpu_inference_session_old.cc

@@ -0,0 +1,217 @@
+/**
+ * Copyright 2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "backend/common/session/gpu_inference_session.h"
+#include <algorithm>
+#include "ir/tensor.h"
+#include "ir/anf.h"
+#include "ir/param_info.h"
+#include "runtime/device/kernel_runtime.h"
+#include "backend/common/session/anf_runtime_algorithm.h"
+#include "include/common/utils/anfalgo.h"
+#include "utils/ms_utils.h"
+#include "runtime/device/ms_device_shape_transfer.h"
+#include "include/common/utils/config_manager.h"
+
+namespace mindspore {
+namespace session {
+void GpuInferenceSession::LoadInputData(const std::shared_ptr<KernelGraph> &kernel_graph,
+                                        const std::vector<tensor::TensorPtr> &inputs_const) const {
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  std::vector<tensor::TensorPtr> inputs(inputs_const);
+  auto input_nodes = kernel_graph->inputs();
+
+  size_t no_weight_input = 0;
+  for (size_t i = 0; i < input_nodes.size(); ++i) {
+    tensor::TensorPtr tensor = nullptr;
+    if (!input_nodes[i]->isa<Parameter>() || !AnfAlgo::OutputAddrExist(input_nodes[i], 0)) {
+      MS_LOG(INFO) << "Kernel graph inputs is not Parameter or without user.";
+      continue;
+    }
+    auto pk_node = input_nodes[i]->cast<ParameterPtr>();
+    MS_EXCEPTION_IF_NULL(pk_node);
+    auto device_address = AnfAlgo::GetMutableOutputAddr(pk_node, 0);
+    MS_EXCEPTION_IF_NULL(device_address);
+    if (!common::AnfAlgo::IsParameterWeight(pk_node)) {
+      tensor = inputs[no_weight_input++];
+      if (!device_address->SyncHostToDevice(trans::GetRuntimePaddingShape(pk_node, 0),
+                                            LongToSize(tensor->data().nbytes()), tensor->data_type(),
+                                            tensor->data_c())) {
+        MS_LOG(EXCEPTION) << "SyncHostToDevice failed.";
+      }
+    }
+  }
+}
+
+GraphId GpuInferenceSession::CompileGraphImpl(NotNull<FuncGraphPtr> func_graph) {
+  auto graph_id = GPUSession::CompileGraphImpl(func_graph);
+  auto kernel_graph = GetGraph(graph_id);
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  // load weight data to device
+  auto input_nodes = kernel_graph->inputs();
+  for (size_t i = 0; i < input_nodes.size(); ++i) {
+    if (!input_nodes[i]->isa<Parameter>() || !AnfAlgo::OutputAddrExist(input_nodes[i], 0)) {
+      MS_LOG(INFO) << "Kernel graph inputs is not Parameter or without user.";
+      continue;
+    }
+    auto pk_node = input_nodes[i]->cast<ParameterPtr>();
+    MS_EXCEPTION_IF_NULL(pk_node);
+    auto device_address = AnfAlgo::GetMutableOutputAddr(pk_node, 0);
+    MS_EXCEPTION_IF_NULL(device_address);
+    if (common::AnfAlgo::IsParameterWeight(pk_node)) {
+      const auto &param_value = pk_node->default_param();
+      MS_EXCEPTION_IF_NULL(param_value);
+      auto tensor = std::dynamic_pointer_cast<tensor::Tensor>(param_value);
+      MS_EXCEPTION_IF_NULL(tensor);
+      if (!device_address->SyncHostToDevice(trans::GetRuntimePaddingShape(pk_node, 0),
+                                            LongToSize(tensor->data().nbytes()), tensor->data_type(),
+                                            tensor->data_c())) {
+        MS_LOG(EXCEPTION) << "SyncHostToDevice failed.";
+      }
+    }
+  }
+  return graph_id;
+}
+
+bool GpuInferenceSession::CheckModelInputs(uint32_t graph_id, const std::vector<tensor::TensorPtr> &inputs,
+                                           std::string *error_msg) const {
+  MS_LOG(INFO) << "Start check client inputs, graph id : " << graph_id;
+  auto kernel_graph = GetGraph(graph_id);
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  auto kernel_graph_inputs = kernel_graph->inputs();
+  size_t no_weight_input = 0;
+  vector<ParameterPtr> paras;
+  // find parameters of graph inputs
+  for (size_t i = 0; i < kernel_graph_inputs.size(); ++i) {
+    if (!kernel_graph_inputs[i]->isa<Parameter>()) {
+      MS_LOG(ERROR) << "Kernel graph inputs have anfnode which is not Parameter.";
+      continue;
+    }
+    auto parameter = kernel_graph_inputs[i]->cast<ParameterPtr>();
+    if (!common::AnfAlgo::IsParameterWeight(parameter)) {
+      paras.push_back(parameter);
+    }
+  }
+
+  // check inputs
+  for (size_t i = 0; i < paras.size(); ++i) {
+    // compare input number
+    if (paras.size() != inputs.size()) {
+      MS_LOG(ERROR) << "Input number is inconsistent. The actual input number [" << inputs.size()
+                    << "] but the graph input number is [" << paras.size() << "]";
+      MS_LOG(ERROR) << "InputsInfo --" << InputsInfo(paras, inputs);
+      if (error_msg != nullptr) {
+        std::stringstream str_stream;
+        str_stream << "Input number is inconsistent. The given input number [" << inputs.size()
+                   << "] but the graph input number is [" << paras.size() << "]\n";
+        str_stream << "InputsInfo --" << InputsInfo(paras, inputs);
+        *error_msg = str_stream.str();
+      }
+      return false;
+    }
+    auto input = inputs[no_weight_input++];
+    if (!CompareInput(input, paras[i])) {
+      MS_LOG(ERROR) << "Please check the input information.";
+      MS_LOG(ERROR) << "InputsInfo --" << InputsInfo(paras, inputs);
+      if (error_msg != nullptr) {
+        std::stringstream str_stream;
+        str_stream << "Please check the input information.\n";
+        str_stream << "InputsInfo --" << InputsInfo(paras, inputs);
+        *error_msg = str_stream.str();
+      }
+      return false;
+    }
+  }
+  return true;
+}
+
+bool GpuInferenceSession::CompareInput(const tensor::TensorPtr &input, const ParameterPtr &parameter) const {
+  MS_EXCEPTION_IF_NULL(input);
+  MS_EXCEPTION_IF_NULL(parameter);
+  // compare dims
+  auto parameter_shape = AnfAlgo::GetOutputDeviceShape(parameter, 0);
+
+  // compare shape
+  auto input_shape = input->shape();
+  vector<size_t> trans_input;
+  (void)std::transform(input_shape.begin(), input_shape.end(), std::back_inserter(trans_input),
+                       [](const int64_t dim) { return static_cast<size_t>(dim); });
+  auto is_scalar_shape = [](const vector<size_t> &shape) {
+    return shape.empty() || (shape.size() == 1 && shape[0] == 1);
+  };
+  if ((!is_scalar_shape(trans_input) || !is_scalar_shape(parameter_shape)) && (trans_input != parameter_shape)) {
+    MS_LOG(ERROR) << "Input shape is inconsistent. The actual shape is " << PrintInputShape(trans_input)
+                  << ", but the parameter shape is " << PrintInputShape(parameter_shape)
+                  << ". parameter : " << parameter->DebugString();
+    return false;
+  }
+
+  // compare data type
+  auto kernel_build_info = AnfAlgo::GetSelectKernelBuildInfo(parameter);
+  if (input->data_type() != kernel_build_info->GetOutputDeviceType(0)) {
+    MS_LOG(ERROR) << "Input data type is inconsistent. The actual data type is " << input->data_type()
+                  << ", but the parameter data type is " << kernel_build_info->GetOutputDeviceType(0)
+                  << ". parameter : " << parameter->DebugString();
+    return false;
+  }
+  return true;
+}
+
+template <typename T>
+std::string GpuInferenceSession::PrintInputShape(std::vector<T> shape) const {
+  string res = "[";
+  for (auto dim : shape) {
+    res += " " + std::to_string(dim);
+  }
+  return res + " ]";
+}
+
+std::string GpuInferenceSession::InputsInfo(const std::vector<ParameterPtr> &paras,
+                                            const std::vector<tensor::TensorPtr> &inputs) const {
+  const std::map<TypeId, std::string> dtype_name_map{
+    {TypeId::kNumberTypeBegin, "Unknown"},   {TypeId::kNumberTypeBool, "Bool"},
+    {TypeId::kNumberTypeFloat64, "Float64"}, {TypeId::kNumberTypeInt8, "Int8"},
+    {TypeId::kNumberTypeUInt8, "Uint8"},     {TypeId::kNumberTypeInt16, "Int16"},
+    {TypeId::kNumberTypeUInt16, "Uint16"},   {TypeId::kNumberTypeInt32, "Int32"},
+    {TypeId::kNumberTypeUInt32, "Uint32"},   {TypeId::kNumberTypeInt64, "Int64"},
+    {TypeId::kNumberTypeUInt64, "Uint64"},   {TypeId::kNumberTypeFloat16, "Float16"},
+    {TypeId::kNumberTypeFloat32, "Float32"},
+  };
+  auto data_type_to_string = [&dtype_name_map](TypeId type_id) {
+    auto it = dtype_name_map.find(type_id);
+    if (it == dtype_name_map.end()) {
+      return std::string("Unknown");
+    }
+    return it->second;
+  };
+
+  std::string graph = "graph inputs:{ ";
+  for (size_t i = 0; i < paras.size(); ++i) {
+    auto &para = paras[i];
+    graph += std::to_string(i) + ": dims " + std::to_string(AnfAlgo::GetOutputDeviceShape(para, 0).size()) +
+             ", shape " + PrintInputShape(AnfAlgo::GetOutputDeviceShape(para, 0)) + ", data type " +
+             data_type_to_string(AnfAlgo::GetSelectKernelBuildInfo(para)->GetOutputDeviceType(0)) + " }";
+  }
+
+  std::string actual = "given inputs:{ ";
+  for (size_t i = 0; i < inputs.size(); ++i) {
+    actual += std::to_string(i) + ": dims " + std::to_string(inputs[i]->shape().size()) + ", shape " +
+              PrintInputShape(inputs[i]->shape()) + ", data type " + data_type_to_string(inputs[i]->data_type()) + " }";
+  }
+  return graph + "   " + actual;
+}
+}  // namespace session
+}  // namespace mindspore

mindspore/ccsrc/backend/common/session/gpu_session_old.cc

@@ -0,0 +1,769 @@
+/**
+ * Copyright 2019-2021 Huawei Technologies Co., Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "backend/common/session/gpu_session.h"
+
+#include <string>
+#include <utility>
+#include "backend/common/optimizer/helper.h"
+#include "backend/common/optimizer/optimizer.h"
+#include "backend/common/optimizer/pass_manager.h"
+#include "backend/common/optimizer/common_backend_optimization.h"
+#include "plugin/device/gpu/optimizer/adam_weight_decay_fusion.h"
+#include "plugin/device/gpu/optimizer/adam_fusion.h"
+#include "plugin/device/gpu/optimizer/alltoall_fusion.h"
+#include "plugin/device/gpu/optimizer/apply_momentum_weight_scale_fusion.h"
+#include "plugin/device/gpu/optimizer/apply_momentum_scale_fusion.h"
+#include "plugin/device/gpu/optimizer/apply_momentum_weight_fusion.h"
+#include "plugin/device/gpu/optimizer/batch_norm_relu_fusion.h"
+#include "plugin/device/gpu/optimizer/batch_norm_relu_grad_fusion.h"
+#include "plugin/device/gpu/optimizer/batch_norm_add_relu_fusion.h"
+#include "plugin/device/gpu/optimizer/post_batch_norm_add_relu_fusion.h"
+#include "plugin/device/gpu/optimizer/batch_norm_add_relu_grad_fusion.h"
+#include "plugin/device/gpu/optimizer/combine_momentum_fusion.h"
+#include "plugin/device/gpu/optimizer/combine_cast_fusion.h"
+#include "plugin/device/gpu/optimizer/cudnn_inplace_fusion.h"
+#include "plugin/device/gpu/optimizer/insert_format_transform_op.h"
+#include "plugin/device/gpu/optimizer/replace_momentum_cast_fusion.h"
+#include "plugin/device/gpu/optimizer/replace_addn_fusion.h"
+#include "plugin/device/gpu/optimizer/print_reduce_fusion.h"
+#include "plugin/device/gpu/optimizer/bce_with_logits_loss_fusion.h"
+#include "plugin/device/gpu/optimizer/remove_format_transform_pair.h"
+#include "plugin/device/gpu/optimizer/remove_redundant_format_transform.h"
+#include "plugin/device/gpu/optimizer/reduce_precision_fusion.h"
+#include "plugin/device/gpu/optimizer/insert_cast_gpu.h"
+#include "plugin/device/gpu/optimizer/relu_v2_pass.h"
+#include "plugin/device/gpu/optimizer/add_relu_v2_fusion.h"
+#include "plugin/device/gpu/optimizer/add_relu_grad_v2_fusion.h"
+#include "plugin/device/gpu/optimizer/matmul_biasadd_fusion.h"
+#include "plugin/device/gpu/optimizer/neighbor_exchange_v2_fusion.h"
+#ifdef ENABLE_GPU_INFER
+#include "plugin/device/gpu/optimizer/trt_pass/graph_converter.h"
+#endif
+#include "common/graph_kernel/adapter/graph_kernel_optimization.h"
+#include "backend/common/pass/communication_op_fusion.h"
+#include "plugin/device/gpu/optimizer/concat_outputs_for_all_gather.h"
+#include "backend/common/pass/getitem_tuple.h"
+#include "backend/common/pass/optimize_updatestate.h"
+#include "backend/common/pass/adjust_depend_for_parallel_optimizer_recompute_all_gather.h"
+#include "runtime/device/ms_device_shape_transfer.h"
+#include "include/common/debug/anf_ir_dump.h"
+#include "include/common/debug/dump_proto.h"
+#ifdef ENABLE_DEBUGGER
+#include "debug/data_dump/e2e_dump.h"
+#include "debug/data_dump/dump_json_parser.h"
+#include "debug/debugger/proto_exporter.h"
+#include "debug/data_dump/dump_utils.h"
+#include "debug/tensor_load.h"
+#else
+#include "debug/debugger/proto_exporter_stub.h"
+#endif
+#include "plugin/device/gpu/hal/device/gpu_kernel_build.h"
+#include "plugin/device/gpu/hal/device/gpu_kernel_runtime.h"
+#include "plugin/device/gpu/hal/device/gpu_stream_assign.h"
+#include "plugin/device/gpu/hal/device/kernel_info_setter.h"
+#include "runtime/device/kernel_runtime_manager.h"
+#include "plugin/device/gpu/hal/device/cuda_driver.h"
+#include "plugin/device/gpu/hal/device/distribution/collective_init.h"
+#include "plugin/device/gpu/hal/device/gpu_bucket.h"
+#include "plugin/device/gpu/hal/device/gpu_device_address.h"
+#include "utils/ms_utils.h"
+#include "include/common/utils/config_manager.h"
+#include "utils/ms_context.h"
+#include "common/graph_kernel/graph_kernel_flags.h"
+#include "include/common/utils/utils.h"
+#include "abstract/utils.h"
+#if ENABLE_CPU && ENABLE_GPU
+#include "ps/util.h"
+#include "ps/ps_cache/ps_cache_manager.h"
+#endif
+
+namespace mindspore {
+namespace session {
+namespace gpu {
+using AnfAlgo = mindspore::session::AnfRuntimeAlgorithm;
+using CollectiveInitializer = device::gpu::CollectiveInitializer;
+using GetLocalRankId = device::gpu::GetLocalRankId;
+using InitNCCLComm = device::gpu::InitNCCLComm;
+
+void GPUSession::Init(uint32_t device_id) {
+  if (CollectiveInitializer::instance().collective_inited()) {
+    device_id = CollectiveInitializer::instance().local_rank_id();
+  }
+  bool ret = device::gpu::CudaDriver::SetDevice(UintToInt(device_id));
+  if (!ret) {
+    MS_LOG(EXCEPTION) << "GPUSession failed to set current device id:" << device_id;
+  }
+  auto ms_context = MsContext::GetInstance();
+  MS_EXCEPTION_IF_NULL(ms_context);
+  ms_context->set_param<uint32_t>(MS_CTX_DEVICE_ID, device_id);
+  if (CollectiveInitializer::instance().collective_inited()) {
+    auto collective_handle = CollectiveInitializer::instance().collective_handle();
+    if (collective_handle != nullptr) {
+      MS_LOG(INFO) << "Start initializing NCCL communicator for device " << device_id;
+      auto init_nccl_comm_funcptr =
+        reinterpret_cast<InitNCCLComm>(dlsym(const_cast<void *>(collective_handle), "InitNCCLComm"));
+      MS_EXCEPTION_IF_NULL(init_nccl_comm_funcptr);
+      (*init_nccl_comm_funcptr)();
+      MS_LOG(INFO) << "End initializing NCCL communicator.";
+      rank_id_ = GetRankId();
+    }
+  }
+#ifndef ENABLE_SECURITY
+  auto &json_parser = DumpJsonParser::GetInstance();
+  // Dump json config file if dump is enabled for GPU old runtime.
+  json_parser.CopyDumpJsonToDir(rank_id_);
+  json_parser.CopyMSCfgJsonToDir(rank_id_);
+#endif
+  MS_LOG(INFO) << "Set device id " << device_id << " for gpu session.";
+  InitExecutor(kGPUDevice, device_id);
+}
+
+void GPUSession::SelectKernel(const std::shared_ptr<KernelGraph> &kernel_graph) const {
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  device::gpu::FormatTransformChecker::GetInstance().CheckSupportFormatTransform(kernel_graph);
+  for (const auto &kernel_node : kernel_graph->execution_order()) {
+    MS_EXCEPTION_IF_NULL(kernel_node);
+    device::gpu::SetKernelInfo(kernel_node);
+  }
+}
+
+void GPUSession::StartKernelRT() const {
+  auto runtime_instance = device::KernelRuntimeManager::Instance().GetSingleKernelRuntime(kGPUDevice, device_id_);
+  MS_EXCEPTION_IF_NULL(runtime_instance);
+  if (!runtime_instance->Init()) {
+    MS_LOG(EXCEPTION) << "GPU start kernel runtime failed";
+  }
+}
+
+void GPUSession::Optimize(const std::shared_ptr<KernelGraph> &kernel_graph) {
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  auto optimizer = std::make_shared<opt::GraphOptimizer>();
+  auto pm = std::make_shared<opt::PassManager>();
+#ifdef ENABLE_GPU_INFER
+  pm->AddPass(std::make_shared<opt::GraphConverter>());
+#endif
+  pm->AddPass(std::make_shared<opt::MatMulBiasAddFusion>());
+  pm->AddPass(std::make_shared<opt::AdamWeightDecayFusion>());
+  pm->AddPass(std::make_shared<opt::AdamFusion>());
+  pm->AddPass(std::make_shared<opt::AllToAllFusion>());
+  pm->AddPass(std::make_shared<opt::ApplyMomentumWeightDecayScaleFusion>());
+  pm->AddPass(std::make_shared<opt::ApplyMomentumScaleFusion>());
+  pm->AddPass(std::make_shared<opt::ApplyMomentumWeightDecayFusion>());
+  if (!graphkernel::GraphKernelFlags::GetInstance().IsEnableGraphKernel()) {
+    pm->AddPass(std::make_shared<opt::CastAllFusion>("cast_all"));
+  }
+  pm->AddPass(std::make_shared<opt::CombineMomentumFusion>("combine_momentum"));
+  pm->AddPass(std::make_shared<opt::ReplaceMomentumCastFusion>());
+  pm->AddPass(std::make_shared<opt::ReplaceAddNFusion>());
+  pm->AddPass(std::make_shared<opt::PrintReduceFusion>("print_reduce"));
+  pm->AddPass(std::make_shared<opt::BCEWithLogitsLossFusion>());
+  pm->AddPass(std::make_shared<opt::InsertCastGPU>("insert_cast_gpu"));
+  pm->AddPass(std::make_shared<opt::NeighborExchangeV2Fusion>());
+  pm->AddPass(std::make_shared<opt::NeighborExchangeV2GradFusion>());
+  optimizer->AddPassManager(pm);
+  (void)optimizer->Optimize(kernel_graph);
+  kernel_graph->SetExecOrderByDefault();
+}
+
+void GPUSession::HardwareOptimize(const std::shared_ptr<KernelGraph> &kernel_graph) {
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  auto optimizer = std::make_shared<opt::GraphOptimizer>();
+  auto pm = std::make_shared<opt::PassManager>();
+  pm->AddPass(std::make_shared<opt::BatchNormReluFusion>());
+  pm->AddPass(std::make_shared<opt::BatchNormReluGradFusion>());
+  pm->AddPass(std::make_shared<opt::BatchNormAddReluFusion>());
+  pm->AddPass(std::make_shared<opt::PostBatchNormAddReluFusion>());
+  pm->AddPass(std::make_shared<opt::BatchNormAddReluGradFusion>());
+  pm->AddPass(std::make_shared<opt::InsertFormatTransformOp>());
+  pm->AddPass(std::make_shared<opt::RemoveFormatTransformPair>());
+  pm->AddPass(std::make_shared<opt::RemoveRedundantFormatTransform>());
+  // Remove node only used by UpdateState, in order to ensure the correct execution sequence in CudnnInplaceAggregate.
+  pm->AddPass(std::make_shared<opt::OptimizeUpdateState>());
+  pm->AddPass(std::make_shared<opt::CudnnInplaceAggregate>());
+  pm->AddPass(std::make_shared<opt::ReluV2Pass>());
+  pm->AddPass(std::make_shared<opt::AddReluV2Fusion>());
+  pm->AddPass(std::make_shared<opt::AddReluGradV2Fusion>());
+  pm->AddPass(std::make_shared<opt::AllReduceFusion>());
+  pm->AddPass(std::make_shared<opt::AdjustDependForParallelOptimizerRecomputeAllGather>());
+  pm->AddPass(std::make_shared<opt::AllGatherFusion>());
+  pm->AddPass(std::make_shared<opt::ConcatOutputsForAllGather>());
+  pm->AddPass(std::make_shared<opt::GetitemTuple>());
+  pm->AddPass(std::make_shared<opt::ReducePrecisionFusion>("reduce_precision"));
+  optimizer->AddPassManager(pm);
+  (void)optimizer->Optimize(kernel_graph);
+  kernel_graph->SetExecOrderByDefault();
+}
+
+void GPUSession::RunOpOptimize(const std::shared_ptr<KernelGraph> &kernel_graph) {
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  auto optimizer = std::make_shared<opt::GraphOptimizer>();
+  auto pm = std::make_shared<opt::PassManager>();
+  pm->AddPass(std::make_shared<opt::BCEWithLogitsLossFusion>());
+  pm->AddPass(std::make_shared<opt::InsertCastGPU>("insert_cast_gpu"));
+  optimizer->AddPassManager(pm);
+  (void)optimizer->Optimize(kernel_graph);
+  kernel_graph->SetExecOrderByDefault();
+}
+
+void GPUSession::RunOpHardwareOptimize(const std::shared_ptr<KernelGraph> &kernel_graph) {
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  auto optimizer = std::make_shared<opt::GraphOptimizer>();
+  auto pm = std::make_shared<opt::PassManager>();
+  pm->AddPass(std::make_shared<opt::ReducePrecisionFusion>("reduce_precision"));
+  optimizer->AddPassManager(pm);
+  (void)optimizer->Optimize(kernel_graph);
+  kernel_graph->SetExecOrderByDefault();
+}
+
+void GPUSession::GraphKernelOptimize(const std::shared_ptr<KernelGraph> &kernel_graph) {
+  if (!graphkernel::GraphKernelFlags::GetInstance().IsEnableGraphKernel()) {
+    return;
+  }
+  graphkernel::GraphKernelOptimize(kernel_graph);
+  kernel_graph->SetExecOrderByDefault();
+}
+
+void GPUSession::AssignStream(const std::shared_ptr<KernelGraph> &kernel_graph) {
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  device::gpu::AssignGpuStream(kernel_graph);
+}
+
+void GPUSession::BuildKernel(const std::shared_ptr<KernelGraph> &kernel_graph) const {
+  auto kernels = kernel_graph->execution_order();
+  device::gpu::CreateGPUKernel(kernels);
+}
+
+void GPUSession::AllocateMemory(const KernelGraph *kernel_graph) const {
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  auto runtime_instance = device::KernelRuntimeManager::Instance().GetSingleKernelRuntime(kGPUDevice, device_id_);
+  MS_EXCEPTION_IF_NULL(runtime_instance);
+  runtime_instance->AssignMemory(*kernel_graph);
+}
+
+void GPUSession::RunOpAllocateMemory(const std::vector<tensor::TensorPtr> &input_tensors,
+                                     const KernelGraph *kernel_graph, bool is_gradient_out) const {
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  auto runtime_instance = device::KernelRuntimeManager::Instance().GetSingleKernelRuntime(kGPUDevice, device_id_);
+  MS_EXCEPTION_IF_NULL(runtime_instance);
+  runtime_instance->RunOpAssignMemory(input_tensors, *kernel_graph, is_gradient_out);
+}
+
+void GPUSession::RunOpGenKernelEvent(const KernelGraph *graph) const {
+  MS_EXCEPTION_IF_NULL(graph);
+  auto runtime_instance = device::KernelRuntimeManager::Instance().GetSingleKernelRuntime(kGPUDevice, device_id_);
+  MS_EXCEPTION_IF_NULL(runtime_instance);
+  runtime_instance->GenKernelEvents(*graph);
+}
+
+void GPUSession::RunOpClearMemory(const KernelGraph *kernel_graph) const {
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  auto runtime_instance = device::KernelRuntimeManager::Instance().GetSingleKernelRuntime(kGPUDevice, device_id_);
+  MS_EXCEPTION_IF_NULL(runtime_instance);
+  runtime_instance->RunOpClearMemory(*kernel_graph);
+}
+
+namespace {
+constexpr auto kAssignInputSize = 3;
+constexpr auto kAssignUpdateIndex = 1;
+bool UpdatedByAssign(const KernelGraphPtr &kernel_graph, const AnfNodePtr &node) {
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  auto manager = kernel_graph->manager();
+  if (manager == nullptr) {
+    return false;
+  }
+  auto &node_users = manager->node_users();
+  auto iter = node_users.find(node);
+  if (iter == node_users.end()) {
+    return false;
+  }
+  auto &users = iter->second;
+  return std::any_of(users.begin(), users.end(), [](const std::pair<AnfNodePtr, int64_t> &user) {
+    MS_EXCEPTION_IF_NULL(user.first);
+    auto output_cnode = user.first->cast<CNodePtr>();
+    return output_cnode != nullptr && IsPrimitiveCNode(output_cnode, prim::kPrimAssign) &&
+           user.second == kAssignUpdateIndex && output_cnode->inputs().size() > kAssignInputSize;
+  });
+}
+
+size_t UpdateGraphInputAbstract(const AnfNodePtr input_node, const tensor::TensorPtr tensor) {
+  MS_EXCEPTION_IF_NULL(input_node);
+  MS_EXCEPTION_IF_NULL(tensor);
+  size_t size = LongToSize(tensor->data().nbytes());
+  if (!input_node->isa<Parameter>()) {
+    return size;
+  }
+  auto input_param = input_node->cast<ParameterPtr>();
+  if (input_param != nullptr && input_param->has_dynamic_shape()) {
+    auto tensor_shape = tensor->shape();
+    std::vector<size_t> shape_tmp;
+    (void)std::transform(tensor_shape.begin(), tensor_shape.end(), std::back_inserter(shape_tmp), LongToSize);
+    common::AnfAlgo::SetOutputInferTypeAndShape({common::AnfAlgo::GetOutputInferDataType(input_node, 0)}, {shape_tmp},
+                                                input_node.get());
+    size = abstract::ShapeSize(shape_tmp) * abstract::TypeIdSize(tensor->data_type());
+  }
+  return size;
+}
+
+bool CheckIfNeedSync(const tensor::TensorPtr &tensor, const DeviceAddressPtr &device_address,
+                     const ParameterPtr &pk_node) {
+  MS_EXCEPTION_IF_NULL(tensor);
+  MS_EXCEPTION_IF_NULL(pk_node);
+  auto tensor_address = std::dynamic_pointer_cast<device::DeviceAddress>(tensor->device_address());
+  bool need_sync = false;
+  auto ms_context = MsContext::GetInstance();
+  MS_EXCEPTION_IF_NULL(ms_context);
+  if (ms_context->get_param<bool>(MS_CTX_ENABLE_PYNATIVE_INFER)) {
+    if (tensor_address == nullptr || tensor_address != device_address) {
+      need_sync = true;
+    }
+  } else if (tensor->NeedSyncHostToDevice() || tensor_address == nullptr) {
+    need_sync = true;
+  } else if (tensor_address != device_address) {
+    if (tensor_address->DeviceType() == device_address->DeviceType()) {
+      AnfAlgo::SetOutputAddr(tensor_address, 0, pk_node.get());
+    } else {
+      need_sync = true;
+    }
+  }
+  return need_sync;
+}
+}  // namespace
+
+void GPUSession::LoadInputData(const std::shared_ptr<KernelGraph> &kernel_graph,
+                               const std::vector<tensor::TensorPtr> &inputs_const) const {
+  std::vector<tensor::TensorPtr> inputs(inputs_const);
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  auto &input_nodes = kernel_graph->input_nodes();
+  auto ms_context = MsContext::GetInstance();
+  MS_EXCEPTION_IF_NULL(ms_context);
+  if (inputs.size() != input_nodes.size()) {
+    MS_LOG(EXCEPTION) << "Tensor input:" << inputs.size() << " is not equal graph inputs:" << input_nodes.size();
+  }
+  for (size_t i = 0; i < inputs.size(); ++i) {
+    auto tensor = inputs[i];
+    MS_EXCEPTION_IF_NULL(tensor);
+    auto input_node = input_nodes[i];
+    MS_EXCEPTION_IF_NULL(input_node);
+    if (input_node->isa<Parameter>() && AnfAlgo::OutputAddrExist(input_node, 0)) {
+#if ENABLE_CPU && ENABLE_GPU
+      const std::string &param_name = input_node->fullname_with_scope();
+      if (ps::ps_cache_instance.IsHashTable(param_name)) {
+        continue;
+      }
+#endif
+      auto pk_node = input_node->cast<ParameterPtr>();
+      auto device_address = AnfAlgo::GetMutableOutputAddr(pk_node, 0);
+      MS_EXCEPTION_IF_NULL(device_address);
+      bool need_sync = CheckIfNeedSync(tensor, device_address, pk_node);
+      if (need_sync) {
+        if (common::AnfAlgo::IsParameterWeight(pk_node) || UpdatedByAssign(kernel_graph, input_node) ||
+            ms_context->get_param<int>(MS_CTX_EXECUTION_MODE) == kPynativeMode) {
+          tensor->set_device_address(device_address);
+        }
+        auto size = UpdateGraphInputAbstract(input_node, tensor);
+        if (!device_address->SyncHostToDevice(trans::GetRuntimePaddingShape(pk_node, 0), size, tensor->data_type(),
+                                              tensor->data_c())) {
+          MS_LOG(EXCEPTION) << "SyncHostToDevice failed.";
+        }
+        if (kernel_graph->IsUpdatedParameter(pk_node)) {
+          tensor->SetIsUpdateByDevice();
+        }
+      }
+    }
+    tensor->set_sync_status(kNoNeedSync);
+  }
+}
+
+GraphId GPUSession::CompileGraphImpl(const AnfNodePtrList &lst, const AnfNodePtrList &outputs) {
+  // Construct graph, if successfully, graph_sum_ + 1
+  auto graph = ConstructKernelGraph(lst, outputs, DeviceAddressType::kGPU);
+  MS_EXCEPTION_IF_NULL(graph);
+  return CompileGraphImpl(graph);
+}
+
+GraphId GPUSession::CompileGraphImpl(NotNull<FuncGraphPtr> func_graph) {
+  std::vector<KernelGraphPtr> all_graphs;
+  auto root_graph = ConstructKernelGraph(func_graph, &all_graphs, DeviceAddressType::kGPU);
+  MS_EXCEPTION_IF_NULL(root_graph);
+  if (all_graphs.size() != 1) {
+    MS_LOG(EXCEPTION) << "Gpu backend does not support multi-graph schedule, graph num is " << all_graphs.size();
+  }
+  // Insert maketuple graph output in case of multi-outputs.
+  // The ConvertTupleOutputToMaketuple pass will insert TupleGetItem.
+  AnfAlgo::InsertMakeTupleForOutput(NOT_NULL(root_graph));
+  opt::BackendCommonOptimization(root_graph);
+  return CompileGraphImpl(root_graph);
+}
+
+GraphId GPUSession::CompileGraphImpl(const KernelGraphPtr &graph) {
+  MS_EXCEPTION_IF_NULL(graph);
+  // Prepare ms context info for dump .pb graph for GPU old runtime.
+  auto context_ptr = MsContext::GetInstance();
+  MS_EXCEPTION_IF_NULL(context_ptr);
+  auto runtime_instance = device::KernelRuntimeManager::Instance().GetSingleKernelRuntime(kGPUDevice, device_id_);
+  MS_EXCEPTION_IF_NULL(runtime_instance);
+#ifndef ENABLE_SECURITY
+  auto &json_parser = DumpJsonParser::GetInstance();
+  json_parser.Parse();
+#endif
+#ifdef ENABLE_DUMP_IR
+  bool save_graphs = context_ptr->get_param<bool>(MS_CTX_SAVE_GRAPHS_FLAG);
+  // Dump .pb graph before graph optimization
+  if (save_graphs) {
+    DumpIRProto(graph, "before_opt_" + std::to_string(graph->graph_id()));
+  }
+#endif
+  // Graph optimization irrelevant to device data format
+  Optimize(graph);
+  // Select kernel build info
+  SelectKernel(graph);
+  // Graph optimization relevant to device data format
+  HardwareOptimize(graph);
+  // Run final optimization
+  FinalOptimize(graph);
+  // Graph kernel fusion optimization
+  GraphKernelOptimize(graph);
+  // Start gpu kernel runtime
+  StartKernelRT();
+#if ENABLE_CPU && ENABLE_GPU
+  InitPsWorker(graph);
+#endif
+  // Assign CUDA streams
+  AssignStream(graph);
+#ifdef ENABLE_DUMP_IR
+  // Dump .pb graph before remove nop nodes
+  if (save_graphs) {
+    DumpIRProto(graph, "before_removeNop_" + std::to_string(graph->graph_id()));
+  }
+#endif
+  opt::AddDynamicShapeAttrPass(graph);
+  const bool pynative_mode = context_ptr->get_param<int>(MS_CTX_EXECUTION_MODE) == kPynativeMode;
+  // Hide NopOp from execution graph in graph mode
+  if (!pynative_mode) {
+    opt::HideNopNode(graph.get());
+  }
+  // Build kernel if node is cnode
+  BuildKernel(graph);
+#ifndef ENABLE_SECURITY
+  // Get summary nodes.
+  SetSummaryNodes(graph.get());
+#endif
+  // Dump .pb graph after graph optimization
+#ifdef ENABLE_DUMP_IR
+  if (save_graphs) {
+    DumpIRProto(graph, "after_opt_" + std::to_string(graph->graph_id()));
+  }
+#endif
+#ifndef ENABLE_SECURITY
+  // GPU old runtime.
+  if (json_parser.e2e_dump_enabled()) {
+    graph->set_root_graph_id(graph->graph_id());
+    std::string final_graph = "trace_code_graph_" + std::to_string(graph->graph_id());
+    std::string root_dir = json_parser.path() + "/rank_" + std::to_string(rank_id_);
+    std::string target_dir = root_dir + "/graphs";
+    std::string ir_file_path = target_dir + "/" + "ms_output_" + final_graph + ".ir";
+    DumpIRProtoWithSrcInfo(graph, final_graph, target_dir, kDebugWholeStack);
+    DumpIR("trace_code_graph", graph, true, kWholeStack, ir_file_path);
+    DumpGraphExeOrder("ms_execution_order_graph_" + std::to_string(graph->graph_id()) + ".csv", root_dir,
+                      graph->execution_order());
+  }
+#endif
+  // Set graph manager.
+  MS_EXCEPTION_IF_NULL(context_);
+  FuncGraphManagerPtr manager = MakeManager({graph});
+  context_->AddManager(manager);
+  if (manager) {
+    manager->AddFuncGraph(graph);
+    graph->set_manager(manager);
+  }
+
+  InitAllBucket(graph);
+  // Alloc memory in graph mode, including static memory and dynamic memory
+  if (!pynative_mode) {
+    AllocateMemory(graph.get());
+  }
+
+  DumpGraphs({graph});
+
+#ifdef ENABLE_DEBUGGER
+  if (debugger_ && debugger_->DebuggerBackendEnabled()) {
+    debugger_->LoadGraphs(graph);
+  }
+#endif
+  MS_LOG(INFO) << "CompileGraph graph_id: " << graph->graph_id();
+  return graph->graph_id();
+}
+
+// GPU old runtime.
+void GPUSession::PreExecuteGraph(const std::shared_ptr<KernelGraph> &kernel_graph,
+                                 const std::vector<tensor::TensorPtr> &inputs, VectorRef *outputs) {
+#ifdef ENABLE_DEBUGGER
+  if (debugger_) {
+    debugger_->PreExecute(kernel_graph);
+  }
+
+  E2eDump::UpdateIterOldRTDump(kernel_graph.get());
+#endif
+
+#if ENABLE_CPU && ENABLE_GPU
+  // Initialize parameter server
+  InitPSParamAndOptim(kernel_graph, inputs);
+#endif
+}
+
+// GPU old runtime.
+void GPUSession::PostExecuteGraph(const std::shared_ptr<KernelGraph> &kernel_graph,
+                                  const std::vector<tensor::TensorPtr> &inputs, VectorRef *outputs) {
+  // Summary
+  auto context_ptr = MsContext::GetInstance();
+  MS_EXCEPTION_IF_NULL(context_ptr);
+#ifndef ENABLE_SECURITY
+  if (context_ptr->get_param<bool>(MS_CTX_ENABLE_GPU_SUMMARY)) {
+    Summary(kernel_graph.get());
+  }
+#endif
+#ifdef ENABLE_DEBUGGER
+  if (debugger_ && debugger_->DebuggerBackendEnabled()) {
+    debugger_->LoadParametersAndConst(kernel_graph);
+  }
+
+  // debug used for dump
+  if (debugger_ && debugger_->CheckDebuggerDumpEnabled()) {
+    Dump(kernel_graph);
+  }
+
+  if (debugger_) {
+    debugger_->PostExecute();
+  }
+#endif
+}
+
+void GPUSession::ExecuteGraph(const std::shared_ptr<KernelGraph> &kernel_graph) {
+  int kernel_num = kernel_graph->execution_order().size();
+  int64_t loopsize = (kernel_num > 1) ? ConfigManager::GetInstance().gpu_loopsink_size() : 1;
+  for (int64_t i = 0; i < loopsize; i++) {
+#if ENABLE_CPU && ENABLE_GPU
+    std::string channel_name;
+    if (ps::PsDataPrefetch::GetInstance().cache_enable() && IsGetNextGraph(kernel_graph, &channel_name)) {
+      ps::ps_cache_instance.IncreaseGraphStep(channel_name);
+    }
+#endif
+    Execute(kernel_graph);
+  }
+}
+
+void GPUSession::UpdateOutputTensors(const VectorRef *outputs,
+                                     const std::map<tensor::TensorPtr, session::KernelWithIndex> &tensor_to_node,
+                                     std::map<DeviceAddressPtr, DeviceAddressPtr> *new_to_old_device_address) {
+  MS_EXCEPTION_IF_NULL(outputs);
+  for (const auto &item : *outputs) {
+    if (utils::isa<VectorRefPtr>(item)) {
+      const auto &vector_ref = utils::cast<VectorRef>(item);
+      UpdateOutputTensors(&vector_ref, tensor_to_node, new_to_old_device_address);
+    } else if (utils::isa<tensor::TensorPtr>(item)) {
+      const auto &tensor = utils::cast<tensor::TensorPtr>(item);
+      MS_EXCEPTION_IF_NULL(tensor);
+      const auto &iter = tensor_to_node.find(tensor);
+      if (iter != tensor_to_node.end()) {
+        const auto &node = iter->second.first;
+        const auto &output_index = iter->second.second;
+        MS_EXCEPTION_IF_NULL(node);
+        // When the parameter does not have a user in the graph and is used as an output, the device address is null,
+        // and there is no need to set the device address for tensor.
+        if (!AnfAlgo::OutputAddrExist(node, output_index, true)) {
+          continue;
+        }
+        auto address = AnfAlgo::GetMutableOutputAddr(node, output_index);
+        // The outputs may have the same tensor, so need skip when the tensor has been set to device address.
+        if ((address == nullptr) || (address->GetPtr() == nullptr)) {
+          // If the device address in the node is invalid, you need to find out whether there is a corresponding
+          // device address in the new to old device address map to check whether the device address in the node
+          // has been replaced with a new one.
+          if ((*new_to_old_device_address).find(address) != (*new_to_old_device_address).end()) {
+            address = (*new_to_old_device_address)[address];
+          } else {
+            continue;
+          }
+        }
+        tensor->set_device_address(address);
+
+        // When the device address of graph cnode output is set in tensor, the graph output need be set new device
+        // address, to avoid that the device address context of tensor be rewritten in the next step or next loop.
+        // But one time memory application scenarios need to be skipped, because the memory is not allocated next step:
+        // 1. Non cnode 2. Communication kernel.
+        bool ps_mode = false;
+#if ((defined ENABLE_CPU) && (!defined _WIN32))
+        ps_mode = ps::PSContext::instance()->is_ps_mode();
+#endif
+        if (node->isa<CNode>() && !common::AnfAlgo::IsCommunicationOp(node) && !ps_mode) {
+          auto new_address = std::make_shared<device::gpu::GPUDeviceAddress>(nullptr, address->GetSize());
+          // If a nop node is output, its previous node should be set.
+          if (common::AnfAlgo::IsNopNode(node)) {
+            auto pre_node = common::AnfAlgo::GetPrevNodeOutput(node, 0, true);
+            if (!pre_node.first->isa<Parameter>()) {
+              AnfAlgo::SetOutputAddr(new_address, pre_node.second, pre_node.first.get());
+            }
+          } else {
+            AnfAlgo::SetOutputAddr(new_address, output_index, node.get());
+          }
+          (*new_to_old_device_address)[new_address] = address;
+          if (graphkernel::GraphKernelFlags::GetInstance().IsEnableGraphKernel()) {
+            auto runtime_instance =
+              device::KernelRuntimeManager::Instance().GetSingleKernelRuntime(kGPUDevice, device_id_);
+            MS_EXCEPTION_IF_NULL(runtime_instance);
+            auto gpu_runtime_instance = dynamic_cast<device::gpu::GPUKernelRuntime *>(runtime_instance);
+            gpu_runtime_instance->SetAddrInvalid(address);
+          }
+        }
+
+        if (common::AnfAlgo::IsDynamicShape(node)) {
+          const auto &updated_shape = common::AnfAlgo::GetOutputInferShape(node, output_index);
+          ShapeVector int_shape;
+          std::transform(updated_shape.begin(), updated_shape.end(), std::back_inserter(int_shape), SizeToInt);
+          tensor->set_shape(int_shape);
+        }
+      }
+      if (tensor->NeedSyncDeviceToHostImmediately()) {
+        tensor->data_sync(false);
+        tensor->set_device_address(nullptr);
+        tensor->set_sync_status(kNeedSyncHostToDevice);
+      }
+    }
+  }
+}
+
+void GPUSession::Execute(const std::shared_ptr<KernelGraph> &kernel_graph) const {
+  auto runtime_instance = device::KernelRuntimeManager::Instance().GetSingleKernelRuntime(kGPUDevice, device_id_);
+  MS_EXCEPTION_IF_NULL(runtime_instance);
+  if (!runtime_instance->Run(*kernel_graph, false)) {
+    MS_LOG(EXCEPTION) << "GPU execute graph failed!";
+  }
+}
+
+KernelGraphPtr GPUSession::BuildOpImpl(const OpRunInfo &op_run_info, const GraphInfo &graph_info,
+                                       const std::vector<tensor::TensorPtr> &input_tensors,
+                                       const std::vector<int64_t> &tensors_mask) {
+  // Check if the graph cache exists.
+  auto it = run_op_graphs_.find(graph_info);
+  if (it != run_op_graphs_.end() && kOpCacheBlackList.find(op_run_info.op_name) == kOpCacheBlackList.end()) {
+    return it->second;
+  }
+
+  // Prepare the graph
+  const auto &kernel_graph = ConstructSingleOpGraph(op_run_info, input_tensors, tensors_mask);
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  RunOpOptimize(kernel_graph);
+  SelectKernel(kernel_graph);
+  RunOpHardwareOptimize(kernel_graph);
+  StartKernelRT();
+  RunOpHideNopNode(kernel_graph);
+  BuildKernel(kernel_graph);
+  auto enable_op_graph_cache = MsContext::GetInstance()->get_param<bool>(MS_CTX_ENABLE_PYNATIVE_OP_GRAPH_CACHE);
+  if (enable_op_graph_cache) {
+    run_op_graphs_[graph_info] = kernel_graph;
+  }
+  return kernel_graph;
+}
+
+void GPUSession::RunOpImplOrigin(const GraphInfo &graph_info, OpRunInfo *op_run_info,
+                                 std::vector<tensor::TensorPtr> *input_tensors, VectorRef *outputs,
+                                 const std::vector<int64_t> &tensors_mask) {
+  RunOpImpl(graph_info, op_run_info, input_tensors, outputs, tensors_mask);
+}
+
+void GPUSession::RunOpImpl(const GraphInfo &graph_info, OpRunInfo *op_run_info,
+                           std::vector<tensor::TensorPtr> *input_tensors, VectorRef *outputs,
+                           const std::vector<int64_t> &tensors_mask) {
+  MS_EXCEPTION_IF_NULL(input_tensors);
+  MS_EXCEPTION_IF_NULL(op_run_info);
+  ProcessInputTensorsForHeterogeneous("GPU", *input_tensors);
+  const auto &kernel_graph = BuildOpImpl(*op_run_info, graph_info, *input_tensors, tensors_mask);
+  EraseValueNodeTensor(tensors_mask, input_tensors);
+  // wait for allreduce
+  for (auto &tensor : *input_tensors) {
+    MS_EXCEPTION_IF_NULL(tensor);
+    if (tensor->NeedWaitDevice()) {
+      tensor->WaitDevice();
+    }
+  }
+
+  // run op
+  MS_EXCEPTION_IF_NULL(kernel_graph);
+  RunOpRemoveNopNode(kernel_graph);
+  RunOpAllocateMemory(*input_tensors, kernel_graph.get(), op_run_info->is_gradient_out);
+  RunOpGenKernelEvent(kernel_graph.get());
+  // Execute the computation
+  LoadInputData(kernel_graph, *input_tensors);
+  Execute(kernel_graph);
+  // Fetch outputs
+  std::map<tensor::TensorPtr, session::KernelWithIndex> tensor_to_node;
+  UpdateOutputs(kernel_graph, outputs, *input_tensors, &tensor_to_node);
+  // update output abstract of dynamic op to op_run_info
+  if (op_run_info->is_dynamic_shape) {
+    UpdateOutputAbstract(kernel_graph, op_run_info);
+  }
+  RunOpClearMemory(kernel_graph.get());
+  if (kOpCacheBlackList.find(op_run_info->op_name) != kOpCacheBlackList.end()) {
+    run_op_graphs_.erase(graph_info);
+  }
+}
+
+#ifdef ENABLE_DEBUGGER
+
+void GPUSession::Dump(const std::shared_ptr<KernelGraph> &kernel_graph) const {
+  // Dump graph and graph history file if e2e_dump is enabled and update cur_dump_iter for GPU old runtime.
+  if (debugger_->DebuggerBackendEnabled()) {
+    MS_EXCEPTION_IF_NULL(kernel_graph);
+    E2eDump::DumpRunIter(kernel_graph, rank_id_);
+    E2eDump::DumpData(kernel_graph.get(), rank_id_, debugger_.get());
+  } else {
+    DumpJsonParser::GetInstance().UpdateDumpIter();
+  }
+}
+
+bool GPUSession::DumpDataEnabledIteration() const {
+  auto runtime_instance = device::KernelRuntimeManager::Instance().GetSingleKernelRuntime(kGPUDevice, device_id_);
+  MS_EXCEPTION_IF_NULL(runtime_instance);
+  return runtime_instance->DumpDataEnabledIteration();
+}
+#endif
+
+void GPUSession::SyncStream() const {
+  auto runtime_instance = device::KernelRuntimeManager::Instance().GetSingleKernelRuntime(kGPUDevice, device_id_);
+  MS_EXCEPTION_IF_NULL(runtime_instance);
+  auto ret = runtime_instance->SyncStream();
+  if (!ret) {
+    MS_LOG(EXCEPTION) << "Sync stream error!";
+  }
+}
+
+std::shared_ptr<device::Bucket> GPUSession::CreateBucket(uint32_t bucket_id, uint32_t bucket_size) {
+  auto bucket = std::make_shared<device::gpu::GPUBucket>(bucket_id, bucket_size);
+
+  auto kernel_runtime = device::KernelRuntimeManager::Instance().GetCurrentKernelRuntime();
+  MS_EXCEPTION_IF_NULL(kernel_runtime);
+  auto compute_stream = kernel_runtime->compute_stream();
+  auto communication_stream = kernel_runtime->communication_stream();
+  MS_EXCEPTION_IF_NULL(compute_stream);
+  MS_EXCEPTION_IF_NULL(communication_stream);
+
+  MS_EXCEPTION_IF_NULL(bucket);
+  bucket->Init({compute_stream}, {communication_stream});
+  return bucket;
+}
+}  // namespace gpu
+}  // namespace session
+}  // namespace mindspore

mindspore/ccsrc/backend/common/somas/somas_solver_alg.cc

@@ -63,84 +63,121 @@ size_t FootPrint::Result() {
 
   return upperbound;
 }
+// Class FootPrint manages memory allocations. The function `findFirst` finds the first suitable
+// offset in the given merged intervals to allocate a BlockTensor, and `Merge` merges overlapping intervals.
+
+// Function: findFirst
+// Purpose: This function aims to find the first available offset within the provided merged intervals to allocate a BlockTensor.
+// If a suitable offset is found, it updates the offset and returns true; otherwise, returns false.
+// Params:
+// - merged: A pointer to a stack of merged intervals representing available memory spaces.
+// - block: The BlockTensor object representing the memory block to be allocated.
+// - offset: A pointer to a size_t variable where the function will store the found offset.
+// Return: A boolean value indicating whether a suitable offset has been found.
 bool FootPrint::findFirst(stack<Interval> *merged, const BlockTensor &block, size_t *offset) {
-  MS_EXCEPTION_IF_NULL(merged);
-  MS_EXCEPTION_IF_NULL(offset);
-  bool bfound = false;
-  std::set<pair<size_t, size_t>, bool (*)(const pair<size_t, size_t> &a, const pair<size_t, size_t> &b)>
-    offsetcandidates(g_pBranching[m_branching_strategy_]);
-  size_t gap;
-
-  Interval a;
-  Interval it;
-
-  size_t block_size;
-  if (block.Alone()) {
-    block_size = block.m_size_;
-  } else {
-    block_size = block.m_start_tensor_->size_;  // consider only first tensor for contiguous block
-  }
-
-  a.ub() = algorithm[m_algorithm_](this);
-  while (!(*merged).empty()) {
-    it = (*merged).top();
-    (*merged).pop();
-    a.lb() = it.ub();
-    if (a.contains(block_size) && a.lb() + block.m_size_ <= algorithm[m_algorithm_](this)) {
-      gap = a.ub() - a.lb() - block_size;
-      offsetcandidates.emplace(pair<size_t, size_t>(a.lb(), gap));
+    MS_EXCEPTION_IF_NULL(merged);
+    MS_EXCEPTION_IF_NULL(offset);
+    bool bfound = false;
+    
+    // Initialize a set to hold candidate offsets, sorted by the given strategy.
+    std::set<pair<size_t, size_t>, bool (*)(const pair<size_t, size_t> &a, const pair<size_t, size_t> &b)>
+        offsetcandidates(g_pBranching[m_branching_strategy_]);
+    
+    size_t gap;
+    Interval a;
+    Interval it;
+    
+    // Calculate the size of the block to be allocated.
+    size_t block_size;
+    if (block.Alone()) {
+        block_size = block.m_size_;
+    } else {
+        block_size = block.m_start_tensor_->size_;  // consider only first tensor for contiguous block
     }
-    a.ub() = it.lb();
-  }
-  a.lb() = m_offset_;
-  if (a.contains(block_size) && a.lb() + block.m_size_ <= algorithm[m_algorithm_](this)) {
-    gap = a.ub() - a.lb() - block_size;
-    offsetcandidates.emplace(pair<size_t, size_t>(a.lb(), gap));
-  }
 
-  if (offsetcandidates.size() > 0) {
-    *offset = (*offsetcandidates.begin()).first;
-    m_foot_print_next_->m_offset_ = std::max(m_foot_print_next_->m_offset_, *offset + block.m_size_);
-    bfound = true;
-  }
+    // Iterate through the merged intervals and find candidate offsets where the block can fit.
+    a.ub() = algorithm[m_algorithm_](this);
+    while (!(*merged).empty()) {
+        it = (*merged).top();
+        (*merged).pop();
+        a.lb() = it.ub();
+        if (a.contains(block_size) && a.lb() + block.m_size_ <= algorithm[m_algorithm_](this)) {
+            gap = a.ub() - a.lb() - block_size;
+            offsetcandidates.emplace(pair<size_t, size_t>(a.lb(), gap));
+        }
+        a.ub() = it.lb();
+    }
+    
+    // If suitable candidates are found, update the offset and return true.
+    if (offsetcandidates.size() > 0) {
+        *offset = (*offsetcandidates.begin()).first;
+        m_foot_print_next_->m_offset_ = std::max(m_foot_print_next_->m_offset_, *offset + block.m_size_);
+        bfound = true;
+    }
 
-  return bfound;
+    return bfound;
 }
+
+// Function: Merge
+// Purpose: This function merges overlapping intervals from the input vector and stores the result in a stack.
+// Params:
+// - interval_v: A pointer to a vector of Interval objects representing the memory spaces to be merged.
+// - s: A pointer to a stack where the function will store the merged intervals.
+// Return: void.
 void FootPrint::Merge(vector<Interval> *interval_v, stack<Interval> *s) {
-  MS_EXCEPTION_IF_NULL(s);
-  MS_EXCEPTION_IF_NULL(interval_v);
-  sort((*interval_v).begin(), (*interval_v).end(),
-       [](Interval &i1, Interval &i2) { return (i1.lb() < i2.lb()) || (i1.lb() == i2.lb() && i1.ub() < i2.ub()); });
-  (*s).push((*interval_v)[0]);
-
-  for (size_t i = 1; i < (*interval_v).size(); i++) {
-    Interval &top = (*s).top();
-    Interval &b = (*interval_v)[i];
-    if (top.ub() < b.lb())
-      (*s).push(b);
-
-    else if (top.ub() < b.ub())
-      top.ub() = b.ub();
-  }
-
-  return;
+    MS_EXCEPTION_IF_NULL(s);
+    MS_EXCEPTION_IF_NULL(interval_v);
+    
+    // Sort the input intervals in ascending order for merging.
+    sort((*interval_v).begin(), (*interval_v).end(),
+         [](Interval &i1, Interval &i2) { return (i1.lb() < i2.lb()) || (i1.lb() == i2.lb() && i1.ub() < i2.ub()); });
+    
+    // Iterate through the sorted intervals and merge overlapping ones.
+    (*s).push((*interval_v)[0]);
+    for (size_t i = 1; i < (*interval_v).size(); i++) {
+        Interval &top = (*s).top();
+        Interval &b = (*interval_v)[i];
+        if (top.ub() < b.lb())
+            (*s).push(b);
+        else if (top.ub() < b.ub())
+            top.ub() = b.ub();
+    }
+    return;
 }
+/**
+ * @brief Find an appropriate offset for placing a BlockTensor in memory considering a set of constraints.
+ *
+ * This function attempts to find a suitable memory offset for a given BlockTensor. It takes into consideration
+ * the constraints defined against the already allocated tensors to find an appropriate location where the
+ * BlockTensor doesn’t violate any constraints.
+ *
+ * @param constraints A vector of DynamicBitSet representing the constraints between tensors.
+ * @param block The BlockTensor for which an offset needs to be found.
+ * @param offset A pointer to a size_t where the found offset will be stored.
+ * @return Returns true if a suitable offset is found, otherwise returns false.
+ */
 bool FootPrint::findOffset(const std::vector<DynamicBitSet> *constraints, const BlockTensor &block, size_t *offset) {
+  
+  // Check for null pointer and initialize local variables
   MS_EXCEPTION_IF_NULL(offset);
   bool bretval = true;
   vector<Interval> l_interval;
-
+  
+  // Reserve space for intervals estimation
   const size_t intervals_estimation = 1000;
   l_interval.reserve(intervals_estimation * sizeof(Interval));
-
+  
   *offset = m_offset_;
-
-  // transform constrained tensors in non eligible intervals
+  
+  // Transform constrained tensors into non-eligible intervals
   if (block.Alone()) {
+    // Handle the case when the block is alone and check specific conditions
     if (m_algorithm_ == static_cast<uint32_t>(kManyObjects) && m_starts_.size() > 0 && m_starts_[0]->Alone() &&
         (*constraints)[block.m_start_tensor_->index_].IsBitTrue(m_starts_[0]->m_start_tensor_->index_) == false) {
       return false;
     }
+    
+    // Iterate over allocated tensors and identify ineligible intervals
     for (size_t i = 0; i < m_starts_.size(); i++) {
       auto allocated_tensor = m_starts_[i]->m_start_tensor_;
       while (allocated_tensor != nullptr) {
@@ -151,6 +188,7 @@ bool FootPrint::findOffset(const std::vector<DynamicBitSet> *constraints, const
       }
     }
   } else {
+    // Handle the case when the block is not alone and compute ineligible intervals differently
     int64_t start_offset = static_cast<int64_t>(m_offset_);
     for (size_t i = 0; i < m_starts_.size(); i++) {
       auto allocated_tensor = m_starts_[i]->m_start_tensor_;
@@ -176,16 +214,17 @@ bool FootPrint::findOffset(const std::vector<DynamicBitSet> *constraints, const
       }
     }
   }
-
-  // merge non-eligible intervals and find a slot to allocate the tensor block
+  
+  // Merge non-eligible intervals and find a slot to allocate the tensor block
   if (!l_interval.empty()) {
     stack<Interval> l_mergedIntervals;
     Merge(&l_interval, &l_mergedIntervals);
     bretval = findFirst(&l_mergedIntervals, block, offset);
   }
-
+  
   return bretval;
 }
+
 void FootPrint::addElem(BlockTensor *block, const size_t &offset) {
   if (m_foot_print_next_ == nullptr) {
     m_foot_print_next_ = std::make_shared<FootPrint>();
@@ -218,18 +257,40 @@ void FootPrint::addElem(BlockTensor *block, const size_t &offset) {
 void FootPrint::printStats() {
   MS_LOG(DEBUG) << "Footprint blocks: " << m_starts_.size() << " \toffset: " << m_offset_;
 }
+
+/**
+ * @brief Evaluates and allocates memory for a series of BlockTensors.
+ *
+ * This function is a part of the FastHeuristic class and is primarily designed to evaluate 
+ * and allocate memory for a set of tensors represented by BlockTensors. The allocation is 
+ * carried out respecting any constraints represented by DynamicBitSet and considering the 
+ * current FootPrint of memory usage. This function is instrumental in optimizing memory 
+ * allocation for tensor blocks.
+ *
+ * @param block_tensors_v Pointer to a vector of BlockTensor, representing the tensor blocks
+ *                        to be allocated memory.
+ * @param foot_print Shared pointer to a FootPrint object representing the current memory footprint.
+ * @param pConstraints Pointer to a vector of DynamicBitSet representing allocation constraints.
+ * @return Returns true if memory allocation for all tensor blocks is successful, else returns false.
+ */
 bool FastHeuristic::Eval(vector<BlockTensor> *block_tensors_v, const std::shared_ptr<FootPrint> &foot_print,
                          const std::vector<DynamicBitSet> *pConstraints) {
+  // Ensure the foot_print pointer is not null.
   MS_EXCEPTION_IF_NULL(foot_print);
+  
+  // Initialize the start time for evaluating the execution time of this method.
   auto start = std::chrono::system_clock::now();
 
+  // Initialize necessary variables.
   std::shared_ptr<FootPrint> p = foot_print;
   bool bpushed = false;
   size_t offset = foot_print->getOffset();
   m_tensors_allocated_ = 0;
   SomasSolverTensorDescPtr tensor = nullptr;
 
+  // Iterate through each tensor block in block_tensors_v.
   for (auto &block : *block_tensors_v) {
+    // If the block does not need reallocation, update its offsets and continue to the next block.
     if (!block.m_bre_allocate_) {
       offset = block.m_start_tensor_->offset_;
       auto aux_id = foot_print->m_solId_;
@@ -241,13 +302,20 @@ bool FastHeuristic::Eval(vector<BlockTensor> *block_tensors_v, const std::shared
       (void)block.offsets_.emplace(aux_id, aux_offset);
       continue;
     }
+    
+    // Reset bpushed and assign the current foot print to p.
     bpushed = false;
     p = foot_print;
     block.m_current_sol_ = foot_print->m_solId_;
+    
+    // Attempt to find a suitable offset in the memory for the current block.
     while (!bpushed) {
       if (p->findOffset(pConstraints, block, &offset)) {
+        // If a suitable offset is found, add the block to the current foot print.
         p->addElem(&block, offset);
         tensor = block.m_start_tensor_;
+        
+        // Update the count of tensors allocated.
         while (tensor) {
           m_tensors_allocated_++;
           tensor = tensor->right_;
@@ -255,19 +323,24 @@ bool FastHeuristic::Eval(vector<BlockTensor> *block_tensors_v, const std::shared
         bpushed = true;
         break;
       }
-      // go to the next footprint slot
+      
+      // If not found, proceed to the next available foot print slot.
       if (p->Next() != nullptr) {
         p = p->Next();
-      } else if (bpushed == false) {  // something went wrong
+      } else if (bpushed == false) {
+        // If the allocation fails, log a warning and return false.
         MS_LOG(WARNING) << "Internal Error: Could not allocate memory for tensor: " << tensor->index_;
         return false;
       }
     }
   }
-
+  
+  // Log the elapsed time of the Fast Heuristic search.
   MS_LOG(DEBUG)
     << "\nElapsed time of Fast Heuristic search: "
     << std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now() - start).count() << " ms";
+  
+  // Return true indicating successful memory allocation for all tensor blocks.
   return true;
 }
 }  // namespace somas

mindspore/ccsrc/backend/common/somas/somas_solver_alg_old.cc

@@ -0,0 +1,347 @@
+/**
+ * Copyright 2020-2021 Huawei Technologies Co., Ltd
+
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+
+ * http://www.apache.org/licenses/LICENSE-2.0
+
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+*/
+
+#include "backend/common/somas/somas_solver_alg.h"
+
+#include <algorithm>
+#include <stack>
+#include <utility>
+
+namespace mindspore {
+namespace somas {
+// offset picking heuristics
+bool SmallestFit(const pair<size_t, size_t> &a, const pair<size_t, size_t> &b) {
+  return a.first < b.first || (a.first == b.first && a.second < b.second);
+}
+bool LargestFit(const pair<size_t, size_t> &a, const pair<size_t, size_t> &b) {
+  return a.first > b.first || (a.first == b.first && a.second < b.second);
+}
+bool BestFit(const pair<size_t, size_t> &a, const pair<size_t, size_t> &b) {
+  return a.second < b.second || (a.second == b.second && a.first < b.first);
+}
+bool WorstFit(const pair<size_t, size_t> &a, const pair<size_t, size_t> &b) {
+  return a.second > b.second || (a.second == b.second && a.first < b.first);
+}
+size_t SharedObjects(FootPrint *p) { return p->Next()->getOffset(); }
+size_t SingleObject(FootPrint *p) { return SIZE_MAX; }
+
+bool (*g_pBranching[kNumFittingTypes])(const pair<size_t, size_t> &a, const pair<size_t, size_t> &b) = {
+  BestFit, SmallestFit
+#ifdef SOMAS_DEBUG
+  ,
+  LargestFit, WorstFit
+#endif
+};
+size_t (*algorithm[kNumAlgorithmTypes])(FootPrint *p) = {SharedObjects, SingleObject};
+
+size_t FootPrint::Result() {
+  std::shared_ptr<FootPrint> foot_print = shared_from_this();
+  size_t upperbound = 0;
+  uint32_t total_footprints = 0;
+  while (foot_print != nullptr) {
+    foot_print->printStats();
+
+    upperbound = foot_print->getOffset();
+    foot_print = foot_print->Next();
+    total_footprints++;
+  }
+
+  MS_LOG(DEBUG) << total_footprints << " footprints allocated";
+
+  return upperbound;
+}
+// Class FootPrint manages memory allocations. The function `findFirst` finds the first suitable
+// offset in the given merged intervals to allocate a BlockTensor, and `Merge` merges overlapping intervals.
+
+// Function: findFirst
+// Purpose: This function aims to find the first available offset within the provided merged intervals to allocate a BlockTensor.
+// If a suitable offset is found, it updates the offset and returns true; otherwise, returns false.
+// Params:
+// - merged: A pointer to a stack of merged intervals representing available memory spaces.
+// - block: The BlockTensor object representing the memory block to be allocated.
+// - offset: A pointer to a size_t variable where the function will store the found offset.
+// Return: A boolean value indicating whether a suitable offset has been found.
+bool FootPrint::findFirst(stack<Interval> *merged, const BlockTensor &block, size_t *offset) {
+    MS_EXCEPTION_IF_NULL(merged);
+    MS_EXCEPTION_IF_NULL(offset);
+    bool bfound = false;
+    
+    // Initialize a set to hold candidate offsets, sorted by the given strategy.
+    std::set<pair<size_t, size_t>, bool (*)(const pair<size_t, size_t> &a, const pair<size_t, size_t> &b)>
+        offsetcandidates(g_pBranching[m_branching_strategy_]);
+    
+    size_t gap;
+    Interval a;
+    Interval it;
+    
+    // Calculate the size of the block to be allocated.
+    size_t block_size;
+    if (block.Alone()) {
+        block_size = block.m_size_;
+    } else {
+        block_size = block.m_start_tensor_->size_;  // consider only first tensor for contiguous block
+    }
+
+    // Iterate through the merged intervals and find candidate offsets where the block can fit.
+    a.ub() = algorithm[m_algorithm_](this);
+    while (!(*merged).empty()) {
+        it = (*merged).top();
+        (*merged).pop();
+        a.lb() = it.ub();
+        if (a.contains(block_size) && a.lb() + block.m_size_ <= algorithm[m_algorithm_](this)) {
+            gap = a.ub() - a.lb() - block_size;
+            offsetcandidates.emplace(pair<size_t, size_t>(a.lb(), gap));
+        }
+        a.ub() = it.lb();
+    }
+    
+    // If suitable candidates are found, update the offset and return true.
+    if (offsetcandidates.size() > 0) {
+        *offset = (*offsetcandidates.begin()).first;
+        m_foot_print_next_->m_offset_ = std::max(m_foot_print_next_->m_offset_, *offset + block.m_size_);
+        bfound = true;
+    }
+
+    return bfound;
+}
+
+// Function: Merge
+// Purpose: This function merges overlapping intervals from the input vector and stores the result in a stack.
+// Params:
+// - interval_v: A pointer to a vector of Interval objects representing the memory spaces to be merged.
+// - s: A pointer to a stack where the function will store the merged intervals.
+// Return: void.
+void FootPrint::Merge(vector<Interval> *interval_v, stack<Interval> *s) {
+    MS_EXCEPTION_IF_NULL(s);
+    MS_EXCEPTION_IF_NULL(interval_v);
+    
+    // Sort the input intervals in ascending order for merging.
+    sort((*interval_v).begin(), (*interval_v).end(),
+         [](Interval &i1, Interval &i2) { return (i1.lb() < i2.lb()) || (i1.lb() == i2.lb() && i1.ub() < i2.ub()); });
+    
+    // Iterate through the sorted intervals and merge overlapping ones.
+    (*s).push((*interval_v)[0]);
+    for (size_t i = 1; i < (*interval_v).size(); i++) {
+        Interval &top = (*s).top();
+        Interval &b = (*interval_v)[i];
+        if (top.ub() < b.lb())
+            (*s).push(b);
+        else if (top.ub() < b.ub())
+            top.ub() = b.ub();
+    }
+    return;
+}
+/**
+ * @brief Find an appropriate offset for placing a BlockTensor in memory considering a set of constraints.
+ *
+ * This function attempts to find a suitable memory offset for a given BlockTensor. It takes into consideration
+ * the constraints defined against the already allocated tensors to find an appropriate location where the
+ * BlockTensor doesn’t violate any constraints.
+ *
+ * @param constraints A vector of DynamicBitSet representing the constraints between tensors.
+ * @param block The BlockTensor for which an offset needs to be found.
+ * @param offset A pointer to a size_t where the found offset will be stored.
+ * @return Returns true if a suitable offset is found, otherwise returns false.
+ */
+bool FootPrint::findOffset(const std::vector<DynamicBitSet> *constraints, const BlockTensor &block, size_t *offset) {
+  
+  // Check for null pointer and initialize local variables
+  MS_EXCEPTION_IF_NULL(offset);
+  bool bretval = true;
+  vector<Interval> l_interval;
+  
+  // Reserve space for intervals estimation
+  const size_t intervals_estimation = 1000;
+  l_interval.reserve(intervals_estimation * sizeof(Interval));
+  
+  *offset = m_offset_;
+  
+  // Transform constrained tensors into non-eligible intervals
+  if (block.Alone()) {
+    // Handle the case when the block is alone and check specific conditions
+    if (m_algorithm_ == static_cast<uint32_t>(kManyObjects) && m_starts_.size() > 0 && m_starts_[0]->Alone() &&
+        (*constraints)[block.m_start_tensor_->index_].IsBitTrue(m_starts_[0]->m_start_tensor_->index_) == false) {
+      return false;
+    }
+    
+    // Iterate over allocated tensors and identify ineligible intervals
+    for (size_t i = 0; i < m_starts_.size(); i++) {
+      auto allocated_tensor = m_starts_[i]->m_start_tensor_;
+      while (allocated_tensor != nullptr) {
+        if ((*constraints)[block.m_start_tensor_->index_].IsBitTrue(allocated_tensor->index_) == false) {
+          l_interval.emplace_back(Interval(allocated_tensor));
+        }
+        allocated_tensor = allocated_tensor->right_;
+      }
+    }
+  } else {
+    // Handle the case when the block is not alone and compute ineligible intervals differently
+    int64_t start_offset = static_cast<int64_t>(m_offset_);
+    for (size_t i = 0; i < m_starts_.size(); i++) {
+      auto allocated_tensor = m_starts_[i]->m_start_tensor_;
+      while (allocated_tensor != nullptr) {
+        int64_t allocated_offset = static_cast<int64_t>(allocated_tensor->offset_);
+        int64_t allocated_size = static_cast<int64_t>(allocated_tensor->size_);
+        int64_t accumulator = 0;
+        for (auto block_tensor = block.m_start_tensor_; block_tensor != nullptr; block_tensor = block_tensor->right_) {
+          if ((*constraints)[block_tensor->index_].IsBitTrue(allocated_tensor->index_) == false) {
+            int64_t start_first_contiguous = allocated_offset - accumulator - SizeToLong(block_tensor->size_);
+            int64_t end_first_contiguous = allocated_offset - accumulator + allocated_size;
+            if (start_first_contiguous > start_offset) {
+              l_interval.emplace_back(Interval(start_first_contiguous, end_first_contiguous));
+            } else {
+              if (end_first_contiguous > start_offset) {
+                l_interval.emplace_back(Interval(start_offset, end_first_contiguous));
+              }
+            }
+          }
+          accumulator += SizeToLong(block_tensor->size_);
+        }
+        allocated_tensor = allocated_tensor->right_;
+      }
+    }
+  }
+  
+  // Merge non-eligible intervals and find a slot to allocate the tensor block
+  if (!l_interval.empty()) {
+    stack<Interval> l_mergedIntervals;
+    Merge(&l_interval, &l_mergedIntervals);
+    bretval = findFirst(&l_mergedIntervals, block, offset);
+  }
+  
+  return bretval;
+}
+
+void FootPrint::addElem(BlockTensor *block, const size_t &offset) {
+  if (m_foot_print_next_ == nullptr) {
+    m_foot_print_next_ = std::make_shared<FootPrint>();
+    size_t newoffset = m_offset_ + block->m_size_;
+    m_foot_print_next_->setOffset(newoffset);
+    m_foot_print_next_->setAlignment(m_alignment_);
+    m_foot_print_next_->m_solId_ = m_solId_;
+    m_starts_.clear();
+    MS_LOG(DEBUG) << "Creating footprint at offset: " << m_offset_;
+  }
+
+  addStart(block);
+  size_t offset1 = offset;
+  SomasSolverTensorDescPtr tensor = block->m_start_tensor_;
+  MS_LOG(DEBUG) << "Allocating block: " << tensor->index_ << " in offset: " << offset;
+  auto sol_id = block->m_current_sol_;
+  if (block->offsets_.find(sol_id) != block->offsets_.end()) {
+    MS_LOG(WARNING) << "Warning addElem: Offset overwritten at solution " << sol_id << " for block "
+                    << block->m_start_tensor_->index_;
+  }
+  (void)block->offsets_.emplace(sol_id, offset);
+  while (tensor) {
+    tensor->offset_ = offset1;
+    offset1 += tensor->size_;
+
+    MS_LOG(DEBUG) << tensor->index_ << " " << tensor->size_ << " " << tensor->offset_;
+    tensor = tensor->right_;
+  }
+}
+void FootPrint::printStats() {
+  MS_LOG(DEBUG) << "Footprint blocks: " << m_starts_.size() << " \toffset: " << m_offset_;
+}
+
+/**
+ * @brief Evaluates and allocates memory for a series of BlockTensors.
+ *
+ * This function is a part of the FastHeuristic class and is primarily designed to evaluate 
+ * and allocate memory for a set of tensors represented by BlockTensors. The allocation is 
+ * carried out respecting any constraints represented by DynamicBitSet and considering the 
+ * current FootPrint of memory usage. This function is instrumental in optimizing memory 
+ * allocation for tensor blocks.
+ *
+ * @param block_tensors_v Pointer to a vector of BlockTensor, representing the tensor blocks
+ *                        to be allocated memory.
+ * @param foot_print Shared pointer to a FootPrint object representing the current memory footprint.
+ * @param pConstraints Pointer to a vector of DynamicBitSet representing allocation constraints.
+ * @return Returns true if memory allocation for all tensor blocks is successful, else returns false.
+ */
+bool FastHeuristic::Eval(vector<BlockTensor> *block_tensors_v, const std::shared_ptr<FootPrint> &foot_print,
+                         const std::vector<DynamicBitSet> *pConstraints) {
+  // Ensure the foot_print pointer is not null.
+  MS_EXCEPTION_IF_NULL(foot_print);
+  
+  // Initialize the start time for evaluating the execution time of this method.
+  auto start = std::chrono::system_clock::now();
+
+  // Initialize necessary variables.
+  std::shared_ptr<FootPrint> p = foot_print;
+  bool bpushed = false;
+  size_t offset = foot_print->getOffset();
+  m_tensors_allocated_ = 0;
+  SomasSolverTensorDescPtr tensor = nullptr;
+
+  // Iterate through each tensor block in block_tensors_v.
+  for (auto &block : *block_tensors_v) {
+    // If the block does not need reallocation, update its offsets and continue to the next block.
+    if (!block.m_bre_allocate_) {
+      offset = block.m_start_tensor_->offset_;
+      auto aux_id = foot_print->m_solId_;
+      auto aux_offset = block.m_start_tensor_->offset_;
+      if (block.offsets_.find(aux_id) != block.offsets_.end()) {
+        MS_LOG(WARNING) << "Warning: Offset overwritten at solution " << aux_id << " for block "
+                        << block.m_start_tensor_->index_;
+      }
+      (void)block.offsets_.emplace(aux_id, aux_offset);
+      continue;
+    }
+    
+    // Reset bpushed and assign the current foot print to p.
+    bpushed = false;
+    p = foot_print;
+    block.m_current_sol_ = foot_print->m_solId_;
+    
+    // Attempt to find a suitable offset in the memory for the current block.
+    while (!bpushed) {
+      if (p->findOffset(pConstraints, block, &offset)) {
+        // If a suitable offset is found, add the block to the current foot print.
+        p->addElem(&block, offset);
+        tensor = block.m_start_tensor_;
+        
+        // Update the count of tensors allocated.
+        while (tensor) {
+          m_tensors_allocated_++;
+          tensor = tensor->right_;
+        }
+        bpushed = true;
+        break;
+      }
+      
+      // If not found, proceed to the next available foot print slot.
+      if (p->Next() != nullptr) {
+        p = p->Next();
+      } else if (bpushed == false) {
+        // If the allocation fails, log a warning and return false.
+        MS_LOG(WARNING) << "Internal Error: Could not allocate memory for tensor: " << tensor->index_;
+        return false;
+      }
+    }
+  }
+  
+  // Log the elapsed time of the Fast Heuristic search.
+  MS_LOG(DEBUG)
+    << "\nElapsed time of Fast Heuristic search: "
+    << std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now() - start).count() << " ms";
+  
+  // Return true indicating successful memory allocation for all tensor blocks.
+  return true;
+}
+}  // namespace somas
+}  // namespace mindspore

mindspore/ccsrc/backend/common/somas/somas_solver_core.cc

@@ -33,40 +33,104 @@ using std::vector;
 
 namespace mindspore {
 namespace somas {
+
+/*
+  Overview:
+    The `MemoryAllocationSolver` function is part of the `SomasSolverCore` class and is responsible for finding
+    optimal memory allocation for tensors. The function explores different algorithms, sorting strategies, and
+    offset strategies to identify the best combination for minimizing memory usage. The function measures the
+    time taken for each combination and logs relevant information for analysis. The result is the establishment
+    of an upper bound on memory usage with the identified best solution.
+
+  Returns:
+    Status: Indicates the success or failure of the memory allocation solver.
+
+  Process:
+    1. Initialize variables and record the start time.
+    2. If the `all_` flag is set, loop over all heuristics.
+       a. Iterate through all combinations of algorithm, sorting strategy, and branching strategy.
+       b. For each combination, clean the blocks, sort tensors, find solutions, and record the memory usage upper bound.
+       c. Update the best solution information based on the least memory usage.
+       d. Verify the solutions and increment the solution count.
+       e. Log the resume of the solver with the best solution information.
+    3. If the `all_` flag is not set, proceed with a single heuristic.
+       a. Build blocks, sort tensors, find solutions, and verify the solutions.
+       b. Log the algorithm and sorting strategy information if not multi-threaded.
+    4. Return the status of the solver.
+*/
 Status SomasSolverCore::MemoryAllocationSolver() {
+  // Record the start time for the memory allocation solver.
   auto start = std::chrono::system_clock::now();
+  
+  // Initialize the return value as SUCCESS.
   Status retval = SUCCESS;
+  
+  // Initialize the best memory usage to the maximum size_t value.
   size_t best = SIZE_MAX;
-  if (all_) {  // loop over all heuristics
+  
+  // Check if all heuristics should be looped over.
+  if (all_) {
+    // Initialize variables to store the best solution information.
     FittingType best_branching = kBest;
     SortingType best_sorting = kGreaterSizeSmallerIndex;
     AlgorithmType best_algorithm = kManyObjects;
     int64_t best_timing = INT64_MAX;
     uint32_t best_sol = 0;
     size_t worst = 0;
+    
+    // Build the blocks required for solving.
     BuildBlocks();
+    
+    // Clean any existing information.
     Clean();
+    
+    // Log the headers for the solver's output.
     MS_LOG(INFO) << "time\tSol#\tResult\t\t\t\tAlgorithm\tSorting Strategy\tOffset Strategy";
+    
+    // Iterate through all available algorithm types.
     for (size_t algorithm = 0; algorithm < static_cast<size_t>(kNumAlgorithmTypes); algorithm++) {
+      // Set the current algorithm type.
       algorithm_ = static_cast<AlgorithmType>(algorithm);
+      
+      // Iterate through all available sorting types.
       for (size_t sort_strategy = 0; sort_strategy < static_cast<size_t>(kNumSortingTypes); sort_strategy++) {
+        // Set the current sorting type.
         sort_strategy_ = static_cast<SortingType>(sort_strategy);
+        
+        // Sort the tensors based on the current sorting strategy.
         SortTensors();
+        
+        // Iterate through all available fitting types (branching strategies).
         for (size_t branching_strategy = 0; branching_strategy < static_cast<size_t>(kNumFittingTypes);
              branching_strategy++) {
+          // Set the current branching strategy.
           branching_strategy_ = static_cast<FittingType>(branching_strategy);
+          
+          // Clean any existing information.
           Clean();
+          
+          // Log the start of timing for the current combination.
           MS_LOG(DEBUG) << "Timing Start " << tensors_.size() << " Tensors";
+          
+          // Record the start time for finding solutions.
           auto start_upper = std::chrono::system_clock::now();
+          
+          // Find the solutions and record the upper bound of memory usage.
           upperbound_ = FindSolutions();
+          
+          // Log the elapsed time for finding the solutions.
           MS_LOG(DEBUG) << "Elapsed time of upper bound testing: "
                         << std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now() -
                                                                                  start_upper)
                              .count()
                         << " ms";
+          
+          // Update the worst memory usage upper bound.
           if (upperbound_ > worst) {
             worst = upperbound_;
           }
+          
+          // Update the best solution information if a better or equal solution is found.
           if (upperbound_ < best || upperbound_ == best) {
             best = upperbound_;
             best_algorithm = algorithm_;
@@ -75,13 +139,23 @@ Status SomasSolverCore::MemoryAllocationSolver() {
             best_sol = sol_count_;
             best_timing = timing_;
           }
+          
+          // Verify the found solutions.
           Verify();
+          
+          // Increment the solution count.
           sol_count_++;
         }
       }
     }
+    
+    // Set the upper bound to the best memory usage found.
     upperbound_ = best;
+    
+    // Record the end time for the solver.
     auto end = std::chrono::system_clock::now();
+    
+    // Calculate the total elapsed time for the solver.
     size_t total_time = std::chrono::duration_cast<std::chrono::milliseconds>((end - start)).count();
     const double giga = 1024. * 1024. * 1024.;
     const double cent = 100.;
@@ -100,12 +174,16 @@ Status SomasSolverCore::MemoryAllocationSolver() {
     best_sol_ = best_sol;
     SetBestSolution();
   } else {
-    // print only for single heuristic no multi thread
+    // If not looping over all heuristics, proceed with a single heuristic.
+    
+    // Log the algorithm and sorting strategy information if not multi-threaded.
     if (!is_multi_thread_valid_) {
       MS_LOG(INFO) << "Algorithm strategy: " << algorithmTypeNames[algorithm_];
       MS_LOG(INFO) << "Sorting strategy: " << sortingNames[sort_strategy_];
       MS_LOG(INFO) << "Offset strategy: " << branchingNames[branching_strategy_];
     }
+    
+    // Build blocks, sort tensors, find solutions, and verify the solutions.
     BuildBlocks();
     SortTensors();
     upperbound_ = FindSolutions();
@@ -114,11 +192,23 @@ Status SomasSolverCore::MemoryAllocationSolver() {
   return retval;
 }
 
+/**
+ * @brief Verifies if the memory allocation solution is correct.
+ * 
+ * This function checks the validity of the memory allocation solution generated 
+ * by the solver by invoking the Verify(const size_t &upperbound) method with the 
+ * calculated upper bound. It logs whether the verification was successful and 
+ * the memory consumed by the solution in gigabytes.
+ *
+ * @return Status - Returns SUCCESS if verification is successful, otherwise returns FAILED.
+ */
 Status SomasSolverCore::Verify() {
   Status retval = SUCCESS;
+  
   if (verify_) {
     MS_LOG(INFO) << "Verifying solution..";
 
+    // Call the overloaded Verify function to perform detailed verification
     if (!Verify(upperbound_)) {
       MS_LOG(WARNING) << "Solver Allocation Memory Check FAILS";
       retval = FAILED;
@@ -133,6 +223,21 @@ Status SomasSolverCore::Verify() {
   return retval;
 }
 
+/**
+ * @brief Verifies the memory allocation solution against various constraints.
+ * 
+ * This function iterates over all tensors and checks the following:
+ * 1. The alignment of each tensor.
+ * 2. Whether tensors with continuous constraints are continuous in memory.
+ * 3. Whether tensors with conflict constraints overlap in memory.
+ * 4. If the calculated upper bound matches the memory footprint of the tensors.
+ * 
+ * Any violations of these constraints are logged, and the overall verification
+ * result is returned as a boolean value.
+ *
+ * @param upperbound - The upper memory bound calculated by the solver.
+ * @return bool - Returns true if all constraints are satisfied, otherwise returns false.
+ */
 bool SomasSolverCore::Verify(const size_t &upperbound) {
   auto start = std::chrono::system_clock::now();
   bool retval = true;
@@ -141,27 +246,35 @@ bool SomasSolverCore::Verify(const size_t &upperbound) {
   SomasSolverTensorDescPtr t2;
 
   for (auto t1_ : tensors_) {
-    // check alignment
+    // Check alignment and update result
     result = std::max(result, t1_.second->size_ + t1_.second->offset_);
+    
     for (auto t2_ : tensors_) {
       t1 = t1_.second;
       t2 = t2_.second;
       if (t1->index_ == t2->index_) continue;
+      
+      // Check if either tensor is lifelong and their indices are different
       bool blifelong = (t1->lifelong_ || t2->lifelong_) && (t1->index_ != t2->index_);
-      if (t2->right_ == t1) {  // continuous constraint
-        // t1 must be continuous to t2
+      
+      // Check continuous constraint violation
+      if (t2->right_ == t1) {
         bool bcontinuous = t1->offset_ == (t2->offset_ + t2->size_);
         if (!bcontinuous) {
           MS_LOG(WARNING) << "Continuous constraint violation in tensors " << t1->index_ << " and" << t2->index_;
           retval = false;
         }
-      } else if (blifelong || constraints_[t1->index_].IsBitTrue(t2->index_) == false) {  // conflict constraint
+      } 
+      // Check conflict constraint violation
+      else if (blifelong || constraints_[t1->index_].IsBitTrue(t2->index_) == false) {
         size_t t1_ub = t1->offset_ + t1->size_;
         size_t t2_ub = t2->offset_ + t2->size_;
         bool b_overlap_lb = ((t2->offset_ >= t1->offset_) && (t2->offset_ < t1_ub));
         bool b_overlap_ub = ((t2_ub > t1->offset_) && (t2_ub < t1_ub));
         bool b_overlap = b_overlap_lb || b_overlap_ub;
         bool biszerosized = t1->size_ == 0 || t2->size_ == 0;
+        
+        // Log if overlapping and not zero-sized
         if (b_overlap && !biszerosized) {
           MS_LOG(WARNING) << "Non-overlap constraint violation in tensors " << t1->index_ << " and" << t2->index_;
           retval = false;
@@ -169,56 +282,89 @@ bool SomasSolverCore::Verify(const size_t &upperbound) {
       }
     }
   }
+  
+  // Check if upperbound matches the memory footprint result
   if (upperbound != result) {
     MS_LOG(WARNING) << "ERROR Invalid upperbound result --> Footprint Result: " << upperbound_
                     << " Tensor Result: " << result + lifelong_memory_;
     retval = false;
   }
+  
+  // Log the time taken for verification
   MS_LOG(DEBUG)
     << "\nElapsed time of Fast Heuristic Check: "
     << std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now() - start).count() << " ms";
+    
   return retval;
 }
 
+/**
+ * @brief Constructs blocks of tensors from the individual tensors.
+ *
+ * This function iterates through all tensors and builds blocks of tensors.
+ * For each tensor, it checks whether it is already part of a block or it is a lifelong tensor.
+ * Lifelong tensors contribute to the `lifelong_memory_`, and their size is accumulated.
+ * For other tensors, the function finds the leftmost tensor in the block, initializes a new block tensor,
+ * and calculates its size by summing up the sizes of the individual tensors in the block.
+ * The created block tensor is then added to the `block_tensors_` list.
+ *
+ * @note Lifelong tensors and their count are logged if found.
+ */
 void SomasSolverCore::BuildBlocks() {
+  // Logging the start of building block tensors
   MS_LOG(DEBUG) << "Building block of tensors";
 
+  // Initializing lifelong memory
   lifelong_memory_ = 0;
+  // Counter for number of tensors in blocks
   uint64_t tensors_block_count = 0;
+  // Iterating through all tensors
   for (auto tensor : tensors_) {
     SomasSolverTensorDescPtr pTensor = tensor.second;
+    // Skipping blocked tensors
     if (pTensor->blocked_) continue;
+    // If tensor is lifelong, adding its size to lifelong memory and continue to next iteration
     if (pTensor->lifelong_) {
       lifelong_memory_ += pTensor->size_;
       continue;
     }
-    // move to the left
+    // Moving tensor to the leftmost in the block
     while (pTensor->left_) pTensor = pTensor->left_;
 
-    // set start tensor
+    // Initializing block tensor and setting start tensor
     BlockTensor bTensor;
     bTensor.m_bre_allocate_ = true;
     bTensor.m_start_tensor_ = pTensor;
-    // find size
+    // Calculating size of block tensor
     bTensor.m_size_ = 0;
-
     do {
+      // Accumulating sizes of tensors in the block
       bTensor.m_size_ += pTensor->size_;
       pTensor->blocked_ = true;
       pTensor = pTensor->right_;
       tensors_block_count++;
     } while (pTensor != nullptr);
 
-    // add to the list
+    // Adding the created block tensor to the list
     this->block_tensors_.emplace_back(bTensor);
   }
 
+  // Logging if any lifelong tensors are found
   if (tensors_block_count != tensors_.size())
     MS_LOG(INFO) << static_cast<int>(tensors_.size() - tensors_block_count) << " lifelong tensors found";
 }
 
+/**
+ * @brief Cleans the state of the block tensors and resets the upperbound.
+ *
+ * This function iterates through each block tensor and resets the allocation flag and offset for each tensor in the block.
+ * It is typically used to clean the state after solving or when the state needs to be re-initialized.
+ * After iterating through all the block tensors, it resets the `upperbound_` to its maximum value.
+ */
 void SomasSolverCore::Clean() {
+  // Iterating through each block tensor
   for (auto &block : block_tensors_) {
+    // Resetting the allocation flag and offset for each tensor in the block
     block.m_bre_allocate_ = true;
     auto pTensor = block.m_start_tensor_;
     while (pTensor) {
@@ -226,10 +372,24 @@ void SomasSolverCore::Clean() {
       pTensor = pTensor->right_;
     }
   }
+  // Resetting the upperbound
   upperbound_ = SIZE_MAX;
 }
 
+/**
+ * @brief Compares two BlockTensors based on their size and index.
+ *
+ * This function is a comparison function used for sorting or ordering BlockTensors.
+ * It returns true if the size of the first BlockTensor `t1` is greater than the size of the second BlockTensor `t2`,
+ * or if they have equal sizes but the index of the start tensor of `t1` is smaller than the index of the start tensor of `t2`.
+ *
+ * @param t1 The first BlockTensor to be compared.
+ * @param t2 The second BlockTensor to be compared.
+ * @return True if `t1` is considered greater than `t2` according to the defined criteria; False otherwise.
+ */
 static bool GreaterSizeSmallerIndex(const BlockTensor &t1, const BlockTensor &t2) {
+  // Comparison function for sorting block tensors
+  // Returns true if size of t1 is greater than t2 or if they have equal sizes but t1 has smaller index
   return t1.m_size_ > t2.m_size_ ||
          (t1.m_size_ == t2.m_size_ && t1.m_start_tensor_->index_ < t2.m_start_tensor_->index_);
 }
@@ -264,6 +424,13 @@ static bool GreaterSizeGreaterConstraintsGreaterIndex(const BlockTensor &t1, con
 }
 #endif
 
+/**
+ * @brief Sorts the block tensors based on the selected sorting strategy.
+ *
+ * This function initializes a map associating each SortingType to its corresponding sorting function.
+ * It then applies the sorting function corresponding to the current `sort_strategy_` to the list of block tensors.
+ * The sorting strategy is logged for debugging purposes.
+ */
 void SomasSolverCore::SortTensors() {  // need to sort the tensors for Fast Heuristic
   MS_LOG(DEBUG) << "Sorting Blocks of tensor, strategy: " << sortingNames[sort_strategy_];
   typedef bool (*SortingFunction)(const BlockTensor &, const BlockTensor &);
@@ -281,6 +448,14 @@ void SomasSolverCore::SortTensors() {  // need to sort the tensors for Fast Heur
   }
 }
 
+/**
+ * @brief Restores the tensor offsets according to a given solution ID.
+ *
+ * Iterates through each block and assigns the offset of each tensor in the block based on the solution ID.
+ * The offset for each tensor is calculated by incrementally adding the size of the tensor to the block offset.
+ *
+ * @param sol_id The solution ID used to restore the solution.
+ */
 void SomasSolverCore::RestoreSolution(uint32_t sol_id) {
   for (auto block : block_tensors_) {
     if (block.offsets_.count(sol_id) == 0) MS_ASSERT(0);
@@ -295,6 +470,17 @@ void SomasSolverCore::RestoreSolution(uint32_t sol_id) {
     }
   }
 }
+
+/**
+ * @brief Searches for an optimized memory footprint using Fast Heuristic algorithm.
+ *
+ * Calls the Fast Heuristic algorithm to evaluate the memory footprint of the current block tensors.
+ * If a solution is found, it logs the details including the timing, result size, and the strategies used.
+ * If the resultant footprint is smaller than the current `upperbound_`, it updates the `upperbound_` and `best_sol_`.
+ *
+ * @param pFootprint Shared pointer to the current FootPrint object.
+ * @return The optimized memory footprint size.
+ */
 size_t SomasSolverCore::Search(const std::shared_ptr<FootPrint> &pFootprint) {
   size_t result = 0;
   FastHeuristic fh;
@@ -325,6 +511,12 @@ size_t SomasSolverCore::Search(const std::shared_ptr<FootPrint> &pFootprint) {
   return upperbound_;
 }
 
+/**
+ * @brief Appends the lifelong tensors to the final solution.
+ *
+ * Calculates the offset for each lifelong tensor and adds it to the upperbound.
+ * The total memory added by the lifelong tensors is logged for debugging.
+ */
 void SomasSolverCore::AppendLifelongTensors() {
   MS_LOG(DEBUG) << "Appending lifelong tensors to solution";
   size_t offset = upperbound_;
@@ -343,6 +535,14 @@ void SomasSolverCore::AppendLifelongTensors() {
   MS_LOG(DEBUG) << lifelong_memory_ << " bytes from lifelong tensors added to solution";
 }
 
+/**
+ * @brief Initiates the process of finding optimized solutions for tensor allocations.
+ *
+ * Creates a new FootPrint object, sets its properties, and initiates a search for optimized solutions.
+ * Appends the lifelong tensors to the solution and cleans up the FootPrint object after the process.
+ *
+ * @return The optimized memory footprint size after appending lifelong tensors.
+ */
 size_t SomasSolverCore::FindSolutions() {
   MS_LOG(DEBUG) << "Start allocating blocks,offset strategy: " << branchingNames[branching_strategy_];
 
@@ -356,6 +556,13 @@ size_t SomasSolverCore::FindSolutions() {
   return upperbound_;
 }
 
+/**
+ * @brief Destroys a linked list of FootPrint objects.
+ *
+ * Iterates through the linked list of FootPrint objects and sets each pointer to null, effectively destroying the list.
+ *
+ * @param pFootprint Shared pointer to the head of the FootPrint linked list.
+ */
 void SomasSolverCore::Destroy(std::shared_ptr<FootPrint> &pFootprint) {
   while (pFootprint != nullptr) {
     if (pFootprint->Next() != nullptr) {

mindspore/ccsrc/backend/common/somas/somas_solver_core_old.cc

@@ -0,0 +1,578 @@
+/**
+ * Copyright 2020-2021 Huawei Technologies Co., Ltd
+
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+
+ * http://www.apache.org/licenses/LICENSE-2.0
+
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+*/
+
+#include <algorithm>
+#include <chrono>
+#include <cstdio>
+#include <ctime>
+#include <memory>
+#include <string>
+#include <vector>
+#include <map>
+#include "utils/hash_map.h"
+#include "backend/common/somas/somas_solver_alg.h"
+#include "backend/common/somas/somas_solver_core.h"
+#include "backend/common/somas/somas_solver_pre.h"
+
+using mindspore::HashMap;
+using std::sort;
+using std::vector;
+
+namespace mindspore {
+namespace somas {
+
+/*
+  Overview:
+    The `MemoryAllocationSolver` function is part of the `SomasSolverCore` class and is responsible for finding
+    optimal memory allocation for tensors. The function explores different algorithms, sorting strategies, and
+    offset strategies to identify the best combination for minimizing memory usage. The function measures the
+    time taken for each combination and logs relevant information for analysis. The result is the establishment
+    of an upper bound on memory usage with the identified best solution.
+
+  Returns:
+    Status: Indicates the success or failure of the memory allocation solver.
+
+  Process:
+    1. Initialize variables and record the start time.
+    2. If the `all_` flag is set, loop over all heuristics.
+       a. Iterate through all combinations of algorithm, sorting strategy, and branching strategy.
+       b. For each combination, clean the blocks, sort tensors, find solutions, and record the memory usage upper bound.
+       c. Update the best solution information based on the least memory usage.
+       d. Verify the solutions and increment the solution count.
+       e. Log the resume of the solver with the best solution information.
+    3. If the `all_` flag is not set, proceed with a single heuristic.
+       a. Build blocks, sort tensors, find solutions, and verify the solutions.
+       b. Log the algorithm and sorting strategy information if not multi-threaded.
+    4. Return the status of the solver.
+*/
+Status SomasSolverCore::MemoryAllocationSolver() {
+  // Record the start time for the memory allocation solver.
+  auto start = std::chrono::system_clock::now();
+  
+  // Initialize the return value as SUCCESS.
+  Status retval = SUCCESS;
+  
+  // Initialize the best memory usage to the maximum size_t value.
+  size_t best = SIZE_MAX;
+  
+  // Check if all heuristics should be looped over.
+  if (all_) {
+    // Initialize variables to store the best solution information.
+    FittingType best_branching = kBest;
+    SortingType best_sorting = kGreaterSizeSmallerIndex;
+    AlgorithmType best_algorithm = kManyObjects;
+    int64_t best_timing = INT64_MAX;
+    uint32_t best_sol = 0;
+    size_t worst = 0;
+    
+    // Build the blocks required for solving.
+    BuildBlocks();
+    
+    // Clean any existing information.
+    Clean();
+    
+    // Log the headers for the solver's output.
+    MS_LOG(INFO) << "time\tSol#\tResult\t\t\t\tAlgorithm\tSorting Strategy\tOffset Strategy";
+    
+    // Iterate through all available algorithm types.
+    for (size_t algorithm = 0; algorithm < static_cast<size_t>(kNumAlgorithmTypes); algorithm++) {
+      // Set the current algorithm type.
+      algorithm_ = static_cast<AlgorithmType>(algorithm);
+      
+      // Iterate through all available sorting types.
+      for (size_t sort_strategy = 0; sort_strategy < static_cast<size_t>(kNumSortingTypes); sort_strategy++) {
+        // Set the current sorting type.
+        sort_strategy_ = static_cast<SortingType>(sort_strategy);
+        
+        // Sort the tensors based on the current sorting strategy.
+        SortTensors();
+        
+        // Iterate through all available fitting types (branching strategies).
+        for (size_t branching_strategy = 0; branching_strategy < static_cast<size_t>(kNumFittingTypes);
+             branching_strategy++) {
+          // Set the current branching strategy.
+          branching_strategy_ = static_cast<FittingType>(branching_strategy);
+          
+          // Clean any existing information.
+          Clean();
+          
+          // Log the start of timing for the current combination.
+          MS_LOG(DEBUG) << "Timing Start " << tensors_.size() << " Tensors";
+          
+          // Record the start time for finding solutions.
+          auto start_upper = std::chrono::system_clock::now();
+          
+          // Find the solutions and record the upper bound of memory usage.
+          upperbound_ = FindSolutions();
+          
+          // Log the elapsed time for finding the solutions.
+          MS_LOG(DEBUG) << "Elapsed time of upper bound testing: "
+                        << std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now() -
+                                                                                 start_upper)
+                             .count()
+                        << " ms";
+          
+          // Update the worst memory usage upper bound.
+          if (upperbound_ > worst) {
+            worst = upperbound_;
+          }
+          
+          // Update the best solution information if a better or equal solution is found.
+          if (upperbound_ < best || upperbound_ == best) {
+            best = upperbound_;
+            best_algorithm = algorithm_;
+            best_branching = branching_strategy_;
+            best_sorting = sort_strategy_;
+            best_sol = sol_count_;
+            best_timing = timing_;
+          }
+          
+          // Verify the found solutions.
+          Verify();
+          
+          // Increment the solution count.
+          sol_count_++;
+        }
+      }
+    }
+    
+    // Set the upper bound to the best memory usage found.
+    upperbound_ = best;
+    
+    // Record the end time for the solver.
+    auto end = std::chrono::system_clock::now();
+    
+    // Calculate the total elapsed time for the solver.
+    size_t total_time = std::chrono::duration_cast<std::chrono::milliseconds>((end - start)).count();
+    const double giga = 1024. * 1024. * 1024.;
+    const double cent = 100.;
+    MS_LOG(INFO) << "SOMAS SOLVER RESUME:";
+    MS_LOG(INFO) << "Best Solution:[" << 1 + best_sol << "/" << sol_count_ << "] ";
+    MS_LOG(INFO) << "Best result:" << best << " Bytes " << (best) / (giga) << " GB ("
+                 << (best - lifelong_memory_) / (giga) << " GB + " << lifelong_memory_ / (giga)
+                 << " GB from lifelong tensors)";
+
+    MS_LOG(INFO) << "Best timing:" << best_timing << " ms";
+    MS_LOG(INFO) << "Best algorithm: " << algorithmTypeNames[best_algorithm];
+    MS_LOG(INFO) << "Best sorting strategy: " << sortingNames[best_sorting];
+    MS_LOG(INFO) << "Best offset strategy: " << branchingNames[best_branching];
+    MS_LOG(INFO) << "Time elapsed: " << total_time << " ms";
+    MS_LOG(INFO) << "Spread:" << static_cast<double>((worst - best) / static_cast<double>(best * cent)) << " %%";
+    best_sol_ = best_sol;
+    SetBestSolution();
+  } else {
+    // If not looping over all heuristics, proceed with a single heuristic.
+    
+    // Log the algorithm and sorting strategy information if not multi-threaded.
+    if (!is_multi_thread_valid_) {
+      MS_LOG(INFO) << "Algorithm strategy: " << algorithmTypeNames[algorithm_];
+      MS_LOG(INFO) << "Sorting strategy: " << sortingNames[sort_strategy_];
+      MS_LOG(INFO) << "Offset strategy: " << branchingNames[branching_strategy_];
+    }
+    
+    // Build blocks, sort tensors, find solutions, and verify the solutions.
+    BuildBlocks();
+    SortTensors();
+    upperbound_ = FindSolutions();
+    Verify();
+  }
+  return retval;
+}
+
+/**
+ * @brief Verifies if the memory allocation solution is correct.
+ * 
+ * This function checks the validity of the memory allocation solution generated 
+ * by the solver by invoking the Verify(const size_t &upperbound) method with the 
+ * calculated upper bound. It logs whether the verification was successful and 
+ * the memory consumed by the solution in gigabytes.
+ *
+ * @return Status - Returns SUCCESS if verification is successful, otherwise returns FAILED.
+ */
+Status SomasSolverCore::Verify() {
+  Status retval = SUCCESS;
+  
+  if (verify_) {
+    MS_LOG(INFO) << "Verifying solution..";
+
+    // Call the overloaded Verify function to perform detailed verification
+    if (!Verify(upperbound_)) {
+      MS_LOG(WARNING) << "Solver Allocation Memory Check FAILS";
+      retval = FAILED;
+    } else {
+      const double giga = 1024. * 1024. * 1024.;
+      MS_LOG(INFO) << "Solver Allocation Memory Check SUCCESS !!";
+      MS_LOG(INFO) << "Result: " << upperbound_ << " (" << (upperbound_) / (giga) << " GB)";
+      retval = SUCCESS;
+    }
+  }
+
+  return retval;
+}
+
+/**
+ * @brief Verifies the memory allocation solution against various constraints.
+ * 
+ * This function iterates over all tensors and checks the following:
+ * 1. The alignment of each tensor.
+ * 2. Whether tensors with continuous constraints are continuous in memory.
+ * 3. Whether tensors with conflict constraints overlap in memory.
+ * 4. If the calculated upper bound matches the memory footprint of the tensors.
+ * 
+ * Any violations of these constraints are logged, and the overall verification
+ * result is returned as a boolean value.
+ *
+ * @param upperbound - The upper memory bound calculated by the solver.
+ * @return bool - Returns true if all constraints are satisfied, otherwise returns false.
+ */
+bool SomasSolverCore::Verify(const size_t &upperbound) {
+  auto start = std::chrono::system_clock::now();
+  bool retval = true;
+  size_t result = 0;
+  SomasSolverTensorDescPtr t1;
+  SomasSolverTensorDescPtr t2;
+
+  for (auto t1_ : tensors_) {
+    // Check alignment and update result
+    result = std::max(result, t1_.second->size_ + t1_.second->offset_);
+    
+    for (auto t2_ : tensors_) {
+      t1 = t1_.second;
+      t2 = t2_.second;
+      if (t1->index_ == t2->index_) continue;
+      
+      // Check if either tensor is lifelong and their indices are different
+      bool blifelong = (t1->lifelong_ || t2->lifelong_) && (t1->index_ != t2->index_);
+      
+      // Check continuous constraint violation
+      if (t2->right_ == t1) {
+        bool bcontinuous = t1->offset_ == (t2->offset_ + t2->size_);
+        if (!bcontinuous) {
+          MS_LOG(WARNING) << "Continuous constraint violation in tensors " << t1->index_ << " and" << t2->index_;
+          retval = false;
+        }
+      } 
+      // Check conflict constraint violation
+      else if (blifelong || constraints_[t1->index_].IsBitTrue(t2->index_) == false) {
+        size_t t1_ub = t1->offset_ + t1->size_;
+        size_t t2_ub = t2->offset_ + t2->size_;
+        bool b_overlap_lb = ((t2->offset_ >= t1->offset_) && (t2->offset_ < t1_ub));
+        bool b_overlap_ub = ((t2_ub > t1->offset_) && (t2_ub < t1_ub));
+        bool b_overlap = b_overlap_lb || b_overlap_ub;
+        bool biszerosized = t1->size_ == 0 || t2->size_ == 0;
+        
+        // Log if overlapping and not zero-sized
+        if (b_overlap && !biszerosized) {
+          MS_LOG(WARNING) << "Non-overlap constraint violation in tensors " << t1->index_ << " and" << t2->index_;
+          retval = false;
+        }
+      }
+    }
+  }
+  
+  // Check if upperbound matches the memory footprint result
+  if (upperbound != result) {
+    MS_LOG(WARNING) << "ERROR Invalid upperbound result --> Footprint Result: " << upperbound_
+                    << " Tensor Result: " << result + lifelong_memory_;
+    retval = false;
+  }
+  
+  // Log the time taken for verification
+  MS_LOG(DEBUG)
+    << "\nElapsed time of Fast Heuristic Check: "
+    << std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now() - start).count() << " ms";
+    
+  return retval;
+}
+
+/**
+ * @brief Constructs blocks of tensors from the individual tensors.
+ *
+ * This function iterates through all tensors and builds blocks of tensors.
+ * For each tensor, it checks whether it is already part of a block or it is a lifelong tensor.
+ * Lifelong tensors contribute to the `lifelong_memory_`, and their size is accumulated.
+ * For other tensors, the function finds the leftmost tensor in the block, initializes a new block tensor,
+ * and calculates its size by summing up the sizes of the individual tensors in the block.
+ * The created block tensor is then added to the `block_tensors_` list.
+ *
+ * @note Lifelong tensors and their count are logged if found.
+ */
+void SomasSolverCore::BuildBlocks() {
+  // Logging the start of building block tensors
+  MS_LOG(DEBUG) << "Building block of tensors";
+
+  // Initializing lifelong memory
+  lifelong_memory_ = 0;
+  // Counter for number of tensors in blocks
+  uint64_t tensors_block_count = 0;
+  // Iterating through all tensors
+  for (auto tensor : tensors_) {
+    SomasSolverTensorDescPtr pTensor = tensor.second;
+    // Skipping blocked tensors
+    if (pTensor->blocked_) continue;
+    // If tensor is lifelong, adding its size to lifelong memory and continue to next iteration
+    if (pTensor->lifelong_) {
+      lifelong_memory_ += pTensor->size_;
+      continue;
+    }
+    // Moving tensor to the leftmost in the block
+    while (pTensor->left_) pTensor = pTensor->left_;
+
+    // Initializing block tensor and setting start tensor
+    BlockTensor bTensor;
+    bTensor.m_bre_allocate_ = true;
+    bTensor.m_start_tensor_ = pTensor;
+    // Calculating size of block tensor
+    bTensor.m_size_ = 0;
+    do {
+      // Accumulating sizes of tensors in the block
+      bTensor.m_size_ += pTensor->size_;
+      pTensor->blocked_ = true;
+      pTensor = pTensor->right_;
+      tensors_block_count++;
+    } while (pTensor != nullptr);
+
+    // Adding the created block tensor to the list
+    this->block_tensors_.emplace_back(bTensor);
+  }
+
+  // Logging if any lifelong tensors are found
+  if (tensors_block_count != tensors_.size())
+    MS_LOG(INFO) << static_cast<int>(tensors_.size() - tensors_block_count) << " lifelong tensors found";
+}
+
+/**
+ * @brief Cleans the state of the block tensors and resets the upperbound.
+ *
+ * This function iterates through each block tensor and resets the allocation flag and offset for each tensor in the block.
+ * It is typically used to clean the state after solving or when the state needs to be re-initialized.
+ * After iterating through all the block tensors, it resets the `upperbound_` to its maximum value.
+ */
+void SomasSolverCore::Clean() {
+  // Iterating through each block tensor
+  for (auto &block : block_tensors_) {
+    // Resetting the allocation flag and offset for each tensor in the block
+    block.m_bre_allocate_ = true;
+    auto pTensor = block.m_start_tensor_;
+    while (pTensor) {
+      pTensor->offset_ = 0;
+      pTensor = pTensor->right_;
+    }
+  }
+  // Resetting the upperbound
+  upperbound_ = SIZE_MAX;
+}
+
+/**
+ * @brief Compares two BlockTensors based on their size and index.
+ *
+ * This function is a comparison function used for sorting or ordering BlockTensors.
+ * It returns true if the size of the first BlockTensor `t1` is greater than the size of the second BlockTensor `t2`,
+ * or if they have equal sizes but the index of the start tensor of `t1` is smaller than the index of the start tensor of `t2`.
+ *
+ * @param t1 The first BlockTensor to be compared.
+ * @param t2 The second BlockTensor to be compared.
+ * @return True if `t1` is considered greater than `t2` according to the defined criteria; False otherwise.
+ */
+static bool GreaterSizeSmallerIndex(const BlockTensor &t1, const BlockTensor &t2) {
+  // Comparison function for sorting block tensors
+  // Returns true if size of t1 is greater than t2 or if they have equal sizes but t1 has smaller index
+  return t1.m_size_ > t2.m_size_ ||
+         (t1.m_size_ == t2.m_size_ && t1.m_start_tensor_->index_ < t2.m_start_tensor_->index_);
+}
+#ifdef SOMAS_DEBUG
+static bool GreaterSizeGreaterIndex(const BlockTensor &t1, const BlockTensor &t2) {
+  return t1.m_size_ > t2.m_size_ ||
+         (t1.m_size_ == t2.m_size_ && t1.m_start_tensor_->index_ > t2.m_start_tensor_->index_);
+}
+static bool GreaterSizeSmallerConstraintsSmallerIndex(const BlockTensor &t1, const BlockTensor &t2) {
+  return t1.m_size_ > t2.m_size_ ||
+         (t1.m_size_ == t2.m_size_ && t1.m_start_tensor_->constraints_ < t2.m_start_tensor_->constraints_) ||
+         (t1.m_size_ == t2.m_size_ && t1.m_start_tensor_->constraints_ == t2.m_start_tensor_->constraints_ &&
+          t1.m_start_tensor_->index_ < t2.m_start_tensor_->index_);
+}
+static bool GreaterSizeSmallerConstraintsGreaterIndex(const BlockTensor &t1, const BlockTensor &t2) {
+  return t1.m_size_ > t2.m_size_ ||
+         (t1.m_size_ == t2.m_size_ && t1.m_start_tensor_->constraints_ < t2.m_start_tensor_->constraints_) ||
+         (t1.m_size_ == t2.m_size_ && t1.m_start_tensor_->constraints_ == t2.m_start_tensor_->constraints_ &&
+          t1.m_start_tensor_->index_ > t2.m_start_tensor_->index_);
+}
+static bool GreaterSizeGreaterConstraintsSmallerIndex(const BlockTensor &t1, const BlockTensor &t2) {
+  return t1.m_size_ > t2.m_size_ ||
+         (t1.m_size_ == t2.m_size_ && t1.m_start_tensor_->constraints_ > t2.m_start_tensor_->constraints_) ||
+         (t1.m_size_ == t2.m_size_ && t1.m_start_tensor_->constraints_ == t2.m_start_tensor_->constraints_ &&
+          t1.m_start_tensor_->index_ < t2.m_start_tensor_->index_);
+}
+static bool GreaterSizeGreaterConstraintsGreaterIndex(const BlockTensor &t1, const BlockTensor &t2) {
+  return t1.m_size_ > t2.m_size_ ||
+         (t1.m_size_ == t2.m_size_ && t1.m_start_tensor_->constraints_ > t2.m_start_tensor_->constraints_) ||
+         (t1.m_size_ == t2.m_size_ && t1.m_start_tensor_->constraints_ == t2.m_start_tensor_->constraints_ &&
+          t1.m_start_tensor_->index_ > t2.m_start_tensor_->index_);
+}
+#endif
+
+/**
+ * @brief Sorts the block tensors based on the selected sorting strategy.
+ *
+ * This function initializes a map associating each SortingType to its corresponding sorting function.
+ * It then applies the sorting function corresponding to the current `sort_strategy_` to the list of block tensors.
+ * The sorting strategy is logged for debugging purposes.
+ */
+void SomasSolverCore::SortTensors() {  // need to sort the tensors for Fast Heuristic
+  MS_LOG(DEBUG) << "Sorting Blocks of tensor, strategy: " << sortingNames[sort_strategy_];
+  typedef bool (*SortingFunction)(const BlockTensor &, const BlockTensor &);
+  mindspore::HashMap<SortingType, SortingFunction> sort_map;
+  sort_map[kGreaterSizeSmallerIndex] = &GreaterSizeSmallerIndex;
+#ifdef SOMAS_DEBUG
+  sort_map[kGreaterSizeGreaterIndex] = &GreaterSizeGreaterIndex;
+  sort_map[kGreaterSizeSmallerConstraintsSmallerIndex] = &GreaterSizeSmallerConstraintsSmallerIndex;
+  sort_map[kGreaterSizeSmallerConstraintsGreaterIndex] = &GreaterSizeSmallerConstraintsGreaterIndex;
+  sort_map[kGreaterSizeGreaterConstraintsSmallerIndex] = &GreaterSizeGreaterConstraintsSmallerIndex;
+  sort_map[kGreaterSizeGreaterConstraintsGreaterIndex] = &GreaterSizeGreaterConstraintsGreaterIndex;
+#endif
+  if (sort_strategy_ < kNumSortingTypes) {
+    sort(block_tensors_.begin(), block_tensors_.end(), *(sort_map[sort_strategy_]));
+  }
+}
+
+/**
+ * @brief Restores the tensor offsets according to a given solution ID.
+ *
+ * Iterates through each block and assigns the offset of each tensor in the block based on the solution ID.
+ * The offset for each tensor is calculated by incrementally adding the size of the tensor to the block offset.
+ *
+ * @param sol_id The solution ID used to restore the solution.
+ */
+void SomasSolverCore::RestoreSolution(uint32_t sol_id) {
+  for (auto block : block_tensors_) {
+    if (block.offsets_.count(sol_id) == 0) MS_ASSERT(0);
+    size_t bestOffset = block.offsets_[sol_id];
+    size_t offset = bestOffset;
+    SomasSolverTensorDescPtr pTensor = block.m_start_tensor_;
+
+    while (pTensor) {
+      pTensor->offset_ = offset;
+      offset += pTensor->size_;
+      pTensor = pTensor->right_;
+    }
+  }
+}
+
+/**
+ * @brief Searches for an optimized memory footprint using Fast Heuristic algorithm.
+ *
+ * Calls the Fast Heuristic algorithm to evaluate the memory footprint of the current block tensors.
+ * If a solution is found, it logs the details including the timing, result size, and the strategies used.
+ * If the resultant footprint is smaller than the current `upperbound_`, it updates the `upperbound_` and `best_sol_`.
+ *
+ * @param pFootprint Shared pointer to the current FootPrint object.
+ * @return The optimized memory footprint size.
+ */
+size_t SomasSolverCore::Search(const std::shared_ptr<FootPrint> &pFootprint) {
+  size_t result = 0;
+  FastHeuristic fh;
+  MS_LOG(INFO) << "Calling FastSolver Search for " << block_tensors_.size() << " tensors ";
+  auto start = std::chrono::system_clock::now();
+  if (fh.Eval(&block_tensors_, pFootprint, &constraints_)) {
+    result = pFootprint->Result();
+    auto end = std::chrono::system_clock::now();
+    timing_ = std::chrono::duration_cast<std::chrono::milliseconds>((end - start)).count();
+    // print for serial all_ or multi thread solver
+    if (all_ || is_multi_thread_valid_) {
+      const double giga = 1073741824.;
+      MS_LOG(INFO) << timing_ << " ms\t" << sol_count_ + 1 << "/"
+                   << static_cast<size_t>(kNumFittingTypes) * static_cast<size_t>(kNumAlgorithmTypes) *
+                        static_cast<size_t>(kNumSortingTypes)
+                   << "\t" << result << " Bytes (" << result / giga << " GB)\t" << algorithmTypeNames[algorithm_]
+                   << "\t" << sortingNames[sort_strategy_] << "\t" << branchingNames[branching_strategy_];
+    }
+  } else {
+    MS_LOG(INFO) << "FastSolver could not find solution";
+  }
+
+  if (result < upperbound_) {
+    upperbound_ = result;
+    best_sol_ = pFootprint->m_solId_;
+  }
+
+  return upperbound_;
+}
+
+/**
+ * @brief Appends the lifelong tensors to the final solution.
+ *
+ * Calculates the offset for each lifelong tensor and adds it to the upperbound.
+ * The total memory added by the lifelong tensors is logged for debugging.
+ */
+void SomasSolverCore::AppendLifelongTensors() {
+  MS_LOG(DEBUG) << "Appending lifelong tensors to solution";
+  size_t offset = upperbound_;
+  std::map<size_t, SomasSolverTensorDescPtr> lifelongTensors;
+  for (auto &t : tensors_) {
+    if (t.second->lifelong_) {
+      (void)lifelongTensors.emplace(t.first, t.second);
+    }
+  }
+  for (auto &t : lifelongTensors) {
+    auto &pTensor = t.second;
+    pTensor->offset_ = offset;
+    offset += pTensor->size_;
+  }
+  upperbound_ += lifelong_memory_;
+  MS_LOG(DEBUG) << lifelong_memory_ << " bytes from lifelong tensors added to solution";
+}
+
+/**
+ * @brief Initiates the process of finding optimized solutions for tensor allocations.
+ *
+ * Creates a new FootPrint object, sets its properties, and initiates a search for optimized solutions.
+ * Appends the lifelong tensors to the solution and cleans up the FootPrint object after the process.
+ *
+ * @return The optimized memory footprint size after appending lifelong tensors.
+ */
+size_t SomasSolverCore::FindSolutions() {
+  MS_LOG(DEBUG) << "Start allocating blocks,offset strategy: " << branchingNames[branching_strategy_];
+
+  std::shared_ptr<FootPrint> pFootprint = std::make_shared<FootPrint>();
+  pFootprint->setBranchingStrategy(static_cast<uint32_t>(branching_strategy_));
+  pFootprint->setCurrentSol(sol_count_);
+  pFootprint->setAlgorithm(static_cast<uint32_t>(algorithm_));
+  Search(pFootprint);
+  AppendLifelongTensors();
+  Destroy(pFootprint);
+  return upperbound_;
+}
+
+/**
+ * @brief Destroys a linked list of FootPrint objects.
+ *
+ * Iterates through the linked list of FootPrint objects and sets each pointer to null, effectively destroying the list.
+ *
+ * @param pFootprint Shared pointer to the head of the FootPrint linked list.
+ */
+void SomasSolverCore::Destroy(std::shared_ptr<FootPrint> &pFootprint) {
+  while (pFootprint != nullptr) {
+    if (pFootprint->Next() != nullptr) {
+      std::shared_ptr<FootPrint> &p = pFootprint;
+      pFootprint = pFootprint->Next();
+      p = nullptr;
+    } else {
+      pFootprint = nullptr;
+    }
+  }
+}
+}  // namespace somas
+}  // namespace mindspore

mindspore/ccsrc/backend/common/somas/somas_solver_pre.cc

@@ -28,52 +28,103 @@
 namespace mindspore {
 namespace somas {
 constexpr auto kSolNumThresholdMultiThread = 8;
+
+/**
+ * @brief Check the validity of tensors at specified indices.
+ *
+ * This function checks whether the tensors at the given indices are non-null 
+ * and logs a warning if they already have a left or right tensor.
+ *
+ * @param pTensors Pointer to the map containing tensors descriptions.
+ * @param index1 The index of the first tensor to be checked.
+ * @param index2 The index of the second tensor to be checked.
+ * @return Returns SUCCESS if the tensors are valid, otherwise returns FAILED.
+ */
 Status SomasSolverPre::CheckTensors(const TensorsDescMap *pTensors, uint32_t index1, uint32_t index2) {
+  // Dereferencing the pointer to TensorsDescMap to obtain the actual map.
   auto tensors = *pTensors;
+
+  // Checking if the tensor at index1 is NULL.
   if (tensors[index1] == nullptr) {
     MS_LOG(WARNING) << "NULL tensor received in continuous constraint (tensor index " << index1 << ")";
     return FAILED;
   }
+
+  // Checking if the tensor at index2 is NULL.
   if (tensors[index2] == nullptr) {
     MS_LOG(WARNING) << "NULL tensor received in continuous constraint (tensor index " << index2 << ")";
     return FAILED;
   }
 
+  // Logging a warning if tensor at index1 already has a right tensor.
   if (tensors[index1]->right_)
     MS_LOG(WARNING) << "Warning:tensor " << index1
                     << " already has a right tensor (id: " << tensors[index1]->right_->index_;
+
+  // Logging a warning if tensor at index2 already has a left tensor.
   if (tensors[index2]->left_)
     MS_LOG(WARNING) << "Warning:tensor " << index2
                     << " already has a left tensor (id: " << tensors[index2]->left_->index_;
+
   return SUCCESS;
 }
+
+/**
+ * @brief Add contiguous information in the tensor map.
+ *
+ * This function iterates through each sublist in continuous_v, and for each pair of indices, 
+ * sets the right_ pointer of the tensor at index1 and the left_ pointer of the tensor at index2.
+ *
+ * @param continuous_v Vector of vectors containing indices information.
+ * @param pTensors Pointer to the map containing tensors descriptions.
+ * @return Returns SUCCESS if the contiguous information is added successfully, otherwise returns FAILED.
+ */
 Status SomasSolverPre::AddContiguousInfoInMap(const vector<vector<size_t>> &continuous_v, TensorsDescMap *pTensors) {
+  // Dereferencing the pointer to TensorsDescMap to access the map.
   auto &tensors = *pTensors;
-  // creating S Lists
+
+  // Iterating through each sublist in continuous_v.
   for (auto &aux : continuous_v) {
+    // Iterating through each element in the sublist, except the last one.
     for (size_t i = 0; i < aux.size() - 1; i++) {
       auto index1 = aux[i];
       auto index2 = aux[i + 1];
+      // Checking if the tensors at index1 and index2 are valid.
       if (CheckTensors(pTensors, SizeToUint(index1), SizeToUint(index2)) == FAILED) {
         return FAILED;
       }
+      // Setting the right_ pointer of tensor at index1 and the left_ pointer of tensor at index2.
       tensors[index1]->right_ = tensors[index2];
       tensors[index2]->left_ = tensors[index1];
     }
   }
   return SUCCESS;
 }
+
+/**
+ * @brief Add contiguous information in multiple tensor maps.
+ *
+ * This function updates the pointers for each tensor in each solution in vecTensorsMap based on contiguous information.
+ *
+ * @param continuous_v Vector of vectors containing indices information.
+ * @param vecTensorsMap Pointer to the vector containing multiple tensor maps.
+ * @param pTensors Pointer to the original map containing tensors descriptions.
+ * @return Returns SUCCESS if the contiguous information is added successfully in all maps, otherwise returns FAILED.
+ */
 Status SomasSolverPre::AddContiguousInfoInMultiMaps(const vector<vector<size_t>> &continuous_v,
                                                     vector<TensorsDescMap> *vecTensorsMap,
                                                     const TensorsDescMap *pTensors) {
-  // creating S Lists
+  // Iterating through each sublist in continuous_v.
   for (auto &aux : continuous_v) {
+    // Iterating through each element in the sublist, except the last one.
     for (size_t i = 0; i < aux.size() - 1; i++) {
       auto index1 = aux[i];
       auto index2 = aux[i + 1];
+      // Checking if the tensors at index1 and index2 are valid.
       if (CheckTensors(pTensors, SizeToUint(index1), SizeToUint(index2)) == FAILED) {
         return FAILED;
       }
+      // Updating the pointers for each solution in vecTensorsMap.
       for (size_t sol = 0; sol < vecTensorsMap->size(); sol++) {
         auto &tensors_sol = (*vecTensorsMap)[sol];
         tensors_sol[index1]->right_ = tensors_sol[index2];
@@ -83,18 +134,67 @@ Status SomasSolverPre::AddContiguousInfoInMultiMaps(const vector<vector<size_t>>
   }
   return SUCCESS;
 }
+
+/**
+ * @brief Create a vector of tensor maps.
+ *
+ * This function creates a vector of TensorsDescMap with total_sol number of elements and initializes them 
+ * based on the input tensor map.
+ *
+ * @param tensors Reference to the input map containing tensors descriptions.
+ * @param total_sol The total number of solutions (size of the vector to be created).
+ * @return Returns a vector of TensorsDescMap.
+ */
 vector<TensorsDescMap> SomasSolverPre::CreateTensorsMaps(const TensorsDescMap &tensors, size_t total_sol) {
+  // Creating a vector of TensorsDescMap with total_sol number of elements.
   vector<TensorsDescMap> vecTensorsMap(total_sol);
+  
+  // Assigning the first element of vecTensorsMap as the input tensors map.
   vecTensorsMap[0] = tensors;
+
+  // Iterating through each tensor in the input tensors map.
   for (auto &pairT : tensors) {
+    // Creating a new TensorsDescMap for each solution except the first one.
     for (size_t sol = 1; sol < total_sol; sol++) {
       SomasSolverTensorDesc newDesc = *(pairT.second.get());
       SomasSolverTensorDescPtr newDescPtr = std::make_shared<SomasSolverTensorDesc>(newDesc);
+      // Adding the new TensorsDescMap to vecTensorsMap.
       (void)vecTensorsMap[sol].emplace(pairT.first, newDescPtr);
     }
   }
   return vecTensorsMap;
 }
+
+/**
+ * @brief Executes memory allocation optimization for tensors in a computational graph.
+ *
+ * The function SomasSolverPre::Solving is responsible for optimizing the memory allocation of tensors within 
+ * a given computational graph. It utilizes various algorithmic strategies, sorting types, and fitting mechanisms
+ * to explore different possible memory allocation solutions. It can operate both in multi-threaded and 
+ * single-threaded modes depending on the system conditions and parameters passed.
+ *
+ * The method initializes the constants representing the number of available strategies and calculates the total
+ * number of solutions based on the Cartesian product of the available types of sorting, fitting, and algorithms.
+ * It assesses the availability of multi-threading and decides the execution path accordingly.
+ *
+ * In a multi-threaded scenario, the method generates various SomasSolverCore instances for every combination of 
+ * strategies. These instances are tasked with finding memory allocation solutions in parallel. Post-execution,
+ * the solutions are evaluated, and the best one is chosen based on the smallest memory footprint.
+ *
+ * In a single-threaded scenario, a single SomasSolverCore instance is generated and tasked with finding the 
+ * memory allocation solution using the specified strategies.
+ *
+ * @param graph Pointer to the computational graph.
+ * @param ptensors Pointer to the map containing tensor descriptions.
+ * @param pConstraints Pointer to the vector containing constraints.
+ * @param continuous_v Vector containing information about continuous memory requirements.
+ * @param bVerifySolution Boolean flag to indicate whether to verify the solution.
+ * @param ball Boolean flag to indicate whether to use all available strategies.
+ * @param sorting Enum representing the sorting type to be used.
+ * @param fitting Enum representing the fitting type to be used.
+ * @param algorithm Enum representing the algorithm type to be used.
+ * @return Status Indicates the success or failure of the function execution.
+ */
 Status SomasSolverPre::Solving(const session::KernelGraph *graph, TensorsDescMap *ptensors,
                                const std::vector<DynamicBitSet> *pConstraints,
                                const vector<vector<size_t>> &continuous_v, bool bVerifySolution, bool ball,
@@ -102,22 +202,28 @@ Status SomasSolverPre::Solving(const session::KernelGraph *graph, TensorsDescMap
   Status ret = SUCCESS;
   try {
     TensorsDescMap &tensors = *ptensors;
+    // Initializing various strategy types and calculating the total number of solutions.
     constexpr size_t numSortingTypes = static_cast<size_t>(kNumSortingTypes);
     constexpr size_t numFittingTypes = static_cast<size_t>(kNumFittingTypes);
     constexpr size_t numAlgorithmTypes = static_cast<size_t>(kNumAlgorithmTypes);
     constexpr size_t total_sol = numSortingTypes * numFittingTypes * numAlgorithmTypes;
+
+    // Assessing multi-threading availability and requirements.
     size_t process_num = common::ThreadPool::GetInstance().GetSyncRunThreadNum();
     bool isMultiThreadPermit = ball && process_num >= total_sol && total_sol > 1;
     bool isMultiThreadValid = isMultiThreadPermit && (total_sol > kSolNumThresholdMultiThread ||
                                                       kParallelComputeSizeThreshold <= tensors.size());
     const double giga = 1024. * 1024. * 1024.;
     if (isMultiThreadValid) {
+      // Multi-threaded Scenario:
+      // Creating SomasSolverCore instances for each combination of strategies and running them in parallel.
       vector<std::shared_ptr<SomasSolverCore>> solvers;
       std::vector<common::Task> tasks;
       vector<TensorsDescMap> vecTensorsMap = CreateTensorsMaps(tensors, total_sol);
       if (AddContiguousInfoInMultiMaps(continuous_v, &vecTensorsMap, ptensors) == FAILED) {
         return FAILED;
       }
+      // Initializing and configuring each solver with respective strategies.
       auto start = std::chrono::system_clock::now();
       for (size_t algorithm_strategy = 0, sol = 0; algorithm_strategy < numAlgorithmTypes; algorithm_strategy++) {
         for (size_t sort_strategy = 0; sort_strategy < numSortingTypes; sort_strategy++) {
@@ -138,6 +244,7 @@ Status SomasSolverPre::Solving(const session::KernelGraph *graph, TensorsDescMap
           }
         }
       }
+      // Executing all tasks in parallel and selecting the best solution post-execution.
       common::ThreadPool::GetInstance().SyncRun(tasks);
       size_t best_sol = 0, worst = 0, best = SIZE_MAX, best_timing = SIZE_MAX;
       for (size_t sol = 0; sol < total_sol; sol++) {
@@ -173,6 +280,8 @@ Status SomasSolverPre::Solving(const session::KernelGraph *graph, TensorsDescMap
       MS_LOG(INFO) << "Spread:" << static_cast<double>((worst - best) / static_cast<double>(best * kFloatPresent))
                    << " %%";
     } else {
+      // Single-threaded Scenario:
+      // Creating a single SomasSolverCore instance and executing it with the specified strategies.
       if (AddContiguousInfoInMap(continuous_v, ptensors) == FAILED) {
         return FAILED;
       }
@@ -195,6 +304,17 @@ Status SomasSolverPre::Solving(const session::KernelGraph *graph, TensorsDescMap
   return ret;
 }
 
+/**
+ * @brief Log information of tensors, constraints, and continuous segments.
+ *
+ * This function logs various information including solver input, solver output, and tensor relations,
+ * if the save_graphs flag is enabled in the context.
+ *
+ * @param graph Pointer to the kernel graph.
+ * @param tensors Reference to the map containing tensors descriptions.
+ * @param pConstraints Pointer to the vector of DynamicBitSet representing constraints.
+ * @param continuous_v Vector of vectors containing continuous segments information.
+ */
 void SomasSolverPre::Log(const session::KernelGraph *graph, const TensorsDescMap &tensors,
                          const std::vector<DynamicBitSet> *pConstraints, const vector<vector<size_t>> &continuous_v) {
   auto context_ptr = MsContext::GetInstance();
@@ -208,6 +328,14 @@ void SomasSolverPre::Log(const session::KernelGraph *graph, const TensorsDescMap
   TensorRelationLog(pConstraints, graph);
 }
 
+/**
+ * @brief Log tensor relations.
+ *
+ * This function logs the relations between tensors by writing the somas_tensor_relation.ir file.
+ *
+ * @param pConstraints Pointer to the vector of DynamicBitSet representing constraints.
+ * @param graph Pointer to the kernel graph.
+ */
 void SomasSolverPre::TensorRelationLog(const std::vector<DynamicBitSet> *pConstraints,
                                        const session::KernelGraph *graph) {
   MS_LOG(INFO) << "SomasSolver::Log Writing somas_tensor_relation.ir..";
@@ -228,6 +356,15 @@ void SomasSolverPre::TensorRelationLog(const std::vector<DynamicBitSet> *pConstr
   MS_LOG(INFO) << "SomasSolver somas_tensor_relation Log done";
 }
 
+/**
+ * @brief Log solver input information.
+ *
+ * This function logs the input information to the solver by writing the somas_solver_input.ir file.
+ *
+ * @param graph Pointer to the kernel graph.
+ * @param tensors Reference to the map containing tensors descriptions.
+ * @param continuous_v Vector of vectors containing continuous segments information.
+ */
 void SomasSolverPre::SolverInputLog(const session::KernelGraph *graph, const TensorsDescMap &tensors,
                                     const vector<vector<size_t>> &continuous_v) {
   MS_LOG(INFO) << "SomasSolver::Log Writing somas_solver_input..";
@@ -252,6 +389,14 @@ void SomasSolverPre::SolverInputLog(const session::KernelGraph *graph, const Ten
   MS_LOG(INFO) << "SomasSolver input Log done";
 }
 
+/**
+ * @brief Log solver output information.
+ *
+ * This function logs the output information from the solver by writing the somas_solver_output.ir file.
+ *
+ * @param graph Pointer to the kernel graph.
+ * @param tensors Reference to the map containing tensors descriptions.
+ */
 void SomasSolverPre::SolverOutputLog(const session::KernelGraph *graph, const TensorsDescMap &tensors) const {
   MS_LOG(INFO) << "SomasSolver::Log Writing somas_solver_output_..";
   auto context_ptr = MsContext::GetInstance();

mindspore/ccsrc/backend/common/somas/somas_solver_pre.h

@@ -66,14 +66,17 @@ enum FittingType {
 };
 
 class DynamicBitSet {
-  const size_t bit_width_ = 64;
+  const size_t bit_width_ = 64;  // The number of bits that each element of the bitset can hold.
 
+  // Helper function to calculate the index of the element in the 'bit_' vector that holds the bit at the given position.
   inline size_t GetIndex(size_t index) const { return index / bit_width_; }
 
+  // Helper function to create a bitmask for the bit at the given position within its corresponding uint64_t element.
   inline uint64_t GetBitMask(size_t index) const {
     return (((uint64_t)0x1) << ((bit_width_ - 1) - (index % bit_width_)));
   }
 
+  // Resets all elements of the 'bit_' vector to the given value.
   inline void Reset(uint64_t val) {
     bit_.clear();
     for (size_t i = 0; i < bit_size_; i++) {
@@ -82,15 +85,18 @@ class DynamicBitSet {
   }
 
  public:
-  size_t bit_size_;
-  std::vector<uint64_t> bit_;
+  size_t bit_size_;  // The number of uint64_t elements required to hold the bitset.
+  std::vector<uint64_t> bit_;  // The vector of uint64_t elements that represents the bitset.
+
+  // Constructor that initializes the bitset with a given number of bits, all set to false.
   explicit DynamicBitSet(size_t count) {
     bit_size_ = (count + bit_width_ - 1) / bit_width_;
     Reset(0x0);
   }
 
-  ~DynamicBitSet() = default;
+  ~DynamicBitSet() = default;  // Default destructor.
 
+  // Sets the bit at the given position to true.
   void SetBitTrue(size_t index, bool log = false) {
     if (log) {
       MS_LOG(INFO) << GetIndex(index) << " " << GetBitMask(index);
@@ -98,10 +104,13 @@ class DynamicBitSet {
     bit_[GetIndex(index)] |= GetBitMask(index);
   }
 
+  // Sets the bit at the given position to false.
   void SetBitFalse(size_t index) { bit_[GetIndex(index)] &= (~GetBitMask(index)); }
 
+  // Returns whether the bit at the given position is true.
   bool IsBitTrue(size_t index) const { return (bit_[GetIndex(index)] & GetBitMask(index)) != 0x0; }
 
+  // Counts and returns the number of true bits in the bitset.
   size_t CountOnesNum() const {
     size_t ret = 0;
     static char ones_num_in_hex[] = "\0\1\1\2\1\2\2\3\1\2\2\3\2\3\3\4";
@@ -120,6 +129,7 @@ class DynamicBitSet {
     return ret;
   }
 
+  // Logs the bitset for debugging purposes.
   void Log() {
     std::cout << "Start Print Bitset ";
     for (size_t i = 0; i < bit_size_; i++) {
@@ -128,6 +138,7 @@ class DynamicBitSet {
     std::cout << std::endl;
   }
 
+  // Performs a bitwise OR operation between two bitsets and stores the result in the first bitset.
   friend void Union(DynamicBitSet *a, DynamicBitSet *b) {
     for (size_t i = 0; i < (*a).bit_size_; i++) {
       (*a).bit_[i] |= (*b).bit_[i];

mindspore/ccsrc/backend/common/somas/somas_solver_pre_old.cc

@@ -0,0 +1,432 @@
+/**
+ * Copyright 2020-2021 Huawei Technologies Co., Ltd
+
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+
+ * http://www.apache.org/licenses/LICENSE-2.0
+
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+*/
+
+#include <cstdio>
+#include <fstream>
+#include <memory>
+#include <string>
+#include <utility>
+#include "include/common/thread_pool.h"
+
+#include "backend/common/somas/somas_solver_core.h"
+#include "backend/common/somas/somas_solver_pre.h"
+#include "include/common/debug/common.h"
+
+namespace mindspore {
+namespace somas {
+constexpr auto kSolNumThresholdMultiThread = 8;
+
+/**
+ * @brief Check the validity of tensors at specified indices.
+ *
+ * This function checks whether the tensors at the given indices are non-null 
+ * and logs a warning if they already have a left or right tensor.
+ *
+ * @param pTensors Pointer to the map containing tensors descriptions.
+ * @param index1 The index of the first tensor to be checked.
+ * @param index2 The index of the second tensor to be checked.
+ * @return Returns SUCCESS if the tensors are valid, otherwise returns FAILED.
+ */
+Status SomasSolverPre::CheckTensors(const TensorsDescMap *pTensors, uint32_t index1, uint32_t index2) {
+  // Dereferencing the pointer to TensorsDescMap to obtain the actual map.
+  auto tensors = *pTensors;
+
+  // Checking if the tensor at index1 is NULL.
+  if (tensors[index1] == nullptr) {
+    MS_LOG(WARNING) << "NULL tensor received in continuous constraint (tensor index " << index1 << ")";
+    return FAILED;
+  }
+
+  // Checking if the tensor at index2 is NULL.
+  if (tensors[index2] == nullptr) {
+    MS_LOG(WARNING) << "NULL tensor received in continuous constraint (tensor index " << index2 << ")";
+    return FAILED;
+  }
+
+  // Logging a warning if tensor at index1 already has a right tensor.
+  if (tensors[index1]->right_)
+    MS_LOG(WARNING) << "Warning:tensor " << index1
+                    << " already has a right tensor (id: " << tensors[index1]->right_->index_;
+
+  // Logging a warning if tensor at index2 already has a left tensor.
+  if (tensors[index2]->left_)
+    MS_LOG(WARNING) << "Warning:tensor " << index2
+                    << " already has a left tensor (id: " << tensors[index2]->left_->index_;
+
+  return SUCCESS;
+}
+
+/**
+ * @brief Add contiguous information in the tensor map.
+ *
+ * This function iterates through each sublist in continuous_v, and for each pair of indices, 
+ * sets the right_ pointer of the tensor at index1 and the left_ pointer of the tensor at index2.
+ *
+ * @param continuous_v Vector of vectors containing indices information.
+ * @param pTensors Pointer to the map containing tensors descriptions.
+ * @return Returns SUCCESS if the contiguous information is added successfully, otherwise returns FAILED.
+ */
+Status SomasSolverPre::AddContiguousInfoInMap(const vector<vector<size_t>> &continuous_v, TensorsDescMap *pTensors) {
+  // Dereferencing the pointer to TensorsDescMap to access the map.
+  auto &tensors = *pTensors;
+
+  // Iterating through each sublist in continuous_v.
+  for (auto &aux : continuous_v) {
+    // Iterating through each element in the sublist, except the last one.
+    for (size_t i = 0; i < aux.size() - 1; i++) {
+      auto index1 = aux[i];
+      auto index2 = aux[i + 1];
+      // Checking if the tensors at index1 and index2 are valid.
+      if (CheckTensors(pTensors, SizeToUint(index1), SizeToUint(index2)) == FAILED) {
+        return FAILED;
+      }
+      // Setting the right_ pointer of tensor at index1 and the left_ pointer of tensor at index2.
+      tensors[index1]->right_ = tensors[index2];
+      tensors[index2]->left_ = tensors[index1];
+    }
+  }
+  return SUCCESS;
+}
+
+/**
+ * @brief Add contiguous information in multiple tensor maps.
+ *
+ * This function updates the pointers for each tensor in each solution in vecTensorsMap based on contiguous information.
+ *
+ * @param continuous_v Vector of vectors containing indices information.
+ * @param vecTensorsMap Pointer to the vector containing multiple tensor maps.
+ * @param pTensors Pointer to the original map containing tensors descriptions.
+ * @return Returns SUCCESS if the contiguous information is added successfully in all maps, otherwise returns FAILED.
+ */
+Status SomasSolverPre::AddContiguousInfoInMultiMaps(const vector<vector<size_t>> &continuous_v,
+                                                    vector<TensorsDescMap> *vecTensorsMap,
+                                                    const TensorsDescMap *pTensors) {
+  // Iterating through each sublist in continuous_v.
+  for (auto &aux : continuous_v) {
+    // Iterating through each element in the sublist, except the last one.
+    for (size_t i = 0; i < aux.size() - 1; i++) {
+      auto index1 = aux[i];
+      auto index2 = aux[i + 1];
+      // Checking if the tensors at index1 and index2 are valid.
+      if (CheckTensors(pTensors, SizeToUint(index1), SizeToUint(index2)) == FAILED) {
+        return FAILED;
+      }
+      // Updating the pointers for each solution in vecTensorsMap.
+      for (size_t sol = 0; sol < vecTensorsMap->size(); sol++) {
+        auto &tensors_sol = (*vecTensorsMap)[sol];
+        tensors_sol[index1]->right_ = tensors_sol[index2];
+        tensors_sol[index2]->left_ = tensors_sol[index1];
+      }
+    }
+  }
+  return SUCCESS;
+}
+
+/**
+ * @brief Create a vector of tensor maps.
+ *
+ * This function creates a vector of TensorsDescMap with total_sol number of elements and initializes them 
+ * based on the input tensor map.
+ *
+ * @param tensors Reference to the input map containing tensors descriptions.
+ * @param total_sol The total number of solutions (size of the vector to be created).
+ * @return Returns a vector of TensorsDescMap.
+ */
+vector<TensorsDescMap> SomasSolverPre::CreateTensorsMaps(const TensorsDescMap &tensors, size_t total_sol) {
+  // Creating a vector of TensorsDescMap with total_sol number of elements.
+  vector<TensorsDescMap> vecTensorsMap(total_sol);
+  
+  // Assigning the first element of vecTensorsMap as the input tensors map.
+  vecTensorsMap[0] = tensors;
+
+  // Iterating through each tensor in the input tensors map.
+  for (auto &pairT : tensors) {
+    // Creating a new TensorsDescMap for each solution except the first one.
+    for (size_t sol = 1; sol < total_sol; sol++) {
+      SomasSolverTensorDesc newDesc = *(pairT.second.get());
+      SomasSolverTensorDescPtr newDescPtr = std::make_shared<SomasSolverTensorDesc>(newDesc);
+      // Adding the new TensorsDescMap to vecTensorsMap.
+      (void)vecTensorsMap[sol].emplace(pairT.first, newDescPtr);
+    }
+  }
+  return vecTensorsMap;
+}
+
+/**
+ * @brief Executes memory allocation optimization for tensors in a computational graph.
+ *
+ * The function SomasSolverPre::Solving is responsible for optimizing the memory allocation of tensors within 
+ * a given computational graph. It utilizes various algorithmic strategies, sorting types, and fitting mechanisms
+ * to explore different possible memory allocation solutions. It can operate both in multi-threaded and 
+ * single-threaded modes depending on the system conditions and parameters passed.
+ *
+ * The method initializes the constants representing the number of available strategies and calculates the total
+ * number of solutions based on the Cartesian product of the available types of sorting, fitting, and algorithms.
+ * It assesses the availability of multi-threading and decides the execution path accordingly.
+ *
+ * In a multi-threaded scenario, the method generates various SomasSolverCore instances for every combination of 
+ * strategies. These instances are tasked with finding memory allocation solutions in parallel. Post-execution,
+ * the solutions are evaluated, and the best one is chosen based on the smallest memory footprint.
+ *
+ * In a single-threaded scenario, a single SomasSolverCore instance is generated and tasked with finding the 
+ * memory allocation solution using the specified strategies.
+ *
+ * @param graph Pointer to the computational graph.
+ * @param ptensors Pointer to the map containing tensor descriptions.
+ * @param pConstraints Pointer to the vector containing constraints.
+ * @param continuous_v Vector containing information about continuous memory requirements.
+ * @param bVerifySolution Boolean flag to indicate whether to verify the solution.
+ * @param ball Boolean flag to indicate whether to use all available strategies.
+ * @param sorting Enum representing the sorting type to be used.
+ * @param fitting Enum representing the fitting type to be used.
+ * @param algorithm Enum representing the algorithm type to be used.
+ * @return Status Indicates the success or failure of the function execution.
+ */
+Status SomasSolverPre::Solving(const session::KernelGraph *graph, TensorsDescMap *ptensors,
+                               const std::vector<DynamicBitSet> *pConstraints,
+                               const vector<vector<size_t>> &continuous_v, bool bVerifySolution, bool ball,
+                               SortingType sorting, FittingType fitting, AlgorithmType algorithm) {
+  Status ret = SUCCESS;
+  try {
+    TensorsDescMap &tensors = *ptensors;
+    // Initializing various strategy types and calculating the total number of solutions.
+    constexpr size_t numSortingTypes = static_cast<size_t>(kNumSortingTypes);
+    constexpr size_t numFittingTypes = static_cast<size_t>(kNumFittingTypes);
+    constexpr size_t numAlgorithmTypes = static_cast<size_t>(kNumAlgorithmTypes);
+    constexpr size_t total_sol = numSortingTypes * numFittingTypes * numAlgorithmTypes;
+
+    // Assessing multi-threading availability and requirements.
+    size_t process_num = common::ThreadPool::GetInstance().GetSyncRunThreadNum();
+    bool isMultiThreadPermit = ball && process_num >= total_sol && total_sol > 1;
+    bool isMultiThreadValid = isMultiThreadPermit && (total_sol > kSolNumThresholdMultiThread ||
+                                                      kParallelComputeSizeThreshold <= tensors.size());
+    const double giga = 1024. * 1024. * 1024.;
+    if (isMultiThreadValid) {
+      // Multi-threaded Scenario:
+      // Creating SomasSolverCore instances for each combination of strategies and running them in parallel.
+      vector<std::shared_ptr<SomasSolverCore>> solvers;
+      std::vector<common::Task> tasks;
+      vector<TensorsDescMap> vecTensorsMap = CreateTensorsMaps(tensors, total_sol);
+      if (AddContiguousInfoInMultiMaps(continuous_v, &vecTensorsMap, ptensors) == FAILED) {
+        return FAILED;
+      }
+      // Initializing and configuring each solver with respective strategies.
+      auto start = std::chrono::system_clock::now();
+      for (size_t algorithm_strategy = 0, sol = 0; algorithm_strategy < numAlgorithmTypes; algorithm_strategy++) {
+        for (size_t sort_strategy = 0; sort_strategy < numSortingTypes; sort_strategy++) {
+          for (size_t branching_strategy = 0; branching_strategy < numFittingTypes; branching_strategy++) {
+            std::shared_ptr<SomasSolverCore> pSolver =
+              std::make_shared<SomasSolverCore>(vecTensorsMap[sol], pConstraints, sol);
+            pSolver->SetAlgorithmStrategy(AlgorithmType(algorithm));
+            pSolver->SetSortingStrategy(SortingType(sort_strategy));
+            pSolver->SetFittingStrategy(FittingType(branching_strategy));
+            pSolver->SetAllStrategies(false);
+            pSolver->VerifySolution(bVerifySolution);
+            auto task = [pSolver]() {
+              return pSolver->MemoryAllocationSolver() == SUCCESS ? common::SUCCESS : common::FAIL;
+            };
+            tasks.emplace_back(task);
+            solvers.emplace_back(pSolver);
+            sol++;
+          }
+        }
+      }
+      // Executing all tasks in parallel and selecting the best solution post-execution.
+      common::ThreadPool::GetInstance().SyncRun(tasks);
+      size_t best_sol = 0, worst = 0, best = SIZE_MAX, best_timing = SIZE_MAX;
+      for (size_t sol = 0; sol < total_sol; sol++) {
+        auto &solver = solvers[sol];
+        auto &upperbound = solver->GetUpperbound();
+        if (upperbound > worst) {
+          worst = upperbound;
+        }
+        if (upperbound <= best) {
+          best = upperbound;
+          best_sol = sol;
+          best_timing = LongToSize(solver->timing_);
+        }
+      }
+      auto end = std::chrono::system_clock::now();
+      size_t total_time = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
+      auto &best_solver = solvers[best_sol];
+      for (auto &tensor : tensors) {
+        *(tensor.second.get()) = *(vecTensorsMap[best_sol][tensor.first]);
+      }
+      max_offset_ = best_solver->GetUpperbound();
+      constexpr float kFloatPresent = 100.0;
+      MS_LOG(INFO) << "SOMAS SOLVER RESUME:";
+      MS_LOG(INFO) << "Best Solution:[" << 1 + best_sol << "/" << total_sol << "] ";
+      MS_LOG(INFO) << "Best result:" << best << " Bytes " << (best) / (giga) << " GB ("
+                   << (best - best_solver->Getlifelongmemory()) / (giga) << " GB + "
+                   << best_solver->Getlifelongmemory() / (giga) << " GB from lifelong tensors)";
+      MS_LOG(INFO) << "Best timing:" << best_timing << " ms";
+      MS_LOG(INFO) << "Best algorithm: " << algorithmTypeNames[best_solver->algorithm_];
+      MS_LOG(INFO) << "Best sorting strategy: " << sortingNames[best_solver->sort_strategy_];
+      MS_LOG(INFO) << "Best offset strategy: " << branchingNames[best_solver->branching_strategy_];
+      MS_LOG(INFO) << "Time elapsed: " << total_time << " ms";
+      MS_LOG(INFO) << "Spread:" << static_cast<double>((worst - best) / static_cast<double>(best * kFloatPresent))
+                   << " %%";
+    } else {
+      // Single-threaded Scenario:
+      // Creating a single SomasSolverCore instance and executing it with the specified strategies.
+      if (AddContiguousInfoInMap(continuous_v, ptensors) == FAILED) {
+        return FAILED;
+      }
+      std::shared_ptr<SomasSolverCore> pSolver = std::make_shared<SomasSolverCore>(tensors, pConstraints, 0, false);
+      pSolver->SetAlgorithmStrategy(algorithm);
+      pSolver->SetSortingStrategy(sorting);
+      pSolver->SetFittingStrategy(fitting);
+      pSolver->SetAllStrategies(ball);
+      pSolver->VerifySolution(bVerifySolution);
+      if (SUCCESS == (pSolver->MemoryAllocationSolver())) {
+        max_offset_ = pSolver->GetUpperbound();
+        MS_LOG(INFO) << "SomasSolver::Solving SUCCESS";
+        MS_LOG(INFO) << "SomasSolver::Solving RESULT: " << max_offset_ << " (" << max_offset_ / (giga) << " GB)";
+      }
+    }
+    Log(graph, tensors, pConstraints, continuous_v);
+  } catch (const std::exception &e) {
+    MS_LOG(EXCEPTION) << "SomasSolver::Solving FAILED: " << e.what();
+  }
+  return ret;
+}
+
+/**
+ * @brief Log information of tensors, constraints, and continuous segments.
+ *
+ * This function logs various information including solver input, solver output, and tensor relations,
+ * if the save_graphs flag is enabled in the context.
+ *
+ * @param graph Pointer to the kernel graph.
+ * @param tensors Reference to the map containing tensors descriptions.
+ * @param pConstraints Pointer to the vector of DynamicBitSet representing constraints.
+ * @param continuous_v Vector of vectors containing continuous segments information.
+ */
+void SomasSolverPre::Log(const session::KernelGraph *graph, const TensorsDescMap &tensors,
+                         const std::vector<DynamicBitSet> *pConstraints, const vector<vector<size_t>> &continuous_v) {
+  auto context_ptr = MsContext::GetInstance();
+  MS_EXCEPTION_IF_NULL(context_ptr);
+  bool save_graphs = context_ptr->get_param<bool>(MS_CTX_SAVE_GRAPHS_FLAG);
+  if (!save_graphs) {
+    return;
+  }
+  SolverInputLog(graph, tensors, continuous_v);
+  SolverOutputLog(graph, tensors);
+  TensorRelationLog(pConstraints, graph);
+}
+
+/**
+ * @brief Log tensor relations.
+ *
+ * This function logs the relations between tensors by writing the somas_tensor_relation.ir file.
+ *
+ * @param pConstraints Pointer to the vector of DynamicBitSet representing constraints.
+ * @param graph Pointer to the kernel graph.
+ */
+void SomasSolverPre::TensorRelationLog(const std::vector<DynamicBitSet> *pConstraints,
+                                       const session::KernelGraph *graph) {
+  MS_LOG(INFO) << "SomasSolver::Log Writing somas_tensor_relation.ir..";
+  auto context_ptr = MsContext::GetInstance();
+  MS_EXCEPTION_IF_NULL(context_ptr);
+  auto save_graphs_path = context_ptr->get_param<std::string>(MS_CTX_SAVE_GRAPHS_PATH);
+  std::string filename =
+    GetSaveGraphsPathName("somas_tensor_relation_" + std::to_string(graph->graph_id()) + ".ir", save_graphs_path);
+  std::ostringstream oss;
+  for (size_t tid1 = 0; tid1 < pConstraints->size(); tid1++) {
+    oss << 't' << tid1 << ' ';
+    for (size_t tid2 = 0; tid2 < (*pConstraints)[tid1].bit_size_; tid2++) {
+      oss << 'H' << std::hex << (*pConstraints)[tid1].bit_[tid2];
+    }
+    oss << std::endl << std::dec;
+  }
+  (void)Common::SaveStringToFile(filename, oss.str());
+  MS_LOG(INFO) << "SomasSolver somas_tensor_relation Log done";
+}
+
+/**
+ * @brief Log solver input information.
+ *
+ * This function logs the input information to the solver by writing the somas_solver_input.ir file.
+ *
+ * @param graph Pointer to the kernel graph.
+ * @param tensors Reference to the map containing tensors descriptions.
+ * @param continuous_v Vector of vectors containing continuous segments information.
+ */
+void SomasSolverPre::SolverInputLog(const session::KernelGraph *graph, const TensorsDescMap &tensors,
+                                    const vector<vector<size_t>> &continuous_v) {
+  MS_LOG(INFO) << "SomasSolver::Log Writing somas_solver_input..";
+  auto context_ptr = MsContext::GetInstance();
+  MS_EXCEPTION_IF_NULL(context_ptr);
+  auto save_graphs_path = context_ptr->get_param<std::string>(MS_CTX_SAVE_GRAPHS_PATH);
+  std::string filename =
+    GetSaveGraphsPathName("somas_solver_input_" + std::to_string(graph->graph_id()) + ".ir", save_graphs_path);
+  std::ostringstream oss;
+  for (auto &t : tensors) {
+    oss << "T " << t.second->index_ << " " << t.second->size_ << " " << t.second->lifelong_ << std::endl;
+  }
+
+  for (auto &s : continuous_v) {
+    oss << "S";
+    for (auto idx : s) {
+      oss << " " << idx;
+    }
+    oss << std::endl;
+  }
+  (void)Common::SaveStringToFile(filename, oss.str());
+  MS_LOG(INFO) << "SomasSolver input Log done";
+}
+
+/**
+ * @brief Log solver output information.
+ *
+ * This function logs the output information from the solver by writing the somas_solver_output.ir file.
+ *
+ * @param graph Pointer to the kernel graph.
+ * @param tensors Reference to the map containing tensors descriptions.
+ */
+void SomasSolverPre::SolverOutputLog(const session::KernelGraph *graph, const TensorsDescMap &tensors) const {
+  MS_LOG(INFO) << "SomasSolver::Log Writing somas_solver_output_..";
+  auto context_ptr = MsContext::GetInstance();
+  MS_EXCEPTION_IF_NULL(context_ptr);
+  auto save_graphs_path = context_ptr->get_param<std::string>(MS_CTX_SAVE_GRAPHS_PATH);
+  std::string out_filename =
+    GetSaveGraphsPathName("somas_solver_output_" + std::to_string(graph->graph_id()) + ".ir", save_graphs_path);
+  std::ostringstream oss;
+  constexpr size_t contiguous_left = 1;
+  constexpr size_t contiguous_mid = 2;
+  constexpr size_t contiguous_right = 3;
+  for (auto &t : tensors) {
+    SomasSolverTensorDescPtr tensor = t.second;
+    int continuous = 0;
+    if (tensor->left_ == nullptr && tensor->right_ != nullptr)
+      continuous = contiguous_left;
+    else if (tensor->left_ != nullptr && tensor->right_ != nullptr)
+      continuous = contiguous_mid;
+    else if (tensor->left_ != nullptr && tensor->right_ == nullptr)
+      continuous = contiguous_right;
+    const size_t alignment = 512;
+    bool size_aligned = tensor->size_ % alignment == 0;
+    bool offset_aligned = tensor->offset_ % alignment == 0;
+
+    oss << std::endl
+        << "tensor_id=" << tensor->index_ << "\tsize=" << tensor->size_ << "\toffset=" << tensor->offset_
+        << "\tcontinuous=" << continuous << "\tsize_aligned=" << size_aligned << "\toffset_aligned=" << offset_aligned;
+  }
+  (void)Common::SaveStringToFile(out_filename, oss.str());
+  MS_LOG(INFO) << "SomasSolver output Log done";
+}
+}  // namespace somas
+}  // namespace mindspore

mindspore/ccsrc/backend/graph_compiler/backend.cc

@@ -54,18 +54,36 @@
 
 namespace mindspore {
 namespace compile {
+// Return the BaseRef value as a bool variable.
 bool Backend::GetCond(const BaseRef &c, bool *value) {
   mindspore::ScopedLongRunning long_running;
   return BaseRefToBool(c, value);
 }
+// Return the BaseRef value as a integer variable.
 bool Backend::GetIndex(const BaseRef &c, int64_t *value) { return BaseRefToInt(utils::cast<ValuePtr>(c), value); }
 
+/**
+ * @brief Construct a new Backend:: Backend object
+ * This function intializes the Backend object.
+ * Turn off multi_graph_sink_, a tech to efficiently reduce data I/O (host-device interaction).
+ * @param name the identify name of the boject.
+ */
 Backend::Backend(const std::string &name) : name_(name) {
   MS_LOG(DEBUG) << "Select backend:" << name;
   convert_fn_ = MsVmConvert;
   is_multi_graph_sink_ = false;
 }
 
+/**
+ * @brief Converts a graph segment to a specific representation and returns the result.
+ *
+ * This function takes a graph segment and converts it to a specific representation(Anf Node in mindspore).
+ * It encapsulates the transformation details and provides an organized output.
+ *
+ * @param segment The graph segment to be converted.
+ * @param target The target device for the conversion.
+ * @return A `LinConvertResult` containing the converted graph information.
+ */
 LinConvertResult MsBackend::MsConvert(const GraphSegmentPtr &segment, const std::string &target) {
   MS_LOG(DEBUG) << "MsConvert";
   MS_EXCEPTION_IF_NULL(segment);
@@ -131,6 +149,14 @@ LinConvertResult MsBackend::MsConvert(const GraphSegmentPtr &segment, const std:
 }
 
 // compile set input output
+/**
+ * @brief Simulates the execution of a graph and returns the outputs.
+ *
+ * This function simulates the execution of the specified graph and returns the outputs as a `VectorRef`.
+ *
+ * @param g The graph ID of the graph to be simulated.
+ * @return A `VectorRef` containing the simulated graph outputs.
+ */
 VectorRef MsBackend::MsSimuRunGraph(const GraphId &g) {
   MS_LOG(DEBUG) << "Set graph input:" << g;
   std::vector<BaseRef> outputs;
@@ -140,6 +166,16 @@ VectorRef MsBackend::MsSimuRunGraph(const GraphId &g) {
 }
 
 namespace {
+  /**
+ * @brief Retrieves a vector of input tensors excluding those associated with value nodes.
+ *
+ * This function takes an `OpRunInfo` object containing input tensors and their associated tensor masks,
+ * and returns a vector of input tensors excluding those have the `kValueNodeTensorMask`.
+ *
+ * @param op_run_info The `OpRunInfo` object containing input tensors and tensor masks.
+ * @return A vector of `tensor::TensorPtr` containing input tensors without value node tensors.
+ * @throws Exception if the size of input tensors and tensors mask are not equal.
+ */
 std::vector<tensor::TensorPtr> GetTensorWithoutValueMask(const OpRunInfo &op_run_info) {
   std::vector<tensor::TensorPtr> tensors_without_value_node;
   const auto &input_tensors = op_run_info.input_tensors;
@@ -148,6 +184,7 @@ std::vector<tensor::TensorPtr> GetTensorWithoutValueMask(const OpRunInfo &op_run
     MS_LOG(EXCEPTION) << "Input tensors size " << input_tensors.size() << " should be equal to tensors mask size "
                       << tensors_mask.size();
   }
+  // traverse the tensors_mark to get all tensors without value nodes.
   for (size_t index = 0; index < tensors_mask.size(); ++index) {
     if (tensors_mask.at(index) != kValueNodeTensorMask) {
       (void)tensors_without_value_node.emplace_back(input_tensors.at(index));
@@ -155,22 +192,38 @@ std::vector<tensor::TensorPtr> GetTensorWithoutValueMask(const OpRunInfo &op_run
   }
   return tensors_without_value_node;
 }
-
+/**
+ * @brief Pushes input tensors(such as return variable in previous function) into a vector.
+ *
+ * This function takes an input argument `arg` and appends corresponding tensors to the `inputs` vector.
+ * It handles various input types including tensor pointers, CSR tensors, value tuples, scalars, monads,
+ * PyObjectRef, and VectorRefPtr.
+ *
+ * @param arg The input argument to be pushed.
+ * @param inputs A pointer to the vector where input tensors will be appended.
+ */
 void PushInputTensor(const BaseRef &arg, std::vector<tensor::TensorPtr> *inputs) {
   MS_EXCEPTION_IF_NULL(inputs);
+
+  // Handle tensor pointer
   if (utils::isa<tensor::TensorPtr>(arg)) {
     auto value = utils::cast<tensor::TensorPtr>(arg);
     inputs->push_back(value);
-  } else if (utils::isa<tensor::CSRTensorPtr>(arg)) {
+  }
+  // Handle CSR tensor
+  else if (utils::isa<tensor::CSRTensorPtr>(arg)) {
     auto csr = utils::cast<tensor::CSRTensorPtr>(arg);
     MS_EXCEPTION_IF_NULL(csr);
     auto csr_values = csr->GetValues();
     MS_EXCEPTION_IF_NULL(csr_values);
     inputs->push_back(csr_values);
     MS_LOG(INFO) << "For CSRTensor, push its values.";
-  } else if (utils::isa<ValuePtr>(arg)) {
+  }
+  // Handle ValuePtr (including ValueTuple, Scalar, and Monad)
+  else if (utils::isa<ValuePtr>(arg)) {
     auto value = utils::cast<ValuePtr>(arg);
     MS_EXCEPTION_IF_NULL(value);
+
     if (value->isa<ValueTuple>()) {
       auto value_tuple = value->cast<ValueTuplePtr>();
       MS_EXCEPTION_IF_NULL(value_tuple);
@@ -186,19 +239,26 @@ void PushInputTensor(const BaseRef &arg, std::vector<tensor::TensorPtr> *inputs)
     } else {
       inputs->push_back(value->cast<tensor::TensorPtr>());
     }
-  } else if (utils::isa<PyObjectRef>(arg)) {
+  }
+  // Handle PyObjectRef
+  else if (utils::isa<PyObjectRef>(arg)) {
     auto value = utils::cast<PyObjectRef>(arg).object_;
     inputs->push_back(py::cast<tensor::TensorPtr>(value));
-  } else if (utils::isa<VectorRefPtr>(arg)) {
+  }
+  // Handle VectorRefPtr
+  else if (utils::isa<VectorRefPtr>(arg)) {
     const auto &args_new = utils::cast<VectorRef>(arg);
     for (const auto &v : args_new) {
       PushInputTensor(v, inputs);
     }
-  } else {
+  }
+  // Handle unsupported input types
+  else {
     MS_LOG(WARNING) << "Invalid input type.";
   }
 }
 
+
 // Insert the front_node related tensor in the input_tensor.
 void PushTensor(const VectorRef &args, const std::vector<AnfNodePtr> &parameters, const AnfNodePtr &front_node,
                 std::vector<tensor::TensorPtr> *input_tensor) {
@@ -211,18 +271,43 @@ void PushTensor(const VectorRef &args, const std::vector<AnfNodePtr> &parameters
   PushInputTensor(args[position], input_tensor);
 }
 
+/**
+ * @brief Updates the output abstract information in the OpRunInfo structure based on the given KernelGraph.
+ *
+ * This function iterates through the execution order of the given KernelGraph and updates the abstract information
+ * in the provided OpRunInfo structure for the specified operation.
+ *
+ * @param kernel_graph The KernelGraph representing the computation graph.
+ * @param op_run_info A pointer to the OpRunInfo structure to be updated.
+ */
 void UpdateOutputAbstract(const KernelGraphPtr &kernel_graph, OpRunInfo *op_run_info) {
   MS_EXCEPTION_IF_NULL(kernel_graph);
   MS_EXCEPTION_IF_NULL(op_run_info);
+
+  // Retrieve the list of kernels in the execution order of the KernelGraph
   const auto &kernels = kernel_graph->execution_order();
+
+  // Iterate through the kernels and update the output abstract information
   for (const auto &kernel : kernels) {
     MS_EXCEPTION_IF_NULL(kernel);
+
+    // Check if the CNode name matches the target operation name
     if (common::AnfAlgo::GetCNodeName(kernel) == op_run_info->op_name) {
+      // Update the abstract information in the OpRunInfo structure
       op_run_info->abstract = kernel->abstract();
     }
   }
 }
-
+/**
+ * @brief Creates an output tensor for a given AnfNode and output index.
+ *
+ * This function creates an output tensor for the specified AnfNode and output index. The tensor is initialized with
+ * the inferred data type and shape of the output, and is associated with the corresponding device tensor.
+ *
+ * @param output_node The AnfNode representing the output.
+ * @param output_index The index of the output in the node.
+ * @return A pointer to the created output tensor.
+ */
 TensorPtr CreateOutputTensor(const AnfNodePtr &output_node, size_t output_index) {
   MS_EXCEPTION_IF_NULL(output_node);
   // Create host tensor, the output tensor should use the infer type, it will be handed correctly by tensor data sync
@@ -442,28 +527,44 @@ MindRTBackend::MindRTBackend(const std::string &backend_name, const std::string
   runtime::GraphScheduler::GetInstance().Initialize();
 }
 
+/**
+ * Compiles graphs and returns information about the generated actors.
+ * @param func_graph The function graph to be compiled.
+ * @return Information about the generated actors.
+ */
 const ActorInfo &MindRTBackend::CompileGraphs(const FuncGraphPtr &func_graph) {
+  // Check for null pointers
   MS_EXCEPTION_IF_NULL(graph_compiler_);
   MS_EXCEPTION_IF_NULL(func_graph);
+
+  // Log start of function graph compilation
   MS_LOG(INFO) << "Status record: start compile function graph: " << func_graph->ToString();
+
+  // Start profiling timer
   PROF_START(compile_func_graph);
+
+  // Wrap the input function graph to create a preprocessed root graph, which contains funcs each graph has(return a funcptr)
   auto root_graph = WrapPrimitives(func_graph);
   MS_EXCEPTION_IF_NULL(root_graph);
   root_graph_ = root_graph.get();
-  // Register a summary callback function, which is called in the final stages of summary.
+
+  // Register a callback function for summary saving
   graph_compiler_->RegisterSummaryCallBackFunc(callbacks::SummarySaveCallback);
 
+  // Get execution mode from context
   auto context_ptr = MsContext::GetInstance();
   MS_EXCEPTION_IF_NULL(context_ptr);
+  // For debug log
   ms_execution_mode_ = context_ptr->get_param<int>(MS_CTX_EXECUTION_MODE);
   real_execution_mode_ = ms_execution_mode_;
 
-  // Compile root graph.
+  // Compile the root graph
   graph_id_to_device_context_.clear();
   func_graph_to_kernel_graph_ids_.clear();
   control_nodes_.clear();
   auto subgraph_need_compile = CompileGraph(root_graph);
-  // Compile sub graphs.
+
+  // Compile sub graphs if needed
   if (subgraph_need_compile) {
     MS_EXCEPTION_IF_NULL(root_graph->manager());
     FuncGraphSet sub_graphs = root_graph->manager()->func_graphs();
@@ -474,27 +575,41 @@ const ActorInfo &MindRTBackend::CompileGraphs(const FuncGraphPtr &func_graph) {
     }
   }
 
-  // Construct the graph compiler info.
+  // Construct graph compiler info
   auto graph_compiler_info = ConstructGraphCompilerInfo(root_graph);
   MS_EXCEPTION_IF_NULL(graph_compiler_info);
+
+  // If in kgraph mode and there are compiled graphs, transform and schedule actor DAG
   if (real_execution_mode_ == kGraphMode && graph_compiler_info->graphs_.size() != 0) {
-    // Transform graph to actor DAG, and schedule the actor DAG.
     const auto &actor_set = runtime::GraphScheduler::GetInstance().Transform(*graph_compiler_info);
     runtime::GraphScheduler::GetInstance().Schedule(actor_set);
   }
+
+  // Retrieve actor information
   const ActorInfo &actor_info = graph_compiler_info->name_;
+
+  // Store graph compiler info
   (void)actor_to_graph_compiler_info_.emplace(graph_compiler_info->name_, std::move(graph_compiler_info));
+
+  // End profiling timer
   PROF_END(compile_func_graph);
 
+  // Reset execution mode if necessary
   if (ms_execution_mode_ != real_execution_mode_) {
     context_ptr->set_param<int>(MS_CTX_EXECUTION_MODE, ms_execution_mode_);
   }
 
+  // Log end of function graph compilation and actor information
   MS_LOG(INFO) << "Status record: end compile function graph: " << func_graph->ToString()
                << ", produce actor: " << actor_info;
   return actor_info;
 }
 
+/**
+ * Compiles the given function graph or its segments, performing necessary partitioning and compilation.
+ * @param func_graph The function graph or a segment of it to be compiled.
+ * @return True if the graph was split into segments and compiled separately, false if compiled as a whole.
+ */
 bool MindRTBackend::CompileGraph(const FuncGraphPtr &func_graph) {
   MS_EXCEPTION_IF_NULL(func_graph);
   MS_EXCEPTION_IF_NULL(graph_partition_);
@@ -523,6 +638,10 @@ bool MindRTBackend::CompileGraph(const FuncGraphPtr &func_graph) {
   return true;
 }
 
+/**
+ * Compiles a specific graph segment, which can be a normal segment or a cut node segment.
+ * @param segment The graph segment to be compiled.
+ */
 void MindRTBackend::CompileGraph(const GraphSegmentPtr &segment) {
   MS_EXCEPTION_IF_NULL(segment);
   // Compile the normal nodes, which doesn't contain the cut node.
@@ -577,6 +696,22 @@ void MindRTBackend::CompileGraph(const GraphSegmentPtr &segment) {
 }
 
 namespace {
+/**
+ * @brief Retrieves input information for the control operation by processing both front and backend CNodes.
+ *
+ * This function processes front and backend CNodes of a control operation, extracts the input information, and
+ * constructs an argument list (args) containing the processed inputs. It also updates the InputTensorInfo structure
+ * based on the input information.
+ *
+ * @param graph_compiler The shared pointer to the GraphCompiler instance.
+ * @param front_cnode The front-end CNode of the control operation.
+ * @param backend_cnode The backend CNode of the control operation.
+ * @param op_output_map The map of KernelWithIndex to tensor::TensorPtr representing operation outputs.
+ * @param parameter_index The map of AnfNodePtr to index representing parameter nodes.
+ * @param graph_inputs The vector of tensor::TensorPtr representing graph inputs.
+ * @param input_tensor_info A pointer to the InputTensorInfo structure to be updated.
+ * @param args A pointer to the VectorRef where the processed input arguments will be stored.
+ */
 void GetControlOpInput(const std::shared_ptr<GraphCompiler> &graph_compiler, const CNodePtr &front_cnode,
                        const CNodePtr &backend_cnode, const std::map<KernelWithIndex, tensor::TensorPtr> &op_output_map,
                        const std::map<AnfNodePtr, size_t> &parameter_index,
@@ -586,20 +721,27 @@ void GetControlOpInput(const std::shared_ptr<GraphCompiler> &graph_compiler, con
   MS_EXCEPTION_IF_NULL(backend_cnode);
   MS_EXCEPTION_IF_NULL(graph_compiler);
   MS_EXCEPTION_IF_NULL(args);
-  size_t front_index = 0;     // Point to front end cnode
-  size_t back_index = 0;      // Point to backend end cnode
-  size_t args_tuple_num = 0;  // Record the input num of maketuple cnode
-  std::vector<ValuePtr> args_tuple;
+
+  // Initialize indices and counters
+  size_t front_index = 0;     // Index pointing to front-end CNode
+  size_t back_index = 0;      // Index pointing to backend CNode
+  size_t args_tuple_num = 0;  // Number of inputs in maketuple CNode
+  std::vector<ValuePtr> args_tuple; // Temporary storage for inputs in maketuple CNode
+
   auto front_size = front_cnode->inputs().size();
   auto back_size = backend_cnode->inputs().size();
+
+  // Loop through the inputs of front and backend CNodes
   while (front_index + 1 < front_size && back_index + 1 < back_size) {
     AnfNodePtr input_node = nullptr;
+
     if (args_tuple_num) {
       input_node = backend_cnode->input(back_index + 1);
     } else {
       input_node = front_cnode->input(front_index + 1);
+
+      // Check if the input node is a primitive make tuple node
       if (IsPrimitiveCNode(input_node, prim::kPrimMakeTuple)) {
-        // Hook multi-input or multi-output.
         MS_LOG(DEBUG) << "The input node of hook op: " << input_node->DebugString() << " is a make tuple node.";
         auto make_tuple = input_node->cast<CNodePtr>();
         MS_EXCEPTION_IF_NULL(make_tuple);
@@ -607,13 +749,16 @@ void GetControlOpInput(const std::shared_ptr<GraphCompiler> &graph_compiler, con
         continue;
       }
     }
-    // Hook single-input or single-output.
+
+    // Extract the real input node (single-input or single-output)
     auto real_input = common::AnfAlgo::VisitKernel(input_node, 0).first;
     MS_EXCEPTION_IF_NULL(real_input);
     ValuePtr value = nullptr;
+
     if (!real_input->isa<ValueNode>()) {
-      value = graph_compiler->GetSingleOpInputTensorByIndex(backend_cnode, op_output_map, parameter_index, graph_inputs,
-                                                            input_tensor_info, back_index);
+      // Process backend input node that is not a ValueNode
+      value = graph_compiler->GetSingleOpInputTensorByIndex(backend_cnode, op_output_map, parameter_index,
+                                                            graph_inputs, input_tensor_info, back_index);
       MS_EXCEPTION_IF_NULL(value);
       ++back_index;
     } else {
@@ -621,7 +766,9 @@ void GetControlOpInput(const std::shared_ptr<GraphCompiler> &graph_compiler, con
       MS_EXCEPTION_IF_NULL(value_node);
       value = value_node->value();
       MS_EXCEPTION_IF_NULL(value);
+
       if (value->isa<ValueSequence>()) {
+        // Process ValueSequence input (multi-input or multi-output)
         const auto &value_sequeue = value->cast<ValueSequencePtr>();
         MS_EXCEPTION_IF_NULL(value_sequeue);
         back_index += value_sequeue->size();
@@ -629,14 +776,19 @@ void GetControlOpInput(const std::shared_ptr<GraphCompiler> &graph_compiler, con
         ++back_index;
       }
     }
+
+    // Add value to temporary args_tuple if processing maketuple input
     if (args_tuple_num) {
       args_tuple.emplace_back(value);
+
       if (args_tuple.size() == args_tuple_num) {
         value = std::make_shared<ValueTuple>(args_tuple);
         args_tuple_num = 0;
         args_tuple.clear();
       }
     }
+
+    // Add value to args if not processing maketuple input
     if (!args_tuple_num) {
       args->emplace_back(value);
       front_index++;
@@ -644,35 +796,54 @@ void GetControlOpInput(const std::shared_ptr<GraphCompiler> &graph_compiler, con
   }
 }
 
+
+/**
+ * @brief Converts a PyObject to a tensor::TensorPtr and adds it to the provided vector of tensors.
+ *
+ * This function converts a PyObject to a tensor::TensorPtr and appends it to the provided vector of tensors. It
+ * supports various data types, including tensor::Tensor, py::float_, py::int_, py::list, and py::tuple.
+ *
+ * @param input_object The input PyObject to be converted.
+ * @param tensors A pointer to the vector of tensor::TensorPtr where the converted tensor will be added.
+ */
 void ConvertPyObjectToTensor(const py::object &input_object, std::vector<tensor::TensorPtr> *tensors) {
   MS_EXCEPTION_IF_NULL(tensors);
   tensor::TensorPtr tensor_ptr = nullptr;
+
   if (py::isinstance<tensor::Tensor>(input_object)) {
     tensor_ptr = py::cast<tensor::TensorPtr>(input_object);
   } else if (py::isinstance<py::float_>(input_object)) {
+    // Convert py::float_ to tensor::Tensor with kFloat32 data type
     double input_value = py::cast<py::float_>(input_object);
     tensor_ptr = std::make_shared<tensor::Tensor>(input_value, kFloat32);
   } else if (py::isinstance<py::int_>(input_object)) {
+    // Convert py::int_ to tensor::Tensor with kInt64 data type
     tensor_ptr = std::make_shared<tensor::Tensor>(py::cast<int64_t>(input_object), kInt64);
   } else if (py::isinstance<py::list>(input_object)) {
+    // Convert py::list to vector of tensors recursively
     auto list_inputs = py::cast<py::list>(input_object);
     for (size_t i = 0; i < list_inputs.size(); ++i) {
       ConvertPyObjectToTensor(list_inputs[i], tensors);
     }
     return;
   } else if (py::isinstance<py::tuple>(input_object)) {
+    // Convert py::tuple to vector of tensors recursively
     auto tuple_inputs = py::cast<py::tuple>(input_object);
     for (size_t i = 0; i < tuple_inputs.size(); ++i) {
       ConvertPyObjectToTensor(tuple_inputs[i], tensors);
     }
     return;
   } else {
+    // Unsupported data type
     MS_EXCEPTION(TypeError) << "Unreasonable data type: " << input_object.get_type() << ".";
   }
+
+  // Add the converted tensor to the vector
   MS_EXCEPTION_IF_NULL(tensor_ptr);
   (void)tensors->emplace_back(tensor_ptr);
 }
 
+
 void RunControlOperator(const std::shared_ptr<GraphCompiler> &graph_compiler, const KernelGraphPtr &graph,
                         const CNodePtr &kernel, const std::map<KernelWithIndex, tensor::TensorPtr> &op_output_map,
                         const std::map<AnfNodePtr, size_t> &parameter_index,
@@ -742,14 +913,29 @@ void TensorValueToVector(const ValuePtr &value, VectorRef *outputs) {
   }
 }
 
+/**
+ * Check if the graph's output is a constant (ValueNode) or a parameter. If it is, there is no need to execute the graph.
+ *
+ * @param graph_output The output node of the computation graph.
+ * @param args Input arguments for the computation graph.
+ * @param outputs Reference to the vector where the output results will be stored if the graph's output is a constant.
+ *
+ * @return True if the graph's output is a constant or a parameter, indicating no need to execute; False otherwise.
+ */
 bool IsGraphOutputValueNodeOrParameter(const AnfNodePtr &graph_output, const VectorRef &args, VectorRef *outputs) {
   MS_EXCEPTION_IF_NULL(graph_output);
   MS_EXCEPTION_IF_NULL(outputs);
+
   if (graph_output->isa<ValueNode>()) {
+    // The graph's output is a constant. No need to execute.
     MS_LOG(INFO) << "Graph's output is a constant. No need to execute.";
+
     VectorRef output_tmp;
     ValuePtr value = GetValueNode(graph_output);
+
+    // Convert the constant tensor value to a vector.
     TensorValueToVector(value, &output_tmp);
+
     if (output_tmp.size() == 1) {
       *outputs = std::move(output_tmp);
     } else if (output_tmp.size() > 1) {
@@ -757,33 +943,43 @@ bool IsGraphOutputValueNodeOrParameter(const AnfNodePtr &graph_output, const Vec
     } else {
       MS_LOG(EXCEPTION) << "Output is empty!";
     }
+
     return true;
   }
 
   if (graph_output->isa<Parameter>()) {
+    // The graph's output is a parameter. If all parameters are inputs, no need to execute.
     MS_LOG(INFO) << "Graph's output is a parameter. If all params are inputs, no need to execute.";
+
     // Find the right parameter as ret_val.
     auto func_graph = graph_output->func_graph();
     MS_EXCEPTION_IF_NULL(func_graph);
     auto params = func_graph->parameters();
+
     if (args.size() != params.size()) {
       MS_LOG(EXCEPTION) << "Input size " << args.size() << " not equal to graph input size " << params.size();
     }
 
     auto it = std::find(params.begin(), params.end(), graph_output);
+
     if (it == params.end()) {
       MS_EXCEPTION(UnknownError) << "When graph output is Parameter, it should be found in graph parameters";
     }
+
     size_t index = it - params.cbegin();
+
     if (index >= args.size()) {
       MS_EXCEPTION(UnknownError) << "Index " << index << " equal or larger than args size " << args.size();
     }
 
     outputs->emplace_back(args[index]);
+
     return true;
   }
+
   return false;
 }
+
 }  // namespace
 
 void FlatValueTupleValue(const ValuePtrList &value, ValuePtrList *flatted_value) {
@@ -826,119 +1022,190 @@ void FlattenValue(const BaseRef &arg, ValuePtrList *flatted_value) {
         (void)flatted_value->emplace_back(value);
       } else {
         FlattenValue(value, flatted_value);
-      }
     }
+      }
   } else {
     MS_LOG(EXCEPTION) << "The value input to flatten should only contains be sequence or dictionary, but it is "
                       << arg.ToString();
   }
 }
 
+/**
+ * Pushes a tensor from the `args` list into the `input_tensor` vector based on the `front_node` and index provided.
+ *
+ * @param args A list of input arguments.
+ * @param parameters A list of AnfNodePtr representing the parameters.
+ * @param front_node The front node representing the parameter in the subgraph.
+ * @param index The index of the tensor within the parameter.
+ * @param input_tensor A vector to store the input tensor.
+ */
 void PushTupleTensor(const VectorRef &args, const std::vector<AnfNodePtr> &parameters, const AnfNodePtr &front_node,
                      size_t index, std::vector<tensor::TensorPtr> *input_tensor) {
+  // Find the position of the `front_node` in the `parameters` list.
   const auto &iter = std::find(parameters.begin(), parameters.end(), front_node);
   const size_t position = iter - parameters.begin();
-  // If the parameter is not found in the parameters of the root graph, it means that it is the input of the subgraph,
-  // and there is no need to input a tensor.
+
+  // If the parameter is not found in the parameters of the root graph, it means it's an input of the subgraph.
+  // In this case, there's no need to input a tensor, so add a nullptr to the `input_tensor` vector.
   if (position >= args.size()) {
     MS_LOG(INFO) << "Position out of args range, position value is " << position << " and args size is " << args.size()
                  << ".";
     (void)input_tensor->emplace_back(nullptr);
     return;
   }
+
+  // Flatten the value of the argument at the specified position.
   ValuePtrList flatted_value_tuple_value;
   FlattenValue(args[position], &flatted_value_tuple_value);
+
+  // Check if the index is within the range of the flattened values.
   if (index >= flatted_value_tuple_value.size()) {
     MS_LOG(EXCEPTION) << "Index out of flatted_value_tuple_value range, index value is " << index
                       << " and flatted_value_tuple_value size is " << flatted_value_tuple_value.size() << ".";
   }
+
+  // Retrieve the tensor input and add it to the `input_tensor` vector.
   auto input = flatted_value_tuple_value[index];
   MS_EXCEPTION_IF_NULL(input);
   auto tensor_input = input->cast<tensor::TensorPtr>();
   input_tensor->push_back(tensor_input);
 }
 
+
+/**
+ * Execute the computation graph represented by a list of KernelGraphPtrs using a single operator approach.
+ *
+ * @param graphs A list of KernelGraphPtrs representing the computation graphs.
+ * @param inputs A list of input tensors for each computation graph.
+ * @param outputs A reference to the vector where the output results will be stored.
+ */
 void MindRTBackend::RunGraphBySingleOp(const std::vector<KernelGraphPtr> &graphs,
                                        const std::vector<std::vector<tensor::TensorPtr>> &inputs, VectorRef *outputs) {
+  // Ensure that previous tasks are finished before starting execution.
   WaitTaskFinish();
+
   MS_EXCEPTION_IF_NULL(graph_compiler_);
   auto &op_executor = runtime::OpExecutor::GetInstance();
+
+  // Register a batch build callback function.
   op_executor.Register([this]() { BatchBuildCallback(); });
+
   for (size_t graph_index = 0; graph_index < graphs.size(); ++graph_index) {
     const auto &graph = graphs[graph_index];
     MS_EXCEPTION_IF_NULL(graph);
+
+    // Initialize data structures for managing operator outputs.
     std::map<KernelWithIndex, tensor::TensorPtr> op_output_map;
     std::map<AnfNodePtr, size_t> parameter_index;
     GraphOutputInfo graph_output_info;
     graph_output_info.graph_outputs = outputs;
+
+    // Get parameter and output indexes for the current graph.
     graph_compiler_->GetParamAndOutputIndex(graph, inputs[graph_index], outputs, &parameter_index,
                                             &graph_output_info.output_indexes);
 
+    // Initialize a reference count map for CNodes in the current graph.
     std::map<KernelWithIndex, size_t> cnode_ref_count;
     auto iter = cnode_ref_counts_.find(graph->graph_id());
+
     if (iter == cnode_ref_counts_.end()) {
+      // Calculate reference counts for CNodes in the graph if not already calculated.
       graph_compiler_->CalculateRefCount(graph, &cnode_ref_count);
       (void)cnode_ref_counts_.emplace(graph->graph_id(), cnode_ref_count);
     } else {
       cnode_ref_count = iter->second;
     }
+
+    // Calculate forward operator output tensor IDs.
     graph_compiler_->CalculateForwardOpOutputCount(graph, inputs[graph_index], &forward_op_output_tensor_id_);
 
     for (const auto &kernel : graph->execution_order()) {
       InputTensorInfo input_tensor_info;
       VectorRef op_outputs;
+
       if (!common::AnfAlgo::IsControlOpExecInBackend(kernel)) {
         OpRunInfo op_run_info;
         GraphInfo graph_info;
+
+        // Get input tensors, operator run info, and graph info for the current kernel.
         graph_compiler_->GetSingleOpInputTensors(kernel, op_output_map, parameter_index, inputs[graph_index],
                                                  &input_tensor_info);
         graph_compiler_->GetSingleOpRunInfoAndGraphInfo(kernel, input_tensor_info, &op_run_info, &graph_info,
                                                         &graph_output_info);
 
+        // Run the operator and store the outputs.
         RunOp(&op_run_info, &op_outputs);
       } else {
+        // If it's a control operator, wait for previous tasks to finish before execution.
         WaitTaskFinish();
+
+        // Run the control operator and update op outputs.
         RunControlOperator(graph_compiler_, graph, kernel, op_output_map, parameter_index, inputs[graph_index],
                            &input_tensor_info, &op_outputs);
-        // Execute remaining lazy tasks before PyNative hook exit.
+
+        // Execute any remaining lazy tasks before exiting the PyNative hook.
         WaitTaskFinish();
       }
 
+      // Update the reference count for input kernel and manage operator outputs.
       graph_compiler_->UpdateRefCount(input_tensor_info.input_kernel, &cnode_ref_count, &op_output_map);
 
+      // Recover graph outputs based on operator outputs and reference counts.
       graph_output_info.graph_output_tensors.clear();
       graph_compiler_->RecoverGraphOutput(kernel, op_outputs, cnode_ref_count, &op_output_map, &graph_output_info);
 
-      // Save grad node to Bucket
+      // Save gradient node addresses to the Bucket if it's a backward graph and not a parallel kernel.
       if (graph->is_bprop() && (!common::AnfAlgo::IsControlOpExecInBackend(kernel)) && !kernel->is_parallel()) {
         graph_compiler_->AddGradAddrToBucket(graph->graph_id(), graph_output_info.graph_output_tensors);
       }
     }
+
+    // Wait for any pending tasks to finish before moving on to the next graph.
     WaitTaskFinish();
-    // Clear bucket resources every step
+
+    // Clear bucket resources at the end of each step.
     if (graph->is_bprop()) {
       graph_compiler_->ClearAllBucket(graph->graph_id());
     }
   }
 }
 
+
+/**
+ * Execute a computation graph using the MindRT backend.
+ *
+ * This function is responsible for executing a computation graph represented by the provided `actor_info` and `args`.
+ * Depending on the execution mode (Pynative or Graph Mode), it processes input tensors and runs the graph. The results
+ * are constructed in the `outputs` parameter.
+ *
+ * @param actor_info Information about the computation graph to be executed.
+ * @param args Input arguments for the computation graph.
+ * @param outputs Reference to the vector where the output results will be stored.
+ */
 void MindRTBackend::RunGraph(const ActorInfo &actor_info, const VectorRef &args, VectorRef *outputs) {
+  // Ensure the root graph is not null.
   MS_EXCEPTION_IF_NULL(root_graph_);
+
+  // Check if the root graph output is a value node or parameter; if yes, no execution is needed.
   if (IsGraphOutputValueNodeOrParameter(root_graph_->output(), args, outputs)) {
     return;
   }
 
+  // Get the context to check for precompile-only mode.
   const auto &context_ptr = MsContext::GetInstance();
   MS_EXCEPTION_IF_NULL(context_ptr);
+  
+  // If in precompile-only mode, stop execution.
   if (context_ptr->get_param<bool>(MS_CTX_PRECOMPILE_ONLY)) {
     MS_LOG(INFO) << "PrecompileOnly, stop run graph";
     return;
   }
 
-  // Open abstract_lock for dynamic_shape
+  // Open the abstract lock for dynamic_shape.
   AnfUtils::OpenAbstractLock();
 
   MS_LOG(INFO) << "Status record: start run actor: " << actor_info;
+
   // Fetch the graph compiler info.
   const auto &graph_iter = actor_to_graph_compiler_info_.find(actor_info);
   if (graph_iter == actor_to_graph_compiler_info_.end()) {
@@ -952,6 +1219,7 @@ void MindRTBackend::RunGraph(const ActorInfo &actor_info, const VectorRef &args,
   WaitTaskFinish();
 
   // Transform args to input tensors.
+
   // Input tensors of the graph.
   std::vector<std::vector<tensor::TensorPtr>> input_tensors;
   for (const auto &kernel_graph : graph_compiler_info.graphs_) {
@@ -972,16 +1240,18 @@ void MindRTBackend::RunGraph(const ActorInfo &actor_info, const VectorRef &args,
   // Input tensors of the control node.
   std::vector<tensor::TensorPtr> input_tensor;
   MS_EXCEPTION_IF_NULL(graph_compiler_info.control_node_parser_);
-  // Get inputs of control node which come from the host actor.
+
+  // Get inputs of the control node which come from the host actor.
   const auto &control_node_parameters = graph_compiler_info.control_node_parser_->control_node_parameters();
   for (const auto &parameter : control_node_parameters) {
     PushTensor(args, origin_parameters, parameter, &input_tensor);
   }
   (void)input_tensors.emplace_back(input_tensor);
 
-  // Run in the pynative mode.
+  // Run in pynative mode.
   MS_EXCEPTION_IF_NULL(outputs);
-  // There will be more than one kernel graph in heterogeneous scenario in a ms function of PyNative Mode.
+
+  // There will be more than one kernel graph in heterogeneous scenarios in a ms function of PyNative Mode.
   if (real_execution_mode_ == kPynativeMode) {
     RunGraphBySingleOp(graph_compiler_info.graphs_, input_tensors, outputs);
     MS_LOG(INFO) << "Status record: end run actor: " << actor_info;
@@ -997,11 +1267,14 @@ void MindRTBackend::RunGraph(const ActorInfo &actor_info, const VectorRef &args,
   MS_EXCEPTION_IF_NULL(graph_compiler_);
   graph_compiler_->Summary(graph_compiler_info.graphs_);
 
-  bool need_contruct_output = !(distributed::recovery::RecoveryContext::GetInstance()->enable_recovery() &&
+  // Construct output results.
+
+  // Ensure the construction of outputs is needed.
+  bool need_construct_output = !(distributed::recovery::RecoveryContext::GetInstance()->enable_recovery() &&
                                 distributed::recovery::RecoveryContext::GetInstance()->need_reset());
-  if (need_contruct_output) {
-    // Update device address for output node of graph.
-    // Summary processing will use the output device address, so must be after the summary processing.
+  if (need_construct_output) {
+    // Update device address for the output node of the graph.
+    // Summary processing will use the output device address, so it must be after the summary processing.
     actor_set->output_actor_->UpdateOutputDeviceAddress();
 
     // Fetch outputs.
@@ -1013,12 +1286,14 @@ void MindRTBackend::RunGraph(const ActorInfo &actor_info, const VectorRef &args,
     }
   }
 
+  // Clear actor data and close the abstract lock for dynamic_shape.
   runtime::GraphScheduler::GetInstance().ClearActorData(actor_set);
-  // Close abstract_lock for dynamic_shape
   AnfUtils::CloseAbstractLock();
+
   MS_LOG(INFO) << "Status record: end run actor: " << actor_info;
 }
 
+
 BaseRef MindRTBackend::ConstructOutputByAbstract(const abstract::AbstractBasePtr &abstract,
                                                  const std::vector<tensor::TensorPtr> &output_tensors,
                                                  size_t *output_position) {
@@ -1167,20 +1442,37 @@ void MindRTBackend::SyncStream() {
   (void)device_context->SyncStream();
 }
 
+/**
+ * @brief Constructs the information needed for graph compilation and execution in the MindRT backend.
+ *
+ * This function collects information about kernel graphs, device contexts, control nodes, and other
+ * relevant details required for the configuration of a graph compiler. It constructs a `GraphCompilerInfo`
+ * object to encapsulate this information for further use in graph compilation and execution.
+ *
+ * @param root_graph The root function graph to be compiled.
+ * @return A unique pointer to a `GraphCompilerInfo` object containing the collected information.
+ */
 std::unique_ptr<GraphCompilerInfo> MindRTBackend::ConstructGraphCompilerInfo(const FuncGraphPtr &root_graph) {
+  // Check for null inputs
   MS_EXCEPTION_IF_NULL(root_graph);
   MS_EXCEPTION_IF_NULL(graph_compiler_);
 
+  // Initialize data structures to store information
   std::vector<KernelGraphPtr> graphs;
   std::vector<DeviceContext *> device_contexts;
   std::string name = "kernel_graph";
+
+  // Loop through graph_id_to_device_context_ and collect graphs and device contexts
   for (const auto &graph_id_to_context : graph_id_to_device_context_) {
     (void)graphs.emplace_back(graph_compiler_->Fetch(graph_id_to_context.first));
     (void)device_contexts.emplace_back(graph_id_to_context.second);
     (void)name.append("_").append(std::to_string(graph_id_to_context.first));
   }
 
+  // Initialize data structures to map function graphs to kernel graphs
   FuncGraphToKernelGraphGroup func_graph_to_kernel_graphs;
+
+  // Loop through func_graph_to_kernel_graph_ids_ and collect kernel graphs for each function graph
   for (const auto &func_graph_to_kernel_graph_ids : func_graph_to_kernel_graph_ids_) {
     const auto &func_graph = func_graph_to_kernel_graph_ids.first;
     for (const auto &sub_kernel_graphs_ids : func_graph_to_kernel_graph_ids.second) {
@@ -1194,9 +1486,11 @@ std::unique_ptr<GraphCompilerInfo> MindRTBackend::ConstructGraphCompilerInfo(con
     }
   }
 
+  // Create a control node parser and parse control nodes
   auto parser = std::make_shared<ControlNodeParser>();
   parser->Parse(control_nodes_, graphs, device_contexts, root_graph, func_graph_to_kernel_graphs);
 
+  // Determine the order of kernel outputs
   runtime::KernelMapPosition outputs_order;
   const auto &root_output =
     common::AnfAlgo::VisitKernelWithReturnType(root_graph->output(), 0, false, {prim::kPrimTupleGetItem}).first;
@@ -1211,6 +1505,7 @@ std::unique_ptr<GraphCompilerInfo> MindRTBackend::ConstructGraphCompilerInfo(con
     }
   }
 
+  // Create and return GraphCompilerInfo
   std::vector<std::vector<int64_t> *> tensors_mask;
   std::vector<std::vector<tensor::TensorPtr> *> input_tensors;
   return std::make_unique<GraphCompilerInfo>(graphs, device_contexts, tensors_mask, input_tensors, control_nodes_,
@@ -1218,6 +1513,7 @@ std::unique_ptr<GraphCompilerInfo> MindRTBackend::ConstructGraphCompilerInfo(con
                                              runtime::GraphExecutionStrategy::kPipeline);
 }
 
+
 std::unique_ptr<GraphCompilerInfo> MindRTBackend::ConstructGraphCompilerInfo(
   const ActorInfo &actor_info, const std::vector<int64_t> *tensors_mask,
   const std::vector<tensor::TensorPtr> *input_tensors, bool need_erase) {

mindspore/ccsrc/backend/graph_compiler/backend.h

@@ -54,6 +54,7 @@ enum SwitchCondStatus {
   kCondAlreadyRun,
 };
 
+// Base Class
 class BACKEND_EXPORT Backend {
  public:
   explicit Backend(const std::string &name);
@@ -76,6 +77,7 @@ class BACKEND_EXPORT Backend {
   bool is_multi_graph_sink_;
 };
 
+// Inherit Ones
 class BACKEND_EXPORT MsBackend : public Backend {
  public:
   MsBackend(const std::string &name, const std::string &target, uint32_t device_id);
@@ -102,6 +104,7 @@ class BACKEND_EXPORT MsBackend : public Backend {
   mindspore::HashMap<GraphId, LinConvertResult> graph_id_map_;
 };
 
+
 class BACKEND_EXPORT MindRTBackend : public Backend {
  public:
   MindRTBackend(const std::string &backend_name, const std::string &device_name, uint32_t device_id);

mindspore/ccsrc/backend/graph_compiler/graph_partition.cc

@@ -37,6 +37,13 @@ const char kMsVm[] = "vm";
 const char kGeVm[] = "ge";
 namespace compile {
 namespace {
+
+/*
+ * @brief Get the other target of the graph. Note that a graph can only have two targets.
+ * 
+ * @param nodes The nodes of the graph.
+ * @return The other target of the graph.
+ */
 std::string GetOtherTarget(const std::vector<AnfNodePtr> &nodes) {
   auto context_ptr = MsContext::GetInstance();
   MS_EXCEPTION_IF_NULL(context_ptr);
@@ -54,6 +61,12 @@ std::string GetOtherTarget(const std::vector<AnfNodePtr> &nodes) {
   return "";
 }
 
+/**
+ * @brief Calculate the reference count of each node in the graph.
+ * 
+ * @param graph The graph to be calculated.
+ * @param nodes_ref The reference count of each node in the graph.
+ */
 void CalcNodeRefCount(const FuncGraphPtr &graph, std::map<AnfNodePtr, size_t> *nodes_ref) {
   MS_EXCEPTION_IF_NULL(graph);
   MS_EXCEPTION_IF_NULL(nodes_ref);
@@ -85,6 +98,13 @@ void CalcNodeRefCount(const FuncGraphPtr &graph, std::map<AnfNodePtr, size_t> *n
   }
 }
 
+/*
+ * @brief Reorder virtual node such as primitive "depend" and "tuple_getitem" to ensure the order of the nodes.
+ * These virtual nodes will be inserted after their parent node.
+ * 
+ * @param nodes The nodes to be reordered.
+ * @param reorder_prim The primitive used to reorder the nodes.
+ */
 std::vector<AnfNodePtr> ReorderVirtualNode(const std::vector<AnfNodePtr> &nodes, const PrimitivePtr &reorder_prim) {
   std::vector<AnfNodePtr> result;
   std::map<size_t, std::vector<AnfNodePtr>> insert_positions;
@@ -136,6 +156,14 @@ std::vector<AnfNodePtr> ReorderVirtualNode(const std::vector<AnfNodePtr> &nodes,
   return result;
 }
 
+/*
+ * @brief Get next nodes to be visited.
+ * 
+ * @param node The current node.
+ * @param nodes_ref The reference count of each node in the graph.
+ * @param result The next nodes to be visited.
+ * @return The next nodes to be visited.
+ */
 std::vector<AnfNodePtr> GetNextNodes(const AnfNodePtr &node, std::map<AnfNodePtr, size_t> *nodes_ref,
                                      std::vector<AnfNodePtr> *result) {
   MS_EXCEPTION_IF_NULL(node);
@@ -226,6 +254,12 @@ struct GraphNodesDependencyInfo {
   std::map<AnfNodePtr, std::vector<AnfNodePtr>> output_edges_;
 };
 
+/*
+ * @brief Get the dependency information of all the nodes in the graph.
+ * 
+ * @param graph The graph to be calculated.
+ * @return The dependency information of the nodes in the graph.
+ */
 GraphNodesDependencyInfo GetNodesDependencyInfo(const FuncGraphPtr &graph) {
   MS_EXCEPTION_IF_NULL(graph);
   GraphNodesDependencyInfo info;
@@ -283,6 +317,16 @@ struct VisitNodesInfo {
   std::map<AnfNodePtr, AnfNodePtr> seed_cast_next_node_;
 };
 
+/*
+ * @brief Get the visit information of all the nodes in the graph.
+ * The information includes the nodes of the default target and another specified target,
+ * and the cast node following a certain node, if exists.
+ * 
+ * @param dependency_info The dependency information of the nodes in the graph.
+ * @param default_target The default target of the graph.
+ * @param other_target The other target of the graph.
+ * @return The visit information of the nodes in the graph.
+ */
 VisitNodesInfo GetVisitNodesInfo(const GraphNodesDependencyInfo &dependency_info, const std::string &default_target,
                                  const std::string &other_target) {
   VisitNodesInfo result;
@@ -366,6 +410,14 @@ void ParallelSortVisitNodeEdges(const std::vector<AnfNodePtr> &output_edges, Gra
   }
 }
 
+/*
+ * @brief Sort the nodes of the graph considering multiple targets to facilitate the parallel execution.
+ * 
+ * @param graph The graph to be sorted.
+ * @param default_target The default target of the graph.
+ * @param other_target The other target of the graph.
+ * @return The sorted nodes of the graph.
+ */
 std::vector<AnfNodePtr> ParallelSort(const FuncGraphPtr &graph, const std::string &default_target,
                                      const std::string &other_target) {
   MS_EXCEPTION_IF_NULL(graph);
@@ -408,6 +460,12 @@ std::vector<AnfNodePtr> ParallelSort(const FuncGraphPtr &graph, const std::strin
   return result;
 }
 
+/*
+ * @brief Add dependency to segments.
+ * 
+ * @param graph The graph to be added.
+ * @param node_to_segment The mapping from node to segment.
+ */
 void AddSegmentDependency(const FuncGraphPtr &graph, const std::map<AnfNodePtr, GraphSegmentPtr> &node_to_segment) {
   MS_EXCEPTION_IF_NULL(graph);
   std::stack<AnfNodePtr> to_visit;
@@ -550,6 +608,14 @@ struct SplitDynamicNodesHelper {
   size_t merge_node_threshold = 6;
 };
 
+/*
+ * @brief Split the nodes into segments according to the dynamic shape.
+ * 
+ * @param segment_nodes The nodes to be split.
+ * @param segments The segments after splitting.
+ * @param node_to_segment The mapping from node to segment.
+ * @param dynamic_nodes_set The set of nodes with dynamic shape.
+ */
 void SplitDynamicNodeSegment(const std::vector<AnfNodePtr> &segment_nodes, std::vector<GraphSegmentPtr> *segments,
                              std::map<AnfNodePtr, GraphSegmentPtr> *node_to_segment,
                              const std::set<AnfNodePtr> &dynamic_nodes_set) {
@@ -590,6 +656,13 @@ void SplitDynamicNodeSegment(const std::vector<AnfNodePtr> &segment_nodes, std::
   helper.AddSegments(segments, node_to_segment);
 }
 
+/*
+ * @brief Convert nodes before the cut node to a segment.
+ * 
+ * @param segment_nodes The nodes to be split.
+ * @param segments The segments after splitting.
+ * @param node_to_segment The mapping from node to segment.
+ */
 void NodesToSegments(const std::vector<AnfNodePtr> &segment_nodes, std::vector<GraphSegmentPtr> *segments,
                      std::map<AnfNodePtr, GraphSegmentPtr> *node_to_segment) {
   if (segment_nodes.empty()) {
@@ -624,22 +697,43 @@ void NodesToSegments(const std::vector<AnfNodePtr> &segment_nodes, std::vector<G
 GraphPartition::GraphPartition(const std::vector<PrimitivePtr> &cut_list, const std::string &backend_name)
     : cut_list_(cut_list), backend_name_(backend_name) {}
 
+/**
+ * @brief Checks if the given AnfNode is a cut point for partitioning.
+ *
+ * This function determines if a given AnfNode should be considered as a cut point for partitioning a graph. It checks
+ * various conditions, such as the node's type and its association with certain primitives or backends.
+ *
+ * @param node The AnfNode to be checked.
+ * @return True if the node is a cut point, otherwise false.
+ */
 bool GraphPartition::IsCut(const AnfNodePtr &node) {
   MS_EXCEPTION_IF_NULL(node);
+
+  // Check if the node is a CNode.
   if (node->isa<CNode>()) {
     auto cnode = node->cast<CNodePtr>();
     auto &inputs = cnode->inputs();
+
+    // Ensure the inputs of the apply node are not empty.
     if (inputs.empty()) {
       MS_LOG(EXCEPTION) << "Inputs of apply node is empty";
     }
+
+    // Get the first input node.
     AnfNodePtr fn = inputs[0];
+
+    // Check if the first input is not a ValueNode of Primitive type, indicating a cut point.
     if (!IsValueNode<Primitive>(fn)) {
       return true;
     }
+
     auto node_prim = GetValueNode<PrimitivePtr>(fn);
+
+    // Check if the node_prim is in the cut_list, indicating a cut point.
     for (auto &prim : cut_list_) {
       MS_EXCEPTION_IF_NULL(prim);
       if (prim->name() == node_prim->name()) {
+        // Handle special cases based on the primitive's name.
         if (prim->name() == prim::kPrimBpropCut->name()) {
           auto ms_context = MsContext::GetInstance();
           MS_EXCEPTION_IF_NULL(ms_context);
@@ -655,6 +749,8 @@ bool GraphPartition::IsCut(const AnfNodePtr &node) {
         return true;
       }
     }
+
+    // Check for backend-specific conditions (e.g., for backend 'kGeVm').
 #ifdef ENABLE_D
     if (backend_name_ == kGeVm) {
       auto name = GetCNodeFuncName(cnode);
@@ -665,38 +761,66 @@ bool GraphPartition::IsCut(const AnfNodePtr &node) {
     }
 #endif
   }
+
+  // If none of the conditions are met, it's not a cut point.
   return false;
 }
 
+
+/**
+ * @brief Partitions a given FuncGraph into multiple segments.
+ *
+ * This function takes a FuncGraph as input and partitions it into multiple segments based on certain criteria. The
+ * segments are returned as a vector of GraphSegmentPtr.
+ *
+ * @param graph The input FuncGraph to be partitioned.
+ * @param multi_target A pointer to a boolean flag indicating whether the graph contains multiple targets. If not needed,
+ *                    you can pass nullptr.
+ * @return A vector of GraphSegmentPtr representing the partitions of the input FuncGraph.
+ */
 std::vector<GraphSegmentPtr> GraphPartition::Partition(const FuncGraphPtr &graph, bool *multi_target) {
   MS_EXCEPTION_IF_NULL(graph);
+  // graph->get_return return the CNode object
+  // TopoSort
   auto nodes = TopoSort(graph->get_return());
   MS_LOG(DEBUG) << "Split all nodes size:" << nodes.size();
   bool contain_multi_target = ContainMultiTarget(nodes);
+
+  // Set the 'multi_target' flag if it's provided.
   if (multi_target != nullptr) {
     *multi_target = contain_multi_target;
   }
 
   auto context_ptr = MsContext::GetInstance();
   MS_EXCEPTION_IF_NULL(context_ptr);
+  // may be use loop sink to reduce operations inside the loop
   auto enable_loop_sink = context_ptr->get_param<bool>(MS_CTX_ENABLE_LOOP_SINK);
   std::string default_target = context_ptr->get_param<std::string>(MS_CTX_DEVICE_TARGET);
+
+  // Perform partitioning based on the criteria.
   if (contain_multi_target || !enable_loop_sink) {
+    // parellize or not
     if (context_ptr->get_param<bool>(MS_CTX_ENABLE_PARALLEL_SPLIT)) {
       auto other_target = GetOtherTarget(nodes);
       nodes = ParallelSort(graph, default_target, other_target);
     } else {
       nodes = SplitSort(graph, default_target);
     }
+    // Reorder the nodes with primitive "depend" and "tuple_getitem" to ensure the correctness of the partitioning.
     nodes = ReorderVirtualNode(nodes, prim::kPrimTupleGetItem);
     nodes = ReorderVirtualNode(nodes, prim::kPrimDepend);
   }
+
+  // Initialize data structures to store segments and nodes.
   std::vector<GraphSegmentPtr> segments;
   std::vector<AnfNodePtr> segment_nodes;
   std::map<AnfNodePtr, GraphSegmentPtr> node_to_segment;
   std::string last_target;
+
+  // Iterate through nodes to create segments.
   for (auto &node : nodes) {
     MS_EXCEPTION_IF_NULL(node);
+
     if (IsCut(node)) {
       NodesToSegments(segment_nodes, &segments, &node_to_segment);
       segment_nodes.clear();
@@ -707,6 +831,8 @@ std::vector<GraphSegmentPtr> GraphPartition::Partition(const FuncGraphPtr &graph
     } else if (node->isa<CNode>()) {
       if (contain_multi_target) {
         std::string cur_target = GetCNodeTarget(node);
+
+        // Check if the target has changed, and if so, start a new segment.
         if (cur_target != last_target && !last_target.empty()) {
           NodesToSegments(segment_nodes, &segments, &node_to_segment);
           segment_nodes.clear();
@@ -716,11 +842,15 @@ std::vector<GraphSegmentPtr> GraphPartition::Partition(const FuncGraphPtr &graph
       segment_nodes.emplace_back(node);
     }
   }
+
   MS_LOG(DEBUG) << "Segment size:" << segments.size();
+
+  // Add segment dependencies and remove useless ones if multiple targets are present.
   if (contain_multi_target) {
     AddSegmentDependency(graph, node_to_segment);
     RemoveUselessDependency(&segments);
   }
+
   return segments;
 }
 }  // namespace compile

mindspore/ccsrc/backend/graph_compiler/graph_partition.h

@@ -38,8 +38,11 @@ class GraphPartition {
   std::vector<GraphSegmentPtr> Partition(const FuncGraphPtr &func_graph, bool *multi_target = nullptr);
 
  private:
+ // To get the point can be cut or not
   bool IsCut(const AnfNodePtr &node);
+  // New Primitives List After cutting
   std::vector<PrimitivePtr> cut_list_;
+  // the key same as the name_ in the backend object
   std::string backend_name_;
 };
 

mindspore/ccsrc/backend/graph_compiler/transform.cc

@@ -421,15 +421,20 @@ void TraverseGraphMap(
   const std::function<std::shared_ptr<FuncGraph>(const PrimitivePtr, const AbstractFunctionPtr)> &get_prim_graph) {
   MS_EXCEPTION_IF_NULL(manager_ptr);
   MS_EXCEPTION_IF_NULL(tr);
+
   for (const auto &fg : fgs) {
+  // traverse all funcgraphs
     MS_EXCEPTION_IF_NULL(fg);
     for (const auto &ct_any : fg->value_nodes()) {
+      // process all value nodes
       AnfNodePtr const_primitive_node = ct_any.first;
       if (const_primitive_node != nullptr && IsValueNode<Primitive>(const_primitive_node)) {
         auto users = manager_ptr->node_users()[const_primitive_node];
+        // traverse CNode
         for (auto &use : users) {
           CNodePtr node = use.first->cast<CNodePtr>();
           MS_EXCEPTION_IF_NULL(node);
+          // coose users use this fg
           if (node->func_graph() != fg) {
             continue;
           }
@@ -445,8 +450,10 @@ void TraverseGraphMap(
                 continue;
               }
             }
+            // lambda func here
             FuncGraphPtr g = get_prim_graph(GetValueNode<PrimitivePtr>(const_primitive_node),
                                             dyn_cast<AbstractFunction>(const_primitive_node->abstract()));
+            // The user index "key" Use this node
             tr->SetEdge(node, key, NewValueNode(g));
           }
         }
@@ -455,13 +462,29 @@ void TraverseGraphMap(
   }
 }
 
+/**
+ * @brief Wraps primitive operations in the given FuncGraph.
+ *
+ * This function takes a FuncGraph as input and wraps primitive operations within it. It creates a new FuncGraph for each
+ * primitive operation and its type, then replaces the original primitive operation with the new FuncGraph call.
+ *
+ * @param graph The input FuncGraph to be processed.
+ * @return The processed FuncGraph with wrapped primitive operations(Prim Graph).
+ **/
 FuncGraphPtr WrapPrimitives(const FuncGraphPtr &graph) {
+
   MS_EXCEPTION_IF_NULL(graph);
+// get root graph
   FuncGraphManagerPtr manager_ptr = graph->manager();
   MS_EXCEPTION_IF_NULL(manager_ptr);
+// using MapPrimTypeFuncGraph = std::map<PrimTypePair, FuncGraphPtr>
   MapPrimTypeFuncGraph prim_graphs;
+
+// lambda function to get FuncGraphs by Type
   auto get_prim_graph = [&prim_graphs](const PrimitivePtr &prim, const AbstractFunctionPtr &type) {
     PrimTypePair prim_type = std::make_pair(prim, type);
+
+// If tis type can't be found, create it.
     if (prim_graphs.end() == prim_graphs.find(prim_type)) {
       FuncGraphPtr g = std::make_shared<FuncGraph>();
       std::vector<AnfNodePtr> args;
@@ -469,15 +492,18 @@ FuncGraphPtr WrapPrimitives(const FuncGraphPtr &graph) {
       MS_EXCEPTION_IF_NULL(prim_ct);
       prim_ct->set_abstract(type);
       args.push_back(prim_ct);
+
       MS_EXCEPTION_IF_NULL(type);
       TypedPrimitiveAbstractClosurePtr tp = dyn_cast<abstract::TypedPrimitiveAbstractClosure>(type->GetUnique());
       MS_EXCEPTION_IF_NULL(tp);
       MS_EXCEPTION_IF_NULL(g);
+// add paras in the para list
       for (auto t : tp->args_spec_list()) {
         ParameterPtr p = g->add_parameter();
         p->set_abstract(t);
         args.push_back(p);
       }
+
       AnfNodePtr out = g->NewCNode(args);
       out->set_abstract(tp->output());
       g->set_output(out);
@@ -489,12 +515,15 @@ FuncGraphPtr WrapPrimitives(const FuncGraphPtr &graph) {
 
   FuncGraphTransaction tr = manager_ptr->Transact();
   auto &fgs = manager_ptr->func_graphs();
+// call the lambda function here
   TraverseGraphMap(manager_ptr, &tr, fgs, get_prim_graph);
+// commit as A FuncGraphTransaction
   tr.Commit();
 
   return graph;
 }
 
+
 CompileGraphs::CompileGraphs(const BackendPtr &backend, const std::vector<PrimitivePtr> &cut_list) : backend_(backend) {
   MS_EXCEPTION_IF_NULL(backend);
   MS_LOG(DEBUG) << "Start vm: " << backend->name();
@@ -543,6 +572,7 @@ FinalVMPtr CompileGraphs::CompileAndLink(const FuncGraphPtr &graph) {
   Reset();
   MS_LOG(DEBUG) << "Begin parameter:" << graph->parameters().size();
 
+// preprocess the graph
   FuncGraphPtr prim_graph = WrapPrimitives(graph);
   Compile(prim_graph);
   MS_EXCEPTION_IF_NULL(prim_graph);

mindspore/ccsrc/backend/graph_compiler/transform.h

@@ -127,6 +127,7 @@ class BACKEND_EXPORT CompileGraphs {
   FinalVMPtr CompileAndLink(const FuncGraphPtr &func_graph);
 
  protected:
+ // A vector stores instructions and their types
   InstSet insts_;
   mindspore::HashMap<FuncGraphPtr, int64_t> mapping_;
   CompileGraphPtr transform_;

mindspore/ccsrc/common/thread_pool.cc

@@ -47,55 +47,108 @@ ThreadPool::ThreadPool() {
   }
 }
 
+/**
+ * @brief Executes tasks in a synchronized loop within a thread.
+ * 
+ * The `SyncRunLoop` function is designed to execute tasks within a thread context. The thread 
+ * checks for the availability of a task in its context, and if found, executes it. 
+ * If there's no task available, the thread yields its execution in favor of other threads until a threshold, 
+ * after which it goes into a wait state.
+ * 
+ * @param context Shared pointer to the thread's context which contains the task and synchronization primitives.
+ * @return Void.
+ *
+ * @note
+ * This function is intended to be run by a worker thread in a thread pool. It continually checks 
+ * for tasks and executes them as they become available.
+ */
 void ThreadPool::SyncRunLoop(const std::shared_ptr<ThreadContext> &context) {
+  // Early exit if the thread context is not provided.
   if (context == nullptr) {
     return;
   }
+
+  // Initialize a counter for the number of times the thread yields without a task.
   size_t yield_count = 0;
+
+  // Infinite loop to keep the thread running and looking for tasks.
   while (true) {
+    // If the thread pool is signaled to exit, terminate this loop/thread.
     if (exit_run_) {
       return;
     }
 
+    // Check if there's no task in the context.
     if (!context->task) {
       ++yield_count;
+      // If the yield count exceeds the threshold, reset it and put the thread into a wait state.
       if (yield_count > kYieldThreshold) {
         yield_count = 0;
         std::unique_lock<std::mutex> lock(context->mutex);
         context->cond_var.wait(lock, [&context, this] { return context->task != nullptr || exit_run_; });
       } else {
+        // If the yield count hasn't reached the threshold, simply yield this thread's execution.
         std::this_thread::yield();
         continue;
       }
     }
 
+    // Recheck exit flag after potentially waiting.
     if (exit_run_) {
       return;
     }
 
+    // Execute the task within the context.
     try {
       auto &task = *(context->task);
       task();
     } catch (std::exception &e) {
+      // Handle any exceptions that might occur during the task execution.
       MsException::Instance().SetException();
     }
+    
+    // Reset the yield counter and clear the task from the context.
     yield_count = 0;
     context->task = nullptr;
   }
 }
 
+/**
+ * @brief Executes multiple tasks using the thread pool in a synchronized manner.
+ * 
+ * The `SyncRun` function takes a list of tasks and runs them using threads in the pool.
+ * If there are more tasks than available threads, the function will dynamically add more threads
+ * (up to a maximum limit). This function ensures that all tasks are executed and will use 
+ * synchronization primitives to efficiently manage task assignment and thread execution.
+ * 
+ * @param tasks A vector of tasks to be executed. Each task is represented as a callable object.
+ * @return Returns `true` if all tasks are successfully executed, otherwise returns `false`.
+ *
+ * @note
+ * - This function assumes that tasks do not throw exceptions.
+ * - It doesn't currently account for tasks' return values other than for a single task case.
+ */
 bool ThreadPool::SyncRun(const std::vector<Task> &tasks) {
+  // If there are no tasks, return true.
   if (tasks.empty()) {
     return true;
   }
+
+  // Special case: If there's only one task, run it and return its status.
   if (tasks.size() == 1) {
     auto ret = tasks[0]();
     return ret == SUCCESS;
   }
+
+  // Acquire a lock for the thread pool's state.
   std::unique_lock<std::mutex> lock(pool_mtx_);
   exit_run_ = false;
+
+  // Determine the number of tasks and currently available threads.
   size_t task_num = tasks.size();
   size_t thread_num = sync_run_threads_.size();
+
+  // If necessary, create additional threads up to the maximum limit.
   if (thread_num < max_thread_num_ && thread_num < task_num) {
     auto new_thread_num = max_thread_num_;
     if (task_num < max_thread_num_) {
@@ -107,13 +160,19 @@ bool ThreadPool::SyncRun(const std::vector<Task> &tasks) {
       sync_run_threads_.emplace_back(std::thread(&ThreadPool::SyncRunLoop, this, contexts_[i]));
     }
   }
+
+  // If there are no available contexts, return true.
   if (contexts_.empty()) {
     return true;
   }
+
+  // Determine the number of threads that will be used to run the tasks.
   size_t used_thread_num = contexts_.size();
   if (task_num < used_thread_num) {
     used_thread_num = task_num;
   }
+
+  // Assign tasks to threads and monitor their execution.
   bool running = true;
   size_t task_index = 0;
   while (running) {

mindspore/ccsrc/runtime/graph_scheduler/graph_compiler.cc

@@ -441,6 +441,7 @@ GraphId GraphCompiler::CompileGraph(const GraphSegmentPtr &segment, const AnfNod
   return graph_id;
 }
 
+// Not splited
 GraphId GraphCompiler::CompileGraph(const FuncGraphPtr &func_graph, const DeviceContext *device_context) {
   MS_EXCEPTION_IF_NULL(session_);
   MS_EXCEPTION_IF_NULL(func_graph);

mindspore/ccsrc/runtime/graph_scheduler/graph_scheduler.cc

@@ -555,28 +555,36 @@ void GraphScheduler::Schedule(const ActorSet *actor_set) {
 #endif
 }
 
+/**
+ * Runs the execution of actors within an actor set, following a specified graph execution strategy.
+ *
+ * @param actor_set Pointer to the ActorSet containing actors to be executed.
+ * @param device_contexts A vector of DeviceContext pointers representing the device contexts for execution.
+ * @param input_tensors A vector of vectors of TensorPtr representing input tensors for each actor.
+ * @param input_tensors_with_value_node A vector of TensorPtr representing input tensors with value nodes.
+ * @param strategy The graph execution strategy (e.g., step or pipeline).
+ */
 void GraphScheduler::Run(ActorSet *const actor_set, const std::vector<DeviceContext *> &device_contexts,
                          const std::vector<std::vector<TensorPtr>> &input_tensors,
                          const std::vector<TensorPtr> &input_tensors_with_value_node, GraphExecutionStrategy strategy) {
+  // Check for null pointers.
   MS_EXCEPTION_IF_NULL(actor_set);
   MS_EXCEPTION_IF_NULL(actor_set->data_prepare_actor_);
-#if !defined(_WIN32) && !defined(_WIN64)
-  SignalGuard sg(IntHandler);
-#endif
 
-  // Construct OpContext.
+  // Initialize OpContext.
   OpContext<DeviceTensor> op_context;
   std::vector<Promise<int>> result(1);
   op_context.sequential_num_ = RandInt::Instance().Get();
   op_context.results_ = &result;
 
+  // Set OpContext for RPC actors if needed.
 #ifdef ENABLE_RPC_ACTOR
-  // Set OpContext to rpc node scheduler.
   auto op_context_setter =
     std::make_shared<RpcActorOpContextSetter>(rpc_node_scheduler_.get(), actor_set->rpc_actors_, &op_context);
   MS_EXCEPTION_IF_NULL(op_context_setter);
 #endif
 
+  // Handle single-op execution for specific cases.
   if ((strategy == GraphExecutionStrategy::kStep) && IsSingleOpActorSet(actor_set)) {
     actor_set->data_prepare_actor_->PrepareData(input_tensors, &op_context, GraphExecutionStrategy::kStep);
     MS_EXCEPTION_IF_NULL(actor_set->kernel_actors_[0]);
@@ -593,36 +601,42 @@ void GraphScheduler::Run(ActorSet *const actor_set, const std::vector<DeviceCont
   ActorDispatcher::Send(actor_set->data_prepare_actor_->GetAID(), &DataPrepareActor::PrepareData, input_tensors,
                         &op_context, GraphExecutionStrategy::kPipeline);
 
-  // Get the run result.
+  // Wait for the data preparation to complete.
   auto result_future = result[0].GetFuture();
   result_future.Wait();
   MsException::Instance().CheckException();
+
+  // Handle potential errors during data preparation.
   if (!result_future.IsOK()) {
 #ifdef ENABLE_DUMP_IR
     mindspore::RDR::TriggerAll();
 #endif
-    // When temporary variable 'op_context' has beed set failed status, the main thread need wait other threads until
-    // they finish respective task, otherwise segmentation fault will happen when these task access 'op_context',
-    // because it has been destroyed.
+    // When the 'op_context' has been set with a failed status, the main thread must wait for other threads to finish
+    // their respective tasks to avoid segmentation faults, as 'op_context' has been destroyed.
     std::mutex mutex;
     std::unique_lock<std::mutex> locker(mutex);
     std::condition_variable thread_blocker;
     const int64_t kTimeToWait = 2;
     (void)thread_blocker.wait_for(locker, std::chrono::seconds(kTimeToWait));
-    // May set exception in the wait time, need throw the exception to avoid affecting the next execution.
+    // Check for exceptions during the wait time and throw them to avoid affecting the next execution.
     MsException::Instance().CheckException();
     MS_LOG(EXCEPTION) << op_context.error_info_;
   }
 
+  // Calculate execution time.
   double end_time = GetTime();
   const size_t kSecondsToMilliseconds = 1000;
+
+  // Set actor execution strategy and timing information.
   SetActorExecutionStrategy(actor_set, strategy, (end_time - start_time) * kSecondsToMilliseconds);
 
 #if ((defined ENABLE_CPU) && (!defined _WIN32) && (!defined _WIN64))
+  // Handle disaster recovery for CPU execution.
   DoDisasterRecovery(actor_set->name_);
 #endif
 }
 
+
 void GraphScheduler::SetActorExecutionStrategy(ActorSet *const actor_set, GraphExecutionStrategy strategy,
                                                double execution_time) const {
   MS_EXCEPTION_IF_NULL(actor_set);

mindspore/ccsrc/utils/cse.cc

@@ -36,6 +36,12 @@ using mindspore::abstract::AbstractBase;
 using mindspore::abstract::AbstractFunction;
 using mindspore::abstract::AbstractFunctionPtr;
 
+/**
+ * @brief Check if the given node has recomputed scope.
+ *
+ * @param node The node to check.
+ * @return true if the node has recomputed scope, otherwise false.
+ */
 bool WithRecomputedScope(const AnfNodePtr &node) {
   MS_EXCEPTION_IF_NULL(node);
   if (!node->isa<CNode>()) {
@@ -45,11 +51,24 @@ bool WithRecomputedScope(const AnfNodePtr &node) {
   return full_name_with_scope.find(kAttrRecompute) == 0;
 }
 
+/**
+ * @brief Check if the given nodes are set as recomputed.
+ *
+ * @param a The first node.
+ * @param b The second node.
+ * @return true if any node is set as recomputed, otherwise false.
+ */
 bool IsSetRecomputed(const CNodePtr &a, const CNodePtr &b) {
   return (WithRecomputedScope(a) && !a->HasAttr(kAttrNeedCseAfterRecompute)) ||
          (WithRecomputedScope(b) && !b->HasAttr(kAttrNeedCseAfterRecompute));
 }
 
+/**
+ * @brief Update debug info and dump flag for the given nodes.
+ *
+ * @param main The main node.
+ * @param node The node to check.
+ */
 void UpdateDebugInfoAndDumpFlag(const AnfNodePtr &main, const AnfNodePtr &node) {
   if (main == nullptr || !main->isa<CNode>()) {
     return;
@@ -61,6 +80,13 @@ void UpdateDebugInfoAndDumpFlag(const AnfNodePtr &main, const AnfNodePtr &node)
   main_cnode->AddFusedDebugInfo(node);
 }
 
+/**
+ * @brief Get abstract representation of a node.
+ *
+ * @param node The node to process.
+ * @param ignore_fg_abs_tracking_id Flag to decide if fg_abs_tracking_id should be ignored.
+ * @return Abstract representation of the node.
+ */
 BasePtr AbsOf(const AnfNodePtr &node, bool ignore_fg_abs_tracking_id) {
   MS_EXCEPTION_IF_NULL(node);
   auto node_abs = node->abstract();
@@ -81,15 +107,33 @@ BasePtr AbsOf(const AnfNodePtr &node, bool ignore_fg_abs_tracking_id) {
   return node_abs;
 }
 
+/*
+ * @brief Build order groups and do replace for one graph.
+ *
+ * This function groups all nodes that might have common subexpressions and
+ * passes them to the DoReplace() function for actual processing. It also
+ * builds order groups based on the nodes in a single graph.
+ *
+ * @param fg The graph to be processed.
+ * @param manager The manager of the graph.
+ * @return Returns true if the replace operation was successful, otherwise false.
+ */
 bool CSE::BuildOrderGroupAndDoReplaceForOneGraph(const FuncGraphPtr &fg, const FuncGraphManagerPtr &manager) const {
   MS_EXCEPTION_IF_NULL(fg);
+  
+  // Lists to store ordering of groups, groupings of nodes based on hash values, and hash values for each node.
   std::vector<std::size_t> order_group;
   mindspore::HashMap<std::size_t, std::vector<AnfNodePtr>> groups;
   mindspore::HashMap<AnfNodePtr, std::size_t> hashes;
 
+  // Topologically sort the nodes in the function graph starting from the return node.
   std::vector<AnfNodePtr> toposet = TopoSort(fg->get_return());
+
+  // Compute the hash value for each node and group nodes based on these values.
   for (auto node : toposet) {
     MS_EXCEPTION_IF_NULL(node);
+    
+    // Skip nodes that have already been hashed.
     if (hashes.find(node) != hashes.end()) {
       continue;
     }
@@ -99,21 +143,26 @@ bool CSE::BuildOrderGroupAndDoReplaceForOneGraph(const FuncGraphPtr &fg, const F
       ValueNodePtr value_node = node->cast<ValueNodePtr>();
       auto value = value_node->value();
       MS_EXCEPTION_IF_NULL(value);
+      // Combine the hash of the node's value with its abstract hash.
       h = hash_combine(value->hash(), (AbsOf(value_node, true)->hash()));
     } else if (node->isa<CNode>()) {
       auto cnode = node->cast<CNodePtr>();
       auto &inputs = cnode->inputs();
       size_t init = 0;
+      // Combine the hash values of all inputs to the compute node.
       h = std::accumulate(inputs.begin(), inputs.end(), init, [&hashes](std::size_t hash, const AnfNodePtr &node_in) {
         return hash_combine(hash, hashes[node_in]);
       });
     } else if (node->isa<Parameter>()) {
+      // For parameter nodes, use the node's hash.
       h = node->hash();
     } else {
       MS_LOG(ERROR) << "Unknown node type";
     }
 
     hashes[node] = h;
+
+    // Group the node based on its hash value.
     if (groups.find(h) == groups.end()) {
       std::vector<AnfNodePtr> innervec({node});
       groups[h] = innervec;
@@ -122,17 +171,42 @@ bool CSE::BuildOrderGroupAndDoReplaceForOneGraph(const FuncGraphPtr &fg, const F
       groups[h].push_back(node);
     }
   }
+
+  // Attempt to replace nodes within each group.
   return DoReplace(manager, order_group, &groups);
 }
 
+/*
+ * @brief Builds order groups and executes replacements for all graphs managed by the given manager.
+ *
+ * Iterates through all the graphs available in the manager and performs
+ * order group building and replacements on each of them. The change status
+ * is aggregated across all the graphs.
+ *
+ * @param manager The manager handling the graphs to be processed.
+ * @return Returns true if there were any changes, otherwise false.
+ */
 bool CSE::BuildOrderGroupAndDoReplace(const FuncGraphManagerPtr manager) const {
   bool changed = false;
+
+  // Iterate over all function graphs managed by the manager.
   for (FuncGraphPtr fg : manager->func_graphs()) {
+    // Attempt to replace nodes for the current function graph and update the 'changed' status.
     changed = BuildOrderGroupAndDoReplaceForOneGraph(fg, manager) || changed;
   }
+
   return changed;
 }
 
+/*
+ * @brief Checks if the node has hidden side effects.
+ *
+ * This function checks if the given node has any attributes indicating
+ * that it has hidden side effects.
+ *
+ * @param node The node to be checked.
+ * @return Returns true if the node has hidden side effects, otherwise false.
+ */
 bool CSE::HasHiddenSideEffect(const AnfNodePtr &node) {
   auto prim = GetCNodePrimitive(node);
   if (prim == nullptr) {
@@ -141,6 +215,17 @@ bool CSE::HasHiddenSideEffect(const AnfNodePtr &node) {
   return prim->HasAttr(GRAPH_FLAG_SIDE_EFFECT_HIDDEN);
 }
 
+/*
+ * @brief Checks whether a node can be replaced by another.
+ *
+ * This function performs various checks to determine if the 'main' node can
+ * be replaced by the 'node'. It takes into account the types and attributes
+ * of the nodes as well as the values they contain.
+ *
+ * @param main The main node that might be replaced.
+ * @param node The node that is considered as a replacement.
+ * @return Returns true if 'main' can be replaced by 'node', otherwise false.
+ */
 bool CSE::CheckReplace(const AnfNodePtr &main, const AnfNodePtr &node) const {
   MS_EXCEPTION_IF_NULL(main);
   MS_EXCEPTION_IF_NULL(node);
@@ -186,30 +271,40 @@ bool CSE::CheckReplace(const AnfNodePtr &main, const AnfNodePtr &node) const {
   return false;
 }
 
+/*
+ * @brief Perform the actual replacement step of CSE.
+ * It inspects all node groups to seek CSE opportunities.
+ *
+ * @param manager The manager of the graph.
+ * @param order_group The order group of the graph.
+ * @param groups The groups of the graph.
+ *
+ */
 bool CSE::DoReplace(const FuncGraphManagerPtr manager, const std::vector<std::size_t> &order_group,
                     mindspore::HashMap<std::size_t, std::vector<AnfNodePtr>> *groups) const {
   bool changes = false;
   std::set<size_t> clear_set;
+
+  // Iterate over each group.
   for (auto &h : order_group) {
     std::vector<AnfNodePtr> &group = (*groups)[h];
-    // If there are more than 2 node in that group, they may be same common expression can be eliminated.
+
+    // If there are more than 1 node in the group, they might represent the same computation.
     if (group.size() > 1) {
+      // Check each node against every other node in the group.
       for (size_t k = 0; k < group.size() - 1; k++) {
         AnfNodePtr main = group[k];
         MS_EXCEPTION_IF_NULL(main);
 
-        // When all node in group has been replaced
-        // or a valuenode node, skip compare in group
+        // Skip nodes that have already been replaced or are value nodes.
         if ((k + 1 + clear_set.size() == group.size()) || (k > 0 && main->isa<ValueNode>())) {
           break;
         }
 
-        // skip node has been replaced
         if (clear_set.find(k) != clear_set.end()) {
           continue;
         }
 
-        // Compare with rest elements in this group.
         for (size_t i = k + 1; i < group.size(); i++) {
           auto node = group[i];
           MS_EXCEPTION_IF_NULL(node);
@@ -217,17 +312,24 @@ bool CSE::DoReplace(const FuncGraphManagerPtr manager, const std::vector<std::si
           if (clear_set.find(i) != clear_set.end()) {
             continue;
           }
+
+          // Nodes must belong to the same function graph.
           if (main->func_graph() != node->func_graph()) {
             continue;
           }
+
+          // Check if the nodes are equivalent.
           if (CheckReplace(node, main)) {
             changes = true;
+            // Optional: Update debug info or dump flags.
             UpdateDebugInfoAndDumpFlag(main, node);
+            // Replace the node.
             (void)manager->Replace(node, main);
             (void)clear_set.insert(i);
           }
         }
       }
+
       clear_set.clear();
     }
   }
@@ -235,6 +337,7 @@ bool CSE::DoReplace(const FuncGraphManagerPtr manager, const std::vector<std::si
   return changes;
 }
 
+
 bool CSE::Cse(const FuncGraphPtr root, const FuncGraphManagerPtr manager) const {
   MS_EXCEPTION_IF_NULL(manager);
   manager->AddFuncGraph(root);