Homec4science
Diffusion tamaas (ci-fix-static-checks) Browse

tamaas/src/core/iterator.hh
7bb4bacf6434ci-fix-static-checks

Last Change
Raw File

iterator.hh
View Options

/*
* SPDX-License-Indentifier: AGPL-3.0-or-later
*
* Copyright (©) 2016-2022 EPFL (École Polytechnique Fédérale de Lausanne),
* Laboratory (LSMS - Laboratoire de Simulation en Mécanique des Solides)
* Copyright (©) 2020-2022 Lucas Frérot
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU Affero General Public License as published
* by the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU Affero General Public License for more details.
*
* You should have received a copy of the GNU Affero General Public License
* along with this program. If not, see <https://www.gnu.org/licenses/>.
*
*/
/* -------------------------------------------------------------------------- */
#ifndef ITERATOR_HH
#define ITERATOR_HH
/* -------------------------------------------------------------------------- */
#include "tamaas.hh"
#include <cstddef>
#include <iterator>
#include <thrust/iterator/iterator_categories.h>
#include <utility>
#include <vector>
/* -------------------------------------------------------------------------- */
namespace tamaas {
namespace iterator_ {
template <typename T>
class iterator {
public:
using value_type = T;
using difference_type = std::ptrdiff_t;
using pointer = T*;
using reference = T&;
using iterator_category = thrust::random_access_device_iterator_tag;
template <typename U>
friend class iterator;
public:
/// constructor
__device__ __host__ iterator(pointer data, difference_type step_size)
: data(data), step(step_size) {}
/// copy constructor
__device__ __host__ iterator(const iterator& o)
: data(o.data), step(o.step) {}
/// move constructor
__device__ __host__ iterator(iterator&& o) noexcept
: data(exchange(o.data, nullptr)), step(exchange(o.step, 0)) {}
/// copy from a different qualified type
template <typename U,
typename = std::enable_if_t<std::is_convertible<T, U>::value>>
__device__ __host__ iterator(const iterator<U>& o)
: data(o.data), step(o.step) {}
/// move from a different qualified type
template <typename U,
typename = std::enable_if_t<std::is_convertible<T, U>::value>>
__device__ __host__ iterator(iterator<U>&& o)
: data(exchange(o.data, nullptr)), step(exchange(o.step, 0)) {}
// assign
__device__ __host__ iterator& operator=(const iterator& o) {
data = o.data;
step = o.step;
return *this;
}
// move
__device__ __host__ iterator& operator=(iterator&& o) noexcept {
data = exchange(o.data, nullptr);
step = exchange(o.step, 0);
return *this;
}
public:
/// dereference iterator
__device__ __host__ inline reference operator*() { return *data; }
/// dereference iterator
__device__ __host__ inline reference operator*() const { return *data; }
/// increment with given offset
__device__ __host__ inline iterator& operator+=(difference_type a) {
data += a * step;
return *this;
}
/// increment iterator
__device__ __host__ inline iterator& operator++() {
*this += 1;
return *this;
}
/// comparison
__device__ __host__ inline bool operator<(const iterator& a) const {
return data < a.data;
}
/// inequality
__device__ __host__ inline bool operator!=(const iterator& a) const {
return data != a.data;
}
/// equality
__device__ __host__ inline bool operator==(const iterator& a) const {
return data == a.data;
}
/// needed for OpenMP range calculations
__device__ __host__ inline difference_type
operator-(const iterator& a) const {
return (data - a.data) / step;
}
void setStep(difference_type s) { step = s; }
protected:
T* data;
difference_type step;
};
template <typename T>
class iterator_view : iterator<T> {
public:
using value_type = T;
using difference_type = std::ptrdiff_t;
using pointer = T*;
using reference = T&;
/// Constructor
iterator_view(pointer data, std::size_t start, ptrdiff_t offset,
std::vector<UInt> strides, std::vector<UInt> sizes)
: iterator<T>(data, start, offset), strides(std::move(strides)),
n(std::move(sizes)), tuple(strides.size()) {}
iterator_view(const iterator_view& o)
: iterator<T>(o), strides(o.strides), n(o.n), tuple(o.tuple) {}
iterator_view& operator=(const iterator_view& o) {
iterator<T>::operator=(o);
strides = o.strides;
n = o.n;
tuple = o.tuple;
}
protected:
void computeAccessOffset() {
UInt index_copy = this->index;
this->tuple.back() = index_copy % this->n.back();
index_copy -= this->tuple.back();
index_copy /= this->n.back();
/// Computing tuple from index
for (UInt d = this->tuple.size() - 2; d > 0; d--) {
this->tuple[d] = index_copy % this->n[d];
index_copy -= this->tuple[d];
index_copy /= this->n[d];
}
this->tuple[0] = index_copy;
this->access_offset = 0;
for (UInt i = 0; i < this->tuple.size(); ++i)
this->access_offset += this->tuple[i] * this->strides[i];
}
std::vector<UInt> strides;
std::vector<UInt> n;
std::vector<UInt> tuple;
};
} // end namespace iterator_
} // namespace tamaas
/* -------------------------------------------------------------------------- */
#endif // ITERATOR_HH
c4science · Help