Add HostEventType to xplane_schema.

PiperOrigin-RevId: 285822505
Change-Id: Ieb1d4af1df36016e6ecb919c32334c84d2327453

tensorflow/core/profiler/utils/tf_op_utils.cc

@@ -19,6 +19,9 @@ limitations under the License.
 
 #include "absl/strings/ascii.h"
 #include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_split.h"
+#include "absl/strings/strip.h"
 #include "tensorflow/core/platform/regexp.h"
 
 namespace tensorflow {

tensorflow/core/profiler/utils/tf_op_utils.h

@@ -16,13 +16,7 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_PROFILER_UTILS_TF_OP_UTILS_H_
 #define TENSORFLOW_CORE_PROFILER_UTILS_TF_OP_UTILS_H_
 
-#include <utility>
-
-#include "absl/strings/match.h"
-#include "absl/strings/str_cat.h"
-#include "absl/strings/str_split.h"
 #include "absl/strings/string_view.h"
-#include "absl/strings/strip.h"
 
 namespace tensorflow {
 namespace profiler {

tensorflow/core/profiler/utils/xplane_schema.cc

@@ -15,15 +15,54 @@ limitations under the License.
 
 #include "tensorflow/core/profiler/utils/xplane_schema.h"
 
+#include "absl/strings/string_view.h"
+
 namespace tensorflow {
 namespace profiler {
 
 const absl::string_view kHostThreads = "Host Threads";
 
-const int kNumStatTypes = static_cast<int>(StatType::kHloModule) + 1;
+constexpr int kNumHostEventTypes =
+    HostEventType::kLastHostEventType - HostEventType::kFirstHostEventType + 1;
 
-static const absl::string_view kStatTypeStrMap[kNumStatTypes] = {
+constexpr int kNumStatTypes =
-    "unknown",         "id",
+    StatType::kLastStatType - StatType::kFirstStatType + 1;
+
+static const absl::string_view kHostEventTypeMetadataMap[] = {
+    "UnknownHostEventType",
+    "TraceContext",
+    "SessionRun",
+    "FunctionRun",
+    "RunGraph",
+    "ExecutorState::Process",
+    "ExecutorDoneCallback",
+    // tf data captured function events.
+    "InstantiatedCapturedFunction::Run",
+    "InstantiatedCapturedFunction::RunWithBorrowedArgs",
+    "InstantiatedCapturedFunction::RunInstantiated",
+    "InstantiatedCapturedFunction::RunAsync",
+    // Functional ops.
+    "CallOp",
+    "ParallelForOp",
+    "ForeverOp",
+    "NumericalGradientOp-EvalRight",
+    "NumericalGradientOp-EvalLeft",
+    "SymbolicGradientOp",
+    "RemoteCallOp",
+    "IfOp",
+    "CaseOp",
+    "WhileOp-EvalCond",
+    "WhileOp-StartBody",
+    "ForOp",
+    "PartitionedCallOp",
+};
+
+static_assert(sizeof(kHostEventTypeMetadataMap) / sizeof(absl::string_view) ==
+                  kNumHostEventTypes,
+              "Mismatch between enum and string map.");
+
+static const absl::string_view kStatTypeStrMap[] = {
+    "UnknownStatType", "id",
     "parent_step_id",  "function_step_id",
     "device_ordinal",  "chip_ordinal",
     "node_ordinal",    "model_id",
@@ -39,6 +78,14 @@ static const absl::string_view kStatTypeStrMap[kNumStatTypes] = {
     "hlo_module",
 };
 
+static_assert(sizeof(kStatTypeStrMap) / sizeof(absl::string_view) ==
+                  kNumStatTypes,
+              "Mismatch between enum and string map.");
+
+absl::Span<const absl::string_view> GetHostEventTypeStrMap() {
+  return absl::MakeConstSpan(kHostEventTypeMetadataMap, kNumHostEventTypes);
+}
+
 absl::Span<const absl::string_view> GetStatTypeStrMap() {
   return absl::MakeConstSpan(kStatTypeStrMap, kNumStatTypes);
 }

tensorflow/core/profiler/utils/xplane_schema.h

@@ -16,6 +16,7 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_PROFILER_UTILS_XPLANE_SCHEMA_H_
 #define TENSORFLOW_CORE_PROFILER_UTILS_XPLANE_SCHEMA_H_
 
+#include "absl/strings/match.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "tensorflow/core/platform/logging.h"
@@ -26,8 +27,41 @@ namespace profiler {
 // Name of XPlane that contains TraceMe events.
 ABSL_CONST_INIT extern const absl::string_view kHostThreads;
 
+// Interesting event types (i.e., TraceMe names).
+enum HostEventType {
+  kFirstHostEventType = 0,
+  kUnknownHostEventType = kFirstHostEventType,
+  kTraceContext,
+  kSessionRun,
+  kFunctionRun,
+  kRunGraph,
+  kExecutorStateProcess,
+  kExecutorDoneCallback,
+  // tf.data captured function events.
+  kTfDataCapturedFunctionRun,
+  kTfDataCapturedFunctionRunWithBorrowedArgs,
+  kTfDataCapturedFunctionRunInstantiated,
+  kTfDataCapturedFunctionRunAsync,
+  // Functional ops.
+  kCallOp,
+  kParallelForOp,
+  kForeverOp,
+  kNumericalGradientOpEvalRight,
+  kNumericalGradientOpEvalLeft,
+  kSymbolicGradientOp,
+  kRemoteCallOp,
+  kIfOp,
+  kCaseOp,
+  kWhileOpEvalCond,
+  kWhileOpStartBody,
+  kForOp,
+  kPartitionedCallOp,
+  kLastHostEventType = kPartitionedCallOp,
+};
+
 enum StatType {
-  kUnknown = 0,
+  kFirstStatType = 0,
+  kUnknownStatType = kFirstStatType,
   // TraceMe arguments.
   kStepId,
   kParentStepId,
@@ -56,14 +90,23 @@ enum StatType {
   kTfOp,
   kHloOp,
   kHloModule,
+  kLastStatType = kHloModule,
 };
 
-ABSL_CONST_INIT extern const int kNumStatTypes;
+absl::Span<const absl::string_view> GetHostEventTypeStrMap();
+
+inline absl::string_view GetHostEventTypeStr(HostEventType event_type) {
+  return GetHostEventTypeStrMap()[event_type];
+}
+
+inline bool IsHostEventType(HostEventType event_type,
+                            absl::string_view event_name) {
+  return GetHostEventTypeStrMap()[event_type] == event_name;
+}
 
 absl::Span<const absl::string_view> GetStatTypeStrMap();
 
 inline absl::string_view GetStatTypeStr(StatType stat_type) {
-  DCHECK_LT(stat_type, kNumStatTypes);
   return GetStatTypeStrMap()[stat_type];
 }