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

[sdesmalen]

Summary

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.

[sdesmalen]

Fixed failing test (streaming-compatible-sve-no-maximize-bandwidth.ll) by forcing fixed-length autovec when -force-streaming-compatible-sve is used.

[david-arm]

LGTM! The change seems sensible to me.

[dmgreen]

llvm/lib/Target/AArch64/AArch64Subtarget.cpp

480

Is there a reason why we don't do this "properly" (with quotes) by having the neon feature disabled when SSVE is being used?

[sdesmalen]

llvm/lib/Target/AArch64/AArch64Subtarget.cpp

480

Yes. We're trying to model two different concepts:
(1) target features regardless of runtime mode. E.g. hasNEON(), which is set by the target attributes such as +simd.
(2) target features for the given runtime mode, which is a subset of (1). These are controlled by the SME attributes such as "aarch64_pstate_sm_compatible".

(1) allows the compiler to parse/emit NEON instructions that are predicated under the HasNEON Predicates in TableGen.

When a function is marked as streaming compatible (i.e. it must be able to run both in streaming mode as well as non-streaming mode) we want to avoid the compiler from using NEON instructions because if the processor is in streaming-SVE mode at runtime, the NEON instructions will be invalid. This is (2), which we represent this by the interface, isNeonAvailable()). We can use this interface to tell SelectionDAG to lower ISD nodes using streaming-compatible SVE instructions instead of using NEON.

The reason we can't use (1), is that a user should still be able to write NEON that compiles, provided that they properly guard the code with if (!__arm_in_streaming_mode()) { /* NEON code here */ }. A user should also be allowed to write inline-asm that disables streaming mode, executes code with NEON instructions, and then re-enables streaming-mode afterwards. If the target being compiled for doesn't allow NEON instructions (i.e. hasNEON() == false), LLVM can't emit them regardless of the runtime mode.

llvm/lib/Target/AArch64/AArch64Subtarget.h

#include "AArch64GenSubtargetInfo.inc"
   // ReserveXRegisterForRA[i] - X#i is not available for register allocator.
   BitVector ReserveXRegisterForRA;
 
   // CustomCallUsedXRegister[i] - X#i call saved.
   BitVector CustomCallSavedXRegs;
 
   bool IsLittle;
 
-  bool StreamingSVEModeDisabled;
+  bool StreamingSVEMode;
+  bool StreamingCompatibleSVEMode;
   unsigned MinSVEVectorSizeInBits;
   unsigned MaxSVEVectorSizeInBits;
   unsigned VScaleForTuning = 2;
   TailFoldingOpts DefaultSVETFOpts = TailFoldingOpts::Disabled;
 
   /// TargetTriple - What processor and OS we're targeting.
   Triple TargetTriple;
 

 
 public:
   /// This constructor initializes the data members to match that
   /// of the specified triple.
   AArch64Subtarget(const Triple &TT, StringRef CPU, StringRef TuneCPU,
                    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)                    \
   bool GETTER() const { return ATTRIBUTE; }
 #include "AArch64GenSubtargetInfo.inc"
 
   const AArch64SelectionDAGInfo *getSelectionDAGInfo() const override {
     return &TSInfo;
#include "AArch64GenSubtargetInfo.inc"
   /// and preferably modeled with SubtargetFeatures or properties in
   /// initializeProperties().
   ARMProcFamilyEnum getProcFamily() const {
     return ARMProcFamily;
   }
 
   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).
   bool isNeonAvailable() const;
 
   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;
   }
 
   bool isXRegisterReserved(size_t i) const { return ReserveXRegister[i]; }
   bool isXRegisterReservedForRA(size_t i) const { return ReserveXRegisterForRA[i]; }
   unsigned getNumXRegisterReserved() const {
     BitVector AllReservedX(AArch64::GPR64commonRegClass.getNumRegs());
const char* getChkStkName() const {
     return "__chkstk";
   }
 
   const char* getSecurityCheckCookieName() const {
     if (isWindowsArm64EC())
       return "__security_check_cookie_arm64ec";
     return "__security_check_cookie";
   }
-
-  bool isStreamingSVEModeDisabled() const { return StreamingSVEModeDisabled; }
 };
 } // End llvm namespace
 
 #endif

llvm/lib/Target/AArch64/AArch64Subtarget.cpp

void AArch64Subtarget::initializeProperties() {
   }
 }
 
 AArch64Subtarget::AArch64Subtarget(const Triple &TT, StringRef CPU,
                                    StringRef TuneCPU, StringRef FS,
                                    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)),
       TLInfo(TM, *this) {
   if (AArch64::isX18ReservedByDefault(TT))
     ReserveXRegister.set(18);
 
   CallLoweringInfo.reset(new AArch64CallLowering(*getTargetLowering()));
void AArch64Subtarget::mirFileLoaded(MachineFunction &MF) const {
   // instructions, specify explicitly if you need it to be correct.
   MachineFrameInfo &MFI = MF.getFrameInfo();
   if (!MFI.isMaxCallFrameSizeComputed())
     MFI.computeMaxCallFrameSize(MF);
 }
 
 bool AArch64Subtarget::useAA() const { return UseAA; }
 
 bool AArch64Subtarget::isNeonAvailable() const {
   if (!hasNEON())
     return false;
 
+  // The 'force-streaming-comaptible-sve' flag overrides the streaming
+  // function attributes.
+  if (ForceStreamingCompatibleSVE.getNumOccurrences() > 0)
     return !ForceStreamingCompatibleSVE;
+
+  return !isStreaming() && !isStreamingCompatible();
 }

llvm/lib/Target/AArch64/AArch64TargetMachine.cpp

AArch64TargetMachine::getSubtargetImpl(const Function &F) const {
   Attribute CPUAttr = F.getFnAttribute("target-cpu");
   Attribute TuneAttr = F.getFnAttribute("tune-cpu");
   Attribute FSAttr = F.getFnAttribute("target-features");
 
   StringRef CPU = CPUAttr.isValid() ? CPUAttr.getValueAsString() : TargetCPU;
   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;
   Attribute VScaleRangeAttr = F.getFnAttribute(Attribute::VScaleRange);
   if (VScaleRangeAttr.isValid()) {
     std::optional<unsigned> VScaleMax = VScaleRangeAttr.getVScaleRangeMax();
     MinSVEVectorSize = VScaleRangeAttr.getVScaleRangeMin() * 128;
     MaxSVEVectorSize = VScaleMax ? *VScaleMax * 128 : 0;
else {
     MinSVEVectorSize =
         (std::min(MinSVEVectorSize, MaxSVEVectorSize) / 128) * 128;
     MaxSVEVectorSize =
         (std::max(MinSVEVectorSize, MaxSVEVectorSize) / 128) * 128;
   }
 
   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) {
     // This needs to be done before we create a new subtarget since any
     // creation will depend on the TM and the code generation flags on the
     // function that reside in TargetOptions.
     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();
 }
 
 void AArch64leTargetMachine::anchor() { }
 
 AArch64leTargetMachine::AArch64leTargetMachine(
     const Target &T, const Triple &TT, StringRef CPU, StringRef FS,
     const TargetOptions &Options, std::optional<Reloc::Model> RM,

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

 }
 
 TypeSize
 AArch64TTIImpl::getRegisterBitWidth(TargetTransformInfo::RegisterKind K) const {
   switch (K) {
   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())
       return TypeSize::getFixed(
           std::max(ST->getMinSVEVectorSizeInBits(), 128u));
 
     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);
   }
   llvm_unreachable("Unsupported register kind");
 }
 
 bool AArch64TTIImpl::isWideningInstruction(Type *DstTy, unsigned Opcode,

llvm/test/CodeGen/AArch64/sme-streaming-compatible-interface.ll

 
 ;
 ; Handle special cases here.
 ;
 
 define <2 x double> @streaming_compatible_with_neon_vectors(<2 x double> %arg) "aarch64_pstate_sm_compatible" nounwind #0 {
 ; CHECK-LABEL: streaming_compatible_with_neon_vectors:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    sub sp, sp, #112
-; CHECK-NEXT:    stp d15, d14, [sp, #32] // 16-byte Folded Spill
-; CHECK-NEXT:    stp d13, d12, [sp, #48] // 16-byte Folded Spill
-; CHECK-NEXT:    stp d11, d10, [sp, #64] // 16-byte Folded Spill
-; CHECK-NEXT:    stp d9, d8, [sp, #80] // 16-byte Folded Spill
-; CHECK-NEXT:    stp x30, x19, [sp, #96] // 16-byte Folded Spill
-; CHECK-NEXT:    str q0, [sp, #16] // 16-byte Folded Spill
+; CHECK-NEXT:    stp d15, d14, [sp, #-96]! // 16-byte Folded Spill
+; CHECK-NEXT:    stp d13, d12, [sp, #16] // 16-byte Folded Spill
+; CHECK-NEXT:    stp d11, d10, [sp, #32] // 16-byte Folded Spill
+; CHECK-NEXT:    stp d9, d8, [sp, #48] // 16-byte Folded Spill
+; CHECK-NEXT:    str x29, [sp, #64] // 8-byte Folded Spill
+; CHECK-NEXT:    stp x30, x19, [sp, #80] // 16-byte Folded Spill
+; CHECK-NEXT:    addvl sp, sp, #-2
+; CHECK-NEXT:    // kill: def $q0 killed $q0 def $z0
+; CHECK-NEXT:    str z0, [sp, #1, mul vl] // 16-byte Folded Spill
 ; CHECK-NEXT:    bl __arm_sme_state
 ; CHECK-NEXT:    and x19, x0, #0x1
 ; CHECK-NEXT:    tbz x19, #0, .LBB4_2
 ; CHECK-NEXT:  // %bb.1:
 ; CHECK-NEXT:    smstop sm
 ; CHECK-NEXT:  .LBB4_2:
-; CHECK-NEXT:    ldr q0, [sp, #16] // 16-byte Folded Reload
+; CHECK-NEXT:    ldr z0, [sp, #1, mul vl] // 16-byte Folded Reload
+; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0
 ; CHECK-NEXT:    bl normal_callee_vec_arg
-; CHECK-NEXT:    str q0, [sp] // 16-byte Folded Spill
+; CHECK-NEXT:    // kill: def $q0 killed $q0 def $z0
+; CHECK-NEXT:    str z0, [sp] // 16-byte Folded Spill
 ; CHECK-NEXT:    tbz x19, #0, .LBB4_4
 ; CHECK-NEXT:  // %bb.3:
 ; CHECK-NEXT:    smstart sm
 ; CHECK-NEXT:  .LBB4_4:
-; CHECK-NEXT:    ldp q1, q0, [sp] // 32-byte Folded Reload
-; CHECK-NEXT:    ldp x30, x19, [sp, #96] // 16-byte Folded Reload
-; CHECK-NEXT:    fadd v0.2d, v1.2d, v0.2d
-; CHECK-NEXT:    ldp d9, d8, [sp, #80] // 16-byte Folded Reload
-; CHECK-NEXT:    ldp d11, d10, [sp, #64] // 16-byte Folded Reload
-; CHECK-NEXT:    ldp d13, d12, [sp, #48] // 16-byte Folded Reload
-; CHECK-NEXT:    ldp d15, d14, [sp, #32] // 16-byte Folded Reload
-; CHECK-NEXT:    add sp, sp, #112
+; CHECK-NEXT:    ldr z0, [sp, #1, mul vl] // 16-byte Folded Reload
+; CHECK-NEXT:    ldr z1, [sp] // 16-byte Folded Reload
+; CHECK-NEXT:    ptrue p0.d, vl2
+; CHECK-NEXT:    fadd z0.d, p0/m, z0.d, z1.d
+; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0
+; CHECK-NEXT:    addvl sp, sp, #2
+; CHECK-NEXT:    ldp x30, x19, [sp, #80] // 16-byte Folded Reload
+; CHECK-NEXT:    ldp d9, d8, [sp, #48] // 16-byte Folded Reload
+; CHECK-NEXT:    ldp d11, d10, [sp, #32] // 16-byte Folded Reload
+; CHECK-NEXT:    ldp d13, d12, [sp, #16] // 16-byte Folded Reload
+; CHECK-NEXT:    ldr x29, [sp, #64] // 8-byte Folded Reload
+; CHECK-NEXT:    ldp d15, d14, [sp], #96 // 16-byte Folded Reload
 ; CHECK-NEXT:    ret
   %res = call <2 x double> @normal_callee_vec_arg(<2 x double> %arg)
   %fadd = fadd <2 x double> %res, %arg
   ret <2 x double> %fadd
 }
 declare <2 x double> @normal_callee_vec_arg(<2 x double>)
 
 define <vscale x 2 x double> @streaming_compatible_with_scalable_vectors(<vscale x 2 x double> %arg) "aarch64_pstate_sm_compatible" nounwind #0 {

llvm/test/Transforms/LoopVectorize/AArch64/streaming-compatible-sve-no-maximize-bandwidth.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt < %s -passes=loop-vectorize -force-streaming-compatible-sve -mattr=+sve -force-target-instruction-cost=1 -scalable-vectorization=off -force-vector-interleave=1 -S 2>&1 | FileCheck %s --check-prefix=SC_SVE
+; RUN: opt < %s -passes=loop-vectorize -force-streaming-compatible-sve -enable-fixedwidth-autovec-in-streaming-mode -mattr=+sve -force-target-instruction-cost=1 -scalable-vectorization=off -force-vector-interleave=1 -S 2>&1 | FileCheck %s --check-prefix=SC_SVE
 ; RUN: opt < %s -passes=loop-vectorize -mattr=+sve -force-target-instruction-cost=1 -scalable-vectorization=off -force-vector-interleave=1 -S 2>&1 | FileCheck %s --check-prefix=NO_SC_SVE
 
 target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
 target triple = "aarch64-unknown-linux-gnu"
 
 @a =  global [32 x i16] zeroinitializer, align 2
 @b =  global [32 x i16] zeroinitializer, align 2
 @c =  global [32 x i16] zeroinitializer, align 2