[llvm] [InstCombine] (uitofp bool X) * Y --> X ? Y : 0 (PR #96216)

[Alex MacLean via llvm-commits]

https://github.com/AlexMaclean created https://github.com/llvm/llvm-project/pull/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/zvSXq7

>From ff28bc6d54e59c42066ccee21c1e8f8a89c9ef05 Mon Sep 17 00:00:00 2001
From: Alex MacLean <amaclean at nvidia.com>
Date: Thu, 13 Jun 2024 23:40:59 +0000
Subject: [PATCH] [InstCombine] (sitofp bool X) * Y --> X ? Y : 0

---
 .../InstCombine/InstCombineMulDivRem.cpp      | 10 ++++++
 llvm/test/Transforms/InstCombine/fmul-bool.ll | 35 +++++++++++++++++++
 2 files changed, 45 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 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
+}