KokkosFunctionWrapper.h

Go to the documentation of this file.

//* This file is part of the MOOSE framework
//* https://mooseframework.inl.gov
//*
//* All rights reserved, see COPYRIGHT for full restrictions
//* https://github.com/idaholab/moose/blob/master/COPYRIGHT
//*
//* Licensed under LGPL 2.1, please see LICENSE for details
//* https://www.gnu.org/licenses/lgpl-2.1.html
 
#pragma once
 
#include "KokkosTypes.h"
 
namespace Moose::Kokkos
{
 
template <typename Object>
class FunctionWrapperHost;
 
class FunctionWrapperDeviceBase
{
public:
  KOKKOS_FUNCTION FunctionWrapperDeviceBase() {}
  KOKKOS_FUNCTION KOKKOS_VIRTUAL ~FunctionWrapperDeviceBase() {}
 
  KOKKOS_FUNCTION KOKKOS_VIRTUAL Real value(Real /* t */, Real3 /* p */) const
  {
    KOKKOS_ASSERT(false);
    return 0;
  }
  KOKKOS_FUNCTION KOKKOS_VIRTUAL Real3 vectorValue(Real /* t */, Real3 /* p */) const
  {
    KOKKOS_ASSERT(false);
    return Real3(0);
  }
  KOKKOS_FUNCTION KOKKOS_VIRTUAL Real3 gradient(Real /* t */, Real3 /* p */) const
  {
    KOKKOS_ASSERT(false);
    return Real3(0);
  }
  KOKKOS_FUNCTION KOKKOS_VIRTUAL Real3 curl(Real /* t */, Real3 /* p */) const
  {
    KOKKOS_ASSERT(false);
    return Real3(0);
  }
  KOKKOS_FUNCTION KOKKOS_VIRTUAL Real div(Real /* t */, Real3 /* p */) const
  {
    KOKKOS_ASSERT(false);
    return 0;
  }
  KOKKOS_FUNCTION KOKKOS_VIRTUAL Real timeDerivative(Real /* t */, Real3 /* p */) const
  {
    KOKKOS_ASSERT(false);
    return 0;
  }
  KOKKOS_FUNCTION KOKKOS_VIRTUAL Real timeIntegral(Real /* t1 */,
                                                   Real /* t2 */,
                                                   Real3 /* p */) const
  {
    KOKKOS_ASSERT(false);
    return 0;
  }
  KOKKOS_FUNCTION KOKKOS_VIRTUAL Real integral() const
  {
    KOKKOS_ASSERT(false);
    return 0;
  }
  KOKKOS_FUNCTION KOKKOS_VIRTUAL Real average() const
  {
    KOKKOS_ASSERT(false);
    return 0;
  }
};
 
template <typename Object>
class FunctionWrapperDevice : public FunctionWrapperDeviceBase
{
  friend class FunctionWrapperHost<Object>;
 
public:
  KOKKOS_FUNCTION FunctionWrapperDevice() {}
 
  KOKKOS_FUNCTION Real value(Real t, Real3 p) const KOKKOS_OVERRIDE
  {
    return _function->value(t, p);
  }
  KOKKOS_FUNCTION Real3 vectorValue(Real t, Real3 p) const KOKKOS_OVERRIDE
  {
    return _function->vectorValue(t, p);
  }
  KOKKOS_FUNCTION Real3 gradient(Real t, Real3 p) const KOKKOS_OVERRIDE
  {
    return _function->gradient(t, p);
  }
  KOKKOS_FUNCTION Real3 curl(Real t, Real3 p) const KOKKOS_OVERRIDE
  {
    return _function->curl(t, p);
  }
  KOKKOS_FUNCTION Real div(Real t, Real3 p) const KOKKOS_OVERRIDE { return _function->div(t, p); }
  KOKKOS_FUNCTION Real timeDerivative(Real t, Real3 p) const KOKKOS_OVERRIDE
  {
    return _function->timeDerivative(t, p);
  }
  KOKKOS_FUNCTION Real timeIntegral(Real t1, Real t2, Real3 p) const KOKKOS_OVERRIDE
  {
    return _function->timeIntegral(t1, t2, p);
  }
  KOKKOS_FUNCTION Real integral() const KOKKOS_OVERRIDE { return _function->integral(); }
  KOKKOS_FUNCTION Real average() const KOKKOS_OVERRIDE { return _function->average(); }
 
protected:
  Object * _function = nullptr;
};
 
class FunctionWrapperHostBase
{
public:
  virtual ~FunctionWrapperHostBase() {}
 
  virtual FunctionWrapperDeviceBase * allocate() = 0;
  virtual void copyFunction() = 0;
  virtual void freeFunction() = 0;
};
 
template <typename Object>
class FunctionWrapperHost : public FunctionWrapperHostBase
{
public:
  FunctionWrapperHost(const void * function)
    : _function_host(*static_cast<const Object *>(function))
  {
  }
  ~FunctionWrapperHost();
 
  FunctionWrapperDeviceBase * allocate() override final;
  void copyFunction() override final;
  void freeFunction() override final;
 
private:
  const Object & _function_host;
  std::unique_ptr<Object> _function_copy;
  Object * _function_device = nullptr;
};
 
template <typename Object>
FunctionWrapperDeviceBase *
FunctionWrapperHost<Object>::allocate()
{
  // Allocate storage for device wrapper on device
  auto wrapper_device = static_cast<FunctionWrapperDevice<Object> *>(
      ::Kokkos::kokkos_malloc<ExecSpace::memory_space>(sizeof(FunctionWrapperDevice<Object>)));
 
  // Allocate device wrapper on device using placement new to populate vtable with device pointers
  ::Kokkos::parallel_for(
      1, KOKKOS_LAMBDA(const int) { new (wrapper_device) FunctionWrapperDevice<Object>(); });
 
  // Allocate storage for function on device
  _function_device =
      static_cast<Object *>(::Kokkos::kokkos_malloc<ExecSpace::memory_space>(sizeof(Object)));
 
  // Let device wrapper point to the copy
  ::Kokkos::Impl::DeepCopy<MemSpace, ::Kokkos::HostSpace>(
      &(wrapper_device->_function), &_function_device, sizeof(Object *));
 
  return wrapper_device;
}
 
template <typename Object>
void
FunctionWrapperHost<Object>::copyFunction()
{
  // Make a copy of function on host to trigger copy constructor
  _function_copy = std::make_unique<Object>(_function_host);
 
  // Copy function to device
  ::Kokkos::Impl::DeepCopy<MemSpace, ::Kokkos::HostSpace>(
      _function_device, _function_copy.get(), sizeof(Object));
}
 
template <typename Object>
void
FunctionWrapperHost<Object>::freeFunction()
{
  _function_copy.reset();
}
 
template <typename Object>
FunctionWrapperHost<Object>::~FunctionWrapperHost()
{
  ::Kokkos::kokkos_free<ExecSpace::memory_space>(_function_device);
}
 
} // namespace Moose::Kokkos