Line data    Source code

       1             : //* This file is part of the MOOSE framework
       2             : //* https://mooseframework.inl.gov
       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             : #pragma once
      11             : 
      12             : #include "KokkosTypes.h"
      13             : 
      14             : namespace Moose
      15             : {
      16             : namespace Kokkos
      17             : {
      18             : 
      19             : template <typename Object>
      20             : class FunctorWrapperHost;
      21             : 
      22             : /**
      23             :  * Base class for device functor wrapper
      24             :  */
      25             : class FunctorWrapperDeviceBase
      26             : {
      27             : public:
      28             :   /**
      29             :    * Constructor
      30             :    */
      31           0 :   KOKKOS_FUNCTION FunctorWrapperDeviceBase() {}
      32             :   /**
      33             :    * Virtual destructor
      34             :    */
      35           0 :   KOKKOS_FUNCTION virtual ~FunctorWrapperDeviceBase() {}
      36             : };
      37             : 
      38             : /**
      39             :  * Device functor wrapper class that provides polymorphic interfaces for a functor. The functor
      40             :  * itself is a static object and does not have any virtual method. Instead, the device wrapper
      41             :  * defines the virtual shims and forwards the calls to the static methods of the stored functor.
      42             :  * @tparam Object The functor class type
      43             :  */
      44             : template <typename Object>
      45             : class FunctorWrapperDevice : public FunctorWrapperDeviceBase
      46             : {
      47             :   friend class FunctorWrapperHost<Object>;
      48             : 
      49             : public:
      50             :   /**
      51             :    * Constructor
      52             :    */
      53           0 :   KOKKOS_FUNCTION FunctorWrapperDevice() {}
      54             : 
      55             : protected:
      56             :   /**
      57             :    * Pointer to the functor on device
      58             :    */
      59             :   Object * _functor = nullptr;
      60             : };
      61             : 
      62             : /**
      63             :  * Base class for host functor wrapper
      64             :  */
      65             : class FunctorWrapperHostBase
      66             : {
      67             : public:
      68             :   /**
      69             :    * Virtual destructor
      70             :    */
      71           0 :   virtual ~FunctorWrapperHostBase() {}
      72             : 
      73             :   /**
      74             :    * Allocate device functor and wrapper
      75             :    * @returns The pointer to the device wrapper
      76             :    */
      77             :   virtual FunctorWrapperDeviceBase * allocate() = 0;
      78             :   /**
      79             :    * Copy functor to device
      80             :    */
      81             :   virtual void copyFunctor() = 0;
      82             :   /**
      83             :    * Free host and device copies of functor
      84             :    */
      85             :   virtual void freeFunctor() = 0;
      86             : };
      87             : 
      88             : /**
      89             :  * Host functor wrapper class that allocates a functor on device and creates its device wrapper.
      90             :  * This class holds the actual device instance of the functor and manages its allocation and
      91             :  * deallocation, and the device wrapper simply keeps a pointer to it.
      92             :  * @tparam Object The functor class type
      93             :  */
      94             : template <typename Object>
      95             : class FunctorWrapperHost : public FunctorWrapperHostBase
      96             : {
      97             : public:
      98             :   /**
      99             :    * Constructor
     100             :    * @param functor Pointer to the functor
     101             :    */
     102           0 :   FunctorWrapperHost(const void * functor) : _functor_host(*static_cast<const Object *>(functor)) {}
     103             :   /**
     104             :    * Desturctor
     105             :    */
     106             :   ~FunctorWrapperHost();
     107             : 
     108             :   FunctorWrapperDeviceBase * allocate() override final;
     109             :   void copyFunctor() override final;
     110             :   void freeFunctor() override final;
     111             : 
     112             : private:
     113             :   /**
     114             :    * Reference of the functor on host
     115             :    */
     116             :   const Object & _functor_host;
     117             :   /**
     118             :    * Copy of the functor on host
     119             :    */
     120             :   std::unique_ptr<Object> _functor_copy;
     121             :   /**
     122             :    * Copy of the functor on device
     123             :    */
     124             :   Object * _functor_device = nullptr;
     125             : };
     126             : 
     127             : template <typename Object>
     128             : FunctorWrapperDeviceBase *
     129           0 : FunctorWrapperHost<Object>::allocate()
     130             : {
     131             :   // Allocate storage for device wrapper on device
     132           0 :   auto wrapper_device = static_cast<FunctorWrapperDevice<Object> *>(
     133           0 :       ::Kokkos::kokkos_malloc<ExecSpace::memory_space>(sizeof(FunctorWrapperDevice<Object>)));
     134             : 
     135             :   // Allocate device wrapper on device using placement new to populate vtable with device pointers
     136           0 :   ::Kokkos::parallel_for(
     137           0 :       1, KOKKOS_LAMBDA(const int) { new (wrapper_device) FunctorWrapperDevice<Object>(); });
     138             : 
     139             :   // Allocate storage for functor on device
     140           0 :   _functor_device =
     141           0 :       static_cast<Object *>(::Kokkos::kokkos_malloc<ExecSpace::memory_space>(sizeof(Object)));
     142             : 
     143             :   // Let device wrapper point to the copy
     144           0 :   ::Kokkos::Impl::DeepCopy<MemSpace, ::Kokkos::HostSpace>(
     145           0 :       &(wrapper_device->_functor), &_functor_device, sizeof(Object *));
     146             : 
     147           0 :   return wrapper_device;
     148             : }
     149             : 
     150             : template <typename Object>
     151             : void
     152           0 : FunctorWrapperHost<Object>::copyFunctor()
     153             : {
     154             :   // Make a copy of functor on host to trigger copy constructor
     155           0 :   _functor_copy = std::make_unique<Object>(_functor_host);
     156             : 
     157             :   // Copy functor to device
     158           0 :   ::Kokkos::Impl::DeepCopy<MemSpace, ::Kokkos::HostSpace>(
     159           0 :       _functor_device, _functor_copy.get(), sizeof(Object));
     160           0 : }
     161             : 
     162             : template <typename Object>
     163             : void
     164           0 : FunctorWrapperHost<Object>::freeFunctor()
     165             : {
     166           0 :   _functor_copy.reset();
     167           0 : }
     168             : 
     169             : template <typename Object>
     170           0 : FunctorWrapperHost<Object>::~FunctorWrapperHost()
     171             : {
     172           0 :   ::Kokkos::kokkos_free<ExecSpace::memory_space>(_functor_device);
     173           0 : }
     174             : 
     175             : } // namespace Kokkos
     176             : } // namespace Moose