KokkosFunctorWrapper.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 FunctorWrapperHost;

class FunctorWrapperDeviceBase
{
public:
  KOKKOS_FUNCTION FunctorWrapperDeviceBase() {}
  KOKKOS_FUNCTION virtual ~FunctorWrapperDeviceBase() {}
};

template <typename Object>
class FunctorWrapperDevice : public FunctorWrapperDeviceBase
{
  friend class FunctorWrapperHost<Object>;

public:
  KOKKOS_FUNCTION FunctorWrapperDevice() {}

protected:
  Object * _functor = nullptr;
};

class FunctorWrapperHostBase
{
public:
  virtual ~FunctorWrapperHostBase() {}

  virtual FunctorWrapperDeviceBase * allocate() = 0;
  virtual void copyFunctor() = 0;
  virtual void freeFunctor() = 0;
};

template <typename Object>
class FunctorWrapperHost : public FunctorWrapperHostBase
{
public:
  FunctorWrapperHost(const void * functor) : _functor_host(*static_cast<const Object *>(functor)) {}
  ~FunctorWrapperHost();

  FunctorWrapperDeviceBase * allocate() override final;
  void copyFunctor() override final;
  void freeFunctor() override final;

private:
  const Object & _functor_host;
  std::unique_ptr<Object> _functor_copy;
  Object * _functor_device = nullptr;
};

template <typename Object>
FunctorWrapperDeviceBase *
FunctorWrapperHost<Object>::allocate()
{
  // Allocate storage for device wrapper on device
  auto wrapper_device = static_cast<FunctorWrapperDevice<Object> *>(
      ::Kokkos::kokkos_malloc<ExecSpace::memory_space>(sizeof(FunctorWrapperDevice<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) FunctorWrapperDevice<Object>(); });

  // Allocate storage for functor on device
  _functor_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->_functor), &_functor_device, sizeof(Object *));

  return wrapper_device;
}

template <typename Object>
void
FunctorWrapperHost<Object>::copyFunctor()
{
  // Make a copy of functor on host to trigger copy constructor
  _functor_copy = std::make_unique<Object>(_functor_host);

  // Copy functor to device
  ::Kokkos::Impl::DeepCopy<MemSpace, ::Kokkos::HostSpace>(
      _functor_device, _functor_copy.get(), sizeof(Object));
}

template <typename Object>
void
FunctorWrapperHost<Object>::freeFunctor()
{
  _functor_copy.reset();
}

template <typename Object>
FunctorWrapperHost<Object>::~FunctorWrapperHost()
{
  ::Kokkos::kokkos_free<ExecSpace::memory_space>(_functor_device);
}

} // namespace Moose::Kokkos