Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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