Tpetra_withLocalAccess.hpp

Go to the documentation of this file.

/*
// @HEADER
// ***********************************************************************
//
//          Tpetra: Templated Linear Algebra Services Package
//                 Copyright (2008) Sandia Corporation
//
// Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
// the U.S. Government retains certain rights in this software.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
// 1. Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the Corporation nor the names of the
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY SANDIA CORPORATION "AS IS" AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL SANDIA CORPORATION OR THE
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// ************************************************************************
// @HEADER
*/

#ifndef TPETRA_WITHLOCALACCESS_HPP
#define TPETRA_WITHLOCALACCESS_HPP

#include "TpetraCore_config.h"
#include "Tpetra_Access.hpp"
#include <functional>
#include <type_traits>


namespace Tpetra {

  // Generic classes needed to implement withLocalAccess

  namespace Details {
    template<class T> struct is_access_mode : public std::false_type {};
    template<> struct is_access_mode<Access::ReadOnly> : public std::true_type {};
    template<> struct is_access_mode<Access::ReadWrite> : public std::true_type {};
    template<> struct is_access_mode<Access::WriteOnly> : public std::true_type {};

    template<class GlobalObjectType>
    struct DefaultMemorySpace {
      using type = typename GlobalObjectType::device_type::memory_space;

      // Given a global object, get its (default) memory space instance.
      static type space (const GlobalObjectType& /* G */) {
        // This stub just assumes that 'type' is default constructible.
        // In Kokkos, default-constructing a memory space instance just
        // gives the default memory space.
        return type ();
      }
    };

    template<class GlobalObjectType>
    struct DefaultExecutionSpace {
      using type = typename GlobalObjectType::device_type::execution_space;

      // Given a global object, get its (default) execution space instance.
      static type space (const GlobalObjectType& /* G */) {
        // This stub just assumes that 'type' is default constructible.
        // In Kokkos, default-constructing a execution space instance just
        // gives the default execution space.
        return type ();
      }
    };

#ifndef DOXYGEN_SHOULD_SKIP_THIS
    template<class GlobalObjectType, class ... Args>
    class LocalAccess; // forward declaration
#endif // DOXYGEN_SHOULD_SKIP_THIS

    template<class LocalAccessType>
    struct GetMasterLocalObject {};

    template<class LocalAccessType>
    typename GetMasterLocalObject<LocalAccessType>::
      master_local_object_type
    getMasterLocalObject (LocalAccessType LA) {
      return GetMasterLocalObject<LocalAccessType>::get (LA);
    }

    template<class LocalAccessType>
    struct GetNonowningLocalObject {};

    template<class LocalAccessType>
    typename GetNonowningLocalObject<LocalAccessType>::
    nonowning_local_object_type
    getNonowningLocalObject(LocalAccessType LA,
      const typename GetMasterLocalObject<LocalAccessType>::
        master_local_object_type& master)
    {
      using impl_type = GetNonowningLocalObject<LocalAccessType>;
      return impl_type::get(LA, master);
    }
  } // namespace Details

  template<class LocalAccessType>
  using with_local_access_function_argument_type =
    typename Details::GetNonowningLocalObject<LocalAccessType>::
      nonowning_local_object_type;

  // Users call readOnly, writeOnly, and readWrite, in order to declare
  // how they intend to access a global object's local data.

  template<class GlobalObjectType>
  Details::LocalAccess<
    GlobalObjectType,
    Access::ReadOnly>
  readOnly (GlobalObjectType&);

  template<class GlobalObjectType>
  Details::LocalAccess<
    GlobalObjectType,
    Access::ReadOnly>
  readOnly (const GlobalObjectType&);

  template<class GlobalObjectType>
  Details::LocalAccess<
    GlobalObjectType,
    Access::WriteOnly>
  writeOnly (GlobalObjectType&);

  template<class GlobalObjectType>
  Details::LocalAccess<
    GlobalObjectType,
    Access::ReadWrite>
  readWrite (GlobalObjectType&);

  // LocalAccess struct

  namespace Details {
    struct unspecified_type {};

    template<class ... Args>
    struct LocalAccessTraits {};

    // An empty list of template arguments constrains nothing.
    template<>
    struct LocalAccessTraits<> {
      using execution_space = unspecified_type;
      using memory_space = unspecified_type;
      using access_mode = unspecified_type;
    };

    // Skip over any instances of unspecified_type.  This makes it
    // easy to define the return types of LocalAccess::on and
    // LocalAccess::at (see below).
    template<class ... Rest>
    struct LocalAccessTraits<unspecified_type, Rest...> {
    private:
      using rest_type = LocalAccessTraits<Rest...>;
    public:
      using execution_space = typename rest_type::execution_space;
      using memory_space = typename rest_type::memory_space;
      using access_mode = typename rest_type::access_mode;
    };

    template<class First, class ... Rest>
    struct LocalAccessTraits<First, Rest...> {
    private:
      using rest_type = LocalAccessTraits<Rest...>;
    public:
      using execution_space = typename std::conditional<
        Kokkos::Impl::is_execution_space<First>::value,
        First,
        typename rest_type::execution_space>::type;
      using memory_space = typename std::conditional<
        Kokkos::Impl::is_memory_space<First>::value,
        First,
        typename rest_type::memory_space>::type;
      using access_mode = typename std::conditional<
        is_access_mode<First>::value,
        First,
        typename rest_type::access_mode>::type;
    };

    template<class GlobalObjectType,
             class Traits,
             bool is_execution_space_specified = ! std::is_same<
               typename Traits::execution_space,
               unspecified_type>::value,
             bool is_memory_space_specified = ! std::is_same<
               typename Traits::memory_space,
               unspecified_type>::value>
    struct SpaceTypes {};

    template<class GlobalObjectType, class Traits>
    struct SpaceTypes<GlobalObjectType, Traits, true, true> {
      using execution_space = typename Traits::execution_space;
      using memory_space = typename Traits::memory_space;
    };

    // If only memory_space was specified, then get execution_space
    // from memory_space's default execution space.
    template<class GlobalObjectType, class Traits>
    struct SpaceTypes<GlobalObjectType, Traits, false, true> {
      using execution_space =
        typename Traits::memory_space::execution_space;
      using memory_space = typename Traits::memory_space;
    };

    // If only execution_space was specified, then get memory_space
    // from execution_space's default memory space.
    template<class GlobalObjectType, class Traits>
    struct SpaceTypes<GlobalObjectType, Traits, true, false> {
      using execution_space = typename Traits::execution_space;
      using memory_space =
        typename Traits::execution_space::memory_space;
    };

    // If neither execution_space nor memory_space were specified,
    // then get them both from their defaults, that depend on
    // GlobalObjectType.
    template<class GlobalObjectType, class Traits>
    struct SpaceTypes<GlobalObjectType, Traits, false, false> {
      using execution_space =
        typename DefaultExecutionSpace<GlobalObjectType>::type;
      using memory_space =
        typename DefaultMemorySpace<GlobalObjectType>::type;
    };

    template<class GlobalObjectType, class ... Args>
    class LocalAccess {
    private:
      using this_type = LocalAccess<GlobalObjectType, Args...>;
      using traits = LocalAccessTraits<Args...>;
      static constexpr bool is_execution_space_specified =
        ! std::is_same<typename traits::execution_space,
                       unspecified_type>::value;
      static constexpr bool is_memory_space_specified =
        ! std::is_same<typename traits::memory_space,
                       unspecified_type>::value;
      static constexpr bool is_access_mode_specified =
        ! std::is_same<typename traits::access_mode,
                       unspecified_type>::value;
      static_assert(is_access_mode_specified, "LocalAccess requires "
        "that you specify the access mode.  You may do this by "
        "always starting construction of a LocalAccess instance "
        "with readOnly, writeOnly, or readWrite.");

    public:
      using global_object_type = GlobalObjectType;

      using execution_space =
        typename SpaceTypes<global_object_type, traits>::execution_space;
      static_assert(! is_execution_space_specified ||
        Kokkos::Impl::is_execution_space<execution_space>::value,
        "Specified execution space is not a Kokkos execution space.");
      static_assert(is_execution_space_specified ||
        Kokkos::Impl::is_execution_space<execution_space>::value,
        "Default execution space is not a Kokkos execution space.");

      using memory_space =
        typename SpaceTypes<global_object_type, traits>::memory_space;
      static_assert(! is_memory_space_specified ||
        Kokkos::Impl::is_memory_space<memory_space>::value,
        "Specified memory space is not a Kokkos memory space.");
      static_assert(is_memory_space_specified ||
        Kokkos::Impl::is_memory_space<memory_space>::value,
        "Default memory space is not a Kokkos memory space.");

      using access_mode = typename traits::access_mode;

      LocalAccess (global_object_type& G,
                   const execution_space& execSpace,
                   const memory_space& memSpace,
                   const bool viewIsValid = true) :
        G_ (G),
        execSpace_ (execSpace),
        memSpace_ (memSpace),
        valid_ (viewIsValid)
      {}

      LocalAccess (global_object_type& G,
                   const bool viewIsValid = true) :
        LocalAccess (G,
          DefaultExecutionSpace<global_object_type>::space (G),
          DefaultMemorySpace<global_object_type>::space (G),
          viewIsValid)
      {}

      LocalAccess (global_object_type& G,
                   const execution_space& execSpace,
                   const bool viewIsValid = true) :
        LocalAccess (G,
          execSpace,
          DefaultMemorySpace<global_object_type>::space (G),
          viewIsValid)
      {}

      LocalAccess (global_object_type& G,
                   const memory_space& memSpace,
                   const bool viewIsValid = true) :
        LocalAccess (G,
          DefaultExecutionSpace<global_object_type>::space (G),
          memSpace,
          viewIsValid)
      {}

      bool isValid () const { return valid_; }

      execution_space getExecutionSpace () const { return execSpace_; }

      memory_space getMemorySpace () const { return memSpace_; }

      this_type
      valid (const bool thisIsValid) const {
        return {G_, getExecutionSpace(), getMemorySpace(),
                thisIsValid};
      }

      template<class NewMemorySpace>
      typename std::conditional<
        is_execution_space_specified,
        LocalAccess<global_object_type, execution_space,
                    NewMemorySpace, access_mode>,
        LocalAccess<global_object_type, NewMemorySpace, access_mode>
        >::type
      on(const NewMemorySpace& memSpace) const {
        using Kokkos::Impl::is_memory_space;
        static_assert(is_memory_space<NewMemorySpace>::value,
          "NewMemorySpace must be a Kokkos memory space.");

        // We can't use std::conditional here, because we need the
        // "select" method.
        using alt_execution_space =
          typename LocalAccess<global_object_type, NewMemorySpace,
                               access_mode>::execution_space;
        auto execSpace = Kokkos::Impl::if_c<
          is_execution_space_specified,
          execution_space,
          alt_execution_space>::select(
            getExecutionSpace(),
            alt_execution_space());
        return {G_, execSpace, memSpace, isValid()};
      }

      template<class NewExecutionSpace>
      typename std::conditional<
        is_memory_space_specified,
        LocalAccess<global_object_type, NewExecutionSpace,
                    memory_space, access_mode>,
        LocalAccess<global_object_type, NewExecutionSpace, access_mode>
        >::type
      at(const NewExecutionSpace& execSpace) const {
        using Kokkos::Impl::is_execution_space;
        static_assert(is_execution_space<NewExecutionSpace>::value,
          "NewExecutionSpace must be a Kokkos execution space.");

        // We can't use std::conditional here, because we need the
        // "select" method.
        using alt_memory_space =
          typename LocalAccess<global_object_type, NewExecutionSpace,
                               access_mode>::memory_space;
        auto memSpace = Kokkos::Impl::if_c<
          is_memory_space_specified,
          memory_space,
          alt_memory_space>::select(
            getMemorySpace(),
            alt_memory_space());
        return {G_, execSpace, memSpace, isValid()};
      }

    public:
      global_object_type& G_;

      execution_space execSpace_;

      memory_space memSpace_;

      bool valid_;
    };
  } // namespace Details

  // Implementations of readOnly, writeOnly, and readWrite

  template<class GlobalObjectType>
  Details::LocalAccess<
    GlobalObjectType,
    Access::ReadOnly>
  readOnly (GlobalObjectType& G)
  {
    return {G};
  }

  template<class GlobalObjectType>
  Details::LocalAccess<
    GlobalObjectType,
    Access::ReadOnly>
  readOnly (const GlobalObjectType& G)
  {
    GlobalObjectType& G_nc = const_cast<GlobalObjectType&> (G);
    return {G_nc};
  }

  template<class GlobalObjectType>
  Details::LocalAccess<
    GlobalObjectType,
    Access::WriteOnly>
  writeOnly (GlobalObjectType& G)
  {
    return {G};
  }

  template<class GlobalObjectType>
  Details::LocalAccess<
    GlobalObjectType,
    Access::ReadWrite>
  readWrite (GlobalObjectType& G)
  {
    return {G};
  }

  // Implementation of withLocalAccess

  namespace Details {

    // Use std::tuple as a compile-time list (Greenspun's 10th Rule)

    // cons<T, std::tuple<Args...>>::type is std::tuple<T, Args...>.
    // This is the usual Lisp CONS function, but for std::tuple.  I
    // got the idea from
    // https://stackoverflow.com/questions/18701798/building-and-accessing-a-list-of-types-at-compile-time
    // but without "struct Void {};", and with head of new list
    // ordered first instead of last (hence "cons").
    template<class ...> struct cons;

    // (CONS SomeType NIL)
    template<class T, template <class ...> class List>
    struct cons<T, List<>> {
      using type = List<T>;
    };

    // (CONS SomeType ListOfTypes)
    template <class T, template <class ...> class List, class ...Types>
    struct cons<T, List<Types...>>
    {
      typedef List<T, Types...> type;
    };

    // Map from std::tuple<Ts...> to std::function<void (Ts...)>.

    // I got inspiration from
    // https://stackoverflow.com/questions/15418841/how-do-i-strip-a-tuple-back-into-a-variadic-template-list-of-types
    //
    // The only significant change was from "using type = T<Ts...>;",
    // to "using type = std::function<void (Ts...)>;".
    template<class T>
    struct tuple_to_function_type { };

    template<typename... Ts>
    struct tuple_to_function_type<std::tuple<Ts...> >
    {
      using type = std::function<void (Ts...)>;
    };

    // Map from a list of zero or more LocalAccess types, to the
    // corresponding list of arguments for the user function to give to
    // withLocalAccess.
    template<class ... Args>
    struct ArgsToFunction {};

    template<>
    struct ArgsToFunction<> {
      using arg_list_type = std::tuple<>;

      // Implementers should only use this.
      using type = std::function<void ()>;
    };

    template<class FirstLocalAccessType, class ... Rest>
    struct ArgsToFunction<FirstLocalAccessType, Rest...> {
      using head_arg_type =
        with_local_access_function_argument_type<FirstLocalAccessType>;
      using tail_arg_list_type =
        typename ArgsToFunction<Rest...>::arg_list_type;
      using arg_list_type =
        typename cons<head_arg_type, tail_arg_list_type>::type;

      // Implementers should only use this.
      using type = typename tuple_to_function_type<arg_list_type>::type;
    };

    template<class ... LocalAccessTypes>
    struct WithLocalAccess {
      using current_user_function_type =
        typename ArgsToFunction<LocalAccessTypes...>::type;

      static void
      withLocalAccess (LocalAccessTypes...,
                       typename ArgsToFunction<LocalAccessTypes...>::type);
    };

    template<>
    struct WithLocalAccess<> {
      using current_user_function_type =
        typename ArgsToFunction<>::type;

      static void
      withLocalAccess (current_user_function_type userFunction)
      {
        userFunction ();
      }
    };

    template<class FirstLocalAccessType>
    struct WithLocalAccess<FirstLocalAccessType> {
      using current_user_function_type =
        typename ArgsToFunction<FirstLocalAccessType>::type;

      static void
      withLocalAccess (current_user_function_type userFunction,
                       FirstLocalAccessType first)
      {
        // The "master" local object is the scope guard for local
        // data.  Its constructor may allocate temporary storage, copy
        // data to the desired memory space, etc.  Its destructor will
        // put everything back.  "Put everything back" could be a
        // no-op, or it could copy data back so where they need to go
        // and/or free temporary storage.
        //
        // Users define this function and the type it returns by
        // specializing GetMasterLocalObject for LocalAccess
        // specializations.
        auto first_lcl_master = getMasterLocalObject (first);

        // The "nonowning" local object is a nonowning view of the
        // "master" local object.  This is the only local object that
        // users see, and they see it as input to their function.
        // Subsequent slices / subviews view this nonowning local
        // object.  All such nonowning views must have lifetime
        // contained within the lifetime of the master local object.
        //
        // Users define this function and the type it returns by
        // specializing GetNonowningLocalObject (see above).
        //
        // Constraining the nonowning views' lifetime to this scope
        // means that master local object types may use low-cost
        // ownership models, like that of std::unique_ptr.  There
        // should be no need for reference counting (in the manner of
        // std::shared_ptr) or Herb Sutter's deferred_heap.
        auto first_lcl_view = getNonowningLocalObject (first, first_lcl_master);

        // Curry the user's function by fixing the first argument.

        WithLocalAccess<>::withLocalAccess
          ([=] () {
             userFunction (first_lcl_view);
           });
      }
    };

    template<class FirstLocalAccessType, class ... Rest>
    struct WithLocalAccess<FirstLocalAccessType, Rest...> {
      using current_user_function_type =
        typename ArgsToFunction<FirstLocalAccessType, Rest...>::type;

      static void
      withLocalAccess (current_user_function_type userFunction,
                       FirstLocalAccessType first,
                       Rest... rest)
      {
        // The "master" local object is the scope guard for local
        // data.  Its constructor may allocate temporary storage, copy
        // data to the desired memory space, etc.  Its destructor will
        // put everything back.  "Put everything back" could be a
        // no-op, or it could copy data back so where they need to go
        // and/or free temporary storage.
        //
        // Users define this function and the type it returns by
        // specializing GetMasterLocalObject for LocalAccess
        // specializations.
        auto first_lcl_master = getMasterLocalObject (first);

        // The "nonowning" local object is a nonowning view of the
        // "master" local object.  This is the only local object that
        // users see, and they see it as input to their function.
        // Subsequent slices / subviews view this nonowning local
        // object.  All such nonowning views must have lifetime
        // contained within the lifetime of the master local object.
        //
        // Users define this function and the type it returns by
        // specializing GetNonowningLocalObject (see above).
        //
        // Constraining the nonowning views' lifetime to this scope
        // means that master local object types may use low-cost
        // ownership models, like that of std::unique_ptr.  There
        // should be no need for reference counting (in the manner of
        // std::shared_ptr) or Herb Sutter's deferred_heap.
        auto first_lcl_view = getNonowningLocalObject (first, first_lcl_master);

        // Curry the user's function by fixing the first argument.

        WithLocalAccess<Rest...>::withLocalAccess
          ([=] (with_local_access_function_argument_type<Rest>... args) {
             userFunction (first_lcl_view, args...);
           },
           rest...);
      }
    };
  } // namespace Details

  // withLocalAccess function declaration and definition

  template<class ... LocalAccessTypes>
  void
  withLocalAccess
    (typename Details::ArgsToFunction<LocalAccessTypes...>::type userFunction,
     LocalAccessTypes... localAccesses)
  {
    using impl_type = Details::WithLocalAccess<LocalAccessTypes...>;
    impl_type::withLocalAccess (userFunction, localAccesses...);
  }

} // namespace Tpetra

#endif // TPETRA_WITHLOCALACCESS_HPP