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 <[email protected]>
iree-org · Feb 4, 2025 · 3fafd30 · 3fafd30
1 parent 002e637
commit 3fafd30
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 4 deletions.
diff --git 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
@@ -14,6 +14,7 @@ namespace mlir::iree_compiler::IREE::HAL {
 
 enum class ResolveCPUAndCPUFeaturesStatus {
   OK,
+  UnknownCPU,
   InconsistentHost,
   UnimplementedMapping,
   ImplicitGenericFallback