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