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