From 9d8e9be4f4e9fc55f0ab1f42e8acd259543b66b5 Mon Sep 17 00:00:00 2001
From: c8ef <c8ef@outlook.com>
Date: Sat, 24 Aug 2024 19:56:59 +0800
Subject: [PATCH] [InstCombine] Update the `select` operand when the `cond` is
 `trunc` and has the `nuw` or `nsw` property.  (#105914)

This patch updates the select operand when the cond has the nuw or nsw
property. Considering the semantics of the nuw and nsw flag, if there is
no poison value in this expression, this code assumes that X can only be
0, 1 or -1.

close: #96765
alive2: https://alive2.llvm.org/ce/z/3n3n2Q
---
 .../InstCombine/InstCombineSelect.cpp         | 19 ++++++
 .../InstCombine/fold-select-trunc.ll          | 68 +++++++++++++++++++
 2 files changed, 87 insertions(+)
 create mode 100644 llvm/test/Transforms/InstCombine/fold-select-trunc.ll

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 18ffc209f259e0b..fcd11126073bf19 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -4201,5 +4201,24 @@ Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
     }
   }
 
+  // select (trunc nuw X to i1), X, Y --> select (trunc nuw X to i1), 1, Y
+  // select (trunc nuw X to i1), Y, X --> select (trunc nuw X to i1), Y, 0
+  // select (trunc nsw X to i1), X, Y --> select (trunc nsw X to i1), -1, Y
+  // select (trunc nsw X to i1), Y, X --> select (trunc nsw X to i1), Y, 0
+  Value *Trunc;
+  if (match(CondVal, m_NUWTrunc(m_Value(Trunc)))) {
+    if (TrueVal == Trunc)
+      return replaceOperand(SI, 1, ConstantInt::get(TrueVal->getType(), 1));
+    if (FalseVal == Trunc)
+      return replaceOperand(SI, 2, ConstantInt::get(FalseVal->getType(), 0));
+  }
+  if (match(CondVal, m_NSWTrunc(m_Value(Trunc)))) {
+    if (TrueVal == Trunc)
+      return replaceOperand(SI, 1,
+                            Constant::getAllOnesValue(TrueVal->getType()));
+    if (FalseVal == Trunc)
+      return replaceOperand(SI, 2, ConstantInt::get(FalseVal->getType(), 0));
+  }
+
   return nullptr;
 }
diff --git a/llvm/test/Transforms/InstCombine/fold-select-trunc.ll b/llvm/test/Transforms/InstCombine/fold-select-trunc.ll
new file mode 100644
index 000000000000000..5567d7d5e1fca93
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/fold-select-trunc.ll
@@ -0,0 +1,68 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+define i8 @fold_select_trunc_nuw_true(i8 %x, i8 %y) {
+; CHECK-LABEL: @fold_select_trunc_nuw_true(
+; 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_nuw_false(i8 %x, i8 %y) {
+; CHECK-LABEL: @fold_select_trunc_nuw_false(
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc nuw i8 [[X:%.*]] to i1
+; CHECK-NEXT:    [[RET:%.*]] = select i1 [[TRUNC]], i8 [[Y:%.*]], i8 0
+; CHECK-NEXT:    ret i8 [[RET]]
+;
+  %trunc = trunc nuw i8 %x to i1
+  %ret = select i1 %trunc, i8 %y, i8 %x
+  ret i8 %ret
+}
+
+define i128 @fold_select_trunc_nsw_true(i128 %x, i128 %y) {
+; CHECK-LABEL: @fold_select_trunc_nsw_true(
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc nsw i128 [[X:%.*]] to i1
+; CHECK-NEXT:    [[RET:%.*]] = select i1 [[TRUNC]], i128 -1, i128 [[Y:%.*]]
+; CHECK-NEXT:    ret i128 [[RET]]
+;
+  %trunc = trunc nsw i128 %x to i1
+  %ret = select i1 %trunc, i128 %x, i128 %y
+  ret i128 %ret
+}
+
+define i8 @fold_select_trunc_nsw_false(i8 %x, i8 %y) {
+; CHECK-LABEL: @fold_select_trunc_nsw_false(
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc nsw i8 [[X:%.*]] to i1
+; CHECK-NEXT:    [[RET:%.*]] = select i1 [[TRUNC]], i8 [[Y:%.*]], i8 0
+; CHECK-NEXT:    ret i8 [[RET]]
+;
+  %trunc = trunc nsw i8 %x to i1
+  %ret = select i1 %trunc, i8 %y, i8 %x
+  ret i8 %ret
+}
+
+define i8 @fold_select_trunc_negative(i8 %x, i8 %y) {
+; CHECK-LABEL: @fold_select_trunc_negative(
+; CHECK-NEXT:    [[TRUNC:%.*]] = trunc i8 [[X:%.*]] to i1
+; CHECK-NEXT:    [[RET:%.*]] = select i1 [[TRUNC]], i8 [[X]], i8 [[Y:%.*]]
+; CHECK-NEXT:    ret i8 [[RET]]
+;
+  %trunc = trunc i8 %x to i1
+  %ret = select i1 %trunc, i8 %x, i8 %y
+  ret i8 %ret
+}
+
+define <2 x i8> @fold_select_trunc_vector(<2 x i8> %x, <2 x i8> %y) {
+; CHECK-LABEL: @fold_select_trunc_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
+}