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 "KokkosVector.h"
      11             : #include "KokkosSystem.h"
      12             : 
      13             : // #define DEVICE_ASSEMBLY
      14             : 
      15             : namespace Moose
      16             : {
      17             : namespace Kokkos
      18             : {
      19             : 
      20             : #ifdef DEVICE_ASSEMBLY
      21             : 
      22             : void
      23             : Vector::DeviceAssembly::create(const Array<Array<libMesh::dof_id_type>> & list)
      24             : {
      25             :   this->list = list;
      26             : 
      27             :   count.createHost(list.size());
      28             :   offset.create(list.size() + 1);
      29             :   offset = 0;
      30             : 
      31             :   for (unsigned int i = 0; i < list.size(); ++i)
      32             :   {
      33             :     count[i] = list[i].size();
      34             :     offset[i + 1] = offset[i] + count[i];
      35             :   }
      36             : 
      37             :   offset.copyToDevice();
      38             : 
      39             :   buffer.createDevice(offset.last());
      40             : }
      41             : 
      42             : void
      43             : Vector::DeviceAssembly::destroy()
      44             : {
      45             :   list.destroy();
      46             :   count.destroy();
      47             :   offset.destroy();
      48             :   buffer.destroy();
      49             : }
      50             : 
      51             : KOKKOS_FUNCTION void
      52             : Vector::operator()(PackBuffer, const PetscCount tid) const
      53             : {
      54             :   _send.buffer[_send.offset[_current_proc] + tid] = _ghost(_send.list[_current_proc][tid]);
      55             : }
      56             : 
      57             : KOKKOS_FUNCTION void
      58             : Vector::operator()(UnpackBuffer, const PetscCount tid) const
      59             : {
      60             :   _local[_recv.list[_current_proc][tid]] += _recv.buffer[_recv.offset[_current_proc] + tid];
      61             : }
      62             : 
      63             : #endif
      64             : 
      65             : void
      66      173933 : Vector::create(libMesh::NumericVector<PetscScalar> & vector, const System & system, bool assemble)
      67             : {
      68      173933 :   auto petsc_vector = dynamic_cast<libMesh::PetscVector<PetscScalar> *>(&vector);
      69             : 
      70             :   mooseAssert(petsc_vector, "Kokkos vector error: provided vector is not a PetscVector.");
      71             : 
      72             :   PetscScalar * array;
      73             :   PetscMemType mtype;
      74      173933 :   LibmeshPetscCallQ(VecGhostGetLocalForm(petsc_vector->vec(), &_local_vector));
      75      173933 :   LibmeshPetscCallQ(
      76             :       VecGetArrayAndMemType(_local_vector ? _local_vector : petsc_vector->vec(), &array, &mtype));
      77             : 
      78      173933 :   auto is_host = mtype == PETSC_MEMTYPE_HOST;
      79      173933 :   auto is_ghosted = vector.type() == libMesh::ParallelType::GHOSTED;
      80      173933 :   bool realloc =
      81      141328 :       !_is_alloc || _assemble != assemble || _is_host != is_host || _is_ghosted != is_ghosted;
      82             : 
      83             : #ifndef MOOSE_ENABLE_KOKKOS_GPU
      84      141328 :   if (!is_host)
      85           0 :     mooseError("PETSc vectors must be on host when Kokkos device capabilities are disabled.");
      86             : #endif
      87             : 
      88      173933 :   _global_vector = petsc_vector->vec();
      89      173933 :   _array = array;
      90      173933 :   _system = &system;
      91      173933 :   _assemble = assemble;
      92      173933 :   _is_host = is_host;
      93      173933 :   _is_ghosted = is_ghosted;
      94             : 
      95      173933 :   if (realloc)
      96             :   {
      97        3604 :     _comm = nullptr;
      98             : 
      99             : #ifdef DEVICE_ASSEMBLY
     100             :     _send.destroy();
     101             :     _recv.destroy();
     102             : #endif
     103             : 
     104        3604 :     _local.destroy();
     105        3604 :     _ghost.destroy();
     106             :   }
     107             : 
     108      173933 :   if (assemble)
     109             :   {
     110       81553 :     if (realloc)
     111             :     {
     112             : #ifdef DEVICE_ASSEMBLY
     113             :       _comm = &_system->getComm();
     114             :       _send.create(_system->getGhostCommList());
     115             :       _recv.create(_system->getLocalCommList());
     116             :       _ghost.createDevice(_system->getNumGhostDofs());
     117             : #else
     118        1662 :       _ghost.create(_system->getNumGhostDofs());
     119             : #endif
     120        1662 :       _ghost.offset(_system->getNumLocalDofs());
     121             :     }
     122             : 
     123       81553 :     if (!_is_host)
     124             :     {
     125          28 :       if (realloc)
     126           8 :         _local.init(_system->getNumLocalDofs());
     127             : 
     128          28 :       _local.aliasDevice(_array);
     129             :     }
     130             :     else
     131             :     {
     132       81525 :       if (realloc)
     133        1654 :         _local.createDevice(_system->getNumLocalDofs());
     134             : 
     135       81525 :       _local.aliasHost(_array);
     136             :     }
     137             :   }
     138             :   else
     139             :   {
     140       92380 :     if (!_is_host)
     141             :     {
     142           4 :       if (realloc)
     143           4 :         _local.init(_is_ghosted ? _system->getNumLocalDofs() + _system->getNumGhostDofs()
     144           0 :                                 : _system->getNumLocalDofs());
     145             : 
     146           4 :       _local.aliasDevice(_array);
     147             :     }
     148             :     else
     149             :     {
     150       92376 :       if (realloc)
     151        2340 :         _local.createDevice(_is_ghosted ? _system->getNumLocalDofs() + _system->getNumGhostDofs()
     152         402 :                                         : _system->getNumLocalDofs());
     153             : 
     154       92376 :       _local.aliasHost(_array);
     155             :     }
     156             :   }
     157             : 
     158      173933 :   _is_alloc = true;
     159      173933 : }
     160             : 
     161             : void
     162       40620 : Vector::destroy()
     163             : {
     164       40620 :   if (!_is_alloc)
     165       37034 :     return;
     166             : 
     167        3586 :   _global_vector = PETSC_NULLPTR;
     168        3586 :   _local_vector = PETSC_NULLPTR;
     169        3586 :   _array = PETSC_NULLPTR;
     170        3586 :   _system = nullptr;
     171        3586 :   _comm = nullptr;
     172             : 
     173             : #ifdef DEVICE_ASSEMBLY
     174             :   _send.destroy();
     175             :   _recv.destroy();
     176             : #endif
     177             : 
     178        3586 :   _local.destroy();
     179        3586 :   _ghost.destroy();
     180             : 
     181        3586 :   _assemble = false;
     182        3586 :   _is_ghosted = false;
     183        3586 :   _is_host = false;
     184        3586 :   _is_alloc = false;
     185             : }
     186             : 
     187             : void
     188       92380 : Vector::copyToDevice()
     189             : {
     190       92380 :   if (_assemble)
     191           0 :     mooseError("Kokkos vector error: copyToDevice() should not be called for an assembled vector.");
     192             : 
     193       92380 :   if (_is_alloc && _is_host)
     194       92376 :     _local.copyToDevice();
     195       92380 : }
     196             : 
     197             : void
     198        1909 : Vector::copyToHost()
     199             : {
     200        1909 :   if (_assemble)
     201           0 :     mooseError("Kokkos vector error: copyToHost() should not be called for an assembled vector.");
     202             : 
     203        1909 :   if (_is_alloc && _is_host)
     204        1905 :     _local.copyToHost();
     205        1909 : }
     206             : 
     207             : void
     208      173933 : Vector::restore()
     209             : {
     210      173933 :   LibmeshPetscCallQ(
     211             :       VecRestoreArrayAndMemType(_local_vector ? _local_vector : _global_vector, &_array));
     212      173933 :   LibmeshPetscCallQ(VecGhostRestoreLocalForm(_global_vector, &_local_vector));
     213      173933 : }
     214             : 
     215             : void
     216       81553 : Vector::close()
     217             : {
     218       81553 :   if (!_assemble)
     219           0 :     mooseError("Kokkos vector error: close() should not be called for a non-assembled vector.");
     220             : 
     221       81553 :   if (!_is_alloc)
     222           0 :     return;
     223             : 
     224             : #ifdef DEVICE_ASSEMBLY
     225             : 
     226             :   // Pack data into send buffer
     227             : 
     228             :   for (_current_proc = 0; _current_proc < _comm->size(); ++_current_proc)
     229             :   {
     230             :     Kokkos::RangePolicy<PackBuffer, Kokkos::IndexType<PetscCount>> policy(
     231             :         0, _send.count[_current_proc]);
     232             :     Kokkos::parallel_for(policy, *this);
     233             :   }
     234             : 
     235             :   Kokkos::fence();
     236             : 
     237             :   // Perform MPI communications
     238             : 
     239             :   MPI_Alltoallv(_send.buffer.deviceData(),
     240             :                 _send.count.data(),
     241             :                 _send.offset.data(),
     242             :                 MPIU_SCALAR,
     243             :                 _recv.buffer.deviceData(),
     244             :                 _recv.count.data(),
     245             :                 _recv.offset.data(),
     246             :                 MPIU_SCALAR,
     247             :                 _comm->get());
     248             : 
     249             :   // Unpack data from receive buffer
     250             : 
     251             :   for (_current_proc = 0; _current_proc < _comm->size(); ++_current_proc)
     252             :   {
     253             :     Kokkos::RangePolicy<UnpackBuffer, Kokkos::IndexType<PetscCount>> policy(
     254             :         0, _recv.count[_current_proc]);
     255             :     Kokkos::parallel_for(policy, *this);
     256             :   }
     257             : 
     258             :   Kokkos::fence();
     259             : 
     260             : #endif
     261             : 
     262             :   // Copy to the host vector if needed
     263             : 
     264       81553 :   if (_is_host)
     265       81525 :     _local.copyToHost();
     266             : 
     267             :   // Restore the PETSc vector
     268             : 
     269       81553 :   restore();
     270             : 
     271             : #ifndef DEVICE_ASSEMBLY
     272             : 
     273       81553 :   _ghost.copyToHost();
     274             : 
     275       81553 :   LibmeshPetscCallQ(
     276             :       VecSetValues(_global_vector,
     277             :                    _system->getNumGhostDofs(),
     278             :                    reinterpret_cast<const PetscInt *>(_system->getDofMap().get_send_list().data()),
     279             :                    _ghost.data(),
     280             :                    ADD_VALUES));
     281             : 
     282             : #endif
     283             : }
     284             : 
     285             : } // namespace Kokkos
     286             : } // namespace Moose