diff options
Diffstat (limited to 'tensorflow/compiler/xla/service/buffer_assignment.h')
| -rw-r--r-- | tensorflow/compiler/xla/service/buffer_assignment.h | 10 |
1 files changed, 9 insertions, 1 deletions
diff --git a/tensorflow/compiler/xla/service/buffer_assignment.h b/tensorflow/compiler/xla/service/buffer_assignment.h index b82acb19b3..ec1375e24d 100644 --- a/tensorflow/compiler/xla/service/buffer_assignment.h +++ b/tensorflow/compiler/xla/service/buffer_assignment.h @@ -465,7 +465,7 @@ class BufferAssigner { // ColocatedBufferSet aggregates a set of related LogicalBuffers from 'module' // which should be colocated in the same buffer allocation. void BuildColocatedBufferSets( - const HloModule* module, const TuplePointsToAnalysis& points_to_analysis, + const HloModule* module, const BufferLiveness& buffer_liveness, std::vector<ColocatedBufferSet>* colocated_buffer_sets); // For each buffer set in 'colocated_buffer_sets', assigns all buffers in the @@ -482,6 +482,14 @@ class BufferAssigner { const std::vector<const LogicalBuffer*>& colocated_set, std::vector<ColocatedBufferSet>* colocated_buffer_sets); + // Conceptually the same as AddSetToColocatedBufferSets, but specific to the + // colocated buffers for while instructions. + void AddWhileSetToColocatedBufferSets( + const std::vector<const LogicalBuffer*>& colocated_set, + const LogicalBuffer* while_init_buffer, const HloInstruction* while_hlo, + const HloComputation& computation, const BufferLiveness& buffer_liveness, + std::vector<ColocatedBufferSet>* colocated_buffer_sets); + const HloModule* module_; // Function which returns the buffer size for a given logical buffer (shape). |