[InstCombine] narrow abs with sign-extended input

[spatel]

Summary

In the motivating cases from https://llvm.org/PR48816 , we have a trailing trunc. But that is not required to reduce the abs width:
https://alive2.llvm.org/ce/z/ECaz-p
...as long as we clear the int-min-is-poison bit (nsw).

We have some existing tests that are affected, and I'm not sure what the overall implications are, but in general I think we favor narrowing operations over preserving nsw/nuw. If not, we could restrict this transform based on type (shouldChangeType() and/or vector vs. scalar).

[lebedev.ri]

LGTM
I think narrowing ops is a win.

[RKSimon]

Cheers - please can you add the PR48816 test case

define i8 @transform_abs_epi8(i8 %0) {
  %2 = sext i8 %0 to i32
  %3 = tail call i32 @llvm.abs.i32(i32 %2, i1 true)
  %4 = trunc i32 %3 to i8
  ret i8 %4
}
declare i32 @llvm.abs.i32(i32, i1 immarg)

(and vector equivalent)

[spatel]

In D95235#2515765, @RKSimon wrote:

Cheers - please can you add the PR48816 test case

define i8 @transform_abs_epi8(i8 %0) {
  %2 = sext i8 %0 to i32
  %3 = tail call i32 @llvm.abs.i32(i32 %2, i1 true)
  %4 = trunc i32 %3 to i8
  ret i8 %4
}
declare i32 @llvm.abs.i32(i32, i1 immarg)

(and vector equivalent)

Sure - will add tests with trunc on pre-commit of tests and this patch.

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

if (Optional<bool> Sign = getKnownSign(IIOperand, II, DL, &AC, &DT)) {
         return replaceInstUsesWith(*II, IIOperand);
 
       // abs(x) -> -x if x < 0
       if (IntMinIsPoison)
         return BinaryOperator::CreateNSWNeg(IIOperand);
       return BinaryOperator::CreateNeg(IIOperand);
     }
 
+    // abs (sext X) --> zext (abs X*)
+    // Clear the IsIntMin (nsw) bit on the abs to allow narrowing.
+    if (match(IIOperand, m_OneUse(m_SExt(m_Value(X))))) {
+      Value *NarrowAbs =
+          Builder.CreateBinaryIntrinsic(Intrinsic::abs, X, Builder.getFalse());
+      return CastInst::Create(Instruction::ZExt, NarrowAbs, II->getType());
+    }
+
     break;
   }
   case Intrinsic::bswap: {
     Value *IIOperand = II->getArgOperand(0);
     Value *X = nullptr;
 
     // bswap(trunc(bswap(x))) -> trunc(lshr(x, c))
     if (match(IIOperand, m_Trunc(m_BSwap(m_Value(X))))) {

llvm/test/Transforms/InstCombine/abs-1.ll

;
   %cmp = icmp slt i8 %x, 1
   %neg = sub i8 0, %x
   %abs = select i1 %cmp, i8 %neg, i8 %x
   ret i8 %abs
 }
 
 define i32 @abs_canonical_5(i8 %x) {
 ; CHECK-LABEL: @abs_canonical_5(
-; CHECK-NEXT:    [[CONV:%.*]] = sext i8 [[X:%.*]] to i32
-; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @llvm.abs.i32(i32 [[CONV]], i1 true)
-; CHECK-NEXT:    ret i32 [[TMP1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.abs.i8(i8 [[X:%.*]], i1 false)
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i8 [[TMP1]] to i32
+; CHECK-NEXT:    ret i32 [[TMP2]]
 ;
   %cmp = icmp sgt i8 %x, 0
   %conv = sext i8 %x to i32
   %neg = sub i32 0, %conv
   %abs = select i1 %cmp, i32 %conv, i32 %neg
   ret i32 %abs
 }
 
;
   %cmp = icmp slt i8 %x, 1
   %neg = sub i8 0, %x
   %abs = select i1 %cmp, i8 %x, i8 %neg
   ret i8 %abs
 }
 
 define i32 @nabs_canonical_5(i8 %x) {
 ; CHECK-LABEL: @nabs_canonical_5(
-; CHECK-NEXT:    [[CONV:%.*]] = sext i8 [[X:%.*]] to i32
-; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @llvm.abs.i32(i32 [[CONV]], i1 false)
-; CHECK-NEXT:    [[ABS:%.*]] = sub nsw i32 0, [[TMP1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.abs.i8(i8 [[X:%.*]], i1 false)
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i8 [[TMP1]] to i32
+; CHECK-NEXT:    [[ABS:%.*]] = sub nsw i32 0, [[TMP2]]
 ; CHECK-NEXT:    ret i32 [[ABS]]
 ;
   %cmp = icmp sgt i8 %x, 0
   %conv = sext i8 %x to i32
   %neg = sub i32 0, %conv
   %abs = select i1 %cmp, i32 %neg, i32 %conv
   ret i32 %abs
 }

llvm/test/Transforms/InstCombine/abs-intrinsic.ll

;
   %and2 = and <4 x i32> %abs, <i32 -4, i32 -4, i32 -4, i32 -4>
   ret <4 x i32> %and2
 }
 
 ; Make sure we infer this add doesn't overflow. The input to the abs has 3
 ; sign bits, the abs reduces this to 2 sign bits.
 define i32 @abs_signbits(i30 %x) {
 ; CHECK-LABEL: @abs_signbits(
-; CHECK-NEXT:    [[EXT:%.*]] = sext i30 [[X:%.*]] to i32
-; CHECK-NEXT:    [[ABS:%.*]] = call i32 @llvm.abs.i32(i32 [[EXT]], i1 false)
-; CHECK-NEXT:    [[ADD:%.*]] = add nuw nsw i32 [[ABS]], 1
+; CHECK-NEXT:    [[TMP1:%.*]] = call i30 @llvm.abs.i30(i30 [[X:%.*]], i1 false)
+; CHECK-NEXT:    [[NARROW:%.*]] = add nuw i30 [[TMP1]], 1
+; CHECK-NEXT:    [[ADD:%.*]] = zext i30 [[NARROW]] to i32
 ; CHECK-NEXT:    ret i32 [[ADD]]
 ;
   %ext = sext i30 %x to i32
   %abs = call i32 @llvm.abs.i32(i32 %ext, i1 false)
   %add = add i32 %abs, 1
   ret i32 %add
 }
 
 define <4 x i32> @abs_signbits_vec(<4 x i30> %x) {
 ; CHECK-LABEL: @abs_signbits_vec(
-; CHECK-NEXT:    [[EXT:%.*]] = sext <4 x i30> [[X:%.*]] to <4 x i32>
-; CHECK-NEXT:    [[ABS:%.*]] = call <4 x i32> @llvm.abs.v4i32(<4 x i32> [[EXT]], i1 false)
-; CHECK-NEXT:    [[ADD:%.*]] = add nuw nsw <4 x i32> [[ABS]], <i32 1, i32 1, i32 1, i32 1>
+; CHECK-NEXT:    [[TMP1:%.*]] = call <4 x i30> @llvm.abs.v4i30(<4 x i30> [[X:%.*]], i1 false)
+; CHECK-NEXT:    [[NARROW:%.*]] = add nuw <4 x i30> [[TMP1]], <i30 1, i30 1, i30 1, i30 1>
+; CHECK-NEXT:    [[ADD:%.*]] = zext <4 x i30> [[NARROW]] to <4 x i32>
 ; CHECK-NEXT:    ret <4 x i32> [[ADD]]
 ;
   %ext = sext <4 x i30> %x to <4 x i32>
   %abs = call <4 x i32> @llvm.abs.v4i32(<4 x i32> %ext, i1 false)
   %add = add <4 x i32> %abs, <i32 1, i32 1, i32 1, i32 1>
   ret <4 x i32> %add
 }
 

 ;
   %abs = call i8 @llvm.abs.i8(i8 %x, i1 false)
   %cmp = icmp ne i8 %abs, -128
   ret i1 %cmp
 }
 
 define i32 @abs_sext(i8 %x) {
 ; CHECK-LABEL: @abs_sext(
-; CHECK-NEXT:    [[S:%.*]] = sext i8 [[X:%.*]] to i32
-; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.abs.i32(i32 [[S]], i1 false)
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.abs.i8(i8 [[X:%.*]], i1 false)
+; CHECK-NEXT:    [[A:%.*]] = zext i8 [[TMP1]] to i32
 ; CHECK-NEXT:    ret i32 [[A]]
 ;
   %s = sext i8 %x to i32
   %a = call i32 @llvm.abs.i32(i32 %s, i1 0)
   ret i32 %a
 }
 
 define <3 x i82> @abs_nsw_sext(<3 x i7> %x) {
 ; CHECK-LABEL: @abs_nsw_sext(
-; CHECK-NEXT:    [[S:%.*]] = sext <3 x i7> [[X:%.*]] to <3 x i82>
-; CHECK-NEXT:    [[A:%.*]] = call <3 x i82> @llvm.abs.v3i82(<3 x i82> [[S]], i1 true)
+; CHECK-NEXT:    [[TMP1:%.*]] = call <3 x i7> @llvm.abs.v3i7(<3 x i7> [[X:%.*]], i1 false)
+; CHECK-NEXT:    [[A:%.*]] = zext <3 x i7> [[TMP1]] to <3 x i82>
 ; CHECK-NEXT:    ret <3 x i82> [[A]]
 ;
   %s = sext <3 x i7> %x to <3 x i82>
   %a = call <3 x i82> @llvm.abs.v3i82(<3 x i82> %s, i1 1)
   ret <3 x i82> %a
 }
 
 define i32 @abs_sext_extra_use(i8 %x, i32* %p) {
 ; CHECK-LABEL: @abs_sext_extra_use(
 ; CHECK-NEXT:    [[S:%.*]] = sext i8 [[X:%.*]] to i32
 ; CHECK-NEXT:    store i32 [[S]], i32* [[P:%.*]], align 4
 ; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.abs.i32(i32 [[S]], i1 false)
 ; CHECK-NEXT:    ret i32 [[A]]
 ;
   %s = sext i8 %x to i32
   store i32 %s, i32* %p
   %a = call i32 @llvm.abs.i32(i32 %s, i1 0)
   ret i32 %a
 }
 
 ; PR48816
 
 define i8 @trunc_abs_sext(i8 %x) {
 ; CHECK-LABEL: @trunc_abs_sext(
-; CHECK-NEXT:    [[S:%.*]] = sext i8 [[X:%.*]] to i32
-; CHECK-NEXT:    [[A:%.*]] = tail call i32 @llvm.abs.i32(i32 [[S]], i1 true)
-; CHECK-NEXT:    [[T:%.*]] = trunc i32 [[A]] to i8
-; CHECK-NEXT:    ret i8 [[T]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8 @llvm.abs.i8(i8 [[X:%.*]], i1 false)
+; CHECK-NEXT:    ret i8 [[TMP1]]
 ;
   %s = sext i8 %x to i32
   %a = tail call i32 @llvm.abs.i32(i32 %s, i1 true)
   %t = trunc i32 %a to i8
   ret i8 %t
 }
 
 define <4 x i8> @trunc_abs_sext_vec(<4 x i8> %x) {
 ; CHECK-LABEL: @trunc_abs_sext_vec(
-; CHECK-NEXT:    [[S:%.*]] = sext <4 x i8> [[X:%.*]] to <4 x i32>
-; CHECK-NEXT:    [[A:%.*]] = tail call <4 x i32> @llvm.abs.v4i32(<4 x i32> [[S]], i1 true)
-; CHECK-NEXT:    [[T:%.*]] = trunc <4 x i32> [[A]] to <4 x i8>
-; CHECK-NEXT:    ret <4 x i8> [[T]]
+; CHECK-NEXT:    [[TMP1:%.*]] = call <4 x i8> @llvm.abs.v4i8(<4 x i8> [[X:%.*]], i1 false)
+; CHECK-NEXT:    ret <4 x i8> [[TMP1]]
 ;
   %s = sext <4 x i8> %x to <4 x i32>
   %a = tail call <4 x i32> @llvm.abs.v4i32(<4 x i32> %s, i1 true)
   %t = trunc <4 x i32> %a to <4 x i8>
   ret <4 x i8> %t
 }