Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
DeviceTensor.cuh
1
2 /**
3 * Copyright (c) 2015-present, Facebook, Inc.
4 * All rights reserved.
5 *
6 * This source code is licensed under the CC-by-NC license found in the
7 * LICENSE file in the root directory of this source tree.
8 */
9
10 // Copyright 2004-present Facebook. All Rights Reserved.
11
12 #pragma once
13
14 #include "Tensor.cuh"
15 #include "DeviceMemory.h"
16
17 namespace faiss { namespace gpu {
18
19 template <typename T,
20 int Dim,
21 bool Contig = false,
22 typename IndexT = int,
23 template <typename U> class PtrTraits = traits::DefaultPtrTraits>
24 class DeviceTensor : public Tensor<T, Dim, Contig, IndexT, PtrTraits> {
25 public:
26 typedef IndexT IndexType;
27 typedef typename PtrTraits<T>::PtrType DataPtrType;
28
29 /// Default constructor
30 __host__ DeviceTensor();
31
32 /// Destructor
33 __host__ ~DeviceTensor();
34
35 /// Move constructor
36 __host__ DeviceTensor(DeviceTensor<T, Dim, Contig, IndexT, PtrTraits>&& t);
37
38 /// Move assignment
39 __host__ DeviceTensor<T, Dim, Contig, IndexT, PtrTraits>&
40 operator=(DeviceTensor<T, Dim, Contig, IndexT, PtrTraits>&& t);
41
42 /// Constructs a tensor of the given size, allocating memory for it
43 /// locally
44 __host__ DeviceTensor(const IndexT sizes[Dim]);
45 __host__ DeviceTensor(std::initializer_list<IndexT> sizes);
46
47 /// Constructs a tensor of the given size, reserving a temporary
48 /// memory reservation via a memory manager.
49 /// The memory reservation should be ordered with respect to the
50 /// given stream.
51 __host__ DeviceTensor(DeviceMemory& m,
52 const IndexT sizes[Dim],
53 cudaStream_t stream);
54 __host__ DeviceTensor(DeviceMemory& m,
55 std::initializer_list<IndexT> sizes,
56 cudaStream_t stream);
57
58 /// Constructs a tensor of the given size and stride, referencing a
59 /// memory region we do not own
60 __host__ DeviceTensor(DataPtrType data,
61 const IndexT sizes[Dim]);
62 __host__ DeviceTensor(DataPtrType data,
63 std::initializer_list<IndexT> sizes);
64
65 /// Constructs a tensor of the given size and stride, referencing a
66 /// memory region we do not own
67 __host__ DeviceTensor(DataPtrType data,
68 const IndexT sizes[Dim],
69 const IndexT strides[Dim]);
70
71 /// Copies a tensor into ourselves, allocating memory for it locally
72 __host__ DeviceTensor(Tensor<T, Dim, Contig, IndexT, PtrTraits>& t,
73 cudaStream_t stream);
74
75 /// Copies a tensor into ourselves, reserving a temporary
76 /// memory reservation via a memory manager.
77 __host__ DeviceTensor(DeviceMemory& m,
78 Tensor<T, Dim, Contig, IndexT, PtrTraits>& t,
79 cudaStream_t stream);
80
81 /// Call to zero out memory
82 __host__ DeviceTensor<T, Dim, Contig, IndexT, PtrTraits>&
83 zero(cudaStream_t stream);
84
85 private:
86 enum AllocState {
87 /// This tensor itself owns the memory, which must be freed via
88 /// cudaFree
89 Owner,
90
91 /// This tensor itself is not an owner of the memory; there is
92 /// nothing to free
93 NotOwner,
94
95 /// This tensor has the memory via a temporary memory reservation
96 Reservation
97 };
98
99 AllocState state_;
100 DeviceMemoryReservation reservation_;
101 };
102
103 } } // namespace
104
105 #include "DeviceTensor-inl.cuh"
faiss::gpu::DeviceTensor::operator=
__host__ DeviceTensor< T, Dim, Contig, IndexT, PtrTraits > & operator=(DeviceTensor< T, Dim, Contig, IndexT, PtrTraits > &&t)
Move assignment.
Definition: DeviceTensor-inl.cuh:38
faiss::gpu::Tensor::data
__host__ __device__ DataPtrType data()
Returns a raw pointer to the start of our data.
Definition: Tensor.cuh:174
faiss::gpu::DeviceTensor::DeviceTensor
__host__ DeviceTensor()
Default constructor.
Definition: DeviceTensor-inl.cuh:19
faiss::gpu::DeviceTensor::zero
__host__ DeviceTensor< T, Dim, Contig, IndexT, PtrTraits > & zero(cudaStream_t stream)
Call to zero out memory.
Definition: DeviceTensor-inl.cuh:199
faiss::gpu::DeviceTensor::~DeviceTensor
__host__ ~DeviceTensor()
Destructor.
Definition: DeviceTensor-inl.cuh:56
faiss::gpu::Tensor::strides
__host__ __device__ const IndexT * strides() const
Returns the stride array.
Definition: Tensor.cuh:248
faiss::gpu::Tensor
Our tensor type.
Definition: Tensor.cuh:31
faiss::gpu::Tensor::sizes
__host__ __device__ const IndexT * sizes() const
Returns the size array.
Definition: Tensor.cuh:243
faiss::gpu::DeviceMemory
Manages temporary memory allocations on a GPU device.
Definition: DeviceMemory.h:47
faiss::gpu::DeviceTensor
Definition: DeviceTensor.cuh:24
faiss::gpu::DeviceMemoryReservation
Definition: DeviceMemory.h:21