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