[InstSimplify] loosen FMF for sqrt(X) * sqrt(X) --> X

[spatel]

Summary

If my fast-math understanding is correct, 'reassoc' alone is not enough because that doesn't give us the freedom to get the negative number cases wrong.

I considered that we might not even need 'reassoc' here, but if we eliminate the sqrt calls, then we may differ in the last bit by eliminating the rounding that occurs in those calls.

[spatel]

Ping.

[scanon]

IIRC Intrinsic::sqrt is undef for negative inputs (unlike the sqrt libcall), so we don't need FMF.noNaNs to license this transformation.

[arsenm]

In D43765#1027289, @scanon wrote:

IIRC Intrinsic::sqrt is undef for negative inputs (unlike the sqrt libcall), so we don't need FMF.noNaNs to license this transformation.

Although IMO if we're fixing all of the math in LLVM I would like to fix this. At least do something like ctlz/cttz for whether negative is undef

[spatel]

In D43765#1027302, @arsenm wrote:

In D43765#1027289, @scanon wrote:

IIRC Intrinsic::sqrt is undef for negative inputs (unlike the sqrt libcall), so we don't need FMF.noNaNs to license this transformation.

Although IMO if we're fixing all of the math in LLVM I would like to fix this. At least do something like ctlz/cttz for whether negative is undef

Sqrt was fixed:
https://reviews.llvm.org/D28797

[spatel]

In D43765#1027330, @spatel wrote:

In D43765#1027302, @arsenm wrote:

In D43765#1027289, @scanon wrote:

IIRC Intrinsic::sqrt is undef for negative inputs (unlike the sqrt libcall), so we don't need FMF.noNaNs to license this transformation.

Although IMO if we're fixing all of the math in LLVM I would like to fix this. At least do something like ctlz/cttz for whether negative is undef

Sqrt was fixed:
https://reviews.llvm.org/D28797

Also, note that we convert sqrt and other libcalls to the LLVM intrinsics when possible in clang:
D39204 (and follow-up commits listed there)
...so we don't have to muddy the code looking for sqrt libcall patterns too. If errno can be set by a libcall, then no amount of FMF should override that; if errno can't be set by the libcall, it should have been converted to the LLVM intrinsic.

[spatel]

Ping * 2.

[efriedma]

test/Transforms/InstSimplify/fast-math.ll

219 ↗

(On Diff #135904)

Maybe add a test that we don't do this without the relevant fast-math flags?

[spatel]

Patch updated:
Include negative tests to show that we're only doing the transform when all 3 of the required FMF are present.

[wristow]

LGTM

llvm/trunk/lib/Analysis/InstructionSimplify.cpp

static Value *SimplifyFMulInst(Value *Op0, Value *Op1, FastMathFlags FMF,
   // fmul X, 1.0 ==> X
   if (match(Op1, m_FPOne()))
     return Op0;
 
   // fmul nnan nsz X, 0 ==> 0
   if (FMF.noNaNs() && FMF.noSignedZeros() && match(Op1, m_AnyZeroFP()))
     return ConstantFP::getNullValue(Op0->getType());
 
-  // sqrt(X) * sqrt(X) --> X
+  // sqrt(X) * sqrt(X) --> X, if we can:
+  // 1. Remove the intermediate rounding (reassociate).
+  // 2. Ignore non-zero negative numbers because sqrt would produce NAN.
+  // 3. Ignore -0.0 because sqrt(-0.0) == -0.0, but -0.0 * -0.0 == 0.0.
   Value *X;
-  if (FMF.isFast() && Op0 == Op1 &&
-      match(Op0, m_Intrinsic<Intrinsic::sqrt>(m_Value(X))))
+  if (Op0 == Op1 && match(Op0, m_Intrinsic<Intrinsic::sqrt>(m_Value(X))) &&
+      FMF.allowReassoc() && FMF.noNaNs() && FMF.noSignedZeros())
     return X;
 
   return nullptr;
 }
 
 Value *llvm::SimplifyFAddInst(Value *Op0, Value *Op1, FastMathFlags FMF,
                               const SimplifyQuery &Q) {
   return ::SimplifyFAddInst(Op0, Op1, FMF, Q, RecursionLimit);

llvm/trunk/test/Transforms/InstSimplify/fast-math.ll

 ; CHECK-NEXT:    ret float -1.000000e+00
 ;
   %neg = fsub float 0.000000e+00, %f
   %div = fdiv nnan float %f, %neg
   ret float %div
 }
 
 ; PR21126: http://llvm.org/bugs/show_bug.cgi?id=21126
-; With unsafe/fast math, sqrt(X) * sqrt(X) is just X.
+; With loose math, sqrt(X) * sqrt(X) is just X.
 
 declare double @llvm.sqrt.f64(double)
 
 define double @sqrt_squared(double %f) {
 ; CHECK-LABEL: @sqrt_squared(
 ; CHECK-NEXT:    ret double [[F:%.*]]
 ;
   %sqrt = call double @llvm.sqrt.f64(double %f)
-  %mul = fmul fast double %sqrt, %sqrt
+  %mul = fmul reassoc nnan nsz double %sqrt, %sqrt
+  ret double %mul
+}
+
+; Negative tests for the above transform: we need all 3 of those flags.
+
+define double @sqrt_squared_not_fast_enough1(double %f) {
+; CHECK-LABEL: @sqrt_squared_not_fast_enough1(
+; CHECK-NEXT:    [[SQRT:%.*]] = call double @llvm.sqrt.f64(double [[F:%.*]])
+; CHECK-NEXT:    [[MUL:%.*]] = fmul nnan nsz double [[SQRT]], [[SQRT]]
+; CHECK-NEXT:    ret double [[MUL]]
+;
+  %sqrt = call double @llvm.sqrt.f64(double %f)
+  %mul = fmul nnan nsz double %sqrt, %sqrt
+  ret double %mul
+}
+
+define double @sqrt_squared_not_fast_enough2(double %f) {
+; CHECK-LABEL: @sqrt_squared_not_fast_enough2(
+; CHECK-NEXT:    [[SQRT:%.*]] = call double @llvm.sqrt.f64(double [[F:%.*]])
+; CHECK-NEXT:    [[MUL:%.*]] = fmul reassoc nnan double [[SQRT]], [[SQRT]]
+; CHECK-NEXT:    ret double [[MUL]]
+;
+  %sqrt = call double @llvm.sqrt.f64(double %f)
+  %mul = fmul reassoc nnan double %sqrt, %sqrt
+  ret double %mul
+}
+
+define double @sqrt_squared_not_fast_enough3(double %f) {
+; CHECK-LABEL: @sqrt_squared_not_fast_enough3(
+; CHECK-NEXT:    [[SQRT:%.*]] = call double @llvm.sqrt.f64(double [[F:%.*]])
+; CHECK-NEXT:    [[MUL:%.*]] = fmul reassoc nsz double [[SQRT]], [[SQRT]]
+; CHECK-NEXT:    ret double [[MUL]]
+;
+  %sqrt = call double @llvm.sqrt.f64(double %f)
+  %mul = fmul reassoc nsz double %sqrt, %sqrt
   ret double %mul
 }