Add sparse core step time breakdown to overview page.

PiperOrigin-RevId: 679197985
openxla · Sep 26, 2024 · 89c60d4 · 89c60d4
1 parent 0e732d6
commit 89c60d4
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 2 deletions.
diff --git a/third_party/tsl/tsl/profiler/utils/xplane_schema.cc b/third_party/tsl/tsl/profiler/utils/xplane_schema.cc
@@ -49,6 +49,7 @@ const absl::string_view kHostCpusPlaneName = "Host CPUs";
 const absl::string_view kSyscallsPlaneName = "Syscalls";
 
 const absl::string_view kStepLineName = "Steps";
+const absl::string_view kSparseCoreStepLineName = "Sparse Core Steps";
 const absl::string_view kTensorFlowNameScopeLineName = "Framework Name Scope";
 const absl::string_view kTensorFlowOpLineName = "Framework Ops";
 const absl::string_view kXlaModuleLineName = "XLA Modules";
@@ -340,6 +341,7 @@ const StatTypeMap& GetStatTypeMap() {
       {"cuda_graph_id", kCudaGraphId},
       {"cuda_graph_exec_id", kCudaGraphExecId},
       {"cuda_graph_orig_id", kCudaGraphOrigId},
+      {"step_idle_time_ps", kStepIdleTimePs},
   });
   DCHECK_EQ(stat_type_map->size(), kNumStatTypes);
   return *stat_type_map;

diff --git a/third_party/tsl/tsl/profiler/utils/xplane_schema.h b/third_party/tsl/tsl/profiler/utils/xplane_schema.h
@@ -70,6 +70,7 @@ TF_CONST_INIT extern const absl::string_view kTensorFlowNameScopeLineName;
 TF_CONST_INIT extern const absl::string_view kTensorFlowOpLineName;
 TF_CONST_INIT extern const absl::string_view kXlaModuleLineName;
 TF_CONST_INIT extern const absl::string_view kXlaOpLineName;
+TF_CONST_INIT extern const absl::string_view kSparseCoreStepLineName;
 TF_CONST_INIT extern const absl::string_view kXlaAsyncOpLineName;
 TF_CONST_INIT extern const absl::string_view kKernelLaunchLineName;
 TF_CONST_INIT extern const absl::string_view kSourceLineName;
@@ -328,7 +329,8 @@ enum StatType {
   // on the GPU device when tracing is in graph level.
   kCudaGraphExecId,
   kCudaGraphOrigId,
-  kLastStatType = kCudaGraphOrigId,
+  kStepIdleTimePs,
+  kLastStatType = kStepIdleTimePs,
 };
 
 enum MegaScaleStatType : uint8_t {

diff --git a/third_party/tsl/tsl/profiler/utils/xplane_utils.cc b/third_party/tsl/tsl/profiler/utils/xplane_utils.cc
@@ -548,7 +548,8 @@ void AggregateXPlane(const XPlane& full_trace, XPlane& aggregated_trace) {
   uint64_t last_op_end_ps = 0;
 
   plane.ForEachLine([&](const XLineVisitor& line) {
-    if (line.Name() == kStepLineName) {
+    if (line.Name() == kStepLineName ||
+        line.Name() == kSparseCoreStepLineName) {
       XLineBuilder aggregated_line =
           aggregated_plane.GetOrCreateLine(line.Id());
       aggregated_line.SetName(kStepLineName);