[ARM][LowOverheadLoops] Handle reductions

[samparker]

Summary

While validating live-out values, record instructions that look like a reduction. This will comprise of a vector op (for now only vadd), a vorr (vmov) which store the previous value of vadd and then a vpsel in the exit block which is predicated upon a vctp. This vctp will combine the last two iterations using the vmov and vadd into a vector which can then be consumed by a vaddv.
Once we have determined that it's safe to perform tail-predication, we need to change this sequence of instructions so that the predication doesn't produce incorrect code. This involves changing the register allocation of the vadd so it updates itself and the predication on the final iteration will not update the falsely predicated lanes. This mimics what the vmov, vctp and vpsel do and so we then don't need any of those instructions.

[SjoerdMeijer]

Big patch: this is just a first scan of the code, and a first round of nits. Now going to look again, to let things sink in.

llvm/include/llvm/CodeGen/ReachingDefAnalysis.h

191–192

you renamed this Defs to Incoming...

193

So rename this one too?

271

and here too?

llvm/lib/Target/ARM/ARMLowOverheadLoops.cpp

205

Can you comment what these members are? I agree that most are self-explanatory, but I am e.g. interested in Init, and if it e.g. can be null (there is a check in the fixup function), and what the meaning is of that.

574

nit: perhaps an assert that VPSEL is a vpsel would be good.

594

nit, I think nicer to read is:

MO.getReg() == 0

->

!MO.getReg().isValid()

602

Nit: just a bit shorter would be:

for (auto &MO : MI->uses()) {
  if (MO.isImm() && MO.getImm() == Imm)
     return true;
return false;

650

nit: was just curious why we expect the first item in the set to be the vpsel. Can we rely on that with a set?

664

Could this be a good candidate for a helper function in ARMBaseInstrInfo.h?

681

I guess you mean VMOV can be an alias for VORR, which you're checking here?

708

Can or should this not be checked much earlier?

[samparker]

llvm/lib/Target/ARM/ARMLowOverheadLoops.cpp

205

'Init' is the possible instruction that maybe initialising our result register, such as a mov #0, but we won't necessarily have an instruction doing this.

650

Because the set only has one member if we get here.

664

I'm not sure how relevant it is to the rest of the backend, but it would be more readable here as a local helper - especially once we add more supported opcodes.

681

Yes, I don't think we actually have a MVE VMOV instruction.

708

Yes, at some point in an unknown patch. I've moved it into one of the legality helpers.

[samparker]

Rebased, which has made checking the VPSEL predicate a much more simple task.

[SjoerdMeijer]

llvm/lib/Target/ARM/ARMLowOverheadLoops.cpp

684

typo: that

808

Perhaps it's good to add an comment here or in the description of the algorithm on line 733 - 753 that reductions need special treatment as they define values that are not used by predicated instructions inside the loop.

1336

Do we have a test case with more than 1 reduction?

[samparker]

• Rebased after adding tests in reductions.ll.
• Re-added support for VADDi8 and VADDi16.
• Added some extra comments and TODOs.

[SjoerdMeijer]

Thanks, nice optimisation.

llvm/lib/Target/ARM/ARMLowOverheadLoops.cpp

277

is it not const anymore?

693

Do we have a (negative) test case with a float reduction?

737

nit: perhaps move this comment down a bit...

739

...and check this earlier.

[samparker]

llvm/lib/Target/ARM/ARMLowOverheadLoops.cpp

693

The vectorizer doesn't seem to want to produce vector float reductions loops, so I've left them out on testing.

llvm/include/llvm/CodeGen/ReachingDefAnalysis.h

public:
   /// Provides the uses, in the same block as MI, of register that MI defines.
   /// This does not consider live-outs.
   void getReachingLocalUses(MachineInstr *MI, int PhysReg,
                             InstSet &Uses) const;
 
   /// Search MBB for a definition of PhysReg and insert it into Defs. If no
   /// definition is found, recursively search the predecessor blocks for them.
   void getLiveOuts(MachineBasicBlock *MBB, int PhysReg, InstSet &Defs,
                    BlockSet &VisitedBBs) const;
+  void getLiveOuts(MachineBasicBlock *MBB, int PhysReg, InstSet &Defs) const;
 
   /// For the given block, collect the instructions that use the live-in
   /// value of the provided register. Return whether the value is still
   /// live on exit.
   bool getLiveInUses(MachineBasicBlock *MBB, int PhysReg,
                      InstSet &Uses) const;
 
   /// Collect the users of the value stored in PhysReg, which is defined
   /// by MI.
private:
   MachineInstr *getInstFromId(MachineBasicBlock *MBB, int InstId) const;
 
   /// Provides the instruction of the closest reaching def instruction of
   /// PhysReg that reaches MI, relative to the begining of MI's basic block.
   MachineInstr *getReachingLocalMIDef(MachineInstr *MI, int PhysReg) const;
 };
 
 } // namespace llvm
 
 #endif // LLVM_CODEGEN_REACHINGDEFSANALYSIS_H

llvm/lib/CodeGen/ReachingDefAnalysis.cpp

while (!ToVisit.empty()) {
       if (getLiveInUses(MBB, PhysReg, Uses))
         ToVisit.insert(ToVisit.end(), MBB->successors().begin(),
                        MBB->successors().end());
       Visited.insert(MBB);
     }
   }
 }
 
+void ReachingDefAnalysis::getLiveOuts(MachineBasicBlock *MBB, int PhysReg,
+                                      InstSet &Defs) const {
+  SmallPtrSet<MachineBasicBlock*, 2> VisitedBBs;
+  getLiveOuts(MBB, PhysReg, Defs, VisitedBBs);
+}
+
 void
 ReachingDefAnalysis::getLiveOuts(MachineBasicBlock *MBB, int PhysReg,
                                  InstSet &Defs, BlockSet &VisitedBBs) const {
   if (VisitedBBs.count(MBB))
     return;
 
   VisitedBBs.insert(MBB);
   LivePhysRegs LiveRegs(*TRI);

llvm/lib/Target/ARM/ARMBaseInstrInfo.h

          Opc == ARM::tSUBi3 || Opc == ARM::tSUBi8 ||
          Opc == ARM::tSUBSi3 || Opc == ARM::tSUBSi8 ||
          Opc == ARM::t2SUBri || Opc == ARM::t2SUBri12 || Opc == ARM::t2SUBSri;
 }
 
 static inline bool isMovRegOpcode(int Opc) {
   return Opc == ARM::MOVr || Opc == ARM::tMOVr || Opc == ARM::t2MOVr;
 }
-
 /// isValidCoprocessorNumber - decide whether an explicit coprocessor
 /// number is legal in generic instructions like CDP. The answer can
 /// vary with the subtarget.
 static inline bool isValidCoprocessorNumber(unsigned Num,
                                             const FeatureBitset& featureBits) {
   // In Armv7 and Armv8-M CP10 and CP11 clash with VFP/NEON, however, the
   // coprocessor is still valid for CDP/MCR/MRC and friends. Allowing it is
   // useful for code which is shared with older architectures which do not know

llvm/lib/Target/ARM/ARMLowOverheadLoops.cpp

public:
     }
 
     unsigned size() const { return Insts.size(); }
     SmallVectorImpl<PredicatedMI> &getInsts() { return Insts; }
     MachineInstr *getPredicateThen() const { return PredicateThen->MI; }
     PredicatedMI *getDivergent() const { return Divergent; }
   };
 
+  struct Reduction {
+    MachineInstr *Init;
+    MachineInstr &Copy;
+    MachineInstr &Reduce;
+    MachineInstr &VPSEL;
+
+    Reduction(MachineInstr *Init, MachineInstr *Mov, MachineInstr *Add,
+              MachineInstr *Sel)
+      : Init(Init), Copy(*Mov), Reduce(*Add), VPSEL(*Sel) { }
+  };
+
   struct LowOverheadLoop {
 
     MachineLoop &ML;
+    MachineBasicBlock *Preheader = nullptr;
     MachineLoopInfo &MLI;
     ReachingDefAnalysis &RDA;
     const TargetRegisterInfo &TRI;
+    const ARMBaseInstrInfo &TII;
     MachineFunction *MF = nullptr;
     MachineInstr *InsertPt = nullptr;
     MachineInstr *Start = nullptr;
     MachineInstr *Dec = nullptr;
     MachineInstr *End = nullptr;
     MachineInstr *VCTP = nullptr;
     SmallPtrSet<MachineInstr*, 4> SecondaryVCTPs;
     VPTBlock *CurrentBlock = nullptr;
     SetVector<MachineInstr*> CurrentPredicate;
     SmallVector<VPTBlock, 4> VPTBlocks;
     SmallPtrSet<MachineInstr*, 4> ToRemove;
+    SmallVector<std::unique_ptr<Reduction>, 1> Reductions;
     SmallPtrSet<MachineInstr*, 4> BlockMasksToRecompute;
     bool Revert = false;
     bool CannotTailPredicate = false;
 
     LowOverheadLoop(MachineLoop &ML, MachineLoopInfo &MLI,
-                    ReachingDefAnalysis &RDA, const TargetRegisterInfo &TRI)
-      : ML(ML), MLI(MLI), RDA(RDA), TRI(TRI) {
+                    ReachingDefAnalysis &RDA, const TargetRegisterInfo &TRI,
+                    const ARMBaseInstrInfo &TII)
+      : ML(ML), MLI(MLI), RDA(RDA), TRI(TRI), TII(TII) {
       MF = ML.getHeader()->getParent();
+      if (auto *MBB = ML.getLoopPreheader())
+        Preheader = MBB;
+      else if (auto *MBB = MLI.findLoopPreheader(&ML, true))
+        Preheader = MBB;
     }
 
     // If this is an MVE instruction, check that we know how to use tail
     // predication with it. Record VPT blocks and return whether the
     // instruction is valid for tail predication.
     bool ValidateMVEInst(MachineInstr *MI);
 
     void AnalyseMVEInst(MachineInstr *MI) {
       CannotTailPredicate = !ValidateMVEInst(MI);
     }
 
     bool IsTailPredicationLegal() const {
       // For now, let's keep things really simple and only support a single
       // block for tail predication.
       return !Revert && FoundAllComponents() && VCTP &&
              !CannotTailPredicate && ML.getNumBlocks() == 1;
     }
 
     // Check that the predication in the loop will be equivalent once we
     // perform the conversion. Also ensure that we can provide the number
     // of elements to the loop start instruction.
     bool ValidateTailPredicate(MachineInstr *StartInsertPt);
 
+    // See whether the live-out instructions are a reduction that we can fixup
+    // later.
+    bool FindValidReduction(InstSet &LiveMIs, InstSet &LiveOutUsers);
+
     // Check that any values available outside of the loop will be the same
     // after tail predication conversion.
-    bool ValidateLiveOuts() const;
+    bool ValidateLiveOuts();
 
     // Is it safe to define LR with DLS/WLS?
     // LR can be defined if it is the operand to start, because it's the same
     // value, or if it's going to be equivalent to the operand to Start.
     MachineInstr *isSafeToDefineLR();
 
     // Check the branch targets are within range and we satisfy our
     // restrictions.
private:
     void RevertWhile(MachineInstr *MI) const;
 
     bool RevertLoopDec(MachineInstr *MI) const;
 
     void RevertLoopEnd(MachineInstr *MI, bool SkipCmp = false) const;
 
     void ConvertVPTBlocks(LowOverheadLoop &LoLoop);
 
+    void FixupReductions(LowOverheadLoop &LoLoop) const;
+
     MachineInstr *ExpandLoopStart(LowOverheadLoop &LoLoop);
 
     void Expand(LowOverheadLoop &LoLoop);
 
     void IterationCountDCE(LowOverheadLoop &LoLoop);
   };
 }
 
auto CannotProvideElements = [this](MachineBasicBlock *MBB,
     // Don't continue searching up through multiple predecessors.
     if (MBB->pred_size() > 1)
       return true;
 
     return false;
   };
 
   // First, find the block that looks like the preheader.
-  MachineBasicBlock *MBB = MLI.findLoopPreheader(&ML, true);
+  MachineBasicBlock *MBB = Preheader;
   if (!MBB) {
     LLVM_DEBUG(dbgs() << "ARM Loops: Didn't find preheader.\n");
     return false;
   }
 
   // Then search backwards for a def, until we get to InsertBB.
   while (MBB != InsertBB) {
     if (CannotProvideElements(MBB, NumElements)) {
static bool isVectorPredicated(MachineInstr *MI) {
   return PIdx != -1 && MI->getOperand(PIdx + 1).getReg() == ARM::VPR;
 }
 
 static bool isRegInClass(const MachineOperand &MO,
                          const TargetRegisterClass *Class) {
   return MO.isReg() && MO.getReg() && Class->contains(MO.getReg());
 }
 
 // MVE 'narrowing' operate on half a lane, reading from half and writing
 // to half, which are referred to has the top and bottom half. The other
 // half retains its previous value.
 static bool retainsPreviousHalfElement(const MachineInstr &MI) {
   const MCInstrDesc &MCID = MI.getDesc();
   uint64_t Flags = MCID.TSFlags;
   return (Flags & ARMII::RetainsPreviousHalfElement) != 0;
 }
 
 // Some MVE instructions read from the top/bottom halves of their operand(s)
 // and generate a vector result with result elements that are double the
 // width of the input.
 static bool producesDoubleWidthResult(const MachineInstr &MI) {
   const MCInstrDesc &MCID = MI.getDesc();
   uint64_t Flags = MCID.TSFlags;
   return (Flags & ARMII::DoubleWidthResult) != 0;
 }
 
 static bool isHorizontalReduction(const MachineInstr &MI) {
   const MCInstrDesc &MCID = MI.getDesc();
   uint64_t Flags = MCID.TSFlags;
   return (Flags & ARMII::HorizontalReduction) != 0;
 }
 
 // Can this instruction generate a non-zero result when given only zeroed
 // operands? This allows us to know that, given operands with false bytes
 // zeroed by masked loads, that the result will also contain zeros in those
 // bytes.
 static bool canGenerateNonZeros(const MachineInstr &MI) {
 
   // Check for instructions which can write into a larger element size,
   // possibly writing into a previous zero'd lane.
   if (producesDoubleWidthResult(MI))
     return true;
 
   switch (MI.getOpcode()) {
   default:
if (!MO.isReg() || !MO.getReg())
       continue;
     if (!isRegInClass(MO, QPRs) && AllowScalars)
       continue;
     if (auto *OpDef = RDA.getMIOperand(&MI, MO))
       if (FalseLanesZero.count(OpDef))
        continue;
     return false;
   }
   LLVM_DEBUG(dbgs() << "ARM Loops: Always False Zeros: " << MI);
   return true;
 }
 
-bool LowOverheadLoop::ValidateLiveOuts() const {
+bool
+LowOverheadLoop::FindValidReduction(InstSet &LiveMIs, InstSet &LiveOutUsers) {
+  // Also check for reductions where the operation needs to be merging values
+  // from the last and previous loop iterations. This means an instruction
+  // producing a value and a vmov storing the value calculated in the previous
+  // iteration. So we can have two live-out regs, one produced by a vmov and
+  // both being consumed by a vpsel.
+  LLVM_DEBUG(dbgs() << "ARM Loops: Looking for reduction live-outs:\n";
+             for (auto *MI : LiveMIs)
+               dbgs() << " - " << *MI);
+
+  // Expect a vmov, a vadd and a single vpsel user.
+  if (LiveMIs.size() != 2 || LiveOutUsers.size() != 1)
+    return false;
+
+  MachineInstr *VPSEL = *LiveOutUsers.begin();
+  if (VPSEL->getOpcode() != ARM::MVE_VPSEL)
+    return false;
+
+  unsigned VPRIdx = llvm::findFirstVPTPredOperandIdx(*VPSEL) + 1;
+  MachineInstr *Pred = RDA.getMIOperand(VPSEL, VPRIdx);
+  if (!Pred || Pred != VCTP) {
+    LLVM_DEBUG(dbgs() << "ARM Loops: Not using equivalent predicate.\n");
+    return false;
+  }
+
+  MachineInstr *Reduce = RDA.getMIOperand(VPSEL, 1);
+  if (!Reduce)
+    return false;
+
+  // TODO: Support more operations that VADD.
+  if (Reduce->getOpcode() != ARM::MVE_VADDi32)
+    return false;
+
+  // Check that the VORR is actually a VMOV.
+  MachineInstr *Copy = RDA.getMIOperand(VPSEL, 2);
+  if (!Copy || Copy->getOpcode() != ARM::MVE_VORR ||
+      !Copy->getOperand(1).isReg() || !Copy->getOperand(2).isReg() ||
+      Copy->getOperand(1).getReg() != Copy->getOperand(2).getReg())
+    return false;
+
+  assert((LiveMIs.count(Reduce) && LiveMIs.count(Copy)) &&
+         "Expected live outs to be consumed by vpsel");
+
+  assert((Reduce->getOperand(0).getReg() == Reduce->getOperand(1).getReg() ||
+          Reduce->getOperand(0).getReg() == Reduce->getOperand(2).getReg()) &&
+         "Expected VADD to be overwriting one of its operands");
+
+  // Check that the vadd and vmov are only used by each other and the vpsel.
+  SmallPtrSet<MachineInstr*, 2> CopyUsers;
+  RDA.getGlobalUses(Copy, Copy->getOperand(0).getReg(), CopyUsers);
+  if (CopyUsers.size() > 2 || !CopyUsers.count(Reduce))
+    return false;
+
+  SmallPtrSet<MachineInstr*, 2> ReduceUsers;
+  RDA.getGlobalUses(Reduce, Reduce->getOperand(0).getReg(), ReduceUsers);
+  if (ReduceUsers.size() > 2 || !ReduceUsers.count(Copy))
+    return false;
+
+  // Then find whether there's an instruction initialising the register that
+  // is storing the reduction.
+  if (!Preheader)
+    return false;
+
+  SmallPtrSet<MachineInstr*, 2> Incoming;
+  RDA.getLiveOuts(Preheader, Copy->getOperand(1).getReg(), Incoming);
+  if (Incoming.size() > 1)
+    return false;
+
+  MachineInstr *Init = Incoming.empty() ? nullptr : *Incoming.begin();
+  LLVM_DEBUG(dbgs() << "ARM Loops: Found a reduction:\n"
+             << " - " << *Copy
+             << " - " << *Reduce
+             << " - " << *VPSEL);
+  Reductions.push_back(std::make_unique<Reduction>(Init, Copy, Reduce, VPSEL));
+  return true;
+}
+
+bool LowOverheadLoop::ValidateLiveOuts() {
   // We want to find out if the tail-predicated version of this loop will
   // produce the same values as the loop in its original form. For this to
   // be true, the newly inserted implicit predication must not change the
   // the (observable) results.
   // We're doing this because many instructions in the loop will not be
   // predicated and so the conversion from VPT predication to tail-predication
   // can result in different values being produced; due to the tail-predication
   // preventing many instructions from updating their falsely predicated
   // lanes. This analysis assumes that all the instructions perform lane-wise
   // operations and don't perform any exchanges.
   // A masked load, whether through VPT or tail predication, will write zeros
   // to any of the falsely predicated bytes. So, from the loads, we know that
   // the false lanes are zeroed and here we're trying to track that those false
   // lanes remain zero, or where they change, the differences are masked away
   // by their user(s).
   // All MVE loads and stores have to be predicated, so we know that any load
   // operands, or stored results are equivalent already. Other explicitly
   // predicated instructions will perform the same operation in the original
   // loop and the tail-predicated form too. Because of this, we can insert
   // loads, stores and other predicated instructions into our Predicated
   // set and build from there.
   const TargetRegisterClass *QPRs = TRI.getRegClass(ARM::MQPRRegClassID);
   SetVector<MachineInstr *> FalseLanesUnknown;
   SmallPtrSet<MachineInstr *, 4> FalseLanesZero;
   SmallPtrSet<MachineInstr *, 4> Predicated;
-  MachineBasicBlock *MBB = ML.getHeader();
+  MachineBasicBlock *Header = ML.getHeader();
 
-  for (auto &MI : *MBB) {
+  for (auto &MI : *Header) {
     const MCInstrDesc &MCID = MI.getDesc();
     uint64_t Flags = MCID.TSFlags;
     if ((Flags & ARMII::DomainMask) != ARMII::DomainMVE)
       continue;
 
     if (isVCTP(&MI) || isVPTOpcode(MI.getOpcode()))
       continue;
 
auto HasPredicatedUsers = [this](MachineInstr *MI, const MachineOperand &MO,
     }
     return true;
   };
 
   // Visit the unknowns in reverse so that we can start at the values being
   // stored and then we can work towards the leaves, hopefully adding more
   // instructions to Predicated. Successfully terminating the loop means that
   // all the unknown values have to found to be masked by predicated user(s).
+  SmallPtrSet<MachineInstr*, 2> NonPredicated;
   for (auto *MI : reverse(FalseLanesUnknown)) {
     for (auto &MO : MI->operands()) {
       if (!isRegInClass(MO, QPRs) || !MO.isDef())
         continue;
       if (!HasPredicatedUsers(MI, MO, Predicated)) {
         LLVM_DEBUG(dbgs() << "ARM Loops: Found an unknown def of : "
                           << TRI.getRegAsmName(MO.getReg()) << " at " << *MI);
-        return false;
+        NonPredicated.insert(MI);
+        continue;
       }
     }
     // Any unknown false lanes have been masked away by the user(s).
     Predicated.insert(MI);
   }
 
-  // Collect Q-regs that are live in the exit blocks. We don't collect scalars
-  // because they won't be affected by lane predication.
-  SmallSet<Register, 2> LiveOuts;
+  SmallPtrSet<MachineInstr *, 2> LiveOutMIs;
+  SmallPtrSet<MachineInstr*, 2> LiveOutUsers;
   SmallVector<MachineBasicBlock *, 2> ExitBlocks;
   ML.getExitBlocks(ExitBlocks);
-  for (auto *MBB : ExitBlocks)
-    for (const MachineBasicBlock::RegisterMaskPair &RegMask : MBB->liveins())
-      if (QPRs->contains(RegMask.PhysReg))
-        LiveOuts.insert(RegMask.PhysReg);
-
-  // Collect the instructions in the loop body that define the live-out values.
-  SmallPtrSet<MachineInstr *, 2> LiveMIs;
   assert(ML.getNumBlocks() == 1 && "Expected single block loop!");
-  for (auto Reg : LiveOuts)
-    if (auto *MI = RDA.getLocalLiveOutMIDef(MBB, Reg))
-      LiveMIs.insert(MI);
+  assert(ExitBlocks.size() == 1 && "Expected a single exit block");
+  MachineBasicBlock *ExitBB = ExitBlocks.front();
+  for (const MachineBasicBlock::RegisterMaskPair &RegMask : ExitBB->liveins()) {
+    // Check Q-regs that are live in the exit blocks. We don't collect scalars
+    // because they won't be affected by lane predication.
+    if (QPRs->contains(RegMask.PhysReg)) {
+      if (auto *MI = RDA.getLocalLiveOutMIDef(Header, RegMask.PhysReg))
+        LiveOutMIs.insert(MI);
+      RDA.getLiveInUses(ExitBB, RegMask.PhysReg, LiveOutUsers);
+    }
+  }
+
+  // If we have any non-predicated live-outs, they need to be part of a
+  // reduction that we can fixup later.
+  if (!NonPredicated.empty() &&
+      !FindValidReduction(NonPredicated, LiveOutUsers))
+    return false;
 
-  LLVM_DEBUG(dbgs() << "ARM Loops: Found loop live-outs:\n";
-             for (auto *MI : LiveMIs)
-               dbgs() << " - " << *MI);
   // We've already validated that any VPT predication within the loop will be
   // equivalent when we perform the predication transformation; so we know that
   // any VPT predicated instruction is predicated upon VCTP. Any live-out
-  // instruction needs to be predicated, so check this here.
-  for (auto *MI : LiveMIs)
-    if (!isVectorPredicated(MI))
+  // instruction needs to be predicated, so check this here. The instructions
+  // in NonPredicated have been found to be a reduction that we can ensure its
+  // legality.
+  for (auto *MI : LiveOutMIs)
+    if (!isVectorPredicated(MI) && !NonPredicated.count(MI))
       return false;
 
   return true;
 }
 
 void LowOverheadLoop::CheckLegality(ARMBasicBlockUtils *BBUtils) {
   if (Revert)
     return;
for (auto &MI : *MBB) {
       if (isLoopStart(MI))
         return &MI;
     }
     if (MBB->pred_size() == 1)
       return SearchForStart(*MBB->pred_begin());
     return nullptr;
   };
 
-  LowOverheadLoop LoLoop(*ML, *MLI, *RDA, *TRI);
+  LowOverheadLoop LoLoop(*ML, *MLI, *RDA, *TRI, *TII);
   // Search the preheader for the start intrinsic.
   // FIXME: I don't see why we shouldn't be supporting multiple predecessors
   // with potentially multiple set.loop.iterations, so we need to enable this.
-  if (auto *Preheader = ML->getLoopPreheader())
-    LoLoop.Start = SearchForStart(Preheader);
-  else if (auto *Preheader = MLI->findLoopPreheader(ML, true))
-    LoLoop.Start = SearchForStart(Preheader);
+  if (LoLoop.Preheader)
+    LoLoop.Start = SearchForStart(LoLoop.Preheader);
   else
     return false;
 
   // Find the low-overhead loop components and decide whether or not to fall
   // back to a normal loop. Also look for a vctp instructions and decide
   // whether we can convert that predicate using tail predication.
   for (auto *MBB : reverse(ML->getBlocks())) {
     for (auto &MI : *MBB) {
MachineInstr* ARMLowOverheadLoops::ExpandLoopStart(LowOverheadLoop &LoLoop) {
   // If we're inserting at a mov lr, then remove it as it's redundant.
   if (InsertPt != Start)
     LoLoop.ToRemove.insert(InsertPt);
   LoLoop.ToRemove.insert(Start);
   LLVM_DEBUG(dbgs() << "ARM Loops: Inserted start: " << *MIB);
   return &*MIB;
 }
 
+void ARMLowOverheadLoops::FixupReductions(LowOverheadLoop &LoLoop) const {
+  LLVM_DEBUG(dbgs() << "ARM Loops: Fixing up reduction(s).\n");
+  auto BuildMov = [this](MachineInstr &InsertPt, Register To, Register From) {
+    MachineBasicBlock *MBB = InsertPt.getParent();
+    MachineInstrBuilder MIB =
+      BuildMI(*MBB, &InsertPt, InsertPt.getDebugLoc(), TII->get(ARM::MVE_VORR));
+    MIB.addDef(To);
+    MIB.addReg(From);
+    MIB.addReg(From);
+    MIB.addImm(0);
+    MIB.addReg(0);
+    MIB.addReg(To);
+    LLVM_DEBUG(dbgs() << "ARM Loops: Inserted VMOV: " << *MIB);
+  };
+
+  for (auto &Reduction : LoLoop.Reductions) {
+    MachineInstr &Copy = Reduction->Copy;
+    MachineInstr &Reduce = Reduction->Reduce;
+    Register DestReg = Copy.getOperand(0).getReg();
+
+    // Change the initialiser if present
+    if (Reduction->Init) {
+      MachineInstr *Init = Reduction->Init;
+
+      for (unsigned i = 0; i < Init->getNumOperands(); ++i) {
+        MachineOperand &MO = Init->getOperand(i);
+        if (MO.isReg() && MO.isUse() && MO.isTied() &&
+            Init->findTiedOperandIdx(i) == 0)
+          Init->getOperand(i).setReg(DestReg);
+      }
+      Init->getOperand(0).setReg(DestReg);
+      LLVM_DEBUG(dbgs() << "ARM Loops: Changed init regs: " << *Init);
+    } else
+      BuildMov(LoLoop.Preheader->instr_back(), DestReg, Copy.getOperand(1).getReg());
+
+    // Change the reducing op to write to the register that is used to copy
+    // its value on the next iteration. Also update the tied-def operand.
+    Reduce.getOperand(0).setReg(DestReg);
+    Reduce.getOperand(5).setReg(DestReg);
+    LLVM_DEBUG(dbgs() << "ARM Loops: Changed reduction regs: " << Reduce);
+
+    // Instead of a vpsel, just copy the register into the necessary one.
+    MachineInstr &VPSEL = Reduction->VPSEL;
+    if (VPSEL.getOperand(0).getReg() != DestReg)
+      BuildMov(VPSEL, VPSEL.getOperand(0).getReg(), DestReg);
+
+    // Remove the unnecessary instructions.
+    LLVM_DEBUG(dbgs() << "ARM Loops: Removing:\n"
+               << " - " << Copy
+               << " - " << VPSEL << "\n");
+    Copy.eraseFromParent();
+    VPSEL.eraseFromParent();
+  }
+}
+
 void ARMLowOverheadLoops::ConvertVPTBlocks(LowOverheadLoop &LoLoop) {
   auto RemovePredicate = [](MachineInstr *MI) {
     LLVM_DEBUG(dbgs() << "ARM Loops: Removing predicate from: " << *MI);
     if (int PIdx = llvm::findFirstVPTPredOperandIdx(*MI)) {
       assert(MI->getOperand(PIdx).getImm() == ARMVCC::Then &&
              "Expected Then predicate!");
       MI->getOperand(PIdx).setImm(ARMVCC::None);
       MI->getOperand(PIdx+1).setReg(0);
else
       LoLoop.Start->eraseFromParent();
     bool FlagsAlreadySet = RevertLoopDec(LoLoop.Dec);
     RevertLoopEnd(LoLoop.End, FlagsAlreadySet);
   } else {
     LoLoop.Start = ExpandLoopStart(LoLoop);
     RemoveDeadBranch(LoLoop.Start);
     LoLoop.End = ExpandLoopEnd(LoLoop);
     RemoveDeadBranch(LoLoop.End);
-    if (LoLoop.IsTailPredicationLegal())
+    if (LoLoop.IsTailPredicationLegal()) {
       ConvertVPTBlocks(LoLoop);
+      FixupReductions(LoLoop);
+    }
     for (auto *I : LoLoop.ToRemove) {
       LLVM_DEBUG(dbgs() << "ARM Loops: Erasing " << *I);
       I->eraseFromParent();
     }
     for (auto *I : LoLoop.BlockMasksToRecompute) {
       LLVM_DEBUG(dbgs() << "ARM Loops: Recomputing VPT/VPST Block Mask: " << *I);
       recomputeVPTBlockMask(*I);
       LLVM_DEBUG(dbgs() << "           ... done: " << *I);

llvm/test/CodeGen/Thumb2/LowOverheadLoops/vector-arith-codegen.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=armv8.1m.main -mattr=+mve -disable-mve-tail-predication=false --verify-machineinstrs %s -o - | FileCheck %s
 
 define dso_local i32 @mul_reduce_add(i32* noalias nocapture readonly %a, i32* noalias nocapture readonly %b, i32 %N) {
 ; CHECK-LABEL: mul_reduce_add:
 ; CHECK:       @ %bb.0: @ %entry
 ; CHECK-NEXT:    cmp r2, #0
 ; CHECK-NEXT:    itt eq
 ; CHECK-NEXT:    moveq r0, #0
 ; CHECK-NEXT:    bxeq lr
 ; CHECK-NEXT:    push {r7, lr}
-; CHECK-NEXT:    adds r3, r2, #3
-; CHECK-NEXT:    vmov.i32 q0, #0x0
-; CHECK-NEXT:    bic r3, r3, #3
-; CHECK-NEXT:    sub.w r12, r3, #4
-; CHECK-NEXT:    movs r3, #1
-; CHECK-NEXT:    add.w lr, r3, r12, lsr #2
+; CHECK-NEXT:    vmov.i32 q1, #0x0
 ; CHECK-NEXT:    movs r3, #0
-; CHECK-NEXT:    dls lr, lr
+; CHECK-NEXT:    dlstp.32 lr, r2
 ; CHECK-NEXT:  .LBB0_1: @ %vector.body
 ; CHECK-NEXT:    @ =>This Inner Loop Header: Depth=1
-; CHECK-NEXT:    vctp.32 r2
-; CHECK-NEXT:    vmov q1, q0
-; CHECK-NEXT:    vpstt
-; CHECK-NEXT:    vldrwt.u32 q0, [r0], #16
-; CHECK-NEXT:    vldrwt.u32 q2, [r1], #16
+; CHECK-NEXT:    vldrw.u32 q0, [r0], #16
+; CHECK-NEXT:    vldrw.u32 q2, [r1], #16
 ; CHECK-NEXT:    adds r3, #4
 ; CHECK-NEXT:    vmul.i32 q0, q2, q0
-; CHECK-NEXT:    subs r2, #4
-; CHECK-NEXT:    vadd.i32 q0, q0, q1
-; CHECK-NEXT:    le lr, .LBB0_1
+; CHECK-NEXT:    vadd.i32 q1, q0, q1
+; CHECK-NEXT:    letp lr, .LBB0_1
 ; CHECK-NEXT:  @ %bb.2: @ %middle.block
-; CHECK-NEXT:    vpsel q0, q0, q1
+; CHECK-NEXT:    vmov q0, q1
 ; CHECK-NEXT:    vaddv.u32 r0, q0
 ; CHECK-NEXT:    pop {r7, pc}
 entry:
   %cmp8 = icmp eq i32 %N, 0
   br i1 %cmp8, label %for.cond.cleanup, label %vector.ph
 
 vector.ph:                                        ; preds = %entry
   %n.rnd.up = add i32 %N, 3

 define dso_local i32 @mul_reduce_add_const(i32* noalias nocapture readonly %a, i32 %b, i32 %N) {
 ; CHECK-LABEL: mul_reduce_add_const:
 ; CHECK:       @ %bb.0: @ %entry
 ; CHECK-NEXT:    cmp r2, #0
 ; CHECK-NEXT:    itt eq
 ; CHECK-NEXT:    moveq r0, #0
 ; CHECK-NEXT:    bxeq lr
 ; CHECK-NEXT:    push {r7, lr}
-; CHECK-NEXT:    adds r1, r2, #3
-; CHECK-NEXT:    movs r3, #1
-; CHECK-NEXT:    bic r1, r1, #3
-; CHECK-NEXT:    vmov.i32 q0, #0x0
-; CHECK-NEXT:    subs r1, #4
-; CHECK-NEXT:    add.w lr, r3, r1, lsr #2
+; CHECK-NEXT:    vmov.i32 q1, #0x0
 ; CHECK-NEXT:    movs r1, #0
-; CHECK-NEXT:    dls lr, lr
+; CHECK-NEXT:    dlstp.32 lr, r2
 ; CHECK-NEXT:  .LBB1_1: @ %vector.body
 ; CHECK-NEXT:    @ =>This Inner Loop Header: Depth=1
-; CHECK-NEXT:    vctp.32 r2
-; CHECK-NEXT:    vmov q1, q0
-; CHECK-NEXT:    vpst
-; CHECK-NEXT:    vldrwt.u32 q0, [r0], #16
+; CHECK-NEXT:    vldrw.u32 q0, [r0], #16
 ; CHECK-NEXT:    adds r1, #4
-; CHECK-NEXT:    subs r2, #4
-; CHECK-NEXT:    vadd.i32 q0, q0, q1
-; CHECK-NEXT:    le lr, .LBB1_1
+; CHECK-NEXT:    vadd.i32 q1, q0, q1
+; CHECK-NEXT:    letp lr, .LBB1_1
 ; CHECK-NEXT:  @ %bb.2: @ %middle.block
-; CHECK-NEXT:    vpsel q0, q0, q1
+; CHECK-NEXT:    vmov q0, q1
 ; CHECK-NEXT:    vaddv.u32 r0, q0
 ; CHECK-NEXT:    pop {r7, pc}
 entry:
   %cmp6 = icmp eq i32 %N, 0
   br i1 %cmp6, label %for.cond.cleanup, label %vector.ph
 
 vector.ph:                                        ; preds = %entry
   %n.rnd.up = add i32 %N, 3

 define dso_local i32 @add_reduce_add_const(i32* noalias nocapture readonly %a, i32 %b, i32 %N) {
 ; CHECK-LABEL: add_reduce_add_const:
 ; CHECK:       @ %bb.0: @ %entry
 ; CHECK-NEXT:    cmp r2, #0
 ; CHECK-NEXT:    itt eq
 ; CHECK-NEXT:    moveq r0, #0
 ; CHECK-NEXT:    bxeq lr
 ; CHECK-NEXT:    push {r7, lr}
-; CHECK-NEXT:    adds r1, r2, #3
-; CHECK-NEXT:    movs r3, #1
-; CHECK-NEXT:    bic r1, r1, #3
-; CHECK-NEXT:    vmov.i32 q0, #0x0
-; CHECK-NEXT:    subs r1, #4
-; CHECK-NEXT:    add.w lr, r3, r1, lsr #2
+; CHECK-NEXT:    vmov.i32 q1, #0x0
 ; CHECK-NEXT:    movs r1, #0
-; CHECK-NEXT:    dls lr, lr
+; CHECK-NEXT:    dlstp.32 lr, r2
 ; CHECK-NEXT:  .LBB2_1: @ %vector.body
 ; CHECK-NEXT:    @ =>This Inner Loop Header: Depth=1
-; CHECK-NEXT:    vctp.32 r2
-; CHECK-NEXT:    vmov q1, q0
-; CHECK-NEXT:    vpst
-; CHECK-NEXT:    vldrwt.u32 q0, [r0], #16
+; CHECK-NEXT:    vldrw.u32 q0, [r0], #16
 ; CHECK-NEXT:    adds r1, #4
-; CHECK-NEXT:    subs r2, #4
-; CHECK-NEXT:    vadd.i32 q0, q0, q1
-; CHECK-NEXT:    le lr, .LBB2_1
+; CHECK-NEXT:    vadd.i32 q1, q0, q1
+; CHECK-NEXT:    letp lr, .LBB2_1
 ; CHECK-NEXT:  @ %bb.2: @ %middle.block
-; CHECK-NEXT:    vpsel q0, q0, q1
+; CHECK-NEXT:    vmov q0, q1
 ; CHECK-NEXT:    vaddv.u32 r0, q0
 ; CHECK-NEXT:    pop {r7, pc}
 entry:
   %cmp6 = icmp eq i32 %N, 0
   br i1 %cmp6, label %for.cond.cleanup, label %vector.ph
 
 vector.ph:                                        ; preds = %entry
   %n.rnd.up = add i32 %N, 3