Skip to content

Commit

Permalink
Split out MemorySpaceAssignmentTest class for re-use.
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 688294711
  • Loading branch information
Google-ML-Automation committed Dec 17, 2024
1 parent d80d15e commit 9b4d49e
Show file tree
Hide file tree
Showing 3 changed files with 477 additions and 397 deletions.
29 changes: 29 additions & 0 deletions xla/service/memory_space_assignment/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,7 @@ xla_cc_test(
":cost_analysis",
":memory_space_assignment",
":memory_space_assignment_proto_cc",
":memory_space_assignment_test_base",
":options",
":prefetch_interval_picker",
":repacking",
Expand Down Expand Up @@ -160,6 +161,34 @@ cc_library(
],
)

cc_library(
name = "memory_space_assignment_test_base",
testonly = True,
hdrs = ["memory_space_assignment_test_base.h"],
deps = [
":buffer_interval_comparator",
":cost_analysis",
":memory_space_assignment",
":options",
":prefetch_interval_picker",
"//xla:shape_util",
"//xla/hlo/analysis:hlo_alias_analysis",
"//xla/hlo/ir:hlo",
"//xla/hlo/transforms:instruction_hoister",
"//xla/hlo/utils:hlo_live_range",
"//xla/service:buffer_value",
"//xla/service:hlo_buffer",
"//xla/service:hlo_cost_analysis",
"//xla/service:hlo_value",
"//xla/tests:hlo_test_base",
"//xla/tsl/lib/core:status_test_util",
"@com_google_absl//absl/status:statusor",
"@tsl//tsl/platform:status",
"@tsl//tsl/platform:statusor",
"@tsl//tsl/platform:test",
],
)

cc_library(
name = "utils",
srcs = ["utils.cc"],
Expand Down
Loading

0 comments on commit 9b4d49e

Please sign in to comment.