MueLu_IndexManager_kokkos_decl.hpp

Go to the documentation of this file.

// @HEADER
// *****************************************************************************
//        MueLu: A package for multigrid based preconditioning
//
// Copyright 2012 NTESS and the MueLu contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER
 
#ifndef MUELU_INDEXMANAGER_KOKKOS_DECL_HPP
#define MUELU_INDEXMANAGER_KOKKOS_DECL_HPP
 
#include "MueLu_ConfigDefs.hpp"
#include "MueLu_Types.hpp"
 
#include <Tpetra_KokkosCompat_ClassicNodeAPI_Wrapper.hpp>
 
#include "Teuchos_OrdinalTraits.hpp"
 
#include "MueLu_BaseClass.hpp"
#include "MueLu_IndexManager_kokkos_fwd.hpp"
 
/*****************************************************************************
 
****************************************************************************/
 
namespace MueLu {
 
template <class LocalOrdinal, class GlobalOrdinal, class Node>
class IndexManager_kokkos : public BaseClass {
#undef MUELU_INDEXMANAGER_KOKKOS_SHORT
#include "MueLu_UseShortNamesOrdinal.hpp"
 
 public:
  using execution_space = typename Node::execution_space;
  using memory_space    = typename Node::memory_space;
  using device_type     = Kokkos::Device<execution_space, memory_space>;
  using intTupleView    = typename Kokkos::View<int[3], device_type>;
  using LOTupleView     = typename Kokkos::View<LO[3], device_type>;
 
 private:
  const int meshLayout = UNCOUPLED;
  int myRank           = -1;
  int numDimensions;        
  int interpolationOrder_;  
  intTupleView coarseRate;  
  intTupleView endRate;     
 
  LO lNumFineNodes;              
  LO lNumFineNodes10;            
  LOTupleView lFineNodesPerDir;  
 
  LO numCoarseNodes;              
  LO numCoarseNodes10;            
  LOTupleView coarseNodesPerDir;  
 
 public:
  IndexManager_kokkos() = default;
 
  IndexManager_kokkos(const int NumDimensions,
                      const int interpolationOrder,
                      const int MyRank,
                      const ArrayView<const LO> LFineNodesPerDir,
                      const ArrayView<const int> CoarseRate);
 
  virtual ~IndexManager_kokkos() {}
 
  void setupIM(const int NumDimensions,
               const int interpolationOrder,
               const ArrayView<const int> coarseRate,
               const ArrayView<const LO> LFineNodesPerDir);
 
  void computeMeshParameters();
 
  int getNumDimensions() const { return numDimensions; }
 
  int getInterpolationOrder() const { return interpolationOrder_; }
 
  LO getNumLocalFineNodes() const { return lNumFineNodes; }
 
  LO getNumCoarseNodes() const { return numCoarseNodes; }
 
  KOKKOS_INLINE_FUNCTION
  intTupleView getCoarseningRates() const { return coarseRate; }
 
  KOKKOS_INLINE_FUNCTION
  intTupleView getCoarseningEndRates() const { return endRate; }
 
  KOKKOS_INLINE_FUNCTION
  LOTupleView getLocalFineNodesPerDir() const { return lFineNodesPerDir; }
 
  KOKKOS_INLINE_FUNCTION
  LOTupleView getCoarseNodesPerDir() const { return coarseNodesPerDir; }
 
  Array<LO> getCoarseNodesPerDirArray() const;
 
  KOKKOS_INLINE_FUNCTION
  void getFineLID2FineTuple(const LO myLID, LO (&tuple)[3]) const {
    LO tmp;
    tuple[2] = myLID / (lFineNodesPerDir(1) * lFineNodesPerDir(0));
    tmp      = myLID % (lFineNodesPerDir(1) * lFineNodesPerDir(0));
    tuple[1] = tmp / lFineNodesPerDir(0);
    tuple[0] = tmp % lFineNodesPerDir(0);
  }  // getFineNodeLocalTuple
 
  KOKKOS_INLINE_FUNCTION
  void getFineTuple2FineLID(const LO tuple[3], LO& myLID) const {
    myLID = tuple[2] * lNumFineNodes10 + tuple[1] * lFineNodesPerDir[0] + tuple[0];
  }  // getFineNodeLID
 
  KOKKOS_INLINE_FUNCTION
  void getCoarseLID2CoarseTuple(const LO myLID, LO (&tuple)[3]) const {
    LO tmp;
    tuple[2] = myLID / numCoarseNodes10;
    tmp      = myLID % numCoarseNodes10;
    tuple[1] = tmp / coarseNodesPerDir[0];
    tuple[0] = tmp % coarseNodesPerDir[0];
  }  // getCoarseNodeLocalTuple
 
  KOKKOS_INLINE_FUNCTION
  void getCoarseTuple2CoarseLID(const LO i, const LO j, const LO k, LO& myLID) const {
    myLID = k * numCoarseNodes10 + j * coarseNodesPerDir[0] + i;
  }  // getCoarseNodeLID
};
 
}  // namespace MueLu
 
#define MUELU_INDEXMANAGER_KOKKOS_SHORT
#endif  // MUELU_INDEXMANAGER_KOKKOS_DECL_HPP