Merge pull request #2018 from IntelPython/resolve-gh-2017

Fix incorrect results from `cumulative_sum` and `cumulative_prod` with `dtype=bool`

Author: ndgrigorian

Diff for: dpctl/tensor/libtensor/include/kernels/accumulators.hpp

@@ -47,6 +47,8 @@ namespace kernels
 namespace accumulators
 {
 
+namespace su_ns = dpctl::tensor::sycl_utils;
+
 using dpctl::tensor::ssize_t;
 using namespace dpctl::tensor::offset_utils;
 
@@ -84,9 +86,18 @@ template <typename srcTy, typename dstTy> struct CastTransformer
     }
 };
 
+template <typename ScanOpT, typename T> struct needs_workaround
+{
+    // workaround needed due to crash in JITing on CPU
+    // remove when CMPLRLLVM-65813 is resolved
+    static constexpr bool value = su_ns::IsSyclLogicalAnd<T, ScanOpT>::value ||
+                                  su_ns::IsSyclLogicalOr<T, ScanOpT>::value;
+};
+
 template <typename BinOpT, typename T> struct can_use_inclusive_scan_over_group
 {
-    static constexpr bool value = sycl::has_known_identity<BinOpT, T>::value;
+    static constexpr bool value = sycl::has_known_identity<BinOpT, T>::value &&
+                                  !needs_workaround<BinOpT, T>::value;
 };
 
 namespace detail
@@ -144,8 +155,6 @@ template <typename T> class stack_strided_t
 
 // Iterative cumulative summation
 
-namespace su_ns = dpctl::tensor::sycl_utils;
-
 using nwiT = std::uint32_t;
 
 template <typename inputT,

Diff for: dpctl/tensor/libtensor/source/accumulators/cumulative_prod.cpp

@@ -70,10 +70,12 @@ template <typename argTy, typename outTy>
 struct TypePairSupportDataForProdAccumulation
 {
     static constexpr bool is_defined = std::disjunction<
+        td_ns::TypePairDefinedEntry<argTy, bool, outTy, bool>,
         td_ns::TypePairDefinedEntry<argTy, bool, outTy, std::int32_t>,
         td_ns::TypePairDefinedEntry<argTy, bool, outTy, std::int64_t>,
 
         // input int8_t
+        td_ns::TypePairDefinedEntry<argTy, std::int8_t, outTy, std::int8_t>,
         td_ns::TypePairDefinedEntry<argTy, std::int8_t, outTy, std::int32_t>,
         td_ns::TypePairDefinedEntry<argTy, std::int8_t, outTy, std::int64_t>,
 
@@ -130,6 +132,11 @@ struct TypePairSupportDataForProdAccumulation
         td_ns::NotDefinedEntry>::is_defined;
 };
 
+template <typename T>
+using CumProdScanOpT = std::conditional_t<std::is_same_v<T, bool>,
+                                          sycl::logical_and<T>,
+                                          sycl::multiplies<T>>;
+
 template <typename fnT, typename srcTy, typename dstTy>
 struct CumProd1DContigFactory
 {
@@ -138,7 +145,7 @@ struct CumProd1DContigFactory
         if constexpr (TypePairSupportDataForProdAccumulation<srcTy,
                                                              dstTy>::is_defined)
         {
-            using ScanOpT = sycl::multiplies<dstTy>;
+            using ScanOpT = CumProdScanOpT<dstTy>;
             constexpr bool include_initial = false;
             if constexpr (std::is_same_v<srcTy, dstTy>) {
                 using dpctl::tensor::kernels::accumulators::NoOpTransformer;
@@ -171,7 +178,7 @@ struct CumProd1DIncludeInitialContigFactory
         if constexpr (TypePairSupportDataForProdAccumulation<srcTy,
                                                              dstTy>::is_defined)
         {
-            using ScanOpT = sycl::multiplies<dstTy>;
+            using ScanOpT = CumProdScanOpT<dstTy>;
             constexpr bool include_initial = true;
             if constexpr (std::is_same_v<srcTy, dstTy>) {
                 using dpctl::tensor::kernels::accumulators::NoOpTransformer;
@@ -204,7 +211,7 @@ struct CumProdStridedFactory
         if constexpr (TypePairSupportDataForProdAccumulation<srcTy,
                                                              dstTy>::is_defined)
         {
-            using ScanOpT = sycl::multiplies<dstTy>;
+            using ScanOpT = CumProdScanOpT<dstTy>;
             constexpr bool include_initial = false;
             if constexpr (std::is_same_v<srcTy, dstTy>) {
                 using dpctl::tensor::kernels::accumulators::NoOpTransformer;
@@ -237,7 +244,7 @@ struct CumProdIncludeInitialStridedFactory
         if constexpr (TypePairSupportDataForProdAccumulation<srcTy,
                                                              dstTy>::is_defined)
         {
-            using ScanOpT = sycl::multiplies<dstTy>;
+            using ScanOpT = CumProdScanOpT<dstTy>;
             constexpr bool include_initial = true;
             if constexpr (std::is_same_v<srcTy, dstTy>) {
                 using dpctl::tensor::kernels::accumulators::NoOpTransformer;

Diff for: dpctl/tensor/libtensor/source/accumulators/cumulative_sum.cpp

@@ -70,10 +70,12 @@ template <typename argTy, typename outTy>
 struct TypePairSupportDataForSumAccumulation
 {
     static constexpr bool is_defined = std::disjunction<
+        td_ns::TypePairDefinedEntry<argTy, bool, outTy, bool>,
         td_ns::TypePairDefinedEntry<argTy, bool, outTy, std::int32_t>,
         td_ns::TypePairDefinedEntry<argTy, bool, outTy, std::int64_t>,
 
         // input int8_t
+        td_ns::TypePairDefinedEntry<argTy, std::int8_t, outTy, std::int8_t>,
         td_ns::TypePairDefinedEntry<argTy, std::int8_t, outTy, std::int32_t>,
         td_ns::TypePairDefinedEntry<argTy, std::int8_t, outTy, std::int64_t>,
 
@@ -130,6 +132,10 @@ struct TypePairSupportDataForSumAccumulation
         td_ns::NotDefinedEntry>::is_defined;
 };
 
+template <typename T>
+using CumSumScanOpT = std::
+    conditional_t<std::is_same_v<T, bool>, sycl::logical_or<T>, sycl::plus<T>>;
+
 template <typename fnT, typename srcTy, typename dstTy>
 struct CumSum1DContigFactory
 {
@@ -138,7 +144,7 @@ struct CumSum1DContigFactory
         if constexpr (TypePairSupportDataForSumAccumulation<srcTy,
                                                             dstTy>::is_defined)
         {
-            using ScanOpT = sycl::plus<dstTy>;
+            using ScanOpT = CumSumScanOpT<dstTy>;
             constexpr bool include_initial = false;
             if constexpr (std::is_same_v<srcTy, dstTy>) {
                 using dpctl::tensor::kernels::accumulators::NoOpTransformer;
@@ -171,7 +177,7 @@ struct CumSum1DIncludeInitialContigFactory
         if constexpr (TypePairSupportDataForSumAccumulation<srcTy,
                                                             dstTy>::is_defined)
         {
-            using ScanOpT = sycl::plus<dstTy>;
+            using ScanOpT = CumSumScanOpT<dstTy>;
             constexpr bool include_initial = true;
             if constexpr (std::is_same_v<srcTy, dstTy>) {
                 using dpctl::tensor::kernels::accumulators::NoOpTransformer;
@@ -204,7 +210,7 @@ struct CumSumStridedFactory
         if constexpr (TypePairSupportDataForSumAccumulation<srcTy,
                                                             dstTy>::is_defined)
         {
-            using ScanOpT = sycl::plus<dstTy>;
+            using ScanOpT = CumSumScanOpT<dstTy>;
             constexpr bool include_initial = false;
             if constexpr (std::is_same_v<srcTy, dstTy>) {
                 using dpctl::tensor::kernels::accumulators::NoOpTransformer;
@@ -237,7 +243,7 @@ struct CumSumIncludeInitialStridedFactory
         if constexpr (TypePairSupportDataForSumAccumulation<srcTy,
                                                             dstTy>::is_defined)
         {
-            using ScanOpT = sycl::plus<dstTy>;
+            using ScanOpT = CumSumScanOpT<dstTy>;
             constexpr bool include_initial = true;
             if constexpr (std::is_same_v<srcTy, dstTy>) {
                 using dpctl::tensor::kernels::accumulators::NoOpTransformer;

Diff for: dpctl/tests/test_tensor_accumulation.py

@@ -421,3 +421,15 @@ def test_cumulative_sum_gh_1901(p):
     inp = dpt.ones(n, dtype=dt)
     r = dpt.cumulative_sum(inp, dtype=dt)
     assert dpt.all(r == dpt.arange(1, n + 1, dtype=dt))
+
+
+@pytest.mark.parametrize(
+    "dt", ["i1", "i2", "i4", "i8", "f2", "f4", "f8", "c8", "c16"]
+)
+def test_gh_2017(dt):
+    "See https://github.com/IntelPython/dpctl/issues/2017"
+    q = get_queue_or_skip()
+    skip_if_dtype_not_supported(dt, q)
+    x = dpt.asarray([-1, 1], dtype=dpt.dtype(dt), sycl_queue=q)
+    r = dpt.cumulative_sum(x, dtype="?")
+    assert dpt.all(r)