From a4ca22506c96cd8508bcfb93a26941f19c7179ef Mon Sep 17 00:00:00 2001
From: Alex MacLean <amaclean@nvidia.com>
Date: Sat, 22 Jun 2024 09:50:21 -0700
Subject: [InstCombine] (uitofp bool X) * Y --> X ? Y : 0 (#96216)

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
---
 .../InstCombine/InstCombineMulDivRem.cpp           | 18 +++++++++++
 llvm/test/Transforms/InstCombine/fmul-bool.ll      | 35 ++++++++++++++++++++++
 2 files changed, 53 insertions(+)
 create mode 100644 llvm/test/Transforms/InstCombine/fmul-bool.ll

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
index 8fcb354..c3f1c12 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
@@ -879,6 +879,24 @@ 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)) {
+      auto *SI = SelectInst::Create(X, Op1, ConstantFP::get(I.getType(), 0.0));
+      SI->copyFastMathFlags(I.getFastMathFlags());
+      return SI;
+    }
+    if (match(Op1, m_UIToFP(m_Value(X))) &&
+        X->getType()->isIntOrIntVectorTy(1)) {
+      auto *SI = SelectInst::Create(X, Op0, ConstantFP::get(I.getType(), 0.0));
+      SI->copyFastMathFlags(I.getFastMathFlags());
+      return SI;
+    }
+  }
+
   // (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 0000000..263ca53
--- /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 nnan nsz 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 nnan nsz <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 nnan nsz <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
+}
-- 
cgit v1.1