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 // 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
Generated by   doxygen 1.8.5