[AArch64][compiler-rt] Add a function returning the current vector length (#92921)

__arm_get_current_vg emits a cntd instruction if in streaming mode or SVE is
available at runtime, otherwise it will return 0.

4 files changed, 135 insertions, 70 deletions

diff --git a/compiler-rt/lib/builtins/CMakeLists.txt b/compiler-rt/lib/builtins/CMakeLists.txt
index c72eb33..0b9e9bd 100644
--- a/compiler-rt/lib/builtins/CMakeLists.txt
+++ b/compiler-rt/lib/builtins/CMakeLists.txt
@@ -562,7 +562,7 @@ set(aarch64_SOURCES
 )
 
 if(COMPILER_RT_HAS_AARCH64_SME AND COMPILER_RT_HAS_FNO_BUILTIN_FLAG AND (COMPILER_RT_HAS_AUXV OR COMPILER_RT_BAREMETAL_BUILD))
-  list(APPEND aarch64_SOURCES aarch64/sme-abi.S aarch64/sme-abi-init.c aarch64/sme-libc-routines.c)
+  list(APPEND aarch64_SOURCES aarch64/sme-abi.S aarch64/sme-abi-init.c aarch64/sme-abi-vg.c aarch64/sme-libc-routines.c)
   message(STATUS "AArch64 SME ABI routines enabled")
   set_source_files_properties(aarch64/sme-libc-routines.c PROPERTIES COMPILE_FLAGS "-fno-builtin")
 else()
diff --git a/compiler-rt/lib/builtins/aarch64/sme-abi-vg.c b/compiler-rt/lib/builtins/aarch64/sme-abi-vg.c
new file mode 100644
index 0000000..e384ab7
--- /dev/null
+++ b/compiler-rt/lib/builtins/aarch64/sme-abi-vg.c
@@ -0,0 +1,45 @@
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "../cpu_model/aarch64.h"
+
+struct FEATURES {
+  long long features;
+};
+
+extern struct FEATURES __aarch64_cpu_features;
+
+struct SME_STATE {
+  long PSTATE;
+  long TPIDR2_EL0;
+};
+
+extern struct SME_STATE __arm_sme_state(void) __arm_streaming_compatible;
+
+extern bool __aarch64_has_sme_and_tpidr2_el0;
+
+#if __GNUC__ >= 9
+#pragma GCC diagnostic ignored "-Wprio-ctor-dtor"
+#endif
+__attribute__((constructor(90))) static void get_aarch64_cpu_features(void) {
+  if (!__aarch64_cpu_features.features)
+    __init_cpu_features();
+}
+
+__attribute__((target("sve"))) long
+__arm_get_current_vg(void) __arm_streaming_compatible {
+  struct SME_STATE State = __arm_sme_state();
+  bool HasSVE = __aarch64_cpu_features.features & (1ULL << FEAT_SVE);
+
+  if (!HasSVE && !__aarch64_has_sme_and_tpidr2_el0)
+    return 0;
+
+  if (HasSVE || (State.PSTATE & 1)) {
+    long vl;
+    __asm__ __volatile__("cntd %0" : "=r"(vl));
+    return vl;
+  }
+
+  return 0;
+}
diff --git a/compiler-rt/lib/builtins/cpu_model/aarch64.c b/compiler-rt/lib/builtins/cpu_model/aarch64.c
index 17bddfc..b868caa 100644
--- a/compiler-rt/lib/builtins/cpu_model/aarch64.c
+++ b/compiler-rt/lib/builtins/cpu_model/aarch64.c
@@ -12,7 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "cpu_model.h"
+#include "aarch64.h"
 
 #if !defined(__aarch64__)
 #error This file is intended only for aarch64-based targets
@@ -53,74 +53,6 @@ _Bool __aarch64_have_lse_atomics
 #endif
 
 #if !defined(DISABLE_AARCH64_FMV)
-// CPUFeatures must correspond to the same AArch64 features in
-// AArch64TargetParser.h
-enum CPUFeatures {
-  FEAT_RNG,
-  FEAT_FLAGM,
-  FEAT_FLAGM2,
-  FEAT_FP16FML,
-  FEAT_DOTPROD,
-  FEAT_SM4,
-  FEAT_RDM,
-  FEAT_LSE,
-  FEAT_FP,
-  FEAT_SIMD,
-  FEAT_CRC,
-  FEAT_SHA1,
-  FEAT_SHA2,
-  FEAT_SHA3,
-  FEAT_AES,
-  FEAT_PMULL,
-  FEAT_FP16,
-  FEAT_DIT,
-  FEAT_DPB,
-  FEAT_DPB2,
-  FEAT_JSCVT,
-  FEAT_FCMA,
-  FEAT_RCPC,
-  FEAT_RCPC2,
-  FEAT_FRINTTS,
-  FEAT_DGH,
-  FEAT_I8MM,
-  FEAT_BF16,
-  FEAT_EBF16,
-  FEAT_RPRES,
-  FEAT_SVE,
-  FEAT_SVE_BF16,
-  FEAT_SVE_EBF16,
-  FEAT_SVE_I8MM,
-  FEAT_SVE_F32MM,
-  FEAT_SVE_F64MM,
-  FEAT_SVE2,
-  FEAT_SVE_AES,
-  FEAT_SVE_PMULL128,
-  FEAT_SVE_BITPERM,
-  FEAT_SVE_SHA3,
-  FEAT_SVE_SM4,
-  FEAT_SME,
-  FEAT_MEMTAG,
-  FEAT_MEMTAG2,
-  FEAT_MEMTAG3,
-  FEAT_SB,
-  FEAT_PREDRES,
-  FEAT_SSBS,
-  FEAT_SSBS2,
-  FEAT_BTI,
-  FEAT_LS64,
-  FEAT_LS64_V,
-  FEAT_LS64_ACCDATA,
-  FEAT_WFXT,
-  FEAT_SME_F64,
-  FEAT_SME_I64,
-  FEAT_SME2,
-  FEAT_RCPC3,
-  FEAT_MOPS,
-  FEAT_MAX,
-  FEAT_EXT = 62, // Reserved to indicate presence of additional features field
-                 // in __aarch64_cpu_features
-  FEAT_INIT      // Used as flag of features initialization completion
-};
 
 // Architecture features used
 // in Function Multi Versioning
diff --git a/compiler-rt/lib/builtins/cpu_model/aarch64.h b/compiler-rt/lib/builtins/cpu_model/aarch64.h
new file mode 100644
index 0000000..15d5300
--- /dev/null
+++ b/compiler-rt/lib/builtins/cpu_model/aarch64.h
@@ -0,0 +1,88 @@
+//===-- cpu_model/aarch64.h --------------------------------------------- -===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "cpu_model.h"
+
+#if !defined(__aarch64__)
+#error This file is intended only for aarch64-based targets
+#endif
+
+#if !defined(DISABLE_AARCH64_FMV)
+
+// CPUFeatures must correspond to the same AArch64 features in
+// AArch64TargetParser.h
+enum CPUFeatures {
+  FEAT_RNG,
+  FEAT_FLAGM,
+  FEAT_FLAGM2,
+  FEAT_FP16FML,
+  FEAT_DOTPROD,
+  FEAT_SM4,
+  FEAT_RDM,
+  FEAT_LSE,
+  FEAT_FP,
+  FEAT_SIMD,
+  FEAT_CRC,
+  FEAT_SHA1,
+  FEAT_SHA2,
+  FEAT_SHA3,
+  FEAT_AES,
+  FEAT_PMULL,
+  FEAT_FP16,
+  FEAT_DIT,
+  FEAT_DPB,
+  FEAT_DPB2,
+  FEAT_JSCVT,
+  FEAT_FCMA,
+  FEAT_RCPC,
+  FEAT_RCPC2,
+  FEAT_FRINTTS,
+  FEAT_DGH,
+  FEAT_I8MM,
+  FEAT_BF16,
+  FEAT_EBF16,
+  FEAT_RPRES,
+  FEAT_SVE,
+  FEAT_SVE_BF16,
+  FEAT_SVE_EBF16,
+  FEAT_SVE_I8MM,
+  FEAT_SVE_F32MM,
+  FEAT_SVE_F64MM,
+  FEAT_SVE2,
+  FEAT_SVE_AES,
+  FEAT_SVE_PMULL128,
+  FEAT_SVE_BITPERM,
+  FEAT_SVE_SHA3,
+  FEAT_SVE_SM4,
+  FEAT_SME,
+  FEAT_MEMTAG,
+  FEAT_MEMTAG2,
+  FEAT_MEMTAG3,
+  FEAT_SB,
+  FEAT_PREDRES,
+  FEAT_SSBS,
+  FEAT_SSBS2,
+  FEAT_BTI,
+  FEAT_LS64,
+  FEAT_LS64_V,
+  FEAT_LS64_ACCDATA,
+  FEAT_WFXT,
+  FEAT_SME_F64,
+  FEAT_SME_I64,
+  FEAT_SME2,
+  FEAT_RCPC3,
+  FEAT_MOPS,
+  FEAT_MAX,
+  FEAT_EXT = 62, // Reserved to indicate presence of additional features field
+                 // in __aarch64_cpu_features
+  FEAT_INIT      // Used as flag of features initialization completion
+};
+
+void __init_cpu_features(void);
+
+#endif // !defined(DISABLE_AARCH64_FMV)