[ARM] Cleanup for the MVETailPrediction pass

This strips out a lot of the code that should no longer be needed from
the MVETailPredictionPass, leaving the important part - find active lane
mask instructions and convert them to VCTP operations.

Differential Revision: https://reviews.llvm.org/D91866

Author: davemgreen

llvm/lib/Target/ARM/MVETailPredication.cpp

@@ -22,23 +22,13 @@
 /// The HardwareLoops pass inserts intrinsics identifying loops that the
 /// backend will attempt to convert into a low-overhead loop. The vectorizer is
 /// responsible for generating a vectorized loop in which the lanes are
-/// predicated upon the iteration counter. This pass looks at these predicated
-/// vector loops, that are targets for low-overhead loops, and prepares it for
-/// code generation. Once the vectorizer has produced a masked loop, there's a
-/// couple of final forms:
-/// - A tail-predicated loop, with implicit predication.
-/// - A loop containing multiple VCPT instructions, predicating multiple VPT
-///   blocks of instructions operating on different vector types.
-///
-/// This pass:
-/// 1) Checks if the predicates of the masked load/store instructions are
-///    generated by intrinsic @llvm.get.active.lanes(). This intrinsic consumes
-///    the the scalar loop tripcount as its second argument, which we extract
-///    to set up the number of elements processed by the loop.
-/// 2) Intrinsic @llvm.get.active.lanes() is then replaced by the MVE target
-///    specific VCTP intrinsic to represent the effect of tail predication.
-///    This will be picked up by the ARM Low-overhead loop pass, which performs
-///    the final transformation to a DLSTP or WLSTP tail-predicated loop.
+/// predicated upon an get.active.lane.mask intrinsic. This pass looks at these
+/// get.active.lane.mask intrinsic and attempts to convert them to VCTP
+/// instructions. This will be picked up by the ARM Low-overhead loop pass later
+/// in the backend, which performs the final transformation to a DLSTP or WLSTP
+/// tail-predicated loop.
+//
+//===----------------------------------------------------------------------===//
 
 #include "ARM.h"
 #include "ARMSubtarget.h"
@@ -57,6 +47,7 @@
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/LoopUtils.h"
 #include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
 
@@ -112,23 +103,18 @@ class MVETailPredication : public LoopPass {
   bool runOnLoop(Loop *L, LPPassManager&) override;
 
 private:
-  /// Perform the relevant checks on the loop and convert if possible.
-  bool TryConvert(Value *TripCount);
-
-  /// Return whether this is a vectorized loop, that contains masked
-  /// load/stores.
-  bool IsPredicatedVectorLoop();
+  /// Perform the relevant checks on the loop and convert active lane masks if
+  /// possible.
+  bool TryConvertActiveLaneMask(Value *TripCount);
 
   /// Perform several checks on the arguments of @llvm.get.active.lane.mask
   /// intrinsic. E.g., check that the loop induction variable and the element
   /// count are of the form we expect, and also perform overflow checks for
   /// the new expressions that are created.
-  bool IsSafeActiveMask(IntrinsicInst *ActiveLaneMask, Value *TripCount,
-                        FixedVectorType *VecTy);
+  bool IsSafeActiveMask(IntrinsicInst *ActiveLaneMask, Value *TripCount);
 
   /// Insert the intrinsic to represent the effect of tail predication.
-  void InsertVCTPIntrinsic(IntrinsicInst *ActiveLaneMask, Value *TripCount,
-                           FixedVectorType *VecTy);
+  void InsertVCTPIntrinsic(IntrinsicInst *ActiveLaneMask, Value *TripCount);
 
   /// Rematerialize the iteration count in exit blocks, which enables
   /// ARMLowOverheadLoops to better optimise away loop update statements inside
@@ -138,25 +124,6 @@ class MVETailPredication : public LoopPass {
 
 } // end namespace
 
-static bool IsDecrement(Instruction &I) {
-  auto *Call = dyn_cast<IntrinsicInst>(&I);
-  if (!Call)
-    return false;
-
-  Intrinsic::ID ID = Call->getIntrinsicID();
-  return ID == Intrinsic::loop_decrement_reg;
-}
-
-static bool IsMasked(Instruction *I) {
-  auto *Call = dyn_cast<IntrinsicInst>(I);
-  if (!Call)
-    return false;
-
-  Intrinsic::ID ID = Call->getIntrinsicID();
-  return ID == Intrinsic::masked_store || ID == Intrinsic::masked_load ||
-         isGatherScatter(Call);
-}
-
 bool MVETailPredication::runOnLoop(Loop *L, LPPassManager&) {
   if (skipLoop(L) || !EnableTailPredication)
     return false;
@@ -207,147 +174,11 @@ bool MVETailPredication::runOnLoop(Loop *L, LPPassManager&) {
       return false;
   }
 
-  // Search for the hardware loop intrinic that decrements the loop counter.
-  IntrinsicInst *Decrement = nullptr;
-  for (auto *BB : L->getBlocks()) {
-    for (auto &I : *BB) {
-      if (IsDecrement(I)) {
-        Decrement = cast<IntrinsicInst>(&I);
-        break;
-      }
-    }
-  }
-
-  if (!Decrement)
-    return false;
-
-  LLVM_DEBUG(dbgs() << "ARM TP: Running on Loop: " << *L << *Setup << "\n"
-             << *Decrement << "\n");
-
-  if (!TryConvert(Setup->getArgOperand(0))) {
-    LLVM_DEBUG(dbgs() << "ARM TP: Can't tail-predicate this loop.\n");
-    return false;
-  }
-
-  return true;
-}
-
-static FixedVectorType *getVectorType(IntrinsicInst *I) {
-  unsigned ID = I->getIntrinsicID();
-  FixedVectorType *VecTy;
-  if (ID == Intrinsic::masked_load || isGather(I)) {
-    if (ID == Intrinsic::arm_mve_vldr_gather_base_wb ||
-        ID == Intrinsic::arm_mve_vldr_gather_base_wb_predicated)
-      // then the type is a StructType
-      VecTy = dyn_cast<FixedVectorType>(I->getType()->getContainedType(0));
-    else
-      VecTy = dyn_cast<FixedVectorType>(I->getType());
-  } else if (ID == Intrinsic::masked_store) {
-    VecTy = dyn_cast<FixedVectorType>(I->getOperand(0)->getType());
-  } else {
-    VecTy = dyn_cast<FixedVectorType>(I->getOperand(2)->getType());
-  }
-  assert(VecTy && "No scalable vectors expected here");
-  return VecTy;
-}
-
-bool MVETailPredication::IsPredicatedVectorLoop() {
-  // Check that the loop contains at least one masked load/store intrinsic.
-  // We only support 'normal' vector instructions - other than masked
-  // load/stores.
-  bool ActiveLaneMask = false;
-  for (auto *BB : L->getBlocks()) {
-    for (auto &I : *BB) {
-      auto *Int = dyn_cast<IntrinsicInst>(&I);
-      if (!Int)
-        continue;
-
-      switch (Int->getIntrinsicID()) {
-      case Intrinsic::get_active_lane_mask:
-        ActiveLaneMask = true;
-        continue;
-      case Intrinsic::sadd_sat:
-      case Intrinsic::uadd_sat:
-      case Intrinsic::ssub_sat:
-      case Intrinsic::usub_sat:
-      case Intrinsic::vector_reduce_add:
-        continue;
-      case Intrinsic::fma:
-      case Intrinsic::trunc:
-      case Intrinsic::rint:
-      case Intrinsic::round:
-      case Intrinsic::floor:
-      case Intrinsic::ceil:
-      case Intrinsic::fabs:
-        if (ST->hasMVEFloatOps())
-          continue;
-        break;
-      default:
-        break;
-      }
-      if (IsMasked(&I)) {
-        auto *VecTy = getVectorType(Int);
-        unsigned Lanes = VecTy->getNumElements();
-        unsigned ElementWidth = VecTy->getScalarSizeInBits();
-        // MVE vectors are 128-bit, but don't support 128 x i1.
-        // TODO: Can we support vectors larger than 128-bits?
-        unsigned MaxWidth = TTI->getRegisterBitWidth(true);
-        if (Lanes * ElementWidth > MaxWidth || Lanes == MaxWidth)
-          return false;
-        MaskedInsts.push_back(cast<IntrinsicInst>(&I));
-        continue;
-      }
-
-      for (const Use &U : Int->args()) {
-        if (isa<VectorType>(U->getType()))
-          return false;
-      }
-    }
-  }
-
-  if (!ActiveLaneMask) {
-    LLVM_DEBUG(dbgs() << "ARM TP: No get.active.lane.mask intrinsic found.\n");
-    return false;
-  }
-  return !MaskedInsts.empty();
-}
-
-// Look through the exit block to see whether there's a duplicate predicate
-// instruction. This can happen when we need to perform a select on values
-// from the last and previous iteration. Instead of doing a straight
-// replacement of that predicate with the vctp, clone the vctp and place it
-// in the block. This means that the VPR doesn't have to be live into the
-// exit block which should make it easier to convert this loop into a proper
-// tail predicated loop.
-static void Cleanup(SetVector<Instruction*> &MaybeDead, Loop *L) {
-  BasicBlock *Exit = L->getUniqueExitBlock();
-  if (!Exit) {
-    LLVM_DEBUG(dbgs() << "ARM TP: can't find loop exit block\n");
-    return;
-  }
-
-  // Drop references and add operands to check for dead.
-  SmallPtrSet<Instruction*, 4> Dead;
-  while (!MaybeDead.empty()) {
-    auto *I = MaybeDead.front();
-    MaybeDead.remove(I);
-    if (I->hasNUsesOrMore(1))
-      continue;
+  LLVM_DEBUG(dbgs() << "ARM TP: Running on Loop: " << *L << *Setup << "\n");
 
-    for (auto &U : I->operands())
-      if (auto *OpI = dyn_cast<Instruction>(U))
-        MaybeDead.insert(OpI);
+  bool Changed = TryConvertActiveLaneMask(Setup->getArgOperand(0));
 
-    Dead.insert(I);
-  }
-
-  for (auto *I : Dead) {
-    LLVM_DEBUG(dbgs() << "ARM TP: removing dead insn: "; I->dump());
-    I->eraseFromParent();
-  }
-
-  for (auto I : L->blocks())
-    DeleteDeadPHIs(I);
+  return Changed;
 }
 
 // The active lane intrinsic has this form:
@@ -368,15 +199,18 @@ static void Cleanup(SetVector<Instruction*> &MaybeDead, Loop *L) {
 // 3) The IV must be an induction phi with an increment equal to the
 //    vector width.
 bool MVETailPredication::IsSafeActiveMask(IntrinsicInst *ActiveLaneMask,
-    Value *TripCount, FixedVectorType *VecTy) {
+                                          Value *TripCount) {
   bool ForceTailPredication =
     EnableTailPredication == TailPredication::ForceEnabledNoReductions ||
     EnableTailPredication == TailPredication::ForceEnabled;
 
   Value *ElemCount = ActiveLaneMask->getOperand(1);
   auto *EC= SE->getSCEV(ElemCount);
   auto *TC = SE->getSCEV(TripCount);
-  int VectorWidth = VecTy->getNumElements();
+  int VectorWidth =
+      cast<FixedVectorType>(ActiveLaneMask->getType())->getNumElements();
+  if (VectorWidth != 4 && VectorWidth != 8 && VectorWidth != 16)
+    return false;
   ConstantInt *ConstElemCount = nullptr;
 
   // 1) Smoke tests that the original scalar loop TripCount (TC) belongs to
@@ -503,21 +337,22 @@ bool MVETailPredication::IsSafeActiveMask(IntrinsicInst *ActiveLaneMask,
   if (VectorWidth == StepValue)
     return true;
 
-  LLVM_DEBUG(dbgs() << "ARM TP: Step value " << StepValue << " doesn't match "
-             "vector width " << VectorWidth << "\n");
+  LLVM_DEBUG(dbgs() << "ARM TP: Step value " << StepValue
+                    << " doesn't match vector width " << VectorWidth << "\n");
 
   return false;
 }
 
 void MVETailPredication::InsertVCTPIntrinsic(IntrinsicInst *ActiveLaneMask,
-    Value *TripCount, FixedVectorType *VecTy) {
+                                             Value *TripCount) {
   IRBuilder<> Builder(L->getLoopPreheader()->getTerminator());
   Module *M = L->getHeader()->getModule();
   Type *Ty = IntegerType::get(M->getContext(), 32);
-  unsigned VectorWidth = VecTy->getNumElements();
+  unsigned VectorWidth =
+      cast<FixedVectorType>(ActiveLaneMask->getType())->getNumElements();
 
   // Insert a phi to count the number of elements processed by the loop.
-  Builder.SetInsertPoint(L->getHeader()->getFirstNonPHI()  );
+  Builder.SetInsertPoint(L->getHeader()->getFirstNonPHI());
   PHINode *Processed = Builder.CreatePHI(Ty, 2);
   Processed->addIncoming(ActiveLaneMask->getOperand(1), L->getLoopPreheader());
 
@@ -553,50 +388,36 @@ void MVETailPredication::InsertVCTPIntrinsic(IntrinsicInst *ActiveLaneMask,
              << "ARM TP: Inserted VCTP: " << *VCTPCall << "\n");
 }
 
-bool MVETailPredication::TryConvert(Value *TripCount) {
-  if (!IsPredicatedVectorLoop()) {
-    LLVM_DEBUG(dbgs() << "ARM TP: no masked instructions in loop.\n");
+bool MVETailPredication::TryConvertActiveLaneMask(Value *TripCount) {
+  SmallVector<IntrinsicInst *, 4> ActiveLaneMasks;
+  for (auto *BB : L->getBlocks())
+    for (auto &I : *BB)
+      if (auto *Int = dyn_cast<IntrinsicInst>(&I))
+        if (Int->getIntrinsicID() == Intrinsic::get_active_lane_mask)
+          ActiveLaneMasks.push_back(Int);
+
+  if (ActiveLaneMasks.empty())
     return false;
-  }
 
   LLVM_DEBUG(dbgs() << "ARM TP: Found predicated vector loop.\n");
-  SetVector<Instruction*> Predicates;
-
-  auto getPredicateOp = [](IntrinsicInst *I) {
-    unsigned IntrinsicID = I->getIntrinsicID();
-    if (IntrinsicID == Intrinsic::arm_mve_vldr_gather_offset_predicated ||
-        IntrinsicID == Intrinsic::arm_mve_vstr_scatter_offset_predicated)
-      return 5;
-    return (IntrinsicID == Intrinsic::masked_load || isGather(I)) ? 2 : 3;
-  };
-
-  // Walk through the masked intrinsics and try to find whether the predicate
-  // operand is generated by intrinsic @llvm.get.active.lane.mask().
-  for (auto *I : MaskedInsts) {
-    Value *PredOp = I->getArgOperand(getPredicateOp(I));
-    auto *Predicate = dyn_cast<Instruction>(PredOp);
-    if (!Predicate || Predicates.count(Predicate))
-      continue;
 
-    auto *ActiveLaneMask = dyn_cast<IntrinsicInst>(Predicate);
-    if (!ActiveLaneMask ||
-        ActiveLaneMask->getIntrinsicID() != Intrinsic::get_active_lane_mask)
-      continue;
-
-    Predicates.insert(Predicate);
+  for (auto *ActiveLaneMask : ActiveLaneMasks) {
     LLVM_DEBUG(dbgs() << "ARM TP: Found active lane mask: "
                       << *ActiveLaneMask << "\n");
 
-    auto *VecTy = getVectorType(I);
-    if (!IsSafeActiveMask(ActiveLaneMask, TripCount, VecTy)) {
+    if (!IsSafeActiveMask(ActiveLaneMask, TripCount)) {
       LLVM_DEBUG(dbgs() << "ARM TP: Not safe to insert VCTP.\n");
       return false;
     }
     LLVM_DEBUG(dbgs() << "ARM TP: Safe to insert VCTP.\n");
-    InsertVCTPIntrinsic(ActiveLaneMask, TripCount, VecTy);
+    InsertVCTPIntrinsic(ActiveLaneMask, TripCount);
   }
 
-  Cleanup(Predicates, L);
+  // Remove dead instructions and now dead phis.
+  for (auto *II : ActiveLaneMasks)
+    RecursivelyDeleteTriviallyDeadInstructions(II);
+  for (auto I : L->blocks())
+    DeleteDeadPHIs(I);
   return true;
 }