Skip to content

Avoid a C-order copy in dpnp.einsum - #3069

Open
vlad-perevezentsev wants to merge 3 commits into
masterfrom
fix_perf_einsum
Open

vlad-perevezentsev wants to merge 3 commits into
masterfrom
fix_perf_einsum

Conversation

@vlad-perevezentsev

Copy link
Copy Markdown
Contributor

This PR improves dpnp.einsum performance by avoiding a copy into C-order.

Commit 039b7f1 (#3058) remaps order="K" to "C" which turned the final dpnp.asarray from a no-op into a full device copy because dpnp contracts operands in reverse order so "ij,jk" computes (a @ b).T and the result is naturally f-contiguous

This PR builds the result directly in the requested layout instead of fixing it afterwards. The last contraction determines the output axis order, so it can choose which operand is on the left side of the matmul and avoid an extra copy.
This is only applied for "C" order, since for other layouts the default operand order already matches the expected order="K" + optimize=True behavior. The swap does not move any data. It only changes the operand order in the matmul

The performance result bench_linalg.MatMul.time_einsum_ij_jk on PVC:

dtype size main fix
float64 1024 0.953 ms 0.612 ms (x0.64)
float64 4096 8.590 ms 7.490 ms (x0.87)
float32 1024 0.574 ms 0.439 ms (x0.76)
float32 4096 7.577 ms 7.005 ms (x0.92)
int32 1024 1.891 ms 1.301 ms (x0.69)
int32 4096 54.324 ms 35.045 ms (x0.65)
int64 1024 4.637 ms 3.735 ms (x0.81)
int64 4096 200.454 ms 162.920 ms (x0.81)
  • Have you provided a meaningful PR description?
  • Have you added a test, reproducer or referred to an issue with a reproducer?
  • Have you tested your changes locally for CPU and GPU devices?
  • Have you made sure that new changes do not introduce compiler warnings?
  • Have you checked performance impact of proposed changes?
  • Have you added documentation for your changes, if necessary?
  • Have you added your changes to the changelog?

@github-actions

Copy link
Copy Markdown
Contributor

View rendered docs @ https://intelpython.github.io/dpnp/pull/3069/index.html

@github-actions

github-actions Bot commented Sep 14, 2026

Copy link
Copy Markdown
Contributor

Array API standard conformance tests for dpnp=0.21.0dev8=py314ha0e2e8e_11 ran successfully.
Passed: 1376
Failed: 0
Skipped: 6

@coveralls

coveralls commented Sep 14, 2026

Copy link
Copy Markdown
Collaborator

Coverage Status

coverage: 78.48% (+0.003%) from 78.477% — fix_perf_einsum into master

@ndgrigorian ndgrigorian left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants