[XLA/GPU] Polishing code style.

- Also remove unnecessary checks.

tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.cc

@@ -1857,8 +1857,7 @@ IrEmitterUnnested::BuildKernelThunkFromBufferSlices(
     absl::string_view name, Thunk::ThunkInfo thunk_info,
     absl::Span<const BufferSlice* const> slices,
     std::function<void(const BufferSlice*, llvm::Value*)>
-        bind_slice_to_ir_value,
-    bool insist_single_temp_buffer) {
+        bind_slice_to_ir_value) {
   const auto& buffer_assn = ir_emitter_context_->buffer_assignment();
 
   // 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()) {
     if (alloc.IsPreallocatedTempBuffer()) {
       if (!temp_buffer.has_value()) {
+        // Retrieve the first seen temp buffer.
         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();
 
         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(

tensorflow/compiler/xla/service/gpu/ir_emitter_unnested.h

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

tensorflow/compiler/xla/service/heap_simulator.cc

@@ -409,10 +409,11 @@ HeapSimulator::Result<HloValue> HeapSimulator::Finish() {
   // 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
   // collecting statistics, e.g. NoFragmentationStatsHeap.
-  size_t total_chunk_count = 0;
-  absl::c_for_each(result.heap_results, [&](const HeapResult<HloValue>& hr) {
-    total_chunk_count += hr.chunk_map.size();
-  });
+  size_t total_chunk_count = absl::c_accumulate(
+      result.heap_results, (size_t)0,
+      [&](size_t lhs, const HeapResult<HloValue>& rhs) -> size_t {
+        return lhs + rhs.chunk_map.size();
+      });
   if (total_chunk_count != 0) {
     // If we were told to assign specific buffers, make sure we've assigned
     // exactly that many buffers.