Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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:343
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:238
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:102
faiss::gpu::DeviceMemory
Manages temporary memory allocations on a GPU device.
Definition: DeviceMemory.h:47
faiss::gpu::DeviceTensor
Definition: DeviceTensor.cuh:24