Home | History | Annotate | Download | only in gpu 
      1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
      2 
      3 Licensed under the Apache License, Version 2.0 (the "License");
      4 you may not use this file except in compliance with the License.
      5 You may obtain a copy of the License at
      6 
      7     http://www.apache.org/licenses/LICENSE-2.0
      8 
      9 Unless required by applicable law or agreed to in writing, software
     10 distributed under the License is distributed on an "AS IS" BASIS,
     11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     12 See the License for the specific language governing permissions and
     13 limitations under the License.
     14 ==============================================================================*/
     15 
     16 #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_GPU_BUFFER_ALLOCATIONS_H_
     17 #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_BUFFER_ALLOCATIONS_H_
     18 
     19 #include <memory>
     20 #include <set>
     21 #include <vector>
     22 
     23 #include "tensorflow/compiler/xla/service/buffer_assignment.h"
     24 #include "tensorflow/compiler/xla/service/device_memory_allocator.h"
     25 #include "tensorflow/compiler/xla/statusor.h"
     26 #include "tensorflow/core/lib/gtl/array_slice.h"
     27 #include "tensorflow/core/platform/stream_executor_no_cuda.h"
     28 
     29 namespace xla {
     30 namespace gpu {
     31 
     32 // A thread-compatible class that encapsulates the base addresses of the
     33 // allocated device buffers.
     34 class BufferAllocations {
     35  public:
     36   // This inner class encapsulates methods that build a BufferAllocations from
     37   // the given buffer assignment.
     38   class Builder {
     39    public:
     40     // Registers preallocated buffers (such as parameter addresses and
     41     // user-specified result buffers) to the given buffer index. The builder
     42     // will skip allocating buffers for registered buffer indices.
     43     void RegisterBuffer(BufferAllocation::Index index,
     44                         perftools::gputools::DeviceMemoryBase address);
     45 
     46     // Builds a BufferAllocations object from the given buffer assignment.
     47     // `memory_allocator` is what this function uses to allocate device memory.
     48     // `device_ordinal` is the number of the device this function allocates
     49     // memory on.
     50     StatusOr<std::unique_ptr<BufferAllocations>> Build(
     51         const BufferAssignment& buffer_assignment, int device_ordinal,
     52         DeviceMemoryAllocator* memory_allocator);
     53 
     54    private:
     55     std::map<BufferAllocation::Index, perftools::gputools::DeviceMemoryBase>
     56         registered_buffers_;
     57   };
     58 
     59   BufferAllocations(const BufferAllocations&) = delete;
     60   BufferAllocations& operator=(const BufferAllocations&) = delete;
     61 
     62   DeviceMemoryAllocator* memory_allocator() const { return memory_allocator_; }
     63   int device_ordinal() const { return device_ordinal_; }
     64 
     65   // Returns the device address of buffer `buffer_index`. `buffer_index` must be
     66   // a valid index, i.e., in [0, buffer_count). This function returns null if
     67   // `buffer_index` is not assigned to a buffer address.
     68   perftools::gputools::DeviceMemoryBase GetDeviceAddress(
     69       BufferAllocation::Index buffer_index) const;
     70 
     71   // Same as above, but also adjusts the returned address for the offset and
     72   // size contained in the given slice.
     73   perftools::gputools::DeviceMemoryBase GetDeviceAddress(
     74       const BufferAllocation::Slice& buffer_slice) const;
     75 
     76   perftools::gputools::DeviceMemoryBase GetTempBufferBase() const {
     77     return temp_buffer_base_;
     78   }
     79 
     80   // Tears down all buffers allocated by this object that are not in
     81   // `live_addresses`.
     82   tensorflow::Status TearDown(
     83       const std::set<perftools::gputools::DeviceMemoryBase>& live_addresses,
     84       const BufferAssignment& buffer_assignment);
     85 
     86  private:
     87   BufferAllocations(BufferAllocation::Index buffer_count, int device_ordinal,
     88                     DeviceMemoryAllocator* memory_allocator)
     89       : buffers_(buffer_count),
     90         device_ordinal_(device_ordinal),
     91         memory_allocator_(memory_allocator) {}
     92 
     93   // Sets the device address of buffer `buffer_index`.
     94   void SetBuffer(BufferAllocation::Index buffer_index,
     95                  perftools::gputools::DeviceMemoryBase buffer);
     96 
     97   // An array of device pointers that stores the address of each buffer
     98   // indexed by Index. Each element can point to a temporary buffer, an
     99   // input buffer, or nullptr if no buffer is needed for that Index.
    100   std::vector<perftools::gputools::DeviceMemoryBase> buffers_;
    101 
    102   // The base address of the memory block that contains all temporary buffers.
    103   perftools::gputools::DeviceMemoryBase temp_buffer_base_;
    104 
    105   int device_ordinal_;
    106 
    107   DeviceMemoryAllocator* memory_allocator_;
    108 };
    109 
    110 }  // namespace gpu
    111 }  // namespace xla
    112 
    113 #endif  // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_BUFFER_ALLOCATIONS_H_
    114