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             : #include "KokkosBlas.hpp"
      11             : 
      12             : #include "KokkosArray.h"
      13             : #include "KokkosThread.h"
      14             : 
      15             : namespace Moose::Kokkos
      16             : {
      17             : 
      18             : void
      19     2193992 : free(void * ptr)
      20             : {
      21     2193992 :   ::Kokkos::kokkos_free<ExecSpace::memory_space>(ptr);
      22     2193992 : }
      23             : 
      24             : template <typename T, typename index_type>
      25             : void
      26          16 : Array<T, 1, index_type, LayoutType::LEFT>::axby(const T a,
      27             :                                                 const Array<T, 1, index_type, LayoutType::LEFT> & x,
      28             :                                                 const char op,
      29             :                                                 const T b,
      30             :                                                 const Array<T, 1, index_type, LayoutType::LEFT> & y,
      31             :                                                 const bool accumulate)
      32             : {
      33          16 :   if (!this->isDeviceAlloc() || !x.isDeviceAlloc() || !y.isDeviceAlloc())
      34           0 :     mooseError("Kokkos array error: all arrays for axby() should be on device.");
      35             : 
      36          16 :   if (!(this->size() == x.size() && (x.size() == y.size())))
      37           0 :     mooseError("Kokkos array error: all arrays for axby() should have the same size.");
      38             : 
      39          16 :   ::Kokkos::RangePolicy<ExecSpace, ::Kokkos::IndexType<index_type>> policy(0, this->size());
      40             : 
      41          16 :   auto data = this->deviceData();
      42             : 
      43          16 :   switch (op)
      44             :   {
      45           4 :     case '+':
      46           4 :       ::Kokkos::parallel_for(
      47          14 :           policy, KOKKOS_LAMBDA(const index_type i) {
      48           6 :             if (accumulate)
      49           3 :               data[i] += (a * x[i]) + (b * y[i]);
      50             :             else
      51           3 :               data[i] = (a * x[i]) + (b * y[i]);
      52             :           });
      53           4 :       break;
      54           4 :     case '-':
      55           4 :       ::Kokkos::parallel_for(
      56          14 :           policy, KOKKOS_LAMBDA(const index_type i) {
      57           6 :             if (accumulate)
      58           3 :               data[i] += (a * x[i]) - (b * y[i]);
      59             :             else
      60           3 :               data[i] = (a * x[i]) - (b * y[i]);
      61             :           });
      62           4 :       break;
      63           4 :     case '*':
      64           4 :       ::Kokkos::parallel_for(
      65          14 :           policy, KOKKOS_LAMBDA(const index_type i) {
      66           6 :             if (accumulate)
      67           3 :               data[i] += (a * x[i]) * (b * y[i]);
      68             :             else
      69           3 :               data[i] = (a * x[i]) * (b * y[i]);
      70             :           });
      71           4 :       break;
      72           4 :     case '/':
      73           4 :       ::Kokkos::parallel_for(
      74          14 :           policy, KOKKOS_LAMBDA(const index_type i) {
      75           6 :             if (accumulate)
      76           3 :               data[i] += (a * x[i]) / (b * y[i]);
      77             :             else
      78           3 :               data[i] = (a * x[i]) / (b * y[i]);
      79             :           });
      80           4 :       break;
      81           0 :     default:
      82           0 :       mooseError("Kokkos array error: unknown operation '", op, "' for axby().");
      83             :   }
      84             : 
      85          16 :   ::Kokkos::fence();
      86          16 : }
      87             : 
      88             : template <typename T, typename index_type>
      89             : void
      90           2 : Array<T, 1, index_type, LayoutType::LEFT>::scal(const T a,
      91             :                                                 const Array<T, 1, index_type, LayoutType::LEFT> & x)
      92             : {
      93           2 :   if (!this->isDeviceAlloc() || !x.isDeviceAlloc())
      94           0 :     mooseError("Kokkos array error: all arrays for scal() should be on device.");
      95             : 
      96           2 :   if (this->size() != x.size())
      97           0 :     mooseError("Kokkos array error: all arrays for scal() should have the same size.");
      98             : 
      99           2 :   auto self = this->deviceView();
     100           2 :   auto other = x.deviceView();
     101             : 
     102           2 :   KokkosBlas::scal(self, a, other);
     103           2 : }
     104             : 
     105             : template <typename T, typename index_type>
     106             : void
     107           2 : Array<T, 1, index_type, LayoutType::LEFT>::scal(const T a)
     108             : {
     109           2 :   return scal(a, *this);
     110             : }
     111             : 
     112             : template <typename T, typename index_type>
     113             : T
     114           2 : Array<T, 1, index_type, LayoutType::LEFT>::dot(const Array<T, 1, index_type, LayoutType::LEFT> & x)
     115             : {
     116           2 :   if (!this->isDeviceAlloc() || !x.isDeviceAlloc())
     117           0 :     mooseError("Kokkos array error: all arrays for dot() should be on device.");
     118             : 
     119           2 :   if (this->size() != x.size())
     120           0 :     mooseError("Kokkos array error: all arrays for dot() should have the same size.");
     121             : 
     122           2 :   auto self = this->deviceView();
     123           2 :   auto other = x.deviceView();
     124             : 
     125           4 :   return KokkosBlas::dot(self, other);
     126           2 : }
     127             : 
     128             : template <typename T, typename index_type>
     129             : T
     130           2 : Array<T, 1, index_type, LayoutType::LEFT>::nrm2()
     131             : {
     132           2 :   if (!this->isDeviceAlloc())
     133           0 :     mooseError("Kokkos array error: array for nrm2() should be on device.");
     134             : 
     135           2 :   auto self = this->deviceView();
     136             : 
     137           4 :   return KokkosBlas::nrm2(self);
     138           2 : }
     139             : 
     140             : template class Array<Real, 1, uint8_t, LayoutType::LEFT>;
     141             : template class Array<Real, 1, uint16_t, LayoutType::LEFT>;
     142             : template class Array<Real, 1, uint32_t, LayoutType::LEFT>;
     143             : template class Array<Real, 1, uint64_t, LayoutType::LEFT>;
     144             : 
     145             : } // namespace Moose::Kokkos