RegisterBase.hpp

Go to the documentation of this file.

 
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
// Copyright (c) Lawrence Livermore National Security, LLC and other
// RAJA Project Developers. See top-level LICENSE and COPYRIGHT
// files for dates and other details. No copyright assignment is required
// to contribute to RAJA.
//
// SPDX-License-Identifier: (BSD-3-Clause)
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
 
#ifndef RAJA_pattern_tensor_RegisterBase_HPP
#define RAJA_pattern_tensor_RegisterBase_HPP
 
#include "RAJA/config.hpp"
 
#include "RAJA/util/macros.hpp"
 
#include "camp/camp.hpp"
#include "RAJA/pattern/tensor/TensorLayout.hpp"
#include "RAJA/pattern/tensor/internal/TensorRef.hpp"
#include "RAJA/util/BitMask.hpp"
 
#include "RAJA/policy/tensor/arch.hpp"
 
namespace RAJA
{
namespace expt
{
template<typename T, typename REGISTER_POLICY>
class Register;
}
 
namespace internal
{
namespace expt
{
class RegisterConcreteBase
{};
 
/*
 * Overload for:    arithmetic + TensorRegister
 
 */
template<
    typename LEFT,
    typename RIGHT,
    typename std::enable_if<std::is_arithmetic<LEFT>::value, bool>::type = true,
    typename std::enable_if<std::is_base_of<RegisterConcreteBase, RIGHT>::value,
                            bool>::type                                  = true>
RAJA_INLINE RAJA_HOST_DEVICE RIGHT operator+(LEFT const& lhs, RIGHT const& rhs)
{
  return RIGHT(lhs).add(rhs);
}
 
/*
 * Overload for:    arithmetic - TensorRegister
 
 */
template<
    typename LEFT,
    typename RIGHT,
    typename std::enable_if<std::is_arithmetic<LEFT>::value, bool>::type = true,
    typename std::enable_if<std::is_base_of<RegisterConcreteBase, RIGHT>::value,
                            bool>::type                                  = true>
RAJA_INLINE RAJA_HOST_DEVICE RIGHT operator-(LEFT const& lhs, RIGHT const& rhs)
{
  return RIGHT(lhs).subtract(rhs);
}
 
/*
 * Overload for:    arithmetic * TensorRegister
 
 */
template<
    typename LEFT,
    typename RIGHT,
    typename std::enable_if<std::is_arithmetic<LEFT>::value, bool>::type = true,
    typename std::enable_if<std::is_base_of<RegisterConcreteBase, RIGHT>::value,
                            bool>::type                                  = true>
RAJA_INLINE RAJA_HOST_DEVICE RIGHT operator*(LEFT const& lhs, RIGHT const& rhs)
{
  return rhs.scale(lhs);
}
 
/*
 * Overload for:    arithmetic / TensorRegister
 
 */
template<
    typename LEFT,
    typename RIGHT,
    typename std::enable_if<std::is_arithmetic<LEFT>::value, bool>::type = true,
    typename std::enable_if<std::is_base_of<RegisterConcreteBase, RIGHT>::value,
                            bool>::type                                  = true>
RAJA_INLINE RAJA_HOST_DEVICE RIGHT operator/(LEFT const& lhs, RIGHT const& rhs)
{
  return RIGHT(lhs).divide(rhs);
}
 
 
template<typename Derived>
class RegisterBase;
 
template<typename T, typename REGISTER_POLICY>
class RegisterBase<RAJA::expt::Register<T, REGISTER_POLICY>>
    : public RegisterConcreteBase
{
public:
  using self_type    = RAJA::expt::Register<T, REGISTER_POLICY>;
  using element_type = camp::decay<T>;
 
  using index_type = camp::idx_t;
 
  using int_element_type =
      typename RegisterTraits<REGISTER_POLICY, T>::int_element_type;
  using int_vector_type =
      RAJA::expt::Register<int_element_type, REGISTER_POLICY>;
 
private:
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type* getThis() { return static_cast<self_type*>(this); }
 
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  constexpr self_type const* getThis() const
  {
    return static_cast<self_type const*>(this);
  }
 
public:
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  static constexpr bool is_root() { return true; }
 
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  constexpr RegisterBase() {}
 
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  ~RegisterBase() {}
 
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  constexpr RegisterBase(RegisterBase const&) {}
 
  RAJA_INLINE
 
  RAJA_HOST_DEVICE
  constexpr RegisterBase(self_type const&) {}
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  static self_type s_broadcast_n(element_type const& value, camp::idx_t N)
  {
    self_type x;
    for (camp::idx_t i = 0; i < N; ++i)
    {
      x.set(value, i);
    }
    return x;
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type get_and_broadcast(int i) const
  {
    self_type x;
    x.broadcast(getThis()->get(i));
    return x;
  }
 
  template<typename T2>
  RAJA_HOST_DEVICE RAJA_INLINE self_type& gather(
      element_type const* ptr,
      RAJA::expt::Register<T2, REGISTER_POLICY> offsets)
  {
#ifdef RAJA_ENABLE_VECTOR_STATS
    RAJA::tensor_stats::num_vector_load_strided_n++;
#endif
    for (camp::idx_t i = 0; i < self_type::s_num_elem; ++i)
    {
      getThis()->set(ptr[offsets.get(i)], i);
    }
    return *getThis();
  }
 
  template<typename T2>
  RAJA_HOST_DEVICE RAJA_INLINE self_type& gather_n(
      element_type const* ptr,
      RAJA::expt::Register<T2, REGISTER_POLICY> const& offsets,
      camp::idx_t N)
  {
#ifdef RAJA_ENABLE_VECTOR_STATS
    RAJA::tensor_stats::num_vector_load_strided_n++;
#endif
    for (camp::idx_t i = 0; i < N; ++i)
    {
      getThis()->set(ptr[offsets.get(i)], i);
    }
    return *getThis();
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type& segmented_load(element_type const* ptr,
                            camp::idx_t segbits,
                            camp::idx_t stride_inner,
                            camp::idx_t stride_outer)
  {
    getThis()->gather(ptr, self_type::s_segmented_offsets(segbits, stride_inner,
                                                          stride_outer));
    return *getThis();
  }
 
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type& segmented_load_nm(element_type const* ptr,
                               camp::idx_t segbits,
                               camp::idx_t stride_inner,
                               camp::idx_t stride_outer,
                               camp::idx_t num_inner,
                               camp::idx_t num_outer)
  {
 
    camp::idx_t num_segments = self_type::s_num_elem >> segbits;
    camp::idx_t seg_size     = 1 << segbits;
 
    camp::idx_t lane = 0;
    for (camp::idx_t seg = 0; seg < num_segments; ++seg)
    {
      for (camp::idx_t i = 0; i < seg_size; ++i)
      {
 
        if (seg >= num_outer || i >= num_inner)
        {
          getThis()->set(element_type(0), lane);
        }
        else
        {
 
          camp::idx_t offset = seg * stride_outer + i * stride_inner;
 
          element_type value = ptr[offset];
 
          getThis()->set(value, lane);
        }
 
        lane++;
      }
    }
 
    return *getThis();
  }
 
  template<typename T2>
  RAJA_HOST_DEVICE RAJA_INLINE self_type const& scatter(
      element_type* ptr,
      RAJA::expt::Register<T2, REGISTER_POLICY> const& offsets) const
  {
#ifdef RAJA_ENABLE_VECTOR_STATS
    RAJA::tensor_stats::num_vector_load_strided_n++;
#endif
    for (camp::idx_t i = 0; i < self_type::s_num_elem; ++i)
    {
      ptr[offsets.get(i)] = getThis()->get(i);
    }
    return *getThis();
  }
 
  template<typename T2>
  RAJA_HOST_DEVICE RAJA_INLINE self_type const& scatter_n(
      element_type* ptr,
      RAJA::expt::Register<T2, REGISTER_POLICY> const& offsets,
      camp::idx_t N) const
  {
#ifdef RAJA_ENABLE_VECTOR_STATS
    RAJA::tensor_stats::num_vector_load_strided_n++;
#endif
    for (camp::idx_t i = 0; i < N; ++i)
    {
      ptr[offsets.get(i)] = getThis()->get(i);
    }
    return *getThis();
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type const& segmented_store(element_type* ptr,
                                   camp::idx_t segbits,
                                   camp::idx_t stride_inner,
                                   camp::idx_t stride_outer) const
  {
    getThis()->scatter(ptr, self_type::s_segmented_offsets(
                                segbits, stride_inner, stride_outer));
    return *getThis();
  }
 
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type const& segmented_store_nm(element_type* ptr,
                                      camp::idx_t segbits,
                                      camp::idx_t stride_inner,
                                      camp::idx_t stride_outer,
                                      camp::idx_t num_inner,
                                      camp::idx_t num_outer) const
  {
 
    camp::idx_t num_segments = self_type::s_num_elem >> segbits;
    camp::idx_t seg_size     = 1 << segbits;
 
    camp::idx_t lane = 0;
    for (camp::idx_t seg = 0; seg < num_segments; ++seg)
    {
      for (camp::idx_t i = 0; i < seg_size; ++i)
      {
 
        if (!(seg >= num_outer || i >= num_inner))
        {
 
          camp::idx_t offset = seg * stride_outer + i * stride_inner;
 
          ptr[offset] = getThis()->get(lane);
        }
 
        lane++;
      }
    }
 
    return *getThis();
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type& operator=(element_type value)
  {
    getThis()->broadcast(value);
    return *getThis();
  }
 
  RAJA_SUPPRESS_HD_WARN
  template<typename T2>
  RAJA_HOST_DEVICE RAJA_INLINE self_type& operator=(
      RAJA::expt::Register<T2, RAJA::expt::scalar_register> const& value)
  {
    getThis()->broadcast(value.get(0));
    return *getThis();
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type& operator=(self_type const& x)
  {
    getThis()->copy(x);
    return *getThis();
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type operator+(self_type const& x) const { return getThis()->add(x); }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type& operator+=(self_type const& x)
  {
    *getThis() = getThis()->add(x);
    return *getThis();
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type operator+(element_type const& x) const { return getThis()->add(x); }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type& operator+=(element_type x)
  {
    *getThis() = getThis()->add(x);
    return *getThis();
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type operator-() const { return self_type(0).subtract(*getThis()); }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type operator-(self_type const& x) const
  {
    return getThis()->subtract(x);
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type& operator-=(self_type const& x)
  {
    *getThis() = getThis()->subtract(x);
    return *getThis();
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type operator-(element_type const& x) const
  {
    return getThis()->subtract(x);
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type& operator-=(element_type const& x)
  {
    *getThis() = getThis()->subtract(x);
    return *getThis();
  }
 
  template<typename RHS>
  RAJA_HOST_DEVICE RAJA_INLINE self_type operator*(RHS const& rhs) const
  {
    return getThis()->multiply(rhs);
  }
 
  template<typename RHS>
  RAJA_HOST_DEVICE RAJA_INLINE self_type& operator*=(RHS const& rhs)
  {
    *getThis() = getThis()->multiply(rhs);
    return *getThis();
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_INLINE
 
  RAJA_HOST_DEVICE
  self_type operator/(self_type const& x) const { return getThis()->divide(x); }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type& operator/=(self_type const& x)
  {
    *getThis() = getThis()->divide(x);
    return *getThis();
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_INLINE
 
  RAJA_HOST_DEVICE
  self_type operator/(element_type const& x) const
  {
    return getThis()->divide(x);
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type& operator/=(element_type const& x)
  {
    *getThis() = getThis()->divide(x);
    return *getThis();
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type divide_n(self_type const& b, camp::idx_t n) const
  {
    self_type q(*getThis());
    for (camp::idx_t i = 0; i < n; ++i)
    {
      q.set(getThis()->get(i) / b.get(i), i);
    }
    return q;
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_HOST_DEVICE
 
  RAJA_INLINE
  self_type divide_n(element_type const& b, camp::idx_t n) const
  {
    self_type q(*getThis());
    for (camp::idx_t i = 0; i < n; ++i)
    {
      q.set(getThis()->get(i) / b, i);
    }
    return q;
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_INLINE
 
  RAJA_HOST_DEVICE
  element_type dot(self_type const& x) const
  {
    return getThis()->multiply(x).sum();
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_INLINE
 
  RAJA_HOST_DEVICE
  self_type multiply_add(self_type const& b, self_type const& c) const
  {
    return (self_type(*getThis()) * self_type(b)) + self_type(c);
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_INLINE
 
  RAJA_HOST_DEVICE
  self_type multiply_subtract(self_type const& b, self_type const& c) const
  {
    return getThis()->multiply_add(b, -c);
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_INLINE
 
  RAJA_HOST_DEVICE
  self_type scale(element_type c) const
  {
    return getThis()->multiply(self_type(c));
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_INLINE
 
  RAJA_HOST_DEVICE
  element_type min_n(camp::idx_t N) const { return getThis()->min(N); }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_INLINE
 
  RAJA_HOST_DEVICE
  element_type max_n(camp::idx_t N) const { return getThis()->max(N); }
 
  RAJA_INLINE
 
  RAJA_HOST_DEVICE
  self_type transpose_shuffle_left(camp::idx_t lvl, self_type const& y) const
  {
    auto const& x = *getThis();
 
    self_type z;
 
    for (camp::idx_t i = 0; i < self_type::s_num_elem; ++i)
    {
 
      // extract value x or y
      camp::idx_t xy_select = (i >> lvl) & 0x1;
 
 
      z.set(xy_select == 0 ? x.get(i) : y.get(i - (1 << lvl)), i);
    }
 
    return z;
  }
 
  RAJA_INLINE
 
  RAJA_HOST_DEVICE
  self_type transpose_shuffle_right(int lvl, self_type const& y) const
  {
    auto const& x = *getThis();
 
    self_type z;
 
    camp::idx_t i0 = 1 << lvl;
 
    for (camp::idx_t i = 0; i < self_type::s_num_elem; ++i)
    {
 
      // extract value x or y
      camp::idx_t xy_select = (i >> lvl) & 0x1;
 
      z.set(xy_select == 0 ? x.get(i0 + i) : y.get(i0 + i - (1 << lvl)), i);
    }
 
    return z;
  }
 
  RAJA_INLINE
  static int_vector_type s_segmented_offsets(camp::idx_t segbits,
                                             camp::idx_t stride_inner,
                                             camp::idx_t stride_outer)
  {
    int_vector_type result;
 
    camp::idx_t num_segments = self_type::s_num_elem >> segbits;
    camp::idx_t seg_size     = 1 << segbits;
 
    camp::idx_t lane = 0;
    for (camp::idx_t seg = 0; seg < num_segments; ++seg)
    {
      for (camp::idx_t i = 0; i < seg_size; ++i)
      {
        result.set(seg * stride_outer + i * stride_inner, lane);
        lane++;
      }
    }
 
    return result;
  }
 
  RAJA_INLINE
  self_type segmented_sum_inner(camp::idx_t segbits,
                                camp::idx_t output_segment) const
  {
    self_type result(0);
 
    // default implementation is dumb, just sum each value into
    // appropriate segment lane
    int output_offset = output_segment * self_type::s_num_elem >> segbits;
 
    for (camp::idx_t i = 0; i < self_type::s_num_elem; ++i)
    {
      auto value =
          getThis()->get(i) + result.get((i >> segbits) + output_offset);
      result.set(value, (i >> segbits) + output_offset);
    }
 
    return result;
  }
 
  RAJA_INLINE
  self_type segmented_sum_outer(camp::idx_t segbits,
                                camp::idx_t output_segment) const
  {
    self_type result(0);
 
    // default implementation is dumb, just sum each value into
    // appropriate segment lane
    int output_offset = output_segment * (1 << segbits);
 
    for (camp::idx_t i = 0; i < self_type::s_num_elem; ++i)
    {
      camp::idx_t output_i = output_offset + (i & ((1 << segbits) - 1));
      auto value           = getThis()->get(i) + result.get(output_i);
      result.set(value, output_i);
    }
 
    return result;
  }
 
  RAJA_INLINE
  self_type segmented_divide_nm(self_type den,
                                camp::idx_t segbits,
                                camp::idx_t num_inner,
                                camp::idx_t num_outer) const
  {
    self_type result;
 
    camp::idx_t num_segments = self_type::s_num_elem >> segbits;
    camp::idx_t seg_size     = 1 << segbits;
 
    camp::idx_t lane = 0;
    for (camp::idx_t seg = 0; seg < num_segments; ++seg)
    {
      for (camp::idx_t i = 0; i < seg_size; ++i)
      {
 
        if (seg >= num_outer || i >= num_inner)
        {
          result.set(element_type(0), lane);
        }
        else
        {
 
          element_type div = getThis()->get(lane) / den.get(lane);
 
          result.set(div, lane);
        }
 
        lane++;
      }
    }
 
    return result;
  }
 
  RAJA_SUPPRESS_HD_WARN
  RAJA_INLINE
 
  RAJA_HOST_DEVICE
  self_type segmented_dot(camp::idx_t segbits,
                          camp::idx_t output_segment,
                          self_type const& x) const
  {
    return getThis()->multiply(x).segmented_sum_inner(segbits, output_segment);
  }
 
  RAJA_INLINE
  self_type segmented_broadcast_inner(camp::idx_t segbits,
                                      camp::idx_t input_segment) const
  {
    self_type result;
 
    camp::idx_t mask   = (1 << segbits) - 1;
    camp::idx_t offset = input_segment << segbits;
 
    // default implementation is dumb, just sum each value into
    // appropriate segment lane
    for (camp::idx_t i = 0; i < self_type::s_num_elem; ++i)
    {
 
      auto off = (i & mask) + offset;
 
      result.set(getThis()->get(off), i);
    }
 
    return result;
  }
 
  RAJA_INLINE
  self_type segmented_broadcast_outer(camp::idx_t segbits,
                                      camp::idx_t input_segment) const
  {
    self_type result;
 
    camp::idx_t offset = input_segment * (self_type::s_num_elem >> segbits);
 
    // default implementation is dumb, just sum each value into
    // appropriate segment lane
    for (camp::idx_t i = 0; i < self_type::s_num_elem; ++i)
    {
 
      auto off = (i >> segbits) + offset;
 
      result.set(getThis()->get(off), i);
    }
 
    return result;
  }
 
  RAJA_INLINE
  std::string to_string() const
  {
    std::string s = "Register(" + std::to_string(self_type::s_num_elem) + ")[ ";
 
    //
    for (camp::idx_t i = 0; i < self_type::s_num_elem; ++i)
    {
      s += std::to_string(getThis()->get(i)) + " ";
    }
 
    s += " ]\n";
 
    return s;
  }
};
 
 
}  // namespace expt
}  // namespace internal
}  // namespace RAJA
 
 
#endif