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

[XLA] Remove an overly pessimistic assertion from memory space assignment.

Browse Source 
The assertion could trigger when the same HloValue is passed multiple times to a
while HLO. We already handle this case in AllocateSegment so the assertion isn't
necessary.

PiperOrigin-RevId: 349591728
Change-Id: I47e513c3d4ab8006ab43dcad59e6935249ab7b31
This commit is contained in:
Berkin Ilbeyi 2020-12-30 12:56:28 -08:00 committed by TensorFlower Gardener
parent 89310df4bc
commit ae6ff84595
2 changed files with 40 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -941,12 +941,8 @@ bool AlternateMemoryBestFitHeap::IsUseAllowedInAlternateMemory(
int64 while_time = instruction_schedule.at(use.instruction);
auto existing_required_assignment =
RequiredMemoryAssignmentAt(while_value, while_time);
if (existing_required_assignment) {
// TODO(berkin): Failing for now when the output is requested to be in
// alternate memory, and the buffer is a while loop output.
CHECK(existing_required_assignment->memory_space == MemorySpace::kDefault)
<< "While loop buffers pinned to alternate memory not "
"currently supported.";
if (existing_required_assignment &&
existing_required_assignment->memory_space == MemorySpace::kDefault) {
VLOG(4) << "While allocation not allowed in alternate memory because "
"there is a required default memory assignment.";
return false;

38
tensorflow/compiler/xla/service/memory_space_assignment_test.cc
View File

@ -4152,6 +4152,44 @@ TEST_P(MemorySpaceAssignmentTest, MoveCopyDoneEarlier) {
find_schedule_index(cos->operand(0)));
}
TEST_P(MemorySpaceAssignmentTest, WhileAliasedArgumentRequiredAssignmentBug) {
// Tests an overly pessimistic assertion when the same HloValue is passed
// multiple times to a while HLO. We already handle this case that the two
// arguments must alias and get the same allocation in AllocateSegment so the
// assertion isn't necessary.
absl::string_view hlo_string = R"(
HloModule bug, is_scheduled=true
while_condition {
param1 = (f32[2,4], f32[2,4], f32[2,4]) parameter(0)
ROOT cond = pred[] constant(true)
}
while_body {
param2 = (f32[2,4], f32[2,4], f32[2,4]) parameter(0)
gte2 = f32[2,4] get-tuple-element(param2), index=0
gte3 = f32[2,4] get-tuple-element(param2), index=1
gte4 = f32[2,4] get-tuple-element(param2), index=2
add = f32[2,4] add(gte2, gte3)
ROOT tuple2 = (f32[2,4], f32[2,4], f32[2,4]) tuple(add, gte3, gte4)
}
ENTRY Entry {
param0 = f32[2,4] parameter(0)
a = f32[2,4] negate(param0)
b = f32[2,4] negate(param0)
tuple = (f32[2,4], f32[2,4], f32[2,4]) tuple(a, b, b)
while = (f32[2,4], f32[2,4], f32[2,4]) while(tuple), condition=while_condition, body=while_body
gte1 = f32[2,4] get-tuple-element(while), index=0
gte2 = f32[2,4] get-tuple-element(while), index=1
ROOT root = f32[2,4] add(gte1, gte2)
}
)";
TF_ASSERT_OK_AND_ASSIGN(auto module,
ParseAndReturnVerifiedModule(hlo_string));
AssignMemorySpace(module.get());
}
TEST_P(MemorySpaceAssignmentTest, DisallowedUseBug) {
// When we have a disallowed use (in this case tanh), we aren't allowed to
// allocate this use in alternate memory. However, if we have another use