Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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