diff --git a/tensorflow/core/profiler/utils/xplane_schema.cc b/tensorflow/core/profiler/utils/xplane_schema.cc index 78fd7af584c..858dd7a99ba 100644 --- a/tensorflow/core/profiler/utils/xplane_schema.cc +++ b/tensorflow/core/profiler/utils/xplane_schema.cc @@ -112,6 +112,10 @@ const HostEventTypeMap& GetHostEventTypeMap() { // Batching related. {"BatchingSessionRun", kBatchingSessionRun}, {"ProcessBatch", kProcessBatch}, + {"ConcatInputTensors", kConcatInputTensors}, + {"MergeInputTensors", kMergeInputTensors}, + {"ScheduleWithoutSplit", kScheduleWithoutSplit}, + {"ScheduleWithSplit", kScheduleWithSplit}, // JAX related. {"LocalExecutable::ExecuteOnLocalDevices", kExecuteOnLocalDevices}, // GPU related. @@ -202,6 +206,10 @@ const StatTypeMap& GetStatTypeMap() { {"memory_size", kDevCapMemorySize}, {"compute_cap_major", kDevCapComputeCapMajor}, {"compute_cap_minor", kDevCapComputeCapMinor}, + // Batching related. + {"batch_size_after_padding", kBatchSizeAfterPadding}, + {"padding_amount", kPaddingAmount}, + {"batching_input_task_size", kBatchingInputTaskSize}, }); DCHECK_EQ(stat_type_map->size(), kNumStatTypes); return *stat_type_map; diff --git a/tensorflow/core/profiler/utils/xplane_schema.h b/tensorflow/core/profiler/utils/xplane_schema.h index 8a9ac4fd278..dd8b4fe5140 100644 --- a/tensorflow/core/profiler/utils/xplane_schema.h +++ b/tensorflow/core/profiler/utils/xplane_schema.h @@ -103,6 +103,10 @@ enum HostEventType { // Batching related. kBatchingSessionRun, kProcessBatch, + kConcatInputTensors, + kMergeInputTensors, + kScheduleWithoutSplit, + kScheduleWithSplit, // JAX related. kExecuteOnLocalDevices, // GPU related. @@ -191,7 +195,11 @@ enum StatType { kDevCapMemorySize, kDevCapComputeCapMajor, kDevCapComputeCapMinor, - kLastStatType = kDevCapComputeCapMinor, + // Batching related. + kBatchSizeAfterPadding, + kPaddingAmount, + kBatchingInputTaskSize, + kLastStatType = kBatchingInputTaskSize, }; inline std::string GpuPlaneName(int32 device_ordinal) {