KokkosScalar.h

Go to the documentation of this file.

//* This file is part of the MOOSE framework
//* https://www.mooseframework.org
//*
//* All rights reserved, see COPYRIGHT for full restrictions
//* https://github.com/idaholab/moose/blob/master/COPYRIGHT
//*
//* Licensed under LGPL 2.1, please see LICENSE for details
//* https://www.gnu.org/licenses/lgpl-2.1.html

#pragma once

#include "KokkosReferenceWrapper.h"

#include <type_traits>
#include <utility>

#ifdef MOOSE_KOKKOS_SCOPE
#define KOKKOS_SCALAR_FUNCTION KOKKOS_FUNCTION
#else
#define KOKKOS_SCALAR_FUNCTION
#endif

namespace Moose::Kokkos
{

template <typename T, typename = typename std::enable_if<std::is_arithmetic<T>::value>::type>
class Scalar : public ReferenceWrapper<T>
{
public:
  using value_type = T;

  Scalar(T & value) : ReferenceWrapper<T>(value) {}
  Scalar(const Scalar & object) : ReferenceWrapper<T>(object) {}

  auto & operator=(const Scalar & value)
  {
    this->_reference = static_cast<const T &>(value);

    return *this;
  }
  template <typename U>
  auto operator=(const U & value) -> decltype(std::declval<T &>() = value, std::declval<Scalar &>())
  {
    this->_reference = value;

    return *this;
  }

  KOKKOS_SCALAR_FUNCTION auto operator+() const -> decltype(+std::declval<const T &>())
  {
    return +value();
  }
  KOKKOS_SCALAR_FUNCTION auto operator-() const -> decltype(-std::declval<const T &>())
  {
    return -value();
  }

  template <typename U>
  auto operator+=(const U & value)
      -> decltype(std::declval<T &>() += value, std::declval<Scalar &>())
  {
    this->_reference += value;

    return *this;
  }
  template <typename U>
  auto operator-=(const U & value)
      -> decltype(std::declval<T &>() -= value, std::declval<Scalar &>())
  {
    this->_reference -= value;

    return *this;
  }
  template <typename U>
  auto operator*=(const U & value)
      -> decltype(std::declval<T &>() *= value, std::declval<Scalar &>())
  {
    this->_reference *= value;

    return *this;
  }
  template <typename U>
  auto operator/=(const U & value)
      -> decltype(std::declval<T &>() /= value, std::declval<Scalar &>())
  {
    this->_reference /= value;

    return *this;
  }
  template <typename U>
  auto operator%=(const U & value)
      -> decltype(std::declval<T &>() %= value, std::declval<Scalar &>())
  {
    this->_reference %= value;

    return *this;
  }

  template <typename U = T>
  auto operator++() -> decltype(++std::declval<U &>(), std::declval<Scalar &>())
  {
    ++this->_reference;

    return *this;
  }
  template <typename U = T>
  auto operator++(int) -> decltype(std::declval<U &>()++)
  {
    return this->_reference++;
  }
  template <typename U = T>
  auto operator--() -> decltype(--std::declval<U &>(), std::declval<Scalar &>())
  {
    --this->_reference;

    return *this;
  }
  template <typename U = T>
  auto operator--(int) -> decltype(std::declval<U &>()--)
  {
    return this->_reference--;
  }

  template <typename U>
  KOKKOS_SCALAR_FUNCTION auto operator+(const U & value) const
      -> decltype(std::declval<const T &>() + value)
  {
    return this->value() + value;
  }
  template <typename U>
  KOKKOS_SCALAR_FUNCTION auto operator-(const U & value) const
      -> decltype(std::declval<const T &>() - value)
  {
    return this->value() - value;
  }
  template <typename U>
  KOKKOS_SCALAR_FUNCTION auto operator*(const U & value) const
      -> decltype(std::declval<const T &>() * value)
  {
    return this->value() * value;
  }
  template <typename U>
  KOKKOS_SCALAR_FUNCTION auto operator/(const U & value) const
      -> decltype(std::declval<const T &>() / value)
  {
    return this->value() / value;
  }
  template <typename U>
  KOKKOS_SCALAR_FUNCTION auto operator%(const U & value) const
      -> decltype(std::declval<const T &>() % value)
  {
    return this->value() % value;
  }

private:
  KOKKOS_SCALAR_FUNCTION const T & value() const { return static_cast<const T &>(*this); }
};

template <typename>
struct is_scalar : std::false_type
{
};

template <typename T, typename Enable>
struct is_scalar<Scalar<T, Enable>> : std::true_type
{
};

template <typename T,
          typename U,
          typename = typename std::enable_if<!is_scalar<typename std::decay<T>::type>::value>::type>
KOKKOS_SCALAR_FUNCTION auto
operator+(const T & left, const Scalar<U> & right) -> decltype(left + static_cast<const U &>(right))
{
  return left + static_cast<const U &>(right);
}

template <typename T,
          typename U,
          typename = typename std::enable_if<!is_scalar<typename std::decay<T>::type>::value>::type>
KOKKOS_SCALAR_FUNCTION auto
operator-(const T & left, const Scalar<U> & right) -> decltype(left - static_cast<const U &>(right))
{
  return left - static_cast<const U &>(right);
}

template <typename T,
          typename U,
          typename = typename std::enable_if<!is_scalar<typename std::decay<T>::type>::value>::type>
KOKKOS_SCALAR_FUNCTION auto
operator*(const T & left, const Scalar<U> & right) -> decltype(left * static_cast<const U &>(right))
{
  return left * static_cast<const U &>(right);
}

template <typename T,
          typename U,
          typename = typename std::enable_if<!is_scalar<typename std::decay<T>::type>::value>::type>
KOKKOS_SCALAR_FUNCTION auto
operator/(const T & left, const Scalar<U> & right) -> decltype(left / static_cast<const U &>(right))
{
  return left / static_cast<const U &>(right);
}

template <typename T,
          typename U,
          typename = typename std::enable_if<!is_scalar<typename std::decay<T>::type>::value>::type>
KOKKOS_SCALAR_FUNCTION auto
operator%(const T & left, const Scalar<U> & right) -> decltype(left % static_cast<const U &>(right))
{
  return left % static_cast<const U &>(right);
}

template <typename T>
struct ArrayDeepCopy<Scalar<T>>
{
  static constexpr bool value = true;
};

// Mimic MOOSE convention
using PostprocessorValue = Scalar<const PostprocessorValue>;

} // namespace Moose::Kokkos