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             : #include "KokkosHeader.h"
      13             : 
      14             : namespace Moose
      15             : {
      16             : namespace Kokkos
      17             : {
      18             : namespace Utils
      19             : {
      20             : 
      21             : /**
      22             :  * Returns the sign of a value
      23             :  * @param x The value
      24             :  * @returns The sign of the value
      25             :  */
      26             : template <typename T>
      27             : KOKKOS_INLINE_FUNCTION T
      28    21599928 : sign(T x)
      29             : {
      30    21599928 :   return x >= 0.0 ? 1.0 : -1.0;
      31             : }
      32             : 
      33             : /**
      34             :  * Find a value in an array
      35             :  * @param target The target value to find
      36             :  * @param begin The pointer to the first element of the array
      37             :  * @param end The pointer next to the last element of the array
      38             :  * @returns The pointer to the target element, \p end if the target element was not found
      39             :  */
      40             : template <typename T>
      41             : KOKKOS_INLINE_FUNCTION const T *
      42     9226076 : find(const T & target, const T * const begin, const T * const end)
      43             : {
      44     9226076 :   if (begin == end)
      45           0 :     return end;
      46             : 
      47     9226076 :   auto left = begin;
      48     9226076 :   auto right = end - 1;
      49             : 
      50    24443213 :   while (left <= right)
      51             :   {
      52    23893123 :     auto mid = left + (right - left) / 2;
      53             : 
      54    23893123 :     if (*mid == target)
      55     8675986 :       return mid;
      56    15217137 :     else if (*mid < target)
      57     9079170 :       left = mid + 1;
      58             :     else
      59     6137967 :       right = mid - 1;
      60             :   }
      61             : 
      62      550090 :   return end;
      63             : }
      64             : 
      65             : /**
      66             :  * Perform an in-place linear solve using Cholesky decomposition
      67             :  * Matrix and right-hand-side vector are modified after this call
      68             :  * @param A The row-major matrix
      69             :  * @param x The solution vector
      70             :  * @param b The right-hand-side vector
      71             :  * @param n The system size
      72             :  */
      73             : KOKKOS_INLINE_FUNCTION void
      74      137000 : choleskySolve(Real * const A, Real * const x, Real * const b, const unsigned int n)
      75             : {
      76      959000 :   for (unsigned int i = 0; i < n; ++i)
      77             :   {
      78     3699000 :     for (unsigned int j = 0; j <= i; ++j)
      79             :     {
      80     2877000 :       Real sum = A[j + n * i];
      81             : 
      82     7672000 :       for (unsigned int k = 0; k < j; ++k)
      83     4795000 :         sum -= A[k + n * i] * A[k + n * j];
      84             : 
      85     2877000 :       if (i == j)
      86      822000 :         A[j + n * i] = ::Kokkos::sqrt(sum);
      87             :       else
      88     2055000 :         A[j + n * i] = sum / A[j + n * j];
      89             :     }
      90             :   }
      91             : 
      92      959000 :   for (unsigned int i = 0; i < n; ++i)
      93             :   {
      94      822000 :     Real sum = b[i];
      95             : 
      96     2877000 :     for (unsigned int j = 0; j < i; ++j)
      97     2055000 :       sum -= A[j + n * i] * b[j];
      98             : 
      99      822000 :     b[i] = sum / A[i + n * i];
     100             :   }
     101             : 
     102      959000 :   for (int i = n - 1; i >= 0; --i)
     103             :   {
     104      822000 :     Real sum = b[i];
     105             : 
     106     2877000 :     for (unsigned int j = i + 1; j < n; ++j)
     107     2055000 :       sum -= A[i + n * j] * b[j];
     108             : 
     109      822000 :     x[i] = sum / A[i + n * i];
     110             :   }
     111      137000 : }
     112             : 
     113             : } // namespace Utils
     114             : } // namespace Kokkos
     115             : } // namespace Moose