Corrected boolean indexing cumsum

- The cumulative sum was being calculated incorrectly -- the offset from stride simplification was unused and the result was incorrect for some cases with negative strides

Author: ndgrigorian

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

@@ -424,7 +424,6 @@ typedef size_t (*mask_positions_strided_impl_fn_ptr_t)(
     size_t,
     const char *,
     int,
-    py::ssize_t,
     const py::ssize_t *,
     char *,
     std::vector<sycl::event> const &);
@@ -434,7 +433,6 @@ size_t mask_positions_strided_impl(sycl::queue q,
                                    size_t n_elems,
                                    const char *mask,
                                    int nd,
-                                   py::ssize_t input_offset,
                                    const py::ssize_t *shape_strides,
                                    char *cumsum,
                                    std::vector<sycl::event> const &depends = {})
@@ -444,7 +442,7 @@ size_t mask_positions_strided_impl(sycl::queue q,
     cumsumT *cumsum_data_ptr = reinterpret_cast<cumsumT *>(cumsum);
     size_t wg_size = 128;
 
-    StridedIndexer strided_indexer{nd, input_offset, shape_strides};
+    StridedIndexer strided_indexer{nd, 0, shape_strides};
     NonZeroIndicator<maskT, cumsumT> non_zero_indicator{};
 
     sycl::event comp_ev =

dpctl/tensor/libtensor/source/boolean_advanced_indexing.cpp

@@ -201,46 +201,27 @@ size_t py_mask_positions(dpctl::tensor::usm_ndarray mask,
         return fn(exec_q, mask_size, mask_data, cumsum_data, depends);
     }
 
-    const py::ssize_t *shape = mask.get_shape_raw();
-    auto const &strides_vector = mask.get_strides_vector();
-
-    using shT = std::vector<py::ssize_t>;
-    shT simplified_shape;
-    shT simplified_strides;
-    py::ssize_t offset(0);
-
+    // Strided implementation
     int mask_nd = mask.get_ndim();
-    int nd = mask_nd;
-
-    dpctl::tensor::py_internal::simplify_iteration_space_1(
-        nd, shape, strides_vector, simplified_shape, simplified_strides,
-        offset);
-
-    if (nd == 1 && simplified_strides[0] == 1) {
-        auto fn = (use_i32)
-                      ? mask_positions_contig_i32_dispatch_vector[mask_typeid]
-                      : mask_positions_contig_i64_dispatch_vector[mask_typeid];
-
-        return fn(exec_q, mask_size, mask_data, cumsum_data, depends);
-    }
+    auto const &shape_vector = mask.get_shape_vector();
+    auto const &strides_vector = mask.get_strides_vector();
 
-    // Strided implementation
     auto strided_fn =
         (use_i32) ? mask_positions_strided_i32_dispatch_vector[mask_typeid]
                   : mask_positions_strided_i64_dispatch_vector[mask_typeid];
-    std::vector<sycl::event> host_task_events;
 
+    std::vector<sycl::event> host_task_events;
     using dpctl::tensor::offset_utils::device_allocate_and_pack;
     const auto &ptr_size_event_tuple = device_allocate_and_pack<py::ssize_t>(
-        exec_q, host_task_events, simplified_shape, simplified_strides);
+        exec_q, host_task_events, shape_vector, strides_vector);
     py::ssize_t *shape_strides = std::get<0>(ptr_size_event_tuple);
     if (shape_strides == nullptr) {
         sycl::event::wait(host_task_events);
         throw std::runtime_error("Unexpected error");
     }
     sycl::event copy_shape_ev = std::get<2>(ptr_size_event_tuple);
 
-    if (2 * static_cast<size_t>(nd) != std::get<1>(ptr_size_event_tuple)) {
+    if (2 * static_cast<size_t>(mask_nd) != std::get<1>(ptr_size_event_tuple)) {
         copy_shape_ev.wait();
         sycl::event::wait(host_task_events);
         sycl::free(shape_strides, exec_q);
@@ -253,7 +234,7 @@ size_t py_mask_positions(dpctl::tensor::usm_ndarray mask,
     dependent_events.insert(dependent_events.end(), depends.begin(),
                             depends.end());
 
-    size_t total_set = strided_fn(exec_q, mask_size, mask_data, nd, offset,
+    size_t total_set = strided_fn(exec_q, mask_size, mask_data, mask_nd,
                                   shape_strides, cumsum_data, dependent_events);
 
     sycl::event::wait(host_task_events);