Layout.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_LAYOUT_HPP
#define RAJA_LAYOUT_HPP
 
#include "RAJA/config.hpp"
 
#include <cassert>
#include <iostream>
#include <limits>
 
#include "RAJA/index/IndexValue.hpp"
 
#include "RAJA/internal/foldl.hpp"
 
#include "RAJA/util/Operators.hpp"
#include "RAJA/util/Permutations.hpp"
 
namespace RAJA
{
 
namespace detail
{
 
 
template<typename Range,
         typename IdxLin        = Index_type,
         ptrdiff_t StrideOneDim = -1>
struct LayoutBase_impl;
 
template<size_t j, size_t n_dims, typename IdxLin = Index_type>
struct stride_calculator
{
  RAJA_INLINE RAJA_HOST_DEVICE constexpr IdxLin operator()(
      IdxLin cur_stride,
      IdxLin const (&sizes)[n_dims]) const
  {
    return stride_calculator<j + 1, n_dims, IdxLin> {}(
        cur_stride * (sizes[j] ? sizes[j] : 1), sizes);
  }
};
 
template<size_t n_dims, typename IdxLin>
struct stride_calculator<n_dims, n_dims, IdxLin>
{
  RAJA_INLINE RAJA_HOST_DEVICE constexpr IdxLin operator()(
      IdxLin cur_stride,
      IdxLin const (&)[n_dims]) const
  {
    return cur_stride;
  }
};
 
template<camp::idx_t... RangeInts, typename IdxLin, ptrdiff_t StrideOneDim>
struct LayoutBase_impl<camp::idx_seq<RangeInts...>, IdxLin, StrideOneDim>
{
public:
  using IndexLinear = IdxLin;
  using IndexRange  = camp::make_idx_seq_t<sizeof...(RangeInts)>;
 
  static inline constexpr size_t n_dims = sizeof...(RangeInts);
  static inline constexpr IdxLin limit = RAJA::operators::limits<IdxLin>::max();
  static inline constexpr ptrdiff_t stride_one_dim = StrideOneDim;
 
  IdxLin sizes[n_dims]       = {0};
  IdxLin strides[n_dims]     = {0};
  IdxLin inv_strides[n_dims] = {0};
  IdxLin inv_mods[n_dims]    = {0};
 
 
  RAJA_INLINE constexpr LayoutBase_impl()                       = default;
  RAJA_INLINE constexpr LayoutBase_impl(LayoutBase_impl const&) = default;
  RAJA_INLINE constexpr LayoutBase_impl(LayoutBase_impl&&)      = default;
  RAJA_INLINE constexpr LayoutBase_impl& operator=(LayoutBase_impl const&) =
      default;
  RAJA_INLINE constexpr LayoutBase_impl& operator=(LayoutBase_impl&&) = default;
 
  template<typename... Types>
  RAJA_INLINE RAJA_HOST_DEVICE constexpr LayoutBase_impl(Types... ns)
      : sizes {static_cast<IdxLin>(stripIndexType(ns))...},
        strides {(detail::stride_calculator<RangeInts + 1, n_dims, IdxLin> {}(
            sizes[RangeInts] ? IdxLin(1) : IdxLin(0),
            sizes))...},
        inv_strides {(strides[RangeInts] ? strides[RangeInts] : IdxLin(1))...},
        inv_mods {(sizes[RangeInts] ? sizes[RangeInts] : IdxLin(1))...}
  {
    static_assert(n_dims == sizeof...(Types),
                  "number of dimensions must match");
  }
 
  template<typename CIdxLin, ptrdiff_t CStrideOneDim>
  RAJA_INLINE RAJA_HOST_DEVICE constexpr LayoutBase_impl(
      const LayoutBase_impl<camp::idx_seq<RangeInts...>,
                            CIdxLin,
                            CStrideOneDim>& rhs)
      : sizes {static_cast<IdxLin>(rhs.sizes[RangeInts])...},
        strides {static_cast<IdxLin>(rhs.strides[RangeInts])...},
        inv_strides {static_cast<IdxLin>(rhs.inv_strides[RangeInts])...},
        inv_mods {static_cast<IdxLin>(rhs.inv_mods[RangeInts])...}
  {}
 
  template<typename... Types>
  RAJA_INLINE constexpr LayoutBase_impl(
      const std::array<IdxLin, n_dims>& sizes_in,
      const std::array<IdxLin, n_dims>& strides_in)
      : sizes {sizes_in[RangeInts]...},
        strides {strides_in[RangeInts]...},
        inv_strides {(strides[RangeInts] ? strides[RangeInts] : IdxLin(1))...},
        inv_mods {(sizes[RangeInts] ? sizes[RangeInts] : IdxLin(1))...}
  {}
 
  template<camp::idx_t N, typename Idx>
  RAJA_INLINE RAJA_HOST_DEVICE void BoundsCheckError(Idx idx) const
  {
    printf("Error at index %d, value %ld is not within bounds [0, %ld] \n",
           static_cast<int>(N), static_cast<long int>(idx),
           static_cast<long int>(sizes[N] - 1));
    RAJA_ABORT_OR_THROW("Out of bounds error \n");
  }
 
  template<camp::idx_t N>
  RAJA_INLINE RAJA_HOST_DEVICE void BoundsCheck() const
  {}
 
  template<camp::idx_t N, typename Idx, typename... Indices>
  RAJA_INLINE RAJA_HOST_DEVICE void BoundsCheck(Idx idx,
                                                Indices... indices) const
  {
    if (sizes[N] > 0 && !(0 <= idx && idx < static_cast<Idx>(sizes[N])))
    {
      BoundsCheckError<N>(idx);
    }
    RAJA_UNUSED_VAR(idx);
    BoundsCheck<N + 1>(indices...);
  }
 
  template<typename... Indices>
  RAJA_INLINE RAJA_HOST_DEVICE RAJA_BOUNDS_CHECK_constexpr IdxLin
  operator()(Indices... indices) const
  {
#if defined(RAJA_BOUNDS_CHECK_INTERNAL)
    BoundsCheck<0>(indices...);
#endif
    // dot product of strides and indices
    return sum<IdxLin>((RangeInts == stride_one_dim
                            ?  // Is this dimension stride-one?
                            indices
                            :  // it's stride one, so dont bother with multiply
                            strides[RangeInts] * indices  // it's not stride one
                        )...);
  }
 
  template<typename... Indices>
  RAJA_INLINE RAJA_HOST_DEVICE void toIndices(IdxLin linear_index,
                                              Indices&&... indices) const
  {
#if defined(RAJA_BOUNDS_CHECK_INTERNAL)
    IdxLin totSize = size_noproj();
    if (totSize > 0 && (linear_index < 0 || linear_index >= totSize))
    {
      printf("Error! Linear index %ld is not within bounds [0, %ld]. \n",
             static_cast<long int>(linear_index),
             static_cast<long int>(totSize - 1));
      RAJA_ABORT_OR_THROW("Out of bounds error \n");
    }
#endif
 
    camp::sink((indices = (camp::decay<Indices>)((linear_index /
                                                  inv_strides[RangeInts]) %
                                                 inv_mods[RangeInts]))...);
  }
 
  RAJA_INLINE RAJA_HOST_DEVICE constexpr IdxLin size() const
  {
    // Multiply together all of the sizes,
    // replacing 1 for any zero-sized dimensions
    return foldl(
        RAJA::operators::multiplies<IdxLin>(),
        (sizes[RangeInts] == IdxLin(0) ? IdxLin(1) : sizes[RangeInts])...);
  }
 
  RAJA_INLINE RAJA_HOST_DEVICE constexpr IdxLin size_noproj() const
  {
    // Multiply together all of the sizes
    return foldl(RAJA::operators::multiplies<IdxLin>(), sizes[RangeInts]...);
  }
 
  template<camp::idx_t DIM>
  RAJA_INLINE RAJA_HOST_DEVICE constexpr IndexLinear get_dim_stride() const
  {
    return strides[DIM];
  }
 
  template<camp::idx_t DIM>
  RAJA_INLINE RAJA_HOST_DEVICE constexpr IndexLinear get_dim_size() const
  {
    return sizes[DIM];
  }
 
  template<camp::idx_t DIM>
  RAJA_INLINE RAJA_HOST_DEVICE constexpr IndexLinear get_dim_begin() const
  {
    return 0;
  }
};
 
}  // namespace detail
 
template<size_t n_dims, typename IdxLin = Index_type, ptrdiff_t StrideOne = -1>
using Layout =
    detail::LayoutBase_impl<camp::make_idx_seq_t<n_dims>, IdxLin, StrideOne>;
 
template<typename IdxLin, typename DimTuple, ptrdiff_t StrideOne = -1>
struct TypedLayout;
 
template<typename IdxLin, typename... DimTypes, ptrdiff_t StrideOne>
struct TypedLayout<IdxLin, camp::tuple<DimTypes...>, StrideOne>
    : public Layout<sizeof...(DimTypes), strip_index_type_t<IdxLin>, StrideOne>
{
 
  using StrippedIdxLin = strip_index_type_t<IdxLin>;
  using Self   = TypedLayout<IdxLin, camp::tuple<DimTypes...>, StrideOne>;
  using Base   = Layout<sizeof...(DimTypes), StrippedIdxLin, StrideOne>;
  using DimArr = std::array<StrippedIdxLin, sizeof...(DimTypes)>;
 
  // Pull in base constructors
  using Base::Base;
 
  RAJA_INLINE RAJA_HOST_DEVICE constexpr IdxLin operator()(
      DimTypes... indices) const
  {
    return IdxLin(Base::operator()(stripIndexType(indices)...));
  }
 
  RAJA_INLINE RAJA_HOST_DEVICE void toIndices(IdxLin linear_index,
                                              DimTypes&... indices) const
  {
    toIndicesHelper(camp::make_idx_seq_t<sizeof...(DimTypes)> {},
                    std::forward<IdxLin>(linear_index),
                    std::forward<DimTypes&>(indices)...);
  }
 
private:
  template<typename... Indices, camp::idx_t... RangeInts>
  RAJA_INLINE RAJA_HOST_DEVICE void toIndicesHelper(camp::idx_seq<RangeInts...>,
                                                    IdxLin linear_index,
                                                    Indices&... indices) const
  {
    StrippedIdxLin locals[sizeof...(DimTypes)];
    Base::toIndices(stripIndexType(linear_index), locals[RangeInts]...);
    camp::sink(
        (indices = Indices {static_cast<Indices>(locals[RangeInts])})...);
  }
};
 
template<ptrdiff_t s1_dim, size_t n_dims, typename IdxLin>
RAJA_INLINE RAJA_HOST_DEVICE constexpr Layout<n_dims, IdxLin, s1_dim>
make_stride_one(Layout<n_dims, IdxLin> const& l)
{
  return Layout<n_dims, IdxLin, s1_dim>(l);
}
 
template<ptrdiff_t s1_dim, typename IdxLin, typename IdxTuple>
RAJA_INLINE RAJA_HOST_DEVICE constexpr TypedLayout<IdxLin, IdxTuple, s1_dim>
make_stride_one(TypedLayout<IdxLin, IdxTuple> const& l)
{
  // strip l to it's base-class type
  using Base    = typename TypedLayout<IdxLin, IdxTuple>::Base;
  Base const& b = (Base const&)l;
 
  // Use non-typed layout to initialize new typed layout
  return TypedLayout<IdxLin, IdxTuple, s1_dim>(b);
}
 
 
}  // namespace RAJA
 
#endif