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

[Yunbo Ni via llvm-commits]

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

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

>From 84761ebf8640d944d1c5dcee13fe432760067c76 Mon Sep 17 00:00:00 2001
From: cardigan1008 <ybni at cse.cuhk.edu.hk>
Date: Sun, 28 Dec 2025 00:46:30 +0800
Subject: [PATCH 1/2] [InstCombine][NFC] Add pre-commit tests for fdiv.

---
 .../select-binop-foldable-floating-point.ll           | 11 +++++++++++
 1 file changed, 11 insertions(+)

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..136c11f1d4370 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 ninf 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
+}

>From 7d1de3f98fecd068e29a2b07c7fad30f923d776f Mon Sep 17 00:00:00 2001
From: cardigan1008 <ybni at cse.cuhk.edu.hk>
Date: Sun, 28 Dec 2025 00:48:16 +0800
Subject: [PATCH 2/2] [InstCombine] Add check for flag progpgation in
 `foldSelectIntoOp`.

---
 llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp    | 9 ++++++++-
 .../InstCombine/select-binop-foldable-floating-point.ll  | 2 +-
 2 files changed, 9 insertions(+), 2 deletions(-)

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 136c11f1d4370..83fa28a406f75 100644
--- a/llvm/test/Transforms/InstCombine/select-binop-foldable-floating-point.ll
+++ b/llvm/test/Transforms/InstCombine/select-binop-foldable-floating-point.ll
@@ -412,7 +412,7 @@ define float @select_nnan_fdiv_invalid(i1 %cond, float %A, float %B) {
 
 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 ninf i1 [[COND:%.*]], float [[B:%.*]], float 1.000000e+00
+; 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]]
 ;