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 <[email protected]>

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();

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 ([email protected])
+//
+// ************************************************************************
+//@HEADER
+*/
+
+#pragma once
+
+#include "common.hpp"
+#include "concepts.hpp"
+#include "pools.hpp"
+#include "traits.hpp"
+
+namespace Space {
+namespace SpaceDim {
+
+template <size_t SpaceIdx>
+void generate_XorShift1024_pool_variant(py::module &_mod) {
+  using space_spec_t = ExecutionSpaceSpecialization<SpaceIdx>;
+  using Sp           = typename space_spec_t::type;
+  using UniformT     = Kokkos::Random_XorShift1024_Pool<Sp>;
+
+  auto name = join("_", "KokkosXorShift1024Pool", space_spec_t::label());
+
+  Common::generate_pool<UniformT, Sp>(_mod, name, demangle<UniformT>());
+}
+}  // namespace SpaceDim
+
+template <size_t SpaceIdx>
+void generate_XorShift1024_pool_variant(
+    py::module &,
+    std::enable_if_t<!is_available<execution_space_t<SpaceIdx>>::value, int> =
+        0) {}
+
+template <size_t SpaceIdx>
+void generate_XorShift1024_pool_variant(
+    py::module &_mod,
+    std::enable_if_t<is_available<execution_space_t<SpaceIdx>>::value, int> =
+        0) {
+  SpaceDim::generate_XorShift1024_pool_variant<SpaceIdx>(_mod);
+}
+}  // namespace Space
+
+namespace {
+// generate data-type, memory-space buffers for concrete dimension
+template <size_t... SpaceIdx>
+void generate_XorShift1024_pool_variant(py::module &_mod,
+                                        std::index_sequence<SpaceIdx...>) {
+  FOLD_EXPRESSION(Space::generate_XorShift1024_pool_variant<SpaceIdx>(_mod));
+}
+}  // namespace

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 ([email protected])
+//
+// ************************************************************************
+//@HEADER
+*/
+
+#pragma once
+
+#include "common.hpp"
+#include "concepts.hpp"
+#include "pools.hpp"
+#include "traits.hpp"
+
+namespace Space {
+namespace SpaceDim {
+
+template <size_t SpaceIdx>
+void generate_XorShift64_pool_variant(py::module &_mod) {
+  using space_spec_t = ExecutionSpaceSpecialization<SpaceIdx>;
+  using Sp           = typename space_spec_t::type;
+  using UniformT     = Kokkos::Random_XorShift64_Pool<Sp>;
+
+  auto name = join("_", "KokkosXorShift64Pool", space_spec_t::label());
+
+  Common::generate_pool<UniformT, Sp>(_mod, name, demangle<UniformT>());
+}
+}  // namespace SpaceDim
+
+template <size_t SpaceIdx>
+void generate_XorShift64_pool_variant(
+    py::module &,
+    std::enable_if_t<!is_available<execution_space_t<SpaceIdx>>::value, int> =
+        0) {}
+
+template <size_t SpaceIdx>
+void generate_XorShift64_pool_variant(
+    py::module &_mod,
+    std::enable_if_t<is_available<execution_space_t<SpaceIdx>>::value, int> =
+        0) {
+  SpaceDim::generate_XorShift64_pool_variant<SpaceIdx>(_mod);
+}
+}  // namespace Space
+
+namespace {
+// generate data-type, memory-space buffers for concrete dimension
+template <size_t... SpaceIdx>
+void generate_XorShift64_pool_variant(py::module &_mod,
+                                      std::index_sequence<SpaceIdx...>) {
+  FOLD_EXPRESSION(Space::generate_XorShift64_pool_variant<SpaceIdx>(_mod));
+}
+}  // namespace

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 ([email protected])
+//
+// ************************************************************************
+//@HEADER
+*/
+
+#pragma once
+
+#include <Kokkos_Core.hpp>
+#include <Kokkos_Random.hpp>
+
+namespace Common {
+template <typename PoolT, typename Sp>
+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<Kokkos::Cuda>;
+  // class decl
+  py::class_<PoolT> _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

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

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)

src/CMakeLists.txt

@@ -1,2 +1,3 @@
 
+ADD_SUBDIRECTORY(pool_variants)
 ADD_SUBDIRECTORY(variants)

src/enumeration.cpp

@@ -145,6 +145,19 @@ void generate_enumeration(py::module &kokkos) {
         std::make_index_sequence<ExecutionSpacesEnd>{});
   }();
 
+  auto _get_execspace_name = [](int idx) {
+    return get_enumeration<ExecutionSpaceSpecialization>(
+        idx, std::make_index_sequence<ExecutionSpacesEnd>{});
+  };
+  auto _get_execspace_idx = [](std::string str) {
+    return get_enumeration<ViewDataTypeSpecialization>(
+        str, std::make_index_sequence<ExecutionSpacesEnd>{});
+  };
+  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

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);
 }

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
+    $<TARGET_OBJECTS:libpykokkos-pool-variants>)

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@
+}