Merge remote-tracking branch 'upstream/main' into patch-1

Author: Changming Sun

BUILD.bazel

@@ -99,6 +99,11 @@ WINDOWS_X86_SRCS = [
     "src/x86/windows/init.c",
 ]
 
+WINDOWS_ARM_SRCS = [
+    "src/arm/windows/init-by-logical-sys-info.c",
+    "src/arm/windows/init.c",
+]
+
 MACH_X86_SRCS = [
     "src/x86/mach/init.c",
 ]
@@ -128,6 +133,7 @@ cc_library(
         ":macos_x86_64_legacy": COMMON_SRCS + X86_SRCS + MACH_SRCS + MACH_X86_SRCS,
         ":macos_arm64": COMMON_SRCS + MACH_SRCS + MACH_ARM_SRCS,
         ":windows_x86_64": COMMON_SRCS + X86_SRCS + WINDOWS_X86_SRCS,
+        ":windows_arm64": COMMON_SRCS + ARM_SRCS + WINDOWS_ARM_SRCS,
         ":android_armv7": COMMON_SRCS + ARM_SRCS + LINUX_SRCS + LINUX_ARM32_SRCS + ANDROID_ARM_SRCS,
         ":android_arm64": COMMON_SRCS + ARM_SRCS + LINUX_SRCS + LINUX_ARM64_SRCS + ANDROID_ARM_SRCS,
         ":android_x86": COMMON_SRCS + X86_SRCS + LINUX_SRCS + LINUX_X86_SRCS,
@@ -149,6 +155,7 @@ cc_library(
     }),
     copts = select({
         ":windows_x86_64": [],
+        ":windows_arm64": [],
         "//conditions:default": C99OPTS,
     }) + [
         "-Iexternal/cpuinfo/include",
@@ -281,6 +288,11 @@ config_setting(
     values = {"cpu": "x64_windows"},
 )
 
+config_setting(
+    name = "windows_arm64",
+    values = {"cpu": "arm64_windows"},
+)
+
 config_setting(
     name = "android_armv7",
     values = {

MODULE.bazel

@@ -0,0 +1 @@
+module(name = "cpuinfo")

WORKSPACE.bazel

@@ -1 +1 @@
-workspace(name = "org_pytorch_cpuinfo")
+workspace(name = "org_pytorch_cpuinfo")

include/cpuinfo.h

@@ -813,6 +813,13 @@ struct cpuinfo_x86_isa {
 	bool avx512vp2intersect;
 	bool avx512_4vnniw;
 	bool avx512_4fmaps;
+	bool amx_bf16;
+	bool amx_tile;
+	bool amx_int8;
+	bool amx_fp16;
+	bool avx_vnni_int8;
+	bool avx_vnni_int16;
+	bool avx_ne_convert;
 	bool hle;
 	bool rtm;
 	bool xtest;
@@ -1329,6 +1336,98 @@ static inline bool cpuinfo_has_x86_avx512_4fmaps(void) {
 #endif
 }
 
+/* [NOTE] Intel Advanced Matrix Extensions (AMX) detection
+ *
+ * I.  AMX is a new extensions to the x86 ISA to work on matrices, consists of
+ *   1) 2-dimentional registers (tiles), hold sub-matrices from larger matrices in memory
+ *   2) Accelerator called Tile Matrix Multiply (TMUL), contains instructions operating on tiles
+ *
+ * II. Platforms that supports AMX:
+ * +-----------------+-----+----------+----------+----------+----------+
+ * |    Platforms    | Gen | amx-bf16 | amx-tile | amx-int8 | amx-fp16 |
+ * +-----------------+-----+----------+----------+----------+----------+
+ * | Sapphire Rapids | 4th |   YES    |   YES    |   YES    |    NO    |
+ * +-----------------+-----+----------+----------+----------+----------+
+ * | Emerald Rapids  | 5th |   YES    |   YES    |   YES    |    NO    |
+ * +-----------------+-----+----------+----------+----------+----------+
+ * | Granite Rapids  | 6th |   YES    |   YES    |   YES    |   YES    |
+ * +-----------------+-----+----------+----------+----------+----------+
+ *
+ * Reference: https://www.intel.com/content/www/us/en/products/docs
+ *    /accelerator-engines/advanced-matrix-extensions/overview.html
+ */
+static inline bool cpuinfo_has_x86_amx_bf16(void) {
+#if CPUINFO_ARCH_X86 || CPUINFO_ARCH_X86_64
+	return cpuinfo_isa.amx_bf16;
+#else
+	return false;
+#endif
+}
+
+static inline bool cpuinfo_has_x86_amx_tile(void) {
+#if CPUINFO_ARCH_X86 || CPUINFO_ARCH_X86_64
+	return cpuinfo_isa.amx_tile;
+#else
+	return false;
+#endif
+}
+
+static inline bool cpuinfo_has_x86_amx_int8(void) {
+#if CPUINFO_ARCH_X86 || CPUINFO_ARCH_X86_64
+	return cpuinfo_isa.amx_int8;
+#else
+	return false;
+#endif
+}
+
+static inline bool cpuinfo_has_x86_amx_fp16(void) {
+#if CPUINFO_ARCH_X86 || CPUINFO_ARCH_X86_64
+	return cpuinfo_isa.amx_fp16;
+#else
+	return false;
+#endif
+}
+
+/*
+ * Intel AVX Vector Neural Network Instructions (VNNI) INT8
+ * Supported Platfroms: Sierra Forest, Arrow Lake, Lunar Lake
+ */
+static inline bool cpuinfo_has_x86_avx_vnni_int8(void) {
+#if CPUINFO_ARCH_X86 || CPUINFO_ARCH_X86_64
+	return cpuinfo_isa.avx_vnni_int8;
+#else
+	return false;
+#endif
+}
+
+/*
+ * Intel AVX Vector Neural Network Instructions (VNNI) INT16
+ * Supported Platfroms: Arrow Lake, Lunar Lake
+ */
+static inline bool cpuinfo_has_x86_avx_vnni_int16(void) {
+#if CPUINFO_ARCH_X86 || CPUINFO_ARCH_X86_64
+	return cpuinfo_isa.avx_vnni_int16;
+#else
+	return false;
+#endif
+}
+
+/*
+ * A new set of instructions, which can convert low precision floating point
+ * like BF16/FP16 to high precision floating point FP32, as well as convert FP32
+ * elements to BF16. This instruction allows the platform to have improved AI
+ * capabilities and better compatibility.
+ *
+ * Supported Platforms: Sierra Forest, Arrow Lake, Lunar Lake
+ */
+static inline bool cpuinfo_has_x86_avx_ne_convert(void) {
+#if CPUINFO_ARCH_X86 || CPUINFO_ARCH_X86_64
+	return cpuinfo_isa.avx_ne_convert;
+#else
+	return false;
+#endif
+}
+
 static inline bool cpuinfo_has_x86_hle(void) {
 #if CPUINFO_ARCH_X86 || CPUINFO_ARCH_X86_64
 	return cpuinfo_isa.hle;

src/arm/windows/init.c

@@ -20,42 +20,56 @@ static struct woa_chip_info woa_chip_unknown = {
 	{{cpuinfo_vendor_unknown, cpuinfo_uarch_unknown, 0}}};
 
 /* Please add new SoC/chip info here! */
-static struct woa_chip_info woa_chips[] = {
+static struct woa_chip_info woa_chips[woa_chip_name_last] = {
 	/* Microsoft SQ1 Kryo 495 4 + 4 cores (3 GHz + 1.80 GHz) */
-	{L"Microsoft SQ1",
-	 woa_chip_name_microsoft_sq_1,
-	 {{
-		  cpuinfo_vendor_arm,
-		  cpuinfo_uarch_cortex_a55,
-		  1800000000,
-	  },
-	  {
-		  cpuinfo_vendor_arm,
-		  cpuinfo_uarch_cortex_a76,
-		  3000000000,
-	  }}},
+	[woa_chip_name_microsoft_sq_1] =
+		{L"Microsoft SQ1",
+		 woa_chip_name_microsoft_sq_1,
+		 {{
+			  cpuinfo_vendor_arm,
+			  cpuinfo_uarch_cortex_a55,
+			  1800000000,
+		  },
+		  {
+			  cpuinfo_vendor_arm,
+			  cpuinfo_uarch_cortex_a76,
+			  3000000000,
+		  }}},
 	/* Microsoft SQ2 Kryo 495 4 + 4 cores (3.15 GHz + 2.42 GHz) */
-	{L"Microsoft SQ2",
-	 woa_chip_name_microsoft_sq_2,
-	 {{
-		  cpuinfo_vendor_arm,
-		  cpuinfo_uarch_cortex_a55,
-		  2420000000,
-	  },
-	  {cpuinfo_vendor_arm, cpuinfo_uarch_cortex_a76, 3150000000}}},
+	[woa_chip_name_microsoft_sq_2] =
+		{L"Microsoft SQ2",
+		 woa_chip_name_microsoft_sq_2,
+		 {{
+			  cpuinfo_vendor_arm,
+			  cpuinfo_uarch_cortex_a55,
+			  2420000000,
+		  },
+		  {cpuinfo_vendor_arm, cpuinfo_uarch_cortex_a76, 3150000000}}},
+	/* Snapdragon (TM) 8cx Gen 3 @ 3.0 GHz */
+	[woa_chip_name_microsoft_sq_3] =
+		{L"Snapdragon (TM) 8cx Gen 3",
+		 woa_chip_name_microsoft_sq_3,
+		 {{
+			  cpuinfo_vendor_arm,
+			  cpuinfo_uarch_cortex_a78,
+			  2420000000,
+		  },
+		  {cpuinfo_vendor_arm, cpuinfo_uarch_cortex_x1, 3000000000}}},
 	/* Microsoft Windows Dev Kit 2023 */
-	{L"Snapdragon Compute Platform",
-	 woa_chip_name_microsoft_sq_3,
-	 {{
-		  cpuinfo_vendor_arm,
-		  cpuinfo_uarch_cortex_a78,
-		  2420000000,
-	  },
-	  {cpuinfo_vendor_arm, cpuinfo_uarch_cortex_x1, 3000000000}}},
+	[woa_chip_name_microsoft_sq_3_devkit] =
+		{L"Snapdragon Compute Platform",
+		 woa_chip_name_microsoft_sq_3_devkit,
+		 {{
+			  cpuinfo_vendor_arm,
+			  cpuinfo_uarch_cortex_a78,
+			  2420000000,
+		  },
+		  {cpuinfo_vendor_arm, cpuinfo_uarch_cortex_x1, 3000000000}}},
 	/* Ampere Altra */
-	{L"Ampere(R) Altra(R) Processor",
-	 woa_chip_name_ampere_altra,
-	 {{cpuinfo_vendor_arm, cpuinfo_uarch_neoverse_n1, 3000000000}}}};
+	[woa_chip_name_ampere_altra] = {
+		L"Ampere(R) Altra(R) Processor",
+		woa_chip_name_ampere_altra,
+		{{cpuinfo_vendor_arm, cpuinfo_uarch_neoverse_n1, 3000000000}}}};
 
 BOOL CALLBACK cpuinfo_arm_windows_init(PINIT_ONCE init_once, PVOID parameter, PVOID* context) {
 	struct woa_chip_info* chip_info = NULL;

src/arm/windows/windows-arm-init.h

@@ -8,8 +8,9 @@ enum woa_chip_name {
 	woa_chip_name_microsoft_sq_1 = 0,
 	woa_chip_name_microsoft_sq_2 = 1,
 	woa_chip_name_microsoft_sq_3 = 2,
-	woa_chip_name_ampere_altra = 3,
-	woa_chip_name_unknown = 4,
+	woa_chip_name_microsoft_sq_3_devkit = 3,
+	woa_chip_name_ampere_altra = 4,
+	woa_chip_name_unknown = 5,
 	woa_chip_name_last = woa_chip_name_unknown
 };
 

src/x86/isa.c

@@ -537,6 +537,48 @@ struct cpuinfo_x86_isa cpuinfo_x86_detect_isa(
 	 */
 	isa.avx512bf16 = avx512_regs && !!(structured_feature_info1.eax & UINT32_C(0x00000020));
 
+	/*
+	 * AMX_BF16 instructions:
+	 * - Intel: edx[bit 22] in structured feature info (ecx = 0).
+	 */
+	isa.amx_bf16 = avx512_regs && !!(structured_feature_info0.edx & UINT32_C(0x00400000));
+
+	/*
+	 * AMX_TILE instructions:
+	 * - Intel: edx[bit 24] in structured feature info (ecx = 0).
+	 */
+	isa.amx_tile = avx512_regs && !!(structured_feature_info0.edx & UINT32_C(0x01000000));
+
+	/*
+	 * AMX_INT8 instructions:
+	 * - Intel: edx[bit 25] in structured feature info (ecx = 0).
+	 */
+	isa.amx_int8 = avx512_regs && !!(structured_feature_info0.edx & UINT32_C(0x02000000));
+
+	/*
+	 * AMX_FP16 instructions:
+	 * - Intel: eax[bit 21] in structured feature info (ecx = 1).
+	 */
+	isa.amx_fp16 = avx512_regs && !!(structured_feature_info1.eax & UINT32_C(0x00200000));
+
+	/*
+	 * AVX_VNNI_INT8 instructions:
+	 * - Intel: edx[bit 4] in structured feature info (ecx = 1).
+	 */
+	isa.avx_vnni_int8 = avx_regs && !!(structured_feature_info1.edx & UINT32_C(0x00000010));
+
+	/*
+	 * AVX_VNNI_INT16 instructions:
+	 * - Intel: edx[bit 10] in structured feature info (ecx = 1).
+	 */
+	isa.avx_vnni_int16 = avx_regs && !!(structured_feature_info1.edx & UINT32_C(0x00000400));
+
+	/*
+	 * AVX_NE_CONVERT instructions:
+	 * - Intel: edx[bit 5] in structured feature info (ecx = 1).
+	 */
+	isa.avx_ne_convert = avx_regs && !!(structured_feature_info1.edx & UINT32_C(0x00000020));
+
 	/*
 	 * HLE instructions:
 	 * - Intel: ebx[bit 4] in structured feature info (ecx = 0).

tools/isa-info.c

@@ -70,7 +70,14 @@ int main(int argc, char** argv) {
 	printf("\tAVX512VP2INTERSECT: %s\n", cpuinfo_has_x86_avx512vp2intersect() ? "yes" : "no");
 	printf("\tAVX512_4VNNIW: %s\n", cpuinfo_has_x86_avx512_4vnniw() ? "yes" : "no");
 	printf("\tAVX512_4FMAPS: %s\n", cpuinfo_has_x86_avx512_4fmaps() ? "yes" : "no");
+	printf("\tAMX_BF16: %s\n", cpuinfo_has_x86_amx_bf16() ? "yes" : "no");
+	printf("\tAMX_TILE: %s\n", cpuinfo_has_x86_amx_tile() ? "yes" : "no");
+	printf("\tAMX_INT8: %s\n", cpuinfo_has_x86_amx_int8() ? "yes" : "no");
+	printf("\tAMX_FP16: %s\n", cpuinfo_has_x86_amx_fp16() ? "yes" : "no");
 	printf("\tAVXVNNI: %s\n", cpuinfo_has_x86_avxvnni() ? "yes" : "no");
+	printf("\tAVX_VNNI_INT8: %s\n", cpuinfo_has_x86_avx_vnni_int8() ? "yes" : "no");
+	printf("\tAVX_VNNI_INT16: %s\n", cpuinfo_has_x86_avx_vnni_int16() ? "yes" : "no");
+	printf("\tAVX_NE_CONVERT: %s\n", cpuinfo_has_x86_avx_ne_convert() ? "yes" : "no");
 
 	printf("Multi-threading extensions:\n");
 	printf("\tMONITOR/MWAIT: %s\n", cpuinfo_has_x86_mwait() ? "yes" : "no");