[llvm] [InstCombine] fold `select (trunc nuw X to i1), X, Y` to `select (trunc nuw X to i1), 1, Y` (PR #105914)

[via llvm-commits]

https://github.com/c8ef created https://github.com/llvm/llvm-project/pull/105914

This patch adds a fold from `select (trunc nuw X to i1), X, Y` to `select (trunc nuw X to i1), 1, Y`. Considering the semantics of the nuw flag, if there is no poison value in this expression, this code assumes that X can only be 0 or 1.

Partially resolved: #96765

>From 5b7752afa5aec5299e986ab9bca26007fbaea810 Mon Sep 17 00:00:00 2001
From: c8ef <c8ef at outlook.com>
Date: Sat, 24 Aug 2024 09:47:29 +0800
Subject: [PATCH] fold select (trunc nuw X to i1), X, Y

---
 .../InstCombine/InstCombineSelect.cpp         | 10 ++++++
 .../InstCombine/fold-select-trunc-if-nuw.ll   | 35 +++++++++++++++++++
 2 files changed, 45 insertions(+)
 create mode 100644 llvm/test/Transforms/InstCombine/fold-select-trunc-if-nuw.ll

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 18ffc209f259e0..49582004d9a4f6 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -4201,5 +4201,15 @@ Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
     }
   }
 
+  // select (trunc nuw X to i1), X, Y --> select (trunc nuw X to i1), 1, Y
+  if (auto *TI = dyn_cast<TruncInst>(CondVal)) {
+    Value *Trunc;
+    if (TI->hasNoUnsignedWrap() && match(CondVal, m_Trunc(m_Value(Trunc))) &&
+        match(TrueVal, m_Specific(Trunc))) {
+      return SelectInst::Create(
+          CondVal, ConstantInt::get(TrueVal->getType(), 1), FalseVal);
+    }
+  }
+
   return nullptr;
 }
diff --git a/llvm/test/Transforms/InstCombine/fold-select-trunc-if-nuw.ll b/llvm/test/Transforms/InstCombine/fold-select-trunc-if-nuw.ll
new file mode 100644
index 00000000000000..09fb2d77f9ef56
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/fold-select-trunc-if-nuw.ll
@@ -0,0 +1,35 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+define i8 @fold_select_trunc_if_nuw(i8 %x, i8 %y) {
+; CHECK-LABEL: @fold_select_trunc_if_nuw(
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc nuw i8 [[X:%.*]] to i1
+; CHECK-NEXT:    [[RET:%.*]] = select i1 [[TRUNC]], i8 1, i8 [[Y:%.*]]
+; CHECK-NEXT:    ret i8 [[RET]]
+;
+  %trunc = trunc nuw i8 %x to i1
+  %ret = select i1 %trunc, i8 %x, i8 %y
+  ret i8 %ret
+}
+
+define i8 @fold_select_trunc_if_nuw_negative(i8 %x, i8 %y) {
+; CHECK-LABEL: @fold_select_trunc_if_nuw_negative(
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc nsw i8 [[X:%.*]] to i1
+; CHECK-NEXT:    [[RET:%.*]] = select i1 [[TRUNC]], i8 [[X]], i8 [[Y:%.*]]
+; CHECK-NEXT:    ret i8 [[RET]]
+;
+  %trunc = trunc nsw i8 %x to i1
+  %ret = select i1 %trunc, i8 %x, i8 %y
+  ret i8 %ret
+}
+
+define <2 x i8> @fold_select_trunc_if_nuw_vector(<2 x i8> %x, <2 x i8> %y) {
+; CHECK-LABEL: @fold_select_trunc_if_nuw_vector(
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc nuw <2 x i8> [[X:%.*]] to <2 x i1>
+; CHECK-NEXT:    [[RET:%.*]] = select <2 x i1> [[TRUNC]], <2 x i8> <i8 1, i8 1>, <2 x i8> [[Y:%.*]]
+; CHECK-NEXT:    ret <2 x i8> [[RET]]
+;
+  %trunc = trunc nuw <2 x i8> %x to <2 x i1>
+  %ret = select <2 x i1> %trunc, <2 x i8> %x, <2 x i8> %y
+  ret <2 x i8> %ret
+}