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