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