diff --git a/op_builder/builder.py b/op_builder/builder.py index 465aa4399a90..ec7340410fc7 100644 --- a/op_builder/builder.py +++ b/op_builder/builder.py @@ -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, diff --git a/tests/unit/ops/test_op_builder.py b/tests/unit/ops/test_op_builder.py index 2bbfa0625f2d..941b2eb88195 100644 --- a/tests/unit/ops/test_op_builder.py +++ b/tests/unit/ops/test_op_builder.py @@ -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