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

Add add new option in HloComputation::AddInstruction to add instruction with different name.

Browse Source 
PiperOrigin-RevId: 336146593
Change-Id: I57ea0152169fc54188cf1002bda772b6153070c9
This commit is contained in:
Yunxing Dai 2020-10-08 12:55:51 -07:00 committed by TensorFlower Gardener
parent d39d5bf305
commit a97f22586d
3 changed files with 10 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tensorflow/compiler/xla/service/hlo_computation.cc
View File

@ -93,10 +93,13 @@ HloComputation::HloComputation(
}
HloInstruction* HloComputation::AddInstruction(
std::unique_ptr<HloInstruction> instruction) {
std::unique_ptr<HloInstruction> instruction, const std::string& new_name) {
CHECK(instruction->opcode() != HloOpcode::kParameter)
<< "Parameter instructions cannot be added to a computation after "
<< "it has been built";
if (!new_name.empty()) {
instruction->SetAndSanitizeName(new_name);
}
return AddInstructionInternal(std::move(instruction));
}

3
tensorflow/compiler/xla/service/hlo_computation.h
View File

@ -122,7 +122,8 @@ class HloComputation {
// Add an instruction to the computation. The computation takes ownership of
// the instruction.
HloInstruction* AddInstruction(std::unique_ptr<HloInstruction> instruction);
HloInstruction* AddInstruction(std::unique_ptr<HloInstruction> instruction,
const std::string& new_name = "");
// Remove the param_no'th parameter from the computation.
// Note this is only applicatable to the computation for the fusion

9
tensorflow/compiler/xla/service/hlo_rematerialization.cc
View File

@ -1521,14 +1521,13 @@ StatusOr<int64> CompressInstruction(MemoryUsageTracker* memory_tracker,
<< ") to" << compact_shape.ToString(true);
HloComputation* computation = best->parent();
HloInstruction* compressed = computation->AddInstruction(
HloInstruction::CreateUnary(compact_shape, HloOpcode::kCopy, best));
compressed->SetAndSanitizeName(best->name() + ".remat_compressed");
HloInstruction::CreateUnary(compact_shape, HloOpcode::kCopy, best),
/*new_name=*/best->name() + ".remat_compressed");
HloInstruction* uncompressed = computation->AddInstruction(
HloInstruction::CreateUnary(best->shape(), HloOpcode::kCopy, compressed));
uncompressed->SetAndSanitizeName(best->name() + ".remat_uncompressed");
HloInstruction::CreateUnary(best->shape(), HloOpcode::kCopy, compressed),
/*new_name=*/best->name() + ".remat_uncompressed");
Item* compressed_item = instruction_list->CreateItem(compressed);
compressed_item->placed = true;