diff --git a/devito/arch/archinfo.py b/devito/arch/archinfo.py index 53de39e892..c294f5def4 100644 --- a/devito/arch/archinfo.py +++ b/devito/arch/archinfo.py @@ -720,6 +720,15 @@ def lscpu(): return {} +def get_apple_platform(brand): + """ + The registered AppleArm matching an Apple-silicon brand string, e.g. + "Apple M2 Pro". Chips not registered yet map to the newest registered one. + """ + key = brand.lower().split('apple ')[1].split()[0] + return platform_registry.get(key, platform_registry['m3']) + + @memoized_func def get_platform(): """Attempt Platform autodetection.""" @@ -753,8 +762,8 @@ def get_platform(): return platform_registry['power8'] elif 'arm' in brand: return platform_registry['arm'] - elif 'm1' in brand: - return platform_registry['m1'] + elif 'apple m' in brand: + return get_apple_platform(brand) elif 'amd' in brand: return platform_registry['amd'] except: diff --git a/tests/test_arch.py b/tests/test_arch.py index 40fa913ab9..22d084cdc1 100644 --- a/tests/test_arch.py +++ b/tests/test_arch.py @@ -6,6 +6,7 @@ from devito import configuration, switchconfig from devito.arch import BLACKWELL, NVIDIAX +from devito.arch.archinfo import AppleArm, get_apple_platform from devito.arch.compiler import ( CudaCompiler, GNUCompiler, compiler_registry, sniff_compiler_version ) @@ -25,6 +26,20 @@ def test_gcc(cc): assert cc in compiler_registry +@pytest.mark.parametrize("brand, expected", [ + ("Apple M1", 'm1'), + ("Apple M1 (Virtual)", 'm1'), + ("Apple M2 Pro", 'm2'), + ("Apple M3 Max", 'm3'), + ("Apple M4 Max", 'm3'), + ("Apple M5", 'm3'), +]) +def test_apple_platform(brand, expected): + platform = get_apple_platform(brand) + assert isinstance(platform, AppleArm) + assert platform.name == expected + + def test_switcharch(): old_compiler = configuration['compiler'] with switchconfig(compiler='gcc-4.9'):