1+ /*
2+ //@HEADER
3+ // ************************************************************************
4+ //
5+ // Kokkos v. 3.0
6+ // Copyright (2020) National Technology & Engineering
7+ // Solutions of Sandia, LLC (NTESS).
8+ //
9+ // Under the terms of Contract DE-NA0003525 with NTESS,
10+ // the U.S. Government retains certain rights in this software.
11+ //
12+ // Redistribution and use in source and binary forms, with or without
13+ // modification, are permitted provided that the following conditions are
14+ // met:
15+ //
16+ // 1. Redistributions of source code must retain the above copyright
17+ // notice, this list of conditions and the following disclaimer.
18+ //
19+ // 2. Redistributions in binary form must reproduce the above copyright
20+ // notice, this list of conditions and the following disclaimer in the
21+ // documentation and/or other materials provided with the distribution.
22+ //
23+ // 3. Neither the name of the Corporation nor the names of the
24+ // contributors may be used to endorse or promote products derived from
25+ // this software without specific prior written permission.
26+ //
27+ // THIS SOFTWARE IS PROVIDED BY NTESS "AS IS" AND ANY
28+ // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
29+ // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
30+ // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL NTESS OR THE
31+ // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
32+ // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
33+ // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
34+ // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
35+ // LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
36+ // NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
37+ // SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
38+ //
39+ // Questions? Contact Christian R. Trott ([email protected] ) 40+ //
41+ // ************************************************************************
42+ //@HEADER
43+ */
44+
45+ #pragma once
46+
47+ #include " common.hpp"
48+ #include " concepts.hpp"
49+ #include " pools.hpp"
50+ #include " traits.hpp"
51+
52+ #include < Kokkos_Core.hpp>
53+
54+ namespace Common {
55+ template <typename Sp, size_t SpaceIdx>
56+ void generate_execution_space_init (py::class_<Sp> &,
57+ enable_if_t <SpaceIdx == Serial_Backend, int > = 0 ) {}
58+
59+ template <typename Sp, size_t SpaceIdx>
60+ void generate_execution_space_init (py::class_<Sp> &,
61+ enable_if_t <SpaceIdx == Threads_Backend, int > = 0 ) {}
62+
63+ template <typename Sp, size_t SpaceIdx>
64+ void generate_execution_space_init (py::class_<Sp> &,
65+ enable_if_t <SpaceIdx == OpenMP_Backend, int > = 0 ) {}
66+
67+ template <typename Sp, size_t SpaceIdx>
68+ void generate_execution_space_init (py::class_<Sp> &_space,
69+ enable_if_t <SpaceIdx == Cuda_Backend, int > = 0 ) {
70+ #if defined(KOKKOS_ENABLE_CUDA) && defined(__CUDACC__)
71+ _space.def (py::init ([](uint64_t stream_ptr, bool manage_stream) {
72+ cudaStream_t s = reinterpret_cast <cudaStream_t>(stream_ptr);
73+ return new Sp{s, manage_stream};
74+ }));
75+ #endif
76+ }
77+
78+ template <typename Sp, size_t SpaceIdx>
79+ void generate_execution_space_init (py::class_<Sp> &_space,
80+ enable_if_t <SpaceIdx == HIP_Backend, int > = 0 ) {
81+ #if defined(KOKKOS_ENABLE_HIP) && defined(__HIPCC__)
82+ _space.def (py::init ([](uint64_t stream_ptr, bool manage_stream) {
83+ hipStream_t s = reinterpret_cast <hipStream_t>(stream_ptr);
84+ return new Sp{s, manage_stream};
85+ }));
86+ #endif
87+ }
88+
89+ template <typename Sp, size_t SpaceIdx>
90+ void generate_execution_space (py::module &_mod, const std::string &_name,
91+ const std::string &_msg) {
92+ if (debug_output ())
93+ std::cerr << " Registering " " as python class '"
94+ << " '..."
95+
96+ py::class_<Sp> _space (_mod, _name.c_str ());
97+ _space.def (py::init ([]() { return new Sp{}; }));
98+
99+ // Add other constructors with arguments if they exist
100+ generate_execution_space_init<Sp, SpaceIdx>(_space);
101+ }
102+ } // namespace Common
103+
104+ namespace Space {
105+ namespace SpaceDim {
106+
107+ template <size_t SpaceIdx>
108+ void generate_execution_space (py::module &_mod) {
109+ using space_spec_t = ExecutionSpaceSpecialization<SpaceIdx>;
110+ using Sp = typename space_spec_t ::type;
111+
112+ auto name = join (" _" " KokkosExecutionSpace" space_spec_t::label ());
113+
114+ Common::generate_execution_space<Sp, SpaceIdx>(_mod, name, demangle<Sp>());
115+ }
116+ } // namespace SpaceDim
117+
118+ template <size_t SpaceIdx>
119+ void generate_execution_space (
120+ py::module &,
121+ std::enable_if_t <!is_available<execution_space_t <SpaceIdx>>::value, int > =
122+ 0 ) {}
123+
124+ template <size_t SpaceIdx>
125+ void generate_execution_space (
126+ py::module &_mod,
127+ std::enable_if_t <is_available<execution_space_t <SpaceIdx>>::value, int > =
128+ 0 ) {
129+ SpaceDim::generate_execution_space<SpaceIdx>(_mod);
130+ }
131+ } // namespace Space
132+
133+ template <size_t ... SpaceIdx>
134+ void generate_execution_spaces (py::module &_mod,
135+ std::index_sequence<SpaceIdx...>) {
136+ FOLD_EXPRESSION (Space::generate_execution_space<SpaceIdx>(_mod));
137+ }
0 commit comments