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 "MooseUtils.h"
      17             : #include "Conversion.h"
      18             : 
      19             : namespace Moose::Kokkos
      20             : {
      21             : 
      22             : using ThreadID = dof_id_type;
      23             : 
      24             : /**
      25             :  * The Kokkos thread object that aids in converting the one-dimensional thread index into
      26             :  * multi-dimensional thread indices
      27             :  */
      28             : template <typename thread_id_type = ThreadID, unsigned int max_dimension = 4>
      29             : class Thread
      30             : {
      31             :   static_assert(std::is_integral_v<thread_id_type>,
      32             :                 "Kokkos thread index type must be an integral type");
      33             :   static_assert(std::is_unsigned_v<thread_id_type>, "Kokkos thread index type must be unsigned");
      34             :   static_assert(!std::is_same_v<bool, thread_id_type>, "Kokkos thread index type must not be bool");
      35             : 
      36             : public:
      37             :   using id_type = thread_id_type;
      38             : 
      39             : #ifdef MOOSE_KOKKOS_SCOPE
      40             :   /**
      41             :    * Set the thread pool size and dimension
      42             :    * @param sizes The size of each dimension
      43             :    */
      44             :   template <typename... size_type>
      45             :   void resize(size_type... sizes);
      46             :   /**
      47             :    * Get the total thread pool size
      48             :    * @returns The total thread pool size
      49             :    */
      50      372066 :   KOKKOS_FUNCTION thread_id_type size() const { return _size; }
      51             :   /**
      52             :    * Get the size of each dimension
      53             :    * @param dim The dimension index
      54             :    * @returns The dimension size
      55             :    */
      56    20773492 :   KOKKOS_FUNCTION thread_id_type size(const unsigned int dim) const
      57             :   {
      58             :     KOKKOS_ASSERT(dim < _dim);
      59             : 
      60    20773492 :     return _dims[dim];
      61             :   }
      62             :   /**
      63             :    * Get the multi-dimensional thread index of a dimension given a one-dimensional thread index
      64             :    * @param tid The one-dimensional thread index
      65             :    * @param dim for which the multi-dimensional thread index is to be returned
      66             :    * @returns The multi-dimensional thread index of the dimension
      67             :    */
      68   313562868 :   KOKKOS_FUNCTION thread_id_type operator()(const thread_id_type tid, const unsigned int dim) const
      69             :   {
      70             :     KOKKOS_ASSERT(dim < _dim);
      71             : 
      72   313562868 :     return (tid / _strides[dim]) % _dims[dim];
      73             :   }
      74             : #endif
      75             : 
      76             : protected:
      77             :   /**
      78             :    * Total thread pool size
      79             :    */
      80             :   thread_id_type _size = 0;
      81             :   /**
      82             :    * Thread pool dimension
      83             :    */
      84             :   unsigned int _dim = 0;
      85             :   /**
      86             :    * Thread pool size of each dimension
      87             :    */
      88             :   thread_id_type _dims[max_dimension];
      89             :   /**
      90             :    * Thread pool stride of each dimension
      91             :    */
      92             :   thread_id_type _strides[max_dimension];
      93             : };
      94             : 
      95             : #ifdef MOOSE_KOKKOS_SCOPE
      96             : template <typename thread_id_type, unsigned int max_dimension>
      97             : template <typename... size_type>
      98             : void
      99      372066 : Thread<thread_id_type, max_dimension>::resize(size_type... sizes)
     100             : {
     101             :   static_assert((std::is_convertible<size_type, thread_id_type>::value && ...),
     102             :                 "All arguments must be convertible to thread_id_type");
     103             :   static_assert(sizeof...(sizes) <= max_dimension, "Number of arguments exceeds maximum dimension");
     104             : 
     105      372066 :   std::vector<thread_id_type> dims;
     106      372066 :   (dims.push_back(sizes), ...);
     107             : 
     108      372066 :   uint64_t size = 1;
     109      372066 :   _dim = dims.size();
     110             : 
     111     1496591 :   for (unsigned int dim = 0; dim < _dim; ++dim)
     112             :   {
     113     1124525 :     _dims[dim] = dims[dim];
     114     1124525 :     _strides[dim] = dim ? _strides[dim - 1] * dims[dim - 1] : 1;
     115     1124525 :     size *= dims[dim];
     116             :   }
     117             : 
     118      372066 :   if (size > std::numeric_limits<thread_id_type>::max())
     119           0 :     mooseError("Kokkos thread error: the dimensions provided (",
     120             :                Moose::stringify(dims),
     121             :                ") has the total size of ",
     122             :                size,
     123             :                " which exceeds the limit of ",
     124             :                MooseUtils::prettyCppType<thread_id_type>(),
     125             :                ".");
     126             : 
     127      372066 :   _size = size;
     128      372066 : }
     129             : #endif
     130             : 
     131             : } // namespace Moose::Kokkos