diff --git a/plugin/xprof/protobuf/BUILD b/plugin/xprof/protobuf/BUILD index 7396c29ff..f72498e45 100644 --- a/plugin/xprof/protobuf/BUILD +++ b/plugin/xprof/protobuf/BUILD @@ -159,7 +159,10 @@ xprof_proto_library( xprof_proto_library( name = "steps_db_proto", srcs = ["steps_db.proto"], - protodeps = [":op_metrics_proto"], + protodeps = [ + ":op_metrics_proto", + ":flat_op_metrics_proto", + ], ) xprof_proto_library( diff --git a/plugin/xprof/protobuf/flat_op_metrics.proto b/plugin/xprof/protobuf/flat_op_metrics.proto index 0a633ac85..1bedcdd9a 100644 --- a/plugin/xprof/protobuf/flat_op_metrics.proto +++ b/plugin/xprof/protobuf/flat_op_metrics.proto @@ -160,6 +160,9 @@ message FlatOpMetricsDb { // Weighted average total op time if running on default pstate. uint64 normalized_total_op_time_ps = 7; + + int64 total_host_infeed_enq_duration_ps = 8; + int64 total_host_infeed_enq_start_timestamp_ps_diff = 9; } message FlatOpMetricsDbMap { diff --git a/plugin/xprof/protobuf/op_stats.proto b/plugin/xprof/protobuf/op_stats.proto index 804792c8a..b633a19e6 100644 --- a/plugin/xprof/protobuf/op_stats.proto +++ b/plugin/xprof/protobuf/op_stats.proto @@ -13,7 +13,7 @@ import "plugin/xprof/protobuf/steps_db.proto"; import "plugin/xprof/protobuf/tf_function.proto"; import "plugin/xprof/protobuf/topology.proto"; -// Performance environment, e.g the peak performance capabilities of the device. +// Performance environment, e.g. the peak performance capabilities of the device. message PerfEnv { // Peak performance of a TPU core or a GPU in TFLOP/s. double peak_tera_flops_per_second = 1; @@ -189,8 +189,13 @@ message OpStats { SourceStats source_stats = 14; // The result for the disaggregated serving latency. DisaggregatedServingLatency disaggregated_serving_latency = 15; - // The new path for device op metrics. device_op_metrics_db will be deprecated - // soon. + // The new path for device op metrics. TODO(b/529857396): device_op_metrics_db + // will be deprecated. tensorflow.profiler.FlatOpMetricsDb flat_device_op_metrics_db = 16; + // The new path for hlo metrics. TODO(b/529857396): + // hlo_metrics_db_complete_steps_only will be deprecated. + tensorflow.profiler.FlatOpMetricsDb flat_hlo_metrics_db_complete_steps_only = + 17; + optional tensorflow.profiler.FlatOpMetricsDb flat_host_op_metrics_db = 18; reserved 7; } diff --git a/plugin/xprof/protobuf/steps_db.proto b/plugin/xprof/protobuf/steps_db.proto index 00da3c6b7..1732018ae 100644 --- a/plugin/xprof/protobuf/steps_db.proto +++ b/plugin/xprof/protobuf/steps_db.proto @@ -4,6 +4,7 @@ package tensorflow.profiler; import "google/protobuf/any.proto"; import "plugin/xprof/protobuf/op_metrics.proto"; +import "plugin/xprof/protobuf/flat_op_metrics.proto"; // Breakdown of step-time on generic hardware. Note that these components are // mutually exclusive so that adding them together is equal to the step time. If @@ -128,7 +129,7 @@ message StepInfoResult { uint64 begin_ps = 3; // Breakdown of the step-time. Can be unpacked into a GenericStepBreakdown. google.protobuf.Any step_breakdown = 4; - // Total time/bytes/occurences for collectives. (All-Reduce, All-to-All etc) + // Total time/bytes/occurrences for collectives. (All-Reduce, All-to-All etc) DeviceMemoryTransfer collectives = 6; } @@ -168,13 +169,16 @@ message PerCoreStepInfo { map core_id_to_replica_id_map = 5; // A map from core_id to all-reduce ops. map all_reduce_db_per_core = 6; - // Information about deivce memory transfers, categoried by source and + // Information about device memory transfers, categoried by source and // destination. Ordered by following categories: // 1. HostToDevice // 2. DeviceToHost // 3. DeviceToDevice // Cores are normally sharing host interfaces (i.e. PCIe). repeated DeviceMemoryTransfer device_memory_transfers = 7; + // The new path for per-step HLO metrics. hlo_metrics_db will be deprecated + // soon. + tensorflow.profiler.FlatOpMetricsDb flat_hlo_metrics_db = 8; reserved 4; } diff --git a/xprof/convert/BUILD b/xprof/convert/BUILD index dd1fea17d..f12498cd0 100644 --- a/xprof/convert/BUILD +++ b/xprof/convert/BUILD @@ -942,6 +942,35 @@ cc_library( ], ) +cc_library( + name = "flat_op_metrics_to_record", + srcs = ["flat_op_metrics_to_record.cc"], + hdrs = ["flat_op_metrics_to_record.h"], + deps = [ + ":op_metrics_to_record", + "@com_google_absl//absl/algorithm:container", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/strings", + "@org_xprof//plugin/xprof/protobuf:flat_op_metrics_proto_cc", + "@org_xprof//plugin/xprof/protobuf:hardware_types_proto_cc", + "@org_xprof//plugin/xprof/protobuf:op_metrics_proto_cc", + "@org_xprof//plugin/xprof/protobuf:op_stats_proto_cc", + "@xla//xla/tsl/profiler/utils:math_utils", + ], +) + +cc_test( + name = "flat_op_metrics_to_record_test", + srcs = ["flat_op_metrics_to_record_test.cc"], + deps = [ + ":flat_op_metrics_to_record", + "//net/proto2/contrib/parse_proto:parse_text_proto", + "@com_google_googletest//:gtest_main", + "@org_xprof//plugin/xprof/protobuf:flat_op_metrics_proto_cc", + "@org_xprof//plugin/xprof/protobuf:roofline_model_proto_cc", + ], +) + cc_library( name = "op_stack", hdrs = ["op_stack.h"], @@ -995,11 +1024,15 @@ cc_library( hdrs = ["op_stats_to_roofline_model.h"], deps = [ ":data_table_utils", + ":flat_op_metrics_db_combiner", + ":flat_op_metrics_to_record", ":op_metrics_db_combiner", ":op_metrics_to_record", ":source_info_utils", + "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", "@com_google_absl//absl/strings", + "@org_xprof//plugin/xprof/protobuf:flat_op_metrics_proto_cc", "@org_xprof//plugin/xprof/protobuf:hardware_types_proto_cc", "@org_xprof//plugin/xprof/protobuf:op_metrics_proto_cc", "@org_xprof//plugin/xprof/protobuf:op_stats_proto_cc", @@ -1272,6 +1305,7 @@ cc_library( srcs = ["step_events_to_steps_db.cc"], hdrs = ["step_events_to_steps_db.h"], deps = [ + ":flat_op_metrics_db_combiner", ":op_metrics_db_combiner", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/base:core_headers", @@ -1279,6 +1313,7 @@ cc_library( "@com_google_absl//absl/log", "@org_xprof//plugin/xprof/protobuf:steps_db_proto_cc", "@org_xprof//xprof/utils:event_span", + "@org_xprof//xprof/utils:flat_op_metrics_db_utils", "@org_xprof//xprof/utils:op_metrics_db_utils", "@xla//xla/tsl/lib/gtl:map_util", "@xla//xla/tsl/platform:logging", @@ -1331,6 +1366,7 @@ cc_library( "@com_google_absl//absl/container:flat_hash_set", "@com_google_absl//absl/log", "@com_google_absl//absl/log:check", + "@com_google_absl//absl/log:vlog_is_on", "@com_google_absl//absl/status", "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", @@ -1440,9 +1476,9 @@ cc_library( "@org_xprof//plugin/xprof/protobuf:op_metrics_proto_cc", "@org_xprof//plugin/xprof/protobuf:steps_db_proto_cc", "@org_xprof//xprof/utils:event_span", + "@org_xprof//xprof/utils:flat_op_metrics_db_utils", "@org_xprof//xprof/utils:op_metrics_db_utils", "@tsl//tsl/profiler/protobuf:xplane_proto_cc", - "@xla//xla/tsl/platform:types", "@xla//xla/tsl/profiler/utils:tf_op_utils", "@xla//xla/tsl/profiler/utils:tf_xplane_visitor", "@xla//xla/tsl/profiler/utils:timespan", @@ -1641,6 +1677,7 @@ cc_test( "@com_google_absl//absl/container:flat_hash_map", "@com_google_googletest//:gtest", "@com_google_googletest//:gtest_main", + "@org_xprof//plugin/xprof/protobuf:flat_op_metrics_proto_cc", "@org_xprof//plugin/xprof/protobuf:hardware_types_proto_cc", "@org_xprof//plugin/xprof/protobuf:op_stats_proto_cc", "@org_xprof//plugin/xprof/protobuf:power_metrics_proto_cc", @@ -3035,6 +3072,7 @@ cc_library( "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log", "@com_google_absl//absl/strings", + "@com_google_absl//absl/strings:string_view", "@org_xprof//plugin/xprof/protobuf:flat_op_metrics_proto_cc", "@org_xprof//xprof/utils:flat_op_metrics_db_utils", "@org_xprof//xprof/utils:gpu_event_stats", diff --git a/xprof/convert/flat_op_metrics_db_combiner.h b/xprof/convert/flat_op_metrics_db_combiner.h index be288f590..bfc9466f6 100644 --- a/xprof/convert/flat_op_metrics_db_combiner.h +++ b/xprof/convert/flat_op_metrics_db_combiner.h @@ -35,6 +35,11 @@ class FlatOpMetricsDbCombiner : public FlatOpMetricsDbBuilder { // occurs. void Combine(const FlatOpMetricsDb& src, bool update_num_cores = true); + // Combines the memory access breakdown for FlatOpMetrics. + static void CombineMemoryAccessedBreakdown( + const tsl::protobuf::RepeatedPtrField& src, + tsl::protobuf::RepeatedPtrField* dst); + private: // Copies FlatOpMetrics metadata (e.g., category, provenance) from src to dst. static void CopyFlatOpMetricsMetadata(const FlatOpMetrics& src, @@ -43,11 +48,6 @@ class FlatOpMetricsDbCombiner : public FlatOpMetricsDbBuilder { // Combines FlatOpMetrics data from src into dst. static void CombineFlatOpMetrics(const FlatOpMetrics& src, FlatOpMetrics* dst, bool update_num_cores); - - // Combines the memory access breakdown for FlatOpMetrics. - static void CombineMemoryAccessedBreakdown( - const tsl::protobuf::RepeatedPtrField& src, - tsl::protobuf::RepeatedPtrField* dst); }; } // namespace profiler diff --git a/xprof/convert/flat_op_metrics_to_record.cc b/xprof/convert/flat_op_metrics_to_record.cc new file mode 100644 index 000000000..473a9f065 --- /dev/null +++ b/xprof/convert/flat_op_metrics_to_record.cc @@ -0,0 +1,58 @@ +/* Copyright 2026 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xprof/convert/flat_op_metrics_to_record.h" + +#include +#include + +#include "absl/algorithm/container.h" + +namespace tensorflow { +namespace profiler { + + +std::vector SortedOpMetricsDb( + const FlatOpMetricsDb& metrics_db, int max_records) { + std::vector result; + // Exclude SparseCore ops to maintain parity with legacy OpMetricsDb behavior. + // SparseCore operations are typically aggregated as children of TensorCore + // operations and are not listed as independent top-level records in this + // view. + result.reserve(absl::c_count_if( + metrics_db.op_instances(), [](const FlatOpMetrics& metrics) { + return metrics.core_type() != FlatOpMetrics::SPARSE_CORE; + })); + for (const FlatOpMetrics& metrics : metrics_db.op_instances()) { + if (metrics.core_type() == FlatOpMetrics::SPARSE_CORE) continue; + result.push_back(&metrics); + } + + auto comp = [](const FlatOpMetrics* a, const FlatOpMetrics* b) { + return std::make_tuple(a->self_time_ps(), a->hlo_name()) > + std::make_tuple(b->self_time_ps(), b->hlo_name()); + }; + int result_size = result.size(); + if (max_records > 0 && result_size > max_records) { + absl::c_partial_sort(result, result.begin() + max_records, comp); + result.resize(max_records); + } else { + absl::c_sort(result, comp); + } + return result; +} + +} // namespace profiler +} // namespace tensorflow diff --git a/xprof/convert/flat_op_metrics_to_record.h b/xprof/convert/flat_op_metrics_to_record.h new file mode 100644 index 000000000..474308236 --- /dev/null +++ b/xprof/convert/flat_op_metrics_to_record.h @@ -0,0 +1,308 @@ +/* Copyright 2026 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef XPROF_CONVERT_FLAT_OP_METRICS_TO_RECORD_H_ +#define XPROF_CONVERT_FLAT_OP_METRICS_TO_RECORD_H_ + +#include +#include +#include +#include +#include + +#include "absl/algorithm/container.h" +#include "absl/log/check.h" +#include "absl/strings/string_view.h" +#include "xla/tsl/profiler/utils/math_utils.h" +#include "xprof/convert/op_metrics_to_record.h" +#include "plugin/xprof/protobuf/flat_op_metrics.pb.h" +#include "plugin/xprof/protobuf/hardware_types.pb.h" +#include "plugin/xprof/protobuf/op_metrics.pb.h" +#include "plugin/xprof/protobuf/op_stats.pb.h" + +namespace tensorflow { +namespace profiler { + +inline const auto& GetMetricsList(const FlatOpMetricsDb& db) { + return db.op_instances(); +} + +inline absl::string_view GetMetricsName(const FlatOpMetrics& metrics) { + return metrics.hlo_name(); +} + +inline double GigaFlopsPerSecondPerCore(const FlatOpMetrics& metrics) { + return tsl::profiler::SafeDivide( + metrics.flops_v2(), tsl::profiler::PicoToNano(metrics.time_ps())); +} + +inline double GigaFlopsPerSecondPerCoreNormalizedOnDvfs( + const FlatOpMetrics& metrics) { + if (metrics.normalized_time_ps() == 0) { + return GigaFlopsPerSecondPerCore(metrics); + } + return GigaFlopsPerSecondPerCore(metrics) * + (tsl::profiler::SafeDivide(metrics.normalized_time_ps(), + metrics.time_ps())); +} + +inline double GigaModelFlopsPerSecondPerCore(const FlatOpMetrics& metrics) { + return tsl::profiler::SafeDivide( + metrics.model_flops_v2(), tsl::profiler::PicoToNano(metrics.time_ps())); +} + +inline double BytesAccessedPerCore( + const FlatOpMetrics& metrics, uint64_t memory_space, + OpMetrics::MemoryAccessed::OperationType operation_type) { + uint64_t bytes = 0; + if (memory_space == MemorySpace::MEMORY_SPACE_ALL) { + bytes = metrics.bytes_accessed(); + } else { + for (const auto& breakdown : metrics.memory_accessed_breakdown()) { + if ((breakdown.operation_type() != operation_type) && + (operation_type != OpMetrics::MemoryAccessed::UNKNOWN)) { + continue; + } + if (((memory_space == MemorySpace::MEMORY_SPACE_HBM) && + (breakdown.memory_space() == MemorySpace::MEMORY_SPACE_HBM)) || + ((memory_space == MemorySpace::MEMORY_SPACE_ON_CHIP) && + (breakdown.memory_space() != MemorySpace::MEMORY_SPACE_HBM))) { + bytes += breakdown.bytes_accessed(); + } + } + } + return bytes; +} + +inline double GigaBytesPerSecondPerCore( + const FlatOpMetrics& metrics, uint64_t memory_space, + OpMetrics::MemoryAccessed::OperationType operation_type) { + return tsl::profiler::SafeDivide( + BytesAccessedPerCore(metrics, memory_space, operation_type), + tsl::profiler::PicoToNano(metrics.time_ps())); +} + +inline double GibiBytesPerSecondPerCore( + const FlatOpMetrics& metrics, uint64_t memory_space, + OpMetrics::MemoryAccessed::OperationType op_type) { + return tsl::profiler::GigaToGibi( + GigaBytesPerSecondPerCore(metrics, memory_space, op_type)); +} + +template +inline void SetExecutionTimes(const FlatOpMetrics& metrics, Record* record) { + record->set_occurrences(metrics.occurrences()); + record->set_total_time_in_us(tsl::profiler::PicoToMicro(metrics.time_ps())); + record->set_avg_time_in_us(tsl::profiler::SafeDivide( + record->total_time_in_us(), metrics.occurrences())); + record->set_total_self_time_in_us( + tsl::profiler::PicoToMicro(metrics.self_time_ps())); + record->set_avg_self_time_in_us(tsl::profiler::SafeDivide( + record->total_self_time_in_us(), metrics.occurrences())); +} + +template +inline void SetTpuUnitFractions(const FlatOpMetrics& metrics, Record* record) { + record->set_dma_stall_fraction( + tsl::profiler::SafeDivide(metrics.dma_stall_ps(), metrics.time_ps())); +} + +// Returns a sorted vector of pointers to FlatOpMetrics in the given database. +// The returned pointers are only valid as long as `metrics_db` exists and is +// not modified. +std::vector SortedOpMetricsDb( + const FlatOpMetricsDb& metrics_db, int max_records = -1); + +template +inline void SetRooflineMetrics(const FlatOpMetrics& metrics, + const PerfEnv& perf_env, + const RunEnvironment& run_env, Record* record, + bool apply_time_scale_factor = false) { + using ::tensorflow::profiler::MemorySpace; + using ::tensorflow::profiler::PerformanceInfo; + + record->set_measured_flop_rate(GigaFlopsPerSecondPerCore(metrics)); + record->set_model_flop_rate(GigaModelFlopsPerSecondPerCore(metrics)); + record->set_measured_memory_bw(GibiBytesPerSecondPerCore( + metrics, tensorflow::profiler::MemorySpace::MEMORY_SPACE_ALL, + OpMetrics::MemoryAccessed::UNKNOWN)); + record->set_flops(metrics.flops()); + record->set_flops_v2(metrics.flops_v2()); + record->set_bytes_accessed(metrics.bytes_accessed()); + record->set_operational_intensity( + tsl::profiler::SafeDivide(metrics.flops_v2(), metrics.bytes_accessed())); + + uint64_t hbm_bytes = 0; + uint64_t cmem_read_bytes = 0; + uint64_t cmem_write_bytes = 0; + uint64_t vmem_read_bytes = 0; + uint64_t vmem_write_bytes = 0; + for (const auto& memory_access : metrics.memory_accessed_breakdown()) { + if (memory_access.memory_space() == PerformanceInfo::MemoryAccessed::HBM) { + hbm_bytes += memory_access.bytes_accessed(); + } else if (memory_access.memory_space() == + PerformanceInfo::MemoryAccessed::CMEM) { + if (memory_access.operation_type() == OpMetrics::MemoryAccessed::READ) { + cmem_read_bytes += memory_access.bytes_accessed(); + } else if (memory_access.operation_type() == + OpMetrics::MemoryAccessed::WRITE) { + cmem_write_bytes += memory_access.bytes_accessed(); + } + } else if (memory_access.memory_space() == + PerformanceInfo::MemoryAccessed::VMEM) { + if (memory_access.operation_type() == OpMetrics::MemoryAccessed::READ) { + vmem_read_bytes += memory_access.bytes_accessed(); + } else if (memory_access.operation_type() == + OpMetrics::MemoryAccessed::WRITE) { + vmem_write_bytes += memory_access.bytes_accessed(); + } + } + } + if (metrics.memory_accessed_breakdown_size() == 0) { + hbm_bytes = metrics.bytes_accessed(); + } + int64_t device_time_ps = apply_time_scale_factor + ? metrics.normalized_time_ps() + : metrics.time_ps(); + record->set_hbm_bw(tsl::profiler::GibibytesPerSecond( + hbm_bytes, tsl::profiler::PicoToNano(metrics.time_ps()))); + record->set_cmem_read_bw(tsl::profiler::GibibytesPerSecond( + cmem_read_bytes, tsl::profiler::PicoToNano(device_time_ps))); + record->set_cmem_write_bw(tsl::profiler::GibibytesPerSecond( + cmem_write_bytes, tsl::profiler::PicoToNano(device_time_ps))); + record->set_vmem_read_bw(tsl::profiler::GibibytesPerSecond( + vmem_read_bytes, tsl::profiler::PicoToNano(device_time_ps))); + record->set_vmem_write_bw(tsl::profiler::GibibytesPerSecond( + vmem_write_bytes, tsl::profiler::PicoToNano(device_time_ps))); + record->set_hbm_operational_intensity( + tsl::profiler::SafeDivide(metrics.flops_v2(), hbm_bytes)); + record->set_cmem_read_operational_intensity( + tsl::profiler::SafeDivide(metrics.flops_v2(), cmem_read_bytes)); + record->set_cmem_write_operational_intensity( + tsl::profiler::SafeDivide(metrics.flops_v2(), cmem_write_bytes)); + record->set_vmem_read_operational_intensity( + tsl::profiler::SafeDivide(metrics.flops_v2(), vmem_read_bytes)); + record->set_vmem_write_operational_intensity( + tsl::profiler::SafeDivide(metrics.flops_v2(), vmem_write_bytes)); + + constexpr absl::string_view kUnknown = "Unknown"; + constexpr absl::string_view kCompute = "Compute"; + constexpr absl::string_view kHbm = "HBM"; + constexpr absl::string_view kCmemRead = "CMEM Read"; + constexpr absl::string_view kCmemWrite = "CMEM Write"; + constexpr absl::string_view kVmemRead = "VMEM Read"; + constexpr absl::string_view kVmemWrite = "VMEM Write"; + constexpr absl::string_view kShmL1 = "Shm/L1"; + + absl::string_view bottleneck_resource = kUnknown; + double bottleneck_utilization = 0; + double bottleneck_operational_intensity = 0; + double peak_flops = + tsl::profiler::TeraToGiga(perf_env.peak_tera_flops_per_second()); + double flops_utilization = + tsl::profiler::SafeDivide(record->measured_flop_rate(), peak_flops); + if (bottleneck_utilization < flops_utilization) { + bottleneck_resource = kCompute; + bottleneck_utilization = flops_utilization; + bottleneck_operational_intensity = record->operational_intensity(); + } + double peak_hbm_bw = GetMemoryPeakBandwidth(perf_env, 0); + double hbm_bw_utilization = tsl::profiler::SafeDivide( + record->hbm_bw(), tsl::profiler::GigaToGibi(peak_hbm_bw)); + if (bottleneck_utilization < hbm_bw_utilization) { + bottleneck_resource = kHbm; + bottleneck_utilization = hbm_bw_utilization; + bottleneck_operational_intensity = record->hbm_operational_intensity(); + } + tensorflow::profiler::HardwareType hardware_type = run_env.hardware_type(); + if (hardware_type == tensorflow::profiler::HardwareType::TPU) { + if (cmem_read_bytes) { + double peak_cmem_read_bw = GetMemoryPeakBandwidth(perf_env, 3); + if (peak_cmem_read_bw) { + double cmem_read_bw_utilization = tsl::profiler::SafeDivide( + record->cmem_read_bw(), + tsl::profiler::GigaToGibi(peak_cmem_read_bw)); + if (bottleneck_utilization < cmem_read_bw_utilization) { + bottleneck_resource = kCmemRead; + bottleneck_utilization = cmem_read_bw_utilization; + bottleneck_operational_intensity = + record->cmem_read_operational_intensity(); + } + } + } + if (cmem_write_bytes) { + double peak_cmem_write_bw = GetMemoryPeakBandwidth(perf_env, 4); + if (peak_cmem_write_bw) { + double cmem_write_bw_utilization = tsl::profiler::SafeDivide( + record->cmem_write_bw(), + tsl::profiler::GigaToGibi(peak_cmem_write_bw)); + if (bottleneck_utilization < cmem_write_bw_utilization) { + bottleneck_resource = kCmemWrite; + bottleneck_utilization = cmem_write_bw_utilization; + bottleneck_operational_intensity = + record->cmem_write_operational_intensity(); + } + } + } + if (vmem_read_bytes) { + double peak_vmem_read_bw = GetMemoryPeakBandwidth(perf_env, 5); + if (peak_vmem_read_bw) { + double vmem_read_bw_utilization = tsl::profiler::SafeDivide( + record->vmem_read_bw(), + tsl::profiler::GigaToGibi(peak_vmem_read_bw)); + if (bottleneck_utilization < vmem_read_bw_utilization) { + bottleneck_resource = kVmemRead; + bottleneck_utilization = vmem_read_bw_utilization; + bottleneck_operational_intensity = + record->vmem_read_operational_intensity(); + } + } + } + if (vmem_write_bytes) { + double peak_vmem_write_bw = GetMemoryPeakBandwidth(perf_env, 6); + if (peak_vmem_write_bw) { + double vmem_write_bw_utilization = tsl::profiler::SafeDivide( + record->vmem_write_bw(), + tsl::profiler::GigaToGibi(peak_vmem_write_bw)); + if (bottleneck_utilization < vmem_write_bw_utilization) { + bottleneck_resource = kVmemWrite; + bottleneck_utilization = vmem_write_bw_utilization; + bottleneck_operational_intensity = + record->vmem_write_operational_intensity(); + } + } + } + } + if (hardware_type == tensorflow::profiler::HardwareType::GPU) { + double peak_shm_l1_bw = GetMemoryPeakBandwidth(perf_env, 2); + if (peak_shm_l1_bw) { + double shm_l1_bw_utilization = tsl::profiler::SafeDivide( + record->hbm_bw(), tsl::profiler::GigaToGibi(peak_shm_l1_bw)); + if (bottleneck_utilization < shm_l1_bw_utilization) { + bottleneck_resource = kShmL1; + bottleneck_utilization = shm_l1_bw_utilization; + bottleneck_operational_intensity = record->hbm_operational_intensity(); + } + } + } + record->set_bound_by(std::string(bottleneck_resource)); + record->set_bottleneck_operational_intensity( + bottleneck_operational_intensity); +} + +} // namespace profiler +} // namespace tensorflow + +#endif // XPROF_CONVERT_FLAT_OP_METRICS_TO_RECORD_H_ diff --git a/xprof/convert/flat_op_metrics_to_record_test.cc b/xprof/convert/flat_op_metrics_to_record_test.cc new file mode 100644 index 000000000..2441f2764 --- /dev/null +++ b/xprof/convert/flat_op_metrics_to_record_test.cc @@ -0,0 +1,116 @@ +/* Copyright 2026 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "xprof/convert/flat_op_metrics_to_record.h" + +#include + +#include "net/proto2/contrib/parse_proto/parse_text_proto.h" +#include "" +#include "plugin/xprof/protobuf/flat_op_metrics.pb.h" +#include "plugin/xprof/protobuf/roofline_model.pb.h" + +namespace tensorflow { +namespace profiler { +namespace { + +using ::google::protobuf::contrib::parse_proto::ParseTextProtoOrDie; +using ::tensorflow::profiler::roofline_model::RooflineModelRecord; + +TEST(FlatOpMetricsToRecordTest, SortedOpMetricsDb) { + FlatOpMetricsDb db = ParseTextProtoOrDie(R"pb( + op_instances { hlo_name: "op1" self_time_ps: 100 core_type: TENSOR_CORE } + op_instances { hlo_name: "op2" self_time_ps: 200 core_type: TENSOR_CORE } + op_instances { hlo_name: "op3" self_time_ps: 150 core_type: TENSOR_CORE } + op_instances { hlo_name: "op4" self_time_ps: 150 core_type: TENSOR_CORE } + op_instances { hlo_name: "op5" self_time_ps: 300 core_type: SPARSE_CORE } + )pb"); + + // Default max_records (-1) + { + std::vector sorted = SortedOpMetricsDb(db); + // Sparse core should be excluded. + ASSERT_EQ(sorted.size(), 4); + // Ordered by self_time_ps descending, then hlo_name descending. + EXPECT_EQ(sorted[0]->hlo_name(), "op2"); // 200 + EXPECT_EQ(sorted[1]->hlo_name(), "op4"); // 150 (op4 > op3) + EXPECT_EQ(sorted[2]->hlo_name(), "op3"); // 150 + EXPECT_EQ(sorted[3]->hlo_name(), "op1"); // 100 + } + + // max_records = 2 + { + std::vector sorted = SortedOpMetricsDb(db, 2); + ASSERT_EQ(sorted.size(), 2); + EXPECT_EQ(sorted[0]->hlo_name(), "op2"); + EXPECT_EQ(sorted[1]->hlo_name(), "op4"); + } +} + +// Minimal PerfEnv and RunEnvironment for testing SetRooflineMetrics. +PerfEnv GetTestPerfEnv() { + return ParseTextProtoOrDie(R"pb( + peak_tera_flops_per_second: 100.0 + peak_bws_giga_bytes_per_second: 1000.0 # Index 0 (HBM) + peak_bws_giga_bytes_per_second: 0.0 + peak_bws_giga_bytes_per_second: 500.0 # Index 2 (Shm/L1 for GPU) + peak_bws_giga_bytes_per_second: 2000.0 # Index 3 (CMEM Read) + peak_bws_giga_bytes_per_second: 2000.0 # Index 4 (CMEM Write) + peak_bws_giga_bytes_per_second: 4000.0 # Index 5 (VMEM Read) + peak_bws_giga_bytes_per_second: 4000.0 # Index 6 (VMEM Write) + )pb"); +} + +RunEnvironment GetTestRunEnv(HardwareType hardware_type) { + RunEnvironment env; + env.set_hardware_type(hardware_type); + return env; +} + +TEST(FlatOpMetricsToRecordTest, SetRooflineMetricsComputeBound) { + FlatOpMetrics metrics = ParseTextProtoOrDie(R"pb( + time_ps: 1000000000000 # 1s + flops: 100000000000000 # 100 TFLOPS + flops_v2: 100000000000000 + bytes_accessed: 1000000000 # 1 GB + )pb"); + + RooflineModelRecord record; + SetRooflineMetrics(metrics, GetTestPerfEnv(), GetTestRunEnv(TPU), &record); + + // Peak flops = 100 TFLOPS. Measured = 100 TFLOPS. Utilization = 1.0 + // Peak HBM BW = 1000 GB/s. Measured = 1 GB/s. Utilization = 0.001 + EXPECT_EQ(record.bound_by(), "Compute"); +} + +TEST(FlatOpMetricsToRecordTest, SetRooflineMetricsHbmBound) { + FlatOpMetrics metrics = ParseTextProtoOrDie(R"pb( + time_ps: 1000000000000 # 1s + flops: 1000000000000 # 1 TFLOPS + flops_v2: 1000000000000 + bytes_accessed: 1000000000000 # 1 TB + )pb"); + + RooflineModelRecord record; + SetRooflineMetrics(metrics, GetTestPerfEnv(), GetTestRunEnv(TPU), &record); + + // Peak flops = 100 TFLOPS. Measured = 1 TFLOPS. Utilization = 0.01 + // Peak HBM BW = 1000 GB/s. Measured = 1000 GB/s. Utilization = 1.0 + EXPECT_EQ(record.bound_by(), "HBM"); +} + +} // namespace +} // namespace profiler +} // namespace tensorflow diff --git a/xprof/convert/op_stats_combiner.cc b/xprof/convert/op_stats_combiner.cc index 3d5457481..8e1eb8c8e 100644 --- a/xprof/convert/op_stats_combiner.cc +++ b/xprof/convert/op_stats_combiner.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include #include +#include #include #include "absl/container/flat_hash_map.h" @@ -44,6 +45,34 @@ namespace profiler { namespace { +// Combines the src PerCoreStepInfo into the dst PerCoreStepInfo. +void CombinePerCoreStepInfo( + int src_host_id, const PerCoreStepInfo& src, bool use_incomplete_step, + PerCoreStepInfo* dst, + FlatOpMetricsDbCombiner* flat_hlo_metrics_db_complete_steps_only_combiner, + FlatOpMetricsDbCombiner* flat_hlo_metrics_db_per_step_combiner) { + CombineCoreIdMap(src_host_id, src.step_info_per_core(), + dst->mutable_step_info_per_core()); + + // Since we have assigned a new step number to the combined result, update + // the step number on each core to this new step number. + uint32_t new_step_num = dst->step_num(); + for (auto& percore_stepinfo : *dst->mutable_step_info_per_core()) { + auto& stepinfo = percore_stepinfo.second; + stepinfo.set_step_num(new_step_num); + } + + if (!use_incomplete_step) { + flat_hlo_metrics_db_complete_steps_only_combiner->Combine( + src.flat_hlo_metrics_db()); + } + flat_hlo_metrics_db_per_step_combiner->Combine(src.flat_hlo_metrics_db()); + CombineCoreIdMap(src_host_id, src.all_reduce_db_per_core(), + dst->mutable_all_reduce_db_per_core()); + CombineCoreIdMap(src_host_id, src.core_id_to_replica_id_map(), + dst->mutable_core_id_to_replica_id_map()); +} + // Combines the src PerCoreStepInfo into the dst PerCoreStepInfo. void CombinePerCoreStepInfo( int src_host_id, const PerCoreStepInfo& src, bool use_incomplete_step, @@ -75,15 +104,27 @@ void CombineStepDatabase( int src_host_id, const StepIntersection& step_intersection, const StepDatabaseResult& src, StepDatabaseResult* dst, OpMetricsDbCombiner* hlo_metrics_db_complete_steps_only_combiner, - std::vector* hlo_metrics_db_per_step_combiners) { + std::vector* hlo_metrics_db_per_step_combiners, + FlatOpMetricsDbCombiner* flat_hlo_metrics_db_complete_steps_only_combiner, + std::vector* + flat_hlo_metrics_db_per_step_combiners, + bool combine_flat_op_metrics_db) { if (src.use_incomplete_step()) dst->set_use_incomplete_step(true); uint32_t src_first_step_idx = step_intersection.FirstStepIndex(src_host_id); for (uint32_t i = 0; i < step_intersection.NumSteps(); i++) { - CombinePerCoreStepInfo( - src_host_id, src.step_sequence(src_first_step_idx + i), - src.use_incomplete_step(), dst->mutable_step_sequence(i), - hlo_metrics_db_complete_steps_only_combiner, - &(*hlo_metrics_db_per_step_combiners)[i]); + if (combine_flat_op_metrics_db) { + CombinePerCoreStepInfo( + src_host_id, src.step_sequence(src_first_step_idx + i), + src.use_incomplete_step(), dst->mutable_step_sequence(i), + flat_hlo_metrics_db_complete_steps_only_combiner, + &(*flat_hlo_metrics_db_per_step_combiners)[i]); + } else { + CombinePerCoreStepInfo( + src_host_id, src.step_sequence(src_first_step_idx + i), + src.use_incomplete_step(), dst->mutable_step_sequence(i), + hlo_metrics_db_complete_steps_only_combiner, + &(*hlo_metrics_db_per_step_combiners)[i]); + } } } @@ -190,20 +231,29 @@ void CombineOpStats( const StepIntersection& step_intersection, const OpStats& src, OpStats* dst, OpMetricsDbCombiner* host_op_metrics_db_combiner, OpMetricsDbCombiner* device_op_metrics_db_combiner, - FlatOpMetricsDbCombiner* flat_op_metrics_db_combiner, + FlatOpMetricsDbCombiner* flat_device_op_metrics_db_combiner, OpMetricsDbCombiner* hlo_metrics_db_complete_steps_only_combiner, - std::vector* hlo_metrics_db_per_step_combiners) { + std::vector* hlo_metrics_db_per_step_combiners, + FlatOpMetricsDbCombiner* flat_hlo_metrics_db_complete_steps_only_combiner, + std::vector* + flat_hlo_metrics_db_per_step_combiners, + bool combine_flat_op_metrics_db) { // Combine host_metrics_db. // Host OpMetricsDb does not need to update the number of cores a certain op // occurs. host_op_metrics_db_combiner->Combine(src.host_op_metrics_db(), /*update_num_cores=*/false); - // Combine device_metrics_db. - device_op_metrics_db_combiner->Combine(src.device_op_metrics_db()); - - // Combine flat_device_op_metrics_db. - if (src.has_flat_device_op_metrics_db()) { - flat_op_metrics_db_combiner->Combine(src.flat_device_op_metrics_db()); + if (combine_flat_op_metrics_db) { + // Combine flat_device_op_metrics_db. + if (src.has_flat_device_op_metrics_db()) { + flat_device_op_metrics_db_combiner->Combine( + src.flat_device_op_metrics_db()); + } + } else { + // Combine device_metrics_db. + if (src.has_device_op_metrics_db()) { + device_op_metrics_db_combiner->Combine(src.device_op_metrics_db()); + } } // Combine step_db. @@ -211,7 +261,10 @@ void CombineOpStats( CombineStepDatabase(src_host_id, step_intersection, src.step_db(), dst->mutable_step_db(), hlo_metrics_db_complete_steps_only_combiner, - hlo_metrics_db_per_step_combiners); + hlo_metrics_db_per_step_combiners, + flat_hlo_metrics_db_complete_steps_only_combiner, + flat_hlo_metrics_db_per_step_combiners, + combine_flat_op_metrics_db); } // Combine run environment info. @@ -316,7 +369,8 @@ StepIntersection ComputeStepIntersectionToMergeOpStats( void CombineAllOpStats(const std::vector& all_op_stats_info, const StepIntersection& step_intersection, - OpStats* combined_op_stats) { + OpStats* combined_op_stats, + bool combine_flat_op_metrics_db) { // A shortcut code path for a single OpStats. There is no need to merge. if (all_op_stats_info.size() == 1) { *combined_op_stats = *all_op_stats_info[0].op_stats; @@ -337,19 +391,49 @@ void CombineAllOpStats(const std::vector& all_op_stats_info, // Initialize all the OpMetricsDbCombiners. OpMetricsDbCombiner host_op_metrics_db_combiner( combined_op_stats->mutable_host_op_metrics_db()); - OpMetricsDbCombiner device_op_metrics_db_combiner( - combined_op_stats->mutable_device_op_metrics_db()); - FlatOpMetricsDbCombiner flat_op_metrics_db_combiner( - combined_op_stats->mutable_flat_device_op_metrics_db()); - OpMetricsDbCombiner hlo_metrics_db_complete_steps_only_combiner( - combined_op_stats->mutable_hlo_metrics_db_complete_steps_only()); + + std::unique_ptr device_op_metrics_db_combiner; + std::unique_ptr flat_device_op_metrics_db_combiner; + std::unique_ptr + hlo_metrics_db_complete_steps_only_combiner; + std::unique_ptr + flat_hlo_metrics_db_complete_steps_only_combiner; + + if (combine_flat_op_metrics_db) { + flat_device_op_metrics_db_combiner = + std::make_unique( + combined_op_stats->mutable_flat_device_op_metrics_db()); + flat_hlo_metrics_db_complete_steps_only_combiner = + std::make_unique( + combined_op_stats + ->mutable_flat_hlo_metrics_db_complete_steps_only()); + } else { + device_op_metrics_db_combiner = std::make_unique( + combined_op_stats->mutable_device_op_metrics_db()); + hlo_metrics_db_complete_steps_only_combiner = + std::make_unique( + combined_op_stats->mutable_hlo_metrics_db_complete_steps_only()); + } + std::vector hlo_metrics_db_per_step_combiners; + std::vector flat_hlo_metrics_db_per_step_combiners; + + if (combine_flat_op_metrics_db) { + flat_hlo_metrics_db_per_step_combiners.reserve( + combined_step_db->step_sequence_size()); + } else { hlo_metrics_db_per_step_combiners.reserve( combined_step_db->step_sequence_size()); + } for (PerCoreStepInfo& step_info : *combined_step_db->mutable_step_sequence()) { + if (combine_flat_op_metrics_db) { + flat_hlo_metrics_db_per_step_combiners.emplace_back( + step_info.mutable_flat_hlo_metrics_db()); + } else { hlo_metrics_db_per_step_combiners.emplace_back( step_info.mutable_hlo_metrics_db()); + } } bool no_accelerator_in_system = NoAcceleratorInSystem(all_op_stats_info); @@ -358,10 +442,14 @@ void CombineAllOpStats(const std::vector& all_op_stats_info, CombineOpStats(no_accelerator_in_system, op_stats_info.src_host_id, op_stats_info.hardware_type, step_intersection, *op_stats_info.op_stats, combined_op_stats, - &host_op_metrics_db_combiner, &device_op_metrics_db_combiner, - &flat_op_metrics_db_combiner, - &hlo_metrics_db_complete_steps_only_combiner, - &hlo_metrics_db_per_step_combiners); + &host_op_metrics_db_combiner, + device_op_metrics_db_combiner.get(), + flat_device_op_metrics_db_combiner.get(), + hlo_metrics_db_complete_steps_only_combiner.get(), + &hlo_metrics_db_per_step_combiners, + flat_hlo_metrics_db_complete_steps_only_combiner.get(), + &flat_hlo_metrics_db_per_step_combiners, + combine_flat_op_metrics_db); } // Sorts all the kernel reports that have been merged by CombineTfOpStats and diff --git a/xprof/convert/op_stats_combiner.h b/xprof/convert/op_stats_combiner.h index 3149df01d..3c447cbb2 100644 --- a/xprof/convert/op_stats_combiner.h +++ b/xprof/convert/op_stats_combiner.h @@ -78,7 +78,8 @@ StepIntersection ComputeStepIntersectionToMergeOpStats( // . The result is stored in . void CombineAllOpStats(const std::vector& all_op_stats_info, const StepIntersection& step_intersection, - OpStats* combined_op_stats); + OpStats* combined_op_stats, + bool combine_flat_op_metrics_db = false); } // namespace profiler } // namespace tensorflow diff --git a/xprof/convert/op_stats_combiner_test.cc b/xprof/convert/op_stats_combiner_test.cc index c02bb2128..38559f8bb 100644 --- a/xprof/convert/op_stats_combiner_test.cc +++ b/xprof/convert/op_stats_combiner_test.cc @@ -22,6 +22,7 @@ limitations under the License. #include "absl/container/flat_hash_map.h" #include "google/protobuf/util/message_differencer.h" #include "xla/tsl/platform/types.h" +#include "plugin/xprof/protobuf/flat_op_metrics.pb.h" #include "plugin/xprof/protobuf/hardware_types.pb.h" #include "plugin/xprof/protobuf/op_stats.pb.h" #include "plugin/xprof/protobuf/power_metrics.pb.h" @@ -234,6 +235,64 @@ TEST(CombineAllOpStatsTest, CombineDisaggregatedServingStats) { 1e-6); } +TEST(CombineAllOpStatsTest, CombineFlatOpMetricsDb) { + OpStats dst_op_stats, src_op_stats_1, src_op_stats_2; + + auto* flat_db_1 = src_op_stats_1.mutable_flat_device_op_metrics_db(); + auto* metric_1 = flat_db_1->add_op_instances(); + metric_1->set_hlo_module_id(1); + metric_1->set_symbol_id(1); + metric_1->set_hlo_name("metric_1"); + metric_1->set_time_ps(100); + + auto* flat_db_2 = src_op_stats_2.mutable_flat_device_op_metrics_db(); + auto* metric_2 = flat_db_2->add_op_instances(); + metric_2->set_hlo_module_id(1); + metric_2->set_symbol_id(2); + metric_2->set_hlo_name("metric_2"); + metric_2->set_time_ps(200); + + auto* step_seq_1 = src_op_stats_1.mutable_step_db()->add_step_sequence(); + step_seq_1->set_step_num(1); + auto* hlo_1 = step_seq_1->mutable_flat_hlo_metrics_db(); + auto* hlo_metric_1 = hlo_1->add_op_instances(); + hlo_metric_1->set_hlo_module_id(2); + hlo_metric_1->set_symbol_id(1); + hlo_metric_1->set_hlo_name("hlo_1"); + hlo_metric_1->set_time_ps(50); + + auto* step_seq_2 = src_op_stats_2.mutable_step_db()->add_step_sequence(); + step_seq_2->set_step_num(1); + auto* hlo_2 = step_seq_2->mutable_flat_hlo_metrics_db(); + auto* hlo_metric_2 = hlo_2->add_op_instances(); + hlo_metric_2->set_hlo_module_id(2); + hlo_metric_2->set_symbol_id(2); + hlo_metric_2->set_hlo_name("hlo_2"); + hlo_metric_2->set_time_ps(150); + + OpStatsInfo op_stats_info_1(&src_op_stats_1, TPU, 0); + OpStatsInfo op_stats_info_2(&src_op_stats_2, TPU, 1); + + std::vector all_op_stats_info = {op_stats_info_1, + op_stats_info_2}; + + absl::flat_hash_map result; + result.insert({0, &src_op_stats_1.step_db()}); + result.insert({1, &src_op_stats_2.step_db()}); + StepIntersection dummy_step_intersection = StepIntersection(1, result); + + CombineAllOpStats(all_op_stats_info, dummy_step_intersection, &dst_op_stats, + /*combine_flat_op_metrics_db=*/true); + + EXPECT_EQ(dst_op_stats.flat_device_op_metrics_db().op_instances_size(), 2); + EXPECT_EQ(dst_op_stats.flat_hlo_metrics_db_complete_steps_only() + .op_instances_size(), + 2); + // Negative invariants: Traditional DBs should be absent/empty. + EXPECT_FALSE(dst_op_stats.has_device_op_metrics_db()); + EXPECT_FALSE(dst_op_stats.has_hlo_metrics_db_complete_steps_only()); +} + } // namespace } // namespace profiler } // namespace tensorflow diff --git a/xprof/convert/op_stats_to_roofline_model.cc b/xprof/convert/op_stats_to_roofline_model.cc index f0fe92915..65ba038e6 100644 --- a/xprof/convert/op_stats_to_roofline_model.cc +++ b/xprof/convert/op_stats_to_roofline_model.cc @@ -28,9 +28,12 @@ limitations under the License. #include "xla/tsl/profiler/utils/math_utils.h" #include "tsl/platform/protobuf.h" #include "xprof/convert/data_table_utils.h" +#include "xprof/convert/flat_op_metrics_db_combiner.h" +#include "xprof/convert/flat_op_metrics_to_record.h" #include "xprof/convert/op_metrics_db_combiner.h" #include "xprof/convert/op_metrics_to_record.h" #include "xprof/convert/source_info_utils.h" +#include "plugin/xprof/protobuf/flat_op_metrics.pb.h" #include "plugin/xprof/protobuf/hardware_types.pb.h" #include "plugin/xprof/protobuf/op_metrics.pb.h" #include "plugin/xprof/protobuf/op_stats.pb.h" @@ -142,6 +145,89 @@ RooflineModelRecord ConvertOpMetricsToRooflineModelRecord( return record; } +RooflineModelRecord ConvertOpMetricsToRooflineModelRecord( + const OpStats& op_stats, const FlatOpMetrics& metrics, + RecordType record_type, uint32_t step_num, uint64_t total_time_ps, + const RooflineModelDatabase& roofline_model_db, bool include_infeed_outfeed, + bool apply_time_scale_multiplier) { + RooflineModelRecord record; + record.set_hlo_name(metrics.hlo_name()); + record.set_hlo_category(metrics.category()); + record.set_hlo_module_id(metrics.hlo_module_id()); + record.set_record_type(record_type); + record.set_step_num(step_num); + *record.mutable_source_info() = metrics.source_info(); + SetExecutionTimes(metrics, &record); + if (record_type == RecordType::AVERAGE_STEP) { + // For RecordType::AVERAGE_STEP, divide by num_steps to show per-step + // numbers when appropriate. + int num_steps = op_stats.step_db().step_sequence_size(); + record.set_total_time_in_us( + tsl::profiler::SafeDivide(record.total_time_in_us(), num_steps)); + record.set_total_self_time_in_us( + tsl::profiler::SafeDivide(record.total_self_time_in_us(), num_steps)); + } + record.set_total_time_per_core_in_us(tsl::profiler::SafeDivide( + record.total_time_in_us(), + op_stats.run_environment().device_core_count())); + record.set_total_time_in_percentage( + tsl::profiler::SafeDivide(metrics.time_ps(), total_time_ps)); + + tensorflow::profiler::SetTpuUnitFractions(metrics, &record); + + // Set the roofline-specific fields. + SetRooflineMetrics(metrics, op_stats.perf_env(), op_stats.run_environment(), + &record, apply_time_scale_multiplier); + const double cmem_wr_utilization = + roofline_model_db.has_cmem() + ? tsl::profiler::SafeDivide(record.cmem_write_bw(), + roofline_model_db.peak_cmem_write_bw()) + : 0; + const double cmem_rd_utilization = + roofline_model_db.has_cmem() + ? tsl::profiler::SafeDivide(record.cmem_read_bw(), + roofline_model_db.peak_cmem_read_bw()) + : 0; + const double vmem_rd_utilization = + roofline_model_db.has_merged_vmem() + ? tsl::profiler::SafeDivide(record.vmem_read_bw(), + roofline_model_db.peak_vmem_read_bw()) + : 0; + const double vmem_wr_utilization = + roofline_model_db.has_merged_vmem() + ? tsl::profiler::SafeDivide(record.vmem_write_bw(), + roofline_model_db.peak_vmem_write_bw()) + : 0; + double flops_utilization = tsl::profiler::SafeDivide( + record.measured_flop_rate(), roofline_model_db.peak_flop_rate()); + if (apply_time_scale_multiplier) { + double measured_flop_rate_normalized = + GigaFlopsPerSecondPerCoreNormalizedOnDvfs(metrics); + flops_utilization = tsl::profiler::SafeDivide( + measured_flop_rate_normalized, roofline_model_db.peak_flop_rate()); + } + const double hbm_utilization = tsl::profiler::SafeDivide( + record.hbm_bw(), roofline_model_db.peak_hbm_bw()); + + const double max_mem_utilization = + std::max({cmem_wr_utilization, cmem_rd_utilization, hbm_utilization, + vmem_wr_utilization, vmem_rd_utilization}); + const double roofline_efficiency = + std::max({max_mem_utilization, flops_utilization}); + // Note, copy-start/done can have utilizations above 1.0 since their + // bytes/time are not accurate as they are asynchronous. + record.set_optimal_flop_rate(tsl::profiler::SafeDivide( + record.measured_flop_rate(), roofline_efficiency)); + record.set_roofline_efficiency(roofline_efficiency); + record.set_flop_rate_relative_to_hw_limit(flops_utilization); + record.set_memory_bw_relative_to_hw_limit(max_mem_utilization); + + record.set_include_infeed_outfeed(include_infeed_outfeed); + record.set_apply_time_scale_multiplier(apply_time_scale_multiplier); + + return record; +} + RooflineModelRecord GenerateRooflineModelProgramRecord( const OpStats& op_stats, const OpMetricsDb& db, RecordType record_type, uint32_t step_num, const RooflineModelDatabase& roofline_model_db, @@ -184,6 +270,90 @@ RooflineModelRecord GenerateRooflineModelProgramRecord( return program_record; } +RooflineModelRecord GenerateRooflineModelProgramRecord( + const OpStats& op_stats, const FlatOpMetricsDb& db, RecordType record_type, + uint32_t step_num, const RooflineModelDatabase& roofline_model_db, + bool include_infeed_outfeed, bool apply_time_scale_multiplier) { + FlatOpMetrics program_metrics; + program_metrics.set_hlo_name("Program"); + program_metrics.set_category("Program"); + program_metrics.set_occurrences(1); + uint64_t infeed_outfeed_time = 0; + for (const FlatOpMetrics& metrics : db.op_instances()) { + // Aggregate innermost ops only to avoid redundant counting. + if (tsl::profiler::MayHaveInnerOps(metrics.category())) continue; + // Skipping sparse core metrics to have equivalence with OpMetricsDb. + if (metrics.core_type() == FlatOpMetrics::SPARSE_CORE) continue; + if (!include_infeed_outfeed && + tsl::profiler::IsInfeedOrOutfeed(metrics.category())) { + infeed_outfeed_time += metrics.time_ps(); + continue; + } + program_metrics.set_flops(program_metrics.flops() + metrics.flops()); + program_metrics.set_flops_v2(program_metrics.flops_v2() + + metrics.flops_v2()); + program_metrics.set_model_flops(program_metrics.model_flops() + + metrics.model_flops()); + program_metrics.set_model_flops_v2(program_metrics.model_flops_v2() + + metrics.model_flops_v2()); + program_metrics.set_bytes_accessed(program_metrics.bytes_accessed() + + metrics.bytes_accessed()); + + FlatOpMetricsDbCombiner::CombineMemoryAccessedBreakdown( + metrics.memory_accessed_breakdown(), + program_metrics.mutable_memory_accessed_breakdown()); + } + uint64_t total_time_ps = db.total_time_ps(); + if (!include_infeed_outfeed) total_time_ps -= infeed_outfeed_time; + program_metrics.set_time_ps(total_time_ps); + RooflineModelRecord program_record = ConvertOpMetricsToRooflineModelRecord( + op_stats, program_metrics, record_type, step_num, total_time_ps, + roofline_model_db, include_infeed_outfeed, apply_time_scale_multiplier); + program_record.set_rank(0); + program_record.set_total_self_time_as_fraction(0.0); + program_record.set_cumulative_total_self_time_as_fraction(0.0); + return program_record; +} + +tsl::protobuf::RepeatedPtrField +ConvertFlatOpMetricsDbToRooflineModelRecords( + const OpStats& op_stats, const FlatOpMetricsDb& db, RecordType record_type, + uint32_t step_num, const RooflineModelDatabase& roofline_model_db, + bool include_infeed_outfeed, bool apply_time_scale_multiplier) { + tsl::protobuf::RepeatedPtrField roofline_model_records; + RooflineModelRecord* program_record = roofline_model_records.Add(); + *program_record = GenerateRooflineModelProgramRecord( + op_stats, db, record_type, step_num, roofline_model_db, + include_infeed_outfeed, apply_time_scale_multiplier); + const RooflineModelRecord* prev_record = program_record; + uint64_t infeed_outfeed_time = 0; + if (!include_infeed_outfeed) { + // Calculate the total time spent on infeed and outfeed ops. + for (const FlatOpMetrics& metrics : db.op_instances()) { + if (metrics.core_type() == FlatOpMetrics::SPARSE_CORE) continue; + if (tsl::profiler::IsInfeedOrOutfeed(metrics.category())) { + infeed_outfeed_time += metrics.time_ps(); + } + } + } + uint64_t total_time_ps = db.total_time_ps() - infeed_outfeed_time; + double total_time_us = tsl::profiler::PicoToMicro(total_time_ps); + for (const auto* metrics : SortedOpMetricsDb(db, kMaxNumRecords)) { + if (metrics->occurrences() == 0) continue; + if (!include_infeed_outfeed && + tsl::profiler::IsInfeedOrOutfeed(metrics->category())) { + continue; + } + RooflineModelRecord* record = roofline_model_records.Add(); + *record = ConvertOpMetricsToRooflineModelRecord( + op_stats, *metrics, record_type, step_num, total_time_ps, + roofline_model_db, include_infeed_outfeed, apply_time_scale_multiplier); + SetRankAndTimeFractions(total_time_us, *prev_record, record); + prev_record = record; + } + return roofline_model_records; +} + tsl::protobuf::RepeatedPtrField ConvertOpMetricsDbToRooflineModelRecords( const OpStats& op_stats, const OpMetricsDb& db, RecordType record_type, @@ -222,6 +392,19 @@ ConvertOpMetricsDbToRooflineModelRecords( return roofline_model_records; } +void SetTimeScaleMultiplier(const OpStats& op_stats, + RooflineModelDatabase& roofline_model_db) { + if (op_stats.has_flat_device_op_metrics_db()) { + roofline_model_db.set_time_scale_multiplier(tsl::profiler::SafeDivide( + op_stats.flat_device_op_metrics_db().normalized_total_op_time_ps(), + op_stats.flat_device_op_metrics_db().total_op_time_ps())); + } else { + roofline_model_db.set_time_scale_multiplier(tsl::profiler::SafeDivide( + op_stats.device_op_metrics_db().normalized_total_op_time_ps(), + op_stats.device_op_metrics_db().total_op_time_ps())); + } +} + RooflineModelDatabase InitializeRooflineModelDatabaseFromOpStats( const OpStats& op_stats, bool include_infeed_outfeed) { tensorflow::profiler::HardwareType hardware_type = @@ -231,9 +414,7 @@ RooflineModelDatabase InitializeRooflineModelDatabaseFromOpStats( RooflineModelDatabase roofline_model_db; const PerfEnv& perf_env = op_stats.perf_env(); roofline_model_db.set_device_type(op_stats.run_environment().device_type()); - roofline_model_db.set_time_scale_multiplier(tsl::profiler::SafeDivide( - op_stats.device_op_metrics_db().normalized_total_op_time_ps(), - op_stats.device_op_metrics_db().total_op_time_ps())); + SetTimeScaleMultiplier(op_stats, roofline_model_db); // Set peak flop rate in GFLOPs/s. roofline_model_db.set_peak_flop_rate( diff --git a/xprof/convert/op_stats_to_roofline_model.h b/xprof/convert/op_stats_to_roofline_model.h index 4f4ab9cce..4ae6293a0 100644 --- a/xprof/convert/op_stats_to_roofline_model.h +++ b/xprof/convert/op_stats_to_roofline_model.h @@ -20,8 +20,10 @@ limitations under the License. #include #include +#include "absl/log/log.h" #include "tsl/platform/protobuf.h" #include "xprof/convert/data_table_utils.h" +#include "plugin/xprof/protobuf/flat_op_metrics.pb.h" #include "plugin/xprof/protobuf/op_metrics.pb.h" #include "plugin/xprof/protobuf/op_stats.pb.h" #include "plugin/xprof/protobuf/roofline_model.pb.h" @@ -52,6 +54,23 @@ ConvertOpMetricsDbToRooflineModelRecords( uint32_t step_num, const RooflineModelDatabase& roofline_model_db, bool include_infeed_outfeed, bool apply_time_scale_multiplier = false); +RooflineModelRecord ConvertOpMetricsToRooflineModelRecord( + const OpStats& op_stats, const FlatOpMetrics& metrics, + RecordType record_type, uint32_t step_num, uint64_t total_time_ps, + const RooflineModelDatabase& roofline_model_db, bool include_infeed_outfeed, + bool apply_time_scale_multiplier = false); + +RooflineModelRecord GenerateRooflineModelProgramRecord( + const OpStats& op_stats, const FlatOpMetricsDb& db, RecordType record_type, + uint32_t step_num, const RooflineModelDatabase& roofline_model_db, + bool include_infeed_outfeed, bool apply_time_scale_multiplier = false); + +tsl::protobuf::RepeatedPtrField +ConvertFlatOpMetricsDbToRooflineModelRecords( + const OpStats& op_stats, const FlatOpMetricsDb& db, RecordType record_type, + uint32_t step_num, const RooflineModelDatabase& roofline_model_db, + bool include_infeed_outfeed, bool apply_time_scale_multiplier = false); + tensorflow::profiler::roofline_model::RooflineModelDatabase ConvertOpStatsToRooflineModel(const tensorflow::profiler::OpStats& tf_op_stats, bool include_infeed_outfeed, @@ -65,18 +84,33 @@ InitializeRooflineModelDatabaseFromOpStats(const OpStats& op_stats, inline void AddRooflineModelRecordForProfileDuration( const OpStats& op_stats, RooflineModelDatabase& roofline_model_db, bool include_infeed_outfeed, bool apply_time_scale_multiplier = false) { - *roofline_model_db.mutable_roofline_model_record() = - ConvertOpMetricsDbToRooflineModelRecords( - op_stats, op_stats.device_op_metrics_db(), RecordType::ALL, - /*step_num=*/0, roofline_model_db, include_infeed_outfeed, - apply_time_scale_multiplier); + if (op_stats.has_flat_device_op_metrics_db() && + !op_stats.flat_device_op_metrics_db().op_instances().empty()) { + *roofline_model_db.mutable_roofline_model_record() = + ConvertFlatOpMetricsDbToRooflineModelRecords( + op_stats, op_stats.flat_device_op_metrics_db(), RecordType::ALL, + /*step_num=*/0, roofline_model_db, include_infeed_outfeed, + apply_time_scale_multiplier); + } else { + *roofline_model_db.mutable_roofline_model_record() = + ConvertOpMetricsDbToRooflineModelRecords( + op_stats, op_stats.device_op_metrics_db(), RecordType::ALL, + /*step_num=*/0, roofline_model_db, include_infeed_outfeed, + apply_time_scale_multiplier); + } } // Generate RooflineModelRecord for the HLO DB over complete steps only. inline void AddRooflineModelRecordsForCompleteSteps( const OpStats& op_stats, RooflineModelDatabase& roofline_model_db, bool include_infeed_outfeed, bool apply_time_scale_multiplier = false) { - if (op_stats.has_hlo_metrics_db_complete_steps_only()) { + if (op_stats.has_flat_hlo_metrics_db_complete_steps_only()) { + *roofline_model_db.add_roofline_model_record() = + GenerateRooflineModelProgramRecord( + op_stats, op_stats.flat_hlo_metrics_db_complete_steps_only(), + RecordType::AVERAGE_STEP, /*step_num=*/0, roofline_model_db, + include_infeed_outfeed, apply_time_scale_multiplier); + } else if (op_stats.has_hlo_metrics_db_complete_steps_only()) { *roofline_model_db.add_roofline_model_record() = GenerateRooflineModelProgramRecord( op_stats, op_stats.hlo_metrics_db_complete_steps_only(), @@ -90,10 +124,19 @@ inline void AddRooflineModelRecordsPerStep( const OpStats& op_stats, RooflineModelDatabase& roofline_model_db, bool include_infeed_outfeed, bool apply_time_scale_multiplier = false) { for (const auto& step_info : op_stats.step_db().step_sequence()) { - *roofline_model_db.add_roofline_model_record() = - GenerateRooflineModelProgramRecord( - op_stats, step_info.hlo_metrics_db(), RecordType::PER_STEP, - step_info.step_num(), roofline_model_db, include_infeed_outfeed); + if (step_info.has_flat_hlo_metrics_db()) { + *roofline_model_db.add_roofline_model_record() = + GenerateRooflineModelProgramRecord( + op_stats, step_info.flat_hlo_metrics_db(), RecordType::PER_STEP, + step_info.step_num(), roofline_model_db, include_infeed_outfeed, + apply_time_scale_multiplier); + } else { + *roofline_model_db.add_roofline_model_record() = + GenerateRooflineModelProgramRecord( + op_stats, step_info.hlo_metrics_db(), RecordType::PER_STEP, + step_info.step_num(), roofline_model_db, include_infeed_outfeed, + apply_time_scale_multiplier); + } } } diff --git a/xprof/convert/step_events_to_steps_db.cc b/xprof/convert/step_events_to_steps_db.cc index b439c475e..0dd02f6bc 100644 --- a/xprof/convert/step_events_to_steps_db.cc +++ b/xprof/convert/step_events_to_steps_db.cc @@ -15,9 +15,9 @@ limitations under the License. #include "xprof/convert/step_events_to_steps_db.h" #include +#include #include #include -#include #include #include @@ -30,9 +30,11 @@ limitations under the License. #include "xla/tsl/platform/logging.h" #include "xla/tsl/platform/types.h" #include "xla/tsl/profiler/utils/timespan.h" +#include "xprof/convert/flat_op_metrics_db_combiner.h" #include "xprof/convert/op_metrics_db_combiner.h" #include "plugin/xprof/protobuf/steps_db.pb.h" #include "xprof/utils/event_span.h" +#include "xprof/utils/flat_op_metrics_db_utils.h" #include "xprof/utils/op_metrics_db_utils.h" namespace tensorflow { @@ -47,7 +49,8 @@ namespace { void StepEventsToPerCoreStepInfo(uint32_t step_num, StepDetails& step_details, PerCoreStepInfo& per_core_step_info) { per_core_step_info.set_step_num(step_num); - OpMetricsDbCombiner combiner(per_core_step_info.mutable_hlo_metrics_db()); + std::unique_ptr combiner; + std::unique_ptr flat_combiner; auto step_time = step_details.StepTime(); if (step_time.duration_ps() == 0) { // In case no step markers are observed for the particular step, Skip the @@ -65,7 +68,13 @@ void StepEventsToPerCoreStepInfo(uint32_t step_num, StepDetails& step_details, AddIdleOp(metrics_db); // TODO(b/397774568): Remove this once the SparseCore OpMetricsDb is // implemented. - if (core_id < kSparseCoreIndexStart) combiner.Combine(metrics_db); + if (core_id < kSparseCoreIndexStart) { + if (!combiner) { + combiner = std::make_unique( + per_core_step_info.mutable_hlo_metrics_db()); + } + combiner->Combine(metrics_db); + } GenericStepBreakdown step_breakdown; auto& category_ps = *(step_breakdown.mutable_category_ps()); @@ -82,6 +91,38 @@ void StepEventsToPerCoreStepInfo(uint32_t step_num, StepDetails& step_details, (*per_core_step_info.mutable_step_info_per_core())[core_id] = std::move(step_info); } + for (auto& [core_id, flat_metrics_db] : + step_details.PerCoreFlatOpMetricsDb()) { + tsl::profiler::Timespan step_time_on_core = + core_id >= kSparseCoreIndexStart ? step_details.StepTimeOnCore(core_id) + : step_time; + SetTotalTimePs(flat_metrics_db, step_time_on_core.duration_ps()); + AddIdleOp(flat_metrics_db); + if (core_id < kSparseCoreIndexStart) { + if (!flat_combiner) { + flat_combiner = std::make_unique( + per_core_step_info.mutable_flat_hlo_metrics_db()); + } + flat_combiner->Combine(flat_metrics_db); + } + + auto [it, inserted] = + per_core_step_info.mutable_step_info_per_core()->try_emplace(core_id); + if (inserted) { + StepInfoResult& step_info = it->second; + GenericStepBreakdown step_breakdown; + auto& category_ps = *(step_breakdown.mutable_category_ps()); + for (auto& metric : flat_metrics_db.op_instances()) { + category_ps[metric.category()] += metric.self_time_ps(); + } + + step_info.set_step_num(step_num); + step_info.set_step_name(step_details.StepName()); + step_info.set_begin_ps(step_time_on_core.begin_ps()); + step_info.set_duration_ps(step_time_on_core.duration_ps()); + step_info.mutable_step_breakdown()->PackFrom(step_breakdown); + } + } auto& all_reduce_db_per_core_map = *per_core_step_info.mutable_all_reduce_db_per_core(); for (const auto& [core_id, all_reduce_db] : step_details.Collectives()) { @@ -178,7 +219,8 @@ StepDatabaseResult ConvertStepEventsToStepDb( if (step_details == nullptr) continue; PerCoreStepInfo per_core_step_info; per_core_step_info.set_step_num(step); - if (!step_details->PerCoreOpMetricsDb().empty()) { + if (!step_details->PerCoreOpMetricsDb().empty() || + !step_details->PerCoreFlatOpMetricsDb().empty()) { StepEventsToPerCoreStepInfo(step, *step_details, per_core_step_info); } else { StepInfoResult step_info = diff --git a/xprof/convert/xplane_to_flat_op_metrics_db.cc b/xprof/convert/xplane_to_flat_op_metrics_db.cc index 67597fe6f..919b19a5d 100644 --- a/xprof/convert/xplane_to_flat_op_metrics_db.cc +++ b/xprof/convert/xplane_to_flat_op_metrics_db.cc @@ -25,6 +25,7 @@ limitations under the License. #include "absl/container/btree_set.h" #include "absl/container/flat_hash_map.h" +#include "absl/strings/string_view.h" #include "xla/tsl/profiler/utils/tf_xplane_visitor.h" #include "xla/tsl/profiler/utils/trace_utils.h" #include "xla/tsl/profiler/utils/timespan.h" @@ -68,12 +69,12 @@ void ConvertSparseCoreDeviceTraceXPlaneToFlatOpMetricsDb( }; struct OpMetricsInProgress { - std::shared_ptr builder; + std::unique_ptr builder; tsl::profiler::AncestorStack event_stack; OpMetricsInProgress() - : builder(std::make_shared()), + : builder(std::make_unique()), event_stack( - [builder = builder](const ParentReference& parent) { + [builder_ptr = builder.get()](const ParentReference& parent) { FlatOpMetrics op_metrics = XEventsFlatOpMetricsDbBuilder::FromXEvent(parent.event); op_metrics.set_time_ps(parent.device_timespan.duration_ps()); @@ -90,11 +91,12 @@ void ConvertSparseCoreDeviceTraceXPlaneToFlatOpMetricsDb( : 1.0; op_metrics.set_normalized_time_ps(op_metrics.time_ps() * factor); - auto key = GetOpKeyFromXEvent(parent.event); + XEventsOpMetricsDbBuilder::OpKey key = + GetOpKeyFromXEvent(parent.event); XEventsFlatOpMetricsDbBuilder::OpKey flat_key; flat_key.program_id = key.program_id; flat_key.symbol_id = key.symbol_id; - builder->AddOpMetric(op_metrics, flat_key); + builder_ptr->AddOpMetric(op_metrics, flat_key); }, [](const ParentReference& parent, const ParentReference& child) { return parent.device_timespan.Includes(child.device_timespan); @@ -152,10 +154,9 @@ void ConvertSparseCoreDeviceTraceXPlaneToFlatOpMetricsDb( // Check if the current module_it encapsulates the event. if (module_it->timespan.Includes(timespan)) { // Insert that event into the stack for that module. - std::string hlo_name = - event.Metadata().HasDisplayName() - ? std::string(event.Metadata().DisplayName()) - : std::string(event.Metadata().Name()); + absl::string_view hlo_name = event.Metadata().HasDisplayName() + ? event.Metadata().DisplayName() + : event.Metadata().Name(); XEventsOpMetricsDbBuilder::OpKey key = GetOpKeyFromXEvent(event); auto module_id = key.program_id.has_value() ? key.program_id.value() : module_it->tc_start_id; diff --git a/xprof/convert/xplane_to_op_stats.cc b/xprof/convert/xplane_to_op_stats.cc index cb7b0cc18..fd66fb7f1 100644 --- a/xprof/convert/xplane_to_op_stats.cc +++ b/xprof/convert/xplane_to_op_stats.cc @@ -30,6 +30,7 @@ limitations under the License. #include "absl/container/flat_hash_set.h" #include "absl/log/check.h" #include "absl/log/log.h" +#include "absl/log/vlog_is_on.h" #include "absl/status/status.h" #include "absl/status/statusor.h" #include "absl/strings/match.h" @@ -41,7 +42,6 @@ limitations under the License. #include "xla/tsl/profiler/utils/tf_xplane_visitor.h" #include "xla/tsl/profiler/utils/timespan.h" #include "xla/tsl/profiler/utils/tpu_xplane_utils.h" -#include "xla/tsl/profiler/utils/xplane_builder.h" #include "xla/tsl/profiler/utils/xplane_schema.h" #include "xla/tsl/profiler/utils/xplane_utils.h" #include "xla/tsl/util/stats_calculator.h" @@ -308,7 +308,8 @@ FlatOpMetrics CreateFusionChildMetrics( } void AddFusionChildrenToFlatOpMetrics( - FlatOpMetricMeta parent_op_metrics, FlatOpMetrics::TpuCoreType core_type, + const FlatOpMetricMeta& parent_op_metrics, + FlatOpMetrics::TpuCoreType core_type, const HloInstructionWrapper* instr_wrapper, FlatOpMetricsDb& children_db) { if (instr_wrapper->FusedChildren().empty()) return; for (const HloInstructionWrapper* child : instr_wrapper->FusedChildren()) { @@ -950,6 +951,40 @@ absl::StatusOr ConvertXSpaceToFlatOpMetricsDb( }); } } + + // StepDb Generation. + std::vector all_step_events; + + // Ensure step_events threads are joined and results combined when the + // function exits. + auto step_events_cleanup = + absl::MakeCleanup([&all_step_events, &step_events, is_tpu]() { + for (auto& device_step_events : all_step_events) { + if (device_step_events.empty()) { + continue; + } + if (is_tpu) { + // In TPU, we take the intersection of step events across cores + // as well as hosts.see b/158249775 and cl/331842545. + IntersectCombineStepEvents(device_step_events, &step_events); + } else { + UnionCombineStepEvents(device_step_events, &step_events); + } + } + }); + + if (options.generate_step_db) { + all_step_events.resize(device_planes.size()); + for (size_t i = 0; i < device_planes.size(); ++i) { + const XPlane* device_trace = device_planes[i]; + auto& current_step_events = all_step_events[i]; + executor->Execute([device_trace, ¤t_step_events]() { + current_step_events = ConvertDeviceTraceXPlaneToStepEvents( + *device_trace, /*for_flat_profile=*/true); + }); + } + } + LOG(INFO) << "ConvertXSpaceToOpStats: Scheduled " << device_planes.size() << " OpMetricsDb generation tasks."; @@ -957,6 +992,29 @@ absl::StatusOr ConvertXSpaceToFlatOpMetricsDb( // The cleanup blocks will execute after this step. } + if (options.generate_step_db) { + if (is_tpu) { + *op_stats.mutable_step_db() = ConvertStepEventsToStepDb( + /*has_device=*/true, /*maybe_drop_incomplete_steps=*/false, + step_events); + *op_stats.mutable_device_op_metrics_db()->mutable_precision_stats() = + ComputePrecisionStats(step_events); + FlatOpMetricsDbCombiner flat_combiner( + op_stats.mutable_flat_hlo_metrics_db_complete_steps_only()); + for (const auto& step_info : op_stats.step_db().step_sequence()) { + flat_combiner.Combine(step_info.flat_hlo_metrics_db()); + } + } else { + StepEvents nonoverlapped_step_events = + ToNonOverlappedStepEvents(step_events); + *op_stats.mutable_step_db() = ConvertStepEventsToStepDb( + /*has_device=*/true, options.maybe_drop_incomplete_steps, + nonoverlapped_step_events); + *op_stats.mutable_device_op_metrics_db()->mutable_precision_stats() = + ComputePrecisionStats(nonoverlapped_step_events); + } + } + // Removed assignment of program_id_to_name_map as it was removed from // FlatOpMetricsDb proto. diff --git a/xprof/convert/xplane_to_step_events.cc b/xprof/convert/xplane_to_step_events.cc index 06169143a..d3e650ecf 100644 --- a/xprof/convert/xplane_to_step_events.cc +++ b/xprof/convert/xplane_to_step_events.cc @@ -28,7 +28,6 @@ limitations under the License. #include "absl/strings/numbers.h" #include "absl/strings/str_split.h" #include "absl/strings/string_view.h" -#include "xla/tsl/platform/types.h" #include "xla/tsl/profiler/utils/tf_op_utils.h" #include "xla/tsl/profiler/utils/tf_xplane_visitor.h" #include "xla/tsl/profiler/utils/timespan.h" @@ -41,6 +40,7 @@ limitations under the License. #include "plugin/xprof/protobuf/op_metrics.pb.h" #include "plugin/xprof/protobuf/steps_db.pb.h" #include "xprof/utils/event_span.h" +#include "xprof/utils/flat_op_metrics_db_utils.h" #include "xprof/utils/op_metrics_db_utils.h" namespace tensorflow { @@ -158,6 +158,59 @@ EventType ClassifyCpuEvent(absl::string_view event_name, bool has_device, } } +bool FindDeviceTraceSpan(const XPlaneVisitor& plane, uint64_t* min_timestamp, + uint64_t* max_timestamp) { + uint64_t min_ts = std::numeric_limits::max(); + uint64_t max_ts = 0; + bool has_events = false; + plane.ForEachLine([&](const XLineVisitor& line) { + if (tsl::profiler::IsDerivedThreadId(line.Id())) return; + if (line.Id() == tsl::profiler::kThreadIdStepInfo) return; + line.ForEachEvent([&](const XEventVisitor& event) { + tsl::profiler::Timespan ts = tsl::profiler::GetDeviceEventTimespan(event); + min_ts = std::min(min_ts, ts.begin_ps()); + max_ts = std::max(max_ts, ts.end_ps()); + has_events = true; + }); + }); + if (has_events) { + *min_timestamp = min_ts; + *max_timestamp = max_ts; + } + return has_events; +} + +void ProcessIncompleteStepInfo(const XPlaneVisitor& plane, + std::optional sc_core_id, + StepEvents* step_markers, + StepEvents* step_events) { + if (!step_markers->empty()) return; + uint64_t min_timestamp = 0; + uint64_t max_timestamp = 0; + if (!FindDeviceTraceSpan(plane, &min_timestamp, &max_timestamp)) return; + if (min_timestamp >= max_timestamp) return; + + uint32_t id = plane.Id(); + if (sc_core_id.has_value()) { + id = kSparseCoreIndexStart + id; + } + const int64_t kIncompleteStepGroupId = 4294967294; + auto timespan = + tsl::profiler::Timespan::FromEndPoints(min_timestamp, max_timestamp); + (*step_markers)[kIncompleteStepGroupId].AddMarker(StepMarker( + StepMarkerType::kDeviceStepMarker, id, "Incomplete Step", timespan)); + (*step_markers)[kIncompleteStepGroupId].SetStepName("Incomplete Step"); + if (!step_events->empty()) { + StepEvents new_step_events; + StepDetails& incomplete_step_details = + new_step_events[kIncompleteStepGroupId]; + for (auto& [group_id, step_details] : *step_events) { + incomplete_step_details.Combine(step_details); + } + *step_events = std::move(new_step_events); + } +} + } // namespace StepEvents ConvertHostThreadsXLineToStepEvents( @@ -356,26 +409,51 @@ StepEvents ConvertDeviceTraceXLineToStepEvents(const uint64_t device_id, return result; } -StepEvents ConvertTpuDeviceTraceXLineToStepEvents(const uint64_t device_id, - const XLineVisitor& line) { +StepEvents ConvertTpuDeviceTraceXLineToStepEvents( + const uint64_t device_id, const XLineVisitor& line, + bool for_flat_profile = false) { StepEvents result; absl::flat_hash_map op_metrics_builder; + absl::flat_hash_map + flat_op_metrics_builder; struct ParentRef { const XEventVisitor event; tsl::profiler::Timespan device_timespan; uint64_t children_duration_ps = 0; int64_t group_id = -1; }; + FlatOpMetrics flat_op_metrics; + OpMetrics op_metrics; tsl::profiler::AncestorStack event_stack( // Adds an OpMetric to the builder based on the provided parent reference. [&](const ParentRef& parent) { - OpMetrics op_metrics = FromXEvent(parent.event); - op_metrics.set_time_ps(parent.device_timespan.duration_ps()); - op_metrics.set_self_time_ps(op_metrics.time_ps() - - parent.children_duration_ps); - op_metrics_builder[parent.group_id].AddOpMetric( - op_metrics, GetOpKeyFromXEvent(parent.event)); + if (for_flat_profile) { + flat_op_metrics = + XEventsFlatOpMetricsDbBuilder::FromXEvent(parent.event); + flat_op_metrics.set_time_ps(parent.device_timespan.duration_ps()); + flat_op_metrics.set_self_time_ps(flat_op_metrics.time_ps() - + parent.children_duration_ps); + std::optional + time_scale_multiplier_stat = + parent.event.GetStat(StatType::kTimeScaleMultiplier); + double factor = time_scale_multiplier_stat.has_value() + ? time_scale_multiplier_stat->DoubleValue() + : 1.0; + flat_op_metrics.set_normalized_time_ps(flat_op_metrics.time_ps() * + factor); + flat_op_metrics_builder[parent.group_id].AddOpMetric( + flat_op_metrics, + XEventsFlatOpMetricsDbBuilder::GetFlatOpKeyFromXEvent( + parent.event)); + } else { + op_metrics = FromXEvent(parent.event); + op_metrics.set_time_ps(parent.device_timespan.duration_ps()); + op_metrics.set_self_time_ps(op_metrics.time_ps() - + parent.children_duration_ps); + op_metrics_builder[parent.group_id].AddOpMetric( + op_metrics, GetOpKeyFromXEvent(parent.event)); + } }, // Checks if the child event is a child of the parent event. [](const ParentRef& parent, const ParentRef& child) { @@ -403,14 +481,20 @@ StepEvents ConvertTpuDeviceTraceXLineToStepEvents(const uint64_t device_id, } }); event_stack.Flush(); - for (auto& [group_id, builder] : op_metrics_builder) { - // Finalize Without the step time now. - result[group_id].SetPerCoreOpMetricsDb(builder.Finalize(), device_id); + if (for_flat_profile) { + for (auto& [group_id, builder] : flat_op_metrics_builder) { + result[group_id].SetPerCoreFlatOpMetricsDb(builder.Finalize(), device_id); + } + } else { + for (auto& [group_id, builder] : op_metrics_builder) { + result[group_id].SetPerCoreOpMetricsDb(builder.Finalize(), device_id); + } } return result; } -StepEvents ConvertDeviceTraceXPlaneToStepEvents(const XPlane& device_trace) { +StepEvents ConvertDeviceTraceXPlaneToStepEvents(const XPlane& device_trace, + bool for_flat_profile) { XPlaneVisitor plane = tsl::profiler::CreateTfXPlaneVisitor(&device_trace); std::optional tpu_core_id = tsl::profiler::GetTensorCoreId(plane.Name()); std::optional sc_core_id = tsl::profiler::GetSparseCoreId(plane.Name()); @@ -437,13 +521,14 @@ StepEvents ConvertDeviceTraceXPlaneToStepEvents(const XPlane& device_trace) { if (!tsl::profiler::IsOpLineName(line.Name())) return; // There should only be a single OpLine per TPU core. DCHECK(step_events.empty()); - step_events = ConvertTpuDeviceTraceXLineToStepEvents(plane.Id(), line); + step_events = ConvertTpuDeviceTraceXLineToStepEvents( + plane.Id(), line, for_flat_profile); } else if (sc_core_id.has_value()) { if (line.Name() != tsl::profiler::kSparseCoreOpLineName) return; // There should only be a single SparseCore StepLine per SparseCore. DCHECK(step_events.empty()); step_events = ConvertTpuDeviceTraceXLineToStepEvents( - kSparseCoreIndexStart + plane.Id(), line); + kSparseCoreIndexStart + plane.Id(), line, for_flat_profile); } else { // There may be multiple streams per GPU device so union the results. StepEvents stream_step_events = @@ -452,6 +537,7 @@ StepEvents ConvertDeviceTraceXPlaneToStepEvents(const XPlane& device_trace) { } } }); + ProcessIncompleteStepInfo(plane, sc_core_id, &step_markers, &step_events); if (!step_events.empty()) { IntersectCombineStepEvents(step_markers, &step_events); } diff --git a/xprof/convert/xplane_to_step_events.h b/xprof/convert/xplane_to_step_events.h index cf83786eb..04215e882 100644 --- a/xprof/convert/xplane_to_step_events.h +++ b/xprof/convert/xplane_to_step_events.h @@ -53,7 +53,8 @@ StepEvents ConvertHostThreadsXPlaneToStepEvents( StepEvents ConvertDeviceTraceXLineToStepEvents(const XLineVisitor& line); // Convert the device trace in XPlane format to StepEvents. -StepEvents ConvertDeviceTraceXPlaneToStepEvents(const XPlane& device_trace); +StepEvents ConvertDeviceTraceXPlaneToStepEvents(const XPlane& device_trace, + bool for_flat_profile = false); } // namespace profiler } // namespace tensorflow diff --git a/xprof/convert/xplane_to_step_events_test.cc b/xprof/convert/xplane_to_step_events_test.cc index dbcb5cb86..7421b06b4 100644 --- a/xprof/convert/xplane_to_step_events_test.cc +++ b/xprof/convert/xplane_to_step_events_test.cc @@ -21,6 +21,7 @@ limitations under the License. #include "" #include "absl/container/flat_hash_map.h" #include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" #include "xla/tsl/profiler/utils/group_events.h" #include "xla/tsl/profiler/utils/preprocess_xplane.h" #include "xla/tsl/profiler/utils/tf_xplane_visitor.h" @@ -45,6 +46,14 @@ using ::tsl::profiler::XLineBuilder; using ::tsl::profiler::XPlaneBuilder; using ::tsl::profiler::XStatsBuilder; +const FlatOpMetrics* FindFlatOp(const FlatOpMetricsDb& db, + absl::string_view name) { + for (const auto& op : db.op_instances()) { + if (op.hlo_name() == name) return &op; + } + return nullptr; +} + // Tests with a sample profile with two steps captured on the host but only one // step on the device. On the host, each step consists of TraceContext -> // FunctionRun -> ExecutorState::Process -> matmul. On the host, each step @@ -199,6 +208,225 @@ TEST(ConvertXPlaneToStepEvents, TpuDevicePlaneToStepEvents) { EXPECT_EQ(step_2.Markers().size(), 1); } +TEST(ConvertXPlaneToStepEvents, TpuDevicePlaneToFlatStepEvents) { + XPlane raw_plane; + XPlaneBuilder plane(&raw_plane); + int64_t device_id = 1; + plane.SetId(device_id); + plane.SetName("/device:TPU:0"); + XLineBuilder op_line = plane.GetOrCreateLine(0); + op_line.SetName(tsl::profiler::kXlaOpLineName); + const XStatMetadata& program_id_stat = + *plane.GetOrCreateStatMetadata(GetStatTypeStr(StatType::kProgramId)); + const XStatMetadata& symbol_id_stat = + *plane.GetOrCreateStatMetadata(GetStatTypeStr(StatType::kSymbolId)); + const XStatMetadata& group_id_stat = + *plane.GetOrCreateStatMetadata(GetStatTypeStr(StatType::kGroupId)); + // Op 1 (symbol 1) + XEventMetadata* op1_meta = plane.GetOrCreateEventMetadata("op1_long"); + op1_meta->set_display_name("op1"); + { + XStatsBuilder stats(op1_meta, &plane); + stats.AddStatValue(program_id_stat, 1); + stats.AddStatValue(symbol_id_stat, 1); + } + + // Op 2 (symbol 2) + XEventMetadata* op2_meta = plane.GetOrCreateEventMetadata("op2_long"); + op2_meta->set_display_name("op2"); + { + XStatsBuilder stats(op2_meta, &plane); + stats.AddStatValue(program_id_stat, 1); + stats.AddStatValue(symbol_id_stat, 2); + } + + // Op 3 (symbol 3) + XEventMetadata* op3_meta = plane.GetOrCreateEventMetadata("op3_long"); + op3_meta->set_display_name("op3"); + { + XStatsBuilder stats(op3_meta, &plane); + stats.AddStatValue(program_id_stat, 1); + stats.AddStatValue(symbol_id_stat, 3); + } + + // Nested 1 (symbol 4) + XEventMetadata* nested1_meta = + plane.GetOrCreateEventMetadata("nested1_long"); + nested1_meta->set_display_name("nested1"); + { + XStatsBuilder stats(nested1_meta, &plane); + stats.AddStatValue(program_id_stat, 1); + stats.AddStatValue(symbol_id_stat, 4); + } + + // Nested 2 (symbol 5) + XEventMetadata* nested2_meta = + plane.GetOrCreateEventMetadata("nested2_long"); + nested2_meta->set_display_name("nested2"); + { + XStatsBuilder stats(nested2_meta, &plane); + stats.AddStatValue(program_id_stat, 1); + stats.AddStatValue(symbol_id_stat, 5); + } + + // Nested 3 (symbol 6) + XEventMetadata* nested3_meta = + plane.GetOrCreateEventMetadata("nested3_long"); + nested3_meta->set_display_name("nested3"); + { + XStatsBuilder stats(nested3_meta, &plane); + stats.AddStatValue(program_id_stat, 1); + stats.AddStatValue(symbol_id_stat, 6); + } + + // Nested 4 (symbol 7) + XEventMetadata* nested4_meta = + plane.GetOrCreateEventMetadata("nested4_long"); + nested4_meta->set_display_name("nested4"); + { + XStatsBuilder stats(nested4_meta, &plane); + stats.AddStatValue(program_id_stat, 1); + stats.AddStatValue(symbol_id_stat, 7); + } + + // Group 1 events + // op1 (0-50) + { + XEventBuilder event = op_line.AddEvent(*op1_meta); + event.SetOffsetPs(0); + event.SetDurationPs(50); + event.AddStatValue(group_id_stat, 1); + } + // nested1 (10-30) + { + XEventBuilder event = op_line.AddEvent(*nested1_meta); + event.SetOffsetPs(10); + event.SetDurationPs(20); + event.AddStatValue(group_id_stat, 1); + } + // op2 (50-100) + { + XEventBuilder event = op_line.AddEvent(*op2_meta); + event.SetOffsetPs(50); + event.SetDurationPs(50); + event.AddStatValue(group_id_stat, 1); + } + // nested2 (60-90) + { + XEventBuilder event = op_line.AddEvent(*nested2_meta); + event.SetOffsetPs(60); + event.SetDurationPs(30); + event.AddStatValue(group_id_stat, 1); + } + + // Group 2 events + // op1 (100-150) + { + XEventBuilder event = op_line.AddEvent(*op1_meta); + event.SetOffsetPs(100); + event.SetDurationPs(50); + event.AddStatValue(group_id_stat, 2); + } + // nested3 (110-140) + { + XEventBuilder event = op_line.AddEvent(*nested3_meta); + event.SetOffsetPs(110); + event.SetDurationPs(30); + event.AddStatValue(group_id_stat, 2); + } + // op3 (150-200) + { + XEventBuilder event = op_line.AddEvent(*op3_meta); + event.SetOffsetPs(150); + event.SetDurationPs(50); + event.AddStatValue(group_id_stat, 2); + } + // nested4 (160-190) + { + XEventBuilder event = op_line.AddEvent(*nested4_meta); + event.SetOffsetPs(160); + event.SetDurationPs(30); + event.AddStatValue(group_id_stat, 2); + } + XLineBuilder step_line = plane.GetOrCreateLine(1); + step_line.SetName(tsl::profiler::kStepLineName); + { + XEventMetadata* event_metadata = plane.CreateEventMetadata(); + XStatsBuilder stats(event_metadata, &plane); + { + XEventBuilder event = step_line.AddEvent(*event_metadata); + event.SetOffsetPs(0); + event.SetDurationPs(100); + event.AddStatValue(group_id_stat, 1); + } + { + XEventBuilder event = step_line.AddEvent(*event_metadata); + event.SetOffsetPs(100); + event.SetDurationPs(100); + event.AddStatValue(group_id_stat, 2); + } + } + + StepEvents step_events = ConvertDeviceTraceXPlaneToStepEvents( + raw_plane, /*for_flat_profile=*/true); + EXPECT_EQ(step_events.size(), 2); + EXPECT_TRUE(step_events.contains(1)); + StepDetails step_1 = step_events[/*group_id=*/1]; + EXPECT_FALSE(step_1.PerCoreOpMetricsDb().contains(device_id)); + ASSERT_TRUE(step_1.PerCoreFlatOpMetricsDb().contains(device_id)); + + const auto& flat_db1 = step_1.PerCoreFlatOpMetricsDb().at(device_id); + EXPECT_EQ(flat_db1.op_instances_size(), 4); + + const auto* op1_g1 = FindFlatOp(flat_db1, "op1"); + ASSERT_NE(op1_g1, nullptr); + EXPECT_EQ(op1_g1->time_ps(), 50); + EXPECT_EQ(op1_g1->self_time_ps(), 30); // 50 - 20 (nested1) + + const auto* nested1_g1 = FindFlatOp(flat_db1, "nested1"); + ASSERT_NE(nested1_g1, nullptr); + EXPECT_EQ(nested1_g1->time_ps(), 20); + EXPECT_EQ(nested1_g1->self_time_ps(), 20); + + const auto* op2_g1 = FindFlatOp(flat_db1, "op2"); + ASSERT_NE(op2_g1, nullptr); + EXPECT_EQ(op2_g1->time_ps(), 50); + EXPECT_EQ(op2_g1->self_time_ps(), 20); // 50 - 30 (nested2) + + const auto* nested2_g1 = FindFlatOp(flat_db1, "nested2"); + ASSERT_NE(nested2_g1, nullptr); + EXPECT_EQ(nested2_g1->time_ps(), 30); + EXPECT_EQ(nested2_g1->self_time_ps(), 30); + + EXPECT_TRUE(step_events.contains(2)); + StepDetails step_2 = step_events[/*group_id=*/2]; + EXPECT_FALSE(step_2.PerCoreOpMetricsDb().contains(device_id)); + ASSERT_TRUE(step_2.PerCoreFlatOpMetricsDb().contains(device_id)); + + const auto& flat_db2 = step_2.PerCoreFlatOpMetricsDb().at(device_id); + EXPECT_EQ(flat_db2.op_instances_size(), 4); + + const auto* op1_g2 = FindFlatOp(flat_db2, "op1"); + ASSERT_NE(op1_g2, nullptr); + EXPECT_EQ(op1_g2->time_ps(), 50); + EXPECT_EQ(op1_g2->self_time_ps(), 20); // 50 - 30 (nested3) + + const auto* nested3_g2 = FindFlatOp(flat_db2, "nested3"); + ASSERT_NE(nested3_g2, nullptr); + EXPECT_EQ(nested3_g2->time_ps(), 30); + EXPECT_EQ(nested3_g2->self_time_ps(), 30); + + const auto* op3_g2 = FindFlatOp(flat_db2, "op3"); + ASSERT_NE(op3_g2, nullptr); + EXPECT_EQ(op3_g2->time_ps(), 50); + EXPECT_EQ(op3_g2->self_time_ps(), 20); // 50 - 30 (nested4) + + const auto* nested4_g2 = FindFlatOp(flat_db2, "nested4"); + ASSERT_NE(nested4_g2, nullptr); + EXPECT_EQ(nested4_g2->time_ps(), 30); + EXPECT_EQ(nested4_g2->self_time_ps(), 30); +} + TEST(ConvertXPlaneToStepEvents, SparseCoreShouldHaveStepMarkers) { XSpace space; XPlane* sc_plane = @@ -230,6 +458,8 @@ TEST(ConvertXPlaneToStepEvents, SparseCoreShouldHaveStepMarkers) { OpMetricsDb op_metrics_db = step_1.PerCoreOpMetricsDb()[/*core_id=*/1 + kSparseCoreIndexStart]; ASSERT_EQ(op_metrics_db.metrics_db_size(), 1); // Sparse op + EXPECT_FALSE(step_1.PerCoreFlatOpMetricsDb().contains( + /*core_id=*/1 + kSparseCoreIndexStart)); const OpMetrics* sparse_core_op = &op_metrics_db.metrics_db(0); EXPECT_EQ(sparse_core_op->name(), "sparse_op"); EXPECT_EQ(sparse_core_op->time_ps(), 880); @@ -237,6 +467,45 @@ TEST(ConvertXPlaneToStepEvents, SparseCoreShouldHaveStepMarkers) { EXPECT_EQ(sparse_core_op->hlo_module_id(), 1); } +TEST(ConvertXPlaneToStepEvents, SparseCoreShouldHaveStepMarkersFlat) { + XSpace space; + XPlane* sc_plane = + tsl::profiler::GetOrCreateTpuXPlane(&space, 0, "TPUv3", 0, 0, 1); + XPlaneBuilder sc_plane_builder(sc_plane); + sc_plane_builder.SetId(1); + XLineBuilder step_line = sc_plane_builder.GetOrCreateLine(0); + step_line.SetName(tsl::profiler::kSparseCoreStepLineName); + tsl::profiler::CreateXEventMetadata(&sc_plane_builder, "Step 1"); + tsl::profiler::CreateXEvent(&sc_plane_builder, &step_line, + /*event_name=*/"Step 1", + /*offset_ps=*/0, /*duration_ps=*/1000, + {{StatType::kGroupId, int64_t{1}}}); + XLineBuilder op_line = sc_plane_builder.GetOrCreateLine(1); + op_line.SetName(tsl::profiler::kSparseCoreOpLineName); + tsl::profiler::CreateXEventMetadata( + &sc_plane_builder, "sparse_op", + {{StatType::kProgramId, int64_t{1}}, {StatType::kSymbolId, int64_t{1}}}); + tsl::profiler::CreateXEvent(&sc_plane_builder, &op_line, + /*event_name=*/"sparse_op", + /*offset_ps=*/100, /*duration_ps=*/880, + {{StatType::kGroupId, int64_t{1}}}); + + StepEvents step_events = ConvertDeviceTraceXPlaneToStepEvents( + *sc_plane, /*for_flat_profile=*/true); + EXPECT_EQ(step_events.size(), 1); + EXPECT_TRUE(step_events.contains(1)); + StepDetails step_1 = step_events[/*group_id=*/1]; + EXPECT_EQ(step_1.Markers().size(), 1); + EXPECT_EQ(step_1.StepTime(), tsl::profiler::Timespan(0, 1000)); + EXPECT_FALSE(step_1.PerCoreOpMetricsDb().contains( + /*core_id=*/1 + kSparseCoreIndexStart)); + ASSERT_TRUE(step_1.PerCoreFlatOpMetricsDb().contains( + /*core_id=*/1 + kSparseCoreIndexStart)); + FlatOpMetricsDb flat_op_metrics_db = + step_1.PerCoreFlatOpMetricsDb().at(/*core_id=*/1 + kSparseCoreIndexStart); + ASSERT_EQ(flat_op_metrics_db.op_instances_size(), 1); +} + TEST(ConvertXPlaneToStepEvents, TpuDevicePlaneNoStepLine) { XPlane raw_plane; XPlaneBuilder plane(&raw_plane); @@ -299,7 +568,11 @@ TEST(ConvertXPlaneToStepEvents, TpuDevicePlaneNoStepLine) { } StepEvents step_events = ConvertDeviceTraceXPlaneToStepEvents(raw_plane); - EXPECT_EQ(step_events.size(), 0); + EXPECT_EQ(step_events.size(), 1); + const int64_t kIncompleteStepGroupId = 4294967294; + ASSERT_TRUE(step_events.contains(kIncompleteStepGroupId)); + const StepDetails& step_details = step_events.at(kIncompleteStepGroupId); + EXPECT_EQ(step_details.StepName(), "Incomplete Step"); } TEST(ConvertXPlaneToOpStats, CpuOnlyPyGrainSingleProcessInputPipelineTest) { diff --git a/xprof/utils/BUILD b/xprof/utils/BUILD index d08016899..5f505bfc0 100644 --- a/xprof/utils/BUILD +++ b/xprof/utils/BUILD @@ -49,6 +49,7 @@ cc_library( "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/log:check", "@com_google_absl//absl/strings", + "@org_xprof//plugin/xprof/protobuf:flat_op_metrics_proto_cc", "@org_xprof//plugin/xprof/protobuf:op_metrics_proto_cc", "@org_xprof//plugin/xprof/protobuf:steps_db_proto_cc", "@xla//xla/tsl/lib/gtl:map_util", @@ -57,6 +58,17 @@ cc_library( ], ) +cc_test( + name = "event_span_test", + srcs = ["event_span_test.cc"], + deps = [ + ":event_span", + "@com_google_googletest//:gtest", + "@com_google_googletest//:gtest_main", + "@org_xprof//plugin/xprof/protobuf:flat_op_metrics_proto_cc", + ], +) + cc_library( name = "hardware_type_utils", srcs = ["hardware_type_utils.cc"], @@ -126,11 +138,13 @@ cc_library( "@com_google_absl//absl/status:statusor", "@com_google_absl//absl/strings", "@org_xprof//plugin/xprof/protobuf:flat_op_metrics_proto_cc", + "@org_xprof//plugin/xprof/protobuf:op_metrics_proto_cc", "@org_xprof//plugin/xprof/protobuf:source_info_proto_cc", "@xla//xla/hlo/ir:hlo", "@xla//xla/tsl/platform:logging", "@xla//xla/tsl/platform:macros", "@xla//xla/tsl/profiler/utils:math_utils", + "@xla//xla/tsl/profiler/utils:tf_op_utils", "@xla//xla/tsl/profiler/utils:xplane_schema", "@xla//xla/tsl/profiler/utils:xplane_visitor", ], diff --git a/xprof/utils/event_span.cc b/xprof/utils/event_span.cc index 89ed9aa9b..1e2c607d6 100644 --- a/xprof/utils/event_span.cc +++ b/xprof/utils/event_span.cc @@ -380,6 +380,8 @@ StepDetails StepDetails::ToNonOverlapped() const { device_memory_transfers_; non_overlapped_step_details.step_name_ = step_name_; non_overlapped_step_details.per_core_op_metrics_db_ = per_core_op_metrics_db_; + non_overlapped_step_details.per_core_flat_op_metrics_db_ = + per_core_flat_op_metrics_db_; return non_overlapped_step_details; } @@ -391,6 +393,10 @@ void StepDetails::Combine(const StepDetails& other) { for (const auto& [core_id, op_metric_db] : other.per_core_op_metrics_db_) { per_core_op_metrics_db_[core_id] = op_metric_db; } + for (const auto& [core_id, flat_op_metric_db] : + other.per_core_flat_op_metrics_db_) { + per_core_flat_op_metrics_db_[core_id] = flat_op_metric_db; + } if (step_name_.empty()) step_name_ = other.step_name_; } diff --git a/xprof/utils/event_span.h b/xprof/utils/event_span.h index bc431aacc..32ea74859 100644 --- a/xprof/utils/event_span.h +++ b/xprof/utils/event_span.h @@ -19,12 +19,14 @@ limitations under the License. #include #include #include +#include #include #include "absl/container/flat_hash_map.h" #include "absl/strings/string_view.h" #include "xla/tsl/platform/types.h" #include "xla/tsl/profiler/utils/timespan.h" +#include "plugin/xprof/protobuf/flat_op_metrics.pb.h" #include "plugin/xprof/protobuf/op_metrics.pb.h" #include "plugin/xprof/protobuf/steps_db.pb.h" @@ -170,6 +172,10 @@ class StepDetails { return device_memory_transfers_; } + const absl::flat_hash_map& PerCoreOpMetricsDb() const { + return per_core_op_metrics_db_; + } + absl::flat_hash_map& PerCoreOpMetricsDb() { return per_core_op_metrics_db_; } @@ -191,7 +197,7 @@ class StepDetails { // Returns the step name. std::string StepName() const { return step_name_; } // Sets the name of this step. - void SetStepName(std::string step_name) { step_name_ = step_name; } + void SetStepName(std::string step_name) { step_name_ = std::move(step_name); } // Converts from overlapped events to non-overlapped events. StepDetails ToNonOverlapped() const; @@ -208,7 +214,20 @@ class StepDetails { std::string DebugString() const; void SetPerCoreOpMetricsDb(OpMetricsDb db, uint32_t core_id) { - per_core_op_metrics_db_[core_id] = db; + per_core_op_metrics_db_[core_id] = std::move(db); + } + + const absl::flat_hash_map& PerCoreFlatOpMetricsDb() + const { + return per_core_flat_op_metrics_db_; + } + + absl::flat_hash_map& PerCoreFlatOpMetricsDb() { + return per_core_flat_op_metrics_db_; + } + + void SetPerCoreFlatOpMetricsDb(FlatOpMetricsDb db, uint32_t core_id) { + per_core_flat_op_metrics_db_[core_id] = std::move(db); } private: @@ -233,6 +252,7 @@ class StepDetails { std::string step_name_; absl::flat_hash_map per_core_op_metrics_db_; + absl::flat_hash_map per_core_flat_op_metrics_db_; }; // Map from step_id to the events happened in that step. diff --git a/xprof/utils/event_span_test.cc b/xprof/utils/event_span_test.cc new file mode 100644 index 000000000..6e1deae6f --- /dev/null +++ b/xprof/utils/event_span_test.cc @@ -0,0 +1,47 @@ +#include "xprof/utils/event_span.h" + +#include "testing/base/public/gmock.h" +#include "" +#include "plugin/xprof/protobuf/flat_op_metrics.pb.h" + +namespace tensorflow { +namespace profiler { +namespace { + +using ::testing::EqualsProto; + +TEST(StepDetailsTest, CombineFlatOpMetricsDb) { + StepDetails step1; + FlatOpMetricsDb db1; + db1.set_total_op_time_ps(100); + step1.SetPerCoreFlatOpMetricsDb(db1, 1); + + StepDetails step2; + FlatOpMetricsDb db2; + db2.set_total_op_time_ps(200); + step2.SetPerCoreFlatOpMetricsDb(db2, 2); + + step1.Combine(step2); + + auto& combined_dbs = step1.PerCoreFlatOpMetricsDb(); + ASSERT_EQ(combined_dbs.size(), 2); + EXPECT_THAT(combined_dbs.at(1), EqualsProto(db1)); + EXPECT_THAT(combined_dbs.at(2), EqualsProto(db2)); +} + +TEST(StepDetailsTest, ToNonOverlappedStepDetailsCopiesFlatOpMetricsDb) { + StepDetails step; + FlatOpMetricsDb db; + db.set_total_op_time_ps(100); + step.SetPerCoreFlatOpMetricsDb(db, 1); + + StepDetails non_overlapped = step.ToNonOverlapped(); + + auto& copied_dbs = non_overlapped.PerCoreFlatOpMetricsDb(); + ASSERT_EQ(copied_dbs.size(), 1); + EXPECT_THAT(copied_dbs.at(1), EqualsProto(db)); +} + +} // namespace +} // namespace profiler +} // namespace tensorflow diff --git a/xprof/utils/flat_op_metrics_db_utils.cc b/xprof/utils/flat_op_metrics_db_utils.cc index 12e010313..b217e03e5 100644 --- a/xprof/utils/flat_op_metrics_db_utils.cc +++ b/xprof/utils/flat_op_metrics_db_utils.cc @@ -36,6 +36,7 @@ limitations under the License. #include "xla/hlo/ir/hlo_opcode.h" #include "xla/tsl/platform/logging.h" #include "xla/tsl/profiler/utils/math_utils.h" +#include "xla/tsl/profiler/utils/tf_op_utils.h" #include "xla/tsl/profiler/utils/xplane_schema.h" #include "xla/tsl/profiler/utils/xplane_visitor.h" #include "plugin/xprof/protobuf/flat_op_metrics.pb.h" @@ -79,8 +80,8 @@ ExtractSourceFileNameAndLineNumber(absl::string_view source_top_line) { // `:`. If the `source_top_line` is not in // the expected format, then `source_info` will not be populated.` void PopulateSourceInfo(absl::string_view source_top_line, - tensorflow::profiler::SourceInfo& source_info) { - const auto file_and_line = + SourceInfo& source_info) { + auto file_and_line = ExtractSourceFileNameAndLineNumber(source_top_line); if (file_and_line.ok()) { source_info.set_file_name(std::move(file_and_line->first)); @@ -88,10 +89,13 @@ void PopulateSourceInfo(absl::string_view source_top_line, } else { LOG(ERROR) << "Failed to extract source filename and line from the input " "source_top_line: '" + << source_top_line << "' with status: " << file_and_line.status(); } } + + // Extracts metadata from an XEventMetadataVisitor and populates the // corresponding fields in a FlatOpMetrics protobuf message. // This function handles names, categories, and various statistics associated @@ -419,15 +423,34 @@ FlatOpMetrics XEventsFlatOpMetricsDbBuilder::FromXEvent( return op_metrics; } +XEventsFlatOpMetricsDbBuilder::OpKey +XEventsFlatOpMetricsDbBuilder::GetFlatOpKeyFromXEvent( + const tsl::profiler::XEventVisitor& event) { + XEventsFlatOpMetricsDbBuilder::OpKey op_key; + if (event.metadata() == nullptr) return op_key; + event.Metadata().ForEachStat([&](const XStatVisitor& stat) { + if (stat.Type().has_value()) { + switch (static_cast(*stat.Type())) { + case StatType::kProgramId: + op_key.program_id = stat.IntOrUintValue(); + break; + case StatType::kSymbolId: + op_key.symbol_id = stat.IntOrUintValue(); + break; + default: + break; + } + } + }); + return op_key; +} + // Adds an operation metric from a tsl::profiler::XEventVisitor to the builder. // It extracts the key and aggregates the metrics. void XEventsFlatOpMetricsDbBuilder::AddOpMetric( const tsl::profiler::XEventVisitor& event) { - auto key = GetOpKeyFromXEvent(event); - XEventsFlatOpMetricsDbBuilder::OpKey flat_key; - flat_key.program_id = key.program_id; - flat_key.symbol_id = key.symbol_id; - AddOpMetric(XEventsFlatOpMetricsDbBuilder::FromXEvent(event), flat_key); + AddOpMetric(XEventsFlatOpMetricsDbBuilder::FromXEvent(event), + GetFlatOpKeyFromXEvent(event)); } // Adds a FlatOpMetrics message with a specific OpKey to the builder. @@ -565,5 +588,76 @@ void SetIdleOp(uint64_t idle_time_ps, FlatOpMetrics& idle_op) { idle_op.set_self_time_ps(idle_time_ps); } +class DeviceTfFlatOpMetricsDbBuilder : public FlatOpMetricsDbBuilder { + public: + explicit DeviceTfFlatOpMetricsDbBuilder(FlatOpMetricsDb* db) + : FlatOpMetricsDbBuilder(db) {} + + void UpdateTfOpMetricsWithDeviceOpMetrics( + absl::string_view tf_op_name, absl::string_view tf_op_type, + const FlatOpMetrics& device_op_metrics) { + FlatOpMetrics* tf_op_metrics = + FlatOpMetricsDbBuilder::LookupOrInsertNewFlatOpMetrics( + /*hlo_module_id=*/0, tf_op_name); + if (tf_op_metrics->category().empty()) { + tf_op_metrics->set_category(tf_op_type == tsl::profiler::kUnknownOp + ? "Unknown" + : std::string(tf_op_type)); + } + tf_op_metrics->set_is_eager(device_op_metrics.is_eager()); + // The occurrences of a TF-op is the maximum among the occurrences of all + // device ops that it contains. + tf_op_metrics->set_occurrences( + std::max(tf_op_metrics->occurrences(), + device_op_metrics.occurrences())); + tf_op_metrics->set_time_ps(tf_op_metrics->time_ps() + + device_op_metrics.time_ps()); + tf_op_metrics->set_self_time_ps(tf_op_metrics->self_time_ps() + + device_op_metrics.self_time_ps()); + tf_op_metrics->set_flops(tf_op_metrics->flops() + + device_op_metrics.flops()); + tf_op_metrics->set_flops_v2(tf_op_metrics->flops_v2() + + device_op_metrics.flops_v2()); + tf_op_metrics->set_model_flops(tf_op_metrics->model_flops() + + device_op_metrics.model_flops()); + tf_op_metrics->set_model_flops_v2(tf_op_metrics->model_flops_v2() + + device_op_metrics.model_flops_v2()); + tf_op_metrics->set_bytes_accessed(tf_op_metrics->bytes_accessed() + + device_op_metrics.bytes_accessed()); + tf_op_metrics->set_core_type(device_op_metrics.core_type()); + } +}; + +FlatOpMetricsDb CreateTfMetricsDbFromDeviceOpMetricsDb( + const FlatOpMetricsDb& device_op_metrics_db, bool with_idle) { + FlatOpMetricsDb tf_op_metrics_db; + DeviceTfFlatOpMetricsDbBuilder builder(&tf_op_metrics_db); + for (const auto& device_op_metrics : device_op_metrics_db.op_instances()) { + if (device_op_metrics.category() == kIdle) { + if (with_idle) { + builder.UpdateTfOpMetricsWithDeviceOpMetrics(kIdle, kIdle, + device_op_metrics); + } + } else if (device_op_metrics.provenance().empty()) { + builder.UpdateTfOpMetricsWithDeviceOpMetrics(device_op_metrics.hlo_name(), + tsl::profiler::kUnknownOp, + device_op_metrics); + } else { + tsl::profiler::TfOp tf_op = + tsl::profiler::ParseTfOpFullname(device_op_metrics.provenance()); + builder.UpdateTfOpMetricsWithDeviceOpMetrics(tf_op.name, tf_op.type, + device_op_metrics); + } + } + tf_op_metrics_db.set_total_op_time_ps( + device_op_metrics_db.total_op_time_ps()); + + tf_op_metrics_db.set_total_time_ps( + with_idle ? device_op_metrics_db.total_time_ps() + : device_op_metrics_db.total_op_time_ps()); + + return tf_op_metrics_db; +} + } // namespace profiler } // namespace tensorflow diff --git a/xprof/utils/flat_op_metrics_db_utils.h b/xprof/utils/flat_op_metrics_db_utils.h index 6eebac4b0..5c795a1c5 100644 --- a/xprof/utils/flat_op_metrics_db_utils.h +++ b/xprof/utils/flat_op_metrics_db_utils.h @@ -26,6 +26,7 @@ limitations under the License. #include "xla/tsl/platform/macros.h" #include "xla/tsl/profiler/utils/xplane_visitor.h" #include "plugin/xprof/protobuf/flat_op_metrics.pb.h" +#include "plugin/xprof/protobuf/op_metrics.pb.h" namespace tensorflow { namespace profiler { @@ -70,6 +71,12 @@ class XEventsFlatOpMetricsDbBuilder { // Constructs a FlatOpMetrics from the provided XEventVisitor. static FlatOpMetrics FromXEvent(const tsl::profiler::XEventVisitor& xevent); + // Returns the OpKey (Program ID and Symbol ID) for the provided + // XEventVisitor. + // This key is used to group and aggregate FlatOpMetrics in the database. + static OpKey GetFlatOpKeyFromXEvent( + const tsl::profiler::XEventVisitor& event); + // Add FlatOpMetrics from XEventVisitor. void AddOpMetric(const tsl::profiler::XEventVisitor& xevent); @@ -151,6 +158,11 @@ uint64_t IdleTimePs(const FlatOpMetricsDb& db); // time spent without any op execution. void SetIdleOp(uint64_t idle_time_ps, FlatOpMetrics& idle_op); + +// Converts from the device flat op metrics to Tf-op metrics. +FlatOpMetricsDb CreateTfMetricsDbFromDeviceOpMetricsDb( + const FlatOpMetricsDb& device_op_metrics_db, bool with_idle = true); + } // namespace profiler } // namespace tensorflow diff --git a/xprof/utils/flat_op_metrics_db_utils_test.cc b/xprof/utils/flat_op_metrics_db_utils_test.cc index 46344ff59..9f42b013c 100644 --- a/xprof/utils/flat_op_metrics_db_utils_test.cc +++ b/xprof/utils/flat_op_metrics_db_utils_test.cc @@ -126,6 +126,31 @@ TEST_F(XEventsFlatOpMetricsDbBuilderTest, BasicEventProcessing) { EXPECT_DOUBLE_EQ(metric.self_time_ps(), 500); } +// ----------------------------------------------------------------------------- +// Test Case: GetFlatOpKeyFromXEvent +// ----------------------------------------------------------------------------- +TEST_F(XEventsFlatOpMetricsDbBuilderTest, GetFlatOpKeyFromXEvent) { + XPlane plane = CreateTestXPlane(); + XPlaneBuilder plane_builder(&plane); + XLineBuilder line_builder = plane_builder.GetOrCreateLine(0); + + // Create event with program_id=123, symbol_id=456 + CreateXEvent(plane_builder, line_builder, "op1", 1000, 500, 123, 456); + + XPlaneVisitor plane_visitor(&plane, {}, {tsl::profiler::FindStatType}); + XEventVisitor event_visitor(&plane_visitor, &plane.lines(0), + &plane.lines(0).events(0)); + + XEventsFlatOpMetricsDbBuilder::OpKey key = + XEventsFlatOpMetricsDbBuilder::GetFlatOpKeyFromXEvent(event_visitor); + + ASSERT_TRUE(key.program_id.has_value()); + EXPECT_EQ(*key.program_id, 123); + ASSERT_TRUE(key.symbol_id.has_value()); + EXPECT_EQ(*key.symbol_id, 456); +} + + // ----------------------------------------------------------------------------- // Test Case 2: Aggregation of Multiple Instances // ----------------------------------------------------------------------------- @@ -717,6 +742,116 @@ TEST_F(XEventsFlatOpMetricsDbBuilderTest, SourceInformationParsing) { EXPECT_EQ(metric.source_info().stack_frame(), "stack_frame_contents"); } +// ----------------------------------------------------------------------------- +// Tests for CreateTfMetricsDbFromDeviceOpMetricsDb +// ----------------------------------------------------------------------------- + +TEST(CreateTfMetricsDbFromDeviceOpMetricsDbTest, AggregationAndIdle) { + FlatOpMetricsDb device_db; + device_db.set_total_op_time_ps(1000); + device_db.set_total_time_ps(1500); + + // Op 1 with provenance + auto* op1 = device_db.add_op_instances(); + op1->set_hlo_name("hlo1"); + op1->set_provenance("TfOp1:TfOpType1"); + op1->set_time_ps(500); + op1->set_self_time_ps(400); + op1->set_flops(100); + op1->set_flops_v2(101); + op1->set_model_flops(102); + op1->set_model_flops_v2(103); + op1->set_bytes_accessed(1000); + op1->set_occurrences(1); + + // Op 2 with same provenance + auto* op2 = device_db.add_op_instances(); + op2->set_hlo_name("hlo2"); + op2->set_provenance("TfOp1:TfOpType1"); + op2->set_time_ps(300); + op2->set_self_time_ps(300); + op2->set_flops(50); + op2->set_flops_v2(51); + op2->set_model_flops(52); + op2->set_model_flops_v2(53); + op2->set_bytes_accessed(500); + op2->set_occurrences(2); + + // Op 3 with empty provenance + auto* op3 = device_db.add_op_instances(); + op3->set_hlo_name("hlo3"); + op3->set_time_ps(200); + op3->set_self_time_ps(200); + op3->set_flops(20); + op3->set_bytes_accessed(200); + op3->set_occurrences(1); + + // Idle Op + auto* idle_op = device_db.add_op_instances(); + idle_op->set_category(kIdle); + idle_op->set_time_ps(500); + idle_op->set_self_time_ps(500); + + // Test with_idle = true + { + FlatOpMetricsDb tf_db = + CreateTfMetricsDbFromDeviceOpMetricsDb(device_db, /*with_idle=*/true); + + EXPECT_EQ(tf_db.total_op_time_ps(), 1000); + EXPECT_EQ(tf_db.total_time_ps(), 1500); + ASSERT_EQ(tf_db.op_instances_size(), 3); + + // Find ops + const FlatOpMetrics* tf_op1 = nullptr; + const FlatOpMetrics* tf_op3 = nullptr; + const FlatOpMetrics* tf_idle = nullptr; + + for (const auto& op : tf_db.op_instances()) { + if (op.hlo_name() == "TfOp1") + tf_op1 = &op; + else if (op.hlo_name() == "hlo3") + tf_op3 = &op; + else if (op.hlo_name() == kIdle) + tf_idle = &op; + } + + ASSERT_NE(tf_op1, nullptr); + EXPECT_EQ(tf_op1->category(), "TfOpType1"); + EXPECT_EQ(tf_op1->time_ps(), 800); // 500 + 300 + EXPECT_EQ(tf_op1->self_time_ps(), 700); // 400 + 300 + EXPECT_EQ(tf_op1->flops(), 150); // 100 + 50 + EXPECT_EQ(tf_op1->flops_v2(), 152); // 101 + 51 + EXPECT_EQ(tf_op1->model_flops(), 154); // 102 + 52 + EXPECT_EQ(tf_op1->model_flops_v2(), 156); // 103 + 53 + EXPECT_EQ(tf_op1->bytes_accessed(), 1500); // 1000 + 500 + EXPECT_EQ(tf_op1->occurrences(), 2); // max(1, 2) + + ASSERT_NE(tf_op3, nullptr); + EXPECT_EQ(tf_op3->category(), "Unknown"); + EXPECT_EQ(tf_op3->time_ps(), 200); + + ASSERT_NE(tf_idle, nullptr); + EXPECT_EQ(tf_idle->category(), kIdle); + EXPECT_EQ(tf_idle->time_ps(), 500); + } + + // Test with_idle = false + { + FlatOpMetricsDb tf_db = + CreateTfMetricsDbFromDeviceOpMetricsDb(device_db, /*with_idle=*/false); + + EXPECT_EQ(tf_db.total_op_time_ps(), 1000); + EXPECT_EQ(tf_db.total_time_ps(), 1000); // Should be total_op_time_ps + ASSERT_EQ(tf_db.op_instances_size(), 2); // Idle should be skipped + + const FlatOpMetrics* tf_idle = nullptr; + for (const auto& op : tf_db.op_instances()) { + if (op.hlo_name() == kIdle) tf_idle = &op; + } + EXPECT_EQ(tf_idle, nullptr); + } +} + } // namespace } // namespace profiler } // namespace tensorflow