Skip to content

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

Merged
metascroy merged 10 commits into
mainfrom
add-bf16-kleidi
Mar 4, 2025
Merged

Optionally enable KleidiAI + clean up setup.py flags#1826
metascroy merged 10 commits into
mainfrom
add-bf16-kleidi

Conversation

@metascroy

Copy link
Copy Markdown
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 Bot commented Mar 4, 2025

Copy link
Copy Markdown

🔗 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
Comment thread setup.py
self.build_experimental_mps = self._os_bool_var(
"TORCHAO_BUILD_EXPERIMENTAL_MPS", default=False
)
if self.build_experimental_mps:

Copy link
Copy Markdown
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
Copy Markdown
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
Jack-Khuu pushed a commit that referenced this pull request Mar 4, 2025
* init

* up

* up

* up

* up

* up

* up

* up

* up

* up
liangel-02 pushed a commit that referenced this pull request Aug 25, 2025
* init

* up

* up

* up

* up

* up

* up

* up

* up

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

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

Development

Successfully merging this pull request may close these issues.

3 participants