[llvm] [InstCombine] Add check for flag propagation in `foldSelectIntoOp` (PR #173735)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Yunbo Ni (cardigan1008)

<details>
<summary>Changes</summary>

Fixes https://github.com/llvm/llvm-project/pull/162003#issuecomment-3693943568. 

The current flag propagation assumes that if a select has both `ninf` and `nnan`, then the operands of the folded operation must be finite. While this assumption holds for `fadd`, `fsub`, and `fmul`, it does not hold for `fdiv`. 

For example, assume we have: 

```
A = 1.0, B = +Inf
A / B = 0.0  (finite, non-NaN)
```

The current transform would turn `fdiv A, B; select ninf nnan cond, A/B, A;` into `A / (select ninf nnan cond, B, 1.0)`. If `cond` is true, the inner select returns `B = +Inf`, and due to the propagated `ninf`, this becomes poison. 

This patch add check for operators before flag propagation to avoid `fdiv` cases. 

Alive2: https://alive2.llvm.org/ce/z/o0MJmS

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


2 Files Affected:

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


``````````diff
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index f52bac5e600cb..67d1845832725 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -555,8 +555,15 @@ Instruction *InstCombinerImpl::foldSelectIntoOp(SelectInst &SI, Value *TrueVal,
       // 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.
+      // Note: This property holds for fadd, fsub, and fmul, but does not
+      // hold for fdiv (e.g. A / Inf == 0.0).
+      bool CanInferFiniteOperandsFromResult =
+          TVI->getOpcode() == Instruction::FAdd ||
+          TVI->getOpcode() == Instruction::FSub ||
+          TVI->getOpcode() == Instruction::FMul;
       NewSelFMF.setNoInfs(TVI->hasNoInfs() ||
-                          (NewSelFMF.noInfs() && NewSelFMF.noNaNs()));
+                          (CanInferFiniteOperandsFromResult &&
+                           NewSelFMF.noInfs() && NewSelFMF.noNaNs()));
       cast<Instruction>(NewSel)->setFastMathFlags(NewSelFMF);
     }
     NewSel->takeName(TVI);
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 c14dd469f1a6e..83fa28a406f75 100644
--- a/llvm/test/Transforms/InstCombine/select-binop-foldable-floating-point.ll
+++ b/llvm/test/Transforms/InstCombine/select-binop-foldable-floating-point.ll
@@ -409,3 +409,14 @@ define float @select_nnan_fdiv_invalid(i1 %cond, float %A, float %B) {
   %D = select nnan i1 %cond, float %C, float %A
   ret float %D
 }
+
+define float @select_fpclass_fdiv_nnan_ninf(i1 %cond, float %A, float %B) {
+; CHECK-LABEL: @select_fpclass_fdiv_nnan_ninf(
+; CHECK-NEXT:    [[C:%.*]] = select nnan i1 [[COND:%.*]], float [[B:%.*]], float 1.000000e+00
+; CHECK-NEXT:    [[D:%.*]] = fdiv float [[A:%.*]], [[C]]
+; CHECK-NEXT:    ret float [[D]]
+;
+  %C = fdiv float %A, %B
+  %D = select ninf nnan i1 %cond, float %C, float %A
+  ret float %D
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/173735