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