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