Add a API to check SVE Length support on ARM CPU.

Signed-off-by: maajidkhann <[email protected]>
pytorch · Aug 6, 2024 · 3208502 · 3208502
1 parent ca67895
commit 3208502
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 0 deletions.
diff --git a/include/cpuinfo.h b/include/cpuinfo.h
@@ -1670,6 +1670,7 @@ struct cpuinfo_arm_isa {
 	bool sve;
 	bool sve2;
 	bool i8mm;
+	int svelen;
 #endif
 	bool rdm;
 	bool fp16arith;
@@ -2042,6 +2043,30 @@ static inline bool cpuinfo_has_arm_sve2(void) {
 #endif
 }
 
+static inline bool cpuinfo_support_arm_sve128(void) {
+#if CPUINFO_ARCH_ARM || CPUINFO_ARCH_ARM64
+    return (cpuinfo_isa.sve || cpuinfo_isa.sve2) && (cpuinfo_isa.svelen >= 16);
+#else
+    return false;
+#endif
+}
+
+static inline bool cpuinfo_support_arm_sve256(void) {
+#if CPUINFO_ARCH_ARM || CPUINFO_ARCH_ARM64
+    return (cpuinfo_isa.sve || cpuinfo_isa.sve2) && (cpuinfo_isa.svelen >= 32);
+#else
+    return false;
+#endif
+}
+
+static inline bool cpuinfo_support_arm_sve512(void) {
+#if CPUINFO_ARCH_ARM || CPUINFO_ARCH_ARM64
+    return (cpuinfo_isa.sve || cpuinfo_isa.sve2) && (cpuinfo_isa.svelen >= 64);
+#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/aarch64-isa.c b/src/arm/linux/aarch64-isa.c
@@ -3,6 +3,9 @@
 #include <arm/linux/api.h>
 #include <cpuinfo/log.h>
 
+#include <sys/prctl.h>
+#include <stdio.h>
+
 void cpuinfo_arm64_linux_decode_isa_from_proc_cpuinfo(
 	uint32_t features,
 	uint32_t features2,
@@ -138,9 +141,11 @@ void cpuinfo_arm64_linux_decode_isa_from_proc_cpuinfo(
 	}
 	if (features & CPUINFO_ARM_LINUX_FEATURE_SVE) {
 		isa->sve = true;
+		isa->svelen = -1;
 	}
 	if (features2 & CPUINFO_ARM_LINUX_FEATURE2_SVE2) {
 		isa->sve2 = true;
+		isa->svelen = -1;
 	}
 	// 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,
@@ -151,4 +156,15 @@ void cpuinfo_arm64_linux_decode_isa_from_proc_cpuinfo(
 	if (features & CPUINFO_ARM_LINUX_FEATURE_ASIMDFHM) {
 		isa->fhm = true;
 	}
+
+    #define PR_SVE_GET_VL 51
+    #define PR_SVE_VL_LEN_MASK 0xffff
+    int ret = prctl(PR_SVE_GET_VL);
+    if (ret < 0) {
+        perror("prctl(PR_SVE_GET_VL) failed");
+        isa->svelen = 0;  // Assume no SVE support if the call fails
+    } else {
+        // Mask out the SVE vector length bits
+        isa->svelen = ret & PR_SVE_VL_LEN_MASK;
+    }
 }
diff --git a/tools/isa-info.c b/tools/isa-info.c
@@ -172,6 +172,11 @@ int main(int argc, char** argv) {
 	printf("\tARM SVE: %s\n", cpuinfo_has_arm_sve() ? "yes" : "no");
 	printf("\tARM SVE 2: %s\n", cpuinfo_has_arm_sve2() ? "yes" : "no");
 
+	printf("ARM SVE Capabilities:\n");
+	printf("\tSVE 128-bit: %s\n", cpuinfo_support_arm_sve128() ? "Yes" : "No");
+	printf("\tSVE 256-bit: %s\n", cpuinfo_support_arm_sve256() ? "Yes" : "No");
+	printf("\tSVE 512-bit: %s\n", cpuinfo_support_arm_sve512() ? "Yes" : "No");
+
 	printf("Cryptography extensions:\n");
 	printf("\tAES: %s\n", cpuinfo_has_arm_aes() ? "yes" : "no");
 	printf("\tSHA1: %s\n", cpuinfo_has_arm_sha1() ? "yes" : "no");