Homec4science
Diffusion tamaas (features/materials) Browse

tamaas/src/core/loop.hh
4278fed63fdefeatures/materials

Last Change
Raw File

loop.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)
*
* 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 LOOP_HH
#define LOOP_HH
/* -------------------------------------------------------------------------- */
#include "loops/apply.hh"
#include "loops/loop_utils.hh"
#include "mpi_interface.hh"
#include "ranges.hh"
#include "tamaas.hh"
#include <thrust/execution_policy.h>
#include <thrust/for_each.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/zip_iterator.h>
#include <thrust/transform_reduce.h>
#include <thrust/tuple.h>
#include <thrust/version.h>
#include <type_traits>
#include <utility>
namespace tamaas {
template <typename T>
struct is_policy : std::false_type {};
template <>
struct is_policy<thrust::detail::host_t> : std::true_type {};
template <>
struct is_policy<const thrust::detail::host_t> : std::true_type {};
template <>
struct is_policy<const thrust::detail::host_t&> : std::true_type {};
// device_t == host_t when thrust backend is CPP
#if TAMAAS_LOOP_BACKEND != TAMAAS_LOOP_BACKEND_CPP
template <>
struct is_policy<thrust::detail::device_t> : std::true_type {};
template <>
struct is_policy<const thrust::detail::device_t> : std::true_type {};
template <>
struct is_policy<const thrust::detail::device_t&> : std::true_type {};
#endif
/**
* @brief Singleton class for automated loops using lambdas
* This class is sweet candy :) It provides abstraction of the paralelism
* paradigm used in loops and allows simple and less error-prone loop syntax,
* with minimum boiler plate. I love it <3
*/
class Loop {
public:
/// Helper class to count iterations within lambda-loop
template <typename T>
class arange {
public:
using it_type = thrust::counting_iterator<T>;
using reference = typename it_type::reference;
arange(T start, T size) : start(start), range_size(size) {}
it_type begin(UInt /*i*/ = 1) const { return it_type(T(start)); }
it_type end(UInt /*i*/ = 1) const { return it_type(range_size); }
UInt getNbComponents() const { return 1; }
private:
T start, range_size;
};
template <typename T>
static arange<T> range(T size) {
return arange<T>(0, size);
}
template <typename T, typename U>
static arange<T> range(U start, T size) {
return arange<T>(start, size);
}
public:
/// Replacement for thrust::transform_iterator which copies values
template <typename Iterator, typename Functor, typename Value>
class transform_iterator
: public thrust::iterator_adaptor<
transform_iterator<Iterator, Functor, Value>, Iterator, Value,
thrust::use_default, thrust::use_default, Value> {
Functor func;
public:
using super_t =
thrust::iterator_adaptor<transform_iterator<Iterator, Functor, Value>,
Iterator, Value, thrust::use_default,
thrust::use_default, Value>;
__host__ __device__ transform_iterator(const Iterator& x,
const Functor& func)
: super_t(x), func(func) {}
friend class thrust::iterator_core_access;
private:
__host__ __device__ auto dereference() const
-> decltype(func(*this->base())) {
return func(*this->base());
}
};
public:
/// Loop functor over ranges
template <typename Functor, typename... Ranges>
static auto loop(Functor&& func, Ranges&&... ranges) ->
typename std::enable_if<not is_policy<Functor>::value, void>::type {
loopImpl(thrust::device, std::forward<Functor>(func),
std::forward<Ranges>(ranges)...);
}
/// Loop over ranges with non-default policy
template <typename DerivedPolicy, typename Functor, typename... Ranges>
static void loop(const thrust::execution_policy<DerivedPolicy>& policy,
Functor&& func, Ranges&&... ranges) {
loopImpl(policy, std::forward<Functor>(func),
std::forward<Ranges>(ranges)...);
}
private:
template <typename T>
using reference_type = typename std::decay<T>::type::reference;
public:
/// Reduce functor over ranges
template <operation op = operation::plus, typename Functor,
typename... Ranges>
static auto reduce(Functor&& func, Ranges&&... ranges) ->
typename std::enable_if<
not is_policy<Functor>::value,
decltype(func(std::declval<reference_type<Ranges>>()...))>::type {
return reduceImpl<op>(thrust::device, std::forward<Functor>(func),
std::forward<Ranges>(ranges)...);
}
/// Reduce over ranges with non-default policy
template <operation op = operation::plus, typename DerivedPolicy,
typename Functor, typename... Ranges>
static auto reduce(const thrust::execution_policy<DerivedPolicy>& policy,
Functor&& func, Ranges&&... ranges)
-> decltype(func(std::declval<reference_type<Ranges>>()...)) {
return reduceImpl<op>(policy, std::forward<Functor>(func),
std::forward<Ranges>(ranges)...);
}
private:
/// Loop over ranges and apply functor
template <typename DerivedPolicy, typename Functor, typename... Ranges>
static void loopImpl(const thrust::execution_policy<DerivedPolicy>& policy,
Functor&& func, Ranges&&... ranges);
/// Loop over ranges, apply functor and reduce result
template <operation op, typename DerivedPolicy, typename Functor,
typename... Ranges>
static auto reduceImpl(const thrust::execution_policy<DerivedPolicy>& policy,
Functor&& func, Ranges&&... ranges)
-> decltype(func(std::declval<reference_type<Ranges>>()...));
public:
/// Constructor
Loop() = delete;
};
/* -------------------------------------------------------------------------- */
/* Template implementation */
/* -------------------------------------------------------------------------- */
template <typename DerivedPolicy, typename Functor, typename... Ranges>
void Loop::loopImpl(const thrust::execution_policy<DerivedPolicy>& policy,
Functor&& func, Ranges&&... ranges) {
auto begin = thrust::make_zip_iterator(thrust::make_tuple(ranges.begin()...));
auto end = thrust::make_zip_iterator(thrust::make_tuple(ranges.end()...));
checkLoopSize(ranges...);
thrust::for_each(policy, begin, end,
detail::ApplyFunctor<Functor>(std::forward<Functor>(func)));
#if (defined(TAMAAS_USE_CUDA) && THRUST_VERSION < 100904)
cudaDeviceSynchronize();
#endif
}
/* -------------------------------------------------------------------------- */
template <operation op, typename DerivedPolicy, typename Functor,
typename... Ranges>
auto Loop::reduceImpl(const thrust::execution_policy<DerivedPolicy>& policy,
Functor&& func, Ranges&&... ranges)
-> decltype(func(std::declval<reference_type<Ranges>>()...)) {
using return_type = decltype(func(std::declval<reference_type<Ranges>>()...));
using apply_type = detail::ApplyFunctor<Functor, return_type>;
checkLoopSize(ranges...);
auto applier = apply_type(func);
detail::reduction_helper<op, return_type> red_helper;
auto begin_zip =
thrust::make_zip_iterator(thrust::make_tuple(ranges.begin()...));
auto end_zip = thrust::make_zip_iterator(thrust::make_tuple(ranges.end()...));
transform_iterator<decltype(begin_zip), apply_type, return_type> begin(
begin_zip, applier);
transform_iterator<decltype(end_zip), apply_type, return_type> end(end_zip,
applier);
auto result =
thrust::reduce(policy, begin, end, red_helper.init(), red_helper);
#if (defined(TAMAAS_USE_CUDA) && THRUST_VERSION < 100904)
cudaDeviceSynchronize();
#endif
return mpi::allreduce<op>(std::move(result));
}
} // namespace tamaas
#endif // LOOP_HH
c4science · Help