From 79cafa530d20de2ff07d53e50c991b77378b5c4e Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Thu, 1 Oct 2026 12:13:50 -0700 Subject: [PATCH 1/3] Fix multi-pass tree reductions and vecdot result strides --- .../kernels/linalg_functions/dot_product.hpp | 12 ++++--- .../include/kernels/linalg_functions/gemm.hpp | 4 +-- .../libtensor/include/kernels/reductions.hpp | 30 ++++++++++------- dpnp/tests/tensor/test_usm_ndarray_linalg.py | 33 +++++++++++++++++++ .../tensor/test_usm_ndarray_reductions.py | 25 ++++++++++++++ 5 files changed, 85 insertions(+), 19 deletions(-) diff --git a/dpnp/tensor/libtensor/include/kernels/linalg_functions/dot_product.hpp b/dpnp/tensor/libtensor/include/kernels/linalg_functions/dot_product.hpp index b3216f382016..40b239e4e127 100644 --- a/dpnp/tensor/libtensor/include/kernels/linalg_functions/dot_product.hpp +++ b/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, diff --git a/dpnp/tensor/libtensor/include/kernels/linalg_functions/gemm.hpp b/dpnp/tensor/libtensor/include/kernels/linalg_functions/gemm.hpp index e93a386ad785..2d973e7749e3 100644 --- a/dpnp/tensor/libtensor/include/kernels/linalg_functions/gemm.hpp +++ b/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, diff --git a/dpnp/tensor/libtensor/include/kernels/reductions.hpp b/dpnp/tensor/libtensor/include/kernels/reductions.hpp index 9628e4028b26..04a988ceba95 100644 --- a/dpnp/tensor/libtensor/include/kernels/reductions.hpp +++ b/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, diff --git a/dpnp/tests/tensor/test_usm_ndarray_linalg.py b/dpnp/tests/tensor/test_usm_ndarray_linalg.py index c28754ca080f..5e508d12ed48 100644 --- a/dpnp/tests/tensor/test_usm_ndarray_linalg.py +++ b/dpnp/tests/tensor/test_usm_ndarray_linalg.py @@ -1028,3 +1028,36 @@ 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] diff --git a/dpnp/tests/tensor/test_usm_ndarray_reductions.py b/dpnp/tests/tensor/test_usm_ndarray_reductions.py index 2c431efa936d..1334d7176892 100644 --- a/dpnp/tests/tensor/test_usm_ndarray_reductions.py +++ b/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] From aa301de13424e9febabc2eac41821aa14c7e4b30 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Thu, 1 Oct 2026 12:17:21 -0700 Subject: [PATCH 2/3] Add gh-3082 to changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index c49bc020827c..63567e7748c8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -108,6 +108,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 and `dpnp.tensor.vecdot` on large rows with some data types [#3082](https://github.com/IntelPython/dpnp/pull/3082) ### Security From d264c751e133d2685198d418a67b9f224d5e4680 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Fri, 2 Oct 2026 10:08:10 -0700 Subject: [PATCH 3/3] Fix bug in `matmul` and add regression tests --- CHANGELOG.md | 2 +- .../include/kernels/linalg_functions/gemm.hpp | 11 +-- dpnp/tests/tensor/test_usm_ndarray_linalg.py | 73 +++++++++++++++++++ 3 files changed, 80 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 63567e7748c8..9bc6a254473f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -108,7 +108,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 and `dpnp.tensor.vecdot` on large rows with some data types [#3082](https://github.com/IntelPython/dpnp/pull/3082) +* 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 diff --git a/dpnp/tensor/libtensor/include/kernels/linalg_functions/gemm.hpp b/dpnp/tensor/libtensor/include/kernels/linalg_functions/gemm.hpp index 2d973e7749e3..b6f5c4249b7f 100644 --- a/dpnp/tensor/libtensor/include/kernels/linalg_functions/gemm.hpp +++ b/dpnp/tensor/libtensor/include/kernels/linalg_functions/gemm.hpp @@ -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; + 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); diff --git a/dpnp/tests/tensor/test_usm_ndarray_linalg.py b/dpnp/tests/tensor/test_usm_ndarray_linalg.py index 5e508d12ed48..ab6ad7c67325 100644 --- a/dpnp/tests/tensor/test_usm_ndarray_linalg.py +++ b/dpnp/tests/tensor/test_usm_ndarray_linalg.py @@ -1061,3 +1061,76 @@ def test_vecdot_tree_broadcast(dt): 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]]]