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