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