fold sext(trunc nsw) to trunc nsw

Author: wenju-he

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

@@ -1561,8 +1561,13 @@ Instruction *InstCombinerImpl::visitSExt(SExtInst &Sext) {
     // If the input has more sign bits than bits truncated, then convert
     // directly to final type.
     unsigned XBitSize = X->getType()->getScalarSizeInBits();
-    if (ComputeNumSignBits(X, &Sext) > XBitSize - SrcBitSize)
-      return CastInst::CreateIntegerCast(X, DestTy, /* isSigned */ true);
+    bool HasNSW = cast<TruncInst>(Src)->hasNoSignedWrap();
+    if (HasNSW || (ComputeNumSignBits(X, &Sext) > XBitSize - SrcBitSize)) {
+      auto *Res = CastInst::CreateIntegerCast(X, DestTy, /* isSigned */ true);
+      if (auto *ResTrunc = dyn_cast<TruncInst>(Res); ResTrunc && HasNSW)
+        ResTrunc->setHasNoSignedWrap(true);
+      return Res;
+    }
 
     // If input is a trunc from the destination type, then convert into shifts.
     if (Src->hasOneUse() && X->getType() == DestTy) {

llvm/test/Transforms/InstCombine/sext-of-trunc-nsw.ll

@@ -144,6 +144,17 @@ define i24 @wide_source_matching_signbits(i32 %x) {
   ret i24 %c
 }
 
+define i32 @wide_source_matching_signbits_has_nsw_flag(i64 %i) {
+; CHECK-LABEL: define i32 @wide_source_matching_signbits_has_nsw_flag(
+; CHECK-SAME: i64 [[I:%.*]]) {
+; CHECK-NEXT:    [[A:%.*]] = trunc nsw i64 [[I]] to i32
+; CHECK-NEXT:    ret i32 [[A]]
+;
+  %a = trunc nsw i64 %i to i16
+  %b = sext i16 %a to i32
+  ret i32 %b
+}
+
 ; negative test - not enough sign-bits
 
 define i24 @wide_source_not_matching_signbits(i32 %x) {