Decoupling dispatching functionality

antonwolfy · antonwolfy · commit f27927ba4e3e · 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
@@ -49,24 +49,22 @@ namespace py = pybind11;
 
 namespace mkl_rng_dev = oneapi::mkl::rng::device;
 
-/*! @brief Functor for unary function evaluation on contiguous array */
 template <typename EngineBuilderT,
-          typename DataT,
-          typename GaussianDistrT,
+          typename DistributorBuilderT,
           unsigned int items_per_wi = 4,
           bool enable_sg_load = true>
 struct RngContigFunctor
 {
 private:
-    // const std::uint32_t seed_;
+    using DataT = typename DistributorBuilderT::result_type;
+
     EngineBuilderT engine_;
-    GaussianDistrT distr_;
+    DistributorBuilderT distr_;
     DataT * const res_ = nullptr;
     const size_t nelems_;
 
 public:
-
-    RngContigFunctor(EngineBuilderT& engine, GaussianDistrT& distr, DataT *res, const size_t n_elems)
+    RngContigFunctor(EngineBuilderT& engine, DistributorBuilderT& distr, DataT *res, const size_t n_elems)
         : engine_(engine), distr_(distr), res_(res), nelems_(n_elems)
     {
     }
@@ -82,7 +80,7 @@ struct RngContigFunctor
         using EngineT = typename EngineBuilderT::EngineType;
         EngineT engine = engine_(nelems_ * global_id); // offset is questionable...
 
-        using DistrT = typename GaussianDistrT::distr_type;
+        using DistrT = typename DistributorBuilderT::distr_type;
         DistrT distr = distr_();
 
         constexpr std::size_t vec_sz = EngineT::vec_size;
diff --git a/dpnp/backend/extensions/rng/device/dispatch/matrix.hpp b/dpnp/backend/extensions/rng/device/dispatch/matrix.hpp
@@ -0,0 +1,56 @@
+//*****************************************************************************
+// Copyright (c) 2024, 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.
+//
+// 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 <oneapi/mkl/rng/device.hpp>
+
+#include "utils/type_dispatch.hpp"
+
+
+namespace dpnp::backend::ext::rng::device::dispatch
+{
+namespace dpctl_td_ns = dpctl::tensor::type_dispatch;
+namespace mkl_rng_dev = oneapi::mkl::rng::device;
+
+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<
+        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;
+};
+} // dpnp::backend::ext::rng::device::dispatch
diff --git a/dpnp/backend/extensions/rng/device/dispatch/table_builder.hpp b/dpnp/backend/extensions/rng/device/dispatch/table_builder.hpp
@@ -0,0 +1,100 @@
+//*****************************************************************************
+// Copyright (c) 2024, 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.
+//
+// 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 <oneapi/mkl/rng/device.hpp>
+
+
+namespace dpnp::backend::ext::rng::device::dispatch
+{
+namespace mkl_rng_dev = oneapi::mkl::rng::device;
+
+template <typename funcPtrT,
+          template <typename fnT, typename E, typename T, typename M> typename factory,
+          int _no_of_engines,
+          int _no_of_types,
+          int _no_of_methods>
+class Dispatch3DTableBuilder
+{
+private:
+    template <typename E, typename T>
+    const std::vector<funcPtrT> row_per_method() const
+    {
+        std::vector<funcPtrT> per_method = {
+            factory<funcPtrT, E, T, mkl_rng_dev::gaussian_method::by_default>{}.get(),
+            factory<funcPtrT, E, T, mkl_rng_dev::gaussian_method::box_muller2>{}.get(),
+        };
+        assert(per_method.size() == _no_of_methods);
+        return per_method;
+    }
+
+    template <typename E>
+    auto table_per_type_and_method() const
+    {
+        std::vector<std::vector<funcPtrT>>
+                   table_by_type = {row_per_method<E, bool>(),
+                                    row_per_method<E, int8_t>(),
+                                    row_per_method<E, uint8_t>(),
+                                    row_per_method<E, int16_t>(),
+                                    row_per_method<E, uint16_t>(),
+                                    row_per_method<E, int32_t>(),
+                                    row_per_method<E, uint32_t>(),
+                                    row_per_method<E, int64_t>(),
+                                    row_per_method<E, uint64_t>(),
+                                    row_per_method<E, sycl::half>(),
+                                    row_per_method<E, float>(),
+                                    row_per_method<E, double>(),
+                                    row_per_method<E, std::complex<float>>(),
+                                    row_per_method<E, std::complex<double>>()};
+        assert(table_by_type.size() == _no_of_types);
+        return table_by_type;
+    }
+
+public:
+    Dispatch3DTableBuilder() = default;
+    ~Dispatch3DTableBuilder() = default;
+
+    void populate(funcPtrT table[][_no_of_types][_no_of_methods]) const
+    {
+        const auto map_by_engine = {table_per_type_and_method<mkl_rng_dev::mrg32k3a<8>>()};
+        assert(map_by_engine.size() == _no_of_engines);
+
+        std::uint16_t engine_id = 0;
+        for (auto &table_by_type : map_by_engine) {
+            std::uint16_t type_id = 0;
+            for (auto &row_by_method : table_by_type) {
+                std::uint16_t method_id = 0;
+                for (auto &fn_ptr : row_by_method) {
+                    table[engine_id][type_id][method_id] = fn_ptr;
+                    ++method_id;
+                }
+                ++type_id;
+            }
+            ++engine_id;
+        }
+    }
+};
+} // dpnp::backend::ext::rng::device::dispatch
diff --git a/dpnp/backend/extensions/rng/device/engine/base_builder.hpp b/dpnp/backend/extensions/rng/device/engine/base_builder.hpp
@@ -67,7 +67,7 @@ class BaseBuilder {
         }
     }
 
-    inline auto operator()() const
+    inline auto operator()(void) const
     {
         switch (no_of_seeds) {
             case 1: {
diff --git a/dpnp/backend/extensions/rng/device/gaussian.cpp b/dpnp/backend/extensions/rng/device/gaussian.cpp
@@ -33,17 +33,15 @@
 // dpctl tensor headers
 #include "kernels/alignment.hpp"
 
-using dpctl::tensor::kernels::alignment_utils::disabled_sg_loadstore_wrapper_krn;
-using dpctl::tensor::kernels::alignment_utils::is_aligned;
-using dpctl::tensor::kernels::alignment_utils::required_alignment;
-
 #include "common_impl.hpp"
 #include "gaussian.hpp"
 
 #include "engine/engine_base.hpp"
 #include "engine/engine_builder.hpp"
 
-// #include "dpnp_utils.hpp"
+#include "dispatch/matrix.hpp"
+#include "dispatch/table_builder.hpp"
+
 
 namespace dpnp
 {
@@ -55,26 +53,31 @@ namespace rng
 {
 namespace device
 {
+namespace dpctl_krn_ns = dpctl::tensor::kernels::alignment_utils;
 namespace dpctl_td_ns = dpctl::tensor::type_dispatch;
 namespace mkl_rng_dev = oneapi::mkl::rng::device;
 namespace py = pybind11;
 namespace type_utils = dpctl::tensor::type_utils;
 
+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
 
 template <typename DataT, typename Method>
-struct GaussianDistr
+struct DistributorBuilder
 {
 private:
     const DataT mean_;
     const DataT stddev_;
 
 public:
-    using method_type = Method;
     using result_type = DataT;
+    using method_type = Method;
     using distr_type = typename mkl_rng_dev::gaussian<DataT, Method>;
 
-    GaussianDistr(const DataT mean, const DataT stddev)
+    DistributorBuilder(const DataT mean, const DataT stddev)
         : mean_(mean), stddev_(stddev)
     {
     }
@@ -128,23 +131,23 @@ static sycl::event gaussian_impl(engine::EngineBase *engine,
             EngineBuilderT eng_builder(engine);
             eng_builder.print(); // TODO: remove
 
-            using GaussianDistrT = GaussianDistr<DataT, Method>;
-            GaussianDistrT distr(mean, stddev);
+            using DistributorBuilderT = DistributorBuilder<DataT, Method>;
+            DistributorBuilderT dist_builder(mean, stddev);
 
             if (is_aligned<required_alignment>(out_ptr)) {
                 constexpr bool enable_sg_load = true;
                 using KernelName = gaussian_kernel<EngineT, DataT, Method, items_per_wi>;
 
                 cgh.parallel_for<KernelName>(sycl::nd_range<1>({global_size}, {local_size}),
-                    details::RngContigFunctor<EngineBuilderT, DataT, GaussianDistrT, items_per_wi, enable_sg_load>(eng_builder, distr, out, n));
+                    details::RngContigFunctor<EngineBuilderT, DistributorBuilderT, items_per_wi, enable_sg_load>(eng_builder, dist_builder, out, n));
             }
             else {
                 constexpr bool disable_sg_load = false;
                 using InnerKernelName = gaussian_kernel<EngineT, DataT, Method, 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<EngineBuilderT, DataT, GaussianDistrT, items_per_wi, disable_sg_load>(eng_builder, distr, out, n));
+                    details::RngContigFunctor<EngineBuilderT, DistributorBuilderT, items_per_wi, disable_sg_load>(eng_builder, dist_builder, out, n));
             }
         });
     } catch (oneapi::mkl::exception const &e) {
@@ -225,97 +228,12 @@ std::pair<sycl::event, sycl::event> gaussian(engine::EngineBase *engine,
      return std::make_pair(ht_ev, gaussian_ev);
 }
 
-template <typename funcPtrT,
-          template <typename fnT, typename E, typename T, typename M> typename factory,
-          int _no_of_engines,
-          int _no_of_types,
-          int _no_of_methods>
-class Dispatch3DTableBuilder
-{
-private:
-    template <typename E, typename T>
-    const std::vector<funcPtrT> row_per_method() const
-    {
-        std::vector<funcPtrT> per_method = {
-            factory<funcPtrT, E, T, mkl_rng_dev::gaussian_method::by_default>{}.get(),
-            factory<funcPtrT, E, T, mkl_rng_dev::gaussian_method::box_muller2>{}.get(),
-        };
-        assert(per_method.size() == _no_of_methods);
-        return per_method;
-    }
-
-    template <typename E>
-    auto table_per_type_and_method() const
-    {
-        std::vector<std::vector<funcPtrT>>
-                   table_by_type = {row_per_method<E, bool>(),
-                                    row_per_method<E, int8_t>(),
-                                    row_per_method<E, uint8_t>(),
-                                    row_per_method<E, int16_t>(),
-                                    row_per_method<E, uint16_t>(),
-                                    row_per_method<E, int32_t>(),
-                                    row_per_method<E, uint32_t>(),
-                                    row_per_method<E, int64_t>(),
-                                    row_per_method<E, uint64_t>(),
-                                    row_per_method<E, sycl::half>(),
-                                    row_per_method<E, float>(),
-                                    row_per_method<E, double>(),
-                                    row_per_method<E, std::complex<float>>(),
-                                    row_per_method<E, std::complex<double>>()};
-        assert(table_by_type.size() == _no_of_types);
-        return table_by_type;
-    }
-
-public:
-    Dispatch3DTableBuilder() = default;
-    ~Dispatch3DTableBuilder() = default;
-
-    void populate(funcPtrT table[][_no_of_types][_no_of_methods]) const
-    {
-        const auto map_by_engine = {table_per_type_and_method<mkl_rng_dev::mrg32k3a<8>>()};
-        assert(map_by_engine.size() == _no_of_engines);
-
-        std::uint16_t engine_id = 0;
-        for (auto &table_by_type : map_by_engine) {
-            std::uint16_t type_id = 0;
-            for (auto &row_by_method : table_by_type) {
-                std::uint16_t method_id = 0;
-                for (auto &fn_ptr : row_by_method) {
-                    table[engine_id][type_id][method_id] = fn_ptr;
-                    ++method_id;
-                }
-                ++type_id;
-            }
-            ++engine_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<
-        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 E, typename T, typename M>
 struct GaussianContigFactory
 {
     fnT get()
     {
-        if constexpr (GaussianTypePairSupportFactory<T, M>::is_defined) {
+        if constexpr (dispatch::GaussianTypePairSupportFactory<T, M>::is_defined) {
             return gaussian_impl<E, T, M>;
         }
         else {
@@ -326,7 +244,7 @@ struct GaussianContigFactory
 
 void init_gaussian_dispatch_table(void)
 {
-    Dispatch3DTableBuilder<gaussian_impl_fn_ptr_t, GaussianContigFactory, engine::no_of_engines, dpctl_td_ns::num_types, no_of_methods> contig;
+    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);
 }
 } // namespace device

Original file line number	Diff line number	Diff line change
`@@ -67,7 +67,7 @@ class BaseBuilder {`
`67`	`67`	`}`
`68`	`68`	`}`
`69`	`69`
`70`		`- inline auto operator()() const`
	`70`	`+ inline auto operator()(void) const`
`71`	`71`	`{`
`72`	`72`	`switch (no_of_seeds) {`
`73`	`73`	`case 1: {`