[AArch64][NFC] Refactor the tail-folding option

[david-arm]

Summary

This patch does simple refactoring of the tail-folding
option in preparation for enabling tail-folding by
default for neoverse-v1. It adds a default tail-folding
option field to the AArch64Subtarget class that can
be set on a per-CPU.

[sdesmalen]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

43–94

There is probably little value in having the copied std::string here with a SmallVector of StringRefs for the comma-separated options. If you're re-parsing the string each time to query what options are set, then you can just as well split the original StringRef TailFoldingOptionLoc at that point and make SmallVector<StringRef> TailFoldTypes local to the parsing function.

96

I'm still not entirely happy with the concept of having to re-parse the full string every time this information is queried, but if that only happens very irregularly, I guess the overhead is only minimal.

Just as a suggestion, one thing you could do is to parse the string for the reductions/recurrences/reverse early on and store that in a EnableBits + DisableBits mask, and have this separate from parsing of simple/disabled/all/default. Then, at the point of querying the tail folding options, you can do something like:

uint8_t InitialBits = ... // either all zero, all ones, or SimpleBits or DefaultBits (queried from the target)
Bits = (InitialBits | EnableBits) ^ DisableBits;

It does mean that it would not be allowed to have a string such as reductions+default, because default/all/disabled/simple all need to come before the other toggle on/off options. Also, it might make make sense to also make it illegal to write disabled+simple+all because only one of those options should be set.

This way the syntax would become:

(disabled|all|default|simple)[+(reductions|recurrences|reverse|noreductions|norecurrences|noreverse)]*

3561–3562

I find this interface a little confusing and wonder if instead you can make TailFoldingOptionLoc a std::string, and then do:

TailFoldingOption Opt(TailFoldingOptionLoc, ST->getSVETailFoldingDefaultOpts());
return Opt.satisfies(Required);

What do you think?

[david-arm]

• Refactored TailFoldingOption class to parse the string at the time of assignment and keep a record of the InitialBits, EnableBits and DisableBits. This requires changing the semantics of the option to require only one of (disabled|all|simple|default), which must also be the first in the list.

[david-arm]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

96

I've tried to implement your suggestion as best I can!

3561–3562

Now that we no longer parse the string each time, I'm not sure if it still makes sense to refactor this way? I could in theory create a small, new wrapper class with a different name just so I can pass in the option, but wasn't sure if it was worth the extra code?

[sdesmalen]

Thanks for refactoring this @david-arm, this seems like a step in the right direction.

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

42–43

nit: All members of a class are private by default, so you can remove this.

43–94

Can you add a comment explaining why these are required?

46–48

Can you make this part of the constructor? e.g.

TailFoldingOption(bool NeedsDefault = true) { ... }
TailFoldingOption(uint8_t InitBits) { ... }

When you make NeedsDefault = true the default, then you can do away with:

if (Val.empty()) {
  setNeedsDefault();
  return;
}

in operator=

110

Could you split the parsing in two steps, such that you first parse only the "disabled|all|simple|default", and only then start iterating the individual toggles, e.g.

if (TailFoldTypes[0] == "disabled")
  ...
else if (TailFoldTypes[0] == "all")
  ...
else if (TailFoldTypes[0] == "simple")
  ...
else if (TailFoldTypes[0] == "default")
  ...
else
  report_fatal_error(..)

for (auto TailFoldType : ArrayRef<StringRef>(TailFoldTypes).drop_front()) {
  if (TailFoldType == "reductions")
    setEnableBit(TFReductions);
  ...
}

llvm/lib/Target/AArch64/Utils/AArch64BaseInfo.h

542

nit: Can you put this in a namespace, e.g.

namespace TailFoldingOpts {
  enum TailFoldingOptsImpl : uint8_t {
    Disabled = 0x00;
    Simple = 0x01;
    ...
  };
};

That makes it a bit easier to use these, e.g.

TailFoldingOpts::Disabled
TailFoldingOpts::Simple
...

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64Subtarget.h

131

There are several of these uint8_t flying round, which will make it awkward to change later. Perhaps worth using a typedef?

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

96

I think the modern way to do this is inline with the variable definitions, i.e

uint8_t InitialBits = 0;
...

[paulwalker-arm]

llvm/lib/Target/AArch64/Utils/AArch64BaseInfo.h

542

The modern way to do this is via an enum class?, i.e.

enum class TailFoldingOpt : uint8_t {
  Disabled = 0x00;
  Simple = 0x01;
  ...
};

[david-arm]

• Address various review comments involving refactoring code.

[david-arm]

I've tried to address the various comments on the patch and maintain the existing functionality!

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

46–48

I don't think this helps because the constructor never gets invoked in any other way except the default TailFoldingOption(). The only time this ever gets set to something other than default values is when setting the option value, which occurs in operator=.

I've changed this code

if (Val.empty()) {
  setNeedsDefault();
  return;
}

to now emit an error because if the user has supplied the option "-msve-tail-folding=" I think we should treat that as user error.

llvm/lib/Target/AArch64/Utils/AArch64BaseInfo.h

542

This turns out to be a huge pain because it prevents casting between the type and an integer. I tried also using LLVM_DECLARE_ENUM_AS_BITMASK, which didn't stop the compiler complaining. :)

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

119–120

As discussed offline I'd rather keep the existing ability to specify the individual options manually (i.e. -sve-tail-folding=reductions+reverse...), which is more natural than forcing the user to first switch everything off before selectively turning options back on.

llvm/lib/Target/AArch64/Utils/AArch64BaseInfo.h

546

Is Disabled used anywhere? If not then I think it's best removed.

[david-arm]

• Removed a reportError when parsing the first option as we want to permit the user doing stuff like -sve-tail-folding=reductions.
• Changed TailFoldingOpts to be an enum class.

[david-arm]

llvm/lib/Target/AArch64/Utils/AArch64BaseInfo.h

542

Thanks for help with this @paulwalker-arm! I've now changed this to be an enum class.

546

It is required now that TailFoldingOpts is an enum class I think.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

121

Should the else case also do setInitialBits(TailFoldingOpts::Disabled);?

llvm/lib/Target/AArch64/AArch64Subtarget.h

#include "AArch64GenSubtargetInfo.inc"
   BitVector CustomCallSavedXRegs;
 
   bool IsLittle;
 
   bool StreamingSVEModeDisabled;
   unsigned MinSVEVectorSizeInBits;
   unsigned MaxSVEVectorSizeInBits;
   unsigned VScaleForTuning = 2;
+  TailFoldingOpts DefaultSVETFOpts = TailFoldingOpts::Disabled;
 
   /// TargetTriple - What processor and OS we're targeting.
   Triple TargetTriple;
 
   AArch64FrameLowering FrameLowering;
   AArch64InstrInfo InstrInfo;
   AArch64SelectionDAGInfo TSInfo;
   AArch64TargetLowering TLInfo;
bool useSVEForFixedLengthVectors(EVT VT) const {
     return VT.getFixedSizeInBits() > AArch64::SVEBitsPerBlock ||
            forceStreamingCompatibleSVE();
   }
 
   bool forceStreamingCompatibleSVE() const;
 
   unsigned getVScaleForTuning() const { return VScaleForTuning; }
 
+  TailFoldingOpts getSVETailFoldingDefaultOpts() const {
+    return DefaultSVETFOpts;
+  }
+
   const char* getChkStkName() const {
     if (isWindowsArm64EC())
       return "__chkstk_arm64ec";
     return "__chkstk";
   }
 
   const char* getSecurityCheckCookieName() const {
     if (isWindowsArm64EC())

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

static cl::opt<bool> EnableFalkorHWPFUnrollFix("enable-falkor-hwpf-unroll-fix",
                                                cl::init(true), cl::Hidden);
 
 static cl::opt<unsigned> SVEGatherOverhead("sve-gather-overhead", cl::init(10),
                                            cl::Hidden);
 
 static cl::opt<unsigned> SVEScatterOverhead("sve-scatter-overhead",
                                             cl::init(10), cl::Hidden);
 
 namespace {
-class TailFoldingKind {
-private:
-  uint8_t Bits = 0; // Currently defaults to disabled.
+class TailFoldingOption {
+  // These bitfields will only ever be set to something non-zero in operator=,
+  // when setting the -sve-tail-folding option. This option should always be of
+  // the form (default|simple|all|disable)[+(Flag1|Flag2|etc)], where here
+  // InitialBits is one of (disabled|all|simple). EnableBits represents
+  // additional flags we're enabling, and DisableBits for those flags we're
+  // disabling. The default flag is tracked in the variable NeedsDefault, since
+  // at the time of setting the option we may not know what the default value
+  // for the CPU is.
+  TailFoldingOpts InitialBits = TailFoldingOpts::Disabled;
+  TailFoldingOpts EnableBits = TailFoldingOpts::Disabled;
+  TailFoldingOpts DisableBits = TailFoldingOpts::Disabled;
+
+  // This value needs to be initialised to true in case the user does not
+  // explicitly set the -sve-tail-folding option.
+  bool NeedsDefault = true;
+
+  void setInitialBits(TailFoldingOpts Bits) { InitialBits = Bits; }
+
+  void setNeedsDefault(bool V) { NeedsDefault = V; }
+
+  void setEnableBit(TailFoldingOpts Bit) {
+    EnableBits |= Bit;
+    DisableBits &= ~Bit;
+  }
+
+  void setDisableBit(TailFoldingOpts Bit) {
+    EnableBits &= ~Bit;
+    DisableBits |= Bit;
+  }
+
+  TailFoldingOpts getBits(TailFoldingOpts DefaultBits) const {
+    TailFoldingOpts Bits = TailFoldingOpts::Disabled;
+
+    assert((InitialBits == TailFoldingOpts::Disabled || !NeedsDefault) &&
+           "Initial bits should only include one of "
+           "(disabled|all|simple|default)");
+    Bits = NeedsDefault ? DefaultBits : InitialBits;
+    Bits |= EnableBits;
+    Bits &= ~DisableBits;
+
+    return Bits;
+  }
+
+  void reportError(std::string Opt) {
+    errs() << "invalid argument '" << Opt
+           << "' to -sve-tail-folding=; the option should be of the form\n"
+              "  (disabled|all|default|simple)[+(reductions|recurrences"
+              "|reverse|noreductions|norecurrences|noreverse)]\n";
+    report_fatal_error("Unrecognised tail-folding option");
+  }
 
 public:
-  enum TailFoldingOpts {
-    TFDisabled = 0x0,
-    TFReductions = 0x01,
-    TFRecurrences = 0x02,
-    TFReverse = 0x04,
-    TFSimple = 0x80,
-    TFAll = TFReductions | TFRecurrences | TFReverse | TFSimple
-  };
 
   void operator=(const std::string &Val) {
-    if (Val.empty())
+    // If the user explicitly sets -sve-tail-folding= then treat as an error.
+    if (Val.empty()) {
+      reportError("");
       return;
-    SmallVector<StringRef, 6> TailFoldTypes;
+    }
+
+    // Since the user is explicitly setting the option we don't automatically
+    // need the default unless they require it.
+    setNeedsDefault(false);
+
+    SmallVector<StringRef, 4> TailFoldTypes;
     StringRef(Val).split(TailFoldTypes, '+', -1, false);
-    for (auto TailFoldType : TailFoldTypes) {
-      if (TailFoldType == "disabled")
-        Bits = 0;
-      else if (TailFoldType == "all")
-        Bits = TFAll;
-      else if (TailFoldType == "default")
-        Bits = 0; // Currently defaults to never tail-folding.
-      else if (TailFoldType == "simple")
-        add(TFSimple);
-      else if (TailFoldType == "reductions")
-        add(TFReductions);
-      else if (TailFoldType == "recurrences")
-        add(TFRecurrences);
-      else if (TailFoldType == "reverse")
-        add(TFReverse);
-      else if (TailFoldType == "noreductions")
-        remove(TFReductions);
-      else if (TailFoldType == "norecurrences")
-        remove(TFRecurrences);
-      else if (TailFoldType == "noreverse")
-        remove(TFReverse);
-      else {
-        errs()
-            << "invalid argument " << TailFoldType.str()
-            << " to -sve-tail-folding=; each element must be one of: disabled, "
-               "all, default, simple, reductions, noreductions, recurrences, "
-               "norecurrences\n";
-      }
+
+    unsigned StartIdx = 1;
+    if (TailFoldTypes[0] == "disabled")
+      setInitialBits(TailFoldingOpts::Disabled);
+    else if (TailFoldTypes[0] == "all")
+      setInitialBits(TailFoldingOpts::All);
+    else if (TailFoldTypes[0] == "default")
+      setNeedsDefault(true);
+    else if (TailFoldTypes[0] == "simple")
+      setInitialBits(TailFoldingOpts::Simple);
+      StartIdx = 0;
+      setInitialBits(TailFoldingOpts::Disabled);
+
+    for (unsigned I = StartIdx; I < TailFoldTypes.size(); I++) {
+      if (TailFoldTypes[I] == "reductions")
+        setEnableBit(TailFoldingOpts::Reductions);
+      else if (TailFoldTypes[I] == "recurrences")
+        setEnableBit(TailFoldingOpts::Recurrences);
+      else if (TailFoldTypes[I] == "reverse")
+        setEnableBit(TailFoldingOpts::Reverse);
+      else if (TailFoldTypes[I] == "noreductions")
+        setDisableBit(TailFoldingOpts::Reductions);
+      else if (TailFoldTypes[I] == "norecurrences")
+        setDisableBit(TailFoldingOpts::Recurrences);
+      else if (TailFoldTypes[I] == "noreverse")
+        setDisableBit(TailFoldingOpts::Reverse);
+      else
+        reportError(Val);
     }
   }
 
-  operator uint8_t() const { return Bits; }
-
-  void add(uint8_t Flag) { Bits |= Flag; }
-  void remove(uint8_t Flag) { Bits &= ~Flag; }
+  bool satisfies(TailFoldingOpts DefaultBits, TailFoldingOpts Required) const {
+    return (getBits(DefaultBits) & Required) == Required;
+  }
 };
 } // namespace
 
-TailFoldingKind TailFoldingKindLoc;
+TailFoldingOption TailFoldingOptionLoc;
 
-cl::opt<TailFoldingKind, true, cl::parser<std::string>> SVETailFolding(
+cl::opt<TailFoldingOption, true, cl::parser<std::string>> SVETailFolding(
     "sve-tail-folding",
     cl::desc(
-        "Control the use of vectorisation using tail-folding for SVE:"
-        "\ndisabled    No loop types will vectorize using tail-folding"
-        "\ndefault     Uses the default tail-folding settings for the target "
-        "CPU"
-        "\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 "
+        "Control the use of vectorisation using tail-folding for SVE where the"
+        " option is specified in the form (Initial)[+(Flag1|Flag2|...)]:"
+        "\ndisabled      (Initial) No loop types will vectorize using "
+        "tail-folding"
+        "\ndefault       (Initial) Uses the default tail-folding settings for "
+        "the target CPU"
+        "\nall           (Initial) All legal loop types will vectorize using "
+        "tail-folding"
+        "\nsimple        (Initial) Use tail-folding for simple loops (not "
+        "reductions or recurrences)"
+        "\nreductions    Use tail-folding for loops containing reductions"
+        "\nnoreductions  Inverse of above"
+        "\nrecurrences   Use tail-folding for loops containing fixed order "
         "recurrences"
-        "\nreverse     Use tail-folding for loops requiring reversed "
-        "predicates"),
-    cl::location(TailFoldingKindLoc));
+        "\nnorecurrences Inverse of above"
+        "\nreverse       Use tail-folding for loops requiring reversed "
+        "predicates"
+        "\nnoreverse     Inverse of above"),
+    cl::location(TailFoldingOptionLoc));
 
 // 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
for (Instruction &I : *BB) {
           return true;
       }
     }
   }
   return false;
 }
 
 bool AArch64TTIImpl::preferPredicateOverEpilogue(TailFoldingInfo *TFI) {
-  if (!ST->hasSVE() || TailFoldingKindLoc == TailFoldingKind::TFDisabled)
+  if (!ST->hasSVE())
     return false;
 
   // We don't currently support vectorisation with interleaving for SVE - with
   // such loops we're better off not using tail-folding. This gives us a chance
   // to fall back on fixed-width vectorisation using NEON's ld2/st2/etc.
   if (TFI->IAI->hasGroups())
     return false;
 
-  TailFoldingKind Required; // Defaults to 0.
+  TailFoldingOpts Required = TailFoldingOpts::Disabled;
   if (TFI->LVL->getReductionVars().size())
-    Required.add(TailFoldingKind::TFReductions);
+    Required |= TailFoldingOpts::Reductions;
   if (TFI->LVL->getFixedOrderRecurrences().size())
-    Required.add(TailFoldingKind::TFRecurrences);
+    Required |= TailFoldingOpts::Recurrences;
 
   // We call this to discover whether any load/store pointers in the loop have
   // negative strides. This will require extra work to reverse the loop
   // predicate, which may be expensive.
   if (containsDecreasingPointers(TFI->LVL->getLoop(),
                                  TFI->LVL->getPredicatedScalarEvolution()))
-    Required.add(TailFoldingKind::TFReverse);
-  if (!Required)
-    Required.add(TailFoldingKind::TFSimple);
+    Required |= TailFoldingOpts::Reverse;
+  if (Required == TailFoldingOpts::Disabled)
+    Required |= TailFoldingOpts::Simple;
 
-  return (TailFoldingKindLoc & Required) == Required;
+  return TailFoldingOptionLoc.satisfies(ST->getSVETailFoldingDefaultOpts(),
+                                        Required);
 }
 
 InstructionCost
 AArch64TTIImpl::getScalingFactorCost(Type *Ty, GlobalValue *BaseGV,
                                      int64_t BaseOffset, bool HasBaseReg,
                                      int64_t Scale, unsigned AddrSpace) const {
   // Scaling factors are not free at all.
   // Operands                     | Rt Latency

llvm/lib/Target/AArch64/Utils/AArch64BaseInfo.h

 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIB_TARGET_AARCH64_UTILS_AARCH64BASEINFO_H
 #define LLVM_LIB_TARGET_AARCH64_UTILS_AARCH64BASEINFO_H
 
 // FIXME: Is it easiest to fix this layering violation by moving the .inc
 // #includes from AArch64MCTargetDesc.h to here?
 #include "MCTargetDesc/AArch64MCTargetDesc.h" // For AArch64::X0 and friends.
+#include "llvm/ADT/BitmaskEnum.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/MC/SubtargetFeature.h"
 #include "llvm/Support/ErrorHandling.h"
 
 namespace llvm {
 
 inline static unsigned getWRegFromXReg(unsigned Reg) {
case 64:
     return AArch64SVEPredPattern::vl64;
   case 128:
     return AArch64SVEPredPattern::vl128;
   case 256:
     return AArch64SVEPredPattern::vl256;
   }
 }
 
+/// An enum to describe what types of loops we should attempt to tail-fold:
+///   Disabled:    None
+///   Reductions:  Loops containing reductions
+///   Recurrences: Loops with first-order recurrences, i.e. that would
+///                  require a SVE splice instruction
+///   Reverse:     Reverse loops
+///   Simple:      Loops that are not reversed and don't contain reductions
+///                  or first-order recurrences.
+///   All:         All
+enum class TailFoldingOpts : uint8_t {
+  Disabled = 0x00,
+  Simple = 0x01,
+  Reductions = 0x02,
+  Recurrences = 0x04,
+  Reverse = 0x08,
+  All = Reductions | Recurrences | Simple | Reverse
+};
+
+LLVM_DECLARE_ENUM_AS_BITMASK(TailFoldingOpts,
+                             /* LargestValue */ (long)TailFoldingOpts::Reverse);
+
 namespace AArch64ExactFPImm {
   struct ExactFPImm {
     const char *Name;
     int Enum;
     const char *Repr;
   };
 #define GET_EXACTFPIMM_DECL
 #include "AArch64GenSystemOperands.inc"

llvm/test/Transforms/LoopVectorize/AArch64/sve-tail-folding-option.ll

 ; RUN: opt < %s -passes=loop-vectorize -sve-tail-folding=disabled -S | FileCheck %s -check-prefix=CHECK-NOTF
 ; RUN: opt < %s -passes=loop-vectorize -sve-tail-folding=default -S | FileCheck %s -check-prefix=CHECK-NOTF
 ; RUN: opt < %s -passes=loop-vectorize -sve-tail-folding=all -S | FileCheck %s -check-prefix=CHECK-TF
-; RUN: opt < %s -passes=loop-vectorize -sve-tail-folding=disabled+simple+reductions+recurrences+reverse -S | FileCheck %s -check-prefix=CHECK-TF
+; RUN: opt < %s -passes=loop-vectorize -sve-tail-folding=simple+reductions+recurrences+reverse -S | FileCheck %s -check-prefix=CHECK-TF
 ; RUN: opt < %s -passes=loop-vectorize -sve-tail-folding=all+noreductions -S | FileCheck %s -check-prefix=CHECK-TF-NORED
 ; RUN: opt < %s -passes=loop-vectorize -sve-tail-folding=all+norecurrences -S | FileCheck %s -check-prefix=CHECK-TF-NOREC
 ; RUN: opt < %s -passes=loop-vectorize -sve-tail-folding=all+noreverse -S | FileCheck %s -check-prefix=CHECK-TF-NOREV
 ; RUN: opt < %s -passes=loop-vectorize -sve-tail-folding=reductions -S | FileCheck %s -check-prefix=CHECK-TF-ONLYRED
 
 target triple = "aarch64-unknown-linux-gnu"
 
 define void @simple_memset(i32 %val, ptr %ptr, i64 %n) #0 {