From 708e5729bb02a1ff17c71ee6dc2ad531a6b18d50 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 16 Dec 2019 12:13:43 -0800 Subject: [PATCH] Add HostEventType to xplane_schema. PiperOrigin-RevId: 285822505 Change-Id: Ieb1d4af1df36016e6ecb919c32334c84d2327453 --- tensorflow/core/profiler/utils/tf_op_utils.cc | 3 ++ tensorflow/core/profiler/utils/tf_op_utils.h | 6 --- .../core/profiler/utils/xplane_schema.cc | 53 +++++++++++++++++-- .../core/profiler/utils/xplane_schema.h | 49 +++++++++++++++-- 4 files changed, 99 insertions(+), 12 deletions(-) diff --git a/tensorflow/core/profiler/utils/tf_op_utils.cc b/tensorflow/core/profiler/utils/tf_op_utils.cc index ae9e8bd9c16..8cf39b43b53 100644 --- a/tensorflow/core/profiler/utils/tf_op_utils.cc +++ b/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 { diff --git a/tensorflow/core/profiler/utils/tf_op_utils.h b/tensorflow/core/profiler/utils/tf_op_utils.h index 5e459ab214b..588dcd4a4a3 100644 --- a/tensorflow/core/profiler/utils/tf_op_utils.h +++ b/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 - -#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 { diff --git a/tensorflow/core/profiler/utils/xplane_schema.cc b/tensorflow/core/profiler/utils/xplane_schema.cc index b05b851ef60..095ef2e1f0e 100644 --- a/tensorflow/core/profiler/utils/xplane_schema.cc +++ b/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(StatType::kHloModule) + 1; +constexpr int kNumHostEventTypes = + HostEventType::kLastHostEventType - HostEventType::kFirstHostEventType + 1; -static const absl::string_view kStatTypeStrMap[kNumStatTypes] = { - "unknown", "id", +constexpr int kNumStatTypes = + 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 GetHostEventTypeStrMap() { + return absl::MakeConstSpan(kHostEventTypeMetadataMap, kNumHostEventTypes); +} + absl::Span GetStatTypeStrMap() { return absl::MakeConstSpan(kStatTypeStrMap, kNumStatTypes); } diff --git a/tensorflow/core/profiler/utils/xplane_schema.h b/tensorflow/core/profiler/utils/xplane_schema.h index 27977858604..0df9e2b21eb 100644 --- a/tensorflow/core/profiler/utils/xplane_schema.h +++ b/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 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 GetStatTypeStrMap(); inline absl::string_view GetStatTypeStr(StatType stat_type) { - DCHECK_LT(stat_type, kNumStatTypes); return GetStatTypeStrMap()[stat_type]; }