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