Make exegesis annotator support AVX512BW registers

We currently abort if we have an unhandled register, which includes the
AVX512BW registers before this patch. This patch changes the SSE/AVX2
register classes to the AVX512 register versions so we pick up the
registers as expected.

gematria/datasets/find_accessed_addrs_exegesis.cc

@@ -169,8 +169,8 @@ Expected<ExecutionAnnotations> ExegesisAnnotator::findAccessedAddrs(
 
     if (MRI.getRegClass(X86::GR64_NOREX2RegClassID).contains(RegisterIndex) ||
         MRI.getRegClass(X86::VR64RegClassID).contains(RegisterIndex) ||
-        MRI.getRegClass(X86::VR128RegClassID).contains(RegisterIndex) ||
-        MRI.getRegClass(X86::VR256RegClassID).contains(RegisterIndex) ||
+        MRI.getRegClass(X86::VR128XRegClassID).contains(RegisterIndex) ||
+        MRI.getRegClass(X86::VR256XRegClassID).contains(RegisterIndex) ||
         MRI.getRegClass(X86::VR512RegClassID).contains(RegisterIndex) ||
         MRI.getRegClass(X86::VK64RegClassID).contains(RegisterIndex)) {
       RegVal.Value = APInt(64, kInitialRegVal);

gematria/datasets/find_accessed_addrs_exegesis_test.cc

@@ -272,5 +272,25 @@ TEST_F(FindAccessedAddrsExegesisTest, STRegister) {
                                 &RegisterAndValue::register_name, "ST1")));
 }
 
+TEST_F(FindAccessedAddrsExegesisTest, AVX512XMMYMMRegisters) {
+  // This test requires AVX512, skip if we are not running on a CPU with
+  // AVX512F.
+  if (!__builtin_cpu_supports("avx512f")) {
+    GTEST_SKIP() << "CPU does not support AVX512, skipping.";
+  }
+
+  auto AddrsOrErr = FindAccessedAddrsExegesis(R"asm(
+    vmovdqu %xmm16, (%rax)
+    vmovdqu %ymm16, (%rax)
+  )asm");
+  ASSERT_TRUE(static_cast<bool>(AddrsOrErr));
+  EXPECT_THAT(
+      AddrsOrErr->initial_registers(),
+      IsSupersetOf(
+          {Property("register_name", &RegisterAndValue::register_name, "XMM16"),
+           Property("register_name", &RegisterAndValue::register_name,
+                    "YMM16")}));
+}
+
 }  // namespace
 }  // namespace gematria