Line data    Source code

       1             : //* This file is part of the MOOSE framework
       2             : //* https://www.mooseframework.org
       3             : //*
       4             : //* All rights reserved, see COPYRIGHT for full restrictions
       5             : //* https://github.com/idaholab/moose/blob/master/COPYRIGHT
       6             : //*
       7             : //* Licensed under LGPL 2.1, please see LICENSE for details
       8             : //* https://www.gnu.org/licenses/lgpl-2.1.html
       9             : 
      10             : #pragma once
      11             : 
      12             : #ifdef MOOSE_KOKKOS_SCOPE
      13             : #include "KokkosHeader.h"
      14             : #endif
      15             : 
      16             : #include "MooseTypes.h"
      17             : 
      18             : using ThreadID = dof_id_type;
      19             : 
      20             : namespace Moose
      21             : {
      22             : namespace Kokkos
      23             : {
      24             : 
      25             : /**
      26             :  * The Kokkos thread object that aids in converting the one-dimensional thread index into
      27             :  * multi-dimensional thread indices
      28             :  */
      29             : class Thread
      30             : {
      31             : #ifdef MOOSE_KOKKOS_SCOPE
      32             : public:
      33             :   /**
      34             :    * Set the thread pool size and dimension
      35             :    * @param dims The vector containing the size of each dimension
      36             :    */
      37             :   void resize(std::vector<ThreadID> dims);
      38             :   /**
      39             :    * Get the total thread pool size
      40             :    * @returns The total thread pool size
      41             :    */
      42       67130 :   ThreadID size() const { return _size; }
      43             :   /**
      44             :    * Get the multi-dimensional thread index of a dimension given a one-dimensional thread index
      45             :    * @param tid The one-dimensional thread index
      46             :    * @param dim for which the multi-dimensional thread index is to be returned
      47             :    * @returns The multi-dimensional thread index of the dimension
      48             :    */
      49    73424320 :   KOKKOS_FUNCTION ThreadID operator()(ThreadID tid, unsigned int dim) const
      50             :   {
      51             :     KOKKOS_ASSERT(dim < _dim);
      52             : 
      53    73424320 :     return (tid / _strides[dim]) % _dims[dim];
      54             :   }
      55             : #endif
      56             : 
      57             : protected:
      58             :   /**
      59             :    * Total thread pool size
      60             :    */
      61             :   ThreadID _size = 0;
      62             :   /**
      63             :    * Thread pool dimension
      64             :    */
      65             :   unsigned int _dim = 0;
      66             :   /**
      67             :    * Thread pool size of each dimension
      68             :    */
      69             :   ThreadID _dims[10];
      70             :   /**
      71             :    * Thread pool stride of each dimension
      72             :    */
      73             :   ThreadID _strides[10];
      74             : };
      75             : 
      76             : } // namespace Kokkos
      77             : } // namespace Moose