From 2a082cd97b1c8511c970365e85d454bde05279fa Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Mon, 14 Sep 2026 10:32:31 -0700 Subject: [PATCH 1/3] Avoid a C-order copy in dpnp.einsum --- dpnp/dpnp_utils/dpnp_utils_einsum.py | 45 +++++++++++++++++++++++----- 1 file changed, 38 insertions(+), 7 deletions(-) diff --git a/dpnp/dpnp_utils/dpnp_utils_einsum.py b/dpnp/dpnp_utils/dpnp_utils_einsum.py index efa669126a1..1f2521ef823 100644 --- a/dpnp/dpnp_utils/dpnp_utils_einsum.py +++ b/dpnp/dpnp_utils/dpnp_utils_einsum.py @@ -880,7 +880,7 @@ def _parse_possible_contraction( return [sort, positions, new_input_sets] -def _reduced_binary_einsum(arr0, sub0, arr1, sub1, sub_others): +def _reduced_binary_einsum(arr0, sub0, arr1, sub1, sub_others, prefer_c=False): """Copied from _reduced_binary_einsum in cupy/core/_einsum.py""" set0 = set(sub0) @@ -916,6 +916,27 @@ def _reduced_binary_einsum(arr0, sub0, arr1, sub1, sub_others): arr1 = _expand_dims_transpose(arr1, sub1, sub_out) return arr0 * arr1, sub_out + if ( + prefer_c + and sub_l + and sub_r + and sub_others.index(sub_l[0]) > sub_others.index(sub_r[0]) + ): + # Swap the operand roles so the free axes of the product already come + # out in `sub_others` order. Otherwise the caller transposes the + # result, which makes it f-contiguous and costs a copy into c-order. + sub_out = sub_b + sub_r + sub_l + arr0, bs0, cs0, ts0, arr1, bs1, cs1, ts1 = ( + arr1, + bs1, + cs1, + ts1, + arr0, + bs0, + cs0, + ts0, + ) + tmp0, shapes0 = _flatten_transpose(arr0, [bs0, ts0, cs0]) tmp1, shapes1 = _flatten_transpose(arr1, [bs1, cs1, ts1]) shapes_out = shapes0[0] + shapes0[1] + shapes1[2] @@ -1211,7 +1232,15 @@ def dpnp_einsum( stacklevel=2, ) - for idx0, idx1 in _iter_path_pairs(path): + # Resolved above the loop because `prefer_c` below needs the final order: + # only a "C" target gains from a product laid out in output order. + if order == "K" and optimize is False and not all_f_contiguous: + # only the unoptimized path of NumPy copies into a c-contiguous + # array, the optimized one is matmul-based, as dpnp always is + order = "C" + + pairs = list(_iter_path_pairs(path)) + for pair_idx, (idx0, idx1) in enumerate(pairs): # "reduced" binary einsum arr0 = operands.pop(idx0) sub0 = input_subscripts.pop(idx0) @@ -1224,7 +1253,13 @@ def dpnp_einsum( ) ) arr_out, sub_out = _reduced_binary_einsum( - arr0, sub0, arr1, sub1, sub_others + arr0, + sub0, + arr1, + sub1, + sub_others, + # only "C" and the last contraction + prefer_c=order == "C" and pair_idx == len(pairs) - 1, ) operands.append(arr_out) input_subscripts.append(sub_out) @@ -1245,10 +1280,6 @@ def dpnp_einsum( # a view is returned for any `order`, the same way NumPy does if not returns_view: - if order == "K" and optimize is False and not all_f_contiguous: - # only the unoptimized path of NumPy copies into a c-contiguous - # array, the optimized one is matmul-based, as dpnp always is - order = "C" arr_out = dpnp.asarray(arr_out, order=order) assert returns_view or arr_out.dtype == result_dtype return dpnp.get_result_array(arr_out, out, casting=casting) From da897622f9957838794c1660d958584872b1655f Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Mon, 14 Sep 2026 10:58:27 -0700 Subject: [PATCH 2/3] Add gh-3069 to changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0f47c903d54..a11a1259406 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -48,6 +48,7 @@ This release is compatible with NumPy 2.5. * `dpnp` uses pybind11 3.1.0 [#3015](https://github.com/IntelPython/dpnp/pull/3015) * Reworked the ASV benchmarks and added end-to-end workload benchmarks derived from dpBench [#2996](https://github.com/IntelPython/dpnp/pull/2996) * Reduced allocations in `dpnp.linalg.norm` by reusing the reduction result as the `sqrt` output buffer in the 2-norm and Frobenius-norm branches [#3062](https://github.com/IntelPython/dpnp/pull/3062) +* Avoided a copy of the `dpnp.einsum` result into C-order by building the product in the requested layout directly [#3069](https://github.com/IntelPython/dpnp/pull/3069) ### Deprecated From 984294f5141f43d8cf89eacac1abd0172aacd408 Mon Sep 17 00:00:00 2001 From: Vladislav Perevezentsev Date: Wed, 16 Sep 2026 01:57:38 -0700 Subject: [PATCH 3/3] Update changelog --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a11a1259406..1b1f37d7629 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -48,7 +48,7 @@ This release is compatible with NumPy 2.5. * `dpnp` uses pybind11 3.1.0 [#3015](https://github.com/IntelPython/dpnp/pull/3015) * Reworked the ASV benchmarks and added end-to-end workload benchmarks derived from dpBench [#2996](https://github.com/IntelPython/dpnp/pull/2996) * Reduced allocations in `dpnp.linalg.norm` by reusing the reduction result as the `sqrt` output buffer in the 2-norm and Frobenius-norm branches [#3062](https://github.com/IntelPython/dpnp/pull/3062) -* Avoided a copy of the `dpnp.einsum` result into C-order by building the product in the requested layout directly [#3069](https://github.com/IntelPython/dpnp/pull/3069) +* Avoided a copy of `dpnp.einsum` result into C-order by building the product in the requested layout directly [#3069](https://github.com/IntelPython/dpnp/pull/3069) ### Deprecated