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