Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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