Skip to content

Commit b599ec1

Browse files
tensorflower-gardenertensorflower-gardener
committed
avoid removing some copy kernels.
PiperOrigin-RevId: 785998543
1 parent a88b21e commit b599ec1

File tree

7 files changed

+723
-2
lines changed

7 files changed

+723
-2
lines changed

‎third_party/xla/xla/hlo/utils/BUILD

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ cc_library(
5252
srcs = ["hlo_live_range.cc"],
5353
hdrs = ["hlo_live_range.h"],
5454
deps = [
55+
":hlo_stack_trace",
5556
"//xla:shape_util",
5657
"//xla/hlo/analysis:hlo_alias_analysis",
5758
"//xla/hlo/ir:hlo",
@@ -62,6 +63,7 @@ cc_library(
6263
"@com_google_absl//absl/log",
6364
"@com_google_absl//absl/log:check",
6465
"@com_google_absl//absl/status:statusor",
66+
"@com_google_absl//absl/strings",
6567
"@com_google_absl//absl/strings:str_format",
6668
"@com_google_absl//absl/types:span",
6769
"@local_tsl//tsl/platform:logging",
@@ -83,9 +85,9 @@ xla_cc_test(
8385
"//xla/hlo/testlib:hlo_hardware_independent_test_base",
8486
"//xla/service:hlo_value",
8587
"//xla/tsl/lib/core:status_test_util",
88+
"//xla/tsl/platform:statusor",
8689
"@com_google_absl//absl/container:flat_hash_map",
8790
"@com_google_googletest//:gtest_main",
88-
"@local_tsl//tsl/platform:statusor",
8991
],
9092
)
9193

@@ -194,6 +196,35 @@ xla_cc_test(
194196
],
195197
)
196198

199+
cc_library(
200+
name = "hlo_stack_trace",
201+
srcs = ["hlo_stack_trace.cc"],
202+
hdrs = ["hlo_stack_trace.h"],
203+
deps = [
204+
"//xla/hlo/ir:hlo",
205+
"//xla/service:hlo_value",
206+
"@com_google_absl//absl/container:flat_hash_map",
207+
"@com_google_absl//absl/container:inlined_vector",
208+
"@com_google_absl//absl/strings",
209+
"@com_google_absl//absl/strings:str_format",
210+
"@com_google_absl//absl/strings:string_view",
211+
"@com_googlesource_code_re2//:re2",
212+
],
213+
)
214+
215+
xla_cc_test(
216+
name = "hlo_stack_trace_test",
217+
srcs = ["hlo_stack_trace_test.cc"],
218+
deps = [
219+
":hlo_stack_trace",
220+
"//xla:shape_util",
221+
"//xla/hlo/analysis:hlo_dataflow_analysis",
222+
"//xla/hlo/testlib:hlo_hardware_independent_test_base",
223+
"//xla/service:hlo_value",
224+
"@com_google_googletest//:gtest_main",
225+
],
226+
)
227+
197228
cc_library(
198229
name = "hlo_query",
199230
srcs = ["hlo_query.cc"],

‎third_party/xla/xla/hlo/utils/hlo_live_range.cc

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,13 +29,15 @@ limitations under the License.
2929
#include "absl/log/check.h"
3030
#include "absl/log/log.h"
3131
#include "absl/status/statusor.h"
32+
#include "absl/strings/str_cat.h"
3233
#include "absl/strings/str_format.h"
3334
#include "absl/types/span.h"
3435
#include "xla/hlo/analysis/hlo_alias_analysis.h"
3536
#include "xla/hlo/ir/dfs_hlo_visitor.h"
3637
#include "xla/hlo/ir/hlo_computation.h"
3738
#include "xla/hlo/ir/hlo_opcode.h"
3839
#include "xla/hlo/ir/hlo_schedule.h"
40+
#include "xla/hlo/utils/hlo_stack_trace.h"
3941
#include "xla/service/hlo_buffer.h"
4042
#include "xla/service/hlo_value.h"
4143
#include "xla/shape_util.h"
@@ -351,6 +353,12 @@ std::string HloLiveRange::ToString() const {
351353
value->index().ToString(), bytes, total_bytes);
352354
}
353355

356+
const HloModule* module =
357+
instructions.empty() ? nullptr : instructions[0]->GetModule();
358+
if (module != nullptr) {
359+
absl::StrAppend(&output, FormatStackTraceBreakdown(sized_buffers, module));
360+
}
361+
354362
return output;
355363
}
356364

‎third_party/xla/xla/hlo/utils/hlo_live_range_test.cc

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,8 +37,8 @@ limitations under the License.
3737
#include "xla/shape.h"
3838
#include "xla/shape_util.h"
3939
#include "xla/tsl/lib/core/status_test_util.h"
40+
#include "xla/tsl/platform/statusor.h"
4041
#include "xla/xla_data.pb.h"
41-
#include "tsl/platform/statusor.h"
4242

4343
namespace xla {
4444
namespace {
@@ -555,6 +555,11 @@ TEST_F(HloLiveRangeTest, ToString) {
555555
multiply{}: 16 bytes (cumulative: 16 bytes)
556556
paramA{}: 16 bytes (cumulative: 32 bytes)
557557
paramX{}: 16 bytes (cumulative: 48 bytes)
558+
Stack trace breakdown for peak usage: 48 bytes
559+
ToString (100.0%, total: 48 bytes, current: 0 bytes, remaining: 48 bytes)
560+
├── multiply (33.3%, total: 16 bytes, current: 16 bytes, remaining: 32 bytes)
561+
├── paramA (33.3%, total: 16 bytes, current: 16 bytes, remaining: 16 bytes)
562+
└── paramX (33.3%, total: 16 bytes, current: 16 bytes, remaining: 0 bytes)
558563
)";
559564
EXPECT_EQ(hlo_live_range_->ToString(), expected_string);
560565
}
@@ -599,6 +604,12 @@ TEST_F(HloLiveRangeTest, ToStringTuple) {
599604
constant{1}: 16 bytes (cumulative: 32 bytes)
600605
paramA{}: 16 bytes (cumulative: 48 bytes)
601606
constant{0}: 4 bytes (cumulative: 52 bytes)
607+
Stack trace breakdown for peak usage: 52 bytes
608+
ToStringTuple (100.0%, total: 52 bytes, current: 0 bytes, remaining: 52 bytes)
609+
├── constant (30.8%, total: 16 bytes, current: 16 bytes, remaining: 36 bytes)
610+
├── constant{1} (30.8%, total: 16 bytes, current: 16 bytes, remaining: 20 bytes)
611+
├── paramA (30.8%, total: 16 bytes, current: 16 bytes, remaining: 4 bytes)
612+
└── constant{0} (7.7%, total: 4 bytes, current: 4 bytes, remaining: 0 bytes)
602613
)";
603614
EXPECT_EQ(hlo_live_range_->ToString(), expected_string);
604615
}

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy