Added destribution method dispatching

antonwolfy · antonwolfy · commit 84074fd0d392 · 2024-05-06T15:04:32.000+02:00
diff --git a/dpnp/backend/extensions/rng/device/common_impl.hpp b/dpnp/backend/extensions/rng/device/common_impl.hpp
@@ -51,8 +51,8 @@ namespace mkl_rng_dev = oneapi::mkl::rng::device;
 
 /*! @brief Functor for unary function evaluation on contiguous array */
 template <typename DataT,
+          typename GaussianDistrT,
           typename ResT = DataT,
-          typename Method = mkl_rng_dev::gaussian_method::by_default,
           typename UnaryOperatorT = ResT,
           unsigned int vec_sz = 8,
           unsigned int items_per_wi = 4,
@@ -61,14 +61,14 @@ struct RngContigFunctor
 {
 private:
     const std::uint32_t seed_;
-    const DataT mean_;
-    const DataT stddev_;
+    GaussianDistrT distr_;
     ResT * const res_ = nullptr;
     const size_t nelems_;
 
 public:
-    RngContigFunctor(const std::uint32_t seed, const DataT mean, const DataT stddev, ResT *res, const size_t n_elems)
-        : seed_(seed), mean_(mean), stddev_(stddev), res_(res), nelems_(n_elems)
+
+    RngContigFunctor(const std::uint32_t seed, GaussianDistrT& distr, ResT *res, const size_t n_elems)
+        : seed_(seed), distr_(distr), res_(res), nelems_(n_elems)
     {
     }
 
@@ -80,8 +80,11 @@ struct RngContigFunctor
         const std::uint8_t sg_size = sg.get_local_range()[0];
         const std::uint8_t max_sg_size = sg.get_max_local_range()[0];
 
-        auto engine = mkl_rng_dev::mrg32k3a<vec_sz>(seed_, nelems_ * global_id);  // offset is questionable...
-        mkl_rng_dev::gaussian<DataT, Method> distr(mean_, stddev_);
+        using EngineT = typename mkl_rng_dev::mrg32k3a<vec_sz>;
+        auto engine = EngineT(seed_, nelems_ * global_id);  // offset is questionable...
+
+        using DistrT = typename GaussianDistrT::distr_type;
+        DistrT distr = distr_();
 
         if constexpr (enable_sg_load) {
             const size_t base = items_per_wi * vec_sz * (nd_it.get_group(0) * nd_it.get_local_range(0) + sg.get_group_id()[0] * max_sg_size);
@@ -92,13 +95,13 @@ struct RngContigFunctor
                     size_t offset = base + static_cast<size_t>(it) * static_cast<size_t>(sg_size);
                     auto out_multi_ptr = sycl::address_space_cast<sycl::access::address_space::global_space, sycl::access::decorated::yes>(&res_[offset]);
 
-                    sycl::vec<DataT, vec_sz> rng_val_vec = mkl_rng_dev::generate(distr, engine);
+                    sycl::vec<DataT, vec_sz> rng_val_vec = mkl_rng_dev::generate<DistrT, EngineT>(distr, engine);
                     sg.store<vec_sz>(out_multi_ptr, rng_val_vec);
                 }
             }
             else {
                 for (size_t offset = base + sg.get_local_id()[0]; offset < nelems_; offset += sg_size) {
-                    res_[offset] = mkl_rng_dev::generate_single(distr, engine);
+                    res_[offset] = mkl_rng_dev::generate_single<DistrT, EngineT>(distr, engine);
                 }
             }
         }
@@ -108,7 +111,7 @@ struct RngContigFunctor
             base = (base / sg_size) * sg_size * items_per_wi * vec_sz + (base % sg_size);
             for (size_t offset = base; offset < std::min(nelems_, base + sg_size * (items_per_wi * vec_sz)); offset += sg_size)
             {
-                res_[offset] = mkl_rng_dev::generate_single(distr, engine);
+                res_[offset] = mkl_rng_dev::generate_single<DistrT, EngineT>(distr, engine);
             }
         }
     }
diff --git a/dpnp/backend/extensions/rng/device/gaussian.cpp b/dpnp/backend/extensions/rng/device/gaussian.cpp
@@ -57,6 +57,40 @@ namespace mkl_rng_dev = oneapi::mkl::rng::device;
 namespace py = pybind11;
 namespace type_utils = dpctl::tensor::type_utils;
 
+constexpr int num_methods = 2; // number of methods of gaussian distribution
+
+// static mkl_rng_dev::gaussian_method get_method(const std::int8_t method) {
+//     switch (method) {
+//         case 0: return mkl_rng_dev::gaussian_method::by_default;
+//         case 1: return mkl_rng_dev::gaussian_method::by_default;
+//         default:
+//             throw py::value_error();
+//     }
+// }
+
+template <typename DataT, typename Method>
+struct GaussianDistr
+{
+private:
+    const DataT mean_;
+    const DataT stddev_;
+
+public:
+    using method_type = Method;
+    using result_type = DataT;
+    using distr_type = typename mkl_rng_dev::gaussian<DataT, Method>;
+
+    GaussianDistr(const DataT mean, const DataT stddev)
+        : mean_(mean), stddev_(stddev)
+    {
+    }
+
+    inline auto operator()(void) const
+    {
+        return distr_type(mean_, stddev_);
+    }
+};
+
 typedef sycl::event (*gaussian_impl_fn_ptr_t)(sycl::queue &,
                                            const std::uint32_t,
                                            const double,
@@ -65,13 +99,12 @@ typedef sycl::event (*gaussian_impl_fn_ptr_t)(sycl::queue &,
                                            char *,
                                            const std::vector<sycl::event> &);
 
-static gaussian_impl_fn_ptr_t gaussian_dispatch_vector[dpctl_td_ns::num_types];
+static gaussian_impl_fn_ptr_t gaussian_dispatch_table[dpctl_td_ns::num_types][num_methods];
 
-// template <typename DataT, typename Method = mkl_rng_dev::gaussian_method::by_default>
-template <typename DataT, unsigned int vec_sz, unsigned int items_per_wi>
+template <typename DataT,  typename Method, unsigned int vec_sz, unsigned int items_per_wi>
 class gaussian_kernel;
 
-template <typename DataT, typename Method = mkl_rng_dev::gaussian_method::by_default>
+template <typename DataT, typename Method>
 static sycl::event gaussian_impl(sycl::queue& exec_q,
                                  const std::uint32_t seed,
                                  const double mean_val,
@@ -98,20 +131,23 @@ static sycl::event gaussian_impl(sycl::queue& exec_q,
         distr_event = exec_q.submit([&](sycl::handler &cgh) {
             cgh.depends_on(depends);
 
+            using GaussianDistrT = GaussianDistr<DataT, Method>;
+            GaussianDistrT distr(mean, stddev);
+
             if (is_aligned<required_alignment>(out_ptr)) {
                 constexpr bool enable_sg_load = true;
-                using KernelName = gaussian_kernel<DataT, vec_sz, items_per_wi>;
+                using KernelName = gaussian_kernel<DataT, Method, vec_sz, items_per_wi>;
 
                 cgh.parallel_for<KernelName>(sycl::nd_range<1>({global_size}, {local_size}),
-                    details::RngContigFunctor<DataT, DataT, Method, DataT, vec_sz, items_per_wi, enable_sg_load>(seed, mean, stddev, out, n));
+                    details::RngContigFunctor<DataT, GaussianDistrT, DataT, DataT, vec_sz, items_per_wi, enable_sg_load>(seed, distr, out, n));
             }
             else {
                 constexpr bool disable_sg_load = false;
-                using InnerKernelName = gaussian_kernel<DataT, vec_sz, items_per_wi>;
+                using InnerKernelName = gaussian_kernel<DataT, Method, vec_sz, items_per_wi>;
                 using KernelName = disabled_sg_loadstore_wrapper_krn<InnerKernelName>;
 
                 cgh.parallel_for<KernelName>(sycl::nd_range<1>({global_size}, {local_size}),
-                    details::RngContigFunctor<DataT, DataT, Method, DataT, vec_sz, items_per_wi, disable_sg_load>(seed, mean, stddev, out, n));
+                    details::RngContigFunctor<DataT, GaussianDistrT, DataT, DataT, vec_sz, items_per_wi, disable_sg_load>(seed, distr, out, n));
             }
         });
     } catch (oneapi::mkl::exception const &e) {
@@ -129,6 +165,7 @@ static sycl::event gaussian_impl(sycl::queue& exec_q,
 }
 
 std::pair<sycl::event, sycl::event> gaussian(sycl::queue exec_q,
+                                             const std::uint8_t method_id,
                                              const std::uint32_t seed,
                                              const double mean,
                                              const double stddev,
@@ -166,10 +203,14 @@ std::pair<sycl::event, sycl::event> gaussian(sycl::queue exec_q,
         throw std::runtime_error("Only population of contiguous array is supported.");
     }
 
+    if (method_id >= num_methods) {
+        throw std::runtime_error("Unknown method=" + std::to_string(method_id) + " for gaussian distribution.");
+    }
+
     auto array_types = dpctl_td_ns::usm_ndarray_types();
     int res_type_id = array_types.typenum_to_lookup_id(res.get_typenum());
 
-    auto gaussian_fn = gaussian_dispatch_vector[res_type_id];
+    auto gaussian_fn = gaussian_dispatch_table[res_type_id][method_id];
     if (gaussian_fn == nullptr) {
         throw py::value_error("No gaussian implementation defined for a required type");
     }
@@ -181,36 +222,95 @@ std::pair<sycl::event, sycl::event> gaussian(sycl::queue exec_q,
      return std::make_pair(ht_ev, gaussian_ev);
 }
 
-template <typename T>
+template <typename funcPtrT,
+          template <typename fnT, typename D, typename S> typename factory,
+          int _num_types,
+          int _num_methods>
+// class DispatchTableBuilder : public dpctl_td_ns::DispatchTableBuilder<funcPtrT, factory, _num_types>
+class DispatchTableBuilder/* : public dpctl_td_ns::DispatchTableBuilder<funcPtrT, factory, _num_types>*/
+{
+private:
+    template <typename dstTy>
+    const std::vector<funcPtrT> row_per_method() const
+    {
+        std::vector<funcPtrT> per_method = {
+            factory<funcPtrT, dstTy, mkl_rng_dev::gaussian_method::by_default>{}.get(),
+            factory<funcPtrT, dstTy, mkl_rng_dev::gaussian_method::box_muller2>{}.get(),
+        };
+        assert(per_method.size() == _num_methods);
+        return per_method;
+    }
+
+public:
+    DispatchTableBuilder() = default;
+    ~DispatchTableBuilder() = default;
+
+    void populate(funcPtrT table[][_num_methods]) const
+    {
+        const auto map_by_dst_type = {row_per_method<bool>(),
+                                      row_per_method<int8_t>(),
+                                      row_per_method<uint8_t>(),
+                                      row_per_method<int16_t>(),
+                                      row_per_method<uint16_t>(),
+                                      row_per_method<int32_t>(),
+                                      row_per_method<uint32_t>(),
+                                      row_per_method<int64_t>(),
+                                      row_per_method<uint64_t>(),
+                                      row_per_method<sycl::half>(),
+                                      row_per_method<float>(),
+                                      row_per_method<double>(),
+                                      row_per_method<std::complex<float>>(),
+                                      row_per_method<std::complex<double>>()};
+        assert(map_by_dst_type.size() == _num_types);
+        int dst_id = 0;
+        for (auto &row : map_by_dst_type) {
+            int src_id = 0;
+            for (auto &fn_ptr : row) {
+                table[dst_id][src_id] = fn_ptr;
+                ++src_id;
+            }
+            ++dst_id;
+        }
+    }
+};
+
+template <typename Ty, typename ArgTy, typename Method, typename argMethod>
+struct TypePairDefinedEntry : std::bool_constant<std::is_same_v<Ty, ArgTy> &&
+                                                 std::is_same_v<Method, argMethod>>
+{
+    static constexpr bool is_defined = true;
+};
+
+template <typename T, typename M>
 struct GaussianTypePairSupportFactory
 {
     static constexpr bool is_defined = std::disjunction<
-        dpctl_td_ns::TypePairDefinedEntry<T, double, T, double>,
-        dpctl_td_ns::TypePairDefinedEntry<T, float, T, float>,
+        TypePairDefinedEntry<T, double, M, mkl_rng_dev::gaussian_method::by_default>,
+        TypePairDefinedEntry<T, double, M, mkl_rng_dev::gaussian_method::box_muller2>,
+        TypePairDefinedEntry<T, float, M, mkl_rng_dev::gaussian_method::by_default>,
+        TypePairDefinedEntry<T, float, M, mkl_rng_dev::gaussian_method::box_muller2>,
         // fall-through
         dpctl_td_ns::NotDefinedEntry>::is_defined;
 };
 
-template <typename fnT, typename T>
+template <typename fnT, typename T, typename M>
 struct GaussianContigFactory
 {
     fnT get()
     {
-        if constexpr (GaussianTypePairSupportFactory<T>::is_defined) {
-            return gaussian_impl<T>;
+        if constexpr (GaussianTypePairSupportFactory<T, M>::is_defined) {
+            return gaussian_impl<T, M>;
         }
         else {
             return nullptr;
         }
     }
 };
 
-void init_gaussian_dispatch_vector(void)
+void init_gaussian_dispatch_table(void)
 {
-    dpctl_td_ns::DispatchVectorBuilder<gaussian_impl_fn_ptr_t, GaussianContigFactory,
-                                       dpctl_td_ns::num_types>
-        contig;
-    contig.populate_dispatch_vector(gaussian_dispatch_vector);
+    DispatchTableBuilder<gaussian_impl_fn_ptr_t, GaussianContigFactory, dpctl_td_ns::num_types, num_methods> contig;
+    contig.populate(gaussian_dispatch_table);
 }
 } // namespace device
 } // namespace rng
diff --git a/dpnp/backend/extensions/rng/device/gaussian.hpp b/dpnp/backend/extensions/rng/device/gaussian.hpp
@@ -42,14 +42,15 @@ namespace rng
 namespace device
 {
 extern std::pair<sycl::event, sycl::event> gaussian(sycl::queue exec_q,
+                                             const std::uint8_t method_id,
                                              const std::uint32_t seed,
                                              const double mean,
                                              const double stddev,
                                              const std::uint64_t n,
                                              dpctl::tensor::usm_ndarray res,
                                              const std::vector<sycl::event> &depends = {});
 
-extern void init_gaussian_dispatch_vector(void);
+extern void init_gaussian_dispatch_table(void);
 } // namespace device
 } // namespace rng
 } // namespace ext
diff --git a/dpnp/backend/extensions/rng/device/rng_py.cpp b/dpnp/backend/extensions/rng/device/rng_py.cpp
@@ -43,13 +43,13 @@ namespace py = pybind11;
 // populate dispatch vectors
 void init_dispatch_vectors(void)
 {
-    rng_dev_ext::init_gaussian_dispatch_vector();
+    // rng_dev_ext::init_gaussian_dispatch_vector();
 }
 
 // populate dispatch tables
 void init_dispatch_tables(void)
 {
-    // lapack_ext::init_heevd_dispatch_table();
+    rng_dev_ext::init_gaussian_dispatch_table();
 }
 
 
@@ -81,7 +81,7 @@ PYBIND11_MODULE(_rng_dev_impl, m)
 
     m.def("_gaussian", &rng_dev_ext::gaussian,
           "",
-          py::arg("sycl_queue"), py::arg("seed"), py::arg("mean"), py::arg("stddev"),
+          py::arg("sycl_queue"), py::arg("method"), py::arg("seed"), py::arg("mean"), py::arg("stddev"),
           py::arg("n"), py::arg("res"),
           py::arg("depends") = py::list());
 }

Original file line number	Diff line number	Diff line change
`@@ -51,8 +51,8 @@ namespace mkl_rng_dev = oneapi::mkl::rng::device;`
`51`	`51`
`52`	`52`	`/! @brief Functor for unary function evaluation on contiguous array /`
`53`	`53`	`template <typename DataT,`
	`54`	`+ typename GaussianDistrT,`
`54`	`55`	`typename ResT = DataT,`
`55`		`- typename Method = mkl_rng_dev::gaussian_method::by_default,`
`56`	`56`	`typename UnaryOperatorT = ResT,`
`57`	`57`	`unsigned int vec_sz = 8,`
`58`	`58`	`unsigned int items_per_wi = 4,`
`@@ -61,14 +61,14 @@ struct RngContigFunctor`
`61`	`61`	`{`
`62`	`62`	`private:`
`63`	`63`	`const std::uint32_t seed_;`
`64`		`- const DataT mean_;`
`65`		`- const DataT stddev_;`
	`64`	`+ GaussianDistrT distr_;`
`66`	`65`	`ResT * const res_ = nullptr;`
`67`	`66`	`const size_t nelems_;`
`68`	`67`
`69`	`68`	`public:`
`70`		`- RngContigFunctor(const std::uint32_t seed, const DataT mean, const DataT stddev, ResT *res, const size_t n_elems)`
`71`		`- : seed_(seed), mean_(mean), stddev_(stddev), res_(res), nelems_(n_elems)`
	`69`	`+`
	`70`	`+ RngContigFunctor(const std::uint32_t seed, GaussianDistrT& distr, ResT *res, const size_t n_elems)`
	`71`	`+ : seed_(seed), distr_(distr), res_(res), nelems_(n_elems)`
`72`	`72`	`{`
`73`	`73`	`}`
`74`	`74`
`@@ -80,8 +80,11 @@ struct RngContigFunctor`
`80`	`80`	`const std::uint8_t sg_size = sg.get_local_range()[0];`
`81`	`81`	`const std::uint8_t max_sg_size = sg.get_max_local_range()[0];`
`82`	`82`
`83`		`- auto engine = mkl_rng_dev::mrg32k3a<vec_sz>(seed_, nelems_ * global_id); // offset is questionable...`
`84`		`- mkl_rng_dev::gaussian<DataT, Method> distr(mean_, stddev_);`
	`83`	`+ using EngineT = typename mkl_rng_dev::mrg32k3a<vec_sz>;`
	`84`	`+ auto engine = EngineT(seed_, nelems_ * global_id); // offset is questionable...`
	`85`	`+`
	`86`	`+ using DistrT = typename GaussianDistrT::distr_type;`
	`87`	`+ DistrT distr = distr_();`
`85`	`88`
`86`	`89`	`if constexpr (enable_sg_load) {`
`87`	`90`	`const size_t base = items_per_wi * vec_sz * (nd_it.get_group(0) * nd_it.get_local_range(0) + sg.get_group_id()[0] * max_sg_size);`
`@@ -92,13 +95,13 @@ struct RngContigFunctor`
`92`	`95`	`size_t offset = base + static_cast<size_t>(it) * static_cast<size_t>(sg_size);`
`93`	`96`	`auto out_multi_ptr = sycl::address_space_cast<sycl::access::address_space::global_space, sycl::access::decorated::yes>(&res_[offset]);`
`94`	`97`
`95`		`- sycl::vec<DataT, vec_sz> rng_val_vec = mkl_rng_dev::generate(distr, engine);`
	`98`	`+ sycl::vec<DataT, vec_sz> rng_val_vec = mkl_rng_dev::generate<DistrT, EngineT>(distr, engine);`
`96`	`99`	`sg.store<vec_sz>(out_multi_ptr, rng_val_vec);`
`97`	`100`	`}`
`98`	`101`	`}`
`99`	`102`	`else {`
`100`	`103`	`for (size_t offset = base + sg.get_local_id()[0]; offset < nelems_; offset += sg_size) {`
`101`		`- res_[offset] = mkl_rng_dev::generate_single(distr, engine);`
	`104`	`+ res_[offset] = mkl_rng_dev::generate_single<DistrT, EngineT>(distr, engine);`
`102`	`105`	`}`
`103`	`106`	`}`
`104`	`107`	`}`
`@@ -108,7 +111,7 @@ struct RngContigFunctor`
`108`	`111`	`base = (base / sg_size) * sg_size * items_per_wi * vec_sz + (base % sg_size);`
`109`	`112`	`for (size_t offset = base; offset < std::min(nelems_, base + sg_size * (items_per_wi * vec_sz)); offset += sg_size)`
`110`	`113`	`{`
`111`		`- res_[offset] = mkl_rng_dev::generate_single(distr, engine);`
	`114`	`+ res_[offset] = mkl_rng_dev::generate_single<DistrT, EngineT>(distr, engine);`
`112`	`115`	`}`
`113`	`116`	`}`
`114`	`117`	`}`
Original file line number	Diff line number	Diff line change
`@@ -43,13 +43,13 @@ namespace py = pybind11;`
`43`	`43`	`// populate dispatch vectors`
`44`	`44`	`void init_dispatch_vectors(void)`
`45`	`45`	`{`
`46`		`- rng_dev_ext::init_gaussian_dispatch_vector();`
	`46`	`+ // rng_dev_ext::init_gaussian_dispatch_vector();`
`47`	`47`	`}`
`48`	`48`
`49`	`49`	`// populate dispatch tables`
`50`	`50`	`void init_dispatch_tables(void)`
`51`	`51`	`{`
`52`		`- // lapack_ext::init_heevd_dispatch_table();`
	`52`	`+ rng_dev_ext::init_gaussian_dispatch_table();`
`53`	`53`	`}`
`54`	`54`
`55`	`55`
`@@ -81,7 +81,7 @@ PYBIND11_MODULE(_rng_dev_impl, m)`
`81`	`81`
`82`	`82`	`m.def("_gaussian", &rng_dev_ext::gaussian,`
`83`	`83`	`"",`
`84`		`- py::arg("sycl_queue"), py::arg("seed"), py::arg("mean"), py::arg("stddev"),`
	`84`	`+ py::arg("sycl_queue"), py::arg("method"), py::arg("seed"), py::arg("mean"), py::arg("stddev"),`
`85`	`85`	`py::arg("n"), py::arg("res"),`
`86`	`86`	`py::arg("depends") = py::list());`
`87`	`87`	`}`