Faiss
 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::Tensor::getSizeInBytes
__host__ __device__ size_t getSizeInBytes() const
Definition: Tensor.cuh:238
Generated by   doxygen 1.8.5