Exposed vector size of an engine

Author: antonwolfy

dpnp/backend/extensions/rng/device/dispatch/table_builder.hpp

@@ -77,12 +77,13 @@ class Dispatch3DTableBuilder
     Dispatch3DTableBuilder() = default;
     ~Dispatch3DTableBuilder() = default;
 
-    void populate(funcPtrT table[][_no_of_types][_no_of_methods]) const
+    template <std::uint8_t... VecSizes>
+    void populate(funcPtrT table[][_no_of_types][_no_of_methods], std::integer_sequence<std::uint8_t, VecSizes...>) const
     {
-        const auto map_by_engine = {table_per_type_and_method<mkl_rng_dev::mrg32k3a<8>>(),
-                                    table_per_type_and_method<mkl_rng_dev::philox4x32x10<8>>(),
-                                    table_per_type_and_method<mkl_rng_dev::mcg31m1<8>>(),
-                                    table_per_type_and_method<mkl_rng_dev::mcg59<8>>()};
+        const auto map_by_engine = {table_per_type_and_method<mkl_rng_dev::mrg32k3a<VecSizes>>()...,
+                                    table_per_type_and_method<mkl_rng_dev::philox4x32x10<VecSizes>>()...,
+                                    table_per_type_and_method<mkl_rng_dev::mcg31m1<VecSizes>>()...,
+                                    table_per_type_and_method<mkl_rng_dev::mcg59<VecSizes>>()...};
         assert(map_by_engine.size() == _no_of_engines);
 
         std::uint16_t engine_id = 0;

dpnp/backend/extensions/rng/device/engine/builder/base_builder.hpp

@@ -111,6 +111,7 @@ class BaseBuilder {
 
     // TODO: remove
     void print() {
+        std::cout << "vector size = " << std::to_string(EngineT::vec_size) << std::endl;
         std::cout << "list_of_seeds: ";
         for (auto &val: seeds) {
             std::cout << std::to_string(val) << ", ";

dpnp/backend/extensions/rng/device/gaussian.cpp

@@ -26,6 +26,7 @@
 #include <pybind11/pybind11.h>
 
 // dpctl tensor headers
+#include "utils/output_validation.hpp"
 #include "utils/type_dispatch.hpp"
 #include "utils/type_utils.hpp"
 #include "kernels/alignment.hpp"
@@ -51,7 +52,22 @@ using dpctl_krn_ns::disabled_sg_loadstore_wrapper_krn;
 using dpctl_krn_ns::is_aligned;
 using dpctl_krn_ns::required_alignment;
 
-constexpr int no_of_methods = 2; // number of methods of gaussian distribution
+constexpr auto no_of_methods = 2; // number of methods of gaussian distribution
+
+constexpr auto seq_of_vec_sizes = std::integer_sequence<std::uint8_t, 2, 4, 8, 16>{};
+constexpr auto vec_sizes_len = seq_of_vec_sizes.size();
+constexpr auto no_of_engines = engine::no_of_engines * vec_sizes_len;
+
+template <typename VecSizeT, VecSizeT ...Ints, auto ...Indices>
+inline auto find_vec_size_impl(const VecSizeT vec_size, std::index_sequence<Indices...>) {
+    return std::min({ ((Ints == vec_size) ? Indices : sizeof...(Indices))... });
+}
+
+template <typename VecSizeT, VecSizeT ...Ints>
+int find_vec_size(const VecSizeT vec_size, std::integer_sequence<VecSizeT, Ints...>) {
+    auto res = find_vec_size_impl<VecSizeT, Ints...>(vec_size, std::make_index_sequence<sizeof...(Ints)>{}); 
+    return (res == sizeof...(Ints)) ? -1 : res;
+}
 
 template <typename DataT, typename Method>
 struct DistributorBuilder
@@ -83,7 +99,7 @@ typedef sycl::event (*gaussian_impl_fn_ptr_t)(engine::EngineBase *engine,
                                               char *,
                                               const std::vector<sycl::event> &);
 
-static gaussian_impl_fn_ptr_t gaussian_dispatch_table[engine::no_of_engines][dpctl_td_ns::num_types][no_of_methods];
+static gaussian_impl_fn_ptr_t gaussian_dispatch_table[no_of_engines][dpctl_td_ns::num_types][no_of_methods];
 
 template <typename EngineT, typename DataT,  typename Method, unsigned int items_per_wi>
 class gaussian_kernel;
@@ -117,7 +133,7 @@ static sycl::event gaussian_impl(engine::EngineBase *engine,
 
             using EngineBuilderT = engine::builder::Builder<EngineT>;
             EngineBuilderT eng_builder(engine);
-            eng_builder.print(); // TODO: remove
+            // eng_builder.print(); // TODO: remove
 
             using DistributorBuilderT = DistributorBuilder<DataT, Method>;
             DistributorBuilderT dist_builder(mean, stddev);
@@ -154,6 +170,7 @@ static sycl::event gaussian_impl(engine::EngineBase *engine,
 
 std::pair<sycl::event, sycl::event> gaussian(engine::EngineBase *engine,
                                              const std::uint8_t method_id,
+                                             const std::uint8_t vec_size,
                                              const double mean,
                                              const double stddev,
                                              const std::uint64_t n,
@@ -176,15 +193,10 @@ std::pair<sycl::event, sycl::event> gaussian(engine::EngineBase *engine,
     }
 
     // ensure that output is ample enough to accommodate all elements
-    auto res_offsets = res.get_minmax_offsets();
-    // destination must be ample enough to accommodate all elements
-    {
-        size_t range =
-            static_cast<size_t>(res_offsets.second - res_offsets.first);
-        if (range + 1 < res_nelems) {
-            throw py::value_error(
-                "Destination array can not accommodate all the elements of source array.");
-        }
+    dpctl::tensor::validation::AmpleMemory::throw_if_not_ample(res, res_nelems);
+
+    if (!dpctl::utils::queues_are_compatible(exec_q, {res})) {
+        throw py::value_error("Execution queue is not compatible with the allocation queue");
     }
 
     bool is_res_c_contig = res.is_c_contiguous();
@@ -201,6 +213,12 @@ std::pair<sycl::event, sycl::event> gaussian(engine::EngineBase *engine,
         throw std::runtime_error("Unknown method=" + std::to_string(method_id) + " for gaussian distribution.");
     }
 
+    int vec_size_id = find_vec_size(vec_size, seq_of_vec_sizes);
+    if (vec_size_id < 0) {
+        throw std::runtime_error("Vector size=" + std::to_string(vec_size) + " is out of supported range");
+    }
+    enginge_id = enginge_id * vec_sizes_len + vec_size_id;
+
     auto array_types = dpctl_td_ns::usm_ndarray_types();
     int res_type_id = array_types.typenum_to_lookup_id(res.get_typenum());
 
@@ -232,7 +250,7 @@ struct GaussianContigFactory
 
 void init_gaussian_dispatch_3d_table(void)
 {
-    dispatch::Dispatch3DTableBuilder<gaussian_impl_fn_ptr_t, GaussianContigFactory, engine::no_of_engines, dpctl_td_ns::num_types, no_of_methods> contig;
-    contig.populate(gaussian_dispatch_table);
+    dispatch::Dispatch3DTableBuilder<gaussian_impl_fn_ptr_t, GaussianContigFactory, no_of_engines, dpctl_td_ns::num_types, no_of_methods> contig;
+    contig.populate(gaussian_dispatch_table, seq_of_vec_sizes);
 }
 } // dpnp::backend::ext::rng::device

dpnp/backend/extensions/rng/device/gaussian.hpp

@@ -34,6 +34,7 @@ namespace dpnp::backend::ext::rng::device
 {
 extern std::pair<sycl::event, sycl::event> gaussian(engine::EngineBase *engine,
                                                     const std::uint8_t method_id,
+                                                    const std::uint8_t vec_size,
                                                     const double mean,
                                                     const double stddev,
                                                     const std::uint64_t n,

dpnp/backend/extensions/rng/device/rng_py.cpp

@@ -100,7 +100,8 @@ PYBIND11_MODULE(_rng_dev_impl, m)
     m.def("_gaussian", &rng_dev_ext::gaussian,
           "",
           py::arg("engine"),
-          py::arg("method"), py::arg("mean"), py::arg("stddev"),
+          py::arg("method_id"), py::arg("vec_size"),
+          py::arg("mean"), py::arg("stddev"),
           py::arg("n"), py::arg("res"),
           py::arg("depends") = py::list());
 }