[TypePromotion] Allow SExtInst as sink

[samparker]

Summary

If these instructions are in the tree, then they're almost certainly going to need a sign extend anyway, so having a trunc in the IR to handle them shouldn't cause any harm. I also noticed some significantly better codegen when targeting arm on one of the IR tests.

llvm/lib/CodeGen/TypePromotion.cpp

if (auto *Return = dyn_cast<ReturnInst>(V))
     return LessOrEqualTypeSize(Return->getReturnValue());
   if (auto *ZExt = dyn_cast<ZExtInst>(V))
     return GreaterThanTypeSize(ZExt);
   if (auto *Switch = dyn_cast<SwitchInst>(V))
     return LessThanTypeSize(Switch->getCondition());
   if (auto *ICmp = dyn_cast<ICmpInst>(V))
     return ICmp->isSigned() || LessThanTypeSize(ICmp->getOperand(0));
 
-  return isa<CallInst>(V);
+  return isa<SExtInst>(V) || isa<CallInst>(V);
 }
 
 /// Return whether this instruction can safely wrap.
 bool TypePromotion::isSafeWrap(Instruction *I) {
   // We can support a, potentially, wrapping instruction (I) if:
   // - It is only used by an unsigned icmp.
   // - The icmp uses a constant.
   // - The wrapping value (I) is decreasing, i.e would underflow - wrapping
if (auto *I = dyn_cast<Instruction>(V)) {
     case Instruction::PHI:
     case Instruction::Select:
     case Instruction::Ret:
     case Instruction::Load:
     case Instruction::Trunc:
     case Instruction::BitCast:
       return isSupportedType(I);
     case Instruction::ZExt:
+    case Instruction::SExt:
       return isSupportedType(I->getOperand(0));
     case Instruction::ICmp:
       // Now that we allow small types than TypeSize, only allow icmp of
       // TypeSize because they will require a trunc to be legalised.
       // TODO: Allow icmp of smaller types, and calculate at the end
       // whether the transform would be beneficial.
       if (isa<PointerType>(I->getOperand(0)->getType()))
         return true;

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

exit:
   ret i32 %ret
 }
 
 define i32 @icmp_i32_zext_i16(i16* %ptr) {
 ; CHECK-LABEL: @icmp_i32_zext_i16(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[GEP:%.*]] = getelementptr inbounds i16, i16* [[PTR:%.*]], i32 0
 ; CHECK-NEXT:    [[TMP0:%.*]] = load i16, i16* [[GEP]], align 1
-; CHECK-NEXT:    [[TMP1:%.*]] = sub nuw nsw i16 [[TMP0]], 1
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i16 [[TMP0]] to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = sub nuw nsw i32 [[TMP1]], 1
 ; CHECK-NEXT:    [[CONV44:%.*]] = zext i16 [[TMP0]] to i32
 ; CHECK-NEXT:    br label [[PREHEADER:%.*]]
 ; CHECK:       preheader:
 ; CHECK-NEXT:    br label [[BODY:%.*]]
 ; CHECK:       body:
-; CHECK-NEXT:    [[TMP2:%.*]] = phi i16 [ [[TMP1]], [[PREHEADER]] ], [ [[TMP3:%.*]], [[IF_END:%.*]] ]
+; CHECK-NEXT:    [[TMP3:%.*]] = phi i32 [ [[TMP2]], [[PREHEADER]] ], [ [[TMP6:%.*]], [[IF_END:%.*]] ]
 ; CHECK-NEXT:    [[SI_0274:%.*]] = phi i32 [ [[CONV44]], [[PREHEADER]] ], [ [[INC:%.*]], [[IF_END]] ]
-; CHECK-NEXT:    [[CONV51266:%.*]] = zext i16 [[TMP2]] to i32
-; CHECK-NEXT:    [[CMP52267:%.*]] = icmp eq i32 [[SI_0274]], [[CONV51266]]
+; CHECK-NEXT:    [[TMP4:%.*]] = trunc i32 [[TMP3]] to i16
+; CHECK-NEXT:    [[CMP52267:%.*]] = icmp eq i32 [[SI_0274]], [[TMP3]]
 ; CHECK-NEXT:    br i1 [[CMP52267]], label [[IF_END]], label [[EXIT:%.*]]
 ; CHECK:       if.end:
 ; CHECK-NEXT:    [[INC]] = add i32 [[SI_0274]], 1
 ; CHECK-NEXT:    [[GEP1:%.*]] = getelementptr inbounds i16, i16* [[PTR]], i32 [[INC]]
-; CHECK-NEXT:    [[TMP3]] = load i16, i16* [[GEP1]], align 1
+; CHECK-NEXT:    [[TMP5:%.*]] = load i16, i16* [[GEP1]], align 1
+; CHECK-NEXT:    [[TMP6]] = zext i16 [[TMP5]] to i32
 ; CHECK-NEXT:    br label [[BODY]]
 ; CHECK:       exit:
-; CHECK-NEXT:    [[RET:%.*]] = sext i16 [[TMP2]] to i32
+; CHECK-NEXT:    [[TMP7:%.*]] = trunc i32 [[TMP3]] to i16
+; CHECK-NEXT:    [[RET:%.*]] = sext i16 [[TMP7]] to i32
 ; CHECK-NEXT:    ret i32 [[RET]]
 ;
 entry:
   %gep = getelementptr inbounds i16, i16* %ptr, i32 0
   %0 = load i16, i16* %gep, align 1
   %1 = sub nuw nsw i16 %0, 1
   %conv44 = zext i16 %0 to i32
   br label %preheader

 @e = common dso_local local_unnamed_addr global i8 0, align 1
 @a = common dso_local local_unnamed_addr global i8 0, align 1
 @d = common dso_local local_unnamed_addr global i32 0, align 4
 
 define void @and_trunc_two_zext() {
 ; CHECK-LABEL: @and_trunc_two_zext(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = load i16, i16* @c, align 2
-; CHECK-NEXT:    [[TMP1:%.*]] = load i16, i16* @b, align 2
-; CHECK-NEXT:    [[CONV:%.*]] = sext i16 [[TMP1]] to i32
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i16 [[TMP0]] to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = load i16, i16* @b, align 2
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i16 [[TMP2]] to i32
+; CHECK-NEXT:    [[TMP4:%.*]] = trunc i32 [[TMP3]] to i16
+; CHECK-NEXT:    [[CONV:%.*]] = sext i16 [[TMP4]] to i32
 ; CHECK-NEXT:    store i32 [[CONV]], i32* @f, align 4
-; CHECK-NEXT:    [[TMP2:%.*]] = trunc i16 [[TMP1]] to i8
-; CHECK-NEXT:    [[CONV1:%.*]] = and i8 [[TMP2]], 1
-; CHECK-NEXT:    store i8 [[CONV1]], i8* @e, align 1
-; CHECK-NEXT:    [[TMP3:%.*]] = load i8, i8* @a, align 1
-; CHECK-NEXT:    [[NARROW:%.*]] = mul nuw i8 [[TMP3]], [[CONV1]]
-; CHECK-NEXT:    [[MUL:%.*]] = zext i8 [[NARROW]] to i32
-; CHECK-NEXT:    store i32 [[MUL]], i32* @d, align 4
-; CHECK-NEXT:    [[TMP4:%.*]] = zext i8 [[NARROW]] to i16
-; CHECK-NEXT:    [[CONV5:%.*]] = or i16 [[TMP0]], [[TMP4]]
-; CHECK-NEXT:    [[TOBOOL:%.*]] = icmp eq i16 [[CONV5]], 0
+; CHECK-NEXT:    [[TMP5:%.*]] = and i32 [[TMP3]], 255
+; CHECK-NEXT:    [[CONV1:%.*]] = and i32 [[TMP5]], 1
+; CHECK-NEXT:    [[TMP6:%.*]] = trunc i32 [[CONV1]] to i8
+; CHECK-NEXT:    store i8 [[TMP6]], i8* @e, align 1
+; CHECK-NEXT:    [[TMP7:%.*]] = load i8, i8* @a, align 1
+; CHECK-NEXT:    [[TMP8:%.*]] = zext i8 [[TMP7]] to i32
+; CHECK-NEXT:    [[NARROW:%.*]] = mul nuw i32 [[TMP8]], [[CONV1]]
+; CHECK-NEXT:    [[TMP9:%.*]] = trunc i32 [[NARROW]] to i8
+; CHECK-NEXT:    store i32 [[NARROW]], i32* @d, align 4
+; CHECK-NEXT:    [[CONV5:%.*]] = or i32 [[TMP1]], [[NARROW]]
+; CHECK-NEXT:    [[TOBOOL:%.*]] = icmp eq i32 [[CONV5]], 0
 ; CHECK-NEXT:    br i1 [[TOBOOL]], label [[IF_END:%.*]], label [[FOR_COND:%.*]]
 ; CHECK:       for.cond:
 ; CHECK-NEXT:    br label [[FOR_COND]]
 ; CHECK:       if.end:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %0 = load i16, i16* @c, align 2
cond.end:
   store i8 %cond, i8* @a, align 1
   ret void
 }
 
 define i1 @dont_replace_trunc_1(i8* %a, i16* %b, i16* %c, i32* %d, i8* %e, i32* %f) {
 ; CHECK-LABEL: @dont_replace_trunc_1(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = load i16, i16* [[C:%.*]], align 2
-; CHECK-NEXT:    [[TMP1:%.*]] = load i16, i16* [[B:%.*]], align 2
-; CHECK-NEXT:    [[CONV:%.*]] = sext i16 [[TMP1]] to i32
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i16 [[TMP0]] to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = load i16, i16* [[B:%.*]], align 2
+; CHECK-NEXT:    [[TMP3:%.*]] = zext i16 [[TMP2]] to i32
+; CHECK-NEXT:    [[TMP4:%.*]] = trunc i32 [[TMP3]] to i16
+; CHECK-NEXT:    [[CONV:%.*]] = sext i16 [[TMP4]] to i32
 ; CHECK-NEXT:    store i32 [[CONV]], i32* [[F:%.*]], align 4
-; CHECK-NEXT:    [[TMP2:%.*]] = trunc i16 [[TMP1]] to i8
-; CHECK-NEXT:    [[CONV1:%.*]] = and i8 [[TMP2]], 1
-; CHECK-NEXT:    store i8 [[CONV1]], i8* [[E:%.*]], align 1
-; CHECK-NEXT:    [[TMP3:%.*]] = load i8, i8* [[A:%.*]], align 1
-; CHECK-NEXT:    [[NARROW:%.*]] = mul nuw i8 [[TMP3]], [[CONV1]]
-; CHECK-NEXT:    [[MUL:%.*]] = zext i8 [[NARROW]] to i32
-; CHECK-NEXT:    store i32 [[MUL]], i32* [[D:%.*]], align 4
-; CHECK-NEXT:    [[TMP4:%.*]] = zext i8 [[NARROW]] to i16
-; CHECK-NEXT:    [[CONV5:%.*]] = or i16 [[TMP0]], [[TMP4]]
-; CHECK-NEXT:    [[TOBOOL:%.*]] = icmp eq i16 [[CONV5]], 0
+; CHECK-NEXT:    [[TMP5:%.*]] = and i32 [[TMP3]], 255
+; CHECK-NEXT:    [[CONV1:%.*]] = and i32 [[TMP5]], 1
+; CHECK-NEXT:    [[TMP6:%.*]] = trunc i32 [[CONV1]] to i8
+; CHECK-NEXT:    store i8 [[TMP6]], i8* [[E:%.*]], align 1
+; CHECK-NEXT:    [[TMP7:%.*]] = load i8, i8* [[A:%.*]], align 1
+; CHECK-NEXT:    [[TMP8:%.*]] = zext i8 [[TMP7]] to i32
+; CHECK-NEXT:    [[NARROW:%.*]] = mul nuw i32 [[TMP8]], [[CONV1]]
+; CHECK-NEXT:    [[TMP9:%.*]] = trunc i32 [[NARROW]] to i8
+; CHECK-NEXT:    store i32 [[NARROW]], i32* [[D:%.*]], align 4
+; CHECK-NEXT:    [[CONV5:%.*]] = or i32 [[TMP1]], [[NARROW]]
+; CHECK-NEXT:    [[TOBOOL:%.*]] = icmp eq i32 [[CONV5]], 0
 ; CHECK-NEXT:    ret i1 [[TOBOOL]]
 ;
 entry:
   %0 = load i16, i16* %c, align 2
   %1 = load i16, i16* %b, align 2
   %conv = sext i16 %1 to i32
   store i32 %conv, i32* %f, align 4
   %2 = trunc i16 %1 to i8