InstCombine: Fold ldexp(ldexp(x, a), b) -> ldexp(x, a + b)

[arsenm]

Summary

The problem here is overflow or underflow which would have occurred in
the inner operation, which the exponent offsetting avoids. We can do
this if we know the two exponents are in the same direction, or
reassoc flags allow unsafe reassociates.

[foad]

LGTM.

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2381–2383

Additionally it would be safe if both exponents are known to be <= 0, but I guess we don't have a helper function for that.

[arsenm]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2381–2383

Right we totally lack FP range tracking

[foad]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2381–2383

I just mean for the integer exponents, you're checking if they're both < 0, but that could be relaxed to both <= 0.

[arsenm]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2381–2383

This checks that they are both positive or both negative so I’m not sure what you mean

[foad]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2381–2383

Your code optimizes if they are both >= 0 or both < 0.
You could optimize if they are both >= 0 or both <= 0. This might optimize more cases.

[arsenm]

4f9aad964f4a8fb364f8069eca1825865ca85b80

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2381–2383

added todo

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

static std::optional<bool> getKnownSign(Value *Op, Instruction *CxtI,
   Value *X, *Y;
   if (match(Op, m_NSWSub(m_Value(X), m_Value(Y))))
     return isImpliedByDomCondition(ICmpInst::ICMP_SLT, X, Y, CxtI, DL);
 
   return isImpliedByDomCondition(
       ICmpInst::ICMP_SLT, Op, Constant::getNullValue(Op->getType()), CxtI, DL);
 }
 
+/// Return true if two values \p Op0 and \p Op1 are known to have the same sign.
+static bool signBitMustBeTheSame(Value *Op0, Value *Op1, Instruction *CxtI,
+                                 const DataLayout &DL, AssumptionCache *AC,
+                                 DominatorTree *DT) {
+  std::optional<bool> Known1 = getKnownSign(Op1, CxtI, DL, AC, DT);
+  if (!Known1)
+    return false;
+  std::optional<bool> Known0 = getKnownSign(Op0, CxtI, DL, AC, DT);
+  if (!Known0)
+    return false;
+  return *Known0 == *Known1;
+}
+
 /// Try to canonicalize min/max(X + C0, C1) as min/max(X, C1 - C0) + C0. This
 /// can trigger other combines.
 static Instruction *moveAddAfterMinMax(IntrinsicInst *II,
                                        InstCombiner::BuilderTy &Builder) {
   Intrinsic::ID MinMaxID = II->getIntrinsicID();
   assert((MinMaxID == Intrinsic::smax || MinMaxID == Intrinsic::smin ||
           MinMaxID == Intrinsic::umax || MinMaxID == Intrinsic::umin) &&
          "Expected a min or max intrinsic");
if (match(II->getArgOperand(0), m_OneUse(m_FNeg(m_Value(X))))) {
       // sin(-x) --> -sin(x)
       Value *NewSin = Builder.CreateUnaryIntrinsic(Intrinsic::sin, X, II);
       Instruction *FNeg = UnaryOperator::CreateFNeg(NewSin);
       FNeg->copyFastMathFlags(II);
       return FNeg;
     }
     break;
   }
+  case Intrinsic::ldexp: {
+    // ldexp(ldexp(x, a), b) -> ldexp(x, a + b)
+    //
+    // The danger is if the first ldexp would overflow to infinity or underflow
+    // to zero, but the combined exponent avoids it. We ignore this with
+    // reassoc.
+    //
+    // It's also safe to fold if we know both exponents have the same sign since
+    // it would just double down on the overflow/underflow which would occur
+    // anyway.
+    Value *Src = II->getArgOperand(0);
+    Value *Exp = II->getArgOperand(1);
+    Value *InnerSrc;
+    Value *InnerExp;
+    if (match(Src, m_OneUse(m_Intrinsic<Intrinsic::ldexp>(
+                       m_Value(InnerSrc), m_Value(InnerExp)))) &&
+        Exp->getType() == InnerExp->getType()) {
+      FastMathFlags FMF = II->getFastMathFlags();
+      FastMathFlags InnerFlags = cast<FPMathOperator>(Src)->getFastMathFlags();
+
+      if ((FMF.allowReassoc() && InnerFlags.allowReassoc()) ||
+          signBitMustBeTheSame(Exp, InnerExp, II, DL, &AC, &DT)) {
+        // TODO: Add nsw/nuw probably safe if integer type exceeds exponent
+        // width.
+        Value *NewExp = Builder.CreateAdd(InnerExp, Exp);
+        II->setArgOperand(1, NewExp);
+        II->setFastMathFlags(InnerFlags); // Or the inner flags.
+        return replaceOperand(*II, 0, InnerSrc);
+      }
+    }
+
+    break;
+  }
   case Intrinsic::ptrauth_auth:
   case Intrinsic::ptrauth_resign: {
     // (sign|resign) + (auth|resign) can be folded by omitting the middle
     // sign+auth component if the key and discriminator match.
     bool NeedSign = II->getIntrinsicID() == Intrinsic::ptrauth_resign;
     Value *Key = II->getArgOperand(1);
     Value *Disc = II->getArgOperand(2);
 

llvm/test/Transforms/InstCombine/ldexp.ll

;
   %ldexp0 = call float @llvm.ldexp.f32.i32(float %x, i32 %a)
   %ldexp1 = call reassoc float @llvm.ldexp.f32.i32(float %ldexp0, i32 %b)
   ret float %ldexp1
 }
 
 define float @ldexp_reassoc_ldexp_reassoc(float %x, i32 %a, i32 %b) {
 ; CHECK-LABEL: define float @ldexp_reassoc_ldexp_reassoc
 ; CHECK-SAME: (float [[X:%.*]], i32 [[A:%.*]], i32 [[B:%.*]]) {
-; CHECK-NEXT:    [[LDEXP0:%.*]] = call reassoc float @llvm.ldexp.f32.i32(float [[X]], i32 [[A]])
-; CHECK-NEXT:    [[LDEXP1:%.*]] = call reassoc float @llvm.ldexp.f32.i32(float [[LDEXP0]], i32 [[B]])
+; CHECK-NEXT:    [[TMP1:%.*]] = add i32 [[A]], [[B]]
+; CHECK-NEXT:    [[LDEXP1:%.*]] = call reassoc float @llvm.ldexp.f32.i32(float [[X]], i32 [[TMP1]])
 ; CHECK-NEXT:    ret float [[LDEXP1]]
 ;
   %ldexp0 = call reassoc float @llvm.ldexp.f32.i32(float %x, i32 %a)
   %ldexp1 = call reassoc float @llvm.ldexp.f32.i32(float %ldexp0, i32 %b)
   ret float %ldexp1
 }
 
 ; Test that we or the inner and outer flags
 define float @ldexp_reassoc_ldexp_reassoc_preserve_flags(float %x, i32 %a, i32 %b) {
 ; CHECK-LABEL: define float @ldexp_reassoc_ldexp_reassoc_preserve_flags
 ; CHECK-SAME: (float [[X:%.*]], i32 [[A:%.*]], i32 [[B:%.*]]) {
-; CHECK-NEXT:    [[LDEXP0:%.*]] = call reassoc ninf float @llvm.ldexp.f32.i32(float [[X]], i32 [[A]])
-; CHECK-NEXT:    [[LDEXP1:%.*]] = call reassoc nnan float @llvm.ldexp.f32.i32(float [[LDEXP0]], i32 [[B]])
+; CHECK-NEXT:    [[TMP1:%.*]] = add i32 [[A]], [[B]]
+; CHECK-NEXT:    [[LDEXP1:%.*]] = call reassoc nnan ninf float @llvm.ldexp.f32.i32(float [[X]], i32 [[TMP1]])
 ; CHECK-NEXT:    ret float [[LDEXP1]]
 ;
   %ldexp0 = call reassoc ninf float @llvm.ldexp.f32.i32(float %x, i32 %a)
   %ldexp1 = call reassoc nnan float @llvm.ldexp.f32.i32(float %ldexp0, i32 %b)
   ret float %ldexp1
 }
 
 define <2 x float> @ldexp_reassoc_ldexp_reassoc_vec(<2 x float> %x, <2 x i32> %a, <2 x i32> %b) {
 ; CHECK-LABEL: define <2 x float> @ldexp_reassoc_ldexp_reassoc_vec
 ; CHECK-SAME: (<2 x float> [[X:%.*]], <2 x i32> [[A:%.*]], <2 x i32> [[B:%.*]]) {
-; CHECK-NEXT:    [[LDEXP0:%.*]] = call reassoc <2 x float> @llvm.ldexp.v2f32.v2i32(<2 x float> [[X]], <2 x i32> [[A]])
-; CHECK-NEXT:    [[LDEXP1:%.*]] = call reassoc <2 x float> @llvm.ldexp.v2f32.v2i32(<2 x float> [[LDEXP0]], <2 x i32> [[B]])
+; CHECK-NEXT:    [[TMP1:%.*]] = add <2 x i32> [[A]], [[B]]
+; CHECK-NEXT:    [[LDEXP1:%.*]] = call reassoc <2 x float> @llvm.ldexp.v2f32.v2i32(<2 x float> [[X]], <2 x i32> [[TMP1]])
 ; CHECK-NEXT:    ret <2 x float> [[LDEXP1]]
 ;
   %ldexp0 = call reassoc <2 x float> @llvm.ldexp.v2f32.v2i32(<2 x float> %x, <2 x i32> %a)
   %ldexp1 = call reassoc <2 x float> @llvm.ldexp.v2f32.v2i32(<2 x float> %ldexp0, <2 x i32> %b)
   ret <2 x float> %ldexp1
 }
 
 define float @ldexp_multi_use_ldexp(float %x, i32 %a, i32 %b, ptr %ptr) {
;
   %ldexp0 = call reassoc float @llvm.ldexp.f32.i32(float %x, i32 %a)
   %ldexp1 = call reassoc float @llvm.ldexp.f32.i64(float %ldexp0, i64 %b)
   ret float %ldexp1
 }
 
 define float @ldexp_ldexp_constants(float %x) {
 ; CHECK-LABEL: define float @ldexp_ldexp_constants
 ; CHECK-SAME: (float [[X:%.*]]) {
-; CHECK-NEXT:    [[LDEXP0:%.*]] = call reassoc float @llvm.ldexp.f32.i32(float [[X]], i32 8)
-; CHECK-NEXT:    [[LDEXP1:%.*]] = call reassoc float @llvm.ldexp.f32.i32(float [[LDEXP0]], i32 24)
+; CHECK-NEXT:    [[LDEXP1:%.*]] = call reassoc float @llvm.ldexp.f32.i32(float [[X]], i32 32)
 ; CHECK-NEXT:    ret float [[LDEXP1]]
 ;
   %ldexp0 = call reassoc float @llvm.ldexp.f32.i32(float %x, i32 8)
   %ldexp1 = call reassoc float @llvm.ldexp.f32.i32(float %ldexp0, i32 24)
   ret float %ldexp1
 }
 
 define float @ldexp_ldexp_opposite_constants(float %x) {
 ; CHECK-LABEL: define float @ldexp_ldexp_opposite_constants
 ; CHECK-SAME: (float [[X:%.*]]) {
-; CHECK-NEXT:    [[LDEXP0:%.*]] = call reassoc float @llvm.ldexp.f32.i32(float [[X]], i32 8)
-; CHECK-NEXT:    [[LDEXP1:%.*]] = call reassoc float @llvm.ldexp.f32.i32(float [[LDEXP0]], i32 -8)
-; CHECK-NEXT:    ret float [[LDEXP1]]
+; CHECK-NEXT:    ret float [[X]]
 ;
   %ldexp0 = call reassoc float @llvm.ldexp.f32.i32(float %x, i32 8)
   %ldexp1 = call reassoc float @llvm.ldexp.f32.i32(float %ldexp0, i32 -8)
   ret float %ldexp1
 }
 
 define float @ldexp_ldexp_negated_variable_reassoc(float %x, i32 %a) {
 ; CHECK-LABEL: define float @ldexp_ldexp_negated_variable_reassoc
 ; CHECK-SAME: (float [[X:%.*]], i32 [[A:%.*]]) {
-; CHECK-NEXT:    [[LDEXP0:%.*]] = call reassoc float @llvm.ldexp.f32.i32(float [[X]], i32 [[A]])
-; CHECK-NEXT:    [[NEG_A:%.*]] = sub i32 0, [[A]]
-; CHECK-NEXT:    [[LDEXP1:%.*]] = call reassoc float @llvm.ldexp.f32.i32(float [[LDEXP0]], i32 [[NEG_A]])
-; CHECK-NEXT:    ret float [[LDEXP1]]
+; CHECK-NEXT:    ret float [[X]]
 ;
   %ldexp0 = call reassoc float @llvm.ldexp.f32.i32(float %x, i32 %a)
   %neg.a = sub i32 0, %a
   %ldexp1 = call reassoc float @llvm.ldexp.f32.i32(float %ldexp0, i32 %neg.a)
   ret float %ldexp1
 }
 
 define float @ldexp_ldexp_negated_variable(float %x, i32 %a) {
;
   ret float %ldexp1
 }
 
 define float @ldexp_ldexp_both_exp_known_positive(float %x, i32 %a.arg, i32 %b.arg) {
 ; CHECK-LABEL: define float @ldexp_ldexp_both_exp_known_positive
 ; CHECK-SAME: (float [[X:%.*]], i32 [[A_ARG:%.*]], i32 [[B_ARG:%.*]]) {
 ; CHECK-NEXT:    [[A:%.*]] = and i32 [[A_ARG]], 127
 ; CHECK-NEXT:    [[B:%.*]] = and i32 [[B_ARG]], 127
-; CHECK-NEXT:    [[LDEXP0:%.*]] = call float @llvm.ldexp.f32.i32(float [[X]], i32 [[A]])
-; CHECK-NEXT:    [[LDEXP1:%.*]] = call float @llvm.ldexp.f32.i32(float [[LDEXP0]], i32 [[B]])
+; CHECK-NEXT:    [[TMP1:%.*]] = add nuw nsw i32 [[A]], [[B]]
+; CHECK-NEXT:    [[LDEXP1:%.*]] = call float @llvm.ldexp.f32.i32(float [[X]], i32 [[TMP1]])
 ; CHECK-NEXT:    ret float [[LDEXP1]]
 ;
   %a = and i32 %a.arg, 127
   %b = and i32 %b.arg, 127
   %ldexp0 = call float @llvm.ldexp.f32.i32(float %x, i32 %a)
   %ldexp1 = call float @llvm.ldexp.f32.i32(float %ldexp0, i32 %b)
   ret float %ldexp1
 }
 
 define float @ldexp_ldexp_both_exp_known_negative(float %x, ptr %a.ptr, ptr %b.ptr) {
 ; CHECK-LABEL: define float @ldexp_ldexp_both_exp_known_negative
 ; CHECK-SAME: (float [[X:%.*]], ptr [[A_PTR:%.*]], ptr [[B_PTR:%.*]]) {
 ; CHECK-NEXT:    [[A:%.*]] = load i32, ptr [[A_PTR]], align 4, !range [[RNG0:![0-9]+]]
 ; CHECK-NEXT:    [[B:%.*]] = load i32, ptr [[B_PTR]], align 4, !range [[RNG0]]
-; CHECK-NEXT:    [[LDEXP0:%.*]] = call float @llvm.ldexp.f32.i32(float [[X]], i32 [[A]])
-; CHECK-NEXT:    [[LDEXP1:%.*]] = call float @llvm.ldexp.f32.i32(float [[LDEXP0]], i32 [[B]])
+; CHECK-NEXT:    [[TMP1:%.*]] = add nsw i32 [[A]], [[B]]
+; CHECK-NEXT:    [[LDEXP1:%.*]] = call float @llvm.ldexp.f32.i32(float [[X]], i32 [[TMP1]])
 ; CHECK-NEXT:    ret float [[LDEXP1]]
 ;
   %a = load i32, ptr %a.ptr, !range !0
   %b = load i32, ptr %b.ptr, !range !0
   %ldexp0 = call float @llvm.ldexp.f32.i32(float %x, i32 %a)
   %ldexp1 = call float @llvm.ldexp.f32.i32(float %ldexp0, i32 %b)
   ret float %ldexp1
 }
;
   %ldexp0 = call float @llvm.ldexp.f32.i32(float %x, i32 %a)
   %ldexp1 = call float @llvm.ldexp.f32.i32(float %ldexp0, i32 %b)
   ret float %ldexp1
 }
 
 define float @ldexp_reassoc_ldexp_reassoc_0(float %x, i32 %y) {
 ; CHECK-LABEL: define float @ldexp_reassoc_ldexp_reassoc_0
 ; CHECK-SAME: (float [[X:%.*]], i32 [[Y:%.*]]) {
-; CHECK-NEXT:    [[LDEXP0:%.*]] = call reassoc float @llvm.ldexp.f32.i32(float [[X]], i32 0)
-; CHECK-NEXT:    [[LDEXP1:%.*]] = call reassoc float @llvm.ldexp.f32.i32(float [[LDEXP0]], i32 [[Y]])
+; CHECK-NEXT:    [[LDEXP1:%.*]] = call reassoc float @llvm.ldexp.f32.i32(float [[X]], i32 [[Y]])
 ; CHECK-NEXT:    ret float [[LDEXP1]]
 ;
   %ldexp0 = call reassoc float @llvm.ldexp.f32.i32(float %x, i32 0)
   %ldexp1 = call reassoc float @llvm.ldexp.f32.i32(float %ldexp0, i32 %y)
   ret float %ldexp1
 }
 
 define float @ldexp_ldexp_0(float %x, i32 %y) {
 ; CHECK-LABEL: define float @ldexp_ldexp_0
 ; CHECK-SAME: (float [[X:%.*]], i32 [[Y:%.*]]) {
-; CHECK-NEXT:    [[LDEXP0:%.*]] = call float @llvm.ldexp.f32.i32(float [[X]], i32 0)
-; CHECK-NEXT:    [[LDEXP1:%.*]] = call float @llvm.ldexp.f32.i32(float [[LDEXP0]], i32 [[Y]])
+; CHECK-NEXT:    [[LDEXP1:%.*]] = call float @llvm.ldexp.f32.i32(float [[X]], i32 [[Y]])
 ; CHECK-NEXT:    ret float [[LDEXP1]]
 ;
   %ldexp0 = call float @llvm.ldexp.f32.i32(float %x, i32 0)
   %ldexp1 = call float @llvm.ldexp.f32.i32(float %ldexp0, i32 %y)
   ret float %ldexp1
 }
 
 !0 = !{i32 -127, i32 0}
 !1 = !{i32 0, i32 127}