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