Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
DeviceTensor-inl.cuh
1 /**
2 * Copyright (c) Facebook, Inc. and its affiliates.
3 *
4 * This source code is licensed under the MIT license found in the
5 * LICENSE file in the root directory of this source tree.
6 */
7
8
9 #include <utility> // std::move
10
11 namespace faiss { namespace gpu {
12
13 template <typename T, int Dim, bool InnerContig,
14 typename IndexT, template <typename U> class PtrTraits>
15 __host__
16 DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor() :
17 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(),
18 state_(AllocState::NotOwner),
19 space_(MemorySpace::Device) {
20 }
21
22 template <typename T, int Dim, bool InnerContig,
23 typename IndexT, template <typename U> class PtrTraits>
24 __host__
25 DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
26 DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t) :
27 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(),
28 state_(AllocState::NotOwner),
29 space_(MemorySpace::Device) {
30 this->operator=(std::move(t));
31 }
32
33 template <typename T, int Dim, bool InnerContig,
34 typename IndexT, template <typename U> class PtrTraits>
35 __host__
36 DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&
37 DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::operator=(
38 DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t) {
39 if (this->state_ == AllocState::Owner) {
40 CUDA_VERIFY(cudaFree(this->data_));
41 }
42
43 this->Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::operator=(
44 std::move(t));
45
46 this->state_ = t.state_; t.state_ = AllocState::NotOwner;
47 this->space_ = t.space_;
48 this->reservation_ = std::move(t.reservation_);
49
50 return *this;
51 }
52
53 template <typename T, int Dim, bool InnerContig,
54 typename IndexT, template <typename U> class PtrTraits>
55 __host__
56 DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::~DeviceTensor() {
57 if (state_ == AllocState::Owner) {
58 FAISS_ASSERT(this->data_ || (this->getSizeInBytes() == 0));
59 CUDA_VERIFY(cudaFree(this->data_));
60 this->data_ = nullptr;
61 }
62
63 // Otherwise, if we have a temporary memory reservation, then its
64 // destructor will return the reservation
65 }
66
67 template <typename T, int Dim, bool InnerContig,
68 typename IndexT, template <typename U> class PtrTraits>
69 __host__
70 DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
71 const IndexT sizes[Dim],
72 MemorySpace space) :
73 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(nullptr, sizes),
74 state_(AllocState::Owner),
75 space_(space) {
76
77 allocMemorySpace(space, &this->data_, this->getSizeInBytes());
78 FAISS_ASSERT(this->data_ || (this->getSizeInBytes() == 0));
79 }
80
81 template <typename T, int Dim, bool InnerContig,
82 typename IndexT, template <typename U> class PtrTraits>
83 __host__
84 DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
85 std::initializer_list<IndexT> sizes,
86 MemorySpace space) :
87 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(nullptr, sizes),
88 state_(AllocState::Owner),
89 space_(space) {
90
91 allocMemorySpace(space, &this->data_, this->getSizeInBytes());
92 FAISS_ASSERT(this->data_ || (this->getSizeInBytes() == 0));
93 }
94
95 // memory reservation constructor
96 template <typename T, int Dim, bool InnerContig,
97 typename IndexT, template <typename U> class PtrTraits>
98 __host__
99 DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
100 DeviceMemory& m,
101 const IndexT sizes[Dim],
102 cudaStream_t stream,
103 MemorySpace space) :
104 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(nullptr, sizes),
105 state_(AllocState::Reservation),
106 space_(space) {
107
108 // FIXME: add MemorySpace to DeviceMemory
109 auto memory = m.getMemory(stream, this->getSizeInBytes());
110
111 this->data_ = (T*) memory.get();
112 FAISS_ASSERT(this->data_ || (this->getSizeInBytes() == 0));
113 reservation_ = std::move(memory);
114 }
115
116 // memory reservation constructor
117 template <typename T, int Dim, bool InnerContig,
118 typename IndexT, template <typename U> class PtrTraits>
119 __host__
120 DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
121 DeviceMemory& m,
122 std::initializer_list<IndexT> sizes,
123 cudaStream_t stream,
124 MemorySpace space) :
125 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(nullptr, sizes),
126 state_(AllocState::Reservation),
127 space_(space) {
128
129 // FIXME: add MemorySpace to DeviceMemory
130 auto memory = m.getMemory(stream, this->getSizeInBytes());
131
132 this->data_ = (T*) memory.get();
133 FAISS_ASSERT(this->data_ || (this->getSizeInBytes() == 0));
134 reservation_ = std::move(memory);
135 }
136
137 template <typename T, int Dim, bool InnerContig,
138 typename IndexT, template <typename U> class PtrTraits>
139 __host__
140 DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
141 DataPtrType data,
142 const IndexT sizes[Dim],
143 MemorySpace space) :
144 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(data, sizes),
145 state_(AllocState::NotOwner),
146 space_(space) {
147 }
148
149 template <typename T, int Dim, bool InnerContig,
150 typename IndexT, template <typename U> class PtrTraits>
151 __host__
152 DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
153 DataPtrType data,
154 std::initializer_list<IndexT> sizes,
155 MemorySpace space) :
156 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(data, sizes),
157 state_(AllocState::NotOwner),
158 space_(space) {
159 }
160
161 template <typename T, int Dim, bool InnerContig,
162 typename IndexT, template <typename U> class PtrTraits>
163 __host__
164 DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
165 DataPtrType data,
166 const IndexT sizes[Dim],
167 const IndexT strides[Dim],
168 MemorySpace space) :
169 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(data, sizes, strides),
170 state_(AllocState::NotOwner),
171 space_(space) {
172 }
173
174 template <typename T, int Dim, bool InnerContig,
175 typename IndexT, template <typename U> class PtrTraits>
176 __host__
177 DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
178 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t,
179 cudaStream_t stream,
180 MemorySpace space) :
181 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(nullptr, t.sizes(), t.strides()),
182 state_(AllocState::Owner),
183 space_(space) {
184
185 allocMemorySpace(space_, &this->data_, this->getSizeInBytes());
186 FAISS_ASSERT(this->data_ || (this->getSizeInBytes() == 0));
187 this->copyFrom(t, stream);
188 }
189
190 template <typename T, int Dim, bool InnerContig,
191 typename IndexT, template <typename U> class PtrTraits>
192 __host__
193 DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::DeviceTensor(
194 DeviceMemory& m,
195 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t,
196 cudaStream_t stream,
197 MemorySpace space) :
198 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(nullptr, t.sizes(), t.strides()),
199 state_(AllocState::Reservation),
200 space_(space) {
201
202 // FIXME: add MemorySpace to DeviceMemory
203 auto memory = m.getMemory(stream, this->getSizeInBytes());
204
205 this->data_ = (T*) memory.get();
206 FAISS_ASSERT(this->data_ || (this->getSizeInBytes() == 0));
207 reservation_ = std::move(memory);
208
209 this->copyFrom(t, stream);
210 }
211
212 template <typename T, int Dim, bool InnerContig,
213 typename IndexT, template <typename U> class PtrTraits>
214 __host__ DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&
215 DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>::zero(
216 cudaStream_t stream) {
217 if (this->data_) {
218 // Region must be contiguous
219 FAISS_ASSERT(this->isContiguous());
220
221 CUDA_VERIFY(cudaMemsetAsync(
222 this->data_, 0, this->getSizeInBytes(), stream));
223 }
224
225 return *this;
226 }
227
228 } } // namespace
faiss::gpu::DeviceTensor::~DeviceTensor
__host__ ~DeviceTensor()
Destructor.
Definition: DeviceTensor-inl.cuh:56
faiss::gpu::DeviceTensor::zero
__host__ DeviceTensor< T, Dim, InnerContig, IndexT, PtrTraits > & zero(cudaStream_t stream)
Call to zero out memory.
Definition: DeviceTensor-inl.cuh:215
faiss::gpu::Tensor::data_
DataPtrType data_
Raw pointer to where the tensor data begins.
Definition: Tensor.cuh:342
faiss::gpu::DeviceMemory::getMemory
virtual DeviceMemoryReservation getMemory(cudaStream_t stream, size_t size)=0
faiss::gpu::Tensor::copyFrom
__host__ void copyFrom(Tensor< T, Dim, InnerContig, IndexT, PtrTraits > &t, cudaStream_t stream)
Copies a tensor into ourselves; sizes must match.
Definition: Tensor-inl.cuh:130
faiss::gpu::Tensor::operator=
__host__ __device__ Tensor< T, Dim, InnerContig, IndexT, PtrTraits > & operator=(Tensor< T, Dim, InnerContig, IndexT, PtrTraits > &t)
Assignment.
Definition: Tensor-inl.cuh:48
faiss::gpu::Tensor
Our tensor type.
Definition: Tensor.cuh:28
faiss::gpu::DeviceTensor::DeviceTensor
__host__ DeviceTensor()
Default constructor.
Definition: DeviceTensor-inl.cuh:16
faiss::gpu::DeviceTensor::operator=
__host__ DeviceTensor< T, Dim, InnerContig, IndexT, PtrTraits > & operator=(DeviceTensor< T, Dim, InnerContig, IndexT, PtrTraits > &&t)
Move assignment.
Definition: DeviceTensor-inl.cuh:37
faiss::gpu::DeviceMemory
Manages temporary memory allocations on a GPU device.
Definition: DeviceMemory.h:44
faiss::gpu::DeviceTensor
Definition: DeviceTensor.cuh:22
faiss::gpu::Tensor::getSizeInBytes
__host__ __device__ size_t getSizeInBytes() const
Definition: Tensor.cuh:238