Create events for TF name scope.

PiperOrigin-RevId: 295003367
Change-Id: Id4fc85f9ad35f968ff4b8b8730c2383e4c8ec214

tensorflow/core/profiler/utils/BUILD

@@ -270,6 +270,7 @@ cc_library(
     hdrs = ["derived_timeline.h"],
     deps = [
         ":group_events",
+        ":tf_op_utils",
         ":tf_xplane_visitor",
         ":trace_utils",
         ":xplane_builder",

tensorflow/core/profiler/utils/derived_timeline.cc

@@ -16,6 +16,8 @@ limitations under the License.
 
 #include "absl/strings/str_split.h"
 #include "tensorflow/core/lib/gtl/map_util.h"
+#include "tensorflow/core/profiler/protobuf/xplane.pb.h"
+#include "tensorflow/core/profiler/utils/tf_op_utils.h"
 #include "tensorflow/core/profiler/utils/tf_xplane_visitor.h"
 #include "tensorflow/core/profiler/utils/trace_utils.h"
 #include "tensorflow/core/profiler/utils/xplane_builder.h"
@@ -106,11 +108,37 @@ class DerivedXLineBuilder {
 };
 
 const absl::string_view kDerivedLineSteps = "Steps";
+const absl::string_view kDerivedLineTensorFlowNameScope =
+    "TensorFlow Name Scope";
 const absl::string_view kDerivedLineTensorFlowOps = "TensorFlow Ops";
 const absl::string_view kDerivedLineXlaModules = "XLA Modules";
 const absl::string_view kDerivedLineXlaOps = "XLA Ops";
 const absl::string_view kAnnotationDelimiter = "::";
 
+void ProcessTfOpEvent(const XEventVisitor& event,
+                      absl::string_view tf_op_full_name,
+                      absl::optional<int64> group_id,
+                      XPlaneBuilder* plane_builder,
+                      DerivedXLineBuilder* tf_name_scope_line_builder,
+                      DerivedXLineBuilder* tf_op_line_builder) {
+  TfOp tf_op = ParseTfOpFullname(tf_op_full_name);
+  if (tf_op.is_tf_op) {
+    std::vector<const XEventMetadata*> tf_name_scope_metadata_per_level;
+    for (const auto& tf_name_scope : ParseTfNameScopes(tf_op)) {
+      tf_name_scope_metadata_per_level.push_back(
+          plane_builder->GetOrCreateEventMetadata(tf_name_scope));
+    }
+    tf_name_scope_line_builder->ExpandOrAddEvents(
+        tf_name_scope_metadata_per_level, event, group_id);
+  }
+  XEventMetadata* event_metadata =
+      plane_builder->GetOrCreateEventMetadata(tf_op_full_name);
+  // Set the display name to op_type so that the events of the same op_type have
+  // the same color in the trace viewer.
+  event_metadata->set_display_name(TfOpEventName(tf_op));
+  tf_op_line_builder->ExpandOrAddEvents({event_metadata}, event, group_id);
+}
+
 }  // namespace
 
 void DeriveEventsFromAnnotations(const SymbolResolver& symbol_resolver,
@@ -131,18 +159,22 @@ void DeriveEventsFromAnnotations(const SymbolResolver& symbol_resolver,
   XPlaneBuilder plane(device_trace);
   DerivedXLineBuilder tf_ops(&plane, kThreadIdTfOp, kDerivedLineTensorFlowOps,
                              start_timestamp_ns, {}, /*try_expand=*/true);
+  DerivedXLineBuilder tf_name_scope(
+      &plane, kThreadIdTfNameScope, kDerivedLineTensorFlowNameScope,
+      start_timestamp_ns, {&tf_ops}, /*try_expand=*/true);
   DerivedXLineBuilder hlo_ops(&plane, kThreadIdHloOp, kDerivedLineXlaOps,
                               start_timestamp_ns, {}, /*try_expand=*/true);
-  DerivedXLineBuilder hlo_modules(&plane, kThreadIdHloModule,
-                                  kDerivedLineXlaModules, start_timestamp_ns,
-                                  {&tf_ops, &hlo_ops}, /*try_expand=*/false);
+  DerivedXLineBuilder hlo_modules(
+      &plane, kThreadIdHloModule, kDerivedLineXlaModules, start_timestamp_ns,
+      {&tf_ops, &tf_name_scope, &hlo_ops}, /*try_expand=*/false);
   DerivedXLineBuilder steps(&plane, kThreadIdStepInfo, kDerivedLineSteps,
-                            start_timestamp_ns, {&tf_ops, &hlo_ops},
+                            start_timestamp_ns,
+                            {&tf_ops, &tf_name_scope, &hlo_ops},
                             /*try_expand=*/true);
 
   // Process events in order by start time.
   for (const XEventVisitor& event : events) {
-    absl::string_view tf_op_fullname;
+    absl::string_view tf_op_full_name;
     absl::string_view hlo_module_name;
     std::vector<absl::string_view> hlo_op_names;
     absl::optional<int64> group_id;
@@ -151,7 +183,7 @@ void DeriveEventsFromAnnotations(const SymbolResolver& symbol_resolver,
       if (stat.Type() == StatType::kGroupId) {
         group_id = stat.IntValue();
       } else if (stat.Type() == StatType::kLevel0) {
-        tf_op_fullname = stat.StrValue();
+        tf_op_full_name = stat.StrValue();
       } else if (stat.Type() == StatType::kHloOp) {
         hlo_op_names = absl::StrSplit(stat.StrValue(), kAnnotationDelimiter);
       } else if (stat.Type() == StatType::kHloModule) {
@@ -188,12 +220,12 @@ void DeriveEventsFromAnnotations(const SymbolResolver& symbol_resolver,
       hlo_ops.ExpandOrAddEvents(hlo_op_metadata_per_level, event, group_id);
       auto tf_op_name = symbol_resolver(hlo_module_name, hlo_op_names.back());
       if (!tf_op_name.empty()) {
-        tf_ops.ExpandOrAddEvents({plane.GetOrCreateEventMetadata(tf_op_name)},
-                                 event, group_id);
+        ProcessTfOpEvent(event, tf_op_name, group_id, &plane, &tf_name_scope,
+                         &tf_ops);
       }
-    } else if (!tf_op_fullname.empty()) {  // GPU kernel not compiled by XLA
-      tf_ops.ExpandOrAddEvents({plane.GetOrCreateEventMetadata(tf_op_fullname)},
-                               event, group_id);
+    } else if (!tf_op_full_name.empty()) {  // GPU kernel not compiled by XLA
+      ProcessTfOpEvent(event, tf_op_full_name, group_id, &plane, &tf_name_scope,
+                       &tf_ops);
     }
   }
   RemoveEmptyLines(device_trace);

tensorflow/core/profiler/utils/derived_timeline_test.cc

@@ -78,7 +78,7 @@ TEST(DerivedTimelineTest, HloModuleNameTest) {
 
 // Checks that the TF op events are expanded.
 TEST(DerivedTimelineTest, TfOpLineTest) {
-  const absl::string_view kTfOpName = "Mul";
+  const absl::string_view kTfOpName = "mul:Mul";
   const absl::string_view kKernelDetails = "kernel_details";
   XSpace space;
   EventGroupNameMap event_group_name_map;
@@ -120,7 +120,7 @@ TEST(DerivedTimelineTest, TfOpLineTest) {
 // Checks that the dependency between the step line and the TF op line prevents
 // TF op events from being expanded.
 TEST(DerivedTimelineTest, DependencyTest) {
-  const absl::string_view kTfOpName = "Mul";
+  const absl::string_view kTfOpName = "mul:Mul";
   const absl::string_view kKernelDetails = "kernel_details";
   XSpace space;
   EventGroupNameMap event_group_name_map({{0, "train 0"}, {1, "train 1"}});
@@ -155,6 +155,56 @@ TEST(DerivedTimelineTest, DependencyTest) {
   });
 }
 
+// Checks that the TF op events are expanded.
+TEST(DerivedTimelineTest, TfOpNameScopeTest) {
+  const absl::string_view kTfOpName = "scope1/scope2/mul:Mul";
+  const absl::string_view kKernelDetails = "kernel_details";
+  XSpace space;
+  EventGroupNameMap event_group_name_map;
+  XPlane* plane = space.add_planes();
+  XPlaneBuilder plane_builder(plane);
+  auto line_builder = plane_builder.GetOrCreateLine(0);
+  auto first_event =
+      CreateXEvent(&plane_builder, &line_builder, "op1", 0, 100, {});
+  first_event.AddStatValue(
+      *plane_builder.GetOrCreateStatMetadata(GetStatTypeStr(StatType::kLevel0)),
+      kTfOpName);
+  first_event.AddStatValue(*plane_builder.GetOrCreateStatMetadata(
+                               GetStatTypeStr(StatType::kKernelDetails)),
+                           kKernelDetails);
+  auto second_event =
+      CreateXEvent(&plane_builder, &line_builder, "op2", 200, 300, {});
+  second_event.AddStatValue(
+      *plane_builder.GetOrCreateStatMetadata(GetStatTypeStr(StatType::kLevel0)),
+      kTfOpName);
+  second_event.AddStatValue(*plane_builder.GetOrCreateStatMetadata(
+                                GetStatTypeStr(StatType::kKernelDetails)),
+                            kKernelDetails);
+  GenerateDerivedTimeLines(event_group_name_map, &space);
+  XPlaneVisitor plane_visitor = CreateTfXPlaneVisitor(plane);
+  // The TF name scope line and the TF op line are added.
+  EXPECT_EQ(plane_visitor.NumLines(), 3);
+  plane_visitor.ForEachLine([&](const XLineVisitor& line_visitor) {
+    int64 line_id = line_visitor.Id();
+    if (line_id == 0) {
+      return;
+    } else if (line_id == kThreadIdTfNameScope) {
+      EXPECT_EQ(line_visitor.NumEvents(), 2);
+      line_visitor.ForEachEvent([&](const XEventVisitor& event_visitor) {
+        EXPECT_EQ(event_visitor.OffsetPs(), 0);
+        EXPECT_EQ(event_visitor.DurationPs(), 500);
+      });
+    } else if (line_id == kThreadIdTfOp) {
+      EXPECT_EQ(line_visitor.NumEvents(), 1);
+      line_visitor.ForEachEvent([&](const XEventVisitor& event_visitor) {
+        EXPECT_EQ(event_visitor.Name(), kTfOpName);
+        EXPECT_EQ(event_visitor.OffsetPs(), 0);
+        EXPECT_EQ(event_visitor.DurationPs(), 500);
+      });
+    }
+  });
+}
+
 }  // namespace
 }  // namespace profiler
 }  // namespace tensorflow

tensorflow/core/profiler/utils/tf_op_utils.cc

@@ -15,8 +15,6 @@ limitations under the License.
 
 #include "tensorflow/core/profiler/utils/tf_op_utils.h"
 
-#include <vector>
-
 #include "absl/strings/ascii.h"
 #include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
@@ -49,7 +47,7 @@ TfOp ParseTfOpFullname(absl::string_view tf_op_fullname) {
   // JAX op types have only lowercase letters and underscores.
   static const LazyRE2 kJaxOpTypeRegEx = {"[a-z_]*"};
 
-  TfOp tf_op = {tf_op_fullname, kUnknownOp};
+  TfOp tf_op = {tf_op_fullname, kUnknownOp, /*is_tf_op=*/false};
   std::vector<absl::string_view> parts =
       absl::StrSplit(tf_op_fullname, absl::MaxSplits(':', 1));
   if (parts.size() != 2) {
@@ -67,29 +65,38 @@ TfOp ParseTfOpFullname(absl::string_view tf_op_fullname) {
     tf_op.type = kDatasetOp;
   } else if (RE2::FullMatch(parts[1], *kTfOpTypeRegEx) &&
              RE2::FullMatch(parts[0], *kTfOpNameRegEx)) {  // TensorFlow
-    tf_op = {parts[0], parts[1]};
+    tf_op = {parts[0], parts[1], /*is_tf_op=*/true};
   } else if (RE2::FullMatch(parts[1], *kJaxOpTypeRegEx)) {  // JAX
-    tf_op = {parts[0], parts[1]};
+    tf_op = {parts[0], parts[1], /*is_tf_op=*/false};
   }
   return tf_op;
 }
 
-std::string TfOpEventName(absl::string_view tf_op_fullname) {
+std::vector<absl::string_view> ParseTfNameScopes(const TfOp& tf_op) {
+  std::vector<absl::string_view> name_scopes = absl::StrSplit(tf_op.name, '/');
+  // The last element is an op name not TF name scope.
+  if (!name_scopes.empty()) name_scopes.pop_back();
+  return name_scopes;
+}
+
+std::string TfOpEventName(const TfOp& tf_op) {
   std::string event_name;
-  TfOp op = ParseTfOpFullname(tf_op_fullname);
-  if (op.type == kUnknownOp) {
+  if (tf_op.type == kUnknownOp) {
     // Some TraceMe names contain trailing whitespace, remove it.
-    event_name =
-        std::string(absl::StripTrailingAsciiWhitespace(tf_op_fullname));
-  } else if (op.type == kDatasetOp) {
+    event_name = std::string(absl::StripTrailingAsciiWhitespace(tf_op.name));
+  } else if (tf_op.type == kDatasetOp) {
     std::vector<absl::string_view> op_parts =
-        absl::StrSplit(tf_op_fullname, kSeparator);
+        absl::StrSplit(tf_op.name, kSeparator);
     event_name = absl::StrCat(kIterator, kSeparator, op_parts.back());
   } else {
-    event_name = std::string(op.type);
+    event_name = std::string(tf_op.type);
   }
   return event_name;
 }
 
+std::string TfOpEventName(absl::string_view tf_op_fullname) {
+  return TfOpEventName(ParseTfOpFullname(tf_op_fullname));
+}
+
 }  // namespace profiler
 }  // namespace tensorflow

tensorflow/core/profiler/utils/tf_op_utils.h

@@ -16,6 +16,8 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_PROFILER_UTILS_TF_OP_UTILS_H_
 #define TENSORFLOW_CORE_PROFILER_UTILS_TF_OP_UTILS_H_
 
+#include <vector>
+
 #include "absl/strings/match.h"
 #include "absl/strings/string_view.h"
 
@@ -32,12 +34,17 @@ ABSL_CONST_INIT extern const absl::string_view kMemcpyDToHOp;
 struct TfOp {
   absl::string_view name;
   absl::string_view type;
+  bool is_tf_op;
 };
 
 TfOp ParseTfOpFullname(absl::string_view tf_op_fullname);
 
+// Returns a vector of TF name scopes extracted from tf_op_full_name.
+std::vector<absl::string_view> ParseTfNameScopes(const TfOp& tf_op);
+
 // Trace event name for TF ops is the op type so they have the same color in
 // trace viewer.
+std::string TfOpEventName(const TfOp& tf_op);
 std::string TfOpEventName(absl::string_view tf_op_fullname);
 
 // Returns true if the given name is not a TensorFlow op.

tensorflow/core/profiler/utils/trace_utils.h

@@ -23,10 +23,12 @@ namespace profiler {
 // First derived stream/thread id.
 constexpr int kThreadIdDerivedMin = 0xdeadbeef;
 constexpr int kThreadIdStepInfo = kThreadIdDerivedMin;
-constexpr int kThreadIdTfOp = kThreadIdDerivedMin + 1;
-constexpr int kThreadIdHloModule = kThreadIdDerivedMin + 2;
-constexpr int kThreadIdHloOp = kThreadIdDerivedMin + 3;
-constexpr int kThreadIdOverhead = kThreadIdDerivedMin + 4;
+constexpr int kThreadIdTfNameScope = kThreadIdDerivedMin + 1;
+constexpr int kThreadIdTfOp = kThreadIdDerivedMin + 2;
+constexpr int kThreadIdHloModule = kThreadIdDerivedMin + 3;
+constexpr int kThreadIdHloOp = kThreadIdDerivedMin + 4;
+constexpr int kThreadIdOverhead = kThreadIdDerivedMin + 5;
+
 // Last derived stream/thread id.
 constexpr int kThreadIdDerivedMax = kThreadIdOverhead;