[TypePromotion] Avoid some unnecessary truncs

[samparker]

Summary

Check for legal zext 'sinks' before inserting a trunc.

I think this resolves some/most/all of the issue in D111237. I won't have the time to fix this if it's broken though.

edit: I setup a stage 2 on my mac and this looks fine enough.

[dmgreen]

Sorry for the delay. This doesn't fix the original case unfortunately, as that version doesn't start with an icmp. The test case is a little different with the loads feeding the loop condition.
But sounds like it should be an improvement, so LGTM (I would suggest committing after the branch ;) )

llvm/lib/CodeGen/TypePromotion.cpp

if (auto *Switch = dyn_cast<SwitchInst>(I)) {
       Type *Ty = TruncTysMap[Switch][0];
       if (Instruction *Trunc = InsertTrunc(Switch->getCondition(), Ty)) {
         Trunc->moveBefore(Switch);
         Switch->setCondition(Trunc);
       }
       continue;
     }
 
+    // Don't insert a trunc for a zext which can still legally promote.
+    if (auto ZExt = dyn_cast<ZExtInst>(I))
+      if (ZExt->getType()->getScalarSizeInBits() > PromotedWidth)
+        continue;
+
     // Now handle the others.
     for (unsigned i = 0; i < I->getNumOperands(); ++i) {
       Type *Ty = TruncTysMap[I][i];
       if (Instruction *Trunc = InsertTrunc(I->getOperand(i), Ty)) {
         Trunc->moveBefore(I);
         I->setOperand(i, Trunc);
       }
     }

llvm/test/Transforms/TypePromotion/AArch64/lit.local.cfg

+if not 'AArch64' in config.root.targets:
+    config.unsupported = True

llvm/test/Transforms/TypePromotion/AArch64/phi-zext-gep.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -mtriple=aarch64 -type-promotion -verify -S %s -o - | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+
+define dso_local i32 @avoid_trunc_gep(i8* nocapture readonly %ip) {
+; CHECK-LABEL: @avoid_trunc_gep(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TAG_0_IN8:%.*]] = load i8, i8* [[IP:%.*]], align 1
+; CHECK-NEXT:    [[TMP0:%.*]] = zext i8 [[TAG_0_IN8]] to i32
+; CHECK-NEXT:    [[CMP9:%.*]] = icmp ult i32 [[TMP0]], 100
+; CHECK-NEXT:    br i1 [[CMP9]], label [[FOR_BODY_PREHEADER:%.*]], label [[FOR_END:%.*]]
+; CHECK:       for.body.preheader:
+; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
+; CHECK:       for.body:
+; CHECK-NEXT:    [[TAG_0_IN10:%.*]] = phi i32 [ [[TMP1:%.*]], [[FOR_BODY]] ], [ [[TMP0]], [[FOR_BODY_PREHEADER]] ]
+; CHECK-NEXT:    [[TAG_0:%.*]] = zext i32 [[TAG_0_IN10]] to i64
+; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds i8, i8* [[IP]], i64 [[TAG_0]]
+; CHECK-NEXT:    [[TAG_0_IN:%.*]] = load i8, i8* [[ARRAYIDX]], align 1
+; CHECK-NEXT:    [[TMP1]] = zext i8 [[TAG_0_IN]] to i32
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[TMP1]], 100
+; CHECK-NEXT:    br i1 [[CMP]], label [[FOR_BODY]], label [[FOR_END_LOOPEXIT:%.*]]
+; CHECK:       for.end.loopexit:
+; CHECK-NEXT:    br label [[FOR_END]]
+; CHECK:       for.end:
+; CHECK-NEXT:    [[TAG_0_IN_LCSSA:%.*]] = phi i32 [ [[TMP0]], [[ENTRY:%.*]] ], [ [[TMP1]], [[FOR_END_LOOPEXIT]] ]
+; CHECK-NEXT:    [[TMP2:%.*]] = trunc i32 [[TAG_0_IN_LCSSA]] to i8
+; CHECK-NEXT:    ret i32 [[TAG_0_IN_LCSSA]]
+;
+entry:
+  %tag.0.in8 = load i8, i8* %ip, align 1
+  %cmp9 = icmp ult i8 %tag.0.in8, 100
+  br i1 %cmp9, label %for.body.preheader, label %for.end
+
+for.body.preheader:                               ; preds = %entry
+  br label %for.body
+
+for.body:                                         ; preds = %for.body.preheader, %for.body
+  %tag.0.in10 = phi i8 [ %tag.0.in, %for.body ], [ %tag.0.in8, %for.body.preheader ]
+  %tag.0 = zext i8 %tag.0.in10 to i64
+  %arrayidx = getelementptr inbounds i8, i8* %ip, i64 %tag.0
+  %tag.0.in = load i8, i8* %arrayidx, align 1
+  %cmp = icmp ult i8 %tag.0.in, 100
+  br i1 %cmp, label %for.body, label %for.end.loopexit
+
+for.end.loopexit:                                 ; preds = %for.body
+  br label %for.end
+
+for.end:                                          ; preds = %for.end.loopexit, %entry
+  %tag.0.in.lcssa = phi i8 [ %tag.0.in8, %entry ], [ %tag.0.in, %for.end.loopexit ]
+  %conv3 = zext i8 %tag.0.in.lcssa to i32
+  ret i32 %conv3
+}