[AArch64][SME] Disable (SLP|Loop)Vectorizer when function may be executed in streaming mode.

[sdesmalen]

Summary

When the SME attributes tell that a function is or may be executed in Streaming
SVE mode, we currently need to be conservative and disable _any_ vectorization
(fixed or scalable) because the code-generator does not yet support generating
streaming-compatible code.

Scalable auto-vec will be gradually enabled in the future when we have
confidence that the loop-vectorizer won't use any SVE or NEON instructions
that are illegal in Streaming SVE mode.

[sdesmalen]

Simplified test using 'sed' macro replacement.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64Subtarget.h

124

Likely doesn't matter much but the specification uses StreamingSVEMode.

[dmgreen]

I agree that the idea here sounds sensible. I feel it might be better to have the variable store isStreamingModeEnabled though. That seems like the standard default, and helps avoids a lot of double negatives.

llvm/lib/Target/AArch64/AArch64TargetMachine.cpp

428

This seems to be in the wrong place, between "SVEMax" and MaxSVEVectorSize.

llvm/test/Transforms/LoopVectorize/AArch64/sme-vectorize.ll

84

Are you sure this tests anything in SLP?

[paulwalker-arm]

If you do go for a rename can I recommend InStreamingSVEMode :)

[sdesmalen]

The reason to go for IsStreamingModeDisabled (rather than Enabled) is to also cover the case of a 'streaming compatible' function, where streaming mode may be enabled, but we don't know at compile-time whether it is.

[paulwalker-arm]

Sorry @sdesmalen, I must keep reminding myself that streaming compatible is a subset of streaming SVE mode :) What about representing each separately? For example IsStreamingSVECompatible and InStreamingSVEMode. The latter likely has no current uses? but will allow that difference between "streaming compatible" and "streaming SVE mode" to be queried if/when necessary.

[sdesmalen]

In D135950#3862297, @paulwalker-arm wrote:

Sorry @sdesmalen, I must keep reminding myself that streaming compatible is a subset of streaming SVE mode :) What about representing each separately? For example IsStreamingSVECompatible and InStreamingSVEMode. The latter likely has no current uses? but will allow that difference between "streaming compatible" and "streaming SVE mode" to be queried if/when necessary.

In that case we'd have the possibility of representing an invalid combination (i.e. a function cannot be both StreamingCompatible and StreamingSVEMode), so then an enum would be a better choice. But since there is (currently) no need for this, is IsStreamingSVEDisabled not better at this point? It makes little sense to add something we don't need and can easily extend later.

llvm/test/Transforms/LoopVectorize/AArch64/sme-vectorize.ll

84

I believe so, the hints passed to LoopVectorize sets the vectorize-width to 1 and the interleave-count to 4 using the attributes. We end up with a fixed-width vector code, so this must be due to SLP.

[paulwalker-arm]

In D135950#3865768, @sdesmalen wrote:

In D135950#3862297, @paulwalker-arm wrote:

Sorry @sdesmalen, I must keep reminding myself that streaming compatible is a subset of streaming SVE mode :) What about representing each separately? For example IsStreamingSVECompatible and InStreamingSVEMode. The latter likely has no current uses? but will allow that difference between "streaming compatible" and "streaming SVE mode" to be queried if/when necessary.

In that case we'd have the possibility of representing an invalid combination (i.e. a function cannot be both StreamingCompatible and StreamingSVEMode), so then an enum would be a better choice. But since there is (currently) no need for this, is IsStreamingSVEDisabled not better at this point? It makes little sense to add something we don't need and can easily extend later.

I was considering StreamingCompatible as being a subset of StreamingSVEMode so both could be true. That said, I do agree, the exact mechanics can be changed at a later point once we understand the interactions more. This is why I previously/already accepted the patch because I'm happy with the current intent.

[sdesmalen]

llvm/lib/Target/AArch64/AArch64TargetMachine.cpp

428

Sorry this comment escaped my attention, I've pushed a fix separately in rG36864d47d6b0

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;
   unsigned MinSVEVectorSizeInBits;
   unsigned MaxSVEVectorSizeInBits;
   unsigned VScaleForTuning = 2;
 
   /// TargetTriple - What processor and OS we're targeting.
   Triple TargetTriple;
 
   AArch64FrameLowering FrameLowering;

 
 public:
   /// This constructor initializes the data members to match that
   /// of the specified triple.
   AArch64Subtarget(const Triple &TT, const std::string &CPU,
                    const std::string &TuneCPU, const std::string &FS,
                    const TargetMachine &TM, bool LittleEndian,
                    unsigned MinSVEVectorSizeInBitsOverride = 0,
-                   unsigned MaxSVEVectorSizeInBitsOverride = 0);
+                   unsigned MaxSVEVectorSizeInBitsOverride = 0,
+                   bool StreamingSVEModeDisabled = true);
 
 // 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"
   /// initializeProperties().
   ARMProcFamilyEnum getProcFamily() const {
     return ARMProcFamily;
   }
 
   bool isXRaySupported() const override { return true; }
 
   unsigned getMinVectorRegisterBitWidth() const {
+    // Don't assume any minimum vector size when PSTATE.SM may not be 0.
+    if (!isStreamingSVEModeDisabled())
+      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());
     AllReservedX |= ReserveXRegister;
#include "AArch64GenSubtargetInfo.inc"
   }
 
   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, const std::string &CPU,
                                    const std::string &TuneCPU,
                                    const std::string &FS,
                                    const TargetMachine &TM, bool LittleEndian,
                                    unsigned MinSVEVectorSizeInBitsOverride,
-                                   unsigned MaxSVEVectorSizeInBitsOverride)
+                                   unsigned MaxSVEVectorSizeInBitsOverride,
+                                   bool StreamingSVEModeDisabled)
     : AArch64GenSubtargetInfo(TT, CPU, TuneCPU, FS),
       ReserveXRegister(AArch64::GPR64commonRegClass.getNumRegs()),
       ReserveXRegisterForRA(AArch64::GPR64commonRegClass.getNumRegs()),
       CustomCallSavedXRegs(AArch64::GPR64commonRegClass.getNumRegs()),
       IsLittle(LittleEndian),
+      StreamingSVEModeDisabled(StreamingSVEModeDisabled),
       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()));

llvm/lib/Target/AArch64/AArch64TargetMachine.cpp

std::string CPU =
       CPUAttr.isValid() ? CPUAttr.getValueAsString().str() : TargetCPU;
   std::string TuneCPU =
       TuneAttr.isValid() ? TuneAttr.getValueAsString().str() : CPU;
   std::string FS =
       FSAttr.isValid() ? FSAttr.getValueAsString().str() : TargetFS;
 
   SmallString<512> Key;
 
+  bool StreamingSVEModeDisabled =
+      !F.hasFnAttribute("aarch64_pstate_sm_enabled") &&
+      !F.hasFnAttribute("aarch64_pstate_sm_compatible") &&
+      !F.hasFnAttribute("aarch64_pstate_sm_body");
+
   unsigned MinSVEVectorSize = 0;
   unsigned MaxSVEVectorSize = 0;
   Attribute VScaleRangeAttr = F.getFnAttribute(Attribute::VScaleRange);
   if (VScaleRangeAttr.isValid()) {
     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;
   }
 
   Key += "SVEMin";
   Key += std::to_string(MinSVEVectorSize);
   Key += "SVEMax";
+  Key += "StreamingSVEModeDisabled=" + std::to_string(StreamingSVEModeDisabled);
   Key += std::to_string(MaxSVEVectorSize);
   Key += CPU;
   Key += TuneCPU;
   Key += 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);
+    I = std::make_unique<AArch64Subtarget>(
+        TargetTriple, CPU, TuneCPU, FS, *this, isLittle, MinSVEVectorSize,
+        MaxSVEVectorSize, StreamingSVEModeDisabled);
   }
   return I.get();
 }
 
 void AArch64leTargetMachine::anchor() { }
 
 AArch64leTargetMachine::AArch64leTargetMachine(
     const Target &T, const Triple &TT, StringRef CPU, StringRef FS,

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.h

Optional<Instruction *> instCombineIntrinsic(InstCombiner &IC,
                                                IntrinsicInst &II) const;
 
   Optional<Value *> simplifyDemandedVectorEltsIntrinsic(
       InstCombiner &IC, IntrinsicInst &II, APInt DemandedElts, APInt &UndefElts,
       APInt &UndefElts2, APInt &UndefElts3,
       std::function<void(Instruction *, unsigned, APInt, APInt &)>
           SimplifyAndSetOp) const;
 
-  TypeSize getRegisterBitWidth(TargetTransformInfo::RegisterKind K) const {
-    switch (K) {
-    case TargetTransformInfo::RGK_Scalar:
-      return TypeSize::getFixed(64);
-    case TargetTransformInfo::RGK_FixedWidthVector:
-      if (ST->hasSVE())
-        return TypeSize::getFixed(
-            std::max(ST->getMinSVEVectorSizeInBits(), 128u));
-      return TypeSize::getFixed(ST->hasNEON() ? 128 : 0);
-    case TargetTransformInfo::RGK_ScalableVector:
-      return TypeSize::getScalable(ST->hasSVE() ? 128 : 0);
-    }
-    llvm_unreachable("Unsupported register kind");
-  }
+  TypeSize getRegisterBitWidth(TargetTransformInfo::RegisterKind K) const;
 
   unsigned getMinVectorRegisterBitWidth() const {
     return ST->getMinVectorRegisterBitWidth();
   }
 
   Optional<unsigned> getVScaleForTuning() const {
     return ST->getVScaleForTuning();
   }

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

cl::desc(
         "\nall         All legal loop types will vectorize using tail-folding"
         "\nsimple      Use tail-folding for simple loops (not reductions or "
         "recurrences)"
         "\nreductions  Use tail-folding for loops containing reductions"
         "\nrecurrences Use tail-folding for loops containing fixed order "
         "recurrences"),
     cl::location(TailFoldingKindLoc));
 
+// Experimental option that will only be fully functional when the
+// code-generator is changed to use SVE instead of NEON for all fixed-width
+// operations.
+static cl::opt<bool> EnableFixedwidthAutovecInStreamingMode(
+    "enable-fixedwidth-autovec-in-streaming-mode", cl::init(false), cl::Hidden);
+
+// Experimental option that will only be fully functional when the cost-model
+// and code-generator have been changed to avoid using scalable vector
+// instructions that are not legal in streaming SVE mode.
+static cl::opt<bool> EnableScalableAutovecInStreamingMode(
+    "enable-scalable-autovec-in-streaming-mode", cl::init(false), cl::Hidden);
+
 bool AArch64TTIImpl::areInlineCompatible(const Function *Caller,
                                          const Function *Callee) const {
   SMEAttrs CallerAttrs(*Caller);
   SMEAttrs CalleeAttrs(*Callee);
   if (CallerAttrs.requiresSMChange(CalleeAttrs,
                                    /*BodyOverridesInterface=*/true) ||
       CallerAttrs.requiresLazySave(CalleeAttrs) ||
       CalleeAttrs.hasNewZAInterface())
Optional<Value *> AArch64TTIImpl::simplifyDemandedVectorEltsIntrinsic(
   case Intrinsic::aarch64_neon_uqxtn:
     SimplifyAndSetOp(&II, 0, OrigDemandedElts, UndefElts);
     break;
   }
 
   return None;
 }
 
+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)
+      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)
+      return TypeSize::getScalable(0);
+
+    return TypeSize::getScalable(ST->hasSVE() ? 128 : 0);
+  }
+  llvm_unreachable("Unsupported register kind");
+}
+
 bool AArch64TTIImpl::isWideningInstruction(Type *DstTy, unsigned Opcode,
                                            ArrayRef<const Value *> Args) {
 
   // A helper that returns a vector type from the given type. The number of
   // elements in type Ty determines the vector width.
   auto toVectorTy = [&](Type *ArgTy) {
     return VectorType::get(ArgTy->getScalarType(),
                            cast<VectorType>(DstTy)->getElementCount());

llvm/test/Transforms/LoopVectorize/AArch64/sme-vectorize.ll

+; RUN: sed -e s/REPLACE_PSTATE_MACRO/aarch64_pstate_sm_enabled/ %s | opt -loop-vectorize -slp-vectorizer -S - | FileCheck %s --check-prefix=CHECK
+; RUN: sed -e s/REPLACE_PSTATE_MACRO/aarch64_pstate_sm_enabled/ %s | opt -loop-vectorize -slp-vectorizer -S -enable-scalable-autovec-in-streaming-mode - | FileCheck %s --check-prefix=CHECK-FORCE-SCALABLE
+; RUN: sed -e s/REPLACE_PSTATE_MACRO/aarch64_pstate_sm_enabled/ %s | opt -loop-vectorize -slp-vectorizer -S -enable-fixedwidth-autovec-in-streaming-mode - | FileCheck %s --check-prefix=CHECK-FORCE-FIXEDWIDTH
+
+; RUN: sed -e s/REPLACE_PSTATE_MACRO/aarch64_pstate_sm_compatible/ %s | opt -loop-vectorize -slp-vectorizer -S - | FileCheck %s --check-prefix=CHECK
+; RUN: sed -e s/REPLACE_PSTATE_MACRO/aarch64_pstate_sm_compatible/ %s | opt -loop-vectorize -slp-vectorizer -S -enable-scalable-autovec-in-streaming-mode - | FileCheck %s --check-prefix=CHECK-FORCE-SCALABLE
+; RUN: sed -e s/REPLACE_PSTATE_MACRO/aarch64_pstate_sm_compatible/ %s | opt -loop-vectorize -slp-vectorizer -S -enable-fixedwidth-autovec-in-streaming-mode - | FileCheck %s --check-prefix=CHECK-FORCE-FIXEDWIDTH
+
+; RUN: sed -e s/REPLACE_PSTATE_MACRO/aarch64_pstate_sm_body/ %s | opt -loop-vectorize -slp-vectorizer -S - | FileCheck %s --check-prefix=CHECK
+; RUN: sed -e s/REPLACE_PSTATE_MACRO/aarch64_pstate_sm_body/ %s | opt -loop-vectorize -slp-vectorizer -S -enable-scalable-autovec-in-streaming-mode - | FileCheck %s --check-prefix=CHECK-FORCE-SCALABLE
+; RUN: sed -e s/REPLACE_PSTATE_MACRO/aarch64_pstate_sm_body/ %s | opt -loop-vectorize -slp-vectorizer -S -enable-fixedwidth-autovec-in-streaming-mode - | FileCheck %s --check-prefix=CHECK-FORCE-FIXEDWIDTH
+
+target triple = "aarch64-unknown-linux-gnu"
+
+attributes #0 = { vscale_range(1,16) "target-features"="+neon,+sme,+sve2" "REPLACE_PSTATE_MACRO" }
+
+define void @test_fixedwidth_loopvec(ptr noalias %dst, ptr readonly %src, i32 %N) #0 {
+; CHECK-LABEL: @test_fixedwidth_loopvec
+; CHECK-NOT: <{{[1-9]+}} x i32>
+; CHECK-FORCE-FIXEDWIDTH-LABEL: @test_fixedwidth_loopvec
+; CHECK-FORCE-FIXEDWIDTH: <{{[1-9]+}} x i32>
+entry:
+  %cmp6 = icmp sgt i32 %N, 0
+  br i1 %cmp6, label %for.body.preheader, label %for.cond.cleanup
+
+for.body.preheader:                               ; preds = %entry
+  %wide.trip.count = zext i32 %N to i64
+  br label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.body, %entry
+  ret void
+
+for.body:                                         ; preds = %for.body.preheader, %for.body
+  %indvars.iv = phi i64 [ 0, %for.body.preheader ], [ %indvars.iv.next, %for.body ]
+  %arrayidx = getelementptr inbounds i32, ptr %src, i64 %indvars.iv
+  %0 = load i32, ptr %arrayidx, align 4
+  %add = add nsw i32 %0, 42
+  %arrayidx2 = getelementptr inbounds i32, ptr %dst, i64 %indvars.iv
+  store i32 %add, ptr %arrayidx2, align 4
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond.not = icmp eq i64 %indvars.iv.next, %wide.trip.count
+  br i1 %exitcond.not, label %for.cond.cleanup, label %for.body, !llvm.loop !0
+}
+
+!0 = distinct !{!0, !1, !2, !3}
+!1 = !{!"llvm.loop.mustprogress"}
+!2 = !{!"llvm.loop.interleave.count", i32 1}
+!3 = !{!"llvm.loop.vectorize.scalable.enable", i1 false}
+
+define void @test_scalable_loopvec(ptr noalias %dst, ptr readonly %src, i32 %N) #0 {
+; CHECK-LABEL: @test_scalable_loopvec
+; CHECK-NOT: <vscale x {{[1-9]+}} x i32>
+; CHECK-FORCE-SCALABLE-LABEL: @test_fixedwidth_loopvec
+; CHECK-FORCE-SCALABLE-LABEL: <vscale x {{[1-9]+}} x i32>
+entry:
+  %cmp6 = icmp sgt i32 %N, 0
+  br i1 %cmp6, label %for.body.preheader, label %for.cond.cleanup
+
+for.body.preheader:                               ; preds = %entry
+  %wide.trip.count = zext i32 %N to i64
+  br label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.body, %entry
+  ret void
+
+for.body:                                         ; preds = %for.body.preheader, %for.body
+  %indvars.iv = phi i64 [ 0, %for.body.preheader ], [ %indvars.iv.next, %for.body ]
+  %arrayidx = getelementptr inbounds i32, ptr %src, i64 %indvars.iv
+  %0 = load i32, ptr %arrayidx, align 4
+  %add = add nsw i32 %0, 42
+  %arrayidx2 = getelementptr inbounds i32, ptr %dst, i64 %indvars.iv
+  store i32 %add, ptr %arrayidx2, align 4
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond.not = icmp eq i64 %indvars.iv.next, %wide.trip.count
+  br i1 %exitcond.not, label %for.cond.cleanup, label %for.body, !llvm.loop !4
+}
+
+!4 = distinct !{!4, !5, !6, !7}
+!5 = !{!"llvm.loop.mustprogress"}
+!6 = !{!"llvm.loop.interleave.count", i32 1}
+!7 = !{!"llvm.loop.vectorize.scalable.enable", i1 true}
+
+define void @test_slp(ptr noalias %dst, ptr readonly %src, i32 %N) #0 {
+; CHECK-LABEL: @test_slp
+; CHECK-NOT: <{{[1-9]+}} x i32>
+; CHECK-FORCE-FIXEDWIDTH-LABEL: @test_slp
+; CHECK-FORCE-FIXEDWIDTH: <{{[1-9]+}} x i32>
+entry:
+  %cmp6 = icmp sgt i32 %N, 0
+  br i1 %cmp6, label %for.body.preheader, label %for.cond.cleanup
+
+for.body.preheader:                               ; preds = %entry
+  %wide.trip.count = zext i32 %N to i64
+  br label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.body, %entry
+  ret void
+
+for.body:                                         ; preds = %for.body.preheader, %for.body
+  %indvars.iv = phi i64 [ 0, %for.body.preheader ], [ %indvars.iv.next, %for.body ]
+  %arrayidx = getelementptr inbounds i32, ptr %src, i64 %indvars.iv
+  %0 = load i32, ptr %arrayidx, align 4
+  %add = add nsw i32 %0, 42
+  %arrayidx2 = getelementptr inbounds i32, ptr %dst, i64 %indvars.iv
+  store i32 %add, ptr %arrayidx2, align 4
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond.not = icmp eq i64 %indvars.iv.next, %wide.trip.count
+  br i1 %exitcond.not, label %for.cond.cleanup, label %for.body, !llvm.loop !8
+}
+
+!8 = distinct !{!8, !9, !10, !11}
+!9 = !{!"llvm.loop.mustprogress"}
+!10 = !{!"llvm.loop.interleave.count", i32 4}
+!11 = !{!"llvm.loop.vectorize.width", i32 1}