logsumexp and reduce_hypot no longer use atomics

This change was made to improve the accuracy of these functions

Author: ndgrigorian

dpctl/tensor/_reduction.py

@@ -288,6 +288,98 @@ def prod(x, axis=None, dtype=None, keepdims=False):
     )
 
 
+def _tree_reduction_over_axis(
+    x,
+    axis,
+    dtype,
+    keepdims,
+    _reduction_fn,
+    _dtype_supported,
+    _default_reduction_type_fn,
+    _identity=None,
+):
+    if not isinstance(x, dpt.usm_ndarray):
+        raise TypeError(f"Expected dpctl.tensor.usm_ndarray, got {type(x)}")
+    nd = x.ndim
+    if axis is None:
+        axis = tuple(range(nd))
+    if not isinstance(axis, (tuple, list)):
+        axis = (axis,)
+    axis = normalize_axis_tuple(axis, nd, "axis")
+    red_nd = len(axis)
+    perm = [i for i in range(nd) if i not in axis] + list(axis)
+    arr2 = dpt.permute_dims(x, perm)
+    res_shape = arr2.shape[: nd - red_nd]
+    q = x.sycl_queue
+    inp_dt = x.dtype
+    if dtype is None:
+        res_dt = _default_reduction_type_fn(inp_dt, q)
+    else:
+        res_dt = dpt.dtype(dtype)
+        res_dt = _to_device_supported_dtype(res_dt, q.sycl_device)
+
+    res_usm_type = x.usm_type
+    if x.size == 0:
+        if _identity is None:
+            raise ValueError("reduction does not support zero-size arrays")
+        else:
+            if keepdims:
+                res_shape = res_shape + (1,) * red_nd
+                inv_perm = sorted(range(nd), key=lambda d: perm[d])
+                res_shape = tuple(res_shape[i] for i in inv_perm)
+            return dpt.astype(
+                dpt.full(
+                    res_shape,
+                    _identity,
+                    dtype=_default_reduction_type_fn(inp_dt, q),
+                    usm_type=res_usm_type,
+                    sycl_queue=q,
+                ),
+                res_dt,
+            )
+    if red_nd == 0:
+        return dpt.astype(x, res_dt, copy=False)
+
+    host_tasks_list = []
+    if _dtype_supported(inp_dt, res_dt):
+        res = dpt.empty(
+            res_shape, dtype=res_dt, usm_type=res_usm_type, sycl_queue=q
+        )
+        ht_e, _ = _reduction_fn(
+            src=arr2, trailing_dims_to_reduce=red_nd, dst=res, sycl_queue=q
+        )
+        host_tasks_list.append(ht_e)
+    else:
+        if dtype is None:
+            raise RuntimeError(
+                "Automatically determined reduction data type does not "
+                "have direct implementation"
+            )
+        tmp_dt = _default_reduction_type_fn(inp_dt, q)
+        tmp = dpt.empty(
+            res_shape, dtype=tmp_dt, usm_type=res_usm_type, sycl_queue=q
+        )
+        ht_e_tmp, r_e = _reduction_fn(
+            src=arr2, trailing_dims_to_reduce=red_nd, dst=tmp, sycl_queue=q
+        )
+        host_tasks_list.append(ht_e_tmp)
+        res = dpt.empty(
+            res_shape, dtype=res_dt, usm_type=res_usm_type, sycl_queue=q
+        )
+        ht_e, _ = ti._copy_usm_ndarray_into_usm_ndarray(
+            src=tmp, dst=res, sycl_queue=q, depends=[r_e]
+        )
+        host_tasks_list.append(ht_e)
+
+    if keepdims:
+        res_shape = res_shape + (1,) * red_nd
+        inv_perm = sorted(range(nd), key=lambda d: perm[d])
+        res = dpt.permute_dims(dpt.reshape(res, res_shape), inv_perm)
+    dpctl.SyclEvent.wait_for(host_tasks_list)
+
+    return res
+
+
 def logsumexp(x, axis=None, dtype=None, keepdims=False):
     """logsumexp(x, axis=None, dtype=None, keepdims=False)
 
@@ -330,7 +422,7 @@ def logsumexp(x, axis=None, dtype=None, keepdims=False):
             array has the data type as described in the `dtype` parameter
             description above.
     """
-    return _reduction_over_axis(
+    return _tree_reduction_over_axis(
         x,
         axis,
         dtype,
@@ -384,7 +476,7 @@ def reduce_hypot(x, axis=None, dtype=None, keepdims=False):
             array has the data type as described in the `dtype` parameter
             description above.
     """
-    return _reduction_over_axis(
+    return _tree_reduction_over_axis(
         x,
         axis,
         dtype,

dpctl/tensor/libtensor/include/kernels/reductions.hpp

@@ -3069,53 +3069,6 @@ struct ProductOverAxis0TempsContigFactory
     }
 };
 
-/* @brief Types supported by hypot-reduction code based on atomic_ref */
-template <typename argTy, typename outTy>
-struct TypePairSupportDataForHypotReductionAtomic
-{
-
-    /* value if true a kernel for <argTy, outTy> must be instantiated, false
-     * otherwise */
-    static constexpr bool is_defined = std::disjunction< // disjunction is C++17
-                                                         // feature, supported
-                                                         // by DPC++ input bool
-        // input bool
-        td_ns::TypePairDefinedEntry<argTy, bool, outTy, float>,
-        td_ns::TypePairDefinedEntry<argTy, bool, outTy, double>,
-        // input int8
-        td_ns::TypePairDefinedEntry<argTy, std::int8_t, outTy, float>,
-        td_ns::TypePairDefinedEntry<argTy, std::int8_t, outTy, double>,
-        // input uint8
-        td_ns::TypePairDefinedEntry<argTy, std::uint8_t, outTy, float>,
-        td_ns::TypePairDefinedEntry<argTy, std::uint8_t, outTy, double>,
-        // input int16
-        td_ns::TypePairDefinedEntry<argTy, std::int16_t, outTy, float>,
-        td_ns::TypePairDefinedEntry<argTy, std::int16_t, outTy, double>,
-        // input uint16
-        td_ns::TypePairDefinedEntry<argTy, std::uint16_t, outTy, float>,
-        td_ns::TypePairDefinedEntry<argTy, std::uint16_t, outTy, double>,
-        // input int32
-        td_ns::TypePairDefinedEntry<argTy, std::int32_t, outTy, float>,
-        td_ns::TypePairDefinedEntry<argTy, std::int32_t, outTy, double>,
-        // input uint32
-        td_ns::TypePairDefinedEntry<argTy, std::uint32_t, outTy, float>,
-        td_ns::TypePairDefinedEntry<argTy, std::uint32_t, outTy, double>,
-        // input int64
-        td_ns::TypePairDefinedEntry<argTy, std::int64_t, outTy, double>,
-        // input uint64
-        td_ns::TypePairDefinedEntry<argTy, std::uint64_t, outTy, double>,
-        // input half
-        td_ns::TypePairDefinedEntry<argTy, sycl::half, outTy, float>,
-        td_ns::TypePairDefinedEntry<argTy, float, outTy, double>,
-        // input float
-        td_ns::TypePairDefinedEntry<argTy, float, outTy, float>,
-        td_ns::TypePairDefinedEntry<argTy, float, outTy, double>,
-        // input double
-        td_ns::TypePairDefinedEntry<argTy, double, outTy, double>,
-        // fall-through
-        td_ns::NotDefinedEntry>::is_defined;
-};
-
 template <typename argTy, typename outTy>
 struct TypePairSupportDataForHypotReductionTemps
 {
@@ -3177,25 +3130,6 @@ struct TypePairSupportDataForHypotReductionTemps
         td_ns::NotDefinedEntry>::is_defined;
 };
 
-template <typename fnT, typename srcTy, typename dstTy>
-struct HypotOverAxisAtomicStridedFactory
-{
-    fnT get() const
-    {
-        if constexpr (TypePairSupportDataForHypotReductionAtomic<
-                          srcTy, dstTy>::is_defined)
-        {
-            using ReductionOpT = su_ns::Hypot<dstTy>;
-            return dpctl::tensor::kernels::
-                reduction_over_group_with_atomics_strided_impl<srcTy, dstTy,
-                                                               ReductionOpT>;
-        }
-        else {
-            return nullptr;
-        }
-    }
-};
-
 template <typename fnT, typename srcTy, typename dstTy>
 struct HypotOverAxisTempsStridedFactory
 {
@@ -3215,44 +3149,6 @@ struct HypotOverAxisTempsStridedFactory
     }
 };
 
-template <typename fnT, typename srcTy, typename dstTy>
-struct HypotOverAxis1AtomicContigFactory
-{
-    fnT get() const
-    {
-        if constexpr (TypePairSupportDataForHypotReductionAtomic<
-                          srcTy, dstTy>::is_defined)
-        {
-            using ReductionOpT = su_ns::Hypot<dstTy>;
-            return dpctl::tensor::kernels::
-                reduction_axis1_over_group_with_atomics_contig_impl<
-                    srcTy, dstTy, ReductionOpT>;
-        }
-        else {
-            return nullptr;
-        }
-    }
-};
-
-template <typename fnT, typename srcTy, typename dstTy>
-struct HypotOverAxis0AtomicContigFactory
-{
-    fnT get() const
-    {
-        if constexpr (TypePairSupportDataForHypotReductionAtomic<
-                          srcTy, dstTy>::is_defined)
-        {
-            using ReductionOpT = su_ns::Hypot<dstTy>;
-            return dpctl::tensor::kernels::
-                reduction_axis0_over_group_with_atomics_contig_impl<
-                    srcTy, dstTy, ReductionOpT>;
-        }
-        else {
-            return nullptr;
-        }
-    }
-};
-
 template <typename fnT, typename srcTy, typename dstTy>
 struct HypotOverAxis1TempsContigFactory
 {
@@ -3291,53 +3187,6 @@ struct HypotOverAxis0TempsContigFactory
     }
 };
 
-/* @brief Types supported by logsumexp-reduction code based on atomic_ref */
-template <typename argTy, typename outTy>
-struct TypePairSupportDataForLogSumExpReductionAtomic
-{
-
-    /* value if true a kernel for <argTy, outTy> must be instantiated, false
-     * otherwise */
-    static constexpr bool is_defined = std::disjunction< // disjunction is C++17
-                                                         // feature, supported
-                                                         // by DPC++ input bool
-        // input bool
-        td_ns::TypePairDefinedEntry<argTy, bool, outTy, float>,
-        td_ns::TypePairDefinedEntry<argTy, bool, outTy, double>,
-        // input int8
-        td_ns::TypePairDefinedEntry<argTy, std::int8_t, outTy, float>,
-        td_ns::TypePairDefinedEntry<argTy, std::int8_t, outTy, double>,
-        // input uint8
-        td_ns::TypePairDefinedEntry<argTy, std::uint8_t, outTy, float>,
-        td_ns::TypePairDefinedEntry<argTy, std::uint8_t, outTy, double>,
-        // input int16
-        td_ns::TypePairDefinedEntry<argTy, std::int16_t, outTy, float>,
-        td_ns::TypePairDefinedEntry<argTy, std::int16_t, outTy, double>,
-        // input uint16
-        td_ns::TypePairDefinedEntry<argTy, std::uint16_t, outTy, float>,
-        td_ns::TypePairDefinedEntry<argTy, std::uint16_t, outTy, double>,
-        // input int32
-        td_ns::TypePairDefinedEntry<argTy, std::int32_t, outTy, float>,
-        td_ns::TypePairDefinedEntry<argTy, std::int32_t, outTy, double>,
-        // input uint32
-        td_ns::TypePairDefinedEntry<argTy, std::uint32_t, outTy, float>,
-        td_ns::TypePairDefinedEntry<argTy, std::uint32_t, outTy, double>,
-        // input int64
-        td_ns::TypePairDefinedEntry<argTy, std::int64_t, outTy, double>,
-        // input uint64
-        td_ns::TypePairDefinedEntry<argTy, std::uint64_t, outTy, double>,
-        // input half
-        td_ns::TypePairDefinedEntry<argTy, sycl::half, outTy, float>,
-        td_ns::TypePairDefinedEntry<argTy, float, outTy, double>,
-        // input float
-        td_ns::TypePairDefinedEntry<argTy, float, outTy, float>,
-        td_ns::TypePairDefinedEntry<argTy, float, outTy, double>,
-        // input double
-        td_ns::TypePairDefinedEntry<argTy, double, outTy, double>,
-        // fall-through
-        td_ns::NotDefinedEntry>::is_defined;
-};
-
 template <typename argTy, typename outTy>
 struct TypePairSupportDataForLogSumExpReductionTemps
 {
@@ -3399,25 +3248,6 @@ struct TypePairSupportDataForLogSumExpReductionTemps
         td_ns::NotDefinedEntry>::is_defined;
 };
 
-template <typename fnT, typename srcTy, typename dstTy>
-struct LogSumExpOverAxisAtomicStridedFactory
-{
-    fnT get() const
-    {
-        if constexpr (TypePairSupportDataForLogSumExpReductionAtomic<
-                          srcTy, dstTy>::is_defined)
-        {
-            using ReductionOpT = su_ns::LogSumExp<dstTy>;
-            return dpctl::tensor::kernels::
-                reduction_over_group_with_atomics_strided_impl<srcTy, dstTy,
-                                                               ReductionOpT>;
-        }
-        else {
-            return nullptr;
-        }
-    }
-};
-
 template <typename fnT, typename srcTy, typename dstTy>
 struct LogSumExpOverAxisTempsStridedFactory
 {
@@ -3437,44 +3267,6 @@ struct LogSumExpOverAxisTempsStridedFactory
     }
 };
 
-template <typename fnT, typename srcTy, typename dstTy>
-struct LogSumExpOverAxis1AtomicContigFactory
-{
-    fnT get() const
-    {
-        if constexpr (TypePairSupportDataForLogSumExpReductionAtomic<
-                          srcTy, dstTy>::is_defined)
-        {
-            using ReductionOpT = su_ns::LogSumExp<dstTy>;
-            return dpctl::tensor::kernels::
-                reduction_axis1_over_group_with_atomics_contig_impl<
-                    srcTy, dstTy, ReductionOpT>;
-        }
-        else {
-            return nullptr;
-        }
-    }
-};
-
-template <typename fnT, typename srcTy, typename dstTy>
-struct LogSumExpOverAxis0AtomicContigFactory
-{
-    fnT get() const
-    {
-        if constexpr (TypePairSupportDataForLogSumExpReductionAtomic<
-                          srcTy, dstTy>::is_defined)
-        {
-            using ReductionOpT = su_ns::LogSumExp<dstTy>;
-            return dpctl::tensor::kernels::
-                reduction_axis0_over_group_with_atomics_contig_impl<
-                    srcTy, dstTy, ReductionOpT>;
-        }
-        else {
-            return nullptr;
-        }
-    }
-};
-
 template <typename fnT, typename srcTy, typename dstTy>
 struct LogSumExpOverAxis1TempsContigFactory
 {