Skip to content

Commit

Permalink
[XLA:MSA] Fixes a bug in GetInefficientAllocationSites(allocation_val…
Browse files Browse the repository at this point in the history
…ues). The function was previously assuming allocation_values can never be empty.

PiperOrigin-RevId: 693424356
  • Loading branch information
mehrdadkhani authored and Google-ML-Automation committed Nov 15, 2024
1 parent 58b5afb commit 8d5fd3a
Show file tree
Hide file tree
Showing 8 changed files with 741 additions and 398 deletions.
16 changes: 15 additions & 1 deletion xla/service/memory_space_assignment/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,7 @@ xla_cc_test(
deps = [
":algorithm",
":allocation",
":allocation_value",
":buffer_interval_comparator",
":cost_analysis",
":memory_space_assignment",
Expand Down Expand Up @@ -107,7 +108,6 @@ xla_cc_test(
"//xla/service/heap_simulator:allocation_block",
"//xla/tests:hlo_test_base",
"//xla/tests:test_utils",
"//xla/tests:verified_hlo_module",
"//xla/tests:xla_internal_test_main",
"//xla/tsl/lib/core:status_test_util",
"@com_google_absl//absl/algorithm:container",
Expand Down Expand Up @@ -273,6 +273,7 @@ cc_library(
srcs = [],
hdrs = ["options.h"],
deps = [
":allocation_value",
":buffer_interval_comparator",
":cost_analysis",
":memory_space_assignment_proto_cc",
Expand Down Expand Up @@ -508,6 +509,7 @@ cc_library(
hdrs = ["algorithm.h"],
deps = [
":allocation",
":allocation_value",
":buffer_interval_comparator",
":cost_analysis",
":memory_bound_loop_optimizer",
Expand Down Expand Up @@ -577,6 +579,18 @@ cc_library(
],
)

cc_library(
name = "allocation_value",
hdrs = ["allocation_value.h"],
deps = [
":allocation",
"//xla/hlo/ir:hlo",
"//xla/service:hlo_value",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/types:span",
],
)

xla_cc_test(
name = "prefetch_interval_picker_test",
srcs = ["prefetch_interval_picker_test.cc"],
Expand Down
218 changes: 117 additions & 101 deletions xla/service/memory_space_assignment/algorithm.cc

Large diffs are not rendered by default.

Loading

0 comments on commit 8d5fd3a

Please sign in to comment.