Merge pull request #262 from gonnet/main

Add `sme2` detection for `aarch64`
pytorch · Sep 26, 2024 · 1e83a2f · 1e83a2f
2 parents a5ff6df + 8f02f6d
commit 1e83a2f
Show file tree

Hide file tree

Showing 7 changed files with 99 additions and 19 deletions.
diff --git a/include/cpuinfo-mock.h b/include/cpuinfo-mock.h
@@ -60,7 +60,7 @@ ssize_t CPUINFO_ABI cpuinfo_mock_read(int fd, void* buffer, size_t capacity);
 void CPUINFO_ABI cpuinfo_set_hwcap(uint32_t hwcap);
 #endif
 #if CPUINFO_ARCH_ARM
-void CPUINFO_ABI cpuinfo_set_hwcap2(uint32_t hwcap2);
+void CPUINFO_ABI cpuinfo_set_hwcap2(uint64_t hwcap2);
 #endif
 #endif
 

diff --git a/include/cpuinfo.h b/include/cpuinfo.h
@@ -1671,6 +1671,12 @@ struct cpuinfo_arm_isa {
 	bool sve2;
 	bool i8mm;
 	bool sme;
+	bool sme2;
+	bool sme2p1;
+	bool sme_i16i32;
+	bool sme_bi32i32;
+	bool sme_b16b16;
+	bool sme_f16f16;
 	uint32_t svelen;
 #endif
 	bool rdm;
@@ -2061,6 +2067,54 @@ static inline bool cpuinfo_has_arm_sme(void) {
 #endif
 }
 
+static inline bool cpuinfo_has_arm_sme2(void) {
+#if CPUINFO_ARCH_ARM64
+	return cpuinfo_isa.sme2;
+#else
+	return false;
+#endif
+}
+
+static inline bool cpuinfo_has_arm_sme2p1(void) {
+#if CPUINFO_ARCH_ARM64
+	return cpuinfo_isa.sme2p1;
+#else
+	return false;
+#endif
+}
+
+static inline bool cpuinfo_has_arm_sme_i16i32(void) {
+#if CPUINFO_ARCH_ARM64
+	return cpuinfo_isa.sme_i16i32;
+#else
+	return false;
+#endif
+}
+
+static inline bool cpuinfo_has_arm_sme_bi32i32(void) {
+#if CPUINFO_ARCH_ARM64
+	return cpuinfo_isa.sme_bi32i32;
+#else
+	return false;
+#endif
+}
+
+static inline bool cpuinfo_has_arm_sme_b16b16(void) {
+#if CPUINFO_ARCH_ARM64
+	return cpuinfo_isa.sme_b16b16;
+#else
+	return false;
+#endif
+}
+
+static inline bool cpuinfo_has_arm_sme_f16f16(void) {
+#if CPUINFO_ARCH_ARM64
+	return cpuinfo_isa.sme_f16f16;
+#else
+	return false;
+#endif
+}
+
 #if CPUINFO_ARCH_RISCV32 || CPUINFO_ARCH_RISCV64
 /* This structure is not a part of stable API. Use cpuinfo_has_riscv_* functions
  * instead. */

diff --git a/src/arm/linux/aarch32-isa.c b/src/arm/linux/aarch32-isa.c
@@ -24,7 +24,7 @@ void cpuinfo_set_wcid(uint32_t wcid) {
 
 void cpuinfo_arm_linux_decode_isa_from_proc_cpuinfo(
 	uint32_t features,
-	uint32_t features2,
+	uint64_t features2,
 	uint32_t midr,
 	uint32_t architecture_version,
 	uint32_t architecture_flags,

diff --git a/src/arm/linux/aarch64-isa.c b/src/arm/linux/aarch64-isa.c
@@ -7,7 +7,7 @@
 
 void cpuinfo_arm64_linux_decode_isa_from_proc_cpuinfo(
 	uint32_t features,
-	uint32_t features2,
+	uint64_t features2,
 	uint32_t midr,
 	const struct cpuinfo_arm_chipset chipset[restrict static 1],
 	struct cpuinfo_arm_isa isa[restrict static 1]) {
@@ -147,6 +147,24 @@ void cpuinfo_arm64_linux_decode_isa_from_proc_cpuinfo(
 	if (features2 & CPUINFO_ARM_LINUX_FEATURE2_SME) {
 		isa->sme = true;
 	}
+	if (features2 & CPUINFO_ARM_LINUX_FEATURE2_SME2) {
+		isa->sme2 = true;
+	}
+	if (features2 & CPUINFO_ARM_LINUX_FEATURE2_SME2P1) {
+		isa->sme2p1 = true;
+	}
+	if (features2 & CPUINFO_ARM_LINUX_FEATURE2_SME_I16I32) {
+		isa->sme_i16i32 = true;
+	}
+	if (features2 & CPUINFO_ARM_LINUX_FEATURE2_SME_BI32I32) {
+		isa->sme_bi32i32 = true;
+	}
+	if (features2 & CPUINFO_ARM_LINUX_FEATURE2_SME_B16B16) {
+		isa->sme_b16b16 = true;
+	}
+	if (features2 & CPUINFO_ARM_LINUX_FEATURE2_SME_F16F16) {
+		isa->sme_f16f16 = true;
+	}
 	// SVEBF16 is set iff SVE and BF16 are both supported, but the SVEBF16
 	// feature flag was added in Linux kernel before the BF16 feature flag,
 	// so we check for either.

diff --git a/src/arm/linux/api.h b/src/arm/linux/api.h
@@ -138,6 +138,12 @@ struct cpuinfo_arm_linux_proc_cpuinfo_cache {
 #define CPUINFO_ARM_LINUX_FEATURE2_RNG UINT32_C(0x00010000)
 #define CPUINFO_ARM_LINUX_FEATURE2_BTI UINT32_C(0x00020000)
 #define CPUINFO_ARM_LINUX_FEATURE2_SME UINT32_C(0x00800000)
+#define CPUINFO_ARM_LINUX_FEATURE2_SME2 UINT64_C(0x0000002000000000)
+#define CPUINFO_ARM_LINUX_FEATURE2_SME2P1 UINT64_C(0x0000004000000000)
+#define CPUINFO_ARM_LINUX_FEATURE2_SME_I16I32 UINT64_C(0x0000008000000000)
+#define CPUINFO_ARM_LINUX_FEATURE2_SME_BI32I32 UINT64_C(0x0000010000000000)
+#define CPUINFO_ARM_LINUX_FEATURE2_SME_B16B16 UINT64_C(0x0000020000000000)
+#define CPUINFO_ARM_LINUX_FEATURE2_SME_F16F16 UINT64_C(0x0000040000000000)
 #endif
 
 #define CPUINFO_ARM_LINUX_VALID_ARCHITECTURE UINT32_C(0x00010000)
@@ -173,7 +179,7 @@ struct cpuinfo_arm_linux_processor {
 	struct cpuinfo_arm_linux_proc_cpuinfo_cache proc_cpuinfo_cache;
 #endif
 	uint32_t features;
-	uint32_t features2;
+	uint64_t features2;
 	/**
 	 * Main ID Register value.
 	 */
@@ -296,14 +302,14 @@ CPUINFO_INTERNAL bool cpuinfo_arm_linux_parse_proc_cpuinfo(
 #if CPUINFO_ARCH_ARM
 CPUINFO_INTERNAL bool cpuinfo_arm_linux_hwcap_from_getauxval(
 	uint32_t hwcap[restrict static 1],
-	uint32_t hwcap2[restrict static 1]);
+	uint64_t hwcap2[restrict static 1]);
 CPUINFO_INTERNAL bool cpuinfo_arm_linux_hwcap_from_procfs(
 	uint32_t hwcap[restrict static 1],
-	uint32_t hwcap2[restrict static 1]);
+	uint64_t hwcap2[restrict static 1]);
 
 CPUINFO_INTERNAL void cpuinfo_arm_linux_decode_isa_from_proc_cpuinfo(
 	uint32_t features,
-	uint32_t features2,
+	uint64_t features2,
 	uint32_t midr,
 	uint32_t architecture_version,
 	uint32_t architecture_flags,
@@ -312,11 +318,11 @@ CPUINFO_INTERNAL void cpuinfo_arm_linux_decode_isa_from_proc_cpuinfo(
 #elif CPUINFO_ARCH_ARM64
 CPUINFO_INTERNAL void cpuinfo_arm_linux_hwcap_from_getauxval(
 	uint32_t hwcap[restrict static 1],
-	uint32_t hwcap2[restrict static 1]);
+	uint64_t hwcap2[restrict static 1]);
 
 CPUINFO_INTERNAL void cpuinfo_arm64_linux_decode_isa_from_proc_cpuinfo(
 	uint32_t features,
-	uint32_t features2,
+	uint64_t features2,
 	uint32_t midr,
 	const struct cpuinfo_arm_chipset chipset[restrict static 1],
 	struct cpuinfo_arm_isa isa[restrict static 1]);

diff --git a/src/arm/linux/hwcap.c b/src/arm/linux/hwcap.c
@@ -31,16 +31,16 @@ void cpuinfo_set_hwcap(uint32_t hwcap) {
 	mock_hwcap = hwcap;
 }
 
-static uint32_t mock_hwcap2 = 0;
-void cpuinfo_set_hwcap2(uint32_t hwcap2) {
+static uint64_t mock_hwcap2 = 0;
+void cpuinfo_set_hwcap2(uint64_t hwcap2) {
 	mock_hwcap2 = hwcap2;
 }
 #endif
 
 #if CPUINFO_ARCH_ARM
 typedef unsigned long (*getauxval_function_t)(unsigned long);
 
-bool cpuinfo_arm_linux_hwcap_from_getauxval(uint32_t hwcap[restrict static 1], uint32_t hwcap2[restrict static 1]) {
+bool cpuinfo_arm_linux_hwcap_from_getauxval(uint32_t hwcap[restrict static 1], uint64_t hwcap2[restrict static 1]) {
 #if CPUINFO_MOCK
 	*hwcap = mock_hwcap;
 	*hwcap2 = mock_hwcap2;
@@ -83,13 +83,13 @@ bool cpuinfo_arm_linux_hwcap_from_getauxval(uint32_t hwcap[restrict static 1], u
 }
 
 #ifdef __ANDROID__
-bool cpuinfo_arm_linux_hwcap_from_procfs(uint32_t hwcap[restrict static 1], uint32_t hwcap2[restrict static 1]) {
+bool cpuinfo_arm_linux_hwcap_from_procfs(uint32_t hwcap[restrict static 1], uint64_t hwcap2[restrict static 1]) {
 #if CPUINFO_MOCK
 	*hwcap = mock_hwcap;
 	*hwcap2 = mock_hwcap2;
 	return true;
 #else
-	uint32_t hwcaps[2] = {0, 0};
+	uint64_t hwcaps[2] = {0, 0};
 	bool result = false;
 	int file = -1;
 
@@ -113,7 +113,7 @@ bool cpuinfo_arm_linux_hwcap_from_procfs(uint32_t hwcap[restrict static 1], uint
 						hwcaps[0] = (uint32_t)elf_auxv.a_un.a_val;
 						break;
 					case AT_HWCAP2:
-						hwcaps[1] = (uint32_t)elf_auxv.a_un.a_val;
+						hwcaps[1] = (uint64_t)elf_auxv.a_un.a_val;
 						break;
 				}
 			} else {
@@ -141,13 +141,13 @@ bool cpuinfo_arm_linux_hwcap_from_procfs(uint32_t hwcap[restrict static 1], uint
 }
 #endif /* __ANDROID__ */
 #elif CPUINFO_ARCH_ARM64
-void cpuinfo_arm_linux_hwcap_from_getauxval(uint32_t hwcap[restrict static 1], uint32_t hwcap2[restrict static 1]) {
+void cpuinfo_arm_linux_hwcap_from_getauxval(uint32_t hwcap[restrict static 1], uint64_t hwcap2[restrict static 1]) {
 #if CPUINFO_MOCK
 	*hwcap = mock_hwcap;
 	*hwcap2 = mock_hwcap2;
 #else
 	*hwcap = (uint32_t)getauxval(AT_HWCAP);
-	*hwcap2 = (uint32_t)getauxval(AT_HWCAP2);
+	*hwcap2 = (uint64_t)getauxval(AT_HWCAP2);
 	return;
 #endif
 }

diff --git a/src/arm/linux/init.c b/src/arm/linux/init.c
@@ -247,7 +247,8 @@ void cpuinfo_arm_linux_init(void) {
 #endif
 
 #if CPUINFO_ARCH_ARM
-	uint32_t isa_features = 0, isa_features2 = 0;
+	uint32_t isa_features = 0;
+	uint64_t isa_features2 = 0;
 #ifdef __ANDROID__
 	/*
 	 * On Android before API 20, libc.so does not provide getauxval
@@ -299,7 +300,8 @@ void cpuinfo_arm_linux_init(void) {
 		&chipset,
 		&cpuinfo_isa);
 #elif CPUINFO_ARCH_ARM64
-	uint32_t isa_features = 0, isa_features2 = 0;
+	uint32_t isa_features = 0;
+	uint64_t isa_features2 = 0;
 	/* getauxval is always available on ARM64 Android */
 	cpuinfo_arm_linux_hwcap_from_getauxval(&isa_features, &isa_features2);
 	cpuinfo_arm64_linux_decode_isa_from_proc_cpuinfo(