Introducing dispatch_table

Author: AlexanderKalistratov

dpnp/backend/extensions/statistics/__init__.py

@@ -24,6 +24,10 @@
 //*****************************************************************************
 
 #include "common.hpp"
+#include "utils/type_dispatch.hpp"
+#include <pybind11/pybind11.h>
+
+namespace dpctl_td_ns = dpctl::tensor::type_dispatch;
 
 namespace statistics
 {
@@ -78,5 +82,43 @@ size_t get_local_mem_size_in_bytes(const sycl::device &device, size_t reserve)
     return local_mem_size - reserve;
 }
 
+pybind11::dtype dtype_from_typenum(int dst_typenum)
+{
+    dpctl_td_ns::typenum_t dst_typenum_t =
+        static_cast<dpctl_td_ns::typenum_t>(dst_typenum);
+    switch (dst_typenum_t) {
+    case dpctl_td_ns::typenum_t::BOOL:
+        return py::dtype("?");
+    case dpctl_td_ns::typenum_t::INT8:
+        return py::dtype("i1");
+    case dpctl_td_ns::typenum_t::UINT8:
+        return py::dtype("u1");
+    case dpctl_td_ns::typenum_t::INT16:
+        return py::dtype("i2");
+    case dpctl_td_ns::typenum_t::UINT16:
+        return py::dtype("u2");
+    case dpctl_td_ns::typenum_t::INT32:
+        return py::dtype("i4");
+    case dpctl_td_ns::typenum_t::UINT32:
+        return py::dtype("u4");
+    case dpctl_td_ns::typenum_t::INT64:
+        return py::dtype("i8");
+    case dpctl_td_ns::typenum_t::UINT64:
+        return py::dtype("u8");
+    case dpctl_td_ns::typenum_t::HALF:
+        return py::dtype("f2");
+    case dpctl_td_ns::typenum_t::FLOAT:
+        return py::dtype("f4");
+    case dpctl_td_ns::typenum_t::DOUBLE:
+        return py::dtype("f8");
+    case dpctl_td_ns::typenum_t::CFLOAT:
+        return py::dtype("c8");
+    case dpctl_td_ns::typenum_t::CDOUBLE:
+        return py::dtype("c16");
+    default:
+        throw py::value_error("Unrecognized dst_typeid");
+    }
+}
+
 } // namespace common
 } // namespace statistics

dpnp/backend/extensions/statistics/common.cpp

@@ -26,13 +26,15 @@
 #pragma once
 
 #include <complex>
-#include <functional>
-#include <tuple>
-#include <type_traits>
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
 
+// clang-format off
+// math_utils.hpp doesn't include sycl header but uses sycl types
+// so sycl.hpp must be included before math_utils.hpp
 #include <sycl/sycl.hpp>
-
 #include "utils/math_utils.hpp"
+// clang-format on
 
 namespace statistics
 {
@@ -180,5 +182,9 @@ sycl::nd_range<Dims> make_ndrange(const sycl::range<Dims> &global_range,
 sycl::nd_range<1>
     make_ndrange(size_t global_size, size_t local_range, size_t work_per_item);
 
+// This function is a copy from dpctl because it is not available in the public
+// headers of dpctl.
+pybind11::dtype dtype_from_typenum(int dst_typenum);
+
 } // namespace common
 } // namespace statistics

dpnp/backend/extensions/statistics/common.hpp

@@ -0,0 +1,292 @@
+//*****************************************************************************
+// 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 <unordered_set>
+#include <vector>
+
+#include "utils/type_dispatch.hpp"
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <sycl/sycl.hpp>
+
+#include "common.hpp"
+
+namespace dpctl_td_ns = dpctl::tensor::type_dispatch;
+namespace py = pybind11;
+
+namespace statistics
+{
+namespace common
+{
+
+template <typename T, typename Rest>
+struct one_of
+{
+    static_assert(std::is_same_v<Rest, std::tuple<>>,
+                  "one_of: second parameter cannot be empty std::tuple");
+    static_assert(false, "one_of: second parameter must be std::tuple");
+};
+
+template <typename T, typename Top, typename... Rest>
+struct one_of<T, std::tuple<Top, Rest...>>
+{
+    static constexpr bool value =
+        std::is_same_v<T, Top> || one_of<T, std::tuple<Rest...>>::value;
+};
+
+template <typename T, typename Top>
+struct one_of<T, std::tuple<Top>>
+{
+    static constexpr bool value = std::is_same_v<T, Top>;
+};
+
+template <typename T, typename Rest>
+constexpr bool one_of_v = one_of<T, Rest>::value;
+
+template <typename FnT>
+using Table = FnT[dpctl_td_ns::num_types];
+template <typename FnT>
+using Table2 = Table<FnT>[dpctl_td_ns::num_types];
+
+using TypeId = int32_t;
+using TypesPair = std::pair<TypeId, TypeId>;
+
+struct int_pair_hash
+{
+    inline size_t operator()(const TypesPair &p) const
+    {
+        std::hash<size_t> hasher;
+        return hasher(size_t(p.first) << (8 * sizeof(TypeId)) |
+                      size_t(p.second));
+    }
+};
+
+using SupportedTypesList = std::vector<TypeId>;
+using SupportedTypesList2 = std::vector<TypesPair>;
+using SupportedTypesSet = std::unordered_set<TypeId>;
+using SupportedTypesSet2 = std::unordered_set<TypesPair, int_pair_hash>;
+
+using DType = py::dtype;
+using DTypePair = std::pair<DType, DType>;
+
+using SupportedDTypeList = std::vector<DType>;
+using SupportedDTypeList2 = std::vector<DTypePair>;
+
+template <typename FnT,
+          typename SupportedTypes,
+          template <typename, typename>
+          typename Func>
+struct TableBuilder2
+{
+    template <typename _FnT, typename T1, typename T2>
+    struct impl
+    {
+        static constexpr bool is_defined =
+            one_of_v<std::tuple<T1, T2>, SupportedTypes>;
+
+        _FnT get()
+        {
+            if constexpr (is_defined) {
+                return Func<T1, T2>::impl;
+            }
+            else {
+                return nullptr;
+            }
+        }
+    };
+
+    using type =
+        dpctl_td_ns::DispatchTableBuilder<FnT, impl, dpctl_td_ns::num_types>;
+};
+
+template <typename FnT>
+class DispatchTable2
+{
+public:
+    DispatchTable2(std::string first_name, std::string second_name)
+        : first_name(first_name), second_name(second_name)
+    {
+    }
+
+    template <typename SupportedTypes,
+              template <typename, typename>
+              typename Func>
+    void populate_dispatch_table()
+    {
+        using TBulder = typename TableBuilder2<FnT, SupportedTypes, Func>::type;
+        TBulder builder;
+
+        builder.populate_dispatch_table(table);
+        populate_supported_types();
+    }
+
+    FnT get_unsafe(int first_typenum, int second_typenum) const
+    {
+        auto array_types = dpctl_td_ns::usm_ndarray_types();
+        const int first_type_id =
+            array_types.typenum_to_lookup_id(first_typenum);
+        const int second_type_id =
+            array_types.typenum_to_lookup_id(second_typenum);
+
+        return table[first_type_id][second_type_id];
+    }
+
+    FnT get(int first_typenum, int second_typenum) const
+    {
+        auto fn = get_unsafe(first_typenum, second_typenum);
+
+        if (fn == nullptr) {
+            auto array_types = dpctl_td_ns::usm_ndarray_types();
+            const int first_type_id =
+                array_types.typenum_to_lookup_id(first_typenum);
+            const int second_type_id =
+                array_types.typenum_to_lookup_id(second_typenum);
+
+            py::dtype first_dtype = dtype_from_typenum(first_type_id);
+            auto first_type_pos =
+                std::find(supported_first_type.begin(),
+                          supported_first_type.end(), first_dtype);
+            if (first_type_pos == supported_first_type.end()) {
+                py::str types = py::str(py::cast(supported_first_type));
+                py::str dtype = py::str(first_dtype);
+
+                py::str err_msg =
+                    py::str("'" + first_name + "' has unsupported type '") +
+                    dtype +
+                    py::str("'."
+                            " Supported types are: ") +
+                    types;
+
+                throw py::value_error(static_cast<std::string>(err_msg));
+            }
+
+            py::dtype second_dtype = dtype_from_typenum(second_type_id);
+            auto second_type_pos =
+                std::find(supported_second_type.begin(),
+                          supported_second_type.end(), second_dtype);
+            if (second_type_pos == supported_second_type.end()) {
+                py::str types = py::str(py::cast(supported_second_type));
+                py::str dtype = py::str(second_dtype);
+
+                py::str err_msg =
+                    py::str("'" + second_name + "' has unsupported type '") +
+                    dtype +
+                    py::str("'."
+                            " Supported types are: ") +
+                    types;
+
+                throw py::value_error(static_cast<std::string>(err_msg));
+            }
+
+            py::str first_dtype_str = py::str(first_dtype);
+            py::str second_dtype_str = py::str(second_dtype);
+            py::str types = py::str(py::cast(all_supported_types));
+
+            py::str err_msg =
+                py::str("'" + first_name + "' and '" + second_name +
+                        "' has unsupported types combination: ('") +
+                first_dtype_str + py::str("', '") + second_dtype_str +
+                py::str("')."
+                        " Supported types combinations are: ") +
+                types;
+
+            throw py::value_error(static_cast<std::string>(err_msg));
+        }
+
+        return fn;
+    }
+
+    const SupportedDTypeList &get_supported_first_type() const
+    {
+        return supported_first_type;
+    }
+
+    const SupportedDTypeList &get_supported_second_type() const
+    {
+        return supported_second_type;
+    }
+
+    const SupportedDTypeList2 &get_all_supported_types() const
+    {
+        return all_supported_types;
+    }
+
+private:
+    void populate_supported_types()
+    {
+        SupportedTypesSet first_supported_types_set;
+        SupportedTypesSet second_supported_types_set;
+        SupportedTypesSet2 all_supported_types_set;
+
+        for (int i = 0; i < dpctl_td_ns::num_types; ++i) {
+            for (int j = 0; j < dpctl_td_ns::num_types; ++j) {
+                if (table[i][j] != nullptr) {
+                    all_supported_types_set.emplace(i, j);
+                    first_supported_types_set.emplace(i);
+                    second_supported_types_set.emplace(j);
+                }
+            }
+        }
+
+        auto to_supported_dtype_list = [](const auto &supported_set,
+                                          auto &supported_list) {
+            SupportedTypesList lst(supported_set.begin(), supported_set.end());
+            std::sort(lst.begin(), lst.end());
+            supported_list.resize(supported_set.size());
+            std::transform(lst.begin(), lst.end(), supported_list.begin(),
+                           [](TypeId i) { return dtype_from_typenum(i); });
+        };
+
+        to_supported_dtype_list(first_supported_types_set,
+                                supported_first_type);
+        to_supported_dtype_list(second_supported_types_set,
+                                supported_second_type);
+
+        SupportedTypesList2 lst(all_supported_types_set.begin(),
+                                all_supported_types_set.end());
+        std::sort(lst.begin(), lst.end());
+        all_supported_types.resize(all_supported_types_set.size());
+        std::transform(lst.begin(), lst.end(), all_supported_types.begin(),
+                       [](TypesPair p) {
+                           return DTypePair(dtype_from_typenum(p.first),
+                                            dtype_from_typenum(p.second));
+                       });
+    }
+
+    std::string first_name;
+    std::string second_name;
+
+    SupportedDTypeList supported_first_type;
+    SupportedDTypeList supported_second_type;
+    SupportedDTypeList2 all_supported_types;
+
+    Table2<FnT> table;
+};
+
+} // namespace common
+} // namespace statistics