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 "KokkosMatrix.h"
      11             : #include "KokkosSystem.h"
      12             : 
      13             : namespace Moose::Kokkos
      14             : {
      15             : 
      16             : void
      17       16153 : Matrix::create(libMesh::SparseMatrix<PetscScalar> & matrix, const System & system)
      18             : {
      19       16153 :   auto petsc_matrix = dynamic_cast<libMesh::PetscMatrix<PetscScalar> *>(&matrix);
      20             : 
      21             :   mooseAssert(petsc_matrix, "Kokkos matrix error: provided matrix is not a PetscMatrix.");
      22             : 
      23             :   PetscMemType mtype;
      24       16153 :   LibmeshPetscCallQ(MatGetCurrentMemType(petsc_matrix->mat(), &mtype));
      25       16153 :   const bool is_host = PetscMemTypeHost(mtype);
      26             : 
      27             : #ifndef MOOSE_ENABLE_KOKKOS_GPU
      28        8962 :   if (!is_host)
      29           0 :     mooseError("PETSc matrices must be on host when Kokkos device capabilities are disabled.");
      30             : #endif
      31             : 
      32       16153 :   if (_is_alloc)
      33       14417 :     return;
      34             : 
      35        1736 :   auto & sparsity = system.getSparsity();
      36             : 
      37        1736 :   _matrix = petsc_matrix->mat();
      38        1736 :   _nr = sparsity.row_ptr.size() - 1;
      39        1736 :   _col_idx = sparsity.col_idx;
      40        1736 :   _row_idx = sparsity.row_idx;
      41        1736 :   _row_ptr = sparsity.row_ptr;
      42        1736 :   _is_host = is_host;
      43             : 
      44        1736 :   if (!_is_host)
      45          14 :     _val.createDevice(_col_idx.size());
      46             :   else
      47        1722 :     _val.create(_col_idx.size());
      48             : 
      49        1736 :   std::vector<PetscInt> col_idx(&_col_idx.begin(), &_col_idx.end());
      50        1736 :   std::vector<PetscInt> row_idx(&_row_idx.begin(), &_row_idx.end());
      51             : 
      52        1736 :   LibmeshPetscCallQ(
      53             :       MatSetPreallocationCOO(_matrix, col_idx.size(), row_idx.data(), col_idx.data()));
      54             : 
      55        1736 :   _is_alloc = true;
      56        1736 : }
      57             : 
      58             : void
      59      134940 : Matrix::destroy()
      60             : {
      61      134940 :   _matrix = PETSC_NULLPTR;
      62      134940 :   _nr = 0;
      63             : 
      64      134940 :   _col_idx.destroy();
      65      134940 :   _row_idx.destroy();
      66      134940 :   _row_ptr.destroy();
      67      134940 :   _val.destroy();
      68             : 
      69      134940 :   _is_host = false;
      70      134940 :   _is_alloc = false;
      71      134940 : }
      72             : 
      73             : void
      74       16153 : Matrix::close()
      75             : {
      76       16153 :   if (_is_host)
      77       16139 :     _val.copyToHost();
      78             : 
      79       16153 :   LibmeshPetscCallQ(
      80             :       MatSetValuesCOO(_matrix, _is_host ? _val.hostData() : _val.deviceData(), ADD_VALUES));
      81             : 
      82       16153 :   ::Kokkos::fence();
      83       16153 : }
      84             : 
      85             : } // namespace Moose::Kokkos