Add pybind class for engines

antonwolfy · antonwolfy · commit 9cff28644bf6 · 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
@@ -50,42 +50,47 @@ namespace py = pybind11;
 namespace mkl_rng_dev = oneapi::mkl::rng::device;
 
 /*! @brief Functor for unary function evaluation on contiguous array */
-template <typename DataT,
+template <typename EngineDistrT,
+          typename DataT,
           typename GaussianDistrT,
-          typename ResT = DataT,
-          typename UnaryOperatorT = ResT,
-          unsigned int vec_sz = 8,
           unsigned int items_per_wi = 4,
           bool enable_sg_load = true>
 struct RngContigFunctor
 {
 private:
-    const std::uint32_t seed_;
+    // const std::uint32_t seed_;
+    EngineDistrT engine_;
     GaussianDistrT distr_;
-    ResT * const res_ = nullptr;
+    DataT * const res_ = nullptr;
     const size_t nelems_;
 
 public:
 
-    RngContigFunctor(const std::uint32_t seed, GaussianDistrT& distr, ResT *res, const size_t n_elems)
-        : seed_(seed), distr_(distr), res_(res), nelems_(n_elems)
+    RngContigFunctor(EngineDistrT& engine, GaussianDistrT& distr, DataT *res, const size_t n_elems)
+        : engine_(engine), distr_(distr), res_(res), nelems_(n_elems)
     {
     }
 
     void operator()(sycl::nd_item<1> nd_it) const
     {
-        auto global_id = nd_it.get_global_id();
+        // auto global_id = nd_it.get_global_id();
+
+        // constexpr std::size_t vec_sz = EngineT::vec_size;
         
         auto sg = nd_it.get_sub_group();
         const std::uint8_t sg_size = sg.get_local_range()[0];
         const std::uint8_t max_sg_size = sg.get_max_local_range()[0];
 
-        using EngineT = typename mkl_rng_dev::mrg32k3a<vec_sz>;
-        auto engine = EngineT(seed_, nelems_ * global_id);  // offset is questionable...
+        // auto engine = EngineT(seed_, nelems_ * global_id);  // offset is questionable...
+
+        using EngineT = typename EngineDistrT::engine_type;
+        EngineT engine = engine_();
 
         using DistrT = typename GaussianDistrT::distr_type;
         DistrT distr = distr_();
 
+        constexpr std::size_t vec_sz = EngineT::vec_size;
+
         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);
 
diff --git a/dpnp/backend/extensions/rng/device/gaussian.cpp b/dpnp/backend/extensions/rng/device/gaussian.cpp
@@ -91,7 +91,26 @@ struct GaussianDistr
     }
 };
 
-typedef sycl::event (*gaussian_impl_fn_ptr_t)(sycl::queue &,
+template <typename EngineBase, typename MklEngineT>
+struct EngineDistr
+{
+private:
+    EngineBase *engine_;
+
+public:
+    using engine_type = MklEngineT;
+
+    EngineDistr(EngineBase *engine) : engine_(engine)
+    {
+    }
+
+    inline auto operator()(void) const
+    {
+        return MklEngineT(engine_->seed_, engine_->offset_);
+    }
+};
+
+typedef sycl::event (*gaussian_impl_fn_ptr_t)(EngineBase *engine,
                                            const std::uint32_t,
                                            const double,
                                            const double,
@@ -101,25 +120,26 @@ typedef sycl::event (*gaussian_impl_fn_ptr_t)(sycl::queue &,
 
 static gaussian_impl_fn_ptr_t gaussian_dispatch_table[dpctl_td_ns::num_types][num_methods];
 
-template <typename DataT,  typename Method, unsigned int vec_sz, unsigned int items_per_wi>
+template <typename EngineT, typename DataT,  typename Method, unsigned int items_per_wi>
 class gaussian_kernel;
 
-template <typename DataT, typename Method>
-static sycl::event gaussian_impl(sycl::queue& exec_q,
+template <typename EngineT, typename DataT, typename Method>
+static sycl::event gaussian_impl(EngineBase *engine,
                                  const std::uint32_t seed,
                                  const double mean_val,
                                  const double stddev_val,
                                  const std::uint64_t n,
                                  char *out_ptr,
                                  const std::vector<sycl::event> &depends)
 {
+    auto exec_q = engine->get_queue();
     type_utils::validate_type_for_device<DataT>(exec_q);
 
     DataT *out = reinterpret_cast<DataT *>(out_ptr);
     DataT mean = static_cast<DataT>(mean_val);
     DataT stddev = static_cast<DataT>(stddev_val);
 
-    constexpr std::size_t vec_sz = 8;
+    constexpr std::size_t vec_sz = EngineT::vec_size;
     constexpr std::size_t items_per_wi = 4;
     constexpr std::size_t local_size = 256;
     const std::size_t wg_items = local_size * vec_sz * items_per_wi;
@@ -131,23 +151,28 @@ static sycl::event gaussian_impl(sycl::queue& exec_q,
         distr_event = exec_q.submit([&](sycl::handler &cgh) {
             cgh.depends_on(depends);
 
+            using EngineDistrT = EngineDistr<MRG32k3a, EngineT>;
+            EngineDistrT eng(static_cast<MRG32k3a*>(engine));
+
+            // EngineT engine = EngineT(seed, 0);
+
             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, Method, vec_sz, items_per_wi>;
+                using KernelName = gaussian_kernel<EngineT, DataT, Method, items_per_wi>;
 
                 cgh.parallel_for<KernelName>(sycl::nd_range<1>({global_size}, {local_size}),
-                    details::RngContigFunctor<DataT, GaussianDistrT, DataT, DataT, vec_sz, items_per_wi, enable_sg_load>(seed, distr, out, n));
+                    details::RngContigFunctor<EngineDistrT, DataT, GaussianDistrT, items_per_wi, enable_sg_load>(eng, distr, out, n));
             }
             else {
                 constexpr bool disable_sg_load = false;
-                using InnerKernelName = gaussian_kernel<DataT, Method, vec_sz, items_per_wi>;
+                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<DataT, GaussianDistrT, DataT, DataT, vec_sz, items_per_wi, disable_sg_load>(seed, distr, out, n));
+                    details::RngContigFunctor<EngineDistrT, DataT, GaussianDistrT, items_per_wi, disable_sg_load>(eng, distr, out, n));
             }
         });
     } catch (oneapi::mkl::exception const &e) {
@@ -164,7 +189,7 @@ static sycl::event gaussian_impl(sycl::queue& exec_q,
     return distr_event;
 }
 
-std::pair<sycl::event, sycl::event> gaussian(sycl::queue exec_q,
+std::pair<sycl::event, sycl::event> gaussian(EngineBase *engine,
                                              const std::uint8_t method_id,
                                              const std::uint32_t seed,
                                              const double mean,
@@ -173,6 +198,9 @@ std::pair<sycl::event, sycl::event> gaussian(sycl::queue exec_q,
                                              dpctl::tensor::usm_ndarray res,
                                              const std::vector<sycl::event> &depends)
 {
+    std::cout << engine->print() << std::endl;
+    auto exec_q = engine->get_queue();
+
     const int res_nd = res.get_ndim();
     const py::ssize_t *res_shape = res.get_shape_raw();
 
@@ -216,7 +244,7 @@ std::pair<sycl::event, sycl::event> gaussian(sycl::queue exec_q,
     }
 
     char *res_data = res.get_data();
-    sycl::event gaussian_ev = gaussian_fn(exec_q, seed, mean, stddev, n, res_data, depends);
+    sycl::event gaussian_ev = gaussian_fn(engine, seed, mean, stddev, n, res_data, depends);
 
     sycl::event ht_ev = dpctl::utils::keep_args_alive(exec_q, {res}, {gaussian_ev});
      return std::make_pair(ht_ev, gaussian_ev);
@@ -299,7 +327,7 @@ struct GaussianContigFactory
     fnT get()
     {
         if constexpr (GaussianTypePairSupportFactory<T, M>::is_defined) {
-            return gaussian_impl<T, M>;
+            return gaussian_impl<mkl_rng_dev::mrg32k3a<8>, T, M>;
         }
         else {
             return nullptr;
diff --git a/dpnp/backend/extensions/rng/device/gaussian.hpp b/dpnp/backend/extensions/rng/device/gaussian.hpp
@@ -31,17 +31,43 @@
 
 #include <dpctl4pybind11.hpp>
 
-namespace dpnp
-{
-namespace backend
-{
-namespace ext
-{
-namespace rng
-{
-namespace device
+class EngineBase {
+public:
+    virtual ~EngineBase() {}
+    virtual sycl::queue get_queue() = 0;
+    virtual std::string print() = 0;
+    // auto get_engine() {
+    //     return nullptr;
+    // }
+};
+
+class MRG32k3a : public EngineBase {
+public:
+    sycl::queue q_;
+    const std::uint32_t seed_;
+    const std::uint64_t offset_;
+
+// public:
+    MRG32k3a(sycl::queue &q, std::uint32_t seed, std::uint64_t offset = 0) : q_(q), seed_(seed), offset_(offset) {}
+
+    sycl::queue get_queue() override {
+        return q_;
+    }
+
+    std::string print() override {
+        return "seed = " + std::to_string(seed_) + ", offset = " + std::to_string(offset_);
+    }
+
+    // auto get_engine() override {
+    //     return oneapi::mkl::rng::device::mrg32k3a<8>(seed_, offset_);
+    // }
+
+    // using engine_type = oneapi::mkl::rng::device::mrg32k3a<8>;
+};
+
+namespace dpnp::backend::ext::rng::device
 {
-extern std::pair<sycl::event, sycl::event> gaussian(sycl::queue exec_q,
+extern std::pair<sycl::event, sycl::event> gaussian(EngineBase *engine,
                                              const std::uint8_t method_id,
                                              const std::uint32_t seed,
                                              const double mean,
@@ -51,8 +77,4 @@ extern std::pair<sycl::event, sycl::event> gaussian(sycl::queue exec_q,
                                              const std::vector<sycl::event> &depends = {});
 
 extern void init_gaussian_dispatch_table(void);
-} // namespace device
-} // namespace rng
-} // namespace ext
-} // namespace backend
-} // namespace dpnp
+} // namespace dpnp::backend::ext::rng::device
diff --git a/dpnp/backend/extensions/rng/device/rng_py.cpp b/dpnp/backend/extensions/rng/device/rng_py.cpp
@@ -52,6 +52,29 @@ void init_dispatch_tables(void)
     rng_dev_ext::init_gaussian_dispatch_table();
 }
 
+class PyEngineBase : public EngineBase {
+public:
+    /* Inherit the constructors */
+    using EngineBase::EngineBase;
+
+    /* Trampoline (need one for each virtual function) */
+    sycl::queue get_queue() override {
+        PYBIND11_OVERRIDE_PURE(
+            sycl::queue, /* Return type */
+            EngineBase,  /* Parent class */
+            get_queue,       /* Name of function in C++ (must match Python name) */
+        );
+    }
+
+    std::string print() override {
+        PYBIND11_OVERRIDE_PURE(
+            std::string, /* Return type */
+            EngineBase,  /* Parent class */
+            print,       /* Name of function in C++ (must match Python name) */
+        );
+    }
+};
+
 
 PYBIND11_MODULE(_rng_dev_impl, m)
 {
@@ -79,9 +102,18 @@ PYBIND11_MODULE(_rng_dev_impl, m)
     //       py::arg("eig_vecs"), py::arg("eig_vals"),
     //       py::arg("depends") = py::list());
 
+    py::class_<EngineBase, PyEngineBase /* <--- trampoline */>(m, "EngineBase")
+        .def(py::init<>())
+        .def("print", &EngineBase::print);
+
+    py::class_<MRG32k3a, EngineBase>(m, "MRG32k3a")
+        .def(py::init<sycl::queue &, std::uint32_t, std::uint64_t>());
+
+
     m.def("_gaussian", &rng_dev_ext::gaussian,
           "",
-          py::arg("sycl_queue"), py::arg("method"), py::arg("seed"), py::arg("mean"), py::arg("stddev"),
+          py::arg("engine"),
+          py::arg("method"), py::arg("seed"), py::arg("mean"), py::arg("stddev"),
           py::arg("n"), py::arg("res"),
           py::arg("depends") = py::list());
 }
