[XLA:HLO] Run HeapSimulator on whole-module if all computations are sequential.
Previously the HeapSimulator was only run on a per-computation basis. This meant that if you had many sub-computations in your module (e.g. many While loops), the space for all of the temporary buffers inside the conditions and bodies of the loops were in distinct memory ranges. This is overly pessimistic if all computations in the module are sequential. This CL changes the HeapSimulator to also run whole-module simulation, calling Alloc and Free on sub-computation buffers at the appropriate nested spot, right next to the calling instruction. The BufferAssigner is updated to take advantage of this when possible, as is MinimumMemoryForSequence. Change: 154908856
This commit is contained in:
A. Unique TensorFlower
TensorFlower Gardener
parent
58196d4bf9
commit
5ad12420e7
@ -666,8 +666,8 @@ cc_library(
|
||||
],
|
||||
deps = [
|
||||
":buffer_liveness",
|
||||
":heap_simulator",
|
||||
":hlo",
|
||||
":hlo_ordering",
|
||||
":logical_buffer",
|
||||
":tuple_points_to_analysis",
|
||||
"//tensorflow/compiler/xla:shape_util",
|
||||
@ -707,51 +707,38 @@ cc_test(
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "heap_simulator",
|
||||
srcs = [
|
||||
"heap_simulator.cc",
|
||||
],
|
||||
hdrs = [
|
||||
"heap_simulator.h",
|
||||
],
|
||||
deps = [
|
||||
":hlo",
|
||||
":liveness_util",
|
||||
":logical_buffer",
|
||||
":tuple_points_to_analysis",
|
||||
"//tensorflow/compiler/xla:statusor",
|
||||
"//tensorflow/compiler/xla:util",
|
||||
"//tensorflow/core:lib",
|
||||
],
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "heap_simulator_test",
|
||||
srcs = ["heap_simulator_test.cc"],
|
||||
deps = [
|
||||
":heap_simulator",
|
||||
":hlo",
|
||||
":hlo_ordering",
|
||||
":logical_buffer",
|
||||
":tuple_points_to_analysis",
|
||||
"//tensorflow/compiler/xla:literal_util",
|
||||
"//tensorflow/compiler/xla:status_macros",
|
||||
"//tensorflow/compiler/xla/tests:hlo_test_base",
|
||||
"//tensorflow/core:lib",
|
||||
"//tensorflow/core:test_main",
|
||||
],
|
||||
)
|
||||
|
||||
# The hlo_ordering library contains both hlo_ordering and heap_simulator because
|
||||
# they are mutually dependent.
|
||||
cc_library(
|
||||
name = "hlo_ordering",
|
||||
srcs = [
|
||||
"heap_simulator.cc",
|
||||
"hlo_ordering.cc",
|
||||
],
|
||||
hdrs = [
|
||||
"heap_simulator.h",
|
||||
"hlo_ordering.h",
|
||||
],
|
||||
deps = [
|
||||
":call_graph",
|
||||
":heap_simulator",
|
||||
":hlo",
|
||||
":liveness_util",
|
||||
":logical_buffer",
|
||||
":tuple_points_to_analysis",
|
||||
"//tensorflow/compiler/xla:shape_util",
|
||||
|
||||
@ -548,6 +548,8 @@ Status BufferAssigner::AssignBuffersForComputation(
|
||||
const FlatSet<const HloInstruction*>* hlos_to_allocate,
|
||||
const FlatSet<const LogicalBuffer*>& colocated_buffers,
|
||||
const FlatSet<BufferAllocation::Index>& colocated_allocations,
|
||||
FlatMap<const HloComputation*, FlatSet<const LogicalBuffer*>>*
|
||||
buffers_to_assign_sequentially,
|
||||
BufferAssignment* assignment) {
|
||||
// Buffers are sorted and assigned to BufferAllocations in decreasing order of
|
||||
// size.
|
||||
@ -578,9 +580,16 @@ Status BufferAssigner::AssignBuffersForComputation(
|
||||
// If there is a sequential instruction ordering, we'll delay assignment of
|
||||
// temp buffers until after the main assignment loop.
|
||||
const BufferLiveness& liveness = assignment->liveness();
|
||||
const std::vector<const HloInstruction*>* sequential_order =
|
||||
liveness.hlo_ordering().SequentialOrder(*computation);
|
||||
FlatSet<const LogicalBuffer*> unassigned_temp_buffers;
|
||||
const bool has_sequential_order =
|
||||
liveness.hlo_ordering().SequentialOrder(*computation) != nullptr;
|
||||
if (has_sequential_order && buffers_to_assign_sequentially != nullptr) {
|
||||
// Every sequential computation must get an entry in the
|
||||
// buffers_to_assign_sequentially map, even if we end up with an empty set
|
||||
// of buffers. This ensures we can correctly determine whether to run
|
||||
// whole-module heap simulation.
|
||||
buffers_to_assign_sequentially->emplace(computation,
|
||||
FlatSet<const LogicalBuffer*>());
|
||||
}
|
||||
|
||||
// Sort the LogicalBuffers first by size. We assign the larger LogicalBuffers
|
||||
// first for simplicity. This means any previously created BufferAllocation is
|
||||
@ -599,7 +608,7 @@ Status BufferAssigner::AssignBuffersForComputation(
|
||||
// important reuse case where an elementwise instruction reuses one of its
|
||||
// operand's buffer. This improves locality.
|
||||
std::sort(sorted_buffers.begin(), sorted_buffers.end(),
|
||||
[this, sequential_order, &liveness, &post_order_position](
|
||||
[this, has_sequential_order, &liveness, &post_order_position](
|
||||
const LogicalBuffer* a, const LogicalBuffer* b) {
|
||||
// Primary sort is by decreasing buffer size.
|
||||
const int64 a_size = buffer_size_(*a);
|
||||
@ -609,7 +618,7 @@ Status BufferAssigner::AssignBuffersForComputation(
|
||||
}
|
||||
// Otherwise live out buffers come before others, if the
|
||||
// instructions are sequentially ordered.
|
||||
if (sequential_order != nullptr) {
|
||||
if (has_sequential_order) {
|
||||
const bool a_live_out = liveness.MaybeLiveOut(*a);
|
||||
const bool b_live_out = liveness.MaybeLiveOut(*b);
|
||||
if (a_live_out != b_live_out) {
|
||||
@ -746,7 +755,7 @@ Status BufferAssigner::AssignBuffersForComputation(
|
||||
}
|
||||
}
|
||||
|
||||
if (!assignment->HasAllocation(*buffer) && sequential_order != nullptr &&
|
||||
if (!assignment->HasAllocation(*buffer) && has_sequential_order &&
|
||||
!liveness.MaybeLiveOut(*buffer)) {
|
||||
// There is a sequential instruction ordering, so we delay assignment of
|
||||
// temp buffers until after the loop. We do this right before we decide to
|
||||
@ -758,7 +767,7 @@ Status BufferAssigner::AssignBuffersForComputation(
|
||||
// for the definition of temp buffers.
|
||||
CHECK(!is_entry_parameter) << *buffer;
|
||||
CHECK(!is_thread_local) << *buffer;
|
||||
unassigned_temp_buffers.insert(buffer);
|
||||
(*buffers_to_assign_sequentially)[computation].insert(buffer);
|
||||
VLOG(3) << "Delaying assignment of temp buffer: " << *buffer;
|
||||
continue;
|
||||
}
|
||||
@ -772,27 +781,68 @@ Status BufferAssigner::AssignBuffersForComputation(
|
||||
}
|
||||
}
|
||||
|
||||
if (!unassigned_temp_buffers.empty()) {
|
||||
TF_RETURN_IF_ERROR(AssignBuffersWithSequentialOrdering(
|
||||
*sequential_order, unassigned_temp_buffers, *computation, assignment));
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status BufferAssigner::AssignBuffersWithSequentialOrdering(
|
||||
const std::vector<const HloInstruction*>& sequence,
|
||||
const FlatSet<const LogicalBuffer*>& buffers_to_assign,
|
||||
const HloComputation& computation, BufferAssignment* assignment) {
|
||||
const FlatMap<const HloComputation*, FlatSet<const LogicalBuffer*>>&
|
||||
buffers_to_assign_sequentially,
|
||||
bool run_whole_module_heap_simulation, BufferAssignment* assignment) {
|
||||
// Run the sequence of instructions through the heap simulator. The heuristic
|
||||
// that seems to give the best results is lazy-best-fit, with all runs of
|
||||
// alloc / free calls sorted in decreasing size order.
|
||||
TF_ASSIGN_OR_RETURN(
|
||||
HeapSimulator::Result result,
|
||||
HeapSimulator::Run(MakeUnique<DecreasingSizeRunsHeap>(
|
||||
MakeUnique<LazyBestFitHeap>(alignment_)),
|
||||
sequence, computation,
|
||||
assignment->points_to_analysis(), buffer_size_,
|
||||
&buffers_to_assign));
|
||||
const HloOrdering& hlo_ordering = assignment->liveness().hlo_ordering();
|
||||
if (run_whole_module_heap_simulation) {
|
||||
// Run the heap simulation over the whole module. This reduces memory usage,
|
||||
// since buffers for kCall and kWhile sub-computations are only live for the
|
||||
// duration of their calling instructions.
|
||||
VLOG(1) << "Running whole-module heap simulation";
|
||||
SequentialHloOrdering::HloModuleSequence module_sequence;
|
||||
FlatSet<const LogicalBuffer*> all_buffers_to_assign;
|
||||
for (const auto& pair : buffers_to_assign_sequentially) {
|
||||
const HloComputation* computation = pair.first;
|
||||
const FlatSet<const LogicalBuffer*>& buffers_to_assign = pair.second;
|
||||
const std::vector<const HloInstruction*>* instruction_sequence =
|
||||
hlo_ordering.SequentialOrder(*computation);
|
||||
CHECK(instruction_sequence != nullptr) << computation->name();
|
||||
module_sequence[computation] = *instruction_sequence;
|
||||
all_buffers_to_assign.insert(buffers_to_assign.begin(),
|
||||
buffers_to_assign.end());
|
||||
}
|
||||
TF_ASSIGN_OR_RETURN(
|
||||
const HeapSimulator::Result result,
|
||||
HeapSimulator::Run(MakeUnique<DecreasingSizeRunsHeap>(
|
||||
MakeUnique<LazyBestFitHeap>(alignment_)),
|
||||
assignment->module(), module_sequence,
|
||||
assignment->points_to_analysis(), buffer_size_,
|
||||
&all_buffers_to_assign));
|
||||
AssignBuffersFromHeapSimulator(result, assignment);
|
||||
} else {
|
||||
// Run the heap-simulation on a per-computation basis. Buffers for
|
||||
// sub-computations are assigned disjoint BufferAllocations, assuming the
|
||||
// worst-case that they may all be live concurrently.
|
||||
VLOG(1) << "Running per-computation heap simulation";
|
||||
for (const auto& pair : buffers_to_assign_sequentially) {
|
||||
const HloComputation* computation = pair.first;
|
||||
const FlatSet<const LogicalBuffer*>& buffers_to_assign = pair.second;
|
||||
const std::vector<const HloInstruction*>* instruction_sequence =
|
||||
hlo_ordering.SequentialOrder(*computation);
|
||||
CHECK(instruction_sequence != nullptr) << computation->name();
|
||||
TF_ASSIGN_OR_RETURN(
|
||||
const HeapSimulator::Result result,
|
||||
HeapSimulator::Run(MakeUnique<DecreasingSizeRunsHeap>(
|
||||
MakeUnique<LazyBestFitHeap>(alignment_)),
|
||||
*computation, *instruction_sequence,
|
||||
assignment->points_to_analysis(), buffer_size_,
|
||||
&buffers_to_assign));
|
||||
AssignBuffersFromHeapSimulator(result, assignment);
|
||||
}
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
void BufferAssigner::AssignBuffersFromHeapSimulator(
|
||||
const HeapSimulator::Result& result, BufferAssignment* assignment) {
|
||||
if (assignment->stats_.preallocated_temp_fragmentation_bytes == -1) {
|
||||
assignment->stats_.preallocated_temp_fragmentation_bytes =
|
||||
result.fragmentation_size;
|
||||
@ -801,8 +851,6 @@ Status BufferAssigner::AssignBuffersWithSequentialOrdering(
|
||||
result.fragmentation_size;
|
||||
}
|
||||
|
||||
// Use the results of the heap simulator to create one allocation per
|
||||
// computation, with LogicalBuffers packed to specific offsets.
|
||||
BufferAllocation* allocation = assignment->NewEmptyAllocation(
|
||||
result.heap_size, /*is_thread_local=*/false, /*is_reusable=*/true);
|
||||
for (const auto& buffer_chunk : result.chunk_map) {
|
||||
@ -810,7 +858,6 @@ Status BufferAssigner::AssignBuffersWithSequentialOrdering(
|
||||
const HeapSimulator::Chunk& chunk = buffer_chunk.second;
|
||||
assignment->AddAssignment(allocation, buffer, chunk.offset, chunk.size);
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
// Adds the 'colocated_set' of buffers to 'colocated_buffer_sets', maintaining
|
||||
@ -1108,8 +1155,6 @@ StatusOr<std::unique_ptr<BufferAssignment>> BufferAssigner::CreateAssignment(
|
||||
TF_ASSIGN_OR_RETURN(std::unique_ptr<BufferLiveness> liveness,
|
||||
BufferLiveness::Run(module, std::move(hlo_ordering)));
|
||||
|
||||
std::vector<const HloComputation*> thread_local_computations;
|
||||
std::vector<const HloComputation*> global_computations;
|
||||
VLOG(1) << "Assigning buffers to module " << module->name();
|
||||
if (hlos_to_allocate != nullptr) {
|
||||
VLOG(3) << "LogicalBuffer assignment restricted to hlos: ";
|
||||
@ -1121,9 +1166,6 @@ StatusOr<std::unique_ptr<BufferAssignment>> BufferAssigner::CreateAssignment(
|
||||
XLA_VLOG_LINES(3, liveness->ToString());
|
||||
XLA_VLOG_LINES(3, liveness->points_to_analysis().ToString());
|
||||
|
||||
TF_RETURN_IF_ERROR(GatherComputationsByAllocationType(
|
||||
module, &thread_local_computations, &global_computations));
|
||||
|
||||
// Set of HLO's to allocate if hlos_to_allocate is given. Passed as a set to
|
||||
// AssignBuffersForComputation for fast membership testing.
|
||||
std::unique_ptr<FlatSet<const HloInstruction*>> hlo_set;
|
||||
@ -1148,16 +1190,38 @@ StatusOr<std::unique_ptr<BufferAssignment>> BufferAssigner::CreateAssignment(
|
||||
AssignColocatedBufferSets(colocated_buffer_sets, assignment.get(),
|
||||
&colocated_buffers, &colocated_allocations);
|
||||
|
||||
std::vector<const HloComputation*> thread_local_computations;
|
||||
std::vector<const HloComputation*> global_computations;
|
||||
TF_RETURN_IF_ERROR(GatherComputationsByAllocationType(
|
||||
module, &thread_local_computations, &global_computations));
|
||||
|
||||
// First assign buffers for global computatations. Temporary buffers for
|
||||
// sequential computations are collected in 'buffers_to_assign_sequentially'.
|
||||
FlatMap<const HloComputation*, FlatSet<const LogicalBuffer*>>
|
||||
buffers_to_assign_sequentially;
|
||||
for (auto* computation : global_computations) {
|
||||
TF_RETURN_IF_ERROR(AssignBuffersForComputation(
|
||||
computation, /*is_thread_local=*/false, hlo_set.get(),
|
||||
colocated_buffers, colocated_allocations, assignment.get()));
|
||||
colocated_buffers, colocated_allocations,
|
||||
&buffers_to_assign_sequentially, assignment.get()));
|
||||
}
|
||||
// Assign buffers with sequential ordering, if any. If all global computations
|
||||
// are sequential, we can run heap simuation on the whole module, which
|
||||
// reduces memory usage.
|
||||
const bool run_whole_module_heap_simulation =
|
||||
buffers_to_assign_sequentially.size() == global_computations.size();
|
||||
TF_RETURN_IF_ERROR(AssignBuffersWithSequentialOrdering(
|
||||
buffers_to_assign_sequentially, run_whole_module_heap_simulation,
|
||||
assignment.get()));
|
||||
|
||||
// Now assign buffers for thread-local computations. All LogicalBuffers get
|
||||
// their own BufferAllocation.
|
||||
for (auto* computation : thread_local_computations) {
|
||||
TF_RET_CHECK(computation != module->entry_computation());
|
||||
TF_RETURN_IF_ERROR(AssignBuffersForComputation(
|
||||
computation, /*is_thread_local=*/true, hlo_set.get(), colocated_buffers,
|
||||
colocated_allocations, assignment.get()));
|
||||
colocated_allocations, /*buffers_to_assign_sequentially=*/nullptr,
|
||||
assignment.get()));
|
||||
}
|
||||
|
||||
// Mark all buffers which may be live out of the entry computation as
|
||||
|
||||
@ -23,6 +23,7 @@ limitations under the License.
|
||||
#include <vector>
|
||||
|
||||
#include "tensorflow/compiler/xla/service/buffer_liveness.h"
|
||||
#include "tensorflow/compiler/xla/service/heap_simulator.h"
|
||||
#include "tensorflow/compiler/xla/service/hlo_computation.h"
|
||||
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
|
||||
#include "tensorflow/compiler/xla/service/hlo_module.h"
|
||||
@ -354,6 +355,9 @@ class BufferAssignment {
|
||||
void AddAssignment(BufferAllocation* allocation, const LogicalBuffer& buffer,
|
||||
int64 offset, int64 size);
|
||||
|
||||
// Returns the HloModule used to construct this assignment.
|
||||
const HloModule& module() { return *module_; }
|
||||
|
||||
// Returns the BufferLiveness object used to construct this assignment.
|
||||
const BufferLiveness& liveness() { return *liveness_; }
|
||||
|
||||
@ -427,14 +431,27 @@ class BufferAssigner {
|
||||
const tensorflow::gtl::FlatSet<const LogicalBuffer*>& colocated_buffers,
|
||||
const tensorflow::gtl::FlatSet<BufferAllocation::Index>&
|
||||
colocated_allocations,
|
||||
tensorflow::gtl::FlatMap<const HloComputation*,
|
||||
tensorflow::gtl::FlatSet<const LogicalBuffer*>>*
|
||||
buffers_to_assign_sequentially,
|
||||
BufferAssignment* assignment);
|
||||
|
||||
// Assigns 'buffers_to_assign' assuming the HLO instructions will be executed
|
||||
// in the given 'sequential_order'.
|
||||
// Assigns 'buffers_to_assign_sequentially' using heap simulation, assuming
|
||||
// the HLO instructions will be executed in the sequential order given by
|
||||
// assignment->liveness().hlo_ordering().SequentialOrder. If
|
||||
// 'run_whole_module_heap_simulation' is true, the heap simulation will be run
|
||||
// assuming all global computations are sequentially ordered.
|
||||
Status AssignBuffersWithSequentialOrdering(
|
||||
const std::vector<const HloInstruction*>& sequential_order,
|
||||
const tensorflow::gtl::FlatSet<const LogicalBuffer*>& buffers_to_assign,
|
||||
const HloComputation& computation, BufferAssignment* assignment);
|
||||
const tensorflow::gtl::FlatMap<
|
||||
const HloComputation*,
|
||||
tensorflow::gtl::FlatSet<const LogicalBuffer*>>&
|
||||
buffers_to_assign_sequentially,
|
||||
bool run_whole_module_heap_simulation, BufferAssignment* assignment);
|
||||
|
||||
// Uses the results of the heap simulator to create a single allocation, with
|
||||
// LogicalBuffers packed to specific offsets.
|
||||
void AssignBuffersFromHeapSimulator(const HeapSimulator::Result& result,
|
||||
BufferAssignment* assignment);
|
||||
|
||||
// Tries to assign the given instruction to the given buffer. Returns if the
|
||||
// assignment was successful.
|
||||
@ -477,8 +494,6 @@ class BufferAssigner {
|
||||
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).
|
||||
LogicalBuffer::SizeFunction buffer_size_;
|
||||
|
||||
|
||||
@ -53,12 +53,44 @@ std::vector<const LogicalBuffer*> UniqueOperandSourceBuffers(
|
||||
|
||||
/*static*/
|
||||
StatusOr<HeapSimulator::Result> HeapSimulator::Run(
|
||||
std::unique_ptr<HeapAlgorithm> algorithm,
|
||||
const std::vector<const HloInstruction*>& instruction_sequence,
|
||||
const HloComputation& computation,
|
||||
std::unique_ptr<HeapAlgorithm> algorithm, const HloModule& module,
|
||||
const SequentialHloOrdering::HloModuleSequence& module_sequence,
|
||||
const TuplePointsToAnalysis& points_to_analysis,
|
||||
const LogicalBuffer::SizeFunction& size_fn,
|
||||
const FlatSet<const LogicalBuffer*>* buffers_to_assign) {
|
||||
HeapSimulator heap(std::move(algorithm), size_fn, buffers_to_assign);
|
||||
const HloComputation* entry_computation = module.entry_computation();
|
||||
const std::vector<const HloInstruction*>& instruction_sequence =
|
||||
FindOrDie(module_sequence, entry_computation);
|
||||
TF_RETURN_IF_ERROR(heap.RunComputation(*entry_computation,
|
||||
instruction_sequence,
|
||||
points_to_analysis, &module_sequence));
|
||||
return heap.Finish();
|
||||
}
|
||||
|
||||
/*static*/
|
||||
StatusOr<HeapSimulator::Result> HeapSimulator::Run(
|
||||
std::unique_ptr<HeapAlgorithm> algorithm, const HloComputation& computation,
|
||||
const std::vector<const HloInstruction*>& instruction_sequence,
|
||||
const TuplePointsToAnalysis& points_to_analysis,
|
||||
const LogicalBuffer::SizeFunction& size_fn,
|
||||
const FlatSet<const LogicalBuffer*>* buffers_to_assign) {
|
||||
HeapSimulator heap(std::move(algorithm), size_fn, buffers_to_assign);
|
||||
TF_RETURN_IF_ERROR(heap.RunComputation(computation, instruction_sequence,
|
||||
points_to_analysis,
|
||||
/*module_sequence=*/nullptr));
|
||||
return heap.Finish();
|
||||
}
|
||||
|
||||
// Runs a heap simulation for the given 'computation', assuming the given
|
||||
// 'instruction_sequence'. If 'module_sequence' is non-null, it is used to find
|
||||
// kCall and kWhile sub-computations, and the heap simulation for those
|
||||
// sub-computations will be run recursively.
|
||||
Status HeapSimulator::RunComputation(
|
||||
const HloComputation& computation,
|
||||
const std::vector<const HloInstruction*>& instruction_sequence,
|
||||
const TuplePointsToAnalysis& points_to_analysis,
|
||||
const SequentialHloOrdering::HloModuleSequence* module_sequence) {
|
||||
// The goal here is to minimize memory usage, assuming the given sequential
|
||||
// ordering of instructions. The strategy is to walk through the instruction
|
||||
// sequence, calling Alloc and Free on the underlying heap algorithm. The
|
||||
@ -67,7 +99,6 @@ StatusOr<HeapSimulator::Result> HeapSimulator::Run(
|
||||
// 'live_buffers' tracks the liveness of each buffer that we assign, by
|
||||
// associating it with a set of HloInstructions that need to be visited. When
|
||||
// the set becomes empty, the buffer is no longer used, and can be freed.
|
||||
HeapSimulator heap(std::move(algorithm), size_fn, buffers_to_assign);
|
||||
FlatMap<const LogicalBuffer*, FlatSet<const HloInstruction*>> live_buffers;
|
||||
|
||||
const HloInstruction* root = computation.root_instruction();
|
||||
@ -90,7 +121,7 @@ StatusOr<HeapSimulator::Result> HeapSimulator::Run(
|
||||
// lifetime of buffers that aren't already connected by a data dependency.
|
||||
std::vector<const LogicalBuffer*> dead_buffers_to_free;
|
||||
for (const LogicalBuffer* buffer : buffers_defined_by_instruction) {
|
||||
if (heap.IgnoreBuffer(buffer)) {
|
||||
if (IgnoreBuffer(buffer)) {
|
||||
continue;
|
||||
}
|
||||
for (const BufferAlias& alias :
|
||||
@ -127,7 +158,7 @@ StatusOr<HeapSimulator::Result> HeapSimulator::Run(
|
||||
std::vector<const LogicalBuffer*> operand_buffers_to_free;
|
||||
for (const LogicalBuffer* operand_buffer :
|
||||
UniqueOperandSourceBuffers(instruction, points_to_analysis)) {
|
||||
if (heap.IgnoreBuffer(operand_buffer)) {
|
||||
if (IgnoreBuffer(operand_buffer)) {
|
||||
continue;
|
||||
}
|
||||
live_buffers[operand_buffer].erase(instruction);
|
||||
@ -142,10 +173,10 @@ StatusOr<HeapSimulator::Result> HeapSimulator::Run(
|
||||
// happen before dead or operand buffers are freed; the instruction reads
|
||||
// the operand buffers to produce its output.
|
||||
//
|
||||
// INVARIANT: Either heap.Alloc or heap.ShareBuffer will be called for each
|
||||
// buffer that we should assign.
|
||||
// INVARIANT: Either Alloc or ShareBuffer will be called for each buffer
|
||||
// that we should assign.
|
||||
for (const LogicalBuffer* buffer : buffers_defined_by_instruction) {
|
||||
if (heap.IgnoreBuffer(buffer)) {
|
||||
if (IgnoreBuffer(buffer)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
@ -159,24 +190,50 @@ StatusOr<HeapSimulator::Result> HeapSimulator::Run(
|
||||
CanShareOperandBufferWithUser(
|
||||
operand_buffer->instruction(), operand_buffer->index(),
|
||||
buffer->instruction(), buffer->index(), points_to_analysis)) {
|
||||
heap.ShareBuffer(buffer, operand_buffer);
|
||||
ShareBuffer(buffer, operand_buffer);
|
||||
shared = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (!shared) {
|
||||
heap.Alloc(buffer);
|
||||
Alloc(buffer);
|
||||
}
|
||||
}
|
||||
|
||||
// If the whole module is sequential, we can save memory by running the
|
||||
// heap-simulation for sub-computations inline. E.g. the buffers for the
|
||||
// condition and body of a kWhile instruction are only live for the duration
|
||||
// of the instruction itself.
|
||||
//
|
||||
// The order that the sub-computations are simulated does not affect
|
||||
// correctness; since the whole module is sequential, we know that the
|
||||
// sub-computations will never be run concurrently.
|
||||
if (module_sequence != nullptr) {
|
||||
if (instruction->opcode() == HloOpcode::kCall ||
|
||||
instruction->opcode() == HloOpcode::kWhile) {
|
||||
for (const HloComputation* called_computation :
|
||||
instruction->called_computations()) {
|
||||
const std::vector<const HloInstruction*>& called_sequence =
|
||||
FindOrDie(*module_sequence, called_computation);
|
||||
TF_RETURN_IF_ERROR(RunComputation(*called_computation,
|
||||
called_sequence, points_to_analysis,
|
||||
module_sequence));
|
||||
}
|
||||
}
|
||||
|
||||
// Other sub-computations (e.g. Map, Reduce, ...) are skipped; they are
|
||||
// assigned "thread-local" allocations, meaning their buffers are not
|
||||
// allocated up-front at the beginning of the computation.
|
||||
}
|
||||
|
||||
// Free buffers that are no longer live. This is the earliest point that we
|
||||
// can de-allocate; right after the last use of the buffer.
|
||||
for (const LogicalBuffer* buffer : dead_buffers_to_free) {
|
||||
heap.Free(buffer);
|
||||
Free(buffer);
|
||||
}
|
||||
for (const LogicalBuffer* buffer : operand_buffers_to_free) {
|
||||
heap.Free(buffer);
|
||||
Free(buffer);
|
||||
}
|
||||
}
|
||||
|
||||
@ -187,10 +244,10 @@ StatusOr<HeapSimulator::Result> HeapSimulator::Run(
|
||||
const FlatSet<const HloInstruction*>& pending = buffer_pending.second;
|
||||
CHECK_EQ(pending.size(), 1) << *buffer;
|
||||
CHECK(*pending.begin() == nullptr) << *buffer;
|
||||
heap.Free(buffer);
|
||||
Free(buffer);
|
||||
}
|
||||
|
||||
return heap.Finish();
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
HeapSimulator::HeapSimulator(
|
||||
@ -309,6 +366,11 @@ HeapSimulator::Result HeapSimulator::Finish() {
|
||||
result.chunk_map.emplace(buffer, chunk);
|
||||
}
|
||||
}
|
||||
// If we were told to assign specific buffers, make sure we've assigned
|
||||
// exactly that many buffers.
|
||||
if (buffers_to_assign_ != nullptr) {
|
||||
CHECK_EQ(buffers_to_assign_->size(), result.chunk_map.size());
|
||||
}
|
||||
}
|
||||
|
||||
// Fragmentation is the difference between the actual and ideal sizes.
|
||||
|
||||
@ -23,6 +23,7 @@ limitations under the License.
|
||||
|
||||
#include "tensorflow/compiler/xla/service/hlo_computation.h"
|
||||
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
|
||||
#include "tensorflow/compiler/xla/service/hlo_ordering.h"
|
||||
#include "tensorflow/compiler/xla/service/logical_buffer.h"
|
||||
#include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h"
|
||||
#include "tensorflow/compiler/xla/statusor.h"
|
||||
@ -63,17 +64,32 @@ class HeapSimulator {
|
||||
};
|
||||
|
||||
// Run the heap simulation with the given algorithm, assuming the given
|
||||
// sequential ordering of instructions. The 'instruction_sequence' must
|
||||
// contain a topologically-consistent total ordering of all instructions in
|
||||
// the computation. The result is invalid if instructions are not run in
|
||||
// exactly this sequence.
|
||||
// module_sequence, which must contain a topologically-consistent total
|
||||
// ordering of all instructions within each computation. The result is invalid
|
||||
// if instructions are not run in exactly this sequence.
|
||||
//
|
||||
// Running heap simulation on the whole module tends to save memory, compared
|
||||
// to running on a per-computation basis, since we can re-use buffer space for
|
||||
// called sub-computations.
|
||||
//
|
||||
// If 'buffers_to_assign' is provided, only those buffers are assigned
|
||||
// offsets, otherwise all buffers defined by the instructions are assigned.
|
||||
static StatusOr<Result> Run(
|
||||
std::unique_ptr<HeapAlgorithm> algorithm, const HloModule& module,
|
||||
const SequentialHloOrdering::HloModuleSequence& module_sequence,
|
||||
const TuplePointsToAnalysis& points_to_analysis,
|
||||
const LogicalBuffer::SizeFunction& size_fn,
|
||||
const tensorflow::gtl::FlatSet<const LogicalBuffer*>* buffers_to_assign =
|
||||
nullptr);
|
||||
|
||||
// Same as above, but runs on a single computation. The 'instruction_sequence'
|
||||
// must contain a topologically-consistent total ordering of all instructions
|
||||
// in the computation. The result is invalid if instructions are not run in
|
||||
// exactly this sequence.
|
||||
static StatusOr<Result> Run(
|
||||
std::unique_ptr<HeapAlgorithm> algorithm,
|
||||
const std::vector<const HloInstruction*>& instruction_sequence,
|
||||
const HloComputation& computation,
|
||||
const std::vector<const HloInstruction*>& instruction_sequence,
|
||||
const TuplePointsToAnalysis& points_to_analysis,
|
||||
const LogicalBuffer::SizeFunction& size_fn,
|
||||
const tensorflow::gtl::FlatSet<const LogicalBuffer*>* buffers_to_assign =
|
||||
@ -86,6 +102,12 @@ class HeapSimulator {
|
||||
const tensorflow::gtl::FlatSet<const LogicalBuffer*>* buffers_to_assign);
|
||||
~HeapSimulator();
|
||||
|
||||
Status RunComputation(
|
||||
const HloComputation& computation,
|
||||
const std::vector<const HloInstruction*>& instruction_sequence,
|
||||
const TuplePointsToAnalysis& points_to_analysis,
|
||||
const SequentialHloOrdering::HloModuleSequence* module_sequence);
|
||||
|
||||
bool IgnoreBuffer(const LogicalBuffer* buffer) const;
|
||||
void Alloc(const LogicalBuffer* buffer);
|
||||
void Free(const LogicalBuffer* buffer);
|
||||
|
||||
@ -19,13 +19,16 @@ limitations under the License.
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "tensorflow/compiler/xla/literal_util.h"
|
||||
#include "tensorflow/compiler/xla/service/hlo_computation.h"
|
||||
#include "tensorflow/compiler/xla/service/hlo_instruction.h"
|
||||
#include "tensorflow/compiler/xla/service/hlo_module.h"
|
||||
#include "tensorflow/compiler/xla/service/hlo_ordering.h"
|
||||
#include "tensorflow/compiler/xla/service/logical_buffer.h"
|
||||
#include "tensorflow/compiler/xla/service/tuple_points_to_analysis.h"
|
||||
#include "tensorflow/compiler/xla/status_macros.h"
|
||||
#include "tensorflow/compiler/xla/tests/hlo_test_base.h"
|
||||
#include "tensorflow/core/lib/gtl/flatmap.h"
|
||||
|
||||
namespace xla {
|
||||
namespace {
|
||||
@ -69,6 +72,7 @@ class HeapCallRecorder : public HeapAlgorithm {
|
||||
// sequence against an expected sequence.
|
||||
class HeapSimulatorTracker {
|
||||
public:
|
||||
// Constructor for testing a single entry computation.
|
||||
HeapSimulatorTracker(
|
||||
const string& name, std::unique_ptr<HloComputation> computation,
|
||||
const std::vector<const HloInstruction*>& instruction_sequence) {
|
||||
@ -83,12 +87,48 @@ class HeapSimulatorTracker {
|
||||
auto zero_size = [](const LogicalBuffer& buffer) { return 0; };
|
||||
auto algorithm = MakeUnique<DecreasingSizeRunsHeap>(
|
||||
MakeUnique<HeapCallRecorder>(&actual_calls_));
|
||||
result_ = HeapSimulator::Run(std::move(algorithm), instruction_sequence,
|
||||
*module_->entry_computation(),
|
||||
*points_to_analysis_, zero_size)
|
||||
result_ = HeapSimulator::Run(
|
||||
std::move(algorithm), *module_->entry_computation(),
|
||||
instruction_sequence, *points_to_analysis_, zero_size)
|
||||
.ConsumeValueOrDie();
|
||||
}
|
||||
|
||||
explicit HeapSimulatorTracker(const string& name) {
|
||||
module_ = MakeUnique<HloModule>(name);
|
||||
}
|
||||
|
||||
// Similar to the single entry computation constructor above, but runs the
|
||||
// simulation over the entire module.
|
||||
void RunWholeModule(
|
||||
const std::vector<const HloInstruction*>& full_module_sequence) {
|
||||
points_to_analysis_ =
|
||||
TuplePointsToAnalysis::Run(module_.get()).ConsumeValueOrDie();
|
||||
|
||||
// Construct the module sequence grouped by computation.
|
||||
SequentialHloOrdering::HloModuleSequence module_sequence;
|
||||
tensorflow::gtl::FlatMap<const HloInstruction*, int> reverse_position;
|
||||
for (int i = 0; i < full_module_sequence.size(); ++i) {
|
||||
const HloInstruction* instruction = full_module_sequence[i];
|
||||
module_sequence[instruction->parent()].push_back(instruction);
|
||||
reverse_position[instruction] = full_module_sequence.size() - i;
|
||||
}
|
||||
|
||||
// Hack the size_fn so that it returns a decreasing value as we step through
|
||||
// the sequence. This lets us ensure the Alloc calls are in the sequence
|
||||
// order. The Free calls are sorted by LogicalBuffer.id, which is at least
|
||||
// deterministic.
|
||||
auto size_fn = [&reverse_position](const LogicalBuffer& buffer) {
|
||||
return reverse_position[buffer.instruction()];
|
||||
};
|
||||
auto algorithm = MakeUnique<DecreasingSizeRunsHeap>(
|
||||
MakeUnique<HeapCallRecorder>(&actual_calls_));
|
||||
result_ = HeapSimulator::Run(std::move(algorithm), *module_,
|
||||
module_sequence, *points_to_analysis_, size_fn)
|
||||
.ConsumeValueOrDie();
|
||||
}
|
||||
|
||||
HloModule* module() { return module_.get(); }
|
||||
|
||||
// Returns the buffer defined at the given instruction and index.
|
||||
const LogicalBuffer* BufferAt(const HloInstruction* instruction,
|
||||
const ShapeIndex& index) const {
|
||||
@ -358,6 +398,86 @@ TEST_F(HeapSimulatorTest, MultiplyDotDotTuple) {
|
||||
});
|
||||
}
|
||||
|
||||
TEST_F(HeapSimulatorTest, WholeModule) {
|
||||
HeapSimulatorTracker tracker(TestName());
|
||||
|
||||
const Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {});
|
||||
const Shape tuple_shape =
|
||||
ShapeUtil::MakeTupleShape({scalar_shape, scalar_shape});
|
||||
|
||||
auto cond_builder = HloComputation::Builder("WhileCond");
|
||||
HloInstruction* cond_param = cond_builder.AddInstruction(
|
||||
HloInstruction::CreateParameter(0, tuple_shape, "cond_param"));
|
||||
HloInstruction* cond_iter = cond_builder.AddInstruction(
|
||||
HloInstruction::CreateGetTupleElement(scalar_shape, cond_param, 0));
|
||||
HloInstruction* cond_data = cond_builder.AddInstruction(
|
||||
HloInstruction::CreateGetTupleElement(scalar_shape, cond_param, 1));
|
||||
HloInstruction* cond_lt = cond_builder.AddInstruction(
|
||||
HloInstruction::CreateBinary(ShapeUtil::MakeShape(PRED, {}),
|
||||
HloOpcode::kLt, cond_iter, cond_data));
|
||||
HloComputation* cond_computation =
|
||||
tracker.module()->AddEmbeddedComputation(cond_builder.Build());
|
||||
|
||||
auto body_builder = HloComputation::Builder("WhileBody");
|
||||
HloInstruction* body_param = body_builder.AddInstruction(
|
||||
HloInstruction::CreateParameter(0, tuple_shape, "body_param"));
|
||||
HloComputation* body_computation =
|
||||
tracker.module()->AddEmbeddedComputation(body_builder.Build());
|
||||
|
||||
auto builder = HloComputation::Builder(TestName());
|
||||
HloInstruction* param = builder.AddInstruction(
|
||||
HloInstruction::CreateParameter(0, tuple_shape, "param"));
|
||||
HloInstruction* while_op = builder.AddInstruction(HloInstruction::CreateWhile(
|
||||
tuple_shape, cond_computation, body_computation, param));
|
||||
tracker.module()->AddEntryComputation(builder.Build());
|
||||
|
||||
tracker.RunWholeModule(
|
||||
{param, while_op, body_param, cond_param, cond_iter, cond_data, cond_lt});
|
||||
tracker.ExpectCallSequence({
|
||||
// The entry computation param and while_op are allocated first.
|
||||
{kAlloc, tracker.BufferAt(param, {})},
|
||||
{kAlloc, tracker.BufferAt(param, {0})},
|
||||
{kAlloc, tracker.BufferAt(param, {1})},
|
||||
{kAlloc, tracker.BufferAt(while_op, {})},
|
||||
{kAlloc, tracker.BufferAt(while_op, {0})},
|
||||
{kAlloc, tracker.BufferAt(while_op, {1})},
|
||||
|
||||
// Now the while body param is allocated and freed.
|
||||
{kAlloc, tracker.BufferAt(body_param, {})},
|
||||
{kAlloc, tracker.BufferAt(body_param, {0})},
|
||||
{kAlloc, tracker.BufferAt(body_param, {1})},
|
||||
{kFree, tracker.BufferAt(body_param, {})},
|
||||
{kFree, tracker.BufferAt(body_param, {0})},
|
||||
{kFree, tracker.BufferAt(body_param, {1})},
|
||||
|
||||
// Now the while cond param is allocated. The GTE instructions just alias
|
||||
// the param elements, so the param tuple can immediately be freed.
|
||||
{kAlloc, tracker.BufferAt(cond_param, {})},
|
||||
{kAlloc, tracker.BufferAt(cond_param, {0})},
|
||||
{kAlloc, tracker.BufferAt(cond_param, {1})},
|
||||
{kFree, tracker.BufferAt(cond_param, {})},
|
||||
|
||||
// Now the final cond less-than buffer is allocated.
|
||||
{kAlloc, tracker.BufferAt(cond_lt, {})},
|
||||
|
||||
// The order of the remaining Free calls is based on the LogicalBuffer.id,
|
||||
// which is deterministic, but not obvious.
|
||||
{kFree, tracker.BufferAt(param, {})},
|
||||
{kFree, tracker.BufferAt(param, {0})},
|
||||
{kFree, tracker.BufferAt(param, {1})},
|
||||
|
||||
{kFree, tracker.BufferAt(while_op, {})},
|
||||
{kFree, tracker.BufferAt(while_op, {0})},
|
||||
{kFree, tracker.BufferAt(while_op, {1})},
|
||||
|
||||
{kFree, tracker.BufferAt(cond_param, {0})},
|
||||
{kFree, tracker.BufferAt(cond_param, {1})},
|
||||
{kFree, tracker.BufferAt(cond_lt, {})},
|
||||
|
||||
{kFinish, nullptr},
|
||||
});
|
||||
}
|
||||
|
||||
// Base class for heap algorithm tests.
|
||||
class HeapAlgorithmTestBase : public ::testing::Test {
|
||||
protected:
|
||||
|
||||
@ -221,23 +221,6 @@ string SequentialHloOrdering::ToString() const {
|
||||
return tensorflow::str_util::Join(pieces, "\n");
|
||||
}
|
||||
|
||||
namespace {
|
||||
StatusOr<int64> MinimumMemoryForSequence(
|
||||
const HloComputation& computation,
|
||||
const std::vector<const HloInstruction*>& sequence,
|
||||
const TuplePointsToAnalysis& points_to_analysis,
|
||||
const LogicalBuffer::SizeFunction& size_function) {
|
||||
// The absolute minimum memory required for a given sequence of instructions
|
||||
// is determined by the sequence of Alloc and Free calls on a simulated heap,
|
||||
// ignoring fragmentation.
|
||||
TF_ASSIGN_OR_RETURN(
|
||||
HeapSimulator::Result result,
|
||||
HeapSimulator::Run(MakeUnique<NoFragmentationStatsHeap>(), sequence,
|
||||
computation, points_to_analysis, size_function));
|
||||
return result.heap_size;
|
||||
}
|
||||
} // namespace
|
||||
|
||||
StatusOr<int64> MinimumMemoryForSequence(
|
||||
const SequentialHloOrdering::HloModuleSequence& module_sequence,
|
||||
const LogicalBuffer::SizeFunction& size_function) {
|
||||
@ -249,17 +232,16 @@ StatusOr<int64> MinimumMemoryForSequence(
|
||||
TF_ASSIGN_OR_RETURN(std::unique_ptr<TuplePointsToAnalysis> points_to_analysis,
|
||||
TuplePointsToAnalysis::Run(module));
|
||||
|
||||
int64 total_memory = 0;
|
||||
for (const auto& pair : module_sequence) {
|
||||
const HloComputation* computation = pair.first;
|
||||
const std::vector<const HloInstruction*>& sequence = pair.second;
|
||||
TF_ASSIGN_OR_RETURN(
|
||||
const int64 memory,
|
||||
MinimumMemoryForSequence(*computation, sequence, *points_to_analysis,
|
||||
size_function));
|
||||
total_memory += memory;
|
||||
}
|
||||
return total_memory;
|
||||
// The absolute minimum memory required for a given sequence of instructions
|
||||
// is determined by the sequence of Alloc and Free calls on a simulated heap,
|
||||
// ignoring fragmentation. We run the heap simulation on the whole module,
|
||||
// rather than summing each computation, since it gives us a better lower
|
||||
// bound, by minimizing the liveness of sub-computations.
|
||||
TF_ASSIGN_OR_RETURN(
|
||||
HeapSimulator::Result result,
|
||||
HeapSimulator::Run(MakeUnique<NoFragmentationStatsHeap>(), *module,
|
||||
module_sequence, *points_to_analysis, size_function));
|
||||
return result.heap_size;
|
||||
}
|
||||
|
||||
namespace {
|
||||
@ -516,6 +498,18 @@ StatusOr<std::vector<const HloInstruction*>> RunDFSMemoryScheduler(
|
||||
return sequence;
|
||||
}
|
||||
|
||||
StatusOr<int64> MinimumMemoryForComputation(
|
||||
const HloComputation& computation,
|
||||
const std::vector<const HloInstruction*>& sequence,
|
||||
const TuplePointsToAnalysis& points_to_analysis,
|
||||
const LogicalBuffer::SizeFunction& size_function) {
|
||||
TF_ASSIGN_OR_RETURN(
|
||||
HeapSimulator::Result result,
|
||||
HeapSimulator::Run(MakeUnique<NoFragmentationStatsHeap>(), computation,
|
||||
sequence, points_to_analysis, size_function));
|
||||
return result.heap_size;
|
||||
}
|
||||
|
||||
StatusOr<std::vector<const HloInstruction*>> CreateMemoryMinimizingSequence(
|
||||
const HloComputation& computation,
|
||||
const TuplePointsToAnalysis& points_to_analysis,
|
||||
@ -523,13 +517,17 @@ StatusOr<std::vector<const HloInstruction*>> CreateMemoryMinimizingSequence(
|
||||
// We try both a list-scheduler based ordering and a DFS based ordering, and
|
||||
// choose whichever returns a lower min-memory, not accounting for
|
||||
// fragmentation.
|
||||
//
|
||||
// Note that this is just a heuristic. One obvious inaccuracy is that the
|
||||
// memory required for sub-computations might be different when considered
|
||||
// within the caller's context. But it's good enough for now.
|
||||
TF_ASSIGN_OR_RETURN(
|
||||
std::vector<const HloInstruction*> list_sequence,
|
||||
ListScheduler::Run(computation, points_to_analysis, size_function));
|
||||
TF_ASSIGN_OR_RETURN(
|
||||
const int64 list_memory,
|
||||
MinimumMemoryForSequence(computation, list_sequence, points_to_analysis,
|
||||
size_function));
|
||||
MinimumMemoryForComputation(computation, list_sequence,
|
||||
points_to_analysis, size_function));
|
||||
VLOG(2) << "Min-memory list sequence: " << list_memory << " bytes";
|
||||
|
||||
TF_ASSIGN_OR_RETURN(
|
||||
@ -537,8 +535,8 @@ StatusOr<std::vector<const HloInstruction*>> CreateMemoryMinimizingSequence(
|
||||
RunDFSMemoryScheduler(computation, points_to_analysis, size_function));
|
||||
TF_ASSIGN_OR_RETURN(
|
||||
const int64 dfs_memory,
|
||||
MinimumMemoryForSequence(computation, dfs_sequence, points_to_analysis,
|
||||
size_function));
|
||||
MinimumMemoryForComputation(computation, dfs_sequence, points_to_analysis,
|
||||
size_function));
|
||||
VLOG(2) << "Min-memory dfs sequence: " << dfs_memory << " bytes";
|
||||
|
||||
if (list_memory <= dfs_memory) {
|
||||
|
||||
@ -155,6 +155,65 @@ TEST_F(HloOrderingTest, InstructionsInDifferentComputations) {
|
||||
EXPECT_FALSE(ordering.ExecutesBefore(y, c));
|
||||
}
|
||||
|
||||
class MinimumMemoryForSequenceTest : public HloTestBase {};
|
||||
|
||||
TEST_F(MinimumMemoryForSequenceTest, MultiComputation) {
|
||||
HloModule module(TestName());
|
||||
const Shape scalar_shape = ShapeUtil::MakeShape(xla::F32, {});
|
||||
const Shape tuple_shape =
|
||||
ShapeUtil::MakeTupleShape({scalar_shape, scalar_shape});
|
||||
|
||||
auto cond_builder = HloComputation::Builder("WhileCond");
|
||||
// Tuple param: 24 bytes (each elem has 8 byte pointer, 4 byte element)
|
||||
HloInstruction* cond_param = cond_builder.AddInstruction(
|
||||
HloInstruction::CreateParameter(0, tuple_shape, "cond_param"));
|
||||
HloInstruction* cond_iter = cond_builder.AddInstruction(
|
||||
HloInstruction::CreateGetTupleElement(scalar_shape, cond_param, 0));
|
||||
HloInstruction* cond_data = cond_builder.AddInstruction(
|
||||
HloInstruction::CreateGetTupleElement(scalar_shape, cond_param, 1));
|
||||
// Free cond_param[] (16 bytes), Alloc PRED[] (1 byte)
|
||||
HloInstruction* cond_lt = cond_builder.AddInstruction(
|
||||
HloInstruction::CreateBinary(ShapeUtil::MakeShape(PRED, {}),
|
||||
HloOpcode::kLt, cond_iter, cond_data));
|
||||
HloComputation* cond_computation =
|
||||
module.AddEmbeddedComputation(cond_builder.Build());
|
||||
|
||||
auto body_builder = HloComputation::Builder("WhileBody");
|
||||
// Tuple param: 24 bytes (each elem has 8 byte pointer, 4 byte element)
|
||||
HloInstruction* body_param = body_builder.AddInstruction(
|
||||
HloInstruction::CreateParameter(0, tuple_shape, "body_param"));
|
||||
HloComputation* body_computation =
|
||||
module.AddEmbeddedComputation(body_builder.Build());
|
||||
|
||||
auto builder = HloComputation::Builder(TestName());
|
||||
// Entry params: 8 bytes (4 bytes per param), TOTAL=8
|
||||
HloInstruction* iter = builder.AddInstruction(
|
||||
HloInstruction::CreateParameter(0, scalar_shape, "param_iter"));
|
||||
HloInstruction* data = builder.AddInstruction(
|
||||
HloInstruction::CreateParameter(1, scalar_shape, "param_data"));
|
||||
// Tuple: 16 bytes (8 bytes per pointer), TOTAL=24
|
||||
HloInstruction* tuple =
|
||||
builder.AddInstruction(HloInstruction::CreateTuple({iter, data}));
|
||||
// While: 8 bytes (4 bytes per element), TOTAL=32
|
||||
// Both cond and body use a max of 24 bytes, TOTAL=56
|
||||
HloInstruction* while_op = builder.AddInstruction(HloInstruction::CreateWhile(
|
||||
tuple_shape, cond_computation, body_computation, tuple));
|
||||
HloComputation* entry_computation =
|
||||
module.AddEntryComputation(builder.Build());
|
||||
|
||||
auto size_fn = [](const LogicalBuffer& buffer) {
|
||||
return ShapeUtil::ByteSizeOf(buffer.shape(), /*pointer_size=*/8);
|
||||
};
|
||||
|
||||
SequentialHloOrdering::HloModuleSequence module_sequence;
|
||||
module_sequence[cond_computation] = {cond_param, cond_iter, cond_data,
|
||||
cond_lt};
|
||||
module_sequence[body_computation] = {body_param};
|
||||
module_sequence[entry_computation] = {iter, data, tuple, while_op};
|
||||
EXPECT_EQ(56,
|
||||
MinimumMemoryForSequence(module_sequence, size_fn).ValueOrDie());
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
} // namespace xla
|
||||
|
||||
Loading…
Reference in New Issue
Block a user