[InstCombine] Fix flag propagation in foldSelectIntoOp

[dtcxzyw]

Consider the following transform:

C = binop float A, nnan OOp
D = select ninf, i1 cond, float C, float A
->
E = select ninf, i1 cond, float OOp, float Identity
F = binop float A, E

We cannot propagate ninf from the original select, because OOp may be inf, and the flag only guarantees that FalseVal (op OOp) is never infinity.
Examples: -inf + +inf = NaN, -inf - -inf = NaN, 0 * inf = NaN
Specifically, if the original select has both ninf and nnan, we can safely propagate the flag.

Alive2:

• fadd: https://alive2.llvm.org/ce/z/TWfktv
• fsub: https://alive2.llvm.org/ce/z/RAsjJb
• fmul: https://alive2.llvm.org/ce/z/8eg4ND

Closes #161634.

[dtcxzyw]

@zyw-bot mfuzz

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Yingwei Zheng (dtcxzyw)

Changes

Alive2: https://alive2.llvm.org/ce/z/NwVqpL
Closes #161634.

---

Full diff: https://github.com/llvm/llvm-project/pull/162003.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp (+4-1)
• (modified) llvm/test/Transforms/InstCombine/select-binop-foldable-floating-point.ll (+44)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 8f60e506e8a33..fb4aaa6992702 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -544,8 +544,11 @@ Instruction *InstCombinerImpl::foldSelectIntoOp(SelectInst &SI, Value *TrueVal,
 
     Value *NewSel = Builder.CreateSelect(SI.getCondition(), Swapped ? C : OOp,
                                          Swapped ? OOp : C, "", &SI);
-    if (isa<FPMathOperator>(&SI))
+    if (isa<FPMathOperator>(&SI)) {
       cast<Instruction>(NewSel)->setFastMathFlags(FMF);
+      if (!TVI->hasNoInfs() && !FMF.noNaNs())
+        cast<Instruction>(NewSel)->setHasNoInfs(false);
+    }
     NewSel->takeName(TVI);
     BinaryOperator *BO =
         BinaryOperator::Create(TVI->getOpcode(), FalseVal, NewSel);
diff --git a/llvm/test/Transforms/InstCombine/select-binop-foldable-floating-point.ll b/llvm/test/Transforms/InstCombine/select-binop-foldable-floating-point.ll
index 253bc9e784c2f..8b5e02acfb02e 100644
--- a/llvm/test/Transforms/InstCombine/select-binop-foldable-floating-point.ll
+++ b/llvm/test/Transforms/InstCombine/select-binop-foldable-floating-point.ll
@@ -23,6 +23,50 @@ define float @select_fpclass_fadd(i1 %cond, float nofpclass(nan) %A, float %B) {
   ret float %D
 }
 
+define float @select_fpclass_fadd_ninf1(i1 %cond, float nofpclass(nan) %A, float %B) {
+; CHECK-LABEL: @select_fpclass_fadd_ninf1(
+; CHECK-NEXT:    [[C:%.*]] = select i1 [[COND:%.*]], float [[B:%.*]], float -0.000000e+00
+; CHECK-NEXT:    [[D:%.*]] = fadd float [[A:%.*]], [[C]]
+; CHECK-NEXT:    ret float [[D]]
+;
+  %C = fadd ninf float %A, %B
+  %D = select i1 %cond, float %C, float %A
+  ret float %D
+}
+
+define float @select_fpclass_fadd_ninf2(i1 %cond, float nofpclass(nan) %A, float %B) {
+; CHECK-LABEL: @select_fpclass_fadd_ninf2(
+; CHECK-NEXT:    [[C:%.*]] = select i1 [[COND:%.*]], float [[B:%.*]], float -0.000000e+00
+; CHECK-NEXT:    [[D:%.*]] = fadd float [[A:%.*]], [[C]]
+; CHECK-NEXT:    ret float [[D]]
+;
+  %C = fadd float %A, %B
+  %D = select ninf i1 %cond, float %C, float %A
+  ret float %D
+}
+
+define float @select_fpclass_fadd_ninf3(i1 %cond, float nofpclass(nan) %A, float %B) {
+; CHECK-LABEL: @select_fpclass_fadd_ninf3(
+; CHECK-NEXT:    [[C:%.*]] = select ninf i1 [[COND:%.*]], float [[B:%.*]], float -0.000000e+00
+; CHECK-NEXT:    [[D:%.*]] = fadd ninf float [[A:%.*]], [[C]]
+; CHECK-NEXT:    ret float [[D]]
+;
+  %C = fadd ninf float %A, %B
+  %D = select ninf i1 %cond, float %C, float %A
+  ret float %D
+}
+
+define float @select_fpclass_fadd_nnan_ninf(i1 %cond, float nofpclass(nan) %A, float %B) {
+; CHECK-LABEL: @select_fpclass_fadd_nnan_ninf(
+; CHECK-NEXT:    [[C:%.*]] = select nnan ninf i1 [[COND:%.*]], float [[B:%.*]], float -0.000000e+00
+; CHECK-NEXT:    [[D:%.*]] = fadd float [[A:%.*]], [[C]]
+; CHECK-NEXT:    ret float [[D]]
+;
+  %C = fadd float %A, %B
+  %D = select nnan ninf i1 %cond, float %C, float %A
+  ret float %D
+}
+
 define float @select_nnan_fadd(i1 %cond, float %A, float %B) {
 ; CHECK-LABEL: @select_nnan_fadd(
 ; CHECK-NEXT:    [[C:%.*]] = select nnan i1 [[COND:%.*]], float [[B:%.*]], float -0.000000e+00

[arsenm]

Can you add a TODO that you can preserve this for operators which cannot introduce infs

[dtcxzyw]

This fold works for fadd, fsub, and fmul. Unfortunately, there are no exceptions. All of these patterns have counterexamples to block the flag propagation. I have updated comments and proofs.