80 lines · python
1"""2Platform-agnostic helper to query for CPU features.3"""4 5import re6 7 8class CPUFeature:9 def __init__(self, linux_cpu_info_flag: str = None, darwin_sysctl_key: str = None):10 self.cpu_info_flag = linux_cpu_info_flag11 self.sysctl_key = darwin_sysctl_key12 13 def __str__(self):14 for arch_class in ALL_ARCHS:15 for feat_var in dir(arch_class):16 if self == getattr(arch_class, feat_var):17 return f"{arch_class.__name__}.{feat_var}"18 raise AssertionError("unreachable")19 20 def is_supported(self, triple, cmd_runner):21 if re.match(".*-.*-linux", triple):22 err_msg, res = self._is_supported_linux(cmd_runner)23 elif re.match(".*-apple-.*", triple):24 err_msg, res = self._is_supported_darwin(cmd_runner)25 else:26 err_msg, res = None, False27 28 if err_msg:29 print(f"CPU feature check failed: {err_msg}")30 31 return res32 33 def _is_supported_linux(self, cmd_runner):34 if not self.cpu_info_flag:35 return f"Unspecified cpuinfo flag for {self}", False36 37 cmd = "cat /proc/cpuinfo"38 err, retcode, output = cmd_runner(cmd)39 if err.Fail() or retcode != 0:40 return output, False41 42 # Assume that every processor presents the same features.43 # Look for the first "Features: ...." line. Features are space separated.44 if m := re.search(r"Features\s*: (.*)\n", output):45 features = m.group(1).split()46 return None, (self.cpu_info_flag in features)47 48 return 'No "Features:" line found in /proc/cpuinfo', False49 50 def _is_supported_darwin(self, cmd_runner):51 if not self.sysctl_key:52 return f"Unspecified sysctl key for {self}", False53 54 cmd = f"sysctl -n {self.sysctl_key}"55 err, retcode, output = cmd_runner(cmd)56 if err.Fail() or retcode != 0:57 return output, False58 59 return None, (output.strip() == "1")60 61 62class AArch64:63 FPMR = CPUFeature("fpmr")64 GCS = CPUFeature("gcs")65 MTE = CPUFeature("mte", "hw.optional.arm.FEAT_MTE4")66 MTE_STORE_ONLY = CPUFeature("mtestoreonly")67 PTR_AUTH = CPUFeature("paca", "hw.optional.arm.FEAT_PAuth2")68 SME = CPUFeature("sme", "hw.optional.arm.FEAT_SME")69 SME_FA64 = CPUFeature("smefa64")70 SME2 = CPUFeature("sme2", "hw.optional.arm.FEAT_SME2")71 SVE = CPUFeature("sve")72 73 74class Loong:75 LASX = CPUFeature("lasx")76 LSX = CPUFeature("lsx")77 78 79ALL_ARCHS = [AArch64, Loong]80