Skip to content

Improve repeat using broadcasting and reshape - #318

Merged
angeloskath merged 1 commit into
mainfrom
repeat
Dec 30, 2023
Merged

angeloskath merged 1 commit into
mainfrom
repeat

Conversation

@angeloskath

Copy link
Copy Markdown
Member

Repeat was unnecessarily using concatenation and launching repeats kernel calls to do the repeat. This improves it by broadcast and reshape which maps to a single copy kernel launch regardless of the number of repeats.

For repeats values of 100 it can be up to 20x faster on my M2 air. It is the same speed for repeats values of 2 across the 1st axis which would be the best case scenario for the concatenation implementation.

@angeloskath
angeloskath requested a review from awni December 30, 2023 04:00
@angeloskath angeloskath mentioned this pull request Dec 30, 2023
4 tasks done
@awni

awni commented Dec 30, 2023

Copy link
Copy Markdown
Member

🤦 thanks for catching that!

@angeloskath
angeloskath merged commit a020a2d into main Dec 30, 2023
@angeloskath
angeloskath deleted the repeat branch December 30, 2023 05:40
JayYun98 pushed a commit to JayYun98/mlx that referenced this pull request Jan 7, 2024
zcbenz pushed a commit to Ritabanm/mlx that referenced this pull request Sep 30, 2026
Silence NonInteractiveExampleWarning in test_unary
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.

2 participants