Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: performance of ufunc resolution for non-nominal signatures #3030

Merged
merged 2 commits into from
Mar 4, 2024

Conversation

agoose77
Copy link
Collaborator

@agoose77 agoose77 commented Feb 21, 2024

@lgray observed that find_ufunc is a significant bottleneck in coffea-dask workflows. Whilst we can optimise this further by changing our internal representation of behaviors (i.e. __getitem__ registers overloads in high-performance lookup structures), for now I suspect that we can improve the performance of the common-path. @jpivarski notes that the behaviour changed in this PR was never intended to be supported, so it represents a correctness improvement regardless.

The case that this PR improves is

np.add(foo, bar, ...) where neither foo nor bar are Awkward Arrays with names (e.g. __record__). This represents most ufunc operations given that even ufunc overloads of records often delegate to Awkwards' ufunc mechanism on arrays of numbers.

@lgray
Copy link
Contributor

lgray commented Feb 21, 2024

This corresponds to a ~20% improvement in execution time (O(1s) out of 5s) building a single task graph in the workflow in question, so this is a rather significant improvement.

Keeping in mind the task graph gets built 180 times or so for the full analysis.

@agoose77 agoose77 requested a review from jpivarski February 22, 2024 10:30
Copy link
Member

@jpivarski jpivarski left a comment

Choose a reason for hiding this comment

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

I think this is ready to be merged. I've check-marked that all of the rules I'd expect are included in the new implementation.

Thanks!

src/awkward/_behavior.py Show resolved Hide resolved
src/awkward/_behavior.py Show resolved Hide resolved
src/awkward/_behavior.py Show resolved Hide resolved
src/awkward/_connect/numpy.py Show resolved Hide resolved
@agoose77 agoose77 enabled auto-merge (squash) February 23, 2024 07:35
@agoose77 agoose77 merged commit ac97e2a into main Mar 4, 2024
38 checks passed
@agoose77 agoose77 deleted the agoose77/fix-performance-ufunc-resolution branch March 4, 2024 19:52
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