[AArch64] Force streaming-compatible codegen when attributes are set.

Before this patch, the only way to generate streaming-compatible code
was to use the `-force-streaming-compatible-sve` flag, but the compiler
should also avoid the use of instructions invalid in streaming mode
when a function has the aarch64_pstate_sm_enabled/compatible attribute.

Reviewed By: paulwalker-arm, david-arm

Differential Revision: https://reviews.llvm.org/D155428

4 files changed, 40 insertions, 18 deletions

diff --git a/llvm/lib/Target/AArch64/AArch64Subtarget.cpp b/llvm/lib/Target/AArch64/AArch64Subtarget.cpp
index d27ddde..450e27b 100644
--- a/llvm/lib/Target/AArch64/AArch64Subtarget.cpp
+++ b/llvm/lib/Target/AArch64/AArch64Subtarget.cpp
@@ -292,13 +292,15 @@ AArch64Subtarget::AArch64Subtarget(const Triple &TT, StringRef CPU,
                                    const TargetMachine &TM, bool LittleEndian,
                                    unsigned MinSVEVectorSizeInBitsOverride,
                                    unsigned MaxSVEVectorSizeInBitsOverride,
-                                   bool StreamingSVEModeDisabled)
+                                   bool StreamingSVEMode,
+                                   bool StreamingCompatibleSVEMode)
     : AArch64GenSubtargetInfo(TT, CPU, TuneCPU, FS),
       ReserveXRegister(AArch64::GPR64commonRegClass.getNumRegs()),
       ReserveXRegisterForRA(AArch64::GPR64commonRegClass.getNumRegs()),
       CustomCallSavedXRegs(AArch64::GPR64commonRegClass.getNumRegs()),
       IsLittle(LittleEndian),
-      StreamingSVEModeDisabled(StreamingSVEModeDisabled),
+      StreamingSVEMode(StreamingSVEMode),
+      StreamingCompatibleSVEMode(StreamingCompatibleSVEMode),
       MinSVEVectorSizeInBits(MinSVEVectorSizeInBitsOverride),
       MaxSVEVectorSizeInBits(MaxSVEVectorSizeInBitsOverride), TargetTriple(TT),
       InstrInfo(initializeSubtargetDependencies(FS, CPU, TuneCPU)),
@@ -479,5 +481,10 @@ bool AArch64Subtarget::isNeonAvailable() const {
   if (!hasNEON())
     return false;
 
-  return !ForceStreamingCompatibleSVE;
+  // The 'force-streaming-comaptible-sve' flag overrides the streaming
+  // function attributes.
+  if (ForceStreamingCompatibleSVE.getNumOccurrences() > 0)
+    return !ForceStreamingCompatibleSVE;
+
+  return !isStreaming() && !isStreamingCompatible();
 }
diff --git a/llvm/lib/Target/AArch64/AArch64Subtarget.h b/llvm/lib/Target/AArch64/AArch64Subtarget.h
index d4c136d..9ab8668 100644
--- a/llvm/lib/Target/AArch64/AArch64Subtarget.h
+++ b/llvm/lib/Target/AArch64/AArch64Subtarget.h
@@ -124,7 +124,8 @@ protected:
 
   bool IsLittle;
 
-  bool StreamingSVEModeDisabled;
+  bool StreamingSVEMode;
+  bool StreamingCompatibleSVEMode;
   unsigned MinSVEVectorSizeInBits;
   unsigned MaxSVEVectorSizeInBits;
   unsigned VScaleForTuning = 2;
@@ -163,7 +164,8 @@ public:
                    StringRef FS, const TargetMachine &TM, bool LittleEndian,
                    unsigned MinSVEVectorSizeInBitsOverride = 0,
                    unsigned MaxSVEVectorSizeInBitsOverride = 0,
-                   bool StreamingSVEModeDisabled = true);
+                   bool StreamingSVEMode = false,
+                   bool StreamingCompatibleSVEMode = false);
 
 // Getters for SubtargetFeatures defined in tablegen
 #define GET_SUBTARGETINFO_MACRO(ATTRIBUTE, DEFAULT, GETTER)                    \
@@ -202,6 +204,12 @@ public:
 
   bool isXRaySupported() const override { return true; }
 
+  /// Returns true if the function has the streaming attribute.
+  bool isStreaming() const { return StreamingSVEMode; }
+
+  /// Returns true if the function has the streaming-compatible attribute.
+  bool isStreamingCompatible() const { return StreamingCompatibleSVEMode; }
+
   /// Returns true if the target has NEON and the function at runtime is known
   /// to have NEON enabled (e.g. the function is known not to be in streaming-SVE
   /// mode, which disables NEON instructions).
@@ -209,7 +217,7 @@ public:
 
   unsigned getMinVectorRegisterBitWidth() const {
     // Don't assume any minimum vector size when PSTATE.SM may not be 0.
-    if (!isStreamingSVEModeDisabled())
+    if (StreamingSVEMode || StreamingCompatibleSVEMode)
       return 0;
     return MinVectorRegisterBitWidth;
   }
@@ -416,8 +424,6 @@ public:
       return "__security_check_cookie_arm64ec";
     return "__security_check_cookie";
   }
-
-  bool isStreamingSVEModeDisabled() const { return StreamingSVEModeDisabled; }
 };
 } // End llvm namespace
 
diff --git a/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp b/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp
index c7a6bb5..0c017bc 100644
--- a/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp
+++ b/llvm/lib/Target/AArch64/AArch64TargetMachine.cpp
@@ -391,10 +391,10 @@ AArch64TargetMachine::getSubtargetImpl(const Function &F) const {
   StringRef TuneCPU = TuneAttr.isValid() ? TuneAttr.getValueAsString() : CPU;
   StringRef FS = FSAttr.isValid() ? FSAttr.getValueAsString() : TargetFS;
 
-  bool StreamingSVEModeDisabled =
-      !F.hasFnAttribute("aarch64_pstate_sm_enabled") &&
-      !F.hasFnAttribute("aarch64_pstate_sm_compatible") &&
-      !F.hasFnAttribute("aarch64_pstate_sm_body");
+  bool StreamingSVEMode = F.hasFnAttribute("aarch64_pstate_sm_enabled") ||
+                          F.hasFnAttribute("aarch64_pstate_sm_body");
+  bool StreamingCompatibleSVEMode =
+      F.hasFnAttribute("aarch64_pstate_sm_compatible");
 
   unsigned MinSVEVectorSize = 0;
   unsigned MaxSVEVectorSize = 0;
@@ -427,8 +427,11 @@ AArch64TargetMachine::getSubtargetImpl(const Function &F) const {
 
   SmallString<512> Key;
   raw_svector_ostream(Key) << "SVEMin" << MinSVEVectorSize << "SVEMax"
-                           << MaxSVEVectorSize << "StreamingSVEModeDisabled="
-                           << StreamingSVEModeDisabled << CPU << TuneCPU << FS;
+                           << MaxSVEVectorSize
+                           << "StreamingSVEMode=" << StreamingSVEMode
+                           << "StreamingCompatibleSVEMode="
+                           << StreamingCompatibleSVEMode << CPU << TuneCPU
+                           << FS;
 
   auto &I = SubtargetMap[Key];
   if (!I) {
@@ -438,8 +441,14 @@ AArch64TargetMachine::getSubtargetImpl(const Function &F) const {
     resetTargetOptions(F);
     I = std::make_unique<AArch64Subtarget>(
         TargetTriple, CPU, TuneCPU, FS, *this, isLittle, MinSVEVectorSize,
-        MaxSVEVectorSize, StreamingSVEModeDisabled);
+        MaxSVEVectorSize, StreamingSVEMode, StreamingCompatibleSVEMode);
   }
+
+  assert((!StreamingSVEMode || I->hasSME()) &&
+         "Expected SME to be available");
+  assert((!StreamingCompatibleSVEMode || I->hasSVEorSME()) &&
+         "Expected SVE or SME to be available");
+
   return I.get();
 }
 
diff --git a/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp b/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
index 72ac539..19fa856 100644
--- a/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp
@@ -1925,8 +1925,7 @@ AArch64TTIImpl::getRegisterBitWidth(TargetTransformInfo::RegisterKind K) const {
   case TargetTransformInfo::RGK_Scalar:
     return TypeSize::getFixed(64);
   case TargetTransformInfo::RGK_FixedWidthVector:
-    if (!ST->isStreamingSVEModeDisabled() &&
-        !EnableFixedwidthAutovecInStreamingMode)
+    if (!ST->isNeonAvailable() && !EnableFixedwidthAutovecInStreamingMode)
       return TypeSize::getFixed(0);
 
     if (ST->hasSVE())
@@ -1935,7 +1934,8 @@ AArch64TTIImpl::getRegisterBitWidth(TargetTransformInfo::RegisterKind K) const {
 
     return TypeSize::getFixed(ST->hasNEON() ? 128 : 0);
   case TargetTransformInfo::RGK_ScalableVector:
-    if (!ST->isStreamingSVEModeDisabled() && !EnableScalableAutovecInStreamingMode)
+    if ((ST->isStreaming() || ST->isStreamingCompatible()) &&
+        !EnableScalableAutovecInStreamingMode)
       return TypeSize::getScalable(0);
 
     return TypeSize::getScalable(ST->hasSVE() ? 128 : 0);