Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
Tensor-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 "../GpuFaissAssert.h"
10 #include "DeviceUtils.h"
11 #include <limits>
12
13 namespace faiss { namespace gpu {
14
15 template <typename T, int Dim, bool InnerContig,
16 typename IndexT, template <typename U> class PtrTraits>
17 __host__ __device__
18 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::Tensor()
19 : data_(nullptr) {
20 static_assert(Dim > 0, "must have > 0 dimensions");
21
22 for (int i = 0; i < Dim; ++i) {
23 size_[i] = 0;
24 stride_[i] = (IndexT) 1;
25 }
26 }
27
28 template <typename T, int Dim, bool InnerContig,
29 typename IndexT, template <typename U> class PtrTraits>
30 __host__ __device__
31 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::Tensor(
32 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t) {
33 this->operator=(t);
34 }
35
36 template <typename T, int Dim, bool InnerContig,
37 typename IndexT, template <typename U> class PtrTraits>
38 __host__ __device__
39 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::Tensor(
40 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t) {
41 this->operator=(std::move(t));
42 }
43
44 template <typename T, int Dim, bool InnerContig,
45 typename IndexT, template <typename U> class PtrTraits>
46 __host__ __device__
47 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>&
48 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::operator=(
49 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t) {
50 data_ = t.data_;
51 for (int i = 0; i < Dim; ++i) {
52 size_[i] = t.size_[i];
53 stride_[i] = t.stride_[i];
54 }
55
56 return *this;
57 }
58
59 template <typename T, int Dim, bool InnerContig,
60 typename IndexT, template <typename U> class PtrTraits>
61 __host__ __device__
62 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>&
63 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::operator=(
64 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t) {
65 data_ = t.data_; t.data_ = nullptr;
66 for (int i = 0; i < Dim; ++i) {
67 stride_[i] = t.stride_[i]; t.stride_[i] = 0;
68 size_[i] = t.size_[i]; t.size_[i] = 0;
69 }
70
71 return *this;
72 }
73
74 template <typename T, int Dim, bool InnerContig,
75 typename IndexT, template <typename U> class PtrTraits>
76 __host__ __device__
77 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::
78 Tensor(DataPtrType data, const IndexT sizes[Dim])
79 : data_(data) {
80 static_assert(Dim > 0, "must have > 0 dimensions");
81
82 for (int i = 0; i < Dim; ++i) {
83 size_[i] = sizes[i];
84 }
85
86 stride_[Dim - 1] = (IndexT) 1;
87 for (int i = Dim - 2; i >= 0; --i) {
88 stride_[i] = stride_[i + 1] * sizes[i + 1];
89 }
90 }
91
92 template <typename T, int Dim, bool InnerContig,
93 typename IndexT, template <typename U> class PtrTraits>
94 __host__ __device__
95 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::
96 Tensor(DataPtrType data, std::initializer_list<IndexT> sizes)
97 : data_(data) {
98 GPU_FAISS_ASSERT(sizes.size() == Dim);
99 static_assert(Dim > 0, "must have > 0 dimensions");
100
101 int i = 0;
102 for (auto s : sizes) {
103 size_[i++] = s;
104 }
105
106 stride_[Dim - 1] = (IndexT) 1;
107 for (int j = Dim - 2; j >= 0; --j) {
108 stride_[j] = stride_[j + 1] * size_[j + 1];
109 }
110 }
111
112
113 template <typename T, int Dim, bool InnerContig,
114 typename IndexT, template <typename U> class PtrTraits>
115 __host__ __device__
116 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::Tensor(
117 DataPtrType data, const IndexT sizes[Dim], const IndexT strides[Dim])
118 : data_(data) {
119 static_assert(Dim > 0, "must have > 0 dimensions");
120
121 for (int i = 0; i < Dim; ++i) {
122 size_[i] = sizes[i];
123 stride_[i] = strides[i];
124 }
125 }
126
127 template <typename T, int Dim, bool InnerContig,
128 typename IndexT, template <typename U> class PtrTraits>
129 __host__ void
130 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::copyFrom(
131 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t,
132 cudaStream_t stream) {
133 // The tensor must be fully contiguous
134 GPU_FAISS_ASSERT(this->isContiguous());
135
136 // Size must be the same (since dimensions are checked and
137 // continuity is assumed, we need only check total number of
138 // elements
139 GPU_FAISS_ASSERT(this->numElements() == t.numElements());
140
141 if (t.numElements() > 0) {
142 GPU_FAISS_ASSERT(this->data_);
143 GPU_FAISS_ASSERT(t.data());
144
145 int ourDev = getDeviceForAddress(this->data_);
146 int tDev = getDeviceForAddress(t.data());
147
148 if (tDev == -1) {
149 CUDA_VERIFY(cudaMemcpyAsync(this->data_,
150 t.data(),
151 this->getSizeInBytes(),
152 ourDev == -1 ? cudaMemcpyHostToHost :
153 cudaMemcpyHostToDevice,
154 stream));
155 } else {
156 CUDA_VERIFY(cudaMemcpyAsync(this->data_,
157 t.data(),
158 this->getSizeInBytes(),
159 ourDev == -1 ? cudaMemcpyDeviceToHost :
160 cudaMemcpyDeviceToDevice,
161 stream));
162 }
163 }
164 }
165
166 template <typename T, int Dim, bool InnerContig,
167 typename IndexT, template <typename U> class PtrTraits>
168 __host__ void
169 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::copyTo(
170 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t,
171 cudaStream_t stream) {
172 // The tensor must be fully contiguous
173 GPU_FAISS_ASSERT(this->isContiguous());
174
175 // Size must be the same (since dimensions are checked and
176 // continuity is assumed, we need only check total number of
177 // elements
178 GPU_FAISS_ASSERT(this->numElements() == t.numElements());
179
180 if (t.numElements() > 0) {
181 GPU_FAISS_ASSERT(this->data_);
182 GPU_FAISS_ASSERT(t.data());
183
184 int ourDev = getDeviceForAddress(this->data_);
185 int tDev = getDeviceForAddress(t.data());
186
187 if (tDev == -1) {
188 CUDA_VERIFY(cudaMemcpyAsync(t.data(),
189 this->data_,
190 this->getSizeInBytes(),
191 ourDev == -1 ? cudaMemcpyHostToHost :
192 cudaMemcpyDeviceToHost,
193 stream));
194 } else {
195 CUDA_VERIFY(cudaMemcpyAsync(t.data(),
196 this->data_,
197 this->getSizeInBytes(),
198 ourDev == -1 ? cudaMemcpyHostToDevice :
199 cudaMemcpyDeviceToDevice,
200 stream));
201 }
202 }
203 }
204
205 template <typename T, int Dim, bool InnerContig,
206 typename IndexT, template <typename U> class PtrTraits>
207 template <typename OtherT, int OtherDim>
208 __host__ __device__ bool
209 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::isSame(
210 const Tensor<OtherT, OtherDim, InnerContig, IndexT, PtrTraits>& rhs) const {
211 if (Dim != OtherDim) {
212 return false;
213 }
214
215 for (int i = 0; i < Dim; ++i) {
216 if (this->getSize(i) != rhs.getSize(i)) {
217 return false;
218 }
219
220 if (this->getStride(i) != rhs.getStride(i)) {
221 return false;
222 }
223 }
224
225 return true;
226 }
227
228 template <typename T, int Dim, bool InnerContig,
229 typename IndexT, template <typename U> class PtrTraits>
230 template <typename OtherT, int OtherDim>
231 __host__ __device__ bool
232 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::isSameSize(
233 const Tensor<OtherT, OtherDim, InnerContig, IndexT, PtrTraits>& rhs) const {
234 if (Dim != OtherDim) {
235 return false;
236 }
237
238 for (int i = 0; i < Dim; ++i) {
239 if (this->getSize(i) != rhs.getSize(i)) {
240 return false;
241 }
242 }
243
244 return true;
245 }
246
247 template <typename T, int Dim, bool InnerContig,
248 typename IndexT, template <typename U> class PtrTraits>
249 template <typename U>
250 __host__ __device__ Tensor<U, Dim, InnerContig, IndexT, PtrTraits>
251 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::cast() {
252 static_assert(sizeof(U) == sizeof(T), "cast must be to same size object");
253
254 return Tensor<U, Dim, InnerContig, IndexT, PtrTraits>(
255 reinterpret_cast<U*>(data_), size_, stride_);
256 }
257
258 template <typename T, int Dim, bool InnerContig,
259 typename IndexT, template <typename U> class PtrTraits>
260 template <typename U>
261 __host__ __device__ const Tensor<U, Dim, InnerContig, IndexT, PtrTraits>
262 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::cast() const {
263 static_assert(sizeof(U) == sizeof(T), "cast must be to same size object");
264
265 return Tensor<U, Dim, InnerContig, IndexT, PtrTraits>(
266 reinterpret_cast<U*>(data_), size_, stride_);
267 }
268
269 template <typename T, int Dim, bool InnerContig,
270 typename IndexT, template <typename U> class PtrTraits>
271 template <typename U>
272 __host__ __device__ Tensor<U, Dim, InnerContig, IndexT, PtrTraits>
273 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::castResize() {
274 static_assert(sizeof(U) >= sizeof(T), "only handles greater sizes");
275 constexpr int kMultiple = sizeof(U) / sizeof(T);
276
277 GPU_FAISS_ASSERT(canCastResize<U>());
278
279 IndexT newSize[Dim];
280 IndexT newStride[Dim];
281
282 for (int i = 0; i < Dim - 1; ++i) {
283 newSize[i] = size_[i];
284 newStride[i] = stride_[i] / kMultiple;
285 }
286
287 newStride[Dim - 1] = 1; // this is the same as the old stride
288 newSize[Dim - 1] = size_[Dim - 1] / kMultiple;
289
290 return Tensor<U, Dim, InnerContig, IndexT, PtrTraits>(
291 reinterpret_cast<U*>(data_), newSize, newStride);
292 }
293
294 template <typename T, int Dim, bool InnerContig,
295 typename IndexT, template <typename U> class PtrTraits>
296 template <typename U>
297 __host__ __device__ const Tensor<U, Dim, InnerContig, IndexT, PtrTraits>
298 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::castResize() const {
299 return const_cast<Tensor<T, Dim, InnerContig, IndexT, PtrTraits>*>(this)->
300 castResize<U>();
301 }
302
303 template <typename T, int Dim, bool InnerContig,
304 typename IndexT, template <typename U> class PtrTraits>
305 template <typename U>
306 __host__ __device__ bool
307 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::canCastResize() const {
308 static_assert(sizeof(U) >= sizeof(T), "only handles greater sizes");
309 constexpr int kMultiple = sizeof(U) / sizeof(T);
310
311 // Ensure that the base pointer is sizeof(U) aligned
312 if (((uintptr_t) data_) % sizeof(U) != 0) {
313 return false;
314 }
315
316 // Check all outer strides
317 for (int i = 0; i < Dim - 1; ++i) {
318 if (stride_[i] % kMultiple != 0) {
319 return false;
320 }
321 }
322
323 // Check inner size
324 if (size_[Dim - 1] % kMultiple != 0) {
325 return false;
326 }
327
328 if (stride_[Dim - 1] != 1) {
329 return false;
330 }
331
332 return true;
333 }
334
335 template <typename T, int Dim, bool InnerContig,
336 typename IndexT, template <typename U> class PtrTraits>
337 template <typename NewIndexT>
338 __host__ Tensor<T, Dim, InnerContig, NewIndexT, PtrTraits>
339 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::castIndexType() const {
340 if (sizeof(NewIndexT) < sizeof(IndexT)) {
341 GPU_FAISS_ASSERT(this->canUseIndexType<NewIndexT>());
342 }
343
344 NewIndexT newSize[Dim];
345 NewIndexT newStride[Dim];
346 for (int i = 0; i < Dim; ++i) {
347 newSize[i] = (NewIndexT) size_[i];
348 newStride[i] = (NewIndexT) stride_[i];
349 }
350
351 return Tensor<T, Dim, InnerContig, NewIndexT, PtrTraits>(
352 data_, newSize, newStride);
353 }
354
355 template <typename T, int Dim, bool InnerContig,
356 typename IndexT, template <typename U> class PtrTraits>
357 template <typename NewIndexT>
358 __host__ bool
359 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::canUseIndexType() const {
360 static_assert(sizeof(size_t) >= sizeof(IndexT),
361 "index size too large");
362 static_assert(sizeof(size_t) >= sizeof(NewIndexT),
363 "new index size too large");
364
365 // Find maximum offset that can be calculated
366 // FIXME: maybe also consider offset in bytes? multiply by sizeof(T)?
367 size_t maxOffset = 0;
368
369 for (int i = 0; i < Dim; ++i) {
370 size_t curMaxOffset = (size_t) size_[i] * (size_t) stride_[i];
371 if (curMaxOffset > maxOffset) {
372 maxOffset = curMaxOffset;
373 }
374 }
375
376 if (maxOffset > (size_t) std::numeric_limits<NewIndexT>::max()) {
377 return false;
378 }
379
380 return true;
381 }
382
383 template <typename T, int Dim, bool InnerContig,
384 typename IndexT, template <typename U> class PtrTraits>
385 __host__ __device__ size_t
386 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::numElements() const {
387 size_t size = (size_t) getSize(0);
388
389 for (int i = 1; i < Dim; ++i) {
390 size *= (size_t) getSize(i);
391 }
392
393 return size;
394 }
395
396 template <typename T, int Dim, bool InnerContig,
397 typename IndexT, template <typename U> class PtrTraits>
398 __host__ __device__ bool
399 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::isContiguous() const {
400 long prevSize = 1;
401
402 for (int i = Dim - 1; i >= 0; --i) {
403 if (getSize(i) != (IndexT) 1) {
404 if (getStride(i) == prevSize) {
405 prevSize *= getSize(i);
406 } else {
407 return false;
408 }
409 }
410 }
411
412 return true;
413 }
414
415 template <typename T, int Dim, bool InnerContig,
416 typename IndexT, template <typename U> class PtrTraits>
417 __host__ __device__ bool
418 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::isConsistentlySized(int i) const {
419 if (i == 0 && getStride(i) > 0 && getSize(i) > 0) {
420 return true;
421 } else if ((i > 0) && (i < Dim) && (getStride(i) > 0) &&
422 ((getStride(i - 1) / getStride(i)) >= getSize(i))) {
423 return true;
424 }
425
426 return false;
427 }
428
429 template <typename T, int Dim, bool InnerContig,
430 typename IndexT, template <typename U> class PtrTraits>
431 __host__ __device__ bool
432 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::isConsistentlySized() const {
433 for (int i = 0; i < Dim; ++i) {
434 if (!isConsistentlySized(i)) {
435 return false;
436 }
437 }
438
439 return true;
440 }
441
442 template <typename T, int Dim, bool InnerContig,
443 typename IndexT, template <typename U> class PtrTraits>
444 __host__ __device__ bool
445 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::isContiguousDim(int i) const {
446 return (i == Dim - 1) || // just in case
447 ((i < Dim - 1) &&
448 ((getStride(i) / getStride(i + 1)) == getSize(i + 1)));
449 }
450
451 template <typename T, int Dim, bool InnerContig,
452 typename IndexT, template <typename U> class PtrTraits>
453 __host__ __device__ Tensor<T, Dim, InnerContig, IndexT, PtrTraits>
454 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::transpose(int dim1,
455 int dim2) const {
456 GPU_FAISS_ASSERT(dim1 >= 0 && dim1 < Dim);
457 GPU_FAISS_ASSERT(dim1 >= 0 && dim2 < Dim);
458
459 // If a tensor is innermost contiguous, one cannot transpose the innermost
460 // dimension
461 if (InnerContig) {
462 GPU_FAISS_ASSERT(dim1 != Dim - 1 && dim2 != Dim - 1);
463 }
464
465 IndexT newSize[Dim];
466 IndexT newStride[Dim];
467
468 for (int i = 0; i < Dim; ++i) {
469 newSize[i] = size_[i];
470 newStride[i] = stride_[i];
471 }
472
473 IndexT tmp = newSize[dim1];
474 newSize[dim1] = newSize[dim2];
475 newSize[dim2] = tmp;
476
477 tmp = newStride[dim1];
478 newStride[dim1] = newStride[dim2];
479 newStride[dim2] = tmp;
480
481 return Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(data_, newSize, newStride);
482 }
483
484 template <typename T, int Dim, bool InnerContig,
485 typename IndexT, template <typename U> class PtrTraits>
486 template <int NewDim>
487 __host__ __device__ Tensor<T, NewDim, InnerContig, IndexT, PtrTraits>
488 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::upcastOuter() {
489 // Can only create tensors of greater dimension
490 static_assert(NewDim > Dim, "Can only upcast to greater dim");
491
492 IndexT newSize[NewDim];
493 IndexT newStride[NewDim];
494
495 int shift = NewDim - Dim;
496
497 for (int i = 0; i < NewDim; ++i) {
498 if (i < shift) {
499 // These are the extended dimensions
500 newSize[i] = (IndexT) 1;
501 newStride[i] = size_[0] * stride_[0];
502 } else {
503 // Shift the remaining dimensions
504 newSize[i] = size_[i - shift];
505 newStride[i] = stride_[i - shift];
506 }
507 }
508
509 return Tensor<T, NewDim, InnerContig, IndexT, PtrTraits>(
510 data_, newSize, newStride);
511 }
512
513 template <typename T, int Dim, bool InnerContig,
514 typename IndexT, template <typename U> class PtrTraits>
515 template <int NewDim>
516 __host__ __device__ Tensor<T, NewDim, InnerContig, IndexT, PtrTraits>
517 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::upcastInner() {
518 // Can only create tensors of greater dimension
519 static_assert(NewDim > Dim, "Can only upcast to greater dim");
520
521 IndexT newSize[NewDim];
522 IndexT newStride[NewDim];
523
524 for (int i = 0; i < NewDim; ++i) {
525 if (i < Dim) {
526 // Existing dimensions get copied over
527 newSize[i] = size_[i];
528 newStride[i] = stride_[i];
529 } else {
530 // Extended dimensions
531 newSize[i] = (IndexT) 1;
532 newStride[i] = (IndexT) 1;
533 }
534 }
535
536 return Tensor<T, NewDim, InnerContig, IndexT, PtrTraits>(
537 data_, newSize, newStride);
538 }
539
540 template <typename T, int Dim, bool InnerContig,
541 typename IndexT, template <typename U> class PtrTraits>
542 template <int NewDim>
543 __host__ __device__ Tensor<T, NewDim, InnerContig, IndexT, PtrTraits>
544 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::downcastOuter() {
545 // Can only create tensors of lesser dimension
546 static_assert(NewDim < Dim, "Can only downcast to lesser dim");
547
548 // We can't downcast non-contiguous tensors, since it leaves
549 // garbage data in the tensor. The tensor needs to be contiguous
550 // in all of the dimensions we are collapsing (no padding in
551 // them).
552 for (int i = 0; i < Dim - NewDim; ++i) {
553 bool cont = isContiguousDim(i);
554 GPU_FAISS_ASSERT(cont);
555 }
556
557 IndexT newSize[NewDim];
558 IndexT newStride[NewDim];
559
560 int ignoredDims = Dim - NewDim;
561 IndexT collapsedSize = 1;
562
563 for (int i = 0; i < Dim; ++i) {
564 if (i < ignoredDims) {
565 // Collapse these dimensions
566 collapsedSize *= getSize(i);
567 } else {
568 // Non-collapsed dimensions
569 if (i == ignoredDims) {
570 // This is the first non-collapsed dimension
571 newSize[i - ignoredDims] = collapsedSize * getSize(i);
572 } else {
573 // Subsequent non-collapsed dimensions
574 newSize[i - ignoredDims] = getSize(i);
575 }
576
577 newStride[i - ignoredDims] = getStride(i);
578 }
579 }
580
581 return Tensor<T, NewDim, InnerContig, IndexT, PtrTraits>(
582 data_, newSize, newStride);
583 }
584
585 template <typename T, int Dim, bool InnerContig,
586 typename IndexT, template <typename U> class PtrTraits>
587 template <int NewDim>
588 __host__ __device__ Tensor<T, NewDim, InnerContig, IndexT, PtrTraits>
589 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::downcastInner() {
590 // Can only create tensors of lesser dimension
591 static_assert(NewDim < Dim, "Can only downcast to lesser dim");
592
593 // We can't downcast non-contiguous tensors, since it leaves
594 // garbage data in the tensor. The tensor needs to be contiguous
595 // in all of the dimensions we are collapsing (no padding in
596 // them).
597 for (int i = NewDim; i < Dim; ++i) {
598 GPU_FAISS_ASSERT(isContiguousDim(i));
599 }
600
601 IndexT newSize[NewDim];
602 IndexT newStride[NewDim];
603
604 IndexT collapsedSize = 1;
605
606 for (int i = Dim - 1; i >= 0; --i) {
607 if (i >= NewDim) {
608 // Collapse these dimensions
609 collapsedSize *= getSize(i);
610 } else {
611 // Non-collapsed dimensions
612 if (i == NewDim - 1) {
613 // This is the first non-collapsed dimension
614 newSize[i] = collapsedSize * getSize(i);
615 newStride[i] = getStride(Dim - 1);
616 } else {
617 // Subsequent non-collapsed dimensions
618 newSize[i] = getSize(i);
619 newStride[i] = getStride(i);
620 }
621 }
622 }
623
624 return Tensor<T, NewDim, InnerContig, IndexT, PtrTraits>(
625 data_, newSize, newStride);
626 }
627
628 template <typename T, int Dim, bool InnerContig,
629 typename IndexT, template <typename U> class PtrTraits>
630 template <int SubDim>
631 __host__ __device__ Tensor<T, SubDim, InnerContig, IndexT, PtrTraits>
632 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::view(DataPtrType at) {
633 static_assert(SubDim >= 1 && SubDim < Dim,
634 "can only create view of lesser dim");
635
636 IndexT viewSizes[SubDim];
637 IndexT viewStrides[SubDim];
638
639 for (int i = 0; i < SubDim; ++i) {
640 viewSizes[i] = size_[Dim - SubDim + i];
641 viewStrides[i] = stride_[Dim - SubDim + i];
642 }
643
644 return Tensor<T, SubDim, InnerContig, IndexT, PtrTraits>(
645 at, viewSizes, viewStrides);
646 }
647
648 template <typename T, int Dim, bool InnerContig,
649 typename IndexT, template <typename U> class PtrTraits>
650 template <int SubDim>
651 __host__ __device__ Tensor<T, SubDim, InnerContig, IndexT, PtrTraits>
652 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::view() {
653 return view<SubDim>(data_);
654 }
655
656 template <typename T, int Dim, bool InnerContig,
657 typename IndexT, template <typename U> class PtrTraits>
658 __host__ __device__ Tensor<T, Dim, InnerContig, IndexT, PtrTraits>
659 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::narrowOutermost(IndexT start,
660 IndexT size) {
661 return this->narrow(0, start, size);
662 }
663
664 template <typename T, int Dim, bool InnerContig,
665 typename IndexT, template <typename U> class PtrTraits>
666 __host__ __device__ Tensor<T, Dim, InnerContig, IndexT, PtrTraits>
667 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::narrow(int dim,
668 IndexT start,
669 IndexT size) {
670 DataPtrType newData = data_;
671
672 GPU_FAISS_ASSERT(start >= 0 &&
673 start < size_[dim] &&
674 (start + size) <= size_[dim]);
675
676 if (start > 0) {
677 newData += (size_t) start * stride_[dim];
678 }
679
680 IndexT newSize[Dim];
681 for (int i = 0; i < Dim; ++i) {
682 if (i == dim) {
683 GPU_FAISS_ASSERT(start + size <= size_[dim]);
684 newSize[i] = size;
685 } else {
686 newSize[i] = size_[i];
687 }
688 }
689
690 // If we were innermost contiguous before, we are still innermost contiguous
691 return Tensor<T, Dim, InnerContig, IndexT, PtrTraits>(newData, newSize, stride_);
692 }
693
694 template <typename T, int Dim, bool InnerContig,
695 typename IndexT, template <typename U> class PtrTraits>
696 template <int NewDim>
697 __host__ __device__ Tensor<T, NewDim, InnerContig, IndexT, PtrTraits>
698 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::view(
699 std::initializer_list<IndexT> sizes) {
700 GPU_FAISS_ASSERT(this->isContiguous());
701
702 GPU_FAISS_ASSERT(sizes.size() == NewDim);
703
704 // The total size of the new view must be the same as the total size
705 // of the old view
706 size_t curSize = numElements();
707 size_t newSize = 1;
708
709 for (auto s : sizes) {
710 newSize *= s;
711 }
712
713 GPU_FAISS_ASSERT(curSize == newSize);
714 return Tensor<T, NewDim, true, IndexT, PtrTraits>(data(), sizes);
715 }
716
717 } } // namespace
faiss::gpu::Tensor::upcastOuter
__host__ __device__ Tensor< T, NewDim, InnerContig, IndexT, PtrTraits > upcastOuter()
Definition: Tensor-inl.cuh:488
faiss::gpu::Tensor::castIndexType
__host__ Tensor< T, Dim, InnerContig, NewIndexT, PtrTraits > castIndexType() const
Definition: Tensor-inl.cuh:339
faiss::gpu::Tensor::isContiguousDim
__host__ __device__ bool isContiguousDim(int i) const
Returns true if the given dimension index has no padding.
Definition: Tensor-inl.cuh:445
faiss::gpu::Tensor::cast
__host__ __device__ Tensor< U, Dim, InnerContig, IndexT, PtrTraits > cast()
Definition: Tensor-inl.cuh:251
faiss::gpu::Tensor::numElements
__host__ __device__ size_t numElements() const
Definition: Tensor-inl.cuh:386
faiss::gpu::Tensor::downcastOuter
__host__ __device__ Tensor< T, NewDim, InnerContig, IndexT, PtrTraits > downcastOuter()
Definition: Tensor-inl.cuh:544
faiss::gpu::Tensor::canCastResize
__host__ __device__ bool canCastResize() const
Returns true if we can castResize() this tensor to the new type.
Definition: Tensor-inl.cuh:307
faiss::gpu::Tensor::data_
DataPtrType data_
Raw pointer to where the tensor data begins.
Definition: Tensor.cuh:342
faiss::gpu::Tensor::Tensor
__host__ __device__ Tensor()
Default constructor.
Definition: Tensor-inl.cuh:18
faiss::gpu::Tensor::upcastInner
__host__ __device__ Tensor< T, NewDim, InnerContig, IndexT, PtrTraits > upcastInner()
Definition: Tensor-inl.cuh:517
faiss::gpu::Tensor::narrowOutermost
__host__ __device__ Tensor< T, Dim, InnerContig, IndexT, PtrTraits > narrowOutermost(IndexT start, IndexT size)
Definition: Tensor-inl.cuh:659
faiss::gpu::Tensor::stride_
IndexT stride_[Dim]
Array of strides (in sizeof(T) terms) per each dimension.
Definition: Tensor.cuh:345
faiss::gpu::Tensor::isContiguous
__host__ __device__ bool isContiguous() const
Definition: Tensor-inl.cuh:399
faiss::gpu::Tensor::sizes
__host__ __device__ const IndexT * sizes() const
Returns the size array.
Definition: Tensor.cuh:243
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::size_
IndexT size_[Dim]
Size per each dimension.
Definition: Tensor.cuh:348
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::strides
__host__ __device__ const IndexT * strides() const
Returns the stride array.
Definition: Tensor.cuh:248
faiss::gpu::Tensor::getSize
__host__ __device__ IndexT getSize(int i) const
Definition: Tensor.cuh:222
faiss::gpu::Tensor::isSameSize
__host__ __device__ bool isSameSize(const Tensor< OtherT, OtherDim, InnerContig, IndexT, PtrTraits > &rhs) const
Returns true if the two tensors are of the same dimensionality and size.
Definition: Tensor-inl.cuh:232
faiss::gpu::Tensor::downcastInner
__host__ __device__ Tensor< T, NewDim, InnerContig, IndexT, PtrTraits > downcastInner()
Definition: Tensor-inl.cuh:589
faiss::gpu::Tensor::narrow
__host__ __device__ Tensor< T, Dim, InnerContig, IndexT, PtrTraits > narrow(int dim, IndexT start, IndexT size)
Definition: Tensor-inl.cuh:667
faiss::gpu::Tensor::data
__host__ __device__ DataPtrType data()
Returns a raw pointer to the start of our data.
Definition: Tensor.cuh:174
faiss::gpu::Tensor::copyTo
__host__ void copyTo(Tensor< T, Dim, InnerContig, IndexT, PtrTraits > &t, cudaStream_t stream)
Copies ourselves into a tensor; sizes must match.
Definition: Tensor-inl.cuh:169
faiss::gpu::Tensor
Our tensor type.
Definition: Tensor.cuh:28
faiss::gpu::Tensor::canUseIndexType
__host__ bool canUseIndexType() const
Definition: Tensor-inl.cuh:359
faiss::gpu::Tensor::transpose
__host__ __device__ Tensor< T, Dim, InnerContig, IndexT, PtrTraits > transpose(int dim1, int dim2) const
Definition: Tensor-inl.cuh:454
faiss::gpu::Tensor::getStride
__host__ __device__ IndexT getStride(int i) const
Definition: Tensor.cuh:228
faiss::gpu::Tensor::castResize
__host__ __device__ Tensor< U, Dim, InnerContig, IndexT, PtrTraits > castResize()
Definition: Tensor-inl.cuh:273
faiss::gpu::Tensor::view
__host__ __device__ Tensor< T, SubDim, InnerContig, IndexT, PtrTraits > view()
Definition: Tensor-inl.cuh:652
faiss::gpu::Tensor::isSame
__host__ __device__ bool isSame(const Tensor< OtherT, OtherDim, InnerContig, IndexT, PtrTraits > &rhs) const
Definition: Tensor-inl.cuh:209