Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
NoTypeTensor.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 "../../FaissAssert.h"
14 #include "Tensor.cuh"
15 #include <initializer_list>
16
17 namespace faiss { namespace gpu {
18
19 template <int Dim, bool Contig = false, typename IndexT = int>
20 class NoTypeTensor {
21 public:
22 NoTypeTensor()
23 : mem_(nullptr),
24 typeSize_(0) {
25 }
26
27 template <typename T>
28 NoTypeTensor(Tensor<T, Dim, Contig, IndexT>& t)
29 : mem_(t.data()),
30 typeSize_(sizeof(T)) {
31 for (int i = 0; i < Dim; ++i) {
32 size_[i] = t.getSize(i);
33 stride_[i] = t.getStride(i);
34 }
35 }
36
37 NoTypeTensor(void* mem, int typeSize, std::initializer_list<IndexT> sizes)
38 : mem_(mem),
39 typeSize_(typeSize) {
40
41 int i = 0;
42 for (auto s : sizes) {
43 size_[i++] = s;
44 }
45
46 stride_[Dim - 1] = (IndexT) 1;
47 for (int j = Dim - 2; j >= 0; --j) {
48 stride_[j] = stride_[j + 1] * size_[j + 1];
49 }
50 }
51
52 NoTypeTensor(void* mem, int typeSize, int sizes[Dim])
53 : mem_(mem),
54 typeSize_(typeSize) {
55 for (int i = 0; i < Dim; ++i) {
56 size_[i] = sizes[i];
57 }
58
59 stride_[Dim - 1] = (IndexT) 1;
60 for (int i = Dim - 2; i >= 0; --i) {
61 stride_[i] = stride_[i + 1] * sizes[i + 1];
62 }
63 }
64
65 NoTypeTensor(void* mem, int typeSize,
66 IndexT sizes[Dim], IndexT strides[Dim])
67 : mem_(mem),
68 typeSize_(typeSize) {
69 for (int i = 0; i < Dim; ++i) {
70 size_[i] = sizes[i];
71 stride_[i] = strides[i];
72 }
73 }
74
75 int getTypeSize() const {
76 return typeSize_;
77 }
78
79 IndexT getSize(int dim) const {
80 FAISS_ASSERT(dim < Dim);
81 return size_[dim];
82 }
83
84 IndexT getStride(int dim) const {
85 FAISS_ASSERT(dim < Dim);
86 return stride_[dim];
87 }
88
89 template <typename T>
90 Tensor<T, Dim, Contig, IndexT> toTensor() {
91 FAISS_ASSERT(sizeof(T) == typeSize_);
92
93 return Tensor<T, Dim, Contig, IndexT>((T*) mem_, size_, stride_);
94 }
95
96 NoTypeTensor<Dim, Contig, IndexT> narrowOutermost(IndexT start, IndexT size) {
97 char* newPtr = (char*) mem_;
98
99 if (start > 0) {
100 newPtr += typeSize_ * start * stride_[0];
101 }
102
103 IndexT newSize[Dim];
104 for (int i = 0; i < Dim; ++i) {
105 if (i == 0) {
106 assert(start + size <= size_[0]);
107 newSize[i] = size;
108 } else {
109 newSize[i] = size_[i];
110 }
111 }
112
113 return NoTypeTensor<Dim, Contig, IndexT>(
114 newPtr, typeSize_, newSize, stride_);
115 }
116
117 private:
118 void* mem_;
119 int typeSize_;
120 IndexT size_[Dim];
121 IndexT stride_[Dim];
122 };
123
124 } } // namespace
faiss::gpu::Tensor::data
__host__ __device__ DataPtrType data()
Returns a raw pointer to the start of our data.
Definition: Tensor.cuh:173
faiss::gpu::NoTypeTensor
Definition: NoTypeTensor.cuh:20
faiss::gpu::Tensor::getStride
__host__ __device__ IndexT getStride(int i) const
Definition: Tensor.cuh:227
faiss::gpu::Tensor
Our tensor type.
Definition: Tensor.cuh:30
faiss::gpu::Tensor::getSize
__host__ __device__ IndexT getSize(int i) const
Definition: Tensor.cuh:221