Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
121 changes: 70 additions & 51 deletions include/yateto/TensorView.h
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,25 @@
#include <limits>
#include <type_traits>

#if defined(__CUDACC__) || defined(__HIP__)
#define YATETO_HOST __host__
#define YATETO_DEVICE __device__
#else
#define YATETO_HOST
#define YATETO_DEVICE
#endif

#define YATETO_HOSTDEVICE YATETO_HOST YATETO_DEVICE

namespace yateto {
template<typename uint_t=unsigned>
class slice {
public:
explicit slice(uint_t start = 0, uint_t stop = std::numeric_limits<uint_t>::max())
#pragma omp declare target begin
YATETO_HOSTDEVICE explicit slice(uint_t start = 0, uint_t stop = std::numeric_limits<uint_t>::max())
: start(start), stop(stop)
{}
#pragma omp declare target end

uint_t start;
uint_t stop;
Expand All @@ -27,61 +39,65 @@ namespace yateto {
template<unsigned Dim, typename real_t, typename uint_t>
class TensorView {
public:
explicit TensorView(std::initializer_list<uint_t> shape) {
#pragma omp declare target
YATETO_HOSTDEVICE explicit TensorView(std::initializer_list<uint_t> shape) {
std::copy(shape.begin(), shape.end(), m_shape);
}

explicit TensorView(uint_t const shape[]) {
YATETO_HOSTDEVICE explicit TensorView(uint_t const shape[]) {
for (uint_t d = 0; d < Dim; ++d) {
m_shape[d] = shape[d];
}
}

constexpr uint_t dim() const {
YATETO_HOSTDEVICE static constexpr uint_t dim() {
return Dim;
}

uint_t shape(uint_t dim) const {
YATETO_HOSTDEVICE uint_t shape(uint_t dim) const {
return m_shape[dim];
}

#pragma omp end declare target
protected:
uint_t m_shape[Dim];
};

template<typename real_t, typename uint_t>
class TensorView<0, real_t, uint_t> {
public:
explicit TensorView(std::initializer_list<uint_t> shape) {}
#pragma omp declare target
YATETO_HOSTDEVICE explicit TensorView(std::initializer_list<uint_t> shape) {}

explicit TensorView(uint_t const shape[]) {}
YATETO_HOSTDEVICE explicit TensorView(uint_t const shape[]) {}

constexpr uint_t dim() const {
YATETO_HOSTDEVICE static constexpr uint_t dim() {
return 0;
}

uint_t shape(uint_t dim) const {
YATETO_HOSTDEVICE uint_t shape(uint_t dim) const {
return 0;
}
#pragma omp end declare target
};

template<unsigned Dim, typename real_t, typename uint_t=unsigned>
class DenseTensorView : public TensorView<Dim, real_t, uint_t> {
public:
explicit DenseTensorView(real_t* values, std::initializer_list<uint_t> shape, std::initializer_list<uint_t> start, std::initializer_list<uint_t> stop)
#pragma omp declare target
YATETO_HOSTDEVICE explicit DenseTensorView(real_t* values, std::initializer_list<uint_t> shape, std::initializer_list<uint_t> start, std::initializer_list<uint_t> stop)
: TensorView<Dim, real_t, uint_t>(shape), m_values(values) {
std::copy(start.begin(), start.end(), m_start);
std::copy(stop.begin(), stop.end(), m_stop);
computeStride();
}

explicit DenseTensorView(real_t* values, std::initializer_list<uint_t> shape)
YATETO_HOSTDEVICE explicit DenseTensorView(real_t* values, std::initializer_list<uint_t> shape)
: TensorView<Dim, real_t, uint_t>(shape), m_values(values), m_start{} {
std::copy(shape.begin(), shape.end(), m_stop);
computeStride();
}

explicit DenseTensorView(real_t* values, uint_t const shape[], uint_t const start[], uint_t const stop[])
YATETO_HOSTDEVICE explicit DenseTensorView(real_t* values, uint_t const shape[], uint_t const start[], uint_t const stop[])
: TensorView<Dim, real_t, uint_t>(shape), m_values(values) {
for (uint_t d = 0; d < Dim; ++d) {
m_start[d] = start[d];
Expand All @@ -90,27 +106,27 @@ namespace yateto {
computeStride();
}

explicit DenseTensorView(real_t* values, uint_t const shape[])
YATETO_HOSTDEVICE explicit DenseTensorView(real_t* values, uint_t const shape[])
: TensorView<Dim, real_t, uint_t>(shape), m_values(values), m_start{} {
for (uint_t d = 0; d < Dim; ++d) {
m_stop[d] = shape[d];
}
computeStride();
}

explicit DenseTensorView(real_t* values, uint_t const shape[], uint_t const stride[])
YATETO_HOSTDEVICE explicit DenseTensorView(real_t* values, uint_t const shape[], uint_t const stride[])
: TensorView<Dim, real_t, uint_t>(shape), m_values(values), m_start{} {
for (uint_t d = 0; d < Dim; ++d) {
m_stop[d] = shape[d];
m_stride[d] = stride[d];
}
}

uint_t size() const {
YATETO_HOSTDEVICE uint_t size() const {
return (m_stop[Dim-1]-m_start[Dim-1]) * m_stride[Dim-1];
}

void setZero() {
YATETO_HOSTDEVICE void setZero() {
uint_t entry[Dim];
std::copy(m_start, m_start + Dim, entry);
while (entry[Dim-1] != m_stop[Dim-1]) {
Expand All @@ -132,41 +148,41 @@ namespace yateto {
}

template<typename Head>
bool isInRange(const uint_t start[Dim], const uint_t stop[Dim], int dim, Head head) const {
YATETO_HOSTDEVICE bool isInRange(const uint_t start[Dim], const uint_t stop[Dim], int dim, Head head) const {
return static_cast<uint_t>(head) >= start[dim] && static_cast<uint_t>(head) < stop[dim];
}

template<typename Head, typename... Tail>
bool isInRange(const uint_t start[Dim], const uint_t stop[Dim], int dim, Head head, Tail... tail) const {
YATETO_HOSTDEVICE bool isInRange(const uint_t start[Dim], const uint_t stop[Dim], int dim, Head head, Tail... tail) const {
return static_cast<uint_t>(head) >= start[dim]
&& static_cast<uint_t>(head) < stop[dim]
&& isInRange(start, stop, dim+1, tail...);
}

template<typename... Entry>
bool isInRange(Entry... entry) const {
YATETO_HOSTDEVICE bool isInRange(Entry... entry) const {
static_assert(sizeof...(entry) == Dim,
"Number of arguments to isInRange(...) does not match the tensor dimension.");
return isInRange(m_start, m_stop, 0, entry...);
}

template<typename... Entry>
real_t& operator()(Entry... entry) {
YATETO_HOSTDEVICE real_t& operator()(Entry... entry) {
static_assert(sizeof...(entry) == Dim,
"Number of arguments to operator() does not match the tensor dimension.");
assert(isInRange(entry...));
return m_values[address(entry...)];
}

template<typename... Entry>
real_t operator()(Entry... entry) const {
YATETO_HOSTDEVICE real_t operator()(Entry... entry) const {
static_assert(sizeof...(entry) == Dim,
"Number of arguments to operator() const does not match the tensor dimension.");
assert(isInRange(entry...));
return m_values[address(entry...)];
}

real_t operator[](uint_t const entry[Dim]) const {
YATETO_HOSTDEVICE real_t operator[](uint_t const entry[Dim]) const {
uint_t addr = 0;
for (uint_t d = 0; d < Dim; ++d) {
assert(entry[d] >= m_start[d] && entry[d] < m_stop[d]);
Expand All @@ -175,7 +191,7 @@ namespace yateto {
return m_values[addr];
}

real_t& operator[](uint_t const entry[Dim]) {
YATETO_HOSTDEVICE real_t& operator[](uint_t const entry[Dim]) {
uint_t addr = 0;
for (uint_t d = 0; d < Dim; ++d) {
assert(entry[d] >= m_start[d] && entry[d] < m_stop[d]);
Expand All @@ -185,7 +201,7 @@ namespace yateto {
}

template<class view_t>
void copyToView(view_t& other) const {
YATETO_HOSTDEVICE void copyToView(view_t& other) const {
assert(Dim == other.dim());

uint_t entry[Dim];
Expand Down Expand Up @@ -217,7 +233,7 @@ namespace yateto {
}

template<typename... Entry>
auto subtensor(Entry... entry) -> DenseTensorView<count_slices<uint_t, Entry...>::value, real_t, uint_t> const {
YATETO_HOSTDEVICE auto subtensor(Entry... entry) -> DenseTensorView<count_slices<uint_t, Entry...>::value, real_t, uint_t> const {
static_assert(sizeof...(entry) == Dim, "Number of arguments to subtensor() does not match tensor dimension.");
constexpr auto nSlices = count_slices<uint_t, Entry...>::value;
uint_t begin[Dim];
Expand All @@ -228,61 +244,62 @@ namespace yateto {
return subtensor;
}

real_t* data() {
YATETO_HOSTDEVICE real_t* data() {
return m_values;
}

const real_t* data() const {
YATETO_HOSTDEVICE const real_t* data() const {
return m_values;
}

protected:
void computeStride() {
YATETO_HOSTDEVICE void computeStride() {
m_stride[0] = 1;
for (uint_t d = 0; d < Dim-1; ++d) {
m_stride[d+1] = m_stride[d] * (m_stop[d] - m_start[d]);
}
}

template<typename Head>
uint_t address(Head head) const {
YATETO_HOSTDEVICE uint_t address(Head head) const {
assert(static_cast<uint_t>(head) >= m_start[Dim-1] && static_cast<uint_t>(head) < m_stop[Dim-1]);
return (head - m_start[Dim-1]) * m_stride[Dim-1];
}

template<typename Head, typename... Tail>
uint_t address(Head head, Tail... tail) const {
YATETO_HOSTDEVICE uint_t address(Head head, Tail... tail) const {
uint_t const d = (Dim-1) - sizeof...(tail);
assert(static_cast<uint_t>(head) >= m_start[d] && static_cast<uint_t>(head) < m_stop[d]);
return (head - m_start[d]) * m_stride[d] + address(tail...);
}

template<typename T, typename std::enable_if<std::is_integral<T>::value, int>::type = 0>
void extractDim(uint_t*& begin, uint_t*&, uint_t*&, uint_t dimNo, T entry) const {
YATETO_HOSTDEVICE void extractDim(uint_t*& begin, uint_t*&, uint_t*&, uint_t dimNo, T entry) const {
assert(static_cast<uint_t>(entry) >= m_start[dimNo] && static_cast<uint_t>(entry) < m_stop[dimNo]);
*begin++ = entry;
}

template<typename T, typename std::enable_if<std::is_same<T, slice<uint_t>>::value, int>::type = 0>
void extractDim(uint_t*& begin, uint_t*& size, uint_t*& stride, uint_t dimNo, T dim) const {
YATETO_HOSTDEVICE void extractDim(uint_t*& begin, uint_t*& size, uint_t*& stride, uint_t dimNo, T dim) const {
*begin = std::max(m_start[dimNo], dim.start);
*size++ = std::min(m_stop[dimNo], dim.stop) - *begin;
++begin;
*stride++ = m_stride[dimNo];
}

template<typename Head>
void extractSubtensor(uint_t* begin, uint_t* size, uint_t* stride, Head head) const {
YATETO_HOSTDEVICE void extractSubtensor(uint_t* begin, uint_t* size, uint_t* stride, Head head) const {
extractDim<Head>(begin, size, stride, Dim-1, head);
}

template<typename Head, typename... Tail>
void extractSubtensor(uint_t* begin, uint_t* size, uint_t* stride, Head head, Tail... tail) const {
YATETO_HOSTDEVICE void extractSubtensor(uint_t* begin, uint_t* size, uint_t* stride, Head head, Tail... tail) const {
uint_t const d = (Dim-1) - sizeof...(tail);
extractDim<Head>(begin, size, stride, d, head);
extractSubtensor(begin, size, stride, tail...);
}

#pragma omp end declare target
real_t* m_values;
uint_t m_start[Dim];
uint_t m_stop[Dim];
Expand All @@ -292,47 +309,49 @@ namespace yateto {
template<typename real_t, typename uint_t>
class DenseTensorView<0,real_t,uint_t> : public TensorView<0, real_t, uint_t> {
public:
explicit DenseTensorView(real_t* values, std::initializer_list<uint_t> shape, std::initializer_list<uint_t> start, std::initializer_list<uint_t> stop)
#pragma omp declare target
YATETO_HOSTDEVICE explicit DenseTensorView(real_t* values, std::initializer_list<uint_t> shape, std::initializer_list<uint_t> start, std::initializer_list<uint_t> stop)
: TensorView<0, real_t, uint_t>(shape), m_values(values) {
}

uint_t size() const {
YATETO_HOSTDEVICE uint_t size() const {
return 1;
}

void setZero() {
YATETO_HOSTDEVICE void setZero() {
m_values[0] = 0.0;
}

template<class view_t>
void copyToView(view_t& other) {
YATETO_HOSTDEVICE void copyToView(view_t& other) {
other.m_values[0] = m_values[0];
}

#pragma omp end declare target
protected:
real_t* m_values;
};

template<typename real_t, typename uint_t>
class CSCMatrixView : public TensorView<2, real_t, uint_t> {
public:
explicit CSCMatrixView(real_t* values, std::initializer_list<uint_t> shape, uint_t const* rowInd, uint_t const* colPtr)
#pragma omp declare target
YATETO_HOSTDEVICE explicit CSCMatrixView(real_t* values, std::initializer_list<uint_t> shape, uint_t const* rowInd, uint_t const* colPtr)
: TensorView<2, real_t, uint_t>(shape), m_values(values), m_rowInd(rowInd), m_colPtr(colPtr) {
}

explicit CSCMatrixView(real_t* values, uint_t const shape[], uint_t const* rowInd, uint_t const* colPtr)
YATETO_HOSTDEVICE explicit CSCMatrixView(real_t* values, uint_t const shape[], uint_t const* rowInd, uint_t const* colPtr)
: TensorView<2, real_t, uint_t>(shape), m_values(values), m_rowInd(rowInd), m_colPtr(colPtr) {
}

uint_t size() const {
YATETO_HOSTDEVICE uint_t size() const {
return m_colPtr[ this->shape(1) ];
}

void setZero() {
YATETO_HOSTDEVICE void setZero() {
memset(m_values, 0, size() * sizeof(real_t));
}

real_t operator()(uint_t row, uint_t col) const {
YATETO_HOSTDEVICE real_t operator()(uint_t row, uint_t col) const {
assert(col >= 0 && col < this->shape(1));
uint_t addr = m_colPtr[ col ];
uint_t stop = m_colPtr[ col+1 ];
Expand All @@ -347,7 +366,7 @@ namespace yateto {
return m_values[addr];
}

real_t& operator()(uint_t row, uint_t col) {
YATETO_HOSTDEVICE real_t& operator()(uint_t row, uint_t col) {
assert(col >= 0 && col < this->shape(1));
uint_t addr = m_colPtr[ col ];
uint_t stop = m_colPtr[ col+1 ];
Expand All @@ -362,7 +381,7 @@ namespace yateto {
return m_values[addr];
}

bool isInRange(uint_t row, uint_t col) const {
YATETO_HOSTDEVICE bool isInRange(uint_t row, uint_t col) const {
assert(col >= 0 && col < this->shape(1));
uint_t addr = m_colPtr[ col ];
uint_t stop = m_colPtr[ col+1 ];
Expand All @@ -376,16 +395,16 @@ namespace yateto {
return false;
}

real_t& operator[](const uint_t entry[2]) {
YATETO_HOSTDEVICE real_t& operator[](const uint_t entry[2]) {
return operator()(entry[0], entry[1]);
}

real_t operator[](const uint_t entry[2]) const {
YATETO_HOSTDEVICE real_t operator[](const uint_t entry[2]) const {
return operator()(entry[0], entry[1]);
}

template<class view_t>
void copyToView(view_t& other) {
YATETO_HOSTDEVICE void copyToView(view_t& other) {
assert(2 == other.dim());
assert(this->shape(0) == other.shape(0) && this->shape(1) == other.shape(1));

Expand All @@ -399,7 +418,7 @@ namespace yateto {
}
}
}

#pragma omp end declare target
protected:
real_t* m_values;
uint_t const* m_rowInd;
Expand Down
Loading
Loading