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.

Sign up for GitHub

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

Jump to bottom

Optionally enable KleidiAI + clean up setup.py flags #1826

Merged
merged 10 commits into from
Mar 4, 2025
Merged

Optionally enable KleidiAI + clean up setup.py flags #1826

merged 10 commits into from
Mar 4, 2025

Conversation

metascroy
Copy link
Contributor

This PR,

  • Guards universal kernels behind TORCHAO_BUILD_CPU_AARCH64
  • Cleans up CMakeExtension to include cmake flags
  • Introduces BuildOptions that control what CPU/MPS kernels are built. Universal kernels are built if TORCHAO_BUILD_CPU_AARCH64 is specified, KleidiAI kernels are built if TORCHAO_BUILD_KLEIDIAI is specified, and MPS kernels are built if TORCHAO_BUILD_EXPERIMENTAL_MPS is specified. On Arm-based Macs, TORCHAO_BUILD_CPU_AARCH64 is enabled by default when USE_CPP=1.

@pytorch-bot PyTorch Bot
Copy link

pytorch-bot bot commented Mar 4, 2025 •
edited
Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1826

Note: Links to docs will display an error until the docs builds have been completed.

⏳ No Failures, 9 Pending

As of commit fe8410b with merge base e767713 (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Mar 4, 2025
@manuelcandales manuelcandales added build topic: improvement Use this tag if this PR is an improvement (doesn't fit into any of the other categories) labels Mar 4, 2025
manuelcandales
manuelcandales reviewed Mar 4, 2025
View reviewed changes
setup.py
self.build_experimental_mps = self._os_bool_var(
"TORCHAO_BUILD_EXPERIMENTAL_MPS", default=False
)
if self.build_experimental_mps:
Copy link
Contributor

Choose a reason for hiding this comment

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

also need to assert _is_arm64 here. the experimental mps ops are intended for apple silicon

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Added arm64

@metascroy metascroy merged commit 9bcd73b into main Mar 4, 2025
51 of 52 checks passed
Jack-Khuu pushed a commit that referenced this pull request Mar 4, 2025
@metascroy @Jack-Khuu
Optionally enable KleidiAI + clean up setup.py flags (#1826)
28ccd73 
* init

* up

* up

* up

* up

* up

* up

* up

* up

* up
@metascroy metascroy mentioned this pull request Mar 4, 2025
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@manuelcandales manuelcandales manuelcandales approved these changes

@digantdesai digantdesai Awaiting requested review from digantdesai

Assignees
No one assigned
Labels
build CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. topic: improvement Use this tag if this PR is an improvement (doesn't fit into any of the other categories)
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@metascroy @manuelcandales @facebook-github-bot