From 944b1d0e86d51a86c1edc0e422155d6caf93ecdc Mon Sep 17 00:00:00 2001 From: "Jonathan R. Madsen" Date: Thu, 28 Apr 2022 13:22:34 -0500 Subject: [PATCH] Random pools (#37) * Add bindings for random pools in Kokkos_Random.hpp * fix indentation warning in pools * remove f-string * Use ExecutionSpaceSpecialization and add get_execution_space function * Formatting Co-authored-by: Nader Al Awar --- include/libpykokkos.hpp | 1 + include/pool_variants/XorShift1024_pool.hpp | 89 +++++++++++++++++++++ include/pool_variants/XorShift64_pool.hpp | 89 +++++++++++++++++++++ include/pools.hpp | 74 +++++++++++++++++ kokkos/__init__.py.in | 1 + kokkos/utility.py | 20 +++++ src/CMakeLists.txt | 1 + src/enumeration.cpp | 13 +++ src/libpykokkos.cpp | 1 + src/pool_variants/CMakeLists.txt | 60 ++++++++++++++ src/pool_variants/pool.cpp.in | 11 +++ src/pool_variants/variant.cpp.in | 9 +++ 12 files changed, 369 insertions(+) create mode 100644 include/pool_variants/XorShift1024_pool.hpp create mode 100644 include/pool_variants/XorShift64_pool.hpp create mode 100644 include/pools.hpp create mode 100644 src/pool_variants/CMakeLists.txt create mode 100644 src/pool_variants/pool.cpp.in create mode 100644 src/pool_variants/variant.cpp.in diff --git a/include/libpykokkos.hpp b/include/libpykokkos.hpp index 5e022d0..891fa6b 100644 --- a/include/libpykokkos.hpp +++ b/include/libpykokkos.hpp @@ -56,4 +56,5 @@ void generate_enumeration(py::module& kokkos); 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 destroy_callbacks(); diff --git a/include/pool_variants/XorShift1024_pool.hpp b/include/pool_variants/XorShift1024_pool.hpp new file mode 100644 index 0000000..57203da --- /dev/null +++ b/include/pool_variants/XorShift1024_pool.hpp @@ -0,0 +1,89 @@ +/* +//@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" + +namespace Space { +namespace SpaceDim { + +template +void generate_XorShift1024_pool_variant(py::module &_mod) { + using space_spec_t = ExecutionSpaceSpecialization; + using Sp = typename space_spec_t::type; + using UniformT = Kokkos::Random_XorShift1024_Pool; + + auto name = join("_", "KokkosXorShift1024Pool", space_spec_t::label()); + + Common::generate_pool(_mod, name, demangle()); +} +} // namespace SpaceDim + +template +void generate_XorShift1024_pool_variant( + py::module &, + std::enable_if_t>::value, int> = + 0) {} + +template +void generate_XorShift1024_pool_variant( + py::module &_mod, + std::enable_if_t>::value, int> = + 0) { + SpaceDim::generate_XorShift1024_pool_variant(_mod); +} +} // namespace Space + +namespace { +// generate data-type, memory-space buffers for concrete dimension +template +void generate_XorShift1024_pool_variant(py::module &_mod, + std::index_sequence) { + FOLD_EXPRESSION(Space::generate_XorShift1024_pool_variant(_mod)); +} +} // namespace diff --git a/include/pool_variants/XorShift64_pool.hpp b/include/pool_variants/XorShift64_pool.hpp new file mode 100644 index 0000000..46d074c --- /dev/null +++ b/include/pool_variants/XorShift64_pool.hpp @@ -0,0 +1,89 @@ +/* +//@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" + +namespace Space { +namespace SpaceDim { + +template +void generate_XorShift64_pool_variant(py::module &_mod) { + using space_spec_t = ExecutionSpaceSpecialization; + using Sp = typename space_spec_t::type; + using UniformT = Kokkos::Random_XorShift64_Pool; + + auto name = join("_", "KokkosXorShift64Pool", space_spec_t::label()); + + Common::generate_pool(_mod, name, demangle()); +} +} // namespace SpaceDim + +template +void generate_XorShift64_pool_variant( + py::module &, + std::enable_if_t>::value, int> = + 0) {} + +template +void generate_XorShift64_pool_variant( + py::module &_mod, + std::enable_if_t>::value, int> = + 0) { + SpaceDim::generate_XorShift64_pool_variant(_mod); +} +} // namespace Space + +namespace { +// generate data-type, memory-space buffers for concrete dimension +template +void generate_XorShift64_pool_variant(py::module &_mod, + std::index_sequence) { + FOLD_EXPRESSION(Space::generate_XorShift64_pool_variant(_mod)); +} +} // namespace diff --git a/include/pools.hpp b/include/pools.hpp new file mode 100644 index 0000000..5548a34 --- /dev/null +++ b/include/pools.hpp @@ -0,0 +1,74 @@ +/* +//@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 +#include + +namespace Common { +template +void generate_pool(py::module &_mod, const std::string &_name, + const std::string &_msg) { + if (debug_output()) + std::cerr << "Registering " << _msg << " as python class '" << _name + << "'..." << std::endl; + + // using PoolT = Kokkos::Random_XorShift64_Pool; + // class decl + py::class_ _pool(_mod, _name.c_str()); + + // default initializer + _pool.def(py::init([]() { return new PoolT{}; })); + + _pool.def(py::init([](uint64_t seed) { return new PoolT{seed}; })); + + _pool.def( + "init", + [](PoolT &_p, uint64_t _seed, int _num_states) { + _p.init(_seed, _num_states); + }, + "Initialize the random pool"); +} +} // namespace Common \ No newline at end of file diff --git a/kokkos/__init__.py.in b/kokkos/__init__.py.in index fc927f5..2ed710e 100644 --- a/kokkos/__init__.py.in +++ b/kokkos/__init__.py.in @@ -184,6 +184,7 @@ try: "DefaultHostMemorySpace", "get_dtype", # get enumeration from string or vice-versa "get_layout", + "get_execution_space", "get_memory_space", "get_memory_trait", "get_host_accessible", # query build support diff --git a/kokkos/utility.py b/kokkos/utility.py index 081f095..7faa2f9 100644 --- a/kokkos/utility.py +++ b/kokkos/utility.py @@ -215,3 +215,23 @@ def create_mirror_view(src, copy=False): def deep_copy(dst, src): """Performs Kokkos::deep_copy""" return dst.deep_copy(src) + + +def random_pool(state, space, seed=None): + """Create a Random_XorShift Pool""" + + if state not in {64, 1024}: + raise ValueError(f"State size {state} not supported, only 64 and 1024.") + + if seed is not None and not isinstance(seed, int): + raise ValueError("Seed must be either None or of type int") + + _space = lib.get_execution_space(space) + _name = f"KokkosXorShift{state}Pool_{_space}" + + _cons = getattr(lib, _name) + + if seed is None: + return _cons() + + return _cons(seed) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 6db1d9e..7c8c4e2 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -1,2 +1,3 @@ +ADD_SUBDIRECTORY(pool_variants) ADD_SUBDIRECTORY(variants) diff --git a/src/enumeration.cpp b/src/enumeration.cpp index 861ea42..4e9c7be 100644 --- a/src/enumeration.cpp +++ b/src/enumeration.cpp @@ -145,6 +145,19 @@ void generate_enumeration(py::module &kokkos) { std::make_index_sequence{}); }(); + auto _get_execspace_name = [](int idx) { + return get_enumeration( + idx, std::make_index_sequence{}); + }; + auto _get_execspace_idx = [](std::string str) { + return get_enumeration( + str, std::make_index_sequence{}); + }; + kokkos.def("get_execution_space", _get_execspace_name, + "Get the execution space"); + kokkos.def("get_execution_space", _get_execspace_idx, + "Get the execution space"); + //----------------------------------------------------------------------------// // // data types diff --git a/src/libpykokkos.cpp b/src/libpykokkos.cpp index 4683d46..118caf0 100644 --- a/src/libpykokkos.cpp +++ b/src/libpykokkos.cpp @@ -112,4 +112,5 @@ PYBIND11_MODULE(libpykokkos, kokkos) { generate_view_variants(kokkos); generate_atomic_variants(kokkos); generate_backend_versions(kokkos); + generate_pool_variants(kokkos); } diff --git a/src/pool_variants/CMakeLists.txt b/src/pool_variants/CMakeLists.txt new file mode 100644 index 0000000..dfa58df --- /dev/null +++ b/src/pool_variants/CMakeLists.txt @@ -0,0 +1,60 @@ + +SET(_functions) + +IF(CMAKE_UNITY_BUILD) + SET(CMAKE_UNITY_BUILD_MODE GROUP) +ENDIF() + +ADD_LIBRARY(libpykokkos-pool-variants OBJECT) + +TARGET_LINK_LIBRARIES(libpykokkos-pool-variants PUBLIC + pybind11::pybind11 + Kokkos::kokkos + libpykokkos::precompiled-headers + libpykokkos::build-options) + +SET(_types XorShift64 XorShift1024) + +MACRO(ADD_VARIANT TYPE_VARIANT) + STRING(TOLOWER "${TYPE_VARIANT}_pool" _TAG) + SET(FUNC "generate_${_TAG}") + IF(NOT ENABLE_QUIET) + MESSAGE(STATUS "Generating '${_TAG}.cpp'...") + ENDIF() + CONFIGURE_FILE( + ${CMAKE_CURRENT_SOURCE_DIR}/variant.cpp.in + ${CMAKE_CURRENT_BINARY_DIR}/${_TAG}.cpp + @ONLY) + SET_SOURCE_FILES_PROPERTIES(${CMAKE_CURRENT_BINARY_DIR}/${_TAG}.cpp + PROPERTIES + UNITY_GROUP "${TYPE_VARIANT}") + TARGET_SOURCES(libpykokkos-pool-variants PUBLIC + ${CMAKE_CURRENT_BINARY_DIR}/${_TAG}.cpp) + LIST(APPEND _functions "${FUNC}") +ENDMACRO() + +# slightly different iterations schemes to improve grouping for unity builds +FOREACH(TYPE_VARIANT ${_types}) + ADD_VARIANT(${TYPE_VARIANT}) +ENDFOREACH() + +FOREACH(_FUNC ${_functions}) + SET(PROTOTYPES "${PROTOTYPES}void ${_FUNC}(py::module&);\n") + SET(INVOCATIONS "${INVOCATIONS} ${_FUNC}(kokkos);\n") +ENDFOREACH() + +CONFIGURE_FILE( + ${CMAKE_CURRENT_SOURCE_DIR}/pool.cpp.in + ${CMAKE_CURRENT_BINARY_DIR}/pools.cpp + @ONLY) + +FILE(GLOB EXISTING_SOURCES + ${CMAKE_CURRENT_SOURCE_DIR}/*.cpp + ${PROJECT_SOURCE_DIR}/include/pool_variants/*.hpp) + +TARGET_SOURCES(libpykokkos-core PUBLIC + ${EXISTING_SOURCES} + ${CMAKE_CURRENT_BINARY_DIR}/pools.cpp) + +TARGET_SOURCES(libpykokkos PUBLIC + $) diff --git a/src/pool_variants/pool.cpp.in b/src/pool_variants/pool.cpp.in new file mode 100644 index 0000000..65e400d --- /dev/null +++ b/src/pool_variants/pool.cpp.in @@ -0,0 +1,11 @@ +// clang-format off + +#include "common.hpp" +#include "fwd.hpp" +#include "defines.hpp" + +@PROTOTYPES@ + +void generate_pool_variants(py::module &kokkos) { +@INVOCATIONS@ +} diff --git a/src/pool_variants/variant.cpp.in b/src/pool_variants/variant.cpp.in new file mode 100644 index 0000000..72cd053 --- /dev/null +++ b/src/pool_variants/variant.cpp.in @@ -0,0 +1,9 @@ +// clang-format off + +#include "fwd.hpp" +#include "pool_variants/@TYPE_VARIANT@_pool.hpp" + +void @FUNC@(py::module &kokkos) { + generate_@TYPE_VARIANT@_pool_variant<@ENUM@>( + kokkos, std::make_index_sequence{}); +}