Faiss
 All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
DeviceTensor-inl.cuh
1 
2 /**
3  * Copyright (c) 2015-present, Facebook, Inc.
4  * All rights reserved.
5  *
6  * This source code is licensed under the CC-by-NC license found in the
7  * LICENSE file in the root directory of this source tree.
8  */
9 
10 // Copyright 2004-present Facebook. All Rights Reserved.
11 
12 #include <utility> // std::move
13 
14 namespace faiss { namespace gpu {
15 
16 template <typename T, int Dim, bool Contig,
17  typename IndexT, template <typename U> class PtrTraits>
18 __host__
19 DeviceTensor<T, Dim, Contig, IndexT, PtrTraits>::DeviceTensor() :
20  Tensor<T, Dim, Contig, IndexT, PtrTraits>(),
21  state_(AllocState::NotOwner) {
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  this->operator=(std::move(t));
32 }
33 
34 template <typename T, int Dim, bool Contig,
35  typename IndexT, template <typename U> class PtrTraits>
36 __host__
37 DeviceTensor<T, Dim, Contig, IndexT, PtrTraits>&
38 DeviceTensor<T, Dim, Contig, IndexT, PtrTraits>::operator=(
39  DeviceTensor<T, Dim, Contig, IndexT, PtrTraits>&& t) {
40  if (this->state_ == AllocState::Owner) {
41  CUDA_VERIFY(cudaFree(this->data_));
42  }
43 
44  this->Tensor<T, Dim, Contig, IndexT, PtrTraits>::operator=(
45  std::move(t));
46 
47  this->state_ = t.state_; t.state_ = AllocState::NotOwner;
48  this->reservation_ = std::move(t.reservation_);
49 
50  return *this;
51 }
52 
53 template <typename T, int Dim, bool Contig,
54  typename IndexT, template <typename U> class PtrTraits>
55 __host__
56 DeviceTensor<T, Dim, Contig, 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 // cudaMalloc constructor
68 template <typename T, int Dim, bool Contig,
69  typename IndexT, template <typename U> class PtrTraits>
70 __host__
71 DeviceTensor<T, Dim, Contig, IndexT, PtrTraits>::DeviceTensor(
72  const IndexT sizes[Dim]) :
73  Tensor<T, Dim, Contig, IndexT, PtrTraits>(nullptr, sizes),
74  state_(AllocState::Owner) {
75 
76  CUDA_VERIFY(cudaMalloc(&this->data_, this->getSizeInBytes()));
77  FAISS_ASSERT(this->data_ || (this->getSizeInBytes() == 0));
78 }
79 
80 template <typename T, int Dim, bool Contig,
81  typename IndexT, template <typename U> class PtrTraits>
82 __host__
83 DeviceTensor<T, Dim, Contig, IndexT, PtrTraits>::DeviceTensor(
84  std::initializer_list<IndexT> sizes) :
85  Tensor<T, Dim, Contig, IndexT, PtrTraits>(nullptr, sizes),
86  state_(AllocState::Owner) {
87 
88  CUDA_VERIFY(cudaMalloc(&this->data_, this->getSizeInBytes()));
89  if(!(this->data_ || this->getSizeInBytes() == 0)) {
90  fprintf(stderr, "could not cudaMalloc %ld bytes!\n", this->getSizeInBytes());
91  abort();
92  }
93 }
94 
95 // memory reservation constructor
96 template <typename T, int Dim, bool Contig,
97  typename IndexT, template <typename U> class PtrTraits>
98 __host__
99 DeviceTensor<T, Dim, Contig, IndexT, PtrTraits>::DeviceTensor(
100  DeviceMemory& m,
101  const IndexT sizes[Dim],
102  cudaStream_t stream) :
103  Tensor<T, Dim, Contig, IndexT, PtrTraits>(nullptr, sizes),
104  state_(AllocState::Reservation) {
105 
106  auto memory = m.getMemory(stream, this->getSizeInBytes());
107 
108  this->data_ = (T*) memory.get();
109  FAISS_ASSERT(this->data_ || (this->getSizeInBytes() == 0));
110  reservation_ = std::move(memory);
111 }
112 
113 // memory reservation constructor
114 template <typename T, int Dim, bool Contig,
115  typename IndexT, template <typename U> class PtrTraits>
116 __host__
117 DeviceTensor<T, Dim, Contig, IndexT, PtrTraits>::DeviceTensor(
118  DeviceMemory& m,
119  std::initializer_list<IndexT> sizes,
120  cudaStream_t stream) :
121  Tensor<T, Dim, Contig, IndexT, PtrTraits>(nullptr, sizes),
122  state_(AllocState::Reservation) {
123 
124  auto memory =
125  m.getMemory(stream, this->getSizeInBytes());
126 
127  this->data_ = (T*) memory.get();
128  FAISS_ASSERT(this->data_ || (this->getSizeInBytes() == 0));
129  reservation_ = std::move(memory);
130 }
131 
132 template <typename T, int Dim, bool Contig,
133  typename IndexT, template <typename U> class PtrTraits>
134 __host__
135 DeviceTensor<T, Dim, Contig, IndexT, PtrTraits>::DeviceTensor(
136  DataPtrType data,
137  const IndexT sizes[Dim]) :
138  Tensor<T, Dim, Contig, IndexT, PtrTraits>(data, sizes),
139  state_(AllocState::NotOwner) {
140 }
141 
142 template <typename T, int Dim, bool Contig,
143  typename IndexT, template <typename U> class PtrTraits>
144 __host__
145 DeviceTensor<T, Dim, Contig, IndexT, PtrTraits>::DeviceTensor(
146  DataPtrType data,
147  std::initializer_list<IndexT> sizes) :
148  Tensor<T, Dim, Contig, IndexT, PtrTraits>(data, sizes),
149  state_(AllocState::NotOwner) {
150 }
151 
152 template <typename T, int Dim, bool Contig,
153  typename IndexT, template <typename U> class PtrTraits>
154 __host__
155 DeviceTensor<T, Dim, Contig, IndexT, PtrTraits>::DeviceTensor(
156  DataPtrType data,
157  const IndexT sizes[Dim],
158  const IndexT strides[Dim]) :
159  Tensor<T, Dim, Contig, IndexT, PtrTraits>(data, sizes, strides),
160  state_(AllocState::NotOwner) {
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  Tensor<T, Dim, Contig, IndexT, PtrTraits>& t,
168  cudaStream_t stream) :
169  Tensor<T, Dim, Contig, IndexT, PtrTraits>(nullptr, t.sizes(), t.strides()),
170  state_(AllocState::Owner) {
171 
172  CUDA_VERIFY(cudaMalloc(&this->data_, this->getSizeInBytes()));
173  FAISS_ASSERT(this->data_ || (this->getSizeInBytes() == 0));
174  this->copyFrom(t, stream);
175 }
176 
177 template <typename T, int Dim, bool Contig,
178  typename IndexT, template <typename U> class PtrTraits>
179 __host__
180 DeviceTensor<T, Dim, Contig, IndexT, PtrTraits>::DeviceTensor(
181  DeviceMemory& m,
182  Tensor<T, Dim, Contig, IndexT, PtrTraits>& t,
183  cudaStream_t stream) :
184  Tensor<T, Dim, Contig, IndexT, PtrTraits>(nullptr, t.sizes(), t.strides()),
185  state_(AllocState::Reservation) {
186 
187  auto memory = m.getMemory(stream, this->getSizeInBytes());
188 
189  this->data_ = (T*) memory.get();
190  FAISS_ASSERT(this->data_ || (this->getSizeInBytes() == 0));
191  reservation_ = std::move(memory);
192 
193  this->copyFrom(t, stream);
194 }
195 
196 template <typename T, int Dim, bool Contig,
197  typename IndexT, template <typename U> class PtrTraits>
198 __host__ DeviceTensor<T, Dim, Contig, IndexT, PtrTraits>&
199 DeviceTensor<T, Dim, Contig, IndexT, PtrTraits>::zero(
200  cudaStream_t stream) {
201  if (this->data_) {
202  // Region must be contiguous
203  FAISS_ASSERT(this->isContiguous());
204 
205  CUDA_VERIFY(cudaMemsetAsync(
206  this->data_, 0, this->getSizeInBytes(), stream));
207  }
208 
209  return *this;
210 }
211 
212 } } // 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:38
faiss::gpu::DeviceTensor::DeviceTensor
__host__ DeviceTensor()
Default constructor.
Definition: DeviceTensor-inl.cuh:19
faiss::gpu::Tensor::data_
DataPtrType data_
Raw pointer to where the tensor data begins.
Definition: Tensor.cuh:331
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:199
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:56
faiss::gpu::Tensor
Our tensor type.
Definition: Tensor.cuh:31
faiss::gpu::Tensor::getSizeInBytes
__host__ __device__ size_t getSizeInBytes() const
Definition: Tensor.cuh:226
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:47
Generated by   doxygen 1.8.5