[TypePromotion] Avoid unnecessary trunc zext pairs

[samparker]

Summary

Following on from D115451, this is the next in a series of improvements for this pass.

[adriantong1024]

This LGTM

llvm/lib/CodeGen/TypePromotion.cpp

 //
 // We collect all instruction connecting sources and sinks in a worklist, so
 // that we can mutate these instruction and perform type promotion when it is
 // legal to do so.
 
 namespace {
 class IRPromoter {
   LLVMContext &Ctx;
-  IntegerType *OrigTy = nullptr;
   unsigned PromotedWidth = 0;
   SetVector<Value *> &Visited;
   SetVector<Value *> &Sources;
   SetVector<Instruction *> &Sinks;
   SmallPtrSetImpl<Instruction *> &SafeWrap;
   IntegerType *ExtTy = nullptr;
   SmallPtrSet<Value *, 8> NewInsts;
   SmallPtrSet<Instruction *, 4> InstsToRemove;
   DenseMap<Value *, SmallVector<Type *, 4>> TruncTysMap;
   SmallPtrSet<Value *, 8> Promoted;
 
   void ReplaceAllUsersOfWith(Value *From, Value *To);
   void ExtendSources();
   void ConvertTruncs();
   void PromoteTree();
   void TruncateSinks();
   void Cleanup();
 
 public:
-  IRPromoter(LLVMContext &C, IntegerType *Ty, unsigned Width,
+  IRPromoter(LLVMContext &C, unsigned Width,
              SetVector<Value *> &visited, SetVector<Value *> &sources,
              SetVector<Instruction *> &sinks,
              SmallPtrSetImpl<Instruction *> &wrap)
-      : Ctx(C), OrigTy(Ty), PromotedWidth(Width), Visited(visited),
+      : Ctx(C), PromotedWidth(Width), Visited(visited),
         Sources(sources), Sinks(sinks), SafeWrap(wrap) {
     ExtTy = IntegerType::get(Ctx, PromotedWidth);
-    assert(OrigTy->getPrimitiveSizeInBits().getFixedSize() <
-               ExtTy->getPrimitiveSizeInBits().getFixedSize() &&
-           "Original type not smaller than extended type");
   }
 
   void Mutate();
 };
 
 class TypePromotion : public FunctionPass {
   unsigned TypeSize = 0;
   LLVMContext *Ctx = nullptr;

 }
 
 /// Return true if V will require any promoted values to be truncated for the
 /// the IR to remain valid. We can't mutate the value type of these
 /// instructions.
 bool TypePromotion::isSink(Value *V) {
   // TODO The truncate also isn't actually necessary because we would already
   // proved that the data value is kept within the range of the original data
-  // type.
+  // type. We currently remove any truncs inserted for handling zext sinks.
 
   // Sinks are:
   // - points where the value in the register is being observed, such as an
   //   icmp, switch or store.
   // - points where value types have to match, such as calls and returns.
   // - zext are included to ease the transformation and are generally removed
   //   later on.
   if (auto *Store = dyn_cast<StoreInst>(V))
for (auto V : Visited) {
     Value *Src = ZExt->getOperand(0);
     if (ZExt->getSrcTy() == ZExt->getDestTy()) {
       LLVM_DEBUG(dbgs() << "IR Promotion: Removing unnecessary cast: " << *ZExt
                         << "\n");
       ReplaceAllUsersOfWith(ZExt, Src);
       continue;
     }
 
-    // Unless they produce a value that is narrower than ExtTy, we can
-    // replace the result of the zext with the input of a newly inserted
-    // trunc.
-    if (NewInsts.count(Src) && isa<TruncInst>(Src) &&
-        Src->getType() == OrigTy) {
+    // We've inserted a trunc for a zext sink, but we already know that the
+    // input is in range, negating the need for the trunc.
+    if (NewInsts.count(Src) && isa<TruncInst>(Src)) {
       auto *Trunc = cast<TruncInst>(Src);
       assert(Trunc->getOperand(0)->getType() == ExtTy &&
              "expected inserted trunc to be operating on i32");
       ReplaceAllUsersOfWith(ZExt, Trunc->getOperand(0));
     }
   }
 
   for (auto *I : InstsToRemove) {
for (auto *V : Visited) {
     if (auto *I = dyn_cast<Instruction>(Masked))
       NewInsts.insert(I);
 
     ReplaceAllUsersOfWith(Trunc, Masked);
   }
 }
 
 void IRPromoter::Mutate() {
-  LLVM_DEBUG(dbgs() << "IR Promotion: Promoting use-def chains from "
-                    << OrigTy->getBitWidth() << " to " << PromotedWidth
-                    << "-bits\n");
+  LLVM_DEBUG(dbgs() << "IR Promotion: Promoting use-def chains to "
+                    << PromotedWidth << "-bits\n");
 
   // Cache original types of the values that will likely need truncating
   for (auto *I : Sinks) {
     if (auto *Call = dyn_cast<CallInst>(I)) {
       for (Value *Arg : Call->args())
         TruncTysMap[Call].push_back(Arg->getType());
     } else if (auto *Switch = dyn_cast<SwitchInst>(I))
       TruncTysMap[I].push_back(Switch->getCondition()->getType());
for (auto *V : CurrentVisited) {
     ++ToPromote;
   }
 
   // DAG optimizations should be able to handle these cases better, especially
   // for function arguments.
   if (ToPromote < 2 || (Blocks.size() == 1 && (NonFreeArgs > SafeWrap.size())))
     return false;
 
-  IRPromoter Promoter(*Ctx, cast<IntegerType>(OrigTy), PromotedWidth,
-                      CurrentVisited, Sources, Sinks, SafeWrap);
+  IRPromoter Promoter(*Ctx, PromotedWidth, CurrentVisited, Sources, Sinks,
+                      SafeWrap);
   Promoter.Mutate();
   return true;
 }
 
 bool TypePromotion::runOnFunction(Function &F) {
   if (skipFunction(F) || DisablePromotion)
     return false;
 

llvm/test/Transforms/TypePromotion/ARM/casts.ll

entry:
   %mul = zext i8 %narrow to i32
   store i32 %mul, i32* %d, align 4
   %4 = zext i8 %narrow to i16
   %conv5 = or i16 %0, %4
   %tobool = icmp eq i16 %conv5, 0
   ret i1 %tobool
 }
 
-define i32 @dont_replace_trunc_2(i16* %a, i8* %b) {
-; CHECK-LABEL: @dont_replace_trunc_2(
+define i32 @dont_return_inserted_trunc(i16* %a, i8* %b) {
+; CHECK-LABEL: @dont_return_inserted_trunc(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = load i16, i16* [[A:%.*]], align 2
 ; CHECK-NEXT:    [[TMP1:%.*]] = zext i16 [[TMP0]] to i32
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i32 [[TMP1]], 8
 ; CHECK-NEXT:    [[NARROW:%.*]] = select i1 [[CMP]], i32 [[TMP1]], i32 0
 ; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[NARROW]], 255
 ; CHECK-NEXT:    [[TMP3:%.*]] = load i8, i8* [[B:%.*]], align 1
 ; CHECK-NEXT:    [[TMP4:%.*]] = zext i8 [[TMP3]] to i32
 ; CHECK-NEXT:    [[OR:%.*]] = or i32 [[TMP4]], [[TMP2]]
 ; CHECK-NEXT:    [[TMP5:%.*]] = trunc i32 [[OR]] to i8
 ; CHECK-NEXT:    store i8 [[TMP5]], i8* [[B]], align 1
 ; CHECK-NEXT:    [[TMP6:%.*]] = trunc i32 [[OR]] to i8
-; CHECK-NEXT:    [[CONV5:%.*]] = zext i8 [[TMP6]] to i32
-; CHECK-NEXT:    ret i32 [[CONV5]]
+; CHECK-NEXT:    ret i32 [[OR]]
 ;
 entry:
   %0 = load i16, i16* %a, align 2
   %cmp = icmp ugt i16 %0, 8
   %narrow = select i1 %cmp, i16 %0, i16 0
   %cond = trunc i16 %narrow to i8
   %1 = load i8, i8* %b, align 1
   %or = or i8 %1, %cond

 ; CHECK-NEXT:    [[TMP1:%.*]] = zext i16 [[TMP0]] to i32
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[TMP1]], 0
 ; CHECK-NEXT:    br i1 [[CMP]], label [[COND_END:%.*]], label [[COND_FALSE:%.*]]
 ; CHECK:       cond.false:
 ; CHECK-NEXT:    [[TMP2:%.*]] = urem i32 535, [[TMP1]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = and i32 [[TMP2]], 255
 ; CHECK-NEXT:    [[TMP4:%.*]] = udiv i32 [[TMP3]], 3
 ; CHECK-NEXT:    [[TMP5:%.*]] = trunc i32 [[TMP4]] to i8
-; CHECK-NEXT:    [[PHITMP:%.*]] = zext i8 [[TMP5]] to i32
 ; CHECK-NEXT:    br label [[COND_END]]
 ; CHECK:       cond.end:
-; CHECK-NEXT:    [[COND:%.*]] = phi i32 [ [[PHITMP]], [[COND_FALSE]] ], [ 0, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[COND:%.*]] = phi i32 [ [[TMP4]], [[COND_FALSE]] ], [ 0, [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    ret i32 [[COND]]
 ;
 entry:
   %0 = load i16, i16* %a
   %cmp = icmp eq i16 %0, 0
   br i1 %cmp, label %cond.end, label %cond.false
 
 cond.false: