spaces.hpp

// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GINKGO_EXTENSIONS_KOKKOS_SPACES_HPP
#define GINKGO_EXTENSIONS_KOKKOS_SPACES_HPP
 
#include <Kokkos_Core.hpp>
 
#include <ginkgo/config.hpp>
#include <ginkgo/core/base/exception.hpp>
#include <ginkgo/core/base/exception_helpers.hpp>
#include <ginkgo/core/base/executor.hpp>
 
 
namespace gko {
namespace ext {
namespace kokkos {
namespace detail {
 
 
template <typename MemorySpace, typename ExecType>
struct compatible_space
    : std::integral_constant<bool, Kokkos::has_shared_space ||
                                       Kokkos::has_shared_host_pinned_space> {};
 
template <>
struct compatible_space<Kokkos::HostSpace, ReferenceExecutor> : std::true_type {
};
 
template <typename MemorySpace>
struct compatible_space<MemorySpace, ReferenceExecutor> {
    // need manual implementation of std::integral_constant because,
    // while compiling for cuda, somehow bool is replaced by __nv_bool
    static constexpr bool value =
        Kokkos::SpaceAccessibility<Kokkos::HostSpace, MemorySpace>::accessible;
};
 
#ifdef KOKKOS_ENABLE_OPENMP
template <typename MemorySpace>
struct compatible_space<MemorySpace, OmpExecutor>
    : compatible_space<MemorySpace, ReferenceExecutor> {};
#endif
 
#ifdef KOKKOS_ENABLE_CUDA
template <typename MemorySpace>
struct compatible_space<MemorySpace, CudaExecutor> {
    static constexpr bool value =
        Kokkos::SpaceAccessibility<Kokkos::Cuda, MemorySpace>::accessible;
};
#endif
 
#ifdef KOKKOS_ENABLE_HIP
template <typename MemorySpace>
struct compatible_space<MemorySpace, HipExecutor> {
    static constexpr bool value =
        Kokkos::SpaceAccessibility<Kokkos::HIP, MemorySpace>::accessible;
};
#endif
 
#ifdef KOKKOS_ENABLE_SYCL
template <typename MemorySpace>
struct compatible_space<MemorySpace, DpcppExecutor> {
    static constexpr bool value =
        Kokkos::SpaceAccessibility<Kokkos::Experimental::SYCL,
                                   MemorySpace>::accessible;
};
#endif
 
 
template <typename MemorySpace, typename ExecType>
inline bool check_compatibility(std::shared_ptr<const ExecType>)
{
    return compatible_space<MemorySpace, ExecType>::value;
}
 
 
template <typename MemorySpace>
inline bool check_compatibility(std::shared_ptr<const Executor> exec)
{
    if (auto p = std::dynamic_pointer_cast<const ReferenceExecutor>(exec)) {
        return check_compatibility<MemorySpace>(p);
    }
    if (auto p = std::dynamic_pointer_cast<const OmpExecutor>(exec)) {
        return check_compatibility<MemorySpace>(p);
    }
    if (auto p = std::dynamic_pointer_cast<const CudaExecutor>(exec)) {
        return check_compatibility<MemorySpace>(p);
    }
    if (auto p = std::dynamic_pointer_cast<const HipExecutor>(exec)) {
        return check_compatibility<MemorySpace>(p);
    }
    if (auto p = std::dynamic_pointer_cast<const DpcppExecutor>(exec)) {
        return check_compatibility<MemorySpace>(p);
    }
    GKO_NOT_IMPLEMENTED;
}
 
 
template <typename MemorySpace, typename T>
inline void assert_compatibility(T&& obj)
{
    GKO_THROW_IF_INVALID(check_compatibility<MemorySpace>(obj.get_executor()),
                         "Executor type and memory space are incompatible");
}
 
 
}  // namespace detail
 
 
inline std::shared_ptr<Executor> create_default_host_executor()
{
#ifdef KOKKOS_ENABLE_SERIAL
    if constexpr (std::is_same_v<Kokkos::DefaultHostExecutionSpace,
                                 Kokkos::Serial>) {
        return ReferenceExecutor::create();
    }
#endif
#ifdef KOKKOS_ENABLE_OPENMP
    if constexpr (std::is_same_v<Kokkos::DefaultHostExecutionSpace,
                                 Kokkos::OpenMP>) {
        return OmpExecutor::create();
    }
#endif
    GKO_NOT_IMPLEMENTED;
}
 
 
template <typename ExecSpace,
          typename MemorySpace = typename ExecSpace::memory_space>
inline std::shared_ptr<Executor> create_executor(ExecSpace ex, MemorySpace = {})
{
    static_assert(
        Kokkos::SpaceAccessibility<ExecSpace, MemorySpace>::accessible);
#ifdef KOKKOS_ENABLE_SERIAL
    if constexpr (std::is_same_v<ExecSpace, Kokkos::Serial>) {
        return ReferenceExecutor::create();
    }
#endif
#ifdef KOKKOS_ENABLE_OPENMP
    if constexpr (std::is_same_v<ExecSpace, Kokkos::OpenMP>) {
        return OmpExecutor::create();
    }
#endif
#ifdef KOKKOS_ENABLE_CUDA
    if constexpr (std::is_same_v<ExecSpace, Kokkos::Cuda>) {
        if constexpr (std::is_same_v<MemorySpace, Kokkos::CudaSpace>) {
            return CudaExecutor::create(
                Kokkos::device_id(), create_default_host_executor(),
                std::make_shared<CudaAllocator>(), ex.cuda_stream());
        }
        if constexpr (std::is_same_v<MemorySpace, Kokkos::CudaUVMSpace>) {
            return CudaExecutor::create(
                Kokkos::device_id(), create_default_host_executor(),
                std::make_shared<CudaUnifiedAllocator>(Kokkos::device_id()),
                ex.cuda_stream());
        }
        if constexpr (std::is_same_v<MemorySpace,
                                     Kokkos::CudaHostPinnedSpace>) {
            return CudaExecutor::create(
                Kokkos::device_id(), create_default_host_executor(),
                std::make_shared<CudaHostAllocator>(Kokkos::device_id()),
                ex.cuda_stream());
        }
    }
#endif
#ifdef KOKKOS_ENABLE_HIP
    if constexpr (std::is_same_v<ExecSpace, Kokkos::HIP>) {
        if constexpr (std::is_same_v<MemorySpace, Kokkos::HIPSpace>) {
            return HipExecutor::create(
                Kokkos::device_id(), create_default_host_executor(),
                std::make_shared<HipAllocator>(), ex.hip_stream());
        }
        if constexpr (std::is_same_v<MemorySpace, Kokkos::HIPManagedSpace>) {
            return HipExecutor::create(
                Kokkos::device_id(), create_default_host_executor(),
                std::make_shared<HipUnifiedAllocator>(Kokkos::device_id()),
                ex.hip_stream());
        }
        if constexpr (std::is_same_v<MemorySpace, Kokkos::HIPHostPinnedSpace>) {
            return HipExecutor::create(
                Kokkos::device_id(), create_default_host_executor(),
                std::make_shared<HipHostAllocator>(Kokkos::device_id()),
                ex.hip_stream());
        }
    }
#endif
#ifdef KOKKOS_ENABLE_SYCL
    if constexpr (std::is_same_v<ExecSpace, Kokkos::Experimental::SYCL>) {
        static_assert(
            std::is_same_v<MemorySpace, Kokkos::Experimental::SYCLSpace>,
            "Ginkgo doesn't support shared memory space allocation for SYCL");
        return DpcppExecutor::create(Kokkos::device_id(),
                                     create_default_host_executor());
    }
#endif
    GKO_NOT_IMPLEMENTED;
}
 
 
inline std::shared_ptr<Executor> create_default_executor(
    Kokkos::DefaultExecutionSpace ex = {})
{
    return create_executor(std::move(ex));
}
 
 
}  // namespace kokkos
}  // namespace ext
}  // namespace gko
 
 
#endif  // GINKGO_EXTENSIONS_KOKKOS_SPACES_HPP