[Hardware][Apple-CPU] Enable native bfloat16 on Apple Silicon (M2 and later) (#24129)

Signed-off-by: ignaciosica <mignacio.sica@gmail.com>
This commit is contained in:
Ignacio Sica 2025-09-10 00:50:21 -03:00 committed by GitHub
parent 7e7db04310
commit 3c2156b3af
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -75,12 +75,12 @@ class CpuPlatform(Platform):
def supported_dtypes(self) -> list[torch.dtype]: def supported_dtypes(self) -> list[torch.dtype]:
if self.get_cpu_architecture() == CpuArchEnum.POWERPC: if self.get_cpu_architecture() == CpuArchEnum.POWERPC:
return [torch.bfloat16, torch.float32] return [torch.bfloat16, torch.float32]
elif sys.platform.startswith( elif (self.get_cpu_architecture() == CpuArchEnum.ARM
"darwin") and self.get_cpu_architecture() == CpuArchEnum.ARM: and sys.platform.startswith("darwin")):
# TODO: change this condition to check if the platform support bf16 if (subprocess.check_output(
# instead of checking the OS. For instance M2 shall supports bf16 ["sysctl -n hw.optional.arm.FEAT_BF16"],
# already. But we need to modify `cpu_extension.cmake` to activate shell=True).strip() == b"1"):
# the feature in the build. return [torch.bfloat16, torch.float16, torch.float32]
return [torch.float16, torch.float32] return [torch.float16, torch.float32]
# x86/aarch64 CPU has supported both bf16 and fp16 natively. # x86/aarch64 CPU has supported both bf16 and fp16 natively.
return [torch.bfloat16, torch.float16, torch.float32] return [torch.bfloat16, torch.float16, torch.float32]