All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
NoTypeTensor.cuh
/**
* Copyright (c) 2015-present, Facebook, Inc.
* All rights reserved.
*
* This source code is licensed under the BSD+Patents license found in the
* LICENSE file in the root directory of this source tree.
*/
// Copyright 2004-present Facebook. All Rights Reserved.
#pragma once
#include "../../FaissAssert.h"
#include "Tensor.cuh"
#include <initializer_list>
namespace faiss { namespace gpu {
template <int Dim, bool InnerContig = false, typename IndexT = int>
class NoTypeTensor {
public:
NoTypeTensor()
: mem_(nullptr),
typeSize_(0) {
}
template <typename T>
NoTypeTensor(Tensor<T, Dim, InnerContig, IndexT>& t)
: mem_(t.data()),
typeSize_(sizeof(T)) {
for (int i = 0; i < Dim; ++i) {
size_[i] = t.getSize(i);
stride_[i] = t.getStride(i);
}
}
NoTypeTensor(void* mem, int typeSize, std::initializer_list<IndexT> sizes)
: mem_(mem),
typeSize_(typeSize) {
int i = 0;
for (auto s : sizes) {
size_[i++] = s;
}
stride_[Dim - 1] = (IndexT) 1;
for (int j = Dim - 2; j >= 0; --j) {
stride_[j] = stride_[j + 1] * size_[j + 1];
}
}
NoTypeTensor(void* mem, int typeSize, int sizes[Dim])
: mem_(mem),
typeSize_(typeSize) {
for (int i = 0; i < Dim; ++i) {
size_[i] = sizes[i];
}
stride_[Dim - 1] = (IndexT) 1;
for (int i = Dim - 2; i >= 0; --i) {
stride_[i] = stride_[i + 1] * sizes[i + 1];
}
}
NoTypeTensor(void* mem, int typeSize,
IndexT sizes[Dim], IndexT strides[Dim])
: mem_(mem),
typeSize_(typeSize) {
for (int i = 0; i < Dim; ++i) {
size_[i] = sizes[i];
stride_[i] = strides[i];
}
}
int getTypeSize() const {
return typeSize_;
}
IndexT getSize(int dim) const {
FAISS_ASSERT(dim < Dim);
return size_[dim];
}
IndexT getStride(int dim) const {
FAISS_ASSERT(dim < Dim);
return stride_[dim];
}
template <typename T>
Tensor<T, Dim, InnerContig, IndexT> toTensor() {
FAISS_ASSERT(sizeof(T) == typeSize_);
return Tensor<T, Dim, InnerContig, IndexT>((T*) mem_, size_, stride_);
}
NoTypeTensor<Dim, InnerContig, IndexT> narrowOutermost(IndexT start,
IndexT size) {
char* newPtr = (char*) mem_;
if (start > 0) {
newPtr += typeSize_ * start * stride_[0];
}
IndexT newSize[Dim];
for (int i = 0; i < Dim; ++i) {
if (i == 0) {
assert(start + size <= size_[0]);
newSize[i] = size;
} else {
newSize[i] = size_[i];
}
}
return NoTypeTensor<Dim, InnerContig, IndexT>(
newPtr, typeSize_, newSize, stride_);
}
private:
void* mem_;
int typeSize_;
IndexT size_[Dim];
IndexT stride_[Dim];
};
} } // namespace
faiss::gpu::NoTypeTensor
Definition: NoTypeTensor.cuh:20
faiss::gpu::Tensor::getSize
__host__ __device__ IndexT getSize(int i) const
Definition: Tensor.cuh:226
faiss::gpu::Tensor::data
__host__ __device__ DataPtrType data()
Returns a raw pointer to the start of our data.
Definition: Tensor.cuh:178
faiss::gpu::Tensor
Our tensor type.
Definition: Tensor.cuh:30
faiss::gpu::Tensor::getStride
__host__ __device__ IndexT getStride(int i) const
Definition: Tensor.cuh:232