From 3fafd3048e28929e6ba1cce62dc0b79280ec47b9 Mon Sep 17 00:00:00 2001
From: Benoit Jacob <jacob.benoit.1@gmail.com>
Date: Tue, 4 Feb 2025 12:25:18 -0500
Subject: [PATCH] Fix mapping of CPU to CPU features on Arm64. (#19900)

The existing mapping only added the features that are default in the ARM
architecture revision. For example, AppleM4 is ArmV8.7-a plus a bunch of
additional features, and we were missing those additional features.

Signed-off-by: Benoit Jacob <jacob.benoit.1@gmail.com>
---
 .../target/LLVMCPU/ResolveCPUAndCPUFeatures.cpp | 17 +++++++++++++----
 .../target/LLVMCPU/ResolveCPUAndCPUFeatures.h   |  1 +
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/compiler/plugins/target/LLVMCPU/ResolveCPUAndCPUFeatures.cpp b/compiler/plugins/target/LLVMCPU/ResolveCPUAndCPUFeatures.cpp
index 628b042b95b1..e6106d57894c 100644
--- a/compiler/plugins/target/LLVMCPU/ResolveCPUAndCPUFeatures.cpp
+++ b/compiler/plugins/target/LLVMCPU/ResolveCPUAndCPUFeatures.cpp
@@ -64,10 +64,16 @@ resolveCPUFeaturesForCPU(const llvm::Triple &triple, std::string &cpu,
     addCpuFeatures(llvm::RISCV::getFeaturesForCPU, cpuFeatureList);
   } else if (triple.isAArch64()) {
     std::vector<llvm::StringRef> cpuFeatureList;
-    const llvm::AArch64::ArchInfo *cpuArch = llvm::AArch64::getArchForCpu(cpu);
-    llvm::AArch64::getExtensionFeatures(cpuArch->DefaultExts, cpuFeatureList);
-    targetCpuFeatures.AddFeature(cpuArch->ArchFeature);
-    for (const auto &feature : cpuFeatureList) {
+    std::optional<llvm::AArch64::CpuInfo> cpuInfo =
+        llvm::AArch64::parseCpu(cpu);
+    if (!cpuInfo) {
+      return ResolveCPUAndCPUFeaturesStatus::UnknownCPU;
+    }
+    llvm::AArch64::ExtensionSet extensions;
+    extensions.addCPUDefaults(*cpuInfo);
+    std::vector<std::string> features;
+    extensions.toLLVMFeatureList(features);
+    for (const auto &feature : features) {
       targetCpuFeatures.AddFeature(feature);
     }
   } else {
@@ -168,6 +174,7 @@ resolveCPUAndCPUFeatures(std::string_view triple_str, std::string &cpu,
                     ResolveCPUAndCPUFeaturesStatus b) {
     return a == ResolveCPUAndCPUFeaturesStatus::OK ? b : a;
   };
+
   return combine(combine(status1, status2), status3);
 }
 
@@ -183,6 +190,8 @@ std::string getMessage(ResolveCPUAndCPUFeaturesStatus status,
     return "Resolution of CPU to CPU-features is not implemented on this "
            "target architecture. Pass explicit "
            "CPU-features, or implement the missing mapping.\n";
+  case ResolveCPUAndCPUFeaturesStatus::UnknownCPU:
+    return "Unknown CPU name for this target architecture.";
   default:
     assert(false);
     return "";
diff --git a/compiler/plugins/target/LLVMCPU/ResolveCPUAndCPUFeatures.h b/compiler/plugins/target/LLVMCPU/ResolveCPUAndCPUFeatures.h
index 51adc2f86a4a..c7f81a893725 100644
--- a/compiler/plugins/target/LLVMCPU/ResolveCPUAndCPUFeatures.h
+++ b/compiler/plugins/target/LLVMCPU/ResolveCPUAndCPUFeatures.h
@@ -14,6 +14,7 @@ namespace mlir::iree_compiler::IREE::HAL {
 
 enum class ResolveCPUAndCPUFeaturesStatus {
   OK,
+  UnknownCPU,
   InconsistentHost,
   UnimplementedMapping,
   ImplicitGenericFallback