Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
Tensor.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 #pragma once
12
13 #include <assert.h>
14 #include <cuda.h>
15 #include <cuda_runtime.h>
16 #include <initializer_list>
17
18 /// Multi-dimensional array class for CUDA device and host usage.
19 /// Originally from Facebook's fbcunn, since added to the Torch GPU
20 /// library cutorch as well.
21
22 namespace faiss { namespace gpu {
23
24 /// Our tensor type
25 template <typename T,
26 int Dim,
27 bool InnerContig,
28 typename IndexT,
29 template <typename U> class PtrTraits>
30 class Tensor;
31
32 /// Type of a subspace of a tensor
33 namespace detail {
34 template <typename TensorType,
35 int SubDim,
36 template <typename U> class PtrTraits>
37 class SubTensor;
38 }
39
40 namespace traits {
41
42 template <typename T>
43 struct RestrictPtrTraits {
44 typedef T* __restrict__ PtrType;
45 };
46
47 template <typename T>
48 struct DefaultPtrTraits {
49 typedef T* PtrType;
50 };
51
52 }
53
54 /**
55 Templated multi-dimensional array that supports strided access of
56 elements. Main access is through `operator[]`; e.g.,
57 `tensor[x][y][z]`.
58
59 - `T` is the contained type (e.g., `float`)
60 - `Dim` is the tensor rank
61 - If `InnerContig` is true, then the tensor is assumed to be innermost
62 - contiguous, and only operations that make sense on contiguous
63 - arrays are allowed (e.g., no transpose). Strides are still
64 - calculated, but innermost stride is assumed to be 1.
65 - `IndexT` is the integer type used for size/stride arrays, and for
66 - all indexing math. Default is `int`, but for large tensors, `long`
67 - can be used instead.
68 - `PtrTraits` are traits applied to our data pointer (T*). By default,
69 - this is just T*, but RestrictPtrTraits can be used to apply T*
70 - __restrict__ for alias-free analysis.
71 */
72 template <typename T,
73 int Dim,
74 bool InnerContig = false,
75 typename IndexT = int,
76 template <typename U> class PtrTraits = traits::DefaultPtrTraits>
77 class Tensor {
78 public:
79 enum { NumDim = Dim };
80 typedef T DataType;
81 typedef IndexT IndexType;
82 enum { IsInnerContig = InnerContig };
83 typedef typename PtrTraits<T>::PtrType DataPtrType;
84 typedef Tensor<T, Dim, InnerContig, IndexT, PtrTraits> TensorType;
85
86 /// Default constructor
87 __host__ __device__ Tensor();
88
89 /// Copy constructor
90 __host__ __device__ Tensor(Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t)
91 = default;
92
93 /// Move constructor
94 __host__ __device__ Tensor(Tensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t)
95 = default;
96
97 /// Assignment
98 __host__ __device__ Tensor<T, Dim, InnerContig, IndexT, PtrTraits>&
99 operator=(Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t) = default;
100
101 /// Move assignment
102 __host__ __device__ Tensor<T, Dim, InnerContig, IndexT, PtrTraits>&
103 operator=(Tensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t);
104
105 /// Constructor that calculates strides with no padding
106 __host__ __device__ Tensor(DataPtrType data,
107 const IndexT sizes[Dim]);
108 __host__ __device__ Tensor(DataPtrType data,
109 std::initializer_list<IndexT> sizes);
110
111 /// Constructor that takes arbitrary size/stride arrays.
112 /// Errors if you attempt to pass non-contiguous strides to a
113 /// contiguous tensor.
114 __host__ __device__ Tensor(DataPtrType data,
115 const IndexT sizes[Dim],
116 const IndexT strides[Dim]);
117
118 /// Copies a tensor into ourselves; sizes must match
119 __host__ void copyFrom(Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t,
120 cudaStream_t stream);
121
122 /// Copies ourselves into a tensor; sizes must match
123 __host__ void copyTo(Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t,
124 cudaStream_t stream);
125
126 /// Returns true if the two tensors are of the same dimensionality,
127 /// size and stride.
128 template <typename OtherT, int OtherDim>
129 __host__ __device__ bool
130 isSame(const Tensor<OtherT, OtherDim, InnerContig, IndexT, PtrTraits>& rhs) const;
131
132 /// Returns true if the two tensors are of the same dimensionality and size
133 template <typename OtherT, int OtherDim>
134 __host__ __device__ bool
135 isSameSize(const Tensor<OtherT, OtherDim, InnerContig, IndexT, PtrTraits>& rhs) const;
136
137 /// Cast to a tensor of a different type of the same size and
138 /// stride. U and our type T must be of the same size
139 template <typename U>
140 __host__ __device__ Tensor<U, Dim, InnerContig, IndexT, PtrTraits> cast();
141
142 /// Const version of `cast`
143 template <typename U>
144 __host__ __device__
145 const Tensor<U, Dim, InnerContig, IndexT, PtrTraits> cast() const;
146
147 /// Cast to a tensor of a different type which is potentially a
148 /// different size than our type T. Tensor must be aligned and the
149 /// innermost dimension must be a size that is a multiple of
150 /// sizeof(U) / sizeof(T), and the stride of the innermost dimension
151 /// must be contiguous. The stride of all outer dimensions must be a
152 /// multiple of sizeof(U) / sizeof(T) as well.
153 template <typename U>
154 __host__ __device__ Tensor<U, Dim, InnerContig, IndexT, PtrTraits> castResize();
155
156 /// Const version of `castResize`
157 template <typename U>
158 __host__ __device__ const Tensor<U, Dim, InnerContig, IndexT, PtrTraits>
159 castResize() const;
160
161 /// Returns true if we can castResize() this tensor to the new type
162 template <typename U>
163 __host__ __device__ bool canCastResize() const;
164
165 /// Attempts to cast this tensor to a tensor of a different IndexT.
166 /// Fails if size or stride entries are not representable in the new
167 /// IndexT.
168 template <typename NewIndexT>
169 __host__ Tensor<T, Dim, InnerContig, NewIndexT, PtrTraits>
170 castIndexType() const;
171
172 /// Returns true if we can use this indexing type to access all elements
173 /// index type
174 template <typename NewIndexT>
175 __host__ bool canUseIndexType() const;
176
177 /// Returns a raw pointer to the start of our data.
178 __host__ __device__ inline DataPtrType data() {
179 return data_;
180 }
181
182 /// Returns a raw pointer to the end of our data, assuming
183 /// continuity
184 __host__ __device__ inline DataPtrType end() {
185 return data() + numElements();
186 }
187
188 /// Returns a raw pointer to the start of our data (const).
189 __host__ __device__ inline
190 const DataPtrType data() const {
191 return data_;
192 }
193
194 /// Returns a raw pointer to the end of our data, assuming
195 /// continuity (const)
196 __host__ __device__ inline DataPtrType end() const {
197 return data() + numElements();
198 }
199
200 /// Cast to a different datatype
201 template <typename U>
202 __host__ __device__ inline
203 typename PtrTraits<U>::PtrType dataAs() {
204 return reinterpret_cast<typename PtrTraits<U>::PtrType>(data_);
205 }
206
207 /// Cast to a different datatype
208 template <typename U>
209 __host__ __device__ inline
210 const typename PtrTraits<const U>::PtrType dataAs() const {
211 return reinterpret_cast<typename PtrTraits<const U>::PtrType>(data_);
212 }
213
214 /// Returns a read/write view of a portion of our tensor.
215 __host__ __device__ inline
216 detail::SubTensor<TensorType, Dim - 1, PtrTraits>
217 operator[](IndexT);
218
219 /// Returns a read/write view of a portion of our tensor (const).
220 __host__ __device__ inline
221 const detail::SubTensor<TensorType, Dim - 1, PtrTraits>
222 operator[](IndexT) const;
223
224 /// Returns the size of a given dimension, `[0, Dim - 1]`. No bounds
225 /// checking.
226 __host__ __device__ inline IndexT getSize(int i) const {
227 return size_[i];
228 }
229
230 /// Returns the stride of a given dimension, `[0, Dim - 1]`. No bounds
231 /// checking.
232 __host__ __device__ inline IndexT getStride(int i) const {
233 return stride_[i];
234 }
235
236 /// Returns the total number of elements contained within our data
237 /// (product of `getSize(i)`)
238 __host__ __device__ size_t numElements() const;
239
240 /// If we are contiguous, returns the total size in bytes of our
241 /// data
242 __host__ __device__ size_t getSizeInBytes() const {
243 return numElements() * sizeof(T);
244 }
245
246 /// Returns the size array.
247 __host__ __device__ inline const IndexT* sizes() const {
248 return size_;
249 }
250
251 /// Returns the stride array.
252 __host__ __device__ inline const IndexT* strides() const {
253 return stride_;
254 }
255
256 /// Returns true if there is no padding within the tensor and no
257 /// re-ordering of the dimensions.
258 /// ~~~
259 /// (stride(i) == size(i + 1) * stride(i + 1)) && stride(dim - 1) == 0
260 /// ~~~
261 __host__ __device__ bool isContiguous() const;
262
263 /// Returns whether a given dimension has only increasing stride
264 /// from the previous dimension. A tensor that was permuted by
265 /// exchanging size and stride only will fail this check.
266 /// If `i == 0` just check `size > 0`. Returns `false` if `stride` is `<= 0`.
267 __host__ __device__ bool isConsistentlySized(int i) const;
268
269 // Returns whether at each dimension `stride <= size`.
270 // If this is not the case then iterating once over the size space will
271 // touch the same memory locations multiple times.
272 __host__ __device__ bool isConsistentlySized() const;
273
274 /// Returns true if the given dimension index has no padding
275 __host__ __device__ bool isContiguousDim(int i) const;
276
277 /// Returns a tensor of the same dimension after transposing the two
278 /// dimensions given. Does not actually move elements; transposition
279 /// is made by permuting the size/stride arrays.
280 /// If the dimensions are not valid, asserts.
281 __host__ __device__ Tensor<T, Dim, InnerContig, IndexT, PtrTraits>
282 transpose(int dim1, int dim2) const;
283
284 /// Upcast a tensor of dimension `D` to some tensor of dimension
285 /// D' > D by padding the leading dimensions by 1
286 /// e.g., upcasting a 2-d tensor `[2][3]` to a 4-d tensor `[1][1][2][3]`
287 template <int NewDim>
288 __host__ __device__ Tensor<T, NewDim, InnerContig, IndexT, PtrTraits>
289 upcastOuter();
290
291 /// Upcast a tensor of dimension `D` to some tensor of dimension
292 /// D' > D by padding the lowest/most varying dimensions by 1
293 /// e.g., upcasting a 2-d tensor `[2][3]` to a 4-d tensor `[2][3][1][1]`
294 template <int NewDim>
295 __host__ __device__ Tensor<T, NewDim, InnerContig, IndexT, PtrTraits>
296 upcastInner();
297
298 /// Downcast a tensor of dimension `D` to some tensor of dimension
299 /// D' < D by collapsing the leading dimensions. asserts if there is
300 /// padding on the leading dimensions.
301 template <int NewDim>
302 __host__ __device__
303 Tensor<T, NewDim, InnerContig, IndexT, PtrTraits> downcastOuter();
304
305 /// Downcast a tensor of dimension `D` to some tensor of dimension
306 /// D' < D by collapsing the leading dimensions. asserts if there is
307 /// padding on the leading dimensions.
308 template <int NewDim>
309 __host__ __device__
310 Tensor<T, NewDim, InnerContig, IndexT, PtrTraits> downcastInner();
311
312 /// Returns a tensor that is a view of the `SubDim`-dimensional slice
313 /// of this tensor, starting at `at`.
314 template <int SubDim>
315 __host__ __device__ Tensor<T, SubDim, InnerContig, IndexT, PtrTraits>
316 view(DataPtrType at);
317
318 /// Returns a tensor that is a view of the `SubDim`-dimensional slice
319 /// of this tensor, starting where our data begins
320 template <int SubDim>
321 __host__ __device__ Tensor<T, SubDim, InnerContig, IndexT, PtrTraits>
322 view();
323
324 /// Returns a tensor of the same dimension that is a view of the
325 /// original tensor with the specified dimension restricted to the
326 /// elements in the range [start, start + size)
327 __host__ __device__ Tensor<T, Dim, InnerContig, IndexT, PtrTraits>
328 narrowOutermost(IndexT start, IndexT size);
329
330 /// Returns a tensor of the same dimension that is a view of the
331 /// original tensor with the specified dimension restricted to the
332 /// elements in the range [start, start + size).
333 /// Can occur in an arbitrary dimension
334 __host__ __device__ Tensor<T, Dim, InnerContig, IndexT, PtrTraits>
335 narrow(int dim, IndexT start, IndexT size);
336
337 /// Returns a view of the given tensor expressed as a tensor of a
338 /// different number of dimensions.
339 /// Only works if we are contiguous.
340 template <int NewDim>
341 __host__ __device__ Tensor<T, NewDim, InnerContig, IndexT, PtrTraits>
342 view(std::initializer_list<IndexT> sizes);
343
344 protected:
345 /// Raw pointer to where the tensor data begins
346 DataPtrType data_;
347
348 /// Array of strides (in sizeof(T) terms) per each dimension
349 IndexT stride_[Dim];
350
351 /// Size per each dimension
352 IndexT size_[Dim];
353 };
354
355 // Utilities for checking a collection of tensors
356 namespace detail {
357
358 template <typename IndexType>
359 bool canUseIndexType() {
360 return true;
361 }
362
363 template <typename IndexType, typename T, typename... U>
364 bool canUseIndexType(const T& arg, const U&... args) {
365 return arg.canUseIndexType<IndexType>() &&
366 canUseIndexType(args...);
367 }
368
369 } // namespace detail
370
371 template <typename IndexType, typename... T>
372 bool canUseIndexType(const T&... args) {
373 return detail::canUseIndexType(args...);
374 }
375
376 namespace detail {
377
378 /// Specialization for a view of a single value (0-dimensional)
379 template <typename TensorType, template <typename U> class PtrTraits>
380 class SubTensor<TensorType, 0, PtrTraits> {
381 public:
382 __host__ __device__ SubTensor<TensorType, 0, PtrTraits>
383 operator=(typename TensorType::DataType val) {
384 *data_ = val;
385 return *this;
386 }
387
388 // operator T&
389 __host__ __device__ operator typename TensorType::DataType&() {
390 return *data_;
391 }
392
393 // const operator T& returning const T&
394 __host__ __device__ operator const typename TensorType::DataType&() const {
395 return *data_;
396 }
397
398 // operator& returning T*
399 __host__ __device__ typename TensorType::DataType* operator&() {
400 return data_;
401 }
402
403 // const operator& returning const T*
404 __host__ __device__ const typename TensorType::DataType* operator&() const {
405 return data_;
406 }
407
408 /// Returns a raw accessor to our slice.
409 __host__ __device__ inline typename TensorType::DataPtrType data() {
410 return data_;
411 }
412
413 /// Returns a raw accessor to our slice (const).
414 __host__ __device__ inline
415 const typename TensorType::DataPtrType data() const {
416 return data_;
417 }
418
419 /// Cast to a different datatype.
420 template <typename T>
421 __host__ __device__ T& as() {
422 return *dataAs<T>();
423 }
424
425 /// Cast to a different datatype (const).
426 template <typename T>
427 __host__ __device__ const T& as() const {
428 return *dataAs<T>();
429 }
430
431 /// Cast to a different datatype
432 template <typename T>
433 __host__ __device__ inline
434 typename PtrTraits<T>::PtrType dataAs() {
435 return reinterpret_cast<typename PtrTraits<T>::PtrType>(data_);
436 }
437
438 /// Cast to a different datatype (const)
439 template <typename T>
440 __host__ __device__ inline
441 typename PtrTraits<const T>::PtrType dataAs() const {
442 return reinterpret_cast<typename PtrTraits<const T>::PtrType>(data_);
443 }
444
445 /// Use the texture cache for reads
446 __device__ inline typename TensorType::DataType ldg() const {
447 #if __CUDA_ARCH__ >= 350
448 return __ldg(data_);
449 #else
450 return *data_;
451 #endif
452 }
453
454 /// Use the texture cache for reads; cast as a particular type
455 template <typename T>
456 __device__ inline T ldgAs() const {
457 #if __CUDA_ARCH__ >= 350
458 return __ldg(dataAs<T>());
459 #else
460 return as<T>();
461 #endif
462 }
463
464 protected:
465 /// One dimension greater can create us
466 friend class SubTensor<TensorType, 1, PtrTraits>;
467
468 /// Our parent tensor can create us
469 friend class Tensor<typename TensorType::DataType,
470 1,
471 TensorType::IsInnerContig,
472 typename TensorType::IndexType,
473 PtrTraits>;
474
475 __host__ __device__ inline SubTensor(
476 TensorType& t,
477 typename TensorType::DataPtrType data)
478 : tensor_(t),
479 data_(data) {
480 }
481
482 /// The tensor we're referencing
483 TensorType& tensor_;
484
485 /// Where our value is located
486 typename TensorType::DataPtrType const data_;
487 };
488
489 /// A `SubDim`-rank slice of a parent Tensor
490 template <typename TensorType,
491 int SubDim,
492 template <typename U> class PtrTraits>
493 class SubTensor {
494 public:
495 /// Returns a view of the data located at our offset (the dimension
496 /// `SubDim` - 1 tensor).
497 __host__ __device__ inline
498 SubTensor<TensorType, SubDim - 1, PtrTraits>
499 operator[](typename TensorType::IndexType index) {
500 if (TensorType::IsInnerContig && SubDim == 1) {
501 // Innermost dimension is stride 1 for contiguous arrays
502 return SubTensor<TensorType, SubDim - 1, PtrTraits>(
503 tensor_, data_ + index);
504 } else {
505 return SubTensor<TensorType, SubDim - 1, PtrTraits>(
506 tensor_,
507 data_ + index * tensor_.getStride(TensorType::NumDim - SubDim));
508 }
509 }
510
511 /// Returns a view of the data located at our offset (the dimension
512 /// `SubDim` - 1 tensor) (const).
513 __host__ __device__ inline
514 const SubTensor<TensorType, SubDim - 1, PtrTraits>
515 operator[](typename TensorType::IndexType index) const {
516 if (TensorType::IsInnerContig && SubDim == 1) {
517 // Innermost dimension is stride 1 for contiguous arrays
518 return SubTensor<TensorType, SubDim - 1, PtrTraits>(
519 tensor_, data_ + index);
520 } else {
521 return SubTensor<TensorType, SubDim - 1, PtrTraits>(
522 tensor_,
523 data_ + index * tensor_.getStride(TensorType::NumDim - SubDim));
524 }
525 }
526
527 // operator& returning T*
528 __host__ __device__ typename TensorType::DataType* operator&() {
529 return data_;
530 }
531
532 // const operator& returning const T*
533 __host__ __device__ const typename TensorType::DataType* operator&() const {
534 return data_;
535 }
536
537 /// Returns a raw accessor to our slice.
538 __host__ __device__ inline typename TensorType::DataPtrType data() {
539 return data_;
540 }
541
542 /// Returns a raw accessor to our slice (const).
543 __host__ __device__ inline
544 const typename TensorType::DataPtrType data() const {
545 return data_;
546 }
547
548 /// Cast to a different datatype.
549 template <typename T>
550 __host__ __device__ T& as() {
551 return *dataAs<T>();
552 }
553
554 /// Cast to a different datatype (const).
555 template <typename T>
556 __host__ __device__ const T& as() const {
557 return *dataAs<T>();
558 }
559
560 /// Cast to a different datatype
561 template <typename T>
562 __host__ __device__ inline
563 typename PtrTraits<T>::PtrType dataAs() {
564 return reinterpret_cast<typename PtrTraits<T>::PtrType>(data_);
565 }
566
567 /// Cast to a different datatype (const)
568 template <typename T>
569 __host__ __device__ inline
570 typename PtrTraits<const T>::PtrType dataAs() const {
571 return reinterpret_cast<typename PtrTraits<const T>::PtrType>(data_);
572 }
573
574 /// Use the texture cache for reads
575 __device__ inline typename TensorType::DataType ldg() const {
576 #if __CUDA_ARCH__ >= 350
577 return __ldg(data_);
578 #else
579 return *data_;
580 #endif
581 }
582
583 /// Use the texture cache for reads; cast as a particular type
584 template <typename T>
585 __device__ inline T ldgAs() const {
586 #if __CUDA_ARCH__ >= 350
587 return __ldg(dataAs<T>());
588 #else
589 return as<T>();
590 #endif
591 }
592
593 /// Returns a tensor that is a view of the SubDim-dimensional slice
594 /// of this tensor, starting where our data begins
595 Tensor<typename TensorType::DataType,
596 SubDim,
597 TensorType::IsInnerContig,
598 typename TensorType::IndexType,
599 PtrTraits> view() {
600 return tensor_.template view<SubDim>(data_);
601 }
602
603 protected:
604 /// One dimension greater can create us
605 friend class SubTensor<TensorType, SubDim + 1, PtrTraits>;
606
607 /// Our parent tensor can create us
608 friend class
609 Tensor<typename TensorType::DataType,
610 TensorType::NumDim,
611 TensorType::IsInnerContig,
612 typename TensorType::IndexType,
613 PtrTraits>;
614
615 __host__ __device__ inline SubTensor(
616 TensorType& t,
617 typename TensorType::DataPtrType data)
618 : tensor_(t),
619 data_(data) {
620 }
621
622 /// The tensor we're referencing
623 TensorType& tensor_;
624
625 /// The start of our sub-region
626 typename TensorType::DataPtrType const data_;
627 };
628
629 } // namespace detail
630
631 template <typename T, int Dim, bool InnerContig,
632 typename IndexT, template <typename U> class PtrTraits>
633 __host__ __device__ inline
634 detail::SubTensor<Tensor<T, Dim, InnerContig, IndexT, PtrTraits>,
635 Dim - 1, PtrTraits>
636 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::operator[](IndexT index) {
637 return detail::SubTensor<TensorType, Dim - 1, PtrTraits>(
638 detail::SubTensor<TensorType, Dim, PtrTraits>(
639 *this, data_)[index]);
640 }
641
642 template <typename T, int Dim, bool InnerContig,
643 typename IndexT, template <typename U> class PtrTraits>
644 __host__ __device__ inline
645 const detail::SubTensor<Tensor<T, Dim, InnerContig, IndexT, PtrTraits>,
646 Dim - 1, PtrTraits>
647 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>::operator[](IndexT index) const {
648 return detail::SubTensor<TensorType, Dim - 1, PtrTraits>(
649 detail::SubTensor<TensorType, Dim, PtrTraits>(
650 const_cast<TensorType&>(*this), data_)[index]);
651 }
652
653 } } // namespace
654
655 #include "Tensor-inl.cuh"
faiss::gpu::Tensor::upcastOuter
__host__ __device__ Tensor< T, NewDim, InnerContig, IndexT, PtrTraits > upcastOuter()
Definition: Tensor-inl.cuh:454
faiss::gpu::Tensor::operator[]
__host__ __device__ detail::SubTensor< TensorType, Dim-1, PtrTraits > operator[](IndexT)
Returns a read/write view of a portion of our tensor.
Definition: Tensor.cuh:636
faiss::gpu::Tensor::castIndexType
__host__ Tensor< T, Dim, InnerContig, NewIndexT, PtrTraits > castIndexType() const
Definition: Tensor-inl.cuh:305
faiss::gpu::detail::SubTensor::view
Tensor< typename TensorType::DataType, SubDim, TensorType::IsInnerContig, typename TensorType::IndexType, PtrTraits > view()
Definition: Tensor.cuh:599
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::detail::SubTensor::dataAs
__host__ __device__ PtrTraits< const T >::PtrType dataAs() const
Cast to a different datatype (const)
Definition: Tensor.cuh:570
faiss::gpu::Tensor::dataAs
__host__ __device__ const PtrTraits< const U >::PtrType dataAs() const
Cast to a different datatype.
Definition: Tensor.cuh:210
faiss::gpu::Tensor::dataAs
__host__ __device__ PtrTraits< U >::PtrType dataAs()
Cast to a different datatype.
Definition: Tensor.cuh:203
faiss::gpu::detail::SubTensor::ldgAs
__device__ T ldgAs() const
Use the texture cache for reads; cast as a particular type.
Definition: Tensor.cuh:585
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::data_
DataPtrType data_
Raw pointer to where the tensor data begins.
Definition: Tensor.cuh:346
faiss::gpu::detail::SubTensor::dataAs
__host__ __device__ PtrTraits< T >::PtrType dataAs()
Cast to a different datatype.
Definition: Tensor.cuh:563
faiss::gpu::Tensor::Tensor
__host__ __device__ Tensor()
Default constructor.
Definition: Tensor-inl.cuh:20
faiss::gpu::Tensor::end
__host__ __device__ DataPtrType end() const
Definition: Tensor.cuh:196
faiss::gpu::detail::SubTensor< TensorType, 0, PtrTraits >::dataAs
__host__ __device__ PtrTraits< const T >::PtrType dataAs() const
Cast to a different datatype (const)
Definition: Tensor.cuh:441
faiss::gpu::Tensor::upcastInner
__host__ __device__ Tensor< T, NewDim, InnerContig, IndexT, PtrTraits > upcastInner()
Definition: Tensor-inl.cuh:483
faiss::gpu::detail::SubTensor::data
__host__ __device__ const TensorType::DataPtrType data() const
Returns a raw accessor to our slice (const).
Definition: Tensor.cuh:544
faiss::gpu::detail::SubTensor< TensorType, 0, PtrTraits >::ldg
__device__ TensorType::DataType ldg() const
Use the texture cache for reads.
Definition: Tensor.cuh:446
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::detail::SubTensor< TensorType, 0, PtrTraits >::as
__host__ __device__ T & as()
Cast to a different datatype.
Definition: Tensor.cuh:421
faiss::gpu::detail::SubTensor::as
__host__ __device__ T & as()
Cast to a different datatype.
Definition: Tensor.cuh:550
faiss::gpu::detail::SubTensor< TensorType, 0, PtrTraits >::tensor_
TensorType & tensor_
The tensor we&#39;re referencing.
Definition: Tensor.cuh:483
faiss::gpu::Tensor::isContiguous
__host__ __device__ bool isContiguous() const
Definition: Tensor-inl.cuh:365
faiss::gpu::Tensor::data
__host__ __device__ const DataPtrType data() const
Returns a raw pointer to the start of our data (const).
Definition: Tensor.cuh:190
faiss::gpu::detail::SubTensor::ldg
__device__ TensorType::DataType ldg() const
Use the texture cache for reads.
Definition: Tensor.cuh:575
faiss::gpu::Tensor::sizes
__host__ __device__ const IndexT * sizes() const
Returns the size array.
Definition: Tensor.cuh:247
faiss::gpu::detail::SubTensor::data_
TensorType::DataPtrType const data_
The start of our sub-region.
Definition: Tensor.cuh:625
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::detail::SubTensor< TensorType, 0, PtrTraits >::ldgAs
__device__ T ldgAs() const
Use the texture cache for reads; cast as a particular type.
Definition: Tensor.cuh:456
faiss::gpu::detail::SubTensor::operator[]
__host__ __device__ const SubTensor< TensorType, SubDim-1, PtrTraits > operator[](typename TensorType::IndexType index) const
Definition: Tensor.cuh:515
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::traits::RestrictPtrTraits
Definition: Tensor.cuh:43
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::detail::SubTensor< TensorType, 0, PtrTraits >::data_
TensorType::DataPtrType const data_
Where our value is located.
Definition: Tensor.cuh:486
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::detail::SubTensor< TensorType, 0, PtrTraits >
Specialization for a view of a single value (0-dimensional)
Definition: Tensor.cuh:380
faiss::gpu::Tensor::end
__host__ __device__ DataPtrType end()
Definition: Tensor.cuh:184
faiss::gpu::traits::DefaultPtrTraits
Definition: Tensor.cuh:48
faiss::gpu::detail::SubTensor::tensor_
TensorType & tensor_
The tensor we&#39;re referencing.
Definition: Tensor.cuh:622
faiss::gpu::detail::SubTensor< TensorType, 0, PtrTraits >::data
__host__ __device__ const TensorType::DataPtrType data() const
Returns a raw accessor to our slice (const).
Definition: Tensor.cuh:415
faiss::gpu::Tensor::operator=
__host__ __device__ Tensor< T, Dim, InnerContig, IndexT, PtrTraits > & operator=(Tensor< T, Dim, InnerContig, IndexT, PtrTraits > &t)=default
Assignment.
faiss::gpu::detail::SubTensor::operator[]
__host__ __device__ SubTensor< TensorType, SubDim-1, PtrTraits > operator[](typename TensorType::IndexType index)
Definition: Tensor.cuh:499
faiss::gpu::detail::SubTensor::as
__host__ __device__ const T & as() const
Cast to a different datatype (const).
Definition: Tensor.cuh:556
faiss::gpu::detail::SubTensor
A SubDim-rank slice of a parent Tensor.
Definition: Tensor.cuh:37
faiss::gpu::detail::SubTensor< TensorType, 0, PtrTraits >::dataAs
__host__ __device__ PtrTraits< T >::PtrType dataAs()
Cast to a different datatype.
Definition: Tensor.cuh:434
faiss::gpu::detail::SubTensor::data
__host__ __device__ TensorType::DataPtrType data()
Returns a raw accessor to our slice.
Definition: Tensor.cuh:538
faiss::gpu::Tensor::castResize
__host__ __device__ Tensor< U, Dim, InnerContig, IndexT, PtrTraits > castResize()
Definition: Tensor-inl.cuh:244
faiss::gpu::detail::SubTensor< TensorType, 0, PtrTraits >::data
__host__ __device__ TensorType::DataPtrType data()
Returns a raw accessor to our slice.
Definition: Tensor.cuh:409
faiss::gpu::detail::SubTensor< TensorType, 0, PtrTraits >::as
__host__ __device__ const T & as() const
Cast to a different datatype (const).
Definition: Tensor.cuh:427
faiss::gpu::Tensor::getSizeInBytes
__host__ __device__ size_t getSizeInBytes() const
Definition: Tensor.cuh:242
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