From b8694f5986eeeb7aacbf5d75430e3d9caa986629 Mon Sep 17 00:00:00 2001 From: Nader Al Awar Date: Thu, 19 Oct 2023 09:44:46 -0500 Subject: [PATCH] ENH: add support for execution space constructors (#57) * ENH: add support for execution space constructors * Add missing iostream include --- CMakeLists.txt | 6 +- examples/user.cpp | 1 + include/execution_spaces.hpp | 137 +++++++++++++++++++++++++++++++++++ include/libpykokkos.hpp | 1 + src/execution_spaces.cpp | 55 ++++++++++++++ src/libpykokkos.cpp | 1 + 6 files changed, 199 insertions(+), 2 deletions(-) create mode 100644 include/execution_spaces.hpp create mode 100644 src/execution_spaces.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 299c057..80482b7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -43,7 +43,8 @@ SET(libpykokkos_SOURCES ${CMAKE_CURRENT_LIST_DIR}/src/enumeration.cpp ${CMAKE_CURRENT_LIST_DIR}/src/available.cpp ${CMAKE_CURRENT_LIST_DIR}/src/common.cpp - ${CMAKE_CURRENT_LIST_DIR}/src/tools.cpp) + ${CMAKE_CURRENT_LIST_DIR}/src/tools.cpp + ${CMAKE_CURRENT_LIST_DIR}/src/execution_spaces.cpp) SET(libpykokkos_HEADERS ${CMAKE_CURRENT_LIST_DIR}/include/libpykokkos.hpp @@ -53,7 +54,8 @@ SET(libpykokkos_HEADERS ${CMAKE_CURRENT_LIST_DIR}/include/common.hpp ${CMAKE_CURRENT_LIST_DIR}/include/traits.hpp ${CMAKE_CURRENT_LIST_DIR}/include/views.hpp - ${CMAKE_CURRENT_LIST_DIR}/include/fwd.hpp) + ${CMAKE_CURRENT_LIST_DIR}/include/fwd.hpp + ${CMAKE_CURRENT_LIST_DIR}/include/execution_spaces.hpp) ADD_LIBRARY(libpykokkos-core OBJECT ${libpykokkos_SOURCES} diff --git a/examples/user.cpp b/examples/user.cpp index a026b48..9deb320 100644 --- a/examples/user.cpp +++ b/examples/user.cpp @@ -4,6 +4,7 @@ #include "Kokkos_Core.hpp" #include +#include struct InitView { explicit InitView(view_type _v) : m_view(_v) {} diff --git a/include/execution_spaces.hpp b/include/execution_spaces.hpp new file mode 100644 index 0000000..669ed88 --- /dev/null +++ b/include/execution_spaces.hpp @@ -0,0 +1,137 @@ +/* +//@HEADER +// ************************************************************************ +// +// Kokkos v. 3.0 +// Copyright (2020) National Technology & Engineering +// Solutions of Sandia, LLC (NTESS). +// +// Under the terms of Contract DE-NA0003525 with NTESS, +// 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 NTESS "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 NTESS 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. +// +// Questions? Contact Christian R. Trott (crtrott@sandia.gov) +// +// ************************************************************************ +//@HEADER +*/ + +#pragma once + +#include "common.hpp" +#include "concepts.hpp" +#include "pools.hpp" +#include "traits.hpp" + +#include + +namespace Common { +template +void generate_execution_space_init(py::class_ &, + enable_if_t = 0) {} + +template +void generate_execution_space_init(py::class_ &, + enable_if_t = 0) {} + +template +void generate_execution_space_init(py::class_ &, + enable_if_t = 0) {} + +template +void generate_execution_space_init(py::class_ &_space, + enable_if_t = 0) { +#if defined(KOKKOS_ENABLE_CUDA) && defined(__CUDACC__) + _space.def(py::init([](uint64_t stream_ptr, bool manage_stream) { + cudaStream_t s = reinterpret_cast(stream_ptr); + return new Sp{s, manage_stream}; + })); +#endif +} + +template +void generate_execution_space_init(py::class_ &_space, + enable_if_t = 0) { +#if defined(KOKKOS_ENABLE_HIP) && defined(__HIPCC__) + _space.def(py::init([](uint64_t stream_ptr, bool manage_stream) { + hipStream_t s = reinterpret_cast(stream_ptr); + return new Sp{s, manage_stream}; + })); +#endif +} + +template +void generate_execution_space(py::module &_mod, const std::string &_name, + const std::string &_msg) { + if (debug_output()) + std::cerr << "Registering " << _msg << " as python class '" << _name + << "'..." << std::endl; + + py::class_ _space(_mod, _name.c_str()); + _space.def(py::init([]() { return new Sp{}; })); + + // Add other constructors with arguments if they exist + generate_execution_space_init(_space); +} +} // namespace Common + +namespace Space { +namespace SpaceDim { + +template +void generate_execution_space(py::module &_mod) { + using space_spec_t = ExecutionSpaceSpecialization; + using Sp = typename space_spec_t::type; + + auto name = join("_", "KokkosExecutionSpace", space_spec_t::label()); + + Common::generate_execution_space(_mod, name, demangle()); +} +} // namespace SpaceDim + +template +void generate_execution_space( + py::module &, + std::enable_if_t>::value, int> = + 0) {} + +template +void generate_execution_space( + py::module &_mod, + std::enable_if_t>::value, int> = + 0) { + SpaceDim::generate_execution_space(_mod); +} +} // namespace Space + +template +void generate_execution_spaces(py::module &_mod, + std::index_sequence) { + FOLD_EXPRESSION(Space::generate_execution_space(_mod)); +} \ No newline at end of file diff --git a/include/libpykokkos.hpp b/include/libpykokkos.hpp index 891fa6b..18aed0e 100644 --- a/include/libpykokkos.hpp +++ b/include/libpykokkos.hpp @@ -57,4 +57,5 @@ void generate_view_variants(py::module& kokkos); void generate_atomic_variants(py::module& kokkos); void generate_backend_versions(py::module& kokkos); void generate_pool_variants(py::module& kokkos); +void generate_execution_spaces(py::module& kokkos); void destroy_callbacks(); diff --git a/src/execution_spaces.cpp b/src/execution_spaces.cpp new file mode 100644 index 0000000..13f9be2 --- /dev/null +++ b/src/execution_spaces.cpp @@ -0,0 +1,55 @@ +/* +//@HEADER +// ************************************************************************ +// +// Kokkos v. 3.0 +// Copyright (2020) National Technology & Engineering +// Solutions of Sandia, LLC (NTESS). +// +// Under the terms of Contract DE-NA0003525 with NTESS, +// 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 NTESS "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 NTESS 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. +// +// Questions? Contact Christian R. Trott (crtrott@sandia.gov) +// +// ************************************************************************ +//@HEADER +*/ + +#include "execution_spaces.hpp" + +#include "common.hpp" +#include "defines.hpp" +#include "fwd.hpp" +#include "traits.hpp" + +void generate_execution_spaces(py::module &kokkos) { + generate_execution_spaces(kokkos, + std::make_index_sequence{}); +} \ No newline at end of file diff --git a/src/libpykokkos.cpp b/src/libpykokkos.cpp index e1fa535..f3305e2 100644 --- a/src/libpykokkos.cpp +++ b/src/libpykokkos.cpp @@ -114,4 +114,5 @@ PYBIND11_MODULE(libpykokkos, kokkos) { generate_atomic_variants(kokkos); generate_backend_versions(kokkos); generate_pool_variants(kokkos); + generate_execution_spaces(kokkos); }