fieldFreeFunctions.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: MIT
// SPDX-FileCopyrightText: 2023 NeoN authors
#pragma once
 
#include <tuple>
 
#include <Kokkos_Core.hpp>
#include "NeoN/core/primitives/label.hpp"
#include "NeoN/helpers/exceptions.hpp"
#include "NeoN/core/parallelAlgorithms.hpp"
#include "NeoN/core/view.hpp"
 
namespace NeoN
{
 
// Forward declaration
template<typename ValueType>
class Vector;
 
 
template<typename T, typename Inner>
void map(Vector<T>& a, const Inner inner, std::pair<localIdx, localIdx> range = {0, 0})
{
    auto [start, end] = range;
    if (end == 0)
    {
        end = a.size();
    }
    auto viewA = a.view();
    parallelFor(
        a.exec(), {start, end}, KOKKOS_LAMBDA(const localIdx i) { viewA[i] = inner(i); }
    );
}
 
template<typename ValueType>
void fill(
    Vector<ValueType>& a,
    const std::type_identity_t<ValueType> value,
    std::pair<localIdx, localIdx> range = {0, 0}
)
{
    auto [start, end] = range;
    if (end == 0)
    {
        end = a.size();
    }
    auto viewA = a.view();
    parallelFor(
        a.exec(), {start, end}, KOKKOS_LAMBDA(const localIdx i) { viewA[i] = value; }
    );
}
 
 
template<typename ValueType>
void setVector(
    Vector<ValueType>& a,
    const View<const std::type_identity_t<ValueType>> b,
    std::pair<localIdx, localIdx> range = {0, 0}
)
{
    auto [start, end] = range;
    if (end == 0)
    {
        end = a.size();
    }
    auto viewA = a.view();
    parallelFor(
        a.exec(), {start, end}, KOKKOS_LAMBDA(const localIdx i) { viewA[i] = b[i]; }
    );
}
 
template<typename ValueType>
void scalarMul(Vector<ValueType>& a, const scalar value)
{
    auto viewA = a.view();
    parallelFor(
        a, KOKKOS_LAMBDA(const localIdx i) { return viewA[i] * value; }
    );
}
 
namespace detail
{
template<typename ValueType, typename BinaryOp>
void fieldBinaryOp(
    Vector<ValueType>& a, const Vector<std::type_identity_t<ValueType>>& b, BinaryOp op
)
{
    NeoN_ASSERT_EQUAL_LENGTH(a, b);
    auto viewA = a.view();
    auto viewB = b.view();
    parallelFor(
        a, KOKKOS_LAMBDA(const localIdx i) { return op(viewA[i], viewB[i]); }
    );
}
}
 
template<typename ValueType>
void add(Vector<ValueType>& a, const Vector<std::type_identity_t<ValueType>>& b)
{
    detail::fieldBinaryOp(
        a, b, KOKKOS_LAMBDA(ValueType va, ValueType vb) { return va + vb; }
    );
}
 
 
template<typename ValueType>
void sub(Vector<ValueType>& a, const Vector<std::type_identity_t<ValueType>>& b)
{
    detail::fieldBinaryOp(
        a, b, KOKKOS_LAMBDA(ValueType va, ValueType vb) { return va - vb; }
    );
}
 
template<typename ValueType>
void mul(Vector<ValueType>& a, const Vector<std::type_identity_t<ValueType>>& b)
{
    detail::fieldBinaryOp(
        a, b, KOKKOS_LAMBDA(ValueType va, ValueType vb) { return va * vb; }
    );
}
 
template<typename... Args>
auto spans(Args&... fields)
{
    return std::make_tuple(fields.view()...);
}
 
template<typename... Args>
auto copyToHosts(Args&... fields)
{
    return std::make_tuple(fields.copyToHost()...);
}
 
template<typename T>
bool equal(Vector<T>& field, T value)
{
    auto hostVector = field.copyToHost();
    auto hostView = hostVector.view();
    for (localIdx i = 0; i < hostView.size(); i++)
    {
        if (hostView[i] != value)
        {
            return false;
        }
    }
    return true;
};
 
template<typename T>
bool equal(const Vector<T>& field, const Vector<T>& field2)
{
    auto [hostVector, hostVector2] = copyToHosts(field, field2);
    auto [hostSpan, hostSpan2] = spans(hostVector, hostVector2);
 
    if (hostSpan.size() != hostSpan2.size())
    {
        return false;
    }
 
    for (localIdx i = 0; i < hostSpan.size(); i++)
    {
        if (hostSpan[i] != hostSpan2[i])
        {
            return false;
        }
    }
 
    return true;
};
 
template<typename T>
bool equal(const Vector<T>& field, View<T> span2)
{
    auto hostView = field.copyToHost().view();
 
    if (hostView.size() != span2.size())
    {
        return false;
    }
 
    for (localIdx i = 0; i < hostView.size(); i++)
    {
        if (hostView[i] != span2[i])
        {
            return false;
        }
    }
 
    return true;
}
 
} // namespace NeoN