mindspore/ccsrc/backend/graph_compiler/backend.cc

@@ -184,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));
@@ -191,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);
@@ -222,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) {
@@ -247,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
@@ -478,28 +527,43 @@ 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 root graph
   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);
   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();
@@ -510,27 +574,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_);
@@ -559,6 +637,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.