indexing.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2021, Lawrence Livermore National Security, LLC and LvArray contributors.
 * All rights reserved.
 * See the LICENSE file for details.
 * SPDX-License-Identifier: (BSD-3-Clause)
 */

#pragma once

// Source includes
#include "Macros.hpp"
#include "typeManipulation.hpp"
#include "limits.hpp"

// TPL includes
#include <RAJA/RAJA.hpp>

namespace LvArray
{

namespace indexing
{

template< bool B_IS_ONE >
struct ConditionalMultiply
{
  template< typename A, typename B >
  static inline LVARRAY_HOST_DEVICE constexpr auto multiply( A const a, B const b )
  { return a * b; }
};

template<>
struct ConditionalMultiply< true >
{
  template< typename A, typename B >
  static inline LVARRAY_HOST_DEVICE constexpr A multiply( A const a, B const & )
  { return a; }
};

template< int SIZE, typename T >
LVARRAY_HOST_DEVICE inline constexpr
std::enable_if_t< (SIZE == 1), T >
multiplyAll( T const * const LVARRAY_RESTRICT values )
{ return values[ 0 ]; }

template< int SIZE, typename T >
LVARRAY_HOST_DEVICE inline constexpr
std::enable_if_t< (SIZE > 1), T >
multiplyAll( T const * const LVARRAY_RESTRICT values )
{ return values[ 0 ] * multiplyAll< SIZE - 1 >( values + 1 ); }

template< int USD, typename INDEX_TYPE, typename INDEX >
LVARRAY_HOST_DEVICE inline constexpr
INDEX_TYPE getLinearIndex( INDEX_TYPE const * const LVARRAY_RESTRICT strides, INDEX const index )
{ return ConditionalMultiply< USD == 0 >::multiply( index, strides[ 0 ] ); }

template< int USD, typename INDEX_TYPE, typename INDEX, typename ... REMAINING_INDICES >
LVARRAY_HOST_DEVICE inline constexpr
INDEX_TYPE getLinearIndex( INDEX_TYPE const * const LVARRAY_RESTRICT strides, INDEX const index, REMAINING_INDICES const ... indices )
{
  return ConditionalMultiply< USD == 0 >::multiply( index, strides[ 0 ] ) +
         getLinearIndex< USD - 1, INDEX_TYPE, REMAINING_INDICES... >( strides + 1, indices ... );
}

inline
std::string getIndexString()
{ return "{}"; }

template< typename INDEX, typename ... REMAINING_INDICES >
std::string getIndexString( INDEX const index, REMAINING_INDICES const ... indices )
{
  std::ostringstream oss;

  oss << "{ " << index;
  using expander = int[];
  (void) expander{ 0, ( void (oss << ", " << indices ), 0 )... };
  oss << " }";

  return oss.str();
}

template< typename INDEX_TYPE, typename ... INDICES >
std::string printDimsAndIndices( INDEX_TYPE const * const LVARRAY_RESTRICT dims, INDICES const... indices )
{
  constexpr int NDIM = sizeof ... (INDICES);
  std::ostringstream oss;
  oss << "dimensions = { " << dims[ 0 ];
  for( int i = 1; i < NDIM; ++i )
  {
    oss << ", " << dims[ i ];
  }

  oss << " } indices = " << getIndexString( indices ... );

  return oss.str();
}

template< typename INDEX_TYPE, typename ... INDICES >
LVARRAY_HOST_DEVICE inline constexpr
bool invalidIndices( INDEX_TYPE const * const LVARRAY_RESTRICT dims, INDICES const ... indices )
{
  int curDim = 0;
  bool invalid = false;
  typeManipulation::forEachArg( [dims, &curDim, &invalid]( auto const index )
  {
    invalid = invalid || ( index < 0 ) || ( index >= dims[ curDim ] );
    ++curDim;
  }, indices ... );

  return invalid;
}

template< typename INDEX_TYPE, typename ... INDICES >
LVARRAY_HOST_DEVICE inline
void checkIndices( INDEX_TYPE const * const LVARRAY_RESTRICT dims, INDICES const ... indices )
{ LVARRAY_ERROR_IF( invalidIndices( dims, indices ... ), "Invalid indices. " << printDimsAndIndices( dims, indices ... ) ); }

template< typename PERMUTATION, typename INDEX_TYPE, camp::idx_t NDIM >
LVARRAY_HOST_DEVICE inline
typeManipulation::CArray< INDEX_TYPE, NDIM > calculateStrides( typeManipulation::CArray< INDEX_TYPE, NDIM > const & dims )
{
  constexpr typeManipulation::CArray< camp::idx_t, NDIM > perm = typeManipulation::asArray( PERMUTATION {} );
  INDEX_TYPE foldedStrides[ NDIM ];

  for( int i = 0; i < NDIM; ++i )
  {
    foldedStrides[ i ] = 1;
    for( int j = i + 1; j < NDIM; ++j )
    {
      foldedStrides[ i ] *= dims[ perm[ j ] ];
    }
  }

  typeManipulation::CArray< INDEX_TYPE, NDIM > strides;
  for( int i = 0; i < NDIM; ++i )
  {
    strides[ perm[ i ] ] = foldedStrides[ i ];
  }

  return strides;
}

} // namespace indexing
} // namespace LvArray