From 8eea4d3b2788223600b5682ad22a3e825a7a473f Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Wed, 30 Sep 2026 11:15:07 -0700 Subject: [PATCH 01/24] Add device RNG Generator API on oneMKL Philox Add a `_rng_impl` pybind11 extension that samples on the device with oneMKL's device Philox4x32x10 engine and SYCL ports of NumPy's distribution algorithms, and expose it as `dpnp.random.Generator`, `dpnp.random.Philox4x3210` and `dpnp.random.default_rng`. Each work-item owns a disjoint block of Philox counters and the counter is advanced on the host, so a seed gives the same stream on any device. The Generator implements random, uniform, standard_normal, normal, lognormal, standard_exponential, exponential, standard_gamma, gamma, chisquare, beta, integers and bytes; other methods raise NotImplementedError. The legacy RNG is left in place for now. --- dpnp/CMakeLists.txt | 1 + dpnp/libtensor/extensions/rng/CMakeLists.txt | 51 ++ dpnp/libtensor/extensions/rng/common.hpp | 220 +++++ .../extensions/rng/distributions.hpp | 42 + dpnp/libtensor/extensions/rng/exponential.cpp | 45 + dpnp/libtensor/extensions/rng/gamma.cpp | 50 ++ dpnp/libtensor/extensions/rng/integers.cpp | 45 + .../extensions/rng/kernels/distributions.hpp | 471 +++++++++++ .../extensions/rng/kernels/engines.hpp | 131 +++ .../extensions/rng/kernels/sample.hpp | 88 ++ dpnp/libtensor/extensions/rng/normal.cpp | 53 ++ dpnp/libtensor/extensions/rng/rng_py.cpp | 58 ++ dpnp/libtensor/extensions/rng/sample.cpp | 149 ++++ dpnp/libtensor/extensions/rng/uniform.cpp | 49 ++ dpnp/random/__init__.py | 7 + dpnp/random/_bit_generator.py | 385 +++++++++ dpnp/random/_generator.py | 781 ++++++++++++++++++ dpnp/tests/test_random_generator.py | 533 ++++++++++++ 18 files changed, 3159 insertions(+) create mode 100644 dpnp/libtensor/extensions/rng/CMakeLists.txt create mode 100644 dpnp/libtensor/extensions/rng/common.hpp create mode 100644 dpnp/libtensor/extensions/rng/distributions.hpp create mode 100644 dpnp/libtensor/extensions/rng/exponential.cpp create mode 100644 dpnp/libtensor/extensions/rng/gamma.cpp create mode 100644 dpnp/libtensor/extensions/rng/integers.cpp create mode 100644 dpnp/libtensor/extensions/rng/kernels/distributions.hpp create mode 100644 dpnp/libtensor/extensions/rng/kernels/engines.hpp create mode 100644 dpnp/libtensor/extensions/rng/kernels/sample.hpp create mode 100644 dpnp/libtensor/extensions/rng/normal.cpp create mode 100644 dpnp/libtensor/extensions/rng/rng_py.cpp create mode 100644 dpnp/libtensor/extensions/rng/sample.cpp create mode 100644 dpnp/libtensor/extensions/rng/uniform.cpp create mode 100644 dpnp/random/_bit_generator.py create mode 100644 dpnp/random/_generator.py create mode 100644 dpnp/tests/test_random_generator.py diff --git a/dpnp/CMakeLists.txt b/dpnp/CMakeLists.txt index e96cf5bbc5a9..116859a4166d 100644 --- a/dpnp/CMakeLists.txt +++ b/dpnp/CMakeLists.txt @@ -229,6 +229,7 @@ add_subdirectory(libtensor/source) add_subdirectory(libtensor/extensions/blas) add_subdirectory(libtensor/extensions/fft) add_subdirectory(libtensor/extensions/lapack) +add_subdirectory(libtensor/extensions/rng) add_subdirectory(libtensor/extensions/vm) add_subdirectory(libtensor/extensions/sparse) diff --git a/dpnp/libtensor/extensions/rng/CMakeLists.txt b/dpnp/libtensor/extensions/rng/CMakeLists.txt new file mode 100644 index 000000000000..89fde0707140 --- /dev/null +++ b/dpnp/libtensor/extensions/rng/CMakeLists.txt @@ -0,0 +1,51 @@ +# -*- coding: utf-8 -*- +# ***************************************************************************** +# Copyright (c) 2026, Intel Corporation +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# - Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# - 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. +# - Neither the name of the copyright holder nor the names of its contributors +# may be used to endorse or promote products derived from this software +# without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "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 THE COPYRIGHT HOLDER OR 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. +# ***************************************************************************** + +set(python_module_name _rng_impl) +set(_module_src + ${CMAKE_CURRENT_SOURCE_DIR}/exponential.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/gamma.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/integers.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/normal.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/rng_py.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/sample.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/uniform.cpp +) + +build_dpnp_pybind11_ext( + ${python_module_name} + SOURCES ${_module_src} +) + +# oneMKL device RNG API is header-only, so no oneMKL library is linked +target_include_directories( + ${python_module_name} + SYSTEM + PRIVATE $ +) diff --git a/dpnp/libtensor/extensions/rng/common.hpp b/dpnp/libtensor/extensions/rng/common.hpp new file mode 100644 index 000000000000..7cf3f262cf31 --- /dev/null +++ b/dpnp/libtensor/extensions/rng/common.hpp @@ -0,0 +1,220 @@ +//***************************************************************************** +// Copyright (c) 2026, Intel Corporation +// All rights reserved. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// - Redistributions of source code must retain the above copyright notice, +// this list of conditions and the following disclaimer. +// - 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. +// - Neither the name of the copyright holder nor the names of its contributors +// may be used to endorse or promote products derived from this software +// without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "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 THE COPYRIGHT HOLDER OR 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. +//***************************************************************************** + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include + +#include "dpnp4pybind11.hpp" +#include +#include + +#include "kernels/distributions.hpp" +#include "kernels/engines.hpp" +#include "kernels/sample.hpp" + +// utils extension header +#include "ext/common.hpp" + +// dpnp tensor headers +#include "utils/output_validation.hpp" +#include "utils/type_dispatch.hpp" +#include "utils/type_utils.hpp" + +namespace dpnp::extensions::rng +{ +namespace py = pybind11; +namespace td_ns = dpnp::tensor::type_dispatch; +namespace rng_ns = dpnp::kernels::rng; + +// A distribution parameter as passed from Python: either a C-contiguous +// usm_ndarray with the data type and the size of the result, whose data +// pointer is `data`, or a Python scalar. +struct ParamArg +{ + const char *data = nullptr; + py::object scalar; +}; + +typedef sycl::event (*sample_fn_ptr_t)(sycl::queue &, + const rng_ns::EngineState &, + const std::size_t, + char *, + const ParamArg &, + const ParamArg &, + const std::vector &); + +using sample_dispatch_table_t = + sample_fn_ptr_t[rng_ns::num_engines][td_ns::num_types]; + +namespace detail +{ +template +rng_ns::Param to_param(const ParamArg &arg) +{ + if (arg.data != nullptr) { + return {reinterpret_cast(arg.data), T{}}; + } + if (arg.scalar.is_none()) { + return {nullptr, T{}}; + } + return {nullptr, py::cast(arg.scalar)}; +} +} // namespace detail + +template class SamplerT> +sycl::event sample_impl(sycl::queue &exec_q, + const rng_ns::EngineState &state, + const std::size_t nelems, + char *result, + const ParamArg &a, + const ParamArg &b, + const std::vector &depends) +{ + dpnp::tensor::type_utils::validate_type_for_device(exec_q); + + T *res = reinterpret_cast(result); + const SamplerT sampler(detail::to_param(a), detail::to_param(b)); + const std::size_t n_wi = rng_ns::num_work_items(nelems); + + sycl::event sample_ev = exec_q.submit([&](sycl::handler &cgh) { + cgh.depends_on(depends); + + using SampleKernel = + rng_ns::SampleFunctor>; + cgh.parallel_for( + sycl::range<1>(n_wi), + SampleKernel(state, sampler, res, nelems, n_wi)); + }); + + return sample_ev; +} + +// Types the floating-point samplers support +template +struct RealTypes + : std::bool_constant || std::is_same_v> +{ +}; + +// Types the integer samplers support +template +struct IntegerTypes : std::bool_constant> +{ +}; + +// The raw output of a bit generator is always 64 bits wide +template +struct RawTypes : std::bool_constant> +{ +}; + +template class SamplerT, + template class SupportedT> +struct SampleFactoryFor +{ + template + struct type + { + fnT get() + { + if constexpr (SupportedT::value) { + return sample_impl; + } + else { + return nullptr; + } + } + }; +}; + +template