Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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