diff --git a/CHANGELOG.md b/CHANGELOG.md index 0f47c903d54..1b1f37d7629 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 `dpnp.einsum` result into C-order by building the product in the requested layout directly [#3069](https://github.com/IntelPython/dpnp/pull/3069) ### Deprecated 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)