[Target][ARM] Replace outdated getARMVPTBlockMask function

[Pierre-vh]

Summary

getARMVPTBlockMask was an outdated function that only handled basic block masks: T, TT, TTT and TTTT.
This worked fine before the MVE VPT Block Insertion Pass improvements as it was the only kind of masks that it could generate, but now it can generate more complex masks that uses E predicates, so it's dangerous to use that function to calculate VPT/VPST block masks.

I replaced it with 2 different functions:

• expandPredBlockMask, in ARMBaseInfo. This adds an "E" or "T" at the end of an existing PredBlockMask.
• recomputeVPTBlockMask, in Thumb2InstrInfo. This takes an iterator to a VPT/VPST instruction and recomputes its block mask by looking at the predicated instructions that follows it. This should be used to recompute a block mask after removing/adding a predicated instruction to the block.

The expandPredBlockMask function is pretty much imported from the MVE VPT Blocks pass.

I also had to change the ARMLowOverheadLoops and MVEVPTBlocks passes as well so they could use these new functions.

[samparker]

llvm/lib/Target/ARM/Thumb2InstrInfo.cpp

742

I think it would be cleaner just to pass a MachineInstr, which you can then derive these iterators.

744

VPST is covered by the isVPTOpcode helper.

[Pierre-vh]

• recomputeVPTBlockMask now takes a reference instead of 2 iterators
• Fixed assertion in recomputeVPTBlockMask

[samparker]

Thanks.

llvm/lib/Target/ARM/ARMBaseInstrInfo.h

 static inline
 bool isUncondBranchOpcode(int Opc) {
   return Opc == ARM::B || Opc == ARM::tB || Opc == ARM::t2B;
 }
 
 // This table shows the VPT instruction variants, i.e. the different
 // mask field encodings, see also B5.6. Predication/conditional execution in
 // the ArmARM.
-
-
-inline static ARM::PredBlockMask getARMVPTBlockMask(unsigned NumInsts) {
-  switch (NumInsts) {
-  case 1:
-    return ARM::PredBlockMask::T;
-  case 2:
-    return ARM::PredBlockMask::TT;
-  case 3:
-    return ARM::PredBlockMask::TTT;
-  case 4:
-    return ARM::PredBlockMask::TTTT;
-  default:
-    break;
-  };
-  llvm_unreachable("Unexpected number of instruction in a VPT block");
-}
-
-
 static inline bool isVPTOpcode(int Opc) {
   return Opc == ARM::MVE_VPTv16i8 || Opc == ARM::MVE_VPTv16u8 ||
          Opc == ARM::MVE_VPTv16s8 || Opc == ARM::MVE_VPTv8i16 ||
          Opc == ARM::MVE_VPTv8u16 || Opc == ARM::MVE_VPTv8s16 ||
          Opc == ARM::MVE_VPTv4i32 || Opc == ARM::MVE_VPTv4u32 ||
          Opc == ARM::MVE_VPTv4s32 || Opc == ARM::MVE_VPTv4f32 ||
          Opc == ARM::MVE_VPTv8f16 || Opc == ARM::MVE_VPTv16i8r ||
          Opc == ARM::MVE_VPTv16u8r || Opc == ARM::MVE_VPTv16s8r ||

llvm/lib/Target/ARM/ARMLowOverheadLoops.cpp

struct LowOverheadLoop {
     MachineInstr *Start = nullptr;
     MachineInstr *Dec = nullptr;
     MachineInstr *End = nullptr;
     MachineInstr *VCTP = nullptr;
     VPTBlock *CurrentBlock = nullptr;
     SetVector<MachineInstr*> CurrentPredicate;
     SmallVector<VPTBlock, 4> VPTBlocks;
     SmallPtrSet<MachineInstr*, 4> ToRemove;
+    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) {
       MF = ML.getHeader()->getParent();
     }
void ARMLowOverheadLoops::ConvertVPTBlocks(LowOverheadLoop &LoLoop) {
   // 4) A VPT block which is not predicated upon a vctp, but contains it and
   //    all instructions within the block are predicated upon in.
 
   for (auto &Block : LoLoop.getVPTBlocks()) {
     SmallVectorImpl<PredicatedMI> &Insts = Block.getInsts();
     if (Block.HasNonUniformPredicate()) {
       PredicatedMI *Divergent = Block.getDivergent();
       if (isVCTP(Divergent->MI)) {
-        // The vctp will be removed, so the size of the vpt block needs to be
-        // modified.
-        uint64_t Size = (uint64_t)getARMVPTBlockMask(Block.size() - 1);
-        Block.getVPST()->getOperand(0).setImm(Size);
-        LLVM_DEBUG(dbgs() << "ARM Loops: Modified VPT block mask.\n");
+        // The vctp will be removed, so the block mask of the VPST/VPT will need
+        // to be recomputed.
+        LoLoop.BlockMasksToRecompute.insert(Block.getVPST());
       } else if (Block.IsOnlyPredicatedOn(LoLoop.VCTP)) {
         // The VPT block has a non-uniform predicate but it's entry is guarded
         // only by a vctp, which means we:
         // - Need to remove the original vpst.
         // - Then need to unpredicate any following instructions, until
         //   we come across the divergent vpr def.
         // - Insert a new vpst to predicate the instruction(s) that following
         //   the divergent vpr def.
         // TODO: We could be producing more VPT blocks than necessary and could
         // fold the newly created one into a proceeding one.
         for (auto I = ++MachineBasicBlock::iterator(Block.getVPST()),
              E = ++MachineBasicBlock::iterator(Divergent->MI); I != E; ++I)
           RemovePredicate(&*I);
 
         unsigned Size = 0;
         auto E = MachineBasicBlock::reverse_iterator(Divergent->MI);
         auto I = MachineBasicBlock::reverse_iterator(Insts.back().MI);
         MachineInstr *InsertAt = nullptr;
         while (I != E) {
           InsertAt = &*I;
           ++Size;
           ++I;
         }
+        // Create a VPST with a null mask, we'll recompute it later.
         MachineInstrBuilder MIB = BuildMI(*InsertAt->getParent(), InsertAt,
                                           InsertAt->getDebugLoc(),
                                           TII->get(ARM::MVE_VPST));
-        MIB.addImm((uint64_t)getARMVPTBlockMask(Size));
+        MIB.addImm(0);
         LLVM_DEBUG(dbgs() << "ARM Loops: Removing VPST: " << *Block.getVPST());
         LLVM_DEBUG(dbgs() << "ARM Loops: Created VPST: " << *MIB);
         LoLoop.ToRemove.insert(Block.getVPST());
+        LoLoop.BlockMasksToRecompute.insert(MIB.getInstr());
       }
     } else if (Block.IsOnlyPredicatedOn(LoLoop.VCTP)) {
       // A vpt block which is only predicated upon vctp and has no internal vpr
       // defs:
       // - Remove vpst.
       // - Unpredicate the remaining instructions.
       LLVM_DEBUG(dbgs() << "ARM Loops: Removing VPST: " << *Block.getVPST());
       LoLoop.ToRemove.insert(Block.getVPST());
if (LoLoop.Revert) {
     LoLoop.End = ExpandLoopEnd(LoLoop);
     RemoveDeadBranch(LoLoop.End);
     if (LoLoop.IsTailPredicationLegal())
       ConvertVPTBlocks(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);
+    }
   }
 
   PostOrderLoopTraversal DFS(LoLoop.ML, *MLI);
   DFS.ProcessLoop();
   const SmallVectorImpl<MachineBasicBlock*> &PostOrder = DFS.getOrder();
   for (auto *MBB : PostOrder) {
     recomputeLiveIns(*MBB);
     // FIXME: For some reason, the live-in print order is non-deterministic for

llvm/lib/Target/ARM/MVEVPTBlockPass.cpp

if (registerDefinedBetween(CmpMI->getOperand(1).getReg(), std::next(CmpMI),
                              MI, TRI))
     return nullptr;
   if (registerDefinedBetween(CmpMI->getOperand(2).getReg(), std::next(CmpMI),
                              MI, TRI))
     return nullptr;
   return &*CmpMI;
 }
 
-static ARM::PredBlockMask ExpandBlockMask(ARM::PredBlockMask BlockMask,
-                                          ARMVCC::VPTCodes Kind) {
-  using PredBlockMask = ARM::PredBlockMask;
-  assert(Kind != ARMVCC::None && "Cannot expand mask with 'None'");
-  assert(countTrailingZeros((unsigned)BlockMask) != 0 &&
-         "Mask is already full");
-
-  auto ChooseMask = [&](PredBlockMask AddedThen, PredBlockMask AddedElse) {
-    return (Kind == ARMVCC::Then) ? AddedThen : AddedElse;
-  };
-
-  switch (BlockMask) {
-  case PredBlockMask::T:
-    return ChooseMask(PredBlockMask::TT, PredBlockMask::TE);
-  case PredBlockMask::TT:
-    return ChooseMask(PredBlockMask::TTT, PredBlockMask::TTE);
-  case PredBlockMask::TE:
-    return ChooseMask(PredBlockMask::TET, PredBlockMask::TEE);
-  case PredBlockMask::TTT:
-    return ChooseMask(PredBlockMask::TTTT, PredBlockMask::TTTE);
-  case PredBlockMask::TTE:
-    return ChooseMask(PredBlockMask::TTET, PredBlockMask::TTEE);
-  case PredBlockMask::TET:
-    return ChooseMask(PredBlockMask::TETT, PredBlockMask::TETE);
-  case PredBlockMask::TEE:
-    return ChooseMask(PredBlockMask::TEET, PredBlockMask::TEEE);
-  default:
-    llvm_unreachable("Unknown Mask");
-  }
-}
-
 // Advances Iter past a block of predicated instructions.
 // Returns true if it successfully skipped the whole block of predicated
 // instructions. Returns false when it stopped early (due to MaxSteps), or if
 // Iter didn't point to a predicated instruction.
 static bool StepOverPredicatedInstrs(MachineBasicBlock::instr_iterator &Iter,
                                      MachineBasicBlock::instr_iterator EndIter,
                                      unsigned MaxSteps,
                                      unsigned &NumInstrsSteppedOver) {

 static bool IsVPRDefinedOrKilledByBlock(MachineBasicBlock::iterator Iter,
                                         MachineBasicBlock::iterator End) {
   for (; Iter != End; ++Iter)
     if (Iter->definesRegister(ARM::VPR) || Iter->killsRegister(ARM::VPR))
       return true;
   return false;
 }
 
+// Creates a T, TT, TTT or TTTT BlockMask depending on BlockSize.
+static ARM::PredBlockMask GetInitialBlockMask(unsigned BlockSize) {
+  switch (BlockSize) {
+  case 1:
+    return ARM::PredBlockMask::T;
+  case 2:
+    return ARM::PredBlockMask::TT;
+  case 3:
+    return ARM::PredBlockMask::TTT;
+  case 4:
+    return ARM::PredBlockMask::TTTT;
+  default:
+    llvm_unreachable("Invalid BlockSize!");
+  }
+}
+
 // Given an iterator (Iter) that points at an instruction with a "Then"
 // predicate, tries to create the largest block of continuous predicated
 // instructions possible, and returns the VPT Block Mask of that block.
 //
 // This will try to perform some minor optimization in order to maximize the
 // size of the block.
 static ARM::PredBlockMask
 CreateVPTBlock(MachineBasicBlock::instr_iterator &Iter,
CreateVPTBlock(MachineBasicBlock::instr_iterator &Iter,
   LLVM_DEBUG(for (MachineBasicBlock::instr_iterator AddedInstIter =
                       std::next(BlockBeg);
                   AddedInstIter != Iter; ++AddedInstIter) {
     dbgs() << "  adding: ";
     AddedInstIter->dump();
   });
 
   // Generate the initial BlockMask
-  ARM::PredBlockMask BlockMask = getARMVPTBlockMask(BlockSize);
+  ARM::PredBlockMask BlockMask = GetInitialBlockMask(BlockSize);
 
   // Remove VPNOTs while there's still room in the block, so we can make the
   // largest block possible.
   ARMVCC::VPTCodes CurrentPredicate = ARMVCC::Else;
   while (BlockSize < 4 && Iter != EndIter &&
          Iter->getOpcode() == ARM::MVE_VPNOT) {
 
     // Try to skip all of the predicated instructions after the VPNOT, stopping
while (BlockSize < 4 && Iter != EndIter &&
     // it at the same time.
     for (; Iter != VPNOTBlockEndIter; ++Iter) {
       // Find the register in which the predicate is
       int OpIdx = findFirstVPTPredOperandIdx(*Iter);
       assert(OpIdx != -1);
 
       // Change the predicate and update the mask
       Iter->getOperand(OpIdx).setImm(CurrentPredicate);
-      BlockMask = ExpandBlockMask(BlockMask, CurrentPredicate);
+      BlockMask = expandPredBlockMask(BlockMask, CurrentPredicate);
 
       LLVM_DEBUG(dbgs() << "  adding : "; Iter->dump());
     }
 
     CurrentPredicate =
         (CurrentPredicate == ARMVCC::Then ? ARMVCC::Else : ARMVCC::Then);
   }
   return BlockMask;

llvm/lib/Target/ARM/Thumb2InstrInfo.h

 // corresponding to MachineInstr::findFirstPredOperandIdx.
 int findFirstVPTPredOperandIdx(const MachineInstr &MI);
 ARMVCC::VPTCodes getVPTInstrPredicate(const MachineInstr &MI,
                                       Register &PredReg);
 inline ARMVCC::VPTCodes getVPTInstrPredicate(const MachineInstr &MI) {
   Register PredReg;
   return getVPTInstrPredicate(MI, PredReg);
 }
+
+// Recomputes the Block Mask of Instr, a VPT or VPST instruction.
+// This rebuilds the block mask of the instruction depending on the predicates
+// of the instructions following it. This should only be used after the
+// MVEVPTBlockInsertion pass has run, and should be used whenever a predicated
+// instruction is added to/removed from the block.
+void recomputeVPTBlockMask(MachineInstr &Instr);
 } // namespace llvm
 
 #endif

llvm/lib/Target/ARM/Thumb2InstrInfo.cpp

ARMVCC::VPTCodes llvm::getVPTInstrPredicate(const MachineInstr &MI,
   if (PIdx == -1) {
     PredReg = 0;
     return ARMVCC::None;
   }
 
   PredReg = MI.getOperand(PIdx+1).getReg();
   return (ARMVCC::VPTCodes)MI.getOperand(PIdx).getImm();
 }
+
+void llvm::recomputeVPTBlockMask(MachineInstr &Instr) {
+  assert(isVPTOpcode(Instr.getOpcode()) && "Not a VPST or VPT Instruction!");
+
+  MachineOperand &MaskOp = Instr.getOperand(0);
+  assert(MaskOp.isImm() && "Operand 0 is not the block mask of the VPT/VPST?!");
+
+  MachineBasicBlock::iterator Iter = ++Instr.getIterator(),
+                              End = Instr.getParent()->end();
+
+  // Verify that the instruction after the VPT/VPST is predicated (it should
+  // be), and skip it.
+  ARMVCC::VPTCodes Pred = getVPTInstrPredicate(*Iter);
+  assert(
+      Pred == ARMVCC::Then &&
+      "VPT/VPST should be followed by an instruction with a 'then' predicate!");
+  ++Iter;
+
+  // Iterate over the predicated instructions, updating the BlockMask as we go.
+  ARM::PredBlockMask BlockMask = ARM::PredBlockMask::T;
+  while (Iter != End) {
+    ARMVCC::VPTCodes Pred = getVPTInstrPredicate(*Iter);
+    if (Pred == ARMVCC::None)
+      break;
+    BlockMask = expandPredBlockMask(BlockMask, Pred);
+    ++Iter;
+  }
+
+  // Rewrite the BlockMask.
+  MaskOp.setImm((int64_t)(BlockMask));
+}

llvm/lib/Target/ARM/Utils/ARMBaseInfo.h

enum class PredBlockMask {
     TTET = 0b0101,
     TEEE = 0b1111,
     TEET = 0b1101,
     TETT = 0b1001,
     TETE = 0b1011
   };
 } // namespace ARM
 
+// Expands a PredBlockMask by adding an E or a T at the end, depending on Kind.
+// e.g ExpandPredBlockMask(T, Then) = TT, ExpandPredBlockMask(TT, Else) = TTE,
+// and so on.
+ARM::PredBlockMask expandPredBlockMask(ARM::PredBlockMask BlockMask,
+                                       ARMVCC::VPTCodes Kind);
+
 inline static const char *ARMVPTPredToString(ARMVCC::VPTCodes CC) {
   switch (CC) {
   case ARMVCC::None:  return "none";
   case ARMVCC::Then:  return "t";
   case ARMVCC::Else:  return "e";
   }
   llvm_unreachable("Unknown VPT code");
 }

llvm/lib/Target/ARM/Utils/ARMBaseInfo.cpp

 //
 //===----------------------------------------------------------------------===//
 #include "ARMBaseInfo.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 
 using namespace llvm;
 namespace llvm {
+ARM::PredBlockMask expandPredBlockMask(ARM::PredBlockMask BlockMask,
+                                       ARMVCC::VPTCodes Kind) {
+  using PredBlockMask = ARM::PredBlockMask;
+  assert(Kind != ARMVCC::None && "Cannot expand a mask with None!");
+  assert(countTrailingZeros((unsigned)BlockMask) != 0 &&
+         "Mask is already full");
+
+  auto ChooseMask = [&](PredBlockMask AddedThen, PredBlockMask AddedElse) {
+    return Kind == ARMVCC::Then ? AddedThen : AddedElse;
+  };
+
+  switch (BlockMask) {
+  case PredBlockMask::T:
+    return ChooseMask(PredBlockMask::TT, PredBlockMask::TE);
+  case PredBlockMask::TT:
+    return ChooseMask(PredBlockMask::TTT, PredBlockMask::TTE);
+  case PredBlockMask::TE:
+    return ChooseMask(PredBlockMask::TET, PredBlockMask::TEE);
+  case PredBlockMask::TTT:
+    return ChooseMask(PredBlockMask::TTTT, PredBlockMask::TTTE);
+  case PredBlockMask::TTE:
+    return ChooseMask(PredBlockMask::TTET, PredBlockMask::TTEE);
+  case PredBlockMask::TET:
+    return ChooseMask(PredBlockMask::TETT, PredBlockMask::TETE);
+  case PredBlockMask::TEE:
+    return ChooseMask(PredBlockMask::TEET, PredBlockMask::TEEE);
+  default:
+    llvm_unreachable("Unknown Mask");
+  }
+}
+
 namespace ARMSysReg {
 
 // lookup system register using 12-bit SYSm value.
 // Note: the search is uniqued using M1 mask
 const MClassSysReg *lookupMClassSysRegBy12bitSYSmValue(unsigned SYSm) {
   return lookupMClassSysRegByM1Encoding12(SYSm);
 }