insert profiling iteration end point after last tbe kernel

mindspore/ccsrc/device/ascend/profiling/profiling_utils.cc

@@ -33,7 +33,7 @@ constexpr char kIterEndNode[] = "PROFILING_ITER_END";
 std::unordered_map<uint32_t, std::vector<std::string>> ProfilingUtils::graph_kernel_name_;
 uint32_t ProfilingUtils::custom_node_index_ = 1;
 
-ProfilingTraceInfo ProfilingUtils::GetProfilingTraceFromEnv(const NotNull<session::KernelGraph *> graph_ptr) {
+ProfilingTraceInfo ProfilingUtils::GetProfilingTraceFromEnv(NotNull<const session::KernelGraph *> graph_ptr) {
   MS_LOG(INFO) << "get env start";
   custom_node_index_ = 1;
   auto &cnode_exec_order = graph_ptr->execution_order();
@@ -148,18 +148,29 @@ std::string ProfilingUtils::GetTraceBpEnd(const std::vector<CNodePtr> &cnode_exe
   }
 
   if (bp_end_str.empty()) {
-    auto last_cnode = cnode_exec_order.back();
+    bp_end_str = GetGraphLastTbeKernelName(cnode_exec_order);
-    MS_EXCEPTION_IF_NULL(last_cnode);
-    bp_end_str = last_cnode->fullname_with_scope();
   }
   return bp_end_str;
 }
 
+std::string ProfilingUtils::GetGraphLastTbeKernelName(const std::vector<CNodePtr> &cnode_exec_order) {
+  std::string last_tbe_kernel_name = "";
+  // find last tbe_kernel
+  for (auto iter = cnode_exec_order.rbegin(); iter != cnode_exec_order.rend(); ++iter) {
+    if (AnfAlgo::GetKernelType(*iter) == TBE_KERNEL) {
+      last_tbe_kernel_name = (*iter)->fullname_with_scope();
+      break;
+    }
+  }
+  if (last_tbe_kernel_name.empty()) {
+    MS_LOG(WARNING) << "tbe kernel not found in graph";
+  }
+  return last_tbe_kernel_name;
+}
+
 std::string ProfilingUtils::GetTraceNetoutput(const std::vector<CNodePtr> &cnode_exec_order) {
   const char *trace_netoutput = std::getenv(kIterEndNode);
-  auto &last_cnode = cnode_exec_order.back();
+  return trace_netoutput == nullptr ? GetGraphLastTbeKernelName(cnode_exec_order) : std::string(trace_netoutput);
-  MS_EXCEPTION_IF_NULL(last_cnode);
-  return trace_netoutput == nullptr ? last_cnode->fullname_with_scope() : std::string(trace_netoutput);
 }
 
 NotNull<CNodePtr> ProfilingUtils::CreateProfilingCNode(const ProfilingContent &profiling_content,

mindspore/ccsrc/device/ascend/profiling/profiling_utils.h

@@ -94,7 +94,7 @@ class ProfilingUtils {
   // And other cnode, like AllReduce, export PROFILING_CUSTOM_1='full name of AllReduce cnode'
   // GetNext, export PROFIFLING_CUSTOM_2='full name fo GetNext cnode'
   // The variable i in PROFILING_CUSTOM_i should start from 1 without interruption.
-  static ProfilingTraceInfo GetProfilingTraceFromEnv(const NotNull<session::KernelGraph *> graph_ptr);
+  static ProfilingTraceInfo GetProfilingTraceFromEnv(NotNull<const session::KernelGraph *> graph_ptr);
 
   // Insert two profiling trace points, one in front and one behind
   static void ProfilingCustomOp(const mindspore::AnfNodePtr &anf_node, const ProfilingTraceInfo &profiling_trace_info,
@@ -114,6 +114,7 @@ class ProfilingUtils {
   static std::string GetTraceBegin(const std::vector<CNodePtr> &cnode_exec_order);
   static std::string GetTraceBpEnd(const std::vector<CNodePtr> &cnode_exec_order);
   static std::string GetTraceNetoutput(const std::vector<CNodePtr> &cnode_exec_order);
+  static std::string GetGraphLastTbeKernelName(const std::vector<CNodePtr> &cnode_exec_order);
   static void GetTraceHccl(const std::vector<CNodePtr> &cnode_exec_order,
                            NotNull<ProfilingTraceInfo *> profiling_trace);
   static void GetCNodeOutputRealNode(const std::string &node_name, const std::vector<CNodePtr> &cnode_exec_order,