[InstCombine] Relaxed constraints of uses for exp(X) * exp(Y) -> exp(X + Y) and exp2(X) * exp2(Y) -> exp2(X + Y)

[vdsered]

Summary

InstCombine didn't perform the transformations when fmul's operands were the same instruction because it required to have one use for each of them which is false in the case. This patch fixes this + adds tests for them and introduces a new function isOnlyUserOfAnyOperand to check these cases in a single place.

This patch is a result of discussion in https://reviews.llvm.org/D102574.

[efriedma]

llvm/include/llvm/Transforms/Utils/Local.h

470

Don't expose isRestrictiveUseConstraintMet in the header if you don't plan to use it anywhere.

474

I don't like the name isRelaxedUseConstraintMet; doesn't provide any indication what it means at first glance. Maybe isOnlyUserOfAnyOperand?

llvm/lib/Transforms/Utils/Local.cpp

3311

llvm::is_splat?

[vdsered]

Updating diff

1. Fixed warnings of clang-tidy
2. Refactored implementation for this patch

[vdsered]

llvm/include/llvm/Transforms/Utils/Local.h

470

I saw how a check like Op0->hasOneUse() || Op1->hasOneUse() is heavily used as a precondition for many patterns and I thought of making it as a helper here because isOnlyUserOfAnyOperand uses it too. However, there are also places where one writes something like match(Op0, m_OneUse(...)) and it can probably be a little bit better in terms of readability. So, I'm not sure if we actually need this function. Removed it here as it is not that important for the goal of the patch.

474

Your suggestion is a good alternative to the previous name. Renamed it.

llvm/lib/Transforms/Utils/Local.cpp

3311

Fixed.

[spatel]

llvm/lib/Transforms/Utils/Local.cpp

3309–3310

Use llvm::any_of() to shorten, and implicit convert operand list to Value?

return any_of(I.operands(), [](Value *V) { return V->hasOneUse();}) || ...

3312

I->operands() ?

[vdsered]

Replaced std::any_of with llvm::any_of and used I->operands() instead of making range directly via begin/end

[vdsered]

llvm/lib/Transforms/Utils/Local.cpp

3309–3310

Fixed.

3312

Fixed.

[vdsered]

Removed irrelevant comments

[spatel]

llvm/lib/Transforms/Utils/Local.cpp

3305

This seems inverted. How about creating an inspection function within the Instruction class itself?

bool Instruction::isOnlyUserOfAnyOperand() const;

[lebedev.ri]

I"m not sure i like isOnlyUserOfAnyOperand() direction,
even though i've encountered this problem before.

I would recommend to simply use hasOneUser() in place of hasOneUse()?

[vdsered]

As I understand it

if (match(Op0, m_Intrinsic<Intrinsic::exp2>(m_Value(X))) &&
        match(Op1, m_Intrinsic<Intrinsic::exp2>(m_Value(Y))) &&
        isOnlyUserOfAnyOperand(&I))

will be

if (match(Op0, m_Intrinsic<Intrinsic::exp2>(m_Value(X))) &&
        match(Op1, m_Intrinsic<Intrinsic::exp2>(m_Value(Y))) &&
        (X->hasOneUser() || Y -> hasOneUser()))

The difference between them is that isOnlyUserOfAnyOperand depends on how many operands an instruction has (a known count, we use it for mul now where it is 2). On the other, hasOneUser depends on the total number of uses of each operand (can be an arbitrary number). I'm not sure, but there can theoretically be situations when a number of uses is more than a number of ops, right?

@lebedev.ri, we could implement isOnlyUserOfAnyOperand like this

bool llvm::isOnlyUserOfAnyOperand(Instruction *I) {
    return llvm::any_of(I->operands(), [](Value *V) { return V->hasOneUser(); })
}

This solution would be good too especially because it is easier to understand, but I'd leave it as-is without hasOneUser in case there are no any performance/logical issues with it.

[vdsered]

llvm/lib/Transforms/Utils/Local.cpp

3305

I'm a little bit confused where to place such a function because here we have, for example, isInstructionTriviallyDead-like functions, but we have more those within Instruction class. I move it to the class.

[vdsered]

Moved isOnlyUserOfAnyOperand to Instruction class

[spatel]

In D102698#2777107, @vdsered wrote:

bool llvm::isOnlyUserOfAnyOperand(Instruction *I) {
    return llvm::any_of(I->operands(), [](Value *V) { return V->hasOneUser(); })
}

This solution would be good too especially because it is easier to understand, but I'd leave it as-is without hasOneUser in case there are no any performance/logical issues with it.

I didn't understand the comment. That implementation definitely looks nicer than what we have in the current patch (nit: no need to explicitly use llvm:: namespace prefix). Is there a reason not to do that (or the alternate suggestion of changing the caller code only with hasOneUser()? I slightly prefer adding this inspection method because it will make the caller code more obviously different than the usual hasOneUse() code pattern, but either way seems fine to me.

[vdsered]

isOnlyUserOfAnyOperand uses hasOneUser to figure out if at least one of ops has only one user

[spatel]

I pushed some baseline tests here:
01120fe5b398

Please rebase/update the patch after that (manually remove TODO comments in the test file).

Other than that, this LGTM now, but let's see if there is agreement on that with other reviewers.

[vdsered]

Rebasing + updated tests with multiple uses in fmul

[vdsered]

Waiting for approval of other reviewers

[efriedma]

LGTM

[vdsered]

I ask someone to commit the patch if there are no any barriers left for that (Roman's agreement?). I don't have commit access.

llvm/include/llvm/Transforms/Utils/Local.h

 //===----------------------------------------------------------------------===//
 //  Transform predicates
 //
 
 /// Given an instruction, is it legal to set operand OpIdx to a non-constant
 /// value?
 bool canReplaceOperandWithVariable(const Instruction *I, unsigned OpIdx);
 
+/// It checks if this instruction is the only user of at least one of
+/// its operands.
+bool isOnlyUserOfAnyOperand(Instruction *I);
+
 //===----------------------------------------------------------------------===//
 //  Value helper functions
 //
 
 /// Invert the given true/false value, possibly reusing an existing copy.
 Value *invertCondition(Value *Condition);
 
 
 //===----------------------------------------------------------------------===//
 //  Assorted
 //
 
 /// If we can infer one attribute from another on the declaration of a
 /// function, explicitly materialize the maximal set in the IR.
 bool inferAttributesFromOthers(Function &F);
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_UTILS_LOCAL_H

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

if (I.hasNoNaNs() && I.hasNoSignedZeros() && Op0 == Op1 &&
       if (match(Op0, m_FDiv(m_Intrinsic<Intrinsic::sqrt>(m_Value(Y)),
                             m_Value(X)))) {
         Value *XX = Builder.CreateFMulFMF(X, X, &I);
         return BinaryOperator::CreateFDivFMF(Y, XX, &I);
       }
     }
 
     // exp(X) * exp(Y) -> exp(X + Y)
-    // Match as long as at least one of exp has only one use.
     if (match(Op0, m_Intrinsic<Intrinsic::exp>(m_Value(X))) &&
         match(Op1, m_Intrinsic<Intrinsic::exp>(m_Value(Y))) &&
-        (Op0->hasOneUse() || Op1->hasOneUse())) {
+        isOnlyUserOfAnyOperand(&I)) {
       Value *XY = Builder.CreateFAddFMF(X, Y, &I);
       Value *Exp = Builder.CreateUnaryIntrinsic(Intrinsic::exp, XY, &I);
       return replaceInstUsesWith(I, Exp);
     }
 
     // exp2(X) * exp2(Y) -> exp2(X + Y)
-    // Match as long as at least one of exp2 has only one use.
     if (match(Op0, m_Intrinsic<Intrinsic::exp2>(m_Value(X))) &&
         match(Op1, m_Intrinsic<Intrinsic::exp2>(m_Value(Y))) &&
-        (Op0->hasOneUse() || Op1->hasOneUse())) {
+        isOnlyUserOfAnyOperand(&I)) {
       Value *XY = Builder.CreateFAddFMF(X, Y, &I);
       Value *Exp2 = Builder.CreateUnaryIntrinsic(Intrinsic::exp2, XY, &I);
       return replaceInstUsesWith(I, Exp2);
     }
 
     // (X*Y) * X => (X*X) * Y where Y != X
     //  The purpose is two-fold:
     //   1) to form a power expression (of X).

llvm/lib/Transforms/Utils/Local.cpp

case Instruction::GetElementPtr:
     gep_type_iterator It = gep_type_begin(I);
     for (auto E = std::next(It, OpIdx); It != E; ++It)
       if (It.isStruct())
         return false;
     return true;
   }
 }
 
+bool llvm::isOnlyUserOfAnyOperand(Instruction *I) {
+  assert(I->getNumOperands() && "Must have at least one operand");
+  auto *FirstOp = I->getOperand(0);
+  auto NumOps = I->getNumOperands();
+  return llvm::any_of(I->operands(), [](Value *V) { return V->hasOneUse(); }) ||
+         (FirstOp->hasNUses(NumOps) && llvm::is_splat(I->operands()));
+}
+
 Value *llvm::invertCondition(Value *Condition) {
   // First: Check if it's a constant
   if (Constant *C = dyn_cast<Constant>(Condition))
     return ConstantExpr::getNot(C);
 
   // Second: If the condition is already inverted, return the original value
   Value *NotCondition;
   if (match(Condition, m_Not(m_Value(NotCondition))))

llvm/test/Transforms/InstCombine/fmul-exp.ll

 ; CHECK-NEXT:    ret double [[TMP1]]
 ;
   %tmp = call double @llvm.exp.f64(double %a)
   %tmp1 = call double @llvm.exp.f64(double %b)
   %mul = fmul reassoc double %tmp, %tmp1
   ret double %mul
 }
 
+define double @exp_a(double %a) {
+; CHECK-LABEL: @exp_a(
+; CHECK-NEXT:    [[TMP1:%.*]] = fadd reassoc double [[A:%.*]], [[A]]
+; CHECK-NEXT:    [[TMP2:%.*]] = call reassoc double @llvm.exp.f64(double [[TMP1]])
+; CHECK-NEXT:    ret double [[TMP2]]
+;
+  %tmp = call double @llvm.exp.f64(double %a)
+  %mul = fmul reassoc double %tmp, %tmp
+  ret double %mul
+}
+
 ; exp(a) * exp(b) * exp(c) * exp(d) => exp(a+b+c+d) with reassoc
 define double @exp_a_exp_b_exp_c_exp_d_fast(double %a, double %b, double %c, double %d) {
 ; CHECK-LABEL: @exp_a_exp_b_exp_c_exp_d_fast(
 ; CHECK-NEXT:    [[TMP:%.*]] = fadd reassoc double [[A:%.*]], [[B:%.*]]
 ; CHECK-NEXT:    [[TMP1:%.*]] = fadd reassoc double [[TMP]], [[C:%.*]]
 ; CHECK-NEXT:    [[TMP2:%.*]] = fadd reassoc double [[TMP1]], [[D:%.*]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = call reassoc double @llvm.exp.f64(double [[TMP2]])
 ; CHECK-NEXT:    ret double [[TMP3]]

llvm/test/Transforms/InstCombine/fmul-exp2.ll

 ;
   %tmp = call double @llvm.exp2.f64(double %a)
   %tmp1 = call double @llvm.exp2.f64(double %b)
   %mul = fmul reassoc double %tmp, %tmp1
   call void @use(double %tmp1)
   ret double %mul
 }
 
+define double @exp2_a_multiple_uses_in_one_user(double %a) {
+; CHECK-LABEL: @exp2_a_multiple_uses_in_one_user(
+; CHECK-NEXT:    [[TMP1:%.*]] = fadd reassoc double [[A:%.*]], [[A]]
+; CHECK-NEXT:    [[TMP2:%.*]] = call reassoc double @llvm.exp2.f64(double [[TMP1]])
+; CHECK-NEXT:    ret double [[TMP2]]
+;
+  %tmp = call double @llvm.exp2.f64(double %a)
+  %mul = fmul reassoc double %tmp, %tmp
+  ret double %mul
+}
+
 ; exp2(a) * exp2(b) reassoc, both with multiple uses
 define double @exp2_a_exp2_b_multiple_uses_both(double %a, double %b) {
 ; CHECK-LABEL: @exp2_a_exp2_b_multiple_uses_both(
 ; CHECK-NEXT:    [[TMP:%.*]] = call double @llvm.exp2.f64(double [[A:%.*]])
 ; CHECK-NEXT:    [[TMP1:%.*]] = call double @llvm.exp2.f64(double [[B:%.*]])
 ; CHECK-NEXT:    [[MUL:%.*]] = fmul reassoc double [[TMP]], [[TMP1]]
 ; CHECK-NEXT:    call void @use(double [[TMP]])
 ; CHECK-NEXT:    call void @use(double [[TMP1]])