OSSForks/tensorflow
1
0
Fork 0
mirror of https://github.com/zebrajr/tensorflow.git synced 2025-12-06 12:20:11 +01:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

Make sure to produce a deterministic Memory usage report.

Browse Source 
PiperOrigin-RevId: 824403476
This commit is contained in:
Adrian Kuegel 2025-10-27 01:32:04 -07:00 committed by TensorFlower Gardener
parent 28054871f6
commit 5d42d91467
2 changed files with 42 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

47
third_party/xla/xla/service/buffer_assignment.cc vendored
View File

@ -990,13 +990,21 @@ std::string BufferAllocation::MemoryUsageReport(const std::string& prefix,
// Group the values by their offset in the allocation.
absl::flat_hash_map<int64_t, OffsetInfo> offset_to_buffers;
std::vector<const HloValue*> sorted_hlo_values;
sorted_hlo_values.reserve(assigned_buffers_.size());
for (const auto& element : assigned_buffers_) {
const HloValue* value = element.first;
OffsetInfo& offset_info = offset_to_buffers[element.second.offset];
sorted_hlo_values.push_back(element.first);
}
absl::c_sort(sorted_hlo_values, [](const HloValue* a, const HloValue* b) {
return a->id() < b->id();
});
for (const HloValue* value : sorted_hlo_values) {
const OffsetSize& offset_size = assigned_buffers_.find(value)->second;
OffsetInfo& offset_info = offset_to_buffers[offset_size.offset];
offset_info.values.push_back(value);
offset_info.offset_size.offset = element.second.offset;
offset_info.offset_size.offset = offset_size.offset;
offset_info.offset_size.size =
std::max(offset_info.offset_size.size, element.second.size);
std::max(offset_info.offset_size.size, offset_size.size);
}
// Sort the offset infos by the max size of the values in the group.
@ -1008,7 +1016,14 @@ std::string BufferAllocation::MemoryUsageReport(const std::string& prefix,
}
absl::c_sort(sorted_offset_infos,
[](const OffsetInfo& a, const OffsetInfo& b) {
return a.offset_size.size > b.offset_size.size;
if (a.offset_size.size != b.offset_size.size) {
return a.offset_size.size > b.offset_size.size;
}
// Each HloValue appears in just one OffsetInfo `values`
// vector. Therefore we can use the id of any element of the
// `values` vector as tie breaker, as long as the `values`
// vector has a deterministic order.
return a.values.back()->id() < b.values.back()->id();
});
StrAppend(&output, prefix,
@ -1030,15 +1045,21 @@ std::string BufferAllocation::MemoryUsageReport(const std::string& prefix,
// of the line.
absl::flat_hash_map<std::string, int64_t> shapes;
for (auto& value : offset_info.values) shapes[value->shape().ToString()]++;
std::vector<std::pair<int64_t, std::string>> sorted_shapes;
sorted_shapes.reserve(shapes.size());
for (const auto& entry : shapes) {
sorted_shapes.emplace_back(entry.second, entry.first);
}
absl::c_sort(sorted_shapes);
StrAppend(
&output,
absl::StrJoin(shapes, ", ", [](std::string* out, const auto& pair) {
if (pair.second == 1) {
return absl::StrAppend(out, pair.first);
}
return absl::StrAppend(out, pair.second, "×", pair.first);
}));
StrAppend(&output,
absl::StrJoin(
sorted_shapes, ", ", [](std::string* out, const auto& pair) {
if (pair.first == 1) {
return absl::StrAppend(out, pair.second);
}
return absl::StrAppend(out, pair.first, "×", pair.second);
}));
StrAppend(&output, "\n");

8
third_party/xla/xla/service/buffer_assignment_test.cc vendored
View File

@ -2215,6 +2215,14 @@ TEST_F(BufferAssignmentTest, PeakBuffers) {
EXPECT_FALSE(buffer.IsInputOrOutput());
EXPECT_TRUE(buffer.IsPreallocatedTempBuffer());
ASSERT_EQ(buffer.assigned_buffers().size(), 4);
const char* const kExpectedMemoryUsageReport =
R"(cumulative_size; size; offset; used_by_n_values; shapes_list
------------------------------------------------------------
800B( 50%); 800B; 0; 2; f32[100], f32[200]
1.2KiB( 75%); 400B; 800; 1; f32[100]
1.6KiB(100%); 400B; 1200; 1; f32[100]
)";
EXPECT_EQ(buffer.MemoryUsageReport(""), kExpectedMemoryUsageReport);
const std::vector<const HloValue*>& peak_buffers =
buffer.PeakMemoryLogicalBuffers();