[InstCombine] (uitofp bool X) * Y --> X ? Y : 0

[AlexMaclean]

Fold mul (uitofp i1 X), Y to select i1 X, Y, 0.0 when the mul is nnan and nsz

Proof: https://alive2.llvm.org/ce/z/_stiPm

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Alex MacLean (AlexMaclean)

Changes

Fold mul (uitofp i1 X), Y to select i1 X, Y, 0.0 when the mul is nnan and nsz

Proof: https://alive2.llvm.org/ce/z/zvSXq7

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp (+10)
• (added) llvm/test/Transforms/InstCombine/fmul-bool.ll (+35)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
index 8fcb3544f682a..dc2c92a84c650 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
@@ -879,6 +879,16 @@ Instruction *InstCombinerImpl::visitFMul(BinaryOperator &I) {
     if (Constant *NegC = ConstantFoldUnaryOpOperand(Instruction::FNeg, C, DL))
       return BinaryOperator::CreateFMulFMF(X, NegC, &I);
 
+  if (I.hasNoNaNs() && I.hasNoSignedZeros()) {
+    // (uitofp bool X) * Y --> X ? Y : 0
+    // Y * (uitofp bool X) --> X ? Y : 0
+    // Note INF * 0 is NaN.
+    if (match(Op0, m_UIToFP(m_Value(X))) && X->getType()->isIntOrIntVectorTy(1))
+      return SelectInst::Create(X, Op1, ConstantFP::get(I.getType(), 0.0));
+    if (match(Op1, m_UIToFP(m_Value(X))) && X->getType()->isIntOrIntVectorTy(1))
+      return SelectInst::Create(X, Op0, ConstantFP::get(I.getType(), 0.0));
+  }
+
   // (select A, B, C) * (select A, D, E) --> select A, (B*D), (C*E)
   if (Value *V = SimplifySelectsFeedingBinaryOp(I, Op0, Op1))
     return replaceInstUsesWith(I, V);
diff --git a/llvm/test/Transforms/InstCombine/fmul-bool.ll b/llvm/test/Transforms/InstCombine/fmul-bool.ll
new file mode 100644
index 0000000000000..73bb4f39b106b
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/fmul-bool.ll
@@ -0,0 +1,35 @@
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+; X * Y (when Y is a boolean) --> Y ? X : 0
+
+define float @fmul_bool(float %x, i1 %y) {
+; CHECK-LABEL: @fmul_bool(
+; CHECK-NEXT:    [[M:%.*]] = select i1 [[Y:%.*]], float [[X:%.*]], float 0.000000e+00
+; CHECK-NEXT:    ret float [[M]]
+;
+  %z = uitofp i1 %y to float
+  %m = fmul nnan nsz float %z, %x
+  ret float %m
+}
+
+define <2 x float> @fmul_bool_vec(<2 x float> %x, <2 x i1> %y) {
+; CHECK-LABEL: @fmul_bool_vec(
+; CHECK-NEXT:    [[M:%.*]] = select <2 x i1> [[Y:%.*]], <2 x float> [[X:%.*]], <2 x float> zeroinitializer
+; CHECK-NEXT:    ret <2 x float> [[M]]
+;
+  %z = uitofp <2 x i1> %y to <2 x float>
+  %m = fmul nnan nsz <2 x float> %z, %x
+  ret <2 x float> %m
+}
+
+define <2 x float> @fmul_bool_vec_commute(<2 x float> %px, <2 x i1> %y) {
+; CHECK-LABEL: @fmul_bool_vec_commute(
+; CHECK-NEXT:    [[X:%.*]] = fmul nnan nsz <2 x float> [[PX:%.*]], [[PX]]
+; CHECK-NEXT:    [[M:%.*]] = select <2 x i1> [[Y:%.*]], <2 x float> [[X]], <2 x float> zeroinitializer
+; CHECK-NEXT:    ret <2 x float> [[M]]
+;
+  %x = fmul nnan nsz <2 x float> %px, %px  ; thwart complexity-based canonicalization
+  %z = uitofp <2 x i1> %y to <2 x float>
+  %m = fmul nnan nsz <2 x float> %x, %z
+  ret <2 x float> %m
+}

[dtcxzyw]

LGTM.