From 2738b30b1b30cbeb3b743cb619ebddd09369205f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mos=C3=A8=20Giordano?= <mose@gnu.org>
Date: Thu, 8 Aug 2024 15:21:02 +0200
Subject: [PATCH] Fix bfloat16 detection on aarch64

---
 src/llvm-demote-float16.cpp | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/llvm-demote-float16.cpp b/src/llvm-demote-float16.cpp
index 5d0d9f5d37c40..6649fd1d1d212 100644
--- a/src/llvm-demote-float16.cpp
+++ b/src/llvm-demote-float16.cpp
@@ -73,12 +73,23 @@ static bool have_fp16(Function &caller, const Triple &TT) {
 }
 
 static bool have_bf16(Function &caller, const Triple &TT) {
-    if (caller.hasFnAttribute("julia.hasbf16")) {
-        return true;
+    Attribute FSAttr = caller.getFnAttribute("target-features");
+    StringRef FS = "";
+    if (FSAttr.isValid()) {
+        FS = FSAttr.getValueAsString();
+    } else if (jl_ExecutionEngine) {
+        FS = jl_ExecutionEngine->getTargetFeatureString();
+    }
+    if (TT.isAArch64()) {
+        if (FS.find("+bf16") != llvm::StringRef::npos) {
+            return true;
+        }
+    } else if (TT.getArch() == Triple::x86_64) {
+        if (caller.hasFnAttribute("julia.hasbf16")) {
+            return true;
+        }
     }
 
-    // there's no targets that fully support bfloat yet;,
-    // AVX512BF16 only provides conversion and dot product instructions.
     return false;
 }