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