Code movement and utility functions

Author: AlexanderKalistratov

dpnp/backend/extensions/statistics/CMakeLists.txt

@@ -26,6 +26,7 @@
 
 set(python_module_name _statistics_impl)
 set(_module_src
+    ${CMAKE_CURRENT_SOURCE_DIR}/common.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/histogram.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/histogram_common.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/statistics_py.cpp

dpnp/backend/extensions/statistics/common.cpp

@@ -0,0 +1,82 @@
+//*****************************************************************************
+// 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.
+//*****************************************************************************
+
+#include "common.hpp"
+
+namespace statistics
+{
+namespace common
+{
+
+size_t get_max_local_size(const sycl::device &device)
+{
+    constexpr const int default_max_cpu_local_size = 256;
+    constexpr const int default_max_gpu_local_size = 0;
+
+    return get_max_local_size(device, default_max_cpu_local_size,
+                              default_max_gpu_local_size);
+}
+
+size_t get_max_local_size(const sycl::device &device,
+                          int cpu_local_size_limit,
+                          int gpu_local_size_limit)
+{
+    int max_work_group_size =
+        device.get_info<sycl::info::device::max_work_group_size>();
+    if (device.is_cpu() && cpu_local_size_limit > 0) {
+        return std::min(cpu_local_size_limit, max_work_group_size);
+    }
+    else if (device.is_gpu() && gpu_local_size_limit > 0) {
+        return std::min(gpu_local_size_limit, max_work_group_size);
+    }
+
+    return max_work_group_size;
+}
+
+sycl::nd_range<1>
+    make_ndrange(size_t global_size, size_t local_range, size_t work_per_item)
+{
+    return make_ndrange(sycl::range<1>(global_size),
+                        sycl::range<1>(local_range),
+                        sycl::range<1>(work_per_item));
+}
+
+size_t get_local_mem_size_in_bytes(const sycl::device &device)
+{
+    // Reserving 1kb for runtime needs
+    constexpr const size_t reserve = 1024;
+
+    return get_local_mem_size_in_bytes(device, reserve);
+}
+
+size_t get_local_mem_size_in_bytes(const sycl::device &device, size_t reserve)
+{
+    size_t local_mem_size =
+        device.get_info<sycl::info::device::local_mem_size>();
+    return local_mem_size - reserve;
+}
+
+} // namespace common
+} // namespace statistics

dpnp/backend/extensions/statistics/common.hpp

@@ -0,0 +1,184 @@
+//*****************************************************************************
+// 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 <complex>
+#include <functional>
+#include <tuple>
+#include <type_traits>
+
+#include <sycl/sycl.hpp>
+
+#include "utils/math_utils.hpp"
+
+namespace statistics
+{
+namespace common
+{
+
+template <typename N, typename D>
+constexpr auto CeilDiv(N n, D d)
+{
+    return (n + d - 1) / d;
+}
+
+template <typename N, typename D>
+constexpr auto Align(N n, D d)
+{
+    return CeilDiv(n, d) * d;
+}
+
+template <typename T, sycl::memory_order Order, sycl::memory_scope Scope>
+struct AtomicOp
+{
+    static void add(T &lhs, const T value)
+    {
+        sycl::atomic_ref<T, Order, Scope> lh(lhs);
+        lh += value;
+    }
+};
+
+template <typename T, sycl::memory_order Order, sycl::memory_scope Scope>
+struct AtomicOp<std::complex<T>, Order, Scope>
+{
+    static void add(std::complex<T> &lhs, const std::complex<T> value)
+    {
+        T *_lhs = reinterpret_cast<T(&)[2]>(lhs);
+        const T *_val = reinterpret_cast<const T(&)[2]>(value);
+        sycl::atomic_ref<T, Order, Scope> lh0(_lhs[0]);
+        lh0 += _val[0];
+        sycl::atomic_ref<T, Order, Scope> lh1(_lhs[1]);
+        lh1 += _val[1];
+    }
+};
+
+template <typename T>
+struct Less
+{
+    bool operator()(const T &lhs, const T &rhs) const
+    {
+        return std::less{}(lhs, rhs);
+    }
+};
+
+template <typename T>
+struct Less<std::complex<T>>
+{
+    bool operator()(const std::complex<T> &lhs,
+                    const std::complex<T> &rhs) const
+    {
+        return dpctl::tensor::math_utils::less_complex(lhs, rhs);
+    }
+};
+
+template <typename T>
+struct IsNan
+{
+    static bool isnan(const T &v)
+    {
+        if constexpr (std::is_floating_point<T>::value) {
+            return sycl::isnan(v);
+        }
+
+        return false;
+    }
+};
+
+template <typename T>
+struct IsNan<std::complex<T>>
+{
+    static bool isnan(const std::complex<T> &v)
+    {
+        T real1 = std::real(v);
+        T imag1 = std::imag(v);
+        return sycl::isnan(real1) || sycl::isnan(imag1);
+    }
+};
+
+size_t get_max_local_size(const sycl::device &device);
+size_t get_max_local_size(const sycl::device &device,
+                          int cpu_local_size_limit,
+                          int gpu_local_size_limit);
+
+inline size_t get_max_local_size(const sycl::queue &queue)
+{
+    return get_max_local_size(queue.get_device());
+}
+
+inline size_t get_max_local_size(const sycl::queue &queue,
+                                 int cpu_local_size_limit,
+                                 int gpu_local_size_limit)
+{
+    return get_max_local_size(queue.get_device(), cpu_local_size_limit,
+                              gpu_local_size_limit);
+}
+
+size_t get_local_mem_size_in_bytes(const sycl::device &device);
+size_t get_local_mem_size_in_bytes(const sycl::device &device, size_t reserve);
+
+inline size_t get_local_mem_size_in_bytes(const sycl::queue &queue)
+{
+    return get_local_mem_size_in_bytes(queue.get_device());
+}
+
+inline size_t get_local_mem_size_in_bytes(const sycl::queue &queue,
+                                          size_t reserve)
+{
+    return get_local_mem_size_in_bytes(queue.get_device(), reserve);
+}
+
+template <typename T>
+size_t get_local_mem_size_in_items(const sycl::device &device)
+{
+    return get_local_mem_size_in_bytes(device) / sizeof(T);
+}
+
+template <typename T>
+size_t get_local_mem_size_in_items(const sycl::device &device, size_t reserve)
+{
+    return get_local_mem_size_in_bytes(device, sizeof(T) * reserve) / sizeof(T);
+}
+
+template <int Dims>
+sycl::nd_range<Dims> make_ndrange(const sycl::range<Dims> &global_range,
+                                  const sycl::range<Dims> &local_range,
+                                  const sycl::range<Dims> &work_per_item)
+{
+    sycl::range<Dims> aligned_global_range;
+
+    for (int i = 0; i < Dims; ++i) {
+        aligned_global_range[i] =
+            Align(CeilDiv(global_range[i], work_per_item[i]), local_range[i]);
+    }
+
+    return sycl::nd_range<Dims>(aligned_global_range, local_range);
+}
+
+sycl::nd_range<1>
+    make_ndrange(size_t global_size, size_t local_range, size_t work_per_item);
+
+} // namespace common
+} // namespace statistics

dpnp/backend/extensions/statistics/histogram.cpp

@@ -45,10 +45,70 @@ namespace dpctl_td_ns = dpctl::tensor::type_dispatch;
 using dpctl::tensor::usm_ndarray;
 
 using namespace statistics::histogram;
+using namespace statistics::common;
 
 namespace
 {
 
+template <typename T, typename DataStorage>
+struct HistogramEdges
+{
+    static constexpr bool const sync_after_init = DataStorage::sync_after_init;
+    using boundsT = std::tuple<T, T>;
+
+    HistogramEdges(const T *global_data, size_t size, sycl::handler &cgh)
+        : data(global_data, sycl::range<1>(size), cgh)
+    {
+    }
+
+    template <int _Dims>
+    void init(const sycl::nd_item<_Dims> &item) const
+    {
+        data.init(item);
+    }
+
+    boundsT get_bounds() const
+    {
+        auto min = data.get_ptr()[0];
+        auto max = data.get_ptr()[data.size() - 1];
+        return {min, max};
+    }
+
+    template <int _Dims, typename dT>
+    size_t get_bin(const sycl::nd_item<_Dims> &,
+                   const dT *val,
+                   const boundsT &) const
+    {
+        uint32_t edges_count = data.size();
+        uint32_t bins_count = edges_count - 1;
+        const auto *bins = data.get_ptr();
+
+        uint32_t bin =
+            std::upper_bound(bins, bins + edges_count, val[0], Less<dT>{}) -
+            bins - 1;
+        bin = std::min(bin, bins_count - 1);
+
+        return bin;
+    }
+
+    template <typename dT>
+    bool in_bounds(const dT *val, const boundsT &bounds) const
+    {
+        Less<dT> _less;
+        return !_less(val[0], std::get<0>(bounds)) &&
+               !_less(std::get<1>(bounds), val[0]) && !IsNan<dT>::isnan(val[0]);
+    }
+
+private:
+    DataStorage data;
+};
+
+template <typename T>
+using CachedEdges = HistogramEdges<T, CachedData<const T, 1>>;
+
+template <typename T>
+using UncachedEdges = HistogramEdges<T, UncachedData<const T, 1>>;
+
 template <typename T, typename BinsT, typename HistType = size_t>
 static sycl::event histogram_impl(sycl::queue &exec_q,
                                   const void *vin,
@@ -65,17 +125,11 @@ static sycl::event histogram_impl(sycl::queue &exec_q,
     HistType *out = static_cast<HistType *>(vout);
 
     auto device = exec_q.get_device();
-
-    uint32_t local_size =
-        device.is_cpu()
-            ? 256
-            : device.get_info<sycl::info::device::max_work_group_size>();
+    const auto local_size = get_max_local_size(device);
 
     constexpr uint32_t WorkPI = 128; // empirically found number
-    auto global_size = Align(CeilDiv(size, WorkPI), local_size);
 
-    auto nd_range =
-        sycl::nd_range(sycl::range<1>(global_size), sycl::range<1>(local_size));
+    const auto nd_range = make_ndrange(size, local_size, WorkPI);
 
     return exec_q.submit([&](sycl::handler &cgh) {
         cgh.depends_on(depends);
@@ -96,21 +150,14 @@ static sycl::event histogram_impl(sycl::queue &exec_q,
         };
 
         auto dispatch_bins = [&](auto &weights) {
-            auto local_mem_size =
-                device.get_info<sycl::info::device::local_mem_size>() /
-                sizeof(T);
+            const auto local_mem_size = get_local_mem_size_in_items<T>(device);
             if (local_mem_size >= bins_count) {
-                uint32_t max_local_copies = local_mem_size / bins_count;
-                uint32_t local_hist_count = std::max(
-                    std::min(
-                        int(std::ceil((float(4 * local_size) / bins_count))),
-                        16),
-                    1);
-                local_hist_count = std::min(local_hist_count, max_local_copies);
+                const auto local_hist_count = get_local_hist_copies_count(
+                    local_mem_size, local_size, bins_count);
 
                 auto hist = HistWithLocalCopies<HistType>(
                     out, bins_count, local_hist_count, cgh);
-                uint32_t free_local_mem = local_mem_size - hist.size();
+                const auto free_local_mem = local_mem_size - hist.size();
 
                 dispatch_edges(free_local_mem, weights, hist);
             }