Skip to content

Detect NEON when py-cpuinfo reports no CPU flags #8765

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

Draft
vineethsaivs wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

Detect NEON when py-cpuinfo reports no CPU flags #8765

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions op_builder/builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -505,12 +505,14 @@ def simd_width(self):
if cpu_info is None:
return '-D__SCALAR__'

# py-cpuinfo reports no 'flags' on some CPUs, e.g. Apple Silicon.
flags = cpu_info.get('flags', [])
if cpu_info['arch'] == 'X86_64':
if 'avx512' in cpu_info['flags'] or 'avx512f' in cpu_info['flags']:
if 'avx512' in flags or 'avx512f' in flags:
return '-D__AVX512__'
elif 'avx2' in cpu_info['flags']:
elif 'avx2' in flags:
return '-D__AVX256__'
elif (cpu_info['arch'] or '').startswith('ARM') and 'sve' in cpu_info['flags']:
elif (cpu_info['arch'] or '').startswith('ARM') and 'sve' in flags:
return '-D__SVE__'
elif cpu_info['arch'] == 'ARM_8':
# NEON is baseline on AArch64; vdivq/vsqrtq used by simd.h are A64-only,
Expand Down
8 changes: 8 additions & 0 deletions tests/unit/ops/test_op_builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -338,3 +338,11 @@ def test_probe_is_compatible_does_not_propagate_a_failed_probe():
@pytest.mark.parametrize("answer", [True, False])
def test_probe_is_compatible_passes_through_a_successful_probe(answer):
assert probe_is_compatible(_AnsweringBuilder(answer)) is answer


@pytest.mark.parametrize("flags, expected", [(None, "-D__NEON__"), (["fp", "asimd", "sve"], "-D__SVE__")])
def test_simd_width_handles_cpu_info_without_flags(flags, expected):
# py-cpuinfo leaves out 'flags' on Apple Silicon, so simd_width must not index it.
cpu_info = {"arch": "ARM_8"} if flags is None else {"arch": "ARM_8", "flags": flags}
with patch.dict(sys.modules, {"cpuinfo": MagicMock(get_cpu_info=MagicMock(return_value=cpu_info))}):
assert make_builder().simd_width() == expected