[ARM][MVE] Tail-predication: clean-up removing unused code

[SjoerdMeijer]

Summary

After the rewrite of this pass (D79175) I missed one thing: the inserted VCTP intrinsic can be cloned to exit blocks if there are instructions present in it that perform the same operation. However, it turns out that for handling reductions, see D75533, it's actually easier not not to have the VCTP in exit blocks. If the exit block was cloned, we rematerialized the trip count in the exit blocks, because that made other dead-code removal easier. But our dead-code removal got more powerful, so I don't think this is still needed and therefore this also removes RematerializeIterCount.

[samparker]

Not sure if we'll have to revisit this again if we're doing, say a reduction, and but also have a VCMP in the loop which writes to the VPR meaning it needs to be spilled and restored for the VPSEL... But we can cross that bridge if we come to it, cheers!

[SjoerdMeijer]

Thanks, and yes, fully agreed. At the moment it isn't doing what it claims to do, so I'd prefer to get rid of it, and revise and change this when we need to.

llvm/lib/Target/ARM/MVETailPredication.cpp

 DisableTailPredication("disable-mve-tail-predication", cl::Hidden,
                        cl::init(true),
                        cl::desc("Disable MVE Tail Predication"));
 namespace {
 
 class MVETailPredication : public LoopPass {
   SmallVector<IntrinsicInst*, 4> MaskedInsts;
   Loop *L = nullptr;
-  LoopInfo *LI = nullptr;
-  const DataLayout *DL;
-  DominatorTree *DT = nullptr;
   ScalarEvolution *SE = nullptr;
   TargetTransformInfo *TTI = nullptr;
-  TargetLibraryInfo *TLI = nullptr;
-  bool ClonedVCTPInExitBlock = false;
 
 public:
   static char ID;
 
   MVETailPredication() : LoopPass(ID) { }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<ScalarEvolutionWrapperPass>();
     AU.addRequired<LoopInfoWrapperPass>();
     AU.addRequired<TargetPassConfig>();
     AU.addRequired<TargetTransformInfoWrapperPass>();
-    AU.addRequired<DominatorTreeWrapperPass>();
-    AU.addRequired<TargetLibraryInfoWrapperPass>();
     AU.addPreserved<LoopInfoWrapperPass>();
     AU.setPreservesCFG();
   }
 
   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 checks on the arguments of @llvm.get.active.lane.mask
   /// intrinsic: check if the first is a loop induction variable, and for the
   /// the second check that no overflow can occur in the expression that use
   /// this backedge-taken count.
   bool IsSafeActiveMask(IntrinsicInst *ActiveLaneMask, Value *TripCount,
                         FixedVectorType *VecTy);
 
   /// Insert the intrinsic to represent the effect of tail predication.
   void InsertVCTPIntrinsic(IntrinsicInst *ActiveLaneMask, Value *TripCount,
-                           FixedVectorType *VecTy,
-                           DenseMap<Instruction *, Instruction *> &NewPredicates);
+                           FixedVectorType *VecTy);
 
   /// Rematerialize the iteration count in exit blocks, which enables
   /// ARMLowOverheadLoops to better optimise away loop update statements inside
   /// hardware-loops.
   void RematerializeIterCount();
 };
 
 } // end namespace
static bool IsMasked(Instruction *I) {
   if (!Call)
     return false;
 
   Intrinsic::ID ID = Call->getIntrinsicID();
   // TODO: Support gather/scatter expand/compress operations.
   return ID == Intrinsic::masked_store || ID == Intrinsic::masked_load;
 }
 
-void MVETailPredication::RematerializeIterCount() {
-  SmallVector<WeakTrackingVH, 16> DeadInsts;
-  SCEVExpander Rewriter(*SE, *DL, "mvetp");
-  ReplaceExitVal ReplaceExitValue = AlwaysRepl;
-
-  formLCSSARecursively(*L, *DT, LI, SE);
-  rewriteLoopExitValues(L, LI, TLI, SE, TTI, Rewriter, DT, ReplaceExitValue,
-                        DeadInsts);
-}
-
 bool MVETailPredication::runOnLoop(Loop *L, LPPassManager&) {
   if (skipLoop(L) || DisableTailPredication)
     return false;
 
   MaskedInsts.clear();
   Function &F = *L->getHeader()->getParent();
   auto &TPC = getAnalysis<TargetPassConfig>();
   auto &TM = TPC.getTM<TargetMachine>();
   auto *ST = &TM.getSubtarget<ARMSubtarget>(F);
-  DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
-  LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
   TTI = &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
   SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
-  auto *TLIP = getAnalysisIfAvailable<TargetLibraryInfoWrapperPass>();
-  TLI = TLIP ? &TLIP->getTLI(*L->getHeader()->getParent()) : nullptr;
-  DL = &L->getHeader()->getModule()->getDataLayout();
   this->L = L;
 
   // The MVE and LOB extensions are combined to enable tail-predication, but
   // there's nothing preventing us from generating VCTP instructions for v8.1m.
   if (!ST->hasMVEIntegerOps() || !ST->hasV8_1MMainlineOps()) {
     LLVM_DEBUG(dbgs() << "ARM TP: Not a v8.1m.main+mve target.\n");
     return false;
   }
for (auto &I : *BB) {
         break;
       }
     }
   }
 
   if (!Decrement)
     return false;
 
-  ClonedVCTPInExitBlock = 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;
   }
 
-  if (ClonedVCTPInExitBlock)
-    RematerializeIterCount();
   return true;
 }
 
 static FixedVectorType *getVectorType(IntrinsicInst *I) {
   unsigned TypeOp = I->getIntrinsicID() == Intrinsic::masked_load ? 0 : 1;
   auto *PtrTy = cast<PointerType>(I->getOperand(TypeOp)->getType());
   auto *VecTy = cast<FixedVectorType>(PtrTy->getElementType());
   assert(VecTy && "No scalable vectors expected here");

 
 // 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 bool Cleanup(DenseMap<Instruction*, Instruction*> &NewPredicates,
-                    SetVector<Instruction*> &MaybeDead, Loop *L) {
+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 false;
-  }
-
-  bool ClonedVCTPInExitBlock = false;
-
-  for (auto &Pair : NewPredicates) {
-    Instruction *OldPred = Pair.first;
-    Instruction *NewPred = Pair.second;
-
-    for (auto &I : *Exit) {
-      if (I.isSameOperationAs(OldPred)) {
-        Instruction *PredClone = NewPred->clone();
-        PredClone->insertBefore(&I);
-        I.replaceAllUsesWith(PredClone);
-        MaybeDead.insert(&I);
-        ClonedVCTPInExitBlock = true;
-        LLVM_DEBUG(dbgs() << "ARM TP: replacing: "; I.dump();
-                   dbgs() << "ARM TP: with:      "; PredClone->dump());
-        break;
-      }
-    }
+    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;
 
     for (auto &U : I->operands())
       if (auto *OpI = dyn_cast<Instruction>(U))
         MaybeDead.insert(OpI);
 
     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 ClonedVCTPInExitBlock;
 }
 
 // The active lane intrinsic has this form:
 //
 //    @llvm.get.active.lane.mask(IV, BTC)
 //
 // Here we perform checks that this intrinsic behaves as expected,
 // which means:
static Value *getNumElements(BasicBlock *Preheader, Value *BTC) {
   IRBuilder<> Builder(Preheader->getTerminator());
   Value *NumElements = Builder.CreateAdd(BTC,
         ConstantInt::get(BTC->getType(), 1), "num.elements");
   LLVM_DEBUG(dbgs() << "ARM TP: Created num elems: " << *NumElements << "\n");
   return NumElements;
 }
 
 void MVETailPredication::InsertVCTPIntrinsic(IntrinsicInst *ActiveLaneMask,
-    Value *TripCount, FixedVectorType *VecTy,
-    DenseMap<Instruction*, Instruction*> &NewPredicates) {
+    Value *TripCount, FixedVectorType *VecTy) {
   IRBuilder<> Builder(L->getLoopPreheader()->getTerminator());
   Module *M = L->getHeader()->getModule();
   Type *Ty = IntegerType::get(M->getContext(), 32);
   unsigned VectorWidth = VecTy->getNumElements();
 
   // The backedge-taken count in @llvm.get.active.lane.mask, its 2nd operand,
   // is one less than the trip count. So we need to find or create
   // %num.elements = %BTC + 1 in the preheader.
case 16: VCTPID = Intrinsic::arm_mve_vctp8; break;
     // vector wants to be <2 x i1>, but v2i1 is not a legal MVE
     // type, so problems happen at isel time.
     // Intrinsic::arm_mve_vctp64 exists for ACLE intrinsics
     // purposes, but takes a v4i1 instead of a v2i1.
   }
   Function *VCTP = Intrinsic::getDeclaration(M, VCTPID);
   Value *VCTPCall = Builder.CreateCall(VCTP, Processed);
   ActiveLaneMask->replaceAllUsesWith(VCTPCall);
-  NewPredicates[ActiveLaneMask] = cast<Instruction>(VCTPCall);
 
   // Add the incoming value to the new phi.
   // TODO: This add likely already exists in the loop.
   Value *Remaining = Builder.CreateSub(Processed, Factor);
   Processed->addIncoming(Remaining, L->getLoopLatch());
   LLVM_DEBUG(dbgs() << "ARM TP: Insert processed elements phi: "
              << *Processed << "\n"
              << "ARM TP: Inserted VCTP: " << *VCTPCall << "\n");
 }
 
 bool MVETailPredication::TryConvert(Value *TripCount) {
   if (!IsPredicatedVectorLoop()) {
     LLVM_DEBUG(dbgs() << "ARM TP: no masked instructions in loop.\n");
     return false;
   }
 
   LLVM_DEBUG(dbgs() << "ARM TP: Found predicated vector loop.\n");
-
   SetVector<Instruction*> Predicates;
-  DenseMap<Instruction*, Instruction*> NewPredicates;
 
   // 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) {
     unsigned PredOp = I->getIntrinsicID() == Intrinsic::masked_load ? 2 : 3;
     auto *Predicate = dyn_cast<Instruction>(I->getArgOperand(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);
     LLVM_DEBUG(dbgs() << "ARM TP: Found active lane mask: "
                       << *ActiveLaneMask << "\n");
 
     auto *VecTy = getVectorType(I);
     if (!IsSafeActiveMask(ActiveLaneMask, TripCount, VecTy)) {
       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, NewPredicates);
+    InsertVCTPIntrinsic(ActiveLaneMask, TripCount, VecTy);
   }
 
-  // Now clean up.
-  ClonedVCTPInExitBlock = Cleanup(NewPredicates, Predicates, L);
+  Cleanup(Predicates, L);
   return true;
 }
 
 Pass *llvm::createMVETailPredicationPass() {
   return new MVETailPredication();
 }
 
 char MVETailPredication::ID = 0;
 
 INITIALIZE_PASS_BEGIN(MVETailPredication, DEBUG_TYPE, DESC, false, false)
 INITIALIZE_PASS_END(MVETailPredication, DEBUG_TYPE, DESC, false, false)

llvm/test/CodeGen/Thumb2/LowOverheadLoops/vector-reduce-mve-tail.ll

 ; CHECK: vector.body:
 ; CHECK: [[ELTS:%[^ ]+]] = phi i32 [ %N, %vector.ph ], [ [[SUB:%[^ ]+]], %vector.body ]
 ; CHECK: [[VCTP:%[^ ]+]] = call <4 x i1> @llvm.arm.mve.vctp32(i32 [[ELTS]])
 ; CHECK: [[SUB]] = sub i32 [[ELTS]], 4
 ; CHECK: call <4 x i32> @llvm.masked.load.v4i32.p0v4i32(<4 x i32>* {{.*}}, i32 4, <4 x i1> [[VCTP]]
 ; CHECK: call <4 x i32> @llvm.masked.load.v4i32.p0v4i32(<4 x i32>* {{.*}}, i32 4, <4 x i1> [[VCTP]],
 
 ; CHECK: middle.block:
-; CHECK: [[VCTP_CLONE:%[^ ]+]] = call <4 x i1> @llvm.arm.mve.vctp32(i32 [[REMAT_ITER:%.*]])
-; CHECK: [[VPSEL:%[^ ]+]] = select <4 x i1> [[VCTP_CLONE]],
+; CHECK: [[VPSEL:%[^ ]+]] = select <4 x i1> [[VCTP]],
 ; CHECK: call i32 @llvm.experimental.vector.reduce.add.v4i32(<4 x i32> [[VPSEL]])
 
 define i32 @vec_mul_reduce_add(i32* noalias nocapture readonly %a, i32* noalias nocapture readonly %b, i32 %N) {
 entry:
   %cmp8 = icmp eq i32 %N, 0
   %0 = add i32 %N, 3
   %1 = lshr i32 %0, 2
   %2 = shl nuw i32 %1, 2
vector.body:                                      ; preds = %vector.body, %vector.ph
   %index.next = add i32 %index, 4
   %scevgep = getelementptr i32, i32* %lsr.iv, i32 4
   %scevgep3 = getelementptr i32, i32* %lsr.iv2, i32 4
   %10 = call i32 @llvm.loop.decrement.reg.i32.i32.i32(i32 %6, i32 1)
   %11 = icmp ne i32 %10, 0
   br i1 %11, label %vector.body, label %middle.block
 
 middle.block:                                     ; preds = %vector.body
-; TODO: check that the intrinsic is also emitted here by the loop vectoriser
-;  %12 = icmp ule <4 x i32> %induction, %broadcast.splat12
-  %12 = call <4 x i1> @llvm.get.active.lane.mask.v4i1.i32(i32 %index, i32 %trip.count.minus.1)
-
-  %13 = select <4 x i1> %12, <4 x i32> %9, <4 x i32> %vec.phi
-  %14 = call i32 @llvm.experimental.vector.reduce.add.v4i32(<4 x i32> %13)
+  %12 = select <4 x i1> %7, <4 x i32> %9, <4 x i32> %vec.phi
+  %13 = call i32 @llvm.experimental.vector.reduce.add.v4i32(<4 x i32> %12)
   br label %for.cond.cleanup
 
 for.cond.cleanup:                                 ; preds = %middle.block, %entry
-  %res.0.lcssa = phi i32 [ 0, %entry ], [ %14, %middle.block ]
+  %res.0.lcssa = phi i32 [ 0, %entry ], [ %13, %middle.block ]
   ret i32 %res.0.lcssa
 }
 
 declare <4 x i32> @llvm.masked.load.v4i32.p0v4i32(<4 x i32>*, i32 immarg, <4 x i1>, <4 x i32>)
 declare i32 @llvm.experimental.vector.reduce.add.v4i32(<4 x i32>)
 declare void @llvm.set.loop.iterations.i32(i32)
 declare i32 @llvm.loop.decrement.reg.i32.i32.i32(i32, i32)
 declare <4 x i1> @llvm.get.active.lane.mask.v4i1.i32(i32, i32)