From 6c3ff46b95eefeacd709c82bdeaba2f04c011407 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Wed, 4 Nov 2026 11:46:54 -0800 Subject: [PATCH] Add dpnp.random.Generator on the oneMKL device RNG API Add the dpnp.libtensor._rng_impl extension, which samples on the device from oneMKL's device philox4x32x10 engine, and the Generator API of NumPy on top of it: BitGenerator, Philox4x3210, SeedSequence (NumPy's), default_rng, and the random, uniform, standard_normal and normal methods. A bit generator keeps its state on the host and every launch rebuilds its engines from it, so the values depend only on the seed and the output size, not on the device. The legacy RandomState and module functions are unchanged. --- CHANGELOG.md | 1 + doc/reference/random.rst | 91 ++- dpnp/AGENTS.md | 2 + dpnp/CMakeLists.txt | 1 + dpnp/libtensor/AGENTS.md | 5 +- dpnp/libtensor/extensions/rng/CMakeLists.txt | 48 ++ dpnp/libtensor/extensions/rng/common.hpp | 284 +++++++++ .../extensions/rng/distributions.hpp | 39 ++ .../extensions/rng/kernels/distributions.hpp | 114 ++++ .../extensions/rng/kernels/engines.hpp | 202 ++++++ .../extensions/rng/kernels/sample.hpp | 133 ++++ .../extensions/rng/kernels/samplers.hpp | 134 ++++ dpnp/libtensor/extensions/rng/normal.cpp | 45 ++ dpnp/libtensor/extensions/rng/rng_py.cpp | 59 ++ dpnp/libtensor/extensions/rng/sample.cpp | 166 +++++ dpnp/libtensor/extensions/rng/uniform.cpp | 50 ++ dpnp/random/__init__.py | 7 + dpnp/random/_bit_generator.py | 391 ++++++++++++ dpnp/random/_generator.py | 575 ++++++++++++++++++ dpnp/tests/AGENTS.md | 1 + dpnp/tests/test_random_generator.py | 433 +++++++++++++ dpnp/tests/test_sycl_queue.py | 73 +++ dpnp/tests/test_usm_type.py | 33 + 23 files changed, 2878 insertions(+), 9 deletions(-) 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/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/kernels/samplers.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/CHANGELOG.md b/CHANGELOG.md index 612b04ceac39..e811916e1b4c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * Added support for boolean input in `dpnp.positive` and the `+` operator of `dpnp.ndarray` [#3128](https://github.com/IntelPython/dpnp/pull/3128) * Added `dpnp.top_k` to return the `k` largest or smallest elements and their indices along an axis, with a `sorted` keyword to return them in sorted order [#3133](https://github.com/IntelPython/dpnp/pull/3133) +* Added the `dpnp.random.Generator` API with `default_rng`, `BitGenerator`, `SeedSequence` and the `Philox4x3210` bit generator, sampling on the device with the oneMKL device RNG API [#3157](https://github.com/IntelPython/dpnp/pull/3157) ### Changed diff --git a/doc/reference/random.rst b/doc/reference/random.rst index 55ac0271307a..d96b52654d27 100644 --- a/doc/reference/random.rst +++ b/doc/reference/random.rst @@ -3,10 +3,90 @@ Random sampling (:mod:`dpnp.random`) ==================================== +.. hint:: `NumPy API Reference: Random sampling `_ + +DPNP draws random numbers on the SYCL device, with the counter-based +Philox4x32x10 generator of the oneMKL device RNG API. It provides the +:class:`Generator` API of NumPy, recommended for new code. The legacy +:class:`RandomState` API and the module-level functions still use the +host-side oneMKL RNG. + +.. code-block:: python + + import dpnp + + rng = dpnp.random.default_rng(seed=42) + x = rng.standard_normal((3, 4)) # on the default device + y = rng.random(5, usm_type="shared") + + rng = dpnp.random.default_rng(42, device="cpu") + +A generator draws all its samples on one SYCL queue, chosen with the +``device`` or ``sycl_queue`` keyword when it is created, and the samples +depend only on the seed, not on the device. All samples are +:class:`dpnp.ndarray`, zero-dimensional when ``size`` is ``None``. + +Differences from NumPy +---------------------- + +* The same seed gives the same stream on every device. + The streams differ from NumPy's. +* The floating-point distributions take a ``dtype`` keyword, as in CuPy, + which defaults to the default floating-point type of the device. All + methods take a ``usm_type`` keyword. +* Parameters on the host (Python scalars, lists and NumPy arrays) are + checked as in NumPy. Parameters in :class:`dpnp.ndarray` are + not, since that would wait for the device. + Invalid ones give undefined values, such as NaN. + +Random Generator +---------------- + +.. autosummary:: + :toctree: generated/ + :nosignatures: + + default_rng + Generator + +Simple random data +^^^^^^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + :nosignatures: + + Generator.random + +Distributions +^^^^^^^^^^^^^ + +.. autosummary:: + :toctree: generated/ + :nosignatures: + + Generator.normal + Generator.standard_normal + Generator.uniform + +Bit generators +-------------- + +.. autosummary:: + :toctree: generated/ + :nosignatures: + + BitGenerator + Philox4x3210 + SeedSequence + +Legacy random generation +------------------------ + .. hint:: `NumPy API Reference: Legacy random generation `_ Simple random data ------------------- +^^^^^^^^^^^^^^^^^^ .. autosummary:: :toctree: generated/ @@ -23,9 +103,8 @@ Simple random data choice bytes - Permutations ------------- +^^^^^^^^^^^^ .. autosummary:: :toctree: generated/ @@ -34,9 +113,8 @@ Permutations shuffle permutation - Distributions -------------- +^^^^^^^^^^^^^ .. autosummary:: :toctree: generated/ @@ -78,9 +156,8 @@ Distributions weibull zipf - Random generator ----------------- +^^^^^^^^^^^^^^^^ .. autosummary:: :toctree: generated/ diff --git a/dpnp/AGENTS.md b/dpnp/AGENTS.md index b7d3956d4134..7efb8ae7b9b7 100644 --- a/dpnp/AGENTS.md +++ b/dpnp/AGENTS.md @@ -6,6 +6,7 @@ - `dpnp_algo/dpnp_elementwise_common.py`: ufunc wrapper classes (`DPNPUnaryFunc`, `DPNPBinaryFunc`, two-output and specialised variants) - `dpnp_utils/`: shared helpers (`get_usm_allocations`, `map_dtype_to_device`, reduction wrappers) - `linalg/`, `fft/`, `random/`, `scipy/`, `lib/`: subpackages, each exporting through its own `__init__.py` +- `random/`: `_bit_generator.py` (`BitGenerator`, `Philox4x3210`; `SeedSequence` is NumPy's) and `_generator.py` (`Generator`, which implements every distribution). The kernels are in `dpnp.libtensor._rng_impl`. - Legacy, don't extend: the Cython `dpnp_algo` path and `dpnp_backend_c` (still used by `dpnp.partition` and `dpnp.random`), and the NumPy fallback `call_origin`. New kernels go into `dpnp/libtensor` (see `dpnp/libtensor/AGENTS.md`). ## Adding a Public Function @@ -20,6 +21,7 @@ - Extension calls are ordered with `_manager = dpctl.utils.SequentialOrderManager[exec_q]`: pass `depends=_manager.submitted_events`, then `_manager.add_event_pair(ht_ev, ev)`. Don't add blocking waits. - Use `dpnp.default_float_type(...)` or `map_dtype_to_device(...)` for the default float type. - Deprecated ufuncs use `DPNPDeprecatedUnaryFunc(..., deprecated_msg=...)`. +- A `dpnp.random` distribution is a `Generator` method: it calls `self._check_queues(...)` on its array parameters (raises `ExecutionPlacementError`), checks host parameters as NumPy does, and draws through `_cont` or `_float_sample` into `BitGenerator._fill`. Parameters on the device are not checked, as the check waits for the device. Temporaries take the output's `usm_type`. - Unsupported NumPy keywords raise `NotImplementedError` (`dpnp.check_limitations`) and are listed under `Limitations` in the docstring. - The Array API spec version is set in three places that change together: `__array_api_version__` in `dpnp/tensor/_array_api.py`, `ARRAY_API_TESTS_VERSION` in `.github/workflows/conda-package.yml` and `doc/reference/array_api.rst` (plus the version test in `dpnp/tests/test_ndarray.py`). Inspection utilities come from `__array_namespace_info__` in `dpnp_array_api_info.py`. - `dpnp.tensor` replaces `dpctl.tensor`, which no longer exists here; `SyclQueue`, `SyclDevice`, `dpctl.memory` and `SequentialOrderManager` still come from `dpctl`. 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/AGENTS.md b/dpnp/libtensor/AGENTS.md index c6cec65c2044..556cd6bf62fa 100644 --- a/dpnp/libtensor/AGENTS.md +++ b/dpnp/libtensor/AGENTS.md @@ -3,7 +3,7 @@ ## Layout - `include/kernels/`: header-only SYCL kernels in namespaces under `dpnp::tensor::kernels`; `include/utils/`: type dispatch, type, offset and memory-overlap helpers (`type_dispatch.hpp`, `type_dispatch_building.hpp`, `type_utils.hpp`, `memory_overlap.hpp`, ...); `include/ext/`: dispatch helpers (`init_dispatch_vector`, `init_dispatch_table`, `DispatchTable`) and validation helpers (`validation_utils.hpp`). The headers ship with dpnp; `dpnp.get_include()` and `python -m dpnp --tensor-include-dir` return `include/`. - `source/`: bindings built by `source/CMakeLists.txt` into `_tensor_impl`, `_tensor_elementwise_impl`, `_tensor_reductions_impl`, `_tensor_sorting_impl`, `_tensor_accumulation_impl`, `_tensor_linalg_impl`, `_statistics_impl` and `_window_impl`. The functions dpnp adds beyond the Array API are listed in the `_dpnp_*_sources` variables there and compiled into the same modules. -- `extensions/{blas,fft,lapack,sparse,vm}/`: oneMKL/oneMath extensions, one module `__impl` per directory. +- `extensions/{blas,fft,lapack,rng,sparse,vm}/`: oneMKL/oneMath extensions, one module `__impl` per directory. - Every module installs into `dpnp/libtensor` and is imported as `dpnp.libtensor.__impl`. A new kernel goes into `source/` and an existing module; only code that needs oneMKL gets a new extension. ## Adding a ufunc Kernel (model: `fabs`) @@ -21,8 +21,9 @@ A new extension's `CMakeLists.txt` lists its sources and calls `build_dpnp_pybin - Validate inputs before launching: queue compatibility, contiguity, writability, memory overlap (including the same logical tensor), dtype and ndim. Reuse the helpers neighbouring files use: `check_queue`, `check_c_contig`, `check_no_overlap`, `check_writable` from `include/ext/validation_utils.hpp` (`putmask`, `statistics`, `isclose`, `interpolate`), or `dpnp::utils::queues_are_compatible`, `dpnp::tensor::overlap::MemoryOverlap` and `dpnp::tensor::validation::CheckWritable` (extensions and the Array API functions). Don't silently accept strided input. - Throw `py::value_error` for bad arguments and `std::runtime_error` for internal failures; catch oneMKL and SYCL exceptions and rethrow them. - Guard double, complex and half with `dpnp::tensor::type_utils::validate_type_for_device(q)`; don't use double literals or casts in kernels that must run without fp64. -- `blas`, `fft`, `lapack` and `sparse` build against either oneMKL or oneMath (`USE_ONEMATH`); `vm` has kernels only with oneMKL: with oneMath it builds empty and `_is_available()` returns `False`. +- `blas`, `fft`, `lapack` and `sparse` build against either oneMKL or oneMath (`USE_ONEMATH`); `vm` has kernels only with oneMKL: with oneMath it builds empty and `_is_available()` returns `False`; `rng` uses the header-only oneMKL device RNG API and links no oneMKL library. - Release the GIL around blocking oneMKL calls (`py::gil_scoped_release`). +- `rng` values depend only on the engine state and the number of elements, never on the device or the layout chosen for it (`extensions/rng/kernels/sample.hpp`); each engine's `items_per_work_item` (`kernels/engines.hpp`) is part of its stream definition. - An extension whose kernels may be launched over more than `INT_MAX` work-items passes `LARGE_RANGES` to `build_dpnp_pybind11_ext`, which adds `-fno-sycl-id-queries-fit-in-int` at compile and link time (DPC++ 2026.2 and later need it at link time). `source/CMakeLists.txt` and `dpnp_backend_c` pass it to both `target_compile_options` and `target_link_options`. - clang-format 23 formats C++ (`.clang-format`: LLVM base, 4-space indent, 80 columns); gersemi formats CMake at 88 columns, except `dpnp/backend/cmake/Modules/`. diff --git a/dpnp/libtensor/extensions/rng/CMakeLists.txt b/dpnp/libtensor/extensions/rng/CMakeLists.txt new file mode 100644 index 000000000000..20f793c978df --- /dev/null +++ b/dpnp/libtensor/extensions/rng/CMakeLists.txt @@ -0,0 +1,48 @@ +# -*- 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}/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..1e72cd95017a --- /dev/null +++ b/dpnp/libtensor/extensions/rng/common.hpp @@ -0,0 +1,284 @@ +//***************************************************************************** +// 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 + +#include "dpnp4pybind11.hpp" +#include +#include + +#include "kernels/engines.hpp" +#include "kernels/sample.hpp" +#include "kernels/samplers.hpp" + +// utils extension header +#include "ext/common.hpp" + +// dpnp tensor headers +#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 size of the result, whose data pointer is `data` and +// whose type lookup id is `type_id`, or a Python scalar. +struct ParamArg +{ + const char *data = nullptr; + int type_id = -1; + py::object scalar; +}; + +typedef sycl::event (*sample_fn_ptr_t)(sycl::queue &, + const rng_ns::EngineState &, + const std::size_t, + const int, + char *, + const std::vector &, + const std::vector &); + +// Samplers of a result type, indexed by the engine and the result type +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, const std::size_t pos) +{ + if (arg.data != nullptr) { + if (arg.type_id != td_ns::GetTypeid{}.get()) { + throw py::value_error("Parameter array " + std::to_string(pos) + + " has an unexpected data type."); + } + return {reinterpret_cast(arg.data), T{}}; + } + return {nullptr, py::cast(arg.scalar)}; +} + +template +SamplerT make_sampler(const std::vector ¶ms, + std::index_sequence) +{ + using param_types = typename SamplerT::param_types; + return SamplerT( + to_param>(params[Is], Is)...); +} + +template +void validate_type_for_device(sycl::queue &exec_q) +{ + dpnp::tensor::type_utils::validate_type_for_device(exec_q); +} + +template +void validate_types_for_device(sycl::queue &exec_q, std::index_sequence) +{ + (validate_type_for_device>(exec_q), ...); +} +template +sycl::event submit_sample(sycl::queue &exec_q, + const rng_ns::EngineState &state, + const SamplerT &sampler, + typename SamplerT::result_type *res, + const std::size_t nelems, + const std::vector &depends) +{ + const std::size_t n_wi = rng_ns::num_work_items(nelems); + + return 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)); + }); +} +} // namespace detail + +// Checks the parameters and the result against SamplerT, then submits the +// kernel, unless there is nothing to sample. +template +sycl::event sample_impl(sycl::queue &exec_q, + const rng_ns::EngineState &state, + const std::size_t nelems, + const int result_type_id, + char *result, + const std::vector ¶ms, + const std::vector &depends) +{ + using T = typename SamplerT::result_type; + using param_types = typename SamplerT::param_types; + constexpr std::size_t n_params = std::tuple_size_v; + + if (result_type_id != td_ns::GetTypeid{}.get()) { + throw py::value_error("The result array has an unexpected data type."); + } + if (params.size() != n_params) { + throw py::value_error("Expected " + std::to_string(n_params) + + " parameters, got " + + std::to_string(params.size()) + "."); + } + + detail::validate_type_for_device(exec_q); + detail::validate_types_for_device( + exec_q, std::make_index_sequence{}); + + const SamplerT sampler = detail::make_sampler( + params, std::make_index_sequence{}); + + if (nelems == 0) { + return sycl::event{}; + } + + T *res = reinterpret_cast(result); + + if constexpr (rng_ns::can_draw_blocked) { + // the blocked layout gives the same values, faster on GPUs but + // slower on CPUs + if (exec_q.get_device().is_gpu()) { + return detail::submit_sample( + exec_q, state, sampler, res, nelems, depends); + } + } + return detail::submit_sample( + exec_q, state, sampler, res, nelems, depends); +} + +// Types the floating-point samplers support +template +struct RealTypes + : std::bool_constant || std::is_same_v> +{ +}; + +// The raw output of a bit generator, whole or as 32-bit halves +template +struct RawTypes : std::bool_constant || + std::is_same_v> +{ +}; + +template class SamplerT, + template class SupportedT> +struct SampleFactoryFor +{ + template + struct type + { + fnT get() + { + if constexpr (SupportedT::value) { + return sample_impl>; + } + else { + return nullptr; + } + } + }; +}; + +template