Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
DeviceTensor.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 "Tensor.cuh"
12 #include "DeviceMemory.h"
13 #include "MemorySpace.h"
14
15 namespace faiss { namespace gpu {
16
17 template <typename T,
18 int Dim,
19 bool InnerContig = false,
20 typename IndexT = int,
21 template <typename U> class PtrTraits = traits::DefaultPtrTraits>
22 class DeviceTensor : public Tensor<T, Dim, InnerContig, IndexT, PtrTraits> {
23 public:
24 typedef IndexT IndexType;
25 typedef typename PtrTraits<T>::PtrType DataPtrType;
26
27 /// Default constructor
28 __host__ DeviceTensor();
29
30 /// Destructor
31 __host__ ~DeviceTensor();
32
33 /// Move constructor
34 __host__ DeviceTensor(DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t);
35
36 /// Move assignment
37 __host__ DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&
38 operator=(DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&& t);
39
40 /// Constructs a tensor of the given size, allocating memory for it
41 /// locally
42 __host__ DeviceTensor(const IndexT sizes[Dim],
43 MemorySpace space = MemorySpace::Device);
44 __host__ DeviceTensor(std::initializer_list<IndexT> sizes,
45 MemorySpace space = MemorySpace::Device);
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 MemorySpace space = MemorySpace::Device);
55 __host__ DeviceTensor(DeviceMemory& m,
56 std::initializer_list<IndexT> sizes,
57 cudaStream_t stream,
58 MemorySpace space = MemorySpace::Device);
59
60 /// Constructs a tensor of the given size and stride, referencing a
61 /// memory region we do not own
62 __host__ DeviceTensor(DataPtrType data,
63 const IndexT sizes[Dim],
64 MemorySpace space = MemorySpace::Device);
65 __host__ DeviceTensor(DataPtrType data,
66 std::initializer_list<IndexT> sizes,
67 MemorySpace space = MemorySpace::Device);
68
69 /// Constructs a tensor of the given size and stride, referencing a
70 /// memory region we do not own
71 __host__ DeviceTensor(DataPtrType data,
72 const IndexT sizes[Dim],
73 const IndexT strides[Dim],
74 MemorySpace space = MemorySpace::Device);
75
76 /// Copies a tensor into ourselves, allocating memory for it locally
77 __host__ DeviceTensor(Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t,
78 cudaStream_t stream,
79 MemorySpace space = MemorySpace::Device);
80
81 /// Copies a tensor into ourselves, reserving a temporary
82 /// memory reservation via a memory manager.
83 __host__ DeviceTensor(DeviceMemory& m,
84 Tensor<T, Dim, InnerContig, IndexT, PtrTraits>& t,
85 cudaStream_t stream,
86 MemorySpace space = MemorySpace::Device);
87
88 /// Call to zero out memory
89 __host__ DeviceTensor<T, Dim, InnerContig, IndexT, PtrTraits>&
90 zero(cudaStream_t stream);
91
92 private:
93 enum AllocState {
94 /// This tensor itself owns the memory, which must be freed via
95 /// cudaFree
96 Owner,
97
98 /// This tensor itself is not an owner of the memory; there is
99 /// nothing to free
100 NotOwner,
101
102 /// This tensor has the memory via a temporary memory reservation
103 Reservation
104 };
105
106 AllocState state_;
107 MemorySpace space_;
108 DeviceMemoryReservation reservation_;
109 };
110
111 } } // namespace
112
113 #include "DeviceTensor-inl.cuh"
faiss::gpu::DeviceTensor::~DeviceTensor
__host__ ~DeviceTensor()
Destructor.
Definition: DeviceTensor-inl.cuh:56
faiss::gpu::DeviceTensor::zero
__host__ DeviceTensor< T, Dim, InnerContig, IndexT, PtrTraits > & zero(cudaStream_t stream)
Call to zero out memory.
Definition: DeviceTensor-inl.cuh:215
faiss::gpu::Tensor::sizes
__host__ __device__ const IndexT * sizes() const
Returns the size array.
Definition: Tensor.cuh:243
faiss::gpu::Tensor::strides
__host__ __device__ const IndexT * strides() const
Returns the stride array.
Definition: Tensor.cuh:248
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
Our tensor type.
Definition: Tensor.cuh:28
faiss::gpu::DeviceTensor::DeviceTensor
__host__ DeviceTensor()
Default constructor.
Definition: DeviceTensor-inl.cuh:16
faiss::gpu::DeviceTensor::operator=
__host__ DeviceTensor< T, Dim, InnerContig, IndexT, PtrTraits > & operator=(DeviceTensor< T, Dim, InnerContig, IndexT, PtrTraits > &&t)
Move assignment.
Definition: DeviceTensor-inl.cuh:37
faiss::gpu::DeviceMemory
Manages temporary memory allocations on a GPU device.
Definition: DeviceMemory.h:44
faiss::gpu::DeviceTensor
Definition: DeviceTensor.cuh:22
faiss::gpu::DeviceMemoryReservation
Definition: DeviceMemory.h:18