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

[XLA/GPU] Polishing code style.

Browse Source 
- Also remove unnecessary checks.
This commit is contained in:
Trent Lo 2020-10-13 11:24:07 -07:00
parent 55a1d7cbbd
commit 12b38f98bb
3 changed files with 9 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc
View File

@ -1857,8 +1857,7 @@ IrEmitterUnnested::BuildKernelThunkFromBufferSlices(
absl::string_view name, Thunk::ThunkInfo thunk_info, absl::string_view name, Thunk::ThunkInfo thunk_info,
absl::Span<const BufferSlice* const> slices, absl::Span<const BufferSlice* const> slices,
std::function<void(const BufferSlice*, llvm::Value*)> std::function<void(const BufferSlice*, llvm::Value*)>
bind_slice_to_ir_value, bind_slice_to_ir_value) {
bool insist_single_temp_buffer) {
const auto& buffer_assn = ir_emitter_context_->buffer_assignment(); const auto& buffer_assn = ir_emitter_context_->buffer_assignment();
// Figure out which buffer allocations need to be passed as arguments to our // Figure out which buffer allocations need to be passed as arguments to our
@ -1874,9 +1873,8 @@ IrEmitterUnnested::BuildKernelThunkFromBufferSlices(
for (const BufferAllocation& alloc : buffer_assn.Allocations()) { for (const BufferAllocation& alloc : buffer_assn.Allocations()) {
if (alloc.IsPreallocatedTempBuffer()) { if (alloc.IsPreallocatedTempBuffer()) {
if (!temp_buffer.has_value()) { if (!temp_buffer.has_value()) {
// Retrieve the first seen temp buffer.
temp_buffer = &alloc; temp_buffer = &alloc;
} else if (insist_single_temp_buffer) {
LOG(FATAL) << "Multiple temp buffers found, but only one is allowed!";
} }
} }
} }
@ -1996,13 +1994,7 @@ std::unique_ptr<KernelThunk> IrEmitterUnnested::BuildKernelThunk(
<< hlo_buffer_slice->gte_index.ToString(); << hlo_buffer_slice->gte_index.ToString();
bindings_.BindHloToIrValue(*instr, value, index); bindings_.BindHloToIrValue(*instr, value, index);
}, });
// Check temp buffer numbers only when the multiheap mode is off.
/*insist_single_temp_buffer=*/inst->parent()
->parent()
->config()
.debug_options()
.xla_multiheap_size_constraint_per_heap() == -1);
} }
std::unique_ptr<KernelThunk> IrEmitterUnnested::BuildKernelThunkForMlir( std::unique_ptr<KernelThunk> IrEmitterUnnested::BuildKernelThunkForMlir(

3
tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h
View File

@ -547,8 +547,7 @@ class IrEmitterUnnested : public IrEmitter,
absl::string_view name, Thunk::ThunkInfo thunk_info, absl::string_view name, Thunk::ThunkInfo thunk_info,
absl::Span<const BufferSlice* const> slices, absl::Span<const BufferSlice* const> slices,
std::function<void(const BufferSlice*, llvm::Value*)> std::function<void(const BufferSlice*, llvm::Value*)>
bind_slice_to_ir_value, bind_slice_to_ir_value);
bool insist_single_temp_buffer = true);
// Returns a KernelThunk that invokes the kernel emitted for `inst`. The // Returns a KernelThunk that invokes the kernel emitted for `inst`. The
// caller needs to make sure `inst` outlives the lifetime of the returned // caller needs to make sure `inst` outlives the lifetime of the returned

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

@ -409,10 +409,11 @@ HeapSimulator::Result<HloValue> HeapSimulator::Finish() {
// Post-process the result to add chunks for shared buffers. An empty chunk // Post-process the result to add chunks for shared buffers. An empty chunk
// map means that either no buffers were allocated, or the heap was only // map means that either no buffers were allocated, or the heap was only
// collecting statistics, e.g. NoFragmentationStatsHeap. // collecting statistics, e.g. NoFragmentationStatsHeap.
size_t total_chunk_count = 0; size_t total_chunk_count = absl::c_accumulate(
absl::c_for_each(result.heap_results, [&](const HeapResult<HloValue>& hr) { result.heap_results, (size_t)0,
total_chunk_count += hr.chunk_map.size(); [&](size_t lhs, const HeapResult<HloValue>& rhs) -> size_t {
}); return lhs + rhs.chunk_map.size();
});
if (total_chunk_count != 0) { if (total_chunk_count != 0) {
// If we were told to assign specific buffers, make sure we've assigned // If we were told to assign specific buffers, make sure we've assigned
// exactly that many buffers. // exactly that many buffers.