experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[XLA] Add mechanism to disable alternative mem allocation at uses

Browse Source 
PiperOrigin-RevId: 347059810
Change-Id: Ie625eecb08ca195163a11e5b65ec776f9e8e2036
This commit is contained in:
Berkin Ilbeyi 2020-12-11 13:07:59 -08:00 committed by TensorFlower Gardener
parent 4a8d9cf405
commit df3f233362
2 changed files with 10 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorflow/compiler/xla/service/memory_space_assignment.cc
View File

@ -893,6 +893,9 @@ AlternateMemoryBestFitHeap::GetSortedColocatedIntervals(
bool AlternateMemoryBestFitHeap::IsUseAllowedInAlternateMemory(
const AllocationValue& value, const HloUse& use) const {
const auto& instruction_schedule = hlo_live_range_.instruction_schedule();
if (!options_.is_use_allowed_in_alternate_mem_fn(use)) {
return false;
}
if (use.instruction->opcode() == HloOpcode::kWhile) {
HloComputation* while_body = use.instruction->while_body();

7
tensorflow/compiler/xla/service/memory_space_assignment.h
View File

@ -400,6 +400,8 @@ class MemorySpaceAssignment {
GlobalDecreasingSizeBestFitHeap<HloValue>::BufferIntervalCompare;
using IsAllowedInAlternateMemoryFunction =
std::function<bool(const HloValue&)>;
using IsUseAllowedInAlternateMemoryFunction =
std::function<bool(const HloUse&)>;
// MemorySpaceAssignment uses a notion of a slow and large default memory
// space and a fast and small alternate memory space.
@ -434,6 +436,11 @@ class MemorySpaceAssignment {
// the opcode) to be placed on the alternate memory.
IsAllowedInAlternateMemoryFunction is_allowed_in_alternate_mem_fn;
// This function can be used to prevent certain HloUses (e.g., based on
// the opcode) to be placed on the alternate memory.
IsUseAllowedInAlternateMemoryFunction is_use_allowed_in_alternate_mem_fn =
[](const HloUse&) { return true; };
// Specifies the upper bound for number of outstanding prefetches and
// evictions, -1 for unlimited.
int64 max_outstanding_prefetches = -1;