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 CC-by-NC 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 "Tensor.cuh"
14 #include "DeviceMemory.h"
15 #include "MemorySpace.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 MemorySpace space = MemorySpace::Device);
46 __host__ DeviceTensor(std::initializer_list<IndexT> sizes,
47 MemorySpace space = MemorySpace::Device);
48
49 /// Constructs a tensor of the given size, reserving a temporary
50 /// memory reservation via a memory manager.
51 /// The memory reservation should be ordered with respect to the
52 /// given stream.
53 __host__ DeviceTensor(DeviceMemory& m,
54 const IndexT sizes[Dim],
55 cudaStream_t stream,
56 MemorySpace space = MemorySpace::Device);
57 __host__ DeviceTensor(DeviceMemory& m,
58 std::initializer_list<IndexT> sizes,
59 cudaStream_t stream,
60 MemorySpace space = MemorySpace::Device);
61
62 /// Constructs a tensor of the given size and stride, referencing a
63 /// memory region we do not own
64 __host__ DeviceTensor(DataPtrType data,
65 const IndexT sizes[Dim],
66 MemorySpace space = MemorySpace::Device);
67 __host__ DeviceTensor(DataPtrType data,
68 std::initializer_list<IndexT> sizes,
69 MemorySpace space = MemorySpace::Device);
70
71 /// Constructs a tensor of the given size and stride, referencing a
72 /// memory region we do not own
73 __host__ DeviceTensor(DataPtrType data,
74 const IndexT sizes[Dim],
75 const IndexT strides[Dim],
76 MemorySpace space = MemorySpace::Device);
77
78 /// Copies a tensor into ourselves, allocating memory for it locally
79 __host__ DeviceTensor(Tensor<T, Dim, Contig, IndexT, PtrTraits>& t,
80 cudaStream_t stream,
81 MemorySpace space = MemorySpace::Device);
82
83 /// Copies a tensor into ourselves, reserving a temporary
84 /// memory reservation via a memory manager.
85 __host__ DeviceTensor(DeviceMemory& m,
86 Tensor<T, Dim, Contig, IndexT, PtrTraits>& t,
87 cudaStream_t stream,
88 MemorySpace space = MemorySpace::Device);
89
90 /// Call to zero out memory
91 __host__ DeviceTensor<T, Dim, Contig, IndexT, PtrTraits>&
92 zero(cudaStream_t stream);
93
94 private:
95 enum AllocState {
96 /// This tensor itself owns the memory, which must be freed via
97 /// cudaFree
98 Owner,
99
100 /// This tensor itself is not an owner of the memory; there is
101 /// nothing to free
102 NotOwner,
103
104 /// This tensor has the memory via a temporary memory reservation
105 Reservation
106 };
107
108 AllocState state_;
109 MemorySpace space_;
110 DeviceMemoryReservation reservation_;
111 };
112
113 } } // namespace
114
115 #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:39
faiss::gpu::Tensor::data
__host__ __device__ DataPtrType data()
Returns a raw pointer to the start of our data.
Definition: Tensor.cuh:173
faiss::gpu::DeviceTensor::DeviceTensor
__host__ DeviceTensor()
Default constructor.
Definition: DeviceTensor-inl.cuh:18
faiss::gpu::DeviceTensor::zero
__host__ DeviceTensor< T, Dim, Contig, IndexT, PtrTraits > & zero(cudaStream_t stream)
Call to zero out memory.
Definition: DeviceTensor-inl.cuh:217
faiss::gpu::DeviceTensor::~DeviceTensor
__host__ ~DeviceTensor()
Destructor.
Definition: DeviceTensor-inl.cuh:58
faiss::gpu::Tensor::strides
__host__ __device__ const IndexT * strides() const
Returns the stride array.
Definition: Tensor.cuh:247
faiss::gpu::Tensor
Our tensor type.
Definition: Tensor.cuh:30
faiss::gpu::Tensor::sizes
__host__ __device__ const IndexT * sizes() const
Returns the size array.
Definition: Tensor.cuh:242
faiss::gpu::DeviceMemory
Manages temporary memory allocations on a GPU device.
Definition: DeviceMemory.h:46
faiss::gpu::DeviceTensor
Definition: DeviceTensor.cuh:24
faiss::gpu::DeviceMemoryReservation
Definition: DeviceMemory.h:20