[InstCombine] Missed optimization in math expression: squashing sqrt functions

[Quolyk]

Summary

This patch enables folding under -ffast-math flag sqrt(a) * sqrt(b) -> sqrt(a*b)

[spatel]

lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

626

Note that we have a fast-math-flag specifically for reassociation now (as well as a FMF for approximation of math functions). I'm not sure what combination would be needed for this transform. It's independent of this patch, but something to think about.

731

What if sqrt(a) or sqrt(b) has another use besides the multiply? I don't think we want to do the transform in that case. Either way, please add a test for that kind of example.

733–739

You can use:

match(Op0, m_Intrinsic<Intrinsic::sqrt>( ...

to simplify this code.

test/Transforms/InstCombine/fmul.ll

185–199 ↗

(On Diff #127243)

Please use the script at utils/update_test_checks.py to auto-generate the checks for any tests you add. The checks will be more exact and have regex for the variable names, so they'll be more robust.

[davide]

test/Transforms/InstCombine/fmul.ll

195–199 ↗

(On Diff #127243)

Please add a negative test showing that this combine doesn't fire unless the conditions are met (allows reassociation/fast-math, etc..)

[hfinkel]

lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

731

I agree, we should have ->hasOneUse() checks here.

740

Needs

BuilderTy::FastMathFlagGuard Guard(Builder);

above this.

[spatel]

LGTM

lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

737

Just a few small bits here:

1. You don’t have to initialize to nullptr; the captured variables aren’t touched if the match fails.
2. There is a matcher for m_OneUse (though it may not save any space here).
3. I prefer to keep the names in the code comments and the actual code variables the same, so Opnd0 would just be ‘A’ etc.

[davide]

LGTM.

[Quolyk]

@davide please accept the revision, if ok. I'll commit.

lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

737

1. I know, Xcode showed warning at that line, so I decided to make it that way.
2. That is useful, didn't know that.
3. I'll take into account, thank you.

[spatel]

This isn't committed yet. Do you have access?

[Quolyk]

In D41322#964867, @spatel wrote:

This isn't committed yet. Do you have access?

I have been granted a few days ago, just need to figure out how to make it.

lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

Instruction *InstCombiner::visitFMul(BinaryOperator &I) {
 
   if (isa<Constant>(Op0))
     std::swap(Op0, Op1);
 
   if (Value *V = SimplifyFMulInst(Op0, Op1, I.getFastMathFlags(),
                                   SQ.getWithInstruction(&I)))
     return replaceInstUsesWith(I, V);
 
   bool AllowReassociate = I.isFast();
 
   // Simplify mul instructions with a constant RHS.
   if (isa<Constant>(Op1)) {
     if (Instruction *FoldedMul = foldOpWithConstantIntoOperand(I))
       return FoldedMul;
 
     // (fmul X, -1.0) --> (fsub -0.0, X)
     if (match(Op1, m_SpecificFP(-1.0))) {
if (OpX && OpY) {
       Log2->setArgOperand(0, OpY);
       Value *FMulVal = Builder.CreateFMul(OpX, Log2);
       Value *FSub = Builder.CreateFSub(FMulVal, OpX);
       FSub->takeName(&I);
       return replaceInstUsesWith(I, FSub);
     }
   }
 
+  // sqrt(a) * sqrt(b) -> sqrt(a * b)
+  if (AllowReassociate &&
+      Op0->hasOneUse() && Op1->hasOneUse()) {
+    Value *Opnd0 = nullptr;
+    Value *Opnd1 = nullptr;
+    if (match(Op0, m_Intrinsic<Intrinsic::sqrt>(m_Value(Opnd0))) &&
+        match(Op1, m_Intrinsic<Intrinsic::sqrt>(m_Value(Opnd1)))) {
+      BuilderTy::FastMathFlagGuard Guard(Builder);
+      Builder.setFastMathFlags(I.getFastMathFlags());
+      Value *FMulVal = Builder.CreateFMul(Opnd0, Opnd1);
+      Value *Sqrt = Intrinsic::getDeclaration(I.getModule(), 
+                                              Intrinsic::sqrt, I.getType());
+      Value *SqrtCall = Builder.CreateCall(Sqrt, FMulVal);
+      return replaceInstUsesWith(I, SqrtCall);
+    }
+  }
+
   // Handle symmetric situation in a 2-iteration loop
   Value *Opnd0 = Op0;
   Value *Opnd1 = Op1;
   for (int i = 0; i < 2; i++) {
     bool IgnoreZeroSign = I.hasNoSignedZeros();
     if (BinaryOperator::isFNeg(Opnd0, IgnoreZeroSign)) {
       BuilderTy::FastMathFlagGuard Guard(Builder);
       Builder.setFastMathFlags(I.getFastMathFlags());

test/Transforms/InstCombine/fmul-sqrt.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -instcombine < %s | FileCheck %s
+
+declare double @llvm.sqrt.f64(double) nounwind readnone speculatable
+declare void @use(double)
+
+; sqrt(a) * sqrt(b) no math flags
+define double @sqrt_a_sqrt_b(double %a, double %b) {
+; CHECK-LABEL: @sqrt_a_sqrt_b(
+; CHECK-NEXT:    [[TMP1:%.*]] = call double @llvm.sqrt.f64(double [[A:%.*]])
+; CHECK-NEXT:    [[TMP2:%.*]] = call double @llvm.sqrt.f64(double [[B:%.*]])
+; CHECK-NEXT:    [[MUL:%.*]] = fmul double [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    ret double [[MUL]]
+;
+  %1 = call double @llvm.sqrt.f64(double %a)
+  %2 = call double @llvm.sqrt.f64(double %b)
+  %mul = fmul double %1, %2
+  ret double %mul
+}
+
+; sqrt(a) * sqrt(b) fast-math, multiple uses
+define double @sqrt_a_sqrt_b_multiple_uses(double %a, double %b) {
+; CHECK-LABEL: @sqrt_a_sqrt_b_multiple_uses(
+; CHECK-NEXT:    [[TMP1:%.*]] = call fast double @llvm.sqrt.f64(double [[A:%.*]])
+; CHECK-NEXT:    [[TMP2:%.*]] = call fast double @llvm.sqrt.f64(double [[B:%.*]])
+; CHECK-NEXT:    [[MUL:%.*]] = fmul fast double [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    call void @use(double [[TMP2]])
+; CHECK-NEXT:    ret double [[MUL]]
+;
+  %1 = call fast double @llvm.sqrt.f64(double %a)
+  %2 = call fast double @llvm.sqrt.f64(double %b)
+  %mul = fmul fast double %1, %2
+  call void @use(double %2)
+  ret double %mul
+}
+
+; sqrt(a) * sqrt(b) => sqrt(a*b) with fast-math
+define double @sqrt_a_sqrt_b_fast(double %a, double %b) {
+; CHECK-LABEL: @sqrt_a_sqrt_b_fast(
+; CHECK-NEXT:    [[TMP1:%.*]] = fmul fast double [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = call fast double @llvm.sqrt.f64(double [[TMP1]])
+; CHECK-NEXT:    ret double [[TMP2]]
+;
+  %1 = call fast double @llvm.sqrt.f64(double %a)
+  %2 = call fast double @llvm.sqrt.f64(double %b)
+  %mul = fmul fast double %1, %2
+  ret double %mul
+}
+
+; sqrt(a) * sqrt(b) * sqrt(c) * sqrt(d) => sqrt(a*b*c+d) with fast-math
+define double @sqrt_a_sqrt_b_sqrt_c_sqrt_d_fast(double %a, double %b, double %c, double %d) {
+; CHECK-LABEL: @sqrt_a_sqrt_b_sqrt_c_sqrt_d_fast(
+; CHECK-NEXT:    [[TMP1:%.*]] = fmul fast double [[A:%.*]], [[B:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = fmul fast double [[TMP1]], [[C:%.*]]
+; CHECK-NEXT:    [[TMP3:%.*]] = fmul fast double [[TMP2]], [[D:%.*]]
+; CHECK-NEXT:    [[TMP4:%.*]] = call fast double @llvm.sqrt.f64(double [[TMP3]])
+; CHECK-NEXT:    ret double [[TMP4]]
+;
+  %1 = call fast double @llvm.sqrt.f64(double %a)
+  %2 = call fast double @llvm.sqrt.f64(double %b)
+  %mul = fmul fast double %1, %2
+  %3 = call fast double @llvm.sqrt.f64(double %c)
+  %mul1 = fmul fast double %mul, %3
+  %4 = call fast double @llvm.sqrt.f64(double %d)
+  %mul2 = fmul fast double %mul1, %4
+  ret double %mul2
+}