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

Adds a max_rematerialized_block_size field.

Browse Source 
PiperOrigin-RevId: 301205071
Change-Id: I9403816bbdc1e079a634b6ada730ecf7983eba6c
This commit is contained in:
A. Unique TensorFlower 2020-03-16 11:46:57 -07:00 committed by TensorFlower Gardener
parent e17dc08a3c
commit 09e0b6cea4
2 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -1648,6 +1648,8 @@ StatusOr<bool> HloRematerialization::RematerializeComputation(
} else {
// Found a valid block. Reset to start looking for single instructions
// again.
max_rematerialized_block_size_ =
std::max(max_rematerialized_block_size_, max_block_size);
changed = true;
min_block_size = 1;
max_block_size = 1;

4
tensorflow/compiler/xla/service/hlo_rematerialization.h
View File

@ -180,6 +180,10 @@ class HloRematerialization : public HloModulePass {
// dead. Hence, no net instructions were added.
int64 net_instructions_added_ = 0;
// Size of the largest block that has been rematerialized. This is actually an
// upper bound (within a factor of 2) on the block size.
int max_rematerialized_block_size_ = 0;
RematerializationMode mode_;
};