Fix multi-pass tree reductions and matmul, vecdot strides in reductions

CHANGELOG.md

@@ -110,6 +110,7 @@ This release is compatible with NumPy 2.5.
 * Fixed the list of events the copy kernels of `dpnp.reshape`, `dpnp.tensor.reshape`, `dpnp.roll` and `dpnp.tensor.roll` wait on being padded with default-constructed events [#3072](https://github.com/IntelPython/dpnp/pull/3072)
 * Fixed `simplify_iteration_three_strides` and `simplify_iteration_four_strides` accumulating into their third and fourth output displacements without zeroing them first, which required the caller to initialize them [#3072](https://github.com/IntelPython/dpnp/pull/3072)
 * Fixed `dpnp.ndarray.flat` indexing and assignment edge cases, adding support for slices, ellipsis, and integer/boolean array indices [#3045](https://github.com/IntelPython/dpnp/pull/3045)
+* Fixed incorrect results of `dpnp.tensor.vecdot` in some cases with strided outputs and of `dpnp.tensor` reductions, `dpnp.tensor.vecdot` and `dpnp.tensor.matmul` on large inputs with some data types [#3082](https://github.com/IntelPython/dpnp/pull/3082)
 
 ### Security
 

dpnp/tensor/libtensor/include/kernels/linalg_functions/dot_product.hpp

@@ -1089,8 +1089,9 @@ sycl::event dot_product_tree_impl(sycl::queue &exec_q,
                     InputIndexerT, ResIndexerT>;
             using ReductionIndexerT = dpnp::tensor::offset_utils::NoOpIndexer;
 
-            const InputIndexerT inp_indexer{/* size */ batches,
-                                            /* step */ reduction_groups_};
+            const InputIndexerT inp_indexer{
+                /* size */ batches,
+                /* step */ remaining_reduction_nelems};
             static constexpr ResIndexerT res_iter_indexer{};
 
             const InputOutputIterIndexerT in_out_iter_indexer{inp_indexer,
@@ -1124,7 +1125,7 @@ sycl::event dot_product_tree_impl(sycl::queue &exec_q,
         const ResIndexerT res_iter_indexer{
             batch_nd, batch_res_offset,
             /* shape */ batch_shape_and_strides,
-            /* strides */ batch_shape_and_strides + 2 * batch_nd};
+            /* strides */ batch_shape_and_strides + 3 * batch_nd};
 
         const InputOutputIterIndexerT in_out_iter_indexer{inp_indexer,
                                                           res_iter_indexer};
@@ -1328,8 +1329,9 @@ sycl::event
                     InputIndexerT, ResIndexerT>;
             using ReductionIndexerT = dpnp::tensor::offset_utils::NoOpIndexer;
 
-            const InputIndexerT inp_indexer{/* size */ batches,
-                                            /* step */ reduction_groups_};
+            const InputIndexerT inp_indexer{
+                /* size */ batches,
+                /* step */ remaining_reduction_nelems};
             static constexpr ResIndexerT res_iter_indexer{};
 
             const InputOutputIterIndexerT in_out_iter_indexer{inp_indexer,

dpnp/tensor/libtensor/include/kernels/linalg_functions/gemm.hpp

@@ -338,7 +338,7 @@ sycl::event tree_reduction_for_gemm(sycl::queue &exec_q,
         using ReductionIndexerT = dpnp::tensor::offset_utils::NoOpIndexer;
 
         const InputIndexerT inp_indexer{/* size */ iter_nelems,
-                                        /* step */ reduction_groups_};
+                                        /* step */ remaining_reduction_nelems};
         static constexpr ResIndexerT res_iter_indexer{};
 
         const InputOutputIterIndexerT in_out_iter_indexer{inp_indexer,
@@ -461,7 +461,7 @@ sycl::event
         // creates a stack of reduction_nelems 2D matrices and we reduce
         // along the stack axis
         const InputIndexerT inp_indexer{/* size */ iter_nelems,
-                                        /* step */ reduction_groups_};
+                                        /* step */ remaining_reduction_nelems};
         static constexpr ResIndexerT res_iter_indexer{};
 
         const InputOutputIterIndexerT in_out_iter_indexer{inp_indexer,
@@ -2432,13 +2432,14 @@ sycl::event
             using dpnp::tensor::offset_utils::Strided1DIndexer;
             using dpnp::tensor::offset_utils::StridedIndexer;
             using dpnp::tensor::offset_utils::ThreeOffsets_CombinedIndexer;
-            using BatchDimsIndexerT =
-                ThreeOffsets_CombinedIndexer<StridedIndexer, StridedIndexer,
-                                             Strided1DIndexer>;
+            using dpnp::tensor::offset_utils::UnpackedStridedIndexer;
+            using BatchDimsIndexerT = ThreeOffsets_CombinedIndexer<
+                StridedIndexer, UnpackedStridedIndexer, Strided1DIndexer>;
             const StridedIndexer lhs_batch_indexer(batch_nd, lhs_batch_offset,
                                                    batch_shape_strides);
-            const StridedIndexer rhs_batch_indexer(
-                batch_nd, rhs_batch_offset, batch_shape_strides + 2 * batch_nd);
+            const UnpackedStridedIndexer rhs_batch_indexer(
+                batch_nd, rhs_batch_offset, batch_shape_strides,
+                batch_shape_strides + 2 * batch_nd);
             const Strided1DIndexer tmp_batch_indexer(
                 /* size   */ batch_nelems,
                 /* step   */ n * m);

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

@@ -1327,8 +1327,9 @@ sycl::event reduction_over_group_temps_strided_impl(
                 using ReductionIndexerT =
                     dpnp::tensor::offset_utils::NoOpIndexer;
 
-                const InputIndexerT inp_indexer{/* size */ iter_nelems,
-                                                /* step */ reduction_groups_};
+                const InputIndexerT inp_indexer{
+                    /* size */ iter_nelems,
+                    /* step */ remaining_reduction_nelems};
                 static constexpr ResIndexerT res_iter_indexer{};
 
                 const InputOutputIterIndexerT in_out_iter_indexer{
@@ -1561,8 +1562,9 @@ sycl::event reduction_axis1_over_group_temps_contig_impl(
                     InputIndexerT, ResIndexerT>;
             using ReductionIndexerT = dpnp::tensor::offset_utils::NoOpIndexer;
 
-            const InputIndexerT inp_indexer{/* size */ iter_nelems,
-                                            /* step */ reduction_groups_};
+            const InputIndexerT inp_indexer{
+                /* size */ iter_nelems,
+                /* step */ remaining_reduction_nelems};
             static constexpr ResIndexerT res_iter_indexer{};
 
             const InputOutputIterIndexerT in_out_iter_indexer{inp_indexer,
@@ -1793,8 +1795,9 @@ sycl::event reduction_axis0_over_group_temps_contig_impl(
                     InputIndexerT, ResIndexerT>;
             using ReductionIndexerT = dpnp::tensor::offset_utils::NoOpIndexer;
 
-            const InputIndexerT inp_indexer{/* size */ iter_nelems,
-                                            /* step */ reduction_groups_};
+            const InputIndexerT inp_indexer{
+                /* size */ iter_nelems,
+                /* step */ remaining_reduction_nelems};
             static constexpr ResIndexerT res_iter_indexer{};
 
             const InputOutputIterIndexerT in_out_iter_indexer{inp_indexer,
@@ -2698,8 +2701,9 @@ sycl::event search_over_group_temps_strided_impl(
                     InputIndexerT, ResIndexerT>;
             using ReductionIndexerT = dpnp::tensor::offset_utils::NoOpIndexer;
 
-            const InputIndexerT inp_indexer{/* size */ iter_nelems,
-                                            /* step */ reduction_groups_};
+            const InputIndexerT inp_indexer{
+                /* size */ iter_nelems,
+                /* step */ remaining_reduction_nelems};
             static constexpr ResIndexerT res_iter_indexer{};
 
             const InputOutputIterIndexerT in_out_iter_indexer{inp_indexer,
@@ -2973,8 +2977,9 @@ sycl::event search_axis1_over_group_temps_contig_impl(
                     InputIndexerT, ResIndexerT>;
             using ReductionIndexerT = dpnp::tensor::offset_utils::NoOpIndexer;
 
-            const InputIndexerT inp_indexer{/* size */ iter_nelems,
-                                            /* step */ reduction_groups_};
+            const InputIndexerT inp_indexer{
+                /* size */ iter_nelems,
+                /* step */ remaining_reduction_nelems};
             static constexpr ResIndexerT res_iter_indexer{};
 
             const InputOutputIterIndexerT in_out_iter_indexer{inp_indexer,
@@ -3240,8 +3245,9 @@ sycl::event search_axis0_over_group_temps_contig_impl(
                     InputIndexerT, ResIndexerT>;
             using ReductionIndexerT = dpnp::tensor::offset_utils::NoOpIndexer;
 
-            const InputIndexerT inp_indexer{/* size */ iter_nelems,
-                                            /* step */ reduction_groups_};
+            const InputIndexerT inp_indexer{
+                /* size */ iter_nelems,
+                /* step */ remaining_reduction_nelems};
             static constexpr ResIndexerT res_iter_indexer{};
 
             const InputOutputIterIndexerT in_out_iter_indexer{inp_indexer,

dpnp/tests/tensor/test_usm_ndarray_linalg.py

@@ -1028,3 +1028,109 @@ def test_matmul_out_appended_axes():
     out = dpt.empty((), dtype="i4")
     dpt.matmul(x1, x2, out=out)
     assert out == n1
+
+
+@pytest.mark.slow
+def test_vecdot_tree_multiple_passes_broadcast():
+    q = get_queue_or_skip()
+    dev = q.sycl_device
+    # number of elements above which the partial results of the work-groups
+    # are reduced
+    wg = 4 * max(dev.sub_group_sizes)
+    max_wg = min(2048, dev.max_work_group_size // 2)
+    n = 8 * wg * 8 * max_wg + 999
+
+    x1 = dpt.zeros((3, n), dtype="i1", sycl_queue=q)
+    for r in range(3):
+        x1[r, n - 1 - r :] = 1
+    # a broadcast x2 makes the strides of the operands differ from those of
+    # the result
+    x2 = dpt.ones(n, dtype="i1", sycl_queue=q)
+    res = dpt.vecdot(x1, x2)
+    assert dpt.asnumpy(res).tolist() == [1, 2, 3]
+
+
+@pytest.mark.parametrize("dt", ["c8", "c16"])
+def test_vecdot_tree_broadcast(dt):
+    q = get_queue_or_skip()
+    skip_if_dtype_not_supported(dt, q)
+
+    n = 100000
+    x1 = dpt.ones((3, n), dtype=dt, sycl_queue=q)
+    x1 *= dpt.asarray([[1], [2], [3]], dtype=dt, sycl_queue=q)
+    x2 = dpt.ones(n, dtype=dt, sycl_queue=q)
+    res = dpt.vecdot(x1, x2)
+    assert dpt.asnumpy(res).tolist() == [n, 2 * n, 3 * n]
+
+
+def _gemm_batch_tree_loop_k(dev):
+    # inner dimension size above which the tree reductions of batched gemm
+    # reduce the partial results
+    wg = 4 * max(dev.sub_group_sizes)
+    max_wg = min(2048, dev.max_work_group_size // 2)
+    return 256 * 4 * wg * 4 * max_wg
+
+
+def _skip_if_too_large(dev, alloc_nbytes):
+    if max(alloc_nbytes) > dev.max_mem_alloc_size:
+        pytest.skip("Allocation exceeds the device's maximum allocation size")
+    if sum(alloc_nbytes) > dev.global_mem_size // 4:
+        pytest.skip("Not enough device memory")
+
+
+@pytest.mark.slow
+def test_matmul_tree_multiple_passes_strided():
+    q = get_queue_or_skip()
+    dev = q.sycl_device
+    k = _gemm_batch_tree_loop_k(dev) + 999
+    _skip_if_too_large(dev, [k])
+
+    x1 = dpt.zeros((1, 1, k), dtype="i1", sycl_queue=q)
+    x1[..., k - 3 :] = 1
+    # a broadcast x2 makes the strided implementation be used
+    x2 = dpt.broadcast_to(
+        dpt.asarray([1, 2], dtype="i1", sycl_queue=q), (2, k, 2)
+    )
+    res = dpt.matmul(x1, x2)
+    assert dpt.asnumpy(res).tolist() == [[[3, 6]], [[3, 6]]]
+
+
+@pytest.mark.slow
+def test_matmul_tree_multiple_passes_contig():
+    q = get_queue_or_skip()
+    dev = q.sycl_device
+    k = _gemm_batch_tree_loop_k(dev) + 999
+    _skip_if_too_large(dev, [2 * k, k])
+
+    x1 = dpt.zeros((1, 2, k), dtype="i1", sycl_queue=q)
+    for r in range(2):
+        x1[0, r, k - 1 - r :] = 1
+    x2 = dpt.ones((1, k, 1), dtype="i1", sycl_queue=q)
+    res = dpt.matmul(x1, x2)
+    assert dpt.asnumpy(res).tolist() == [[[1], [2]]]
+
+
+def test_matmul_tree_broadcast_batch_x2():
+    q = get_queue_or_skip()
+    dev = q.sycl_device
+    # inner dimension size above which batched gemm reduces the partial
+    # results with multiple work-groups
+    max_wg = min(2048, dev.max_work_group_size // 2)
+    k = 256 * 4 * max_wg + 999
+
+    x1 = dpt.zeros((2, 1, k), dtype="i1", sycl_queue=q)
+    x1[0, :, k - 3 :] = 1
+    x1[1, :, k - 2 :] = 1
+    # x2 is broadcast along the batch dimension
+    x2 = dpt.broadcast_to(
+        dpt.asarray([1, 2], dtype="i1", sycl_queue=q), (2, k, 2)
+    )
+    res = dpt.matmul(x1, x2)
+    assert dpt.asnumpy(res).tolist() == [[[3, 6]], [[2, 4]]]
+
+    # a broadcast x1 keeps a non-broadcast x2 on the strided implementation
+    x1 = dpt.broadcast_to(x1[:1], (2, 1, k))
+    x2 = dpt.ones((2, k, 2), dtype="i1", sycl_queue=q)
+    x2[1] = 2
+    res = dpt.matmul(x1, x2)
+    assert dpt.asnumpy(res).tolist() == [[[3, 3]], [[6, 6]]]

dpnp/tests/tensor/test_usm_ndarray_reductions.py

@@ -702,3 +702,28 @@ def test_count_nonzero(dt):
     res = dpt.count_nonzero(x)
     assert res == 7
     assert res.dtype == expected_dt
+
+
+def _tree_reduction_loop_size(dev):
+    # number of elements of a row above which the tree reductions reduce the
+    # partial results of the work-groups
+    wg = 4 * max(dev.sub_group_sizes)
+    max_wg = min(2048, dev.max_work_group_size // 2)
+    return 8 * wg * 8 * max_wg
+
+
+@pytest.mark.slow
+def test_tree_reductions_multiple_passes():
+    q = get_queue_or_skip()
+    n = _tree_reduction_loop_size(q.sycl_device) + 12345
+
+    x = dpt.zeros((3, n), dtype="i1", sycl_queue=q)
+    pos = [7, n - 5, n // 2]
+    for r, p in enumerate(pos):
+        x[r, p] = 1
+    assert dpt.asnumpy(dpt.argmax(x, axis=1)).tolist() == pos
+
+    x = dpt.ones((3, n), dtype="i1", sycl_queue=q)
+    for r, p in enumerate(pos):
+        x[r, p : p + r + 1] = 2
+    assert dpt.asnumpy(dpt.prod(x, axis=1)).tolist() == [2, 4, 8]