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 "KokkosArray.h"
      13             : 
      14             : #ifdef MOOSE_KOKKOS_SCOPE
      15             : #include "KokkosUtils.h"
      16             : #endif
      17             : 
      18             : #include "libmesh/petsc_matrix.h"
      19             : 
      20             : namespace Moose
      21             : {
      22             : namespace Kokkos
      23             : {
      24             : 
      25             : class System;
      26             : 
      27             : /**
      28             :  * The Kokkos wrapper class for PETSc matrix
      29             :  */
      30             : class Matrix
      31             : {
      32             : public:
      33             :   /**
      34             :    * Default constructor
      35             :    */
      36       41580 :   Matrix() = default;
      37             :   /**
      38             :    * Destructor
      39             :    */
      40       40620 :   ~Matrix() { destroy(); }
      41             :   /**
      42             :    * Free all data and reset
      43             :    */
      44             :   void destroy();
      45             : 
      46             : #ifdef MOOSE_KOKKOS_SCOPE
      47             :   /**
      48             :    * Get whether the matrix was allocated
      49             :    * @returns Whether the matrix was allocated
      50             :    */
      51             :   bool isAlloc() const { return _is_alloc; }
      52             :   /**
      53             :    * Create the matrix from a libMesh PetscMatrix
      54             :    * @param matrix The libMesh PetscMatrix
      55             :    * @param system The Kokkos system
      56             :    */
      57             :   void create(libMesh::SparseMatrix<PetscScalar> & matrix, const System & system);
      58             :   /**
      59             :    * Assemble the underlying PETSc matrix
      60             :    */
      61             :   void close();
      62             : 
      63             :   /**
      64             :    * Zero a row
      65             :    * @param i The row index local to this process to be zeroed
      66             :    */
      67       40902 :   KOKKOS_FUNCTION void zero(PetscInt i)
      68             :   {
      69      350410 :     for (PetscInt j = _row_ptr[i]; j < _row_ptr[i + 1]; ++j)
      70      309508 :       _val[j] = 0;
      71       40902 :   }
      72             :   /**
      73             :    * Get an entry with given row and column indices
      74             :    * @param i The row index local to this process
      75             :    * @param j The global column index
      76             :    * @returns The reference of the element
      77             :    */
      78     6062628 :   KOKKOS_FUNCTION PetscScalar & operator()(PetscInt i, PetscInt j) const
      79             :   {
      80     6062628 :     auto idx = find(i, j);
      81             : 
      82             :     KOKKOS_ASSERT(idx != -1)
      83             : 
      84     6062628 :     return _val[idx];
      85             :   }
      86             :   /**
      87             :    * Assign a scalar value uniformly
      88             :    * @param scalar The scalar value to be assigned
      89             :    */
      90        4654 :   auto & operator=(PetscScalar scalar)
      91             :   {
      92        4654 :     _val = scalar;
      93             : 
      94        4654 :     return *this;
      95             :   }
      96             : #endif
      97             : 
      98             : private:
      99             : #ifdef MOOSE_KOKKOS_SCOPE
     100             :   /**
     101             :    * Get the index of given row and column indices
     102             :    * @param i The row index local to this process
     103             :    * @param j The global column index
     104             :    * @returns The index into the nonzero vector, and -1 if not found
     105             :    */
     106             :   KOKKOS_FUNCTION PetscInt find(PetscInt i, PetscInt j) const;
     107             : #endif
     108             : 
     109             :   /**
     110             :    * PETSc matrix
     111             :    */
     112             :   Mat _matrix = PETSC_NULLPTR;
     113             :   /**
     114             :    * Number of rows local to this process
     115             :    */
     116             :   PetscCount _nr = 0;
     117             :   /**
     118             :    * CSR vectors on device
     119             :    */
     120             :   ///@{
     121             :   Array<PetscInt> _col_idx;
     122             :   Array<PetscInt> _row_idx;
     123             :   Array<PetscInt> _row_ptr;
     124             :   Array<PetscScalar> _val;
     125             :   ///@}
     126             :   /**
     127             :    * Flag whether the PETSc matrix is a host matrix
     128             :    */
     129             :   bool _is_host = false;
     130             :   /**
     131             :    * Flag whether the matrix was allocated
     132             :    */
     133             :   bool _is_alloc = false;
     134             : };
     135             : 
     136             : #ifdef MOOSE_KOKKOS_SCOPE
     137             : KOKKOS_FUNCTION inline PetscInt
     138     6062628 : Matrix::find(PetscInt i, PetscInt j) const
     139             : {
     140             :   KOKKOS_ASSERT(i < _nr);
     141             : 
     142     6062628 :   auto begin = _col_idx.data() + _row_ptr[i];
     143     6062628 :   auto end = _col_idx.data() + _row_ptr[i + 1];
     144     6062628 :   auto target = Utils::find(j, begin, end);
     145             : 
     146     6062628 :   if (target == end)
     147           0 :     return -1;
     148             :   else
     149     6062628 :     return target - _col_idx.data();
     150             : }
     151             : #endif
     152             : 
     153             : } // namespace Kokkos
     154             : } // namespace Moose