[AArch64][compiler-rt] Add a function returning the current vector length

compiler-rt/lib/builtins/aarch64/sme-abi-init.c

@@ -2,6 +2,8 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+#include "../cpu_model/aarch64.h"
+
 __attribute__((visibility("hidden"), nocommon))
 _Bool __aarch64_has_sme_and_tpidr2_el0;
 
@@ -50,3 +52,35 @@ __attribute__((constructor(90)))
 static void init_aarch64_has_sme(void) {
   __aarch64_has_sme_and_tpidr2_el0 = has_sme();
 }
+
+#if __GNUC__ >= 9
+#pragma GCC diagnostic ignored "-Wprio-ctor-dtor"
+#endif
+__attribute__((constructor(90))) static void get_aarch64_cpu_features(void) {
+  if (!__get_aarch64_features())
+    __init_cpu_features();
+}
+
+struct SME_STATE {
+  long PSTATE;
+  long TPIDR2_EL0;
+};
+
+extern struct SME_STATE __arm_sme_state(void) __arm_streaming_compatible;
+
+__attribute__((target("sve"))) long
+__arm_get_current_vg(void) __arm_streaming_compatible {
+  struct SME_STATE State = __arm_sme_state();
+  bool HasSVE = __get_aarch64_features() & (1ULL << FEAT_SVE);
+
+  if (!HasSVE && !has_sme())
+    return 0;
+
+  if (HasSVE || (State.PSTATE & 1)) {
+    long vl;
+    __asm__ __volatile__("cntd %0" : "=r"(vl));
+    return vl;
+  }
+
+  return 0;
+}

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
@@ -129,6 +61,10 @@ struct {
   // As features grows new fields could be added
 } __aarch64_cpu_features __attribute__((visibility("hidden"), nocommon));
 
+long long __get_aarch64_features(void) {
+  return __aarch64_cpu_features.features;
+}
+
 // The formatter wants to re-order these includes, but doing so is incorrect:
 // clang-format off
 #if defined(__APPLE__)

compiler-rt/lib/builtins/cpu_model/aarch64.h

@@ -0,0 +1,90 @@
+//===-- 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
+};
+
+long long __get_aarch64_features(void);
+
+void CONSTRUCTOR_ATTRIBUTE __init_cpu_features(void);
+
+#endif // !defined(DISABLE_AARCH64_FMV)